nunchaku

Nunchaku is a high-performance inference engine optimized for 4-bit neural networks using the SVDQuant post-training quantization technique. It is designed to maintain high visual fidelity while significantly reducing memory usage and increasing inference speed.

Key benefits include:

• Memory Reduction: Achieves up to 3.6× memory reduction (e.g., on 12B FLUX.1-dev compared to BF16).
• Speedup: Offers significant speed improvements (e.g., 8.7× speedup over 16-bit models on specific hardware) by eliminating the need for CPU offloading.
• High Fidelity: Maintains superior visual quality compared to standard W4A4 or W4A8 baselines.

The nunchaku package is organized into several functional subpackages. Depending on your task, you should import from the following modules:

• nunchaku.models: Core model definitions and loading.
• nunchaku.lora: LoRA (Low-Rank Adaptation) support and management.
• nunchaku.pipeline: High-level inference pipelines.
• nunchaku.ops: Specialized operations and kernels.
• nunchaku.caching: Caching mechanisms for performance optimization.
• nunchaku.utils: General utility functions.

SVDQuant is the quantization method used by Nunchaku to enable efficient 4-bit quantization of diffusion models by managing outliers:

1. Outlier Migration: Outliers are migrated from activations to weights. This makes the updated activation $\hat{\mathbf{X}}$ easier to quantize, while the updated weights $\hat{\mathbf{W}}$ become more difficult.
2. Low-Rank Decomposition: The difficult weights $\hat{\mathbf{W}}$ are decomposed using SVD into a low-rank component $\mathbf{L}_1\mathbf{L}_2$ and a residual $\hat{\mathbf{W}}-\mathbf{L}_1\mathbf{L}_2$.
3. Precision Handling: The low-rank branch runs at 16-bit precision to alleviate quantization difficulty.

The Nunchaku engine optimizes the overhead of running the low-rank branch through kernel fusion. Instead of running components naively (which introduces significant latency due to extra reads/writes), Nunchaku fuses kernels:

• Down Projection & Quantize: These kernels are fused because they use the same input.
• Up Projection & 4-Bit Compute: These kernels are fused because they share the same output.

This reduces data movement overhead and improves inference speed.

To begin using Nunchaku for 4-bit FLUX.1 models, follow these resources:

• Installation: Follow the Installation Guide.
• Basic Usage: Refer to the Usage Tutorial.
• ComfyUI Integration: Use the ComfyUI-nunchaku plugin.
• Custom Quantization: Use DeepCompressor for custom model quantization.
• API Reference: For programmatic access, consult the Python API Reference.

To significantly reduce GPU memory usage when running FLUX.1 models, you can use Nunchaku's built-in CPU offload in combination with Diffusers' sequential offload.

1. Nunchaku CPU Offload: Set offload=True during the transformer initialization. This offloads inactive model components to CPU memory.
2. Diffusers Sequential Offload: Call pipeline.enable_sequential_cpu_offload() on your pipeline instance. This enables automatic device management.

Important: When using pipeline.enable_sequential_cpu_offload(), do not manually move the pipeline to a device using .to('cuda'), as the sequential offload mechanism handles device placement automatically.

Nunchaku uses pytest for testing. You must provide a Hugging Face access token via the HF_TOKEN environment variable to download necessary models and datasets, unless you have already authenticated via huggingface-cli login.

To speed up test runs that generate images, you can set the NUNCHAKU_TEST_CACHE_ROOT environment variable to cache results. If not set, images are saved to test_results/ref.

Use latency.py to measure the end-to-end latency or individual DiT forward step latency for INT4 models.

CLI Arguments:

• -t <steps>: Number of inference steps (default: 20).
• -g <scale>: Guidance scale (default: 5).
• --pag-scale <scale>: PAG guidance scale (default: 2).
• --mode <mode>: Set to step to measure the latency of a single DiT forward step instead of end-to-end generation.
• --warmup-times <count>: Number of warmup runs (default: 2).
• --test-times <count>: Number of test runs (default: 10).

python latency.py

Qwen-Image-Edit-2509 is a monthly iteration of the model. To run the 4-bit quantized version, use NunchakuQwenImageTransformer2DModel.

Requirement: This requires diffusers version 0.36.0 or higher.

Use the generate.py script to generate images from text prompts directly via the CLI. By default, the script uses the INT4 model and saves the result to output.png.

CLI Arguments:

• --prompt <text>: The text prompt for image generation.
• -o or --output-path <path>: Specify a custom path for the generated image.
• -p <precision>: Select model precision. Options include int4 (default) or bf16.
• -t <steps>: Number of inference steps (default: 20).
• -g <scale>: Classifier-free guidance scale (default: 5).
• --pag-scale <scale>: Adjust the PAG guidance scale (default: 2).

python generate.py --prompt "You Text Prompt"

Nunchaku supports Cache-DiT to enable advanced acceleration strategies like DB Cache. To use this feature, you must first install the cache-dit package.

This feature is compatible with models such as FLUX.1 and Qwen-Image.

pip install cache-dit