Replace with SVDQuant+GPTQ NVFP4 DiT and text encoder, native in Diffusers

NOTICE

@@ -2,22 +2,32 @@ Qwen is licensed under the Qwen RESEARCH LICENSE AGREEMENT, Copyright (c) 2026 H
 
 Built with Qwen.
 
-This repository redistributes, for non-commercial research and evaluation use only, files derived from:
+This repository redistributes, for non-commercial research and evaluation use only, files derived from
+Qwen/Qwen-Image-2.1 (revision 790c92633540aa0cb11d9abf19eb46d861714758).
 
-- Qwen/Qwen-Image-2.1 (revision 790c92633540aa0cb11d9abf19eb46d861714758)
-  model_index.json, LICENSE, processor/, scheduler/, vae/, text_encoder/config.json,
-  text_encoder/generation_config.json. Unmodified.
+Unmodified files: LICENSE, processor/, scheduler/, text_encoder/generation_config.json.
 
-- HangGlidersRule/Darkstar-Qwen-Image-2.1-Base-ModelOpt-W4A4-NVFP4 (revision 3b7aa842dd9c1ec8f90d005979f8b2a0d7d7e95b)
-  transformer/config.json, transformer/diffusion_pytorch_model.safetensors. Unmodified
-  (sha256 a6c3080b764344729ef95846faa3a0a01eb36bdef93f785ed07f4ac5751ad76a).
+MODIFIED FILES (changed by joseplcam):
 
-- BennyDaBall/Qwen-Image-2.1-NVFP4 (revision 1a38d44a3a2f35cb0b543a25b04da0a963e7b5e6)
-  text_encoders/qwen3vl_8b_nvfp4.safetensors
-  (sha256 cdd9b85bcb60d5cb358ddecfd45a48fa15cb728b70d88bd2028840eaef095d16).
+- transformer/diffusion_pytorch_model.safetensors, transformer/config.json
+  Quantized from the official BF16 transformer with SVDQuant + GPTQ (NVFP4 weights and activations,
+  rank-32 BF16 low-rank branch) using diffuse-compressor. Blocks 0, 1, 30, 31 and the global
+  modulation are the official BF16 tensors. config.json gained a `quantization_config` entry.
+  sha256 0304fc3e84b13663b439601380066754ef9b92025e159275abe163bc5f1ceeb4
 
-MODIFIED FILE: text_encoder/model.safetensors was changed by joseplcam from the BennyDaBall file above.
-The only change is tensor names: the prefixes "model.layers.", "model.embed_tokens." and "model.norm."
-were renamed to "model.language_model.layers.", "model.language_model.embed_tokens." and
-"model.language_model.norm." to match the Hugging Face Qwen3-VL layout. Tensor data, dtypes, shapes and
-file metadata are unchanged (sha256 of the result 42882503fe37fbc72f374fb9e55e79f020f9cecca02463417176fdd902be470e).
+- text_encoder/model.safetensors, text_encoder/config.json
+  Quantized from the official BF16 Qwen3-VL text encoder with the same method: the MLP projections
+  (gate/up/down) of decoder layers 4-31, rank-128 low-rank branch. All attention projections, decoder
+  layers 0-3 and 32-35, the vision tower, embeddings and lm_head are the official BF16 tensors.
+  config.json gained a `nunchaku_lite` entry.
+  sha256 83c4462b5ce6b98e2708ca53031e2f9a60a8e402c4c4a0ccd19cb2df8e817ae2
+
+- vae/diffusion_pytorch_model.safetensors, vae/config.json
+  The official VAE weights cast from FP32 to BF16.
+  sha256 71879ffd5321e6d10c3c87513e2b474b1252efa7f3dec2969214a9bf06a6dd5c
+
+- model_index.json
+  The `text_encoder` and `transformer` entries point to the custom classes below.
+
+NEW FILES: text_encoder/modeling_nunchaku_qwen3vl.py, transformer/modeling_nunchaku_qwenimage21.py,
+text_encoder/nunchaku_kernels.py, transformer/nunchaku_kernels.py.

README.md

@@ -10,129 +10,150 @@ tags:
   - qwen-image
   - qwen-image-2.1
   - nvfp4
-  - sglang
+  - svdquant
+  - nunchaku
   - blackwell
   - text-to-image
   - image-editing
 ---
 
-# Qwen-Image-2.1-NVFP4 (DiT + text encoder) for SGLang Diffusion
+# Qwen-Image-2.1-NVFP4 (SVDQuant, DiT + text encoder, native in Diffusers)
 
 **Built with Qwen.** Non-commercial research and evaluation use only (Qwen Research License, see `LICENSE` and `NOTICE`).
 
