lite-infer/flux.2-klein-9b-nunchaku-lite-nvfp4_r32-bnb4-text-encoder

FLUX.2 Klein 9B Nunchaku Lite NVFP4 r32

Diffusers-loadable NVFP4 conversion of black-forest-labs/FLUX.2-klein-9B, quantized with `diffuse-compressor`. It loads with a plain from_pretrained call — no runtime graph patches and no extra runtime package.

The transformer uses quant_method: nunchaku_lite, NVFP4 SVDQ with group size 16, rank 32, 144 SVDQ targets and 3 AWQ W4A16 targets. Six outer linears (embedders, norm_out.linear, proj_out) stay in bf16, and the text_encoder component is BitsAndBytes 4-bit NF4 with bf16 compute. QKV projections are not fused, so this trades some speed for loading through the stock Diffusers graph. Calibrated on 128 prompts at 4 steps, 1024x1024.

Benchmark

NVIDIA RTX PRO 4000 Blackwell, settings as in the Run snippet below, one warmup and three measured runs, everything resident on the GPU with no offload. Latency covers denoise and decode, with text embeddings precomputed; VRAM is peak device usage.

That is 1.78x faster than dense bf16 using 2.42x less VRAM.

Output Comparison

[image]

Dense reference (left) and this build (right). Same prompt, seed, scheduler, resolution, and step count. Pixel MAE 31.21 / RMSE 45.71, versus 30.14 / 45.32 for the INT4 build.

Run

Requires the Hugging Face kernels package and DIFFUSERS_TRUST_REMOTE_KERNELS=true, plus a Diffusers build with the nunchaku_lite quantizer. NVFP4 needs a Blackwell or newer NVIDIA GPU; Hopper is unsupported. On Turing through Ada use the INT4 build instead.

import torch
from diffusers import Flux2KleinPipeline

pipe = Flux2KleinPipeline.from_pretrained(
    "lite-infer/flux.2-klein-9b-nunchaku-lite-nvfp4_r32-bnb4-text-encoder",
    torch_dtype=torch.bfloat16,
).to("cuda")

image = pipe(
    prompt="A glass robot tending orchids in a sunlit greenhouse, cinematic lighting, highly detailed",
    generator=torch.Generator("cuda").manual_seed(12345),
    width=1024,
    height=1024,
    num_inference_steps=4,
    guidance_scale=1.0,
).images[0]
image.save("output.png")