Attention explorer on LTX: depth groups and clip regions (prompt focus stays H3)

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

CHANGELOG.md

@@ -14,6 +14,7 @@
 
 - **Temporal style: auto** picks a style for each shot from its own words (motion words punch, calm words hold, loop words loop).
 - **Randomizer, REINS and Block repeat on LTX** (new traits `block_scope` and `block_rows`); the other block-hook features stay H3-only. Whole-span repeat runs as each-block on LTX.
+- **Attention explorer on LTX**: the depth groups and the clip regions (four time bands of the picture, and the sound) work; Prompt focus stays H3-only and is left out on LTX.
 - **Identity bank** (LTX, needs a Taste key): keeps the attention look of the clips you like and leans on it in the next ones of the same size (BachVid), harder on steps where the picture is still changing (KV-Lock). Untested on a real LTX run.
 - **Guide frames** (LTX, Settings > Engine): lays the middle picture of the clip before, or one picture from each of the last clips (needs the JoyAI-Echo LoRA), into the shot as weak guides through ComfyUI's own LTX guide attention. The Image to Video preset has the new slot and crops the guides again before decoding.
 - **LTX 2.3 is tested** on a real model, with two presets: **Video to Video, direct** (the video is encoded and sampled again from part-way down the schedule; the sigmas slot, default `0.909375, 0.725, 0.421875, 0.0`, says how much changes) and **Image to Video, from References** (reference picture 1 is the start frame; nothing is picked twice). The earlier LTX presets (Text to Video, Image to Video, Anchor as guide, IC-LoRA) are gone. **Text to video is the fallback in both**: with no picture or video fed in, the anchor is bypassed (image to video) or a grey clip of the project's size and length is sampled with the whole schedule (video to video), and the run says so. Presets load separate files (diffusion model, video VAE, audio VAE, text encoder). LoRAs and Noise merging work on LTX; other enhancements are not tested there yet. Video to Video shows Width, Height and Length "(no video)" in Project settings for that fallback.

modules/sampling/attention_bandit/__init__.py

@@ -44,7 +44,7 @@ MOUNT = "generation.sampling"
 STAGE = "sampling"
 CATEGORY = "sampling"
 STATUS = "experimental"
