fix(mm): stop the SDNQ guards from aborting folder-encoder identification

invokeai/backend/model_manager/configs/qwen3_encoder.py

@@ -1,4 +1,5 @@
 import json
+from collections.abc import Iterable
 from typing import Any, Literal, Optional, Self
 
 from pydantic import Field
@@ -14,6 +15,11 @@
 from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
 from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType, Qwen3VariantType
 from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor
+from invokeai.backend.quantization.sdnq.detection import (
+    folder_has_sdnq_marker,
+    safetensors_have_sdnq_keys,
+    safetensors_tensor_names,
+)
 from invokeai.backend.quantization.sdnq.sdnq_tensor import SDNQTensor
 
 
@@ -55,12 +61,34 @@ def _has_sdnq_tensors(state_dict: dict[str | int, Any]) -> bool:
 def _has_sdnq_keys(state_dict: dict[str | int, Any]) -> bool:
     """Check if state dict has SDNQ-style keys (weight + scale pairs)."""
     keys = {k for k in state_dict.keys() if isinstance(k, str)}
-    for key in keys:
-        if key.endswith(".weight"):
-            base = key[:-7]
-            if f"{base}.scale" in keys:
-                return True
-    return False
+    return any(key.endswith(".weight") and f"{key[: -len('.weight')]}.scale" in keys for key in keys)
+
+
+def _folder_tensor_names(mod: ModelOnDisk) -> set[str]:
+    """Tensor names declared by every safetensors shard under `mod.path`.
+
+    Folder probes must not go through ``ModelOnDisk.load_state_dict()``. It refuses to pick a file
+    when a folder holds more than one weight file and raises ``ValueError`` — not a
+    ``NotAMatchError`` — which aborts that config's probe rather than declining the model, so a
+    sharded encoder ends up stored as ``unknown``. Every folder-layout Qwen3 encoder we ship is
+    sharded: the FLUX.2 Klein 4B/9B and Z-Image ``text_encoder`` downloads are 2-4 shards each.
+    """
+    return safetensors_tensor_names(mod.weight_files())
+
+
+def _folder_is_sdnq_quantized(mod: ModelOnDisk) -> bool:
+    """True if the Qwen3 encoder folder at `mod.path` holds SDNQ-quantized weights.
+
+    `Qwen3Encoder_Qwen3Encoder_Config` and `Qwen3Encoder_SDNQ_Folder_Config` must be mutually
+    exclusive: one rejects what the other requires. That only holds if both ask the *same* question,
+    so both call this. They used to ask different ones — the unquantized config looked for the marker
+    in `text_encoder/` as well and for keys across every shard, while the SDNQ config looked only at
+    the root and only at safetensors sitting directly in it. A markerless SDNQ encoder in the nested
+    `text_encoder/` layout was therefore rejected by *both* and stored as `unknown`.
+    """
+    if any(folder_has_sdnq_marker(folder) for folder in (mod.path, mod.path / "text_encoder")):
+        return True
+    return safetensors_have_sdnq_keys(mod.weight_files())
 
 
 def _has_t5_encoder_keys(state_dict: dict[str | int, Any]) -> bool:
@@ -92,8 +120,8 @@ def _has_gemma2_keys(state_dict: dict[str | int, Any]) -> bool:
     return False
 
 
