From 7cd6b6be79b04b6041000ff7b2fed649a6a21c8c Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 25 May 2026 05:47:56 +0100 Subject: [PATCH 01/27] refactor(lora): rename native_loader to native_adapter Frees the name for pipelines/native_transformer. Module covers the full LyCORIS adapter family (LoRA/LoKR/LoHA/OFT/IA3/GLoRA/Norm/Full), not just LoRA. --- modules/lora/{native_loader.py => native_adapter.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename modules/lora/{native_loader.py => native_adapter.py} (100%) diff --git a/modules/lora/native_loader.py b/modules/lora/native_adapter.py similarity index 100% rename from modules/lora/native_loader.py rename to modules/lora/native_adapter.py From 5b321d2071faeacfb2c78b606a7794dab6627ab9 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 25 May 2026 05:47:56 +0100 Subject: [PATCH 02/27] refactor(lora): update importers for native_adapter rename 5 pipeline files + 2 test files. Mechanical substitution. 141/141 adapter tests pass. --- modules/lora/lora_load.py | 2 +- pipelines/anima/anima_lora.py | 74 ++++++++++++++-------------- pipelines/chroma/chroma_lora.py | 48 +++++++++--------- pipelines/ernie/ernie_lora.py | 46 +++++++++--------- pipelines/flux/flux2_lora.py | 78 +++++++++++++++--------------- pipelines/z_image/zimage_lora.py | 48 +++++++++--------- test/test-anima-native-adapters.py | 4 +- test/test-flux2-native-adapters.py | 2 +- 8 files changed, 151 insertions(+), 151 deletions(-) diff --git a/modules/lora/lora_load.py b/modules/lora/lora_load.py index a454eff91..d87654c95 100644 --- a/modules/lora/lora_load.py +++ b/modules/lora/lora_load.py @@ -19,7 +19,7 @@ exclude_errors = [ # shared.sd_model_type -> dotted module path of a pipeline native loader # exposing ``try_load(name, network_on_disk, lora_scale)``. New archs add an -# entry here and ship a per-arch ``try_load`` (either binding native_loader's +# entry here and ship a per-arch ``try_load`` (either binding native_adapter's # generic helpers via try_load_chain, or rolling their own). _NATIVE_DISPATCH = { 'zimage': 'pipelines.z_image.zimage_lora', diff --git a/pipelines/anima/anima_lora.py b/pipelines/anima/anima_lora.py index 39530c243..cc9c01dbf 100644 --- a/pipelines/anima/anima_lora.py +++ b/pipelines/anima/anima_lora.py @@ -25,23 +25,23 @@ already stamped on the network_layer_mapping. Adapter and TE paths bypass the rename and are flattened verbatim. Network-key construction (transformer vs llm_adapter vs te) is parameterized -in :mod:`modules.lora.native_loader` via the ``network_prefix`` kwarg; this +in :mod:`modules.lora.native_adapter` via the ``network_prefix`` kwarg; this module supplies :func:`network_prefix_for` to pick per ``prefix_used``. Family-specific dispatch (LoRA, LoHA, LoKR, OFT, IA3, GLoRA, Norm, Full) is -inherited from native_loader's generics; alpha / scale / DoRA flow through +inherited from native_adapter's generics; alpha / scale / DoRA flow through the standard ``NetworkWeights.w`` slots rather than being baked into the factor weights at load time. """ from collections import OrderedDict -from modules.lora import native_loader +from modules.lora import native_adapter # === Arch-specific prefix configuration === # # Order matters: longer / more-specific prefixes must precede shorter ones, -# because :func:`native_loader.parse_key` returns the first match. Both +# because :func:`native_adapter.parse_key` returns the first match. Both # ``diffusion_model.llm_adapter.`` and ``text_encoders.qwen3_06b.transformer.model.`` # start with ``diffusion_model.`` / ``text_encoders.`` so they must be listed first. @@ -58,37 +58,37 @@ ANIMA_PREFIXES = ( # Tests address these through the anima_lora module surface; sibling pipelines # do the same (see flux2_lora / zimage_lora / ernie_lora). -LORA_SUFFIXES = native_loader.LORA_SUFFIXES -LOKR_SUFFIXES = native_loader.LOKR_SUFFIXES -LOHA_SUFFIXES = native_loader.LOHA_SUFFIXES -OFT_SUFFIXES = native_loader.OFT_SUFFIXES -IA3_SUFFIXES = native_loader.IA3_SUFFIXES -GLORA_SUFFIXES = native_loader.GLORA_SUFFIXES -NORM_SUFFIXES = native_loader.NORM_SUFFIXES -FULL_SUFFIXES = native_loader.FULL_SUFFIXES +LORA_SUFFIXES = native_adapter.LORA_SUFFIXES +LOKR_SUFFIXES = native_adapter.LOKR_SUFFIXES +LOHA_SUFFIXES = native_adapter.LOHA_SUFFIXES +OFT_SUFFIXES = native_adapter.OFT_SUFFIXES +IA3_SUFFIXES = native_adapter.IA3_SUFFIXES +GLORA_SUFFIXES = native_adapter.GLORA_SUFFIXES +NORM_SUFFIXES = native_adapter.NORM_SUFFIXES +FULL_SUFFIXES = native_adapter.FULL_SUFFIXES -LORA_MARKERS = native_loader.LORA_MARKERS -LOKR_MARKERS = native_loader.LOKR_MARKERS -LOHA_MARKERS = native_loader.LOHA_MARKERS -OFT_MARKERS = native_loader.OFT_MARKERS -IA3_MARKERS = native_loader.IA3_MARKERS -GLORA_MARKERS = native_loader.GLORA_MARKERS -NORM_MARKERS = native_loader.NORM_MARKERS -FULL_MARKERS = native_loader.FULL_MARKERS +LORA_MARKERS = native_adapter.LORA_MARKERS +LOKR_MARKERS = native_adapter.LOKR_MARKERS +LOHA_MARKERS = native_adapter.LOHA_MARKERS +OFT_MARKERS = native_adapter.OFT_MARKERS +IA3_MARKERS = native_adapter.IA3_MARKERS +GLORA_MARKERS = native_adapter.GLORA_MARKERS +NORM_MARKERS = native_adapter.NORM_MARKERS +FULL_MARKERS = native_adapter.FULL_MARKERS -SUFFIX_NORMALIZE = native_loader.SUFFIX_NORMALIZE -BARE_DIFFUSERS_PREFIX_USED = native_loader.BARE_DIFFUSERS_PREFIX_USED -has_marker = native_loader.has_marker +SUFFIX_NORMALIZE = native_adapter.SUFFIX_NORMALIZE +BARE_DIFFUSERS_PREFIX_USED = native_adapter.BARE_DIFFUSERS_PREFIX_USED +has_marker = native_adapter.has_marker def parse_key(key, suffixes): - """Anima-bound :func:`native_loader.parse_key`.""" - return native_loader.parse_key(key, suffixes, prefixes=ANIMA_PREFIXES) + """Anima-bound :func:`native_adapter.parse_key`.""" + return native_adapter.parse_key(key, suffixes, prefixes=ANIMA_PREFIXES) def group_by_suffixes(state_dict, suffixes): - """Anima-bound :func:`native_loader.group_by_suffixes`.""" - return native_loader.group_by_suffixes(state_dict, suffixes, prefixes=ANIMA_PREFIXES) + """Anima-bound :func:`native_adapter.group_by_suffixes`.""" + return native_adapter.group_by_suffixes(state_dict, suffixes, prefixes=ANIMA_PREFIXES) # === Cosmos 2.0 path rename (transformer only) === @@ -174,7 +174,7 @@ def network_prefix_for(prefix_used): return "lora_transformer_" -# === Native loaders (thin wrappers over native_loader generics) === +# === Native loaders (thin wrappers over native_adapter generics) === _BIND_KWARGS = dict( resolve_targets=resolve_targets, @@ -185,40 +185,40 @@ _BIND_KWARGS = dict( def try_load_lora(name, network_on_disk, lora_scale): - return native_loader.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_lokr(name, network_on_disk, lora_scale): - return native_loader.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_loha(name, network_on_disk, lora_scale): - return native_loader.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_oft(name, network_on_disk, lora_scale): - return native_loader.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_ia3(name, network_on_disk, lora_scale): - return native_loader.try_load_ia3(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_ia3(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_glora(name, network_on_disk, lora_scale): - return native_loader.try_load_glora(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_glora(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_norm(name, network_on_disk, lora_scale): - return native_loader.try_load_norm(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_norm(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_full(name, network_on_disk, lora_scale): - return native_loader.try_load_full(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_full(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load(name, network_on_disk, lora_scale): """Run every Anima family loader, merge any that match.""" - return native_loader.try_load_chain( + return native_adapter.try_load_chain( name, network_on_disk, lora_scale, family_loaders=( try_load_lora, try_load_lokr, try_load_loha, try_load_oft, diff --git a/pipelines/chroma/chroma_lora.py b/pipelines/chroma/chroma_lora.py index 8354a98b5..965d5bf52 100644 --- a/pipelines/chroma/chroma_lora.py +++ b/pipelines/chroma/chroma_lora.py @@ -30,13 +30,13 @@ any ``_mod_lin`` / ``_modulation_lin`` keys land in ``unmapped``. LoRAs targeting the approximator pass through unchanged. """ -from modules.lora import native_loader -from modules.lora.native_loader import ChunkSpec +from modules.lora import native_adapter +from modules.lora.native_adapter import ChunkSpec # === Arch-specific prefix configuration === -KNOWN_PREFIXES = native_loader.KNOWN_PREFIXES_DEFAULT +KNOWN_PREFIXES = native_adapter.KNOWN_PREFIXES_DEFAULT BARE_FLUX_PREFIXES = ("double_blocks.", "single_blocks.") @@ -57,24 +57,24 @@ LINEAR1_DIMS = [3072, 3072, 3072, 12288] # === Re-exports for test/back-compat === -LORA_SUFFIXES = native_loader.LORA_SUFFIXES -LOKR_SUFFIXES = native_loader.LOKR_SUFFIXES -LOHA_SUFFIXES = native_loader.LOHA_SUFFIXES -OFT_SUFFIXES = native_loader.OFT_SUFFIXES +LORA_SUFFIXES = native_adapter.LORA_SUFFIXES +LOKR_SUFFIXES = native_adapter.LOKR_SUFFIXES +LOHA_SUFFIXES = native_adapter.LOHA_SUFFIXES +OFT_SUFFIXES = native_adapter.OFT_SUFFIXES -LORA_MARKERS = native_loader.LORA_MARKERS -LOKR_MARKERS = native_loader.LOKR_MARKERS -LOHA_MARKERS = native_loader.LOHA_MARKERS -OFT_MARKERS = native_loader.OFT_MARKERS +LORA_MARKERS = native_adapter.LORA_MARKERS +LOKR_MARKERS = native_adapter.LOKR_MARKERS +LOHA_MARKERS = native_adapter.LOHA_MARKERS +OFT_MARKERS = native_adapter.OFT_MARKERS -SUFFIX_NORMALIZE = native_loader.SUFFIX_NORMALIZE -BARE_DIFFUSERS_PREFIX_USED = native_loader.BARE_DIFFUSERS_PREFIX_USED -has_marker = native_loader.has_marker +SUFFIX_NORMALIZE = native_adapter.SUFFIX_NORMALIZE +BARE_DIFFUSERS_PREFIX_USED = native_adapter.BARE_DIFFUSERS_PREFIX_USED +has_marker = native_adapter.has_marker def parse_key(key, suffixes): - """Chroma-bound :func:`native_loader.parse_key`.""" - return native_loader.parse_key( + """Chroma-bound :func:`native_adapter.parse_key`.""" + return native_adapter.parse_key( key, suffixes, prefixes=KNOWN_PREFIXES, bare_prefixes=BARE_FLUX_PREFIXES, @@ -83,8 +83,8 @@ def parse_key(key, suffixes): def group_by_suffixes(state_dict, suffixes): - """Chroma-bound :func:`native_loader.group_by_suffixes`.""" - return native_loader.group_by_suffixes( + """Chroma-bound :func:`native_adapter.group_by_suffixes`.""" + return native_adapter.group_by_suffixes( state_dict, suffixes, prefixes=KNOWN_PREFIXES, bare_prefixes=BARE_FLUX_PREFIXES, @@ -202,7 +202,7 @@ def _split_single_linear1(block_idx): return targets -# === Native loaders (thin wrappers over native_loader generics) === +# === Native loaders (thin wrappers over native_adapter generics) === _BIND_KWARGS = dict( @@ -215,24 +215,24 @@ _BIND_KWARGS = dict( def try_load_lora(name, network_on_disk, lora_scale): - return native_loader.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_lokr(name, network_on_disk, lora_scale): - return native_loader.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_loha(name, network_on_disk, lora_scale): - return native_loader.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_oft(name, network_on_disk, lora_scale): - return native_loader.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load(name, network_on_disk, lora_scale): """Run every Chroma family loader, merge any that match.""" - return native_loader.try_load_chain( + return native_adapter.try_load_chain( name, network_on_disk, lora_scale, family_loaders=(try_load_lora, try_load_lokr, try_load_loha, try_load_oft), ) diff --git a/pipelines/ernie/ernie_lora.py b/pipelines/ernie/ernie_lora.py index c7ad8d497..e76d716aa 100644 --- a/pipelines/ernie/ernie_lora.py +++ b/pipelines/ernie/ernie_lora.py @@ -16,12 +16,12 @@ modules (no fused QKV) and ``ErnieImageFeedForward`` exposes ``gate_proj``, straight passthrough; no chunking, no renames, no dispatch table. """ -from modules.lora import native_loader +from modules.lora import native_adapter # === Arch-specific prefix configuration === -KNOWN_PREFIXES = native_loader.KNOWN_PREFIXES_DEFAULT +KNOWN_PREFIXES = native_adapter.KNOWN_PREFIXES_DEFAULT BARE_DIFFUSERS_PREFIXES = ( "layers.", "adaLN_modulation.", "final_norm.", "final_linear.", @@ -30,24 +30,24 @@ BARE_DIFFUSERS_PREFIXES = ( # === Re-exports for test/back-compat === -LORA_SUFFIXES = native_loader.LORA_SUFFIXES -LOKR_SUFFIXES = native_loader.LOKR_SUFFIXES -LOHA_SUFFIXES = native_loader.LOHA_SUFFIXES -OFT_SUFFIXES = native_loader.OFT_SUFFIXES +LORA_SUFFIXES = native_adapter.LORA_SUFFIXES +LOKR_SUFFIXES = native_adapter.LOKR_SUFFIXES +LOHA_SUFFIXES = native_adapter.LOHA_SUFFIXES +OFT_SUFFIXES = native_adapter.OFT_SUFFIXES -LORA_MARKERS = native_loader.LORA_MARKERS -LOKR_MARKERS = native_loader.LOKR_MARKERS -LOHA_MARKERS = native_loader.LOHA_MARKERS -OFT_MARKERS = native_loader.OFT_MARKERS +LORA_MARKERS = native_adapter.LORA_MARKERS +LOKR_MARKERS = native_adapter.LOKR_MARKERS +LOHA_MARKERS = native_adapter.LOHA_MARKERS +OFT_MARKERS = native_adapter.OFT_MARKERS -SUFFIX_NORMALIZE = native_loader.SUFFIX_NORMALIZE -BARE_DIFFUSERS_PREFIX_USED = native_loader.BARE_DIFFUSERS_PREFIX_USED -has_marker = native_loader.has_marker +SUFFIX_NORMALIZE = native_adapter.SUFFIX_NORMALIZE +BARE_DIFFUSERS_PREFIX_USED = native_adapter.BARE_DIFFUSERS_PREFIX_USED +has_marker = native_adapter.has_marker def parse_key(key, suffixes): - """ERNIE-bound :func:`native_loader.parse_key`.""" - return native_loader.parse_key( + """ERNIE-bound :func:`native_adapter.parse_key`.""" + return native_adapter.parse_key( key, suffixes, prefixes=KNOWN_PREFIXES, bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, @@ -55,8 +55,8 @@ def parse_key(key, suffixes): def group_by_suffixes(state_dict, suffixes): - """ERNIE-bound :func:`native_loader.group_by_suffixes`.""" - return native_loader.group_by_suffixes( + """ERNIE-bound :func:`native_adapter.group_by_suffixes`.""" + return native_adapter.group_by_suffixes( state_dict, suffixes, prefixes=KNOWN_PREFIXES, bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, @@ -75,7 +75,7 @@ def resolve_targets(prefix_used, base): return [] -# === Native loaders (thin wrappers over native_loader generics) === +# === Native loaders (thin wrappers over native_adapter generics) === _BIND_KWARGS = dict( @@ -87,24 +87,24 @@ _BIND_KWARGS = dict( def try_load_lora(name, network_on_disk, lora_scale): - return native_loader.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_lokr(name, network_on_disk, lora_scale): - return native_loader.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_loha(name, network_on_disk, lora_scale): - return native_loader.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_oft(name, network_on_disk, lora_scale): - return native_loader.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load(name, network_on_disk, lora_scale): """Run every ERNIE family loader, merge any that match.""" - return native_loader.try_load_chain( + return native_adapter.try_load_chain( name, network_on_disk, lora_scale, family_loaders=(try_load_lora, try_load_lokr, try_load_loha, try_load_oft), ) diff --git a/pipelines/flux/flux2_lora.py b/pipelines/flux/flux2_lora.py index 81994cf82..1c69a66b4 100644 --- a/pipelines/flux/flux2_lora.py +++ b/pipelines/flux/flux2_lora.py @@ -16,12 +16,12 @@ produced by ``Flux2Transformer2DModel.save_lora_adapter()``). Diffusers-PEFT BFL/kohya keys are mapped to diffusers paths via ``F2_SINGLE_MAP`` / ``F2_DOUBLE_MAP`` / ``F2_QKV_MAP``. Fused QKV in double_blocks emits three -Q/K/V targets each carrying a :class:`modules.lora.native_loader.ChunkSpec` +Q/K/V targets each carrying a :class:`modules.lora.native_adapter.ChunkSpec` that the generic loaders use to chunk the up-weight or instantiate the appropriate ``NetworkModule*Chunk`` variant. Per-family fused-QKV handling is inherited from -:mod:`modules.lora.native_loader`; see the loader-by-loader notes there. +:mod:`modules.lora.native_adapter`; see the loader-by-loader notes there. LyCORIS algorithm coverage relative to upstream ``KohakuBlueleaf/LyCORIS/lycoris/modules/``: @@ -49,13 +49,13 @@ to inject the ``diffusion_model.`` prefix for bare-BFL keys and bake kohya import os from modules.logger import log -from modules.lora import native_loader -from modules.lora.native_loader import ChunkSpec +from modules.lora import native_adapter +from modules.lora.native_adapter import ChunkSpec # === Arch-specific prefix configuration === -KNOWN_PREFIXES = native_loader.KNOWN_PREFIXES_DEFAULT + ("lycoris_",) +KNOWN_PREFIXES = native_adapter.KNOWN_PREFIXES_DEFAULT + ("lycoris_",) BARE_FLUX_PREFIXES = ( "single_blocks.", "double_blocks.", "img_in.", "txt_in.", @@ -108,34 +108,34 @@ KOHYA_SUFFIX_MAP = { # === Re-exports for backward compatibility === # The offline test suite addresses these via the flux2_lora module surface. -# Re-export rather than asking tests to import native_loader directly. +# Re-export rather than asking tests to import native_adapter directly. -LORA_SUFFIXES = native_loader.LORA_SUFFIXES -LOKR_SUFFIXES = native_loader.LOKR_SUFFIXES -LOHA_SUFFIXES = native_loader.LOHA_SUFFIXES -OFT_SUFFIXES = native_loader.OFT_SUFFIXES -IA3_SUFFIXES = native_loader.IA3_SUFFIXES -GLORA_SUFFIXES = native_loader.GLORA_SUFFIXES -NORM_SUFFIXES = native_loader.NORM_SUFFIXES -FULL_SUFFIXES = native_loader.FULL_SUFFIXES +LORA_SUFFIXES = native_adapter.LORA_SUFFIXES +LOKR_SUFFIXES = native_adapter.LOKR_SUFFIXES +LOHA_SUFFIXES = native_adapter.LOHA_SUFFIXES +OFT_SUFFIXES = native_adapter.OFT_SUFFIXES +IA3_SUFFIXES = native_adapter.IA3_SUFFIXES +GLORA_SUFFIXES = native_adapter.GLORA_SUFFIXES +NORM_SUFFIXES = native_adapter.NORM_SUFFIXES +FULL_SUFFIXES = native_adapter.FULL_SUFFIXES -LORA_MARKERS = native_loader.LORA_MARKERS -LOKR_MARKERS = native_loader.LOKR_MARKERS -LOHA_MARKERS = native_loader.LOHA_MARKERS -OFT_MARKERS = native_loader.OFT_MARKERS -IA3_MARKERS = native_loader.IA3_MARKERS -GLORA_MARKERS = native_loader.GLORA_MARKERS -NORM_MARKERS = native_loader.NORM_MARKERS -FULL_MARKERS = native_loader.FULL_MARKERS +LORA_MARKERS = native_adapter.LORA_MARKERS +LOKR_MARKERS = native_adapter.LOKR_MARKERS +LOHA_MARKERS = native_adapter.LOHA_MARKERS +OFT_MARKERS = native_adapter.OFT_MARKERS +IA3_MARKERS = native_adapter.IA3_MARKERS +GLORA_MARKERS = native_adapter.GLORA_MARKERS +NORM_MARKERS = native_adapter.NORM_MARKERS +FULL_MARKERS = native_adapter.FULL_MARKERS -SUFFIX_NORMALIZE = native_loader.SUFFIX_NORMALIZE -BARE_DIFFUSERS_PREFIX_USED = native_loader.BARE_DIFFUSERS_PREFIX_USED -has_marker = native_loader.has_marker +SUFFIX_NORMALIZE = native_adapter.SUFFIX_NORMALIZE +BARE_DIFFUSERS_PREFIX_USED = native_adapter.BARE_DIFFUSERS_PREFIX_USED +has_marker = native_adapter.has_marker def parse_key(key, suffixes): - """Flux2-bound :func:`native_loader.parse_key`. Returns ``(prefix_used, base, suffix)`` or ``None``.""" - return native_loader.parse_key( + """Flux2-bound :func:`native_adapter.parse_key`. Returns ``(prefix_used, base, suffix)`` or ``None``.""" + return native_adapter.parse_key( key, suffixes, prefixes=KNOWN_PREFIXES, bare_prefixes=BARE_FLUX_PREFIXES, @@ -144,8 +144,8 @@ def parse_key(key, suffixes): def group_by_suffixes(state_dict, suffixes): - """Flux2-bound :func:`native_loader.group_by_suffixes`.""" - return native_loader.group_by_suffixes( + """Flux2-bound :func:`native_adapter.group_by_suffixes`.""" + return native_adapter.group_by_suffixes( state_dict, suffixes, prefixes=KNOWN_PREFIXES, bare_prefixes=BARE_FLUX_PREFIXES, @@ -231,7 +231,7 @@ def _bfl_to_diffusers_targets(base): return targets -# === Native loaders (thin wrappers over native_loader generics) === +# === Native loaders (thin wrappers over native_adapter generics) === _BIND_KWARGS = dict( @@ -244,40 +244,40 @@ _BIND_KWARGS = dict( def try_load_lora(name, network_on_disk, lora_scale): - return native_loader.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_lokr(name, network_on_disk, lora_scale): - return native_loader.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_loha(name, network_on_disk, lora_scale): - return native_loader.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_oft(name, network_on_disk, lora_scale): - return native_loader.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_ia3(name, network_on_disk, lora_scale): - return native_loader.try_load_ia3(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_ia3(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_glora(name, network_on_disk, lora_scale): - return native_loader.try_load_glora(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_glora(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_norm(name, network_on_disk, lora_scale): - return native_loader.try_load_norm(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_norm(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_full(name, network_on_disk, lora_scale): - return native_loader.try_load_full(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_full(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load(name, network_on_disk, lora_scale): """Single dispatcher entry point: run every family loader, merge any that match.""" - return native_loader.try_load_chain( + return native_adapter.try_load_chain( name, network_on_disk, lora_scale, family_loaders=( try_load_lora, try_load_lokr, try_load_loha, try_load_oft, diff --git a/pipelines/z_image/zimage_lora.py b/pipelines/z_image/zimage_lora.py index 49c35f9d8..f3d39e536 100644 --- a/pipelines/z_image/zimage_lora.py +++ b/pipelines/z_image/zimage_lora.py @@ -22,37 +22,37 @@ OFT block structure is tied to the target module's ``out_features`` so a Q/K/V split is not a drop-in. """ -from modules.lora import native_loader -from modules.lora.native_loader import ChunkSpec +from modules.lora import native_adapter +from modules.lora.native_adapter import ChunkSpec # === Arch-specific prefix configuration === -KNOWN_PREFIXES = native_loader.KNOWN_PREFIXES_DEFAULT +KNOWN_PREFIXES = native_adapter.KNOWN_PREFIXES_DEFAULT BARE_DIFFUSERS_PREFIXES = ("layers.", "noise_refiner.", "context_refiner.") # === Re-exports for test/back-compat === -LORA_SUFFIXES = native_loader.LORA_SUFFIXES -LOKR_SUFFIXES = native_loader.LOKR_SUFFIXES -LOHA_SUFFIXES = native_loader.LOHA_SUFFIXES -OFT_SUFFIXES = native_loader.OFT_SUFFIXES +LORA_SUFFIXES = native_adapter.LORA_SUFFIXES +LOKR_SUFFIXES = native_adapter.LOKR_SUFFIXES +LOHA_SUFFIXES = native_adapter.LOHA_SUFFIXES +OFT_SUFFIXES = native_adapter.OFT_SUFFIXES -LORA_MARKERS = native_loader.LORA_MARKERS -LOKR_MARKERS = native_loader.LOKR_MARKERS -LOHA_MARKERS = native_loader.LOHA_MARKERS -OFT_MARKERS = native_loader.OFT_MARKERS +LORA_MARKERS = native_adapter.LORA_MARKERS +LOKR_MARKERS = native_adapter.LOKR_MARKERS +LOHA_MARKERS = native_adapter.LOHA_MARKERS +OFT_MARKERS = native_adapter.OFT_MARKERS -SUFFIX_NORMALIZE = native_loader.SUFFIX_NORMALIZE -BARE_DIFFUSERS_PREFIX_USED = native_loader.BARE_DIFFUSERS_PREFIX_USED -has_marker = native_loader.has_marker +SUFFIX_NORMALIZE = native_adapter.SUFFIX_NORMALIZE +BARE_DIFFUSERS_PREFIX_USED = native_adapter.BARE_DIFFUSERS_PREFIX_USED +has_marker = native_adapter.has_marker def parse_key(key, suffixes): - """Z-Image-bound :func:`native_loader.parse_key`.""" - return native_loader.parse_key( + """Z-Image-bound :func:`native_adapter.parse_key`.""" + return native_adapter.parse_key( key, suffixes, prefixes=KNOWN_PREFIXES, bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, @@ -60,8 +60,8 @@ def parse_key(key, suffixes): def group_by_suffixes(state_dict, suffixes): - """Z-Image-bound :func:`native_loader.group_by_suffixes`.""" - return native_loader.group_by_suffixes( + """Z-Image-bound :func:`native_adapter.group_by_suffixes`.""" + return native_adapter.group_by_suffixes( state_dict, suffixes, prefixes=KNOWN_PREFIXES, bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, @@ -127,7 +127,7 @@ def _underscore_to_diffusers_targets(base): return [(base, None)] -# === Native loaders (thin wrappers over native_loader generics) === +# === Native loaders (thin wrappers over native_adapter generics) === _BIND_KWARGS = dict( @@ -139,24 +139,24 @@ _BIND_KWARGS = dict( def try_load_lora(name, network_on_disk, lora_scale): - return native_loader.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_lokr(name, network_on_disk, lora_scale): - return native_loader.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_loha(name, network_on_disk, lora_scale): - return native_loader.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_oft(name, network_on_disk, lora_scale): - return native_loader.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) + return native_adapter.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load(name, network_on_disk, lora_scale): """Run every Z-Image family loader, merge any that match.""" - return native_loader.try_load_chain( + return native_adapter.try_load_chain( name, network_on_disk, lora_scale, family_loaders=(try_load_lora, try_load_lokr, try_load_loha, try_load_oft), ) diff --git a/test/test-anima-native-adapters.py b/test/test-anima-native-adapters.py index ad77366a1..f017512e4 100644 --- a/test/test-anima-native-adapters.py +++ b/test/test-anima-native-adapters.py @@ -5,10 +5,10 @@ Offline unit tests for Anima native adapter loaders. Anima is the only native arch with a multi-component network namespace: keys route into ``lora_transformer_*`` (Cosmos 2.0 DiT), ``lora_llm_adapter_*`` (a custom Qwen3-projection MLP), or ``lora_te_*`` (Qwen3 text encoder). Routing -is parameterized in ``modules.lora.native_loader`` via the ``network_prefix`` +is parameterized in ``modules.lora.native_adapter`` via the ``network_prefix`` callable that ``pipelines.anima.anima_lora`` supplies. -Covers the eight families exposed through native_loader's generics (LoRA, +Covers the eight families exposed through native_adapter's generics (LoRA, LoKR, LoHA, OFT, IA3, GLoRA, Norm, Full), focused on: - LoRA across all five recognized prefixes (BFL transformer / BFL llm_adapter / diff --git a/test/test-flux2-native-adapters.py b/test/test-flux2-native-adapters.py index 9400033a8..f7bd043b0 100644 --- a/test/test-flux2-native-adapters.py +++ b/test/test-flux2-native-adapters.py @@ -580,7 +580,7 @@ def test_parse_key_all_prefixes(): def test_resolve_targets_qkv_chunking(): - from modules.lora.native_loader import ChunkSpec + from modules.lora.native_adapter import ChunkSpec # Kohya double_blocks fused QKV → three chunks targeting Q/K/V. targets = F.resolve_targets('lora_unet_', 'double_blocks_0_img_attn_qkv') assert targets == [ From 08ad6196d06b7f46d0e8e35bc85512b35e566ced Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 25 May 2026 05:47:56 +0100 Subject: [PATCH 03/27] feat(pipelines): add native_transformer for single-file DiT loading Read .safetensors, strip prefix, partition siblings, run optional converter, from_config + load_state_dict + validate, dtype/quant/offload. Per-arch knobs via TransformerSpec(cls, subfolder, prefixes, converter, siblings, acceptable_missing, forbidden_markers). SiblingSpec covers bundled components like Anima's llm_adapter. 36 offline tests. --- pipelines/native_transformer.py | 459 +++++++++++++++++++ test/test-native-transformer.py | 761 ++++++++++++++++++++++++++++++++ 2 files changed, 1220 insertions(+) create mode 100644 pipelines/native_transformer.py create mode 100644 test/test-native-transformer.py diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py new file mode 100644 index 000000000..757fbbe4c --- /dev/null +++ b/pipelines/native_transformer.py @@ -0,0 +1,459 @@ +"""Generic native loader for DiT transformers and bundled sibling components. + +Loads a single-safetensors file into a diffusers (or custom) transformer class +when the user selects an override via the UNET dropdown (``shared.opts.sd_unet``). +Bypasses :func:`diffusers.loaders.FromOriginalModelMixin.from_single_file` so +sdnext owns prefix detection, optional sibling partitioning, dtype/quant/offload +handling, and explicit validation of missing/unexpected keys. + +The per-arch knobs are captured in :class:`TransformerSpec`. Arches register a +spec at import time via :func:`register`; arches without a registration get a +default spec that handles BFL-style ``model.diffusion_model.`` prefix stripping +and opportunistically picks up a diffusers converter from +``SINGLE_FILE_LOADABLE_CLASSES`` if the class has one. + +Algorithm: + +1. Read the safetensors state dict (.gguf and .pth are rejected up front). +2. Detect and strip one of the spec's known prefixes (raises on mixed prefixes). +3. Check forbidden markers (catches structural mismatches like Cosmos 1.0 keys + in a Cosmos 2.0 loader). +4. Partition off sibling component keys (e.g. Anima's bundled ``llm_adapter.*``). +5. Run the spec's converter if present (else pass through unchanged). +6. Fetch ``/config.json`` from the base repo, instantiate via + ``cls.from_config``, ``load_state_dict(strict=False)``, validate, dtype-cast, + quantize, and offload-place. +7. Repeat the build for each populated sibling (no converter, no quant by + default; sibling weights are read raw from the bundled file). + +Returns ``(transformer, sibling_components_dict)``. The dict is empty for +arches with no siblings; for Anima it carries the ``llm_adapter`` if the +community file bundled one, else the empty dict and the caller falls back to +loading the adapter from the base repo. +""" + +import os +import time +from dataclasses import dataclass, field +from typing import Callable + +import huggingface_hub as hf + +from modules import shared, devices, sd_models, model_quant, errors +from modules.logger import log + + +DEFAULT_PREFIXES: tuple[str, ...] = ( + "model.diffusion_model.", + "diffusion_model.", + "net.", +) +DEFAULT_ACCEPTABLE_MISSING: tuple[str, ...] = ( + "rope.", + "pos_embedder.", + "learnable_pos_embed.", +) + + +@dataclass(frozen=True) +class SiblingSpec: + """Describes a non-transformer component that may ship inline in the same + safetensors as the transformer (e.g. Anima's ``llm_adapter``). + + ``subfolder`` names the base repo subfolder holding the canonical config and + weights when the sibling is NOT bundled inline; ``inline_prefix`` is the + key prefix that identifies the sibling's weights within the bundled file + (after the transformer's prefix has already been stripped). + """ + + subfolder: str + inline_prefix: str + acceptable_missing: tuple[str, ...] = () + + +@dataclass(frozen=True) +class TransformerSpec: + """Per-arch configuration for the native loader. + + Most arches only need to override ``cls`` (and rely on the default + prefixes, no converter, no siblings). Arches with bundled-sibling + components (Anima) or unusual key conventions (custom converters, + Cosmos-style structural markers) override the relevant fields. + """ + + cls: type + subfolder: str = "transformer" + prefixes: tuple[str, ...] = DEFAULT_PREFIXES + converter: Callable[[dict], dict] | None = None + siblings: dict[str, SiblingSpec] = field(default_factory=dict) + acceptable_missing: tuple[str, ...] = DEFAULT_ACCEPTABLE_MISSING + forbidden_markers: tuple[tuple[str, str], ...] = () + + +REGISTRY: dict[type, TransformerSpec] = {} + + +def register(cls: type, spec: TransformerSpec | None = None) -> None: + """Register a transformer class with an explicit spec, or with the default + spec if ``spec`` is None. Idempotent: re-registering the same class + replaces the previous entry. + """ + if spec is None: + spec = TransformerSpec(cls=cls) + if spec.cls is not cls: + raise ValueError(f"register: spec.cls ({spec.cls.__name__}) does not match cls ({cls.__name__})") + REGISTRY[cls] = spec + + +def lookup(cls: type) -> TransformerSpec: + """Return the registered spec for ``cls``, or synthesize a default one. + + The synthesized default opportunistically pulls a converter from diffusers' + ``SINGLE_FILE_LOADABLE_CLASSES`` table if one exists for the class name and + is not a pass-through no-op lambda. + """ + if cls in REGISTRY: + return REGISTRY[cls] + converter = auto_pickup_converter(cls) + return TransformerSpec(cls=cls, converter=converter) + + +def auto_pickup_converter(cls: type) -> Callable[[dict], dict] | None: + """Pull a checkpoint converter from diffusers for ``cls`` when one exists. + + Skipped for the no-op identity lambda some classes register + (notably ``QwenImageTransformer2DModel``), because using it would silently + accept whatever key naming the file happens to have. + """ + try: + from diffusers.loaders.single_file_model import SINGLE_FILE_LOADABLE_CLASSES + except ImportError: + return None + entry = SINGLE_FILE_LOADABLE_CLASSES.get(cls.__name__) + if entry is None: + return None + fn = entry.get("checkpoint_mapping_fn") + if fn is None or is_noop_converter(fn): + return None + return fn + + +def is_noop_converter(fn: Callable) -> bool: + """Detect ``lambda checkpoint, **kwargs: checkpoint`` and equivalents. + + Strips inline ``#`` comments from the source line before inspecting the + body, so that diagnostic markers like ``# noqa`` on the lambda's source + line do not defeat the detection. + """ + try: + import inspect + src = inspect.getsource(fn).strip() + except (OSError, TypeError): + return False + if "lambda" not in src: + return False + if "#" in src: + src = src.split("#", 1)[0].rstrip() + body = src.split(":", 1)[-1].strip().rstrip(",").rstrip(")") + return body.endswith("checkpoint") + + +def resolve_path() -> str | None: + """Return the absolute path of the UNET dropdown selection, or None if + no selection is active or the file is unresolvable. + """ + sel = shared.opts.sd_unet + if sel is None or sel in ("Default", "None"): + return None + from modules import sd_unet + if sel not in list(sd_unet.unet_dict): + log.error(f'Load module: type=transformer file="{sel}" not found') + return None + path = sd_unet.unet_dict[sel] + if not os.path.exists(path): + log.error(f'Load module: type=transformer path="{path}" does not exist') + return None + return path + + +def load( + local_file: str, + repo_id: str, + spec: TransformerSpec, + diffusers_cfg: dict | None = None, + sibling_classes: dict[str, type] | None = None, +) -> tuple[object, dict[str, object]]: + """Load the transformer (and any bundled siblings) from ``local_file``. + + ``sibling_classes`` supplies the runtime class for each sibling named in + ``spec.siblings``. Required when a sibling has a dynamic class (e.g. + Anima's ``AnimaLLMAdapter`` is loaded from remote_code at runtime). + Missing sibling classes raise ``ValueError`` if the corresponding sibling + keys are present in the bundled file. + + Returns ``(transformer, siblings_dict)``. ``siblings_dict`` is keyed by + sibling name and is empty for non-sibling specs, or for sibling specs + whose keys are absent from the bundled file. + """ + if diffusers_cfg is None: + diffusers_cfg = {} + if sibling_classes is None: + sibling_classes = {} + + t0 = time.time() + if not local_file.lower().endswith(".safetensors"): + raise ValueError( + f"Load model: type={spec.cls.__name__} custom transformer requires .safetensors, " + f'got "{local_file}"' + ) + + _, quant_args = model_quant.get_dit_args( + diffusers_cfg, module="Model", device_map=True, allow_quant=True, + ) + quant_type = model_quant.get_quant_type(quant_args) + + state_dict = sd_models.read_state_dict(local_file, what="transformer") + state_dict = strip_prefix(state_dict, spec.prefixes, spec.cls.__name__) + check_forbidden_markers(state_dict, spec.forbidden_markers, spec.cls.__name__, local_file) + transformer_sd, sibling_sds = partition_siblings(state_dict, spec.siblings) + del state_dict + + sibling_counts = {name: len(sd) for name, sd in sibling_sds.items() if sd} + log.info( + f'Load model: type={spec.cls.__name__} custom="{os.path.basename(local_file)}" ' + f"transformer_keys={len(transformer_sd)} siblings={sibling_counts or '{}'}" + ) + + transformer_cfg = fetch_component_config(repo_id, spec.subfolder) + transformer = build_component( + component_name="transformer", + state_dict=transformer_sd, + config=transformer_cfg, + cls=spec.cls, + converter=spec.converter, + acceptable_missing=spec.acceptable_missing, + quant_args=quant_args, + quant_type=quant_type, + ) + del transformer_sd + devices.torch_gc() + + loaded_siblings: dict[str, object] = {} + for name, sibling_sd in sibling_sds.items(): + if not sibling_sd: + continue + sibling_spec = spec.siblings[name] + sibling_cls = sibling_classes.get(name) + if sibling_cls is None: + raise ValueError( + f"Load model: type={spec.cls.__name__} bundled sibling '{name}' present in " + f"file but no class was supplied via sibling_classes" + ) + sibling_cfg = fetch_component_config(repo_id, sibling_spec.subfolder) + loaded_siblings[name] = build_component( + component_name=name, + state_dict=sibling_sd, + config=sibling_cfg, + cls=sibling_cls, + converter=None, + acceptable_missing=sibling_spec.acceptable_missing, + quant_args={}, + quant_type=None, + ) + + sd_models.allow_post_quant = False + devices.torch_gc() + log.debug(f"Load model: type={spec.cls.__name__} native_transformer time={time.time() - t0:.2f}") + return transformer, loaded_siblings + + +def strip_prefix(state_dict: dict, prefixes: tuple[str, ...], type_name: str) -> dict: + """Detect and uniformly strip the most common known prefix from every key. + + Order matters: longer prefixes win over shorter ones with the same suffix + (e.g. ``model.diffusion_model.`` beats ``diffusion_model.``). If some keys + match the dominant prefix and others do not, raises ValueError because + mixed prefixes indicate a malformed file rather than a recoverable export + quirk. + """ + sorted_prefixes = sorted(prefixes, key=len, reverse=True) + counts: dict[str, int] = {} + seen = 0 + for key in state_dict: + for prefix in sorted_prefixes: + if key.startswith(prefix): + counts[prefix] = counts.get(prefix, 0) + 1 + seen += 1 + break + total = len(state_dict) + if seen == 0: + log.debug(f"Load model: type={type_name} native_transformer prefix=bare") + return state_dict + dominant = max(counts, key=counts.get) + if counts[dominant] != total: + raise ValueError( + f"Load model: type={type_name} native_transformer has mixed prefixes " + f"(total={total} {dominant}={counts[dominant]})" + ) + log.debug(f'Load model: type={type_name} native_transformer prefix="{dominant}"') + offset = len(dominant) + return {key[offset:]: value for key, value in state_dict.items()} + + +def check_forbidden_markers( + state_dict: dict, + forbidden_markers: tuple[tuple[str, str], ...], + type_name: str, + local_file: str, +) -> None: + """Raise if any forbidden marker key is present in the state_dict. + + Catches structural mismatches that pass prefix detection but indicate the + file is from an incompatible architecture variant (e.g. Cosmos 1.0 keys + showing up in a Cosmos 2.0 loader path). + """ + for marker, description in forbidden_markers: + if marker in state_dict: + raise ValueError( + f"Load model: type={type_name} native_transformer rejects " + f'"{os.path.basename(local_file)}" ({description}; marker key {marker!r})' + ) + + +def partition_siblings( + state_dict: dict, + siblings: dict[str, SiblingSpec], +) -> tuple[dict, dict[str, dict]]: + """Split state_dict into (transformer_sd, {sibling_name: sibling_sd}). + + Keys matching a sibling's ``inline_prefix`` go into that sibling's dict + with the prefix stripped; everything else stays in the transformer dict. + Sibling names with no matching keys still appear in the output dict but + map to an empty dict, so the caller can iterate uniformly. + """ + sibling_sds: dict[str, dict] = {name: {} for name in siblings} + transformer_sd: dict = {} + if not siblings: + return state_dict, sibling_sds + sibling_lookups = [(name, siblings[name].inline_prefix) for name in siblings] + for key, value in state_dict.items(): + matched = False + for name, prefix in sibling_lookups: + if key.startswith(prefix): + sibling_sds[name][key[len(prefix):]] = value + matched = True + break + if not matched: + transformer_sd[key] = value + return transformer_sd, sibling_sds + + +def fetch_component_config(repo_id: str, subfolder: str) -> dict: + """Download and parse ``/config.json`` from the base repo.""" + relative_path = f"{subfolder}/config.json" + try: + local = hf.hf_hub_download( + repo_id, filename=relative_path, cache_dir=shared.opts.diffusers_dir, + ) + except Exception as e: + raise RuntimeError( + f'Load model: native_transformer failed to download {relative_path} ' + f'from repo="{repo_id}": {e}' + ) from e + return shared.readfile(local, as_type="dict") + + +def build_component( + *, + component_name: str, + state_dict: dict, + config: dict, + cls: type, + converter: Callable[[dict], dict] | None, + acceptable_missing: tuple[str, ...], + quant_args: dict, + quant_type: str | None, +) -> object: + """Convert (if needed), instantiate, load weights, dtype-cast, quantize, + and offload-place a single component. Raises on any hard failure. + """ + try: + sd = converter(state_dict) if converter is not None else state_dict + component = cls.from_config(config) + missing, unexpected = component.load_state_dict(sd, strict=False) + validate_state_dict_load(component_name, missing, unexpected, acceptable_missing) + del sd + devices.torch_gc() + component = component.to(dtype=devices.dtype) + except Exception as e: + log.error(f"Load model: native_transformer {component_name} load failed: {e}") + errors.display(e, "Load") + raise + + if component_name == "transformer": + apply_quant(component, quant_type) + + if shared.opts.diffusers_offload_mode != "none": + sd_models.move_model(component, devices.cpu) + + if not hasattr(component, "quantization_config"): + if hasattr(component, "config") and hasattr(component.config, "quantization_config"): + component.quantization_config = component.config.quantization_config + elif quant_type is not None and quant_args.get("quantization_config") is not None: + component.quantization_config = quant_args.get("quantization_config") + return component + + +def validate_state_dict_load( + component_name: str, + missing: list[str], + unexpected: list[str], + acceptable_missing: tuple[str, ...], +) -> None: + """Raise ValueError if load_state_dict produced unexpected keys or + non-acceptable missing keys. Buffer-only missing keys matching the + ``acceptable_missing`` prefix list are logged at debug level and ignored. + """ + if unexpected: + sample = ", ".join(unexpected[:5]) + raise ValueError( + f"Load model: native_transformer {component_name} has {len(unexpected)} " + f"unexpected keys (sample: {sample})" + ) + hard_missing = [ + k for k in missing if not any(k.startswith(p) for p in acceptable_missing) + ] + if hard_missing: + sample = ", ".join(hard_missing[:5]) + raise ValueError( + f"Load model: native_transformer {component_name} missing " + f"{len(hard_missing)} required keys (sample: {sample})" + ) + if missing: + log.debug( + f"Load model: native_transformer {component_name} ignored " + f"{len(missing)} buffer-only missing keys" + ) + + +def apply_quant(transformer: object, quant_type: str | None) -> None: + """Apply SDNQ / layerwise quantization to the bare transformer. + + SDNQ 'pre' and 'auto' would normally route through ``quantization_config`` + at ``from_pretrained`` time; the native path bypasses that boundary, so + we call the per-module quant path directly. SDNQ 'post' and + ``layerwise_quantization`` go through ``do_post_load_quant`` as usual. + NVIDIAModelOptConfig (TRT) is not supported on this path. + """ + if quant_type == "NVIDIAModelOptConfig": + log.warning( + "Load model: native_transformer quant=TRT not supported on native path, skipping" + ) + elif quant_type == "SDNQConfig": + if shared.opts.sdnq_quantize_mode == "pre": + log.info( + "Load model: native_transformer quant=SDNQ pre-mode applied post-load " + "on native path" + ) + model_quant.sdnq_quantize_model(transformer, op="transformer") + model_quant.do_post_load_quant(transformer, allow=False) diff --git a/test/test-native-transformer.py b/test/test-native-transformer.py new file mode 100644 index 000000000..bac70ad6f --- /dev/null +++ b/test/test-native-transformer.py @@ -0,0 +1,761 @@ +#!/usr/bin/env python +""" +Offline unit tests for pipelines.native_transformer. + +Covers the pure helpers that own per-arch knob handling: + +- ``strip_prefix`` for single/multi prefix detection and mixed-prefix rejection +- ``partition_siblings`` for inline-sibling key partitioning +- ``check_forbidden_markers`` for structural-mismatch rejection +- ``is_noop_converter`` for diffusers no-op lambda detection +- ``validate_state_dict_load`` for unexpected / missing key handling +- ``register`` / ``lookup`` registry behavior and default spec synthesis +- ``auto_pickup_converter`` for diffusers ``SINGLE_FILE_LOADABLE_CLASSES`` integration +- ``TransformerSpec`` / ``SiblingSpec`` defaults + +Plus one end-to-end ``load`` test against a tiny mock module that exercises +the read -> strip -> convert -> from_config -> load_state_dict -> validate +pipeline without needing a real diffusers transformer or hf_hub_download. + +No running server required. + +Usage: + python test/test-native-transformer.py +""" + +import os +import sys +import tempfile + +import torch +import safetensors.torch + +script_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, script_dir) +os.chdir(script_dir) + +os.environ['SD_INSTALL_QUIET'] = '1' + +# Bootstrap cmd_args before any module that pulls in shared.py. +import modules.cmd_args # pylint: disable=wrong-import-position +import installer # pylint: disable=wrong-import-position +_orig_argv = sys.argv +sys.argv = [sys.argv[0]] +try: + modules.cmd_args.parse_args() +finally: + sys.argv = _orig_argv +installer.add_args(modules.cmd_args.parser) +modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([]) + +from modules.errors import log # pylint: disable=wrong-import-position +from pipelines import native_transformer as nt # pylint: disable=wrong-import-position + + +# ============================================================ +# Test infrastructure +# ============================================================ + +results: dict[str, dict] = {} + + +def category(name: str): + if name not in results: + results[name] = {'passed': 0, 'failed': 0, 'tests': []} + return name + + +def record(cat: str, passed: bool, name: str, detail: str = ''): + status = 'PASS' if passed else 'FAIL' + results[cat]['passed' if passed else 'failed'] += 1 + results[cat]['tests'].append((status, name)) + msg = f' {status}: {name}' + if detail: + msg += f' ({detail})' + if passed: + log.info(msg) + else: + log.error(msg) + + +def run_test(cat: str, fn): + name = fn.__name__ + try: + ok = fn() + if ok is False: + record(cat, False, name) + else: + record(cat, True, name) + except AssertionError as e: + record(cat, False, name, str(e)) + except Exception as e: # pylint: disable=broad-except + record(cat, False, name, f'exception: {e}') + import traceback + traceback.print_exc() + + +# ============================================================ +# strip_prefix +# ============================================================ + +def test_strip_prefix_bare_keys_pass_through(): + sd = {'layers.0.weight': 1, 'layers.0.bias': 2} + out = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test') + assert out == sd, 'bare keys must pass through unchanged' + + +def test_strip_prefix_dominant_single_variant(): + sd = {f'model.diffusion_model.layers.{i}.weight': i for i in range(10)} + out = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test') + assert all(k.startswith('layers.') for k in out) + assert len(out) == 10 + + +def test_strip_prefix_picks_longest_match_first(): + """``model.diffusion_model.`` must beat ``diffusion_model.`` when both match.""" + sd = { + 'model.diffusion_model.layers.0.weight': 1, + 'model.diffusion_model.layers.1.weight': 2, + } + out = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test') + # If shorter prefix matched, keys would start with 'model.' + assert 'layers.0.weight' in out + assert 'layers.1.weight' in out + assert not any(k.startswith('model.') for k in out) + + +def test_strip_prefix_mixed_prefixes_raises(): + sd = { + 'model.diffusion_model.layers.0.weight': 1, + 'net.layers.0.weight': 2, + } + try: + nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test') + raise AssertionError('expected ValueError') + except ValueError as e: + assert 'mixed prefixes' in str(e) + + +def test_strip_prefix_net_variant(): + sd = {'net.layers.0.weight': 1, 'net.layers.1.bias': 2} + out = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test') + assert set(out.keys()) == {'layers.0.weight', 'layers.1.bias'} + + +def test_strip_prefix_diffusion_model_variant(): + sd = {'diffusion_model.layers.0.weight': 1, 'diffusion_model.layers.1.bias': 2} + out = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test') + assert set(out.keys()) == {'layers.0.weight', 'layers.1.bias'} + + +def test_strip_prefix_custom_prefix_set(): + sd = {'lora_unet_blocks_0.weight': 1, 'lora_unet_blocks_1.weight': 2} + out = nt.strip_prefix(sd, ('lora_unet_',), 'Test') + assert set(out.keys()) == {'blocks_0.weight', 'blocks_1.weight'} + + +# ============================================================ +# partition_siblings +# ============================================================ + +def test_partition_siblings_empty_spec_returns_state_dict_unchanged(): + sd = {'a': 1, 'b': 2} + transformer_sd, siblings = nt.partition_siblings(sd, {}) + assert transformer_sd == sd + assert siblings == {} + + +def test_partition_siblings_no_matches_keeps_all_in_transformer(): + sd = {'layers.0.weight': 1, 'layers.1.weight': 2} + siblings_spec = {'llm_adapter': nt.SiblingSpec(subfolder='llm_adapter', inline_prefix='llm_adapter.')} + transformer_sd, siblings = nt.partition_siblings(sd, siblings_spec) + assert transformer_sd == sd + assert siblings == {'llm_adapter': {}} + + +def test_partition_siblings_single_sibling_split(): + sd = { + 'layers.0.weight': 'tx0', + 'layers.1.weight': 'tx1', + 'llm_adapter.input_proj.weight': 'ad0', + 'llm_adapter.output_proj.weight': 'ad1', + } + siblings_spec = {'llm_adapter': nt.SiblingSpec(subfolder='llm_adapter', inline_prefix='llm_adapter.')} + transformer_sd, siblings = nt.partition_siblings(sd, siblings_spec) + assert set(transformer_sd.keys()) == {'layers.0.weight', 'layers.1.weight'} + assert set(siblings['llm_adapter'].keys()) == {'input_proj.weight', 'output_proj.weight'} + assert siblings['llm_adapter']['input_proj.weight'] == 'ad0' + + +def test_partition_siblings_multiple_siblings(): + sd = { + 'layers.0.weight': 'tx', + 'sibling_a.x.weight': 'a0', + 'sibling_b.y.weight': 'b0', + 'sibling_b.z.weight': 'b1', + } + siblings_spec = { + 'sibling_a': nt.SiblingSpec(subfolder='a', inline_prefix='sibling_a.'), + 'sibling_b': nt.SiblingSpec(subfolder='b', inline_prefix='sibling_b.'), + } + transformer_sd, siblings = nt.partition_siblings(sd, siblings_spec) + assert list(transformer_sd.keys()) == ['layers.0.weight'] + assert set(siblings['sibling_a'].keys()) == {'x.weight'} + assert set(siblings['sibling_b'].keys()) == {'y.weight', 'z.weight'} + + +# ============================================================ +# check_forbidden_markers +# ============================================================ + +def test_forbidden_markers_passes_when_absent(): + sd = {'layers.0.weight': 1} + markers = (('legacy.marker.weight', 'old format'),) + nt.check_forbidden_markers(sd, markers, 'Test', '/tmp/x.safetensors') + # no exception = pass + + +def test_forbidden_markers_raises_when_present(): + sd = {'layers.0.weight': 1, 'legacy.marker.weight': 2} + markers = (('legacy.marker.weight', 'old Cosmos 1.0 structure'),) + try: + nt.check_forbidden_markers(sd, markers, 'Test', '/tmp/x.safetensors') + raise AssertionError('expected ValueError') + except ValueError as e: + msg = str(e) + assert 'old Cosmos 1.0 structure' in msg + assert 'legacy.marker.weight' in msg + + +def test_forbidden_markers_empty_tuple_no_op(): + sd = {'layers.0.weight': 1} + nt.check_forbidden_markers(sd, (), 'Test', '/tmp/x.safetensors') + + +# ============================================================ +# is_noop_converter +# ============================================================ + +def test_noop_converter_identity_lambda(): + fn = lambda checkpoint, **kwargs: checkpoint # pylint: disable=unnecessary-lambda-assignment + assert nt.is_noop_converter(fn) is True + + +def test_noop_converter_real_function(): + def real(checkpoint, **kwargs): # pylint: disable=unused-argument + return {k.replace('a.', 'b.'): v for k, v in checkpoint.items()} + assert nt.is_noop_converter(real) is False + + +def test_noop_converter_lambda_with_modification(): + fn = lambda checkpoint, **kwargs: {k: v.float() for k, v in checkpoint.items()} # pylint: disable=unnecessary-lambda-assignment + assert nt.is_noop_converter(fn) is False + + +# ============================================================ +# validate_state_dict_load +# ============================================================ + +def test_validate_accepts_buffer_only_missing(): + nt.validate_state_dict_load( + component_name='transformer', + missing=['rope.freqs', 'pos_embedder.pos'], + unexpected=[], + acceptable_missing=('rope.', 'pos_embedder.'), + ) + + +def test_validate_rejects_unexpected(): + try: + nt.validate_state_dict_load( + component_name='transformer', + missing=[], + unexpected=['some.junk.weight'], + acceptable_missing=(), + ) + raise AssertionError('expected ValueError') + except ValueError as e: + assert 'unexpected' in str(e) + assert 'some.junk.weight' in str(e) + + +def test_validate_rejects_hard_missing(): + try: + nt.validate_state_dict_load( + component_name='transformer', + missing=['layers.0.weight', 'rope.freqs'], + unexpected=[], + acceptable_missing=('rope.',), + ) + raise AssertionError('expected ValueError') + except ValueError as e: + msg = str(e) + assert 'missing' in msg + assert 'layers.0.weight' in msg + # Buffer-only missing must not show up in the hard-missing list + assert msg.count('rope.freqs') == 0 + + +def test_validate_empty_passes(): + nt.validate_state_dict_load( + component_name='transformer', + missing=[], + unexpected=[], + acceptable_missing=(), + ) + + +# ============================================================ +# register / lookup +# ============================================================ + +class FakeTransformer: + """Minimal stand-in for a diffusers transformer class.""" + + +class FakeTransformer2: + """Second stand-in for register/lookup tests.""" + + +def test_register_with_explicit_spec(): + nt.REGISTRY.clear() + spec = nt.TransformerSpec(cls=FakeTransformer, subfolder='custom_sub') + nt.register(FakeTransformer, spec) + assert nt.lookup(FakeTransformer) is spec + + +def test_register_with_default_spec(): + nt.REGISTRY.clear() + nt.register(FakeTransformer) + spec = nt.lookup(FakeTransformer) + assert spec.cls is FakeTransformer + assert spec.subfolder == 'transformer' + assert spec.prefixes == nt.DEFAULT_PREFIXES + assert spec.converter is None + assert spec.siblings == {} + + +def test_register_idempotent_replaces(): + nt.REGISTRY.clear() + spec1 = nt.TransformerSpec(cls=FakeTransformer, subfolder='sub_one') + spec2 = nt.TransformerSpec(cls=FakeTransformer, subfolder='sub_two') + nt.register(FakeTransformer, spec1) + nt.register(FakeTransformer, spec2) + assert nt.lookup(FakeTransformer) is spec2 + + +def test_register_rejects_mismatched_cls(): + try: + nt.register(FakeTransformer, nt.TransformerSpec(cls=FakeTransformer2)) + raise AssertionError('expected ValueError') + except ValueError as e: + assert 'does not match' in str(e) + + +def test_lookup_synthesizes_default_for_unregistered(): + nt.REGISTRY.clear() + spec = nt.lookup(FakeTransformer) + assert spec.cls is FakeTransformer + assert spec.subfolder == 'transformer' # default + assert spec.converter is None # FakeTransformer has no diffusers entry + + +# ============================================================ +# auto_pickup_converter +# ============================================================ + +def test_auto_pickup_returns_none_for_unknown_class(): + assert nt.auto_pickup_converter(FakeTransformer) is None + + +def test_auto_pickup_returns_real_diffusers_converter(): + import diffusers + fn = nt.auto_pickup_converter(diffusers.FluxTransformer2DModel) + assert fn is not None + assert callable(fn) + assert fn.__name__ == 'convert_flux_transformer_checkpoint_to_diffusers' + + +def test_auto_pickup_skips_noop_converter_qwen(): + """QwenImageTransformer2DModel registers a no-op lambda in diffusers; + auto_pickup_converter must return None so the spec falls back to no + converter (the user-registered spec can override with a real converter).""" + import diffusers + assert nt.auto_pickup_converter(diffusers.QwenImageTransformer2DModel) is None + + +# ============================================================ +# TransformerSpec / SiblingSpec defaults +# ============================================================ + +def test_transformer_spec_defaults(): + spec = nt.TransformerSpec(cls=FakeTransformer) + assert spec.subfolder == 'transformer' + assert spec.prefixes == ('model.diffusion_model.', 'diffusion_model.', 'net.') + assert spec.converter is None + assert spec.siblings == {} + assert spec.acceptable_missing == ('rope.', 'pos_embedder.', 'learnable_pos_embed.') + assert spec.forbidden_markers == () + + +def test_sibling_spec_defaults(): + spec = nt.SiblingSpec(subfolder='llm_adapter', inline_prefix='llm_adapter.') + assert spec.subfolder == 'llm_adapter' + assert spec.inline_prefix == 'llm_adapter.' + assert spec.acceptable_missing == () + + +def test_transformer_spec_is_frozen(): + spec = nt.TransformerSpec(cls=FakeTransformer) + try: + spec.subfolder = 'changed' # type: ignore[misc] + except Exception as e: # pylint: disable=broad-except + assert 'FrozenInstanceError' in type(e).__name__ or 'frozen' in str(e).lower() + return + raise AssertionError('expected FrozenInstanceError') + + +# ============================================================ +# Integration: end-to-end load() with a tiny mock module +# ============================================================ +# We sidestep diffusers + hf_hub_download by patching: +# - ``fetch_component_config`` to return a hand-rolled config dict +# - ``model_quant.get_dit_args`` / ``model_quant.get_quant_type`` / quant +# application to no-ops (we only want to test the load path itself). +# The mock cls is a torch.nn.Module subclass whose ``from_config`` constructs +# a fresh module of the expected shape; ``load_state_dict`` is the standard +# PyTorch method. + +class MockMiniTransformer(torch.nn.Module): + """Tiny stand-in: linear in -> linear out, plus a nested rope sub-module + holding a buffer the trainer state dict won't carry. Nested mirrors how + real DiTs structure rope / pos_embedder buffers.""" + + @classmethod + def from_config(cls, config: dict) -> 'MockMiniTransformer': + return cls(dim=config['dim']) + + def __init__(self, dim: int): + super().__init__() + self.in_proj = torch.nn.Linear(dim, dim) + self.out_proj = torch.nn.Linear(dim, dim) + self.rope = torch.nn.Module() + self.rope.register_buffer('freqs', torch.zeros(dim)) + + +def write_fixture(state_dict_keys: dict, fd: int, path: str) -> str: + os.close(fd) + safetensors.torch.save_file(state_dict_keys, path) + return path + + +def test_load_end_to_end_with_bfl_prefix_no_converter(): + """Exercise the full load pipeline: read .safetensors, strip prefix, + no converter, instantiate via from_config, load weights, validate. + """ + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + dim = 8 + # Save with model.diffusion_model. prefix; in_proj.* and out_proj.* + # are the real weights the mock cls expects after the strip. + raw = { + 'model.diffusion_model.in_proj.weight': torch.randn(dim, dim), + 'model.diffusion_model.in_proj.bias': torch.zeros(dim), + 'model.diffusion_model.out_proj.weight': torch.randn(dim, dim), + 'model.diffusion_model.out_proj.bias': torch.zeros(dim), + } + write_fixture(raw, fd, path) + + # Patch fetch_component_config to return our hand-rolled config. + orig_fetch = nt.fetch_component_config + nt.fetch_component_config = lambda repo, sub: {'dim': dim} + + # Patch quant helpers (we only care about the load path). + from modules import model_quant + orig_get_dit = model_quant.get_dit_args + orig_get_qtype = model_quant.get_quant_type + orig_do_post = model_quant.do_post_load_quant + model_quant.get_dit_args = lambda *a, **k: ({}, {}) + model_quant.get_quant_type = lambda *a, **k: None + model_quant.do_post_load_quant = lambda *a, **k: None + + try: + spec = nt.TransformerSpec(cls=MockMiniTransformer) + transformer, siblings = nt.load( + local_file=path, + repo_id='fake/repo', + spec=spec, + diffusers_cfg={}, + ) + finally: + nt.fetch_component_config = orig_fetch + model_quant.get_dit_args = orig_get_dit + model_quant.get_quant_type = orig_get_qtype + model_quant.do_post_load_quant = orig_do_post + + assert isinstance(transformer, MockMiniTransformer) + assert transformer.in_proj.weight.shape == (dim, dim) + # Weights from the fixture should match what was loaded. + loaded_in_w = transformer.in_proj.weight.detach().cpu() + fixture_in_w = raw['model.diffusion_model.in_proj.weight'].to(loaded_in_w.dtype) + assert torch.allclose(loaded_in_w, fixture_in_w) + assert siblings == {} + finally: + if os.path.exists(path): + os.unlink(path) + + +def test_load_end_to_end_with_sibling_partition(): + """Bundled-sibling case: file carries both transformer and sibling weights, + sibling_classes supplies the runtime sibling class, partition routes each + half into its target.""" + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + dim = 8 + sibling_dim = 4 + raw = { + # Transformer half (after strip). + 'model.diffusion_model.in_proj.weight': torch.randn(dim, dim), + 'model.diffusion_model.in_proj.bias': torch.zeros(dim), + 'model.diffusion_model.out_proj.weight': torch.randn(dim, dim), + 'model.diffusion_model.out_proj.bias': torch.zeros(dim), + # Sibling half (after strip + sibling partition). + 'model.diffusion_model.sibling.in_proj.weight': torch.randn(sibling_dim, sibling_dim), + 'model.diffusion_model.sibling.in_proj.bias': torch.zeros(sibling_dim), + 'model.diffusion_model.sibling.out_proj.weight': torch.randn(sibling_dim, sibling_dim), + 'model.diffusion_model.sibling.out_proj.bias': torch.zeros(sibling_dim), + } + write_fixture(raw, fd, path) + + orig_fetch = nt.fetch_component_config + + def patched_fetch(_repo, sub): + return {'dim': dim if sub == 'transformer' else sibling_dim} + + nt.fetch_component_config = patched_fetch + + from modules import model_quant + orig_get_dit = model_quant.get_dit_args + orig_get_qtype = model_quant.get_quant_type + orig_do_post = model_quant.do_post_load_quant + model_quant.get_dit_args = lambda *a, **k: ({}, {}) + model_quant.get_quant_type = lambda *a, **k: None + model_quant.do_post_load_quant = lambda *a, **k: None + + try: + spec = nt.TransformerSpec( + cls=MockMiniTransformer, + siblings={ + 'sibling': nt.SiblingSpec( + subfolder='sibling', + inline_prefix='sibling.', + acceptable_missing=('rope.',), + ), + }, + ) + transformer, siblings = nt.load( + local_file=path, + repo_id='fake/repo', + spec=spec, + diffusers_cfg={}, + sibling_classes={'sibling': MockMiniTransformer}, + ) + finally: + nt.fetch_component_config = orig_fetch + model_quant.get_dit_args = orig_get_dit + model_quant.get_quant_type = orig_get_qtype + model_quant.do_post_load_quant = orig_do_post + + assert isinstance(transformer, MockMiniTransformer) + assert transformer.in_proj.weight.shape == (dim, dim) + assert 'sibling' in siblings + assert isinstance(siblings['sibling'], MockMiniTransformer) + assert siblings['sibling'].in_proj.weight.shape == (sibling_dim, sibling_dim) + finally: + if os.path.exists(path): + os.unlink(path) + + +def test_load_raises_on_missing_sibling_class(): + """Sibling keys present in file but caller forgot to supply the class.""" + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + dim = 8 + sibling_dim = 4 + raw = { + 'model.diffusion_model.in_proj.weight': torch.randn(dim, dim), + 'model.diffusion_model.in_proj.bias': torch.zeros(dim), + 'model.diffusion_model.out_proj.weight': torch.randn(dim, dim), + 'model.diffusion_model.out_proj.bias': torch.zeros(dim), + 'model.diffusion_model.sibling.in_proj.weight': torch.randn(sibling_dim, sibling_dim), + } + write_fixture(raw, fd, path) + + orig_fetch = nt.fetch_component_config + nt.fetch_component_config = lambda repo, sub: {'dim': dim} + + from modules import model_quant + orig_get_dit = model_quant.get_dit_args + orig_get_qtype = model_quant.get_quant_type + orig_do_post = model_quant.do_post_load_quant + model_quant.get_dit_args = lambda *a, **k: ({}, {}) + model_quant.get_quant_type = lambda *a, **k: None + model_quant.do_post_load_quant = lambda *a, **k: None + + try: + spec = nt.TransformerSpec( + cls=MockMiniTransformer, + siblings={'sibling': nt.SiblingSpec(subfolder='s', inline_prefix='sibling.')}, + ) + raised = False + try: + nt.load( + local_file=path, + repo_id='fake/repo', + spec=spec, + diffusers_cfg={}, + sibling_classes={}, # missing! + ) + except ValueError as e: + raised = True + assert "'sibling'" in str(e) + assert 'sibling_classes' in str(e) + assert raised, 'expected ValueError' + finally: + nt.fetch_component_config = orig_fetch + model_quant.get_dit_args = orig_get_dit + model_quant.get_quant_type = orig_get_qtype + model_quant.do_post_load_quant = orig_do_post + finally: + if os.path.exists(path): + os.unlink(path) + + +def test_load_rejects_non_safetensors(): + spec = nt.TransformerSpec(cls=MockMiniTransformer) + try: + nt.load( + local_file='/tmp/some.gguf', + repo_id='fake/repo', + spec=spec, + diffusers_cfg={}, + ) + raise AssertionError('expected ValueError') + except ValueError as e: + assert '.safetensors' in str(e) + + +# ============================================================ +# Run +# ============================================================ + +def run_all(): + log.warning('=== strip_prefix ===') + cat = category('strip') + for fn in [ + test_strip_prefix_bare_keys_pass_through, + test_strip_prefix_dominant_single_variant, + test_strip_prefix_picks_longest_match_first, + test_strip_prefix_mixed_prefixes_raises, + test_strip_prefix_net_variant, + test_strip_prefix_diffusion_model_variant, + test_strip_prefix_custom_prefix_set, + ]: + run_test(cat, fn) + + log.warning('=== partition_siblings ===') + cat = category('partition') + for fn in [ + test_partition_siblings_empty_spec_returns_state_dict_unchanged, + test_partition_siblings_no_matches_keeps_all_in_transformer, + test_partition_siblings_single_sibling_split, + test_partition_siblings_multiple_siblings, + ]: + run_test(cat, fn) + + log.warning('=== forbidden_markers ===') + cat = category('forbidden') + for fn in [ + test_forbidden_markers_passes_when_absent, + test_forbidden_markers_raises_when_present, + test_forbidden_markers_empty_tuple_no_op, + ]: + run_test(cat, fn) + + log.warning('=== noop_converter detection ===') + cat = category('noop') + for fn in [ + test_noop_converter_identity_lambda, + test_noop_converter_real_function, + test_noop_converter_lambda_with_modification, + ]: + run_test(cat, fn) + + log.warning('=== validate_state_dict_load ===') + cat = category('validate') + for fn in [ + test_validate_accepts_buffer_only_missing, + test_validate_rejects_unexpected, + test_validate_rejects_hard_missing, + test_validate_empty_passes, + ]: + run_test(cat, fn) + + log.warning('=== register / lookup ===') + cat = category('registry') + for fn in [ + test_register_with_explicit_spec, + test_register_with_default_spec, + test_register_idempotent_replaces, + test_register_rejects_mismatched_cls, + test_lookup_synthesizes_default_for_unregistered, + ]: + run_test(cat, fn) + + log.warning('=== auto_pickup_converter ===') + cat = category('autopickup') + for fn in [ + test_auto_pickup_returns_none_for_unknown_class, + test_auto_pickup_returns_real_diffusers_converter, + test_auto_pickup_skips_noop_converter_qwen, + ]: + run_test(cat, fn) + + log.warning('=== TransformerSpec / SiblingSpec ===') + cat = category('specs') + for fn in [ + test_transformer_spec_defaults, + test_sibling_spec_defaults, + test_transformer_spec_is_frozen, + ]: + run_test(cat, fn) + + log.warning('=== end-to-end load ===') + cat = category('load') + for fn in [ + test_load_end_to_end_with_bfl_prefix_no_converter, + test_load_end_to_end_with_sibling_partition, + test_load_raises_on_missing_sibling_class, + test_load_rejects_non_safetensors, + ]: + run_test(cat, fn) + + log.warning('=== Results ===') + total_passed = 0 + total_failed = 0 + for cat_name, info in results.items(): + ok = info['failed'] == 0 + status = 'PASS' if ok else 'FAIL' + log.info(f" {cat_name}: {info['passed']} passed, {info['failed']} failed [{status}]") + total_passed += info['passed'] + total_failed += info['failed'] + log.warning(f'Total: {total_passed} passed, {total_failed} failed') + return total_failed == 0 + + +if __name__ == '__main__': + import time + t0 = time.time() + ok = run_all() + log.warning(f'Total time: {time.time() - t0:.2f}s') + sys.exit(0 if ok else 1) From cec6d0dce523b64c03fa886e497744ac6e5880ec Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 25 May 2026 05:47:56 +0100 Subject: [PATCH 04/27] refactor(anima): port anima_transformer to native_transformer 217-line bespoke loader collapses to a 40-line ANIMA_SPEC in pipelines/anima/__init__.py (Cosmos converter + llm_adapter sibling + Cosmos 1.0 forbidden marker). Drop the class-keyed REGISTRY: Anima and raw Cosmos share CosmosTransformer3DModel but need different specs. Specs pass via explicit native_spec= kwarg; make_default_spec(cls) covers the auto-converter case. --- pipelines/anima/__init__.py | 40 +++++ pipelines/anima/anima_transformer.py | 217 --------------------------- pipelines/model_anima.py | 38 ++--- pipelines/native_transformer.py | 47 +++--- test/test-native-transformer.py | 69 +++------ 5 files changed, 94 insertions(+), 317 deletions(-) delete mode 100644 pipelines/anima/anima_transformer.py diff --git a/pipelines/anima/__init__.py b/pipelines/anima/__init__.py index e69de29bb..6d7953090 100644 --- a/pipelines/anima/__init__.py +++ b/pipelines/anima/__init__.py @@ -0,0 +1,40 @@ +"""Anima pipeline package. + +Exports :data:`ANIMA_SPEC` for use by :mod:`pipelines.model_anima` together +with :mod:`pipelines.native_transformer`. The spec captures the Anima-specific +knobs that differ from the native-loader defaults: + +- The bundled ``llm_adapter`` sibling: Anima community files frequently inline + the custom AnimaLLMAdapter weights in the same safetensors as the + transformer. The resolved adapter class is supplied at load time via + ``sibling_classes`` because AnimaLLMAdapter is loaded dynamically through + ``trust_remote_code`` and is not available at import time. +- Cosmos 1.0 structural marker: any community file whose state dict contains + a Cosmos 1.0 nested key (``net.blocks.block1.*``) is rejected with a clear + error since Anima is Cosmos 2.0 only. +- All other knobs (prefixes, ``acceptable_missing`` buffers) use the defaults + from :mod:`pipelines.native_transformer`. +""" + +import diffusers +from diffusers.loaders.single_file_utils import convert_cosmos_transformer_checkpoint_to_diffusers + +from pipelines.native_transformer import TransformerSpec, SiblingSpec + + +ANIMA_SPEC = TransformerSpec( + cls=diffusers.CosmosTransformer3DModel, + converter=convert_cosmos_transformer_checkpoint_to_diffusers, + siblings={ + 'llm_adapter': SiblingSpec( + subfolder='llm_adapter', + inline_prefix='llm_adapter.', + ), + }, + forbidden_markers=( + ( + 'net.blocks.block1.blocks.0.block.attn.to_q.0.weight', + 'unsupported Cosmos 1.0 structure', + ), + ), +) diff --git a/pipelines/anima/anima_transformer.py b/pipelines/anima/anima_transformer.py deleted file mode 100644 index 06692dc4e..000000000 --- a/pipelines/anima/anima_transformer.py +++ /dev/null @@ -1,217 +0,0 @@ -"""Anima custom-transformer loader. - -Called from :func:`pipelines.model_anima.load_anima` when the user has selected -a transformer file via the UNET dropdown (``shared.opts.sd_unet``). Reads the -safetensors directly, strips the BFL-style prefix, splits off the bundled -``llm_adapter.*`` keys, and routes the two halves into the diffusers -``CosmosTransformer3DModel`` and the remote ``AnimaLLMAdapter`` respectively. - -The transformer half is run through diffusers' -``convert_cosmos_transformer_checkpoint_to_diffusers`` (Cosmos 2.0 branch), -whose rename table covers Anima's native key fragments exactly, so the -converted state dict drops cleanly into ``CosmosTransformer3DModel`` with no -ad-hoc renames needed here. The adapter half matches the base repo's -``llm_adapter/diffusion_pytorch_model.safetensors`` exactly, so it loads -as-is. - -Supported input formats (safetensors only; GGUF and .pth are rejected early): - -- Bare BFL keys: ``blocks.0.self_attn.q_proj.weight`` (e.g. ``rdbtAnima_v027``) -- ``model.diffusion_model.`` prefix (e.g. ``animaika_v35``) -- ``diffusion_model.`` prefix (ComfyUI-style export) -- ``net.`` prefix (NVIDIA/Cosmos native export, e.g. ``animayume_v04``) - -Quantization: SDNQ (pre/post/auto) and ``layerwise_quantization`` are honored. -SDNQ pre-mode is applied post-load here because this path bypasses -``from_pretrained``, where ``quantization_config`` normally takes effect. -TensorRT (``NVIDIAModelOptConfig``) is not supported and is skipped with a -warning. GGUF would require a separate converter and is not supported. -""" - -import os -import time -import diffusers -import huggingface_hub as hf -from modules import shared, devices, sd_models, model_quant, errors -from modules.logger import log - - -KNOWN_PREFIXES = ("model.diffusion_model.", "diffusion_model.", "net.") -ADAPTER_PREFIX = "llm_adapter." -COSMOS_1_MARKER = "net.blocks.block1.blocks.0.block.attn.to_q.0.weight" - -# Buffer keys that CosmosTransformer3DModel creates at __init__ time and do -# not appear in trainer state dicts. Acceptable in the "missing" set. -ACCEPTABLE_MISSING = ("rope.", "pos_embedder.", "learnable_pos_embed.") - - -def load_custom_transformer(repo_id, local_file, diffusers_load_config, adapter_cls): - """Load a custom Anima transformer (and optional bundled adapter) from a safetensors file. - - Returns ``(transformer, llm_adapter_or_none)``. If the file does not bundle - an adapter, the second element is ``None`` and the caller should fall back - to the base repo's adapter via ``AnimaLLMAdapter.from_pretrained``. - Raises on any hard failure (prefix mix, shape mismatch, missing configs). - """ - t0 = time.time() - - if not local_file.lower().endswith('.safetensors'): - raise ValueError(f'Load model: type=Anima custom transformer requires .safetensors, got "{local_file}"') - - # from_config + load_state_dict does not consume load_args (device_map, - # torch_dtype, etc.); dtype is applied via explicit .to() below. Only - # quant_type is read from this call. - _, quant_args = model_quant.get_dit_args( - diffusers_load_config, module='Model', device_map=True, allow_quant=True, - ) - quant_type = model_quant.get_quant_type(quant_args) - - transformer_cfg = fetch_component_config(repo_id, 'transformer/config.json') - adapter_cfg = fetch_component_config(repo_id, 'llm_adapter/config.json') - - state_dict = sd_models.read_state_dict(local_file, what='transformer') - state_dict = strip_prefix(state_dict) - transformer_sd, adapter_sd = partition_adapter(state_dict) - del state_dict - - if COSMOS_1_MARKER in transformer_sd: - raise ValueError(f'Load model: type=Anima custom transformer has unsupported Cosmos 1.0 structure (file="{local_file}")') - - log.info(f'Load model: type=Anima custom="{os.path.basename(local_file)}" transformer_keys={len(transformer_sd)} adapter_keys={len(adapter_sd)}') - - transformer = build_transformer(transformer_sd, transformer_cfg, quant_args, quant_type) - del transformer_sd - devices.torch_gc() - - if adapter_sd: - llm_adapter = build_adapter(adapter_sd, adapter_cfg, adapter_cls) - else: - log.info('Load model: type=Anima custom transformer has no bundled adapter, caller will load from base repo') - llm_adapter = None - - sd_models.allow_post_quant = False # transformer already quantized above - devices.torch_gc() - log.debug(f'Load model: type=Anima custom transformer time={time.time()-t0:.2f}') - return transformer, llm_adapter - - -def fetch_component_config(repo_id, relative_path): - """Download and parse a component config.json from the base repo.""" - try: - local = hf.hf_hub_download(repo_id, filename=relative_path, cache_dir=shared.opts.diffusers_dir) - except Exception as e: - raise RuntimeError(f'Load model: type=Anima failed to download {relative_path} from repo="{repo_id}": {e}') from e - return shared.readfile(local, as_type='dict') - - -def strip_prefix(state_dict): - """Detect and uniformly strip the BFL-style prefix from all keys. - - Supported prefixes (longest first, so ``model.diffusion_model.`` beats ``diffusion_model.``): - ``model.diffusion_model.``, ``diffusion_model.``, or no prefix. Raises - ValueError if some keys match the dominant prefix and others do not, - since mixed prefixes indicate a malformed file. - """ - counts = {p: sum(1 for k in state_dict if k.startswith(p)) for p in KNOWN_PREFIXES} - total = len(state_dict) - dominant = max(counts, key=counts.get) - if counts[dominant] == 0: - log.debug('Load model: type=Anima custom transformer prefix=bare') - return state_dict - if counts[dominant] != total: - raise ValueError( - f'Load model: type=Anima custom transformer has mixed prefixes ' - f'(total={total} {dominant}={counts[dominant]})' - ) - log.debug(f'Load model: type=Anima custom transformer prefix="{dominant}"') - offset = len(dominant) - return {k[offset:]: v for k, v in state_dict.items()} - - -def partition_adapter(state_dict): - """Split into (transformer_sd, adapter_sd) by the ``llm_adapter.`` prefix.""" - transformer_sd = {} - adapter_sd = {} - for key, value in state_dict.items(): - if key.startswith(ADAPTER_PREFIX): - adapter_sd[key[len(ADAPTER_PREFIX):]] = value - else: - transformer_sd[key] = value - return transformer_sd, adapter_sd - - -def build_transformer(transformer_sd, transformer_cfg, quant_args, quant_type): - """Convert, instantiate, load, dtype-cast, quantize, and (if offloading) move to CPU.""" - from diffusers.loaders.single_file_utils import convert_cosmos_transformer_checkpoint_to_diffusers - try: - converted = convert_cosmos_transformer_checkpoint_to_diffusers(transformer_sd) - transformer = diffusers.CosmosTransformer3DModel.from_config(transformer_cfg) - missing, unexpected = transformer.load_state_dict(converted, strict=False) - validate_state_dict_load('transformer', missing, unexpected) - del converted - devices.torch_gc() - transformer = transformer.to(dtype=devices.dtype) - except Exception as e: - log.error(f'Load model: type=Anima transformer load failed: {e}') - errors.display(e, 'Load') - raise - - apply_quant(transformer, quant_type) - - if shared.opts.diffusers_offload_mode != 'none': - sd_models.move_model(transformer, devices.cpu) - - if not hasattr(transformer, 'quantization_config'): - if hasattr(transformer, 'config') and hasattr(transformer.config, 'quantization_config'): - transformer.quantization_config = transformer.config.quantization_config - elif (quant_type is not None) and (quant_args.get('quantization_config', None) is not None): - transformer.quantization_config = quant_args.get('quantization_config', None) - return transformer - - -def build_adapter(adapter_sd, adapter_cfg, adapter_cls): - """Instantiate AnimaLLMAdapter from the base repo config and load bundled weights.""" - try: - adapter = adapter_cls.from_config(adapter_cfg) - missing, unexpected = adapter.load_state_dict(adapter_sd, strict=False) - validate_state_dict_load('adapter', missing, unexpected) - adapter = adapter.to(dtype=devices.dtype) - except Exception as e: - log.error(f'Load model: type=Anima adapter load failed: {e}') - errors.display(e, 'Load') - raise - if shared.opts.diffusers_offload_mode != 'none': - sd_models.move_model(adapter, devices.cpu) - return adapter - - -def validate_state_dict_load(component, missing, unexpected): - """Raise ValueError if load_state_dict produced unexpected keys or non-buffer missing keys.""" - if unexpected: - sample = ', '.join(unexpected[:5]) - raise ValueError(f'Load model: type=Anima {component} has {len(unexpected)} unexpected keys (sample: {sample})') - hard_missing = [k for k in missing if not any(k.startswith(p) for p in ACCEPTABLE_MISSING)] - if hard_missing: - sample = ', '.join(hard_missing[:5]) - raise ValueError(f'Load model: type=Anima {component} missing {len(hard_missing)} required keys (sample: {sample})') - if missing: - log.debug(f'Load model: type=Anima {component} ignored {len(missing)} buffer-only missing keys') - - -def apply_quant(transformer, quant_type): - """Apply SDNQ / layerwise quantization to the bare transformer. - - SDNQ 'pre' and 'auto' would normally route through ``quantization_config`` - at ``from_pretrained`` time; since we bypass that boundary, we call the - per-module quant path directly. SDNQ 'post' and ``layerwise_quantization`` - go through ``do_post_load_quant`` as usual. - """ - if quant_type == 'NVIDIAModelOptConfig': - log.warning('Load model: type=Anima quant=TRT not supported on custom transformer path, skipping') - elif quant_type == 'SDNQConfig': - if shared.opts.sdnq_quantize_mode == 'pre': - log.info('Load model: type=Anima quant=SDNQ pre-mode applied post-load on custom transformer path') - model_quant.sdnq_quantize_model(transformer, op='transformer') - # allow=False avoids double-applying SDNQ in auto mode (applied directly - # above); post mode fires regardless of allow, and layerwise always fires. - model_quant.do_post_load_quant(transformer, allow=False) diff --git a/pipelines/model_anima.py b/pipelines/model_anima.py index 0e9908eed..4ed2228fb 100644 --- a/pipelines/model_anima.py +++ b/pipelines/model_anima.py @@ -16,37 +16,25 @@ def _import_from_file(module_name, file_path): return mod -def resolve_custom_transformer_path(): - """Return an absolute path if the user selected a transformer in the UNET - dropdown and the file is resolvable, else ``None``. - """ - sel = shared.opts.sd_unet - if sel is None or sel in ('Default', 'None'): - return None - from modules import sd_unet - if sel not in list(sd_unet.unet_dict): - log.error(f'Load module: type=transformer file="{sel}" not found') - return None - path = sd_unet.unet_dict[sel] - if not os.path.exists(path): - log.error(f'Load module: type=transformer path="{path}" does not exist') - return None - return path - - def load_transformer_components(repo_id, diffusers_load_config, adapter_cls): """Load (transformer, llm_adapter_or_none). - If the UNET dropdown points at a valid safetensors, route through the - custom-transformer helper, which also extracts the bundled adapter - weights. Otherwise fall back to ``generic.load_transformer`` and return - ``None`` for the adapter so the caller loads it from the base repo. + If the UNET dropdown points at a valid safetensors, route through + :mod:`pipelines.native_transformer` with :data:`pipelines.anima.ANIMA_SPEC`, + which extracts any bundled ``llm_adapter`` weights inline with the + transformer. Otherwise fall back to :func:`generic.load_transformer` and + return ``None`` for the adapter so the caller loads it from the base repo. """ - local_file = resolve_custom_transformer_path() + from pipelines import native_transformer + local_file = native_transformer.resolve_path() if local_file is not None: - from pipelines.anima import anima_transformer + from pipelines.anima import ANIMA_SPEC try: - return anima_transformer.load_custom_transformer(repo_id, local_file, diffusers_load_config, adapter_cls) + transformer, siblings = native_transformer.load( + local_file, repo_id, ANIMA_SPEC, diffusers_load_config, + sibling_classes={'llm_adapter': adapter_cls}, + ) + return transformer, siblings.get('llm_adapter') except Exception as e: log.error(f'Load model: type=Anima custom transformer="{local_file}": {e}') errors.display(e, 'Load') diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index 757fbbe4c..5f46de9fa 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -6,11 +6,15 @@ Bypasses :func:`diffusers.loaders.FromOriginalModelMixin.from_single_file` so sdnext owns prefix detection, optional sibling partitioning, dtype/quant/offload handling, and explicit validation of missing/unexpected keys. -The per-arch knobs are captured in :class:`TransformerSpec`. Arches register a -spec at import time via :func:`register`; arches without a registration get a -default spec that handles BFL-style ``model.diffusion_model.`` prefix stripping -and opportunistically picks up a diffusers converter from -``SINGLE_FILE_LOADABLE_CLASSES`` if the class has one. +The per-arch knobs are captured in :class:`TransformerSpec`. Each pipeline +defines its spec in ``pipelines//__init__.py`` and passes it explicitly +to :func:`load` (or to :func:`pipelines.generic.load_transformer` via the +``native_spec`` kwarg). No class-keyed registry: two pipelines may share a +transformer class but need different specs (e.g. Anima vs raw Cosmos both +use ``CosmosTransformer3DModel`` but Anima has a bundled ``llm_adapter`` +sibling). Pipelines without a custom spec fall back to +:func:`make_default_spec`, which opportunistically picks up a real converter +from diffusers' ``SINGLE_FILE_LOADABLE_CLASSES`` table. Algorithm: @@ -90,32 +94,17 @@ class TransformerSpec: forbidden_markers: tuple[tuple[str, str], ...] = () -REGISTRY: dict[type, TransformerSpec] = {} +def make_default_spec(cls: type) -> TransformerSpec: + """Synthesize a default spec for ``cls``: default prefixes, no siblings, + no forbidden markers, and a converter picked up automatically from + diffusers' ``SINGLE_FILE_LOADABLE_CLASSES`` table if one exists (and is + not the no-op identity lambda that ``QwenImageTransformer2DModel`` and + a few other classes register). - -def register(cls: type, spec: TransformerSpec | None = None) -> None: - """Register a transformer class with an explicit spec, or with the default - spec if ``spec`` is None. Idempotent: re-registering the same class - replaces the previous entry. + Used by callers (notably :func:`pipelines.generic.load_transformer`) when + a pipeline does not supply a custom ``TransformerSpec`` of its own. """ - if spec is None: - spec = TransformerSpec(cls=cls) - if spec.cls is not cls: - raise ValueError(f"register: spec.cls ({spec.cls.__name__}) does not match cls ({cls.__name__})") - REGISTRY[cls] = spec - - -def lookup(cls: type) -> TransformerSpec: - """Return the registered spec for ``cls``, or synthesize a default one. - - The synthesized default opportunistically pulls a converter from diffusers' - ``SINGLE_FILE_LOADABLE_CLASSES`` table if one exists for the class name and - is not a pass-through no-op lambda. - """ - if cls in REGISTRY: - return REGISTRY[cls] - converter = auto_pickup_converter(cls) - return TransformerSpec(cls=cls, converter=converter) + return TransformerSpec(cls=cls, converter=auto_pickup_converter(cls)) def auto_pickup_converter(cls: type) -> Callable[[dict], dict] | None: diff --git a/test/test-native-transformer.py b/test/test-native-transformer.py index bac70ad6f..b0a524b29 100644 --- a/test/test-native-transformer.py +++ b/test/test-native-transformer.py @@ -9,7 +9,7 @@ Covers the pure helpers that own per-arch knob handling: - ``check_forbidden_markers`` for structural-mismatch rejection - ``is_noop_converter`` for diffusers no-op lambda detection - ``validate_state_dict_load`` for unexpected / missing key handling -- ``register`` / ``lookup`` registry behavior and default spec synthesis +- ``make_default_spec`` default-spec synthesis with diffusers converter pickup - ``auto_pickup_converter`` for diffusers ``SINGLE_FILE_LOADABLE_CLASSES`` integration - ``TransformerSpec`` / ``SiblingSpec`` defaults @@ -306,58 +306,37 @@ def test_validate_empty_passes(): # ============================================================ -# register / lookup +# make_default_spec # ============================================================ class FakeTransformer: """Minimal stand-in for a diffusers transformer class.""" -class FakeTransformer2: - """Second stand-in for register/lookup tests.""" - - -def test_register_with_explicit_spec(): - nt.REGISTRY.clear() - spec = nt.TransformerSpec(cls=FakeTransformer, subfolder='custom_sub') - nt.register(FakeTransformer, spec) - assert nt.lookup(FakeTransformer) is spec - - -def test_register_with_default_spec(): - nt.REGISTRY.clear() - nt.register(FakeTransformer) - spec = nt.lookup(FakeTransformer) +def test_make_default_spec_for_unknown_class(): + spec = nt.make_default_spec(FakeTransformer) assert spec.cls is FakeTransformer assert spec.subfolder == 'transformer' assert spec.prefixes == nt.DEFAULT_PREFIXES - assert spec.converter is None + assert spec.converter is None # no diffusers entry for FakeTransformer assert spec.siblings == {} + assert spec.forbidden_markers == () -def test_register_idempotent_replaces(): - nt.REGISTRY.clear() - spec1 = nt.TransformerSpec(cls=FakeTransformer, subfolder='sub_one') - spec2 = nt.TransformerSpec(cls=FakeTransformer, subfolder='sub_two') - nt.register(FakeTransformer, spec1) - nt.register(FakeTransformer, spec2) - assert nt.lookup(FakeTransformer) is spec2 +def test_make_default_spec_picks_up_real_diffusers_converter(): + import diffusers + spec = nt.make_default_spec(diffusers.FluxTransformer2DModel) + assert spec.converter is not None + assert spec.converter.__name__ == 'convert_flux_transformer_checkpoint_to_diffusers' -def test_register_rejects_mismatched_cls(): - try: - nt.register(FakeTransformer, nt.TransformerSpec(cls=FakeTransformer2)) - raise AssertionError('expected ValueError') - except ValueError as e: - assert 'does not match' in str(e) - - -def test_lookup_synthesizes_default_for_unregistered(): - nt.REGISTRY.clear() - spec = nt.lookup(FakeTransformer) - assert spec.cls is FakeTransformer - assert spec.subfolder == 'transformer' # default - assert spec.converter is None # FakeTransformer has no diffusers entry +def test_make_default_spec_skips_qwen_image_noop(): + """QwenImageTransformer2DModel's diffusers entry is a no-op lambda; the + default spec must NOT pick it up, leaving converter=None so the caller + sees only their own (potentially absent) override.""" + import diffusers + spec = nt.make_default_spec(diffusers.QwenImageTransformer2DModel) + assert spec.converter is None # ============================================================ @@ -701,14 +680,12 @@ def run_all(): ]: run_test(cat, fn) - log.warning('=== register / lookup ===') - cat = category('registry') + log.warning('=== make_default_spec ===') + cat = category('default_spec') for fn in [ - test_register_with_explicit_spec, - test_register_with_default_spec, - test_register_idempotent_replaces, - test_register_rejects_mismatched_cls, - test_lookup_synthesizes_default_for_unregistered, + test_make_default_spec_for_unknown_class, + test_make_default_spec_picks_up_real_diffusers_converter, + test_make_default_spec_skips_qwen_image_noop, ]: run_test(cat, fn) From 98fba8ddc3a7af53780dfd8ae8573dbde1f38ff9 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 25 May 2026 05:47:56 +0100 Subject: [PATCH 05/27] feat(pipelines): wire native_spec dispatch in generic.load_transformer New native_spec=None kwarg. When set and the UNET dropdown points at a .safetensors, dispatches to native_transformer.load (threading allow_quant/dtype/modules_to_not_convert/modules_dtype_dict). Pipelines without a spec stay on cls.from_single_file unchanged. --- pipelines/generic.py | 22 +++++++++++++++++++++- pipelines/native_transformer.py | 24 ++++++++++++++++++++++-- 2 files changed, 43 insertions(+), 3 deletions(-) diff --git a/pipelines/generic.py b/pipelines/generic.py index b6aa3e92b..768077ff7 100644 --- a/pipelines/generic.py +++ b/pipelines/generic.py @@ -19,7 +19,17 @@ def _loader(component): return 'runai' if shared.opts.runai_streamer_transformers else 'default' -def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer", allow_quant=True, variant=None, dtype=None, modules_to_not_convert=None, modules_dtype_dict=None, **kwargs): +def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer", allow_quant=True, variant=None, dtype=None, modules_to_not_convert=None, modules_dtype_dict=None, native_spec=None, **kwargs): + """Load a DiT transformer from the base repo, or from a user-selected + single file when the UNET dropdown (``shared.opts.sd_unet``) is set. + + When ``native_spec`` is supplied and a .safetensors override is selected, + dispatches to :func:`pipelines.native_transformer.load` so the per-arch + spec (multi-prefix detection, optional converter, optional sibling + partitioning, forbidden markers) drives the load. Pipelines without a + spec continue to use the legacy ``from_single_file`` path; this preserves + behavior for Mode D arches until they explicitly opt in. + """ if shared.state.interrupted: return None transformer = None @@ -55,6 +65,16 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer **load_args, ) transformer = model_quant.do_post_load_quant(transformer, allow=quant_type is not None) + elif local_file is not None and local_file.lower().endswith('.safetensors') and native_spec is not None: + from pipelines import native_transformer + log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader=native_transformer args={load_args}') + transformer, _ = native_transformer.load( + local_file, repo_id, native_spec, load_config, + allow_quant=allow_quant, + dtype=dtype, + modules_to_not_convert=modules_to_not_convert, + modules_dtype_dict=modules_dtype_dict, + ) elif local_file is not None and local_file.lower().endswith('.safetensors'): log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("diffusers")} args={load_args}') if dtype is not None: diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index 5f46de9fa..57d2e3ed4 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -171,6 +171,11 @@ def load( spec: TransformerSpec, diffusers_cfg: dict | None = None, sibling_classes: dict[str, type] | None = None, + *, + allow_quant: bool = True, + dtype=None, + modules_to_not_convert: list | None = None, + modules_dtype_dict: dict | None = None, ) -> tuple[object, dict[str, object]]: """Load the transformer (and any bundled siblings) from ``local_file``. @@ -180,6 +185,11 @@ def load( Missing sibling classes raise ``ValueError`` if the corresponding sibling keys are present in the bundled file. + Keyword-only arguments ``allow_quant``, ``dtype``, ``modules_to_not_convert``, + and ``modules_dtype_dict`` mirror the corresponding kwargs of + :func:`pipelines.generic.load_transformer` so the dispatch from there can + plumb the caller's intent through unchanged. + Returns ``(transformer, siblings_dict)``. ``siblings_dict`` is keyed by sibling name and is empty for non-sibling specs, or for sibling specs whose keys are absent from the bundled file. @@ -197,7 +207,10 @@ def load( ) _, quant_args = model_quant.get_dit_args( - diffusers_cfg, module="Model", device_map=True, allow_quant=True, + diffusers_cfg, module="Model", device_map=True, + allow_quant=allow_quant, + modules_to_not_convert=modules_to_not_convert, + modules_dtype_dict=modules_dtype_dict, ) quant_type = model_quant.get_quant_type(quant_args) @@ -213,6 +226,7 @@ def load( f"transformer_keys={len(transformer_sd)} siblings={sibling_counts or '{}'}" ) + effective_dtype = dtype if dtype is not None else devices.dtype transformer_cfg = fetch_component_config(repo_id, spec.subfolder) transformer = build_component( component_name="transformer", @@ -223,6 +237,7 @@ def load( acceptable_missing=spec.acceptable_missing, quant_args=quant_args, quant_type=quant_type, + dtype=effective_dtype, ) del transformer_sd devices.torch_gc() @@ -248,6 +263,7 @@ def load( acceptable_missing=sibling_spec.acceptable_missing, quant_args={}, quant_type=None, + dtype=effective_dtype, ) sd_models.allow_post_quant = False @@ -362,9 +378,13 @@ def build_component( acceptable_missing: tuple[str, ...], quant_args: dict, quant_type: str | None, + dtype=None, ) -> object: """Convert (if needed), instantiate, load weights, dtype-cast, quantize, and offload-place a single component. Raises on any hard failure. + + ``dtype`` overrides ``devices.dtype`` when supplied; otherwise the global + default is used. """ try: sd = converter(state_dict) if converter is not None else state_dict @@ -373,7 +393,7 @@ def build_component( validate_state_dict_load(component_name, missing, unexpected, acceptable_missing) del sd devices.torch_gc() - component = component.to(dtype=devices.dtype) + component = component.to(dtype=dtype if dtype is not None else devices.dtype) except Exception as e: log.error(f"Load model: native_transformer {component_name} load failed: {e}") errors.display(e, "Load") From b681cccccbe7f3df2c3816245457c7a35292be7b Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 25 May 2026 05:47:56 +0100 Subject: [PATCH 06/27] feat(ernie): register ERNIE_SPEC for native dispatch ErnieImageTransformer2DModel has no from_single_file; selecting an Ernie finetune in the UNET dropdown previously crashed with "is not a valid JSON file" from from_pretrained. Probe of jibMixErnie_v20.safetensors: 409/409 keys overlap with the model state dict after stripping model.diffusion_model., zero missing or unexpected. Spec is the minimum TransformerSpec(cls=...). --- pipelines/ernie/__init__.py | 29 +++++++++++++++++++++++++++++ pipelines/model_ernie.py | 2 ++ 2 files changed, 31 insertions(+) diff --git a/pipelines/ernie/__init__.py b/pipelines/ernie/__init__.py index e69de29bb..dfed6ff84 100644 --- a/pipelines/ernie/__init__.py +++ b/pipelines/ernie/__init__.py @@ -0,0 +1,29 @@ +"""ERNIE-Image pipeline package. + +Exports :data:`ERNIE_SPEC` for use by :mod:`pipelines.model_ernie` together +with :mod:`pipelines.native_transformer`. The spec captures the +ERNIE-Image-specific knobs that differ from the native-loader defaults: + +- No converter is needed: community Ernie trainer dumps use BFL-style keys + (``model.diffusion_model.``-prefixed) whose names match diffusers' + ``ErnieImageTransformer2DModel.state_dict()`` verbatim after prefix strip. + Probed against ``jibMixErnie_v20.safetensors`` (the upstream community + finetune): 409/409 keys overlap with zero missing or unexpected. +- No siblings; no forbidden markers; default prefixes + (``model.diffusion_model.``, ``diffusion_model.``, ``net.``) cover every + exporter seen in the wild. + +Before this spec, selecting an Ernie finetune via the UNET dropdown crashed +in :func:`diffusers.loaders.ModelMixin.from_pretrained` with a misleading +``OSError: ... is not a valid JSON file`` because +``ErnieImageTransformer2DModel`` lacks ``from_single_file`` support and the +fallback path treats the safetensors as a directory looking for +``config.json``. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +ERNIE_SPEC = TransformerSpec(cls=diffusers.ErnieImageTransformer2DModel) diff --git a/pipelines/model_ernie.py b/pipelines/model_ernie.py index 382b1fe80..471c081be 100644 --- a/pipelines/model_ernie.py +++ b/pipelines/model_ernie.py @@ -14,10 +14,12 @@ def load_ernie_image(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=ERNIE-Image repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args} pe={shared.opts.model_ernie_enable_pe}') + from pipelines.ernie import ERNIE_SPEC transformer = generic.load_transformer( repo_id, cls_name=diffusers.ErnieImageTransformer2DModel, load_config=diffusers_load_config, + native_spec=ERNIE_SPEC, ) text_encoder = generic.load_text_encoder( repo_id, From d01d637d9433abe0756bc71fae3f81ad820819d3 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 25 May 2026 05:47:56 +0100 Subject: [PATCH 07/27] refactor(wanai): retire bespoke load_transformer, route through generic The 40-line hand-rolled loader was generic.load_transformer plus a VACE/standard class switch and a hardcoded GGUF rejection. Switch moves to load_wan; GGUF rejection removed (generic handles it via GGUFQuantizationConfig). No native_spec passed: WanTransformer3DModel has a working diffusers converter, so from_single_file via generic stays correct. --- pipelines/model_wanai.py | 56 +++++----------------------------------- 1 file changed, 7 insertions(+), 49 deletions(-) diff --git a/pipelines/model_wanai.py b/pipelines/model_wanai.py index f7a4413a8..e62fb165d 100644 --- a/pipelines/model_wanai.py +++ b/pipelines/model_wanai.py @@ -1,51 +1,8 @@ -import os import transformers import diffusers from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae from modules.logger import log - - -def load_transformer(repo_id, diffusers_load_config=None, subfolder='transformer'): - if diffusers_load_config is None: - diffusers_load_config = {} - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True) - fn = None - - if 'VACE' in repo_id: - transformer_cls = diffusers.WanVACETransformer3DModel - else: - transformer_cls = diffusers.WanTransformer3DModel - - if shared.opts.sd_unet is not None and shared.opts.sd_unet != 'Default': - from modules import sd_unet - if shared.opts.sd_unet not in list(sd_unet.unet_dict): - log.error(f'Load module: type=Transformer not found: {shared.opts.sd_unet}') - return None - fn = sd_unet.unet_dict[shared.opts.sd_unet] if os.path.exists(sd_unet.unet_dict[shared.opts.sd_unet]) else None - - if fn is not None and 'gguf' in fn.lower(): - log.error('Load model: type=WanAI format="gguf" unsupported') - transformer = None - elif fn is not None and 'safetensors' in fn.lower(): - log.debug(f'Load model: type=WanAI {subfolder}="{fn}" quant="{model_quant.get_quant(repo_id)}" args={load_args}') - transformer = transformer_cls.from_single_file( - fn, - cache_dir=shared.opts.hfcache_dir, - **load_args, - **quant_args, - ) - else: - log.debug(f'Load model: type=WanAI {subfolder}="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') - transformer = transformer_cls.from_pretrained( - repo_id, - subfolder=subfolder, - cache_dir=shared.opts.hfcache_dir, - **load_args, - **quant_args, - ) - if shared.opts.diffusers_offload_mode != 'none' and transformer is not None: - sd_models.move_model(transformer, devices.cpu) - return transformer +from pipelines import generic def load_text_encoder(repo_id, diffusers_load_config=None): @@ -71,26 +28,27 @@ def load_wan(checkpoint_info, diffusers_load_config=None): diffusers_load_config = {} repo_id = sd_models.path_to_repo(checkpoint_info) sd_models.hf_auth_check(checkpoint_info) + transformer_cls = diffusers.WanVACETransformer3DModel if 'VACE' in repo_id else diffusers.WanTransformer3DModel boundary_ratio = None if 'a14b' in repo_id.lower() or 'fun-14b' in repo_id.lower(): if shared.opts.model_wan_stage == 'high noise' or shared.opts.model_wan_stage == 'first': - transformer = load_transformer(repo_id, diffusers_load_config, 'transformer') + transformer = generic.load_transformer(repo_id, cls_name=transformer_cls, load_config=diffusers_load_config, subfolder='transformer') transformer_2 = None boundary_ratio = 0.0 elif shared.opts.model_wan_stage == 'low noise' or shared.opts.model_wan_stage == 'second': transformer = None - transformer_2 = load_transformer(repo_id, diffusers_load_config, 'transformer_2') + transformer_2 = generic.load_transformer(repo_id, cls_name=transformer_cls, load_config=diffusers_load_config, subfolder='transformer_2') boundary_ratio = 1000.0 elif shared.opts.model_wan_stage == 'combined' or shared.opts.model_wan_stage == 'both': - transformer = load_transformer(repo_id, diffusers_load_config, 'transformer') - transformer_2 = load_transformer(repo_id, diffusers_load_config, 'transformer_2') + transformer = generic.load_transformer(repo_id, cls_name=transformer_cls, load_config=diffusers_load_config, subfolder='transformer') + transformer_2 = generic.load_transformer(repo_id, cls_name=transformer_cls, load_config=diffusers_load_config, subfolder='transformer_2') boundary_ratio = shared.opts.model_wan_boundary else: log.error(f'Load model: type=WanAI stage="{shared.opts.model_wan_stage}" unsupported') return None else: - transformer = load_transformer(repo_id, diffusers_load_config, 'transformer') + transformer = generic.load_transformer(repo_id, cls_name=transformer_cls, load_config=diffusers_load_config, subfolder='transformer') transformer_2 = None text_encoder = load_text_encoder(repo_id, diffusers_load_config) From e7d2e28d3d68fdce3b24d82c1dfede327d3cb15d Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 25 May 2026 05:47:56 +0100 Subject: [PATCH 08/27] feat(pipelines): wire native_spec for arches missing single_file_model entry ChronoEdit, HunyuanImage, Kandinsky5, LongCat, Ovis, NucleusMoE, Bria, BriaFibo. Each inherits from_single_file via FromOriginalModelMixin but has no SINGLE_FILE_LOADABLE_CLASSES entry; calling it crashed with "FromOriginalModelMixin is currently only compatible with [...]". TransformerSpec(cls=...) for each. Bria specs in pipelines/bria/__init__.py; the rest are module-level in model_.py since they have no package directory. --- pipelines/bria/__init__.py | 29 +++++++++++++++++++++++++++++ pipelines/model_bria.py | 5 ++++- pipelines/model_chrono.py | 6 +++++- pipelines/model_hyimage.py | 6 +++++- pipelines/model_kandinsky.py | 6 +++++- pipelines/model_longcat.py | 6 +++++- pipelines/model_nucleus.py | 5 +++++ pipelines/model_ovis.py | 6 +++++- 8 files changed, 63 insertions(+), 6 deletions(-) diff --git a/pipelines/bria/__init__.py b/pipelines/bria/__init__.py index e69de29bb..4dcb2d97a 100644 --- a/pipelines/bria/__init__.py +++ b/pipelines/bria/__init__.py @@ -0,0 +1,29 @@ +"""Bria pipeline package. + +Exports :data:`BRIA_SPEC` and :data:`BRIA_FIBO_SPEC` for use by +:mod:`pipelines.model_bria` together with :mod:`pipelines.native_transformer`. + +Bria ships two transformer variants: + +- The original Bria family uses a custom :class:`BriaTransformer2DModel` + imported from :mod:`pipelines.bria.transformer_bria`. The custom class + has no entry in diffusers' ``SINGLE_FILE_LOADABLE_CLASSES`` table, so a + community single-file selection previously crashed in + ``ModelMixin.from_pretrained``. +- Bria FIBO (and FIBO Edit) use the upstream + :class:`diffusers.BriaFiboTransformer2DModel`, which also lacks a + ``SINGLE_FILE_LOADABLE_CLASSES`` entry. Same crash. + +Both specs use the default 3-prefix detection +(``model.diffusion_model.``, ``diffusion_model.``, ``net.``), no +converter, and no siblings. +""" + +import diffusers + +from pipelines.bria.transformer_bria import BriaTransformer2DModel +from pipelines.native_transformer import TransformerSpec + + +BRIA_SPEC = TransformerSpec(cls=BriaTransformer2DModel) +BRIA_FIBO_SPEC = TransformerSpec(cls=diffusers.BriaFiboTransformer2DModel) diff --git a/pipelines/model_bria.py b/pipelines/model_bria.py index ae9c9fa7a..961f9c26f 100644 --- a/pipelines/model_bria.py +++ b/pipelines/model_bria.py @@ -20,11 +20,13 @@ def load_bria(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=BriaFibo repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + from pipelines.bria import BRIA_FIBO_SPEC transformer = generic.load_transformer( repo_id, cls_name=diffusers.BriaFiboTransformer2DModel, load_config=diffusers_load_config, allow_quant=False, + native_spec=BRIA_FIBO_SPEC, ) text_encoder = generic.load_text_encoder( repo_id, @@ -66,7 +68,8 @@ def load_bria(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=Bria repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - transformer = generic.load_transformer(repo_id, cls_name=BriaTransformer2DModel, load_config=diffusers_load_config) + from pipelines.bria import BRIA_SPEC + transformer = generic.load_transformer(repo_id, cls_name=BriaTransformer2DModel, load_config=diffusers_load_config, native_spec=BRIA_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config) pipe = BriaPipeline.from_pretrained( diff --git a/pipelines/model_chrono.py b/pipelines/model_chrono.py index 502e7d89d..1bab65c04 100644 --- a/pipelines/model_chrono.py +++ b/pipelines/model_chrono.py @@ -3,6 +3,10 @@ import transformers from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae from modules.logger import log from pipelines import generic +from pipelines.native_transformer import TransformerSpec + + +CHRONOEDIT_SPEC = TransformerSpec(cls=diffusers.ChronoEditTransformer3DModel) def postprocess(p, result): # pylint: disable=unused-argument @@ -21,7 +25,7 @@ def load_chrono(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=ChronoEdit repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - transformer = generic.load_transformer(repo_id, cls_name=diffusers.ChronoEditTransformer3DModel, load_config=diffusers_load_config, subfolder="transformer") + transformer = generic.load_transformer(repo_id, cls_name=diffusers.ChronoEditTransformer3DModel, load_config=diffusers_load_config, subfolder="transformer", native_spec=CHRONOEDIT_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.UMT5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder") try: diff --git a/pipelines/model_hyimage.py b/pipelines/model_hyimage.py index 4cf9776ee..2bed9851b 100644 --- a/pipelines/model_hyimage.py +++ b/pipelines/model_hyimage.py @@ -5,6 +5,10 @@ import diffusers from modules import shared, sd_models, devices, model_quant, sd_hijack_te, sd_hijack_vae from modules.logger import log from pipelines import generic +from pipelines.native_transformer import TransformerSpec + + +HUNYUANIMAGE_SPEC = TransformerSpec(cls=diffusers.HunyuanImageTransformer2DModel) def load_hyimage(checkpoint_info, diffusers_load_config=None): # pylint: disable=unused-argument @@ -16,7 +20,7 @@ def load_hyimage(checkpoint_info, diffusers_load_config=None): # pylint: disable load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) log.debug(f'Load model: type=HunyuanImage21 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - transformer = generic.load_transformer(repo_id, cls_name=diffusers.HunyuanImageTransformer2DModel, load_config=diffusers_load_config, subfolder="transformer") + transformer = generic.load_transformer(repo_id, cls_name=diffusers.HunyuanImageTransformer2DModel, load_config=diffusers_load_config, subfolder="transformer", native_spec=HUNYUANIMAGE_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen2_5_VLForConditionalGeneration, load_config=diffusers_load_config, subfolder="text_encoder") text_encoder_2 = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder_2", allow_shared=False) diff --git a/pipelines/model_kandinsky.py b/pipelines/model_kandinsky.py index d8e950177..4f06e185d 100644 --- a/pipelines/model_kandinsky.py +++ b/pipelines/model_kandinsky.py @@ -3,6 +3,10 @@ import diffusers from modules import shared, sd_models, devices, model_quant, sd_hijack_te, sd_hijack_vae from modules.logger import log from pipelines import generic +from pipelines.native_transformer import TransformerSpec + + +KANDINSKY5_SPEC = TransformerSpec(cls=diffusers.Kandinsky5Transformer3DModel) def load_kandinsky21(checkpoint_info, diffusers_load_config=None): @@ -83,7 +87,7 @@ def load_kandinsky5(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) log.debug(f'Load model: type=Kandinsky50 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - transformer = generic.load_transformer(repo_id, cls_name=diffusers.Kandinsky5Transformer3DModel, load_config=diffusers_load_config) + transformer = generic.load_transformer(repo_id, cls_name=diffusers.Kandinsky5Transformer3DModel, load_config=diffusers_load_config, native_spec=KANDINSKY5_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen2_5_VLForConditionalGeneration, load_config=diffusers_load_config) if 'I2I' in repo_id: diff --git a/pipelines/model_longcat.py b/pipelines/model_longcat.py index ac8a27322..8af13fffe 100644 --- a/pipelines/model_longcat.py +++ b/pipelines/model_longcat.py @@ -3,6 +3,10 @@ import diffusers from modules import shared, devices, sd_models, model_quant, sd_hijack_te from modules.logger import log from pipelines import generic +from pipelines.native_transformer import TransformerSpec + + +LONGCAT_SPEC = TransformerSpec(cls=diffusers.LongCatImageTransformer2DModel) def load_longcat(checkpoint_info, diffusers_load_config=None): @@ -14,7 +18,7 @@ def load_longcat(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=LongCat repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={diffusers_load_config}') - transformer = generic.load_transformer(repo_id, cls_name=diffusers.LongCatImageTransformer2DModel, load_config=diffusers_load_config) + transformer = generic.load_transformer(repo_id, cls_name=diffusers.LongCatImageTransformer2DModel, load_config=diffusers_load_config, native_spec=LONGCAT_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen2_5_VLForConditionalGeneration, load_config=diffusers_load_config) text_processor = transformers.Qwen2VLProcessor.from_pretrained(repo_id, subfolder='tokenizer', cache_dir=shared.opts.hfcache_dir) diff --git a/pipelines/model_nucleus.py b/pipelines/model_nucleus.py index b7087fa59..8625bc7db 100644 --- a/pipelines/model_nucleus.py +++ b/pipelines/model_nucleus.py @@ -3,6 +3,10 @@ import transformers from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae from modules.logger import log from pipelines import generic +from pipelines.native_transformer import TransformerSpec + + +NUCLEUS_SPEC = TransformerSpec(cls=diffusers.NucleusMoEImageTransformer2DModel) def load_nucleus(checkpoint_info, diffusers_load_config=None): @@ -18,6 +22,7 @@ def load_nucleus(checkpoint_info, diffusers_load_config=None): repo_id, cls_name=diffusers.NucleusMoEImageTransformer2DModel, load_config=diffusers_load_config, + native_spec=NUCLEUS_SPEC, ) text_encoder = generic.load_text_encoder( repo_id, diff --git a/pipelines/model_ovis.py b/pipelines/model_ovis.py index e4e65270f..1cd26a63c 100644 --- a/pipelines/model_ovis.py +++ b/pipelines/model_ovis.py @@ -3,6 +3,10 @@ import diffusers from modules import shared, devices, sd_models, model_quant, sd_hijack_te from modules.logger import log from pipelines import generic +from pipelines.native_transformer import TransformerSpec + + +OVIS_SPEC = TransformerSpec(cls=diffusers.OvisImageTransformer2DModel) def load_ovis(checkpoint_info, diffusers_load_config=None): @@ -14,7 +18,7 @@ def load_ovis(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=OvisImage repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={diffusers_load_config}') - transformer = generic.load_transformer(repo_id, cls_name=diffusers.OvisImageTransformer2DModel, load_config=diffusers_load_config) + transformer = generic.load_transformer(repo_id, cls_name=diffusers.OvisImageTransformer2DModel, load_config=diffusers_load_config, native_spec=OVIS_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3Model, load_config=diffusers_load_config) pipe = diffusers.OvisImagePipeline.from_pretrained( From e302518d7374a9d7229962bb96dfa344c409e127 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 25 May 2026 05:47:56 +0100 Subject: [PATCH 09/27] feat(pipelines): wire native_spec for arches without from_single_file CogView3+, CogView4, GLM-Image, HunyuanDiT, Joy, Kandinsky3 UNet, PixArt, PRX. These classes lack from_single_file entirely; the UNET dropdown fell through to cls.from_pretrained(local_file), which decodes the .safetensors header as UTF-8 and crashes "is not a valid JSON file". TransformerSpec(cls=...) module-level in each model_.py. Kandinsky3 also needs subfolder='unet' since it loads a UNet not a transformer. --- pipelines/model_cogview.py | 9 +++++++-- pipelines/model_glm.py | 7 ++++++- pipelines/model_hunyuandit.py | 6 +++++- pipelines/model_joy.py | 5 +++++ pipelines/model_kandinsky.py | 3 ++- pipelines/model_pixart.py | 6 +++++- pipelines/model_prx.py | 6 +++++- 7 files changed, 35 insertions(+), 7 deletions(-) diff --git a/pipelines/model_cogview.py b/pipelines/model_cogview.py index 673c41565..146177063 100644 --- a/pipelines/model_cogview.py +++ b/pipelines/model_cogview.py @@ -3,6 +3,11 @@ import diffusers from modules import shared, devices, sd_models, model_quant, sd_hijack_te from modules.logger import log from pipelines import generic +from pipelines.native_transformer import TransformerSpec + + +COGVIEW3_SPEC = TransformerSpec(cls=diffusers.CogView3PlusTransformer2DModel) +COGVIEW4_SPEC = TransformerSpec(cls=diffusers.CogView4Transformer2DModel) def load_cogview3(checkpoint_info, diffusers_load_config=None): @@ -14,7 +19,7 @@ def load_cogview3(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) log.debug(f'Load model: type=CogView3 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - transformer = generic.load_transformer(repo_id, cls_name=diffusers.CogView3PlusTransformer2DModel, load_config=diffusers_load_config, subfolder="transformer") + transformer = generic.load_transformer(repo_id, cls_name=diffusers.CogView3PlusTransformer2DModel, load_config=diffusers_load_config, subfolder="transformer", native_spec=COGVIEW3_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder") pipe = diffusers.CogView3PlusPipeline.from_pretrained( @@ -40,7 +45,7 @@ def load_cogview4(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) log.debug(f'Load model: type=CogView4 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - transformer = generic.load_transformer(repo_id, cls_name=diffusers.CogView4Transformer2DModel, load_config=diffusers_load_config, subfolder="transformer") + transformer = generic.load_transformer(repo_id, cls_name=diffusers.CogView4Transformer2DModel, load_config=diffusers_load_config, subfolder="transformer", native_spec=COGVIEW4_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.GlmModel, load_config=diffusers_load_config, subfolder="text_encoder", allow_quant=True) pipe = diffusers.CogView4Pipeline.from_pretrained( diff --git a/pipelines/model_glm.py b/pipelines/model_glm.py index c58358f88..97f218bce 100644 --- a/pipelines/model_glm.py +++ b/pipelines/model_glm.py @@ -5,6 +5,10 @@ import diffusers from modules import shared, devices, sd_models, model_quant, sd_hijack_te from modules.logger import log, console from pipelines import generic +from pipelines.native_transformer import TransformerSpec + + +GLM_IMAGE_SPEC = TransformerSpec(cls=diffusers.GlmImageTransformer2DModel) class GLMTokenProgressProcessor(transformers.LogitsProcessor): @@ -98,7 +102,8 @@ def load_glm_image(checkpoint_info, diffusers_load_config=None): transformer = generic.load_transformer( repo_id, cls_name=diffusers.GlmImageTransformer2DModel, - load_config=diffusers_load_config + load_config=diffusers_load_config, + native_spec=GLM_IMAGE_SPEC, ) # Load text encoder (ByT5 for glyph) - cannot use shared T5 as GLM-Image requires specific ByT5 encoder (1472 hidden size) diff --git a/pipelines/model_hunyuandit.py b/pipelines/model_hunyuandit.py index b2933a94f..71e75bf48 100644 --- a/pipelines/model_hunyuandit.py +++ b/pipelines/model_hunyuandit.py @@ -3,6 +3,10 @@ import diffusers from modules import shared, sd_models, devices, model_quant from modules.logger import log from pipelines import generic +from pipelines.native_transformer import TransformerSpec + + +HUNYUANDIT_SPEC = TransformerSpec(cls=diffusers.HunyuanDiT2DModel) def load_hunyuandit(checkpoint_info, diffusers_load_config=None): @@ -19,7 +23,7 @@ def load_hunyuandit(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) log.debug(f'Load model: type=HunyuanDiT repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - transformer = generic.load_transformer(repo_id, cls_name=diffusers.HunyuanDiT2DModel, load_config=diffusers_load_config) + transformer = generic.load_transformer(repo_id, cls_name=diffusers.HunyuanDiT2DModel, load_config=diffusers_load_config, native_spec=HUNYUANDIT_SPEC) repo_te = 'Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers' if 'HunyuanDiT-v1' in repo_id else repo_id text_encoder_2 = generic.load_text_encoder(repo_te, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder_2", allow_shared=False) # this is not normal t5 diff --git a/pipelines/model_joy.py b/pipelines/model_joy.py index 5ed363315..99b8de829 100644 --- a/pipelines/model_joy.py +++ b/pipelines/model_joy.py @@ -3,6 +3,10 @@ import transformers from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae from modules.logger import log from pipelines import generic +from pipelines.native_transformer import TransformerSpec + + +JOY_SPEC = TransformerSpec(cls=diffusers.JoyImageEditTransformer3DModel) def load_joy(checkpoint_info, diffusers_load_config=None): @@ -18,6 +22,7 @@ def load_joy(checkpoint_info, diffusers_load_config=None): repo_id, cls_name=diffusers.JoyImageEditTransformer3DModel, load_config=diffusers_load_config, + native_spec=JOY_SPEC, ) text_encoder = generic.load_text_encoder( repo_id, diff --git a/pipelines/model_kandinsky.py b/pipelines/model_kandinsky.py index 4f06e185d..41a1bb85c 100644 --- a/pipelines/model_kandinsky.py +++ b/pipelines/model_kandinsky.py @@ -6,6 +6,7 @@ from pipelines import generic from pipelines.native_transformer import TransformerSpec +KANDINSKY3_UNET_SPEC = TransformerSpec(cls=diffusers.Kandinsky3UNet, subfolder='unet') KANDINSKY5_SPEC = TransformerSpec(cls=diffusers.Kandinsky5Transformer3DModel) @@ -54,7 +55,7 @@ def load_kandinsky3(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) log.debug(f'Load model: type=Kandinsky30 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - unet = generic.load_transformer(repo_id, cls_name=diffusers.Kandinsky3UNet, load_config=diffusers_load_config, subfolder="unet", variant="fp16") + unet = generic.load_transformer(repo_id, cls_name=diffusers.Kandinsky3UNet, load_config=diffusers_load_config, subfolder="unet", variant="fp16", native_spec=KANDINSKY3_UNET_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder", variant="fp16", allow_shared=False) pipe = diffusers.Kandinsky3Pipeline.from_pretrained( diff --git a/pipelines/model_pixart.py b/pipelines/model_pixart.py index 1ced8b659..0a79b3178 100644 --- a/pipelines/model_pixart.py +++ b/pipelines/model_pixart.py @@ -4,6 +4,10 @@ from huggingface_hub import file_exists from modules import shared, devices, sd_models, model_quant, sd_hijack_te from modules.logger import log from pipelines import generic +from pipelines.native_transformer import TransformerSpec + + +PIXART_SPEC = TransformerSpec(cls=diffusers.PixArtTransformer2DModel) def load_pixart(checkpoint_info, diffusers_load_config=None): @@ -24,7 +28,7 @@ def load_pixart(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=PixArtSigma repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - transformer = generic.load_transformer(repo_id, cls_name=diffusers.PixArtTransformer2DModel, load_config=diffusers_load_config) + transformer = generic.load_transformer(repo_id, cls_name=diffusers.PixArtTransformer2DModel, load_config=diffusers_load_config, native_spec=PIXART_SPEC) text_encoder = generic.load_text_encoder(repo_id_tenc, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config) pipe = diffusers.PixArtSigmaPipeline.from_pretrained( diff --git a/pipelines/model_prx.py b/pipelines/model_prx.py index 52088444c..8b3fc3a53 100644 --- a/pipelines/model_prx.py +++ b/pipelines/model_prx.py @@ -2,6 +2,10 @@ import diffusers from modules import shared, devices, sd_models, model_quant, sd_hijack_te from modules.logger import log from pipelines import generic +from pipelines.native_transformer import TransformerSpec + + +PRX_SPEC = TransformerSpec(cls=diffusers.PRXTransformer2DModel) def load_prx(checkpoint_info, diffusers_load_config=None): @@ -14,7 +18,7 @@ def load_prx(checkpoint_info, diffusers_load_config=None): log.debug(f'Load model: type=PRX repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') from transformers.models.t5gemma.modeling_t5gemma import T5GemmaEncoder - transformer = generic.load_transformer(repo_id, cls_name=diffusers.PRXTransformer2DModel, load_config=diffusers_load_config) + transformer = generic.load_transformer(repo_id, cls_name=diffusers.PRXTransformer2DModel, load_config=diffusers_load_config, native_spec=PRX_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=T5GemmaEncoder, load_config=diffusers_load_config) pipe = diffusers.PRXPipeline.from_pretrained( From 3368ce95f7adf98bc9fbdc1f62669f9e61742074 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 25 May 2026 05:47:56 +0100 Subject: [PATCH 10/27] feat(pipelines): wire native_spec for Qwen-Image and custom-class arches Qwen-Image: diffusers registers a no-op identity lambda for QwenImageTransformer2DModel; from_single_file silently accepted any key naming and loaded with mismatches. QWEN_SPEC sets converter=None to short-circuit the no-op. Custom transformer classes (Lens, Step1X-Edit, VIBE Sana, f-lite): runtime-registered onto diffusers, never in SINGLE_FILE_LOADABLE_CLASSES. Specs land in each arch's existing pipelines//__init__.py. NextStep skipped: loader returns None up front. --- pipelines/f_lite/__init__.py | 6 +++++- pipelines/lens/__init__.py | 5 +++++ pipelines/model_flite.py | 2 +- pipelines/model_lens.py | 2 +- pipelines/model_qwen.py | 2 ++ pipelines/model_step1x_edit.py | 3 ++- pipelines/model_vibe.py | 2 ++ pipelines/qwen/__init__.py | 22 ++++++++++++++++++++++ pipelines/step1x/__init__.py | 6 ++++++ pipelines/vibe/__init__.py | 7 ++++++- 10 files changed, 52 insertions(+), 5 deletions(-) diff --git a/pipelines/f_lite/__init__.py b/pipelines/f_lite/__init__.py index e4ddb20f5..5c4a23513 100644 --- a/pipelines/f_lite/__init__.py +++ b/pipelines/f_lite/__init__.py @@ -1,5 +1,9 @@ +from pipelines.native_transformer import TransformerSpec from .pipeline import FLitePipeline, FLitePipelineOutput, APGConfig from .model import DiT -__all__ = ["APGConfig", "DiT", "FLitePipeline", "FLitePipelineOutput"] +FLITE_SPEC = TransformerSpec(cls=DiT, subfolder='dit_model') + + +__all__ = ["APGConfig", "DiT", "FLITE_SPEC", "FLitePipeline", "FLitePipelineOutput"] diff --git a/pipelines/lens/__init__.py b/pipelines/lens/__init__.py index bcfca204f..cae35a683 100644 --- a/pipelines/lens/__init__.py +++ b/pipelines/lens/__init__.py @@ -2,6 +2,7 @@ import diffusers as _diffusers import transformers as _transformers +from pipelines.native_transformer import TransformerSpec from .pipeline import LensPipeline, LensPipelineOutput from .pipeline_image import LensImg2ImgPipeline, LensInpaintPipeline from .reasoner import PromptReasoner @@ -9,6 +10,9 @@ from .resolution import RESOLUTION_BUCKETS, resolve_resolution from .text_encoder import LensGptOssEncoder from .transformer import LensTransformer2DModel + +LENS_SPEC = TransformerSpec(cls=LensTransformer2DModel) + # --------------------------------------------------------------------------- # Make our custom subclasses discoverable by ``diffusers.DiffusionPipeline``. # @@ -38,6 +42,7 @@ if not hasattr(_diffusers, "LensInpaintPipeline"): del _diffusers, _transformers __all__ = [ + "LENS_SPEC", "LensPipeline", "LensPipelineOutput", "LensImg2ImgPipeline", diff --git a/pipelines/model_flite.py b/pipelines/model_flite.py index 52e7f2f81..aa63969c3 100644 --- a/pipelines/model_flite.py +++ b/pipelines/model_flite.py @@ -19,7 +19,7 @@ def load_flite(checkpoint_info, diffusers_load_config=None): diffusers.FLitePipeline = f_lite.FLitePipeline sys.modules['f_lite'] = f_lite - dit_model = generic.load_transformer(repo_id, cls_name=f_lite.DiT, load_config=diffusers_load_config, subfolder="dit_model") + dit_model = generic.load_transformer(repo_id, cls_name=f_lite.DiT, load_config=diffusers_load_config, subfolder="dit_model", native_spec=f_lite.FLITE_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder") pipe = f_lite.FLitePipeline.from_pretrained( diff --git a/pipelines/model_lens.py b/pipelines/model_lens.py index a4d07a4d9..f6e64abfb 100644 --- a/pipelines/model_lens.py +++ b/pipelines/model_lens.py @@ -15,7 +15,7 @@ def load_lens(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=Lens repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} reasoner={shared.opts.model_lens_enable_pe} args={load_args}') - transformer = generic.load_transformer(repo_id, cls_name=lens.LensTransformer2DModel, load_config=diffusers_load_config) + transformer = generic.load_transformer(repo_id, cls_name=lens.LensTransformer2DModel, load_config=diffusers_load_config, native_spec=lens.LENS_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=lens.LensGptOssEncoder, load_config=diffusers_load_config, allow_quant=False) pipe = lens.LensPipeline.from_pretrained( diff --git a/pipelines/model_qwen.py b/pipelines/model_qwen.py index 52c9fd251..094d23ed3 100644 --- a/pipelines/model_qwen.py +++ b/pipelines/model_qwen.py @@ -57,12 +57,14 @@ def load_qwen(checkpoint_info, diffusers_load_config=None): transformer_subfolder = "transformer" if transformer is None: + from pipelines.qwen import QWEN_SPEC transformer = generic.load_transformer( repo_transformer, subfolder=transformer_subfolder, cls_name=diffusers.QwenImageTransformer2DModel, load_config=diffusers_load_config, modules_to_not_convert=["transformer_blocks.0.img_mod.1.weight"], + native_spec=QWEN_SPEC, ) repo_te = 'Qwen/Qwen-Image' diff --git a/pipelines/model_step1x_edit.py b/pipelines/model_step1x_edit.py index aba2bd5ab..47d0c8adb 100644 --- a/pipelines/model_step1x_edit.py +++ b/pipelines/model_step1x_edit.py @@ -22,7 +22,8 @@ def load_step1x_edit(checkpoint_info, diffusers_load_config=None): text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen2_5_VLForConditionalGeneration, load_config=diffusers_load_config) processor = transformers.Qwen2_5_VLProcessor.from_pretrained(repo_id, cache_dir=shared.opts.hfcache_dir, subfolder='processor') - transformer = generic.load_transformer(repo_id, cls_name=Step1XEditTransformer2DModel, load_config=diffusers_load_config) + from pipelines.step1x import STEP1X_SPEC + transformer = generic.load_transformer(repo_id, cls_name=Step1XEditTransformer2DModel, load_config=diffusers_load_config, native_spec=STEP1X_SPEC) pipe = Step1XEditPipeline.from_pretrained( repo_id, diff --git a/pipelines/model_vibe.py b/pipelines/model_vibe.py index ecac08867..999ab9dd7 100644 --- a/pipelines/model_vibe.py +++ b/pipelines/model_vibe.py @@ -21,11 +21,13 @@ def load_vibe(checkpoint_info, diffusers_load_config=None): sys.modules['vibe.transformer.vibe_sana_editing'] = diffusers # monkey patch since hf model_index.json points to custom class path + from pipelines.vibe import VIBE_SPEC transformer = generic.load_transformer( repo_id, cls_name=VIBESanaEditingModel, load_config=diffusers_load_config, allow_quant=False, + native_spec=VIBE_SPEC, ) text_encoder = generic.load_text_encoder( repo_id, diff --git a/pipelines/qwen/__init__.py b/pipelines/qwen/__init__.py index e1455f988..6dffd74b0 100644 --- a/pipelines/qwen/__init__.py +++ b/pipelines/qwen/__init__.py @@ -1,2 +1,24 @@ +"""Qwen-Image pipeline package. + +Exports :data:`QWEN_SPEC` for use by :mod:`pipelines.model_qwen` together +with :mod:`pipelines.native_transformer`. Qwen-Image is the lone Mode C +arch: diffusers registers a no-op identity lambda as its +``SINGLE_FILE_LOADABLE_CLASSES`` converter, so ``from_single_file`` silently +accepts whatever key naming the community file uses and loads with mangled +weights instead of raising. + +The spec explicitly sets ``converter=None`` to short-circuit the no-op +pickup; if a real Qwen-Image converter is needed for some trainer dump +in the wild, it can be plugged in here. Until then, validation surfaces a +clear error listing unexpected/missing keys instead of letting a malformed +file load silently. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec from pipelines.qwen.qwen_nunchaku import load_qwen_nunchaku from pipelines.qwen.qwen_pruning import check_qwen_pruning + + +QWEN_SPEC = TransformerSpec(cls=diffusers.QwenImageTransformer2DModel, converter=None) diff --git a/pipelines/step1x/__init__.py b/pipelines/step1x/__init__.py index f7983b96f..2de77fcb0 100644 --- a/pipelines/step1x/__init__.py +++ b/pipelines/step1x/__init__.py @@ -1,8 +1,14 @@ +from pipelines.native_transformer import TransformerSpec from pipelines.step1x.pipeline_output import Step1XEditPipelineOutput from pipelines.step1x.pipeline_step1x_edit import Step1XEditPipeline from pipelines.step1x.transformer_step1x_edit import Step1XEditTransformer2DModel + +STEP1X_SPEC = TransformerSpec(cls=Step1XEditTransformer2DModel) + + __all__ = [ + "STEP1X_SPEC", "Step1XEditPipeline", "Step1XEditPipelineOutput", "Step1XEditTransformer2DModel", diff --git a/pipelines/vibe/__init__.py b/pipelines/vibe/__init__.py index 49b0390cb..5aa63c70c 100644 --- a/pipelines/vibe/__init__.py +++ b/pipelines/vibe/__init__.py @@ -1,6 +1,11 @@ """VIBE pipeline components for SD.Next.""" +from pipelines.native_transformer import TransformerSpec from .vibe_sana_editing import VIBESanaEditingModel from .vibe_sana_pipeline import VIBESanaEditingPipeline, VIBESanaImagePipeline -__all__ = ["VIBESanaEditingModel", "VIBESanaEditingPipeline", "VIBESanaImagePipeline"] + +VIBE_SPEC = TransformerSpec(cls=VIBESanaEditingModel) + + +__all__ = ["VIBE_SPEC", "VIBESanaEditingModel", "VIBESanaEditingPipeline", "VIBESanaImagePipeline"] From a68ae8991d67e7a7cb900acd59826ca472ef8fd4 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 25 May 2026 05:47:56 +0100 Subject: [PATCH 11/27] fix(flux2_klein): route through native_transformer with explicit converter Klein shares Flux2Transformer2DModel with full Flux 2 but uses a smaller config. from_single_file instantiates with the class default (= Flux 2 full size); loading a Klein-shaped file crashed at load_model_dict_into_meta with: Cannot load because double_stream_modulation_img.linear.weight expected shape (36864, 6144), but got (24576, 4096) FLUX2_KLEIN_SPEC fetches Klein's transformer/config.json from the base repo and uses convert_flux2_transformer_checkpoint_to_diffusers explicitly. Probed against snofsSexNudesAndOtherFunStuff_v14Base: 233/233 keys overlap. --- pipelines/model_flux2_klein.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/pipelines/model_flux2_klein.py b/pipelines/model_flux2_klein.py index 3a238805b..a6bae4a51 100644 --- a/pipelines/model_flux2_klein.py +++ b/pipelines/model_flux2_klein.py @@ -1,8 +1,25 @@ import transformers import diffusers +from diffusers.loaders.single_file_utils import convert_flux2_transformer_checkpoint_to_diffusers from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae from modules.logger import log from pipelines import generic +from pipelines.native_transformer import TransformerSpec + + +# Klein shares Flux2Transformer2DModel with full Flux 2, but uses a smaller +# config (hidden_size and friends). diffusers' from_single_file picks the +# class default (= Flux 2 full), so loading a Klein-shaped community file +# crashes at load_model_dict_into_meta with a shape mismatch like +# "expected (36864, 6144), got (24576, 4096)". Routing through +# native_transformer pulls the Klein transformer/config.json from the base +# repo first and instantiates Flux2Transformer2DModel at the right size, +# then runs the diffusers Flux 2 converter to split fused QKV blocks and +# rename BFL keys into the diffusers-expected names. +FLUX2_KLEIN_SPEC = TransformerSpec( + cls=diffusers.Flux2Transformer2DModel, + converter=convert_flux2_transformer_checkpoint_to_diffusers, +) def load_flux2_klein(checkpoint_info, diffusers_load_config=None): @@ -15,7 +32,7 @@ def load_flux2_klein(checkpoint_info, diffusers_load_config=None): log.debug(f'Load model: type=Flux2Klein repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') # Load transformer - Klein uses Flux2Transformer2DModel (same class as Flux2, different size) - transformer = generic.load_transformer(repo_id, cls_name=diffusers.Flux2Transformer2DModel, load_config=diffusers_load_config) + transformer = generic.load_transformer(repo_id, cls_name=diffusers.Flux2Transformer2DModel, load_config=diffusers_load_config, native_spec=FLUX2_KLEIN_SPEC) # Load text encoder - Klein uses Qwen3 (4B for Klein-4B, 8B for Klein-9B) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3ForCausalLM, load_config=diffusers_load_config) From 55db15c2139ca9bcf7883113660db28d20fb32af Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 25 May 2026 06:09:02 +0100 Subject: [PATCH 12/27] docs: cleanup --- modules/lora/native_adapter.py | 15 ++++++++------- pipelines/ernie/__init__.py | 28 +++++++++++----------------- pipelines/qwen/__init__.py | 19 +++++++------------ test/test-anima-native-adapters.py | 2 +- test/test-flux2-native-adapters.py | 4 ++-- 5 files changed, 29 insertions(+), 39 deletions(-) diff --git a/modules/lora/native_adapter.py b/modules/lora/native_adapter.py index d0a1679ef..f2bed9ec3 100644 --- a/modules/lora/native_adapter.py +++ b/modules/lora/native_adapter.py @@ -1,7 +1,6 @@ """Shared scaffolding for native adapter loaders. -The four native adapter loaders (z-image, chroma, ernie, flux2) all implement -the same algorithm: +Each per-arch native adapter loader implements the same algorithm: 1. Read the safetensors state dict 2. Test for family-specific markers; bail out if absent @@ -25,7 +24,7 @@ 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. Loader business logic itself lands in subsequent commits. +generic helpers. """ import os @@ -45,7 +44,8 @@ from modules.lora import lora_common as l # Universal prefix list shared by every native arch loader. Per-arch loaders -# extend this with arch-specific entries (e.g. flux2 adds ``"lycoris_"``). +# extend this with arch-specific entries when their files use additional +# vendor-specific naming conventions. KNOWN_PREFIXES_DEFAULT = ("diffusion_model.", "transformer.", "lora_unet_") @@ -56,9 +56,10 @@ KNOWN_PREFIXES_DEFAULT = ("diffusion_model.", "transformer.", "lora_unet_") BARE_DIFFUSERS_PREFIX_USED = "bare_diffusers" -# Default network-key prefix. Single-component arches (flux2, zimage, chroma, -# ernie) keep this default; multi-component arches (anima: transformer plus -# llm_adapter plus text_encoder) pass a callable that picks per ``prefix_used``. +# Default network-key prefix. Single-component arches keep this default; +# multi-component arches (those with separate text-encoder or adapter +# components alongside the transformer) pass a callable that picks per +# ``prefix_used``. NETWORK_PREFIX_DEFAULT = "lora_transformer_" diff --git a/pipelines/ernie/__init__.py b/pipelines/ernie/__init__.py index dfed6ff84..d73a4f7ef 100644 --- a/pipelines/ernie/__init__.py +++ b/pipelines/ernie/__init__.py @@ -1,23 +1,17 @@ """ERNIE-Image pipeline package. -Exports :data:`ERNIE_SPEC` for use by :mod:`pipelines.model_ernie` together -with :mod:`pipelines.native_transformer`. The spec captures the -ERNIE-Image-specific knobs that differ from the native-loader defaults: +Exports :data:`ERNIE_SPEC`. The minimum +``TransformerSpec(cls=ErnieImageTransformer2DModel)`` works because Ernie +trainer dumps use BFL-style ``model.diffusion_model.``-prefixed keys +whose names match the diffusers state_dict verbatim after prefix strip +(probed against a community finetune: full key overlap with zero +missing or unexpected). No siblings, no converter, no forbidden markers. -- No converter is needed: community Ernie trainer dumps use BFL-style keys - (``model.diffusion_model.``-prefixed) whose names match diffusers' - ``ErnieImageTransformer2DModel.state_dict()`` verbatim after prefix strip. - Probed against ``jibMixErnie_v20.safetensors`` (the upstream community - finetune): 409/409 keys overlap with zero missing or unexpected. -- No siblings; no forbidden markers; default prefixes - (``model.diffusion_model.``, ``diffusion_model.``, ``net.``) cover every - exporter seen in the wild. - -Before this spec, selecting an Ernie finetune via the UNET dropdown crashed -in :func:`diffusers.loaders.ModelMixin.from_pretrained` with a misleading -``OSError: ... is not a valid JSON file`` because -``ErnieImageTransformer2DModel`` lacks ``from_single_file`` support and the -fallback path treats the safetensors as a directory looking for +Without this spec, selecting an Ernie file via the UNET dropdown crashes +in :func:`diffusers.loaders.ModelMixin.from_pretrained` with a +misleading ``OSError: ... is not a valid JSON file``, because +``ErnieImageTransformer2DModel`` lacks ``from_single_file`` and the +fallback treats the safetensors as a directory looking for ``config.json``. """ diff --git a/pipelines/qwen/__init__.py b/pipelines/qwen/__init__.py index 6dffd74b0..aac930654 100644 --- a/pipelines/qwen/__init__.py +++ b/pipelines/qwen/__init__.py @@ -1,17 +1,12 @@ """Qwen-Image pipeline package. -Exports :data:`QWEN_SPEC` for use by :mod:`pipelines.model_qwen` together -with :mod:`pipelines.native_transformer`. Qwen-Image is the lone Mode C -arch: diffusers registers a no-op identity lambda as its -``SINGLE_FILE_LOADABLE_CLASSES`` converter, so ``from_single_file`` silently -accepts whatever key naming the community file uses and loads with mangled -weights instead of raising. - -The spec explicitly sets ``converter=None`` to short-circuit the no-op -pickup; if a real Qwen-Image converter is needed for some trainer dump -in the wild, it can be plugged in here. Until then, validation surfaces a -clear error listing unexpected/missing keys instead of letting a malformed -file load silently. +Exports :data:`QWEN_SPEC`. diffusers registers a no-op identity lambda +for ``QwenImageTransformer2DModel`` in ``SINGLE_FILE_LOADABLE_CLASSES``, +so ``from_single_file`` silently accepts whatever key naming the file +uses and loads with mismatched weights. The spec sets ``converter=None`` +explicitly to skip that no-op; validation then surfaces mismatches as +clear errors. A real converter can be plugged in here if a trainer +format that needs one is encountered. """ import diffusers diff --git a/test/test-anima-native-adapters.py b/test/test-anima-native-adapters.py index f017512e4..abbe25213 100644 --- a/test/test-anima-native-adapters.py +++ b/test/test-anima-native-adapters.py @@ -8,7 +8,7 @@ custom Qwen3-projection MLP), or ``lora_te_*`` (Qwen3 text encoder). Routing is parameterized in ``modules.lora.native_adapter`` via the ``network_prefix`` callable that ``pipelines.anima.anima_lora`` supplies. -Covers the eight families exposed through native_adapter's generics (LoRA, +Covers the families exposed through native_adapter's generics (LoRA, LoKR, LoHA, OFT, IA3, GLoRA, Norm, Full), focused on: - LoRA across all five recognized prefixes (BFL transformer / BFL llm_adapter / diff --git a/test/test-flux2-native-adapters.py b/test/test-flux2-native-adapters.py index f7bd043b0..196a95763 100644 --- a/test/test-flux2-native-adapters.py +++ b/test/test-flux2-native-adapters.py @@ -2,8 +2,8 @@ """ Offline unit tests for Flux2/Klein native adapter loaders. -Covers the nine native families (LoRA, LoKR, LoHA, OFT, BOFT, IA3, -GLoRA, Norm, Full) plus DoRA threading via the universal +Covers the native families (LoRA, LoKR, LoHA, OFT, BOFT, IA3, GLoRA, +Norm, Full) plus DoRA threading via the universal NetworkModule.finalize_updown hook, and ex_bias accumulation across stacked Norm adapters. From 9326804d6942325a249364578b181227ba8da8bf Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 25 May 2026 16:19:42 +0100 Subject: [PATCH 13/27] fix(sd_unet): reload model when reverting UNET dropdown to Default load_unet returned early when sd_unet was set to 'Default'/'None', leaving any previously-loaded custom UNET in memory. The model stayed on the override and the dropdown looked stuck. Now detects the transition from a custom selection to Default and triggers sd_models.load_diffuser() to reload via the canonical from_pretrained path. --- modules/sd_unet.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/modules/sd_unet.py b/modules/sd_unet.py index db0ffbe22..f8c42bb45 100644 --- a/modules/sd_unet.py +++ b/modules/sd_unet.py @@ -46,6 +46,15 @@ def load_unet(model, repo_id: str | None = None): return if shared.opts.sd_unet == 'Default' or shared.opts.sd_unet == 'None': + # If a custom UNET / transformer was previously loaded, switching back + # to 'Default' has to actually revert the override. Without this, the + # model keeps the prior finetune transformer in memory and the + # dropdown looks "stuck". + if loaded_unet in (None, 'Default', 'None'): + return + log.info(f'Load module: type=UNet name="Default" (was="{loaded_unet}") reverting to base transformer') + loaded_unet = shared.opts.sd_unet + sd_models.load_diffuser() return if shared.opts.sd_unet not in list(unet_dict): From bd60372cc5e231c0a4b628b7ea616e21b8415f7c Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 25 May 2026 17:12:33 +0100 Subject: [PATCH 14/27] refactor(pipelines): move per-arch transformer specs into packages Specs for chrono, cogview, flux2_klein, glm, hunyuandit, hyimage, joy, kandinsky, longcat, nucleus, ovis, pixart and prx lived module-level in model_.py. Move each into pipelines//__init__.py to match the layout used by anima, bria, ernie, f_lite, lens, qwen, step1x and vibe. Each model_.py now imports its spec lazily inside the load function, so the package only gets pulled in when that arch is actually loaded (kandinsky 2.x never goes through native dispatch and stays untouched). Also drop the "without this spec..." crash paragraph from the new package docstrings plus ernie/__init__.py and bria/__init__.py. --- pipelines/bria/__init__.py | 6 ++---- pipelines/chrono/__init__.py | 15 +++++++++++++++ pipelines/cogview/__init__.py | 15 +++++++++++++++ pipelines/ernie/__init__.py | 7 ------- pipelines/flux2_klein/__init__.py | 26 ++++++++++++++++++++++++++ pipelines/glm/__init__.py | 15 +++++++++++++++ pipelines/hunyuandit/__init__.py | 14 ++++++++++++++ pipelines/hyimage/__init__.py | 19 +++++++++++++++++++ pipelines/joy/__init__.py | 15 +++++++++++++++ pipelines/kandinsky/__init__.py | 23 +++++++++++++++++++++++ pipelines/longcat/__init__.py | 15 +++++++++++++++ pipelines/model_chrono.py | 5 +---- pipelines/model_cogview.py | 7 ++----- pipelines/model_flux2_klein.py | 18 +----------------- pipelines/model_glm.py | 5 +---- pipelines/model_hunyuandit.py | 5 +---- pipelines/model_hyimage.py | 5 +---- pipelines/model_joy.py | 5 +---- pipelines/model_kandinsky.py | 7 ++----- pipelines/model_longcat.py | 5 +---- pipelines/model_nucleus.py | 5 +---- pipelines/model_ovis.py | 5 +---- pipelines/model_pixart.py | 5 +---- pipelines/model_prx.py | 5 +---- pipelines/nucleus/__init__.py | 15 +++++++++++++++ pipelines/ovis/__init__.py | 15 +++++++++++++++ pipelines/pixart/__init__.py | 14 ++++++++++++++ pipelines/prx/__init__.py | 14 ++++++++++++++ 28 files changed, 232 insertions(+), 78 deletions(-) create mode 100644 pipelines/chrono/__init__.py create mode 100644 pipelines/cogview/__init__.py create mode 100644 pipelines/flux2_klein/__init__.py create mode 100644 pipelines/glm/__init__.py create mode 100644 pipelines/hunyuandit/__init__.py create mode 100644 pipelines/hyimage/__init__.py create mode 100644 pipelines/joy/__init__.py create mode 100644 pipelines/kandinsky/__init__.py create mode 100644 pipelines/longcat/__init__.py create mode 100644 pipelines/nucleus/__init__.py create mode 100644 pipelines/ovis/__init__.py create mode 100644 pipelines/pixart/__init__.py create mode 100644 pipelines/prx/__init__.py diff --git a/pipelines/bria/__init__.py b/pipelines/bria/__init__.py index 4dcb2d97a..e6eb9aaee 100644 --- a/pipelines/bria/__init__.py +++ b/pipelines/bria/__init__.py @@ -7,12 +7,10 @@ Bria ships two transformer variants: - The original Bria family uses a custom :class:`BriaTransformer2DModel` imported from :mod:`pipelines.bria.transformer_bria`. The custom class - has no entry in diffusers' ``SINGLE_FILE_LOADABLE_CLASSES`` table, so a - community single-file selection previously crashed in - ``ModelMixin.from_pretrained``. + has no entry in diffusers' ``SINGLE_FILE_LOADABLE_CLASSES`` table. - Bria FIBO (and FIBO Edit) use the upstream :class:`diffusers.BriaFiboTransformer2DModel`, which also lacks a - ``SINGLE_FILE_LOADABLE_CLASSES`` entry. Same crash. + ``SINGLE_FILE_LOADABLE_CLASSES`` entry. Both specs use the default 3-prefix detection (``model.diffusion_model.``, ``diffusion_model.``, ``net.``), no diff --git a/pipelines/chrono/__init__.py b/pipelines/chrono/__init__.py new file mode 100644 index 000000000..b9fe64887 --- /dev/null +++ b/pipelines/chrono/__init__.py @@ -0,0 +1,15 @@ +"""ChronoEdit pipeline package. + +Exports :data:`CHRONOEDIT_SPEC`. The minimum +``TransformerSpec(cls=ChronoEditTransformer3DModel)`` works because +ChronoEdit community files use BFL-style ``model.diffusion_model.`` +prefixed keys whose names match the diffusers state_dict verbatim after +prefix strip. No siblings, no converter, no forbidden markers. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +CHRONOEDIT_SPEC = TransformerSpec(cls=diffusers.ChronoEditTransformer3DModel) diff --git a/pipelines/cogview/__init__.py b/pipelines/cogview/__init__.py new file mode 100644 index 000000000..7163f543f --- /dev/null +++ b/pipelines/cogview/__init__.py @@ -0,0 +1,15 @@ +"""CogView pipeline package. + +Exports :data:`COGVIEW3_SPEC` and :data:`COGVIEW4_SPEC`. Both default +specs work because CogView community files use BFL-style +``model.diffusion_model.``-prefixed keys whose names match the diffusers +state_dict verbatim after prefix strip. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +COGVIEW3_SPEC = TransformerSpec(cls=diffusers.CogView3PlusTransformer2DModel) +COGVIEW4_SPEC = TransformerSpec(cls=diffusers.CogView4Transformer2DModel) diff --git a/pipelines/ernie/__init__.py b/pipelines/ernie/__init__.py index d73a4f7ef..f01ec13ff 100644 --- a/pipelines/ernie/__init__.py +++ b/pipelines/ernie/__init__.py @@ -6,13 +6,6 @@ trainer dumps use BFL-style ``model.diffusion_model.``-prefixed keys whose names match the diffusers state_dict verbatim after prefix strip (probed against a community finetune: full key overlap with zero missing or unexpected). No siblings, no converter, no forbidden markers. - -Without this spec, selecting an Ernie file via the UNET dropdown crashes -in :func:`diffusers.loaders.ModelMixin.from_pretrained` with a -misleading ``OSError: ... is not a valid JSON file``, because -``ErnieImageTransformer2DModel`` lacks ``from_single_file`` and the -fallback treats the safetensors as a directory looking for -``config.json``. """ import diffusers diff --git a/pipelines/flux2_klein/__init__.py b/pipelines/flux2_klein/__init__.py new file mode 100644 index 000000000..72d898a92 --- /dev/null +++ b/pipelines/flux2_klein/__init__.py @@ -0,0 +1,26 @@ +"""Flux 2 Klein pipeline package. + +Exports :data:`FLUX2_KLEIN_SPEC`. Klein shares +:class:`Flux2Transformer2DModel` with full Flux 2 but uses a smaller +config (hidden_size and friends). diffusers' ``from_single_file`` picks +the class default (= Flux 2 full), so loading a Klein-shaped community +file crashes at ``load_model_dict_into_meta`` with a shape mismatch like +``expected (36864, 6144), got (24576, 4096)``. + +Routing through :mod:`pipelines.native_transformer` pulls the Klein +``transformer/config.json`` from the base repo first and instantiates +``Flux2Transformer2DModel`` at the right size, then runs the diffusers +Flux 2 converter to split fused QKV blocks and rename BFL keys into the +diffusers-expected names. +""" + +import diffusers +from diffusers.loaders.single_file_utils import convert_flux2_transformer_checkpoint_to_diffusers + +from pipelines.native_transformer import TransformerSpec + + +FLUX2_KLEIN_SPEC = TransformerSpec( + cls=diffusers.Flux2Transformer2DModel, + converter=convert_flux2_transformer_checkpoint_to_diffusers, +) diff --git a/pipelines/glm/__init__.py b/pipelines/glm/__init__.py new file mode 100644 index 000000000..9de6a1cc6 --- /dev/null +++ b/pipelines/glm/__init__.py @@ -0,0 +1,15 @@ +"""GLM-Image pipeline package. + +Exports :data:`GLM_IMAGE_SPEC`. The minimum +``TransformerSpec(cls=GlmImageTransformer2DModel)`` works because +GLM-Image community files use BFL-style ``model.diffusion_model.`` +prefixed keys whose names match the diffusers state_dict verbatim after +prefix strip. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +GLM_IMAGE_SPEC = TransformerSpec(cls=diffusers.GlmImageTransformer2DModel) diff --git a/pipelines/hunyuandit/__init__.py b/pipelines/hunyuandit/__init__.py new file mode 100644 index 000000000..da6fbe23d --- /dev/null +++ b/pipelines/hunyuandit/__init__.py @@ -0,0 +1,14 @@ +"""HunyuanDiT pipeline package. + +Exports :data:`HUNYUANDIT_SPEC`. The minimum +``TransformerSpec(cls=HunyuanDiT2DModel)`` works because HunyuanDiT +community files use BFL-style ``model.diffusion_model.``-prefixed keys +whose names match the diffusers state_dict verbatim after prefix strip. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +HUNYUANDIT_SPEC = TransformerSpec(cls=diffusers.HunyuanDiT2DModel) diff --git a/pipelines/hyimage/__init__.py b/pipelines/hyimage/__init__.py new file mode 100644 index 000000000..42cc55c6e --- /dev/null +++ b/pipelines/hyimage/__init__.py @@ -0,0 +1,19 @@ +"""HunyuanImage pipeline package. + +Exports :data:`HUNYUANIMAGE_SPEC`. The minimum +``TransformerSpec(cls=HunyuanImageTransformer2DModel)`` works because +HunyuanImage 2.1 community files use BFL-style +``model.diffusion_model.``-prefixed keys whose names match the diffusers +state_dict verbatim after prefix strip. + +The HunyuanImage 3 path is a transformers ``AutoModelForCausalLM`` and +does not go through the native transformer loader, so it does not need a +spec here. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +HUNYUANIMAGE_SPEC = TransformerSpec(cls=diffusers.HunyuanImageTransformer2DModel) diff --git a/pipelines/joy/__init__.py b/pipelines/joy/__init__.py new file mode 100644 index 000000000..1ee025cbe --- /dev/null +++ b/pipelines/joy/__init__.py @@ -0,0 +1,15 @@ +"""Joy-Image-Edit pipeline package. + +Exports :data:`JOY_SPEC`. The minimum +``TransformerSpec(cls=JoyImageEditTransformer3DModel)`` works because +Joy community files use BFL-style ``model.diffusion_model.``-prefixed +keys whose names match the diffusers state_dict verbatim after prefix +strip. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +JOY_SPEC = TransformerSpec(cls=diffusers.JoyImageEditTransformer3DModel) diff --git a/pipelines/kandinsky/__init__.py b/pipelines/kandinsky/__init__.py new file mode 100644 index 000000000..1dcd1e3cb --- /dev/null +++ b/pipelines/kandinsky/__init__.py @@ -0,0 +1,23 @@ +"""Kandinsky pipeline package. + +Exports :data:`KANDINSKY3_UNET_SPEC` and :data:`KANDINSKY5_SPEC`. + +Kandinsky ships several generations under one family: + +- Kandinsky 2.1 / 2.2 are unet-based and load through diffusers' + combined pipelines without going through the native loader, so no + spec is needed. +- Kandinsky 3 uses :class:`Kandinsky3UNet` in the ``unet`` subfolder of + the repo, hence ``subfolder='unet'`` instead of the default + ``'transformer'``. +- Kandinsky 5 uses the new :class:`Kandinsky5Transformer3DModel` and + follows the default layout. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +KANDINSKY3_UNET_SPEC = TransformerSpec(cls=diffusers.Kandinsky3UNet, subfolder='unet') +KANDINSKY5_SPEC = TransformerSpec(cls=diffusers.Kandinsky5Transformer3DModel) diff --git a/pipelines/longcat/__init__.py b/pipelines/longcat/__init__.py new file mode 100644 index 000000000..f68f49c9c --- /dev/null +++ b/pipelines/longcat/__init__.py @@ -0,0 +1,15 @@ +"""LongCat pipeline package. + +Exports :data:`LONGCAT_SPEC`. The minimum +``TransformerSpec(cls=LongCatImageTransformer2DModel)`` works because +LongCat community files use BFL-style ``model.diffusion_model.`` +prefixed keys whose names match the diffusers state_dict verbatim after +prefix strip. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +LONGCAT_SPEC = TransformerSpec(cls=diffusers.LongCatImageTransformer2DModel) diff --git a/pipelines/model_chrono.py b/pipelines/model_chrono.py index 1bab65c04..be66af407 100644 --- a/pipelines/model_chrono.py +++ b/pipelines/model_chrono.py @@ -3,10 +3,6 @@ import transformers from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae from modules.logger import log from pipelines import generic -from pipelines.native_transformer import TransformerSpec - - -CHRONOEDIT_SPEC = TransformerSpec(cls=diffusers.ChronoEditTransformer3DModel) def postprocess(p, result): # pylint: disable=unused-argument @@ -25,6 +21,7 @@ def load_chrono(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=ChronoEdit repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + from pipelines.chrono import CHRONOEDIT_SPEC transformer = generic.load_transformer(repo_id, cls_name=diffusers.ChronoEditTransformer3DModel, load_config=diffusers_load_config, subfolder="transformer", native_spec=CHRONOEDIT_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.UMT5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder") diff --git a/pipelines/model_cogview.py b/pipelines/model_cogview.py index 146177063..a48489e1f 100644 --- a/pipelines/model_cogview.py +++ b/pipelines/model_cogview.py @@ -3,11 +3,6 @@ import diffusers from modules import shared, devices, sd_models, model_quant, sd_hijack_te from modules.logger import log from pipelines import generic -from pipelines.native_transformer import TransformerSpec - - -COGVIEW3_SPEC = TransformerSpec(cls=diffusers.CogView3PlusTransformer2DModel) -COGVIEW4_SPEC = TransformerSpec(cls=diffusers.CogView4Transformer2DModel) def load_cogview3(checkpoint_info, diffusers_load_config=None): @@ -19,6 +14,7 @@ def load_cogview3(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) log.debug(f'Load model: type=CogView3 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + from pipelines.cogview import COGVIEW3_SPEC transformer = generic.load_transformer(repo_id, cls_name=diffusers.CogView3PlusTransformer2DModel, load_config=diffusers_load_config, subfolder="transformer", native_spec=COGVIEW3_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder") @@ -45,6 +41,7 @@ def load_cogview4(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) log.debug(f'Load model: type=CogView4 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + from pipelines.cogview import COGVIEW4_SPEC transformer = generic.load_transformer(repo_id, cls_name=diffusers.CogView4Transformer2DModel, load_config=diffusers_load_config, subfolder="transformer", native_spec=COGVIEW4_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.GlmModel, load_config=diffusers_load_config, subfolder="text_encoder", allow_quant=True) diff --git a/pipelines/model_flux2_klein.py b/pipelines/model_flux2_klein.py index a6bae4a51..00c666df6 100644 --- a/pipelines/model_flux2_klein.py +++ b/pipelines/model_flux2_klein.py @@ -1,25 +1,8 @@ import transformers import diffusers -from diffusers.loaders.single_file_utils import convert_flux2_transformer_checkpoint_to_diffusers from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae from modules.logger import log from pipelines import generic -from pipelines.native_transformer import TransformerSpec - - -# Klein shares Flux2Transformer2DModel with full Flux 2, but uses a smaller -# config (hidden_size and friends). diffusers' from_single_file picks the -# class default (= Flux 2 full), so loading a Klein-shaped community file -# crashes at load_model_dict_into_meta with a shape mismatch like -# "expected (36864, 6144), got (24576, 4096)". Routing through -# native_transformer pulls the Klein transformer/config.json from the base -# repo first and instantiates Flux2Transformer2DModel at the right size, -# then runs the diffusers Flux 2 converter to split fused QKV blocks and -# rename BFL keys into the diffusers-expected names. -FLUX2_KLEIN_SPEC = TransformerSpec( - cls=diffusers.Flux2Transformer2DModel, - converter=convert_flux2_transformer_checkpoint_to_diffusers, -) def load_flux2_klein(checkpoint_info, diffusers_load_config=None): @@ -32,6 +15,7 @@ def load_flux2_klein(checkpoint_info, diffusers_load_config=None): log.debug(f'Load model: type=Flux2Klein repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') # Load transformer - Klein uses Flux2Transformer2DModel (same class as Flux2, different size) + from pipelines.flux2_klein import FLUX2_KLEIN_SPEC transformer = generic.load_transformer(repo_id, cls_name=diffusers.Flux2Transformer2DModel, load_config=diffusers_load_config, native_spec=FLUX2_KLEIN_SPEC) # Load text encoder - Klein uses Qwen3 (4B for Klein-4B, 8B for Klein-9B) diff --git a/pipelines/model_glm.py b/pipelines/model_glm.py index 97f218bce..c46b0ae63 100644 --- a/pipelines/model_glm.py +++ b/pipelines/model_glm.py @@ -5,10 +5,6 @@ import diffusers from modules import shared, devices, sd_models, model_quant, sd_hijack_te from modules.logger import log, console from pipelines import generic -from pipelines.native_transformer import TransformerSpec - - -GLM_IMAGE_SPEC = TransformerSpec(cls=diffusers.GlmImageTransformer2DModel) class GLMTokenProgressProcessor(transformers.LogitsProcessor): @@ -99,6 +95,7 @@ def load_glm_image(checkpoint_info, diffusers_load_config=None): log.debug(f'Load model: type=GLM-Image repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') # Load transformer (DiT decoder - 7B) with quantization support + from pipelines.glm import GLM_IMAGE_SPEC transformer = generic.load_transformer( repo_id, cls_name=diffusers.GlmImageTransformer2DModel, diff --git a/pipelines/model_hunyuandit.py b/pipelines/model_hunyuandit.py index 71e75bf48..1caa06f16 100644 --- a/pipelines/model_hunyuandit.py +++ b/pipelines/model_hunyuandit.py @@ -3,10 +3,6 @@ import diffusers from modules import shared, sd_models, devices, model_quant from modules.logger import log from pipelines import generic -from pipelines.native_transformer import TransformerSpec - - -HUNYUANDIT_SPEC = TransformerSpec(cls=diffusers.HunyuanDiT2DModel) def load_hunyuandit(checkpoint_info, diffusers_load_config=None): @@ -23,6 +19,7 @@ def load_hunyuandit(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) log.debug(f'Load model: type=HunyuanDiT repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + from pipelines.hunyuandit import HUNYUANDIT_SPEC transformer = generic.load_transformer(repo_id, cls_name=diffusers.HunyuanDiT2DModel, load_config=diffusers_load_config, native_spec=HUNYUANDIT_SPEC) repo_te = 'Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers' if 'HunyuanDiT-v1' in repo_id else repo_id text_encoder_2 = generic.load_text_encoder(repo_te, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder_2", allow_shared=False) # this is not normal t5 diff --git a/pipelines/model_hyimage.py b/pipelines/model_hyimage.py index 2bed9851b..83b1f9c1c 100644 --- a/pipelines/model_hyimage.py +++ b/pipelines/model_hyimage.py @@ -5,10 +5,6 @@ import diffusers from modules import shared, sd_models, devices, model_quant, sd_hijack_te, sd_hijack_vae from modules.logger import log from pipelines import generic -from pipelines.native_transformer import TransformerSpec - - -HUNYUANIMAGE_SPEC = TransformerSpec(cls=diffusers.HunyuanImageTransformer2DModel) def load_hyimage(checkpoint_info, diffusers_load_config=None): # pylint: disable=unused-argument @@ -20,6 +16,7 @@ def load_hyimage(checkpoint_info, diffusers_load_config=None): # pylint: disable load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) log.debug(f'Load model: type=HunyuanImage21 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + from pipelines.hyimage import HUNYUANIMAGE_SPEC transformer = generic.load_transformer(repo_id, cls_name=diffusers.HunyuanImageTransformer2DModel, load_config=diffusers_load_config, subfolder="transformer", native_spec=HUNYUANIMAGE_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen2_5_VLForConditionalGeneration, load_config=diffusers_load_config, subfolder="text_encoder") text_encoder_2 = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder_2", allow_shared=False) diff --git a/pipelines/model_joy.py b/pipelines/model_joy.py index 99b8de829..e0c25f6c7 100644 --- a/pipelines/model_joy.py +++ b/pipelines/model_joy.py @@ -3,10 +3,6 @@ import transformers from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae from modules.logger import log from pipelines import generic -from pipelines.native_transformer import TransformerSpec - - -JOY_SPEC = TransformerSpec(cls=diffusers.JoyImageEditTransformer3DModel) def load_joy(checkpoint_info, diffusers_load_config=None): @@ -18,6 +14,7 @@ def load_joy(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=JoyImageEdit repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + from pipelines.joy import JOY_SPEC transformer = generic.load_transformer( repo_id, cls_name=diffusers.JoyImageEditTransformer3DModel, diff --git a/pipelines/model_kandinsky.py b/pipelines/model_kandinsky.py index 41a1bb85c..e9f81d811 100644 --- a/pipelines/model_kandinsky.py +++ b/pipelines/model_kandinsky.py @@ -3,11 +3,6 @@ import diffusers from modules import shared, sd_models, devices, model_quant, sd_hijack_te, sd_hijack_vae from modules.logger import log from pipelines import generic -from pipelines.native_transformer import TransformerSpec - - -KANDINSKY3_UNET_SPEC = TransformerSpec(cls=diffusers.Kandinsky3UNet, subfolder='unet') -KANDINSKY5_SPEC = TransformerSpec(cls=diffusers.Kandinsky5Transformer3DModel) def load_kandinsky21(checkpoint_info, diffusers_load_config=None): @@ -55,6 +50,7 @@ def load_kandinsky3(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) log.debug(f'Load model: type=Kandinsky30 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + from pipelines.kandinsky import KANDINSKY3_UNET_SPEC unet = generic.load_transformer(repo_id, cls_name=diffusers.Kandinsky3UNet, load_config=diffusers_load_config, subfolder="unet", variant="fp16", native_spec=KANDINSKY3_UNET_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder", variant="fp16", allow_shared=False) @@ -88,6 +84,7 @@ def load_kandinsky5(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) log.debug(f'Load model: type=Kandinsky50 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + from pipelines.kandinsky import KANDINSKY5_SPEC transformer = generic.load_transformer(repo_id, cls_name=diffusers.Kandinsky5Transformer3DModel, load_config=diffusers_load_config, native_spec=KANDINSKY5_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen2_5_VLForConditionalGeneration, load_config=diffusers_load_config) diff --git a/pipelines/model_longcat.py b/pipelines/model_longcat.py index 8af13fffe..397d9c0de 100644 --- a/pipelines/model_longcat.py +++ b/pipelines/model_longcat.py @@ -3,10 +3,6 @@ import diffusers from modules import shared, devices, sd_models, model_quant, sd_hijack_te from modules.logger import log from pipelines import generic -from pipelines.native_transformer import TransformerSpec - - -LONGCAT_SPEC = TransformerSpec(cls=diffusers.LongCatImageTransformer2DModel) def load_longcat(checkpoint_info, diffusers_load_config=None): @@ -18,6 +14,7 @@ def load_longcat(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=LongCat repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={diffusers_load_config}') + from pipelines.longcat import LONGCAT_SPEC transformer = generic.load_transformer(repo_id, cls_name=diffusers.LongCatImageTransformer2DModel, load_config=diffusers_load_config, native_spec=LONGCAT_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen2_5_VLForConditionalGeneration, load_config=diffusers_load_config) text_processor = transformers.Qwen2VLProcessor.from_pretrained(repo_id, subfolder='tokenizer', cache_dir=shared.opts.hfcache_dir) diff --git a/pipelines/model_nucleus.py b/pipelines/model_nucleus.py index 8625bc7db..e465f2216 100644 --- a/pipelines/model_nucleus.py +++ b/pipelines/model_nucleus.py @@ -3,10 +3,6 @@ import transformers from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae from modules.logger import log from pipelines import generic -from pipelines.native_transformer import TransformerSpec - - -NUCLEUS_SPEC = TransformerSpec(cls=diffusers.NucleusMoEImageTransformer2DModel) def load_nucleus(checkpoint_info, diffusers_load_config=None): @@ -18,6 +14,7 @@ def load_nucleus(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=NucleusMoEImage repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + from pipelines.nucleus import NUCLEUS_SPEC transformer = generic.load_transformer( repo_id, cls_name=diffusers.NucleusMoEImageTransformer2DModel, diff --git a/pipelines/model_ovis.py b/pipelines/model_ovis.py index 1cd26a63c..753237610 100644 --- a/pipelines/model_ovis.py +++ b/pipelines/model_ovis.py @@ -3,10 +3,6 @@ import diffusers from modules import shared, devices, sd_models, model_quant, sd_hijack_te from modules.logger import log from pipelines import generic -from pipelines.native_transformer import TransformerSpec - - -OVIS_SPEC = TransformerSpec(cls=diffusers.OvisImageTransformer2DModel) def load_ovis(checkpoint_info, diffusers_load_config=None): @@ -18,6 +14,7 @@ def load_ovis(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=OvisImage repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={diffusers_load_config}') + from pipelines.ovis import OVIS_SPEC transformer = generic.load_transformer(repo_id, cls_name=diffusers.OvisImageTransformer2DModel, load_config=diffusers_load_config, native_spec=OVIS_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3Model, load_config=diffusers_load_config) diff --git a/pipelines/model_pixart.py b/pipelines/model_pixart.py index 0a79b3178..0d2ff1127 100644 --- a/pipelines/model_pixart.py +++ b/pipelines/model_pixart.py @@ -4,10 +4,6 @@ from huggingface_hub import file_exists from modules import shared, devices, sd_models, model_quant, sd_hijack_te from modules.logger import log from pipelines import generic -from pipelines.native_transformer import TransformerSpec - - -PIXART_SPEC = TransformerSpec(cls=diffusers.PixArtTransformer2DModel) def load_pixart(checkpoint_info, diffusers_load_config=None): @@ -28,6 +24,7 @@ def load_pixart(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=PixArtSigma repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + from pipelines.pixart import PIXART_SPEC transformer = generic.load_transformer(repo_id, cls_name=diffusers.PixArtTransformer2DModel, load_config=diffusers_load_config, native_spec=PIXART_SPEC) text_encoder = generic.load_text_encoder(repo_id_tenc, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config) diff --git a/pipelines/model_prx.py b/pipelines/model_prx.py index 8b3fc3a53..4c886f231 100644 --- a/pipelines/model_prx.py +++ b/pipelines/model_prx.py @@ -2,10 +2,6 @@ import diffusers from modules import shared, devices, sd_models, model_quant, sd_hijack_te from modules.logger import log from pipelines import generic -from pipelines.native_transformer import TransformerSpec - - -PRX_SPEC = TransformerSpec(cls=diffusers.PRXTransformer2DModel) def load_prx(checkpoint_info, diffusers_load_config=None): @@ -18,6 +14,7 @@ def load_prx(checkpoint_info, diffusers_load_config=None): log.debug(f'Load model: type=PRX repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') from transformers.models.t5gemma.modeling_t5gemma import T5GemmaEncoder + from pipelines.prx import PRX_SPEC transformer = generic.load_transformer(repo_id, cls_name=diffusers.PRXTransformer2DModel, load_config=diffusers_load_config, native_spec=PRX_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=T5GemmaEncoder, load_config=diffusers_load_config) diff --git a/pipelines/nucleus/__init__.py b/pipelines/nucleus/__init__.py new file mode 100644 index 000000000..798c82edb --- /dev/null +++ b/pipelines/nucleus/__init__.py @@ -0,0 +1,15 @@ +"""Nucleus MoE-Image pipeline package. + +Exports :data:`NUCLEUS_SPEC`. The minimum +``TransformerSpec(cls=NucleusMoEImageTransformer2DModel)`` works because +Nucleus community files use BFL-style ``model.diffusion_model.`` +prefixed keys whose names match the diffusers state_dict verbatim after +prefix strip. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +NUCLEUS_SPEC = TransformerSpec(cls=diffusers.NucleusMoEImageTransformer2DModel) diff --git a/pipelines/ovis/__init__.py b/pipelines/ovis/__init__.py new file mode 100644 index 000000000..7a7b86c8e --- /dev/null +++ b/pipelines/ovis/__init__.py @@ -0,0 +1,15 @@ +"""Ovis-Image pipeline package. + +Exports :data:`OVIS_SPEC`. The minimum +``TransformerSpec(cls=OvisImageTransformer2DModel)`` works because +Ovis community files use BFL-style ``model.diffusion_model.``-prefixed +keys whose names match the diffusers state_dict verbatim after prefix +strip. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +OVIS_SPEC = TransformerSpec(cls=diffusers.OvisImageTransformer2DModel) diff --git a/pipelines/pixart/__init__.py b/pipelines/pixart/__init__.py new file mode 100644 index 000000000..fb0dc043a --- /dev/null +++ b/pipelines/pixart/__init__.py @@ -0,0 +1,14 @@ +"""PixArt pipeline package. + +Exports :data:`PIXART_SPEC`. The minimum +``TransformerSpec(cls=PixArtTransformer2DModel)`` works because PixArt +community files use BFL-style ``model.diffusion_model.``-prefixed keys +whose names match the diffusers state_dict verbatim after prefix strip. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +PIXART_SPEC = TransformerSpec(cls=diffusers.PixArtTransformer2DModel) diff --git a/pipelines/prx/__init__.py b/pipelines/prx/__init__.py new file mode 100644 index 000000000..85d28ea23 --- /dev/null +++ b/pipelines/prx/__init__.py @@ -0,0 +1,14 @@ +"""PRX pipeline package. + +Exports :data:`PRX_SPEC`. The minimum +``TransformerSpec(cls=PRXTransformer2DModel)`` works because PRX +community files use BFL-style ``model.diffusion_model.``-prefixed keys +whose names match the diffusers state_dict verbatim after prefix strip. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +PRX_SPEC = TransformerSpec(cls=diffusers.PRXTransformer2DModel) From d55af33d369523e5ce315609427020ae25da0c0b Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 25 May 2026 19:28:39 +0100 Subject: [PATCH 15/27] refactor(pipelines): split kandinsky into kandinsky3 and kandinsky5 K3 and K5 share only the marketing name: different diffusers pipeline classes (Kandinsky3Pipeline vs Kandinsky5T2I/I2IPipeline), different denoiser classes (Kandinsky3UNet vs Kandinsky5Transformer3DModel), no shared scaffolding. Separate packages so each spec lives next to the arch it serves. K2.1 and K2.2 stay outside both packages since they go through KandinskyCombinedPipeline.from_pretrained directly with no spec. --- pipelines/kandinsky/__init__.py | 23 ----------------------- pipelines/kandinsky3/__init__.py | 13 +++++++++++++ pipelines/kandinsky5/__init__.py | 12 ++++++++++++ pipelines/model_kandinsky.py | 4 ++-- 4 files changed, 27 insertions(+), 25 deletions(-) delete mode 100644 pipelines/kandinsky/__init__.py create mode 100644 pipelines/kandinsky3/__init__.py create mode 100644 pipelines/kandinsky5/__init__.py diff --git a/pipelines/kandinsky/__init__.py b/pipelines/kandinsky/__init__.py deleted file mode 100644 index 1dcd1e3cb..000000000 --- a/pipelines/kandinsky/__init__.py +++ /dev/null @@ -1,23 +0,0 @@ -"""Kandinsky pipeline package. - -Exports :data:`KANDINSKY3_UNET_SPEC` and :data:`KANDINSKY5_SPEC`. - -Kandinsky ships several generations under one family: - -- Kandinsky 2.1 / 2.2 are unet-based and load through diffusers' - combined pipelines without going through the native loader, so no - spec is needed. -- Kandinsky 3 uses :class:`Kandinsky3UNet` in the ``unet`` subfolder of - the repo, hence ``subfolder='unet'`` instead of the default - ``'transformer'``. -- Kandinsky 5 uses the new :class:`Kandinsky5Transformer3DModel` and - follows the default layout. -""" - -import diffusers - -from pipelines.native_transformer import TransformerSpec - - -KANDINSKY3_UNET_SPEC = TransformerSpec(cls=diffusers.Kandinsky3UNet, subfolder='unet') -KANDINSKY5_SPEC = TransformerSpec(cls=diffusers.Kandinsky5Transformer3DModel) diff --git a/pipelines/kandinsky3/__init__.py b/pipelines/kandinsky3/__init__.py new file mode 100644 index 000000000..da46fa862 --- /dev/null +++ b/pipelines/kandinsky3/__init__.py @@ -0,0 +1,13 @@ +"""Kandinsky 3.0 pipeline package. + +Exports :data:`KANDINSKY3_UNET_SPEC`. Kandinsky 3 uses +:class:`Kandinsky3UNet` in the ``unet`` subfolder of the repo, hence +``subfolder='unet'`` instead of the default ``'transformer'``. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +KANDINSKY3_UNET_SPEC = TransformerSpec(cls=diffusers.Kandinsky3UNet, subfolder='unet') diff --git a/pipelines/kandinsky5/__init__.py b/pipelines/kandinsky5/__init__.py new file mode 100644 index 000000000..a2927f1c9 --- /dev/null +++ b/pipelines/kandinsky5/__init__.py @@ -0,0 +1,12 @@ +"""Kandinsky 5.0 pipeline package. + +Exports :data:`KANDINSKY5_SPEC`. Kandinsky 5 uses +:class:`Kandinsky5Transformer3DModel` and follows the default layout. +""" + +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +KANDINSKY5_SPEC = TransformerSpec(cls=diffusers.Kandinsky5Transformer3DModel) diff --git a/pipelines/model_kandinsky.py b/pipelines/model_kandinsky.py index e9f81d811..3fec7f510 100644 --- a/pipelines/model_kandinsky.py +++ b/pipelines/model_kandinsky.py @@ -50,7 +50,7 @@ def load_kandinsky3(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) log.debug(f'Load model: type=Kandinsky30 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - from pipelines.kandinsky import KANDINSKY3_UNET_SPEC + from pipelines.kandinsky3 import KANDINSKY3_UNET_SPEC unet = generic.load_transformer(repo_id, cls_name=diffusers.Kandinsky3UNet, load_config=diffusers_load_config, subfolder="unet", variant="fp16", native_spec=KANDINSKY3_UNET_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder", variant="fp16", allow_shared=False) @@ -84,7 +84,7 @@ def load_kandinsky5(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) log.debug(f'Load model: type=Kandinsky50 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - from pipelines.kandinsky import KANDINSKY5_SPEC + from pipelines.kandinsky5 import KANDINSKY5_SPEC transformer = generic.load_transformer(repo_id, cls_name=diffusers.Kandinsky5Transformer3DModel, load_config=diffusers_load_config, native_spec=KANDINSKY5_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen2_5_VLForConditionalGeneration, load_config=diffusers_load_config) From b29178743c0db3c98b0694afb55942f2db1b2239 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 25 May 2026 23:00:47 +0100 Subject: [PATCH 16/27] fix(native_transformer): forward modules_to_not_convert to sdnq_quantize_model apply_quant called sdnq_quantize_model(transformer, op="transformer") without the per-call skip lists, dropping modules_to_not_convert and modules_dtype_dict that the caller passed into generic.load_transformer. The from_pretrained path receives them via quantization_config; the native path now threads them directly into the explicit quantizer call. Plumbed load -> build_component -> apply_quant as keyword args. No behavior change for callers that did not set the skip lists. --- pipelines/native_transformer.py | 33 +++++++++++++++++++++++++++++---- 1 file changed, 29 insertions(+), 4 deletions(-) diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index 57d2e3ed4..df682d051 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -238,6 +238,8 @@ def load( quant_args=quant_args, quant_type=quant_type, dtype=effective_dtype, + modules_to_not_convert=modules_to_not_convert, + modules_dtype_dict=modules_dtype_dict, ) del transformer_sd devices.torch_gc() @@ -379,12 +381,16 @@ def build_component( quant_args: dict, quant_type: str | None, dtype=None, + modules_to_not_convert: list | None = None, + modules_dtype_dict: dict | None = None, ) -> object: """Convert (if needed), instantiate, load weights, dtype-cast, quantize, and offload-place a single component. Raises on any hard failure. ``dtype`` overrides ``devices.dtype`` when supplied; otherwise the global - default is used. + default is used. ``modules_to_not_convert`` and ``modules_dtype_dict`` + are forwarded to :func:`apply_quant` for the transformer component + (ignored for siblings). """ try: sd = converter(state_dict) if converter is not None else state_dict @@ -400,7 +406,12 @@ def build_component( raise if component_name == "transformer": - apply_quant(component, quant_type) + apply_quant( + component, + quant_type, + modules_to_not_convert=modules_to_not_convert, + modules_dtype_dict=modules_dtype_dict, + ) if shared.opts.diffusers_offload_mode != "none": sd_models.move_model(component, devices.cpu) @@ -445,7 +456,12 @@ def validate_state_dict_load( ) -def apply_quant(transformer: object, quant_type: str | None) -> None: +def apply_quant( + transformer: object, + quant_type: str | None, + modules_to_not_convert: list | None = None, + modules_dtype_dict: dict | None = None, +) -> None: """Apply SDNQ / layerwise quantization to the bare transformer. SDNQ 'pre' and 'auto' would normally route through ``quantization_config`` @@ -453,6 +469,10 @@ def apply_quant(transformer: object, quant_type: str | None) -> None: we call the per-module quant path directly. SDNQ 'post' and ``layerwise_quantization`` go through ``do_post_load_quant`` as usual. NVIDIAModelOptConfig (TRT) is not supported on this path. + + ``modules_to_not_convert`` and ``modules_dtype_dict`` are forwarded + to :func:`sdnq_quantize_model` so per-call skip lists set by the + caller are honored on the native path. """ if quant_type == "NVIDIAModelOptConfig": log.warning( @@ -464,5 +484,10 @@ def apply_quant(transformer: object, quant_type: str | None) -> None: "Load model: native_transformer quant=SDNQ pre-mode applied post-load " "on native path" ) - model_quant.sdnq_quantize_model(transformer, op="transformer") + model_quant.sdnq_quantize_model( + transformer, + op="transformer", + modules_to_not_convert=modules_to_not_convert, + modules_dtype_dict=modules_dtype_dict, + ) model_quant.do_post_load_quant(transformer, allow=False) From cda4822ca0693664d2dc8b8282a0c975e8eadd9e Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 25 May 2026 23:00:56 +0100 Subject: [PATCH 17/27] fix(chroma): route through native_transformer to avoid SDNQ pre-mode crash Loading a Chroma transformer override via UNET dropdown crashed at inference with a shape mismatch in unpack_uint4: from_single_file does not integrate quantization_config the way from_pretrained does, so the dequantizer state was set up but the .weight tensor was never packed. CHROMA_SPEC routes through native_transformer with the diffusers Chroma converter, which loads bare and quantizes explicitly via sdnq_quantize_model. --- pipelines/chroma/__init__.py | 18 ++++++++++++++++++ pipelines/model_chroma.py | 3 ++- 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/pipelines/chroma/__init__.py b/pipelines/chroma/__init__.py index e69de29bb..b8eaa68c6 100644 --- a/pipelines/chroma/__init__.py +++ b/pipelines/chroma/__init__.py @@ -0,0 +1,18 @@ +"""Chroma pipeline package. + +Exports :data:`CHROMA_SPEC`. Chroma community files use BFL-style +``model.diffusion_model.``-prefixed keys that need renaming into the +diffusers naming convention, so the spec plugs in +:func:`convert_chroma_transformer_checkpoint_to_diffusers` explicitly. +""" + +import diffusers +from diffusers.loaders.single_file_utils import convert_chroma_transformer_checkpoint_to_diffusers + +from pipelines.native_transformer import TransformerSpec + + +CHROMA_SPEC = TransformerSpec( + cls=diffusers.ChromaTransformer2DModel, + converter=convert_chroma_transformer_checkpoint_to_diffusers, +) diff --git a/pipelines/model_chroma.py b/pipelines/model_chroma.py index 1e7f9ea8a..ee33d2109 100644 --- a/pipelines/model_chroma.py +++ b/pipelines/model_chroma.py @@ -14,7 +14,8 @@ def load_chroma(checkpoint_info, diffusers_load_config=None): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) log.debug(f'Load model: type=Chroma repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - transformer = generic.load_transformer(repo_id, cls_name=diffusers.ChromaTransformer2DModel, load_config=diffusers_load_config, modules_to_not_convert=["distilled_guidance_layer"]) + from pipelines.chroma import CHROMA_SPEC + transformer = generic.load_transformer(repo_id, cls_name=diffusers.ChromaTransformer2DModel, load_config=diffusers_load_config, modules_to_not_convert=["distilled_guidance_layer"], native_spec=CHROMA_SPEC) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config) pipe = diffusers.ChromaPipeline.from_pretrained( From bab6c253f44d56701bcb6a0a21574885627371bb Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 25 May 2026 23:31:33 +0100 Subject: [PATCH 18/27] feat(sd_unet): reset UNET override on cross-arch base model swap A custom UNET selected via the UNET dropdown carried over silently when the user swapped to a base model of a different arch, then crashed inside the per-arch loader because shared.opts.sd_unet was still pointing at the previous arch's file. reload_model_weights now runs sd_detect.detect_pipeline on the new checkpoint before unloading the old model, compares the detected pipeline class against the loaded model's class, and resets shared.opts.sd_unet to Default when they differ. --- modules/sd_models.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/modules/sd_models.py b/modules/sd_models.py index cc36d57ea..e4f48eb70 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -1412,6 +1412,17 @@ def reload_model_weights(sd_model=None, info=None, op='model', force=False, revi jobid = shared.state.begin('Load model') if sd_model is None: sd_model = model_data.sd_model if op == 'model' or op == 'dict' else model_data.sd_refiner + if op == 'model' and sd_model is not None and shared.opts.sd_unet not in (None, 'Default', 'None'): + old_class = type(sd_model).__name__ + try: + new_pipeline, _ = sd_detect.detect_pipeline(checkpoint_info.path, op) + except Exception: + new_pipeline = None + new_class = getattr(new_pipeline, '__name__', None) + if new_class is not None and new_class != old_class: + log.info(f'Load model: pipeline changed {old_class}->{new_class}, resetting sd_unet from "{shared.opts.sd_unet}" to Default') + shared.opts.data["sd_unet"] = 'Default' + sd_unet.loaded_unet = None if sd_model is None: # previous model load failed current_checkpoint_info = None else: From 75933bfaee0c20efb54c0c91bed6d33d31abff8e Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 25 May 2026 23:58:45 +0100 Subject: [PATCH 19/27] feat(native_transformer): log per-phase progress to reduce load silence The native loader was emitting one line at dispatch, a single line after prefix detection, and the total time at the end. For an 18 GB transformer that meant ~60s of silence with no indication anything was happening. Add DEBUG lines before each long phase: read_state_dict, optional converter, load_state_dict into the model, and dtype cast. Total-time log at the end is unchanged. --- pipelines/native_transformer.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index df682d051..9fb6ef2c1 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -214,6 +214,7 @@ def load( ) quant_type = model_quant.get_quant_type(quant_args) + log.debug(f'Load model: native_transformer reading state_dict cls={spec.cls.__name__} file="{os.path.basename(local_file)}"') state_dict = sd_models.read_state_dict(local_file, what="transformer") state_dict = strip_prefix(state_dict, spec.prefixes, spec.cls.__name__) check_forbidden_markers(state_dict, spec.forbidden_markers, spec.cls.__name__, local_file) @@ -393,13 +394,20 @@ def build_component( (ignored for siblings). """ try: - sd = converter(state_dict) if converter is not None else state_dict + if converter is not None: + log.debug(f'Load model: native_transformer {component_name} converter={converter.__name__} keys={len(state_dict)}') + sd = converter(state_dict) + else: + sd = state_dict + log.debug(f'Load model: native_transformer {component_name} loading keys={len(sd)} cls={cls.__name__}') component = cls.from_config(config) missing, unexpected = component.load_state_dict(sd, strict=False) validate_state_dict_load(component_name, missing, unexpected, acceptable_missing) del sd devices.torch_gc() - component = component.to(dtype=dtype if dtype is not None else devices.dtype) + target_dtype = dtype if dtype is not None else devices.dtype + log.debug(f'Load model: native_transformer {component_name} cast dtype={target_dtype}') + component = component.to(dtype=target_dtype) except Exception as e: log.error(f"Load model: native_transformer {component_name} load failed: {e}") errors.display(e, "Load") From 09c6229b1ce41e49d519285cc7ac0912a3d40680 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Tue, 26 May 2026 00:05:54 +0100 Subject: [PATCH 20/27] fix(ui): refresh UNET dropdown after cross-arch model swap When reload_model_weights resets shared.opts.sd_unet to Default on a cross-arch swap, the backend value changed but the dropdown component kept showing the previous selection because sd_unet was not in the change_checkpoint button's outputs. Add the sd_unet component to the outputs and return its post-reload value via gr.update. --- modules/ui_settings.py | 24 ++++++++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/modules/ui_settings.py b/modules/ui_settings.py index 6441c2e90..09d2e83a4 100644 --- a/modules/ui_settings.py +++ b/modules/ui_settings.py @@ -391,6 +391,10 @@ def create_quicksettings(interfaces): if shared.opts.notification_audio_enable and os.path.exists(os.path.join(paths.script_path, shared.opts.notification_audio_path)): gr.Audio(interactive=False, value=os.path.join(paths.script_path, shared.opts.notification_audio_path), elem_id="audio_notification", visible=False) + def change_checkpoint_with_unet_sync(value, progress=False, force=False): + checkpoint_update, settings_text = run_settings_single(value, key='sd_model_checkpoint', progress=progress, force=force) + return checkpoint_update, get_value_for_setting('sd_unet'), settings_text + for k, _item in quicksettings_list: component = shared.settings_components[k] info = shared.opts.data_labels[k] @@ -405,20 +409,32 @@ def create_quicksettings(interfaces): change_handlers = [component.blur] else: change_handlers = [component.release if hasattr(component, 'release') else component.change] + progress_flag = info.refresh is not None + if k == 'sd_model_checkpoint': + def fn(value, progress=progress_flag): + return change_checkpoint_with_unet_sync(value, progress=progress) + outputs = [component, shared.settings_components['sd_unet'], text_settings] + else: + def fn(value, k=k, progress=progress_flag): + return run_settings_single(value, key=k, progress=progress) + outputs = [component, text_settings] for change_handler in change_handlers: change_handler( - fn=lambda value, k=k, progress=info.refresh is not None: run_settings_single(value, key=k, progress=progress), + fn=fn, inputs=[component], - outputs=[component, text_settings], + outputs=outputs, show_progress='full' if info.refresh is not None else 'hidden', ) + def set_checkpoint_sync_unet(value, _dummy): + return change_checkpoint_with_unet_sync(value, force=True) + button_set_checkpoint = gr.Button('Change model', elem_id='change_checkpoint', visible=False) button_set_checkpoint.click( - fn=lambda value, _: run_settings_single(value, key='sd_model_checkpoint', force=True), + fn=set_checkpoint_sync_unet, _js="consumeDesiredCheckpointName", inputs=[shared.settings_components['sd_model_checkpoint'], dummy_component], - outputs=[shared.settings_components['sd_model_checkpoint'], text_settings], + outputs=[shared.settings_components['sd_model_checkpoint'], shared.settings_components['sd_unet'], text_settings], ) button_set_refiner = gr.Button('Change refiner', elem_id='change_refiner', visible=False) button_set_refiner.click( From 7b9c5d5d8ca08de3b735a8fe7c10b24b62ed8ecc Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Tue, 26 May 2026 21:08:02 +0100 Subject: [PATCH 21/27] fix(sd_models): hide misleading state_dict progress bar for safetensors safetensors.torch.load_file opens the file by path internally, so the progress.open wrapper around it never advanced from 0/N - it just left a stuck '0.0/18.2 GB' bar on screen for the duration of the load. Skip the wrapper on that path; .ckpt and stream_load paths still get the working bar. --- modules/sd_models_utils.py | 35 +++++++++++++++++++---------------- 1 file changed, 19 insertions(+), 16 deletions(-) diff --git a/modules/sd_models_utils.py b/modules/sd_models_utils.py index de241c8ca..773d955b0 100644 --- a/modules/sd_models_utils.py +++ b/modules/sd_models_utils.py @@ -100,26 +100,29 @@ def read_state_dict(checkpoint_file, map_location=None, what:str='model'): # pyl if not os.path.isfile(checkpoint_file): log.error(f'Load dict: path="{checkpoint_file}" not a file') return None + _, extension = os.path.splitext(checkpoint_file) + if extension.lower() == ".ckpt" and shared.opts.sd_disable_ckpt: + log.warning(f"Checkpoint loading disabled: {checkpoint_file}") + return None try: pl_sd = None - with progress.open(checkpoint_file, 'rb', description=f'[cyan]Load {what}: [yellow]{checkpoint_file}', auto_refresh=True, console=console) as f: - _, extension = os.path.splitext(checkpoint_file) - if extension.lower() == ".ckpt" and shared.opts.sd_disable_ckpt: - log.warning(f"Checkpoint loading disabled: {checkpoint_file}") - return None - if shared.opts.stream_load: - if extension.lower() == ".safetensors": - buffer = f.read() - pl_sd = safetensors.torch.load(buffer) - else: - buffer = io.BytesIO(f.read()) - pl_sd = torch.load(buffer, map_location='cpu') - else: - if extension.lower() == ".safetensors": - pl_sd = safetensors.torch.load_file(checkpoint_file, device='cpu') + # safetensors.torch.load_file opens its own handle by path, so wrapping + # with progress.open leaves the bar stuck at 0/total. Skip the wrapper + # on that path; other paths actually read through f and update. + if extension.lower() == ".safetensors" and not shared.opts.stream_load: + pl_sd = safetensors.torch.load_file(checkpoint_file, device='cpu') + else: + with progress.open(checkpoint_file, 'rb', description=f'[cyan]Load {what}: [yellow]{checkpoint_file}', auto_refresh=True, console=console) as f: + if shared.opts.stream_load: + if extension.lower() == ".safetensors": + buffer = f.read() + pl_sd = safetensors.torch.load(buffer) + else: + buffer = io.BytesIO(f.read()) + pl_sd = torch.load(buffer, map_location='cpu') else: pl_sd = torch.load(f, map_location='cpu') - sd = get_state_dict_from_checkpoint(pl_sd) + sd = get_state_dict_from_checkpoint(pl_sd) del pl_sd except Exception as e: errors.display(e, f'Load model: {checkpoint_file}') From 7ae45c59d423babb85b47b9656f2a1673866990d Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Tue, 26 May 2026 21:08:10 +0100 Subject: [PATCH 22/27] refactor(native_transformer): accept precomputed quant_args from caller generic.load_transformer already computes quant_args and quant_type at the top of the function, but native_transformer.load was redoing the same work and emitting a duplicate 'Quantization: module=Model ...' log line milliseconds later. Pass them through instead; fall back to computing them when called directly (model_anima still does so). --- pipelines/generic.py | 2 ++ pipelines/native_transformer.py | 21 +++++++++++++-------- 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/pipelines/generic.py b/pipelines/generic.py index 768077ff7..007d31ad6 100644 --- a/pipelines/generic.py +++ b/pipelines/generic.py @@ -74,6 +74,8 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer dtype=dtype, modules_to_not_convert=modules_to_not_convert, modules_dtype_dict=modules_dtype_dict, + quant_args=quant_args, + quant_type=quant_type, ) elif local_file is not None and local_file.lower().endswith('.safetensors'): log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("diffusers")} args={load_args}') diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index 9fb6ef2c1..ceb485ede 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -176,6 +176,8 @@ def load( dtype=None, modules_to_not_convert: list | None = None, modules_dtype_dict: dict | None = None, + quant_args: dict | None = None, + quant_type: str | None = None, ) -> tuple[object, dict[str, object]]: """Load the transformer (and any bundled siblings) from ``local_file``. @@ -188,7 +190,9 @@ def load( Keyword-only arguments ``allow_quant``, ``dtype``, ``modules_to_not_convert``, and ``modules_dtype_dict`` mirror the corresponding kwargs of :func:`pipelines.generic.load_transformer` so the dispatch from there can - plumb the caller's intent through unchanged. + plumb the caller's intent through unchanged. ``quant_args`` and + ``quant_type`` are precomputed by the caller; when ``None`` they are + derived here via ``model_quant.get_dit_args``. Returns ``(transformer, siblings_dict)``. ``siblings_dict`` is keyed by sibling name and is empty for non-sibling specs, or for sibling specs @@ -206,13 +210,14 @@ def load( f'got "{local_file}"' ) - _, quant_args = model_quant.get_dit_args( - diffusers_cfg, module="Model", device_map=True, - allow_quant=allow_quant, - modules_to_not_convert=modules_to_not_convert, - modules_dtype_dict=modules_dtype_dict, - ) - quant_type = model_quant.get_quant_type(quant_args) + if quant_args is None: + _, quant_args = model_quant.get_dit_args( + diffusers_cfg, module="Model", device_map=True, + allow_quant=allow_quant, + modules_to_not_convert=modules_to_not_convert, + modules_dtype_dict=modules_dtype_dict, + ) + quant_type = model_quant.get_quant_type(quant_args) log.debug(f'Load model: native_transformer reading state_dict cls={spec.cls.__name__} file="{os.path.basename(local_file)}"') state_dict = sd_models.read_state_dict(local_file, what="transformer") From c29168a76602c8ae466c36f2c7ad360ef3d8b50d Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Tue, 26 May 2026 21:54:19 +0100 Subject: [PATCH 23/27] feat(native_transformer): implement per-tensor pre-mode SDNQ quantization The transformer load path now mirrors diffusers' load_model_dict_into_meta: build the component inside init_empty_weights(include_buffers=False) so parameter slots stay meta while computed buffers keep their __init__ values, then iterate the converted state_dict tensor by tensor. Each tensor is dispatched through SDNQQuantizer.check_if_quantized_param and create_quantized_param so Linear/Conv/Embed weights are packed to uint4 in flight, while non-quantizable parameters go through accelerate.set_module_tensor_to_device. Peak memory during transformer load drops from the full bf16 model size to one layer's worth of bf16 plus the running uint4 result. The prior load_state_dict + .to(dtype) + post-walk sdnq_quantize_model sequence remains as the fallback path for siblings and non-SDNQ quant types. --- pipelines/native_transformer.py | 141 ++++++++++++++++++++++++++++---- 1 file changed, 126 insertions(+), 15 deletions(-) diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index ceb485ede..20f3b8145 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -42,9 +42,10 @@ from dataclasses import dataclass, field from typing import Callable import huggingface_hub as hf +import torch from modules import shared, devices, sd_models, model_quant, errors -from modules.logger import log +from modules.logger import log, console DEFAULT_PREFIXES: tuple[str, ...] = ( @@ -376,6 +377,99 @@ def fetch_component_config(repo_id: str, subfolder: str) -> dict: return shared.readfile(local, as_type="dict") +def build_component_quantized( + *, + component_name: str, + state_dict: dict, + config: dict, + cls: type, + quant_args: dict, + dtype, + acceptable_missing: tuple[str, ...], +) -> object: + """Build a component with per-tensor SDNQ quantization during load. + + Mirrors the per-tensor loop in + :func:`diffusers.models.model_loading_utils.load_model_dict_into_meta`: + iterates the (already-converted) state_dict and dispatches each tensor + through ``SDNQQuantizer.check_if_quantized_param`` / + ``create_quantized_param`` so Linear/Conv/Embed weights are packed to + uint4 as they land, while biases, LayerNorms, and other non-quantizable + parameters go through ``accelerate.set_module_tensor_to_device``. + + The component is constructed inside ``init_empty_weights(include_buffers= + False)`` so parameter slots are meta tensors (no full bf16 instantiation + upfront) while computed buffers like ``rope.freqs`` materialize normally + during ``cls.from_config(config)``. After the loop, the quantizer's + ``_process_model_after_weight_loading`` hook attaches + ``quantization_config`` and handles the CPU offload move. + + Caller is responsible for running the converter and prefix stripping + before passing ``state_dict``. + """ + import rich.progress as rp + from accelerate import init_empty_weights + from accelerate.utils import set_module_tensor_to_device + from modules.sdnq.quantizer import SDNQQuantizer + + quantization_config = quant_args.get("quantization_config") + if quantization_config is None: + raise ValueError( + f"Load model: native_transformer {component_name} " + f"per-tensor quantization requires quant_args['quantization_config']" + ) + + target_dtype = dtype if dtype is not None else devices.dtype + quantizer = SDNQQuantizer(quantization_config, pre_quantized=False) + quantizer.torch_dtype = target_dtype + + with init_empty_weights(include_buffers=False): + component = cls.from_config(config) + + quantizer._process_model_before_weight_loading(component, device_map=None) # pylint: disable=protected-access + + target_device = ( + devices.cpu if shared.opts.diffusers_offload_mode != "none" + else devices.device + ) + + expected_keys = set(component.state_dict().keys()) + loaded_keys: set[str] = set() + unexpected: list[str] = [] + total = len(state_dict) + + pbar = rp.Progress( + rp.TextColumn(f'[cyan]Load {component_name}:'), + rp.BarColumn(), + rp.MofNCompleteColumn(), + rp.TaskProgressColumn(), + rp.TimeRemainingColumn(), + rp.TimeElapsedColumn(), + rp.TextColumn('[cyan]{task.description}'), + console=console, + ) + with pbar: + task = pbar.add_task(total=total, description=cls.__name__) + for name, value in state_dict.items(): + if name in expected_keys: + if torch.is_floating_point(value): + value = value.to(target_dtype) + if quantizer.check_if_quantized_param(component, value, name): + quantizer.create_quantized_param(component, value, name, target_device, dtype=target_dtype) + else: + set_module_tensor_to_device(component, name, target_device, value=value, dtype=target_dtype) + loaded_keys.add(name) + else: + unexpected.append(name) + pbar.update(task, advance=1) + + missing = sorted(expected_keys - loaded_keys) + validate_state_dict_load(component_name, missing, unexpected, acceptable_missing) + + component = quantizer._process_model_after_weight_loading(component) # pylint: disable=protected-access + return component + + def build_component( *, component_name: str, @@ -393,10 +487,16 @@ def build_component( """Convert (if needed), instantiate, load weights, dtype-cast, quantize, and offload-place a single component. Raises on any hard failure. + For the transformer component under SDNQ, the per-tensor pre-mode path + in :func:`build_component_quantized` is used so quantization is applied + in flight (one layer's worth of bf16 in memory at a time). All other + cases (siblings, non-quantized loads, NVIDIAModelOptConfig, layerwise + quant) go through the standard load_state_dict + post-quantize path. + ``dtype`` overrides ``devices.dtype`` when supplied; otherwise the global default is used. ``modules_to_not_convert`` and ``modules_dtype_dict`` - are forwarded to :func:`apply_quant` for the transformer component - (ignored for siblings). + are forwarded to :func:`apply_quant` for the post-mode path; pre-mode + receives them via the SDNQConfig in ``quant_args``. """ try: if converter is not None: @@ -404,6 +504,21 @@ def build_component( sd = converter(state_dict) else: sd = state_dict + + if component_name == "transformer" and quant_type == "SDNQConfig": + component = build_component_quantized( + component_name=component_name, + state_dict=sd, + config=config, + cls=cls, + quant_args=quant_args, + dtype=dtype, + acceptable_missing=acceptable_missing, + ) + del sd + devices.torch_gc() + return component + log.debug(f'Load model: native_transformer {component_name} loading keys={len(sd)} cls={cls.__name__}') component = cls.from_config(config) missing, unexpected = component.load_state_dict(sd, strict=False) @@ -475,28 +590,24 @@ def apply_quant( modules_to_not_convert: list | None = None, modules_dtype_dict: dict | None = None, ) -> None: - """Apply SDNQ / layerwise quantization to the bare transformer. + """Apply post-load quantization to a fully-loaded transformer. - SDNQ 'pre' and 'auto' would normally route through ``quantization_config`` - at ``from_pretrained`` time; the native path bypasses that boundary, so - we call the per-module quant path directly. SDNQ 'post' and - ``layerwise_quantization`` go through ``do_post_load_quant`` as usual. - NVIDIAModelOptConfig (TRT) is not supported on this path. + Used as a fallback for cases that don't go through the per-tensor + pre-mode path in :func:`build_component_quantized`: ``layerwise_quantization`` + (via ``do_post_load_quant``), and the no-quant case (no-op). SDNQ post + mode also reaches this path because pre-mode dispatch only triggers + when an SDNQConfig is present (which itself only happens under modes + ``pre`` or ``auto``). ``modules_to_not_convert`` and ``modules_dtype_dict`` are forwarded to :func:`sdnq_quantize_model` so per-call skip lists set by the - caller are honored on the native path. + caller are honored. """ if quant_type == "NVIDIAModelOptConfig": log.warning( "Load model: native_transformer quant=TRT not supported on native path, skipping" ) elif quant_type == "SDNQConfig": - if shared.opts.sdnq_quantize_mode == "pre": - log.info( - "Load model: native_transformer quant=SDNQ pre-mode applied post-load " - "on native path" - ) model_quant.sdnq_quantize_model( transformer, op="transformer", From 18656d1d799cf1ea698264adc00b419529a3b958 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Fri, 29 May 2026 00:07:58 +0100 Subject: [PATCH 24/27] fix(native_transformer): forward load kwargs to from_config load_transformer threads **kwargs into its from_pretrained and from_single_file branches. The native dispatch branch bypasses both loaders and builds via cls.from_config, so those kwargs were dropped there. Thread them through load, build_component, and build_component_quantized into from_config so callers passing extra args alongside native_spec are honored instead of silently dropped. Siblings do not receive them. --- pipelines/generic.py | 1 + pipelines/native_transformer.py | 18 +++++++--- test/test-native-transformer.py | 63 +++++++++++++++++++++++++++++++++ 3 files changed, 78 insertions(+), 4 deletions(-) diff --git a/pipelines/generic.py b/pipelines/generic.py index 007d31ad6..93c6fbf1b 100644 --- a/pipelines/generic.py +++ b/pipelines/generic.py @@ -76,6 +76,7 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer modules_dtype_dict=modules_dtype_dict, quant_args=quant_args, quant_type=quant_type, + **kwargs, ) elif local_file is not None and local_file.lower().endswith('.safetensors'): log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("diffusers")} args={load_args}') diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index 20f3b8145..b598caf99 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -179,6 +179,7 @@ def load( modules_dtype_dict: dict | None = None, quant_args: dict | None = None, quant_type: str | None = None, + **kwargs, ) -> tuple[object, dict[str, object]]: """Load the transformer (and any bundled siblings) from ``local_file``. @@ -193,7 +194,11 @@ def load( :func:`pipelines.generic.load_transformer` so the dispatch from there can plumb the caller's intent through unchanged. ``quant_args`` and ``quant_type`` are precomputed by the caller; when ``None`` they are - derived here via ``model_quant.get_dit_args``. + derived here via ``model_quant.get_dit_args``. Extra ``**kwargs`` are + forwarded to the transformer's ``cls.from_config``, the native path's + construction step (it bypasses ``from_pretrained``/``from_single_file``, + where ``load_transformer`` otherwise routes them). Siblings do not + receive them. Returns ``(transformer, siblings_dict)``. ``siblings_dict`` is keyed by sibling name and is empty for non-sibling specs, or for sibling specs @@ -247,6 +252,7 @@ def load( dtype=effective_dtype, modules_to_not_convert=modules_to_not_convert, modules_dtype_dict=modules_dtype_dict, + **kwargs, ) del transformer_sd devices.torch_gc() @@ -386,6 +392,7 @@ def build_component_quantized( quant_args: dict, dtype, acceptable_missing: tuple[str, ...], + **kwargs, ) -> object: """Build a component with per-tensor SDNQ quantization during load. @@ -424,7 +431,7 @@ def build_component_quantized( quantizer.torch_dtype = target_dtype with init_empty_weights(include_buffers=False): - component = cls.from_config(config) + component = cls.from_config(config, **kwargs) quantizer._process_model_before_weight_loading(component, device_map=None) # pylint: disable=protected-access @@ -483,6 +490,7 @@ def build_component( dtype=None, modules_to_not_convert: list | None = None, modules_dtype_dict: dict | None = None, + **kwargs, ) -> object: """Convert (if needed), instantiate, load weights, dtype-cast, quantize, and offload-place a single component. Raises on any hard failure. @@ -496,7 +504,8 @@ def build_component( ``dtype`` overrides ``devices.dtype`` when supplied; otherwise the global default is used. ``modules_to_not_convert`` and ``modules_dtype_dict`` are forwarded to :func:`apply_quant` for the post-mode path; pre-mode - receives them via the SDNQConfig in ``quant_args``. + receives them via the SDNQConfig in ``quant_args``. Extra ``**kwargs`` + reach ``cls.from_config`` for both construction paths. """ try: if converter is not None: @@ -514,13 +523,14 @@ def build_component( quant_args=quant_args, dtype=dtype, acceptable_missing=acceptable_missing, + **kwargs, ) del sd devices.torch_gc() return component log.debug(f'Load model: native_transformer {component_name} loading keys={len(sd)} cls={cls.__name__}') - component = cls.from_config(config) + component = cls.from_config(config, **kwargs) missing, unexpected = component.load_state_dict(sd, strict=False) validate_state_dict_load(component_name, missing, unexpected, acceptable_missing) del sd diff --git a/test/test-native-transformer.py b/test/test-native-transformer.py index b0a524b29..d9f9f06d4 100644 --- a/test/test-native-transformer.py +++ b/test/test-native-transformer.py @@ -422,6 +422,19 @@ class MockMiniTransformer(torch.nn.Module): self.rope.register_buffer('freqs', torch.zeros(dim)) +class MockKwargsTransformer(MockMiniTransformer): + """Records the kwargs from_config received. Mirrors diffusers from_config, + which accepts **kwargs (config overrides); lets a test assert the native + load path forwards caller kwargs to construction instead of dropping them.""" + + last_kwargs: dict = {} + + @classmethod + def from_config(cls, config: dict, **kwargs) -> 'MockKwargsTransformer': + cls.last_kwargs = dict(kwargs) + return cls(dim=config['dim']) + + def write_fixture(state_dict_keys: dict, fd: int, path: str) -> str: os.close(fd) safetensors.torch.save_file(state_dict_keys, path) @@ -484,6 +497,55 @@ def test_load_end_to_end_with_bfl_prefix_no_converter(): os.unlink(path) +def test_load_forwards_kwargs_to_from_config(): + """Caller **kwargs reach cls.from_config (the native path's only + construction step, since it bypasses from_pretrained/from_single_file). + Guards against silently dropping args a caller passes alongside + native_spec.""" + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + dim = 8 + raw = { + 'model.diffusion_model.in_proj.weight': torch.randn(dim, dim), + 'model.diffusion_model.in_proj.bias': torch.zeros(dim), + 'model.diffusion_model.out_proj.weight': torch.randn(dim, dim), + 'model.diffusion_model.out_proj.bias': torch.zeros(dim), + } + write_fixture(raw, fd, path) + + orig_fetch = nt.fetch_component_config + nt.fetch_component_config = lambda repo, sub: {'dim': dim} + from modules import model_quant + orig_get_dit = model_quant.get_dit_args + orig_get_qtype = model_quant.get_quant_type + orig_do_post = model_quant.do_post_load_quant + model_quant.get_dit_args = lambda *a, **k: ({}, {}) + model_quant.get_quant_type = lambda *a, **k: None + model_quant.do_post_load_quant = lambda *a, **k: None + + MockKwargsTransformer.last_kwargs = {} + try: + spec = nt.TransformerSpec(cls=MockKwargsTransformer) + transformer, _ = nt.load( + local_file=path, + repo_id='fake/repo', + spec=spec, + diffusers_cfg={}, + low_cpu_mem_usage=True, + ) + finally: + nt.fetch_component_config = orig_fetch + model_quant.get_dit_args = orig_get_dit + model_quant.get_quant_type = orig_get_qtype + model_quant.do_post_load_quant = orig_do_post + + assert MockKwargsTransformer.last_kwargs == {'low_cpu_mem_usage': True} + assert isinstance(transformer, MockKwargsTransformer) + finally: + if os.path.exists(path): + os.unlink(path) + + def test_load_end_to_end_with_sibling_partition(): """Bundled-sibling case: file carries both transformer and sibling weights, sibling_classes supplies the runtime sibling class, partition routes each @@ -711,6 +773,7 @@ def run_all(): cat = category('load') for fn in [ test_load_end_to_end_with_bfl_prefix_no_converter, + test_load_forwards_kwargs_to_from_config, test_load_end_to_end_with_sibling_partition, test_load_raises_on_missing_sibling_class, test_load_rejects_non_safetensors, From 25fa8e0f15a9f96d9731cad6a8365486b0561259 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Fri, 29 May 2026 00:20:11 +0100 Subject: [PATCH 25/27] refactor(ui): shorten unet-sync handlers, align load logs - ui_settings: rename to sync_checkpoint_unet / sync_checkpoint_unet_forced - sd_models: cross-arch unet reset log uses key=value house style - generic: native transformer dispatch logs loader=native --- modules/sd_models.py | 2 +- modules/ui_settings.py | 10 +++++----- pipelines/generic.py | 2 +- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/modules/sd_models.py b/modules/sd_models.py index e4f48eb70..d61059a4e 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -1420,7 +1420,7 @@ def reload_model_weights(sd_model=None, info=None, op='model', force=False, revi new_pipeline = None new_class = getattr(new_pipeline, '__name__', None) if new_class is not None and new_class != old_class: - log.info(f'Load model: pipeline changed {old_class}->{new_class}, resetting sd_unet from "{shared.opts.sd_unet}" to Default') + log.info(f'Load model: pipeline cls={old_class} changed={new_class} unet="{shared.opts.sd_unet}" set to default') shared.opts.data["sd_unet"] = 'Default' sd_unet.loaded_unet = None if sd_model is None: # previous model load failed diff --git a/modules/ui_settings.py b/modules/ui_settings.py index 09d2e83a4..aa9e6a4a9 100644 --- a/modules/ui_settings.py +++ b/modules/ui_settings.py @@ -391,7 +391,7 @@ def create_quicksettings(interfaces): if shared.opts.notification_audio_enable and os.path.exists(os.path.join(paths.script_path, shared.opts.notification_audio_path)): gr.Audio(interactive=False, value=os.path.join(paths.script_path, shared.opts.notification_audio_path), elem_id="audio_notification", visible=False) - def change_checkpoint_with_unet_sync(value, progress=False, force=False): + def sync_checkpoint_unet(value, progress=False, force=False): checkpoint_update, settings_text = run_settings_single(value, key='sd_model_checkpoint', progress=progress, force=force) return checkpoint_update, get_value_for_setting('sd_unet'), settings_text @@ -412,7 +412,7 @@ def create_quicksettings(interfaces): progress_flag = info.refresh is not None if k == 'sd_model_checkpoint': def fn(value, progress=progress_flag): - return change_checkpoint_with_unet_sync(value, progress=progress) + return sync_checkpoint_unet(value, progress=progress) outputs = [component, shared.settings_components['sd_unet'], text_settings] else: def fn(value, k=k, progress=progress_flag): @@ -426,12 +426,12 @@ def create_quicksettings(interfaces): show_progress='full' if info.refresh is not None else 'hidden', ) - def set_checkpoint_sync_unet(value, _dummy): - return change_checkpoint_with_unet_sync(value, force=True) + def sync_checkpoint_unet_forced(value, _dummy): + return sync_checkpoint_unet(value, force=True) button_set_checkpoint = gr.Button('Change model', elem_id='change_checkpoint', visible=False) button_set_checkpoint.click( - fn=set_checkpoint_sync_unet, + fn=sync_checkpoint_unet_forced, _js="consumeDesiredCheckpointName", inputs=[shared.settings_components['sd_model_checkpoint'], dummy_component], outputs=[shared.settings_components['sd_model_checkpoint'], shared.settings_components['sd_unet'], text_settings], diff --git a/pipelines/generic.py b/pipelines/generic.py index 93c6fbf1b..4b8cc61ec 100644 --- a/pipelines/generic.py +++ b/pipelines/generic.py @@ -67,7 +67,7 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer transformer = model_quant.do_post_load_quant(transformer, allow=quant_type is not None) elif local_file is not None and local_file.lower().endswith('.safetensors') and native_spec is not None: from pipelines import native_transformer - log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader=native_transformer args={load_args}') + log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader=native args={load_args}') transformer, _ = native_transformer.load( local_file, repo_id, native_spec, load_config, allow_quant=allow_quant, From 23c06c39d6886ad981d82990710312c57472a058 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Fri, 29 May 2026 01:06:43 +0100 Subject: [PATCH 26/27] refactor(sd_unet): route transformer override through reload_model_weights The custom-transformer apply (DiT branch) and the revert-to-default path now call reload_model_weights(force=True) instead of load_diffuser, so both go through the same managed reload: unload-before-rebuild for a lower VRAM peak and consistent job/checkpoint bookkeeping. Gate the cross-arch unet reset on an actual checkpoint change so a same-checkpoint reload skips the pipeline-class comparison, which could false-positive and clear the override when the model was left in an img2img/inpaint variant after an interrupted generation. --- modules/sd_models.py | 4 +++- modules/sd_unet.py | 4 ++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/modules/sd_models.py b/modules/sd_models.py index d61059a4e..206bdea09 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -1412,7 +1412,9 @@ def reload_model_weights(sd_model=None, info=None, op='model', force=False, revi jobid = shared.state.begin('Load model') if sd_model is None: sd_model = model_data.sd_model if op == 'model' or op == 'dict' else model_data.sd_refiner - if op == 'model' and sd_model is not None and shared.opts.sd_unet not in (None, 'Default', 'None'): + loaded_ckpt = getattr(sd_model, 'sd_checkpoint_info', None) if sd_model is not None else None + changed_checkpoint = loaded_ckpt is None or checkpoint_info is None or loaded_ckpt.filename != checkpoint_info.filename + if op == 'model' and sd_model is not None and changed_checkpoint and shared.opts.sd_unet not in (None, 'Default', 'None'): old_class = type(sd_model).__name__ try: new_pipeline, _ = sd_detect.detect_pipeline(checkpoint_info.path, op) diff --git a/modules/sd_unet.py b/modules/sd_unet.py index f8c42bb45..4143acc52 100644 --- a/modules/sd_unet.py +++ b/modules/sd_unet.py @@ -54,7 +54,7 @@ def load_unet(model, repo_id: str | None = None): return log.info(f'Load module: type=UNet name="Default" (was="{loaded_unet}") reverting to base transformer') loaded_unet = shared.opts.sd_unet - sd_models.load_diffuser() + sd_models.reload_model_weights(force=True) return if shared.opts.sd_unet not in list(unet_dict): @@ -83,7 +83,7 @@ def load_unet(model, repo_id: str | None = None): model.prior_pipe.text_encoder = prior_text_encoder.to(devices.device, dtype=devices.dtype) elif any([m in model.__class__.__name__ for m in dit_models]) or hasattr(model, 'transformer'): # noqa: C419 # pylint: disable=use-a-generator loaded_unet = shared.opts.sd_unet - sd_models.load_diffuser() # TODO model load: force-reloading entire model as loading transformers only leads to massive memory usage + sd_models.reload_model_weights(force=True) # full reload: in-place transformer swap leaks memory else: if not hasattr(model, 'unet') or model.unet is None: log.error('Load module: type=UNET not found in current model') From 77913c7c0a7e4e7410c78589069ea1a6acdc70d5 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Fri, 29 May 2026 02:52:18 +0100 Subject: [PATCH 27/27] docs: cleanup - generic.load_transformer: shorter dispatch description, no internal taxonomy - native_transformer.load: trim kwargs forwarding wording - sd_unet revert: one-line comment - test: shorten kwargs-forwarding docstrings --- modules/sd_unet.py | 5 +---- pipelines/generic.py | 9 +++------ pipelines/native_transformer.py | 11 ++++------- test/test-native-transformer.py | 12 +++++------- 4 files changed, 13 insertions(+), 24 deletions(-) diff --git a/modules/sd_unet.py b/modules/sd_unet.py index 4143acc52..a371fe8ea 100644 --- a/modules/sd_unet.py +++ b/modules/sd_unet.py @@ -46,10 +46,7 @@ def load_unet(model, repo_id: str | None = None): return if shared.opts.sd_unet == 'Default' or shared.opts.sd_unet == 'None': - # If a custom UNET / transformer was previously loaded, switching back - # to 'Default' has to actually revert the override. Without this, the - # model keeps the prior finetune transformer in memory and the - # dropdown looks "stuck". + # Switching back to Default reverts a previously-loaded custom transformer. if loaded_unet in (None, 'Default', 'None'): return log.info(f'Load module: type=UNet name="Default" (was="{loaded_unet}") reverting to base transformer') diff --git a/pipelines/generic.py b/pipelines/generic.py index 4b8cc61ec..8ba75214b 100644 --- a/pipelines/generic.py +++ b/pipelines/generic.py @@ -23,12 +23,9 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer """Load a DiT transformer from the base repo, or from a user-selected single file when the UNET dropdown (``shared.opts.sd_unet``) is set. - When ``native_spec`` is supplied and a .safetensors override is selected, - dispatches to :func:`pipelines.native_transformer.load` so the per-arch - spec (multi-prefix detection, optional converter, optional sibling - partitioning, forbidden markers) drives the load. Pipelines without a - spec continue to use the legacy ``from_single_file`` path; this preserves - behavior for Mode D arches until they explicitly opt in. + With ``native_spec`` set and a .safetensors override selected, dispatches + to :func:`pipelines.native_transformer.load`. Without a spec, a single-file + override falls back to ``from_single_file``. """ if shared.state.interrupted: return None diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index b598caf99..b52326e5a 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -191,13 +191,10 @@ def load( Keyword-only arguments ``allow_quant``, ``dtype``, ``modules_to_not_convert``, and ``modules_dtype_dict`` mirror the corresponding kwargs of - :func:`pipelines.generic.load_transformer` so the dispatch from there can - plumb the caller's intent through unchanged. ``quant_args`` and - ``quant_type`` are precomputed by the caller; when ``None`` they are - derived here via ``model_quant.get_dit_args``. Extra ``**kwargs`` are - forwarded to the transformer's ``cls.from_config``, the native path's - construction step (it bypasses ``from_pretrained``/``from_single_file``, - where ``load_transformer`` otherwise routes them). Siblings do not + :func:`pipelines.generic.load_transformer` and are forwarded unchanged. + ``quant_args`` and ``quant_type`` are precomputed by the caller; when + ``None`` they are derived here via ``model_quant.get_dit_args``. Extra + ``**kwargs`` reach the transformer's ``cls.from_config``; siblings do not receive them. Returns ``(transformer, siblings_dict)``. ``siblings_dict`` is keyed by diff --git a/test/test-native-transformer.py b/test/test-native-transformer.py index d9f9f06d4..538ce84a5 100644 --- a/test/test-native-transformer.py +++ b/test/test-native-transformer.py @@ -423,9 +423,9 @@ class MockMiniTransformer(torch.nn.Module): class MockKwargsTransformer(MockMiniTransformer): - """Records the kwargs from_config received. Mirrors diffusers from_config, - which accepts **kwargs (config overrides); lets a test assert the native - load path forwards caller kwargs to construction instead of dropping them.""" + """Records the kwargs from_config received, so a test can assert the native + path forwards caller kwargs to construction. Mirrors diffusers from_config, + which accepts **kwargs.""" last_kwargs: dict = {} @@ -498,10 +498,8 @@ def test_load_end_to_end_with_bfl_prefix_no_converter(): def test_load_forwards_kwargs_to_from_config(): - """Caller **kwargs reach cls.from_config (the native path's only - construction step, since it bypasses from_pretrained/from_single_file). - Guards against silently dropping args a caller passes alongside - native_spec.""" + """Caller **kwargs reach cls.from_config through the native load path + rather than being dropped.""" fd, path = tempfile.mkstemp(suffix='.safetensors') try: dim = 8