fix(lora): support ai-toolkit DoRA magnitudes for Krea-2 LoRAs

[Pfannkuchensack]

Summary

fix: Krea-2 DoRA LoRAs trained with ai-toolkit fail to load with ValueError: Unsupported lora format: dict_keys(['to_gate.magnitude', 'to_k.magnitude', 'to_q.magnitude', 'to_v.magnitude']).

Why. ai-toolkit stores the DoRA magnitude as <layer>.magnitude, a suffix the Krea-2 converter's suffix table didn't know. The key fell through to the generic rsplit fallback, so every magnitude in an attention block got grouped into a bogus parent layer (...attn) whose sub-keys were to_q.magnitude, to_k.magnitude, … — exactly the dict shown in the error.

While fixing that, a second, pre-existing bug surfaced. Two mutually incompatible DoRA magnitude conventions exist, and they index different axes of the (out_features, in_features) weight:

Trainer	Key	Indexes	Shape
LyCORIS / kohya / OneTrainer	.dora_scale	in_features	(1, in)
PEFT / Diffusers / ai-toolkit	.lora_magnitude_vector.weight, .magnitude	out_features	(out,)

DoRALayer only implemented the LyCORIS math (norm over the output dim). The existing .lora_magnitude_vector.weight -> dora_scale mapping therefore applied a PEFT magnitude with LyCORIS math: silently wrong weights on square layers, and a broadcast error on non-square ones (e.g. Krea-2's ff.down, 6144 × 16384).

How.

• DoRALayer carries an explicit magnitude_is_out_dim flag instead of guessing the orientation from tensor shape (shape alone can't disambiguate a square layer). The out-dim branch normalizes over the input dim, matching PEFT/ai-toolkit; the LyCORIS path is untouched.
• any_lora_layer_from_state_dict also dispatches on the new dora_magnitude value key.
• The Krea-2 converter maps .magnitude and routes both PEFT-style magnitude keys to dora_magnitude.

The Krea-2 model-config probe already accepted these files, so no identification changes were needed.

Related Issues / Discussions

Closes #9515

Distinct from #9424 (Krea-2 LoKr/LyCORIS support), which is not addressed here.

QA Instructions

Test adapter: https://huggingface.co/khronex/krea2-dora-ehrmantraut-test (ai-toolkit 0.10.17, network.type: dora, native/ComfyUI key layout, 256 modules / 768 tensors)

End-to-end (this is the reproducer from the issue)

1. Install the LoRA above (it probes as lora / krea-2 / lycoris).
2. Generate with Krea-2-Raw + the LoRA applied at weight 1.0.
3. Before this PR: generation fails during text encoding with Unsupported lora format. After: the image generates and the trained subject is clearly present.

Verified on an RTX 4090, Krea-2-Raw (fp8_storage: true), 1024×1024, 28 steps, CFG 4.5, seed 1234, identical prompt with and without the LoRA. The subject identity changes as expected while the composition stays put. Server log for the LoRA run shows krea2_lora_loader executing, and zero Failed to find module for LoRA layer key warnings — all 256 layers patch onto real modules.

Structural check

All 256 converted patch keys were resolved against a meta-device Krea2Transformer2DModel (stock config): 0 missing modules, 0 shape mismatches, including magnitude length == out_features.

Numeric check

The merged weight was compared against ai-toolkit's own forward pass (ToolkitModuleMixin.forward + DoRAModule.apply_dora) for out < in, square, and out > in layers — max absolute deviation 1.4e-6.

Unit tests

pytest tests/backend/patches tests/backend/model_manager/configs/test_krea2_lora_config.py

New tests/backend/patches/layers/test_dora_layer.py covers both magnitude conventions (the out-dim test asserts equivalence with the ai-toolkit reference formula); test_krea2_lora_conversion_utils.py gains an ai-toolkit .magnitude regression test plus orientation assertions on the existing DoRA tests.

Merge Plan

Normal merge — backend only, no DB or schema changes.

Note the behavior change for adapters that were already loading: Krea-2 LoRAs with .lora_magnitude_vector.weight (PEFT/Diffusers DoRA) previously had their magnitude applied along the wrong axis. Square layers will now produce different — correct — output; non-square ones previously raised a broadcast error.

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)