from diffusers import DiffusionPipeline, TorchAoConfig, PipelineQuantizationConfig
from torchao.quantization import Int8WeightOnlyConfig
import torch 
ckpt_id = "black-forest-labs/FLUX.1-dev"
pipeline_quant_config = PipelineQuantizationConfig(
 quant_mapping={"transformer": TorchAoConfig(Int8WeightOnlyConfig())}
)
pipe = DiffusionPipeline.from_pretrained(
 ckpt_id, quantization_config=pipeline_quant_config, torch_dtype=torch.bfloat16
).to("cuda")
_ = pipe("dog", num_inference_steps=2)