from diffusers import DiffusionPipeline
import torch
import time

model_path = "/mnt/contest_ceph/zhanghao/models/stable-diffusion-v1-5"
# model_path = "/mnt/contest_ceph/zhanghao/models/stable-diffusion-3.5-medium"
pipeline = DiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float16)
device = "cuda" if torch.cuda.is_available() else "npu" if torch.npu.is_available() else "cpu"
pipeline.to(device)
start = time.time()
image = pipeline("An image of a squirrel in Picasso style").images[0]
end = time.time()
print(f"elapsed: {end - start}")
image.save("squirrel_picasso.png")