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

refactor(native-loader): replace required_markers with reactive fallback

Browse Source 
A converter error or load_state_dict mismatch now raises OverrideArchMismatch, which load_transformer catches to drop the override and load the base transformer. No per-arch markers to maintain.
This commit is contained in:
CalamitousFelicitousness
2026-06-04 21:20:10 +01:00
parent 596cae1b90
commit 32081ade7c
6 changed files with 118 additions and 421 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pipelines/anima/__init__.py
-7
View File
@@ -12,9 +12,6 @@ knobs that differ from the native-loader defaults:
- 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.
- Cosmos 2.0 required markers: the self-attention and adaptive layer-norm
modulation key families every Cosmos block carries. A UNET/DiT override
lacking them is not a Cosmos checkpoint and is dropped before load.
- All other knobs (prefixes, ``acceptable_missing`` buffers) use the defaults
from :mod:`pipelines.native_transformer`.
"""
@@ -40,8 +37,4 @@ ANIMA_SPEC = TransformerSpec(
'unsupported Cosmos 1.0 structure',
),
),
required_markers=(
('self_attn.', 'Cosmos transformer self-attention blocks'),
('adaln_modulation_', 'Cosmos adaptive layer-norm modulation'),
),
)
pipelines/chroma/__init__.py
-9
View File
@@ -4,11 +4,6 @@ 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.
``required_markers`` names the two key families the converter dereferences
unconditionally (the double-stream blocks and the distilled guidance layer).
A UNET/DiT override missing them is not a Chroma checkpoint and is rejected
before the converter, rather than crashing inside it.
"""
import diffusers
@@ -20,8 +15,4 @@ from pipelines.native_transformer import TransformerSpec
CHROMA_SPEC = TransformerSpec(
cls=diffusers.ChromaTransformer2DModel,
converter=convert_chroma_transformer_checkpoint_to_diffusers,
required_markers=(
("double_blocks.", "Chroma double-stream transformer blocks"),
("distilled_guidance_layer.", "Chroma distilled guidance layer"),
),
)
pipelines/flux2_klein/__init__.py
-8
View File
@@ -12,10 +12,6 @@ Routing through :mod:`pipelines.native_transformer` pulls the Klein
``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.
``required_markers`` names the dual-stream block families the converter
expects; a UNET/DiT override lacking them is not a Flux 2 checkpoint and is
rejected before the converter.
"""
import diffusers
@@ -27,8 +23,4 @@ from pipelines.native_transformer import TransformerSpec
FLUX2_KLEIN_SPEC = TransformerSpec(
cls=diffusers.Flux2Transformer2DModel,
converter=convert_flux2_transformer_checkpoint_to_diffusers,
required_markers=(
("double_blocks.", "Flux 2 double-stream transformer blocks"),
("single_blocks.", "Flux 2 single-stream transformer blocks"),
),
)
pipelines/generic_transformer.py
+41 -43
View File
@@ -30,6 +30,25 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer
quant_type = model_quant.get_quant_type(quant_args)
dtype = dtype or devices.dtype
def load_from_repo():
nonlocal quant_args
log.debug(f'Load model: transformer="{repo_id}" cls={cls_name.__name__} subfolder={subfolder} quant="{quant_type}" loader={get_loader("diffusers")} args={load_args}')
if 'sdnq-' in repo_id.lower():
quant_args = {}
if dtype is not None:
load_args['torch_dtype'] = dtype
if subfolder is not None:
load_args['subfolder'] = subfolder
if variant is not None:
load_args['variant'] = variant
return cls_name.from_pretrained(
repo_id,
cache_dir=shared.opts.hfcache_dir,
**load_args,
**quant_args,
**kwargs,
)
local_file = None
from modules import sd_unet
if shared.opts.sd_unet is not None and shared.opts.sd_unet != 'Default':
@@ -38,19 +57,6 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer
elif os.path.exists(sd_unet.unet_dict[shared.opts.sd_unet]):
local_file = sd_unet.unet_dict[shared.opts.sd_unet]
# drop a UNET/DiT override whose architecture does not match this model's
# transformer and load the base transformer instead of crashing inside the
# arch-specific converter; header-only check, so a large mismatched file is
# rejected before the eager state-dict read
if local_file is not None and local_file.lower().endswith('.safetensors') and native_spec is not None:
from pipelines import native_transformer
compatible, reason = native_transformer.check_override_compatible(local_file, native_spec)
if not compatible:
log.warning(f'Load model: transformer override="{shared.opts.sd_unet}" incompatible with cls={cls_name.__name__} ({reason}); ignoring override and loading base transformer')
shared.opts.data['sd_unet'] = 'Default'
sd_unet.loaded_unet = None
local_file = None
# 1. load gguf
if local_file is not None and local_file.lower().endswith('.gguf'):
log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader={get_loader("diffusers")} args={load_args}')
@@ -62,19 +68,25 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer
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,
)
try:
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,
)
except native_transformer.OverrideArchMismatch as e:
log.warning(f'Load model: transformer override="{shared.opts.sd_unet}" incompatible with cls={cls_name.__name__} ({e}); ignoring override and loading base transformer')
shared.opts.data['sd_unet'] = 'Default'
sd_unet.loaded_unet = None
transformer = load_from_repo()
# 3. load safetensors with diffusers loader
elif local_file is not None and local_file.lower().endswith('.safetensors'):
@@ -91,24 +103,10 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer
**kwargs,
)
# 4. default loading from diffusers repo
# 4. default loading from diffusers repo (also the fallback when an
# incompatible override is dropped above)
else:
log.debug(f'Load model: transformer="{repo_id}" cls={cls_name.__name__} subfolder={subfolder} quant="{quant_type}" loader={get_loader("diffusers")} args={load_args}')
if 'sdnq-' in repo_id.lower():
quant_args = {}
if dtype is not None:
load_args['torch_dtype'] = dtype
if subfolder is not None:
load_args['subfolder'] = subfolder
if variant is not None:
load_args['variant'] = variant
transformer = cls_name.from_pretrained(
repo_id,
cache_dir=shared.opts.hfcache_dir,
**load_args,
**quant_args,
**kwargs,
)
transformer = load_from_repo()
sd_models.allow_post_quant = False # we already handled it
if shared.opts.diffusers_offload_mode != 'none' and transformer is not None:
pipelines/native_transformer.py
+36 -161
View File
@@ -21,9 +21,7 @@ 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) and required markers (catches a file from a wholly
different architecture, e.g. an Anima checkpoint selected as a Chroma
override, before it reaches the arch-specific converter).
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
@@ -63,15 +61,15 @@ DEFAULT_ACCEPTABLE_MISSING: tuple[str, ...] = (
class OverrideArchMismatch(Exception):
"""Raised when a user-selected UNET/DiT override file does not match the
architecture of the spec's transformer class: it is missing the arch's
defining key families (``required_markers``) or carries a forbidden one.
"""Raised when a user-selected UNET/DiT override cannot be loaded as the
spec's transformer class: the arch-specific converter rejects its keys, or
``load_state_dict`` reports a structural mismatch (unexpected or missing
required keys).
Callers that can recover catch this to drop the override and fall back to
the base repo transformer (see
:func:`pipelines.generic_transformer.load_transformer`). Where it is not
caught it surfaces as a clear load error instead of an opaque crash inside
an arch-specific converter.
:func:`pipelines.generic_transformer.load_transformer` catches this to drop
the override and load the base repo transformer instead, so a stale or
wrong-arch UNET selection degrades to the base model rather than crashing
inside a converter.
"""
@@ -99,16 +97,6 @@ class TransformerSpec:
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.
``required_markers`` and ``forbidden_markers`` are the two halves of the
structural arch check. Each is a tuple of ``(marker, description)``.
Forbidden markers are exact keys that must be absent (catches an
incompatible variant of the same family). Required markers are key-family
substrings that must each match at least one key (catches a file from a
different architecture entirely, e.g. an Anima checkpoint selected as a
Chroma override). Specs with a converter should set ``required_markers``,
since the converter runs before the normal load-time key validation and a
wrong-arch file would otherwise crash inside it.
"""
cls: type
@@ -118,7 +106,6 @@ class TransformerSpec:
siblings: dict[str, SiblingSpec] = field(default_factory=dict)
acceptable_missing: tuple[str, ...] = DEFAULT_ACCEPTABLE_MISSING
forbidden_markers: tuple[tuple[str, str], ...] = ()
required_markers: tuple[tuple[str, str], ...] = ()
def make_default_spec(cls: type) -> TransformerSpec:
@@ -252,7 +239,6 @@ def load(
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)
check_required_markers(state_dict, spec.required_markers, spec.cls.__name__, local_file)
transformer_sd, sibling_sds = partition_siblings(state_dict, spec.siblings)
del state_dict
@@ -311,82 +297,39 @@ def load(
return transformer, loaded_siblings
def detect_prefix(keys, prefixes: tuple[str, ...], type_name: str) -> str:
"""Return the single known prefix shared by every key, or ``""`` when the
keys are bare (no known prefix present).
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.
Operates on key names only (any iterable of strings), so it can run against
a header-only key list without reading tensor data.
"""
sorted_prefixes = sorted(prefixes, key=len, reverse=True)
counts: dict[str, int] = {}
total = 0
seen = 0
for key in keys:
total += 1
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:
return ""
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]})"
)
return dominant
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.
Thin wrapper over :func:`detect_prefix` that rewrites the dict once a single
dominant prefix is confirmed; bare key sets pass through unchanged.
"""
dominant = detect_prefix(state_dict, prefixes, type_name)
if dominant == "":
log.debug(f"Load model: type={type_name} native_transformer prefix=bare")
return state_dict
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 first_present_marker(keys, markers: tuple[tuple[str, str], ...]) -> tuple[str, str] | None:
"""Return the first ``(marker, description)`` whose exact key is present in
``keys``, else None. ``keys`` is any container supporting ``in`` over key
names (dict or set). Backs the forbidden-marker check.
"""
for marker, description in markers:
if marker in keys:
return marker, description
return None
def first_missing_marker(keys, markers: tuple[tuple[str, str], ...]) -> tuple[str, str] | None:
"""Return the first ``(marker, description)`` whose substring matches no key
in ``keys``, else None. Each marker names a key family (e.g.
``"double_blocks."``) that a valid file of the arch must contain. Backs the
required-marker check.
"""
if not markers:
return None
key_list = keys if isinstance(keys, (list, tuple)) else list(keys)
for marker, description in markers:
if not any(marker in k for k in key_list):
return marker, description
return None
def check_forbidden_markers(
state_dict: dict,
forbidden_markers: tuple[tuple[str, str], ...],
@@ -399,90 +342,12 @@ def check_forbidden_markers(
file is from an incompatible architecture variant (e.g. Cosmos 1.0 keys
showing up in a Cosmos 2.0 loader path).
"""
hit = first_present_marker(state_dict, forbidden_markers)
if hit is not None:
marker, description = hit
raise ValueError(
f"Load model: type={type_name} native_transformer rejects "
f'"{os.path.basename(local_file)}" ({description}; marker key {marker!r})'
)
def check_required_markers(
state_dict: dict,
required_markers: tuple[tuple[str, str], ...],
type_name: str,
local_file: str,
) -> None:
"""Raise :class:`OverrideArchMismatch` if any required arch marker matches
no key. Each marker names a key family the arch's checkpoint must carry
(e.g. Chroma's ``"double_blocks."``); absence means the file is not a
``type_name`` checkpoint. No-op when ``required_markers`` is empty.
Runs before the arch-specific converter, which would otherwise crash on a
wrong-arch file with an opaque error instead of this actionable one.
"""
miss = first_missing_marker(state_dict, required_markers)
if miss is not None:
marker, description = miss
raise OverrideArchMismatch(
f"Load model: type={type_name} native_transformer rejects "
f'"{os.path.basename(local_file)}" (missing {description}; '
f"no key contains {marker!r})"
)
def peek_keys(local_file: str) -> list[str]:
"""Read only the safetensors header key list (no tensor data).
Lets :func:`check_override_compatible` reject a multi-GB mismatched override
before the eager full-tensor read in
:func:`modules.sd_models.read_state_dict`.
"""
import safetensors.torch
with safetensors.torch.safe_open(local_file, framework="pt", device="cpu") as f:
return list(f.keys())
def check_override_compatible(local_file: str, spec: TransformerSpec) -> tuple[bool, str]:
"""Header-only architecture check for a user-selected UNET/DiT override.
Reads just the safetensors key list, detects and strips the spec's prefix,
then applies the spec's forbidden and required markers. Returns
``(True, "")`` when the file looks loadable as ``spec.cls``, else
``(False, reason)`` with a short human-readable reason. Never reads tensor
data.
On a header read error returns ``(True, "")`` to defer to the full loader
rather than misclassifying an I/O problem as an arch mismatch; the full
load then raises with the real error.
"""
try:
raw_keys = peek_keys(local_file)
except Exception as e: # pylint: disable=broad-except
log.debug(
f"Load model: native_transformer compatibility peek failed "
f'file="{os.path.basename(local_file)}": {e}'
)
return True, ""
try:
dominant = detect_prefix(raw_keys, spec.prefixes, spec.cls.__name__)
except ValueError as e:
return False, str(e)
if dominant:
offset = len(dominant)
stripped = [k[offset:] for k in raw_keys]
else:
stripped = raw_keys
forbidden = first_present_marker(set(stripped), spec.forbidden_markers)
if forbidden is not None:
marker, description = forbidden
return False, f"{description} (marker key {marker!r})"
missing = first_missing_marker(stripped, spec.required_markers)
if missing is not None:
marker, description = missing
return False, f"missing {description} (no key contains {marker!r})"
return True, ""
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(
@@ -655,7 +520,14 @@ def build_component(
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)
try:
sd = converter(state_dict)
except Exception as e:
raise OverrideArchMismatch(
f"Load model: type={cls.__name__} native_transformer converter "
f"{converter.__name__} rejected the override ({type(e).__name__}: {e}); "
f"file does not look like a {cls.__name__} checkpoint"
) from e
else:
sd = state_dict
@@ -683,6 +555,8 @@ def build_component(
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 OverrideArchMismatch:
raise
except Exception as e:
log.error(f"Load model: native_transformer {component_name} load failed: {e}")
errors.display(e, "Load")
@@ -713,13 +587,14 @@ def validate_state_dict_load(
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
"""Raise :class:`OverrideArchMismatch` if load_state_dict produced
unexpected keys or non-acceptable missing keys, which means the override's
weights do not fit the target class. 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(
raise OverrideArchMismatch(
f"Load model: native_transformer {component_name} has {len(unexpected)} "
f"unexpected keys (sample: {sample})"
)
@@ -728,7 +603,7 @@ def validate_state_dict_load(
]
if hard_missing:
sample = ", ".join(hard_missing[:5])
raise ValueError(
raise OverrideArchMismatch(
f"Load model: native_transformer {component_name} missing "
f"{len(hard_missing)} required keys (sample: {sample})"
)
test/test-native-transformer.py
+41 -193
View File
@@ -273,8 +273,8 @@ def test_validate_rejects_unexpected():
unexpected=['some.junk.weight'],
acceptable_missing=(),
)
raise AssertionError('expected ValueError')
except ValueError as e:
raise AssertionError('expected OverrideArchMismatch')
except nt.OverrideArchMismatch as e:
assert 'unexpected' in str(e)
assert 'some.junk.weight' in str(e)
@@ -287,8 +287,8 @@ def test_validate_rejects_hard_missing():
unexpected=[],
acceptable_missing=('rope.',),
)
raise AssertionError('expected ValueError')
except ValueError as e:
raise AssertionError('expected OverrideArchMismatch')
except nt.OverrideArchMismatch as e:
msg = str(e)
assert 'missing' in msg
assert 'layers.0.weight' in msg
@@ -375,7 +375,6 @@ def test_transformer_spec_defaults():
assert spec.siblings == {}
assert spec.acceptable_missing == ('rope.', 'pos_embedder.', 'learnable_pos_embed.')
assert spec.forbidden_markers == ()
assert spec.required_markers == ()
def test_sibling_spec_defaults():
@@ -685,31 +684,57 @@ def test_load_rejects_non_safetensors():
assert '.safetensors' in str(e)
def test_load_raises_override_arch_mismatch():
"""A file missing the spec's required markers raises OverrideArchMismatch
from load() (the defense-in-depth contract), before the converter runs."""
def crashing_converter(sd):
"""diffusers-style layer count that blows up when the block family is
absent, mirroring convert_chroma_..._to_diffusers on a wrong-arch file."""
return list(set(int(k.split('.')[1]) for k in sd if 'double_blocks.' in k))[-1]
def test_build_component_converter_crash_raises_mismatch():
"""A converter that crashes on wrong-arch keys is wrapped as
OverrideArchMismatch (chaining the original), not the raw IndexError."""
try:
nt.build_component(
component_name='transformer',
state_dict={'blocks.0.self_attn.weight': torch.zeros(2)},
config={'dim': 8},
cls=MockMiniTransformer,
converter=crashing_converter,
acceptable_missing=(),
quant_args={},
quant_type=None,
)
raise AssertionError('expected OverrideArchMismatch')
except nt.OverrideArchMismatch as e:
assert 'MockMiniTransformer' in str(e)
assert isinstance(e.__cause__, IndexError), 'original error must be chained'
def test_load_converter_crash_raises_mismatch():
"""End-to-end: a crashing converter surfaces from load() as
OverrideArchMismatch so load_transformer can drop the override."""
fd, path = tempfile.mkstemp(suffix='.safetensors')
try:
raw = {'model.diffusion_model.blocks.0.self_attn.q_proj.weight': torch.zeros(4)}
raw = {'model.diffusion_model.blocks.0.self_attn.weight': torch.zeros(8, 8)}
write_fixture(raw, fd, path)
orig_fetch = nt.fetch_component_config
nt.fetch_component_config = lambda repo, sub: {'dim': 8}
from modules import model_quant
orig_get_dit = model_quant.get_dit_args
orig_get_qtype = model_quant.get_quant_type
model_quant.get_dit_args = lambda *a, **k: ({}, {})
model_quant.get_quant_type = lambda *a, **k: None
try:
spec = nt.TransformerSpec(
cls=MockMiniTransformer,
required_markers=(('double_blocks.', 'Chroma double-stream blocks'),),
)
spec = nt.TransformerSpec(cls=MockMiniTransformer, converter=crashing_converter)
raised = False
try:
nt.load(local_file=path, repo_id='fake/repo', spec=spec, diffusers_cfg={})
except nt.OverrideArchMismatch as e:
raised = True
assert 'double_blocks.' in str(e)
assert 'MockMiniTransformer' in str(e)
assert raised, 'expected OverrideArchMismatch'
finally:
nt.fetch_component_config = orig_fetch
model_quant.get_dit_args = orig_get_dit
model_quant.get_quant_type = orig_get_qtype
finally:
@@ -717,153 +742,6 @@ def test_load_raises_override_arch_mismatch():
os.unlink(path)
# ============================================================
# detect_prefix
# ============================================================
def test_detect_prefix_bare_returns_empty():
keys = ['layers.0.weight', 'layers.0.bias']
assert nt.detect_prefix(keys, nt.DEFAULT_PREFIXES, 'Test') == ''
def test_detect_prefix_returns_dominant():
keys = [f'model.diffusion_model.layers.{i}.weight' for i in range(5)]
assert nt.detect_prefix(keys, nt.DEFAULT_PREFIXES, 'Test') == 'model.diffusion_model.'
def test_detect_prefix_mixed_raises():
keys = ['model.diffusion_model.layers.0.weight', 'net.layers.0.weight']
try:
nt.detect_prefix(keys, nt.DEFAULT_PREFIXES, 'Test')
raise AssertionError('expected ValueError')
except ValueError as e:
assert 'mixed prefixes' in str(e)
def test_detect_prefix_accepts_iterable_keys():
"""Header-only key lists (no tensors) must work, not just dicts."""
keys = ['diffusion_model.a.weight', 'diffusion_model.b.weight']
assert nt.detect_prefix(keys, nt.DEFAULT_PREFIXES, 'Test') == 'diffusion_model.'
# ============================================================
# check_required_markers
# ============================================================
def test_required_markers_passes_when_present():
sd = {'double_blocks.0.x': 1, 'distilled_guidance_layer.in_proj.weight': 2}
markers = (('double_blocks.', 'double blocks'), ('distilled_guidance_layer.', 'guidance'))
nt.check_required_markers(sd, markers, 'Chroma', '/tmp/x.safetensors')
# no exception = pass
def test_required_markers_raises_when_missing():
sd = {'blocks.0.self_attn.q_proj.weight': 1}
markers = (('double_blocks.', 'Chroma double-stream blocks'),)
try:
nt.check_required_markers(sd, markers, 'Chroma', '/tmp/x.safetensors')
raise AssertionError('expected OverrideArchMismatch')
except nt.OverrideArchMismatch as e:
msg = str(e)
assert 'Chroma double-stream blocks' in msg
assert 'double_blocks.' in msg
def test_required_markers_requires_all():
"""All listed markers must match; one missing raises."""
sd = {'double_blocks.0.x': 1}
markers = (('double_blocks.', 'double blocks'), ('distilled_guidance_layer.', 'guidance'))
try:
nt.check_required_markers(sd, markers, 'Chroma', '/tmp/x.safetensors')
raise AssertionError('expected OverrideArchMismatch')
except nt.OverrideArchMismatch as e:
assert 'guidance' in str(e)
def test_required_markers_empty_no_op():
sd = {'anything.weight': 1}
nt.check_required_markers(sd, (), 'Test', '/tmp/x.safetensors')
# ============================================================
# check_override_compatible (header-only)
# ============================================================
def write_keys_fixture(keys: list) -> str:
fd, path = tempfile.mkstemp(suffix='.safetensors')
os.close(fd)
safetensors.torch.save_file({k: torch.zeros(2) for k in keys}, path)
return path
def test_compat_accepts_matching_arch():
path = write_keys_fixture([
'model.diffusion_model.double_blocks.0.img_attn.qkv.weight',
'model.diffusion_model.distilled_guidance_layer.in_proj.weight',
])
try:
spec = nt.TransformerSpec(
cls=MockMiniTransformer,
converter=lambda sd: sd,
required_markers=(('double_blocks.', 'double blocks'), ('distilled_guidance_layer.', 'guidance')),
)
ok, reason = nt.check_override_compatible(path, spec)
assert ok is True, reason
assert reason == ''
finally:
os.unlink(path)
def test_compat_rejects_missing_required():
path = write_keys_fixture([
'model.diffusion_model.blocks.0.self_attn.q_proj.weight',
'model.diffusion_model.blocks.0.adaln_modulation_self_attn.1.weight',
])
try:
spec = nt.TransformerSpec(
cls=MockMiniTransformer,
converter=lambda sd: sd,
required_markers=(('double_blocks.', 'Chroma double-stream blocks'),),
)
ok, reason = nt.check_override_compatible(path, spec)
assert ok is False
assert 'double_blocks.' in reason
finally:
os.unlink(path)
def test_compat_rejects_forbidden_marker():
path = write_keys_fixture(['double_blocks.0.x.weight', 'legacy.bad.key.weight'])
try:
spec = nt.TransformerSpec(
cls=MockMiniTransformer,
forbidden_markers=(('legacy.bad.key.weight', 'legacy format'),),
required_markers=(('double_blocks.', 'double blocks'),),
)
ok, reason = nt.check_override_compatible(path, spec)
assert ok is False
assert 'legacy format' in reason
finally:
os.unlink(path)
def test_compat_defers_on_unreadable_file():
spec = nt.TransformerSpec(cls=MockMiniTransformer, required_markers=(('double_blocks.', 'd'),))
ok, reason = nt.check_override_compatible('/nonexistent/path/model.safetensors', spec)
assert ok is True
assert reason == ''
def test_compat_no_required_markers_accepts_any():
path = write_keys_fixture(['whatever.weight'])
try:
spec = nt.TransformerSpec(cls=MockMiniTransformer)
ok, reason = nt.check_override_compatible(path, spec)
assert ok is True, reason
finally:
os.unlink(path)
# ============================================================
# Run
# ============================================================
@@ -947,37 +825,6 @@ def run_all():
]:
run_test(cat, fn)
log.warning('=== detect_prefix ===')
cat = category('detect')
for fn in [
test_detect_prefix_bare_returns_empty,
test_detect_prefix_returns_dominant,
test_detect_prefix_mixed_raises,
test_detect_prefix_accepts_iterable_keys,
]:
run_test(cat, fn)
log.warning('=== check_required_markers ===')
cat = category('required')
for fn in [
test_required_markers_passes_when_present,
test_required_markers_raises_when_missing,
test_required_markers_requires_all,
test_required_markers_empty_no_op,
]:
run_test(cat, fn)
log.warning('=== check_override_compatible ===')
cat = category('compat')
for fn in [
test_compat_accepts_matching_arch,
test_compat_rejects_missing_required,
test_compat_rejects_forbidden_marker,
test_compat_defers_on_unreadable_file,
test_compat_no_required_markers_accepts_any,
]:
run_test(cat, fn)
log.warning('=== end-to-end load ===')
cat = category('load')
for fn in [
@@ -986,7 +833,8 @@ def run_all():
test_load_end_to_end_with_sibling_partition,
test_load_raises_on_missing_sibling_class,
test_load_rejects_non_safetensors,
test_load_raises_override_arch_mismatch,
test_build_component_converter_crash_raises_mismatch,
test_load_converter_crash_raises_mismatch,
]:
run_test(cat, fn)