Update app.py

app.py

@@ -27,20 +27,40 @@ 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(model, processor, raw_image, input_points=None, input_boxes=None):
+def predict_masks_and_scores_sam_hq(raw_image, input_points=None, input_boxes=None):
     if input_boxes is not None:
         input_boxes = [input_boxes]
 
-    inputs = processor(raw_image, input_boxes=input_boxes, input_points=input_points, return_tensors="pt")
+    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(model.device)
+    inputs = inputs.to(sam_hq_model.device)
 
     with torch.no_grad():
-        outputs = model(**inputs)
+        outputs = sam_hq_model(**inputs)
 
-    masks = processor.image_processor.post_process_masks(
+    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
@@ -68,8 +88,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_model, sam_processor, user_image, input_boxes=input_boxes, input_points=input_points)
-    sam_hq_masks, sam_hq_scores = predict_masks_and_scores(sam_hq_model, sam_hq_processor, 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')