File size: 527 Bytes
c37dede
 
db0d777
c37dede
db0d777
c37dede
db0d777
c37dede
db0d777
c37dede
db0d777
c37dede
 
db0d777
c37dede
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)

pipe = pipe.to("cuda")

prompt = "a photograph of an astronaut riding a horse"

image = pipe(prompt).images[0]  # image here is in [PIL format](https://pillow.readthedocs.io/en/stable/)

# Now to display an image you can either save it such as:
image.save(f"astronaut_rides_horse.png")

# or if you're in a google colab you can directly display it with 
image