stable diffusion 01

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')