flux_fp8.py

import torch
from optimum.quanto import QuantizedDiffusersModel, freeze, qfloat8, quantize

from diffusers import FluxPipeline, FluxTransformer2DModel


class QuantizedFluxTransformer2DModel(QuantizedDiffusersModel):
    base_class = FluxTransformer2DModel


dtype = torch.bfloat16

transformer = FluxTransformer2DModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev", subfolder="transformer", torch_dtype=dtype
)

quantize(transformer, weights=qfloat8)
freeze(transformer)

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev", transformer=None, torch_dtype=dtype
)

pipe.transformer = transformer
pipe.enable_model_cpu_offload()

generator = torch.Generator().manual_seed(12345)
image = pipe(
    prompt="high quality photo of a dog sitting beside a tree with a red soccer ball at its side, the background it's a lake with a boat sailing in it and an airplane flying in the cloudy sky.",
    width=1024,
    height=1024,
    num_inference_steps=20,
    generator=generator,
    guidance_scale=3.5,
).images[0]
image.save("fluxfp8_image.png")

Thanks a lot for investigating into this, I really appreciate it.