milo01/FLUX.2-klein-9b-INT8-transformer

FLUX.2 [klein] 9B (step-distilled) — INT8 (W8A8) Transformer

Quantized transformer checkpoint for [FLUX.2 [klein] 9B](https://huggingface.co/black-forest-labs/FLUX.2-klein-9B) (step-distilled).

INT8 weight and activation quantization via NVIDIA ModelOpt with calibrated input scales.

Note: This repo contains only the quantized transformer weights. The text encoder, VAE, tokenizer, and scheduler are loaded from the base model: `black-forest-labs/FLUX.2-klein-9B`.

Model Details

Quantization Details

Evaluation

Evaluated on 48 prompts (T2I (16 each), editing (16 each), composition (16 each)). Both BF16 baseline and INT8 outputs are generated with identical prompts and seeds, then scored independently.

<details> <summary><strong>📊 Understanding the Metrics</strong></summary>

We report two categories of metrics:

Text-Image Alignment — measures output quality independently:

• —CLIP Score ↑: Uses OpenAI's CLIP model to score how well each generated image matches its text prompt. Both BF16 and quantized models are evaluated independently against the same prompts — this is not a comparison between the two outputs, but an independent quality measure for each. Higher is better (typical range: 0.25–0.35).

Fidelity — measures how closely the quantized output matches the BF16 baseline:

• —LPIPS ↓ (Learned Perceptual Image Patch Similarity): Uses a neural network to judge perceptual similarity the way a human would. Unlike pixel-level metrics, LPIPS captures structural and textural differences. 0 = perceptually identical, 1 = completely different. Values below 0.1 indicate very high fidelity.
• —PSNR ↑ (Peak Signal-to-Noise Ratio): Measures pixel-level accuracy in decibels. Higher values mean less error. 20–30 dB is typical for quantized model comparisons; 30+ dB is excellent.
• —FID ↓ (Fréchet Inception Distance): Compares the statistical distribution of all generated images (not individual pairs). Lower means the quantized model produces images from the same visual distribution as BF16. Sensitive to sample size — our 48-image evaluation provides a directional signal rather than a definitive score.

</details>

Text-Image Alignment (CLIP Score ↑)

CLIP score measures how well the generated image matches the text prompt (higher = better). Both models are evaluated independently:

Fidelity vs BF16 Baseline

These metrics measure how closely the quantized output matches the BF16 reference:

Per-Task Breakdown

Comparison with FP8 (Reference)

Black Forest Labs officially provides FP8 quantized checkpoints for FLUX.2 Klein. However, FP8 (float8_e4m3fn) requires hardware support introduced with NVIDIA Ada Lovelace (RTX 40-series / L4 / L40). INT8 offers a quantized alternative at the same ~2× compression ratio for GPUs that lack native FP8 support (e.g., Ampere, Turing, or non-NVIDIA hardware with INT8 acceleration).

The table below compares both formats against the same BF16 baseline (CLIP 0.6426), evaluated with identical prompts and seeds:

Per-Task Breakdown

Visual Comparison (BF16 vs INT8)

All images generated with identical prompts and seeds (4 denoising steps, 1024×1024).

Text-to-Image

"Oil painting of a stormy seascape in the style of J.M.W. Turner, violent waves crashing against rocks, ship barely visible in mist, thick impasto texture"

[image]

Image Editing

Base: "A red sports car parked in a garage" Edit: "Change the car color to yellow and make the garage look like a futuristic space hangar"

[image]

Multi-Reference Composition (2 references)

Ref 1: "A weathered bronze statue of a Greek philosopher" Ref 2: "A lush tropical rainforest canopy" Compose: "The statue is being reclaimed by the jungle, with vines and flowers growing over its features"

[image]

INT8 Performance Benchmarks

Measured on NVIDIA RTX 5090 with PyTorch 2.10.0+cu130 and CUDA 13.0. Full INT8 stack (INT8 transformer + INT8 text encoder). Resolution: 1024×1024.

<table> <tr><th>Model</th><th>Steps</th><th>Eager</th><th>Compiled</th><th>Throughput</th><th>VRAM</th></tr> <tr><td>klein-4b</td><td>4</td><td>1.77s</td><td>0.72s</td><td>1.387 img/s</td><td>11.25 GB</td></tr> <tr><td>klein-base-4b</td><td>50</td><td>33.25s</td><td>9.92s</td><td>0.101 img/s</td><td>11.26 GB</td></tr> <tr style="background-color: rgba(128,128,128,0.15);"><td><strong>klein-9b ◂</strong></td><td><strong>4</strong></td><td><strong>3.04s</strong></td><td><strong>1.09s</strong></td><td><strong>0.917 img/s</strong></td><td><strong>20.15 GB</strong></td></tr> <tr><td>klein-base-9b</td><td>50</td><td>62.49s</td><td>18.70s</td><td>0.053 img/s</td><td>20.16 GB</td></tr> </table>

torch.compile speedup: 2.5× (klein-4b), 3.4× (klein-base-4b), 2.8× (klein-9b), 3.3× (klein-base-9b) Why the large speedup? Our pipeline loads INT8 weights using TorchAO, which represents linear layers as W8A8 quantized tensors. In eager mode, each quantized matmul dispatches separate CUDA kernels for dequantization and computation. With torch.compile, the full graph is traced and these operations are fused into optimized Triton kernels that perform dequantize + matmul in a single pass, eliminating kernel launch overhead and intermediate memory traffic.

Usage

🚧 Code release coming soon. A pip-installable loader library is in preparation.

Compatibility

This checkpoint uses the official FLUX.2 single-file safetensors format — the same key layout and structure used by Black Forest Labs for their official FP8 and NVFP4 quantized models. Any loader that supports quantized FLUX.2 single-file checkpoints can load this INT8 checkpoint.

License

This model inherits the license from the base model: [FLUX Non-Commercial](https://huggingface.co/black-forest-labs/FLUX.2-klein-9B/blob/main/LICENSE.md).

Acknowledgments

• —Black Forest Labs for FLUX.2
• —NVIDIA for ModelOpt quantization tools
• —TorchAO for quantized tensor runtime