-This repository is a ready-to-run combination of existing NVFP4 quantizations of
-[Qwen/Qwen-Image-2.1](https://huggingface.co/Qwen/Qwen-Image-2.1), packaged in the diffusers layout so that
-[SGLang Diffusion](https://github.com/sgl-project/sglang) can load it from one model id and run both large
-components on native Blackwell FP4 kernels. No new quantization was done here: the value of this repo is the tested
-combination, the small key-rename that makes the encoder loadable by SGLang, and the measurements below.
+An NVFP4 quantization of [Qwen/Qwen-Image-2.1](https://huggingface.co/Qwen/Qwen-Image-2.1), made
+from the official BF16 weights with SVDQuant + GPTQ and calibrated on both text-to-image prompts and
+image edits. Both large components run on native Blackwell FP4 tensor cores, and the pipeline loads with a
+single `DiffusionPipeline.from_pretrained` call.
 
-Text-to-image, image editing and RGBA output all work as in the original pipeline.
+At 512x512 on an RTX PRO 6000 Blackwell it is about **1.7x faster than BF16 and needs about 40% less VRAM**,
+and it matches BF16 closely on text-to-image, typography and edits, including a hard recolor edit that
+earlier 4-bit text encoders failed.
 
-## What is inside, and where each part comes from
-
-| Path | Precision | Source | Changes |
-|---|---|---|---|
-| `transformer/` (7B single-stream DiT) | NVFP4 W4A4, group 16, FP8 block scales. Blocks 0, 1, 30, 31 and the embeddings, modulation and output layers stay BF16 | [HangGlidersRule/Darkstar-Qwen-Image-2.1-Base-ModelOpt-W4A4-NVFP4](https://huggingface.co/HangGlidersRule/Darkstar-Qwen-Image-2.1-Base-ModelOpt-W4A4-NVFP4) @ `3b7aa84`. NVIDIA ModelOpt export of the base model | None (byte-identical) |
-| `text_encoder/model.safetensors` (Qwen3-VL 8B) | NVFP4 for the 252 language-model linear layers; vision tower, embeddings and `lm_head` stay BF16 | [BennyDaBall/Qwen-Image-2.1-NVFP4](https://huggingface.co/BennyDaBall/Qwen-Image-2.1-NVFP4) @ `1a38d44`, file `text_encoders/qwen3vl_8b_nvfp4.safetensors`. ComfyUI-style NVFP4 converted from the official BF16 weights | Tensor names renamed to the Hugging Face Qwen3-VL layout (`model.layers.*` → `model.language_model.layers.*`, same for `embed_tokens` and `norm`). Data unchanged |
-| `vae/`, `processor/`, `scheduler/`, `model_index.json`, `text_encoder/*.json`, `LICENSE` | BF16 / FP32 as released | [Qwen/Qwen-Image-2.1](https://huggingface.co/Qwen/Qwen-Image-2.1) @ `790c926` | None |
-
-Why these two parts:
-
-- The DiT is a ModelOpt NVFP4 export. SGLang detects it from `transformer/config.json` and runs it with FlashInfer's
-  CUTLASS FP4 GEMM, which has an SM120 path (RTX 50 series, RTX PRO 6000). Keeping the first two and last two blocks in
-  BF16 preserves quality: its author reports parity with BF16 on a 25-case evaluation.
-- The encoder is the only published NVFP4 Qwen3-VL 8B built from the official weights (not an abliterated variant).
-  SGLang runs its Comfy NVFP4 layers natively via `comfy-kitchen` FP4 matmuls. The rename was needed because SGLang's
-  native Qwen3-VL encoder expects the Hugging Face tensor names.
-
-Total size is about 14 GB, against about 32 GB for the BF16 release.
-
-## Usage (SGLang Diffusion)
-
-Tested with SGLang commit `ffac53d779c08dcdab2d07e5e2a41dba83f0e65c`, PyTorch 2.13.0+cu130, FlashInfer 0.6.18 and
-`comfy-kitchen==0.2.35`. You need a Blackwell GPU (compute capability 10.0 or 12.0).
-
-SGLang only recognizes the Comfy NVFP4 encoder when its weights path points at the file, so pass it explicitly:
+## Usage
 
 ```python
-from sglang.multimodal_gen import DiffGenerator
-
-repo = "joseplcam/Qwen-Image-2.1-NVFP4"
-with DiffGenerator.from_pretrained(
-    model_path=repo,
-    component_weights_paths={"text_encoder": f"{repo}/text_encoder/model.safetensors"},
-    performance_mode="speed",       # keep every component resident on the GPU
-    attention_backend="torch_sdpa",  # "sage_attn" is faster at 1024x1024 and above
-) as generator:
-    result = generator.generate(sampling_params_kwargs=dict(
-        prompt="A capybara reading a book by candlelight",
-        width=512, height=512, num_inference_steps=40, guidance_scale=1.0, seed=42,
-    ))
-```
-
-Add `image_path="input.png"` for editing. The same arguments work with a local clone of this repository.
+import torch
+from diffusers import DiffusionPipeline
 
-Serving with batching of concurrent text-to-image requests:
+pipe = DiffusionPipeline.from_pretrained(
+    "joseplcam/Qwen-Image-2.1-NVFP4", dtype=torch.bfloat16, trust_remote_code=True
+).to("cuda")
 
-```bash
-sglang serve --model-path joseplcam/Qwen-Image-2.1-NVFP4 \
-  --component-weights-paths.text_encoder joseplcam/Qwen-Image-2.1-NVFP4/text_encoder/model.safetensors \
-  --performance-mode speed --batching-max-size 8 --batching-delay-ms 50
+image = pipe("A capybara reading a book by candlelight", height=512, width=512, num_inference_steps=40).images[0]
+edited = pipe("Make it night time with moonlight", image=image, output_resolution=512, num_inference_steps=40).images[0]
 ```
 
-Tips:
-
-- Set `performance_mode="speed"`. Without it, SGLang's automatic memory policy may stream the text encoder layer by
-  layer from host memory even on a 96 GB GPU, which made every request about 10x slower in our tests.
-- The first run JIT-compiles FlashInfer's FP4 kernels (about 5 minutes; cached in `~/.cache/sglang`). On machines with
-  about 32 GB of RAM, set `MAX_JOBS=4`, or the compiler may run out of memory.
-- On NixOS, the JIT linker needs `libcuda` and `libcudart`, for example
-  `LIBRARY_PATH=/run/opengl-driver/lib:<cuda>/lib`.
-- torch.compile made this model slower in every configuration we tried, so leave it off.
-
-## Measurements (RTX PRO 6000 Blackwell 96 GB, SGLang serve, 512x512, 40 steps)
-
-Average of 3 runs after a warmup. Batch size is the number of concurrent requests, each with a different prompt
-(and a different reference image for edits). VRAM is the peak for the whole GPU.
-
-Text-to-image:
-
-| Batch | SDPA | SageAttention2 | SDPA + torch.compile | SageAttention2 + torch.compile |
-|---|---|---|---|---|
-| 1 | **2.42 s** (0.41 img/s, 16.7 GB) | 2.54 s (0.39 img/s, 16.7 GB) | 6.39 s (0.16 img/s) | 6.80 s (0.15 img/s) |
-| 2 | **2.86 s** (0.70 img/s, 17.9 GB) | 3.02 s (0.66 img/s, 17.9 GB) | 5.75 s (0.35 img/s) | 6.17 s (0.32 img/s) |
-| 4 | **5.58 s** (0.72 img/s, 20.1 GB) | 5.82 s (0.69 img/s, 20.1 GB) | 6.39 s (0.63 img/s) | 7.77 s (0.52 img/s) |
-| 8 | **11.33 s** (0.71 img/s, 25.0 GB) | 11.70 s (0.68 img/s, 25.1 GB) | 12.34 s (0.65 img/s) | 14.35 s (0.56 img/s) |
-
-Image editing runs at about 2.4 s per image (0.41 img/s, 16.8 GB) at any batch size, because SGLang does not merge
-edit requests that carry reference images; they run one after another.
-
-Comparison with other precisions (single request, SDPA):
-
-| Setup | 512x512 | 1024x1024 | Peak VRAM at 512x512 |
-|---|---|---|---|
-| BF16 DiT + BF16 encoder (official release) | 2.85 s | 11.45 s | 34.7 GB |
-| NVFP4 DiT + BF16 encoder | 2.26 s | 6.00 s | 26.7 GB |
-| NVFP4 DiT + NVFP4 encoder (this repo) | 2.29 s | not measured | 16.9 GB |
-
-At 1024x1024 with the NVFP4 DiT, SageAttention2 took 5.53 s against 6.00 s for SDPA.
-
-Quality: we saw no artifacts on our sample prompts or edits. Quantization changes composition details compared with
-BF16 at the same seed (LPIPS against BF16: 0.26 at 512x512 and 0.11 at 1024x1024 for the NVFP4 DiT; 0.19 for adding
-the NVFP4 encoder). Check quality on your own prompts, especially for text rendering.
+`trust_remote_code=True` loads two small files from this repository:
+
+- `text_encoder/modeling_nunchaku_qwen3vl.py`: a `Qwen3VLForConditionalGeneration` subclass that swaps the
+  quantized linears for Diffusers' own `SVDQW4A4Linear` before loading the weights.
+- `nunchaku_kernels.py` (in `text_encoder/` and `transformer/`): Diffusers loads the NVFP4 kernels from
+  `rootonchair/nunchaku-lite-kernels`, which is no longer downloadable. This file points that name at
+  [joseplcam/nunchaku-lite-kernels](https://huggingface.co/joseplcam/nunchaku-lite-kernels), an unmodified
+  build of the same open-source kernels, and sets `DIFFUSERS_TRUST_REMOTE_KERNELS=true` unless you already set
+  it. Set `LOCAL_KERNELS` yourself to use a different build.
+
+Use a different seed for an edit than the one that generated its input image. Qwen-Image-2.1 returns an
+over-sharpened copy that ignores the prompt when the edit starts from the same noise
+([diffusers #14824](https://github.com/huggingface/diffusers/issues/14824)); this is a base-model behaviour.
+
+### Requirements
+
+- NVIDIA Blackwell GPU with compute capability 12.0 (RTX 50 series, RTX PRO 6000). The published kernel build
+  targets `sm_120a` only.
+- Linux x86_64, Python 3.12, PyTorch 2.13 with CUDA 13.0, which the kernel build targets.
+- Diffusers from `main` with Qwen-Image-2.1 and Nunchaku Lite support (tested at commit `0377f0c`),
+  `transformers>=5.12`, `accelerate`, `kernels>=0.14`.
+
+## What is inside
+
+| Component | Precision | Details |
+|---|---|---|
+| `transformer/` (7B DiT) | NVFP4 W4A4, group 16, FP8 block scales + BF16 rank-32 low-rank branch | All attention and MLP projections of blocks 2-29 (196 layers). Blocks 0, 1, 30, 31 and the global modulation stay BF16. |
+| `text_encoder/` (Qwen3-VL 8B) | NVFP4 W4A4 + BF16 rank-128 low-rank branch | MLP projections (gate/up/down) of decoder layers 4-31 (84 layers). All attention projections, layers 0-3 and 32-35, the vision tower, embeddings and `lm_head` stay BF16. |
+| `vae/` | BF16 | Official weights cast from FP32. |
+| `processor/`, `scheduler/` | as released | Unmodified. |
+
+Total download: about 17 GB, against about 32 GB for the BF16 release.
+
+## How it was made
+
+- **Method:** SVDQuant with GPTQ residual rounding, via [diffuse-compressor](https://github.com/rootonchair/diffuse-compressor)
+  (commit `0965874`). A low-rank BF16 branch absorbs the outliers of each weight and SmoothQuant-style
+  scaling migrates activation outliers; GPTQ then rounds the 4-bit residual using calibration statistics.
+  Activation scales are dynamic, so nothing is fixed to the calibration inputs.
+- **Calibration data (128 samples per component):** 64 prompts from the qdiff prompt set and 64 image edits
+  from the train split of [VyoJ/NHR-Edit-Change_Only](https://huggingface.co/datasets/VyoJ/NHR-Edit-Change_Only).
+  Three of every four samples at 512x512, the rest at 1024x1024. The DiT saw 20 denoising steps per sample,
+  with the prefix KV cache disabled so the prompt and reference-image tokens pass through every step.
+- **Sensitive layers:** the first and last two DiT blocks and the modulation stay BF16. For the text encoder,
+  quantizing attention hurt edits most. With every encoder linear quantized, the "turn the parrots blue"
+  recolor below worked in 1 of 8 seeds (rank 32) or 4-5 of 8 (rank 128); quantizing only the MLPs with rank
+  128 brought it to 7 of 8, the same as BF16.
+- **Runtime:** Diffusers' built-in Nunchaku Lite quantizer for the DiT and the same `SVDQW4A4Linear` layers
+  for the text encoder, running the [nunchaku-lite](https://github.com/rootonchair/nunchaku-lite) CUDA kernels.
+
+The scripts that produced this repository are in `tools/`: `quantize_dit.py`,
+`quantize_text_encoder.py --rank 128 --edge-layers 4 --skip-attention`, `package.py`, `evaluate.py` and
+`parrot_test.py`. They need a checkout of diffuse-compressor (its `examples/` package) at
+`$DIFFUSE_COMPRESSOR`.
+
+## Results
+
+RTX PRO 6000 Blackwell (96 GB), Diffusers, 40 steps, guidance 1, no `torch.compile`, after a warmup.
+Peak VRAM is PyTorch's peak allocation.
+
+| | BF16 (official) | This repo |
+|---|---|---|
+| Text-to-image, 512x512 | 3.32 s | **1.91 s** (1.74x) |
+| Edit, 512x512 | 3.73 s | **2.18 s** (1.71x) |
+| Text-to-image, 1024x1024 | 14.06 s | **8.12 s** (1.73x) |
+| Peak VRAM, 512x512 | 31.9-32.5 GiB | **18.6-19.3 GiB** |
+| Throughput, 512x512 batch 1-8 | 0.28-0.30 img/s | **0.49-0.55 img/s** |
+
+Batching several prompts gives no extra throughput in Diffusers on this GPU; it is saturated at batch 1.
+The pipeline shares one reference-image list across a batch, so edits with different images run one at a time.
+
+Closeness to BF16 (LPIPS with AlexNet, same seeds, 512x512; lower is closer, around 0.1 is hard to tell apart):
+
+| Set | LPIPS vs BF16 |
+|---|---|
+| 8 text-to-image prompts (portraits, typography, counting, scenes) | 0.148 |
+| 12 held-out NHR-Edit test edits + 1 parrot recolor | 0.046 |
+
+Instruction following on a hard recolor ("Turn the parrots blue", 8 seeds, tiles with at least one blue parrot),
+with this repo's DiT:
+
+| Text encoder | Followed |
+|---|---|
+| BF16 | 7 / 8 |
+| NVFP4, every linear, rank 32 | 1 / 8 |
+| NVFP4, every linear, rank 128, 4 BF16 edge layers | 4-5 / 8 |
+| **NVFP4, MLPs only, rank 128, 4 BF16 edge layers (this repo)** | **7 / 8** |
+
+With the BF16 text encoder, this repo's DiT followed the edit in 8 of 8 seeds, the same as the BF16 DiT.
 
 ## Limitations
 
+- Measured on one GPU model and one software stack; quantization shifts details of individual images.
+- The evaluation is small (8 prompts, 13 edits, one 8-seed recolor test). It is not a benchmark.
+- The kernel build covers `sm_120a`, PyTorch 2.13, CUDA 13 and CPython 3.12 only.
 - Non-commercial research and evaluation use only.
-- Measured on one GPU model only (RTX PRO 6000 Blackwell). Other Blackwell GPUs are untested.
-- The encoder needs the explicit `component_weights_paths` argument shown above.
-- diffusers `from_pretrained` does not load the NVFP4 components; use SGLang.
+
+## Previous version
+
+Until 2026-09-25 this repository held a different combination for SGLang Diffusion: the ModelOpt NVFP4 DiT from
+HangGlidersRule/Darkstar-Qwen-Image-2.1-Base-ModelOpt-W4A4-NVFP4 and the NVFP4 Qwen3-VL encoder from
+BennyDaBall/Qwen-Image-2.1-NVFP4. It is still available in this repository's commit history. It was
+replaced because its 4-bit text encoder lost hard edits and it did not load in Diffusers.
 
 ## Credits and license
 
-- Qwen team: Qwen-Image 2.1 model, VAE, processor and configuration.
-- HangGlidersRule (Model Forge): ModelOpt NVFP4 DiT.
-- BennyDaBall: NVFP4 Qwen3-VL 8B text encoder.
+- Qwen team: Qwen-Image-2.1.
+- [SVDQuant](https://arxiv.org/abs/2411.05007) (MIT Han Lab), [nunchaku-lite](https://github.com/rootonchair/nunchaku-lite)
+  and [diffuse-compressor](https://github.com/rootonchair/diffuse-compressor) (rootonchair).
+- Calibration edits: [VyoJ/NHR-Edit-Change_Only](https://huggingface.co/datasets/VyoJ/NHR-Edit-Change_Only).
 
-Qwen is licensed under the Qwen RESEARCH LICENSE AGREEMENT, Copyright (c) 2026 Hangzhou Tongyi Laboratory Technology
-Co., Ltd. All Rights Reserved. See `LICENSE` and `NOTICE`, which also lists the one modified file.
+Qwen is licensed under the Qwen RESEARCH LICENSE AGREEMENT, Copyright (c) 2026 Hangzhou Tongyi Laboratory
+Technology Co., Ltd. All Rights Reserved. See `LICENSE` and `NOTICE`, which lists every modified file.

model_index.json

@@ -10,15 +10,15 @@
     "FlowMatchEulerDiscreteScheduler"
   ],
   "text_encoder": [
-    "transformers",
-    "Qwen3VLForConditionalGeneration"
+    "modeling_nunchaku_qwen3vl",
+    "NunchakuQwen3VLForConditionalGeneration"
   ],
   "transformer": [
-    "diffusers",
-    "QwenImage21Transformer2DModel"
+    "modeling_nunchaku_qwenimage21",
+    "NunchakuQwenImage21Transformer2DModel"
   ],
   "vae": [
     "diffusers",
     "AutoencoderKLQwenImage21"
   ]
-}
+}

text_encoder/config.json

@@ -5,6 +5,101 @@
   "dtype": "bfloat16",
   "image_token_id": 151655,
   "model_type": "qwen3_vl",
+  "nunchaku_lite": {
+    "compute_dtype": "bfloat16",
+    "quant_method": "nunchaku_lite",
+    "svdq_w4a4": {
+      "group_size": 16,
+      "precision": "nvfp4",
+      "rank": 128,
+      "targets": [
+        "model.language_model.layers.4.mlp.gate_proj",
+        "model.language_model.layers.4.mlp.up_proj",
+        "model.language_model.layers.4.mlp.down_proj",
+        "model.language_model.layers.5.mlp.gate_proj",
+        "model.language_model.layers.5.mlp.up_proj",
+        "model.language_model.layers.5.mlp.down_proj",
+        "model.language_model.layers.6.mlp.gate_proj",
+        "model.language_model.layers.6.mlp.up_proj",
+        "model.language_model.layers.6.mlp.down_proj",
+        "model.language_model.layers.7.mlp.gate_proj",
+        "model.language_model.layers.7.mlp.up_proj",
+        "model.language_model.layers.7.mlp.down_proj",
+        "model.language_model.layers.8.mlp.gate_proj",
+        "model.language_model.layers.8.mlp.up_proj",
+        "model.language_model.layers.8.mlp.down_proj",
+        "model.language_model.layers.9.mlp.gate_proj",
+        "model.language_model.layers.9.mlp.up_proj",
+        "model.language_model.layers.9.mlp.down_proj",
+        "model.language_model.layers.10.mlp.gate_proj",
+        "model.language_model.layers.10.mlp.up_proj",
+        "model.language_model.layers.10.mlp.down_proj",
+        "model.language_model.layers.11.mlp.gate_proj",
+        "model.language_model.layers.11.mlp.up_proj",
+        "model.language_model.layers.11.mlp.down_proj",
+        "model.language_model.layers.12.mlp.gate_proj",
+        "model.language_model.layers.12.mlp.up_proj",
+        "model.language_model.layers.12.mlp.down_proj",
+        "model.language_model.layers.13.mlp.gate_proj",
+        "model.language_model.layers.13.mlp.up_proj",
+        "model.language_model.layers.13.mlp.down_proj",
+        "model.language_model.layers.14.mlp.gate_proj",
+        "model.language_model.layers.14.mlp.up_proj",
+        "model.language_model.layers.14.mlp.down_proj",
+        "model.language_model.layers.15.mlp.gate_proj",
+        "model.language_model.layers.15.mlp.up_proj",
+        "model.language_model.layers.15.mlp.down_proj",
+        "model.language_model.layers.16.mlp.gate_proj",
+        "model.language_model.layers.16.mlp.up_proj",
+        "model.language_model.layers.16.mlp.down_proj",
+        "model.language_model.layers.17.mlp.gate_proj",
+        "model.language_model.layers.17.mlp.up_proj",
+        "model.language_model.layers.17.mlp.down_proj",
+        "model.language_model.layers.18.mlp.gate_proj",
+        "model.language_model.layers.18.mlp.up_proj",
+        "model.language_model.layers.18.mlp.down_proj",
+        "model.language_model.layers.19.mlp.gate_proj",
+        "model.language_model.layers.19.mlp.up_proj",
+        "model.language_model.layers.19.mlp.down_proj",
+        "model.language_model.layers.20.mlp.gate_proj",
+        "model.language_model.layers.20.mlp.up_proj",
+        "model.language_model.layers.20.mlp.down_proj",
+        "model.language_model.layers.21.mlp.gate_proj",
+        "model.language_model.layers.21.mlp.up_proj",
+        "model.language_model.layers.21.mlp.down_proj",
+        "model.language_model.layers.22.mlp.gate_proj",
+        "model.language_model.layers.22.mlp.up_proj",
+        "model.language_model.layers.22.mlp.down_proj",
+        "model.language_model.layers.23.mlp.gate_proj",
+        "model.language_model.layers.23.mlp.up_proj",
+        "model.language_model.layers.23.mlp.down_proj",
+        "model.language_model.layers.24.mlp.gate_proj",
+        "model.language_model.layers.24.mlp.up_proj",
+        "model.language_model.layers.24.mlp.down_proj",
+        "model.language_model.layers.25.mlp.gate_proj",
+        "model.language_model.layers.25.mlp.up_proj",
+        "model.language_model.layers.25.mlp.down_proj",
+        "model.language_model.layers.26.mlp.gate_proj",
+        "model.language_model.layers.26.mlp.up_proj",
+        "model.language_model.layers.26.mlp.down_proj",
+        "model.language_model.layers.27.mlp.gate_proj",
+        "model.language_model.layers.27.mlp.up_proj",
+        "model.language_model.layers.27.mlp.down_proj",
+        "model.language_model.layers.28.mlp.gate_proj",
+        "model.language_model.layers.28.mlp.up_proj",
+        "model.language_model.layers.28.mlp.down_proj",
+        "model.language_model.layers.29.mlp.gate_proj",
+        "model.language_model.layers.29.mlp.up_proj",
+        "model.language_model.layers.29.mlp.down_proj",
+        "model.language_model.layers.30.mlp.gate_proj",
+        "model.language_model.layers.30.mlp.up_proj",
+        "model.language_model.layers.30.mlp.down_proj",
+        "model.language_model.layers.31.mlp.gate_proj",
+        "model.language_model.layers.31.mlp.up_proj",
+        "model.language_model.layers.31.mlp.down_proj"
+      ]
+    }
+  },
   "text_config": {
     "attention_bias": false,
     "attention_dropout": 0.0,

text_encoder/model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:42882503fe37fbc72f374fb9e55e79f020f9cecca02463417176fdd902be470e
-size 7549900080
+oid sha256:83c4462b5ce6b98e2708ca53031e2f9a60a8e402c4c4a0ccd19cb2df8e817ae2
+size 11812006400

text_encoder/modeling_nunchaku_qwen3vl.py

@@ -0,0 +1,47 @@
+"""Qwen3-VL text encoder with SVDQuant NVFP4 (Nunchaku Lite) linears.
+
+Loaded by Diffusers as a custom pipeline component (`trust_remote_code=True`).
+`config.json` holds the compact Nunchaku Lite config under `nunchaku_lite`; the
+listed linears are swapped for Diffusers' own `SVDQW4A4Linear` before the
+quantized state dict is loaded, so inference uses the same NVFP4 kernels as the
+transformer. Everything else (vision tower, embeddings, the BF16 layers) loads
+as in the stock model.
+"""
+
+from pathlib import Path
+
+import torch
+from accelerate import init_empty_weights
+from huggingface_hub import snapshot_download
+
+from .nunchaku_kernels import KERNELS  # noqa: F401  (sets up the kernels before Diffusers imports them)
+
+# isort: off  -- must stay below `.nunchaku_kernels`, which prepares the kernels this import loads
+from diffusers.quantizers.nunchaku.utils import replace_with_nunchaku_linear  # noqa: E402
+from safetensors.torch import load_file  # noqa: E402
+from transformers import Qwen3VLForConditionalGeneration  # noqa: E402
+# isort: on
+
+
+class NunchakuQwen3VLForConditionalGeneration(Qwen3VLForConditionalGeneration):
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, subfolder: str = "", **kwargs):
+        path = Path(pretrained_model_name_or_path) / subfolder
+        if not path.is_dir():
+            path = Path(snapshot_download(str(pretrained_model_name_or_path), allow_patterns=[f"{subfolder}/*"])) / subfolder
+        dtype = kwargs.get("dtype") or kwargs.get("torch_dtype") or torch.bfloat16
+        if dtype == "auto":
+            dtype = torch.bfloat16
+
+        config = cls.config_class.from_pretrained(path)
+        with init_empty_weights():  # parameters on meta, buffers (e.g. rotary tables) materialized
+            model = cls(config)
+        replace_with_nunchaku_linear(model, config.nunchaku_lite, dtype)
+        # strict=False only because non-persistent buffers are absent; both checks below stay strict
+        result = model.load_state_dict(load_file(path / "model.safetensors"), strict=False, assign=True)
+        if result.unexpected_keys:
+            raise ValueError(f"Checkpoint has {len(result.unexpected_keys)} unexpected tensors, e.g. {result.unexpected_keys[:3]}")
+        missing = [name for name, p in model.named_parameters() if p.device.type == "meta"]
+        if missing:
+            raise ValueError(f"Checkpoint is missing {len(missing)} parameters, e.g. {missing[:3]}")
+        return model.eval()

text_encoder/nunchaku_kernels.py

@@ -0,0 +1,19 @@
+# Shared by text_encoder/ and transformer/: whichever component Diffusers loads first sets this up.
+#
+# Diffusers fetches the Nunchaku Lite kernels from `rootonchair/nunchaku-lite-kernels`, which is
+# no longer downloadable. Loading this repo with trust_remote_code=True already runs this code,
+# so point that kernel name at the rebuilt copy before Diffusers imports its Nunchaku utilities.
+import os
+
+from huggingface_hub import snapshot_download
+
+KERNELS = "rootonchair/nunchaku-lite-kernels"
+if KERNELS not in os.environ.get("LOCAL_KERNELS", ""):
+    local = f"{KERNELS}={snapshot_download('joseplcam/nunchaku-lite-kernels')}"
+    os.environ["LOCAL_KERNELS"] = ":".join(filter(None, [os.environ.get("LOCAL_KERNELS"), local]))
+if "DIFFUSERS_TRUST_REMOTE_KERNELS" not in os.environ:
+    # Diffusers read this variable into a constant at import time, so set both.
+    import diffusers.utils.constants
+
+    os.environ["DIFFUSERS_TRUST_REMOTE_KERNELS"] = "true"
+    diffusers.utils.constants.DIFFUSERS_TRUST_REMOTE_KERNELS = True

tools/evaluate.py

@@ -0,0 +1,103 @@
+"""Compare a Qwen-Image-2.1 pipeline against the BF16 original on text-to-image and edits.
+
+Writes, per pipeline: images, per-item LPIPS vs BF16 (same seeds), timings and peak VRAM.
+Edits come from the NHR-Edit test split (held out from calibration) plus the parrot
+recolor that 4-bit models failed before. Instruction following is judged by eye from
+the saved grids; LPIPS only measures closeness to BF16.
+
+Usage:
+    python tools/qwen21_nvfp4/evaluate.py <pipeline path or id> <tag> [--trust-remote-code] [--baseline <bf16 out dir>]
+"""
+
+import argparse
+import json
+import time
+import urllib.request
+from io import BytesIO
+from pathlib import Path
+
+import datasets
+import torch
+from diffusers import DiffusionPipeline
+from PIL import Image
+
+PROMPTS = [
+    "A capybara reading a book by candlelight",
+    "A red vintage bicycle leaning against a blue door",
+    "A bowl of ramen with a soft boiled egg, top view",
+    "A neon sign that says OPEN in a rainy alley",
+    "A glass teapot with blooming flower tea on a table",
+    "A portrait of an elderly fisherman with a knitted cap, soft window light",
+    'A chalkboard menu that reads "SOUP OF THE DAY: TOMATO"',
+    "Three red apples and two green pears on a wooden table",
+]
+PARROTS = "https://r0k.us/graphics/kodak/kodak/kodim23.png"
+
+
+def edit_items(n: int) -> list[tuple[str, Image.Image, str]]:
+    data = datasets.load_dataset("VyoJ/NHR-Edit-Change_Only", split="test")
+    items = [(f"nhr{i}", data[i]["source"].convert("RGB"), data[i]["edit_instruction"]) for i in range(n)]
+    parrots = Image.open(BytesIO(urllib.request.urlopen(PARROTS).read())).convert("RGB")
+    return items + [("parrots", parrots, "Turn the parrots blue")]
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("pipeline")
+    parser.add_argument("tag")
+    parser.add_argument("--trust-remote-code", action="store_true")
+    parser.add_argument("--edits", type=int, default=12)
+    parser.add_argument("--size", type=int, default=512)
+    parser.add_argument("--out", default="outputs/qwen21_eval")
+    parser.add_argument("--baseline", help="output dir of the BF16 run, for LPIPS")
+    args = parser.parse_args()
+
+    out = Path(args.out) / args.tag
+    out.mkdir(parents=True, exist_ok=True)
+    pipe = DiffusionPipeline.from_pretrained(args.pipeline, dtype=torch.bfloat16, trust_remote_code=args.trust_remote_code).to("cuda")
+    pipe.set_progress_bar_config(disable=True)
+    gen = lambda seed: torch.Generator("cuda").manual_seed(seed)  # noqa: E731
+
+    def run(name, seed, **kwargs):
+        torch.cuda.synchronize()
+        start = time.perf_counter()
+        image = pipe(num_inference_steps=40, generator=gen(seed), **kwargs).images[0]
+        torch.cuda.synchronize()
+        image.save(out / f"{name}.png")
+        return time.perf_counter() - start
+
+    run("warmup", 0, prompt="warmup", height=args.size, width=args.size)
+    run("warmup", 1, prompt="warmup", image=Image.open(out / "warmup.png"), output_resolution=args.size)
+    torch.cuda.reset_peak_memory_stats()
+    times = {"t2i": [], "edit": []}
+    for i, prompt in enumerate(PROMPTS):
+        times["t2i"].append(run(f"t2i{i}", 1000 + i, prompt=prompt, height=args.size, width=args.size))
+    for name, image, prompt in edit_items(args.edits):
+        times["edit"].append(run(name, 2000, prompt=prompt, image=image, output_resolution=args.size))
+    report = {
+        "tag": args.tag,
+        "size": args.size,
+        "t2i_s": round(sum(times["t2i"]) / len(times["t2i"]), 3),
+        "edit_s": round(sum(times["edit"]) / len(times["edit"]), 3),
+        "peak_vram_gib": round(torch.cuda.max_memory_allocated() / 2**30, 2),
+    }
+    if args.baseline:
+        import lpips
+        import numpy as np
+
+        metric = lpips.LPIPS(net="alex", verbose=False)
+        load = lambda p: torch.from_numpy(np.array(Image.open(p).convert("RGB"))).permute(2, 0, 1)[None].float() / 127.5 - 1  # noqa: E731
+        scores = {
+            p.stem: round(metric(load(Path(args.baseline) / p.name), load(p)).item(), 4)
+            for p in sorted(out.glob("*.png"))
+            if p.stem != "warmup" and (Path(args.baseline) / p.name).exists()
+        }
+        report["lpips"] = scores
+        report["lpips_t2i_mean"] = round(sum(v for k, v in scores.items() if k.startswith("t2i")) / len(PROMPTS), 4)
+        report["lpips_edit_mean"] = round(sum(v for k, v in scores.items() if not k.startswith("t2i")) / (args.edits + 1), 4)
+    (out / "report.json").write_text(json.dumps(report, indent=1))
+    print(json.dumps(report))
+
+
+if __name__ == "__main__":
+    main()

tools/modeling_nunchaku_qwen3vl.py

@@ -0,0 +1,47 @@
+"""Qwen3-VL text encoder with SVDQuant NVFP4 (Nunchaku Lite) linears.
+
+Loaded by Diffusers as a custom pipeline component (`trust_remote_code=True`).
+`config.json` holds the compact Nunchaku Lite config under `nunchaku_lite`; the
+listed linears are swapped for Diffusers' own `SVDQW4A4Linear` before the
+quantized state dict is loaded, so inference uses the same NVFP4 kernels as the
+transformer. Everything else (vision tower, embeddings, the BF16 layers) loads
+as in the stock model.
+"""
+
+from pathlib import Path
+
+import torch
+from accelerate import init_empty_weights
+from huggingface_hub import snapshot_download
+
+from .nunchaku_kernels import KERNELS  # noqa: F401  (sets up the kernels before Diffusers imports them)
+
+# isort: off  -- must stay below `.nunchaku_kernels`, which prepares the kernels this import loads
+from diffusers.quantizers.nunchaku.utils import replace_with_nunchaku_linear  # noqa: E402
+from safetensors.torch import load_file  # noqa: E402
+from transformers import Qwen3VLForConditionalGeneration  # noqa: E402
+# isort: on
+
+
+class NunchakuQwen3VLForConditionalGeneration(Qwen3VLForConditionalGeneration):
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, subfolder: str = "", **kwargs):
+        path = Path(pretrained_model_name_or_path) / subfolder
+        if not path.is_dir():
+            path = Path(snapshot_download(str(pretrained_model_name_or_path), allow_patterns=[f"{subfolder}/*"])) / subfolder
+        dtype = kwargs.get("dtype") or kwargs.get("torch_dtype") or torch.bfloat16
+        if dtype == "auto":
+            dtype = torch.bfloat16
+
+        config = cls.config_class.from_pretrained(path)
+        with init_empty_weights():  # parameters on meta, buffers (e.g. rotary tables) materialized
+            model = cls(config)
+        replace_with_nunchaku_linear(model, config.nunchaku_lite, dtype)
+        # strict=False only because non-persistent buffers are absent; both checks below stay strict
+        result = model.load_state_dict(load_file(path / "model.safetensors"), strict=False, assign=True)
+        if result.unexpected_keys:
+            raise ValueError(f"Checkpoint has {len(result.unexpected_keys)} unexpected tensors, e.g. {result.unexpected_keys[:3]}")
+        missing = [name for name, p in model.named_parameters() if p.device.type == "meta"]
+        if missing:
+            raise ValueError(f"Checkpoint is missing {len(missing)} parameters, e.g. {missing[:3]}")
+        return model.eval()

tools/modeling_nunchaku_qwenimage21.py

@@ -0,0 +1,14 @@
+"""Qwen-Image-2.1 transformer, unchanged, loaded as a custom component.
+
+Its only job is importing `nunchaku_kernels`, so the NVFP4 kernels are set up even when
+Diffusers loads this component before the text encoder. The Nunchaku Lite quantization
+itself is handled by Diffusers' built-in quantizer from `config.json`.
+"""
+
+from diffusers import QwenImage21Transformer2DModel
+
+from .nunchaku_kernels import KERNELS  # noqa: F401
+
+
+class NunchakuQwenImage21Transformer2DModel(QwenImage21Transformer2DModel):
+    pass

tools/nunchaku_kernels.py

@@ -0,0 +1,19 @@
+# Shared by text_encoder/ and transformer/: whichever component Diffusers loads first sets this up.
+#
+# Diffusers fetches the Nunchaku Lite kernels from `rootonchair/nunchaku-lite-kernels`, which is
+# no longer downloadable. Loading this repo with trust_remote_code=True already runs this code,
+# so point that kernel name at the rebuilt copy before Diffusers imports its Nunchaku utilities.
+import os
+
+from huggingface_hub import snapshot_download
+
+KERNELS = "rootonchair/nunchaku-lite-kernels"
+if KERNELS not in os.environ.get("LOCAL_KERNELS", ""):
+    local = f"{KERNELS}={snapshot_download('joseplcam/nunchaku-lite-kernels')}"
+    os.environ["LOCAL_KERNELS"] = ":".join(filter(None, [os.environ.get("LOCAL_KERNELS"), local]))
+if "DIFFUSERS_TRUST_REMOTE_KERNELS" not in os.environ:
+    # Diffusers read this variable into a constant at import time, so set both.
+    import diffusers.utils.constants
+
+    os.environ["DIFFUSERS_TRUST_REMOTE_KERNELS"] = "true"
+    diffusers.utils.constants.DIFFUSERS_TRUST_REMOTE_KERNELS = True

tools/package.py

@@ -0,0 +1,81 @@
+"""Package the quantized DiT and text encoder into one Diffusers pipeline folder.
+
+- transformer/: Nunchaku Lite NVFP4 checkpoint (loaded by Diffusers' built-in quantizer)
+- text_encoder/: Nunchaku Lite NVFP4 checkpoint + custom component class (trust_remote_code)
+- vae/: official weights stored in BF16
+- processor/, scheduler/, model_index.json, LICENSE: from Qwen/Qwen-Image-2.1
+
+Usage:
+    python tools/qwen21_nvfp4/package.py --dit <dit.safetensors> --text-encoder <te.safetensors> --out <dir>
+"""
+
+import argparse
+import json
+import os
+import shutil
+import sys
+from pathlib import Path
+
+import torch
+from diffusers import AutoencoderKLQwenImage21
+
+# examples/ is not part of the installed package, so a checkout of the repo is needed
+sys.path.insert(0, os.environ.get("DIFFUSE_COMPRESSOR", str(Path.home() / "src" / "diffuse-compressor")))
+from examples.convert_nunchaku_lite_diffusers import (  # noqa: E402
+    build_diffusers_quantization_config,
+    package_diffusers_pipeline,
+)
+
+HERE = Path(__file__).parent
+# component -> (module file, class); both import nunchaku_kernels, which ships next to each
+CUSTOM_COMPONENTS = {
+    "text_encoder": ("modeling_nunchaku_qwen3vl", "NunchakuQwen3VLForConditionalGeneration"),
+    "transformer": ("modeling_nunchaku_qwenimage21", "NunchakuQwenImage21Transformer2DModel"),
+}
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dit", required=True)
+    parser.add_argument("--text-encoder", required=True)
+    parser.add_argument("--out", required=True)
+    parser.add_argument("--model-id", default="Qwen/Qwen-Image-2.1")
+    args = parser.parse_args()
+
+    out = package_diffusers_pipeline(args.dit, args.model_id, args.out, compute_dtype="bfloat16")
+
+    te = out / "text_encoder"
+    for dense in [*te.glob("*.safetensors"), *te.glob("*.safetensors.index.json")]:
+        dense.unlink()
+    shutil.copy2(args.text_encoder, te / "model.safetensors")
+    config = json.loads((te / "config.json").read_text())
+    config["nunchaku_lite"] = build_diffusers_quantization_config(args.text_encoder, compute_dtype="bfloat16")
+    (te / "config.json").write_text(json.dumps(config, indent=2, sort_keys=True) + "\n")
+
+    index = json.loads((out / "model_index.json").read_text())
+    for component, (module, cls) in CUSTOM_COMPONENTS.items():
+        shutil.copy2(HERE / f"{module}.py", out / component / f"{module}.py")
+        shutil.copy2(HERE / "nunchaku_kernels.py", out / component / "nunchaku_kernels.py")
+        index[component] = [module, cls]
+    (out / "model_index.json").write_text(json.dumps(index, indent=2) + "\n")
+
+    vae = AutoencoderKLQwenImage21.from_pretrained(out / "vae", dtype=torch.bfloat16)
+    shutil.rmtree(out / "vae")
+    vae.save_pretrained(out / "vae")
+    vae_config = json.loads((out / "vae" / "config.json").read_text())
+    vae_config.pop("_name_or_path", None)  # a local path, meaningless on the Hub
+    (out / "vae" / "config.json").write_text(json.dumps(vae_config, indent=2, sort_keys=True) + "\n")
+
+    for leftover in ("assets", ".gitattributes"):
+        path = out / leftover
+        shutil.rmtree(path) if path.is_dir() else path.unlink(missing_ok=True)
+    for doc in ("README.md", "NOTICE"):  # model card and license notice
+        shutil.copy2(HERE / doc, out / doc)
+    (out / "tools").mkdir()
+    for script in sorted(HERE.glob("*.py")):  # the recipe, for reproducibility
+        shutil.copy2(script, out / "tools" / script.name)
+    print(out)
+
+
+if __name__ == "__main__":
+    main()

tools/parrot_test.py

@@ -0,0 +1,69 @@
+"""Hard-edit check: "Turn the parrots blue" over 8 seeds at 512x512, per text encoder variant.
+
+4-bit text encoders tend to lose this recolor (the model keeps the reference colors), so it
+separates encoder variants that LPIPS alone does not. Each row of the output grid is one
+variant; count the tiles with at least one blue parrot.
+
+Usage:
+    python tools/qwen21_nvfp4/parrot_test.py <packaged pipeline> <out.jpg> bf16 <encoder checkpoint> ...
+"""
+
+import io
+import json
+import shutil
+import sys
+import tempfile
+import urllib.request
+from pathlib import Path
+
+import torch
+from diffusers import DiffusionPipeline
+from PIL import Image
+from transformers import Qwen3VLForConditionalGeneration
+
+sys.path.insert(0, str(Path(__file__).parent))
+from package import build_diffusers_quantization_config  # noqa: E402
+
+PARROTS = "https://r0k.us/graphics/kodak/kodak/kodim23.png"
+
+
+def load_encoder(encoder_cls, pipeline_dir: Path, variant: str):
+    if variant == "bf16":
+        return Qwen3VLForConditionalGeneration.from_pretrained("Qwen/Qwen-Image-2.1", subfolder="text_encoder", dtype=torch.bfloat16)
+    with tempfile.TemporaryDirectory() as tmp:  # the packaged text_encoder folder with this checkpoint swapped in
+        folder = Path(tmp)
+        for name in ("generation_config.json",):
+            shutil.copy2(pipeline_dir / "text_encoder" / name, folder / name)
+        config = json.loads((pipeline_dir / "text_encoder" / "config.json").read_text())
+        config["nunchaku_lite"] = build_diffusers_quantization_config(variant, compute_dtype="bfloat16")
+        (folder / "config.json").write_text(json.dumps(config))
+        (folder / "model.safetensors").symlink_to(Path(variant).resolve())
+        return encoder_cls.from_pretrained(folder, dtype=torch.bfloat16)
+
+
+def main():
+    pipeline_dir, out, variants = Path(sys.argv[1]), sys.argv[2], sys.argv[3:]
+    pipe = DiffusionPipeline.from_pretrained(pipeline_dir, dtype=torch.bfloat16, trust_remote_code=True).to("cuda")
+    pipe.set_progress_bar_config(disable=True)
+    parrots = Image.open(io.BytesIO(urllib.request.urlopen(PARROTS).read())).convert("RGB")
+    encoder_cls = type(pipe.text_encoder)
+    rows = []
+    for variant in variants:
+        pipe.text_encoder = None
+        torch.cuda.empty_cache()
+        pipe.text_encoder = load_encoder(encoder_cls, pipeline_dir, variant).to("cuda")
+        rows.append([
+            pipe(prompt="Turn the parrots blue", image=parrots, output_resolution=512, num_inference_steps=40,
+                 generator=torch.Generator("cuda").manual_seed(seed)).images[0].convert("RGB").resize((128, 128))
+            for seed in range(1, 9)
+        ])
+        print(Path(variant).name, "done", flush=True)
+    grid = Image.new("RGB", (128 * 8, 128 * len(rows)))
+    for y, row in enumerate(rows):
+        for x, tile in enumerate(row):
+            grid.paste(tile, (128 * x, 128 * y))
+    grid.save(out)
+
+
+if __name__ == "__main__":
+    main()

tools/quantize_dit.py

@@ -0,0 +1,144 @@
+"""SVDQuant + GPTQ NVFP4 quantization of the Qwen-Image-2.1 DiT with diffuse-compressor.
+
+Calibration mixes text-to-image prompts and image edits (NHR-Edit train split),
+mostly at 512x512 with some 1024x1024, so both workloads shape the scales.
+The first two and last two blocks and the global modulation stay in BF16.
+
+Usage (from the repo root, envs/base_final interpreter):
+    python tools/qwen21_nvfp4/quantize_dit.py --num-samples 128 --rank 32
+"""
+
+import argparse
+import os
+import random
+import sys
+from dataclasses import replace
+from pathlib import Path
+
+import torch
+from diffusers import QwenImage21Pipeline
+
+# examples/ is not part of the installed package, so a checkout of the repo is needed
+sys.path.insert(0, os.environ.get("DIFFUSE_COMPRESSOR", str(Path.home() / "src" / "diffuse-compressor")))
+
+from diffuse_compressor import (  # noqa: E402
+    CalibrationSpec,
+    ExportSpec,
+    LoggingConfig,
+    QuantizationCacheSpec,
+    quantize_and_export,
+)
+from diffuse_compressor.config import GptqSpec  # noqa: E402
+from examples.image_to_image.quantize_hf import scan_linear_targets  # noqa: E402
+from examples.text_to_image.utils import (  # noqa: E402
+    make_generator,
+    save_diffusers_images,
+    standard_prompt_records,
+    svdquant_spec,
+)
+
+SKIP = [
+    "transformer_blocks.0.*",
+    "transformer_blocks.1.*",
+    "transformer_blocks.30.*",
+    "transformer_blocks.31.*",
+    "modulation*",
+]
+
+
+def records(num_samples: int, edit_dataset: str) -> list[dict]:
+    import datasets
+
+    half = num_samples // 2
+    rng = random.Random(0)
+    # 3 of every 4 samples at 512x512 (the working resolution), the rest at 1024x1024
+    size = lambda i: 1024 if i % 4 == 3 else 512  # noqa: E731
+
+    t2i = [
+        {"filename": r["filename"], "prompt": r["prompt"], "seed": r["seed"], "size": size(i)}
+        for i, r in enumerate(standard_prompt_records(num_samples - half))
+    ]
+    data = datasets.load_dataset(edit_dataset, split="train")
+    edits = [
+        {
+            "filename": f"edit-{j}",
+            "prompt": data[j]["edit_instruction"],
+            "image": data[j]["source"].convert("RGB"),
+            "seed": 100_000 + j,
+            "size": size(i),
+        }
+        for i, j in enumerate(sorted(rng.sample(range(len(data)), half)))
+    ]
+    return t2i + edits
+
+
+def forward_fn(pipe, steps: int, device: str):
+    def forward(sample: dict):
+        kwargs = {
+            "prompt": sample["prompt"],
+            "num_inference_steps": steps,
+            "generator": make_generator(sample["seed"], device=device),
+            # the prefix KV cache is a stateful object that calibration replay cannot move or
+            # re-run; without it the prefix tokens pass through the same linears every step
+            "use_kv_cache": False,
+        }
+        if sample.get("image") is not None:
+            kwargs.update(image=sample["image"], output_resolution=sample["size"])
+        else:
+            kwargs.update(height=sample["size"], width=sample["size"])
+        return pipe(**kwargs)
+
+    return forward
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-id", default="Qwen/Qwen-Image-2.1")
+    parser.add_argument("--edit-dataset", default="VyoJ/NHR-Edit-Change_Only")
+    parser.add_argument("--num-samples", type=int, default=128)
+    parser.add_argument("--steps", type=int, default=20)
+    parser.add_argument("--rank", type=int, default=32)
+    parser.add_argument("--no-gptq", action="store_true")
+    parser.add_argument("--out-dir", default="outputs/qwen21_nvfp4")
+    parser.add_argument("--cache-mode", choices=("reuse", "refresh", "disabled"), default="reuse")
+    parser.add_argument("--inspect", action="store_true")
+    args = parser.parse_args()
+
+    device = "cuda"
+    pipe = QwenImage21Pipeline.from_pretrained(args.model_id, torch_dtype=torch.bfloat16).to(device)
+    pipe.set_progress_bar_config(disable=True)
+    scan = scan_linear_targets(pipe.transformer, precision="nvfp4", rank=args.rank, skip=SKIP)
+    print(scan.format_text().splitlines()[0])
+    if args.inspect:
+        print(scan.format_text())
+        return
+
+    tag = f"svdq-nvfp4_r{args.rank}{'' if args.no_gptq else '-gptq'}-qwen-image-2.1"
+    out = Path(args.out_dir)
+    cache = out / "calibration" / tag
+    spec = replace(svdquant_spec("nvfp4", compute_device=device), rank=args.rank, gptq=GptqSpec(enabled=not args.no_gptq))
+    quantize_and_export(
+        pipe.transformer,
+        spec,
+        scan.target_config,
+        CalibrationSpec(
+            samples=records(args.num_samples, args.edit_dataset),
+            num_samples=args.num_samples,
+            cache_num_samples=args.num_samples,
+            batch_size=1,
+            cache_dir=cache / "inputs",
+            cache_mode=args.cache_mode,
+            forward_fn=forward_fn(pipe, args.steps, device),
+            max_rows_per_target=4096,
+            artifact_cache=QuantizationCacheSpec(cache / "artifacts", args.cache_mode),
+            output_dir=cache / "inputs" / "samples",
+            output_save_fn=save_diffusers_images,
+            scope_capture_mode="all_targets",
+        ),
+        ExportSpec(output=out / "checkpoints" / f"{tag}.safetensors"),
+        LoggingConfig(name=tag),
+    )
+
+
+if __name__ == "__main__":
+    main()

tools/quantize_text_encoder.py

@@ -0,0 +1,121 @@
+"""SVDQuant + GPTQ NVFP4 quantization of the Qwen-Image-2.1 text encoder (Qwen3-VL 8B).
+
+Only the language-model linears are quantized. The vision tower, embeddings,
+lm_head (unused by the pipeline) and the first/last `--edge-layers` decoder
+layers stay BF16: the pipeline feeds the last decoder layer's output straight
+into the DiT. Calibration uses the same text-to-image and edit records as the DiT, run
+through the full pipeline so condition images are resized exactly as in use.
+
+Usage (from the repo root, envs/base_final interpreter):
+    python tools/qwen21_nvfp4/quantize_text_encoder.py --num-samples 128 --rank 32
+"""
+
+import argparse
+import sys
+from dataclasses import replace
+from pathlib import Path
+
+import torch
+from diffusers import QwenImage21Pipeline
+
+sys.path.insert(0, str(Path(__file__).parent))
+from quantize_dit import records  # noqa: E402  (also puts diffuse-compressor on sys.path)
+
+from diffuse_compressor import (  # noqa: E402
+    CalibrationSpec,
+    ExportSpec,
+    LoggingConfig,
+    QuantizationCacheSpec,
+    quantize_and_export,
+)
+from diffuse_compressor.calibration import replay  # noqa: E402
+from diffuse_compressor.config import GptqSpec  # noqa: E402
+from examples.image_to_image.quantize_hf import scan_linear_targets  # noqa: E402
+from examples.text_to_image.utils import make_generator, svdquant_spec  # noqa: E402
+
+# Full-model replay CPU-offloads the model through Accelerate, which puts the 16 GB encoder in
+# host RAM and trips the 90% RAM guard on a 32 GB machine. The GPU has room for it, and
+# returning False makes the library fall back to `model.to(device)`.
+replay._accelerate_cpu_offload_for_full_replay = lambda *args, **kwargs: False
+
+NUM_LAYERS = 36
+
+
+def skip_patterns(edge_layers: int, skip_attention: bool) -> list[str]:
+    """Vision tower, lm_head, the first/last `edge_layers` decoder layers and optionally all attention stay BF16."""
+    edges = [*range(edge_layers), *range(NUM_LAYERS - edge_layers, NUM_LAYERS)]
+    patterns = ["model.visual.*", "lm_head", *(f"model.language_model.layers.{i}.*" for i in edges)]
+    return patterns + (["model.language_model.layers.*.self_attn.*"] if skip_attention else [])
+
+
+def forward_fn(pipe, device: str):
+    def forward(sample: dict):
+        kwargs = {
+            "prompt": sample["prompt"],
+            "num_inference_steps": 1,  # the encoder runs once per call; one DiT step is enough
+            "output_type": "latent",
+            "generator": make_generator(sample["seed"], device=device),
+        }
+        if sample.get("image") is not None:
+            kwargs.update(image=sample["image"], output_resolution=sample["size"])
+        else:
+            kwargs.update(height=sample["size"], width=sample["size"])
+        return pipe(**kwargs)
+
+    return forward
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-id", default="Qwen/Qwen-Image-2.1")
+    parser.add_argument("--edit-dataset", default="VyoJ/NHR-Edit-Change_Only")
+    parser.add_argument("--num-samples", type=int, default=128)
+    parser.add_argument("--rank", type=int, default=32)
+    parser.add_argument("--edge-layers", type=int, default=2, help="first/last decoder layers kept in BF16")
+    parser.add_argument("--skip-attention", action="store_true", help="keep all attention projections in BF16")
+    parser.add_argument("--no-gptq", action="store_true")
+    parser.add_argument("--out-dir", default="outputs/qwen21_nvfp4")
+    parser.add_argument("--cache-mode", choices=("reuse", "refresh", "disabled"), default="reuse")
+    parser.add_argument("--inspect", action="store_true")
+    args = parser.parse_args()
+
+    device = "cuda"
+    pipe = QwenImage21Pipeline.from_pretrained(args.model_id, dtype=torch.bfloat16).to(device)
+    pipe.set_progress_bar_config(disable=True)
+    encoder = pipe.text_encoder
+    scan = scan_linear_targets(
+        encoder, precision="nvfp4", rank=args.rank, skip=skip_patterns(args.edge_layers, args.skip_attention)
+    )
+    print(scan.format_text().splitlines()[0])
+    if args.inspect:
+        print(scan.format_text())
+        return
+
+    variant = f"r{args.rank}_e{args.edge_layers}{'_mlp' if args.skip_attention else ''}{'' if args.no_gptq else '-gptq'}"
+    tag = f"svdq-nvfp4_{variant}-qwen-image-2.1-text-encoder"
+    out = Path(args.out_dir)
+    cache = out / "calibration" / tag
+    spec = replace(svdquant_spec("nvfp4", compute_device=device), rank=args.rank, gptq=GptqSpec(enabled=not args.no_gptq))
+    quantize_and_export(
+        encoder,
+        spec,
+        scan.target_config,
+        CalibrationSpec(
+            samples=records(args.num_samples, args.edit_dataset),
+            num_samples=args.num_samples,
+            cache_num_samples=args.num_samples,
+            batch_size=1,
+            cache_dir=cache / "inputs",
+            cache_mode=args.cache_mode,
+            forward_fn=forward_fn(pipe, device),
+            max_rows_per_target=4096,
+            artifact_cache=QuantizationCacheSpec(cache / "artifacts", args.cache_mode),
+            scope_capture_mode="all_targets",
+        ),
+        ExportSpec(output=out / "checkpoints" / f"{tag}.safetensors"),
+        LoggingConfig(name=tag),
+    )
+
+
+if __name__ == "__main__":
+    main()

transformer/config.json

@@ -1,7 +1,6 @@
 {
   "_class_name": "QwenImage21Transformer2DModel",
-  "_diffusers_version": "0.41.0.dev0",
-  "_name_or_path": "/mnt/hf_cache/hub/models--Qwen--Qwen-Image-2.1/snapshots/b3179ad355be050328e483a9dfdd9e60cd62adfa/transformer",
+  "_diffusers_version": "0.37.0.dev0",
   "attention_head_dim": 128,
   "axes_dims_rope": [
     16,
@@ -18,43 +17,210 @@
   "out_channels": 64,
   "patch_size": 1,
   "quantization_config": {
-    "config_groups": {
-      "group_0": {
-        "input_activations": {
-          "dynamic": false,
-          "num_bits": 4,
-          "type": "float",
-          "group_size": 16
-        },
-        "weights": {
-          "dynamic": false,
-          "num_bits": 4,
-          "type": "float",
-          "group_size": 16
-        },
-        "targets": [
-          "Linear"
-        ]
-      }
-    },
-    "ignore": [
-      "*img_in",
-      "*modulation*",
-      "*norm_out*",
-      "*proj_out",
-      "*time_text_embed*",
-      "*transformer_blocks.0.*",
-      "*transformer_blocks.1.*",
-      "*transformer_blocks.30.*",
-      "*transformer_blocks.31.*",
-      "*txt_in*"
-    ],
-    "quant_algo": "NVFP4",
-    "producer": {
-      "name": "modelopt",
-      "version": "0.46.0rc2"
-    },
-    "quant_method": "modelopt",
-    "quant_type": "nvfp4"
+    "compute_dtype": "bfloat16",
+    "quant_method": "nunchaku_lite",
+    "svdq_w4a4": {
+      "group_size": 16,
+      "precision": "nvfp4",
+      "rank": 32,
+      "targets": [
+        "transformer_blocks.2.attn.to_q",
+        "transformer_blocks.2.attn.to_k",
+        "transformer_blocks.2.attn.to_v",
+        "transformer_blocks.2.attn.to_out.0",
+        "transformer_blocks.2.img_mlp.proj",
+        "transformer_blocks.2.img_mlp.out",
+        "transformer_blocks.2.img_mlp.gate_layer",
+        "transformer_blocks.3.attn.to_q",
+        "transformer_blocks.3.attn.to_k",
+        "transformer_blocks.3.attn.to_v",
+        "transformer_blocks.3.attn.to_out.0",
+        "transformer_blocks.3.img_mlp.proj",
+        "transformer_blocks.3.img_mlp.out",
+        "transformer_blocks.3.img_mlp.gate_layer",
+        "transformer_blocks.4.attn.to_q",
+        "transformer_blocks.4.attn.to_k",
+        "transformer_blocks.4.attn.to_v",
+        "transformer_blocks.4.attn.to_out.0",
+        "transformer_blocks.4.img_mlp.proj",
+        "transformer_blocks.4.img_mlp.out",
+        "transformer_blocks.4.img_mlp.gate_layer",
+        "transformer_blocks.5.attn.to_q",
+        "transformer_blocks.5.attn.to_k",
+        "transformer_blocks.5.attn.to_v",
+        "transformer_blocks.5.attn.to_out.0",
+        "transformer_blocks.5.img_mlp.proj",
+        "transformer_blocks.5.img_mlp.out",
+        "transformer_blocks.5.img_mlp.gate_layer",
+        "transformer_blocks.6.attn.to_q",
+        "transformer_blocks.6.attn.to_k",
+        "transformer_blocks.6.attn.to_v",
+        "transformer_blocks.6.attn.to_out.0",
+        "transformer_blocks.6.img_mlp.proj",
+        "transformer_blocks.6.img_mlp.out",
+        "transformer_blocks.6.img_mlp.gate_layer",
+        "transformer_blocks.7.attn.to_q",
+        "transformer_blocks.7.attn.to_k",
+        "transformer_blocks.7.attn.to_v",
+        "transformer_blocks.7.attn.to_out.0",
+        "transformer_blocks.7.img_mlp.proj",
+        "transformer_blocks.7.img_mlp.out",
+        "transformer_blocks.7.img_mlp.gate_layer",
+        "transformer_blocks.8.attn.to_q",
+        "transformer_blocks.8.attn.to_k",
+        "transformer_blocks.8.attn.to_v",
+        "transformer_blocks.8.attn.to_out.0",
+        "transformer_blocks.8.img_mlp.proj",
+        "transformer_blocks.8.img_mlp.out",
+        "transformer_blocks.8.img_mlp.gate_layer",
+        "transformer_blocks.9.attn.to_q",
+        "transformer_blocks.9.attn.to_k",
+        "transformer_blocks.9.attn.to_v",
+        "transformer_blocks.9.attn.to_out.0",
+        "transformer_blocks.9.img_mlp.proj",
+        "transformer_blocks.9.img_mlp.out",
+        "transformer_blocks.9.img_mlp.gate_layer",
+        "transformer_blocks.10.attn.to_q",
+        "transformer_blocks.10.attn.to_k",
+        "transformer_blocks.10.attn.to_v",
+        "transformer_blocks.10.attn.to_out.0",
+        "transformer_blocks.10.img_mlp.proj",
+        "transformer_blocks.10.img_mlp.out",
+        "transformer_blocks.10.img_mlp.gate_layer",
+        "transformer_blocks.11.attn.to_q",
+        "transformer_blocks.11.attn.to_k",
+        "transformer_blocks.11.attn.to_v",
+        "transformer_blocks.11.attn.to_out.0",
+        "transformer_blocks.11.img_mlp.proj",
+        "transformer_blocks.11.img_mlp.out",
+        "transformer_blocks.11.img_mlp.gate_layer",
+        "transformer_blocks.12.attn.to_q",
+        "transformer_blocks.12.attn.to_k",
+        "transformer_blocks.12.attn.to_v",
+        "transformer_blocks.12.attn.to_out.0",
+        "transformer_blocks.12.img_mlp.proj",
+        "transformer_blocks.12.img_mlp.out",
+        "transformer_blocks.12.img_mlp.gate_layer",
+        "transformer_blocks.13.attn.to_q",
+        "transformer_blocks.13.attn.to_k",
+        "transformer_blocks.13.attn.to_v",
+        "transformer_blocks.13.attn.to_out.0",
+        "transformer_blocks.13.img_mlp.proj",
+        "transformer_blocks.13.img_mlp.out",
+        "transformer_blocks.13.img_mlp.gate_layer",
+        "transformer_blocks.14.attn.to_q",
+        "transformer_blocks.14.attn.to_k",
+        "transformer_blocks.14.attn.to_v",
+        "transformer_blocks.14.attn.to_out.0",
+        "transformer_blocks.14.img_mlp.proj",
+        "transformer_blocks.14.img_mlp.out",
+        "transformer_blocks.14.img_mlp.gate_layer",
+        "transformer_blocks.15.attn.to_q",
+        "transformer_blocks.15.attn.to_k",
+        "transformer_blocks.15.attn.to_v",
+        "transformer_blocks.15.attn.to_out.0",
+        "transformer_blocks.15.img_mlp.proj",
+        "transformer_blocks.15.img_mlp.out",
+        "transformer_blocks.15.img_mlp.gate_layer",
+        "transformer_blocks.16.attn.to_q",
+        "transformer_blocks.16.attn.to_k",
+        "transformer_blocks.16.attn.to_v",
+        "transformer_blocks.16.attn.to_out.0",
+        "transformer_blocks.16.img_mlp.proj",
+        "transformer_blocks.16.img_mlp.out",
+        "transformer_blocks.16.img_mlp.gate_layer",
+        "transformer_blocks.17.attn.to_q",
+        "transformer_blocks.17.attn.to_k",
+        "transformer_blocks.17.attn.to_v",
+        "transformer_blocks.17.attn.to_out.0",
+        "transformer_blocks.17.img_mlp.proj",
+        "transformer_blocks.17.img_mlp.out",
+        "transformer_blocks.17.img_mlp.gate_layer",
+        "transformer_blocks.18.attn.to_q",
+        "transformer_blocks.18.attn.to_k",
+        "transformer_blocks.18.attn.to_v",
+        "transformer_blocks.18.attn.to_out.0",
+        "transformer_blocks.18.img_mlp.proj",
+        "transformer_blocks.18.img_mlp.out",
+        "transformer_blocks.18.img_mlp.gate_layer",
+        "transformer_blocks.19.attn.to_q",
+        "transformer_blocks.19.attn.to_k",
+        "transformer_blocks.19.attn.to_v",
+        "transformer_blocks.19.attn.to_out.0",
+        "transformer_blocks.19.img_mlp.proj",
+        "transformer_blocks.19.img_mlp.out",
+        "transformer_blocks.19.img_mlp.gate_layer",
+        "transformer_blocks.20.attn.to_q",
+        "transformer_blocks.20.attn.to_k",
+        "transformer_blocks.20.attn.to_v",
+        "transformer_blocks.20.attn.to_out.0",
+        "transformer_blocks.20.img_mlp.proj",
+        "transformer_blocks.20.img_mlp.out",
+        "transformer_blocks.20.img_mlp.gate_layer",
+        "transformer_blocks.21.attn.to_q",
+        "transformer_blocks.21.attn.to_k",
+        "transformer_blocks.21.attn.to_v",
+        "transformer_blocks.21.attn.to_out.0",
+        "transformer_blocks.21.img_mlp.proj",
+        "transformer_blocks.21.img_mlp.out",
+        "transformer_blocks.21.img_mlp.gate_layer",
+        "transformer_blocks.22.attn.to_q",
+        "transformer_blocks.22.attn.to_k",
+        "transformer_blocks.22.attn.to_v",
+        "transformer_blocks.22.attn.to_out.0",
+        "transformer_blocks.22.img_mlp.proj",
+        "transformer_blocks.22.img_mlp.out",
+        "transformer_blocks.22.img_mlp.gate_layer",
+        "transformer_blocks.23.attn.to_q",
+        "transformer_blocks.23.attn.to_k",
+        "transformer_blocks.23.attn.to_v",
+        "transformer_blocks.23.attn.to_out.0",
+        "transformer_blocks.23.img_mlp.proj",
+        "transformer_blocks.23.img_mlp.out",
+        "transformer_blocks.23.img_mlp.gate_layer",
+        "transformer_blocks.24.attn.to_q",
+        "transformer_blocks.24.attn.to_k",
+        "transformer_blocks.24.attn.to_v",
+        "transformer_blocks.24.attn.to_out.0",
+        "transformer_blocks.24.img_mlp.proj",
+        "transformer_blocks.24.img_mlp.out",
+        "transformer_blocks.24.img_mlp.gate_layer",
+        "transformer_blocks.25.attn.to_q",
+        "transformer_blocks.25.attn.to_k",
+        "transformer_blocks.25.attn.to_v",
+        "transformer_blocks.25.attn.to_out.0",
+        "transformer_blocks.25.img_mlp.proj",
+        "transformer_blocks.25.img_mlp.out",
+        "transformer_blocks.25.img_mlp.gate_layer",
+        "transformer_blocks.26.attn.to_q",
+        "transformer_blocks.26.attn.to_k",
+        "transformer_blocks.26.attn.to_v",
+        "transformer_blocks.26.attn.to_out.0",
+        "transformer_blocks.26.img_mlp.proj",
+        "transformer_blocks.26.img_mlp.out",
+        "transformer_blocks.26.img_mlp.gate_layer",
+        "transformer_blocks.27.attn.to_q",
+        "transformer_blocks.27.attn.to_k",
+        "transformer_blocks.27.attn.to_v",
+        "transformer_blocks.27.attn.to_out.0",
+        "transformer_blocks.27.img_mlp.proj",
+        "transformer_blocks.27.img_mlp.out",
+        "transformer_blocks.27.img_mlp.gate_layer",
+        "transformer_blocks.28.attn.to_q",
+        "transformer_blocks.28.attn.to_k",
+        "transformer_blocks.28.attn.to_v",
+        "transformer_blocks.28.attn.to_out.0",
+        "transformer_blocks.28.img_mlp.proj",
+        "transformer_blocks.28.img_mlp.out",
+        "transformer_blocks.28.img_mlp.gate_layer",
+        "transformer_blocks.29.attn.to_q",
+        "transformer_blocks.29.attn.to_k",
+        "transformer_blocks.29.attn.to_v",
+        "transformer_blocks.29.attn.to_out.0",
+        "transformer_blocks.29.img_mlp.proj",
+        "transformer_blocks.29.img_mlp.out",
+        "transformer_blocks.29.img_mlp.gate_layer"
+      ]
+    }
   }
-}
+}

transformer/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a6c3080b764344729ef95846faa3a0a01eb36bdef93f785ed07f4ac5751ad76a
-size 5451672424
+oid sha256:0304fc3e84b13663b439601380066754ef9b92025e159275abe163bc5f1ceeb4
+size 5603299808

transformer/modeling_nunchaku_qwenimage21.py

@@ -0,0 +1,14 @@
+"""Qwen-Image-2.1 transformer, unchanged, loaded as a custom component.
+
+Its only job is importing `nunchaku_kernels`, so the NVFP4 kernels are set up even when
+Diffusers loads this component before the text encoder. The Nunchaku Lite quantization
+itself is handled by Diffusers' built-in quantizer from `config.json`.
+"""
+
+from diffusers import QwenImage21Transformer2DModel
+
+from .nunchaku_kernels import KERNELS  # noqa: F401
+
+
+class NunchakuQwenImage21Transformer2DModel(QwenImage21Transformer2DModel):
+    pass

transformer/nunchaku_kernels.py

@@ -0,0 +1,19 @@
+# Shared by text_encoder/ and transformer/: whichever component Diffusers loads first sets this up.
+#
+# Diffusers fetches the Nunchaku Lite kernels from `rootonchair/nunchaku-lite-kernels`, which is
+# no longer downloadable. Loading this repo with trust_remote_code=True already runs this code,
+# so point that kernel name at the rebuilt copy before Diffusers imports its Nunchaku utilities.
+import os
+
+from huggingface_hub import snapshot_download
+
+KERNELS = "rootonchair/nunchaku-lite-kernels"
+if KERNELS not in os.environ.get("LOCAL_KERNELS", ""):
+    local = f"{KERNELS}={snapshot_download('joseplcam/nunchaku-lite-kernels')}"
+    os.environ["LOCAL_KERNELS"] = ":".join(filter(None, [os.environ.get("LOCAL_KERNELS"), local]))
+if "DIFFUSERS_TRUST_REMOTE_KERNELS" not in os.environ:
+    # Diffusers read this variable into a constant at import time, so set both.
+    import diffusers.utils.constants
+
+    os.environ["DIFFUSERS_TRUST_REMOTE_KERNELS"] = "true"
+    diffusers.utils.constants.DIFFUSERS_TRUST_REMOTE_KERNELS = True

vae/config.json

@@ -1,6 +1,6 @@
 {
   "_class_name": "AutoencoderKLQwenImage21",
-  "_diffusers_version": "0.37.0.dev0",
+  "_diffusers_version": "0.41.0.dev0",
   "attn_scales": [],
   "base_dim": 96,
   "decoder_base_dim": 144,

vae/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a07a1b7c4ee2966a1b3bdc37de9b4f983d56937e46619f709a80b6e490675417
-size 1350989512
+oid sha256:71879ffd5321e6d10c3c87513e2b474b1252efa7f3dec2969214a9bf06a6dd5c
+size 675508656