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

@@ -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",

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,