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:
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.
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.