diffusers
78bf2c62 - [LoRA] add LoKr adapter support (Z-Image, Flux2/Klein) (#14163)

Commit
7 days ago
[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>
Parents
Loading