add ascend support

This commit is contained in:
root
2025-09-03 10:16:24 +08:00
parent 33c57acb1b
commit 31c16831b1
4 changed files with 26 additions and 12 deletions

View File

@@ -5,9 +5,10 @@ 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)
pipeline.to("cuda")
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")
image.save("squirrel_picasso.png")