MiniMax H3-Image: REINS and the Attention explorer work on it; the explorer skips the sound and clip regions on a model without sound

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

modules/models/minimax_h3_image/__init__.py

@@ -13,7 +13,7 @@ Everything the model needs at run time (the image latent, the picture labels, th
 shipped inside the file) is ComfyUI's own.
 """
 
-from ..._core import traits as _traits
+from ..._core import dit_hooks as _dit_hooks, traits as _traits
 
 ID = "model_minimax_h3_image"
 TITLE = "MiniMax H3-Image"
@@ -39,12 +39,33 @@ def detect(keys) -> bool:
     return any(all(f"{p}{k}" in keyset for k in _PRESENT) and f"{p}{_ABSENT}" not in keyset for p in _PREFIXES)
 
 
+MODEL_CLASS = "comfy.ldm.minimax.model.MiniMaxH3ImageModel"
+
+#: What this model announces beyond core's own reading: its blocks are the H3 blocks (REINS reads the picture's rows
+#: through them, the Attention explorer scales their attention). Block repeat and Attention temperature work the same way
+#: but nobody has tried them on this model, so they stay hidden.
+ANNOUNCED = ["block_rows", "block_scope", "hides:block_repeat", "hides:attention_temperature"]
+
+
+def is_image_model(model) -> bool:
+    return _traits.has_block(model, MODEL_CLASS)
+
+
+def traits(model):
+    return ANNOUNCED if is_image_model(model) else ()
+
+
+def probe_blocks(keys):
+    """How many transformer blocks the file has, read off its header."""
+    return _dit_hooks.count_indexed(keys, tuple(f"{p}blocks." for p in _PREFIXES))
+
+
 def probe_traits(keys) -> list:
     """What core reads off the loaded model (a 3-axis latent with a length-1 time axis, flow prediction), said before it loads."""
     if not detect(keys):
         return []
     from comfy.latent_formats import MiniMaxH3Video
-    found = ["predict_flow"]
+    found = ["predict_flow", *ANNOUNCED]
     rank = _traits.LATENT_RANK.get(MiniMaxH3Video.latent_dimensions)
     if rank:
         found.append(rank)
@@ -113,4 +134,5 @@ def presets():
              "slots": text_to_image(), "project": PROJECT}]
 
 
-PROVIDES = {"detect": detect, "probe_traits": probe_traits, "pipeline_presets": presets}
+TRAITS = traits
+PROVIDES = {"detect": detect, "probe_traits": probe_traits, "probe_blocks": probe_blocks, "pipeline_presets": presets}

modules/models/minimax_h3_image/tests/test_features.py

@@ -0,0 +1,130 @@
+"""REINS and the Attention explorer on a real (tiny, random-weight) H3-Image forward: they run, and they do not touch a sound that is not there."""
+
+import random
+
+import pytest
+import torch
+
+
+@pytest.fixture(autouse=True)
+def root(comfyui, tmp_path, monkeypatch):
+    from modules.system.taste import store
+    monkeypatch.setattr(store, "ROOT", tmp_path / "taste")
+    monkeypatch.setattr(store, "current_prompt_id", lambda: "run-1")
+
+
+class TinyImage:
+    """ComfyUI's own MiniMaxH3ImageModel, four blocks wide, on the CPU."""
+
+    def __init__(self, blocks=4):
+        import comfy.ops
+        from comfy.latent_formats import MiniMaxH3Video
+        from comfy.ldm.minimax.model import MiniMaxH3ImageModel
+        from comfy.model_patcher import ModelPatcher
+
+        torch.manual_seed(0)
+        dm = MiniMaxH3ImageModel(
+            pdd_sigmas=[1.0, 0.5, 0.0], hidden_size=64, num_layers=blocks, token_refiner_num_layers=1,
+            num_attention_heads=2, attention_head_dim=32, ffn_hidden_size=64, text_dim=48, timestep_input_dim=32,
+            time_embed_hidden_size=64, time_embed_dim=32, rope_inv_freq_len=4, dtype=torch.float32,
+            operations=comfy.ops.manual_cast)
+        for p in dm.parameters():
+            torch.nn.init.normal_(p, std=0.05)
+            p.requires_grad_(False)
+        dm.rope.inv_freq.copy_(torch.rand(4))
+
+        class Config:
+            latent_format = MiniMaxH3Video
+
+        class Base(torch.nn.Module):
+            pass
+
+        base = Base()
+        base.diffusion_model, base.model_config = dm, Config()
+        self.patcher = ModelPatcher(base, load_device=torch.device("cpu"), offload_device=torch.device("cpu"))
+        self.latent, self.context = torch.randn(1, 24, 1, 4, 4), torch.randn(1, 5, 48)
+
+    def run(self, patcher=None, sigma=0.5):
+        import comfy.patcher_extension as ext
+        patcher = patcher or self.patcher
+        to = dict(patcher.model_options.get("transformer_options", {}))
+        to["wrappers"] = ext.copy_nested_dicts(to.get("wrappers", {}))
+        ext.merge_nested_dicts(to["wrappers"], patcher.wrappers, copy_dict1=False)
+        to.setdefault("sample_sigmas", torch.tensor([1.0, sigma, 0.0]))
+        to["sigmas"] = torch.tensor([sigma])
+        with torch.inference_mode():
+            return patcher.model.diffusion_model(self.latent, torch.tensor([sigma * 1000.0]), self.context, transformer_options=to)
+
+
+def _load(tiny, **settings):
+    from modules.sampling.modifiers.nodes import FunUILoadModifiers
+    return FunUILoadModifiers.execute(tiny.patcher, {"taste": {"key": "fox"}, **settings}).result
+
+
+def _through(patched, fn):
+    from comfy.patcher_extension import WrappersMP
+    call = fn
+    for ws in reversed([w for ws in patched.wrappers.get(WrappersMP.OUTER_SAMPLE, {}).values() for w in ws]):
+        call = (lambda w, inner: (lambda: w(lambda: inner())))(ws, call)
+    return call()
+
+
+def test_the_loaded_model_announces_what_the_two_features_need():
+    from core import traits
+    from modules.models import minimax_h3, minimax_h3_image as m
+
+    tiny = TinyImage()
+    assert m.is_image_model(tiny.patcher) and not minimax_h3.is_h3(tiny.patcher)
+    found = set(traits.announced(tiny.patcher)) if hasattr(traits, "announced") else set(m.traits(tiny.patcher))
+    assert {"block_rows", "block_scope"} <= found and "audio_stream" not in found
+
+
+def test_reins_captures_the_picture_rows_and_a_learned_direction_moves_the_picture():
+    from modules.system.taste import store
+    tiny = TinyImage()
+    base = tiny.run()
+    patched, status = _load(tiny, reins={"enabled": True, "blocks": "3", "strength": 0.5})
+    assert "reins" in status.splitlines()[0]
+    _through(patched, lambda: tiny.run(patched, sigma=0.5))        # the last step: captured
+    assert store.rate("run-1", "liked")["recorded"] == ["reins"]
+    hidden = tiny.patcher.model.diffusion_model.blocks[0].attn.out_proj.weight.shape[0]
+    liked = torch.zeros(hidden)
+    liked[0] = 1.0
+    for i, (vec, rating) in enumerate([(liked, "liked"), (-liked, "disliked")] * 2):
+        store.capture("fox", "reins", {3: vec}, prompt_id=f"t{i}")
+        store.rate(f"t{i}", rating)
+    moved, _ = _load(tiny, reins={"enabled": True, "blocks": "3", "strength": 0.5})
+    assert not torch.allclose(_through(moved, lambda: tiny.run(moved)), base)
+    quiet, _ = _load(tiny, reins={"enabled": True, "blocks": "3", "strength": 0.0})
+    assert torch.equal(_through(quiet, lambda: tiny.run(quiet)), base)
+
+
+def test_the_attention_explorer_changes_the_picture_with_no_sound_and_no_clip_regions(monkeypatch):
+    from core import log
+    from modules.sampling import attention_bandit as m
+    tiny = TinyImage()
+    base = tiny.run()
+    patched, _ = _load(tiny, attention_bandit={"enabled": True, "mode": "manual", "range": 0.4, "blocks": "0-3", "sound": False})
+    monkeypatch.setattr(m, "draw", lambda table, rng=random, **kw: [2] * len(table))
+    log.new_run()
+    out = _through(patched, lambda: tiny.run(patched))
+    assert not torch.equal(out, base), "the explorer ran and changed nothing"
+    said = " ".join(e["message"] for e in log.history())
+    assert "clip regions" not in said and "sound could not be told apart" not in said and "changed nothing" not in said
+    monkeypatch.setattr(m, "draw", lambda table, rng=random, **kw: [m.STAY] * len(table))
+    plain, _ = _load(tiny, attention_bandit={"enabled": True, "blocks": "0-3"})
+    assert torch.equal(_through(plain, lambda: tiny.run(plain)), base)
+
+
+def test_what_was_rated_with_a_sound_stream_does_not_teach_a_model_without_one():
+    from modules.sampling import attention_bandit as m
+    tiny = TinyImage()
+    sig = []
+    orig = m.make_plan
+    m.make_plan = lambda rows, cond, signature, *a, **k: (sig.append(list(signature)), orig(rows, cond, signature, *a, **k))[1]
+    try:
+        patched, _ = _load(tiny, attention_bandit={"enabled": True, "blocks": "0-3"})
+        _through(patched, lambda: tiny.run(patched))
+    finally:
+        m.make_plan = orig
+    assert sig and sig[0][-1] == 2

modules/sampling/attention_bandit/__init__.py

@@ -400,6 +400,10 @@ def install(patcher, values, key):
     video_width, sound_width = getattr(dm, "inner_dim", None), getattr(dm, "audio_inner_dim", None)
     if paired:
         signature.append(1)            # ratings made on a packed-sequence model (H3) must not teach a paired one
+    silent = not paired and getattr(dm, "audio_latents_dim", 1) is None      # a picture-only model (H3-Image): no sound to spare, no clip time to cut into regions
+    if silent:
+        regions, sound = False, True
+        signature.append(2)            # and what was rated on a model with a sound stream must not teach one without
     if focus and paired:
         focus = False                  # the prompt's rows are not one run of the sequence here: that part is H3's
     if not sound and not paired: