2025-08-20 17:53:15 +08:00
|
|
|
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)
|
2025-09-03 10:16:24 +08:00
|
|
|
device = "cuda" if torch.cuda.is_available() else "npu" if torch.npu.is_available() else "cpu"
|
|
|
|
|
pipeline.to(device)
|
2025-08-20 17:53:15 +08:00
|
|
|
start = time.time()
|
|
|
|
|
image = pipeline("An image of a squirrel in Picasso style").images[0]
|
|
|
|
|
end = time.time()
|
|
|
|
|
print(f"elapsed: {end - start}")
|
2025-09-03 10:16:24 +08:00
|
|
|
image.save("squirrel_picasso.png")
|
