summaryrefslogtreecommitdiff
path: root/diffusion/tutorials/01_tut.py
blob: beb6819c2910d86ecd5775be7d18ee5c3eac7b98 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

from torch import autocast
from diffusers import StableDiffusionPipeline
from PIL import Image
prompt = 'a photo of an astronaut riding a horse on mars'

pipeline = StableDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4',
                                                   use_auth_token=True,
                                                   revision='fp16'
                                                   ).to('cuda')
with autocast('cuda'):
    output = pipeline(prompt)

print(output)
img = output['images'][0]
img.show()
img.save('./output/fp16.png')