About

nunchaku_lite brings Nunchaku-quantized image generation models to standard Diffusers pipelines with a lean runtime, native CUDA kernels, and drop-in pipeline loading. It is designed for fast SVDQ W4A4 inference with a small integration surface: load a normal Diffusers pipeline, replace only the quantized component, and keep the rest of the workflow unchanged.

Its core features include:

• Efficient pipeline loading: constructs the patched component up front so Diffusers does not first materialize dense BF16 transformer or UNet weights that are immediately replaced.
• Native quantized kernels: packages the CUDA kernels and Python wrappers needed for INT4 and FP4 Nunchaku checkpoints.
• AWQ GEMM acceleration: unlike original Nunchaku, nunchaku_lite ships a native W4A16 group-size-64/int32 GEMM path for large AWQ projections, turning chunked GEMV bottlenecks into high-throughput CUDA matmuls.
• Broad image-model coverage: includes built-in adapters for FLUX.1, FLUX.2 Klein, Qwen-Image, Qwen-Image-Edit, SDXL, SDXL-Turbo, and Z-Image Turbo.
• Runtime LoRA support: exposes Diffusers-style LoRA loading, adapter strength control, multi-LoRA composition, and unload/reset flows for supported model families.
• Extensible adapter registry: keeps model-specific graph rewrites isolated in small adapters without requiring pipeline subclasses or a dependency on the full nunchaku Python package.

Benchmarks

On an NVIDIA RTX PRO 6000, nunchaku_lite delivers up to 1.79x lower latency, up to 46% lower peak CUDA memory, and up to 71% smaller transformer storage versus unmodified Diffusers pipelines in the current benchmark set.

See docs/benchmarks.md for charts, generated samples, run settings, and reproduction commands.

Supported Models

nunchaku_lite currently supports Diffusers pipelines whose transformer or UNet component matches one of the built-in adapters below. Checkpoints should be SVDQ W4A4 Nunchaku v2 .safetensors files in INT4 or FP4 format. Full runnable model guides live under docs/models/.

The Qwen low-VRAM guides use enable_model_cpu_offload(), which requires accelerate.

Additional model families should be added through the adapter registry rather than pipeline subclasses. See docs/roadmap.md for planned coverage and remaining feature work.

Requirements

• Python 3.10 or newer
• PyTorch 2.7 or newer with CUDA
• Diffusers 0.36 or newer

Native source builds also require CUDA toolkit 12.6 or newer with nvcc. The default runtime path uses prebuilt kernels from rootonchair/nunchaku-lite-kernels through the Hugging Face kernels library instead.

The native build detects the local GPU architecture by default. Supported targets are sm75, sm80, sm86, sm89, sm120a, and sm121a, subject to the installed CUDA toolkit version. CUDA 12.6 or newer is the documented minimum; sm120a requires CUDA 12.8 or newer, and sm121a requires CUDA 13.0 or newer. The Hugging Face kernels package provides prebuilt kernels for CUDA 7.5, 8.0, 8.6, 8.9, 12.0a, and 12.1a.

The root Python package does not compile CUDA code. CUDA toolkit requirements apply only when installing the local nunchaku_lite_kernels package from ./nunchaku-lite-kernels.

CUDA version note: nunchaku_lite first tries a locally installed nunchaku_lite_kernels package, then falls back to rootonchair/nunchaku-lite-kernels through Hugging Face kernels.

Installation

Install from source:

pip install .

This installs the root Python runtime and uses Hugging Face prebuilt kernels by default.

To build and use the local CUDA kernels package:

pip install --no-build-isolation ./nunchaku-lite-kernels

--no-build-isolation ensures the extension is compiled against the PyTorch version installed in the active environment.

When nunchaku_lite_kernels is installed locally, nunchaku_lite prefers it over the Hugging Face fallback.

Build and install a wheel:

python setup.py bdist_wheel
pip install dist/nunchaku_lite-*.whl

By default, the local kernels package build uses NUNCHAKU_INSTALL_MODE=FAST and compiles for visible local CUDA devices. To build all supported architectures:

