Add three adapter families to the z-image native loader and chain them
with the existing lora path through load_safetensors. Mixed-family
files (for example gta6_amateur_photography_zimagebase_v2.safetensors,
which carries lora and lokr groups in the same file) now load
completely instead of having one family silently dropped.
Shared helpers in pipelines/z_image/zimage_lora.py parse keys by
suffix list, rename legacy attention.out and attention.wo to
attention.to_out.0, and split fused attention.qkv into to_q/k/v. For
lora the split chunks the up weight along dim 0. For lokr the split
emits three NetworkModuleLokrChunk entries that share the tensors and
slice the kronecker product at apply time.
Fused attention.qkv for loha and oft is skipped with a warning. No
NetworkModuleHadaChunk exists, oft rotations are tied to out_features
and cannot be cleanly split across q/k/v, and no real z-image adapter
in that layout exists today.
load_safetensors for zimage chains try_load_lora, try_load_lokr,
try_load_loha and try_load_oft and merges their module dicts into a
single Network so mixed files load every module.
Add zimage to allow_native so lora_force_diffusers picks between
native and diffusers. Before this, zimage always took the diffusers
path regardless of the setting.
pipelines/z_image/zimage_lora.py reads the safetensors and writes
directly into network_layer_mapping, so Z-Image LoRAs no longer go
through the diffusers PEFT converter that raised KeyError on
state dicts with partial alpha keys.
Key formats handled: ai-toolkit, kohya lora_unet_, bare transformer.
and no-prefix. Pre-refactor fused attention.qkv is split into
to_q/k/v; attention.out and attention.wo are renamed to
attention.to_out.0. Alpha and dora_scale are preserved.
Load Flux2/Klein LoRAs as native NetworkModuleLora objects, bypassing
diffusers PEFT. Handles kohya (lora_unet_), AI toolkit (diffusion_model.),
diffusers PEFT (transformer.), and bare BFL key formats with automatic
QKV splitting for double block fused attention weights.
Includes shape validation to reject architecture-mismatched LoRAs early.
Respects lora_force_diffusers setting to fall back to PEFT when needed.
- detect f2 model type for LoRAs via metadata, architecture, and filename/folder
- preprocess bare BFL-format keys with diffusion_model prefix for Flux2LoraLoaderMixin
- handle LoKR format via native NetworkModuleLokr with on-the-fly kron(w1, w2)
- add NetworkModuleLokrChunk for fused QKV split into separate Q/K/V modules
- activate native modules loaded via diffusers path
- improve error message for Flux1/Flux2 architecture mismatch
Vladimir Mandic
CalamitousFelicitousness
awsr
Disty0