Randomizer (attention temperature) now works on LTX: new block_scope trait, LTX blocks counted

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

core/dit_hooks.py

@@ -121,6 +121,8 @@ def block_count(model) -> int:
     """How many hookable blocks the model has, or 0 when it cannot be read."""
     root = getattr(getattr(model, "model", None), "diffusion_model", None)
     blocks = getattr(root, "blocks", None)
+    if blocks is None:
+        blocks = getattr(root, "transformer_blocks", None)       # LTX names them so
     try:
         return len(blocks)
     except TypeError:

modules/models/ltx/__init__.py

@@ -6,10 +6,11 @@ folder mentions LTX.
 
 What it contributes, and what it deliberately does NOT:
 
-* **Traits** -- a joint audio+video model with a compressed time axis. It does NOT announce
-  `dit_block_hooks`: the block hooks other modules install (attention temperature, block repeat...)
-  are written for H3's packed block layout, and LTX's blocks take `(video, audio)` as a pair. Until
-  a module is written for that layout it is absent here, said, rather than half-working.
+* **Traits** -- a joint audio+video model with a compressed time axis. It announces `block_scope` (a
+  hook may hand a block changed options: the Randomizer) but NOT `dit_block_hooks`: the hooks that edit
+  rows (REINS, block repeat...) are written for H3's packed block layout, and LTX's blocks take
+  `(video, audio)` as a pair. Until a module is written for that layout it is absent here, said,
+  rather than half-working.
 * **An empty latent** -- video 128 channels at /32 and /8, plus an audio latent whose length comes
   from the AUDIO VAE and the frame rate. That is why this provider asks for `audio_vae` and
   `frame_rate`: without them the audio half cannot be sized, and a video-only latent would make a
@@ -66,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"]
+    return ["audio_stream", "temporal_latent", "temporal_compression", "ltx_av", "block_scope"]
 
 
 def traits(model):
@@ -74,7 +75,11 @@ def traits(model):
         return ()
     # ltx_av: transformer_blocks that take (video, audio) as a pair and have the video<->audio
     # cross-attention submodules. What the LTX-only modules require.
-    return ["audio_stream", "ltx_av"]
+    # block_scope: its forward honours patches_replace["dit"][("double_block", i)] and routes attention
+    # 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"]
 
 
 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", "adaln_modalities", *H3_ONLY]
+    found = ["audio_stream", "dit_block_hooks", "block_scope", "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", *H3_ONLY]
+    found = ["audio_stream", "dit_block_hooks", "block_scope", *H3_ONLY]
     if has_block(model, ADALN_CLASS):
         found.append("adaln_modalities")
     return found

modules/sampling/attention_temperature/__init__.py

@@ -22,7 +22,7 @@ MOUNT = "generation.sampling"
 STAGE = "sampling"
 CATEGORY = "sampling"
 STATUS = "proven"
-REQUIRES = ["dit_block_hooks"]
+REQUIRES = ["block_scope"]
 
 SETTINGS = {
     "enabled": {

modules/sampling/attention_temperature/tests/test_on_ltx.py

@@ -0,0 +1,36 @@
+"""The randomizer on a real (tiny) LTX forward: its blocks take (video, audio) as a pair."""
+
+import pytest
+import torch
+
+
+@pytest.fixture(autouse=True)
+def _needs_comfy(comfyui):
+    """Imports comfy."""
+
+
+def test_it_is_offered_on_ltx_and_changes_the_output(tiny_ltx):
+    from modules.sampling.modifiers.nodes import FunUILoadModifiers
+    base_v, base_a = tiny_ltx.run()
+    patched, status = FunUILoadModifiers.execute(
+        tiny_ltx.patcher, {"attention_temperature": {"enabled": True, "amount": -0.5, "blocks": "0-1"}}).result
+    assert "attention_temperature: 0.5x (sharper) at blocks 0-1" in status
+    v, a = tiny_ltx.run(patched)
+    assert not torch.allclose(v, base_v), "temperature ran and changed nothing"
+    assert v.shape == base_v.shape and a.shape == base_a.shape
+
+
+def test_blocks_past_the_models_last_are_said_not_silently_dropped(tiny_ltx):
+    from modules.sampling.modifiers.nodes import FunUILoadModifiers
+    from core import dit_hooks
+    n = dit_hooks.block_count(tiny_ltx.patcher)
+    assert n > 0, "LTX's blocks were not counted"
+    _p, status = FunUILoadModifiers.execute(
+        tiny_ltx.patcher, {"attention_temperature": {"enabled": True, "amount": -0.5, "blocks": f"0-{n + 3}"}}).result
+    assert f"0-{n - 1}" in status
+
+
+def test_row_editing_modules_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