fix(krea2-lora): validate kohya keys by un-flattening, not by prefix

[Pfannkuchensack]

Summary

Follow-up to #9449, addressing the two non-blocking notes from review 4888833569.

1. lora_unet_blocks_ swept Wan and Anima kohya LoRAs into the Krea-2 override.

#9449 admitted the kohya/LyCORIS layout by matching per-module prefixes such as
lora_unet_blocks_. That spelling is not Krea-2's alone — Wan writes it verbatim
(wan_lora_conversion_utils._KOHYA_KEY_REGEX) and Anima writes
lora_unet_[llm_adapter_]blocks_<idx>_ (anima_lora_constants._KOHYA_ANIMA_RE). A mislabeled
Wan or Anima file with an explicit Krea-2 base override was therefore accepted, installed, and
then silently no-op'd at generation time: the un-flattener rejects self_attn / cross_attn /
mlp_layer0, so every layer warn-skipped. "Installs, then does nothing" is exactly the failure
mode install-time validation exists to prevent.

2. The doubled-separator spelling could not be installed.

The converter deliberately tolerates lora_unet__blocks_... (the lstrip, plus a dedicated
test), but no prefix in the list spelled out the doubled underscore, so a transformer-only
adapter written that way was rejected by the override — the very class of file that part of
#9449 was meant to make installable.

The fix replaces prefix matching for the kohya layout with the converter's own
reconstruction: a kohya key counts as Krea-2 only if its flattened path un-flattens to a leaf of
the native Krea-2 module vocabulary. Wan/Anima keys do not reach a leaf and are rejected; the
doubled-separator spelling goes through the identical lstrip the converter uses and is
accepted. The dotted layouts (transformer., diffusion_model., text_encoder.) still match
by prefix as before.

Both notes fall out of one change: split_kohya_krea2_key() is now the single place that splits
a kohya key into (flat module path, separator, weight suffix). The converter's per-module gate,
the rewrite it guards, and identification all call it, so they cannot disagree about which module
a key belongs to — the invariant round 3 of the review checked by hand is now structural.

Where the code lives. configs/lora.py cannot import krea2_lora_conversion_utils: it pulls
in the patch layers, which import model_manager.load, closing an import cycle back into
model_manager.configs. The reconstruction helpers therefore moved to krea2_lora_constants,
the same reason anima_lora_constants exists. The move is verbatim apart from three names
becoming public; no logic changed.

One intended behavior change beyond the two notes. An adapter targeting an nn.Sequential
position that holds no Linear (lora_unet_tmlp_1, lora_unet_tproj_0, lora_unet_txtmlp_2)
previously matched by prefix, installed, and then warn-skipped at apply time. The parsing tree
lists the Linear-bearing positions literally, so these no longer install. Same rationale as
note 1, covered by its own test.

No user-facing behavior changes for correctly-labeled Krea-2 LoRAs: every layout #9449 accepted
is still accepted.

Related Issues / Discussions

Follow-up to #9449 (merged as 77f00f1) — non-blocking notes 2 and 3 of review
#9449 (review), left as optional
follow-ups in the round-3 approval
#9449 (review).

QA Instructions

Automated:

pytest tests/backend/patches tests/backend/model_manager/configs

818 passed, 3 skipped locally.

The 22 new parametrized cases in tests/backend/model_manager/configs/test_krea2_lora_config.py
are load-bearing: checked out against #9449 as merged, 11 of them fail — 5 foreign
lora_unet_blocks_ modules that were wrongly accepted, 3 doubled-separator modules that were
wrongly rejected, and the 3 non-Linear Sequential indices. (The llm_adapter_blocks_ and
doubled-separator txtfusion cases pass on both sides — the first matched no prefix, the second
reaches auto-detection via its txtfusion substring — they pin the behavior rather than the
regression.)

Manual, if you have the files:

