diff --git a/conftest.py b/conftest.py index bacfb569..59de0475 100644 --- a/conftest.py +++ b/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 = [] diff --git a/modules/models/ltx/__init__.py b/modules/models/ltx/__init__.py index f8e86101..b66bc297 100644 --- a/modules/models/ltx/__init__.py +++ b/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: diff --git a/modules/models/minimax_h3/__init__.py b/modules/models/minimax_h3/__init__.py index 280712f0..839f0850 100644 --- a/modules/models/minimax_h3/__init__.py +++ b/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 diff --git a/modules/sampling/attention_temperature/tests/test_on_ltx.py b/modules/sampling/attention_temperature/tests/test_on_ltx.py index 76a3a388..cb2ccd98 100644 --- a/modules/sampling/attention_temperature/tests/test_on_ltx.py +++ b/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 diff --git a/modules/sampling/block_repeat/__init__.py b/modules/sampling/block_repeat/__init__.py index 39b66536..d354eed2 100644 --- a/modules/sampling/block_repeat/__init__.py +++ b/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} diff --git a/modules/sampling/block_repeat/tests/test_on_ltx.py b/modules/sampling/block_repeat/tests/test_on_ltx.py new file mode 100644 index 00000000..c3471f0f --- /dev/null +++ b/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 diff --git a/modules/sampling/reins/__init__.py b/modules/sampling/reins/__init__.py index e2871524..4715aeea 100644 --- a/modules/sampling/reins/__init__.py +++ b/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 diff --git a/modules/sampling/reins/tests/test_on_ltx.py b/modules/sampling/reins/tests/test_on_ltx.py new file mode 100644 index 00000000..ad064f91 --- /dev/null +++ b/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