diff --git a/modules/attention/__init__.py b/modules/attention/__init__.py index 78994780c..8ae75345b 100644 --- a/modules/attention/__init__.py +++ b/modules/attention/__init__.py @@ -1,12 +1,12 @@ """Attention backends: one scaled_dot_product_attention router over the registered backends, the per-generation context, and the diffusers-side processor and dispatcher setup.""" from modules.attention.registry import AttentionBackend, AttentionCall, Constraints, Platform, Registry, registry from modules.attention.router import Plan, PlanEntry, build_plan, get_plan, install_router, reapply, reapply_options, report -from modules.attention.dispatcher import set_diffusers_attention, set_attention_dispatcher, hijack_kernels, get_kernel_hijack, get_hf_api_hijack +from modules.attention.dispatcher import set_diffusers_attention, set_attention_dispatcher, list_dispatcher_backends, hijack_kernels, get_kernel_hijack, get_hf_api_hijack from modules.attention import backends, context, debug __all__ = [ 'AttentionBackend', 'AttentionCall', 'Constraints', 'Platform', 'Registry', 'registry', 'Plan', 'PlanEntry', 'build_plan', 'get_plan', 'install_router', 'reapply', 'reapply_options', 'report', - 'set_diffusers_attention', 'set_attention_dispatcher', 'hijack_kernels', 'get_kernel_hijack', 'get_hf_api_hijack', + 'set_diffusers_attention', 'set_attention_dispatcher', 'list_dispatcher_backends', 'hijack_kernels', 'get_kernel_hijack', 'get_hf_api_hijack', 'backends', 'context', 'debug', ] diff --git a/modules/attention/dispatcher.py b/modules/attention/dispatcher.py index 19d58e99b..ea7466373 100644 --- a/modules/attention/dispatcher.py +++ b/modules/attention/dispatcher.py @@ -94,3 +94,13 @@ def set_attention_dispatcher(pipe): log.warning(f'Attention dispatcher: active={prev[0].value} list={backends} target={attn} not found') else: log.debug(f'Attention dispatcher: active={prev[0].value} list={backends}') + + +def list_dispatcher_backends() -> list: + """The kernels diffusers can dispatch attention to, for anything that offers hf_attention as a choice.""" + try: + from diffusers.models import attention_dispatch as a + return sorted(b.value for b in a._AttentionBackendRegistry.list_backends()) # pylint: disable=protected-access + except Exception as e: + log.error(f'Attention dispatcher: {e}') + return [] diff --git a/scripts/xyz/xyz_grid_classes.py b/scripts/xyz/xyz_grid_classes.py index 2da29d7b3..5ce14b422 100644 --- a/scripts/xyz/xyz_grid_classes.py +++ b/scripts/xyz/xyz_grid_classes.py @@ -3,6 +3,12 @@ from scripts.xyz.xyz_grid_shared import ( # pylint: disable=no-name-in-module, u apply_task_arg, apply_task_args, apply_setting, + apply_attention, + apply_attention_overrides, + apply_attention_dispatcher, + list_sdp_overrides, + save_attention, + restore_attention, apply_prompt_primary, apply_prompt_refine, apply_prompt_detailer, @@ -40,7 +46,7 @@ from scripts.xyz.xyz_grid_shared import ( # pylint: disable=no-name-in-module, u format_nothing, str_permutations, ) -from modules import shared, shared_items, sd_samplers, ipadapter, sd_models, sd_vae, sd_unet +from modules import shared, shared_items, sd_samplers, ipadapter, sd_models, sd_vae, sd_unet, attention from modules.control.units import controlnet, t2iadapter from modules.control import processor @@ -107,6 +113,7 @@ class SharedSettingsStackHelper(): disable_apply_metadata = None disable_apply_params = None sdnq_quant_mode = None + attention_settings = None def __enter__(self): # Save overridden settings so they can be restored later @@ -140,6 +147,7 @@ class SharedSettingsStackHelper(): self.disable_apply_metadata = shared.opts.disable_apply_metadata self.disable_apply_params = shared.opts.disable_apply_params self.sdnq_quant_mode = shared.opts.sdnq_quantize_weights_mode + self.attention_settings = save_attention() shared.opts.data["disable_apply_metadata"] = [] shared.opts.data["disable_apply_params"] = '' @@ -188,6 +196,7 @@ class SharedSettingsStackHelper(): if self.sdnq_quant_mode != shared.opts.sdnq_quantize_weights_mode: shared.opts.data["sdnq_quantize_weights_mode"] = self.sdnq_quant_mode sd_models.reload_model_weights(op='model') + restore_attention(self.attention_settings) axis_options = [ @@ -250,6 +259,20 @@ axis_options = [ AxisOption("[Postprocess] Detailer strength", str, apply_field("detailer_strength")), AxisOption("[Quant] SDNQ quant mode", str, apply_sdnq_quant, cost=0.9, fmt=format_value_add_label, choices=lambda: ['none'] + sorted(shared_items.sdnq_quant_modes)), AxisOption("[Quant] SDNQ quant mode TE", str, apply_sdnq_quant_te, cost=0.9, fmt=format_value_add_label, choices=lambda: ['none'] + sorted(shared_items.sdnq_quant_modes)), + AxisOption("[Attention] Method", str, apply_setting('cross_attention_optimization'), cost=0.2, choices=shared_items.list_crossattention), + AxisOption("[Attention] SDP override", str, apply_attention_overrides, cost=0.2, choices=list_sdp_overrides), + AxisOption("[Attention] Dispatcher", str, apply_attention_dispatcher, cost=0.2, choices=lambda: ['None'] + attention.list_dispatcher_backends()), + AxisOption("[Attention] SDNQ matmul", str, apply_attention('sdnq_attention_matmul_type'), cost=0.2, choices=lambda: list(shared_items.sdnq_matmul_modes)), + AxisOption("[Attention] SDNQ PV matmul", str, apply_attention('sdnq_attention_pv_matmul_type'), cost=0.2, choices=lambda: list(shared_items.sdnq_matmul_modes)), + AxisOption("[Attention] SDNQ smooth K", str, apply_attention('sdnq_attention_smooth_k'), cost=0.2, choices=lambda: ['False', 'True']), + AxisOption("[Attention] SDNQ hadamard", str, apply_attention('sdnq_attention_use_hadamard'), cost=0.2, choices=lambda: ['False', 'True']), + AxisOption("[Attention] SDNQ fp16 accumulation", str, apply_attention('sdnq_attention_use_fp16_accum'), cost=0.2, choices=lambda: ['False', 'True']), + AxisOption("[Sparse] Enabled", str, apply_attention('sparse_attention_enabled'), cost=0.2, choices=lambda: ['False', 'True']), + AxisOption("[Sparse] KV budget", int, apply_attention('sparse_attention_budget'), cost=0.2), + AxisOption("[Sparse] Minimum sequence", int, apply_attention('sparse_attention_min_tokens'), cost=0.2), + AxisOption("[Sparse] Dense steps", int, apply_attention('sparse_attention_schedule_steps'), cost=0.2), + AxisOption("[Sparse] Dense step bonus", int, apply_attention('sparse_attention_schedule_bump'), cost=0.2), + AxisOption("[Sparse] Shared heads", str, apply_attention('sparse_attention_head_shared'), cost=0.2, choices=lambda: ['False', 'True']), AxisOption("[HDR] Mode", int, apply_field("hdr_mode")), AxisOption("[HDR] Brightness", float, apply_field("hdr_brightness")), AxisOption("[HDR] Color", float, apply_field("hdr_color")), diff --git a/scripts/xyz/xyz_grid_shared.py b/scripts/xyz/xyz_grid_shared.py index b90d37848..f84696f41 100644 --- a/scripts/xyz/xyz_grid_shared.py +++ b/scripts/xyz/xyz_grid_shared.py @@ -81,6 +81,75 @@ def apply_setting(field): return fun +def attention_options() -> list: + """Attention settings an axis can change; the stack helper restores exactly this set.""" + from modules import attention + return ['cross_attention_optimization', 'hf_attention', *attention.reapply_options()] + + +def list_sdp_overrides() -> list: + item = shared.opts.data_labels.get('sdp_overrides', None) + args = item.component_args if item is not None else None + args = args() if callable(args) else args + return ['None'] + list((args or {}).get('choices', None) or []) + + +def apply_attention(field): + def fun(p, x, xs): + from modules import attention + apply_setting(field)(p, x, xs) + attention.reapply() # backends read their settings when the chain is built, so a write on its own changes nothing + owner = next((backend for backend in attention.registry.backends.values() if field in backend.options), None) + plan = attention.get_plan() + if owner is not None and plan is not None and owner.name not in plan.chain(): + log.warning(f'XYZ grid apply attention: {field} is read by "{owner.label}" which is not in the active chain={plan.chain()}') + return fun + + +def apply_attention_overrides(p, x, xs): + from modules import attention + labels = [label.strip() for label in str(x).split('+') if len(label.strip()) > 0 and label.strip().lower() != 'none'] + unknown = [label for label in labels if attention.registry.by_label(label) is None] + if len(unknown) > 0: + log.warning(f'XYZ grid apply attention: unknown overrides={unknown} available={attention.registry.labels()}') + shared.opts.data['sdp_overrides'] = labels + attention.reapply() + log.debug(f'XYZ grid apply attention: overrides={labels}') + + +def apply_attention_dispatcher(p, x, xs): + from modules import attention + value = '' if str(x).strip().lower() in ['none', 'default'] else str(x).strip() + shared.opts.data['hf_attention'] = value + if shared.sd_loaded: + attention.set_attention_dispatcher(shared.sd_model) + log.debug(f'XYZ grid apply attention: dispatcher="{value}"') + + +def save_attention() -> dict: + return {field: shared.opts.data[field] for field in attention_options() if field in shared.opts.data} + + +def restore_attention(saved: dict): + """Put back whatever an attention axis changed, keys it introduced included, then rebuild what reads them.""" + from modules import attention + changed = [] + for field in attention_options(): + if (field in saved) == (field in shared.opts.data) and saved.get(field, None) == shared.opts.data.get(field, None): + continue + changed.append(field) + if field in saved: + shared.opts.data[field] = saved[field] + else: + shared.opts.data.pop(field, None) + if len(changed) == 0: + return + attention.reapply() + if 'hf_attention' in changed and shared.sd_loaded: + attention.set_attention_dispatcher(shared.sd_model) + log.debug(f'XYZ grid restore attention: {changed}') + + def apply_seed(p, x, xs): p.seed = x p.all_seeds = None diff --git a/test/test-xyz-attention.py b/test/test-xyz-attention.py new file mode 100644 index 000000000..3a16754f2 --- /dev/null +++ b/test/test-xyz-attention.py @@ -0,0 +1,252 @@ +#!/usr/bin/env python +""" +Offline unit tests for the attention and sparse axes of the xyz grid. + +Covers: + +- every [Attention] and [Sparse] axis resolves its choices and targets a registered option +- applying an axis and then leaving the grid restores shared.opts.data exactly, keys the axis + introduced included, so a grid never leaks its last cell into the session +- the axes backed by boolean options take the string the dropdown hands them +- the sdp override axis turns a label, a plus joined pair or None into the option's list +- an override axis apply reaches the router: the rebuilt chain contains the backend it named +- the restore set covers every setting the sparse stage reads + +No running server required. Nothing is moved to the accelerator, and no axis value that would +install a package is used. + +Usage: + python test/test-xyz-attention.py +""" + +import os +import sys + +import torch + +script_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, script_dir) +os.chdir(script_dir) + +os.environ['SD_INSTALL_QUIET'] = '1' + +# Bootstrap cmd_args before any module that pulls in shared.py. +import modules.cmd_args # pylint: disable=wrong-import-position +import installer # pylint: disable=wrong-import-position +orig_argv = sys.argv +sys.argv = [sys.argv[0]] +try: + modules.cmd_args.parse_args() +finally: + sys.argv = orig_argv +installer.add_args(modules.cmd_args.parser) +modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([]) + +stock_sdpa = torch.nn.functional.scaled_dot_product_attention # importing shared installs the configured hijacks in-process + +from modules.errors import log # pylint: disable=wrong-import-position +from modules import attention, shared # pylint: disable=wrong-import-position +from modules.attention.sparse import stage as sparse_stage # pylint: disable=wrong-import-position +from scripts.xyz import xyz_grid_shared as xyz # pylint: disable=wrong-import-position +from scripts.xyz.xyz_grid_classes import axis_options # pylint: disable=wrong-import-position + + +results: dict[str, dict] = {} + + +def category(name: str): + if name not in results: + results[name] = {'passed': 0, 'failed': 0, 'tests': []} + return name + + +def record(cat: str, passed: bool, name: str, detail: str = ''): + status = 'PASS' if passed else 'FAIL' + results[cat]['passed' if passed else 'failed'] += 1 + results[cat]['tests'].append((status, name)) + msg = f' {status}: {name}' + if detail: + msg += f' ({detail})' + if passed: + log.info(msg) + else: + log.error(msg) + + +def run_test(cat: str, fn): + try: + ok = fn() + record(cat, ok is not False, fn.__name__) + except AssertionError as e: + record(cat, False, fn.__name__, str(e)) + except Exception as e: # pylint: disable=broad-except + record(cat, False, fn.__name__, f'exception: {e}') + import traceback + traceback.print_exc() + + +def attention_axes(): + return [axis for axis in axis_options if axis.label.startswith('[Attention]') or axis.label.startswith('[Sparse]')] + + +def sample_value(axis): + """A value the axis accepts that is not the current one, avoiding backends whose prepare installs a package.""" + if axis.choices is None: + return int(shared.opts.get(option_of(axis) or 'sparse_attention_budget') or 0) + 5 + choices = [choice for choice in axis.choices() if choice not in ['Sage attention', 'Flash attention', 'Triton Flash attention']] + current = str(shared.opts.get(option_of(axis)) if option_of(axis) else '') + return next((choice for choice in choices if str(choice) != current), choices[0]) + + +def option_of(axis): + """The option an axis writes, read back from the closure the axis factory built.""" + closure = getattr(axis.apply, '__closure__', None) or () + for cell in closure: + if isinstance(cell.cell_contents, str) and cell.cell_contents in shared.opts.data_labels: + return cell.cell_contents + return {'[Attention] SDP override': 'sdp_overrides', '[Attention] Dispatcher': 'hf_attention'}.get(axis.label, None) + + +def bool_axes(): + return [axis for axis in attention_axes() if isinstance(shared.opts.get(option_of(axis) or ''), bool)] + + +# ============================================================ +# Tests +# ============================================================ + +def test_axes_are_registered(): + axes = attention_axes() + assert len(axes) >= 10, f'only {len(axes)} attention axes' + for axis in axes: + option = option_of(axis) + assert option is not None, f'{axis.label} names no option' + assert option in shared.opts.data_labels, f'{axis.label} targets unknown option {option}' + return True + + +def test_axis_choices_resolve(): + for axis in attention_axes(): + if axis.choices is None: + assert axis.type is int, f'{axis.label} has no choices and is not numeric' + continue + choices = axis.choices() + assert isinstance(choices, list) and len(choices) > 0, f'{axis.label} resolved {choices}' + return True + + +def test_axes_write_and_restore_exactly(): + for axis in attention_axes(): + before = dict(shared.opts.data) + saved = xyz.save_attention() + axis.apply(None, sample_value(axis), []) + assert dict(shared.opts.data) != before, f'{axis.label} wrote nothing' + xyz.restore_attention(saved) + after = dict(shared.opts.data) + leaked = {key for key in set(before) | set(after) if before.get(key, '') != after.get(key, '')} + assert not leaked, f'{axis.label} leaked {sorted(leaked)}' + return True + + +def test_bool_axes_coerce_the_string_the_dropdown_sends(): + axes = bool_axes() + assert len(axes) >= 3, f'only {len(axes)} boolean axes' + for axis in axes: + option = option_of(axis) + saved = xyz.save_attention() + axis.apply(None, 'True', []) + assert shared.opts.data[option] is True, f'{axis.label} took "True" as {shared.opts.data[option]!r}' + axis.apply(None, 'False', []) + assert shared.opts.data[option] is False, f'{axis.label} took "False" as {shared.opts.data[option]!r}' + xyz.restore_attention(saved) + return True + + +def test_override_axis_parses_labels(): + saved = xyz.save_attention() + try: + xyz.apply_attention_overrides(None, 'None', []) + assert shared.opts.data['sdp_overrides'] == [], shared.opts.data['sdp_overrides'] + xyz.apply_attention_overrides(None, 'Flex attention', []) + assert shared.opts.data['sdp_overrides'] == ['Flex attention'], shared.opts.data['sdp_overrides'] + xyz.apply_attention_overrides(None, 'Flex attention+SDNQ attention', []) + assert shared.opts.data['sdp_overrides'] == ['Flex attention', 'SDNQ attention'], shared.opts.data['sdp_overrides'] + finally: + xyz.restore_attention(saved) + return True + + +def test_override_axis_rebuilds_the_chain(): + saved = xyz.save_attention() + try: + xyz.apply_attention_overrides(None, 'Flex attention', []) + assert 'flex' in attention.get_plan().chain(), attention.get_plan().chain() + xyz.apply_attention_overrides(None, 'None', []) + assert 'flex' not in attention.get_plan().chain(), attention.get_plan().chain() + finally: + xyz.restore_attention(saved) + return True + + +def test_dispatcher_axis_clears_on_none(): + saved = xyz.save_attention() + try: + xyz.apply_attention_dispatcher(None, 'native', []) + assert shared.opts.data['hf_attention'] == 'native', shared.opts.data['hf_attention'] + xyz.apply_attention_dispatcher(None, 'None', []) + assert shared.opts.data['hf_attention'] == '', repr(shared.opts.data['hf_attention']) + finally: + xyz.restore_attention(saved) + return True + + +def test_restore_set_covers_every_attention_setting(): + covered = set(xyz.attention_options()) + missing = set(sparse_stage.OPTION_NAMES) - covered + assert not missing, f'sparse settings outside the restore set: {sorted(missing)}' + for backend in attention.registry.backends.values(): + outside = set(backend.options) - covered + assert not outside, f'{backend.label} settings outside the restore set: {sorted(outside)}' + return True + + +def run_all(): + log.warning('=== xyz attention axes ===') + cat = category('axes') + for fn in [ + test_axes_are_registered, + test_axis_choices_resolve, + ]: + run_test(cat, fn) + + log.warning('=== apply and restore ===') + cat = category('apply') + for fn in [ + test_axes_write_and_restore_exactly, + test_bool_axes_coerce_the_string_the_dropdown_sends, + test_override_axis_parses_labels, + test_override_axis_rebuilds_the_chain, + test_dispatcher_axis_clears_on_none, + test_restore_set_covers_every_attention_setting, + ]: + run_test(cat, fn) + + log.warning('=== Results ===') + total_passed = 0 + total_failed = 0 + for cat_name, info in results.items(): + status = 'PASS' if info['failed'] == 0 else 'FAIL' + log.info(f" {cat_name}: {info['passed']} passed, {info['failed']} failed [{status}]") + total_passed += info['passed'] + total_failed += info['failed'] + log.warning(f'Total: {total_passed} passed, {total_failed} failed') + return total_failed == 0 + + +if __name__ == '__main__': + import time + t0 = time.time() + ok = run_all() + torch.nn.functional.scaled_dot_product_attention = stock_sdpa + log.warning(f'Total time: {time.time() - t0:.2f}s') + sys.exit(0 if ok else 1)