A bare key that matches no known prefix is parsed with the
bare-diffusers sentinel and handed to the resolver instead of being
dropped at parse time. The per-arch lists of bare diffusers prefixes are
gone, and a path that names no live module counts as unmapped.
A full-weight extraction on Z-Image bound 308 modules and applied 172 of
them, silently dropping the rest, and left 71 more unmapped.
assign_network_names_to_compvis_modules puts every transformer module in
network_layer_mapping but skips stamping network_layer_name on norms,
which is the attribute the apply pass keys off. try_load_full bound those
modules through the mapping and they then never applied; stamp them
loader-locally, as try_load_norm already does.
Z-Image also names three module groups differently from the diffusers
tree: the qk-norms (q_norm/k_norm vs norm_q/norm_k), and the patch
embedder and final layer, which live in ModuleDicts keyed by
"{patch_size}-{f_patch_size}" and so carry a key the checkpoint has no
notion of. Read that key from the live model rather than hardcoding it.
The counts close exactly: 68 qk-norms plus 3 non-block targets are the 71
that went unmapped.
zimage, chroma, ernie and krea2 chained only lora/lokr/loha/oft while
flux2 and anima ran all eight families, so ia3/glora/norm/full files
(e.g. full-diff extractions with diff/diff_b keys) reported not loaded
on the short-chain arches. The generic family loaders are arch-agnostic;
wire the missing four into each chain.
- add ia3/glora/norm/full wrappers and chain entries to the four arch
modules, with matching suffix/marker re-exports
- add a chain-level full-diff test per suite: zimage covers the legacy
attention.out alias and the fused-qkv skip, chroma the proj rename,
ernie the passthrough
transformer., bare-diffusers, and lora_transformer_ bases are already in
network-key form for every arch, yet each per-arch resolve_targets repeated the
same passthrough branch for them. Move that into a shared
PASSTHROUGH_PREFIXES_DEFAULT set consulted by resolve_group_targets, leaving each
arch's resolve_targets to only the prefixes it actually rewrites (kohya / BFL).
lycoris_ stays in flux2, the one arch that recognizes it.
Pure refactor: the same keys resolve to the same modules.
Replaces zimage's four family loaders with thin wrappers binding
native_loader's generics to z-image's prefix tuples and resolve_targets.
resolve_targets folds the legacy attention.qkv split and attention.out
alias rename into the path-resolution step. Fused qkv now emits three
ChunkSpec(idx, total=3) entries; attention.out / attention.out.0 /
attention.wo aliases collapse to attention.to_out.0.
BARE_DIFFUSERS_PREFIXES allows bare paths starting with layers. /
noise_refiner. / context_refiner. to pass through to the loader. This
matches real Z-Image LoRAs exported via
ZImageTransformer2DModel.save_lora_adapter().
LoHA on fused qkv now binds via NetworkModuleHadaChunk (added to shared
infra by the flux2 PR) instead of being skipped. Test renamed to
test_loha_legacy_fused_qkv_chunked.
parse_key returns (prefix_used, base, suffix) instead of the old
(network_key, suffix); parse test updated.
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.
CalamitousFelicitousness
vladmandic
resonantsky