From b434eb0c1b6a9739b46b6294f91ee0671944a76c Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 6 Sep 2026 00:10:07 +0100 Subject: [PATCH] refactor(lora): route every bare key through the resolver The parser no longer takes reference-name prefixes to tell bare reference keys from bare diffusers keys. Any bare key carries the sentinel and the arch resolver renames what it knows and passes the rest through. Flux2 keeps its list for file-format detection only. --- modules/lora/native_adapter.py | 51 ++++++++++------------------- pipelines/chroma/chroma_lora.py | 18 +++------- pipelines/flux/flux2_lora.py | 10 +++--- pipelines/minimax/minimax_lora.py | 10 ++---- test/test-chroma-native-adapters.py | 8 ++--- test/test-flux2-native-adapters.py | 10 +++--- 6 files changed, 37 insertions(+), 70 deletions(-) diff --git a/modules/lora/native_adapter.py b/modules/lora/native_adapter.py index e62b9ddb9..289dbaf2f 100644 --- a/modules/lora/native_adapter.py +++ b/modules/lora/native_adapter.py @@ -22,8 +22,8 @@ of a fused weight is described) and the per-arch ``resolve_targets`` callable 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`` and ``resolve_targets`` to the generic helpers. +Per-arch loader modules import this module and pass their own ``prefixes`` +and ``resolve_targets`` to the generic helpers. """ import os @@ -54,11 +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` 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, and -# a path naming no live module counts as unmapped instead of vanishing. +# Sentinel ``prefix_used`` value emitted by :func:`parse_key` for a bare path, +# one that matched no arch prefix. 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, and a path naming no +# live module counts as unmapped instead of vanishing. BARE_DIFFUSERS_PREFIX_USED = "bare_diffusers" @@ -389,16 +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=()): +def parse_key(key, suffixes, *, prefixes=KNOWN_PREFIXES_DEFAULT): """Return ``(prefix_used, base, suffix_normalized)`` or ``None``. - ``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``). + ``prefix_used`` is the matched element of ``prefixes``, or + ``BARE_DIFFUSERS_PREFIX_USED`` for a 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` to the raw key before format detection so callers do not have to opt in. @@ -412,7 +411,7 @@ def parse_key(key, suffixes, *, prefixes=KNOWN_PREFIXES_DEFAULT, bare_prefixes=( prefix_used = p stripped = key[len(p):] break - if prefix_used is None and not any(key.startswith(p) for p in bare_prefixes): + if prefix_used is None: prefix_used = BARE_DIFFUSERS_PREFIX_USED matched_suffix = None @@ -433,7 +432,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=()): +def group_by_suffixes(state_dict, suffixes, *, prefixes=KNOWN_PREFIXES_DEFAULT): """Group state-dict entries by ``(prefix_used, base)``. Returns ``{(prefix_used, base): {suffix: tensor, ...}}`` where each suffix @@ -443,7 +442,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) + parsed = parse_key(key, suffixes, prefixes=prefixes) if parsed is None: continue prefix_used, base, suffix = parsed @@ -570,7 +569,6 @@ 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=(), network_prefix=NETWORK_PREFIX_DEFAULT, group_by_suffixes_fn=group_by_suffixes, network_alpha=None, @@ -593,7 +591,6 @@ def try_load_lora(name, network_on_disk, lora_scale, *, groups = group_by_suffixes_fn( state_dict, LORA_SUFFIXES, prefixes=prefixes, - bare_prefixes=bare_prefixes, ) if network_alpha is not None and any("alpha" in w for w in groups.values()): network_alpha = None @@ -673,7 +670,6 @@ 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=(), network_prefix=NETWORK_PREFIX_DEFAULT, group_by_suffixes_fn=group_by_suffixes, arch_name="generic"): @@ -696,7 +692,6 @@ def try_load_lokr(name, network_on_disk, lora_scale, *, groups = group_by_suffixes_fn( state_dict, LOKR_SUFFIXES, prefixes=prefixes, - bare_prefixes=bare_prefixes, ) unmapped = 0 @@ -758,7 +753,6 @@ 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=(), network_prefix=NETWORK_PREFIX_DEFAULT, group_by_suffixes_fn=group_by_suffixes, arch_name="generic"): @@ -780,7 +774,6 @@ def try_load_loha(name, network_on_disk, lora_scale, *, groups = group_by_suffixes_fn( state_dict, LOHA_SUFFIXES, prefixes=prefixes, - bare_prefixes=bare_prefixes, ) unmapped = 0 @@ -831,7 +824,6 @@ 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=(), network_prefix=NETWORK_PREFIX_DEFAULT, group_by_suffixes_fn=group_by_suffixes, arch_name="generic"): @@ -859,7 +851,6 @@ def try_load_oft(name, network_on_disk, lora_scale, *, groups = group_by_suffixes_fn( state_dict, OFT_SUFFIXES, prefixes=prefixes, - bare_prefixes=bare_prefixes, ) unmapped = 0 @@ -891,7 +882,6 @@ 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=(), network_prefix=NETWORK_PREFIX_DEFAULT, group_by_suffixes_fn=group_by_suffixes, arch_name="generic"): @@ -918,7 +908,6 @@ def try_load_ia3(name, network_on_disk, lora_scale, *, groups = group_by_suffixes_fn( state_dict, IA3_SUFFIXES, prefixes=prefixes, - bare_prefixes=bare_prefixes, ) unmapped = 0 @@ -946,7 +935,6 @@ 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=(), network_prefix=NETWORK_PREFIX_DEFAULT, group_by_suffixes_fn=group_by_suffixes, arch_name="generic"): @@ -969,7 +957,6 @@ def try_load_glora(name, network_on_disk, lora_scale, *, groups = group_by_suffixes_fn( state_dict, GLORA_SUFFIXES, prefixes=prefixes, - bare_prefixes=bare_prefixes, ) unmapped = 0 @@ -997,7 +984,6 @@ 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=(), network_prefix=NETWORK_PREFIX_DEFAULT, group_by_suffixes_fn=group_by_suffixes, arch_name="generic"): # pylint: disable=unused-argument @@ -1023,7 +1009,6 @@ def try_load_norm(name, network_on_disk, lora_scale, *, groups = group_by_suffixes_fn( state_dict, NORM_SUFFIXES, prefixes=prefixes, - bare_prefixes=bare_prefixes, ) unmapped = 0 @@ -1062,7 +1047,6 @@ 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=(), network_prefix=NETWORK_PREFIX_DEFAULT, group_by_suffixes_fn=group_by_suffixes, arch_name="generic"): @@ -1084,7 +1068,6 @@ def try_load_full(name, network_on_disk, lora_scale, *, groups = group_by_suffixes_fn( state_dict, FULL_SUFFIXES, prefixes=prefixes, - bare_prefixes=bare_prefixes, ) unmapped = 0 diff --git a/pipelines/chroma/chroma_lora.py b/pipelines/chroma/chroma_lora.py index 2459318d5..c84faf036 100644 --- a/pipelines/chroma/chroma_lora.py +++ b/pipelines/chroma/chroma_lora.py @@ -44,15 +44,6 @@ from modules.lora.native_adapter import ChunkSpec KNOWN_PREFIXES = native_adapter.KNOWN_PREFIXES_DEFAULT -# distilled_guidance_layer. is deliberately a bare-BFL prefix, not a -# bare-diffusers one: the resolver renames BFL MLP leaves (in_layer/out_layer) -# and passes diffusers-named leaves (linear_1/linear_2, in_proj, out_proj, -# norms.N) through verbatim, so both namings route correctly. -BARE_FLUX_PREFIXES = ( - "double_blocks.", "single_blocks.", - "img_in.", "txt_in.", "final_layer.", "distilled_guidance_layer.", -) - # === Fused weight dims === # Defaults match Chroma1-HD (``inner_dim = num_attention_heads * @@ -93,7 +84,6 @@ def parse_key(key, suffixes): return native_adapter.parse_key( key, suffixes, prefixes=KNOWN_PREFIXES, - bare_prefixes=BARE_FLUX_PREFIXES, ) @@ -102,7 +92,6 @@ def group_by_suffixes(state_dict, suffixes): return native_adapter.group_by_suffixes( state_dict, suffixes, prefixes=KNOWN_PREFIXES, - bare_prefixes=BARE_FLUX_PREFIXES, ) @@ -116,14 +105,16 @@ def resolve_targets(prefix_used, base): - ``lora_unet_``: kohya underscore-flat Flux path; parse block type/index and module suffix, rename to diffusers. - - ``diffusion_model.`` or bare BFL (None): dotted Flux path; same rewrite. + - ``diffusion_model.`` or bare (the sentinel): dotted Flux path; same + rewrite. BFL MLP leaves under ``distilled_guidance_layer`` rename and + diffusers-named leaves pass through verbatim, so both namings route. Universal passthrough prefixes are handled upstream by :func:`native_adapter.resolve_group_targets`. """ if prefix_used == "lora_unet_": return _kohya_to_diffusers(base) - if prefix_used in (None, "diffusion_model."): + if prefix_used in (BARE_DIFFUSERS_PREFIX_USED, "diffusion_model."): return _bfl_to_diffusers(base) return [] @@ -256,7 +247,6 @@ def _split_single_linear1(block_idx): _BIND_KWARGS = dict( resolve_targets=resolve_targets, prefixes=KNOWN_PREFIXES, - bare_prefixes=BARE_FLUX_PREFIXES, arch_name="chroma", ) diff --git a/pipelines/flux/flux2_lora.py b/pipelines/flux/flux2_lora.py index 9785ed4a9..95651d859 100644 --- a/pipelines/flux/flux2_lora.py +++ b/pipelines/flux/flux2_lora.py @@ -159,7 +159,6 @@ def parse_key(key, suffixes): return native_adapter.parse_key( key, suffixes, prefixes=KNOWN_PREFIXES, - bare_prefixes=BARE_FLUX_PREFIXES, ) @@ -168,7 +167,6 @@ def group_by_suffixes(state_dict, suffixes): return native_adapter.group_by_suffixes( state_dict, suffixes, prefixes=KNOWN_PREFIXES, - bare_prefixes=BARE_FLUX_PREFIXES, ) @@ -179,15 +177,16 @@ def resolve_targets(prefix_used, base): """Return ``[(diffusers_path, ChunkSpec | None), ...]`` for a parsed group key. For ``lora_unet_`` prefix, applies ``KOHYA_SUFFIX_MAP`` then ``F2_*_MAP``. - For BFL / bare-BFL, applies ``F2_*_MAP`` directly. Unrecognized prefixes - return an empty list. + For BFL and bare keys, applies ``F2_*_MAP`` directly; a bare path the maps + do not know binds verbatim upstream. Unrecognized prefixes return an empty + list. Universal passthrough prefixes (including ``lycoris_``) are handled upstream by :func:`native_adapter.resolve_group_targets`. """ if prefix_used == "lora_unet_": return _kohya_to_diffusers_targets(base) - if prefix_used in (None, "diffusion_model."): + if prefix_used in (BARE_DIFFUSERS_PREFIX_USED, "diffusion_model."): return _bfl_to_diffusers_targets(base) return [] @@ -252,7 +251,6 @@ def _bfl_to_diffusers_targets(base): _BIND_KWARGS = dict( resolve_targets=resolve_targets, prefixes=KNOWN_PREFIXES, - bare_prefixes=BARE_FLUX_PREFIXES, arch_name="f2", ) diff --git a/pipelines/minimax/minimax_lora.py b/pipelines/minimax/minimax_lora.py index 42b81c396..5709772f4 100644 --- a/pipelines/minimax/minimax_lora.py +++ b/pipelines/minimax/minimax_lora.py @@ -34,9 +34,6 @@ KNOWN_PREFIXES = ( "token_refiner.", ) + native_adapter.KNOWN_PREFIXES_DEFAULT -# 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.") - STANDALONE_RENAMES = { "video_patch_proj": "proj_in", @@ -135,14 +132,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) + parsed = native_adapter.parse_key(key, suffixes, prefixes=KNOWN_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=()): # pylint: disable=unused-argument +def group_by_suffixes(state_dict, suffixes, *, prefixes=None): # 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(): @@ -169,7 +166,7 @@ def _block_targets(target_stack, base): def resolve_targets(prefix_used, base): """Return ``[(diffusers_path, ChunkSpec | None), ...]`` for MiniMax keys.""" - if prefix_used == "diffusion_model." or prefix_used is None: + if prefix_used in ("diffusion_model.", BARE_DIFFUSERS_PREFIX_USED): if base.startswith("transformer."): return [(base[len("transformer."):], None)] if base.startswith("text_encoder."): @@ -230,7 +227,6 @@ def file_alpha(network_on_disk): _BIND_KWARGS = dict( resolve_targets=resolve_targets, prefixes=KNOWN_PREFIXES, - bare_prefixes=BARE_PREFIXES, network_prefix=network_prefix_for, group_by_suffixes_fn=group_by_suffixes, arch_name="minimaxh3", diff --git a/test/test-chroma-native-adapters.py b/test/test-chroma-native-adapters.py index 696746694..b21114a8f 100644 --- a/test/test-chroma-native-adapters.py +++ b/test/test-chroma-native-adapters.py @@ -591,7 +591,7 @@ def test_parse_key_all_prefixes(): # Bare BFL path (no prefix) ('double_blocks.0.img_attn.proj.lora_A.weight', C.LORA_SUFFIXES, - (None, 'double_blocks.0.img_attn.proj', 'lora_down.weight')), + (C.BARE_DIFFUSERS_PREFIX_USED, 'double_blocks.0.img_attn.proj', 'lora_down.weight')), ('random.unrelated.key', C.LORA_SUFFIXES, None), ] for key, suffixes, expected in cases: @@ -649,7 +649,7 @@ def test_resolve_targets_extra_and_guidance(): for bfl_base, diffusers_path in C.CHROMA_EXTRA_MAP.items(): for prefix, base in [ ('diffusion_model.', bfl_base), - (None, bfl_base), + (C.BARE_DIFFUSERS_PREFIX_USED, bfl_base), ('lora_unet_', bfl_base.replace('.', '_')), ]: targets = C.resolve_targets(prefix, base) @@ -657,12 +657,12 @@ def test_resolve_targets_extra_and_guidance(): cases = [ # BFL MLP leaves rename to the PixArt projection names. (('diffusion_model.', 'distilled_guidance_layer.layers.0.in_layer'), 'distilled_guidance_layer.layers.0.linear_1'), - ((None, 'distilled_guidance_layer.layers.1.out_layer'), 'distilled_guidance_layer.layers.1.linear_2'), + ((C.BARE_DIFFUSERS_PREFIX_USED, 'distilled_guidance_layer.layers.1.out_layer'), 'distilled_guidance_layer.layers.1.linear_2'), (('lora_unet_', 'distilled_guidance_layer_layers_0_in_layer'), 'distilled_guidance_layer_layers_0_linear_1'), (('lora_unet_', 'distilled_guidance_layer_layers_1_out_layer'), 'distilled_guidance_layer_layers_1_linear_2'), # Verbatim leaves are untouched in either naming. (('diffusion_model.', 'distilled_guidance_layer.in_proj'), 'distilled_guidance_layer.in_proj'), - ((None, 'distilled_guidance_layer.layers.0.linear_1'), 'distilled_guidance_layer.layers.0.linear_1'), + ((C.BARE_DIFFUSERS_PREFIX_USED, 'distilled_guidance_layer.layers.0.linear_1'), 'distilled_guidance_layer.layers.0.linear_1'), ] for (prefix, base), expected in cases: targets = C.resolve_targets(prefix, base) diff --git a/test/test-flux2-native-adapters.py b/test/test-flux2-native-adapters.py index 1375e25a1..816573570 100644 --- a/test/test-flux2-native-adapters.py +++ b/test/test-flux2-native-adapters.py @@ -670,7 +670,7 @@ def test_parse_key_all_prefixes(): ('transformer.', 'transformer_blocks.0.attn.to_q', 'lora_up.weight')), ('double_blocks.10.img_mlp.0.lora_A.weight', F.LORA_SUFFIXES, - (None, 'double_blocks.10.img_mlp.0', 'lora_down.weight')), + (F.BARE_DIFFUSERS_PREFIX_USED, 'double_blocks.10.img_mlp.0', 'lora_down.weight')), ('random.unrelated.key', F.LORA_SUFFIXES, None), ] for key, suffixes, expected in cases: @@ -712,13 +712,13 @@ def test_resolve_targets_extra_modules(): for bfl_base, diffusers_path in F.F2_EXTRA_MAP.items(): targets = F.resolve_targets('diffusion_model.', bfl_base) assert targets == [(diffusers_path, None)], f'{bfl_base} -> {targets}' - targets = F.resolve_targets(None, bfl_base) + targets = F.resolve_targets(F.BARE_DIFFUSERS_PREFIX_USED, bfl_base) assert targets == [(diffusers_path, None)], f'bare {bfl_base} -> {targets}' targets = F.resolve_targets('lora_unet_', bfl_base.replace('.', '_')) assert targets == [(diffusers_path, None)], f'kohya {bfl_base} -> {targets}' - # guidance_in is a bare BFL prefix in its own right. + # guidance_in is a bare BFL path. got = F.parse_key('guidance_in.in_layer.lora_A.weight', F.LORA_SUFFIXES) - assert got == (None, 'guidance_in.in_layer', 'lora_down.weight'), f'bare guidance_in parse -> {got}' + assert got == (F.BARE_DIFFUSERS_PREFIX_USED, 'guidance_in.in_layer', 'lora_down.weight'), f'bare guidance_in parse -> {got}' return True @@ -734,7 +734,7 @@ def test_parse_key_peft_wrapper_unwrap(): cases = [ ('base_model.model.double_blocks.1.img_attn.proj.lora_A.weight', F.LORA_SUFFIXES, - (None, 'double_blocks.1.img_attn.proj', 'lora_down.weight')), + (F.BARE_DIFFUSERS_PREFIX_USED, 'double_blocks.1.img_attn.proj', 'lora_down.weight')), ('base_model.model.transformer.transformer_blocks.0.attn.to_q.lora_A.weight', F.LORA_SUFFIXES, ('transformer.', 'transformer_blocks.0.attn.to_q', 'lora_down.weight')),