import torch
import time
from diffusers import StableDiffusion3Pipeline
#from xformers.ops import MemoryEfficientAttentionFlashAttentionOp, MemoryEfficientAttentionTritonFwdFlashBwOp

pipe = StableDiffusion3Pipeline.from_pretrained("/public/home/lijian/model/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)
pipe = pipe.to("cuda")

prompt = "A cat holding a sign that says hello world"
for i in range(3):
    start = time.time()
    image = pipe(prompt, negative_prompt="", num_inference_steps=28, guidance_scale=7.0,).images[0]
    print(f"All time: {time.time()-start}")
    image.save(f"SD3-medium.png")