From 596cae1b90008ecfd83f96cc5bba88c7845dcdec Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Wed, 3 Jun 2026 03:23:24 +0100 Subject: [PATCH 1/4] fix(native-loader): handle arch-mismatched UNET/DiT override on cold start --- pipelines/anima/__init__.py | 7 + pipelines/chroma/__init__.py | 9 ++ pipelines/flux2_klein/__init__.py | 8 ++ pipelines/generic_transformer.py | 15 ++- pipelines/native_transformer.py | 174 ++++++++++++++++++++++-- test/test-native-transformer.py | 212 ++++++++++++++++++++++++++++++ 6 files changed, 411 insertions(+), 14 deletions(-) diff --git a/pipelines/anima/__init__.py b/pipelines/anima/__init__.py index 6d7953090..506ff32ae 100644 --- a/pipelines/anima/__init__.py +++ b/pipelines/anima/__init__.py @@ -12,6 +12,9 @@ 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`. """ @@ -37,4 +40,8 @@ ANIMA_SPEC = TransformerSpec( 'unsupported Cosmos 1.0 structure', ), ), + required_markers=( + ('self_attn.', 'Cosmos transformer self-attention blocks'), + ('adaln_modulation_', 'Cosmos adaptive layer-norm modulation'), + ), ) diff --git a/pipelines/chroma/__init__.py b/pipelines/chroma/__init__.py index b8eaa68c6..771e15cf8 100644 --- a/pipelines/chroma/__init__.py +++ b/pipelines/chroma/__init__.py @@ -4,6 +4,11 @@ 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 @@ -15,4 +20,8 @@ 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"), + ), ) diff --git a/pipelines/flux2_klein/__init__.py b/pipelines/flux2_klein/__init__.py index 72d898a92..68df69745 100644 --- a/pipelines/flux2_klein/__init__.py +++ b/pipelines/flux2_klein/__init__.py @@ -12,6 +12,10 @@ 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 @@ -23,4 +27,8 @@ 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"), + ), ) diff --git a/pipelines/generic_transformer.py b/pipelines/generic_transformer.py index 6f6569708..c67975eb4 100644 --- a/pipelines/generic_transformer.py +++ b/pipelines/generic_transformer.py @@ -31,13 +31,26 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer dtype = dtype or devices.dtype local_file = None + from modules import sd_unet if shared.opts.sd_unet is not None and shared.opts.sd_unet != 'Default': - from modules import sd_unet if shared.opts.sd_unet not in list(sd_unet.unet_dict): log.error(f'Load module: type=transformer file="{shared.opts.sd_unet}" not found') 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}') diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index b52326e5a..af43b71da 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -21,7 +21,9 @@ 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). + 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). 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 @@ -60,6 +62,19 @@ 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. + + 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. + """ + + @dataclass(frozen=True) class SiblingSpec: """Describes a non-transformer component that may ship inline in the same @@ -84,6 +99,16 @@ 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 @@ -93,6 +118,7 @@ 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: @@ -226,6 +252,7 @@ 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 @@ -284,39 +311,82 @@ def load( 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. +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). 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 state_dict: + for key in keys: + total += 1 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 + return "" 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], ...], @@ -329,12 +399,90 @@ 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). """ - 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})' - ) + 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, "" def partition_siblings( diff --git a/test/test-native-transformer.py b/test/test-native-transformer.py index 538ce84a5..fbb43f3d9 100644 --- a/test/test-native-transformer.py +++ b/test/test-native-transformer.py @@ -375,6 +375,7 @@ 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(): @@ -684,6 +685,185 @@ 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.""" + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + raw = {'model.diffusion_model.blocks.0.self_attn.q_proj.weight': torch.zeros(4)} + write_fixture(raw, fd, path) + 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'),), + ) + 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 raised, 'expected OverrideArchMismatch' + finally: + model_quant.get_dit_args = orig_get_dit + model_quant.get_quant_type = orig_get_qtype + finally: + if os.path.exists(path): + 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 # ============================================================ @@ -767,6 +947,37 @@ 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 [ @@ -775,6 +986,7 @@ 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, ]: run_test(cat, fn) From 32081ade7c967a709093e6a258f3071d40d9dfdc Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Thu, 4 Jun 2026 21:20:10 +0100 Subject: [PATCH 2/4] refactor(native-loader): replace required_markers with reactive fallback 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. --- pipelines/anima/__init__.py | 7 - pipelines/chroma/__init__.py | 9 -- pipelines/flux2_klein/__init__.py | 8 - pipelines/generic_transformer.py | 84 ++++++----- pipelines/native_transformer.py | 197 +++++-------------------- test/test-native-transformer.py | 234 ++++++------------------------ 6 files changed, 118 insertions(+), 421 deletions(-) diff --git a/pipelines/anima/__init__.py b/pipelines/anima/__init__.py index 506ff32ae..6d7953090 100644 --- a/pipelines/anima/__init__.py +++ b/pipelines/anima/__init__.py @@ -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'), - ), ) diff --git a/pipelines/chroma/__init__.py b/pipelines/chroma/__init__.py index 771e15cf8..b8eaa68c6 100644 --- a/pipelines/chroma/__init__.py +++ b/pipelines/chroma/__init__.py @@ -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"), - ), ) diff --git a/pipelines/flux2_klein/__init__.py b/pipelines/flux2_klein/__init__.py index 68df69745..72d898a92 100644 --- a/pipelines/flux2_klein/__init__.py +++ b/pipelines/flux2_klein/__init__.py @@ -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"), - ), ) diff --git a/pipelines/generic_transformer.py b/pipelines/generic_transformer.py index c67975eb4..6bd598e48 100644 --- a/pipelines/generic_transformer.py +++ b/pipelines/generic_transformer.py @@ -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: diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index af43b71da..1961142f3 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -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 ``/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})" ) diff --git a/test/test-native-transformer.py b/test/test-native-transformer.py index fbb43f3d9..4bcd1335a 100644 --- a/test/test-native-transformer.py +++ b/test/test-native-transformer.py @@ -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) From 5f50117bf945c2a2d68ac2c6dc50e4977853deca Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Fri, 5 Jun 2026 20:08:43 +0100 Subject: [PATCH 3/4] docs(native-loader): document shape-mismatch override behavior --- pipelines/native_transformer.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index 1961142f3..a8ebe0629 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -63,13 +63,16 @@ DEFAULT_ACCEPTABLE_MISSING: tuple[str, ...] = ( class OverrideArchMismatch(Exception): """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). + ``load_state_dict`` reports unexpected or missing keys. :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. + + A tensor shape mismatch on otherwise-matching keys is not raised as this; it + surfaces as the native ``load_state_dict`` error, which already names the + conflicting shapes, so it stays a hard load error rather than a fall back. """ From e12ae85d06654d4ed5f9c4dfeb7bd862dea40c14 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 6 Jun 2026 01:08:31 +0100 Subject: [PATCH 4/4] test(native-loader): add shape-mismatch unit and override self-heal API test --- test/test-native-transformer.py | 27 ++++ test/test-override-mismatch-api.py | 215 +++++++++++++++++++++++++++++ 2 files changed, 242 insertions(+) create mode 100644 test/test-override-mismatch-api.py diff --git a/test/test-native-transformer.py b/test/test-native-transformer.py index 4bcd1335a..3120b6ddd 100644 --- a/test/test-native-transformer.py +++ b/test/test-native-transformer.py @@ -710,6 +710,32 @@ def test_build_component_converter_crash_raises_mismatch(): assert isinstance(e.__cause__, IndexError), 'original error must be chained' +def test_build_component_shape_mismatch_is_hard_error(): + """A tensor shape mismatch on otherwise-matching keys stays a hard + RuntimeError (the native size-mismatch message), it is NOT converted to + OverrideArchMismatch and so does not silently fall back to base.""" + # MockMiniTransformer(dim=8) expects (8, 8) projections; feed (4, 4). + sd = { + 'in_proj.weight': torch.randn(4, 4), 'in_proj.bias': torch.zeros(4), + 'out_proj.weight': torch.randn(4, 4), 'out_proj.bias': torch.zeros(4), + } + orig_display = nt.errors.display + nt.errors.display = lambda *a, **k: None # silence the expected traceback dump + try: + nt.build_component( + component_name='transformer', state_dict=sd, config={'dim': 8}, + cls=MockMiniTransformer, converter=None, acceptable_missing=(), + quant_args={}, quant_type=None, + ) + raise AssertionError('expected RuntimeError') + except nt.OverrideArchMismatch: + raise AssertionError('shape mismatch must not be OverrideArchMismatch') from None + except RuntimeError as e: + assert 'size mismatch' in str(e).lower() + finally: + nt.errors.display = orig_display + + 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.""" @@ -834,6 +860,7 @@ def run_all(): test_load_raises_on_missing_sibling_class, test_load_rejects_non_safetensors, test_build_component_converter_crash_raises_mismatch, + test_build_component_shape_mismatch_is_hard_error, test_load_converter_crash_raises_mismatch, ]: run_test(cat, fn) diff --git a/test/test-override-mismatch-api.py b/test/test-override-mismatch-api.py new file mode 100644 index 000000000..da902cb60 --- /dev/null +++ b/test/test-override-mismatch-api.py @@ -0,0 +1,215 @@ +#!/usr/bin/env python +""" +API integration tests for UNET/DiT override architecture-mismatch self-heal. + +Verifies the reactive behavior end-to-end against a running SD.Next instance: +when a UNET/DiT override does not match the base model's architecture, the load +drops the override, resets the UNET dropdown to ``Default``, and the base model +still generates, instead of crashing inside an arch-specific converter. + +Covers: +- GET /sdapi/v1/unets, /sdapi/v1/sd-models (discovery / sanity) +- baseline: base model with no override loads and generates +- mismatch self-heal: base model + wrong-arch override -> override dropped, + sd_unet reset to Default, generation still succeeds +- match preserved (optional): a correct-arch override stays applied + +Requires a running SD.Next instance with the relevant models on disk. Model and +UNET names are environment-specific, so pass them explicitly. Run with no model +args to just list what the server has available. + +Usage: + python test/test-override-mismatch-api.py \ + --url http://127.0.0.1:7860 \ + --base-model "Diffusers/lodestones/Chroma1-HD [0e0c60ece1]" \ + --mismatch-unet "novaOrangeAM_v15" \ + [--match-unet ""] \ + [--steps 4] +""" + +import sys +import time +import argparse +import requests +import urllib3 + +urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) + + +class OverrideMismatchAPITest: + """Drives the override-mismatch self-heal scenarios over the HTTP API.""" + + def __init__(self, base_url, base_model=None, mismatch_unet=None, match_unet=None, steps=4): + self.base_url = base_url.rstrip('/') + self.base_model = base_model + self.mismatch_unet = mismatch_unet + self.match_unet = match_unet + self.steps = steps + self.load_timeout = 900 + self.gen_timeout = 600 + self.passed = 0 + self.failed = 0 + self.skipped = 0 + + # ---- low-level helpers ------------------------------------------------- + + def _get(self, endpoint): + r = requests.get(f'{self.base_url}{endpoint}', timeout=60, verify=False) + r.raise_for_status() + return r.json() + + def _post(self, endpoint, data=None, params=None, timeout=60): + r = requests.post(f'{self.base_url}{endpoint}', json=data, params=params, timeout=timeout, verify=False) + return r + + def record(self, ok, name, detail=''): + tag = 'PASS' if ok else 'FAIL' + self.passed += 1 if ok else 0 + self.failed += 0 if ok else 1 + line = f' {tag}: {name}' + if detail: + line += f' ({detail})' + print(line, flush=True) + + def skip(self, name, reason): + self.skipped += 1 + print(f' SKIP: {name} ({reason})', flush=True) + + # ---- mid-level operations ---------------------------------------------- + + def set_options(self, **kwargs): + r = self._post('/sdapi/v1/options', data=kwargs, timeout=self.load_timeout) + return r.status_code == 200, (r.text[:200] if r.status_code != 200 else '') + + def reload(self, force=True): + r = self._post('/sdapi/v1/reload-checkpoint', params={'force': str(force).lower()}, timeout=self.load_timeout) + return r.status_code == 200, (r.text[:200] if r.status_code != 200 else '') + + def get_sd_unet(self): + return self._get('/sdapi/v1/options').get('sd_unet') + + def generate(self): + payload = {'prompt': 'a photo of a cat', 'steps': self.steps, 'width': 512, 'height': 512, 'save_images': False} + t0 = time.time() + r = self._post('/sdapi/v1/txt2img', data=payload, timeout=self.gen_timeout) + elapsed = time.time() - t0 + if r.status_code != 200: + return False, f'http {r.status_code}: {r.text[:160]}' + body = r.json() + images = body.get('images') or [] + if not images: + return False, f'no images returned ({elapsed:.1f}s)' + return True, f'{elapsed:.1f}s' + + # ---- scenarios --------------------------------------------------------- + + def test_discovery(self): + print('=== discovery ===', flush=True) + try: + unets = self._get('/sdapi/v1/unets') + self.record(isinstance(unets, list), 'GET /sdapi/v1/unets', f'{len(unets)} unets') + models = self._get('/sdapi/v1/sd-models') + self.record(isinstance(models, list), 'GET /sdapi/v1/sd-models', f'{len(models)} models') + if not (self.base_model and self.mismatch_unet): + print(' available UNET names:', flush=True) + for u in unets: + print(f' - {u.get("name")}', flush=True) + print(' available model titles:', flush=True) + for m in models[:40]: + print(f' - {m.get("title")}', flush=True) + except Exception as e: + self.record(False, 'discovery', f'exception: {e}') + + def test_baseline(self): + print('=== baseline (base model, no override) ===', flush=True) + if not self.base_model: + self.skip('baseline', 'no --base-model') + return False + ok, err = self.set_options(sd_unet='Default', sd_model_checkpoint=self.base_model) + if not ok: + self.record(False, 'set base model + Default unet', err) + return False + ok, err = self.reload(force=True) + if not ok: + self.record(False, 'reload base model', err) + return False + ok, detail = self.generate() + self.record(ok, 'generate with base model', detail) + return ok + + def test_mismatch_self_heal(self): + print('=== mismatch self-heal ===', flush=True) + if not (self.base_model and self.mismatch_unet): + self.skip('mismatch self-heal', 'needs --base-model and --mismatch-unet') + return + # configure base model with the wrong-arch override, then force a clean reload + ok, err = self.set_options(sd_model_checkpoint=self.base_model, sd_unet=self.mismatch_unet) + if not ok: + self.record(False, 'set base model + mismatch override', err) + return + ok, err = self.reload(force=True) + # the reload itself must not error out (the whole point of the fix) + self.record(ok, 'reload does not error on mismatched override', err) + # override must self-heal back to Default + healed = self.get_sd_unet() + self.record(healed == 'Default', 'sd_unet reset to Default', f'sd_unet={healed!r}') + # base model must still be usable + gen_ok, detail = self.generate() + self.record(gen_ok, 'generate after self-heal', detail) + + def test_match_preserved(self): + print('=== match preserved (no false drop) ===', flush=True) + if not (self.base_model and self.match_unet): + self.skip('match preserved', 'no --match-unet') + return + ok, err = self.set_options(sd_model_checkpoint=self.base_model, sd_unet=self.match_unet) + if not ok: + self.record(False, 'set base model + matching override', err) + return + ok, err = self.reload(force=True) + self.record(ok, 'reload with matching override', err) + kept = self.get_sd_unet() + self.record(kept == self.match_unet, 'matching override kept (not dropped)', f'sd_unet={kept!r}') + gen_ok, detail = self.generate() + self.record(gen_ok, 'generate with matching override', detail) + + def cleanup(self): + self.set_options(sd_unet='Default') + + def run(self): + try: + self.test_discovery() + self.test_baseline() + self.test_mismatch_self_heal() + self.test_match_preserved() + finally: + self.cleanup() + print('=== results ===', flush=True) + print(f' passed={self.passed} failed={self.failed} skipped={self.skipped}', flush=True) + return self.failed == 0 + + +def main(): + ap = argparse.ArgumentParser(description='Override arch-mismatch self-heal API tests') + ap.add_argument('--url', default='http://127.0.0.1:7860', help='SD.Next base URL') + ap.add_argument('--base-model', default=None, help='checkpoint title to load as the base (DiT arch)') + ap.add_argument('--mismatch-unet', default=None, help='UNET name whose arch does NOT match the base') + ap.add_argument('--match-unet', default=None, help='optional UNET name whose arch DOES match the base') + ap.add_argument('--steps', type=int, default=4) + args = ap.parse_args() + + try: + requests.get(f'{args.url.rstrip("/")}/sdapi/v1/sd-models', timeout=10, verify=False) + except Exception as e: + print(f'cannot reach SD.Next at {args.url}: {e}', flush=True) + return 2 + + ok = OverrideMismatchAPITest( + args.url, base_model=args.base_model, mismatch_unet=args.mismatch_unet, + match_unet=args.match_unet, steps=args.steps, + ).run() + return 0 if ok else 1 + + +if __name__ == '__main__': + sys.exit(main())