feat(model-manager): support LyCORIS LoKr LoRAs for Krea-2

[Balrog994]

Summary

Krea-2 LoRAs in LyCORIS LoKr format can't be installed at all at the moment. The probe turns them away with:

model does not match Krea-2 LoRA heuristics (no complete lora_A/B or lora_down/up pair)

I hit this with an adapter trained in ai-toolkit, which emits LoKr by default, so I'd guess a fair few of the Krea-2 LoRAs going around are in the same boat.

The Krea-2 probe only knows about lora_A/lora_B and lora_down/lora_up. LoKr carries Kronecker factors instead (lokr_w1/lokr_w2, or the factored _a/_b pairs plus an optional lokr_t2), so _has_complete_lora_pair() says no and the adapter is never identified. Setting the base by hand doesn't get you anywhere either, that path goes through the same helper.

There's a second one further down: _SUFFIX_TO_VALUE_KEY doesn't list the lokr_* suffixes, so even past the probe _group_by_layer() would drop the weights and leave every layer holding nothing but its alpha.

The rest was already there. LoKRLayer exists, any_lora_layer_from_state_dict() reaches for it the moment it sees lokr_w1, and the Qwen-Image converter handles LoKr fine. Only the Krea-2 side was never wired up, which fits it being marked WIP in #9304.

So this adds the two missing pieces: a _has_lokr_layer() helper that both config gates consult, and the lokr_* suffixes in the converter's suffix map. I left the signature-module check alone, so a LoKr that only touches generic transformer blocks still won't be claimed as Krea-2.

Related Issues / Discussions

Follows on from #9304.

QA Instructions

Take any ai-toolkit LoKr for Krea-2 (keys look like diffusion_model.txtfusion.<block>.attn.wq.lokr_w1) and install it through the Model Manager. Before, it's rejected with the error above. After, it installs as base=krea-2, format=lycoris and turns up in the LoRA picker.

I tested with a 32-module adapter on the text-fusion stage against a GGUF-quantised Krea-2 Turbo. It converts to 32 LoKRLayers, every target key resolves to a real nn.Linear in Krea2Transformer2DModel, the reconstructed weight (4×4 ⊗ 640×640 = 2560×2560) matches the layer it patches, and generation comes out visibly different with the LoRA switched on.

8 new tests, covering full and factored LoKr, the native→diffusers key renaming, and that a LoKr without the Krea-2 signature modules still gets turned away. tests/backend/patches/ and tests/backend/model_manager/ pass (1030 passed, 133 skipped, 1 xfailed), as do ruff check and format. I only ran them on 3.12/Windows locally, so the rest of the matrix is on CI.

Merge Plan

Nothing special, backend only.

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)

[Balrog994]

Merged main in — thanks @Pfannkuchensack for keeping the branch current while this sat.

#9449 landed just after your last merge and touched the same four files, which is where the
conflict came from. The textual part was trivial (we both added imports). The part worth flagging is
that the two changes overlap behaviourally, and I've changed one of your test expectations because
of it.

Your per-module gate converts a flattened module only when every suffix on it is one
_group_by_layer knows, and LyCORIS algorithms stayed verbatim because lokr_* wasn't such a
suffix. This PR makes it one. So a kohya-flattened LoKr module now un-flattens and lands on its real
Linear instead of degrading to the per-layer "Failed to find module" warning:

lora_unet_txtfusion_layerwise_blocks_0_attn_wq.lokr_w1
  before: lora_transformer-lora_unet_txtfusion_layerwise_blocks_0_attn_wq   (no-op)
  after:  lora_transformer-text_fusion.layerwise_blocks.0.attn.to_q          (LoKRLayer)

The hazard your comment guards against — the grouper's blind rsplit(".", 2) cutting inside a
dotted module path — doesn't apply to a suffix the grouper recognises, so I don't think anything is
lost here.

