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)