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>
94 lines
3.7 KiB
Python
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
|