WaveCut/FLUX.2-klein-9B-OrbitQuant-W4A4

FLUX.2 Klein 9B OrbitQuant W4A4

This is a complete Diffusers pipeline derived from black-forest-labs/FLUX.2-klein-9B at revision 92196c8e11f7b6cf2b7493e037d8c5345c559216.

Both compute-heavy components are compressed:

• transformer: 144 OrbitQuant W4A4 projections plus 3 AdaLN INT4 weight-only projections;
• Qwen3 text encoder: 252 OrbitQuant W4A4 projections; lm_head remains BF16.

The text encoder is quantized to permit a controlled comparison with the published FLUX.2-klein-9B-SDNQ-uint4-static pipeline. This is an extension of OrbitQuant's architecture-independent adapter; the OrbitQuant paper itself leaves text encoders in BF16.

Install

pip install "orbitquant[hf,kernels]>=0.6.0"

OrbitQuant uses packed low-bit inference by default and does not silently materialize all weights in BF16. The optimized native kernel package is provisioned automatically at first model load: OrbitQuant derives the exact runtime variant (torch minor and CUDA version for CUDA, the torch stable ABI for CPU, plus OS and architecture), downloads the matching prebuilt wheel from the OrbitQuant kernels release with checksum verification, and caches it under ~/.cache/orbitquant/kernels. When no variant matches, CUDA inference falls back to the Triton packed path. Provision explicitly (or inspect the resolution) with:

orbitquant kernels-install
orbitquant kernels-status

For ComfyUI, install the ComfyUI-OrbitQuant node pack; its install hook provisions the same kernels automatically. The full contract lives in the provisioning documentation.

With PyTorch 2.9, launch CUDA inference with expandable allocator segments to minimize reserved memory:

PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python generate.py

Diffusers

import torch
import orbitquant  # registers the quantizer
from diffusers import Flux2KleinPipeline

pipe = Flux2KleinPipeline.from_pretrained(
    "WaveCut/FLUX.2-klein-9B-OrbitQuant-W4A4",
    torch_dtype=torch.bfloat16,
).to("cuda")

image = pipe(
    prompt="An intricate orbital conservatory above Earth, documentary realism",
    height=1024,
    width=1024,
    num_inference_steps=4,
    guidance_scale=1.0,
    generator=torch.Generator(device="cuda").manual_seed(0),
).images[0]
image.save("orbitquant.png")

Convert the source checkpoint on load

To create a fresh transformer-only OrbitQuant pipeline directly from the source safetensors checkpoint, use the normal Diffusers loader. This row-streams source weights into packed tensors instead of keeping the complete BF16 transformer and the quantized transformer resident together:

import torch
import orbitquant
from diffusers import DiffusionPipeline
from orbitquant import (
    OrbitQuantConfig,
    build_diffusers_pipeline_quantization_config,
)

qconfig = build_diffusers_pipeline_quantization_config(
    OrbitQuantConfig(target_policy="auto"),
    components="transformer",
)
pipe = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.2-klein-9B",
    quantization_config=qconfig,
    torch_dtype=torch.bfloat16,
)
pipe.enable_model_cpu_offload()

Use pipe.enable_sequential_cpu_offload() instead for sequential offload. The published artifact also quantizes its text encoder for the SDNQ comparison; on-the-fly conversion keeps text encoders in source precision unless they are explicitly included in components. Guaranteed bounded-memory conversion requires a safetensors source checkpoint.

Transformers Component

The quantized Qwen3 component can also be loaded through Transformers:

import torch
import orbitquant
from transformers import AutoModelForCausalLM, AutoTokenizer

repo_id = "WaveCut/FLUX.2-klein-9B-OrbitQuant-W4A4"
tokenizer = AutoTokenizer.from_pretrained(repo_id, subfolder="tokenizer")
text_encoder = AutoModelForCausalLM.from_pretrained(
    repo_id,
    subfolder="text_encoder",
    torch_dtype=torch.bfloat16,
).to("cuda")

Quantization

Setting	Value
Weight / activation bits	W4A4
Rotation	RPBH, seed 0
Block policy	Largest power-of-two divisor of the input dimension
Codebook	Lloyd-Max version 2
Row norms	BF16
AdaLN	INT4 RTN, group size 64, BF16 activations
Runtime	auto_fused; native RPBH/INT8 surrogate plus CUTLASS W4A4 on supported CUDA GPUs
Calibration data	None

