[LoRA] add LoKr adapter support (Z-Image, Flux2/Klein)

[christopher5106]

What does this PR do?

Adds support for loading LoKr (LyCORIS Kronecker product) adapters, for Z-Image and Flux2/Klein checkpoints in the wild, plus a generic path that works for any model whose LoKr state dict already uses diffusers module names.

Fixes #13221. Also addresses #13261 and #13137 (Flux 2 Klein LoKr).

Follows up on the discussion in #13326 — see "Why not split the fused QKV" below for measurements that motivated a different approach for the fused-QKV part. cc @CalamitousFelicitousness whose prototype this builds on.

How it works

• PeftAdapterMixin.load_lora_adapter detects lokr_ keys and injects a peft LoKrConfig instead of a LoraConfig. The new _create_lokr_config infers the config from tensor shapes: it reconstructs each module's Kronecker factorization and finds the decompose_factor under which peft's own factorization() reproduces every shape; per-module rank_pattern/alpha_pattern make peft recreate full vs. rank-decomposed factors exactly as stored (for full-matrix factors the rank is set to max(lokr_w2.shape) so peft also creates full matrices).
• Alpha follows the LyCORIS/ComfyUI convention: alpha / rank scaling applies only when a factor is rank-decomposed and is baked into lokr_w1 at conversion; full-rank modules ignore alpha (ai-toolkit stores a ~1e10 placeholder there). The config sets alpha = r so the runtime scaling is 1.0.
• State dict conversions cover the four formats found in the wild, each validated against a real published checkpoint:

