Views
No views yet
This repo does not contain the full pipeline. Text encoders, VAE, and scheduler are loaded from black-forest-labs/FLUX.2-klein-4B.
| Subfolder | Precision | Format | Size | Use case |
|---|---|---|---|---|
transformer_bf16/ | bfloat16 | safetensors | ~7.7 GB | LoRA training, evaluation baselines, re-quantization |
transformer_fp8_static/ | float8_e4m3fn | torchao .pt | ~3.9 GB | Production inference (~2x memory saving) |
input_scale values from BFL's calibration. The checkpoint
is a torch.save dict containing:state_dict — torchao AffineQuantizedTensor weightsact_scales — per-Linear static activation scales (float32)fp8_dtype — "float8_e4m3fn"1from diffusers import Flux2Transformer2DModel, Flux2KleinPipeline
2from PIL import Image
3import torch
4
5# Load transformer (bf16)
6transformer = Flux2Transformer2DModel.from_pretrained(
7 "photoroom/FLUX.2-klein-4b-fp8-diffusers",
8 subfolder="transformer_bf16",
9 torch_dtype=torch.bfloat16,
10).to("cuda")
11
12# Load pipeline (text encoders, VAE, scheduler from BFL)
13pipe = Flux2KleinPipeline.from_pretrained(
14 "black-forest-labs/FLUX.2-klein-4B",
15 transformer=transformer,
16 torch_dtype=torch.bfloat16,
17)
18
19# Run inference
20image = Image.open("input.png").convert("RGB")
21result = pipe(
22 prompt="a product on a marble countertop",
23 image=[image],
24 height=1024,
25 width=1024,
26 guidance_scale=1.0,
27 num_inference_steps=4,
28 generator=torch.Generator(device="cuda").manual_seed(42),
29).images[0]
30result.save("output.png")1from diffusers import Flux2Transformer2DModel, Flux2KleinPipeline
2from huggingface_hub import hf_hub_download
3from load_torchao import load_torchao_fp8_static_model
4from PIL import Image
5import torch
6
7# Load FP8 static transformer
8ckpt_path = hf_hub_download(
9 "photoroom/FLUX.2-klein-4b-fp8-diffusers",
10 filename="transformer_fp8_static/model_fp8_static.pt",
11)
12
13transformer = load_torchao_fp8_static_model(
14 ckpt_path=ckpt_path,
15 base_model_or_factory=lambda: Flux2Transformer2DModel.from_pretrained(
16 "photoroom/FLUX.2-klein-4b-fp8-diffusers",
17 subfolder="transformer_bf16",
18 torch_dtype=torch.bfloat16,
19 ),
20 device="cuda",
21)
22
23# Load pipeline (text encoders, VAE, scheduler from BFL)
24pipe = Flux2KleinPipeline.from_pretrained(
25 "black-forest-labs/FLUX.2-klein-4B",
26 transformer=transformer,
27 torch_dtype=torch.bfloat16,
28)
29
30# Run inference
31# image = Image.open("input.png").convert("RGB")
32result = pipe(
33 prompt="a cat holding a frame with FP8 writing on it",
34 image=[None],
35 height=1024,
36 width=1024,
37 guidance_scale=1.0,
38 num_inference_steps=4,
39 generator=torch.Generator(device="cuda").manual_seed(42),
40).images[0]
41result.save("output.png")