From bc0cca324c62baf385f89f195c8a7bf8a58bff8d Mon Sep 17 00:00:00 2001 From: Jason Hoku Date: Wed, 4 Nov 2026 17:29:46 -0800 Subject: [PATCH] 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 --- README.md | 10 +-- generation_orchestrator.py | 13 +++- image_generation.py | 40 +++++++++++- tests/conftest.py | 1 + tests/test_issue16_attention.py | 109 ++++++++++++++++++++++++++++++++ 5 files changed, 164 insertions(+), 9 deletions(-) create mode 100644 tests/test_issue16_attention.py diff --git a/README.md b/README.md index ddaf971..5762b6f 100644 --- a/README.md +++ b/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. diff --git a/generation_orchestrator.py b/generation_orchestrator.py index afd39ef..1c9f303 100644 --- a/generation_orchestrator.py +++ b/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) diff --git a/image_generation.py b/image_generation.py index 17ec4c5..c78418a 100644 --- a/image_generation.py +++ b/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: diff --git a/tests/conftest.py b/tests/conftest.py index 5914197..9400858 100644 --- a/tests/conftest.py +++ b/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: {}, diff --git a/tests/test_issue16_attention.py b/tests/test_issue16_attention.py new file mode 100644 index 0000000..dc39c66 --- /dev/null +++ b/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