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