Concretely: two of the five parametrizations of
test_kohya_lycoris_algorithm_keys_do_not_abort_the_load were failing, both the LoKr ones. Rather
than special-case them out I moved the expected layer into the parametrization — the LoKr cases
expect the converted module, hada_* and diff still expect the verbatim key. The invariant the
test is actually about, that a mixed-algorithm file never takes the whole adapter down, holds either
way. Also added a test for the kohya + LoKr path end to end.

If you'd rather this PR left the kohya path completely alone, the alternative is to give that gate
its own suffix set instead of reusing _SUFFIX_TO_VALUE_KEY, so LoKr stays unsupported there and
your test is untouched. It's a smaller blast radius but it means a kohya-flattened LoKr keeps
silently doing nothing, which seemed like the worse outcome. Happy to switch if you disagree.

tests/backend/patches/ and tests/backend/model_manager/ are green locally (1415 passed) and ruff
is clean.

[Balrog994]

Merged main again to clear the conflicts — #9517 and #9518 both landed in these files. The only
real adjustment was _has_lokr_layer taking the same key_filter callable that #9518 switched
_has_complete_lora_pair to; the lokr_* and dora_magnitude suffix entries merged cleanly.

Tests and ruff green, nothing else changed on this side.

[Balrog994]

Merged main again to bring the branch up to date with the 7.0 line (it had fallen ~3.7k commits behind). No conflicts this time: of the four files this PR touches, only configs/lora.py moved on main, and that was the new MiniMax H3 and LTX-2 LyCORIS configs, which are independent of the Krea-2 path.

Krea-2 tests (test_krea2_lora_config.py, test_krea2_lora_conversion_utils.py) are green, and ruff is clean on the touched files. Nothing changed on this side.

[Pfannkuchensack]

Thanks for this. The routing works, and it fixes the Krea-2 LoKr files that are actually out there. Nothing blocking: three things I'd fix before merge, plus a few nits.

What I checked

Real-world layouts. I read only the safetensors headers of published Krea-2 LoKr adapters on HF and built files with identical keys and shapes. Between them they cover four layouts:

Source	Layout
ai-toolkit (eviligh3e/chle_krea2_lokr_v3, azikimbowe/krea2_lokr_stile)	diffusion_model.blocks.N.attn.gate.{lokr_w1,lokr_w2,alpha}
ai-toolkit, factored w2 (another public adapter)	native, lokr_w1 + lokr_w2_a/lokr_w2_b
NeedAHugNOW/City-Pop-LoKr	kohya flat, lora_unet_blocks_0_attn_gate.lokr_w1 + lokr_w2_a/lokr_w2_b (one module with a full lokr_w2)
Kentus/Naz_Ferra_KREA2_LoKr	native, 264 modules including txtmlp.3, tproj.1, last.linear

On main all five classify as Unknown. On this branch each matches LoRA_LyCORIS_Krea2_Config and nothing else. Every layer converts to a LoKRLayer whose reconstructed weight has the shape of a real nn.Linear in Krea2Transformer2DModel. LoKr files shaped like Qwen-Image, FLUX, Anima, Wan and Z-Image adapters classify exactly as they do on main; none is claimed as Krea-2.

End to end. I installed the real city-pop-texture-v2.safetensors through the model manager; it comes in as krea-2 / lora / lycoris. I then generated on Krea-2 Turbo Q4_K_M (GGUF) at 1024², 8 steps, fixed seed, on an RTX 4090. There were no Failed to find module warnings, the image changes clearly (mean absolute pixel difference 40.5 against no LoRA), and repeated runs are bit-identical.

The focused tests pass (104), and ruff check / ruff format --check are clean.

Should fix

1. An incomplete LoKr module installs, then fails at generation.

