From 5875cdd29e8ab634d9b48ca5ce63c890756fe524 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 5 Sep 2026 02:24:04 +0100 Subject: [PATCH] feat(lora): remap anima 1.0 lora block indices onto depth-expanded checkpoints Anima 2.9B interleaves twelve new blocks among the 28 of Anima 1.0, so a block index trained against 1.0 names a different block on the expanded model. Every such key still resolves, since blocks 0 to 27 exist either way, so the mismatch was silent. The Anima loader now shifts base-depth indices onto the blocks that carry those weights, keyed by (base depth, expanded depth) and applied only when the transformer is expanded and the LoRA stays inside the base depth. Transformer keys move; llm_adapter and text encoder keys keep their own numbering. --- pipelines/anima/anima_lora.py | 68 ++++++++++++++++ test/test-anima-native-adapters.py | 125 +++++++++++++++++++++++++++-- 2 files changed, 188 insertions(+), 5 deletions(-) diff --git a/pipelines/anima/anima_lora.py b/pipelines/anima/anima_lora.py index cea98faa8..bc7d9c7ab 100644 --- a/pipelines/anima/anima_lora.py +++ b/pipelines/anima/anima_lora.py @@ -33,8 +33,10 @@ the standard ``NetworkWeights.w`` slots rather than being baked into the factor weights at load time. """ +import re from collections import OrderedDict +from modules.logger import log from modules.lora import native_adapter @@ -178,12 +180,78 @@ def network_prefix_for(prefix_used): return "lora_transformer_" +# === Depth-expanded checkpoint block remap === +# +# Anima 2.9B grows the 28-block DiT to 40 by interleaving new blocks among the +# originals, so a block index trained against 1.0 no longer names the same +# block. Keyed by (base depth, expanded depth) rather than by model name, so +# any Anima transformer of a listed depth picks up the table. + +BLOCK_EXPANSIONS = { + (28, 40): (2, 5, 8, 11, 14, 17, 21, 24, 27, 30, 33, 36), +} + +BLOCK_INDEX = re.compile(r"^blocks([._])(\d+)(?=[._])") + + +def expansion_map(base_depth, expanded_depth): + """Map each base-model block index to its position in the expanded model.""" + inserted = set(BLOCK_EXPANSIONS[(base_depth, expanded_depth)]) + return dict(enumerate(n for n in range(expanded_depth) if n not in inserted)) + + +def transformer_depth(): + """Block count of the loaded transformer, or 0 when there is nothing to read.""" + from modules import shared + pipe = getattr(shared.sd_model, "pipe", shared.sd_model) + blocks = getattr(getattr(pipe, "transformer", None), "transformer_blocks", None) + return len(blocks) if blocks is not None else 0 + + +def block_index(base): + """Leading ``blocks.N`` / ``blocks_N`` index of a parsed key, else None.""" + m = BLOCK_INDEX.match(base) + return int(m.group(2)) if m is not None else None + + +def remap_blocks(groups): + """Shift base-depth block indices onto the blocks that carry those weights. + + A LoRA whose highest block index fits inside the base depth was trained on + the unexpanded model. Only transformer keys move: llm_adapter and text + encoder keys carry their own unrelated numbering. + """ + indices = [i for (prefix, base) in groups + if network_prefix_for(prefix) == "lora_transformer_" and (i := block_index(base)) is not None] + if not indices: + return groups + depth = transformer_depth() + match = next(((b, e) for (b, e) in BLOCK_EXPANSIONS if e == depth and max(indices) < b), None) + if match is None: + return groups + table = expansion_map(*match) + out = {} + for (prefix, base), w in groups.items(): + i = block_index(base) if network_prefix_for(prefix) == "lora_transformer_" else None + if i is not None: + base = BLOCK_INDEX.sub(rf"blocks\g<1>{table[i]}", base) + out[(prefix, base)] = w + log.info(f'Network load: arch=anima block remap {match[0]}->{match[1]} keys={len(indices)}') + return out + + +def group_by_suffixes_remapped(state_dict, suffixes, **kwargs): + """Group by suffix, then remap block indices for depth-expanded transformers.""" + return remap_blocks(native_adapter.group_by_suffixes(state_dict, suffixes, **kwargs)) + + # === Native loaders (thin wrappers over native_adapter generics) === _BIND_KWARGS = dict( resolve_targets=resolve_targets, prefixes=ANIMA_PREFIXES, network_prefix=network_prefix_for, + group_by_suffixes_fn=group_by_suffixes_remapped, arch_name="anima", ) diff --git a/test/test-anima-native-adapters.py b/test/test-anima-native-adapters.py index 08eae0963..dcdff5059 100644 --- a/test/test-anima-native-adapters.py +++ b/test/test-anima-native-adapters.py @@ -205,10 +205,10 @@ def build_cosmos_block(): return block -def build_mock_transformer(): +def build_mock_transformer(n_blocks=N_BLOCKS): """Mirror diffusers' Cosmos2TransformerModel top-level layout.""" transformer = _Holder() - transformer.transformer_blocks = torch.nn.ModuleList([build_cosmos_block() for _ in range(N_BLOCKS)]) + transformer.transformer_blocks = torch.nn.ModuleList([build_cosmos_block() for _ in range(n_blocks)]) # Time embedding transformer.time_embed = _Holder() @@ -338,13 +338,14 @@ class _MockAnimaSdModel: self.__class__.__name__ = 'AnimaTextToImagePipeline' -def install_mock_pipe(): +def install_mock_pipe(n_blocks=N_BLOCKS): """Set shared.sd_model to a mock exposing an Anima-shaped 3-component pipeline. Each test re-installs so stamped ``network_layer_name`` attributes from - prior tests do not leak across runs. + prior tests do not leak across runs. ``n_blocks`` sizes the DiT, which the + depth-expansion block remap reads. """ - transformer = build_mock_transformer() + transformer = build_mock_transformer(n_blocks) llm_adapter = build_mock_llm_adapter() text_encoder = build_mock_text_encoder() pipe = _MockAnimaPipeline(transformer, llm_adapter, text_encoder) @@ -1342,6 +1343,107 @@ def test_dora_square_weight_1d_defaults_to_per_input(): return True +# ============================================================ +# Tests - depth-expanded checkpoint block remap +# ============================================================ + +CAT_REMAP = category('remap') + +BASE_DEPTH = 28 +EXPANDED_DEPTH = 40 + + +def sd_lora_kohya_block(index): + """Kohya transformer LoRA on self_attn.q_proj at an arbitrary block index.""" + stem = f'lora_unet_blocks_{index}_self_attn_q_proj' + return { + f'{stem}.lora_down.weight': torch.randn(RANK, HIDDEN), + f'{stem}.lora_up.weight': torch.randn(HIDDEN, RANK), + f'{stem}.alpha': torch.tensor(float(RANK)), + } + + +def _load_at_depth(try_fn, state_dict, n_blocks, name='test'): + install_mock_pipe(n_blocks) + with TempLora(state_dict, name=name) as nod: + return try_fn(name, nod, lora_scale=1.0) + + +def test_expansion_map_matches_manifest(): + """The 28->40 table skips the author's insertion positions in order.""" + table = A.expansion_map(BASE_DEPTH, EXPANDED_DEPTH) + inserted = set(A.BLOCK_EXPANSIONS[(BASE_DEPTH, EXPANDED_DEPTH)]) + assert len(table) == BASE_DEPTH + assert set(table.values()).isdisjoint(inserted), 'no base block may land on an inserted block' + assert (table[0], table[1], table[2], table[27]) == (0, 1, 3, 39) + assert list(table.values()) == sorted(table.values()), 'order must be preserved' + return True + + +def test_expansion_map_covers_every_base_block_uniquely(): + """All base indices land on distinct blocks inside the expanded depth.""" + table = A.expansion_map(BASE_DEPTH, EXPANDED_DEPTH) + assert len(set(table.values())) == BASE_DEPTH + assert max(table.values()) < EXPANDED_DEPTH + return True + + +def test_remap_binds_last_base_block_to_top_of_expanded(): + """A 1.0 LoRA on block 27 binds to block 39 of a 40-block transformer.""" + net = _load_at_depth(A.try_load_lora, sd_lora_kohya_block(27), EXPANDED_DEPTH) + assert net is not None and len(net.modules) == 1, f'got {net.modules if net else None}' + assert 'lora_transformer_transformer_blocks_39_attn1_to_q' in net.modules + return True + + +def test_remap_leaves_blocks_before_first_insertion(): + """Blocks 0 and 1 precede the first insertion, so they do not move.""" + for i in (0, 1): + net = _load_at_depth(A.try_load_lora, sd_lora_kohya_block(i), EXPANDED_DEPTH) + assert f'lora_transformer_transformer_blocks_{i}_attn1_to_q' in net.modules + return True + + +def test_no_remap_on_base_depth_model(): + """A 28-block transformer matches no expansion entry, so indices pass through.""" + net = _load_at_depth(A.try_load_lora, sd_lora_kohya_block(1), BASE_DEPTH) + assert 'lora_transformer_transformer_blocks_1_attn1_to_q' in net.modules + return True + + +def test_no_remap_when_lora_is_native_to_expanded_depth(): + """A LoRA reaching past the base depth was trained on the expanded model.""" + sd = sd_lora_kohya_block(5) + sd.update(sd_lora_kohya_block(39)) + net = _load_at_depth(A.try_load_lora, sd, EXPANDED_DEPTH) + assert 'lora_transformer_transformer_blocks_5_attn1_to_q' in net.modules + assert 'lora_transformer_transformer_blocks_39_attn1_to_q' in net.modules + return True + + +def test_remap_leaves_adapter_and_te_indices_alone(): + """llm_adapter blocks and text encoder layers carry unrelated numbering.""" + install_mock_pipe(EXPANDED_DEPTH) + out = A.remap_blocks({ + ('lora_unet_', 'blocks_27_self_attn_q_proj'): {}, + ('diffusion_model.llm_adapter.', 'blocks.1.self_attn.q_proj'): {}, + ('lora_te_', 'layers_1_self_attn_q_proj'): {}, + }) + assert ('lora_unet_', 'blocks_39_self_attn_q_proj') in out + assert ('diffusion_model.llm_adapter.', 'blocks.1.self_attn.q_proj') in out + assert ('lora_te_', 'layers_1_self_attn_q_proj') in out + return True + + +def test_remap_handles_dotted_bfl_keys(): + """BFL keys arrive dotted rather than underscore-flattened.""" + install_mock_pipe(EXPANDED_DEPTH) + out = A.remap_blocks({('diffusion_model.', 'blocks.27.self_attn.q_proj'): {}}) + assert ('diffusion_model.', 'blocks.39.self_attn.q_proj') in out + return True + + + # ============================================================ # Test runner # ============================================================ @@ -1417,6 +1519,19 @@ def run_tests(): ]: run_test(CAT_MATH, fn) + log.warning('=== depth-expanded block remap ===') + for fn in [ + test_expansion_map_matches_manifest, + test_expansion_map_covers_every_base_block_uniquely, + test_remap_binds_last_base_block_to_top_of_expanded, + test_remap_leaves_blocks_before_first_insertion, + test_no_remap_on_base_depth_model, + test_no_remap_when_lora_is_native_to_expanded_depth, + test_remap_leaves_adapter_and_te_indices_alone, + test_remap_handles_dotted_bfl_keys, + ]: + run_test(CAT_REMAP, fn) + elapsed = time.time() - t0 log.warning('=== Results ===') total_pass = 0