1. Take a Wan or Anima kohya LoRA, set the base explicitly to Krea-2 during install.
   Before: installs, then produces no visible effect and logs
   Failed to find module for LoRA layer key: per layer. After: install is refused.
2. Take a Krea-2 kohya LoRA whose keys use the doubled separator (lora_unet__blocks_...),
   install with an explicit Krea-2 base. Before: rejected. After: installs and applies.
3. Any normal Krea-2 LoRA (diffusers PEFT, native/ComfyUI, or single-separator kohya) —
   unchanged in both install and generation.

Merge Plan

Nothing special. Backend only, no DB schema, no redux slice, no API surface change.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration — n/a, backend only
• Documentation added / updated (if applicable) — n/a
• Updated What's New copy (if doing a release after this PR) — n/a

[lstein]

Approving. Reviewed adversarially at 3c8db2472f in a clean worktree with a full venv — I assumed the change was broken and tried to prove it. Everything below is execution-backed, not read-through.

Verification

1. Acceptance-set diff — is anything that actually converts newly rejected?

I enumerated ~190 kohya module spellings (Krea-2 native Linear + non-Linear, Wan in both diffusers and native naming, Anima with and without llm_adapter_, Flux, SDXL, Qwen), each in single- and doubled-separator form, and cross-tabulated old prefix accepts × new filter accepts × converter actually rewrites the key × resulting patch target.

Every newly-rejected key has converted = 0 — the converter leaves it verbatim today, so it installs and warn-skips at apply time. No key that actually converts is newly rejected. Every newly-accepted key (the doubled-separator spellings) converts to a real target. No row was accepted-but-not-converted.

2. Is the parsing tree the right vocabulary?

Instantiated the real Krea2Transformer2DModel (diffusers 0.39.0) on the meta device with the loader's KREA2_TRANSFORMER_CONFIG, then checked both directions:

• all 32/32 tree leaves resolve to a real nn.Linear (txtfusion_projector → text_fusion.projector included);
• all 64/64 nn.Linear modules of the model are reachable from a kohya spelling the filter accepts.

Exact match both ways — the tree neither over- nor under-approximates the model. That is the strongest evidence for the central claim: the accepted vocabulary is precisely the set of modules a LoRA can actually patch.

3. The import-cycle justification. Not taken on faith. Patching configs/lora.py to import krea2_lora_conversion_utils fails with ImportError: cannot import name 'ControlLoRA_LyCORIS_FLUX_Config' from partially initialized module ... (most likely due to a circular import), via patches.layers → model_manager.load.load_base → configs.factory → configs.lora. The move to krea2_lora_constants is necessary, and the new chain (krea2_lora_constants → kohya_key_utils → typing) closes no cycle; both package __init__.py files are empty.

4. Tests. 818 passed, 3 skipped locally, matching the PR description. The load-bearing claim holds exactly: swapping in main's lora.py under the new test file gives precisely 11 failures — the 5 foreign lora_unet_blocks_ modules, the 3 doubled-separator modules, and the 3 non-Linear Sequential indices. CI green on all 17 checks.

5. End-to-end through ModelConfigFactory with real .safetensors files and allow_unknown=False:

file	explicit base=Krea2
Wan kohya (self_attn_q, cross_attn_k, ffn_0)	no match ✅
Anima kohya (mlp_layer0, adaln_modulation_1)	no match ✅
Krea-2 kohya, single separator	installs as LoRA_LyCORIS_Krea2_Config ✅
Krea-2 kohya, doubled separator	installs ✅
Krea-2 kohya, Linears + norms mixed	installs ✅
Krea-2 kohya, norms only	no match ✅

Attacks that found nothing