_has_lokr_layer (lora.py#L959-L984) accepts any single lokr_* key, so the "no orphaned half" comment doesn't hold. LoKRLayer needs a complete w1 side and a complete w2 side. With keys under diffusion_model.txtfusion.layerwise_blocks.0.attn.wq.:

Module carries	Probe on this branch	At generation
only lokr_w1	krea-2 / lycoris	AssertionError in LoKRLayer.__init__
only lokr_w2	krea-2 / lycoris	ValueError: Unsupported lora format
lokr_w1 + lokr_w2_a (no _b)	krea-2 / lycoris	AssertionError
lokr_w1 / lokr_w2 + dora_scale	krea-2 / lycoris	KeyError: 'lora_up.weight' (dispatched to DoRALayer)

This config otherwise rejects that kind of file at identification time (_lora_weight_keys_are_all_paired). A per-module check, grouped by key.rsplit(".", 1)[0], would close the gap:

def _lokr_module_is_complete(suffixes: set[str]) -> bool:
    has_w1 = "lokr_w1" in suffixes or {"lokr_w1_a", "lokr_w1_b"} <= suffixes
    has_w2 = "lokr_w2" in suffixes or {"lokr_w2_a", "lokr_w2_b"} <= suffixes
    return has_w1 and has_w2

For dora_scale on a LoKr module, either reject it at probe time or strip it the way anima_lora_conversion_utils._make_layer_patch already does.

2. The explicit-override LoKr test doesn't cover the override change.

test_explicit_krea2_override_accepts_lokr_lora (test_krea2_lora_config.py#L395-L401) uses a fixture containing txtfusion, which the automatic path already accepts. With the or _has_lokr_layer(state_dict, _key_is_supported_krea2_layout) term removed, the test still passes.

Two cases would actually exercise the override:

• A LoKr without the Krea-2 signature, e.g. transformer.transformer_blocks.0.attn.to_q.lokr_w1/w2, or kohya lora_unet_blocks_0_attn_wq.lokr_w1/w2 (which auto-probes as Unknown).
• A rejection case: a Wan kohya LoKr (lora_unet_blocks_0_self_attn_q.lokr_w1) must still be refused under the override.

3. The conversion tests can't tell the factors apart.

The new LoKr conversion tests (test_krea2_lora_conversion_utils.py#L497-L587) use torch.ones and assert only shapes or is not None. kron(ones(2, 2), ones(3, 4)) and kron(ones(3, 4), ones(2, 2)) are both 6×8. So swapping lokr_w1↔lokr_w2 or _a↔_b in _SUFFIX_TO_VALUE_KEY leaves all of them green.

To pin the mapping:

• Use distinct random tensors and compare get_weight against torch.kron(w1, w2) and torch.kron(w1_a @ w1_b, w2_a @ w2_b).
• For the factored case, also assert scale() == alpha / rank.

The code itself gets this right today; the tests just wouldn't notice a regression.

Nits

• is_state_dict_likely_krea2_lora (krea2_lora_conversion_utils.py#L165-L191) has no caller outside its tests. This PR adds a third copy of the LoKr suffix list there, which already lacks .lokr_t2, plus the function's first two tests. I'd leave it untouched (or delete it) and drop those tests.
• The LoKr suffixes now live in three places: _LOKR_WEIGHT_SUFFIXES, the tuple above, and _SUFFIX_TO_VALUE_KEY. One constant in krea2_lora_constants.py would keep them in sync.
• Some comments are stale:
  • krea2_lora_conversion_utils.py#L119-L125 and #L146-L151 still say lokr_w1 modules stay verbatim.
  • The _validate_looks_like_lora docstring and comment (lora.py#L1049-L1053) still describe only A/B pairs.

Not this PR, but worth knowing

On GGUF Krea-2, a LoKr nearly doubles denoise time. Same setup as above, one warm-up then three interleaved runs:

	Denoise (mean)
no LoRA	10.65 s
PEFT LoRA (A/B)	11.66 s (1.09×)
City-Pop LoKr	20.25 s (1.90×)

Sidecar patching has a low-rank fast path only for LoRALayer. Every other layer type materializes its full delta weight on every forward. That's pre-existing and shared with the other bases, so it's a candidate for a separate follow-up.

[Balrog994]

Thanks for the thorough review, and especially for testing against the published adapters and running it end to end. All three should-fix points and the nits are addressed in 6de2785.

1. Incomplete LoKr modules. You're right: the "no orphaned half" comment was wrong. Completeness is now checked per module (grouped by key.rpartition(".")[0]), on both the automatic and the explicit-override path, next to _lora_weight_keys_are_all_paired. Each operand must be present exactly once, whole or as a complete _a/_b pair. That is slightly stricter than your sketch, because it also rejects a module that spells an operand both ways (lokr_w1 plus lokr_w1_a/_b), which trips the same assertion in LoKRLayer.__init__. All four rows of your table, plus that case and a lone lokr_t2, are now refused at identification.

For dora_scale on a LoKr module I went with stripping in the converter, the way anima_lora_conversion_utils._make_layer_patch does, with a warning. The ai-toolkit magnitude spelling is stripped too, since it dispatches to DoRALayer the same way. This keeps such files usable, though only as an approximation of what was trained, and it stays consistent with Anima. Happy to switch to rejecting if you'd rather keep the probe strict.

2. Override test. Replaced it with two cases the automatic probe returns as Unknown, dotted transformer.transformer_blocks.0.attn.to_q.lokr_* and kohya lora_unet_blocks_0_attn_wq.lokr_*. Each case asserts the automatic probe raises NotAMatchError, then that the override accepts it. Also added the Wan kohya LoKr (lora_unet_blocks_0_self_attn_q), which is still refused under the override.

3. Conversion tests. They now use random factors and compare get_weight against torch.kron(w1, w2) / torch.kron(w1_a @ w1_b, w2_a @ w2_b). The factored cases also check scale() == alpha / rank.

To make sure the new tests actually guard what they claim, I injected each regression temporarily and confirmed they go red:

Injected regression	Failing tests
override or _has_complete_lokr_layer(...) term removed	2
completeness gate removed (automatic / override)	12 / 6
operand check loosened to "any factor present"	6
lokr_w1 ↔ lokr_w2 swapped in _SUFFIX_TO_VALUE_KEY	5
lokr_w1_a ↔ _b / lokr_w2_a ↔ _b swapped	1 / 2
DoRA-magnitude stripping removed	3

Nits.

• is_state_dict_likely_krea2_lora is back to exactly what main has, and its two tests are gone.
• The LoKr suffixes now live once, as LOKR_WEIGHT_KEYS in krea2_lora_constants.py. Both _LOKR_WEIGHT_SUFFIXES in configs/lora.py and _SUFFIX_TO_VALUE_KEY derive from it.
• The stale comments in the kohya pass and the _validate_looks_like_lora docstring are updated.

Krea-2 focused tests: 123 passed. ruff check / ruff format --check are clean.

Thanks also for the GGUF timing numbers. Agreed that the missing low-rank fast path for non-LoRALayer sidecars is a separate follow-up.

[jimmybrancaccio]

The suffix mappings in this PR also address the generation-time error ValueError: Unsupported lora format: dict_keys(['alpha']) for dotted native LoKr keys, independently of the model-identification changes.

A minimal synthetic example of the grouping failure on main is:

state_dict = {
    'diffusion_model.blocks.0.attn.wq.alpha': torch.tensor(1.0),
    'diffusion_model.blocks.0.attn.wq.lokr_w1': torch.randn(2, 2),
    'diffusion_model.blocks.0.attn.wq.lokr_w2': torch.randn(3, 3),
}

_group_by_layer() explicitly splits .alpha, but the unrecognized one-part .lokr_w1 / .lokr_w2 suffixes fall through to rsplit('.', maxsplit=2). This puts the alpha and its factors into different groups, so conversion reaches any_lora_layer_from_state_dict() with only alpha.

The seven lokr_* suffix entries proposed here keep those values together and route them to the existing LoKRLayer; they do not require new Kronecker-product math. Posting this additional reproduction here because this PR already contains the converter fix, rather than opening a duplicate PR.