Files
olivv-cs-ComfyUI-FunPack/tests/test_per_block_any_model.py
T
digital-garbageandClaude Opus 5 71f41259b3 feat(lora): per-block application works on any model, not just LTX
Per-block scaling is derived from the LoRA's own deltas — which blocks it put
energy into — so it is a property of the adapter, not of the architecture. Two
things made it LTX-only anyway:

The gate required image_model in {ltxv, ltxav} and mode ltx2, so a model
FunPack had not heard of got the global fallback without anything ever reading
its weights. Whether per-block can run is now discovered from the patches:
fewer than two numbered blocks and the existing fallback takes over, as before.

The key pattern only matched transformer_blocks.N. — MiniMax H3 names its
stack blocks.N., so even with the gate open every H3 patch counted as global
and per-block would have been a silent no-op. The pattern now covers blocks,
transformer_blocks, double_blocks, single_blocks, joint_blocks and layers.

The semantic anchors (PAG 14, STG 14/19) were absolute indices measured on
LTX's 28-block DiT and meant nothing on a stack of another depth. They are
positions now; a 28-block model still resolves to exactly 14 and 19, so LTX
behaviour is unchanged, and a deeper stack gets the equivalent blocks.

Declines, loudly, when a model numbers blocks in two separate stacks
(Flux-style double/single): both start at 0, so one scale would land on two
unrelated blocks.

Untested on a real model — local runs cannot load one. 13 new tests, 11 of
which fail against the previous code.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-08-25 13:10:49 +03:00

94 lines
3.7 KiB
Python

"""Per-block LoRA application is driven by the adapter's deltas, not by the architecture.
It used to run only for image_model in {ltxv, ltxav} and only for keys named
`transformer_blocks.N.`, so every other model got the global fallback without anything ever
looking at its weights. MiniMax H3 names its stack `blocks.N.`, which the old pattern did
not match at all — per-block would have been a silent no-op there even with the gate open.
"""
import sys
import pytest
import torch
sys.path.insert(0, ".")
@pytest.fixture
def mm():
import model_management
return model_management
@pytest.mark.parametrize("key,expected", [
("diffusion_model.transformer_blocks.7.attn.to_q.weight", 7), # LTX
("diffusion_model.blocks.7.attn.qkv.weight", 7), # MiniMax H3
("diffusion_model.double_blocks.3.img_attn.qkv.weight", 3),
("diffusion_model.single_blocks.11.linear1.weight", 11),
("model.layers.4.self_attn.q_proj.weight", 4),
("diffusion_model.final_layer.weight", None),
])
def test_a_numbered_block_is_found_whatever_the_stack_is_called(mm, key, expected):
assert mm.transformer_block_index(key) == expected
def test_transformer_blocks_is_not_read_as_blocks(mm):
"""`transformer_blocks` must match as itself, not as a suffix of the generic name."""
assert mm.block_container_name(
"diffusion_model.transformer_blocks.2.ff.net.0.proj.weight") == "transformer_blocks"
def test_ltx_anchors_are_unchanged_on_a_28_block_model(mm):
"""PAG/STG's 14 and 19 were measured on LTX's 28 blocks. Expressing them as positions
must resolve back to exactly those two blocks, or this generalisation silently retunes
the model it was derived from."""
assert mm.semantic_anchor_blocks(range(28)) == frozenset({14, 19})
def test_anchors_scale_to_a_deeper_stack(mm):
"""A 48-block model gets the proportionally equivalent blocks, not 14 and 19."""
anchors = sorted(mm.semantic_anchor_blocks(range(48)))
assert anchors == [round(14 / 27 * 47), round(19 / 27 * 47)]
assert anchors != [14, 19]
def test_a_stack_too_short_to_compare_yields_what_it_has(mm):
assert mm.semantic_anchor_blocks([0]) == frozenset({0})
assert mm.semantic_anchor_blocks([]) == frozenset()
class _Entry(dict):
pass
def test_per_block_is_no_longer_gated_on_the_model(mm):
"""The gate is 'did the user ask for it'. Whether it can actually run is discovered
from the deltas afterwards."""
loader = mm.FunPackLoraLoader()
stack = {"mode": "minimax_h3", "per_block": True, "loras": []}
assert loader._per_block_supported(object(), stack, _Entry()) is True
assert loader._per_block_supported(object(), {"mode": "ltx2", "per_block": False},
_Entry()) is False
def _deltas(prefix, n, energy=1.0):
return {f"diffusion_model.{prefix}.{i}.attn.qkv.weight": torch.full((4, 4), energy * (i + 1))
for i in range(n)}
def test_an_h3_style_stack_gets_a_real_block_profile(mm):
"""The end of the change that matters: `blocks.N.` deltas now produce per-block scales
instead of falling through to global."""
loader = mm.FunPackLoraLoader()
template = loader._block_profile_template(_deltas("blocks", 8))
assert template is not None
assert template["block_count"] == 8
assert len(template["base_scales"]) == 8
def test_two_block_stacks_decline_rather_than_collide(mm, capsys):
"""double_blocks.0 and single_blocks.0 are different blocks with the same index."""
loader = mm.FunPackLoraLoader()
mixed = {**_deltas("double_blocks", 4), **_deltas("single_blocks", 4)}
assert loader._block_profile_template(mixed) is None
assert "per-block declined" in capsys.readouterr().out