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)

[lstein]

Thanks for this — the core change is solid. I ran an adversarial review (fresh-context reviewers instructed to assume the change is broken and prove it, findings verified by execution against the branch), and the un-flattening itself survives every attack we threw at it:

• All 264 real native module paths round-trip through the greedy parse tree with zero mismatches; no level of the tree can mis-commit (layerwise_blocks/refiner_blocks have no sibling prefix keys).
• No foreign architecture's flattened keys (FLUX, Wan, Qwen-Image, Z-Image, SDXL, Anima) parse against the tree or get claimed by/steal dispatch from this converter; install-time classification of the kohya fixture uniquely matches LoRA_LyCORIS_Krea2_Config.
• All 40 fixture layers resolve against the real Krea2Transformer2DModel in both the direct-patch and sidecar/quantized paths, including to_out.0; leftover unmappable keys degrade to the per-layer warning in both paths.
• Conversion is idempotent, doesn't mutate the input dict, and the collision guard fires in both dict orders; doubled-separator aliases collide loudly rather than silently dropping.
• lora_te_* keys, alpha-only groups, and orphaned halves behave identically pre/post-PR; fixture shapes independently recomputed against KREA2_TRANSFORMER_CONFIG with no discrepancies.

One regression did survive verification, though, and I'd like it fixed before merge.

1. Blocking: kohya keys with non-LoRA LyCORIS suffixes now abort the entire load (pre-PR they loaded)

_maybe_convert_kohya_krea2_state_dict rewrites any lora_unet_<parseable-path>.<suffix> key regardless of suffix, but _group_by_layer can only re-split a dotted key whose suffix is in _SUFFIX_TO_VALUE_KEY. A LyCORIS suffix such as .lokr_w1, .hada_w1_a, or .diff falls into the rsplit(".", maxsplit=2) fallback, which now splits inside the module path itself:

lora_unet_blocks_6_attn_wq.lokr_w1
  → transformer_blocks.6.attn.to_q.lokr_w1          (kohya + native passes)
  → layer "transformer_blocks.6.attn" = {"to_q.lokr_w1", "to_q.lokr_w2"}   (grouping)
  → ValueError: Unsupported lora format: dict_keys(['to_q.lokr_w1', 'to_q.lokr_w2'])

…and the whole adapter fails to load at generation time.

Pre-PR the identical file loaded: the verbatim flattened key rsplit at its only dot, grouped into a valid LoKRLayer/FullLayer, and was merely skipped with the per-layer "Failed to find module" warning at apply time. The rsplit fallback was safe precisely because flattened kohya paths contain no dots; the rewrite introduces dots without teaching the grouper the single-component LyCORIS suffixes.

Trigger: a mixed-algorithm LyCORIS file (per-module algorithms are a supported LyCORIS preset feature) containing at least one ordinary lora_down/up pair plus one LoKr/LoHa/full module. It passes install-time validation because the orphan-pair guard only checks lora_A/B/down/up suffixes. This also interacts badly with #9424: once LoKr keys install, a pure-LoKr kohya adapter would hit this on every layer.

Fix is small — either gate the rewrite so a key whose weight suffix isn't one the grouper can re-split stays verbatim, or add the LyCORIS single-component suffixes (lokr_w1, lokr_w2, hada_w1_a…hada_w2_b, diff, diff_b, on_input) to _SUFFIX_TO_VALUE_KEY. The first option preserves this PR's "left untouched unless we're sure" philosophy.

Non-blocking notes

2. "Left untouched verbatim" contract violated for over-accepted Sequential indices. The tree uses a bare INDEX_PLACEHOLDER for tmlp/tproj/txtmlp, so e.g. lora_unet_tmlp_1 (an activation — no weights exist there) is rewritten to tmlp.1.*, a half-converted key the native pass doesn't recognize, rather than being left verbatim as the module docstring and test_unrecognized_kohya_flattened_keys_are_left_untouched promise. End behavior is still warn-and-skip, so this is diagnostics only. Enumerating the literal indices ({"0": {}, "2": {}}-style) would close it.

3. New TypeError for state dicts containing a non-string key (reachable from .pt/.ckpt sources only). A dict with an int key + PEFT text_fusion keys + kohya keys loaded fine pre-PR (nothing looked native, so the native pass early-returned). Post-PR the kohya rewrite flips _looks_like_native_krea2_lora to True, and the per-key _looks_like_native_krea2_key(key) call then evaluates "txtfusion" in 0 → TypeError: argument of type 'int' is not iterable. The missing isinstance(key, str) guard is pre-existing, but this PR makes it fire for a dict class that previously loaded — worth the one-line guard while you're here.

4. Gap, fine to defer: a transformer-only kohya adapter (no txtfusion keys) can't install even with an explicit Krea-2 base override — _has_krea2_lora_keys matches kohya keys only via the txtfusion substring (the blocks-only fallback requires dotted .attn.wq.), and _KREA2_SUPPORTED_LORA_PREFIXES has no lora_unet_ spelling. You already declared auto-detection out of scope, but the override escape hatch — whose comment exists precisely for adapters lacking the auto-detection keys — doesn't cover the new layout either, so the converter's new capability is unreachable for that adapter class. Adding a lora_unet_ entry to the override prefixes would close it, here or in a follow-up.

