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')),