forked from EngineX-Ascend/enginex-ascend-910-diffusers
add ascend support
This commit is contained in:
5
test.py
5
test.py
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user