Update app.py

app.py

@@ -17,31 +17,41 @@ from PIL import Image
 from gradio_image_prompter import ImagePrompter
 
 
-#sam_hq_model = SamHQModel.from_pretrained("syscv-community/sam-hq-vit-base", device_map="auto", torch_dtype="auto")
-#sam_hq_processor = SamHQProcessor.from_pretrained("syscv-community/sam-hq-vit-base")
-#sam_model = SamModel.from_pretrained("facebook/sam-vit-base", device_map="auto", torch_dtype="auto")
-#sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
+sam_hq_model = SamHQModel.from_pretrained("syscv-community/sam-hq-vit-base", device_map="auto", torch_dtype="auto")
+sam_hq_processor = SamHQProcessor.from_pretrained("syscv-community/sam-hq-vit-base")
+sam_model = SamModel.from_pretrained("facebook/sam-vit-base", device_map="auto", torch_dtype="auto")
+sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
 
 @spaces.GPU
 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]
+
     if model_id == 'sam':
-        model = SamModel.from_pretrained("facebook/sam-vit-base").to("cuda")
-        processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
+        inputs = sam_processor(raw_image, input_boxes=input_boxes, input_points=input_points, return_tensors="pt")
     else:
-        model = SamHQModel.from_pretrained("syscv-community/sam-hq-vit-base").to("cuda")
-        processor = SamHQProcessor.from_pretrained("syscv-community/sam-hq-vit-base")
-
-    inputs = processor(raw_image, input_boxes=[input_boxes] if input_boxes else None,
-                       input_points=[input_points] if input_points else None, 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("cuda")
 
-    with torch.no_grad():
-        outputs = model(**inputs)
+    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)
 
-    masks = processor.image_processor.post_process_masks(outputs.pred_masks.cpu(), original_sizes, reshaped_sizes)
+    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