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 116 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.
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.
…ation # Conflicts: # invokeai/app/invocations/flux_vae_decode.py # invokeai/backend/model_manager/configs/factory.py # invokeai/backend/model_manager/configs/qwen3_encoder.py # invokeai/backend/model_manager/configs/t5_encoder.py # invokeai/backend/model_manager/load/model_loaders/flux.py # invokeai/backend/model_manager/load/model_loaders/z_image.py # invokeai/frontend/web/openapi.json # invokeai/frontend/web/src/services/api/schema.ts
…/InvokeAI into feature/svd-quantization
JPPhoto
left a comment
JPPhoto
left a comment
There was a problem hiding this comment.
Here's what I hope is a complete list. I am going to check in failing tests for the merge blockers.
Merge blockers
-
invokeai/backend/model_manager/configs/main.py_get_submodelsonly checksis_dir(). Empty or incompletetransformer,text_encoder,tokenizer, orvaedirectories are recorded as present, soinvokeai/app/invocations/model.pytreats the pipeline as self-contained and readiness allows it. Loading then fails at runtime.Test:create a pipeline with valid metadata and empty component directories and assert it is rejected or not self-contained. -
invokeai/app/invocations/model.pyis_self_contained_sdnq_pipelinedoes not require theTransformersubmodel. A malformed index can expose VAE, text encoder, and tokenizer entries while omitting the transformer, and still pass the self-contained check even though every loader requests the transformer.Test:keep the transformer directory but remove its model-index entry and assert the pipeline is not self-contained. -
invokeai/backend/model_manager/configs/qwen3_encoder.py_validate_is_qwen3_encoderuses genericmodel.layers.*andmodel.embed_tokens.weightkeys for config-less folders. Those keys also identify Qwen2 and Qwen2-VL folders, whileinvokeai/backend/model_manager/load/model_loaders/z_image.pyreconstructsQwen3ForCausalLMand rejects their missing or visual-tower keys.Test:identify config-less SDNQ Qwen2 and Qwen2-VL folders and assert they are rejected while Qwen3 remains accepted. -
invokeai/frontend/web/src/services/api/types.tsisZImageDiffusersMainModelConfigandisFlux2DiffusersMainModelConfigreturn true for any nonempty SDNQsubmodelsmap. A partial pipeline containing only a transformer is therefore offered as a valid source;invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.tscan auto-select it, and backend invocation validation fails only after graph construction.Test:expose a partial SDNQ pipeline with only a transformer and assert it is absent from source selectors and cannot satisfy readiness.
Follow-on PR
-
invokeai/backend/model_manager/configs/main.pyMain_SDNQ_Diffusers_FLUX_Config._get_submodelsrecords every component advertised bymodel_index.jsonwithout checking that its directory exists. Partial FLUX.1 downloads can register missing CLIP, T5, tokenizer, or VAE submodels.Test:remove one advertised FLUX.1 component directory and assert discovery omits it. -
invokeai/app/invocations/flux_model_loader.pyFluxModelLoaderInvocationstill requires separate T5, CLIP, and VAE inputs even though the new FLUX.1 SDNQ starter model is documented as a complete pipeline with no dependencies.Test:invoke the loader with only a complete FLUX.1 SDNQ pipeline and assert all component outputs are derived from it. -
invokeai/backend/model_manager/configs/main.pyfrom_model_on_diskallows a serializedsubmodelsmap to override fresh filesystem discovery. After a component is deleted, rehydrating the stored config can still mark the model self-contained.Test:serialize a complete config, remove one component directory, rehydrate it, and assert self-contained detection fails. -
invokeai/backend/model_manager/configs/main.py_get_submodelstrustsmodel_index.jsonclass names without validating each component's actual config or weights. A stale index can label a Qwen2, vision encoder, or unrelated VAE as a loadable Qwen3 pipeline component.Test:pair a valid index with mismatched component configs and assert identification rejects the bundle. -
invokeai/backend/quantization/sdnq/utils.pydequantize_uint4_per_groupperforms an unconditional full-tensorunique()and several reductions during the first uint4 dequantization.invokeai/backend/quantization/sdnq/sdnq_tensor.pyalso performs diagnostic reductions and prints during the first few dequantizations. This adds synchronization, allocations, and stdout noise to model inference, with avoidable OOM risk for large tensors.Test:dequantize a large uint4 tensor and assert diagnostics do not scan the full tensor or write to stdout unless explicitly enabled. -
invokeai/backend/model_manager/load/model_loaders/vae.py_load_sdnq_vae,invokeai/backend/model_manager/load/model_loaders/flux.py_load_sdnq_clipand_load_sdnq_t5, and the custom Qwen loader construct inference models without consistently callingeval(). Any nonzero dropout or training-sensitive configuration remains in training mode.Test:load each custom SDNQ model and assert the returned module and descendants are in evaluation mode. -
invokeai/app/invocations/flux_vae_decode.py_vae_decodeaddsvae.config.shift_factorunconditionally forAutoencoderKL. Standard AutoencoderKL configurations may haveshift_factor=None, producing a runtimeTypeError.Test:decode with an AutoencoderKL whose config has no shift factor and assert decoding succeeds with zero shift. -
invokeai/frontend/web/src/services/api/types.tssource-model predicates andinvokeai/frontend/web/src/features/queue/store/readiness.tssource availability checks do not validate the specific required submodels. A partial SDNQ source can make VAE and Qwen3 sources appear available even though the backend will reject it.Test:use a single-file main model plus a partial SDNQ source and assert readiness remains blocked until a complete source or explicit components are selected.
|
Looks like I don't have write access. Attaching a test patch (with one new file and two updated files) that you can apply: |
Addresses the four merge blockers from the PR invoke-ai#9228 review and checks in the reviewer's failing tests. - _get_submodels (FLUX.2 + Z-Image) only checked is_dir(), so the empty component folders an interrupted download leaves behind were recorded as present. Require the files the component's loader actually needs: config + weights for transformer/text_encoder/vae, a vocab/config file for the weightless tokenizer folder. - is_self_contained_sdnq_pipeline() did not require the Transformer submodel, so a malformed model_index.json advertising only VAE/encoder/ tokenizer passed the check even though every loader requests the transformer. - Qwen3Encoder_SDNQ_Folder_Config's config-less fallback matched the generic Qwen keys (model.layers.* / model.embed_tokens.weight), which Qwen2 and Qwen2-VL folders carry too. Mirror the single-file path: reject a bundled visual tower and require the Qwen3-only q_norm/k_norm weights. - isZImage/isFlux2DiffusersMainModelConfig returned true for any nonempty SDNQ submodels map, offering a transformer-only pipeline in the source pickers. Added isSelfContainedSDNQPipeline() mirroring the backend's required-submodel set and reused it in readiness.ts and buildZImageGraph.ts so both sides agree on what "complete" means. Test fixtures now write real component content, since discovery no longer accepts bare directories; the empty-dir case is requested explicitly.
JPPhoto
left a comment
JPPhoto
left a comment
There was a problem hiding this comment.
Next round:
Merge blockers
-
invokeai/app/invocations/flux_vae_decode.py:68raisesTypeErrorwhenAutoencoderKL.config.shift_factorisNone. Test: decode with a standard config lackingshift_factor. -
invokeai/backend/model_manager/configs/main.py:2237trusts model-index class names plus generic files, so mismatched Qwen2, Qwen-VL, or unrelated components can pass self-contained detection. Test: pair a Qwen3 index with mismatched component config and weights; assert rejection. Alternative: dispatch loading from the component's actual architecture instead of the advertised class.
Follow-on PR candidates
-
invokeai/backend/model_manager/configs/main.py:2617records FLUX.1 components without checking directory contents. Test: remove each advertised component and assert discovery omits it. -
invokeai/app/invocations/flux_model_loader.py:71still requires separate T5, CLIP, and VAE inputs for a complete FLUX.1 SDNQ pipeline. Test: invoke with only the pipeline and assert all outputs use its submodels. Alternative: add a dedicated self-contained FLUX.1 loader mode with optional component inputs. -
invokeai/backend/model_manager/configs/t5_encoder.py:_safetensors_dir_has_sdnq_keysrejects weight/scale pairs split across shards. Test: split a pair across two shards without a quantization marker. Alternative: use a manifest-based component inventory shared by identification and loading. -
invokeai/backend/model_manager/load/model_loaders/flux.py:1632andinvokeai/backend/model_manager/load/model_loaders/vae.py:330do not consistently put custom SDNQ models in evaluation mode. Test: load each component and assert all modules are evaluation-mode. -
invokeai/backend/quantization/sdnq/utils.py:242performs full-tensor diagnostics, whileinvokeai/backend/quantization/sdnq/loaders.py:367writes directly to stdout. Test: dequantize a large tensor and assert no unsolicited output or full scan. -
invokeai/backend/model_manager/configs/main.py:2617stores index-derived paths while loaders reconstruct fixed directory names. Test: use valid components under nonstandard index keys and assert loading follows discovered paths. -
invokeai/backend/model_manager/configs/main.py:2278lets serialized submodel maps bypass fresh filesystem state. Test: serialize a complete config, delete a component, rehydrate, and assert it is not self-contained.
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)