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