-def _has_qwen_vl_visual_tower(state_dict: dict[str | int, Any]) -> bool:
-    """Check if state dict bundles a Qwen-VL vision tower (Qwen2-VL / Qwen2.5-VL / Qwen3-VL).
+def _has_qwen_vl_visual_tower(tensor_names: Iterable[str | int]) -> bool:
+    """Check if the tensor names bundle a Qwen-VL vision tower (Qwen2-VL / Qwen2.5-VL / Qwen3-VL).
 
     VL encoders ship a visual tower alongside the language model, whereas a text-only Qwen3 encoder
     never does. A VL file otherwise satisfies the Qwen3 key heuristic (it has ``model.layers.*`` /
@@ -105,23 +133,29 @@ def _has_qwen_vl_visual_tower(state_dict: dict[str | int, Any]) -> bool:
     layout that ComfyUI single-file Qwen3-VL checkpoints use. Matching only bare ``visual.blocks.*``
     missed that layout, letting a single-file Qwen3-VL 4B encoder match both configs and get misrouted to
     the text-only Qwen3 type - silently breaking the single-file/GGUF Krea-2 encoder install path.
+
+    Takes any iterable of names so a folder probe can pass the union of its safetensors headers; a
+    state dict iterates over its keys, so existing call sites are unaffected.
     """
-    for key in state_dict.keys():
+    for key in tensor_names:
         if isinstance(key, str) and (key.startswith(("visual.", "model.visual.")) or ".visual." in key):
             return True
     return False
 
 
-def _has_qwen3_specific_keys(state_dict: dict[str | int, Any]) -> bool:
+def _has_qwen3_specific_keys(tensor_names: Iterable[str | int]) -> bool:
     """Check for Qwen3-only QK-normalization weights (``q_norm``/``k_norm`` per attention block).
 
     Qwen3 adds an RMSNorm on the query and key projections that Qwen2 does not have. A Qwen2
     causal LM otherwise satisfies the generic ``_has_qwen3_keys`` heuristic (same ``model.layers.*``
     / ``model.embed_tokens.weight`` layout), so this is the discriminator that keeps a Qwen2 file
     from being accepted as a Qwen3 encoder the loader would fail to build. Covers both the
     PyTorch/diffusers naming and the llama.cpp/GGUF naming.
+
+    Takes any iterable of names, so a folder probe can pass safetensors header keys instead of a
+    state dict it cannot load.
     """
-    for key in state_dict.keys():
+    for key in tensor_names:
         if not isinstance(key, str):
             continue
         if ".self_attn.q_norm." in key or ".self_attn.k_norm." in key:
@@ -342,23 +376,10 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -
 
     @classmethod
     def _reject_if_sdnq_quantized(cls, mod: ModelOnDisk) -> None:
-        # Primary signal: quantization_config.json with quant_method="sdnq" (at root or in
-        # text_encoder/). Fallback: SDNQ-style weight+scale key pairs in the state dict. This mirrors
-        # the detection in Qwen3Encoder_SDNQ_Folder_Config so the two stay mutually exclusive.
-        for folder in (mod.path, mod.path / "text_encoder"):
-            quant_config_path = folder / "quantization_config.json"
-            if not quant_config_path.exists():
-                continue
-            try:
-                with open(quant_config_path, "r", encoding="utf-8") as f:
-                    quant_config = json.load(f)
-            except (json.JSONDecodeError, OSError):
-                continue
-            if quant_config.get("quant_method") == "sdnq":
-                raise NotAMatchError("folder is SDNQ-quantized; use Qwen3Encoder_SDNQ_Folder_Config")
-
-        if _has_sdnq_keys(mod.load_state_dict()):
-            raise NotAMatchError("state dict looks SDNQ-quantized; use Qwen3Encoder_SDNQ_Folder_Config")
+        # Shared with Qwen3Encoder_SDNQ_Folder_Config so the two configs cannot both reject (or both
+        # accept) the same folder.
+        if _folder_is_sdnq_quantized(mod):
+            raise NotAMatchError("folder is SDNQ-quantized; use Qwen3Encoder_SDNQ_Folder_Config")
 
     @classmethod
     def _get_variant_from_config(cls, config_path) -> Qwen3VariantType:
@@ -527,25 +548,10 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -
 
         raise_for_override_fields(cls, override_fields)
 
-        matched = False
-
-        # Check for quantization_config.json with quant_method="sdnq"
-        quant_config_path = mod.path / "quantization_config.json"
-        if quant_config_path.exists():
-            with open(quant_config_path, "r", encoding="utf-8") as f:
-                quant_config = json.load(f)
-            if quant_config.get("quant_method") == "sdnq":
-                matched = True
-
-        # Fallback: check if safetensors files have SDNQ-style keys
-        if not matched:
-            safetensors_files = list(mod.path.glob("*.safetensors"))
-            if safetensors_files:
-                state_dict = mod.load_state_dict()
-                if _has_sdnq_keys(state_dict):
-                    matched = True
-
-        if not matched:
+        # Shared with the rejection guard in Qwen3Encoder_Qwen3Encoder_Config: exactly one of the two
+        # configs claims a given folder. A corrupt or foreign quantization_config.json falls through
+        # to the key check there instead of aborting this probe with a JSONDecodeError.
+        if not _folder_is_sdnq_quantized(mod):
             raise NotAMatchError("directory does not look like an SDNQ-quantized Qwen3 encoder")
 
         # A root config.json next to tokenizer files is a complete causal LM (TextLLM), not a Qwen3
@@ -608,25 +614,22 @@ def _validate_is_qwen3_encoder(cls, mod: ModelOnDisk) -> None:
                 "(only Qwen3ForCausalLM is supported)"
             )
 
-        # Fallback for folders without a usable config.json architecture: check the state dict. The
+        # Fallback for folders without a usable config.json architecture: check the tensor names. The
         # generic Qwen keys (model.layers. / model.embed_tokens.weight) are NOT enough — Qwen2 and
         # Qwen2-VL folders carry exactly the same ones, and the loader reconstructs a text-only
         # Qwen3ForCausalLM that fails on Qwen2's missing q/k-norm params and on Qwen-VL's visual
         # tower. Mirror the single-file path: reject a bundled visual tower and require the
         # Qwen3-only q/k-norm weights. An SDNQ transformer/VAE folder has transformer_blocks. /
-        # decoder. keys instead and is rejected by both checks. Loading the state dict can raise for
-        # sharded folders (multiple weight files), so treat that as "no usable signal" rather than
-        # letting it abort identification.
-        try:
-            state_dict = mod.load_state_dict()
-        except Exception:
-            state_dict = {}
+        # decoder. keys instead and is rejected by both checks. Names come from the safetensors
+        # headers because a sharded folder has no single state dict to load — that path used to
+        # yield no signal at all and reject every sharded markerless SDNQ encoder.
+        tensor_names = _folder_tensor_names(mod)
 
-        if _has_qwen_vl_visual_tower(state_dict):
+        if _has_qwen_vl_visual_tower(tensor_names):
             raise NotAMatchError(
                 "state dict bundles a Qwen-VL visual tower; this is a Qwen-VL encoder, not a text-only Qwen3 encoder"
             )
-        if _has_qwen3_specific_keys(state_dict):
+        if _has_qwen3_specific_keys(tensor_names):
             return
 
         raise NotAMatchError(

invokeai/backend/quantization/sdnq/detection.py

@@ -10,18 +10,52 @@
 """
 
 import json
+from collections.abc import Iterable
 from pathlib import Path
 
 from safetensors import safe_open
 
 _QUANTIZATION_CONFIG_FILENAME = "quantization_config.json"
 
 
-def folder_has_sdnq_keys(folder_path: Path) -> bool:
-    """True if the safetensors in `folder_path` carry an SDNQ ``<name>.weight`` / ``<name>.scale`` pair.
+def safetensors_tensor_names(files: Iterable[Path]) -> set[str]:
+    """Union of the tensor names declared by `files`, read from their headers.
+
+    Identification only ever needs the *names* — "is there a `<name>.scale` next to this
+    `<name>.weight`", "is there a visual tower" — never the tensor data, and the header carries them
+    per file. That is what makes this shard-safe, which `ModelOnDisk.load_state_dict()` is not: it
+    refuses to pick a file when a folder holds more than one weight file and raises `ValueError` —
+    not a `NotAMatchError` — so it aborts that config's probe instead of declining the model. Every
+    sharded encoder folder (the FLUX.2 Klein and Z-Image `text_encoder` downloads are all 2-4 shards)
+    hit that and fell back to `unknown`.
+
+    Non-safetensors entries are skipped: they are the set `sdnq_sd_loader` reads, and a file whose
+    header cannot be read contributes nothing rather than failing the whole check.
+    """
+    names: set[str] = set()
+    for file in sorted(files):
+        if file.suffix != ".safetensors":
+            continue
+        try:
+            with safe_open(file, framework="pt", device="cpu") as f:
+                names.update(f.keys())
+        except Exception:
+            continue
+    return names
+
+
+def safetensors_have_sdnq_keys(files: Iterable[Path]) -> bool:
+    """True if `files` carry an SDNQ ``<name>.weight`` / ``<name>.scale`` pair.
 
     The pair is resolved across the union of every shard, never within a single file: sharding splits
     a checkpoint by tensor order, so a weight and its scale routinely land in different files.
+    """
+    names = safetensors_tensor_names(files)
+    return any(name.endswith(".weight") and f"{name[: -len('.weight')]}.scale" in names for name in names)
+
+
+def folder_has_sdnq_keys(folder_path: Path) -> bool:
+    """True if the safetensors directly in `folder_path` carry an SDNQ weight/scale pair.
 
     Only safetensors are inspected, which is the same set `sdnq_sd_loader` reads — a `.bin` holding
     SDNQ-shaped tensors is not something we could load anyway, so calling it SDNQ would only move the
@@ -30,15 +64,23 @@ def folder_has_sdnq_keys(folder_path: Path) -> bool:
     if not folder_path.is_dir():
         return False
 
-    keys: set[str] = set()
-    for shard in sorted(folder_path.glob("*.safetensors")):
-        try:
-            with safe_open(shard, framework="pt", device="cpu") as f:
-                keys.update(f.keys())
-        except Exception:
-            continue
+    return safetensors_have_sdnq_keys(folder_path.glob("*.safetensors"))
+
+
+def folder_has_sdnq_marker(folder_path: Path) -> bool:
+    """True if `folder_path` holds a `quantization_config.json` naming SDNQ as the quant method.
 
-    return any(key.endswith(".weight") and f"{key[: -len('.weight')]}.scale" in keys for key in keys)
+    A marker that is missing, unreadable or names another method is not evidence *against* SDNQ
+    weights — callers fall through to the key shape rather than treating False as "not SDNQ".
+    """
+    marker = folder_path / _QUANTIZATION_CONFIG_FILENAME
+    if not marker.is_file():
+        return False
+    try:
+        with open(marker, "r", encoding="utf-8") as f:
+            return json.load(f).get("quant_method") == "sdnq"
+    except (json.JSONDecodeError, OSError):
+        return False
 
 
 def is_sdnq_folder(folder_path: Path) -> bool:
@@ -48,15 +90,4 @@ def is_sdnq_folder(folder_path: Path) -> bool:
     back to the key shape. The fallback is what covers exports that ship no marker — without it such
     a folder reads as plain diffusers to identification and as SDNQ to nothing at all.
     """
-    marker = folder_path / _QUANTIZATION_CONFIG_FILENAME
-    if marker.is_file():
-        try:
-            with open(marker, "r", encoding="utf-8") as f:
-                if json.load(f).get("quant_method") == "sdnq":
-                    return True
-        except (json.JSONDecodeError, OSError):
-            pass
-        # A marker that exists but names another method is not evidence *against* SDNQ keys, so fall
-        # through rather than returning False here.
-
-    return folder_has_sdnq_keys(folder_path)
+    return folder_has_sdnq_marker(folder_path) or folder_has_sdnq_keys(folder_path)