diff --git a/modules/lora/native_adapter.py b/modules/lora/native_adapter.py index dc20b20f2..ff2b6c934 100644 --- a/modules/lora/native_adapter.py +++ b/modules/lora/native_adapter.py @@ -65,6 +65,13 @@ BARE_DIFFUSERS_PREFIX_USED = "bare_diffusers" NETWORK_PREFIX_DEFAULT = "lora_transformer_" +# Prefixes whose parsed ``base`` is already a network-key tail (``arch_prefix + +# base.replace(".", "_")`` matches the stamped module name), so the loader binds +# them directly with no per-arch rewrite. Arch-local already-resolved prefixes +# (e.g. flux2's ``lycoris_``) stay in that arch's ``resolve_targets``. +PASSTHROUGH_PREFIXES_DEFAULT = ("transformer.", BARE_DIFFUSERS_PREFIX_USED, "lora_transformer_") + + def _resolve_prefix(network_prefix, prefix_used): """Return the network-key prefix for one parsed group. @@ -364,19 +371,11 @@ read_state_dict = sd_models.read_state_dict def resolve_group_targets(resolve_targets, prefix_used, base): """Map a parsed ``(prefix_used, base)`` group to ``[(diffusers_path, chunk), ...]``. - Handles the universal ``lora_transformer_`` passthrough before deferring to - the arch's ``resolve_targets`` for Flux-layout, kohya, and bare prefixes. - - ``lora_transformer_`` is the namespace - ``lora_convert.assign_network_names_to_compvis_modules`` stamps on every - arch's transformer modules (``lora_transformer_`` + module path with dots - replaced by underscores). A file already saved in that form (e.g. - OneTrainer, which trains against the diffusers layout with QKV pre-split) - carries the network-key tail verbatim, so its base needs no rename or - chunking and binds for any arch whose ``network_prefix`` resolves to - ``lora_transformer_``. + Passthrough prefixes (:data:`PASSTHROUGH_PREFIXES_DEFAULT`) bind verbatim; + everything else defers to the arch's ``resolve_targets``. Centralizing the + passthrough keeps each arch's resolver to the prefixes it actually rewrites. """ - if prefix_used == "lora_transformer_": + if prefix_used in PASSTHROUGH_PREFIXES_DEFAULT: return [(base, None)] return resolve_targets(prefix_used, base) diff --git a/pipelines/chroma/chroma_lora.py b/pipelines/chroma/chroma_lora.py index 965d5bf52..8a95a69eb 100644 --- a/pipelines/chroma/chroma_lora.py +++ b/pipelines/chroma/chroma_lora.py @@ -103,12 +103,10 @@ 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. - - ``transformer.`` or bare-diffusers: already a diffusers path; passthrough. + + Universal passthrough prefixes are handled upstream by + :func:`native_adapter.resolve_group_targets`. """ - if prefix_used == "transformer.": - return [(base, None)] - if prefix_used == BARE_DIFFUSERS_PREFIX_USED: - return [(base, None)] if prefix_used == "lora_unet_": return _kohya_to_diffusers(base) if prefix_used in (None, "diffusion_model."): diff --git a/pipelines/ernie/ernie_lora.py b/pipelines/ernie/ernie_lora.py index e76d716aa..e3eb3d842 100644 --- a/pipelines/ernie/ernie_lora.py +++ b/pipelines/ernie/ernie_lora.py @@ -67,10 +67,14 @@ def group_by_suffixes(state_dict, suffixes): def resolve_targets(prefix_used, base): - """Passthrough for every recognized prefix. ERNIE has no fused targets or path - renames; the base path is already the diffusers module path.""" - if prefix_used in ("diffusion_model.", "transformer.", "lora_unet_", - BARE_DIFFUSERS_PREFIX_USED, None): + """Passthrough: ERNIE has fully split attention and no path renames, so its + ``diffusion_model.`` / ``lora_unet_`` / bare-BFL keys are already diffusers + module paths. + + Universal passthrough prefixes are handled upstream by + :func:`native_adapter.resolve_group_targets`. + """ + if prefix_used in ("diffusion_model.", "lora_unet_", None): return [(base, None)] return [] diff --git a/pipelines/flux/flux2_lora.py b/pipelines/flux/flux2_lora.py index 1c69a66b4..622fe838e 100644 --- a/pipelines/flux/flux2_lora.py +++ b/pipelines/flux/flux2_lora.py @@ -160,18 +160,17 @@ 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. For ``transformer.``, - ``lycoris_``, and bare-diffusers, returns the base verbatim with no chunking. - Unrecognized prefixes return an empty list. + For BFL / bare-BFL, applies ``F2_*_MAP`` directly. ``lycoris_`` is an + already-underscored diffusers path, returned verbatim. Unrecognized + prefixes return an empty list. + + Universal passthrough prefixes 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."): return _bfl_to_diffusers_targets(base) - if prefix_used == "transformer.": - return [(base, None)] - if prefix_used == BARE_DIFFUSERS_PREFIX_USED: - return [(base, None)] if prefix_used == "lycoris_": # base is an already-underscored diffusers path (e.g. # 'transformer_blocks_0_attn_add_k_proj'). The caller's network_key diff --git a/pipelines/z_image/zimage_lora.py b/pipelines/z_image/zimage_lora.py index f3d39e536..7a3ad7a69 100644 --- a/pipelines/z_image/zimage_lora.py +++ b/pipelines/z_image/zimage_lora.py @@ -83,13 +83,12 @@ def resolve_targets(prefix_used, base): Everything else (modern split-attention paths, MLP, norms, embedders) is returned verbatim. + + Universal passthrough prefixes are handled upstream by + :func:`native_adapter.resolve_group_targets`. """ if prefix_used == "lora_unet_": return _underscore_to_diffusers_targets(base) - if prefix_used == "transformer.": - return [(base, None)] - if prefix_used == BARE_DIFFUSERS_PREFIX_USED: - return [(base, None)] if prefix_used in (None, "diffusion_model."): return _dotted_to_diffusers_targets(base) return [] diff --git a/test/test-flux2-native-adapters.py b/test/test-flux2-native-adapters.py index 943120ad2..be755e724 100644 --- a/test/test-flux2-native-adapters.py +++ b/test/test-flux2-native-adapters.py @@ -599,7 +599,7 @@ def test_resolve_targets_qkv_chunking(): targets = F.resolve_targets('diffusion_model.', 'single_blocks.7.linear1') assert targets == [('single_transformer_blocks.7.attn.to_qkv_mlp_proj', None)] - targets = F.resolve_targets('transformer.', 'transformer_blocks.0.attn.to_q') + targets = F.native_adapter.resolve_group_targets(F.resolve_targets, 'transformer.', 'transformer_blocks.0.attn.to_q') assert targets == [('transformer_blocks.0.attn.to_q', None)] targets = F.resolve_targets('weird_prefix.', 'whatever') @@ -690,8 +690,8 @@ def test_parse_key_bare_diffusers_and_peft_default(): got = F.parse_key(key, suffixes) assert got == expected, f'parse_key({key!r}) = {got}, expected {expected}' - # resolve_targets passes the bare-diffusers path through verbatim. - targets = F.resolve_targets(bd, 'single_transformer_blocks.5.attn.to_out') + # the shared resolver passes the bare-diffusers path through verbatim. + targets = F.native_adapter.resolve_group_targets(F.resolve_targets, bd, 'single_transformer_blocks.5.attn.to_out') assert targets == [('single_transformer_blocks.5.attn.to_out', None)], f'targets={targets}' return True