[LoRA] add LoKr adapter support (Z-Image, Flux2/Klein) (#14163)
* [LoRA] add LoKr adapter support (Z-Image, Flux2/Klein)
Adds loading of LoKr (LyCORIS Kronecker product) adapters:
- `load_lora_adapter` detects `lokr_` keys and injects a peft `LoKrConfig`,
inferred from the tensor shapes via `_create_lokr_config` (decompose
factor, per-module rank/alpha patterns).
- State dict conversions for the formats in the wild: ai-toolkit Z-Image
(dotted diffusers paths under `diffusion_model.`), ai-toolkit BFL Flux2
(fused qkv), LyCORIS underscore format, and bare dotted diffusers paths.
- BFL fused-QKV LoKr cannot be split exactly into separate Q/K/V Kronecker
factors, so `Flux2LoraLoaderMixin.load_lora_weights` fuses the model's
QKV projections and maps the adapter 1:1 (exact).
- Alpha follows the LyCORIS convention: scaling applies only to
rank-decomposed factors and is baked into the weights at conversion.
Fixes #13221
* [LoRA] LoKr review follow-ups: adapter-agnostic error message, refuse fused-QKV load over an unfused adapter
- "Invalid adapter checkpoint. We currently support LoRA and LoKr." replaces the
message that still said "LoRA checkpoint" and described the substring check.
- Loading a fused-QKV LoKr checkpoint now refuses when an adapter is already
injected on to_q/to_k/to_v or add_{q,k,v}_proj: fuse_qkv_projections() would
replace those modules and orphan it. Covered by a test.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
* [LoRA] rename the LoKr tests to test_lokr_conversion.py
The file exercises checkpoint conversion and loading for LoKr adapters, not LoRA.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* [LoRA] LoKr review follow-ups: scope to Z-Image/Flux2, model-level fused QKV, LoKrTesterMixin
- Only the Z-Image and Flux2 loader mixins accept LoKr checkpoints, the two that
convert them; the other mixins are back to their LoRA-only check.
- The fused-QKV handling moves to `_maybe_fuse_qkv_projections_for_lokr` in
peft_utils, called from `load_lora_adapter`, so the model fuses its projections
whichever entry point loads the adapter. The quantized-model error explains why.
- `_bake_lokr_alpha` becomes `_bake_lokr_alpha_` as it edits the dict in place.
- The BFL Flux2 converter no longer accepts expanded diffusers block names under
`diffusion_model.`: no LoKr checkpoint uses that layout; diffusers-named
checkpoints have no prefix and take the generic converter.
- The LyCORIS converter names the keys it does not recognize.
- Tests move from tests/lora to a `LoKrTesterMixin` in tests/models/testing_utils,
with the checkpoint formats tested on the Flux2 and Z-Image model test classes.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* [LoRA] mark ZImageLoraLoaderMixin.load_lora_weights as copied from Flux2
---------
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>