Everything above was verified by executing the actual conversion pipeline on the triggering state dicts (and the pre-PR behavior confirmed by bypassing the new pass). The PR's 38 new tests plus the full tests/backend/patches and tests/backend/model_manager suites pass locally.

[lstein]

Thanks for the quick turnaround — I re-ran the adversarial review against 664d288550 (fresh-context reviewer attacking the fix itself, every finding re-verified by executing the conversion pipeline on the branch). Most of it lands cleanly:

• The suffix gate fixes the original repro: mixed files with lokr_w1/w2, hada_*, or diff/diff_b modules now load, with the LyCORIS module left verbatim and warn-skipped at apply time.
• The literal Sequential indices work as claimed: tmlp_1/tproj_0/txtmlp_2 stay verbatim, tmlp_0/2, tproj_1, txtmlp_1/3 still convert to the right diffusers names, and the new exact-match-before-placeholder precedence in _kohya_module_path_is_leaf correctly mirrors insert_periods_into_kohya_key (round-trip fuzz over every valid native path found no disagreement).
• Every key iteration in the file is now int-key safe (kohya pass, native pass, detection, grouping — all verified with a non-string key present).
• A transformer-only kohya adapter now installs via the explicit Krea-2 override, and FLUX's lora_unet_double_blocks_ spelling is still rejected.

But the blocking finding is only narrowed, not fixed — the most realistic trigger still aborts the load.

1. Blocking: the per-key gate orphans .alpha (and .dora_scale) from their verbatim LyCORIS siblings

LyCORIS saves an alpha tensor per module for LoKr/LoHa — see this repo's own captured fixtures, e.g. qwen_image_lora_kohya_format.py, where every lokr_w1/lokr_w2 pair has a sibling .alpha in exactly this flattened spelling. .alpha is in _SUFFIX_TO_VALUE_KEY, so the gate rewrites it while its lokr_w1/lokr_w2 siblings stay verbatim, splitting one module into two groups:

lora_unet_blocks_6_attn_wq.lokr_w1      → verbatim
lora_unet_blocks_6_attn_wq.lokr_w2      → verbatim
lora_unet_blocks_6_attn_wq.alpha        → transformer_blocks.6.attn.to_q.alpha
  → group "transformer_blocks.6.attn.to_q" = {"alpha"}
  → ValueError: Unsupported lora format: dict_keys(['alpha'])   # whole load aborts

The weight-decomposed variant (wd=True) fails the same way through .dora_scale: the orphaned {"dora_scale"} group dispatches to DoRALayer and dies with KeyError: 'lora_up.weight'. Pre-PR both dicts loaded (the flat module grouped as {lokr_w1, lokr_w2, alpha} → valid LoKRLayer, warn-skipped at apply).

test_kohya_lycoris_algorithm_keys_do_not_abort_the_load passes only because its suffix tuples omit alpha; adding "alpha" to any tuple reproduces the abort.

Fix: make the gate per-module rather than per-key — pre-scan the flat paths and rewrite a path's keys only if every suffix sharing that flat path is in _SUFFIX_TO_VALUE_KEY. That keeps each module atomic: fully converted or fully verbatim, never split. (And please add "alpha" to the test tuples.)

Non-blocking notes

2. lora_unet_blocks_ sweeps Wan and Anima kohya LoRAs into the override. The comment says the prefixes are spelled per-module so they "would [not] match every other architecture's kohya LoRA too", but Wan's kohya spelling is literally lora_unet_blocks_<idx>_... (wan_lora_conversion_utils._KOHYA_KEY_REGEX) and Anima's is lora_unet_(llm_adapter_)?blocks_<idx>_.... Verified: a Wan-keyed or Anima-keyed pair with an explicit Krea-2 override is ACCEPTED, installs, and then silently no-ops at generation (the unflattener rejects self_attn/cross_attn, so every layer warn-skips). It needs the explicit override to trigger, so a mislabeled file is required — but that override path is exactly where a user with a mislabeled file ends up, and "installs then does nothing" is the failure mode the install-time validation exists to prevent. A cheap tightening: during override validation, accept a lora_unet_blocks_ file only if at least one of its flat paths actually unflattens to a Krea-2 leaf.

3. The doubled-separator spelling can't install. The converter deliberately tolerates lora_unet__blocks_... (the lstrip, plus its dedicated test), but none of the new prefixes match the doubled underscore, so that variant of a transformer-only adapter is still rejected by the override (NotAMatchError) — the class of file this part of the fix was meant to make installable.

Attacks that failed, for the record: precedence disagreement between the two tree-walkers (none — verified over all valid paths), rewritten-vs-verbatim key collisions (impossible by construction; the source_keys guard still fires for real aliases), suffix-less kohya keys (partition yields "." → verbatim, same as pre-PR), .lora_magnitude_vector.weight in kohya spelling (partition splits at the first dot, so the full two-segment suffix matches and groups as dora_scale correctly), int keys anywhere in the pipeline. Both touched test files plus the full tests/backend/patches and tests/backend/model_manager suites pass locally (1079 passed).