Format	Example checkpoint	Conversion
ai-toolkit Z-Image (diffusion_model.layers.0.attention.to_q.lokr_w1)	F16/z-image-turbo-flow-dpo (the checkpoint from #13221)	_convert_non_diffusers_lokr_to_diffusers (paths already match the model)
ai-toolkit BFL Flux2 (diffusion_model.double_blocks.0.img_attn.qkv.lokr_w1)	puttmorbidly233/lora klein_snofs_v1_2.safetensors (the checkpoint from #13261)	_convert_non_diffusers_flux2_lokr_to_diffusers
LyCORIS underscore (lycoris_transformer_blocks_0_attn_to_q.lokr_w1)	gattaplayer/besch-flux2-klein-9b-lokr-lion-3e-6-bs2-ga2-v02	_convert_lycoris_flux2_lokr_to_diffusers
Bare dotted diffusers paths with .alpha keys	bghira/flux2-klein-9b-distillation-lokr	_convert_non_diffusers_lokr_to_diffusers

Why not split the fused QKV

BFL-format checkpoints apply LoKr to the fused QKV projection of the double blocks, while the diffusers Flux2 model has separate Q/K/V. For LoRA this split is exact (shared lora_A, chunked lora_B), but a Kronecker product delta over the fused projection is mathematically not a Kronecker product per chunk. #13326 approximated each chunk with a rank-1 Van Loan re-factorization; measuring that on the real klein_snofs_v1_2 checkpoint gives a mean relative delta error of 69% on the 48 QKV projections (max 79%, min 32%), i.e. the adapter loads but its QKV contribution is mostly destroyed — the same class of problem that got #13997 rejected.

Instead, the converter maps these modules to the model's fused attn.to_qkv / attn.to_added_qkv projections, and Flux2LoraLoaderMixin.load_lora_weights calls transformer.fuse_qkv_projections() (idempotent) before injection, with an info log. Loading is then exact. Quantized transformers raise a clear error for such checkpoints since fusing would concatenate packed quantized weights.

Tests / verification

• New ZImageLoKrTests and Flux2LoKrTests: injected get_delta_weight equals kron(w1, w2) exactly, covering fused QKV (img + txt streams), rank-decomposed factors with meaningful alpha, full-rank factors with placeholder alpha, and the LyCORIS format.
• Real layer-0 tensors from F16/z-image-turbo-flow-dpo loaded into a real-dimension (dim=3840) Z-Image transformer: all deltas exact, scaling 1.0.
• All four real checkpoints above convert with every key consumed and a valid inferred config.
• Existing suites pass: tests/lora/test_lora_layers_z_image.py (36 passed) and tests/lora/test_lora_layers_flux2.py (39 passed).

Who can review?

@sayakpaul @asomoza

[sayakpaul]

Will review next week but triggering an AI review as well.

[sergereview]

🤗 Serge says:

Solid, well-tested addition of LoKr adapter loading for Z-Image and Flux2/Klein. The design is coherent: shape-inferred LoKrConfig, alpha baked into the left Kronecker factor per LyCORIS convention, and a fuse-then-map strategy for the fused-QKV case that avoids the lossy per-chunk re-factorization. I verified the key integration points and found no blocking issues.

Correctness (verified, no defects)

• _create_lokr_config — collections is imported; factorization/LoKrConfig are imported lazily from peft; the decompose_factor search and rank/rank_pattern derivation are internally consistent and reproduced exactly by the tests.
• _bake_lokr_alpha — rank is read from the inner dimension of the decomposed factor (w2_b.shape[0] / w1_b.shape[0]), and full-matrix modules correctly drop alpha via the rank is None guard. Scaling matches the test expectations ((alpha/rank) * kron(w1, w2_a@w2_b)).
• Flux2 fused-QKV path — Flux2Transformer2DModel inherits fuse_qkv_projections from AttentionMixin; the "Added" in processor name guard does not spuriously trip (Flux2 processors are Flux2AttnProcessor/Flux2KVAttnProcessor), and fuse_projections creates to_qkv/to_added_qkv used by _get_fused_projections at inference. The .attn.to_qkv. substring check does not falsely match to_qkv_mlp_proj. is_quantized is a real attribute on the model.
• re is imported in lora_conversion_utils.py; logger exists in lora_pipeline.py.

Minor (non-blocking)

• In _bake_lokr_alpha, a malformed module that has a lokr_w*_b but neither lokr_w1 nor lokr_w1_a would raise KeyError at the final assignment. Not reachable for the four in-the-wild formats, so purely defensive — fine to leave.

Tests

• Flux2LoKrTests and ZImageLoKrTests cover fused QKV (img+txt), rank-decomposed factors with meaningful alpha, full-rank factors with placeholder alpha, and the LyCORIS underscore format, asserting get_delta_weight equals kron(w1, w2). All referenced imports resolve.

LGTM.

serge v0.1.0 · model: claude-opus-4-8 · 19 LLM turns · 23 tool calls · 110.2s · 880940 in / 7039 out tokens

[BenjaminBossan]

From the PEFT perspective, this looks correct to me. I have two small comments, but I'll leave it to the Diffusers maintainers to decide if they're worth addressing.

[BenjaminBossan]

The message could be confusing because it still mentions "LoRA checkpoint".

How about "Invalid adapter checkpoint. We currently support LoRA and LoKr."

I wouldn't really mention substring matching, as it's an implementation detail. Same applies to the error messages below.

[christopher5106]

Done in 96bb95c: "Invalid adapter checkpoint. We currently support LoRA and LoKr." at all 25 sites this PR touches, and no mention of the substring check anymore. The three pre-existing messages this PR did not touch keep their old text.

[BenjaminBossan]

Another potential error case would be if there is already an adapter loaded onto a non-fused q, k, or v layer, right? Maybe not terribly likely, but could be worth checking.

[christopher5106]

Good catch, done in 96bb95c: before calling fuse_qkv_projections() the loader now walks transformer.named_modules() and raises if any to_q/to_k/to_v or add_{q,k,v}_proj is already a PEFT BaseTunerLayer, naming the first one and pointing at unload_lora_weights(). Fusing would have replaced those modules and silently orphaned the adapter. Covered by test_lokr_fused_qkv_checkpoint_refuses_when_unfused_projections_are_adapted, which injects a plain LoRA on to_q first and checks the model is left unfused.

[christopher5106]

Both review comments are addressed in 96bb95c (message wording, and refusing a fused-QKV LoKr load when an adapter already sits on the unfused projections, with a test). The branch was also rebased onto current main in September; tests/lora/test_lora_lokr.py is green locally. @sayakpaul @BenjaminBossan could you take another look when you have bandwidth?

[BenjaminBossan]

Thanks for addressing my previous comments. There is no issue there, so from that point of view, it's a 👍 from me.

When going over the PR again, I saw that the unit tests specifically check the fused QKV part of the conversion step, but never the whole pipeline. This is probably not easy to achieve, as it would require the reference output. Just flagging this here for @sayakpaul, not sure what the requirements are in Diffusers. At a glance, I couldn't find LoRA conversion unit tests to compare to.

[BenjaminBossan]

Nit about naming: The file name should mention "conversion", whereas "lora" isn't really tested here.

[christopher5106]

Renamed to tests/lora/test_lokr_conversion.py in c6d5763, thanks!