REINS and Block repeat work on LTX's (video, audio) blocks: new block_rows trait

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Tw9fjYZhHqk196LykrWy3T

conftest.py

@@ -241,6 +241,9 @@ class TinyLTX:
                       torch.tensor([sigma * 1000.0]), self.context, frame_rate=25, transformer_options=to)
 
 
+TinyLTX.sample = TinyH3.sample           # `run` inside the patcher's OUTER_SAMPLE wrappers, the same way
+
+
 _TINY_LTX = []
 
 

modules/models/ltx/__init__.py

@@ -67,7 +67,7 @@ def probe_traits(keys) -> list:
     """What `traits()` would say, read before anything loads."""
     if not detect(keys):
         return []
-    return ["audio_stream", "temporal_latent", "temporal_compression", "ltx_av", "block_scope"]
+    return ["audio_stream", "temporal_latent", "temporal_compression", "ltx_av", "block_scope", "block_rows"]
 
 
 def traits(model):
@@ -79,7 +79,9 @@ def traits(model):
     # through optimized_attention with the block's transformer_options, so a hook that only SCOPES a
     # block (hands it changed transformer_options) works. Not dit_block_hooks: a block here takes and
     # returns (video, audio), where the hooks that edit rows expect one packed tensor.
-    return ["audio_stream", "ltx_av", "block_scope"]
+    # block_rows: the hooks that read or push the video rows of a block's output (REINS, Block repeat)
+    # handle the pair too (see those modules).
+    return ["audio_stream", "ltx_av", "block_scope", "block_rows"]
 
 
 def frames_for(length: int) -> int:

modules/models/minimax_h3/__init__.py

@@ -122,7 +122,7 @@ def probe_traits(keys) -> list:
         return []
     from comfy.latent_formats import MiniMaxH3AV
 
-    found = ["audio_stream", "dit_block_hooks", "block_scope", "adaln_modalities", *H3_ONLY]
+    found = ["audio_stream", "dit_block_hooks", "block_scope", "block_rows", "adaln_modalities", *H3_ONLY]
     rank = _traits.LATENT_RANK.get(MiniMaxH3AV.latent_dimensions)
     if rank:
         found.append(rank)
@@ -138,7 +138,7 @@ def traits(model):
     # dit_block_hooks: its forward honours patches_replace["dit"][("double_block",
     # i)] and routes every attention call through optimized_attention, which is
     # what core/dit_hooks.py hooks.
-    found = ["audio_stream", "dit_block_hooks", "block_scope", *H3_ONLY]
+    found = ["audio_stream", "dit_block_hooks", "block_scope", "block_rows", *H3_ONLY]
     if has_block(model, ADALN_CLASS):
         found.append("adaln_modalities")
     return found

modules/sampling/attention_temperature/tests/test_on_ltx.py

@@ -30,7 +30,7 @@ def test_blocks_past_the_models_last_are_said_not_silently_dropped(tiny_ltx):
     assert f"0-{n - 1}" in status
 
 
-def test_row_editing_modules_stay_off_ltx(tiny_ltx):
+def test_modules_written_for_h3s_packed_rows_stay_off_ltx(tiny_ltx):
     from modules.sampling.modifiers.nodes import FunUILoadModifiers
-    _p, status = FunUILoadModifiers.execute(tiny_ltx.patcher, {"reins": {"enabled": True}}).result
-    assert "reins: needs dit_block_hooks" in status
+    _p, status = FunUILoadModifiers.execute(tiny_ltx.patcher, {"stas": {"enabled": True}}).result
+    assert "stas: needs dit_block_hooks" in status

modules/sampling/block_repeat/__init__.py

@@ -35,7 +35,7 @@ MOUNT = "generation.sampling"
 STAGE = "sampling"
 CATEGORY = "sampling"
 STATUS = "experimental"
