fix(attention): attention modes never worked - wrap the override the way ComfyUI calls it (#16)

ComfyUI's wrap_attn() calls transformer_options["optimized_attention_override"] as
override(original_func, q, k, v, heads, ...). We installed the registered attention
function itself, so q was the original function: "'function' object has no attribute
'shape'" for every installed non-default mode (sage, flash, xformers, pytorch, ...).
A mode that was not installed fell back to default silently, which is why "sage3 worked".

- make_attention_override() drops the leading function and forwards the rest
- generate_image_attn_safe(): one clear log line naming the mode, one retry with the
  default attention; the item records attention_fallback
- README lists the real mode names (sageattn was never one)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

README.md

@@ -196,7 +196,7 @@ If you'd rather build the graph yourself:
 * **Live Dashboard Updates:** Configure `flush_batch_every` to update the dashboard incrementally (e.g., every 4 images) instead of waiting for the entire batch to complete.
 * **Remote VAE Support:** Offload VAE decoding to remote servers (HuggingFace endpoints or local) for 20-30% faster generation and lower VRAM usage.
 * **VAE Selection & Iteration:** Test multiple VAEs per config — add individual VAEs from a searchable dropdown or expand entire VAE folders. Each VAE becomes a dimension in the Cartesian grid.
-* **Attention Mode Testing:** Grid-test different attention implementations (`default`, `sageattn`, `xformers`, etc.) as a combinatorial dimension. Find the optimal attention backend for your hardware.
+* **Attention Mode Testing:** Grid-test different attention implementations (`default`, `sage`, `sage3`, `flash`, `xformers`, etc.) as a combinatorial dimension. Find the optimal attention backend for your hardware.
 * **CLIP Encoding Combinators:** Advanced prompt syntax with `AND`, `CAT`, `AVG(weight)`, and `BREAK` keywords for multi-segment conditioning. `AND` creates separate conditioning entries, `CAT` concatenates token contexts, `AVG(0.5)` blends prompts, and `BREAK` forces new 77-token chunks.
 * **Model-Specific Prompts:** Per-config prompt prefix and suffix fields for model-appropriate quality tags (e.g., prepend `"masterpiece, best quality"` for anime models, append `"4k, photorealistic"` for realistic models).
 * **Recursive Cartesian Prompt Builder:** Visual chip-based prompt editor with `{option1|option2|option3}` syntax that expands into all combinations. Nest multiple groups for exponential prompt variation.
@@ -623,7 +623,7 @@ Config Arrays are the building blocks of your testing grid. Each array defines a
 * **CFG:** Comma-separated list of CFG scale values (e.g., `"6.0, 7.0, 8.0"`).
 * **Samplers Dropdown:** Searchable dropdown auto-populated with all available samplers from the connected sampler node. Select multiple to add them as a comma-separated list.
 * **Schedulers Dropdown:** Searchable dropdown auto-populated with all available schedulers. Works the same as the samplers dropdown.
-* **Attention Mode:** Dropdown selector for attention implementations (`default`, `sageattn`, `xformers`, etc.). Select multiple to test as a grid dimension.
+* **Attention Mode:** Dropdown selector for attention implementations (`default`, `sage`, `sage3`, `flash`, `xformers`, etc.). Select multiple to test as a grid dimension.
 
 #### Iteration Count Display
 
@@ -782,7 +782,7 @@ VAEs are treated as a Cartesian dimension just like samplers or models. If you a
 
 ### Attention Modes
 
-Select which attention backend implementations to test. Available options depend on your system (e.g., `default`, `sageattn`, `xformers`). Like other parameters, you can select multiple modes and they become a grid dimension.
+Select which attention backend implementations to test. Available options depend on your system (e.g., `default`, `xformers`, `pytorch`, `flash`, `sage`, `sage3`, `sub_quad`, `split`; a backend that is not installed is skipped with a warning and the job runs with default attention). If a backend is installed but fails while sampling, the job is retried once with default attention and the item records `attention_fallback` (the mode that failed). Like other parameters, you can select multiple modes and they become a grid dimension.
 
 ### Skip Layer Guidance (DiT models)
 
@@ -1339,7 +1339,7 @@ The Config Builder generates JSON in the exact format required by the Ultimate S
     "lora": "style.safetensors:0.8:0.8 + detail.safetensors:0.5:0.5",
     "model": "base_xl.safetensors",
     "vae": ["vae_v1.safetensors", "vae_v2.safetensors"],
-    "attention_mode": ["default", "sageattn"],
+    "attention_mode": ["default", "sage"],
     "prompt_prefix": "masterpiece, best quality, ",
     "prompt_suffix": ", 4k, detailed",
     "lora_omit_triggers": ["masterpiece", "best quality"],
@@ -1844,7 +1844,7 @@ After modifying backend Python files, restart ComfyUI. After modifying `web/` JS
 
 ### Update 2/16/26 — Config Builder & Dashboard Overhaul
 * 🎛️ **VAE Selection & Iteration:** Test multiple VAEs per config with searchable dropdown and folder expansion. VAEs are a full Cartesian dimension.
-* 🧠 **Attention Mode Testing:** Grid-test different attention implementations (`default`, `sageattn`, `xformers`) as a combinatorial dimension.
+* 🧠 **Attention Mode Testing:** Grid-test different attention implementations (`default`, `sage`, `xformers`) as a combinatorial dimension.
 * 🔗 **CLIP Encoding Combinators:** `AND`, `CAT`, `AVG(weight)`, and `BREAK` keywords for multi-segment conditioning and prompt blending.
 * 📝 **Model-Specific Prompts:** Per-config prompt prefix and suffix fields for model-appropriate quality tags.
 * 🔄 **Recursive Cartesian Prompt Builder:** Visual chip-based editor with `{option1|option2}` syntax expanding into all combinations.

generation_orchestrator.py

@@ -26,7 +26,7 @@ from .model_loader import (
 )
 from .lora_utils import expand_lora_folder
 from .image_generation import (
-    generate_image, flush_batch_with_vae, flush_batch_with_remote_vae,
+    generate_image, generate_image_attn_safe, flush_batch_with_vae, flush_batch_with_remote_vae,
     create_image_metadata, decode_latent_with_vae, pil_to_image_array, calculate_eta, print_generation_progress
 )
 from .config_utils import sanitize_session_name
@@ -2091,6 +2091,7 @@ def run_generation_loop(
             latent_in = make_empty_latent(latent_channels, gen_w, gen_h, latent_ratio)
             
         result_latent = None
+        attn_fallback = None
         try:
             attention_mode = conf.get("attention_mode", "default")
             if attention_mode != "default":
@@ -2107,7 +2108,7 @@ def run_generation_loop(
             if extra_opts:
                 print(f"[GridTester] ⚙️ Extra options: {', '.join(extra_opts)}")
 
-            result_latent, duration = generate_image(
+            result_latent, duration, attn_fallback = generate_image_attn_safe(
                 patched_model, current_seed, eff_conf["steps"], eff_conf["cfg"],
                 conf["sampler"], conf["scheduler"], final_positive, final_negative,
                 latent_in, conf["denoise"], attention_mode=attention_mode,
@@ -2732,6 +2733,9 @@ def run_generation_loop(
                     gen_index=gen_index_offset + total_generated
                 )
                 meta["config_fp"] = job_fp
+                if attn_fallback:
+                    meta["attention_mode"] = "default"
+                    meta["attention_fallback"] = attn_fallback
                 if _is_pv:
                     meta.update({
                         "preview": True, "preview_of": job_fp,
@@ -3758,7 +3762,7 @@ def _run_distributed_generation(
                 if extra_opts:
                     print(f"[GridTester] ⚙️ Extra options: {', '.join(extra_opts)}")
 
-                result_latent, duration = generate_image(
+                result_latent, duration, attn_fallback = generate_image_attn_safe(
                     patched_model, current_seed, conf["steps"], conf["cfg"],
                     conf["sampler"], conf["scheduler"], final_positive, final_negative,
                     latent_in, conf["denoise"], attention_mode=attention_mode,
@@ -3802,6 +3806,9 @@ def _run_distributed_generation(
                 if ux.get("stamp"):
                     meta.update(ux["stamp"])
                 meta["dist_worker"] = "master"
+                if attn_fallback:
+                    meta["attention_mode"] = "default"
+                    meta["attention_fallback"] = attn_fallback
 
                 pending_batch.append((result_latent["samples"].clone(), meta))
                 pending_job_ids.append(job_id)

image_generation.py

@@ -138,6 +138,44 @@ def pil_to_image_array(pil_image):
 # NOTE: torch.inference_mode() decorator — see module docstring above for why
 # this is REQUIRED on all out-of-graph inference entry points. Removing it
 # will reintroduce the same 6× VRAM blowup we fixed in May 2026.
+def make_attention_override(attn_func, name=""):
+    """Wrap a registered ComfyUI attention function as an `optimized_attention_override`.
+
+    ComfyUI's wrap_attn() calls the override as `override(original_func, q, k, v, heads, ...)`
+    (the ORIGINAL attention function comes first), whereas the functions in
+    REGISTERED_ATTENTION_FUNCTIONS take `(q, k, v, heads, ...)`. Installing the registered
+    function directly made q the original function ("'function' object has no attribute
+    'shape'", issue #16). Drop the leading function and forward the rest.
+    """
+    def _override(func, *args, **kwargs):
+        return attn_func(*args, **kwargs)
+    _override.uscg_attention = name
+    return _override
+
+
+def generate_image_attn_safe(*args, attention_mode="default", **kwargs):
+    """generate_image() with a one-shot fallback to the default attention.
+
+    A non-default attention backend can fail at sample time (not installed for this
+    torch/CUDA build, unsupported head dim, ...). Instead of losing the job, print one
+    clear line and retry it once with attention_mode="default".
+
+    Returns (latent, duration, fallback_from) where fallback_from is the attention mode
+    that failed, or None when the requested mode worked.
+    """
+    try:
+        latent, duration = generate_image(*args, attention_mode=attention_mode, **kwargs)
+        return latent, duration, None
+    except Exception as e:
+        if not attention_mode or attention_mode == "default" or type(e).__name__ == "InterruptProcessingException":
+            raise
+        print(f"[GridTester] ❌ Attention mode '{attention_mode}' failed ({type(e).__name__}: {e}): "
+              f"check that the backend is installed for this torch/CUDA build; "
+              f"this job will be retried with the default attention")
+        latent, duration = generate_image(*args, attention_mode="default", **kwargs)
+        return latent, duration, attention_mode
+
+
 @torch.inference_mode()
 def generate_image(
     patched_model,
@@ -229,7 +267,7 @@ def generate_image(
                 original_attn_override = patched_model.model_options.get("transformer_options", {}).get("optimized_attention_override")
                 if "transformer_options" not in patched_model.model_options:
                     patched_model.model_options["transformer_options"] = {}
-                patched_model.model_options["transformer_options"]["optimized_attention_override"] = attn_func
+                patched_model.model_options["transformer_options"]["optimized_attention_override"] = make_attention_override(attn_func, attention_mode)
             else:
                 print(f"[GridTester] ⚠️ Attention mode '{attention_mode}' not available, using default")
         except Exception as e:

tests/conftest.py

@@ -129,6 +129,7 @@ _GEN_ORCH_DEPS = {
     },
     "image_generation": {
         "generate_image": lambda *a, **kw: None,
+        "generate_image_attn_safe": lambda *a, **kw: (None, None, None),
         "flush_batch_with_vae": lambda *a, **kw: None,
         "flush_batch_with_remote_vae": lambda *a, **kw: None,
         "create_image_metadata": lambda *a, **kw: {},

tests/test_issue16_attention.py

@@ -0,0 +1,109 @@
+"""Issue #16: registered attention functions were installed as the override directly.
+
+ComfyUI calls `optimized_attention_override(original_func, q, k, v, heads, ...)`, so the
+registered `(q, k, v, heads, ...)` function received the original function as q and died
+with "'function' object has no attribute 'shape'".
+"""
+import importlib.util
+import os
+import sys
+import types
+
+import pytest
+
+_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, _ROOT)
+
+
+@pytest.fixture
+def ig(monkeypatch):
+    """The real image_generation.py (conftest registers a stub under the package name)."""
+    import torch
+    if not hasattr(torch, "inference_mode"):
+        monkeypatch.setattr(torch, "inference_mode", lambda *a, **k: (lambda f: f), raising=False)
+    pkg = os.path.basename(_ROOT)
+    spec = importlib.util.spec_from_file_location(f"{pkg}.image_generation_real", os.path.join(_ROOT, "image_generation.py"))
+    mod = importlib.util.module_from_spec(spec)
+    mod.__package__ = pkg
+    try:
+        spec.loader.exec_module(mod)
+    except Exception as e:  # relative imports needing real ComfyUI
+        pytest.skip(f"image_generation.py not importable under the stubs: {e}")
+    return mod
+
+
+def test_override_drops_original_function(ig):
+    calls = []
+
+    def fake_sage(*args, **kwargs):
+        calls.append((args, kwargs))
+        return "out"
+
+    override = ig.make_attention_override(fake_sage, "sage")
+    orig = lambda *a, **k: "orig"
+    q, k, v = object(), object(), object()
+    assert override(orig, q, k, v, 8, mask=None, _inside_attn_wrapper=True) == "out"
+    args, kwargs = calls[0]
+    assert args == (q, k, v, 8)
+    assert orig not in args
+    assert kwargs == {"mask": None, "_inside_attn_wrapper": True}
+    assert override.uscg_attention == "sage"
+
+
+def _fake_comfy(monkeypatch, funcs):
+    mod = types.ModuleType("comfy.ldm.modules.attention")
+    mod.get_attention_function = lambda name, default=...: funcs.get(name, default)
+    for n in ("comfy.ldm", "comfy.ldm.modules"):
+        monkeypatch.setitem(sys.modules, n, types.ModuleType(n))
+    monkeypatch.setitem(sys.modules, "comfy.ldm.modules.attention", mod)
+
+
+def test_generate_image_installs_and_restores_override(ig, monkeypatch):
+    """The install/restore block of generate_image, exercised via source-level contract:
+    the installed object is the wrapper, not the registered function."""
+    import inspect
+    src = inspect.getsource(ig.generate_image)
+    assert "make_attention_override(attn_func, attention_mode)" in src
+    assert 'del patched_model.model_options["transformer_options"]["optimized_attention_override"]' in src
+
+
+def test_fallback_retries_once_with_default(ig, monkeypatch, capsys):
+    seen = []
+
+    def fake_generate(*a, attention_mode="default", **kw):
+        seen.append(attention_mode)
+        if attention_mode == "sage":
+            raise AttributeError("'function' object has no attribute 'shape'")
+        return {"samples": 1}, 1.5
+
+    monkeypatch.setattr(ig, "generate_image", fake_generate)
+    lat, dur, fb = ig.generate_image_attn_safe("m", attention_mode="sage")
+    assert seen == ["sage", "default"]
+    assert (lat, dur, fb) == ({"samples": 1}, 1.5, "sage")
+    out = capsys.readouterr().out
+    assert "Attention mode 'sage' failed" in out and "default attention" in out
+
+
+def test_no_fallback_for_default_or_success(ig, monkeypatch):
+    monkeypatch.setattr(ig, "generate_image", lambda *a, attention_mode="default", **kw: ("L", 1.0))
+    assert ig.generate_image_attn_safe(attention_mode="sage") == ("L", 1.0, None)
+
+    def boom(*a, **kw):
+        raise RuntimeError("x")
+    monkeypatch.setattr(ig, "generate_image", boom)
+    with pytest.raises(RuntimeError):
+        ig.generate_image_attn_safe(attention_mode="default")
+
+
+def test_interrupt_is_not_retried(ig, monkeypatch):
+    class InterruptProcessingException(Exception):
+        pass
+    n = []
+
+    def boom(*a, **kw):
+        n.append(1)
+        raise InterruptProcessingException()
+    monkeypatch.setattr(ig, "generate_image", boom)
+    with pytest.raises(InterruptProcessingException):
+        ig.generate_image_attn_safe(attention_mode="sage")
+    assert len(n) == 1