README.md

---
license: other
license_name: flux-non-commercial-license
license_link: https://huggingface.co/black-forest-labs/FLUX.2-klein-9B/blob/main/LICENSE.md
base_model: black-forest-labs/FLUX.2-klein-9B
pipeline_tag: text-to-image
library_name: diffusers
tags:
  - diffusers
  - flux
  - orbitquant
  - quantization
  - text-to-image
---

# FLUX.2 Klein 9B OrbitQuant W4A4

This is a complete Diffusers pipeline derived from
[`black-forest-labs/FLUX.2-klein-9B`](https://huggingface.co/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`](https://huggingface.co/WaveCut/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

```bash
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](https://github.com/iamwavecut/OrbitQuant/releases/tag/kernels-v1)
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:

```bash
orbitquant kernels-install
orbitquant kernels-status
```

For ComfyUI, install the
[ComfyUI-OrbitQuant](https://github.com/iamwavecut/ComfyUI-OrbitQuant) node
pack; its install hook provisions the same kernels automatically. The full
contract lives in the
[provisioning documentation](https://github.com/iamwavecut/OrbitQuant/blob/main/docs/kernel-audit.md#native-package-provisioning).

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

```bash
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python generate.py
```

## Diffusers

```python
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:

```python
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:

```python
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.

![BF16, SDNQ UINT4 and OrbitQuant W4A4 across ten difficult prompts](assets/image_generation_comparison_matrix.webp)

### 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](https://arxiv.org/abs/2607.02461)
- [OrbitQuant library](https://github.com/iamwavecut/OrbitQuant)
- [Source FLUX.2 Klein 9B model](https://huggingface.co/black-forest-labs/FLUX.2-klein-9B)
- [Controlled SDNQ comparison checkpoint](https://huggingface.co/WaveCut/FLUX.2-klein-9B-SDNQ-uint4-static)