diff --git a/samplers.py b/samplers.py index eda95cfa..cc3bfd03 100644 --- a/samplers.py +++ b/samplers.py @@ -600,6 +600,62 @@ def _apply_quality_sharpness(denoised, prev_denoised, sharpness): return denoised +# --------------------------------------------------------------------------- +# Blackwell (sm_120) xformers masked-attention fallback +# --------------------------------------------------------------------------- +# On GPUs newer than capability (9, 0), xformers' memory_efficient_attention has no +# kernel that accepts a tensor attn_bias — every backend is rejected ("No operator +# found ... too new ... attn_bias type is "). Unmasked attention +# still works, which is why i2v ANCHOR scenes generate fine but GUIDE scenes crash: the +# LTX guide path (comfy/ldm/lightricks/model.py _attention_with_guide_mask) passes a +# tensor mask. PyTorch SDPA handles masks on every GPU, so when the active backend is +# xformers on such a device we route only the MASKED calls to attention_pytorch via +# ComfyUI's per-call optimized_attention_override hook (wrap_attn). Unmasked calls stay +# on the fast xformers path. This is the programmatic equivalent of the user's +# --use-pytorch-cross-attention, scoped to masked attention and installed ONLY when this +# exact failure condition is detected, so other setups are untouched. +def _funpack_mask_safe_attention_override(func, *args, **kwargs): + # `func` is the raw backend ComfyUI selected (attention_xformers). wrap_attn has + # already stamped _inside_attn_wrapper into kwargs before calling us, so invoking + # attention_pytorch here cannot re-enter the override (no recursion). + mask = kwargs.get("mask") + if mask is None and len(args) >= 5: + mask = args[4] # (q, k, v, heads, mask, ...) + if mask is not None: + import comfy.ldm.modules.attention as _am + return _am.attention_pytorch(*args, **kwargs) + return func(*args, **kwargs) + + +def _funpack_install_mask_safe_attention(model): + """When the active attention backend is xformers on a GPU too new for its tensor-bias + kernels (capability > (9, 0), e.g. Blackwell sm_120), install a per-call override so + MASKED attention (the LTX guide path) falls back to SDPA. No-op on supported GPUs, when + ComfyUI already uses a mask-capable backend, or when another override is present. + Idempotent; never raises.""" + try: + import comfy.ldm.modules.attention as _am + if getattr(_am, "optimized_attention", None) is not getattr(_am, "attention_xformers", None): + return # SDPA / sage / flash / split selected — masks already handled + if not torch.cuda.is_available(): + return + major, minor = torch.cuda.get_device_capability() + if major <= 9: + return # xformers cutlass/flash tensor-bias kernels support cap <= (9, 0) + to = model.model_options.setdefault("transformer_options", {}) + existing = to.get("optimized_attention_override") + if existing is _funpack_mask_safe_attention_override: + return # already installed (idempotent across scenes/runs) + if existing is not None: + return # respect another tool's override; don't stomp it + to["optimized_attention_override"] = _funpack_mask_safe_attention_override + print(f"[FunPack AV] xformers active on capability {major}.{minor} GPU (too new for " + "its tensor-bias kernels): routing MASKED attention to PyTorch SDPA so guide " + "scenes don't crash. Unmasked attention stays on xformers.") + except Exception as _e: + print(f"[FunPack AV] mask-safe attention install skipped: {_e}") + + # --------------------------------------------------------------------------- # KV-Lock: variance-gated scheduler for BachVid K/V injection (arxiv 2603.09657) # @@ -3768,6 +3824,12 @@ class FunPackLTXAVSceneChainSampler: except Exception as _e: print(f"[FunPackLTXAVSceneChainSampler] hook strip failed: {_e}") + # Blackwell (sm_120) GPUs can't run xformers attention with a tensor mask; the LTX + # guide path uses one, so anchor scenes generate but guide scenes crash. Route masked + # attention to SDPA when that exact combo is detected (no-op otherwise). Threaded via + # transformer_options so it reaches every scene's model forward. + _funpack_install_mask_safe_attention(model) + # Decode-time noise (folded in from LTXV's Set VAE Decoder Noise per the boundary law: # the Chain Sampler owns IMAGES decode, so this lives here, not on a separate node). # Stamp settings onto a shallow copy of the VAE so we never mutate the shared input. diff --git a/tests/test_mask_safe_attention.py b/tests/test_mask_safe_attention.py new file mode 100644 index 00000000..0def0c0f --- /dev/null +++ b/tests/test_mask_safe_attention.py @@ -0,0 +1,138 @@ +"""Unit tests for the Blackwell (sm_120) xformers masked-attention fallback in samplers.py. + +On GPUs newer than capability (9, 0) xformers has no kernel that accepts a tensor attn_bias, +so i2v ANCHOR scenes (unmasked attention) generate fine but GUIDE scenes (LTX guide path +passes a tensor mask) crash. _funpack_install_mask_safe_attention installs a per-call override +(ComfyUI wrap_attn) that routes only MASKED attention to SDPA when that exact combo is present. + +Covers: override routing (masked→pytorch, unmasked→original, mask as kwarg or positional), +and the install gating (only xformers + too-new GPU; respects/ignores other backends, existing +overrides, missing CUDA; idempotent). +""" +import sys +import types +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +# Minimal comfy stubs so `import samplers` works without a full ComfyUI env. +for _name in ( + "comfy", "comfy.k_diffusion", "comfy.k_diffusion.sampling", + "comfy.model_sampling", "comfy.nested_tensor", "comfy.sample", + "comfy.samplers", "comfy.utils", +): + sys.modules.setdefault(_name, types.ModuleType(_name)) +sys.modules["comfy.nested_tensor"].NestedTensor = object + +import samplers # noqa: E402 + + +def _install_attn_module(*, backend="xformers", with_pytorch_sentinel=True): + """Register a fake comfy.ldm.modules.attention with selectable active backend.""" + am = types.ModuleType("comfy.ldm.modules.attention") + + def attention_xformers(*a, **kw): + return "xformers" + + def attention_pytorch(*a, **kw): + return "pytorch" + + am.attention_xformers = attention_xformers + am.attention_pytorch = attention_pytorch + am.optimized_attention = attention_xformers if backend == "xformers" else attention_pytorch + for parent in ("comfy", "comfy.ldm", "comfy.ldm.modules", "comfy.ldm.modules.attention"): + sys.modules.setdefault(parent, types.ModuleType(parent)) + sys.modules["comfy.ldm.modules.attention"] = am + return am + + +def _fake_model(): + return types.SimpleNamespace(model_options={}) + + +def _force_capability(monkeypatch, cap, available=True): + cuda = types.SimpleNamespace( + is_available=lambda: available, + get_device_capability=lambda *a, **k: cap, + ) + monkeypatch.setattr(samplers.torch, "cuda", cuda, raising=False) + + +# ── override routing ────────────────────────────────────────────────────────── +def test_override_routes_masked_to_pytorch_kwarg(): + _install_attn_module() + sentinel = object() + out = samplers._funpack_mask_safe_attention_override( + lambda *a, **k: "xformers", "q", "k", "v", 8, mask=sentinel, + ) + assert out == "pytorch" + + +def test_override_routes_masked_to_pytorch_positional(): + _install_attn_module() + # (q, k, v, heads, mask) — mask as the 5th positional arg. + out = samplers._funpack_mask_safe_attention_override( + lambda *a, **k: "xformers", "q", "k", "v", 8, object(), + ) + assert out == "pytorch" + + +def test_override_passes_unmasked_to_original(): + _install_attn_module() + out = samplers._funpack_mask_safe_attention_override( + lambda *a, **k: "xformers", "q", "k", "v", 8, mask=None, + ) + assert out == "xformers" + + +# ── install gating ──────────────────────────────────────────────────────────── +def test_installs_on_xformers_and_blackwell(monkeypatch): + _install_attn_module(backend="xformers") + _force_capability(monkeypatch, (12, 0)) + m = _fake_model() + samplers._funpack_install_mask_safe_attention(m) + assert (m.model_options["transformer_options"]["optimized_attention_override"] + is samplers._funpack_mask_safe_attention_override) + + +def test_no_install_on_supported_capability(monkeypatch): + _install_attn_module(backend="xformers") + _force_capability(monkeypatch, (9, 0)) # Hopper — xformers tensor-bias kernels OK + m = _fake_model() + samplers._funpack_install_mask_safe_attention(m) + assert "optimized_attention_override" not in m.model_options.get("transformer_options", {}) + + +def test_no_install_when_backend_not_xformers(monkeypatch): + _install_attn_module(backend="pytorch") # user already on SDPA + _force_capability(monkeypatch, (12, 0)) + m = _fake_model() + samplers._funpack_install_mask_safe_attention(m) + assert "optimized_attention_override" not in m.model_options.get("transformer_options", {}) + + +def test_no_install_without_cuda(monkeypatch): + _install_attn_module(backend="xformers") + _force_capability(monkeypatch, (12, 0), available=False) + m = _fake_model() + samplers._funpack_install_mask_safe_attention(m) + assert "optimized_attention_override" not in m.model_options.get("transformer_options", {}) + + +def test_does_not_stomp_existing_override(monkeypatch): + _install_attn_module(backend="xformers") + _force_capability(monkeypatch, (12, 0)) + other = lambda func, *a, **k: func(*a, **k) + m = types.SimpleNamespace(model_options={"transformer_options": {"optimized_attention_override": other}}) + samplers._funpack_install_mask_safe_attention(m) + assert m.model_options["transformer_options"]["optimized_attention_override"] is other + + +def test_idempotent(monkeypatch): + _install_attn_module(backend="xformers") + _force_capability(monkeypatch, (12, 0)) + m = _fake_model() + samplers._funpack_install_mask_safe_attention(m) + samplers._funpack_install_mask_safe_attention(m) + assert (m.model_options["transformer_options"]["optimized_attention_override"] + is samplers._funpack_mask_safe_attention_override)