abhishekchohan/flux2-klein-4b-fp8

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

An FP8 (float8wo) weight-only quantized version of `black-forest-labs/FLUX.2-klein-4B`. 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 (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 (100 of 109 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 & speed footprint

Measured quality (paired BF16 vs FP8)

Paired comparison with identical prompts and seeds (16 images: 8 prompts × 2 seeds), measuring the full denoising trajectory. The FP8 run quantizes both transformer and text encoder, so these numbers reflect the shipped model:

Interpreting LPIPS: < 0.1 ≈ near-lossless, ~0.18 = noticeable, ≥ 0.27 = visibly degraded (calibration from FLUX quantization literature).

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

Identical prompts and seeds; the two columns are visually indistinguishable at a glance.

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-4b-fp8",
    torch_dtype=torch.bfloat16,
)
pipe.enable_model_cpu_offload()  # or .to("cuda") if it fits

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

Evaluation protocol

Measured on an NVIDIA RTX 4090 with diffusers 0.39 / torchao 0.18 / torch 2.11 (cu128): 16 paired images (8 prompts × 2 seeds), identical prompts & seeds for BF16 and FP8, 1024×1024, 8 steps, guidance 1.0. Preference via PickScore, alignment via CLIP-ViT-L/14.

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; latency gains depend on native FP8 tensor cores (Ada/Hopper/Blackwell) and were measured on an RTX 4090.
• —Quality metrics are measured on a 16-image paired set; treat them as indicative.
• —Base model license/usage restrictions of black-forest-labs/FLUX.2-klein-4B still apply.