fallpt-chat / app.py
lelafav502's picture
Update app.py
c37dede
raw
history blame
527 Bytes
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