fix(model-manager): apply Krea-2 LoRAs in kohya key layout

[Pfannkuchensack]

Summary

Krea-2 LoRAs trained with sd-scripts / LyCORIS install without complaint but have no effect on the image. They flatten the module path and prefix it with lora_unet_, over the native (ComfyUI) module names:

lora_unet_blocks_6_attn_wv.lora_down.weight

The Krea-2 converter only knew the dotted layouts (diffusers PEFT and native), so every key missed its module and the adapter silently did nothing. The only trace is a per-layer warning that is easy to miss in a busy log:

WARNING --> Failed to find module for LoRA layer key: lora_transformer-lora_unet_blocks_6_attn_wv

This was reported as "LoRAs don't work on GGUF Krea-2 models", but GGUF is not involved — I verified the quantised path separately (see QA below) and it applies LoRAs correctly. The layout is simply what a kohya-trained adapter looks like, on any Krea-2 model.

The fix un-flattens those keys to the dotted native layout before the existing native→diffusers step, reusing the repo's kohya_key_utils parsing tree — the same approach flux_onetrainer_lora_conversion_utils already uses. The tree walks the native module vocabulary, which resolves the flattened form's only genuine ambiguity: layerwise_blocks and refiner_blocks are the native components that themselves contain an underscore.

The tree doubles as a whitelist. insert_periods_into_kohya_key only rejects leftover tokens, so a prefix of a real path (blocks.0.attn) parses cleanly without naming a module — an added leaf check rejects those. Anything that cannot be reconstructed with certainty is left untouched rather than rewritten into a plausible-looking key that still matches nothing.

Non-Linear natives (mod.lin, prenorm/postnorm, attn.qknorm.*, last.norm/last.modulation) are deliberately excluded from the tree: they have no Linear counterpart in the diffusers layout — mod.lin for instance is folded into the scale_shift_table parameter — so an adapter targeting them cannot be applied, and renaming it anyway would turn "unsupported" into a silent no-op.

Detection is untouched. This layout is already recognised as Krea-2, because txtfusion matches as a substring of the flattened key. A transformer-only kohya adapter still won't auto-detect; that is a separate concern on the config side and out of scope here.

Related Issues / Discussions

Follows on from #9304 (Krea-2 support, LoRA marked WIP).

Orthogonal to #9424 (LyCORIS LoKr support for Krea-2), which touches the same two areas but solves a different half: that one is about weight-factor suffixes (lokr_w1/lokr_w2) on dotted keys and fixes installability; this one is about the flattened module path and fixes applicability. Expect a small merge overlap in krea2_lora_conversion_utils.py; the changes do not conflict functionally.

QA Instructions

Verified against a real adapter, Krea2_Don-Martin_LoRA-step00001200.safetensors (sd-scripts, ss_base_model_version=krea2, ss_network_dim=32, ss_network_alpha=32): 264 modules / 792 tensors, all lora_unet_*, suffixes lora_down.weight / lora_up.weight / alpha.

Before, every one of its 264 layers logged Failed to find module. After, all 264 convert and resolve against a full-size Krea2Transformer2DModel — 0 unresolved, 0 non-Linear, 0 in/out-feature mismatches, alpha preserved at 32.

To reproduce the user-visible behaviour: install a kohya-format Krea-2 LoRA, generate with it enabled at a high weight, and compare against the same seed with it disabled. Before this change the two images are identical and the log is full of Failed to find module for LoRA layer key: lora_transformer-lora_unet_*; after, the LoRA takes effect and the warnings are gone.

I also confirmed the quantised path is not implicated, since the report blamed GGUF. Building the real Krea2Transformer2DModel, quantising every Linear to GGML Q8_0, wrapping as GGMLTensor the way gguf_sd_loader and load_state_dict(assign=True) do, and running on CUDA/bf16 with the memory-efficient attention processor and the exact apply_smart_model_patches(..., force_sidecar_patching=True) call from krea2_denoise.py: the LoRA effect matches the bf16 direct-patch path (max |patched − base| 9.945 vs 9.938 over all 40 Linears), with unpatching restoring the output exactly.

Tests. New fixture krea2_lora_kohya_format.py captures 120 real keys with real shapes (first and last transformer block, one layerwise and one refiner text-fusion block, every top-level module), following the existing lora_state_dicts convention. test_kohya_krea2_lora_layers_match_the_real_transformer builds Krea2Transformer2DModel on the meta device and asserts each of the 40 layers lands on an actual nn.Linear whose in/out features agree with the LoRA's own down/up shapes — that catches a swapped rename such as ff.gate ↔ ff.down, whose SwiGLU shapes are transposed and which a name-only check would let through. Plus 16 parametrised mapping cases (including a multi-digit block index), and coverage for alpha, a doubled separator after the prefix, unrecognised keys being left untouched, and the alias-collision guard.

tests/backend/patches and tests/backend/model_manager pass (1047 passed, 132 skipped, 1 xfailed), as do ruff check and ruff format. Run locally on Python 3.11 / Windows / CUDA; the rest of the matrix is on CI.

Merge Plan

Nothing special, backend only. If #9424 lands first, rebase — both edit krea2_lora_conversion_utils.py in different places (suffix map vs. key normalisation).

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
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)