import torch,time
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained("./FLUX.1-dev/", torch_dtype=torch.bfloat16)
pipe.enable_model_cpu_offload()

prompt = "a tiny astronaut hatching from an egg on the moon"
for i in range(3):
    start = time.time()
    out = pipe(
        prompt=prompt,
        guidance_scale=3.5,
        height=768,
        width=1360,
        num_inference_steps=50,
    ).images[0]
    print(f"Time:{time.time()-start}")
out.save("image.png")
