From 18bbe288b084b411a779d49de3aa575bae40cb14 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Tue, 14 Jul 2026 03:51:14 +0100 Subject: [PATCH] fix(lora): apply full-diff norm targets and map z-image's renamed modules 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. --- modules/lora/native_adapter.py | 6 ++ pipelines/z_image/zimage_lora.py | 57 +++++++++++++++++ test/test-zimage-native-adapters.py | 96 ++++++++++++++++++++++++++++- 3 files changed, 158 insertions(+), 1 deletion(-) diff --git a/modules/lora/native_adapter.py b/modules/lora/native_adapter.py index d6dc8537e..0ab038c73 100644 --- a/modules/lora/native_adapter.py +++ b/modules/lora/native_adapter.py @@ -1039,6 +1039,12 @@ def try_load_full(name, network_on_disk, lora_scale, *, if sd_module is None: unmapped += 1 continue + # Loader-local stamping, same as try_load_norm: a full-weight extraction carries the + # norm weights too, and assign_network_names_to_compvis_modules puts norms in the + # mapping but never stamps network_layer_name on them, which is what the apply pass + # keys off. Without this they bind and then silently never apply. + if not getattr(sd_module, "network_layer_name", None): + sd_module.network_layer_name = network_key nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module) net.modules[network_key] = network_full.NetworkModuleFull(net, nw) diff --git a/pipelines/z_image/zimage_lora.py b/pipelines/z_image/zimage_lora.py index 2c29b5552..e8c6d7ec4 100644 --- a/pipelines/z_image/zimage_lora.py +++ b/pipelines/z_image/zimage_lora.py @@ -12,6 +12,14 @@ Recognized key prefixes: ``diffusion_model.``, ``transformer.``, ``lora_unet_``, or bare paths starting with the known block-level prefixes (``layers.``, ``noise_refiner.``, ``context_refiner.``). +Checkpoint names that differ from the diffusers module tree are rewritten by +:func:`resolve_targets`: qk-norms (``attention.q_norm`` / ``k_norm`` -> +``attention.norm_q`` / ``norm_k``) and the non-block targets in +``ZIMAGE_EXTRA_MAP`` (``x_embedder`` and ``final_layer.*`` live in ModuleDicts +keyed by ``"{patch_size}-{f_patch_size}"``, read from the live model by +:func:`patch_keys`). ``t_embedder.mlp.N`` and ``cap_embedder.N`` already match +and pass through. + Pre-refactor Z-Image attention layouts (fused ``attention.qkv``, bare ``attention.out`` / ``attention.wo``) are rewritten by :func:`resolve_targets` to the current diffusers ``to_q``/``to_k``/``to_v`` and ``to_out.0``. For @@ -34,6 +42,42 @@ KNOWN_PREFIXES = native_adapter.KNOWN_PREFIXES_DEFAULT BARE_DIFFUSERS_PREFIXES = ("layers.", "noise_refiner.", "context_refiner.") +# Checkpoint qk-norm names vs the diffusers attention module names. +ZIMAGE_NORM_ALIASES = { + ".attention.q_norm": ".attention.norm_q", + ".attention.k_norm": ".attention.norm_k", +} + +# Non-block targets. The patch embedder and the final layer live in ModuleDicts keyed by +# "{patch_size}-{f_patch_size}", so the diffusers path carries a key the checkpoint does not: +# {key} is filled from the live model. t_embedder.mlp.N and cap_embedder.N already match and +# pass through verbatim. +ZIMAGE_EXTRA_MAP = { + "x_embedder": "all_x_embedder.{key}", + "final_layer.linear": "all_final_layer.{key}.linear", + "final_layer.adaLN_modulation.1": "all_final_layer.{key}.adaLN_modulation.1", +} + +ZIMAGE_EXTRA_KOHYA_MAP = {k.replace(".", "_"): v for k, v in ZIMAGE_EXTRA_MAP.items()} + +# Both shipped Z-Image variants (Base, Turbo) build a single entry; read from the live model and +# fall back to it. Tests patch this directly. +PATCH_KEYS_DEFAULT = ["2-1"] + + +def patch_keys(): + """ModuleDict keys of the live transformer's ``all_x_embedder``.""" + try: + from modules import shared + pipe = getattr(shared.sd_model, "pipe", shared.sd_model) + embedder = getattr(getattr(pipe, "transformer", None), "all_x_embedder", None) + keys = list(embedder.keys()) if embedder is not None else [] + if keys: + return keys + except Exception: + pass + return PATCH_KEYS_DEFAULT + # === Re-exports for test/back-compat === @@ -106,6 +150,9 @@ def resolve_targets(prefix_used, base): def _dotted_to_diffusers_targets(base): """For BFL / bare-BFL keys like ``layers.0.attention.qkv``.""" + extra = ZIMAGE_EXTRA_MAP.get(base) + if extra is not None: + return [(extra.format(key=key), None) for key in patch_keys()] if base.endswith(".attention.qkv"): stem = base[:-len(".attention.qkv")] return [ @@ -113,6 +160,9 @@ def _dotted_to_diffusers_targets(base): (f"{stem}.attention.to_k", ChunkSpec(idx=1, total=3)), (f"{stem}.attention.to_v", ChunkSpec(idx=2, total=3)), ] + for alias, renamed in ZIMAGE_NORM_ALIASES.items(): + if base.endswith(alias): + return [(base[:-len(alias)] + renamed, None)] for alias in (".attention.out.0", ".attention.out", ".attention.wo"): if base.endswith(alias): stem = base[:-len(alias)] @@ -122,6 +172,9 @@ def _dotted_to_diffusers_targets(base): def _underscore_to_diffusers_targets(base): """For kohya flat-underscore keys like ``layers_0_attention_qkv``.""" + extra = ZIMAGE_EXTRA_KOHYA_MAP.get(base) + if extra is not None: + return [(extra.format(key=key), None) for key in patch_keys()] if base.endswith("_attention_qkv"): stem = base[:-len("_attention_qkv")] return [ @@ -129,6 +182,10 @@ def _underscore_to_diffusers_targets(base): (f"{stem}_attention_to_k", ChunkSpec(idx=1, total=3)), (f"{stem}_attention_to_v", ChunkSpec(idx=2, total=3)), ] + for alias, renamed in ZIMAGE_NORM_ALIASES.items(): + underscored = alias.replace(".", "_") + if base.endswith(underscored): + return [(base[:-len(underscored)] + renamed.replace(".", "_"), None)] for alias in ("_attention_out_0", "_attention_out", "_attention_wo"): if base.endswith(alias): stem = base[:-len(alias)] diff --git a/test/test-zimage-native-adapters.py b/test/test-zimage-native-adapters.py index 2a8df2eb2..77004d18e 100644 --- a/test/test-zimage-native-adapters.py +++ b/test/test-zimage-native-adapters.py @@ -125,6 +125,7 @@ MLP_HIDDEN = int(HIDDEN / 3 * 8) # 256, matches ZImageTransformerBlock FeedFor ADALN_OUT = 4 * HIDDEN # 384, matches Z-Image adaLN_modulation N_LAYERS = 2 # main transformer blocks N_REFINER = 1 # noise_refiner / context_refiner blocks +PATCH_KEY = "2-1" # the single all_x_embedder / all_final_layer key both variants build # pylint: disable=attribute-defined-outside-init @@ -184,6 +185,26 @@ def build_mock_transformer(): transformer.layers = torch.nn.ModuleList([build_zimage_block(modulation=False) for _ in range(N_LAYERS)]) transformer.noise_refiner = torch.nn.ModuleList([build_zimage_block(modulation=True) for _ in range(N_REFINER)]) transformer.context_refiner = torch.nn.ModuleList([build_zimage_block(modulation=True) for _ in range(N_REFINER)]) + # Non-block targets. all_x_embedder / all_final_layer are ModuleDicts keyed by + # "{patch_size}-{f_patch_size}"; both shipped variants build the single key below. + final_layer = _Holder() + final_layer.linear = torch.nn.Linear(HIDDEN, HIDDEN, bias=True) + final_layer.adaLN_modulation = torch.nn.Sequential( + torch.nn.SiLU(), + torch.nn.Linear(HIDDEN, ADALN_OUT, bias=True), + ) + transformer.all_x_embedder = torch.nn.ModuleDict({PATCH_KEY: torch.nn.Linear(HIDDEN, HIDDEN, bias=True)}) + transformer.all_final_layer = torch.nn.ModuleDict({PATCH_KEY: final_layer}) + transformer.t_embedder = _Holder() + transformer.t_embedder.mlp = torch.nn.Sequential( + torch.nn.Linear(HIDDEN, HIDDEN, bias=True), + torch.nn.SiLU(), + torch.nn.Linear(HIDDEN, HIDDEN, bias=True), + ) + transformer.cap_embedder = torch.nn.Sequential( + torch.nn.RMSNorm(HIDDEN), + torch.nn.Linear(HIDDEN, HIDDEN, bias=True), + ) return transformer @@ -653,6 +674,76 @@ def test_full_diff_chain(): return True +def test_full_stamps_norm_targets(): + """A Full adapter on a norm target binds AND applies. + + assign_network_names_to_compvis_modules puts norms in the mapping but never + stamps network_layer_name on them, and the apply pass keys off that + attribute: without loader-local stamping the module binds and then silently + never applies. Regression for the RedLSP extraction, which bound 308 + modules and applied only 172. + """ + sd = { + 'diffusion_model.layers.0.attention_norm1.diff': torch.randn(HIDDEN), + 'diffusion_model.layers.0.attention.q_norm.diff': torch.randn(HEAD_DIM), + } + net = _load_via(Z.try_load, sd) + assert net is not None and len(net.modules) == 2, f'got {sorted(net.modules) if net else None}' + expected = {'lora_transformer_layers_0_attention_norm1', 'lora_transformer_layers_0_attention_norm_q'} + assert set(net.modules) == expected, f'got {set(net.modules)}' + for key, mod in net.modules.items(): + assert getattr(mod.sd_module, 'network_layer_name', None) == key, \ + f'{key}: host module not stamped, the apply pass will skip it' + return True + + +def test_resolve_targets_norm_aliases_and_extra(): + """qk-norm renames and the ModuleDict-keyed non-block targets.""" + cases = [ + (('diffusion_model.', 'layers.0.attention.q_norm'), 'layers.0.attention.norm_q'), + (('diffusion_model.', 'layers.3.attention.k_norm'), 'layers.3.attention.norm_k'), + (('lora_unet_', 'layers_0_attention_q_norm'), 'layers_0_attention_norm_q'), + ((None, 'noise_refiner.0.attention.k_norm'), 'noise_refiner.0.attention.norm_k'), + ] + for (prefix, base), expected in cases: + targets = Z.resolve_targets(prefix, base) + assert targets == [(expected, None)], f'({prefix}, {base}) -> {targets}' + # the patch key comes from the live model's all_x_embedder ModuleDict + install_mock_pipe() + for base, expected in [ + ('x_embedder', f'all_x_embedder.{PATCH_KEY}'), + ('final_layer.linear', f'all_final_layer.{PATCH_KEY}.linear'), + ('final_layer.adaLN_modulation.1', f'all_final_layer.{PATCH_KEY}.adaLN_modulation.1'), + ]: + targets = Z.resolve_targets('diffusion_model.', base) + assert targets == [(expected, None)], f'{base} -> {targets}' + kohya = Z.resolve_targets('lora_unet_', base.replace('.', '_')) + assert kohya == [(expected, None)], f'kohya {base} -> {kohya}' + return True + + +def test_full_extra_modules_bind(): + """x_embedder and final_layer bind through the ModuleDict paths.""" + sd = { + 'diffusion_model.x_embedder.diff': torch.randn(HIDDEN, HIDDEN), + 'diffusion_model.final_layer.linear.diff': torch.randn(HIDDEN, HIDDEN), + 'diffusion_model.final_layer.adaLN_modulation.1.diff': torch.randn(ADALN_OUT, HIDDEN), + 'diffusion_model.t_embedder.mlp.0.diff': torch.randn(HIDDEN, HIDDEN), + 'diffusion_model.cap_embedder.1.diff': torch.randn(HIDDEN, HIDDEN), + } + net = _load_via(Z.try_load, sd) + assert net is not None and len(net.modules) == 5, f'got {sorted(net.modules) if net else None}' + expected = { + f'lora_transformer_all_x_embedder_{PATCH_KEY}', + f'lora_transformer_all_final_layer_{PATCH_KEY}_linear', + f'lora_transformer_all_final_layer_{PATCH_KEY}_adaLN_modulation_1', + 'lora_transformer_t_embedder_mlp_0', + 'lora_transformer_cap_embedder_1', + } + assert set(net.modules) == expected, f'got {set(net.modules)}' + return True + + def test_loha_bfl_proj(): """BFL LoHA on a non-fused proj target binds via NetworkModuleHada.""" net = _load_via(Z.try_load_loha, sd_loha_bfl_proj()) @@ -768,7 +859,8 @@ def run_tests(): t0 = time.time() log.warning('=== Parsing primitives ===') - for fn in [test_parse_key_all_prefixes, test_marker_disambiguation]: + for fn in [test_parse_key_all_prefixes, test_marker_disambiguation, + test_resolve_targets_norm_aliases_and_extra]: run_test(CAT_PARSE, fn) log.warning('=== Loaders ===') @@ -785,6 +877,8 @@ def run_tests(): test_lokr_legacy_fused_qkv_chunked, test_lokr_lycoris_prefix_passthrough, test_full_diff_chain, + test_full_stamps_norm_targets, + test_full_extra_modules_bind, test_loha_bfl_proj, test_loha_legacy_fused_qkv_chunked, test_oft_lycoris_no_npe,