| #!/usr/bin/env python3 | |
| from diffusers import DiffusionPipeline | |
| import torch | |
| path = "runwayml/stable-diffusion-v1-5" | |
| run_compile = False # Set True / False | |
| pipe = DiffusionPipeline.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 | |