[LoRA] add LoKr adapter support (Z-Image, Flux2/Klein) - #3
Open
christopher5106 wants to merge 3 commits into
Open
christopher5106 wants to merge 3 commits into
christopher5106 wants to merge 3 commits into
Conversation
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 huggingface#13221
… 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>
The file exercises checkpoint conversion and loading for LoKr adapters, not LoRA. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Upstream: huggingface#14163 (open since 2026-07-10). Both review comments were addressed in 96bb95c on 2026-09-26, and a re-review was requested.
Problem: diffusers only loads LoRA adapters, so LoKr (LyCORIS Kronecker-product) checkpoints for Z-Image and FLUX.2 / Klein cannot be loaded.
Change:
load_lora_adapterdetectslokr_keys and injects a peftLoKrConfiginferred from the tensor shapes.Test:
tests/lora/test_lora_lokr.py.In
scenario: yes, cherry-picked as one commit. The branch has two commits, to squash at the next re-sync.