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

Merge pull request #4876 from vladmandic/feat/native-transformer-loader

Browse Source 
Feat/native transformer loader
This commit is contained in:
Vladimir Mandic
2026-06-01 06:59:47 +02:00
committed by GitHub
parent 76763da29d 77913c7c0a
commit 24aced130d
63 changed files with 2088 additions and 493 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/lora/lora_load.py
+1 -1
View File
@@ -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',
modules/lora/native_loader.py → modules/lora/native_adapter.py
+8 -7
View File
@@ -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_"
modules/sd_models.py
+13
View File
@@ -1412,6 +1412,19 @@ 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
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)
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 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
current_checkpoint_info = None
else:
modules/sd_models_utils.py
+19 -16
View File
@@ -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}')
modules/sd_unet.py
+7 -1
View File
@@ -46,6 +46,12 @@ def load_unet(model, repo_id: str | None = None):
return
if shared.opts.sd_unet == 'Default' or shared.opts.sd_unet == 'None':
# 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')
loaded_unet = shared.opts.sd_unet
sd_models.reload_model_weights(force=True)
return
if shared.opts.sd_unet not in list(unet_dict):
@@ -74,7 +80,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')
modules/ui_settings.py
+20 -4
View File
@@ -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 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
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 sync_checkpoint_unet(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 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=lambda value, _: run_settings_single(value, key='sd_model_checkpoint', force=True),
fn=sync_checkpoint_unet_forced,
_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(
pipelines/anima/__init__.py
+40
View File
@@ -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',
),
),
)
pipelines/anima/anima_lora.py
+37 -37
View File
@@ -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,
pipelines/anima/anima_transformer.py
-217
View File
@@ -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)
pipelines/bria/__init__.py
+27
View File
@@ -0,0 +1,27 @@
"""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.
- Bria FIBO (and FIBO Edit) use the upstream
:class:`diffusers.BriaFiboTransformer2DModel`, which also lacks a
``SINGLE_FILE_LOADABLE_CLASSES`` entry.
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)
pipelines/chroma/__init__.py
+18
View File
@@ -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,
)
pipelines/chroma/chroma_lora.py
+24 -24
View File
@@ -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),
)
pipelines/chrono/__init__.py
+15
View File
@@ -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)
pipelines/cogview/__init__.py
+15
View File
@@ -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)
pipelines/ernie/__init__.py
+16
View File
@@ -0,0 +1,16 @@
"""ERNIE-Image pipeline package.
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.
"""
import diffusers
from pipelines.native_transformer import TransformerSpec
ERNIE_SPEC = TransformerSpec(cls=diffusers.ErnieImageTransformer2DModel)
pipelines/ernie/ernie_lora.py
+23 -23
View File
@@ -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),
)
pipelines/f_lite/__init__.py
+5 -1
View File
@@ -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"]
pipelines/flux/flux2_lora.py
+39 -39
View File
@@ -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,
pipelines/flux2_klein/__init__.py
+26
View File
@@ -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,
)
pipelines/generic.py
+21 -1
View File
@@ -19,7 +19,14 @@ 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.
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
transformer = None
@@ -55,6 +62,19 @@ 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 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,
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}')
if dtype is not None:
pipelines/glm/__init__.py
+15
View File
@@ -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)
pipelines/hunyuandit/__init__.py
+14
View File
@@ -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)
pipelines/hyimage/__init__.py
+19
View File
@@ -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)
pipelines/joy/__init__.py
+15
View File
@@ -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)
pipelines/kandinsky3/__init__.py
+13
View File
@@ -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')
pipelines/kandinsky5/__init__.py
+12
View File
@@ -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)
pipelines/lens/__init__.py
+5
View File
@@ -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",
pipelines/longcat/__init__.py
+15
View File
@@ -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)
pipelines/model_anima.py
+13 -25
View File
@@ -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')
pipelines/model_bria.py
+4 -1
View File
@@ -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(
pipelines/model_chroma.py
+2 -1
View File
@@ -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(
pipelines/model_chrono.py
+2 -1
View File
@@ -21,7 +21,8 @@ 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")
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")
try:
pipelines/model_cogview.py
+4 -2
View File
@@ -14,7 +14,8 @@ 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")
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")
pipe = diffusers.CogView3PlusPipeline.from_pretrained(
@@ -40,7 +41,8 @@ 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")
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)
pipe = diffusers.CogView4Pipeline.from_pretrained(
pipelines/model_ernie.py
+2
View File
@@ -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,
pipelines/model_flite.py
+1 -1
View File
@@ -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(
pipelines/model_flux2_klein.py
+2 -1
View File
@@ -15,7 +15,8 @@ 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)
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)
text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3ForCausalLM, load_config=diffusers_load_config)
pipelines/model_glm.py
+3 -1
View File
@@ -95,10 +95,12 @@ 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,
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)
pipelines/model_hunyuandit.py
+2 -1
View File
@@ -19,7 +19,8 @@ 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)
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
pipelines/model_hyimage.py
+2 -1
View File
@@ -16,7 +16,8 @@ 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")
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)
pipelines/model_joy.py
+2
View File
@@ -14,10 +14,12 @@ 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,
load_config=diffusers_load_config,
native_spec=JOY_SPEC,
)
text_encoder = generic.load_text_encoder(
repo_id,
pipelines/model_kandinsky.py
+4 -2
View File
@@ -50,7 +50,8 @@ 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")
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)
pipe = diffusers.Kandinsky3Pipeline.from_pretrained(
@@ -83,7 +84,8 @@ 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)
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)
if 'I2I' in repo_id:
pipelines/model_lens.py
+1 -1
View File
@@ -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(
pipelines/model_longcat.py
+2 -1
View File
@@ -14,7 +14,8 @@ 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)
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)
pipelines/model_nucleus.py
+2
View File
@@ -14,10 +14,12 @@ 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,
load_config=diffusers_load_config,
native_spec=NUCLEUS_SPEC,
)
text_encoder = generic.load_text_encoder(
repo_id,
pipelines/model_ovis.py
+2 -1
View File
@@ -14,7 +14,8 @@ 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)
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)
pipe = diffusers.OvisImagePipeline.from_pretrained(
pipelines/model_pixart.py
+2 -1
View File
@@ -24,7 +24,8 @@ 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)
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)
pipe = diffusers.PixArtSigmaPipeline.from_pretrained(
pipelines/model_prx.py
+2 -1
View File
@@ -14,7 +14,8 @@ 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)
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)
pipe = diffusers.PRXPipeline.from_pretrained(
pipelines/model_qwen.py
+2
View File
@@ -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'
pipelines/model_step1x_edit.py
+2 -1
View File
@@ -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,
pipelines/model_vibe.py
+2
View File
@@ -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,
pipelines/model_wanai.py
+7 -49
View File
@@ -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)
pipelines/native_transformer.py
+624
View File
@@ -0,0 +1,624 @@
"""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`. Each pipeline
defines its spec in ``pipelines/<arch>/__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:
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 ``<subfolder>/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
import torch
from modules import shared, devices, sd_models, model_quant, errors
from modules.logger import log, console
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], ...] = ()
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).
Used by callers (notably :func:`pipelines.generic.load_transformer`) when
a pipeline does not supply a custom ``TransformerSpec`` of its own.
"""
return TransformerSpec(cls=cls, converter=auto_pickup_converter(cls))
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,
*,
allow_quant: bool = True,
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,
**kwargs,
) -> 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.
Keyword-only arguments ``allow_quant``, ``dtype``, ``modules_to_not_convert``,
and ``modules_dtype_dict`` mirror the corresponding kwargs of
: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
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}"'
)
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")
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 '{}'}"
)
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",
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,
dtype=effective_dtype,
modules_to_not_convert=modules_to_not_convert,
modules_dtype_dict=modules_dtype_dict,
**kwargs,
)
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,
dtype=effective_dtype,
)
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 ``<subfolder>/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_quantized(
*,
component_name: str,
state_dict: dict,
config: dict,
cls: type,
quant_args: dict,
dtype,
acceptable_missing: tuple[str, ...],
**kwargs,
) -> 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, **kwargs)
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,
state_dict: dict,
config: dict,
cls: type,
converter: Callable[[dict], dict] | None,
acceptable_missing: tuple[str, ...],
quant_args: dict,
quant_type: str | None,
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.
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 post-mode path; pre-mode
receives them via the SDNQConfig in ``quant_args``. Extra ``**kwargs``
reach ``cls.from_config`` for both construction paths.
"""
try:
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
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,
**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, **kwargs)
missing, unexpected = component.load_state_dict(sd, strict=False)
validate_state_dict_load(component_name, missing, unexpected, acceptable_missing)
del sd
devices.torch_gc()
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")
raise
if component_name == "transformer":
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)
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,
modules_to_not_convert: list | None = None,
modules_dtype_dict: dict | None = None,
) -> None:
"""Apply post-load quantization to a fully-loaded transformer.
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.
"""
if quant_type == "NVIDIAModelOptConfig":
log.warning(
"Load model: native_transformer quant=TRT not supported on native path, skipping"
)
elif quant_type == "SDNQConfig":
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)
pipelines/nucleus/__init__.py
+15
View File
@@ -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)
pipelines/ovis/__init__.py
+15
View File
@@ -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)
pipelines/pixart/__init__.py
+14
View File
@@ -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)
pipelines/prx/__init__.py
+14
View File
@@ -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)
pipelines/qwen/__init__.py
+17
View File
@@ -1,2 +1,19 @@
"""Qwen-Image pipeline package.
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
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)
pipelines/step1x/__init__.py
+6
View File
@@ -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",
pipelines/vibe/__init__.py
+6 -1
View File
@@ -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"]
pipelines/z_image/zimage_lora.py
+24 -24
View File
@@ -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),
)
test/test-anima-native-adapters.py
+2 -2
View File
@@ -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 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 /
test/test-flux2-native-adapters.py
+3 -3
View File
@@ -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.
@@ -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 == [
test/test-native-transformer.py
+799
View File
@@ -0,0 +1,799 @@
#!/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
- ``make_default_spec`` default-spec synthesis with diffusers converter pickup
- ``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=(),
)
# ============================================================
# make_default_spec
# ============================================================
class FakeTransformer:
"""Minimal stand-in for a diffusers transformer class."""
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 # no diffusers entry for FakeTransformer
assert spec.siblings == {}
assert spec.forbidden_markers == ()
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_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
# ============================================================
# 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))
class MockKwargsTransformer(MockMiniTransformer):
"""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 = {}
@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)
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_forwards_kwargs_to_from_config():
"""Caller **kwargs reach cls.from_config through the native load path
rather than being dropped."""
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
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('=== make_default_spec ===')
cat = category('default_spec')
for fn in [
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)
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_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,
]:
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)