olivv-cs-ComfyUI-FunPack/tests/test_mask_safe_attention.py

"""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)