Update app.py

app.py

@@ -27,44 +27,37 @@ sam_model = SamModel.from_pretrained("facebook/sam-vit-base", device_map="auto",
 sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
 
 @spaces.GPU
-def predict_masks_and_scores_sam_hq(raw_image, input_points=None, input_boxes=None):
+def predict_masks_and_scores(model_id, raw_image, input_points=None, input_boxes=None):
     if input_boxes is not None:
         input_boxes = [input_boxes]
 
-    inputs = sam_hq_processor(raw_image, input_boxes=input_boxes, input_points=input_points, return_tensors="pt")
+    if model_id == 'sam':
+        inputs = sam_processor(raw_image, input_boxes=input_boxes, input_points=input_points, return_tensors="pt")
+    else:
+        inputs = sam_hq_processor(raw_image, input_boxes=input_boxes, input_points=input_points, return_tensors="pt")
     original_sizes = inputs["original_sizes"]
     reshaped_sizes = inputs["reshaped_input_sizes"]
 
-    inputs = inputs.to(sam_hq_model.device)
-
-    with torch.no_grad():
-        outputs = sam_hq_model(**inputs)
-
-    masks = sam_hq_processor.image_processor.post_process_masks(
-        outputs.pred_masks.cpu(), original_sizes, reshaped_sizes
-    )
+    if model_id == 'sam':
+        inputs = inputs.to(sam_model.device)
+        with torch.no_grad():
+            outputs = sam_model(**inputs)
+    else:
+        inputs = inputs.to(sam_hq_model.device)
+        with torch.no_grad():
+            outputs = sam_hq_model(**inputs)
+
+    if model_id == 'sam':
+        masks = sam_processor.image_processor.post_process_masks(
+            outputs.pred_masks.cpu(), original_sizes, reshaped_sizes
+        )
+    else:
+        masks = sam_hq_processor.image_processor.post_process_masks(
+            outputs.pred_masks.cpu(), original_sizes, reshaped_sizes
+        )
     scores = outputs.iou_scores
     return masks, scores
 
-@spaces.GPU
-def predict_masks_and_scores_sam(raw_image, input_points=None, input_boxes=None):
-    if input_boxes is not None:
-        input_boxes = [input_boxes]
-
-    inputs = sam_processor(raw_image, input_boxes=input_boxes, input_points=input_points, return_tensors="pt")
-    original_sizes = inputs["original_sizes"]
-    reshaped_sizes = inputs["reshaped_input_sizes"]
-
-    inputs = inputs.to(sam_model.device)
-
-    with torch.no_grad():
-        outputs = sam_model(**inputs)
-
-    masks = sam_processor.image_processor.post_process_masks(
-        outputs.pred_masks.cpu(), original_sizes, reshaped_sizes
-    )
-    scores = outputs.iou_scores
-    return masks, scores
 
 
 def process_inputs(prompts):
@@ -88,8 +81,8 @@ def process_inputs(prompts):
     input_points = [input_points] if input_points else None
     user_image = prompts['image']
 
-    sam_masks, sam_scores = predict_masks_and_scores_sam(user_image, input_boxes=input_boxes, input_points=input_points)
-    sam_hq_masks, sam_hq_scores = predict_masks_and_scores_sam_hq(user_image, input_boxes=input_boxes, input_points=input_points)
+    sam_masks, sam_scores = predict_masks_and_scores('sam', user_image, input_boxes=input_boxes, input_points=input_points)
+    sam_hq_masks, sam_hq_scores = predict_masks_and_scores('sam_hq', user_image, input_boxes=input_boxes, input_points=input_points)
 
     if input_boxes and input_points:
         img1_b64 = show_all_annotations_on_image_base64(user_image, sam_masks[0][0], sam_scores[:, 0, :], input_boxes[0], input_points[0], model_name='SAM')