Update app.py

app.py

@@ -4,17 +4,25 @@ import torch
 from PIL import Image
 from transformers import SamModel, SamProcessor
 from gradio_image_prompter import ImagePrompter
-
+import spaces
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-sam_model = SamModel.from_pretrained("facebook/sam-vit-huge").to(device)
+sam_model = SamModel.from_pretrained("facebook/sam-vit-huge").to("cuda")
 sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
-slimsam_model = SamModel.from_pretrained("nielsr/slimsam-50-uniform").to(device)
+slimsam_model = SamModel.from_pretrained("nielsr/slimsam-50-uniform").to("cuda")
 slimsam_processor = SamProcessor.from_pretrained("nielsr/slimsam-50-uniform")
 
-def sam_box_inference(image, model, x_min, y_min, x_max, y_max):
+def get_processor_and_model(slim: bool):
+    if slim:
+        return slimsam_processor, slimsam_model
+    return sam_processor, sam_model
+
+@spaces.GPU
+def sam_box_inference(image, x_min, y_min, x_max, y_max, *, slim=False):
 
-    inputs = sam_processor(
+    processor, model = get_processor_and_model(slim)
+    
+    inputs = processor(
         Image.fromarray(image),
         input_boxes=[[[[x_min, y_min, x_max, y_max]]]],
         return_tensors="pt"
@@ -23,7 +31,7 @@ def sam_box_inference(image, model, x_min, y_min, x_max, y_max):
     with torch.no_grad():
         outputs = model(**inputs)
 
-    mask = sam_processor.image_processor.post_process_masks(
+    mask = processor.image_processor.post_process_masks(
         outputs.pred_masks.cpu(),
         inputs["original_sizes"].cpu(),
         inputs["reshaped_input_sizes"].cpu()
@@ -33,17 +41,20 @@ def sam_box_inference(image, model, x_min, y_min, x_max, y_max):
     print(mask.shape)
     return [(mask, "mask")]
 
+@spaces.GPU
+def sam_point_inference(image, x, y, *, slim=False):
 
-def sam_point_inference(image, model, x, y):
-    inputs = sam_processor(
+    processor, model = get_processor_and_model(slim)
+    
+    inputs = processor(
         image,
         input_points=[[[x, y]]],
         return_tensors="pt").to(device)
 
     with torch.no_grad():
-        outputs = sam_model(**inputs)
+        outputs = model(**inputs)
 
-    mask = sam_processor.post_process_masks(
+    mask = processor.post_process_masks(
         outputs.pred_masks.cpu(),
         inputs["original_sizes"].cpu(),
         inputs["reshaped_input_sizes"].cpu()
@@ -72,8 +83,8 @@ def infer_point(img):
         center_x = int(np.mean(nonzero_indices[1]))
         center_y = int(np.mean(nonzero_indices[0]))
     print("Point inference returned.")
-    return ((image, sam_point_inference(image, slimsam_model, center_x, center_y)),
-    (image, sam_point_inference(image, sam_model, center_x, center_y)))
+    return ((image, sam_point_inference(image, center_x, center_y, slim=True)),
+    (image, sam_point_inference(image, center_x, center_y)))
 
 def infer_box(prompts):
     # background (original image) layers[0] ( point prompt) composite (total image)
@@ -86,8 +97,8 @@ def infer_box(prompts):
     print(points)
 
     # x_min = points[0] x_max = points[3] y_min = points[1] y_max = points[4]
-    return ((image, sam_box_inference(image, slimsam_model, points[0], points[1], points[3], points[4])),
-    (image, sam_box_inference(image, sam_model, points[0], points[1], points[3], points[4])))
+    return ((image, sam_box_inference(image, points[0], points[1], points[3], points[4], slim=True)),
+    (image, sam_box_inference(image, points[0], points[1], points[3], points[4])))
 with gr.Blocks(title="SlimSAM") as demo:
   gr.Markdown("# SlimSAM")
   gr.Markdown("SlimSAM is the pruned-distilled version of SAM that is smaller.")