diff --git a/modules/loaders/sla_attention.py b/modules/loaders/sla_attention.py index 143b863d..2f2eaf49 100644 --- a/modules/loaders/sla_attention.py +++ b/modules/loaders/sla_attention.py @@ -24,10 +24,22 @@ below exists and why a run that never sparsified says so in the log. """ import logging +from .._core import traits as _traits + SLA_NAME = "sla_h3" _H3_HEAD_DIM = 128 -_H3_MODEL_NAMES = ("MiniMaxH3Model",) +# Qualified on purpose, per core/traits.py's has_block: "a bare name matches on +# the class name alone... actively wrong" for a false positive that silently +# sparsifies a model this was never validated on -- exactly the LTX-collision +# risk this file's own docstring already warns about, just not guarded against +# strongly enough before this fix (v4 matched the bare name; found in extensive_ +# testing round 1 on this port). Kept as a local string matching modules/models/ +# minimax_h3's own MODEL_CLASS rather than importing that module -- it pulls in +# comfy_extras at import time, which this file deliberately does not require +# just to be importable on a machine with no ComfyUI on the path. Keep the two +# in sync if the H3 model class is ever renamed. +_H3_MODEL_CLASS = "comfy.ldm.minimax.model.MiniMaxH3Model" # Validated on an RTX 5090 at 768p/15s (ComfyUI-H3-SLA-Attention's own measurements): # 3.7x the attention throughput of stock ComfyUI at sparsity 0.90 / block 64. @@ -67,10 +79,12 @@ def is_h3_model(model): SLA is only sound on H3: the sparsity ratio it runs at is the one the turbo LoRA was distilled to tolerate, and H3's packed [text | cond | audio | video] sequence is what the prefix pinning is for. Head shape alone would also match LTX, and - silently sparsifying LTX attention is a quality loss with no LoRA compensating it. + silently sparsifying LTX attention is a quality loss with no LoRA compensating it + -- which is also why the class name is matched qualified (module + name), not + bare: a bare name is exactly as blind to an unrelated same-named class as it is + to LTX's different head shape. """ - diffusion = getattr(getattr(model, "model", None), "diffusion_model", None) - return type(diffusion).__name__ in _H3_MODEL_NAMES + return _traits.has_block(model, _H3_MODEL_CLASS) def new_state(): diff --git a/modules/loaders/tests/test_sla_attention.py b/modules/loaders/tests/test_sla_attention.py index 22aa6e97..39f5bbfa 100644 --- a/modules/loaders/tests/test_sla_attention.py +++ b/modules/loaders/tests/test_sla_attention.py @@ -190,8 +190,12 @@ def test_a_non_h3_model_is_not_handed_a_minimax_kwarg(): # ── the H3 gate ─────────────────────────────────────────────────────────────── class _Patcher: - def __init__(self, class_name): - inner = type(class_name, (), {})() + def __init__(self, class_name, module="comfy.ldm.minimax.model"): + # has_block (core/traits.py) matches a QUALIFIED class name -- module and + # name both -- so the stub needs a real __module__ and a modules() walk + # the way an nn.Module provides, not just a bare type() with a name. + cls = type(class_name, (), {"__module__": module, "modules": lambda self: [self]}) + inner = cls() self.model = types.SimpleNamespace(diffusion_model=inner) self.model_options = {} self.wrappers = [] @@ -214,6 +218,18 @@ def test_sla_refuses_a_model_that_is_not_h3_and_says_so(): assert "not a MiniMax H3 model" in note +def test_a_same_named_class_from_a_different_module_is_not_treated_as_h3(): + """A bare class-name match would treat any unrelated MiniMaxH3Model -- built by a + different custom node, or a future FunPack model reusing the name by accident -- + as the real thing, and sparsify it at H3's ratio with H3's prefix-pinning logic + on a model that was never validated for either. The qualified check requires the + DEFINING module too, not just the name.""" + model = _Patcher("MiniMaxH3Model", module="some_other_package.model") + out, note, installed = sla.install_sla(model) + assert out is model and installed is False + assert "not a MiniMax H3 model" in note + + def test_sla_reports_when_the_machine_cannot_run_it(): model = _Patcher("MiniMaxH3Model") out, note, installed = sla.install_sla(model)