Merge branch 'dev' into feat/lora-sdnq-stack

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,
@@ -43,7 +49,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
 
@@ -114,6 +120,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
@@ -151,6 +158,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"] = ''
 
@@ -203,6 +211,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 = [
@@ -270,6 +279,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")),

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