changes for gpu usage

README.md

@@ -4,7 +4,7 @@ emoji: 🏢
 colorFrom: green
 colorTo: red
 sdk: gradio
-sdk_version: 3.48.0
+sdk_version: 4.0.0
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -1,5 +1,5 @@
+import spaces
 import gradio as gr
-import sys
 import torch
 from omegaconf import OmegaConf
 from PIL import Image
@@ -45,6 +45,7 @@ def combine_masks(uploaded_mask, sketched_mask):
     else:
         raise ValueError("Please provide a mask")
 
+@spaces.GPU
 def remove_obj(image, uploaded_mask, seed):
     image_pil, sketched_mask = image["image"], image["mask"]
     mask = dilate_mask(combine_masks(uploaded_mask, sketched_mask))