Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 09:14:35 +02:00
Code Issues Packages Projects Releases Wiki Activity

refactor(lora): offer unknown bare keys to the resolver

Browse Source 
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.
This commit is contained in:
CalamitousFelicitousness
2026-09-06 00:00:21 +01:00
parent 1c91fb2047
commit 783b66c3be
9 changed files with 41 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/lora/native_adapter.py
+24 -39
View File
@@ -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
pipelines/chroma/chroma_lora.py
-7
View File
@@ -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",
)
pipelines/ernie/ernie_lora.py
-7
View File
@@ -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",
)
pipelines/flux/flux2_lora.py
-5
View File
@@ -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",
)
pipelines/krea2/krea2_lora.py
-7
View File
@@ -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",
)
pipelines/minimax/minimax_lora.py
+2 -5
View File
@@ -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",
pipelines/z_image/zimage_lora.py
-4
View File
@@ -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",
)
test/test-minimax-native-adapters.py
+13
View File
@@ -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)
test/test-zimage-native-adapters.py
+2 -2
View File
@@ -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 = [