fix(mm): identify sharded Qwen3 encoder folders instead of falling back to unknown

[Pfannkuchensack]

Summary

Two users reported that the Qwen3 encoder for FLUX.2 Klein 4B and 9B is registered as an unknown model after downloading it from the starter models list. This is a probing bug, and it is reproducible.

The starter-model source black-forest-labs/FLUX.2-klein-{4B,9B}::text_encoder+tokenizer lands the encoder on disk as several shards:

<model>/text_encoder/config.json
<model>/text_encoder/model-00001-of-00004.safetensors
<model>/text_encoder/model-00002-of-00004.safetensors
<model>/text_encoder/model-00003-of-00004.safetensors
<model>/text_encoder/model-00004-of-00004.safetensors
<model>/text_encoder/model.safetensors.index.json
<model>/tokenizer/...

Qwen3Encoder_Qwen3Encoder_Config._reject_if_sdnq_quantized() calls mod.load_state_dict() as its fallback signal. ModelOnDisk.resolve_weight_file() refuses to pick a file when a folder holds more than one weight file and raises a plain ValueError — not a NotAMatchError:

Qwen3Encoder_Qwen3Encoder_Config:
  ValueError("Multiple weight files found for this model: [.../model-00001-of-00004.safetensors, ...].
             Please specify the intended file using the 'path' argument")

Because that is not a NotAMatchError, it aborts this config's probe instead of being recorded as a non-match, no other candidate matches the folder, and identification falls through to Unknown_Config. Both the 4B and the 9B encoder are multi-shard, so both are affected.

This was introduced together with the SDNQ support in #9228 — before that, this config never touched the state dict.

Fix

The SDNQ fallback only needs tensor names, so read them from the safetensors headers per shard rather than loading a state dict:

• new _files_look_sdnq_quantized() helper — iterates the model's safetensors files, collects keys via safe_open(...).keys(), and looks for SDNQ weight + scale pairs across the union of all shards. No tensor data is materialized, so this is cheap even for a 16 GB encoder, and it works for any number of shards.
• _has_sdnq_keys() now delegates to a shared key predicate so the state-dict and file-based paths cannot drift apart.
• Qwen3Encoder_SDNQ_Folder_Config's mirrored fallback had the identical crash for sharded SDNQ folders and now uses the same helper. Its file scope is deliberately unchanged (root-level *.safetensors only), so this is not a widening of SDNQ detection.

Verification

Against a real black-forest-labs/FLUX.2-klein-9B::text_encoder+tokenizer install:

Before:

Unknown_Config (base=unknown, type=unknown, format=unknown)

After:

Qwen3Encoder_Qwen3Encoder_Config | type=qwen3_encoder | format=qwen3_encoder | variant=qwen3_8b

Tests

New TestShardedQwen3EncoderFolder in tests/backend/model_manager/configs/test_qwen3_encoder_config.py:

• a sharded text_encoder/ folder (plus a sibling tokenizer/ folder, matching the real download layout) identifies as a Qwen3 encoder with the correct variant — this fails on main with the ValueError above.
• a sharded SDNQ encoder is still rejected, with the weight and scale tensors deliberately placed in different shards to prove the union-of-headers check works.

tests/backend/model_manager/configs/test_qwen3_encoder_config.py and test_qwen3_encoder_sdnq_folder_identification.py pass (24 tests), ruff clean.

Related

Two other folder-based configs call mod.load_state_dict() without an explicit path and would hit the same ValueError on sharded folders — IPAdapter_InvokeAI_Config_Base and T5Encoder_BnBLLMint8_Config. Left untouched here to keep this PR scoped to the reported regression.

🤖 Generated with Claude Code