NUNCHAKU_INSTALL_MODE=ALL pip install --no-build-isolation ./nunchaku-lite-kernels

Quick Start

Run from the repository root:

from pathlib import Path

import torch
from diffusers import FluxPipeline

from nunchaku_lite import load_nunchaku_pipeline


model_id = "black-forest-labs/FLUX.1-schnell"
precision = "fp4"  # "int4" or "fp4"
checkpoints = {
    "int4": "nunchaku-ai/nunchaku-flux.1-schnell/svdq-int4_r32-flux.1-schnell.safetensors",
    "fp4": "nunchaku-ai/nunchaku-flux.1-schnell/svdq-fp4_r32-flux.1-schnell.safetensors",
}
output_path = Path(f"outputs/flux_schnell_nunchaku_lite_{precision}.png")

pipe = load_nunchaku_pipeline(
    model_id,
    pipeline_cls=FluxPipeline,
    checkpoint=checkpoints[precision],
    target="flux",
    precision=precision,
    torch_dtype=torch.bfloat16,
    device="cuda",
)
pipe = pipe.to("cuda")

image = pipe(
    "A cat holding a sign that says hello world",
    height=1024,
    width=1024,
    num_inference_steps=4,
    guidance_scale=0.0,
    generator=torch.Generator(device="cuda").manual_seed(12345),
).images[0]

output_path.parent.mkdir(parents=True, exist_ok=True)
image.save(output_path)
print(f"saved {output_path}")

Runtime LoRA

Pipelines loaded with load_nunchaku_pipeline expose Diffusers-style LoRA methods when the selected adapter supports runtime LoRA:

from pathlib import Path

import torch
from diffusers import FluxPipeline

from nunchaku_lite import load_nunchaku_pipeline


model_id = "black-forest-labs/FLUX.1-dev"
precision = "fp4"
checkpoint = "nunchaku-tech/nunchaku-flux.1-dev/svdq-fp4_r32-flux.1-dev.safetensors"
output_path = Path("outputs/flux_dev_ghibsky_lora_fp4.png")

pipe = load_nunchaku_pipeline(
    model_id,
    pipeline_cls=FluxPipeline,
    checkpoint=checkpoint,
    target="flux",
    precision=precision,
    torch_dtype=torch.bfloat16,
)
pipe.enable_model_cpu_offload()

pipe.load_lora_weights(
    "aleksa-codes/flux-ghibsky-illustration",
    weight_name="lora.safetensors",
    adapter_name="ghibsky",
)
pipe.set_adapters("ghibsky", adapter_weights=0.75)

image = pipe(
    "GHIBSKY style painting of a cozy mountain cabin beside a clear lake at sunset",
    height=1024,
    width=1024,
    num_inference_steps=28,
    guidance_scale=3.5,
    generator=torch.Generator(device="cpu").manual_seed(12345),
).images[0]

output_path.parent.mkdir(parents=True, exist_ok=True)
image.save(output_path)
print(f"saved {output_path}")

Checkpoint paths can be local .safetensors files or Hugging Face paths of the form:

org-or-user/repo-name/path/to/checkpoint.safetensors

Documentation

Public API

from nunchaku_lite import (
    TransformerAdapter,
    list_adapters,
    load_nunchaku_pipeline,
    patch_transformer,
    register_adapter,
)

load_nunchaku_pipeline(...) is the preferred entry point for normal pipeline loading. Use patch_transformer(...) only when a Diffusers component has already been constructed and needs in-place patching.

See docs/api.md for argument details, runtime LoRA examples, and adapter registry usage.

Development

Run the unit tests:

pytest -q tests

Build the extension in place:

python setup.py build_ext --inplace

See docs/development.md for full inference tests, adapter authoring guidance, runtime LoRA implementation notes, and repository layout.

License

nunchaku_lite is licensed under the Apache License, Version 2.0. See LICENSE.

Acknowledgements

nunchaku_lite builds on the Nunchaku project and uses selected native kernel code for the lite runtime. We are grateful to the maintainers and contributors of these projects.