Flux-Schnell Optimum Quanto

run.py

from optimum.quanto import freeze, qfloat8, quantize
from diffusers import FluxPipeline
import torch
import time

seed=1337
generator = torch.Generator("cuda").manual_seed(seed)

pipeline = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to("cuda")

# Uncomment to quantize the model
# quantize(pipeline.transformer, weights=qfloat8)
# freeze(pipeline.transformer)
# quantize(pipeline.text_encoder, weights=qfloat8)
# freeze(pipeline.text_encoder)

t1 = time.time()
image = pipeline(
    prompt="ghibli style, a fantasy landscape with castles",
    guidance_scale=0.0,
    output_type="pil",
    num_inference_steps=4,
    max_sequence_length=256,
    generator=generator
).images[0]
t2 = time.time()
print(f"Time taken: {t2-t1} seconds")

image.save("output.png")