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

refactor(lora): route every bare key through the resolver

Browse Source 
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.
This commit is contained in:
CalamitousFelicitousness
2026-09-06 00:10:07 +01:00
parent 8b93d20b1a
commit b434eb0c1b
6 changed files with 37 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/lora/native_adapter.py
+17 -34
View File
@@ -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
pipelines/chroma/chroma_lora.py
+4 -14
View File
@@ -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",
)
pipelines/flux/flux2_lora.py
+4 -6
View File
@@ -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",
)
pipelines/minimax/minimax_lora.py
+3 -7
View File
@@ -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",
test/test-chroma-native-adapters.py
+4 -4
View File
@@ -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)
test/test-flux2-native-adapters.py
+5 -5
View File
@@ -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')),