From 3e4888d05fd01b725365eeeae603d69394daaf6f Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 12 Jul 2026 22:55:08 +0100 Subject: [PATCH] feat(lora): map chroma non-block and guidance approximator modules BFL-format chroma adapters targeting the embedders, final projection or the distilled guidance layer MLPs resolved verbatim and unmapped: the embedder/final-layer names differ from diffusers outright, and the approximator MLP leaves are in_layer/out_layer in BFL but linear_1/ linear_2 in the diffusers PixArt projection. Only in_proj, out_proj and norms.N shared names and bound. - add CHROMA_EXTRA_MAP (kohya form derived) and GUIDANCE_LEAF_MAP to both target resolvers - route bare distilled_guidance_layer keys through the resolver instead of the diffusers passthrough so both leaf namings resolve; add img_in/txt_in/final_layer bare prefixes - cover all key forms and end-to-end binding in the offline suite --- pipelines/chroma/chroma_lora.py | 60 ++++++++++++++++++--- test/test-chroma-native-adapters.py | 83 ++++++++++++++++++++++++++++- 2 files changed, 134 insertions(+), 9 deletions(-) diff --git a/pipelines/chroma/chroma_lora.py b/pipelines/chroma/chroma_lora.py index 8a95a69eb..e1d679fba 100644 --- a/pipelines/chroma/chroma_lora.py +++ b/pipelines/chroma/chroma_lora.py @@ -7,9 +7,10 @@ Entry points, one per family: :func:`try_load_lora` (plus DoRA), :func:`try_load_lokr`, :func:`try_load_loha`, :func:`try_load_oft`. Recognized key prefixes: ``diffusion_model.``, ``transformer.``, -``lora_unet_``, plus bare BFL paths (``double_blocks.`` / ``single_blocks.``) -and bare diffusers paths (``transformer_blocks.`` / -``single_transformer_blocks.`` / ``distilled_guidance_layer.``). +``lora_unet_``, ``lycoris_``, plus bare BFL paths (``double_blocks.`` / +``single_blocks.`` / ``img_in.`` / ``txt_in.`` / ``final_layer.`` / +``distilled_guidance_layer.``) and bare diffusers paths +(``transformer_blocks.`` / ``single_transformer_blocks.``). Chroma LoRAs are trained against the Flux block layout regardless of save format. The diffusers ``ChromaTransformer2DModel`` exposes split-attention @@ -26,8 +27,11 @@ Fused weight handling: Chroma's modulation generator is the central ``distilled_guidance_layer`` approximator (replacing Flux's per-block ``norm1.linear``). The pruned ``ChromaAdaLayerNormZeroPruned`` classes have no ``.linear`` submodule, so -any ``_mod_lin`` / ``_modulation_lin`` keys land in ``unmapped``. LoRAs -targeting the approximator pass through unchanged. +any ``_mod_lin`` / ``_modulation_lin`` keys land in ``unmapped``. Approximator +keys pass through verbatim except the MLP leaves, where BFL +``layers.{i}.in_layer`` / ``out_layer`` rename to diffusers ``linear_1`` / +``linear_2``. Non-block embedder and final-layer keys rename via +``CHROMA_EXTRA_MAP``. """ from modules.lora import native_adapter @@ -38,11 +42,17 @@ from modules.lora.native_adapter import ChunkSpec KNOWN_PREFIXES = native_adapter.KNOWN_PREFIXES_DEFAULT -BARE_FLUX_PREFIXES = ("double_blocks.", "single_blocks.") +# distilled_guidance_layer. is deliberately a bare-BFL prefix, not a +# bare-diffusers one: the resolver renames BFL MLP leaves (in_layer/out_layer) +# and passes diffusers-named leaves (linear_1/linear_2, in_proj, out_proj, +# norms.N) through verbatim, so both namings route correctly. +BARE_FLUX_PREFIXES = ( + "double_blocks.", "single_blocks.", + "img_in.", "txt_in.", "final_layer.", "distilled_guidance_layer.", +) BARE_DIFFUSERS_PREFIXES = ( "transformer_blocks.", "single_transformer_blocks.", - "distilled_guidance_layer.", ) @@ -116,6 +126,13 @@ def resolve_targets(prefix_used, base): def _kohya_to_diffusers(base): """For kohya keys like ``double_blocks_0_img_attn_qkv`` or ``single_blocks_5_linear1``.""" + extra = CHROMA_EXTRA_KOHYA_MAP.get(base) + if extra is not None: + return [(extra, None)] + if base.startswith("distilled_guidance_layer_layers_"): + for bfl_leaf, dif_leaf in GUIDANCE_LEAF_MAP.items(): + if base.endswith("_" + bfl_leaf): + return [(base[:-len(bfl_leaf)] + dif_leaf, None)] if base.startswith("double_blocks_"): rest = base[len("double_blocks_"):] idx, _, suffix = rest.partition("_") @@ -129,6 +146,14 @@ def _kohya_to_diffusers(base): def _bfl_to_diffusers(base): """For BFL dotted keys like ``double_blocks.0.img_attn.qkv``.""" + extra = CHROMA_EXTRA_MAP.get(base) + if extra is not None: + return [(extra, None)] + if base.startswith("distilled_guidance_layer.layers."): + stem, _, leaf = base.rpartition(".") + mapped = GUIDANCE_LEAF_MAP.get(leaf) + if mapped is not None: + return [(f"{stem}.{mapped}", None)] parts = base.split(".") if len(parts) < 3: return [(base, None)] @@ -142,6 +167,27 @@ def _bfl_to_diffusers(base): return [(base, None)] +# Non-block BFL targets. Chroma prunes time_in/guidance_in/vector_in and the +# final-layer adaLN modulation, so only the embedders and the final projection +# need renames. None carry a block index, so the kohya form is derivable by +# underscoring the BFL path. +CHROMA_EXTRA_MAP = { + "img_in": "x_embedder", + "txt_in": "context_embedder", + "final_layer.linear": "proj_out", +} + +CHROMA_EXTRA_KOHYA_MAP = {k.replace(".", "_"): v for k, v in CHROMA_EXTRA_MAP.items()} + +# distilled_guidance_layer MLP leaves: BFL in_layer/out_layer vs diffusers +# PixArtAlphaTextProjection linear_1/linear_2. The other approximator leaves +# (in_proj, out_proj, norms.N) share names on both sides and pass verbatim. +GUIDANCE_LEAF_MAP = { + "in_layer": "linear_1", + "out_layer": "linear_2", +} + + # Static non-fused renames (underscore-keyed for dispatch from either kohya or BFL paths). _DOUBLE_STATIC = { "img_attn_proj": "attn.to_out.0", diff --git a/test/test-chroma-native-adapters.py b/test/test-chroma-native-adapters.py index df8f2238c..7f29f8ec3 100644 --- a/test/test-chroma-native-adapters.py +++ b/test/test-chroma-native-adapters.py @@ -218,10 +218,24 @@ def build_mock_transformer(): transformer = _Holder() transformer.transformer_blocks = torch.nn.ModuleList([build_double_block() for _ in range(N_DOUBLE)]) transformer.single_transformer_blocks = torch.nn.ModuleList([build_single_block() for _ in range(N_SINGLE)]) - # distilled_guidance_layer - Chroma's central modulation approximator - # Minimal stand-in: just one linear submodule the tests can target + # distilled_guidance_layer - Chroma's central modulation approximator. + # Mirrors ChromaApproximator: in_proj / out_proj Linears, PixArt-shaped + # MLP layers (linear_1 / linear_2) and RMSNorms. transformer.distilled_guidance_layer = _Holder() transformer.distilled_guidance_layer.in_proj = torch.nn.Linear(HIDDEN, HIDDEN, bias=True) + transformer.distilled_guidance_layer.out_proj = torch.nn.Linear(HIDDEN, HIDDEN, bias=True) + transformer.distilled_guidance_layer.layers = torch.nn.ModuleList() + transformer.distilled_guidance_layer.norms = torch.nn.ModuleList() + for _ in range(2): + mlp = _Holder() + mlp.linear_1 = torch.nn.Linear(HIDDEN, HIDDEN, bias=True) + mlp.linear_2 = torch.nn.Linear(HIDDEN, HIDDEN, bias=True) + transformer.distilled_guidance_layer.layers.append(mlp) + transformer.distilled_guidance_layer.norms.append(torch.nn.RMSNorm(HIDDEN)) + # Non-block CHROMA_EXTRA_MAP targets. + transformer.x_embedder = torch.nn.Linear(HIDDEN, HIDDEN) + transformer.context_embedder = torch.nn.Linear(HIDDEN, HIDDEN) + transformer.proj_out = torch.nn.Linear(HIDDEN, HIDDEN, bias=True) return transformer @@ -630,6 +644,32 @@ def test_resolve_targets_static_renames(): return True +def test_resolve_targets_extra_and_guidance(): + """Non-block extra-map renames and guidance-layer MLP leaf renames, all key forms.""" + for bfl_base, diffusers_path in C.CHROMA_EXTRA_MAP.items(): + for prefix, base in [ + ('diffusion_model.', bfl_base), + (None, bfl_base), + ('lora_unet_', bfl_base.replace('.', '_')), + ]: + targets = C.resolve_targets(prefix, base) + assert targets == [(diffusers_path, None)], f'({prefix}, {base}) -> {targets}' + cases = [ + # BFL MLP leaves rename to the PixArt projection names. + (('diffusion_model.', 'distilled_guidance_layer.layers.0.in_layer'), 'distilled_guidance_layer.layers.0.linear_1'), + ((None, 'distilled_guidance_layer.layers.1.out_layer'), 'distilled_guidance_layer.layers.1.linear_2'), + (('lora_unet_', 'distilled_guidance_layer_layers_0_in_layer'), 'distilled_guidance_layer_layers_0_linear_1'), + (('lora_unet_', 'distilled_guidance_layer_layers_1_out_layer'), 'distilled_guidance_layer_layers_1_linear_2'), + # Verbatim leaves are untouched in either naming. + (('diffusion_model.', 'distilled_guidance_layer.in_proj'), 'distilled_guidance_layer.in_proj'), + ((None, 'distilled_guidance_layer.layers.0.linear_1'), 'distilled_guidance_layer.layers.0.linear_1'), + ] + for (prefix, base), expected in cases: + targets = C.resolve_targets(prefix, base) + assert targets == [(expected, None)], f'({prefix}, {base}) -> {targets}' + return True + + def test_resolve_targets_onetrainer_passthrough(): """The ``lora_transformer_`` passthrough lives in the shared resolver. @@ -805,6 +845,43 @@ def test_lora_distilled_guidance(): return True +def sd_lokr_bfl_extra_modules(): + """BFL LoKR spanning the non-block extra targets and guidance MLP leaves. + + Full-matrix factors with the ai-toolkit placeholder alpha, mirroring the + layout of real full-preset checkpoints. + """ + bases = [ + 'img_in', 'txt_in', 'final_layer.linear', + 'distilled_guidance_layer.layers.0.in_layer', + 'distilled_guidance_layer.layers.1.out_layer', + ] + sd = {} + for base in bases: + sd[f'diffusion_model.{base}.lokr_w1'] = torch.randn(LOKR_W1_DIM, LOKR_W1_DIM) + sd[f'diffusion_model.{base}.lokr_w2'] = torch.randn(HIDDEN // LOKR_W1_DIM, HIDDEN // LOKR_W1_DIM) + sd[f'diffusion_model.{base}.alpha'] = torch.tensor(9999220736.0) + return sd + + +def test_lokr_bfl_extra_and_guidance(): + """Embedder/final-layer renames and guidance MLP leaf renames all bind.""" + net = _load_via(C.try_load_lokr, sd_lokr_bfl_extra_modules()) + assert net is not None and len(net.modules) == 5, f'got {sorted(net.modules) if net else None}' + expected = { + 'lora_transformer_x_embedder', + 'lora_transformer_context_embedder', + 'lora_transformer_proj_out', + 'lora_transformer_distilled_guidance_layer_layers_0_linear_1', + 'lora_transformer_distilled_guidance_layer_layers_1_linear_2', + } + assert set(net.modules) == expected, f'got {set(net.modules)}' + # Full-matrix factors: the placeholder alpha must not scale. + for nk, mod in net.modules.items(): + assert mod.dim is None and mod.calc_scale() == 1.0, f'{nk}: dim={mod.dim} scale={mod.calc_scale()}' + return True + + def test_lora_dora_threading(): """dora_scale flows into NetworkModuleLora.""" net = _load_via(C.try_load_lora, sd_lora_with_dora_scale()) @@ -979,6 +1056,7 @@ def run_tests(): log.warning('=== Parsing primitives ===') for fn in [test_parse_key_all_prefixes, test_marker_disambiguation, test_resolve_targets_static_renames, + test_resolve_targets_extra_and_guidance, test_resolve_targets_onetrainer_passthrough]: run_test(CAT_PARSE, fn) @@ -1000,6 +1078,7 @@ def run_tests(): test_lokr_bfl_img_attn_proj, test_lokr_bfl_img_attn_qkv_chunked, test_lokr_bfl_single_linear1_unequal_chunks, + test_lokr_bfl_extra_and_guidance, test_loha_bfl_img_attn_proj, test_loha_bfl_img_attn_qkv_chunked, test_oft_bfl_img_attn_proj,