fix(mm): stop the SDNQ guards from aborting folder-encoder identification - #9561
Pfannkuchensack wants to merge 6 commits into
Conversation
…ck to unknown
The starter-model download `black-forest-labs/FLUX.2-klein-{4B,9B}::text_encoder+tokenizer`
lands the Qwen3 encoder as several `model-0000N-of-0000M.safetensors` shards. The SDNQ
rejection guard in `Qwen3Encoder_Qwen3Encoder_Config` called `mod.load_state_dict()`, which
raises `ValueError("Multiple weight files found for this model")` - not a `NotAMatchError` -
when a folder holds more than one weight file. That aborted this config's probe entirely, so
no candidate matched and the encoder was stored as `unknown`.
Make the SDNQ key check shard-safe: read tensor *names* from the safetensors headers per
shard instead of loading a state dict. That is cheap (no tensor data is materialized), works
for any number of shards, and keeps detecting SDNQ weight+scale pairs even when they are
split across shards. The mirrored fallback in `Qwen3Encoder_SDNQ_Folder_Config` had the same
crash for sharded SDNQ folders and now uses the same helper, with its file scope unchanged.
Verified against a real `FLUX.2-klein-9B::text_encoder+tokenizer` install, which now
identifies as `Qwen3Encoder_Qwen3Encoder_Config` / variant `qwen3_8b`.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…tion The SDNQ support added an "is this folder SDNQ-quantized?" guard to Qwen3Encoder_Qwen3Encoder_Config - a folder config that until then never touched weights - and implemented it with `mod.load_state_dict()`. That call refuses to pick a file when a directory holds more than one weight file and raises ValueError, which is not a NotAMatchError, so the factory records it as an error, no candidate matches and the model falls back to Unknown_Config. Every folder-layout Qwen3 encoder we ship as a starter model is sharded (FLUX.2 Klein 9B: 4, Klein 4B: 2, Z-Image text_encoder: 3), so all three install as "Unable to identify model". The SDNQ checks in identification now read tensor names from the safetensors headers via safetensors_tensor_names() / safetensors_have_sdnq_keys() in the sdnq detection module - the module the SDNQ PR created so this question has one implementation. Reading headers is per file, so it works for any number of shards, and it resolves weight/scale pairs across the union of all of them: sharding splits a checkpoint by tensor order and routinely separates a weight from its scale. Three further gaps closed while here: - Qwen3Encoder_Qwen3Encoder_Config and Qwen3Encoder_SDNQ_Folder_Config must partition folders but asked different questions: the unquantized one checked the marker in text_encoder/ too and keys across every shard, the SDNQ one only the root and only safetensors directly in it. A markerless SDNQ encoder in the nested layout was rejected by both. Both now call one shared predicate. - A corrupt quantization_config.json aborted the SDNQ probe with a JSONDecodeError instead of falling through to the key check. - The Qwen3-only q_norm/k_norm fallback read the state dict inside a bare except, so a sharded folder yielded no signal and was rejected. Single-file and GGUF configs keep using the state dict, which is correct for them. Verified against a real FLUX.2-klein-9B::text_encoder+tokenizer install: it now identifies as Qwen3Encoder_Qwen3Encoder_Config / variant qwen3_8b, and still does with ModelOnDisk.load_state_dict patched to raise. Closes invoke-ai#9567
lstein
left a comment
There was a problem hiding this comment.
Adversarial review at fc753bd.
The sharded-encoder fix itself is right, and the tests are load-bearing: all 7 new tests fail against main and pass here, the 115 tests in the SDNQ / Qwen3 identification suites are green, ruff is clean. Two things the shared predicate now lets through, though, both reproduced with on-disk fixtures against this branch and against main.
Blocker: the nested text_encoder/ SDNQ layout is now identified, but the loader cannot open it
The PR deliberately makes Qwen3Encoder_SDNQ_Folder_Config claim a folder whose weights live under text_encoder/ (test_nested_markerless_sdnq_encoder_is_claimed, test_marker_in_text_encoder_subfolder_is_honored). But Qwen3EncoderSDNQLoader._load_from_sdnq hands config.path (the root) straight to sdnq_sd_loader, which globs *.safetensors in that directory only and reads only the root quantization_config.json (invokeai/backend/quantization/sdnq/loaders.py:214-218). Nothing in the loader resolves text_encoder/, unlike the unquantized Qwen3EncoderLoader, which does (z_image.py:774-785).
Triggering sequence:
- Install
Disty0/Z-Image-Turbo-SDNQ-uint4-svd-r32::text_encoder+tokenizer(or build the PR's own_make_sdnq_encoderfixture):text_encoder/model-0000N-of-0000M.safetensors,text_encoder/quantization_config.json,tokenizer/. main:Unknown_Config, the user sees "unable to identify" at install time. This branch:Qwen3Encoder_SDNQ_Folder_Config,variant=qwen3_4b, selectable in the picker.- Generate →
ValueError: No safetensors files found in <root>fromsdnq_sd_loader.
I ran step 3 directly against the PR's fixture shape, with and without the marker in text_encoder/, and both raise. The rule identification already states for itself (the comment in _validate_is_qwen3_encoder, and Main_SDNQ_Diffusers_ZImage_Config._validate_has_sdnq_transformer at configs/main.py:3028-3034, which exists precisely because "the model installed and only failed when a loader opened the transformer path at generation time") is that it must not accept a folder its loader cannot open. Before this PR the nested layout was unknown-at-install; after it, it is installed-but-broken, which is the worse outcome.
Fix: give _load_from_sdnq the same layout resolution the unquantized loader has (weights and marker from text_encoder/ when that directory exists), with a test that runs sdnq_sd_loader path resolution against the nested fixture. If that is out of scope for this PR, the SDNQ config has to keep rejecting the nested layout so it stays "unknown" rather than "installed".
A related consequence of the same gap: a marker in text_encoder/ is honoured by identification but the loader parses the root one, so a real group_size / modules_dtype_dict would be ignored and defaulted. Same fix.
Should fix: full SDNQ pipeline bundles now match Qwen3Encoder_SDNQ_Folder_Config as well
_folder_is_sdnq_quantized looks for the marker in text_encoder/ and scans mod.weight_files(), which is rglob over the whole tree (model_on_disk.py:61-65), and Qwen3Encoder_SDNQ_Folder_Config.from_model_on_disk has no pipeline exclusion: the model_index.json / transformer guard in the unquantized config (qwen3_encoder.py:330-337) is not mirrored. For a Z-Image or FLUX.2 Klein SDNQ bundle (the three SDNQ starter models), all_matches is now [Main_SDNQ_Diffusers_ZImage_Config, Qwen3Encoder_SDNQ_Folder_Config]; on main it is the Main config alone. The factory's type sort hides this while the Main config matches. When it declines, the whole bundle registers as a Qwen3 encoder at the bundle root, which the SDNQ loader cannot open either:
- transformer folder with marker +
config.jsonbut no shards yet (interrupted download):main→Unknown_Config; this branch →Qwen3Encoder_SDNQ_Folder_Config. - corrupt
model_index.json: same.
It also breaks the property the PR description advertises, that exactly one config claims a folder. Fix: mirror the pipeline exclusion in Qwen3Encoder_SDNQ_Folder_Config.from_model_on_disk, and scope the header scan in _folder_is_sdnq_quantized / _folder_tensor_names to the root and text_encoder/ (the set the loaders actually read) instead of rglob. Once the loader fix above lands, "root + text_encoder/" is exactly that set.
Attacks that did not land
load_state_dicton the folder paths:test_folder_probe_never_loads_a_state_dictproves it for the unquantized config, and_validate_is_qwen3_encoder/_get_variant_from_dirin the SDNQ config never call it either.- Weight/scale pairs split across shards: resolved on the union, verified with the pair in different files.
- Test sensitivity: 7/7 new tests fail on
main(ValueError: Multiple weight files, and the four nested/markerless/corrupt-marker cases raiseNotAMatchError). - Single-file and GGUF Qwen3 paths: untouched, still on the state dict.
- Corrupt shard:
safetensors_tensor_namesswallows it per file, same behaviour as the pre-existingfolder_has_sdnq_keys, so no regression.
…can open Review follow-up on the sharded-encoder fix. Claiming the nested `text_encoder/` layout without teaching the loader about it traded "unknown at install" for "installed, then ValueError at generation time", and the shared SDNQ predicate scanned the whole tree, so a pipeline bundle answered yes to it too. Blocker: `Qwen3EncoderSDNQLoader._load_from_sdnq` handed `config.path` straight to `sdnq_sd_loader`, which globs one directory and reads the quantization_config.json beside it. For a `text_encoder+tokenizer` install that root holds no shards, so generation failed with "No safetensors files found", and a real group_size in text_encoder/quantization_config.json was ignored in favour of the default. The layout resolution now lives in one place, resolve_qwen3_encoder_dir(), exposed on the config as resolve_text_encoder_dir() the way T5Encoder_SDNQ_Config already does for its two layouts, and the loader calls it. Identification cannot accept a folder the loader is unable to open. Should fix: `_folder_is_sdnq_quantized` scanned `mod.weight_files()`, an rglob over the whole tree, and Qwen3Encoder_SDNQ_Folder_Config had no pipeline exclusion. An SDNQ bundle's transformer and VAE are SDNQ-quantized too, so the bundle root matched as a Qwen3 encoder alongside Main_SDNQ_Diffusers_*; the factory's type sort hid that until the Main config declined - an interrupted download, a corrupt model_index.json - and then the whole bundle registered as an encoder at a path the SDNQ loader cannot open. The scan is now scoped to the root and `text_encoder/`, which is exactly the set the loaders read, and the model_index.json / transformer exclusion is mirrored from the unquantized config. Six new tests, all failing before this commit: the loader opening the nested layout (and the root raising), standalone resolution, complete bundle, interrupted download, corrupt model_index.json, and SDNQ weights that live only in `vae/`.
Summary
fix(mm): stop the SDNQ guards from aborting folder-encoder identification
The SDNQ support in #9228 added an "is this folder SDNQ-quantized?" guard to
Qwen3Encoder_Qwen3Encoder_Config— a folder config that until then never touched weights — andimplemented it with
mod.load_state_dict(). That call refuses to pick a file when a directory holdsmore than one weight file and raises
ValueError, which is not aNotAMatchError, so thefactory records it as an error and the config never matches. With no candidate left, the model falls
back to
Unknown_Config.Every folder-layout Qwen3 encoder we ship as a starter model is sharded, so all three install as
"Unable to identify model":
black-forest-labs/FLUX.2-klein-9B::text_encoder+tokenizerblack-forest-labs/FLUX.2-klein-4B::text_encoder+tokenizerTongyi-MAI/Z-Image-Turbo::text_encoder+tokenizerWhat changed
The SDNQ checks in identification read tensor names from safetensors headers.
safetensors_tensor_names()/safetensors_have_sdnq_keys()live inbackend/quantization/sdnq/detection.py— the module the SDNQ PR created precisely so thisquestion has one implementation — and
folder_has_sdnq_keys()now delegates to them. Readingheaders is per file, so it works for any number of shards, and it resolves weight/scale pairs
across the union of all of them: sharding splits a checkpoint by tensor order and routinely
separates a weight from its scale.
The two Qwen3 folder configs now ask the same question.
Qwen3Encoder_Qwen3Encoder_Config(rejects SDNQ) and
Qwen3Encoder_SDNQ_Folder_Config(requires it) must partition folders, but theylooked at different places: the unquantized one checked the marker in
text_encoder/too and keysacross every shard, the SDNQ one only the root and only safetensors sitting directly in it. A
markerless SDNQ encoder in the nested
text_encoder/layout was therefore rejected by both —the exact shape that lands a model in
unknown. Both now call one shared predicate.A corrupt
quantization_config.jsonno longer aborts the probe.Qwen3Encoder_SDNQ_Folder_Configcalledjson.load()unguarded; it now falls through to the keycheck, like every other marker read.
The Qwen3-only
q_norm/k_normfallback works for sharded folders. It read the state dictinside a bare
except, so a sharded folder yielded no signal at all and was rejected. It readsheader names now.
No behaviour changes for single-file checkpoints or GGUFs — those paths still use the state dict,
which is correct for them.
Related Issues / Discussions
Closes #9567
QA Instructions
Verified against a real
black-forest-labs/FLUX.2-klein-9B::text_encoder+tokenizerinstall (4 shards)on Windows:
ModelConfigFactory.from_model_on_disk()→Unknown_Config, withQwen3Encoder_Qwen3Encoder_Config: ValueError: Multiple weight files found for this modelin thedetails.
Qwen3Encoder_Qwen3Encoder_Config,variant=qwen3_8b, and it is selectable for FLUX.2Klein 9B.
ModelOnDisk.load_state_dictpatched to raise, identification still succeeds.To reproduce in the UI: Model Manager → Starter Models → install "FLUX.2 Klein Qwen3 8B Encoder"
(or the 4B / Z-Image encoders). It should register as a Qwen3 encoder, not as Unknown.
New tests:
load_state_dict()is called at alltext_encoder/is honouredq_normkeysMerge Plan
Normal merge. Supersedes the narrower first pass on this branch.
Checklist
What's Newcopy (if doing a release after this PR)