Update app.py

app.py

@@ -15,9 +15,8 @@ import torch
 from diffusers import StableDiffusionBrushNetPipeline, BrushNetModel, UniPCMultistepScheduler
 import random
 import gradio as gr
-import spaces
 
-mobile_sam = sam_model_registry['vit_h'](checkpoint='data/ckpt/sam_vit_h_4b8939.pth')
+mobile_sam = sam_model_registry['vit_h'](checkpoint='data/ckpt/sam_vit_h_4b8939.pth').to("cuda")
 mobile_sam.eval()
 mobile_predictor = SamPredictor(mobile_sam)
 colors = [(255, 0, 0), (0, 255, 0)]
@@ -74,7 +73,6 @@ def resize_image(input_image, resolution):
     img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
     return img
 
-@spaces.GPU
 def process(input_image, 
     original_image, 
     original_mask, 
@@ -275,7 +273,7 @@ with block:
         for p, l in sel_pix:
             points.append(p)
             labels.append(l)
-        mobile_predictor=mobile_predictor.to("cuda")
+        mobile_predictor=mobile_predictor
         mobile_predictor.set_image(img if isinstance(img, np.ndarray) else np.array(img))
         with torch.no_grad():
             masks, _, _ = mobile_predictor.predict(point_coords=np.array(points), point_labels=np.array(labels), multimask_output=False)