From 783b66c3be88254e792c57d3b86d63bacd9eff31 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 6 Sep 2026 00:00:21 +0100 Subject: [PATCH] refactor(lora): offer unknown bare keys to the resolver 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. --- modules/lora/native_adapter.py | 63 +++++++++++----------------- pipelines/chroma/chroma_lora.py | 7 ---- pipelines/ernie/ernie_lora.py | 7 ---- pipelines/flux/flux2_lora.py | 5 --- pipelines/krea2/krea2_lora.py | 7 ---- pipelines/minimax/minimax_lora.py | 7 +--- pipelines/z_image/zimage_lora.py | 4 -- test/test-minimax-native-adapters.py | 13 ++++++ test/test-zimage-native-adapters.py | 4 +- 9 files changed, 41 insertions(+), 76 deletions(-) diff --git a/modules/lora/native_adapter.py b/modules/lora/native_adapter.py index 20f771ed7..e62b9ddb9 100644 --- a/modules/lora/native_adapter.py +++ b/modules/lora/native_adapter.py @@ -23,8 +23,7 @@ each loader passes in (how a parsed ``(prefix, base)`` maps to one or more diffusers paths plus optional chunk descriptors). Per-arch loader modules import this module and pass their own ``prefixes``, -``bare_prefixes``, ``bare_diffusers_prefixes``, and ``resolve_targets`` to the -generic helpers. +``bare_prefixes`` and ``resolve_targets`` to the generic helpers. """ import os @@ -55,10 +54,11 @@ from modules.lora import lora_common as l KNOWN_PREFIXES_DEFAULT = ("diffusion_model.", "transformer.", "lora_unet_", "lora_transformer_", "lycoris_") -# Sentinel ``prefix_used`` value emitted by :func:`parse_key` when a bare path -# starting with a member of ``bare_diffusers_prefixes`` matches. A loader +# Sentinel ``prefix_used`` value emitted by :func:`parse_key` for a bare path +# that matched no arch prefix and no ``bare_prefixes`` member. A loader # ``resolve_targets`` may dispatch on this string to rewrite the base path; -# when it declines, :func:`resolve_group_targets` binds the path verbatim. +# when it declines, :func:`resolve_group_targets` binds the path verbatim, and +# a path naming no live module counts as unmapped instead of vanishing. BARE_DIFFUSERS_PREFIX_USED = "bare_diffusers" @@ -389,14 +389,15 @@ def lokr_shapes_match(sd_module, kron_shape, chunk: ChunkSpec | None) -> bool: # === Parsing primitives === -def parse_key(key, suffixes, *, prefixes=KNOWN_PREFIXES_DEFAULT, bare_prefixes=(), bare_diffusers_prefixes=()): +def parse_key(key, suffixes, *, prefixes=KNOWN_PREFIXES_DEFAULT, bare_prefixes=()): """Return ``(prefix_used, base, suffix_normalized)`` or ``None``. - ``prefix_used`` is the matched element of ``prefixes``, ``BARE_DIFFUSERS_PREFIX_USED`` - if a member of ``bare_diffusers_prefixes`` matched, or ``None`` for a key - that matched a member of ``bare_prefixes``. ``base`` is the path with prefix - and suffix removed. ``suffix_normalized`` is the suffix (without the leading - dot) after applying :data:`SUFFIX_NORMALIZE` (e.g. ``lora_A.weight`` becomes + ``prefix_used`` is the matched element of ``prefixes``, ``None`` for a key + that matched a member of ``bare_prefixes``, or ``BARE_DIFFUSERS_PREFIX_USED`` + for any other bare key, which the loader offers to the resolver and counts + as unmapped when nothing binds. ``base`` is the path with prefix and suffix + removed. ``suffix_normalized`` is the suffix (without the leading dot) after + applying :data:`SUFFIX_NORMALIZE` (e.g. ``lora_A.weight`` becomes ``lora_down.weight``). Always applies :func:`unwrap_peft_wrapper` and :func:`strip_peft_adapter_name` @@ -411,11 +412,8 @@ def parse_key(key, suffixes, *, prefixes=KNOWN_PREFIXES_DEFAULT, bare_prefixes=( prefix_used = p stripped = key[len(p):] break - if prefix_used is None: - if any(key.startswith(p) for p in bare_diffusers_prefixes): - prefix_used = BARE_DIFFUSERS_PREFIX_USED - elif not any(key.startswith(p) for p in bare_prefixes): - return None + if prefix_used is None and not any(key.startswith(p) for p in bare_prefixes): + prefix_used = BARE_DIFFUSERS_PREFIX_USED matched_suffix = None split_at = -1 @@ -435,7 +433,7 @@ def parse_key(key, suffixes, *, prefixes=KNOWN_PREFIXES_DEFAULT, bare_prefixes=( return prefix_used, base, suffix -def group_by_suffixes(state_dict, suffixes, *, prefixes=KNOWN_PREFIXES_DEFAULT, bare_prefixes=(), bare_diffusers_prefixes=()): +def group_by_suffixes(state_dict, suffixes, *, prefixes=KNOWN_PREFIXES_DEFAULT, bare_prefixes=()): """Group state-dict entries by ``(prefix_used, base)``. Returns ``{(prefix_used, base): {suffix: tensor, ...}}`` where each suffix @@ -445,12 +443,7 @@ def group_by_suffixes(state_dict, suffixes, *, prefixes=KNOWN_PREFIXES_DEFAULT, """ groups: dict[tuple, dict[str, torch.Tensor]] = {} for key, value in state_dict.items(): - parsed = parse_key( - key, suffixes, - prefixes=prefixes, - bare_prefixes=bare_prefixes, - bare_diffusers_prefixes=bare_diffusers_prefixes, - ) + parsed = parse_key(key, suffixes, prefixes=prefixes, bare_prefixes=bare_prefixes) if parsed is None: continue prefix_used, base, suffix = parsed @@ -577,7 +570,7 @@ def slice_bias_delta(w, chunk: ChunkSpec, fused_out): def try_load_lora(name, network_on_disk, lora_scale, *, resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT, - bare_prefixes=(), bare_diffusers_prefixes=(), + bare_prefixes=(), network_prefix=NETWORK_PREFIX_DEFAULT, group_by_suffixes_fn=group_by_suffixes, network_alpha=None, @@ -601,7 +594,6 @@ def try_load_lora(name, network_on_disk, lora_scale, *, state_dict, LORA_SUFFIXES, prefixes=prefixes, bare_prefixes=bare_prefixes, - bare_diffusers_prefixes=bare_diffusers_prefixes, ) if network_alpha is not None and any("alpha" in w for w in groups.values()): network_alpha = None @@ -681,7 +673,7 @@ def try_load_lora(name, network_on_disk, lora_scale, *, def try_load_lokr(name, network_on_disk, lora_scale, *, resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT, - bare_prefixes=(), bare_diffusers_prefixes=(), + bare_prefixes=(), network_prefix=NETWORK_PREFIX_DEFAULT, group_by_suffixes_fn=group_by_suffixes, arch_name="generic"): @@ -705,7 +697,6 @@ def try_load_lokr(name, network_on_disk, lora_scale, *, state_dict, LOKR_SUFFIXES, prefixes=prefixes, bare_prefixes=bare_prefixes, - bare_diffusers_prefixes=bare_diffusers_prefixes, ) unmapped = 0 @@ -767,7 +758,7 @@ def try_load_lokr(name, network_on_disk, lora_scale, *, def try_load_loha(name, network_on_disk, lora_scale, *, resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT, - bare_prefixes=(), bare_diffusers_prefixes=(), + bare_prefixes=(), network_prefix=NETWORK_PREFIX_DEFAULT, group_by_suffixes_fn=group_by_suffixes, arch_name="generic"): @@ -790,7 +781,6 @@ def try_load_loha(name, network_on_disk, lora_scale, *, state_dict, LOHA_SUFFIXES, prefixes=prefixes, bare_prefixes=bare_prefixes, - bare_diffusers_prefixes=bare_diffusers_prefixes, ) unmapped = 0 @@ -841,7 +831,7 @@ def try_load_loha(name, network_on_disk, lora_scale, *, def try_load_oft(name, network_on_disk, lora_scale, *, resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT, - bare_prefixes=(), bare_diffusers_prefixes=(), + bare_prefixes=(), network_prefix=NETWORK_PREFIX_DEFAULT, group_by_suffixes_fn=group_by_suffixes, arch_name="generic"): @@ -870,7 +860,6 @@ def try_load_oft(name, network_on_disk, lora_scale, *, state_dict, OFT_SUFFIXES, prefixes=prefixes, bare_prefixes=bare_prefixes, - bare_diffusers_prefixes=bare_diffusers_prefixes, ) unmapped = 0 @@ -902,7 +891,7 @@ def try_load_oft(name, network_on_disk, lora_scale, *, def try_load_ia3(name, network_on_disk, lora_scale, *, resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT, - bare_prefixes=(), bare_diffusers_prefixes=(), + bare_prefixes=(), network_prefix=NETWORK_PREFIX_DEFAULT, group_by_suffixes_fn=group_by_suffixes, arch_name="generic"): @@ -930,7 +919,6 @@ def try_load_ia3(name, network_on_disk, lora_scale, *, state_dict, IA3_SUFFIXES, prefixes=prefixes, bare_prefixes=bare_prefixes, - bare_diffusers_prefixes=bare_diffusers_prefixes, ) unmapped = 0 @@ -958,7 +946,7 @@ def try_load_ia3(name, network_on_disk, lora_scale, *, def try_load_glora(name, network_on_disk, lora_scale, *, resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT, - bare_prefixes=(), bare_diffusers_prefixes=(), + bare_prefixes=(), network_prefix=NETWORK_PREFIX_DEFAULT, group_by_suffixes_fn=group_by_suffixes, arch_name="generic"): @@ -982,7 +970,6 @@ def try_load_glora(name, network_on_disk, lora_scale, *, state_dict, GLORA_SUFFIXES, prefixes=prefixes, bare_prefixes=bare_prefixes, - bare_diffusers_prefixes=bare_diffusers_prefixes, ) unmapped = 0 @@ -1010,7 +997,7 @@ def try_load_glora(name, network_on_disk, lora_scale, *, def try_load_norm(name, network_on_disk, lora_scale, *, resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT, - bare_prefixes=(), bare_diffusers_prefixes=(), + bare_prefixes=(), network_prefix=NETWORK_PREFIX_DEFAULT, group_by_suffixes_fn=group_by_suffixes, arch_name="generic"): # pylint: disable=unused-argument @@ -1037,7 +1024,6 @@ def try_load_norm(name, network_on_disk, lora_scale, *, state_dict, NORM_SUFFIXES, prefixes=prefixes, bare_prefixes=bare_prefixes, - bare_diffusers_prefixes=bare_diffusers_prefixes, ) unmapped = 0 @@ -1076,7 +1062,7 @@ def try_load_norm(name, network_on_disk, lora_scale, *, def try_load_full(name, network_on_disk, lora_scale, *, resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT, - bare_prefixes=(), bare_diffusers_prefixes=(), + bare_prefixes=(), network_prefix=NETWORK_PREFIX_DEFAULT, group_by_suffixes_fn=group_by_suffixes, arch_name="generic"): @@ -1099,7 +1085,6 @@ def try_load_full(name, network_on_disk, lora_scale, *, state_dict, FULL_SUFFIXES, prefixes=prefixes, bare_prefixes=bare_prefixes, - bare_diffusers_prefixes=bare_diffusers_prefixes, ) unmapped = 0 diff --git a/pipelines/chroma/chroma_lora.py b/pipelines/chroma/chroma_lora.py index 3cd0395a6..2459318d5 100644 --- a/pipelines/chroma/chroma_lora.py +++ b/pipelines/chroma/chroma_lora.py @@ -53,10 +53,6 @@ BARE_FLUX_PREFIXES = ( "img_in.", "txt_in.", "final_layer.", "distilled_guidance_layer.", ) -BARE_DIFFUSERS_PREFIXES = ( - "transformer_blocks.", "single_transformer_blocks.", -) - # === Fused weight dims === # Defaults match Chroma1-HD (``inner_dim = num_attention_heads * @@ -98,7 +94,6 @@ def parse_key(key, suffixes): key, suffixes, prefixes=KNOWN_PREFIXES, bare_prefixes=BARE_FLUX_PREFIXES, - bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, ) @@ -108,7 +103,6 @@ def group_by_suffixes(state_dict, suffixes): state_dict, suffixes, prefixes=KNOWN_PREFIXES, bare_prefixes=BARE_FLUX_PREFIXES, - bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, ) @@ -263,7 +257,6 @@ _BIND_KWARGS = dict( resolve_targets=resolve_targets, prefixes=KNOWN_PREFIXES, bare_prefixes=BARE_FLUX_PREFIXES, - bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, arch_name="chroma", ) diff --git a/pipelines/ernie/ernie_lora.py b/pipelines/ernie/ernie_lora.py index b0b072be3..2255851f3 100644 --- a/pipelines/ernie/ernie_lora.py +++ b/pipelines/ernie/ernie_lora.py @@ -25,10 +25,6 @@ from modules.lora import native_adapter KNOWN_PREFIXES = native_adapter.KNOWN_PREFIXES_DEFAULT -BARE_DIFFUSERS_PREFIXES = ( - "layers.", "adaLN_modulation.", "final_norm.", "final_linear.", -) - # === Re-exports for test/back-compat === @@ -60,7 +56,6 @@ def parse_key(key, suffixes): return native_adapter.parse_key( key, suffixes, prefixes=KNOWN_PREFIXES, - bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, ) @@ -69,7 +64,6 @@ def group_by_suffixes(state_dict, suffixes): return native_adapter.group_by_suffixes( state_dict, suffixes, prefixes=KNOWN_PREFIXES, - bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, ) @@ -95,7 +89,6 @@ def resolve_targets(prefix_used, base): _BIND_KWARGS = dict( resolve_targets=resolve_targets, prefixes=KNOWN_PREFIXES, - bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, arch_name="ernieimage", ) diff --git a/pipelines/flux/flux2_lora.py b/pipelines/flux/flux2_lora.py index 681c65087..9785ed4a9 100644 --- a/pipelines/flux/flux2_lora.py +++ b/pipelines/flux/flux2_lora.py @@ -64,8 +64,6 @@ BARE_FLUX_PREFIXES = ( "double_stream_modulation_", ) -BARE_DIFFUSERS_PREFIXES = ("single_transformer_blocks.", "transformer_blocks.") - # === BFL to diffusers mapping === @@ -162,7 +160,6 @@ def parse_key(key, suffixes): key, suffixes, prefixes=KNOWN_PREFIXES, bare_prefixes=BARE_FLUX_PREFIXES, - bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, ) @@ -172,7 +169,6 @@ def group_by_suffixes(state_dict, suffixes): state_dict, suffixes, prefixes=KNOWN_PREFIXES, bare_prefixes=BARE_FLUX_PREFIXES, - bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, ) @@ -257,7 +253,6 @@ _BIND_KWARGS = dict( resolve_targets=resolve_targets, prefixes=KNOWN_PREFIXES, bare_prefixes=BARE_FLUX_PREFIXES, - bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, arch_name="f2", ) diff --git a/pipelines/krea2/krea2_lora.py b/pipelines/krea2/krea2_lora.py index af3e54681..30d5df0e4 100644 --- a/pipelines/krea2/krea2_lora.py +++ b/pipelines/krea2/krea2_lora.py @@ -23,10 +23,6 @@ KNOWN_PREFIXES = native_adapter.KNOWN_PREFIXES_DEFAULT # Top-level module names that a bare LoRA key can start with: the transformer's own # checkpoint-style names plus the upstream-diffusers names (as saved by # ``Krea2Transformer2DModel.save_lora_adapter()``). -BARE_DIFFUSERS_PREFIXES = ( - "blocks.", "txtfusion.", "first.", "last.", "tmlp.", "tproj.", "txtmlp.", - "transformer_blocks.", "text_fusion.", "img_in.", "txt_in.", "time_embed.", "time_mod_proj.", "final_layer.", -) # Upstream-diffusers attention/ff leaves -> checkpoint leaves (block-level modules). DIFFUSERS_LEAF_MAP = { @@ -83,7 +79,6 @@ def parse_key(key, suffixes): return native_adapter.parse_key( key, suffixes, prefixes=KNOWN_PREFIXES, - bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, ) @@ -92,7 +87,6 @@ def group_by_suffixes(state_dict, suffixes): return native_adapter.group_by_suffixes( state_dict, suffixes, prefixes=KNOWN_PREFIXES, - bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, ) @@ -145,7 +139,6 @@ def _underscore_to_dotted(base): _BIND_KWARGS = dict( resolve_targets=resolve_targets, prefixes=KNOWN_PREFIXES, - bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, arch_name="krea2", ) diff --git a/pipelines/minimax/minimax_lora.py b/pipelines/minimax/minimax_lora.py index 52ca14d25..e9186ccf9 100644 --- a/pipelines/minimax/minimax_lora.py +++ b/pipelines/minimax/minimax_lora.py @@ -37,8 +37,6 @@ KNOWN_PREFIXES = ( # Reference keys outside the block stacks carry no arch prefix in reference saves; the base is the whole module path. BARE_PREFIXES = ("video_patch_proj.", "audio_patch_proj.", "condition_proj.", "time_embedder.", "final_layer.") -# Diffusers module names saved without a component prefix (peft dumps, kohya-suffixed exports) bind verbatim. -BARE_DIFFUSERS_PREFIXES = ("transformer_blocks.", "token_refiner.refiner_blocks.", "proj_in.", "audio_proj_in.", "context_embedder.", "time_embedder.linear_", "norm_out.", "proj_out.", "audio_proj_out.") STANDALONE_RENAMES = { "video_patch_proj": "proj_in", @@ -126,14 +124,14 @@ def parse_key(key, suffixes): key = native_adapter.unwrap_peft_wrapper(key) if key.startswith("dit."): key = "diffusion_model." + key[len("dit."):] - parsed = native_adapter.parse_key(key, suffixes, prefixes=KNOWN_PREFIXES, bare_prefixes=BARE_PREFIXES, bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES) + parsed = native_adapter.parse_key(key, suffixes, prefixes=KNOWN_PREFIXES, bare_prefixes=BARE_PREFIXES) if parsed is None: return None prefix_used, base, suffix = parsed return prefix_used, base, normalize_mini_max_suffix(suffix) -def group_by_suffixes(state_dict, suffixes, *, prefixes=None, bare_prefixes=(), bare_diffusers_prefixes=()): # pylint: disable=unused-argument +def group_by_suffixes(state_dict, suffixes, *, prefixes=None, bare_prefixes=()): # pylint: disable=unused-argument """MiniMax-bound :func:`native_adapter.group_by_suffixes`.""" groups: dict[tuple, dict[str, object]] = {} for key, value in state_dict.items(): @@ -238,7 +236,6 @@ _BIND_KWARGS = dict( resolve_targets=resolve_targets, prefixes=KNOWN_PREFIXES, bare_prefixes=BARE_PREFIXES, - bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, network_prefix=network_prefix_for, group_by_suffixes_fn=group_by_suffixes, arch_name="minimaxh3", diff --git a/pipelines/z_image/zimage_lora.py b/pipelines/z_image/zimage_lora.py index e8c6d7ec4..534a09551 100644 --- a/pipelines/z_image/zimage_lora.py +++ b/pipelines/z_image/zimage_lora.py @@ -40,7 +40,6 @@ from modules.lora.native_adapter import ChunkSpec 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 = { @@ -109,7 +108,6 @@ def parse_key(key, suffixes): return native_adapter.parse_key( key, suffixes, prefixes=KNOWN_PREFIXES, - bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, ) @@ -118,7 +116,6 @@ def group_by_suffixes(state_dict, suffixes): return native_adapter.group_by_suffixes( state_dict, suffixes, prefixes=KNOWN_PREFIXES, - bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, ) @@ -199,7 +196,6 @@ def _underscore_to_diffusers_targets(base): _BIND_KWARGS = dict( resolve_targets=resolve_targets, prefixes=KNOWN_PREFIXES, - bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, arch_name="zimage", ) diff --git a/test/test-minimax-native-adapters.py b/test/test-minimax-native-adapters.py index 616f32caf..9d8e7ff30 100644 --- a/test/test-minimax-native-adapters.py +++ b/test/test-minimax-native-adapters.py @@ -489,6 +489,18 @@ def test_diffusers_peft_keys_bind_verbatim(): return True +def test_unknown_bare_key_reaches_the_resolver(): + """A bare key that names no module is offered verbatim and left unbound, not dropped at parse time.""" + sd = lora_pair('transformer_blocks.0.attn.to_q', LINEAR_SHAPES['transformer_blocks.0.attn.to_q']) + sd.update(lora_pair('nowhere.proj', (4, 4))) + mapping = native_mapping(sd) + assert set(mapping) == {'transformer_blocks.0.attn.to_q', 'nowhere.proj'}, sorted(mapping) + net = load_native(sd, name='stray') + assert set(net.modules) == {'lora_transformer_transformer_blocks_0_attn_to_q'}, sorted(net.modules) + assert net.mismatch == 0 + return True + + # ============================================================ # Tests - loader against the reference loaders # ============================================================ @@ -637,6 +649,7 @@ def run_tests(): test_chunk_reorder_composes_with_slice, test_qkv_split_order, test_diffusers_peft_keys_bind_verbatim, + test_unknown_bare_key_reaches_the_resolver, ]: run_test(CAT_RESOLVE, fn) diff --git a/test/test-zimage-native-adapters.py b/test/test-zimage-native-adapters.py index 77004d18e..d7d799694 100644 --- a/test/test-zimage-native-adapters.py +++ b/test/test-zimage-native-adapters.py @@ -456,8 +456,8 @@ def test_parse_key_all_prefixes(): """parse_key recognizes BFL, PEFT, kohya, and bare-diffusers keys. Returns (prefix_used, base, suffix) - prefix_used is the matched - KNOWN_PREFIXES element, BARE_DIFFUSERS_PREFIX_USED for bare paths - matching BARE_DIFFUSERS_PREFIXES, or None when no prefix is recognized. + KNOWN_PREFIXES element, BARE_DIFFUSERS_PREFIX_USED for any other bare + path, or None when no suffix is recognized. """ bd = Z.BARE_DIFFUSERS_PREFIX_USED cases = [