File size: 3,959 Bytes
a4957b7
 
7ad44bc
a4957b7
 
 
 
7ad44bc
a4957b7
 
 
 
7ad44bc
 
 
 
 
 
 
 
 
 
59c59db
7ad44bc
 
 
 
 
 
 
59c59db
7ad44bc
 
 
 
 
 
 
 
 
 
 
59c59db
7ad44bc
59c59db
 
 
 
 
 
7ad44bc
 
a4957b7
 
 
 
 
d16e497
16cd752
 
d16e497
a4957b7
 
 
 
 
 
 
7ad44bc
 
 
a4957b7
 
 
7ad44bc
 
 
 
b51c896
59c59db
7ad44bc
 
 
 
 
 
a4957b7
 
 
 
 
 
 
7ad44bc
 
 
 
 
 
 
 
a4957b7
 
7ad44bc
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
---
base_model: black-forest-labs/FLUX.2-klein-4B
license: apache-2.0
tags:
- orbitquant
- quantized
- diffusers
- diffusion-transformer
---

# black-forest-labs/FLUX.2-klein-4B OrbitQuant W2A4

This repository contains a compact OrbitQuant transformer-component artifact for the source Diffusers model listed above. It is intended to be loaded into the original pipeline, not used as a standalone Diffusers pipeline repository.

OrbitQuant is a calibration-free post-training quantization method for image and video diffusion transformers. This artifact keeps the text encoders, VAE, embeddings, timestep MLP, and final heads in the source precision by default and replaces the transformer linear projections with OrbitQuant modules.

## Usage

Install the package from this repository, then load the base pipeline and patch its transformer component with this artifact:

```python
import torch
from diffusers import Flux2KleinPipeline
from huggingface_hub import snapshot_download
from orbitquant import load_quantized_pipeline_component

base_model = "black-forest-labs/FLUX.2-klein-4B"
artifact_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W2A4"

artifact_dir = snapshot_download(artifact_id, repo_type="model")
pipe = Flux2KleinPipeline.from_pretrained(
    base_model,
    torch_dtype=torch.bfloat16,
)
load_quantized_pipeline_component(
    pipe,
    artifact_dir,
    component="transformer",
    device="cuda",
)
pipe.to("cuda")

image = pipe(
    prompt="A precise product photo of a red ceramic mug on a wooden desk",
    height=1024,
    width=1024,
    num_inference_steps=4,
    guidance_scale=1.0,
).images[0]
image.save("flux2-klein-orbitquant.png")
```

## Quantization

- Method: `orbitquant`
- Bits: `W2A4`
- Runtime mode: `dequant_bf16`
- Activation kernel backend: `triton_cuda`
- Quantization device: `cuda`
- Weight quantization backend: `triton_cuda`
- Quantization staging: `component`
- Target policy: `flux2`
- Rotation: `rpbh`
- Rotation seed: `0`
- Block size: `paper`
- Block size policy: `largest_power_of_two_dividing_dim`
- Codebook: `lloyd_max`
- Codebook version: `1`
- Quantized transformer modules: `100`
- AdaLN INT4 modules: `3`
- Skipped modules: `6`
- Calibration data: none
- Text encoders and VAE: left in source precision by default

## Visual Comparison

The following assets are stored in this artifact and compare the BF16 base generation against the OrbitQuant generation with the same prompt and seed.

![assets/image_generation_comparison_matrix.webp](assets/image_generation_comparison_matrix.webp)

![assets/original_vs_orbitquant_flux2-native_seed0_W2A4_color-binding.webp](assets/original_vs_orbitquant_flux2-native_seed0_W2A4_color-binding.webp)

![assets/original_vs_orbitquant_flux2-native_seed0_W2A4_counting.webp](assets/original_vs_orbitquant_flux2-native_seed0_W2A4_counting.webp)

![assets/original_vs_orbitquant_flux2-native_seed0_W2A4_cyrillic-text-rendering.webp](assets/original_vs_orbitquant_flux2-native_seed0_W2A4_cyrillic-text-rendering.webp)

## Source

- Model: `black-forest-labs/FLUX.2-klein-4B`
- Revision: `e7b7dc27f91deacad38e78976d1f2b499d76a294`
- Source license: `apache-2.0`
- OrbitQuant paper: https://arxiv.org/abs/2607.02461

## Artifact Files

- `model.safetensors`: packed OrbitQuant/INT4 module tensors.
- `quantization_config.json`: serialized OrbitQuant runtime settings.
- `orbitquant_manifest.json`: source provenance, policies, module lists, and checksums.
- `orbitquant_codebooks.safetensors`: Lloyd-Max codebooks.
- `orbitquant_rotations.safetensors`: deterministic RPBH rotation metadata.

## Limitations

- This is a transformer-component artifact; load it into the source pipeline as shown above.
- Runtime mode may dequantize packed weights before BF16 matmul. Disk artifacts are compact, while runtime VRAM depends on the selected backend.
- Quality depends on the source model and bit setting. Very low-bit settings can degrade prompt following or visual detail.