abhishekchohan/flux2-klein-9b-fp8

FLUX.2 Klein 9B — FP8 Weight-Only Quantized (torchao)

An FP8 (float8wo) weight-only quantized version of black-forest-labs/FLUX.2-klein-9B. Both the diffusion transformer and the text encoder have their projection weights stored in FP8 to cut memory roughly in half; the sensitive I/O layers are kept in BF16 to preserve quality. Quantized with torchao and validated end-to-end on an NVIDIA L40S (quantize → save → reload → inference, plus a paired BF16-vs-FP8 fidelity eval).

What was quantized (and what was protected)

Transformer — only the attention/FFN nn.Linear layers inside the transformer blocks (transformer_blocks.* / single_transformer_blocks.*) were quantized (144 of 153 linears). The I/O boundary stays BF16: x_embedder / context_embedder, time & guidance embedders, adaLN/modulation, norm_out / proj_out, and all norms.

Text encoder — the attention/MLP projection linears were quantized to FP8, while the token embeddings and output head are protected in BF16.

VAE — left in BF16.

Memory footprint

Component	BF16	FP8 (this repo)
Transformer weights	~16.9 GB	~8.8 GB (1.92×)
Text encoder	~15.3 GB	~8.8 GB (1.74×)
Full pipeline dir	—	~17.75 GB

Quality

The quantized pipeline was validated to reload and generate coherent images. A paired BF16-vs-FP8 fidelity eval (transformer and text encoder quantized) on this model measured:

Metric	Value	Reading
LPIPS (perceptual)	0.131	low perceptual drift
PSNR	23.3 dB	good fidelity
SSIM	0.847	good structural similarity

The smaller 4B sibling (abhishekchohan/flux2-klein-4b-fp8) measured LPIPS 0.0995 / PSNR 23.76 dB / SSIM 0.902 on the same setup. If you need the absolute highest fidelity, prefer the 4B, or keep the 9B text encoder in BF16.

Side-by-side: BF16 (left) vs FP8 (right)

Identical prompts and seeds; the two columns are visually near-identical.

Prompt: a photo of a cat sitting on a windowsill at golden hour

Prompt: a cinematic portrait of an elderly fisherman, dramatic lighting

Prompt: a watercolor painting of a mountain lake at sunrise

Usage

import torch
from diffusers import Flux2KleinPipeline

pipe = Flux2KleinPipeline.from_pretrained(
    "abhishekchohan/flux2-klein-9b-fp8",
    torch_dtype=torch.bfloat16,
)
pipe.enable_model_cpu_offload()  # recommended at this size

image = pipe(
    prompt="a photo of a cat sitting on a windowsill at golden hour",
    num_inference_steps=8,
    guidance_scale=1.0,
    height=1024,
    width=1024,
).images[0]

⚠️ Loading note (pickle / FP8 tensor subclasses)

The quantized weights are stored as pickle .bin (not safetensors) because torchao's FP8 tensor subclasses cannot be serialized to safetensors. Loading therefore uses torch.load deserialization:

• HF Hub flags pickle files; your client may prompt you to trust/allow them.
• Load with the standard from_pretrained above — diffusers re-materializes the FP8 weights automatically. Requires a torchao install (pip install torchao).

Details

• Backend: torchao weight-only quantization (Float8WeightOnlyConfig, float8_e4m3fn).
• Compute dtype: bfloat16 activations, FP8 (e4m3) weights, dequantized on the fly.
• Scheduler/tokenizer/VAE: identical to the base model.

Limitations

• Weight-only FP8 reduces memory; it does not speed up inference on GPUs without native FP8 GEMM.
• Base model license/usage restrictions of black-forest-labs/FLUX.2-klein-9B still apply.