Upload app.py

app.py

@@ -6,6 +6,9 @@ import tempfile
 import numpy as np
 import torch as torch
 torch.backends.cuda.matmul.allow_tf32 = True
+import spaces
+import functools
+
 
 
 from diffusers import (
@@ -83,6 +86,12 @@ if __name__ == "__main__":
         t_start=0,
     ).to(device)
 
+    try:
+        import xformers
+        pipe.enable_xformers_memory_efficient_attention()
+    except:
+        pass  # run without xformers
+
     # Cache example images in memory
     example_images_dir = "files/image"
     example_images = []
@@ -93,7 +102,7 @@ if __name__ == "__main__":
 
     # Create a Gradio interface
     interface = gr.Interface(
-        fn=lambda image: process_image(pipe, vae_2, image),
+        fn=spaces.GPU(functools.partial(process_image, pipe, vae_2)),
         inputs=gr.Image(type="pil"),
         outputs=gr.Image(type="pil"),
         title="Dereflection Any Image",