-REQUIRES = ["dit_block_hooks"]
+REQUIRES = ["block_rows"]
 
 SETTINGS = {
     "enabled": {
@@ -121,6 +121,10 @@ def install(patcher, values, key):
     steps = int(values.get("steps", 0))
     video_only = bool(values.get("video_only"))
     span = values.get("shape") == "span"
+    paired = hasattr(getattr(patcher.model, "diffusion_model", None), "transformer_blocks")      # audio+video blocks: (video, audio) in and out
+    if span and paired:
+        _say("this model's blocks take picture and sound as a pair, so 'whole span' runs as 'each block'")
+        span = False
 
     if span:
         lo, hi = blocks[0], blocks[-1]
@@ -154,7 +158,20 @@ def install(patcher, values, key):
             dit_hooks.add_block_hook(patcher, key, i, skip)
         what = f"blocks {lo}-{hi} loop {times + 1}x as one span"
     else:
+        def repeat_pair(args, extra):
+            once = extra["original_block"](args)["img"]
+            if not _in_window(args, steps):
+                return {"img": once}
+            if video_only:         # these blocks add their residual into their input: keep the single pass whole
+                once = tuple(t.clone() for t in once)
+            h = once
+            for _ in range(times):
+                h = extra["original_block"]({**args, "img": h})["img"]
+            return {"img": (h[0], once[1]) if video_only else h}
+
         def repeat(args, extra):
+            if paired:
+                return repeat_pair(args, extra)
             once = extra["original_block"](args)["img"]
             if not _in_window(args, steps):
                 return {"img": once}

modules/sampling/block_repeat/tests/test_on_ltx.py

@@ -0,0 +1,49 @@
+"""Block repeat on a (tiny) LTX forward: picture and sound go through each block as a pair."""
+
+import pytest
+import torch
+
+
+@pytest.fixture(autouse=True)
+def _needs_comfy(comfyui):
+    """Imports comfy."""
+
+
+def _load(tiny, **values):
+    from modules.sampling.modifiers.nodes import FunUILoadModifiers
+    return FunUILoadModifiers.execute(tiny.patcher, {"block_repeat": {"enabled": True, **values}}).result
+
+
+def _calls(tiny, patched):
+    counts, handles = {}, []
+    for i, block in enumerate(tiny.patcher.model.diffusion_model.transformer_blocks):
+        handles.append(block.register_forward_pre_hook(lambda m, a, kw=None, i=i: counts.__setitem__(i, counts.get(i, 0) + 1)))
+    try:
+        tiny.run(patched)
+    finally:
+        for h in handles:
+            h.remove()
+    return counts
+
+
+def test_each_named_block_runs_times_plus_one(tiny_ltx):
+    patched, status = _load(tiny_ltx, blocks="1", times=2)
+    assert "block_repeat: blocks 1 run 3x each" in status
+    counts = _calls(tiny_ltx, patched)
+    assert counts[1] == 3 and counts[0] == 1
+
+
+def test_it_changes_the_picture_and_video_only_leaves_the_sound_single_pass(tiny_ltx):
+    base_v, base_a = tiny_ltx.run()
+    both, _ = _load(tiny_ltx, blocks="1")
+    v, a = tiny_ltx.run(both)
+    assert not torch.allclose(v, base_v) and not torch.allclose(a, base_a)
+    only, _ = _load(tiny_ltx, blocks="1", video_only=True)
+    v2, a2 = tiny_ltx.run(only)
+    assert not torch.allclose(v2, base_v)
+    assert v2.shape == base_v.shape and a2.shape == base_a.shape
+
+
+def test_whole_span_runs_as_each_block_here(tiny_ltx):
+    patched, status = _load(tiny_ltx, blocks="1-2", shape="span")
+    assert "run 2x each" in status

modules/sampling/reins/__init__.py

@@ -23,7 +23,7 @@ MOUNT = "generation.sampling"
 STAGE = "guidance"
 CATEGORY = "guidance"
 STATUS = "experimental"
-REQUIRES = ["dit_block_hooks"]
+REQUIRES = ["block_rows"]
 USES = ["taste_store"]
 
 SETTINGS = {
@@ -87,27 +87,32 @@ def install(patcher, values, key):
     def make(block):
         def hook(args, extra):
             out = extra["original_block"](args)["img"]
-            span = video_span(args, out)
-            if span is None:
-                return {"img": out}
-            rows = out[span[0]:span[1]]
+            if isinstance(out, (tuple, list)):         # an audio+video block hands back (video, audio): the video is all picture
+                rows, width = out[0].reshape(-1, out[0].shape[-1]), out[0].shape[-1]
+            else:
+                span = video_span(args, out)
+                if span is None:
+                    return {"img": out}
+                rows, width = out[span[0]:span[1]], out.shape[-1]
             if dit_hooks.last_step(args.get("transformer_options")):
                 captured[block] = rows.mean(0, dtype=torch.float32).detach()
             direction = live["dirs"].get(block)
             if direction is None or strength <= 0.0:
                 return {"img": out}
-            if direction.numel() != out.shape[-1]:
+            if direction.numel() != width:
                 _say(f"block {block}'s learned direction is from a different model; "
                      "learning only there")
                 return {"img": out}
-            d = cast.get((block, out.dtype))
+            d = cast.get((block, rows.dtype))
             if d is None:
-                d = cast[(block, out.dtype)] = direction.to(device=out.device, dtype=out.dtype)
+                d = cast[(block, rows.dtype)] = direction.to(device=rows.device, dtype=rows.dtype)
             norm = torch.linalg.vector_norm(rows, dim=-1, dtype=torch.float32).mean()
             # No host sync: a bad norm makes the push zero instead of a branch.
             scale = torch.where(torch.isfinite(norm) & (norm > 0), norm * strength,
                                 torch.zeros_like(norm))
-            rows += d * scale.to(out.dtype)
+            rows += d * scale.to(rows.dtype)
+            if isinstance(out, (tuple, list)) and rows.data_ptr() != out[0].data_ptr():
+                out[0].copy_(rows.view_as(out[0]))         # the reshape had to copy: put the push back
             return {"img": out}
         return hook
 

modules/sampling/reins/tests/test_on_ltx.py

@@ -0,0 +1,41 @@
+"""Taste steering on a (tiny) LTX forward: its blocks hand back (video, audio)."""
+
+import pytest
+import torch
+
+
+@pytest.fixture(autouse=True)
+def _needs_comfy(comfyui):
+    """Imports comfy."""
+
+
+@pytest.fixture(autouse=True)
+def root(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")
+
+
+def _width(tiny):
+    return tiny.patcher.model.diffusion_model.transformer_blocks[0].attn1.to_out[0].weight.shape[0]
+
+
+def _load(tiny, **values):
+    from modules.sampling.modifiers.nodes import FunUILoadModifiers
+    return FunUILoadModifiers.execute(
+        tiny.patcher, {"reins": {"enabled": True, **values}, "taste": {"key": "fox"}}).result
+
+
+def test_it_learns_at_every_block_and_steers_only_the_picture(tiny_ltx):
+    from modules.system.taste import store
+    liked = torch.zeros(_width(tiny_ltx))
+    liked[0] = 1.0
+    for i, (vec, rating) in enumerate([(liked, "liked"), (-liked, "disliked")] * 2):
+        store.capture("fox", "reins", {1: vec}, prompt_id=f"t{i}")
+        store.rate(f"t{i}", rating)
+    base_v, base_a = tiny_ltx.run()
+    patched, status = _load(tiny_ltx, blocks="1", strength=0.5)
+    assert "reins: learning at every block; steers 1" in status
+    v, a = tiny_ltx.sample(patched)
+    assert not torch.allclose(v, base_v), "the learned direction changed nothing"
+    assert v.shape == base_v.shape and a.shape == base_a.shape