feat: Adding Support for SD.Next Quantization Engine (SDNQ) (Flux1&Flux2klein4B/9B&Z-Image)#9228
feat: Adding Support for SD.Next Quantization Engine (SDNQ) (Flux1&Flux2klein4B/9B&Z-Image)#9228Pfannkuchensack wants to merge 41 commits into
Conversation
Add support for loading SDNQ-quantized models with on-the-fly CPU dequantization, similar to existing GGUF support. New features: - SDNQTensor class with __torch_dispatch__ for automatic dequantization - Support for symmetric/asymmetric int8/uint8/fp8 quantization - Optional SVD correction (low-rank approximation) - Model loaders for Flux and Z-Image SDNQ models - Automatic format detection via weight+scale key pairs New files: - invokeai/backend/quantization/sdnq/ (core module) - tests/backend/quantization/sdnq/ (unit tests) Modified files: - taxonomy.py: Add ModelFormat.SDNQQuantized - configs/main.py: Add Main_SDNQ_FLUX_Config, Main_SDNQ_ZImage_Config - configs/factory.py: Register SDNQ configs - model_loaders/flux.py: Add FluxSDNQCheckpointModel - model_loaders/z_image.py: Add ZImageSDNQCheckpointModel
- Add uint4 per-group quantization with packed weight unpacking - Handle 1D flattened weights (reshape to 2D before unpacking) - Support SDNQ diffusers format for FLUX transformer and T5 - Add SDNQ VAE loading with AutoencoderKL - Add diagnostic logging for debugging dequantization - Fix bit order in uint4 unpacking (lower, upper)
…tion The test was checking `(weight - zero_point) * scale`, but SDNQ (Disty0/sdnq) defines asymmetric dequantization as `zero_point + weight * scale` (via torch.addcmul), where zero_point is a post-scale bias rather than a pre-scale integer offset. The implementation already follows this convention; only the test expectation was wrong.
…tion The test was checking `(weight - zero_point) * scale`, but SDNQ (Disty0/sdnq) defines asymmetric dequantization as `zero_point + weight * scale` (via torch.addcmul), where zero_point is a post-scale bias rather than a pre-scale integer offset. The implementation already follows this convention; only the test expectation was wrong. feat(sdnq): support sidecar LoRA application on SDNQ-quantized layers Bring SDNQ to feature parity with GGUF in the sidecar patching path so LoRA, LoKr, DoRA, FullLayer, and FluxControl patches apply correctly to SDNQ-quantized Linear and Conv2d modules. Without this, the sidecar aggregate replaced the SDNQTensor weight with a meta tensor and patches silently produced wrong results. - Add SDNQTensor branch in CustomModuleMixin._aggregate_patch_parameters mirroring the GGMLTensor branch. - Extend the (GGMLTensor) dtype-cast exclusion to also cover SDNQTensor in CustomLinear, CustomConv2d, CustomInvokeLinearNF4, and CustomInvokeLinear8bitLt. - Add `linear_with_sdnq_quantized_tensor` and `linear_sdnq_quantized` fixtures so the existing custom-module test matrix exercises SDNQ alongside GGUF, BnB-8bit, and NF4.
Add T5Encoder_SDNQ_Config for diffusers-style T5 bundles whose text_encoder_2/ folder holds SDNQ-quantized safetensors (detected via quantization_config.json's quant_method or via the SDNQ-style weight+scale key pairs). Add T5EncoderSDNQLoader that materializes the T5EncoderModel on meta, then loads the SDNQ state dict, and re-shares the embed_tokens/shared weight per HuggingFace's tied- weight convention.
Add Main_SDNQ_Flux2_Config covering Klein 4B/9B and their Base variants (detected via _get_flux2_variant on the dequantized SDNQTensor shapes plus the existing filename heuristic), and Flux2SDNQCheckpointModel that loads diffusers-layout SDNQ FLUX.2 checkpoints straight into Flux2Transformer2DModel. Architecture (num_layers, hidden_size, attention head count, guidance presence) is detected from state-dict shapes the same way the fp16 loader does, since SDNQTensor.shape reports the dequantized shape. BFL-layout SDNQ FLUX.2 checkpoints are not supported here — that would require an SDNQTensor-aware port of the _convert_flux2_bfl_to_diffusers fuse logic.
Add Main_SDNQ_Diffusers_ZImage_Config so a complete SDNQ ZImagePipeline folder (model_index.json + transformer/ + text_encoder/ + tokenizer/ + vae/) is recognised on install and its submodels are wired up. Extend ZImageSDNQCheckpointModel to load the transformer from the subfolder using ZImageTransformer2DModel.from_config() so non-default architecture parameters (e.g. axes_lens [1536,512,512] in newer Z-Image Turbo SDNQ exports) are honoured instead of the single-file path's hardcoded [1024,512,512]. Verified end-to-end against Tongyi-MAI/Z-Image-Turbo-SDNQ-uint4-svd-r32: 269 quantized + 252 regular tensors load into a 6.15B-param model with 0 missing / 0 unexpected keys.
T5Encoder_SDNQ_Config originally only looked for text_encoder_2/ as a subfolder of mod.path, which works for standalone T5 bundles but misses the case where a parent FluxPipeline / similar config registers its T5 submodel with path_or_prefix pointing straight at the text_encoder_2 folder. Allow both layouts in both the config's detection logic and T5EncoderSDNQLoader's te_dir resolution. Verified end-to-end with Disty0/FLUX.1-schnell-SDNQ-uint4-svd-r32.
The diffusers→BFL state-dict converter renamed norm_out.linear.{weight,bias}
to final_layer.adaLN_modulation.1.{weight,bias} but did not swap the
two halves along dim 0. diffusers' AdaLayerNormContinuous packs the
linear output as (scale, shift); BFL's LastLayer packs as (shift, scale).
Without the swap, the final adaLN modulation runs with scale and shift
permuted, which produces structured-but-very-noisy output for every
pixel. Reuse the same pattern the FLUX.2 converter applies for the
analogous adaLN_modulation key.
ZImageSDNQCheckpointModel only handled the Transformer submodel, so attempts to use an SDNQ ZImagePipeline as the "Qwen3 & VAE source model" (which triggers loads for TextEncoder / Tokenizer / VAE) crashed with "Only Transformer submodels are currently supported". Add per-submodel handlers that load text_encoder/ via sdnq_sd_loader into an empty Qwen3ForCausalLM (re-sharing lm_head with embed_tokens when tied), tokenizer/ via AutoTokenizer, and vae/ via AutoencoderKL.from_pretrained. The single-file SDNQ checkpoint path keeps its transformer-only behaviour but now raises a clearer error when asked for a different submodel.
Add support for SDNQ-quantized Flux2KleinPipeline folders, which mix uint4 and int5 dtypes across layers (chosen dynamically by SDNQ during quantization to stay under a per-group loss budget). Core changes: - Add INT5_ASYM quantization type + unpack_uint5 + dequantize_int5_per_group. Sign-extension matches Disty0/sdnq's unpack_int convention (raw 0..31 - 16). zero_point is optional (dynamic-mixed sometimes emits scale-only int5 tensors). - _infer_quantization_type now takes a per_tensor_dtype override; the loader builds an inverted map from quantization_config.json's modules_dtype_dict. - _get_original_shape uses the packed weight size as the authoritative source for in_features, fixing a bug where Klein 4B's group_size=64 layers were misread as group_size=128 (the previous fallback). Pipeline integration: - Add Main_SDNQ_Diffusers_Flux2_Config matching Flux2Pipeline / Flux2KleinPipeline folders with quantized transformer. - Flux2SDNQCheckpointModel now dispatches all pipeline submodels: transformer (Flux2Transformer2DModel.from_config + sdnq state dict), text_encoder (Qwen3ForCausalLM SDNQ + lm_head/embed_tokens tie), tokenizer (AutoTokenizer), vae (AutoencoderKLFlux2 / AutoencoderKL). - Extend flux2_klein_model_loader._validate_diffusers_format and the isFlux2DiffusersMainModelConfig FE filter to also accept SDNQ pipeline configs (when submodels is populated). Verified against Disty0/FLUX.2-klein-4B-SDNQ-4bit-dynamic: 98 uint4 + 2 int5 tensors load into a 3.88B-param Flux2Transformer2DModel with 0 missing / 0 unexpected keys; both dequant paths produce reasonable zero-centred weight distributions.
Main_Diffusers_Flux2_Config so identification routes them to the SDNQ configs instead. Without this both configs accept the folder and the plain diffusers loader wins, then crashes when reading packed uint8 weights as bf16.
diffusion_pytorch_model-{00001,00002}-of-00002.safetensors and FLUX.2
dev's sharded transformer both load. Detect cross-shard key collisions
as a corruption signal.
"main_is_diffusers" in z_image_model_loader and flux2_klein_model_loader so the auto-extract-submodels branch handles them. Without this the loader demanded a separate VAE/Qwen3 source even though the SDNQ pipeline carries those submodels itself. - Drop the ui_model_format=Diffusers hint on Klein's qwen3_source_model field so the FE combobox can also show SDNQ pipeline configs (the FE filter already accepts them).
Loading the Klein 4B SDNQ pipeline as the main model errored with "No Qwen3 Encoder selected" in the UI even though the pipeline carries its own Qwen3 + VAE submodels, and the Model Manager showed no format badge at all on SDNQ models. - flux2_klein_model_loader now treats SDNQ-with-submodels as main_is_diffusers, so the auto-extract-submodels branch handles SDNQ pipelines exactly like plain diffusers. Drop the ui_model_format=Diffusers hint on qwen3_source_model so the combobox can also show SDNQ pipeline configs. - readiness.ts no longer demands a standalone VAE/Qwen3 for FLUX.2 Klein when the main model is itself a pipeline (diffusers or SDNQ-with-submodels). Without this the Invoke button stayed disabled with "Non-diffusers FLUX.2 Klein models require a standalone Qwen3 Encoder" even when the SDNQ pipeline could self-source everything. - Register sdnq_quantized in zModelFormat, the manually-edited OpenAPI schema, ModelFormatBadge, and MODEL_FORMAT_TO_LONG_NAME so SDNQ models render an "sdnq" badge instead of an empty placeholder.
- 4 new starter models covering all SDNQ pipelines verified end-to-end in this branch: FLUX.1 schnell, Z-Image Turbo, FLUX.2 Klein 4B (dynamic mixed), FLUX.2 Klein 9B (dynamic mixed + SVD). Each entry is self-contained (no separate encoder/VAE dependencies because the SDNQ pipeline folder bundles them). - New /configuration/sdnq-quantization/ page: support matrix, VRAM footprints, install steps (Starter Models + HF + Folder), LoRA compatibility notes, SDNQ-vs-SVDQuant/Nunchaku disambiguation, comparison with GGUF/NF4/FP8, troubleshooting. - Cross-link from fp8-storage.mdx's "no-op on quantized" caution.
Pfannkuchensack
requested review from
JPPhoto,
blessedcoolant,
dunkeroni and
lstein
as code owners
github-actions
Bot
added
python
Z-Image and Qwen3 SDNQ configs were missing `variant` (and `cpu_only` on Qwen3) fields that exist on the other variants of the same union, breaking TypeScript narrowing on the FE. - Main_SDNQ_ZImage_Config: add variant (default Turbo) - Main_SDNQ_Diffusers_ZImage_Config: add variant, detect from scheduler_config.json shift value - Qwen3Encoder_SDNQ_Config: add cpu_only + variant, detect from embed_tokens shape - Qwen3Encoder_SDNQ_Folder_Config: add cpu_only + variant, detect from config.json hidden_size - Regenerate FE schema.ts Discriminator tags are unchanged since variant has no default.
Resolve conflicts in qwen3_encoder.py (keep both SDNQ tensor/key detection and the Qwen-VL visual-tower guard) and regenerate openapi.json + schema.ts from the merged backend. Fix T5TokenizerFast under transformers 5.x: it is now an alias of T5Tokenizer. Use T5Tokenizer in flux.py (was used unimported -> F821) and also match the new "T5Tokenizer" class name in the SDNQ FLUX model_index.json submodel probe.
JPPhoto
left a comment
JPPhoto
left a comment
There was a problem hiding this comment.
-
invokeai/backend/quantization/sdnq/loaders.py:257: SDNQ dequantization can use the wrong group size for mixed-group tensors. The loader trustsquantization_config.jsongroup_sizewhenever it is greater than zero, and only infers fromscale.shapewhen the config says zero. But the same file notes that Klein 4B has layers whose actual group size differs from the nominal config value. When such a tensor reachesdequantize_uint4_per_group()ordequantize_int5_per_group(),num_groups = in_features // group_sizeno longer matches the scale tensor's group dimension, so the reshape/broadcast path fails or dequantizes incorrectly. To expose this issue, add a test that loads a uint4 SDNQ tensor with globalgroup_size=128butscale.shape[1]corresponding to 64-wide groups, then callsget_dequantized_tensor(). -
invokeai/app/invocations/z_image_model_loader.py:90: SDNQ Z-Image pipeline folders are not treated as self-contained by the invocation/graph path. The backend addsMain_SDNQ_Diffusers_ZImage_Configwithsubmodels, and the docs describe SDNQ pipeline installs as one install containing transformer, encoders, and VAE, butZImageModelLoaderInvocationstill only uses an explicitly selectedvae_modelorqwen3_source_model; it never falls back to the main model when the selected main config is an SDNQ pipeline with submodels. The frontend repeats that requirement ininvokeai/frontend/web/src/features/queue/store/readiness.ts:322andinvokeai/frontend/web/src/features/nodes/util/graph/generation/buildZImageGraph.ts:52, so a freshly installed SDNQ Z-Image pipeline remains blocked unless the user manually selects a component source, possibly the same model again. To expose this issue, add a graph/readiness test for az-imagemain model withformat: sdnq_quantizedand non-emptysubmodels, with no separate VAE/Qwen3 selections, and assert generation is allowed and the model loader uses the main model for those submodels. -
invokeai/backend/model_manager/configs/factory.py:210: Single-file SDNQ FLUX.2 identification is fragile and can route to the FLUX.1 config first. The union checksMain_SDNQ_FLUX_ConfigbeforeMain_SDNQ_Flux2_Config, despite the nearby comment requiring FLUX.2 before FLUX.1 for ambiguous transformer checkpoints.Main_SDNQ_FLUX_Configonly checks for generic main keys and SDNQ keys before deriving a FLUX.1 variant, so a prefixed FLUX.2 SDNQ state dict can be accepted asbase=fluxand later loaded by the FLUX.1 loader. The alternate bare diffusers-style single-file path is also at risk becauseMain_SDNQ_Flux2_Configcalls_has_main_keys(), whose accepted prefixes do not include baretransformer_blocks.orcontext_embedder.even though the FLUX.2 single-file loader later looks for those keys. To expose this issue, add model-identification tests for SDNQ FLUX.2 single-file state dicts in both prefixed and bare diffusers layouts and assert they identify asMain_SDNQ_Flux2_Config, not FLUX.1 or unknown.
…LUX.2 single-file identification Address three SDNQ review findings (PR invoke-ai#9228): - loaders.py: derive per-tensor group_size from the scale tensor's group dimension instead of trusting quantization_config.json. Dynamic-mixed models (e.g. FLUX.2 Klein 4B) quantize some layers with a group size that differs from the nominal config value; trusting the config produced a num_groups that disagreed with scale.shape[1] and broke the reshape/ broadcast in dequantize_uint4/int5_per_group. - z_image_model_loader: fall back to the main model for the VAE/Qwen3 submodels when it is a self-contained SDNQ pipeline (format=sdnq_quantized with submodels). Mirror the same self-contained handling in the frontend readiness checks and buildZImageGraph, so a freshly installed SDNQ Z-Image pipeline generates without manually selecting a component source. - FLUX.2 single-file identification: list SDNQ FLUX.2 before FLUX.1 in the config union, reject FLUX.2 state dicts in Main_SDNQ_FLUX_Config so the two stay mutually exclusive regardless of iteration order, and accept the bare diffusers transformer layout (transformer_blocks./context_embedder.) in Main_SDNQ_Flux2_Config so the real single-file SDNQ FLUX.2 checkpoint is identified instead of falling through to FLUX.1 or unknown. Add regression tests for all three: mixed-group uint4 dequant, the Z-Image loader/readiness self-contained fallback, and SDNQ FLUX.2 identification.
|
…LUX.2/Z-Image identification Address five SDNQ review findings (PR invoke-ai#9228): - loaders.py: derive per-tensor group_size from the scale tensor's group dimension instead of trusting quantization_config.json. Dynamic-mixed models (e.g. FLUX.2 Klein 4B) quantize some layers with a group size that differs from the nominal config value; trusting the config produced a num_groups that disagreed with scale.shape[1] and broke the reshape/ broadcast in dequantize_uint4/int5_per_group. - z_image_model_loader: fall back to the main model for the VAE/Qwen3 submodels when it is a self-contained SDNQ pipeline (format=sdnq_quantized with submodels). Mirror the same self-contained handling in the frontend readiness checks and buildZImageGraph, so a freshly installed SDNQ Z-Image pipeline generates without manually selecting a component source. - FLUX.2 single-file identification: list SDNQ FLUX.2 before FLUX.1 in the config union, reject FLUX.2 state dicts in Main_SDNQ_FLUX_Config so the two stay mutually exclusive regardless of iteration order, and accept the bare diffusers transformer layout (transformer_blocks./context_embedder.) in Main_SDNQ_Flux2_Config so the real single-file SDNQ FLUX.2 checkpoint is identified instead of falling through to FLUX.1 or unknown. - Main_Diffusers_ZImage_Config: reject SDNQ-quantized ZImagePipeline folders so Main_SDNQ_Diffusers_ZImage_Config matches them instead (mirrors the FLUX and FLUX.2 guards). Without this a full SDNQ Z-Image pipeline could be classified as plain diffusers, mis-reading packed weights and breaking the self-contained path. - z_image loader: check missing/unexpected keys after load_state_dict in the SDNQ transformer folder path and fail fast with the offending key list, instead of silently returning a model with required params on meta tensors that fails later during device movement or inference. Add regression tests for all five: mixed-group uint4 dequant, the Z-Image loader/readiness self-contained fallback, SDNQ FLUX.2 identification, and SDNQ ZImagePipeline-folder identification.
…chensack/InvokeAI into feature/svd-quantization
There was a problem hiding this comment.
-
invokeai/app/invocations/z_image_model_loader.py:74:qwen3_source_modelstill declaresui_model_format=ModelFormat.Diffusers, while_validate_diffusers_format()now accepts SDNQ pipeline configs with submodels atinvokeai/app/invocations/z_image_model_loader.py:159. Generic node and workflow model pickers enforce that schema hint, filtering out configs whoseconfig.formatis not listed atinvokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/fields/inputs/ModelIdentifierFieldInputComponent.tsx:58andinvokeai/frontend/web/src/features/controlLayers/components/CanvasWorkflowIntegration/WorkflowFieldRenderer.tsx:257. A user building a Z-Image graph manually, or filling a workflow field, cannot select an SDNQ Z-Image pipeline as the component source even though the backend now supports it. This was fixed for FLUX.2 by removing the format hint, but not for Z-Image. To expose this issue, add a test that builds thez_image_model_loader.qwen3_source_modelfield template and asserts an SDNQz-imagemain model withsubmodelsis not filtered out. -
invokeai/backend/model_manager/configs/qwen3_encoder.py:388:Qwen3Encoder_SDNQ_Folder_Config.from_model_on_disk()accepts any directory withquantization_config.jsonwherequant_method == "sdnq"and never verifies the folder is actually a Qwen3 encoder. It does not checkconfig.jsonclass names or Qwen3 state-dict keys before returning a Qwen3 config, and the fallback only checks generic SDNQweight+scalekeys. That means importing a standalone SDNQ transformer, VAE, or other SDNQ component folder can be stored astype=qwen3_encoder; the loader then fails later atinvokeai/backend/model_manager/load/model_loaders/z_image.py:1527orinvokeai/backend/model_manager/load/model_loaders/z_image.py:1541when Qwen-specific weights are missing. To expose this issue, add a model-identification test that pointsModelConfigFactory.from_model_on_disk()at an SDNQtransformer/orvae/folder withquant_method: "sdnq"and asserts it does not resolve toQwen3Encoder_SDNQ_Folder_Config. -
invokeai/backend/model_manager/load/model_loaders/flux.py:1229: the FLUX.2 SDNQ pipeline transformer is created underaccelerate.init_empty_weights()and loaded withmodel.load_state_dict(sd, assign=True, strict=False)atinvokeai/backend/model_manager/load/model_loaders/flux.py:1235, but the missing and unexpected key lists are ignored. If an SDNQ export is partial, mismatched, or missing a shard key, required parameters remain on the meta device and the model is returned anyway, failing later during device movement or inference. The same unchecked pattern exists in newly added SDNQ component loaders atinvokeai/backend/model_manager/load/model_loaders/flux.py:2022,invokeai/backend/model_manager/load/model_loaders/flux.py:2047,invokeai/backend/model_manager/load/model_loaders/flux.py:2077, andinvokeai/backend/model_manager/load/model_loaders/vae.py:128. To expose this issue, add a loader test that removes one required SDNQ transformer weight from a folder load and asserts the loader raises with the missing key instead of returning a model with meta parameters.
Address the SDNQ review findings on the feature branch: Dequantization / loading correctness - sdnq/loaders.py: derive per-tensor group_size from the scale tensor's group dimension instead of trusting quantization_config.json. Dynamic-mixed models (e.g. FLUX.2 Klein 4B) quantize some layers with a group size that differs from the nominal config value; trusting the config produced a num_groups that disagreed with scale.shape[1] and broke the reshape/broadcast in dequantize_uint4/int5_per_group. - Add shared raise_on_incomplete_sdnq_load() and apply it to every SDNQ folder loader that used load_state_dict(strict=False) without checking the result: FLUX.2 pipeline transformer, SDNQ CLIP/T5/VAE (flux.py), standalone SDNQ VAE (vae.py), and both Z-Image Qwen3 text-encoder paths (z_image.py). A partial or mismatched export now fails fast with the offending key list instead of returning a model with parameters left on the meta device. Tied/re-shared weights (T5 encoder.embed_tokens, Qwen3 lm_head, CLIP position_ids) are allow-listed. Model identification - Reject FLUX.2 in Main_SDNQ_FLUX_Config and accept the bare diffusers transformer layout (transformer_blocks./context_embedder.) in Main_SDNQ_Flux2_Config, and list FLUX.2 before FLUX.1 in the config union, so a single-file SDNQ FLUX.2 checkpoint identifies as FLUX.2 instead of falling through to FLUX.1 or unknown. - Main_Diffusers_ZImage_Config: reject SDNQ-quantized ZImagePipeline folders so Main_SDNQ_Diffusers_ZImage_Config matches them instead (mirrors the FLUX and FLUX.2 guards). - Qwen3Encoder_SDNQ_Folder_Config: verify the folder is actually a Qwen3 encoder (Qwen3 config class name or Qwen3 state-dict keys) before claiming it, so an SDNQ transformer/VAE folder is not stored as type=qwen3_encoder. - Qwen3Encoder_Qwen3Encoder_Config: reject SDNQ folders so it no longer competes with the SDNQ folder config (both share the Qwen3Encoder type, making the factory tiebreak non-deterministic). Z-Image self-contained SDNQ pipelines - z_image_model_loader: fall back to the main model for the VAE/Qwen3 submodels when it is a self-contained SDNQ pipeline (format=sdnq_quantized with submodels), and drop the ui_model_format=diffusers hint from qwen3_source_model so SDNQ pipelines are not filtered out of the node/workflow model pickers. - Mirror the self-contained handling in the frontend readiness checks and buildZImageGraph so a freshly installed SDNQ Z-Image pipeline generates without a manually selected component source. Add regression tests for every fix: mixed-group uint4 dequant, the load-guard helper, SDNQ FLUX.2 identification, SDNQ ZImagePipeline-folder identification, Qwen3 encoder folder identification, the Z-Image self-contained loader/readiness path, and the qwen3_source_model field template.
JPPhoto
left a comment
JPPhoto
left a comment
There was a problem hiding this comment.
-
invokeai/backend/model_manager/configs/qwen3_encoder.py:452:Qwen3Encoder_SDNQ_Folder_Configlacks the complete causal LM guard that exists inQwen3Encoder_Qwen3Encoder_Configatinvokeai/backend/model_manager/configs/qwen3_encoder.py:227. An SDNQ root folder withconfig.json, tokenizer files, andQwen3ForCausalLMwill be accepted asqwen3_encoderinstead of being left to the TextLLM path, bypassing the disambiguation already documented intests/backend/model_manager/configs/test_qwen3_encoder_config.py. -
invokeai/app/invocations/z_image_model_loader.py:150: self-contained SDNQ Z-Image detection treats any truthysubmodelsdict as sufficient.Main_SDNQ_Diffusers_ZImage_Configonly requires an SDNQ transformer and then builds whatever submodels it recognizes atinvokeai/backend/model_manager/configs/main.py:1836andinvokeai/backend/model_manager/configs/main.py:1896, so a partial or partially recognized pipeline can bypass VAE/Qwen3 readiness and fail later when the loader tries fixedvae,text_encoder, ortokenizerfolders. -
invokeai/frontend/web/src/features/nodes/util/graph/generation/buildZImageGraph.ts:55: the same truthy-submodelscheck is duplicated in frontend graph validation, andinvokeai/frontend/web/src/features/queue/store/readiness.ts:327/invokeai/frontend/web/src/features/queue/store/readiness.ts:796use it to suppress missing VAE/Qwen3 source errors. This should require the specific submodels needed for the workflow, not just any submodel entry. -
invokeai/backend/model_manager/load/model_loaders/z_image.py:1589: the standaloneQwen3EncoderSDNQLoaderstill ignores themissing_keysandunexpected_keysreturned byload_state_dict(strict=False). The later meta-parameter guard catches some missing required parameters, but unexpected keys from an incompatible or contaminated SDNQ export are silently accepted, unlike the other SDNQ load paths that now callraise_on_incomplete_sdnq_load. -
invokeai/backend/model_manager/load/model_loaders/flux.py:558:T5EncoderSDNQLoaderstill relies onassertto enforce missing/unexpected key checks for standalone SDNQ T5 loads. Running with optimized Python strips those assertions, which reintroduces silent partial-load behavior for this path instead of using the new explicit helper.
Address the SDNQ review findings on the feature branch: Dequantization / loading correctness - sdnq/loaders.py: derive per-tensor group_size from the scale tensor's group dimension instead of trusting quantization_config.json. Dynamic-mixed models (e.g. FLUX.2 Klein 4B) quantize some layers with a group size that differs from the nominal config value; trusting the config produced a num_groups that disagreed with scale.shape[1] and broke the reshape/broadcast in dequantize_uint4/int5_per_group. - Add shared raise_on_incomplete_sdnq_load() and use it in every SDNQ folder loader that called load_state_dict(strict=False) without checking the result: FLUX.2 pipeline transformer, SDNQ CLIP/T5/VAE (flux.py), standalone SDNQ VAE (vae.py), and both Z-Image Qwen3 text-encoder paths (z_image.py). A partial or mismatched export now fails fast with the offending key list instead of returning a model with parameters left on the meta device. Tied/re-shared weights (T5 encoder.embed_tokens, Qwen3 lm_head, CLIP position_ids) are allow-listed. Replace the T5 SDNQ assert checks (stripped by python -O) and add the missing unexpected-key check to the standalone Qwen3 SDNQ loader. Model identification - Reject FLUX.2 in Main_SDNQ_FLUX_Config and accept the bare diffusers transformer layout (transformer_blocks./context_embedder.) in Main_SDNQ_Flux2_Config, and list FLUX.2 before FLUX.1 in the config union, so a single-file SDNQ FLUX.2 checkpoint identifies as FLUX.2 instead of falling through to FLUX.1 or unknown. - Main_Diffusers_ZImage_Config: reject SDNQ-quantized ZImagePipeline folders so Main_SDNQ_Diffusers_ZImage_Config matches them instead (mirrors the FLUX/FLUX.2 guards). - Qwen3Encoder_SDNQ_Folder_Config: verify the folder is actually a Qwen3 encoder (Qwen3 config class or Qwen3 state-dict keys) and reject complete causal LMs (root config.json + tokenizer files) so an SDNQ transformer/VAE/TextLLM folder is not stored as type=qwen3_encoder. - Qwen3Encoder_Qwen3Encoder_Config: reject SDNQ folders so it no longer competes with the SDNQ folder config (both share the Qwen3Encoder type, making the factory tiebreak non-deterministic). Z-Image self-contained SDNQ pipelines - z_image_model_loader: fall back to the main model for the VAE/Qwen3 submodels when it is a self-contained SDNQ pipeline, but only when the pipeline actually exposes the required vae + text_encoder + tokenizer submodels (a truthy submodels dict is not enough — a partial pipeline would fail later on missing folders). Drop the ui_model_format=diffusers hint from qwen3_source_model so SDNQ pipelines are not filtered out of the node/workflow model pickers. - Mirror both the self-contained handling and the specific-submodels requirement in the frontend readiness checks and buildZImageGraph, so a freshly installed SDNQ Z-Image pipeline generates without a manually selected component source. Add regression tests for every fix: mixed-group uint4 dequant, the load-guard helper, SDNQ FLUX.2 identification, SDNQ ZImagePipeline-folder identification, Qwen3 encoder folder identification (SDNQ rejection, causal-LM rejection), the Z-Image self-contained loader/readiness path incl. partial-pipeline handling, and the qwen3_source_model field template.
|
@Pfannkuchensack I think we're getting closer. Here are some more findings and suggested testing approaches:
|
…ine support Address the SDNQ review findings on the feature branch: Dequantization / loading correctness - sdnq/loaders.py: derive per-tensor group_size from the scale tensor's group dimension instead of trusting quantization_config.json. Dynamic-mixed models (e.g. FLUX.2 Klein 4B) quantize some layers with a group size that differs from the nominal config value; trusting the config produced a num_groups that disagreed with scale.shape[1] and broke the reshape/broadcast in dequantize_uint4/int5_per_group. - Add shared raise_on_incomplete_sdnq_load() and use it in every SDNQ folder loader that called load_state_dict(strict=False) without checking the result: FLUX.2 pipeline transformer, SDNQ CLIP/T5/VAE (flux.py), standalone SDNQ VAE (vae.py), and both Z-Image Qwen3 text-encoder paths (z_image.py). A partial or mismatched export now fails fast with the offending key list instead of returning a model with parameters left on the meta device. Tied/re-shared weights (T5 encoder.embed_tokens, Qwen3 lm_head, CLIP position_ids) are allow-listed. Replace the T5 SDNQ assert checks (stripped by python -O) and add the missing unexpected-key check to the standalone Qwen3 SDNQ loader. Model identification - Reject FLUX.2 in Main_SDNQ_FLUX_Config and accept the bare diffusers transformer layout (transformer_blocks./context_embedder.) in Main_SDNQ_Flux2_Config, and list FLUX.2 before FLUX.1 in the config union, so a single-file SDNQ FLUX.2 checkpoint identifies as FLUX.2 instead of falling through to FLUX.1 or unknown. - Main_Diffusers_ZImage_Config: reject SDNQ-quantized ZImagePipeline folders so Main_SDNQ_Diffusers_ZImage_Config matches them instead (mirrors the FLUX/FLUX.2 guards). - Qwen3Encoder_SDNQ_Folder_Config: verify the folder is actually a Qwen3 encoder, reject complete causal LMs (root config.json + tokenizer files), accept the same compatible Qwen architectures as the unquantized config (Qwen2VLForConditionalGeneration / Qwen2ForCausalLM / Qwen3ForCausalLM via a shared constant), and make the state-dict fallback resilient to sharded folders, so an SDNQ transformer/VAE/TextLLM folder is not stored as type=qwen3_encoder and a compatible sharded encoder is not wrongly rejected. - Qwen3Encoder_Qwen3Encoder_Config: reject SDNQ folders so it no longer competes with the SDNQ folder config (both share the Qwen3Encoder type, making the factory tiebreak non-deterministic). Z-Image / FLUX.2 self-contained SDNQ pipelines - Add a shared is_self_contained_sdnq_pipeline() helper requiring the specific VAE + Qwen3 (text_encoder + tokenizer) submodels a pipeline install must ship. A truthy submodels dict is not enough: Main_SDNQ_Diffusers_* configs record whichever submodels they recognize, so a partial pipeline can expose only the transformer and would fail at runtime on missing fixed subfolders. - z_image_model_loader and flux2_klein_model_loader: use the helper for both the main-model self-contained fallback and the explicit qwen3_source validator, so a partial pipeline (as main model or selected source) requires an explicit VAE/Qwen3 source instead of emitting requests against missing folders. - Drop the ui_model_format=diffusers hint from z_image qwen3_source_model so SDNQ pipelines are not filtered out of the node/workflow model pickers. - Mirror the same specific-submodels requirement in the frontend readiness checks (Z-Image + FLUX.2, generate and canvas tabs) and buildZImageGraph, and bring the FLUX.2 canvas readiness block to parity with the generate tab (it previously used format !== 'diffusers' and never recognized SDNQ pipelines). Add regression tests for every fix: mixed-group uint4 dequant, the load-guard helper, SDNQ FLUX.2 identification, SDNQ ZImagePipeline-folder identification, Qwen3 encoder folder identification (SDNQ rejection, causal-LM rejection, compatible Qwen2 class names, sharded folders), the Z-Image and FLUX.2 self-contained loader/readiness paths incl. partial-pipeline handling, and the qwen3_source_model field template.
|
@Pfannkuchensack I see one more issue (but reserve the right to find more):
|
…ine support Address the SDNQ review findings on the feature branch: Dequantization / loading correctness - sdnq/loaders.py: derive per-tensor group_size from the scale tensor's group dimension instead of trusting quantization_config.json. Dynamic-mixed models (e.g. FLUX.2 Klein 4B) quantize some layers with a group size that differs from the nominal config value; trusting the config produced a num_groups that disagreed with scale.shape[1] and broke the reshape/broadcast in dequantize_uint4/int5_per_group. - Add shared raise_on_incomplete_sdnq_load() and use it in every SDNQ folder loader that called load_state_dict(strict=False) without checking the result: FLUX.2 pipeline transformer, SDNQ CLIP/T5/VAE (flux.py), standalone SDNQ VAE (vae.py), and both Z-Image Qwen3 text-encoder paths (z_image.py). A partial or mismatched export now fails fast with the offending key list instead of returning a model with parameters left on the meta device. Tied/re-shared weights (T5 encoder.embed_tokens, Qwen3 lm_head, CLIP position_ids) are allow-listed. Replace the T5 SDNQ assert checks (stripped by python -O) and add the missing unexpected-key check to the standalone Qwen3 SDNQ loader. Model identification - Reject FLUX.2 in Main_SDNQ_FLUX_Config and accept the bare diffusers transformer layout (transformer_blocks./context_embedder.) in Main_SDNQ_Flux2_Config, and list FLUX.2 before FLUX.1 in the config union, so a single-file SDNQ FLUX.2 checkpoint identifies as FLUX.2 instead of falling through to FLUX.1 or unknown. - Main_Diffusers_ZImage_Config: reject SDNQ-quantized ZImagePipeline folders so Main_SDNQ_Diffusers_ZImage_Config matches them instead (mirrors the FLUX/FLUX.2 guards). - Qwen3Encoder_SDNQ_Folder_Config: verify the folder is actually a Qwen3 encoder, reject complete causal LMs (root config.json + tokenizer files), accept the same compatible Qwen architectures as the unquantized config (Qwen2VLForConditionalGeneration / Qwen2ForCausalLM / Qwen3ForCausalLM via a shared _QWEN3_ENCODER_ARCHITECTURES constant), and make the state-dict fallback resilient to sharded folders — so an SDNQ transformer/VAE/TextLLM folder is not stored as type=qwen3_encoder and a compatible sharded encoder is not wrongly rejected. - Qwen3Encoder_Qwen3Encoder_Config: reject SDNQ folders so it no longer competes with the SDNQ folder config (both share the Qwen3Encoder type, making the factory tiebreak non-deterministic). SDNQ pipeline submodel discovery - Main_SDNQ_Diffusers_Flux2_Config / Main_SDNQ_Diffusers_ZImage_Config _get_submodels(): record the TextEncoder for any compatible encoder class (_QWEN3_ENCODER_ARCHITECTURES) and the Tokenizer for the slow and fast Qwen2 tokenizer classes, not just Qwen3ForCausalLM / Qwen2Tokenizer. Otherwise a valid pipeline advertising e.g. Qwen2ForCausalLM or Qwen2TokenizerFast was recorded as partial and forced to use separate VAE/Qwen3 sources. Z-Image / FLUX.2 self-contained SDNQ pipelines - Add a shared is_self_contained_sdnq_pipeline() helper requiring the specific VAE + Qwen3 (text_encoder + tokenizer) submodels a pipeline install must ship. A truthy submodels dict is not enough: a partial (or partially recognized) pipeline can expose only the transformer and would fail at runtime on missing fixed subfolders. - z_image_model_loader and flux2_klein_model_loader: use the helper for both the main-model self-contained fallback and the explicit qwen3_source validator, so a partial pipeline (as main model or selected source) requires an explicit VAE/Qwen3 source instead of emitting requests against missing folders. - Drop the ui_model_format=diffusers hint from z_image qwen3_source_model so SDNQ pipelines are not filtered out of the node/workflow model pickers. - Mirror the same specific-submodels requirement in the frontend readiness checks (Z-Image + FLUX.2, generate and canvas tabs) and buildZImageGraph, and bring the FLUX.2 canvas readiness block to parity with the generate tab (it previously used format !== 'diffusers' and never recognized SDNQ pipelines). Add regression tests for every fix: mixed-group uint4 dequant, the load-guard helper, SDNQ FLUX.2 identification, SDNQ ZImagePipeline-folder identification, Qwen3 encoder folder identification (SDNQ rejection, causal-LM rejection, compatible Qwen2 class names, sharded folders), SDNQ pipeline submodel discovery across compatible encoder/tokenizer classes, the Z-Image and FLUX.2 self-contained loader/readiness paths incl. partial-pipeline handling, and the qwen3_source_model field template.
|
3 participants
Summary
Adds support for SDNQ (SD.Next Quantization) as a new quantization format in InvokeAI, enabling memory-efficient inference for large models on consumer GPUs.
What's included:
sdnqquantization backend (invokeai/backend/quantization/sdnq/) withSDNQTensor, dequant utils, and safetensors loaders (incl. multi-shard support)norm_outscale/shift fix)ZImagePipelinediffusers folders (all submodels dispatched via SDNQ loader)ZImagePipeline/Flux2KleinPipelinefolders asmain_is_diffusersso submodels auto-extract (no separate VAE/Qwen3 source required)SDNQmodel format badge, schema/types regeneration, readiness updates, Klein FE combobox now accepts SDNQ pipeline configsdocs/src/content/docs/configuration/sdnq-quantization.mdxtests/backend/quantization/sdnq/covering tensor dequant + loader behavior; custom-modules tests extendedWhy: SDNQ enables running FLUX, FLUX.2, and Z-Image on lower-VRAM GPUs by loading pre-quantized weight folders directly, without runtime conversion overhead.
Related Issues / Discussions
Closes #8789
QA Instructions
diffusion_pytorch_model-*-of-*.safetensorsfiles merge correctly (Klein 9B, FLUX.2 dev)bf16reads on packeduint8weightsnorm_outscale/shift swap).uv run --extra cuda pytest tests/backend/quantization/sdnq/.Merge Plan
Needs Testing
Checklist
What's Newcopy (if doing a release after this PR)