Update app.py

app.py

@@ -10,6 +10,27 @@ 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__)))
 ASSETS_DIR = os.path.join(os.path.dirname(__file__), "assets")