The packed transformer and text-encoder weight payload is 10.67 GB; including the VAE, the weight payload is 10.84 GB. The complete pipeline before this card asset is 10.85 GB.

Benchmark (RTX PRO 6000, OrbitQuant 0.5.0)

All three variants ran back-to-back on one NVIDIA RTX PRO 6000 Blackwell Workstation Edition (96 GB) in one session, one separate process per variant, same protocol and pinned revisions (BF16 92196c8, SDNQ ed71b3f, OrbitQuant ee3a38f), Torch 2.9.1+cu128, Diffusers 0.39.0, Transformers 5.13.0, BF16 arithmetic, no CPU offload, 1024x1024, four steps, guidance 1.0, seed 0, ten prompts, PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, orbitquant==0.5.0 with a locally built sm_120 native kernel package:

Variant	Load	Cold image	Hot mean	Hot median	Hot NVML peak	CUDA allocated peak	CUDA reserved peak
BF16 (no quantization)	3.06 s	1.54 s	1.2014 s	1.2019 s	36.40 GB	34.72 GB	35.12 GB
SDNQ UINT4	3.02 s	5.81 s	1.2735 s	1.2730 s	15.42 GB	13.75 GB	14.14 GB
OrbitQuant W4A4 (0.5.0)	1.95 s	4.96 s	1.1633 s	1.1640 s	14.17 GB	12.51 GB	12.89 GB

OrbitQuant's hot median is 3.2% faster than unquantized BF16 and 8.6% faster than SDNQ, it loads fastest of the three, and it uses 22.2 GB less peak NVML memory than BF16 and 1.25 GB less than SDNQ. Machine-readable results: benchmark/summary.json.

Visual Comparison

The matrix uses the complete ten-prompt stress pack with full 1024x1024 tiles and WebP quality 95. BF16 is the full-precision reference from the controlled visual run; the SDNQ and OrbitQuant columns use the recorded benchmark outputs. Every row uses the same prompt, seed, resolution, step count and guidance.

Visual Assessment

• No collapse: all thirty outputs are finite, coherent and detailed. OrbitQuant did not produce blank, noisy or structurally broken images.
• Micro-detail and materials: all three variants preserve gears, filigree, architectural interiors, paper grain, metal, resin and reflected surfaces. OrbitQuant remains competitive with BF16 and SDNQ in these cases.
• Dense composition: all variants retain foreground/background separation and the main hierarchy in the architectural cutaway and orbital-banquet prompts. Individual requested objects move or disappear because quantization changes the denoising trajectory.
• Counting: none of the variants reliably renders exactly nine performers or every exact repeated motif. This is a base-model limitation in the tested setting rather than an OrbitQuant-only collapse.
• English typography: OrbitQuant is strongest on this row: it preserves the headline, subtitle and all four specification lines. SDNQ preserves the headline and three table lines but omits or corrupts some requested text.
• Russian typography: all variants render the large headline, subtitle and archive stamp well; small contents text contains errors in every column.
• Japanese and Chinese typography: visual glyph quality is plausible, but exact requested strings are not reliably reproduced by any variant.
• Trajectory fidelity: both quantizers change the denoising trajectory at the same seed; neither remains consistently closer to BF16 across all ten prompts.

This subjective paired inspection demonstrates non-collapse and exposes concrete failure modes; it is not a substitute for GenEval or another task-specific objective metric. In this controlled FLUX.2 Klein 9B comparison, OrbitQuant produces the smaller complete 4-bit pipeline, reaches SDNQ hot-generation parity with lower runtime memory, and produces the strongest English fine-print result in the prompt pack.

Limitations

• Optimized native CUDA measurements require a locally built matching ABI3 kernel package.
• Triton packed matmul is a compatible fallback, not the fastest measured backend.
• The optimized CUDA W4A4 tensor-core path maps each fixed Lloyd-Max codebook to symmetric INT8 surrogate values and one scalar. Use runtime_mode="dequant_bf16" for exact-centroid reference evaluation.
• W4A4 activation quantization can move the denoising trajectory farther from BF16 than weight-only UINT4.
• Quantizing the text encoder is outside the paper's default layer policy.
• This model inherits the FLUX Non-Commercial License from the source checkpoint.

References

• OrbitQuant paper
• OrbitQuant library
• Source FLUX.2 Klein 9B model
• Controlled SDNQ comparison checkpoint