Update app.py

app.py

@@ -10,7 +10,25 @@ from matplotlib.patches import Patch
 from PIL import Image, ImageOps
 from torchvision import transforms
 
+import spaces
+from diffusers import DiffusionPipeline
 
+pipe = DiffusionPipeline.from_pretrained(...)
+pipe.to('cuda')
+
+[email protected]
+def generate(prompt):
+    return pipe(prompt).images
+
+gr.Interface(
+    fn=generate,
+    inputs=gr.Text(),
+    outputs=gr.Gallery(),
+).launch()
+
+@spaces.GPU(duration=120)
+def generate(prompt):
+   return pipe(prompt).images
 
 # ----------------- HELPER FUNCTIONS ----------------- #
 os.chdir(os.path.dirname(os.path.abspath(__file__)))