[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