|
|
|
from diffusers import StableDiffusionPipeline |
|
import torch |
|
|
|
path = "runwayml/stable-diffusion-v1-5" |
|
|
|
run_compile = True |
|
|
|
pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16) |
|
pipe = pipe.to("cuda:0") |
|
pipe.unet.to(memory_format=torch.channels_last) |
|
|
|
if run_compile: |
|
print("Run torch compile") |
|
pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) |
|
|
|
prompt = "ghibli style, a fantasy landscape with castles" |
|
|
|
for _ in range(3): |
|
images = pipe(prompt=prompt).images |
|
|