diffuse_compressor is a model-agnostic SVDQuant toolkit for diffusion transformers. It prepares user-selected projection targets, runs calibration replay, quantizes diffusion backbones to INT4 or NVFP4-style weights, and exports Nunchaku-compatible safetensors checkpoints.

The core package deliberately avoids hard-coding Flux, PixArt, Sana, ERNIE, LongCat, video, or image-edit architecture details. Those choices live in TargetConfig examples and downstream user configuration, so the library stays small enough to adapt to new diffusion models.

Highlights

• Model-agnostic target collection with wildcard paths, class scans, grouped QKV/KV projections, skips, and target-level overrides.
• Generic model rewrites such as splitting fused linear or convolution projections before target collection.
• SVDQuant for linear projections and pointwise Conv2d projections.
• Optional GPTQ residual rounding for INT4 and FP4/NVFP4 quantization.
• INT4 and FP4/NVFP4-style residual weight export paths, including DeepCompressor-style scale dtype metadata.
• Disk-backed calibration replay with scoped activation capture, cache reuse, artifact caching, and memory-sensitive replay.
• Nunchaku Lite and Nunchaku-style safetensors export with adjacent config metadata and optional runtime manifests.
• Runnable Diffusers examples for text-to-image, image-to-image, and text-to-video configuration sketches.

Installation

Install the package in editable mode:

python -m pip install -e .

Install development tools:

python -m pip install -e ".[dev]"

Install example data-loading extras for examples that download calibration images from Hugging Face datasets:

python -m pip install -e ".[examples]"

Nunchaku Lite runtime patching still requires installing nunchaku_lite from its release or private package channel. The nunchaku-lite extra is only an explicit optional-runtime marker; it does not install a public PyPI package.

Guidelines

• Quantize a new Hugging Face model: inspect a generic model, calibrate it, package the checkpoint, and load it directly with DiffusionPipeline.from_pretrained.
• Text-to-image end-to-end guide: quantize, evaluate, and run inference with FLUX.2 Klein 4B.
• Image-to-image end-to-end guide: quantize, evaluate, and run inference with LongCat Image Edit.
• Adding a new model: adapt target configs, patches, calibration scopes, inspection, quantization, evaluation, and inference for another architecture.

Example CLIs

Run one of the Diffusers-backed examples:

python examples/text_to_image/quantize_flux1_schnell.py --precision int4
python examples/text_to_image/quantize_flux1_schnell.py --precision nvfp4

Override defaults for a larger run:

python examples/text_to_image/quantize_flux2_klein_4b.py \
  --precision int4 \
  --model-id black-forest-labs/FLUX.2-klein-4B \
  --num-samples 128 \
  --batch-size 1 \
  --output outputs/checkpoints/svdq-int4_r32-flux2-klein-4b.safetensors

Use a lower-memory preset when VRAM or calibration RAM is tight:

python examples/text_to_image/quantize_flux2_klein_4b.py \
  --precision int4 \
  --num-samples 64 \
  --cache-num-samples 64 \
  --batch-size 1 \
  --sample-batch-size 32 \
  --scope-capture-mode one-target \
  --pipeline-offload sequential \
  --offload-model \
  --compute-device cuda

For GPU VRAM, the main knobs are --batch-size 1, --pipeline-offload model or sequential, --offload-model, and --compute-device cuda. The --cache-num-samples, --sample-batch-size, and --scope-capture-mode one-target options primarily reduce calibration memory and replay working set. See docs/low_memory_quantization.md for the full tradeoffs.

Supported example families:

Generic Diffusers repositories can be scanned without a model-specific target map. When a denoiser contains homogeneous repeated ModuleList block stacks, the scanner keeps outer embeddings and projections dense, applies SVDQ only inside those blocks, and derives safe block calibration scopes. Models without such a stack retain the broad compatible-linear fallback. Normalization and modulation linears use the single Diffusers-supported plain AWQ W4A16 layout:

python examples/text_to_image/quantize_hf.py MODEL_ID --inspect-config
python examples/text_to_image/quantize_hf.py MODEL_ID --precision int4
python examples/image_to_image/quantize_hf.py MODEL_ID --dataset DATASET_ID
python examples/text_to_video/quantize_hf.py MODEL_ID --precision nvfp4

Use --skip globs to exclude targets from discovery. Generic mode does not fuse QKV projections or emit the original Nunchaku AdaNorm-interleaved AWQ layout; use a model-specific script when structural patches or fused runtime targets are required.

Package a generic-manifest checkpoint as a complete pipeline for Diffusers' Nunchaku Lite backend:

python examples/convert_nunchaku_lite_diffusers.py \
  --checkpoint outputs/checkpoints/svdq-int4_r32-flux-2-klein-4b.safetensors \
  --model-id black-forest-labs/FLUX.2-klein-4B

The full example table, command matrix, output paths, defaults, and offload notes are preserved in docs/examples.md.

Configuration

A target config answers which model modules become quantized runtime projections. It can describe:

• structural patches to expose targetable child modules;
• single-module and grouped projection targets;
• pointwise Conv2d projector targets;
• skipped modules and unquantized state-dict patterns;
• calibration scopes that replay and clear activations by block;
• runtime-specific tensor layouts such as Nunchaku SVDQ, AWQ W4A16, and AdaNorm AWQ W4A16.

Inspect a config before running a full quantization job:

from diffuse_compressor import inspect_target_config

report = inspect_target_config(model, target_config)
print(report.format_text())
assert report.ok

Example scripts also support:

python examples/text_to_image/quantize_flux1_schnell.py --inspect-config

See docs/configuration.md for target rules, skips, calibration scopes, inspection output, and small runnable recipes.

Checkpoint Export

The Nunchaku exporter writes one safetensors file containing:

• quantized target parameters under configured export_name prefixes;
• untouched non-target model parameters required for strict runtime loading;
• compact quantization_config.* compatibility metadata;
• optional quantization_config.runtime_manifest metadata when the checkpoint can declare a generic Nunchaku Lite runtime ABI.

Config metadata is written beside the checkpoint as <checkpoint-stem>.config.yaml. The schema is documented in docs/checkpoint_metadata.md, and the Nunchaku Lite runtime manifest is documented in docs/nunchaku_lite_manifest_v1.md.

For Nunchaku-style SVDQuant, quantized linear targets use keys such as:

transformer_blocks.0.attn.to_qkv.qweight
transformer_blocks.0.attn.to_qkv.wscales
transformer_blocks.0.attn.to_qkv.smooth_factor
transformer_blocks.0.attn.to_qkv.smooth_factor_orig
transformer_blocks.0.attn.to_qkv.proj_down
transformer_blocks.0.attn.to_qkv.proj_up

Documentation

Development

Install in editable mode with test/build tools:

python -m pip install -e ".[dev]"

Run the test suite:

pytest

Run a focused test while iterating on quantization or export behavior:

pytest tests/test_quantize_export.py

Build source and wheel distributions:

python -m build

Acknowledgements

This project builds on ideas and compatibility targets from DeepCompressor, the original Nunchaku SVDQuant diffusion compression repository.

License

See LICENSE.