summaryrefslogtreecommitdiff
path: root/diffusion/tutorials/01_tut.py
diff options
context:
space:
mode:
Diffstat (limited to 'diffusion/tutorials/01_tut.py')
-rw-r--r--diffusion/tutorials/01_tut.py17
1 files changed, 17 insertions, 0 deletions
diff --git a/diffusion/tutorials/01_tut.py b/diffusion/tutorials/01_tut.py
new file mode 100644
index 0000000..beb6819
--- /dev/null
+++ b/diffusion/tutorials/01_tut.py
@@ -0,0 +1,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')