• Wan's full kohya vocabulary — diffusers naming (attn1_to_q/k/v, attn1_to_out_0, attn1_norm_q/k, attn2_*, ffn_net_0_proj, ffn_net_2) and native naming (self_attn_q/k/v/o, cross_attn_*, cross_attn_k_img/v_img, ffn_0, ffn_2, norm3, modulation). None reconstructs to a Krea-2 leaf.
• Anima's, with and without llm_adapter_: mlp_layer0, mlp_layer1, mlp_layer_0, adaln_modulation_1/3, {self,cross}_attn_{q,k,v}_proj, output_proj.
• Flux double_blocks_/single_blocks_/single_transformer_blocks_, SDXL UNet (down_blocks_0_attentions_0_transformer_blocks_0_attn1_to_q), Qwen transformer_blocks_.
• The greedy, non-backtracking un-flattener against the two underscore-bearing components (txtfusion_layerwise_blocks_*, txtfusion_refiner_blocks_*) — resolves correctly, and txtfusion_projector does not get mis-consumed.
• Degenerate inputs: empty and all-underscore flat paths (lora_unet_.lora_down.weight, lora_unet___.lora_down.weight) — _kohya_module_path_is_leaf("") returns False; suffix-less kohya keys; the non-string keys .pt/.ckpt sources carry; a dotted spelling like lora_unet_last.linear.lora_down.weight.
• Precedence divergence between insert_periods_into_kohya_key and _kohya_module_path_is_leaf. Both check exact match before INDEX_PLACEHOLDER, and the reconstructed path re-splits into exactly the components the forward walk matched (no component can contain a .), so the re-walk always takes the same branch. tmlp's literal "0"/"2" children are the case that would break under the opposite precedence, and they resolve correctly.
• Whether lstrip("_") opens a new door for a foreign file: after stripping, the path still has to reach a Krea-2 leaf, so it does not.
• The auto-detection path (_has_krea2_lora_keys) and the dotted prefix list — untouched, and _KREA2_SUPPORTED_LORA_PREFIXES has no remaining consumer outside _key_is_supported_krea2_layout.

Non-blocking notes

1. The description understates the behavior change. Besides tmlp_1 / tproj_0 / txtmlp_2, these Krea-2-native kohya spellings also stop installing: blocks_N_mod_lin, blocks_N_prenorm, blocks_N_postnorm, blocks_N_attn_qknorm_{qnorm,knorm}, txtfusion_*_mod_lin, and txtmlp_0. All of them warn-skip today (the converter leaves them verbatim), so refusing them follows exactly the stated rationale — but it is a wider set than "three nn.Sequential positions", and a file patching only those modules now fails to install. Worth a sentence in the PR body rather than a code change.

2. Identification still checks less than conversion does. It accepts on module-path validity alone; the converter additionally requires the module's entire suffix set to be convertible. So a genuine Krea-2 module carrying an extra unrecognized suffix alongside its pair still installs and then warn-skips — verified with .lora_mid.weight, .scale and .lokr_w1, each producing lora_transformer-lora_unet_blocks_0_attn_wq, which matches no module. This is pre-existing (#9449's prefix accepted it too), and tightening it per key is exactly what re-opens the alpha-orphan split that round 2 of #9449 fixed — so leaving it alone is the right call. Noting it only because it is the last residual of the class this PR closes.

3. _has_complete_lora_pair needs only one accepted pair, so a hybrid file (Wan kohya plus a single blocks_3_attn_wq module) with an explicit Krea-2 base still installs and no-ops on the Wan layers — verified end-to-end. Pre-existing any-not-all property of the gate, and it takes a deliberately mislabeled file to reach.

4. _lora_weight_keys_are_all_paired's key_filter parameter has no caller that passes it — both call sites hand it the whole state dict. Pre-existing dead parameter that this PR generalizes rather than drops; could go.

Verdict

The change is a strict narrowing on foreign keys and a strict widening on exactly the doubled-separator spelling the converter already tolerated. split_kohya_krea2_key() as the single splitting point makes the gate/rewrite/identification agreement structural rather than checked by hand, and the accepted vocabulary now coincides exactly with the model's real Linear layers. Nothing blocking.