diff --git a/core/dit_hooks.py b/core/dit_hooks.py index 044a8449..40b30df8 100644 --- a/core/dit_hooks.py +++ b/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: diff --git a/modules/models/ltx/__init__.py b/modules/models/ltx/__init__.py index 3025bb3b..f8e86101 100644 --- a/modules/models/ltx/__init__.py +++ b/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: diff --git a/modules/models/minimax_h3/__init__.py b/modules/models/minimax_h3/__init__.py index 6e2f7149..280712f0 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", "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 diff --git a/modules/sampling/attention_temperature/__init__.py b/modules/sampling/attention_temperature/__init__.py index 5f6d1593..632f804b 100644 --- a/modules/sampling/attention_temperature/__init__.py +++ b/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": { diff --git a/modules/sampling/attention_temperature/tests/test_on_ltx.py b/modules/sampling/attention_temperature/tests/test_on_ltx.py new file mode 100644 index 00000000..76a3a388 --- /dev/null +++ b/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