From c35adc4d0d98e3a67058d8fb9cf1a372d1d72582 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Tue, 14 Apr 2026 19:41:26 +0100 Subject: [PATCH 1/2] feat(lora): native z-image lora loader 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. --- modules/lora/lora_load.py | 7 + modules/lora/lora_overrides.py | 1 + pipelines/z_image/zimage_lora.py | 212 +++++++++++++++++++++++++++++++ 3 files changed, 220 insertions(+) create mode 100644 pipelines/z_image/zimage_lora.py diff --git a/modules/lora/lora_load.py b/modules/lora/lora_load.py index 021697272..9790b5cc7 100644 --- a/modules/lora/lora_load.py +++ b/modules/lora/lora_load.py @@ -49,6 +49,13 @@ def load_safetensors(name, network_on_disk: network.NetworkOnDisk) -> network.Ne log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" type=lora {"cached" if cached else ""}') if cached is not None: return cached + if shared.sd_model_type == 'zimage': + from pipelines.z_image import zimage_lora + lora_scale = shared.opts.extra_networks_default_multiplier + zimage_net = zimage_lora.try_load_lora(name, network_on_disk, lora_scale) + if zimage_net is not None: + lora_cache[name] = zimage_net + return zimage_net net = network.Network(name, network_on_disk) net.mtime = os.path.getmtime(network_on_disk.filename) state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network') diff --git a/modules/lora/lora_overrides.py b/modules/lora/lora_overrides.py index d877889ca..49a9569bc 100644 --- a/modules/lora/lora_overrides.py +++ b/modules/lora/lora_overrides.py @@ -27,6 +27,7 @@ allow_native = [ 'sd3', 'f1', 'chroma', + 'zimage', ] diff --git a/pipelines/z_image/zimage_lora.py b/pipelines/z_image/zimage_lora.py new file mode 100644 index 000000000..c84be3c3b --- /dev/null +++ b/pipelines/z_image/zimage_lora.py @@ -0,0 +1,212 @@ +"""Z-Image native LoRA loader. + +Bypasses the diffusers PEFT pipeline entirely. Reads the safetensors directly, +maps each target module into sdnext's existing ``network_layer_mapping``, and +returns a ``Network`` populated with ``NetworkModuleLora`` entries that +``network_activate`` will apply. + +Recognized LoRA key formats: + +- AI Toolkit diffusers: ``diffusion_model.layers.0.attention.to_q.lora_A.weight`` +- AI Toolkit down/up: ``diffusion_model.layers.0.attention.to_q.lora_down.weight`` +- Kohya (lora_unet): ``lora_unet_layers_0_attention_to_q.lora_down.weight`` +- Bare transformer: ``transformer.layers.0.attention.to_q.lora_down.weight`` +- Bare (no prefix): ``layers.0.attention.to_q.lora_A.weight`` + +Alpha and dora_scale keys are preserved and passed to ``NetworkModuleLora`` +which applies the alpha/rank scale at weight-apply time. +""" + +import os +import time +import torch +from modules import shared, sd_models +from modules.logger import log +from modules.lora import network, network_lora, lora_convert +from modules.lora import lora_common as l + + +KNOWN_PREFIXES = ("diffusion_model.", "transformer.", "lora_unet_") + +WEIGHT_MARKERS = ( + ".lora_down.weight", ".lora_up.weight", + ".lora_A.weight", ".lora_B.weight", +) +META_MARKERS = (".alpha", ".dora_scale") + +SUFFIX_NORMALIZE = { + "lora_A.weight": "lora_down.weight", + "lora_B.weight": "lora_up.weight", +} + +# Pre-refactor Z-Image module names map to current diffusers Attention submodule +# names. LoRAs trained against the original Lumina/Z-Image attention layout use +# fused `qkv` and bare `out` / `wo`; diffusers now exposes `to_q`/`to_k`/`to_v` +# and `to_out.0`. The fused qkv up-weight is chunked into three along dim 0. +ATTENTION_OUT_ALIASES = ("attention_out", "attention_out_0", "attention_wo") +ATTENTION_OUT_TARGET = "attention_to_out_0" +ATTENTION_QKV_SUFFIX = "attention_qkv" +ATTENTION_QKV_TARGETS = ("attention_to_q", "attention_to_k", "attention_to_v") + + +def try_load_lora(name, network_on_disk, lora_scale): + """Try loading a Z-Image LoRA as native modules. + + Returns a ``Network`` if at least one module matched, otherwise ``None``. + """ + t0 = time.time() + state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network') + + if not any(any(m in k for m in WEIGHT_MARKERS) for k in state_dict): + return None + + sd_model = getattr(shared.sd_model, "pipe", shared.sd_model) + lora_convert.assign_network_names_to_compvis_modules(sd_model) + mapping = getattr(shared.sd_model, 'network_layer_mapping', {}) or {} + + net = network.Network(name, network_on_disk) + net.mtime = os.path.getmtime(network_on_disk.filename) + + groups = group_state_dict(state_dict) + groups = expand_legacy_attention(groups) + unmapped = 0 + shape_mismatch = 0 + for network_key, w in groups.items(): + if 'lora_down.weight' not in w or 'lora_up.weight' not in w: + continue + sd_module = mapping.get(network_key) + if sd_module is None: + unmapped += 1 + continue + if not shapes_match(sd_module, w['lora_down.weight'], w['lora_up.weight']): + log.warning(f'Network load: type=LoRA name="{name}" key={network_key} shape mismatch') + shape_mismatch += 1 + continue + nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module) + net.modules[network_key] = network_lora.NetworkModuleLora(net, nw) + + if len(net.modules) == 0: + if unmapped or shape_mismatch: + log.debug(f'Network load: type=LoRA name="{name}" native no-match unmapped={unmapped} mismatch={shape_mismatch}') + return None + + log.debug( + f'Network load: type=LoRA name="{name}" native modules={len(net.modules)}' + f' unmapped={unmapped} mismatch={shape_mismatch} scale={lora_scale}' + ) + l.timer.activate += time.time() - t0 + return net + + +def shapes_match(sd_module, down_w: torch.Tensor, up_w: torch.Tensor) -> bool: + if not hasattr(sd_module, 'weight'): + return False + if hasattr(sd_module, 'sdnq_dequantizer'): + mod_shape = sd_module.sdnq_dequantizer.original_shape + else: + mod_shape = sd_module.weight.shape + if len(mod_shape) < 2 or len(down_w.shape) < 2 or len(up_w.shape) < 2: + return False + return down_w.shape[1] == mod_shape[1] and up_w.shape[0] == mod_shape[0] + + +def group_state_dict(state_dict): + """Group state_dict entries by target module. + + Returns ``{network_key: {suffix: tensor, ...}}`` with suffixes normalized to + ``lora_down.weight`` / ``lora_up.weight`` (plus optional ``alpha`` / ``dora_scale``). + ``network_key`` matches the sdnext ``network_layer_name`` convention + (``lora_transformer_``). + """ + groups: dict[str, dict[str, torch.Tensor]] = {} + for key, value in state_dict.items(): + parsed = parse_key(key) + if parsed is None: + continue + network_key, suffix = parsed + slot = groups.get(network_key) + if slot is None: + slot = {} + groups[network_key] = slot + slot[suffix] = value + return groups + + +def expand_legacy_attention(groups): + """Rewrite pre-refactor Z-Image attention module names. + + - ``..._attention_qkv`` groups are split into three separate + ``..._attention_to_q/k/v`` groups. The down weight is shared and the up + weight is chunked along dim 0 (which carries the concatenated q/k/v + outputs in the fused layout). + - ``..._attention_out``, ``..._attention_out_0`` and ``..._attention_wo`` + are renamed to ``..._attention_to_out_0``. + """ + out: dict[str, dict[str, torch.Tensor]] = {} + for key, w in groups.items(): + new_key = None + for alias in ATTENTION_OUT_ALIASES: + if key.endswith("_" + alias): + new_key = key[: -len(alias)] + ATTENTION_OUT_TARGET + break + if new_key is not None: + out[new_key] = w + continue + + if key.endswith("_" + ATTENTION_QKV_SUFFIX): + stem = key[: -len(ATTENTION_QKV_SUFFIX)] + down = w.get("lora_down.weight") + up = w.get("lora_up.weight") + alpha = w.get("alpha") + dora = w.get("dora_scale") + if down is None or up is None or up.shape[0] % 3 != 0: + out[key] = w + continue + chunks = torch.chunk(up, 3, dim=0) + for target, chunk in zip(ATTENTION_QKV_TARGETS, chunks): + split_key = stem + target + split = { + "lora_down.weight": down, + "lora_up.weight": chunk.contiguous(), + } + if alpha is not None: + split["alpha"] = alpha + if dora is not None: + split["dora_scale"] = dora + out[split_key] = split + continue + + out[key] = w + return out + + +def parse_key(key: str): + """Parse one state_dict key into ``(network_key, suffix)`` or ``None``. + + The suffix is one of: ``lora_down.weight``, ``lora_up.weight``, ``alpha``, + ``dora_scale``. Diffusers ``lora_A/lora_B`` aliases are normalized to + ``lora_down/lora_up`` here. + """ + stripped = key + for p in KNOWN_PREFIXES: + if key.startswith(p): + stripped = key[len(p):] + break + + split_at = -1 + raw_suffix = None + for marker in WEIGHT_MARKERS + META_MARKERS: + if stripped.endswith(marker): + split_at = len(stripped) - len(marker) + raw_suffix = marker.lstrip('.') + break + if split_at < 0: + return None + + base = stripped[:split_at] + if not base: + return None + + suffix = SUFFIX_NORMALIZE.get(raw_suffix, raw_suffix) + network_key = 'lora_transformer_' + base.replace('.', '_') + return network_key, suffix From 65e59e72249f6a568004db209077879577b43892 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Tue, 14 Apr 2026 20:37:52 +0100 Subject: [PATCH 2/2] feat(lora): z-image native lokr, loha and oft loaders 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. --- modules/lora/lora_load.py | 10 +- pipelines/z_image/zimage_lora.py | 373 ++++++++++++++++++++++--------- 2 files changed, 279 insertions(+), 104 deletions(-) diff --git a/modules/lora/lora_load.py b/modules/lora/lora_load.py index 9790b5cc7..d2be416cb 100644 --- a/modules/lora/lora_load.py +++ b/modules/lora/lora_load.py @@ -52,7 +52,15 @@ def load_safetensors(name, network_on_disk: network.NetworkOnDisk) -> network.Ne if shared.sd_model_type == 'zimage': from pipelines.z_image import zimage_lora lora_scale = shared.opts.extra_networks_default_multiplier - zimage_net = zimage_lora.try_load_lora(name, network_on_disk, lora_scale) + zimage_net = None + for try_fn in (zimage_lora.try_load_lora, zimage_lora.try_load_lokr, zimage_lora.try_load_loha, zimage_lora.try_load_oft): + sub = try_fn(name, network_on_disk, lora_scale) + if sub is None: + continue + if zimage_net is None: + zimage_net = sub + else: + zimage_net.modules.update(sub.modules) if zimage_net is not None: lora_cache[name] = zimage_net return zimage_net diff --git a/pipelines/z_image/zimage_lora.py b/pipelines/z_image/zimage_lora.py index c84be3c3b..2a138859e 100644 --- a/pipelines/z_image/zimage_lora.py +++ b/pipelines/z_image/zimage_lora.py @@ -1,20 +1,33 @@ -"""Z-Image native LoRA loader. +"""Z-Image native adapter loader. -Bypasses the diffusers PEFT pipeline entirely. Reads the safetensors directly, -maps each target module into sdnext's existing ``network_layer_mapping``, and -returns a ``Network`` populated with ``NetworkModuleLora`` entries that -``network_activate`` will apply. +Runs when :func:`modules.lora.lora_overrides.get_method` returns ``'native'`` +(``lora_force_diffusers`` off and ``zimage`` in ``allow_native``). Reads the +safetensors directly and writes into sdnext's existing +``network_layer_mapping``, returning a ``Network`` populated with +``NetworkModule*`` entries that ``network_activate`` will apply. If the +setting is on, the diffusers PEFT path handles the file instead. -Recognized LoRA key formats: +Entry points, one per family: -- AI Toolkit diffusers: ``diffusion_model.layers.0.attention.to_q.lora_A.weight`` -- AI Toolkit down/up: ``diffusion_model.layers.0.attention.to_q.lora_down.weight`` -- Kohya (lora_unet): ``lora_unet_layers_0_attention_to_q.lora_down.weight`` -- Bare transformer: ``transformer.layers.0.attention.to_q.lora_down.weight`` -- Bare (no prefix): ``layers.0.attention.to_q.lora_A.weight`` +- LoRA (+ DoRA) via :func:`try_load_lora` +- LoKR via :func:`try_load_lokr` +- LoHA via :func:`try_load_loha` (not yet validated against a real Z-Image adapter) +- OFT via :func:`try_load_oft` (not yet validated against a real Z-Image adapter) -Alpha and dora_scale keys are preserved and passed to ``NetworkModuleLora`` -which applies the alpha/rank scale at weight-apply time. +Recognized key prefixes for every family: ``diffusion_model.``, +``transformer.``, ``lora_unet_``, or bare. Diffusers-PEFT ``lora_A``/``lora_B`` +are normalized to ``lora_down``/``lora_up``. + +Pre-refactor Z-Image attention layouts (fused ``attention.qkv``, bare +``attention.out`` / ``attention.wo``) are rewritten to the current diffusers +``to_q``/``to_k``/``to_v`` and ``to_out.0``. For LoRA the fused qkv up-weight +is chunked along dim 0 at load time. For LoKR the split is deferred to apply +time via :class:`NetworkModuleLokrChunk`, which materializes ``kron(w1, w2)`` +once per forward pass and returns the designated slice. + +Fused ``attention.qkv`` for LoHA and OFT is skipped with a warning: no +``NetworkModuleHadaChunk`` exists, and an OFT block structure is tied to the +target module's ``out_features``, so a Q/K/V split is not a drop-in. """ import os @@ -22,27 +35,48 @@ import time import torch from modules import shared, sd_models from modules.logger import log -from modules.lora import network, network_lora, lora_convert +from modules.lora import network, network_lora, network_lokr, network_hada, network_oft, lora_convert from modules.lora import lora_common as l KNOWN_PREFIXES = ("diffusion_model.", "transformer.", "lora_unet_") -WEIGHT_MARKERS = ( +# Every family also picks up the universal optional keys +# (alpha, scale, bias, dora_scale) via base NetworkModule.__init__. +LORA_SUFFIXES = ( ".lora_down.weight", ".lora_up.weight", ".lora_A.weight", ".lora_B.weight", + ".alpha", ".dora_scale", ".bias", ".scale", ) -META_MARKERS = (".alpha", ".dora_scale") +LOKR_SUFFIXES = ( + ".lokr_w1", ".lokr_w2", + ".lokr_w1_a", ".lokr_w1_b", + ".lokr_w2_a", ".lokr_w2_b", + ".lokr_t2", + ".alpha", ".dora_scale", ".bias", ".scale", +) +LOHA_SUFFIXES = ( + ".hada_w1_a", ".hada_w1_b", + ".hada_w2_a", ".hada_w2_b", + ".hada_t1", ".hada_t2", + ".alpha", ".dora_scale", ".bias", ".scale", +) +OFT_SUFFIXES = ( + ".oft_blocks", ".oft_diag", + ".alpha", ".dora_scale", ".bias", ".scale", +) + +# Presence of any of these substrings anywhere in a key marks a file as belonging to that family. +LORA_MARKERS = (".lora_down.weight", ".lora_up.weight", ".lora_A.weight", ".lora_B.weight") +LOKR_MARKERS = (".lokr_w1", ".lokr_w2") +LOHA_MARKERS = (".hada_w1_a", ".hada_w1_b", ".hada_w2_a", ".hada_w2_b") +OFT_MARKERS = (".oft_blocks", ".oft_diag") SUFFIX_NORMALIZE = { "lora_A.weight": "lora_down.weight", "lora_B.weight": "lora_up.weight", } -# Pre-refactor Z-Image module names map to current diffusers Attention submodule -# names. LoRAs trained against the original Lumina/Z-Image attention layout use -# fused `qkv` and bare `out` / `wo`; diffusers now exposes `to_q`/`to_k`/`to_v` -# and `to_out.0`. The fused qkv up-weight is chunked into three along dim 0. ATTENTION_OUT_ALIASES = ("attention_out", "attention_out_0", "attention_wo") ATTENTION_OUT_TARGET = "attention_to_out_0" ATTENTION_QKV_SUFFIX = "attention_qkv" @@ -50,25 +84,18 @@ ATTENTION_QKV_TARGETS = ("attention_to_q", "attention_to_k", "attention_to_v") def try_load_lora(name, network_on_disk, lora_scale): - """Try loading a Z-Image LoRA as native modules. - - Returns a ``Network`` if at least one module matched, otherwise ``None``. - """ + """Try loading a Z-Image LoRA (plus DoRA) as native modules.""" t0 = time.time() state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network') - - if not any(any(m in k for m in WEIGHT_MARKERS) for k in state_dict): + if not has_marker(state_dict, LORA_MARKERS): return None - sd_model = getattr(shared.sd_model, "pipe", shared.sd_model) - lora_convert.assign_network_names_to_compvis_modules(sd_model) - mapping = getattr(shared.sd_model, 'network_layer_mapping', {}) or {} + mapping = resolve_mapping() + net = new_network(name, network_on_disk) - net = network.Network(name, network_on_disk) - net.mtime = os.path.getmtime(network_on_disk.filename) + groups = group_by_suffixes(state_dict, LORA_SUFFIXES) + groups = expand_legacy_attention_lora(groups) - groups = group_state_dict(state_dict) - groups = expand_legacy_attention(groups) unmapped = 0 shape_mismatch = 0 for network_key, w in groups.items(): @@ -85,14 +112,134 @@ def try_load_lora(name, network_on_disk, lora_scale): nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module) net.modules[network_key] = network_lora.NetworkModuleLora(net, nw) - if len(net.modules) == 0: - if unmapped or shape_mismatch: - log.debug(f'Network load: type=LoRA name="{name}" native no-match unmapped={unmapped} mismatch={shape_mismatch}') + return finalize_network(net, name, 'LoRA', lora_scale, t0, unmapped=unmapped, mismatch=shape_mismatch) + + +def try_load_lokr(name, network_on_disk, lora_scale): + """Try loading a Z-Image LoKR as native modules.""" + t0 = time.time() + state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network') + if not has_marker(state_dict, LOKR_MARKERS): return None + mapping = resolve_mapping() + net = new_network(name, network_on_disk) + + groups = group_by_suffixes(state_dict, LOKR_SUFFIXES) + groups, chunk_info = expand_legacy_attention_lokr(groups) + + unmapped = 0 + for network_key, w in groups.items(): + has_1 = "lokr_w1" in w or ("lokr_w1_a" in w and "lokr_w1_b" in w) + has_2 = "lokr_w2" in w or ("lokr_w2_a" in w and "lokr_w2_b" in w) + if not (has_1 and has_2): + continue + sd_module = mapping.get(network_key) + if sd_module is None: + unmapped += 1 + continue + nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module) + chunked = chunk_info.get(network_key) + if chunked is not None: + idx, num = chunked + net.modules[network_key] = network_lokr.NetworkModuleLokrChunk(net, nw, idx, num) + else: + net.modules[network_key] = network_lokr.NetworkModuleLokr(net, nw) + + return finalize_network(net, name, 'LoKR', lora_scale, t0, unmapped=unmapped) + + +def try_load_loha(name, network_on_disk, lora_scale): + """Try loading a Z-Image LoHA as native modules. Fused attention.qkv groups are skipped.""" + t0 = time.time() + state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network') + if not has_marker(state_dict, LOHA_MARKERS): + return None + + mapping = resolve_mapping() + net = new_network(name, network_on_disk) + + groups = group_by_suffixes(state_dict, LOHA_SUFFIXES) + groups = rename_attention_out(groups) + + unmapped = 0 + skipped_qkv = 0 + for network_key, w in groups.items(): + if network_key.endswith("_" + ATTENTION_QKV_SUFFIX): + log.warning(f'Network load: type=LoHA name="{name}" key={network_key} fused qkv skipped (unsupported)') + skipped_qkv += 1 + continue + if not all(k in w for k in ("hada_w1_a", "hada_w1_b", "hada_w2_a", "hada_w2_b")): + continue + sd_module = mapping.get(network_key) + if sd_module is None: + unmapped += 1 + continue + nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module) + net.modules[network_key] = network_hada.NetworkModuleHada(net, nw) + + return finalize_network(net, name, 'LoHA', lora_scale, t0, unmapped=unmapped, skipped=skipped_qkv) + + +def try_load_oft(name, network_on_disk, lora_scale): + """Try loading a Z-Image OFT adapter as native modules. Fused attention.qkv groups are skipped.""" + t0 = time.time() + state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network') + if not has_marker(state_dict, OFT_MARKERS): + return None + + mapping = resolve_mapping() + net = new_network(name, network_on_disk) + + groups = group_by_suffixes(state_dict, OFT_SUFFIXES) + groups = rename_attention_out(groups) + + unmapped = 0 + skipped_qkv = 0 + for network_key, w in groups.items(): + if network_key.endswith("_" + ATTENTION_QKV_SUFFIX): + log.warning(f'Network load: type=OFT name="{name}" key={network_key} fused qkv skipped (unsupported)') + skipped_qkv += 1 + continue + if not ("oft_blocks" in w or "oft_diag" in w): + continue + sd_module = mapping.get(network_key) + if sd_module is None: + unmapped += 1 + continue + nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module) + net.modules[network_key] = network_oft.NetworkModuleOFT(net, nw) + + return finalize_network(net, name, 'OFT', lora_scale, t0, unmapped=unmapped, skipped=skipped_qkv) + + +def has_marker(state_dict, markers): + return any(any(m in k for m in markers) for k in state_dict) + + +def resolve_mapping(): + sd_model = getattr(shared.sd_model, "pipe", shared.sd_model) + lora_convert.assign_network_names_to_compvis_modules(sd_model) + return getattr(shared.sd_model, 'network_layer_mapping', {}) or {} + + +def new_network(name, network_on_disk): + net = network.Network(name, network_on_disk) + net.mtime = os.path.getmtime(network_on_disk.filename) + return net + + +def finalize_network(net, name, family, lora_scale, t0, unmapped=0, mismatch=0, skipped=0): + if len(net.modules) == 0: + if unmapped or mismatch or skipped: + log.debug( + f'Network load: type={family} name="{name}" native no-match' + f' unmapped={unmapped} mismatch={mismatch} skipped={skipped}' + ) + return None log.debug( - f'Network load: type=LoRA name="{name}" native modules={len(net.modules)}' - f' unmapped={unmapped} mismatch={shape_mismatch} scale={lora_scale}' + f'Network load: type={family} name="{name}" native modules={len(net.modules)}' + f' unmapped={unmapped} mismatch={mismatch} skipped={skipped} scale={lora_scale}' ) l.timer.activate += time.time() - t0 return net @@ -110,17 +257,17 @@ def shapes_match(sd_module, down_w: torch.Tensor, up_w: torch.Tensor) -> bool: return down_w.shape[1] == mod_shape[1] and up_w.shape[0] == mod_shape[0] -def group_state_dict(state_dict): +def group_by_suffixes(state_dict, suffixes): """Group state_dict entries by target module. - Returns ``{network_key: {suffix: tensor, ...}}`` with suffixes normalized to - ``lora_down.weight`` / ``lora_up.weight`` (plus optional ``alpha`` / ``dora_scale``). - ``network_key`` matches the sdnext ``network_layer_name`` convention - (``lora_transformer_``). + Returns ``{network_key: {suffix: tensor, ...}}`` where ``network_key`` follows + the sdnext convention ``lora_transformer_``. Only keys + whose suffix appears in ``suffixes`` are kept; ``lora_A``/``lora_B`` are + normalized to ``lora_down``/``lora_up``. """ groups: dict[str, dict[str, torch.Tensor]] = {} for key, value in state_dict.items(): - parsed = parse_key(key) + parsed = parse_key(key, suffixes) if parsed is None: continue network_key, suffix = parsed @@ -132,73 +279,19 @@ def group_state_dict(state_dict): return groups -def expand_legacy_attention(groups): - """Rewrite pre-refactor Z-Image attention module names. - - - ``..._attention_qkv`` groups are split into three separate - ``..._attention_to_q/k/v`` groups. The down weight is shared and the up - weight is chunked along dim 0 (which carries the concatenated q/k/v - outputs in the fused layout). - - ``..._attention_out``, ``..._attention_out_0`` and ``..._attention_wo`` - are renamed to ``..._attention_to_out_0``. - """ - out: dict[str, dict[str, torch.Tensor]] = {} - for key, w in groups.items(): - new_key = None - for alias in ATTENTION_OUT_ALIASES: - if key.endswith("_" + alias): - new_key = key[: -len(alias)] + ATTENTION_OUT_TARGET - break - if new_key is not None: - out[new_key] = w - continue - - if key.endswith("_" + ATTENTION_QKV_SUFFIX): - stem = key[: -len(ATTENTION_QKV_SUFFIX)] - down = w.get("lora_down.weight") - up = w.get("lora_up.weight") - alpha = w.get("alpha") - dora = w.get("dora_scale") - if down is None or up is None or up.shape[0] % 3 != 0: - out[key] = w - continue - chunks = torch.chunk(up, 3, dim=0) - for target, chunk in zip(ATTENTION_QKV_TARGETS, chunks): - split_key = stem + target - split = { - "lora_down.weight": down, - "lora_up.weight": chunk.contiguous(), - } - if alpha is not None: - split["alpha"] = alpha - if dora is not None: - split["dora_scale"] = dora - out[split_key] = split - continue - - out[key] = w - return out - - -def parse_key(key: str): - """Parse one state_dict key into ``(network_key, suffix)`` or ``None``. - - The suffix is one of: ``lora_down.weight``, ``lora_up.weight``, ``alpha``, - ``dora_scale``. Diffusers ``lora_A/lora_B`` aliases are normalized to - ``lora_down/lora_up`` here. - """ +def parse_key(key, suffixes): stripped = key for p in KNOWN_PREFIXES: if key.startswith(p): stripped = key[len(p):] break + matched_suffix = None split_at = -1 - raw_suffix = None - for marker in WEIGHT_MARKERS + META_MARKERS: + for marker in suffixes: if stripped.endswith(marker): split_at = len(stripped) - len(marker) - raw_suffix = marker.lstrip('.') + matched_suffix = marker.lstrip('.') break if split_at < 0: return None @@ -207,6 +300,80 @@ def parse_key(key: str): if not base: return None - suffix = SUFFIX_NORMALIZE.get(raw_suffix, raw_suffix) + suffix = SUFFIX_NORMALIZE.get(matched_suffix, matched_suffix) network_key = 'lora_transformer_' + base.replace('.', '_') return network_key, suffix + + +def rename_attention_out(groups): + """Rename legacy ``attention.out`` / ``attention.wo`` keys to ``attention.to_out.0``.""" + out: dict[str, dict[str, torch.Tensor]] = {} + for key, w in groups.items(): + new_key = None + for alias in ATTENTION_OUT_ALIASES: + if key.endswith("_" + alias): + new_key = key[: -len(alias)] + ATTENTION_OUT_TARGET + break + out[new_key or key] = w + return out + + +def expand_legacy_attention_lora(groups): + """Rename attention.out, then split fused attention.qkv into three per-projection groups. + + The down-weight is shared across Q/K/V and the up-weight is chunked along + dim 0 (the concatenated output dim in the fused layout). + """ + groups = rename_attention_out(groups) + out: dict[str, dict[str, torch.Tensor]] = {} + for key, w in groups.items(): + if not key.endswith("_" + ATTENTION_QKV_SUFFIX): + out[key] = w + continue + stem = key[: -len(ATTENTION_QKV_SUFFIX)] + down = w.get("lora_down.weight") + up = w.get("lora_up.weight") + if down is None or up is None or up.shape[0] % 3 != 0: + out[key] = w + continue + chunks = torch.chunk(up, 3, dim=0) + alpha = w.get("alpha") + dora = w.get("dora_scale") + for target, chunk in zip(ATTENTION_QKV_TARGETS, chunks): + split_key = stem + target + split = { + "lora_down.weight": down, + "lora_up.weight": chunk.contiguous(), + } + if alpha is not None: + split["alpha"] = alpha + if dora is not None: + split["dora_scale"] = dora + out[split_key] = split + return out + + +def expand_legacy_attention_lokr(groups): + """Rename attention.out, then split fused attention.qkv for LoKR. + + For LoKR the split is deferred to apply time via ``NetworkModuleLokrChunk`` + (returned via a parallel ``chunk_info`` dict keyed by the split network key). + The three resulting groups share the same tensor dict by shallow copy; the + chunk module computes ``kron(w1, w2)`` once and slices the designated row + range. + + Returns ``(groups, chunk_info)`` where ``chunk_info[network_key] == (index, num)``. + """ + groups = rename_attention_out(groups) + out: dict[str, dict[str, torch.Tensor]] = {} + chunk_info: dict[str, tuple[int, int]] = {} + for key, w in groups.items(): + if not key.endswith("_" + ATTENTION_QKV_SUFFIX): + out[key] = w + continue + stem = key[: -len(ATTENTION_QKV_SUFFIX)] + for i, target in enumerate(ATTENTION_QKV_TARGETS): + split_key = stem + target + out[split_key] = dict(w) + chunk_info[split_key] = (i, 3) + return out, chunk_info