mirror of
https://github.com/vladmandic/automatic
synced 2026-09-18 16:54:33 +02:00
feat(pipelines): add native_transformer for single-file DiT loading
Read .safetensors, strip prefix, partition siblings, run optional converter, from_config + load_state_dict + validate, dtype/quant/offload. Per-arch knobs via TransformerSpec(cls, subfolder, prefixes, converter, siblings, acceptable_missing, forbidden_markers). SiblingSpec covers bundled components like Anima's llm_adapter. 36 offline tests.
This commit is contained in:
CalamitousFelicitousness
@@ -0,0 +1,459 @@
|
||||
"""Generic native loader for DiT transformers and bundled sibling components.
|
||||
|
||||
Loads a single-safetensors file into a diffusers (or custom) transformer class
|
||||
when the user selects an override via the UNET dropdown (``shared.opts.sd_unet``).
|
||||
Bypasses :func:`diffusers.loaders.FromOriginalModelMixin.from_single_file` so
|
||||
sdnext owns prefix detection, optional sibling partitioning, dtype/quant/offload
|
||||
handling, and explicit validation of missing/unexpected keys.
|
||||
|
||||
The per-arch knobs are captured in :class:`TransformerSpec`. Arches register a
|
||||
spec at import time via :func:`register`; arches without a registration get a
|
||||
default spec that handles BFL-style ``model.diffusion_model.`` prefix stripping
|
||||
and opportunistically picks up a diffusers converter from
|
||||
``SINGLE_FILE_LOADABLE_CLASSES`` if the class has one.
|
||||
|
||||
Algorithm:
|
||||
|
||||
1. Read the safetensors state dict (.gguf and .pth are rejected up front).
|
||||
2. Detect and strip one of the spec's known prefixes (raises on mixed prefixes).
|
||||
3. Check forbidden markers (catches structural mismatches like Cosmos 1.0 keys
|
||||
in a Cosmos 2.0 loader).
|
||||
4. Partition off sibling component keys (e.g. Anima's bundled ``llm_adapter.*``).
|
||||
5. Run the spec's converter if present (else pass through unchanged).
|
||||
6. Fetch ``<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
|
||||
|
||||
from modules import shared, devices, sd_models, model_quant, errors
|
||||
from modules.logger import log
|
||||
|
||||
|
||||
DEFAULT_PREFIXES: tuple[str, ...] = (
|
||||
"model.diffusion_model.",
|
||||
"diffusion_model.",
|
||||
"net.",
|
||||
)
|
||||
DEFAULT_ACCEPTABLE_MISSING: tuple[str, ...] = (
|
||||
"rope.",
|
||||
"pos_embedder.",
|
||||
"learnable_pos_embed.",
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SiblingSpec:
|
||||
"""Describes a non-transformer component that may ship inline in the same
|
||||
safetensors as the transformer (e.g. Anima's ``llm_adapter``).
|
||||
|
||||
``subfolder`` names the base repo subfolder holding the canonical config and
|
||||
weights when the sibling is NOT bundled inline; ``inline_prefix`` is the
|
||||
key prefix that identifies the sibling's weights within the bundled file
|
||||
(after the transformer's prefix has already been stripped).
|
||||
"""
|
||||
|
||||
subfolder: str
|
||||
inline_prefix: str
|
||||
acceptable_missing: tuple[str, ...] = ()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TransformerSpec:
|
||||
"""Per-arch configuration for the native loader.
|
||||
|
||||
Most arches only need to override ``cls`` (and rely on the default
|
||||
prefixes, no converter, no siblings). Arches with bundled-sibling
|
||||
components (Anima) or unusual key conventions (custom converters,
|
||||
Cosmos-style structural markers) override the relevant fields.
|
||||
"""
|
||||
|
||||
cls: type
|
||||
subfolder: str = "transformer"
|
||||
prefixes: tuple[str, ...] = DEFAULT_PREFIXES
|
||||
converter: Callable[[dict], dict] | None = None
|
||||
siblings: dict[str, SiblingSpec] = field(default_factory=dict)
|
||||
acceptable_missing: tuple[str, ...] = DEFAULT_ACCEPTABLE_MISSING
|
||||
forbidden_markers: tuple[tuple[str, str], ...] = ()
|
||||
|
||||
|
||||
REGISTRY: dict[type, TransformerSpec] = {}
|
||||
|
||||
|
||||
def register(cls: type, spec: TransformerSpec | None = None) -> None:
|
||||
"""Register a transformer class with an explicit spec, or with the default
|
||||
spec if ``spec`` is None. Idempotent: re-registering the same class
|
||||
replaces the previous entry.
|
||||
"""
|
||||
if spec is None:
|
||||
spec = TransformerSpec(cls=cls)
|
||||
if spec.cls is not cls:
|
||||
raise ValueError(f"register: spec.cls ({spec.cls.__name__}) does not match cls ({cls.__name__})")
|
||||
REGISTRY[cls] = spec
|
||||
|
||||
|
||||
def lookup(cls: type) -> TransformerSpec:
|
||||
"""Return the registered spec for ``cls``, or synthesize a default one.
|
||||
|
||||
The synthesized default opportunistically pulls a converter from diffusers'
|
||||
``SINGLE_FILE_LOADABLE_CLASSES`` table if one exists for the class name and
|
||||
is not a pass-through no-op lambda.
|
||||
"""
|
||||
if cls in REGISTRY:
|
||||
return REGISTRY[cls]
|
||||
converter = auto_pickup_converter(cls)
|
||||
return TransformerSpec(cls=cls, converter=converter)
|
||||
|
||||
|
||||
def auto_pickup_converter(cls: type) -> Callable[[dict], dict] | None:
|
||||
"""Pull a checkpoint converter from diffusers for ``cls`` when one exists.
|
||||
|
||||
Skipped for the no-op identity lambda some classes register
|
||||
(notably ``QwenImageTransformer2DModel``), because using it would silently
|
||||
accept whatever key naming the file happens to have.
|
||||
"""
|
||||
try:
|
||||
from diffusers.loaders.single_file_model import SINGLE_FILE_LOADABLE_CLASSES
|
||||
except ImportError:
|
||||
return None
|
||||
entry = SINGLE_FILE_LOADABLE_CLASSES.get(cls.__name__)
|
||||
if entry is None:
|
||||
return None
|
||||
fn = entry.get("checkpoint_mapping_fn")
|
||||
if fn is None or is_noop_converter(fn):
|
||||
return None
|
||||
return fn
|
||||
|
||||
|
||||
def is_noop_converter(fn: Callable) -> bool:
|
||||
"""Detect ``lambda checkpoint, **kwargs: checkpoint`` and equivalents.
|
||||
|
||||
Strips inline ``#`` comments from the source line before inspecting the
|
||||
body, so that diagnostic markers like ``# noqa`` on the lambda's source
|
||||
line do not defeat the detection.
|
||||
"""
|
||||
try:
|
||||
import inspect
|
||||
src = inspect.getsource(fn).strip()
|
||||
except (OSError, TypeError):
|
||||
return False
|
||||
if "lambda" not in src:
|
||||
return False
|
||||
if "#" in src:
|
||||
src = src.split("#", 1)[0].rstrip()
|
||||
body = src.split(":", 1)[-1].strip().rstrip(",").rstrip(")")
|
||||
return body.endswith("checkpoint")
|
||||
|
||||
|
||||
def resolve_path() -> str | None:
|
||||
"""Return the absolute path of the UNET dropdown selection, or None if
|
||||
no selection is active or the file is unresolvable.
|
||||
"""
|
||||
sel = shared.opts.sd_unet
|
||||
if sel is None or sel in ("Default", "None"):
|
||||
return None
|
||||
from modules import sd_unet
|
||||
if sel not in list(sd_unet.unet_dict):
|
||||
log.error(f'Load module: type=transformer file="{sel}" not found')
|
||||
return None
|
||||
path = sd_unet.unet_dict[sel]
|
||||
if not os.path.exists(path):
|
||||
log.error(f'Load module: type=transformer path="{path}" does not exist')
|
||||
return None
|
||||
return path
|
||||
|
||||
|
||||
def load(
|
||||
local_file: str,
|
||||
repo_id: str,
|
||||
spec: TransformerSpec,
|
||||
diffusers_cfg: dict | None = None,
|
||||
sibling_classes: dict[str, type] | None = None,
|
||||
) -> tuple[object, dict[str, object]]:
|
||||
"""Load the transformer (and any bundled siblings) from ``local_file``.
|
||||
|
||||
``sibling_classes`` supplies the runtime class for each sibling named in
|
||||
``spec.siblings``. Required when a sibling has a dynamic class (e.g.
|
||||
Anima's ``AnimaLLMAdapter`` is loaded from remote_code at runtime).
|
||||
Missing sibling classes raise ``ValueError`` if the corresponding sibling
|
||||
keys are present in the bundled file.
|
||||
|
||||
Returns ``(transformer, siblings_dict)``. ``siblings_dict`` is keyed by
|
||||
sibling name and is empty for non-sibling specs, or for sibling specs
|
||||
whose keys are absent from the bundled file.
|
||||
"""
|
||||
if diffusers_cfg is None:
|
||||
diffusers_cfg = {}
|
||||
if sibling_classes is None:
|
||||
sibling_classes = {}
|
||||
|
||||
t0 = time.time()
|
||||
if not local_file.lower().endswith(".safetensors"):
|
||||
raise ValueError(
|
||||
f"Load model: type={spec.cls.__name__} custom transformer requires .safetensors, "
|
||||
f'got "{local_file}"'
|
||||
)
|
||||
|
||||
_, quant_args = model_quant.get_dit_args(
|
||||
diffusers_cfg, module="Model", device_map=True, allow_quant=True,
|
||||
)
|
||||
quant_type = model_quant.get_quant_type(quant_args)
|
||||
|
||||
state_dict = sd_models.read_state_dict(local_file, what="transformer")
|
||||
state_dict = strip_prefix(state_dict, spec.prefixes, spec.cls.__name__)
|
||||
check_forbidden_markers(state_dict, spec.forbidden_markers, spec.cls.__name__, local_file)
|
||||
transformer_sd, sibling_sds = partition_siblings(state_dict, spec.siblings)
|
||||
del state_dict
|
||||
|
||||
sibling_counts = {name: len(sd) for name, sd in sibling_sds.items() if sd}
|
||||
log.info(
|
||||
f'Load model: type={spec.cls.__name__} custom="{os.path.basename(local_file)}" '
|
||||
f"transformer_keys={len(transformer_sd)} siblings={sibling_counts or '{}'}"
|
||||
)
|
||||
|
||||
transformer_cfg = fetch_component_config(repo_id, spec.subfolder)
|
||||
transformer = build_component(
|
||||
component_name="transformer",
|
||||
state_dict=transformer_sd,
|
||||
config=transformer_cfg,
|
||||
cls=spec.cls,
|
||||
converter=spec.converter,
|
||||
acceptable_missing=spec.acceptable_missing,
|
||||
quant_args=quant_args,
|
||||
quant_type=quant_type,
|
||||
)
|
||||
del transformer_sd
|
||||
devices.torch_gc()
|
||||
|
||||
loaded_siblings: dict[str, object] = {}
|
||||
for name, sibling_sd in sibling_sds.items():
|
||||
if not sibling_sd:
|
||||
continue
|
||||
sibling_spec = spec.siblings[name]
|
||||
sibling_cls = sibling_classes.get(name)
|
||||
if sibling_cls is None:
|
||||
raise ValueError(
|
||||
f"Load model: type={spec.cls.__name__} bundled sibling '{name}' present in "
|
||||
f"file but no class was supplied via sibling_classes"
|
||||
)
|
||||
sibling_cfg = fetch_component_config(repo_id, sibling_spec.subfolder)
|
||||
loaded_siblings[name] = build_component(
|
||||
component_name=name,
|
||||
state_dict=sibling_sd,
|
||||
config=sibling_cfg,
|
||||
cls=sibling_cls,
|
||||
converter=None,
|
||||
acceptable_missing=sibling_spec.acceptable_missing,
|
||||
quant_args={},
|
||||
quant_type=None,
|
||||
)
|
||||
|
||||
sd_models.allow_post_quant = False
|
||||
devices.torch_gc()
|
||||
log.debug(f"Load model: type={spec.cls.__name__} native_transformer time={time.time() - t0:.2f}")
|
||||
return transformer, loaded_siblings
|
||||
|
||||
|
||||
def strip_prefix(state_dict: dict, prefixes: tuple[str, ...], type_name: str) -> dict:
|
||||
"""Detect and uniformly strip the most common known prefix from every key.
|
||||
|
||||
Order matters: longer prefixes win over shorter ones with the same suffix
|
||||
(e.g. ``model.diffusion_model.`` beats ``diffusion_model.``). If some keys
|
||||
match the dominant prefix and others do not, raises ValueError because
|
||||
mixed prefixes indicate a malformed file rather than a recoverable export
|
||||
quirk.
|
||||
"""
|
||||
sorted_prefixes = sorted(prefixes, key=len, reverse=True)
|
||||
counts: dict[str, int] = {}
|
||||
seen = 0
|
||||
for key in state_dict:
|
||||
for prefix in sorted_prefixes:
|
||||
if key.startswith(prefix):
|
||||
counts[prefix] = counts.get(prefix, 0) + 1
|
||||
seen += 1
|
||||
break
|
||||
total = len(state_dict)
|
||||
if seen == 0:
|
||||
log.debug(f"Load model: type={type_name} native_transformer prefix=bare")
|
||||
return state_dict
|
||||
dominant = max(counts, key=counts.get)
|
||||
if counts[dominant] != total:
|
||||
raise ValueError(
|
||||
f"Load model: type={type_name} native_transformer has mixed prefixes "
|
||||
f"(total={total} {dominant}={counts[dominant]})"
|
||||
)
|
||||
log.debug(f'Load model: type={type_name} native_transformer prefix="{dominant}"')
|
||||
offset = len(dominant)
|
||||
return {key[offset:]: value for key, value in state_dict.items()}
|
||||
|
||||
|
||||
def check_forbidden_markers(
|
||||
state_dict: dict,
|
||||
forbidden_markers: tuple[tuple[str, str], ...],
|
||||
type_name: str,
|
||||
local_file: str,
|
||||
) -> None:
|
||||
"""Raise if any forbidden marker key is present in the state_dict.
|
||||
|
||||
Catches structural mismatches that pass prefix detection but indicate the
|
||||
file is from an incompatible architecture variant (e.g. Cosmos 1.0 keys
|
||||
showing up in a Cosmos 2.0 loader path).
|
||||
"""
|
||||
for marker, description in forbidden_markers:
|
||||
if marker in state_dict:
|
||||
raise ValueError(
|
||||
f"Load model: type={type_name} native_transformer rejects "
|
||||
f'"{os.path.basename(local_file)}" ({description}; marker key {marker!r})'
|
||||
)
|
||||
|
||||
|
||||
def partition_siblings(
|
||||
state_dict: dict,
|
||||
siblings: dict[str, SiblingSpec],
|
||||
) -> tuple[dict, dict[str, dict]]:
|
||||
"""Split state_dict into (transformer_sd, {sibling_name: sibling_sd}).
|
||||
|
||||
Keys matching a sibling's ``inline_prefix`` go into that sibling's dict
|
||||
with the prefix stripped; everything else stays in the transformer dict.
|
||||
Sibling names with no matching keys still appear in the output dict but
|
||||
map to an empty dict, so the caller can iterate uniformly.
|
||||
"""
|
||||
sibling_sds: dict[str, dict] = {name: {} for name in siblings}
|
||||
transformer_sd: dict = {}
|
||||
if not siblings:
|
||||
return state_dict, sibling_sds
|
||||
sibling_lookups = [(name, siblings[name].inline_prefix) for name in siblings]
|
||||
for key, value in state_dict.items():
|
||||
matched = False
|
||||
for name, prefix in sibling_lookups:
|
||||
if key.startswith(prefix):
|
||||
sibling_sds[name][key[len(prefix):]] = value
|
||||
matched = True
|
||||
break
|
||||
if not matched:
|
||||
transformer_sd[key] = value
|
||||
return transformer_sd, sibling_sds
|
||||
|
||||
|
||||
def fetch_component_config(repo_id: str, subfolder: str) -> dict:
|
||||
"""Download and parse ``<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(
|
||||
*,
|
||||
component_name: str,
|
||||
state_dict: dict,
|
||||
config: dict,
|
||||
cls: type,
|
||||
converter: Callable[[dict], dict] | None,
|
||||
acceptable_missing: tuple[str, ...],
|
||||
quant_args: dict,
|
||||
quant_type: str | None,
|
||||
) -> object:
|
||||
"""Convert (if needed), instantiate, load weights, dtype-cast, quantize,
|
||||
and offload-place a single component. Raises on any hard failure.
|
||||
"""
|
||||
try:
|
||||
sd = converter(state_dict) if converter is not None else state_dict
|
||||
component = cls.from_config(config)
|
||||
missing, unexpected = component.load_state_dict(sd, strict=False)
|
||||
validate_state_dict_load(component_name, missing, unexpected, acceptable_missing)
|
||||
del sd
|
||||
devices.torch_gc()
|
||||
component = component.to(dtype=devices.dtype)
|
||||
except Exception as e:
|
||||
log.error(f"Load model: native_transformer {component_name} load failed: {e}")
|
||||
errors.display(e, "Load")
|
||||
raise
|
||||
|
||||
if component_name == "transformer":
|
||||
apply_quant(component, quant_type)
|
||||
|
||||
if shared.opts.diffusers_offload_mode != "none":
|
||||
sd_models.move_model(component, devices.cpu)
|
||||
|
||||
if not hasattr(component, "quantization_config"):
|
||||
if hasattr(component, "config") and hasattr(component.config, "quantization_config"):
|
||||
component.quantization_config = component.config.quantization_config
|
||||
elif quant_type is not None and quant_args.get("quantization_config") is not None:
|
||||
component.quantization_config = quant_args.get("quantization_config")
|
||||
return component
|
||||
|
||||
|
||||
def validate_state_dict_load(
|
||||
component_name: str,
|
||||
missing: list[str],
|
||||
unexpected: list[str],
|
||||
acceptable_missing: tuple[str, ...],
|
||||
) -> None:
|
||||
"""Raise ValueError if load_state_dict produced unexpected keys or
|
||||
non-acceptable missing keys. Buffer-only missing keys matching the
|
||||
``acceptable_missing`` prefix list are logged at debug level and ignored.
|
||||
"""
|
||||
if unexpected:
|
||||
sample = ", ".join(unexpected[:5])
|
||||
raise ValueError(
|
||||
f"Load model: native_transformer {component_name} has {len(unexpected)} "
|
||||
f"unexpected keys (sample: {sample})"
|
||||
)
|
||||
hard_missing = [
|
||||
k for k in missing if not any(k.startswith(p) for p in acceptable_missing)
|
||||
]
|
||||
if hard_missing:
|
||||
sample = ", ".join(hard_missing[:5])
|
||||
raise ValueError(
|
||||
f"Load model: native_transformer {component_name} missing "
|
||||
f"{len(hard_missing)} required keys (sample: {sample})"
|
||||
)
|
||||
if missing:
|
||||
log.debug(
|
||||
f"Load model: native_transformer {component_name} ignored "
|
||||
f"{len(missing)} buffer-only missing keys"
|
||||
)
|
||||
|
||||
|
||||
def apply_quant(transformer: object, quant_type: str | None) -> None:
|
||||
"""Apply SDNQ / layerwise quantization to the bare transformer.
|
||||
|
||||
SDNQ 'pre' and 'auto' would normally route through ``quantization_config``
|
||||
at ``from_pretrained`` time; the native path bypasses that boundary, so
|
||||
we call the per-module quant path directly. SDNQ 'post' and
|
||||
``layerwise_quantization`` go through ``do_post_load_quant`` as usual.
|
||||
NVIDIAModelOptConfig (TRT) is not supported on this path.
|
||||
"""
|
||||
if quant_type == "NVIDIAModelOptConfig":
|
||||
log.warning(
|
||||
"Load model: native_transformer quant=TRT not supported on native path, skipping"
|
||||
)
|
||||
elif quant_type == "SDNQConfig":
|
||||
if shared.opts.sdnq_quantize_mode == "pre":
|
||||
log.info(
|
||||
"Load model: native_transformer quant=SDNQ pre-mode applied post-load "
|
||||
"on native path"
|
||||
)
|
||||
model_quant.sdnq_quantize_model(transformer, op="transformer")
|
||||
model_quant.do_post_load_quant(transformer, allow=False)
|
||||
@@ -0,0 +1,761 @@
|
||||
#!/usr/bin/env python
|
||||
"""
|
||||
Offline unit tests for pipelines.native_transformer.
|
||||
|
||||
Covers the pure helpers that own per-arch knob handling:
|
||||
|
||||
- ``strip_prefix`` for single/multi prefix detection and mixed-prefix rejection
|
||||
- ``partition_siblings`` for inline-sibling key partitioning
|
||||
- ``check_forbidden_markers`` for structural-mismatch rejection
|
||||
- ``is_noop_converter`` for diffusers no-op lambda detection
|
||||
- ``validate_state_dict_load`` for unexpected / missing key handling
|
||||
- ``register`` / ``lookup`` registry behavior and default spec synthesis
|
||||
- ``auto_pickup_converter`` for diffusers ``SINGLE_FILE_LOADABLE_CLASSES`` integration
|
||||
- ``TransformerSpec`` / ``SiblingSpec`` defaults
|
||||
|
||||
Plus one end-to-end ``load`` test against a tiny mock module that exercises
|
||||
the read -> strip -> convert -> from_config -> load_state_dict -> validate
|
||||
pipeline without needing a real diffusers transformer or hf_hub_download.
|
||||
|
||||
No running server required.
|
||||
|
||||
Usage:
|
||||
python test/test-native-transformer.py
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
|
||||
import torch
|
||||
import safetensors.torch
|
||||
|
||||
script_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, script_dir)
|
||||
os.chdir(script_dir)
|
||||
|
||||
os.environ['SD_INSTALL_QUIET'] = '1'
|
||||
|
||||
# Bootstrap cmd_args before any module that pulls in shared.py.
|
||||
import modules.cmd_args # pylint: disable=wrong-import-position
|
||||
import installer # pylint: disable=wrong-import-position
|
||||
_orig_argv = sys.argv
|
||||
sys.argv = [sys.argv[0]]
|
||||
try:
|
||||
modules.cmd_args.parse_args()
|
||||
finally:
|
||||
sys.argv = _orig_argv
|
||||
installer.add_args(modules.cmd_args.parser)
|
||||
modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([])
|
||||
|
||||
from modules.errors import log # pylint: disable=wrong-import-position
|
||||
from pipelines import native_transformer as nt # pylint: disable=wrong-import-position
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Test infrastructure
|
||||
# ============================================================
|
||||
|
||||
results: dict[str, dict] = {}
|
||||
|
||||
|
||||
def category(name: str):
|
||||
if name not in results:
|
||||
results[name] = {'passed': 0, 'failed': 0, 'tests': []}
|
||||
return name
|
||||
|
||||
|
||||
def record(cat: str, passed: bool, name: str, detail: str = ''):
|
||||
status = 'PASS' if passed else 'FAIL'
|
||||
results[cat]['passed' if passed else 'failed'] += 1
|
||||
results[cat]['tests'].append((status, name))
|
||||
msg = f' {status}: {name}'
|
||||
if detail:
|
||||
msg += f' ({detail})'
|
||||
if passed:
|
||||
log.info(msg)
|
||||
else:
|
||||
log.error(msg)
|
||||
|
||||
|
||||
def run_test(cat: str, fn):
|
||||
name = fn.__name__
|
||||
try:
|
||||
ok = fn()
|
||||
if ok is False:
|
||||
record(cat, False, name)
|
||||
else:
|
||||
record(cat, True, name)
|
||||
except AssertionError as e:
|
||||
record(cat, False, name, str(e))
|
||||
except Exception as e: # pylint: disable=broad-except
|
||||
record(cat, False, name, f'exception: {e}')
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
|
||||
# ============================================================
|
||||
# strip_prefix
|
||||
# ============================================================
|
||||
|
||||
def test_strip_prefix_bare_keys_pass_through():
|
||||
sd = {'layers.0.weight': 1, 'layers.0.bias': 2}
|
||||
out = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
|
||||
assert out == sd, 'bare keys must pass through unchanged'
|
||||
|
||||
|
||||
def test_strip_prefix_dominant_single_variant():
|
||||
sd = {f'model.diffusion_model.layers.{i}.weight': i for i in range(10)}
|
||||
out = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
|
||||
assert all(k.startswith('layers.') for k in out)
|
||||
assert len(out) == 10
|
||||
|
||||
|
||||
def test_strip_prefix_picks_longest_match_first():
|
||||
"""``model.diffusion_model.`` must beat ``diffusion_model.`` when both match."""
|
||||
sd = {
|
||||
'model.diffusion_model.layers.0.weight': 1,
|
||||
'model.diffusion_model.layers.1.weight': 2,
|
||||
}
|
||||
out = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
|
||||
# If shorter prefix matched, keys would start with 'model.'
|
||||
assert 'layers.0.weight' in out
|
||||
assert 'layers.1.weight' in out
|
||||
assert not any(k.startswith('model.') for k in out)
|
||||
|
||||
|
||||
def test_strip_prefix_mixed_prefixes_raises():
|
||||
sd = {
|
||||
'model.diffusion_model.layers.0.weight': 1,
|
||||
'net.layers.0.weight': 2,
|
||||
}
|
||||
try:
|
||||
nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
|
||||
raise AssertionError('expected ValueError')
|
||||
except ValueError as e:
|
||||
assert 'mixed prefixes' in str(e)
|
||||
|
||||
|
||||
def test_strip_prefix_net_variant():
|
||||
sd = {'net.layers.0.weight': 1, 'net.layers.1.bias': 2}
|
||||
out = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
|
||||
assert set(out.keys()) == {'layers.0.weight', 'layers.1.bias'}
|
||||
|
||||
|
||||
def test_strip_prefix_diffusion_model_variant():
|
||||
sd = {'diffusion_model.layers.0.weight': 1, 'diffusion_model.layers.1.bias': 2}
|
||||
out = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
|
||||
assert set(out.keys()) == {'layers.0.weight', 'layers.1.bias'}
|
||||
|
||||
|
||||
def test_strip_prefix_custom_prefix_set():
|
||||
sd = {'lora_unet_blocks_0.weight': 1, 'lora_unet_blocks_1.weight': 2}
|
||||
out = nt.strip_prefix(sd, ('lora_unet_',), 'Test')
|
||||
assert set(out.keys()) == {'blocks_0.weight', 'blocks_1.weight'}
|
||||
|
||||
|
||||
# ============================================================
|
||||
# partition_siblings
|
||||
# ============================================================
|
||||
|
||||
def test_partition_siblings_empty_spec_returns_state_dict_unchanged():
|
||||
sd = {'a': 1, 'b': 2}
|
||||
transformer_sd, siblings = nt.partition_siblings(sd, {})
|
||||
assert transformer_sd == sd
|
||||
assert siblings == {}
|
||||
|
||||
|
||||
def test_partition_siblings_no_matches_keeps_all_in_transformer():
|
||||
sd = {'layers.0.weight': 1, 'layers.1.weight': 2}
|
||||
siblings_spec = {'llm_adapter': nt.SiblingSpec(subfolder='llm_adapter', inline_prefix='llm_adapter.')}
|
||||
transformer_sd, siblings = nt.partition_siblings(sd, siblings_spec)
|
||||
assert transformer_sd == sd
|
||||
assert siblings == {'llm_adapter': {}}
|
||||
|
||||
|
||||
def test_partition_siblings_single_sibling_split():
|
||||
sd = {
|
||||
'layers.0.weight': 'tx0',
|
||||
'layers.1.weight': 'tx1',
|
||||
'llm_adapter.input_proj.weight': 'ad0',
|
||||
'llm_adapter.output_proj.weight': 'ad1',
|
||||
}
|
||||
siblings_spec = {'llm_adapter': nt.SiblingSpec(subfolder='llm_adapter', inline_prefix='llm_adapter.')}
|
||||
transformer_sd, siblings = nt.partition_siblings(sd, siblings_spec)
|
||||
assert set(transformer_sd.keys()) == {'layers.0.weight', 'layers.1.weight'}
|
||||
assert set(siblings['llm_adapter'].keys()) == {'input_proj.weight', 'output_proj.weight'}
|
||||
assert siblings['llm_adapter']['input_proj.weight'] == 'ad0'
|
||||
|
||||
|
||||
def test_partition_siblings_multiple_siblings():
|
||||
sd = {
|
||||
'layers.0.weight': 'tx',
|
||||
'sibling_a.x.weight': 'a0',
|
||||
'sibling_b.y.weight': 'b0',
|
||||
'sibling_b.z.weight': 'b1',
|
||||
}
|
||||
siblings_spec = {
|
||||
'sibling_a': nt.SiblingSpec(subfolder='a', inline_prefix='sibling_a.'),
|
||||
'sibling_b': nt.SiblingSpec(subfolder='b', inline_prefix='sibling_b.'),
|
||||
}
|
||||
transformer_sd, siblings = nt.partition_siblings(sd, siblings_spec)
|
||||
assert list(transformer_sd.keys()) == ['layers.0.weight']
|
||||
assert set(siblings['sibling_a'].keys()) == {'x.weight'}
|
||||
assert set(siblings['sibling_b'].keys()) == {'y.weight', 'z.weight'}
|
||||
|
||||
|
||||
# ============================================================
|
||||
# check_forbidden_markers
|
||||
# ============================================================
|
||||
|
||||
def test_forbidden_markers_passes_when_absent():
|
||||
sd = {'layers.0.weight': 1}
|
||||
markers = (('legacy.marker.weight', 'old format'),)
|
||||
nt.check_forbidden_markers(sd, markers, 'Test', '/tmp/x.safetensors')
|
||||
# no exception = pass
|
||||
|
||||
|
||||
def test_forbidden_markers_raises_when_present():
|
||||
sd = {'layers.0.weight': 1, 'legacy.marker.weight': 2}
|
||||
markers = (('legacy.marker.weight', 'old Cosmos 1.0 structure'),)
|
||||
try:
|
||||
nt.check_forbidden_markers(sd, markers, 'Test', '/tmp/x.safetensors')
|
||||
raise AssertionError('expected ValueError')
|
||||
except ValueError as e:
|
||||
msg = str(e)
|
||||
assert 'old Cosmos 1.0 structure' in msg
|
||||
assert 'legacy.marker.weight' in msg
|
||||
|
||||
|
||||
def test_forbidden_markers_empty_tuple_no_op():
|
||||
sd = {'layers.0.weight': 1}
|
||||
nt.check_forbidden_markers(sd, (), 'Test', '/tmp/x.safetensors')
|
||||
|
||||
|
||||
# ============================================================
|
||||
# is_noop_converter
|
||||
# ============================================================
|
||||
|
||||
def test_noop_converter_identity_lambda():
|
||||
fn = lambda checkpoint, **kwargs: checkpoint # pylint: disable=unnecessary-lambda-assignment
|
||||
assert nt.is_noop_converter(fn) is True
|
||||
|
||||
|
||||
def test_noop_converter_real_function():
|
||||
def real(checkpoint, **kwargs): # pylint: disable=unused-argument
|
||||
return {k.replace('a.', 'b.'): v for k, v in checkpoint.items()}
|
||||
assert nt.is_noop_converter(real) is False
|
||||
|
||||
|
||||
def test_noop_converter_lambda_with_modification():
|
||||
fn = lambda checkpoint, **kwargs: {k: v.float() for k, v in checkpoint.items()} # pylint: disable=unnecessary-lambda-assignment
|
||||
assert nt.is_noop_converter(fn) is False
|
||||
|
||||
|
||||
# ============================================================
|
||||
# validate_state_dict_load
|
||||
# ============================================================
|
||||
|
||||
def test_validate_accepts_buffer_only_missing():
|
||||
nt.validate_state_dict_load(
|
||||
component_name='transformer',
|
||||
missing=['rope.freqs', 'pos_embedder.pos'],
|
||||
unexpected=[],
|
||||
acceptable_missing=('rope.', 'pos_embedder.'),
|
||||
)
|
||||
|
||||
|
||||
def test_validate_rejects_unexpected():
|
||||
try:
|
||||
nt.validate_state_dict_load(
|
||||
component_name='transformer',
|
||||
missing=[],
|
||||
unexpected=['some.junk.weight'],
|
||||
acceptable_missing=(),
|
||||
)
|
||||
raise AssertionError('expected ValueError')
|
||||
except ValueError as e:
|
||||
assert 'unexpected' in str(e)
|
||||
assert 'some.junk.weight' in str(e)
|
||||
|
||||
|
||||
def test_validate_rejects_hard_missing():
|
||||
try:
|
||||
nt.validate_state_dict_load(
|
||||
component_name='transformer',
|
||||
missing=['layers.0.weight', 'rope.freqs'],
|
||||
unexpected=[],
|
||||
acceptable_missing=('rope.',),
|
||||
)
|
||||
raise AssertionError('expected ValueError')
|
||||
except ValueError as e:
|
||||
msg = str(e)
|
||||
assert 'missing' in msg
|
||||
assert 'layers.0.weight' in msg
|
||||
# Buffer-only missing must not show up in the hard-missing list
|
||||
assert msg.count('rope.freqs') == 0
|
||||
|
||||
|
||||
def test_validate_empty_passes():
|
||||
nt.validate_state_dict_load(
|
||||
component_name='transformer',
|
||||
missing=[],
|
||||
unexpected=[],
|
||||
acceptable_missing=(),
|
||||
)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# register / lookup
|
||||
# ============================================================
|
||||
|
||||
class FakeTransformer:
|
||||
"""Minimal stand-in for a diffusers transformer class."""
|
||||
|
||||
|
||||
class FakeTransformer2:
|
||||
"""Second stand-in for register/lookup tests."""
|
||||
|
||||
|
||||
def test_register_with_explicit_spec():
|
||||
nt.REGISTRY.clear()
|
||||
spec = nt.TransformerSpec(cls=FakeTransformer, subfolder='custom_sub')
|
||||
nt.register(FakeTransformer, spec)
|
||||
assert nt.lookup(FakeTransformer) is spec
|
||||
|
||||
|
||||
def test_register_with_default_spec():
|
||||
nt.REGISTRY.clear()
|
||||
nt.register(FakeTransformer)
|
||||
spec = nt.lookup(FakeTransformer)
|
||||
assert spec.cls is FakeTransformer
|
||||
assert spec.subfolder == 'transformer'
|
||||
assert spec.prefixes == nt.DEFAULT_PREFIXES
|
||||
assert spec.converter is None
|
||||
assert spec.siblings == {}
|
||||
|
||||
|
||||
def test_register_idempotent_replaces():
|
||||
nt.REGISTRY.clear()
|
||||
spec1 = nt.TransformerSpec(cls=FakeTransformer, subfolder='sub_one')
|
||||
spec2 = nt.TransformerSpec(cls=FakeTransformer, subfolder='sub_two')
|
||||
nt.register(FakeTransformer, spec1)
|
||||
nt.register(FakeTransformer, spec2)
|
||||
assert nt.lookup(FakeTransformer) is spec2
|
||||
|
||||
|
||||
def test_register_rejects_mismatched_cls():
|
||||
try:
|
||||
nt.register(FakeTransformer, nt.TransformerSpec(cls=FakeTransformer2))
|
||||
raise AssertionError('expected ValueError')
|
||||
except ValueError as e:
|
||||
assert 'does not match' in str(e)
|
||||
|
||||
|
||||
def test_lookup_synthesizes_default_for_unregistered():
|
||||
nt.REGISTRY.clear()
|
||||
spec = nt.lookup(FakeTransformer)
|
||||
assert spec.cls is FakeTransformer
|
||||
assert spec.subfolder == 'transformer' # default
|
||||
assert spec.converter is None # FakeTransformer has no diffusers entry
|
||||
|
||||
|
||||
# ============================================================
|
||||
# auto_pickup_converter
|
||||
# ============================================================
|
||||
|
||||
def test_auto_pickup_returns_none_for_unknown_class():
|
||||
assert nt.auto_pickup_converter(FakeTransformer) is None
|
||||
|
||||
|
||||
def test_auto_pickup_returns_real_diffusers_converter():
|
||||
import diffusers
|
||||
fn = nt.auto_pickup_converter(diffusers.FluxTransformer2DModel)
|
||||
assert fn is not None
|
||||
assert callable(fn)
|
||||
assert fn.__name__ == 'convert_flux_transformer_checkpoint_to_diffusers'
|
||||
|
||||
|
||||
def test_auto_pickup_skips_noop_converter_qwen():
|
||||
"""QwenImageTransformer2DModel registers a no-op lambda in diffusers;
|
||||
auto_pickup_converter must return None so the spec falls back to no
|
||||
converter (the user-registered spec can override with a real converter)."""
|
||||
import diffusers
|
||||
assert nt.auto_pickup_converter(diffusers.QwenImageTransformer2DModel) is None
|
||||
|
||||
|
||||
# ============================================================
|
||||
# TransformerSpec / SiblingSpec defaults
|
||||
# ============================================================
|
||||
|
||||
def test_transformer_spec_defaults():
|
||||
spec = nt.TransformerSpec(cls=FakeTransformer)
|
||||
assert spec.subfolder == 'transformer'
|
||||
assert spec.prefixes == ('model.diffusion_model.', 'diffusion_model.', 'net.')
|
||||
assert spec.converter is None
|
||||
assert spec.siblings == {}
|
||||
assert spec.acceptable_missing == ('rope.', 'pos_embedder.', 'learnable_pos_embed.')
|
||||
assert spec.forbidden_markers == ()
|
||||
|
||||
|
||||
def test_sibling_spec_defaults():
|
||||
spec = nt.SiblingSpec(subfolder='llm_adapter', inline_prefix='llm_adapter.')
|
||||
assert spec.subfolder == 'llm_adapter'
|
||||
assert spec.inline_prefix == 'llm_adapter.'
|
||||
assert spec.acceptable_missing == ()
|
||||
|
||||
|
||||
def test_transformer_spec_is_frozen():
|
||||
spec = nt.TransformerSpec(cls=FakeTransformer)
|
||||
try:
|
||||
spec.subfolder = 'changed' # type: ignore[misc]
|
||||
except Exception as e: # pylint: disable=broad-except
|
||||
assert 'FrozenInstanceError' in type(e).__name__ or 'frozen' in str(e).lower()
|
||||
return
|
||||
raise AssertionError('expected FrozenInstanceError')
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Integration: end-to-end load() with a tiny mock module
|
||||
# ============================================================
|
||||
# We sidestep diffusers + hf_hub_download by patching:
|
||||
# - ``fetch_component_config`` to return a hand-rolled config dict
|
||||
# - ``model_quant.get_dit_args`` / ``model_quant.get_quant_type`` / quant
|
||||
# application to no-ops (we only want to test the load path itself).
|
||||
# The mock cls is a torch.nn.Module subclass whose ``from_config`` constructs
|
||||
# a fresh module of the expected shape; ``load_state_dict`` is the standard
|
||||
# PyTorch method.
|
||||
|
||||
class MockMiniTransformer(torch.nn.Module):
|
||||
"""Tiny stand-in: linear in -> linear out, plus a nested rope sub-module
|
||||
holding a buffer the trainer state dict won't carry. Nested mirrors how
|
||||
real DiTs structure rope / pos_embedder buffers."""
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict) -> 'MockMiniTransformer':
|
||||
return cls(dim=config['dim'])
|
||||
|
||||
def __init__(self, dim: int):
|
||||
super().__init__()
|
||||
self.in_proj = torch.nn.Linear(dim, dim)
|
||||
self.out_proj = torch.nn.Linear(dim, dim)
|
||||
self.rope = torch.nn.Module()
|
||||
self.rope.register_buffer('freqs', torch.zeros(dim))
|
||||
|
||||
|
||||
def write_fixture(state_dict_keys: dict, fd: int, path: str) -> str:
|
||||
os.close(fd)
|
||||
safetensors.torch.save_file(state_dict_keys, path)
|
||||
return path
|
||||
|
||||
|
||||
def test_load_end_to_end_with_bfl_prefix_no_converter():
|
||||
"""Exercise the full load pipeline: read .safetensors, strip prefix,
|
||||
no converter, instantiate via from_config, load weights, validate.
|
||||
"""
|
||||
fd, path = tempfile.mkstemp(suffix='.safetensors')
|
||||
try:
|
||||
dim = 8
|
||||
# Save with model.diffusion_model. prefix; in_proj.* and out_proj.*
|
||||
# are the real weights the mock cls expects after the strip.
|
||||
raw = {
|
||||
'model.diffusion_model.in_proj.weight': torch.randn(dim, dim),
|
||||
'model.diffusion_model.in_proj.bias': torch.zeros(dim),
|
||||
'model.diffusion_model.out_proj.weight': torch.randn(dim, dim),
|
||||
'model.diffusion_model.out_proj.bias': torch.zeros(dim),
|
||||
}
|
||||
write_fixture(raw, fd, path)
|
||||
|
||||
# Patch fetch_component_config to return our hand-rolled config.
|
||||
orig_fetch = nt.fetch_component_config
|
||||
nt.fetch_component_config = lambda repo, sub: {'dim': dim}
|
||||
|
||||
# Patch quant helpers (we only care about the load path).
|
||||
from modules import model_quant
|
||||
orig_get_dit = model_quant.get_dit_args
|
||||
orig_get_qtype = model_quant.get_quant_type
|
||||
orig_do_post = model_quant.do_post_load_quant
|
||||
model_quant.get_dit_args = lambda *a, **k: ({}, {})
|
||||
model_quant.get_quant_type = lambda *a, **k: None
|
||||
model_quant.do_post_load_quant = lambda *a, **k: None
|
||||
|
||||
try:
|
||||
spec = nt.TransformerSpec(cls=MockMiniTransformer)
|
||||
transformer, siblings = nt.load(
|
||||
local_file=path,
|
||||
repo_id='fake/repo',
|
||||
spec=spec,
|
||||
diffusers_cfg={},
|
||||
)
|
||||
finally:
|
||||
nt.fetch_component_config = orig_fetch
|
||||
model_quant.get_dit_args = orig_get_dit
|
||||
model_quant.get_quant_type = orig_get_qtype
|
||||
model_quant.do_post_load_quant = orig_do_post
|
||||
|
||||
assert isinstance(transformer, MockMiniTransformer)
|
||||
assert transformer.in_proj.weight.shape == (dim, dim)
|
||||
# Weights from the fixture should match what was loaded.
|
||||
loaded_in_w = transformer.in_proj.weight.detach().cpu()
|
||||
fixture_in_w = raw['model.diffusion_model.in_proj.weight'].to(loaded_in_w.dtype)
|
||||
assert torch.allclose(loaded_in_w, fixture_in_w)
|
||||
assert siblings == {}
|
||||
finally:
|
||||
if os.path.exists(path):
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
def test_load_end_to_end_with_sibling_partition():
|
||||
"""Bundled-sibling case: file carries both transformer and sibling weights,
|
||||
sibling_classes supplies the runtime sibling class, partition routes each
|
||||
half into its target."""
|
||||
fd, path = tempfile.mkstemp(suffix='.safetensors')
|
||||
try:
|
||||
dim = 8
|
||||
sibling_dim = 4
|
||||
raw = {
|
||||
# Transformer half (after strip).
|
||||
'model.diffusion_model.in_proj.weight': torch.randn(dim, dim),
|
||||
'model.diffusion_model.in_proj.bias': torch.zeros(dim),
|
||||
'model.diffusion_model.out_proj.weight': torch.randn(dim, dim),
|
||||
'model.diffusion_model.out_proj.bias': torch.zeros(dim),
|
||||
# Sibling half (after strip + sibling partition).
|
||||
'model.diffusion_model.sibling.in_proj.weight': torch.randn(sibling_dim, sibling_dim),
|
||||
'model.diffusion_model.sibling.in_proj.bias': torch.zeros(sibling_dim),
|
||||
'model.diffusion_model.sibling.out_proj.weight': torch.randn(sibling_dim, sibling_dim),
|
||||
'model.diffusion_model.sibling.out_proj.bias': torch.zeros(sibling_dim),
|
||||
}
|
||||
write_fixture(raw, fd, path)
|
||||
|
||||
orig_fetch = nt.fetch_component_config
|
||||
|
||||
def patched_fetch(_repo, sub):
|
||||
return {'dim': dim if sub == 'transformer' else sibling_dim}
|
||||
|
||||
nt.fetch_component_config = patched_fetch
|
||||
|
||||
from modules import model_quant
|
||||
orig_get_dit = model_quant.get_dit_args
|
||||
orig_get_qtype = model_quant.get_quant_type
|
||||
orig_do_post = model_quant.do_post_load_quant
|
||||
model_quant.get_dit_args = lambda *a, **k: ({}, {})
|
||||
model_quant.get_quant_type = lambda *a, **k: None
|
||||
model_quant.do_post_load_quant = lambda *a, **k: None
|
||||
|
||||
try:
|
||||
spec = nt.TransformerSpec(
|
||||
cls=MockMiniTransformer,
|
||||
siblings={
|
||||
'sibling': nt.SiblingSpec(
|
||||
subfolder='sibling',
|
||||
inline_prefix='sibling.',
|
||||
acceptable_missing=('rope.',),
|
||||
),
|
||||
},
|
||||
)
|
||||
transformer, siblings = nt.load(
|
||||
local_file=path,
|
||||
repo_id='fake/repo',
|
||||
spec=spec,
|
||||
diffusers_cfg={},
|
||||
sibling_classes={'sibling': MockMiniTransformer},
|
||||
)
|
||||
finally:
|
||||
nt.fetch_component_config = orig_fetch
|
||||
model_quant.get_dit_args = orig_get_dit
|
||||
model_quant.get_quant_type = orig_get_qtype
|
||||
model_quant.do_post_load_quant = orig_do_post
|
||||
|
||||
assert isinstance(transformer, MockMiniTransformer)
|
||||
assert transformer.in_proj.weight.shape == (dim, dim)
|
||||
assert 'sibling' in siblings
|
||||
assert isinstance(siblings['sibling'], MockMiniTransformer)
|
||||
assert siblings['sibling'].in_proj.weight.shape == (sibling_dim, sibling_dim)
|
||||
finally:
|
||||
if os.path.exists(path):
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
def test_load_raises_on_missing_sibling_class():
|
||||
"""Sibling keys present in file but caller forgot to supply the class."""
|
||||
fd, path = tempfile.mkstemp(suffix='.safetensors')
|
||||
try:
|
||||
dim = 8
|
||||
sibling_dim = 4
|
||||
raw = {
|
||||
'model.diffusion_model.in_proj.weight': torch.randn(dim, dim),
|
||||
'model.diffusion_model.in_proj.bias': torch.zeros(dim),
|
||||
'model.diffusion_model.out_proj.weight': torch.randn(dim, dim),
|
||||
'model.diffusion_model.out_proj.bias': torch.zeros(dim),
|
||||
'model.diffusion_model.sibling.in_proj.weight': torch.randn(sibling_dim, sibling_dim),
|
||||
}
|
||||
write_fixture(raw, fd, path)
|
||||
|
||||
orig_fetch = nt.fetch_component_config
|
||||
nt.fetch_component_config = lambda repo, sub: {'dim': dim}
|
||||
|
||||
from modules import model_quant
|
||||
orig_get_dit = model_quant.get_dit_args
|
||||
orig_get_qtype = model_quant.get_quant_type
|
||||
orig_do_post = model_quant.do_post_load_quant
|
||||
model_quant.get_dit_args = lambda *a, **k: ({}, {})
|
||||
model_quant.get_quant_type = lambda *a, **k: None
|
||||
model_quant.do_post_load_quant = lambda *a, **k: None
|
||||
|
||||
try:
|
||||
spec = nt.TransformerSpec(
|
||||
cls=MockMiniTransformer,
|
||||
siblings={'sibling': nt.SiblingSpec(subfolder='s', inline_prefix='sibling.')},
|
||||
)
|
||||
raised = False
|
||||
try:
|
||||
nt.load(
|
||||
local_file=path,
|
||||
repo_id='fake/repo',
|
||||
spec=spec,
|
||||
diffusers_cfg={},
|
||||
sibling_classes={}, # missing!
|
||||
)
|
||||
except ValueError as e:
|
||||
raised = True
|
||||
assert "'sibling'" in str(e)
|
||||
assert 'sibling_classes' in str(e)
|
||||
assert raised, 'expected ValueError'
|
||||
finally:
|
||||
nt.fetch_component_config = orig_fetch
|
||||
model_quant.get_dit_args = orig_get_dit
|
||||
model_quant.get_quant_type = orig_get_qtype
|
||||
model_quant.do_post_load_quant = orig_do_post
|
||||
finally:
|
||||
if os.path.exists(path):
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
def test_load_rejects_non_safetensors():
|
||||
spec = nt.TransformerSpec(cls=MockMiniTransformer)
|
||||
try:
|
||||
nt.load(
|
||||
local_file='/tmp/some.gguf',
|
||||
repo_id='fake/repo',
|
||||
spec=spec,
|
||||
diffusers_cfg={},
|
||||
)
|
||||
raise AssertionError('expected ValueError')
|
||||
except ValueError as e:
|
||||
assert '.safetensors' in str(e)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Run
|
||||
# ============================================================
|
||||
|
||||
def run_all():
|
||||
log.warning('=== strip_prefix ===')
|
||||
cat = category('strip')
|
||||
for fn in [
|
||||
test_strip_prefix_bare_keys_pass_through,
|
||||
test_strip_prefix_dominant_single_variant,
|
||||
test_strip_prefix_picks_longest_match_first,
|
||||
test_strip_prefix_mixed_prefixes_raises,
|
||||
test_strip_prefix_net_variant,
|
||||
test_strip_prefix_diffusion_model_variant,
|
||||
test_strip_prefix_custom_prefix_set,
|
||||
]:
|
||||
run_test(cat, fn)
|
||||
|
||||
log.warning('=== partition_siblings ===')
|
||||
cat = category('partition')
|
||||
for fn in [
|
||||
test_partition_siblings_empty_spec_returns_state_dict_unchanged,
|
||||
test_partition_siblings_no_matches_keeps_all_in_transformer,
|
||||
test_partition_siblings_single_sibling_split,
|
||||
test_partition_siblings_multiple_siblings,
|
||||
]:
|
||||
run_test(cat, fn)
|
||||
|
||||
log.warning('=== forbidden_markers ===')
|
||||
cat = category('forbidden')
|
||||
for fn in [
|
||||
test_forbidden_markers_passes_when_absent,
|
||||
test_forbidden_markers_raises_when_present,
|
||||
test_forbidden_markers_empty_tuple_no_op,
|
||||
]:
|
||||
run_test(cat, fn)
|
||||
|
||||
log.warning('=== noop_converter detection ===')
|
||||
cat = category('noop')
|
||||
for fn in [
|
||||
test_noop_converter_identity_lambda,
|
||||
test_noop_converter_real_function,
|
||||
test_noop_converter_lambda_with_modification,
|
||||
]:
|
||||
run_test(cat, fn)
|
||||
|
||||
log.warning('=== validate_state_dict_load ===')
|
||||
cat = category('validate')
|
||||
for fn in [
|
||||
test_validate_accepts_buffer_only_missing,
|
||||
test_validate_rejects_unexpected,
|
||||
test_validate_rejects_hard_missing,
|
||||
test_validate_empty_passes,
|
||||
]:
|
||||
run_test(cat, fn)
|
||||
|
||||
log.warning('=== register / lookup ===')
|
||||
cat = category('registry')
|
||||
for fn in [
|
||||
test_register_with_explicit_spec,
|
||||
test_register_with_default_spec,
|
||||
test_register_idempotent_replaces,
|
||||
test_register_rejects_mismatched_cls,
|
||||
test_lookup_synthesizes_default_for_unregistered,
|
||||
]:
|
||||
run_test(cat, fn)
|
||||
|
||||
log.warning('=== auto_pickup_converter ===')
|
||||
cat = category('autopickup')
|
||||
for fn in [
|
||||
test_auto_pickup_returns_none_for_unknown_class,
|
||||
test_auto_pickup_returns_real_diffusers_converter,
|
||||
test_auto_pickup_skips_noop_converter_qwen,
|
||||
]:
|
||||
run_test(cat, fn)
|
||||
|
||||
log.warning('=== TransformerSpec / SiblingSpec ===')
|
||||
cat = category('specs')
|
||||
for fn in [
|
||||
test_transformer_spec_defaults,
|
||||
test_sibling_spec_defaults,
|
||||
test_transformer_spec_is_frozen,
|
||||
]:
|
||||
run_test(cat, fn)
|
||||
|
||||
log.warning('=== end-to-end load ===')
|
||||
cat = category('load')
|
||||
for fn in [
|
||||
test_load_end_to_end_with_bfl_prefix_no_converter,
|
||||
test_load_end_to_end_with_sibling_partition,
|
||||
test_load_raises_on_missing_sibling_class,
|
||||
test_load_rejects_non_safetensors,
|
||||
]:
|
||||
run_test(cat, fn)
|
||||
|
||||
log.warning('=== Results ===')
|
||||
total_passed = 0
|
||||
total_failed = 0
|
||||
for cat_name, info in results.items():
|
||||
ok = info['failed'] == 0
|
||||
status = 'PASS' if ok else 'FAIL'
|
||||
log.info(f" {cat_name}: {info['passed']} passed, {info['failed']} failed [{status}]")
|
||||
total_passed += info['passed']
|
||||
total_failed += info['failed']
|
||||
log.warning(f'Total: {total_passed} passed, {total_failed} failed')
|
||||
return total_failed == 0
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
import time
|
||||
t0 = time.time()
|
||||
ok = run_all()
|
||||
log.warning(f'Total time: {time.time() - t0:.2f}s')
|
||||
sys.exit(0 if ok else 1)
|
||||
Reference in New Issue
Block a user