-REQUIRES = ["dit_block_hooks"]
+REQUIRES = ["block_scope"]
 USES = ["taste_store"]
 
 SETTINGS = {
@@ -316,6 +316,28 @@ def region_spans(options):
     return [(edges[i], edges[i + 1]) for i in range(BANDS)] + [seg["audio"]]
 
 
+def paired_regions(q, factors, video_width, sound_width):
+    """q with each clip region's rows asking differently, for a model whose picture and sound attend in
+    calls of their own (rows run frame by frame): the picture's self-attention is cut into BANDS equal
+    time bands, the sound's self-attention is the last region. Any other call comes back as it was."""
+    width = int(q.shape[-1])
+    if width == video_width:
+        edges = [int(q.shape[1]) * i // BANDS for i in range(BANDS + 1)]
+        spans = list(zip(edges[:-1], edges[1:]))
+        mine = factors[:BANDS]
+    elif width == sound_width:
+        spans, mine = [(0, int(q.shape[1]))], factors[BANDS:BANDS + 1]
+    else:
+        return q
+    if all(f == 1.0 for f in mine):
+        return q
+    q = q.clone()
+    for (a, b), factor in zip(spans, mine):
+        if factor != 1.0:
+            q[:, a:b] /= factor
+    return q
+
+
 # Carried on the per-block copy of transformer_options, as Randomizer does, so an
 # attention call knows its block's factors without a shared flag.
 _SCOPE = "funui_attention_bandit"
@@ -343,6 +365,11 @@ def install(patcher, values, key):
             "text": 0, "text_wanted": False, "region_hits": 0, "region_wanted": False}
 
     captured = taste.collect(patcher, key, KIND, keep=MAX_ROWS)
+    dm = getattr(patcher.model, "diffusion_model", None)
+    paired = hasattr(dm, "transformer_blocks")                 # picture and sound attend in separate calls, one stream each
+    video_width, sound_width = getattr(dm, "inner_dim", None), getattr(dm, "audio_inner_dim", None)
+    if focus and paired:
+        focus = False                  # the prompt's rows are not one run of the sequence here: that part is H3's
 
     def scoped_for(index):
         def scoped(args, extra):
@@ -367,7 +394,10 @@ def install(patcher, values, key):
         live["scaled"] += 1
         if temperature != 1.0:
             q = q / temperature
-        if live["region_wanted"] and q.ndim == 4:
+        if live["region_wanted"] and paired and q.ndim == 3 and q.shape == k.shape:
+            q = paired_regions(q, live["regions"], video_width, sound_width)
+            live["region_hits"] += 1
+        elif live["region_wanted"] and q.ndim == 4:
             spans = region_spans(options)
             if spans and spans[-1][1] <= q.shape[2]:
                 q = q.clone()

modules/sampling/attention_bandit/tests/test_on_ltx.py

@@ -0,0 +1,71 @@
+"""The attention explorer on a (tiny) LTX forward: depth groups and clip regions; the prompt focus is H3's."""
+
+import random
+
+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 _load(tiny, **values):
+    from modules.sampling.modifiers.nodes import FunUILoadModifiers
+    settings = {"taste": {"key": "fox"}, "attention_bandit": {"enabled": True, "mode": "manual", "range": 0.5, "blocks": "0-3", **values}}
+    return FunUILoadModifiers.execute(tiny.patcher, settings).result
+
+
+def _outer(patched):
+    from comfy.patcher_extension import WrappersMP
+    return [w for ws in patched.wrappers[WrappersMP.OUTER_SAMPLE].values() for w in ws]
+
+
+def _run(outer, fn):
+    call = fn
+    for o in reversed(outer):
+        call = (lambda w, inner: (lambda: w(lambda: inner())))(o, call)
+    return call()
+
+
+def _draw(monkeypatch, groups, regions):
+    from modules.sampling import attention_bandit as m
+    picks = iter([groups, regions])
+    monkeypatch.setattr(m, "draw", lambda table, rng=random, **kw: [next(picks)] * len(table))
+
+
+def test_depth_groups_change_the_picture(tiny_ltx, monkeypatch):
+    base = tiny_ltx.run()
+    patched, status = _load(tiny_ltx, prompt_focus=False, regions=False)
+    assert "4 depth groups over 4 blocks" in status
+    monkeypatch.setattr("modules.sampling.attention_bandit.draw", lambda table, rng=random, **kw: [2] * len(table))
+    out = _run(_outer(patched), lambda: tiny_ltx.run(patched))
+    assert not torch.equal(out[0], base[0]), "the explorer ran and changed nothing"
+
+
+def test_clip_regions_change_the_picture_and_the_sound_apart(tiny_ltx, monkeypatch):
+    from core import log
+    base = tiny_ltx.run()
+    patched, _ = _load(tiny_ltx, prompt_focus=False)
+    _draw(monkeypatch, groups=1, regions=0)          # groups as is, every region sharper
+    log.new_run()
+    out = _run(_outer(patched), lambda: tiny_ltx.run(patched))
+    assert not torch.equal(out[0], base[0]) and not torch.equal(out[1], base[1])
+    assert any("clip regions in" in e["message"] for e in log.history())
+
+
+def test_the_band_helper_leaves_other_calls_alone():
+    from modules.sampling.attention_bandit import paired_regions
+    q = torch.ones(1, 8, 32)
+    other = torch.ones(1, 8, 7)
+    assert paired_regions(other, [0.5] * 5, 32, 16) is other
+    cut = paired_regions(q, [0.5, 1.0, 1.0, 1.0, 1.0], 32, 16)
+    assert torch.equal(cut[:, :2], torch.full((1, 2, 32), 2.0)) and torch.equal(cut[:, 2:], q[:, 2:])