Update app.py

app.py

@@ -1,3 +1,18 @@
+# pip install diffusers transformers torch
+
+from diffusers import StableDiffusionPipeline
+import torch
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+checkpoint = "Deci/DeciDiffusion-v1-0"
+pipeline = StableDiffusionPipeline.from_pretrained(checkpoint, custom_pipeline=checkpoint, torch_dtype=torch.float16)
+pipeline.unet = pipeline.unet.from_pretrained(checkpoint, subfolder='flexible_unet', torch_dtype=torch.float16)
+
+pipeline = pipeline.to(device)
+
+img = pipeline(prompt=['A photo of an astronaut riding a horse on Mars']).images[0]
+
 import gradio as gr
 
 gr.Interface.load("models/Deci/DeciDiffusion-v1-0").launch()