Update app.py

app.py

@@ -16,45 +16,32 @@ from utils import *
 from PIL import Image
 from gradio_image_prompter import ImagePrompter
 
-#sam_hq_model = SamHQModel.from_pretrained("syscv-community/sam-hq-vit-huge")
-#sam_hq_processor = SamHQProcessor.from_pretrained("syscv-community/sam-hq-vit-huge")
-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-huge")
-#sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
-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':
-        inputs = sam_processor(raw_image, input_boxes=input_boxes, input_points=input_points, return_tensors="pt")
+        model = SamModel.from_pretrained("facebook/sam-vit-base").to("cuda")
+        processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
     else:
-        inputs = sam_hq_processor(raw_image, input_boxes=input_boxes, input_points=input_points, return_tensors="pt")
+        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")
+
     original_sizes = inputs["original_sizes"]
     reshaped_sizes = inputs["reshaped_input_sizes"]
+    inputs = inputs.to("cuda")
 
-    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)
+    with torch.no_grad():
+        outputs = 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
-        )
+    masks = processor.image_processor.post_process_masks(outputs.pred_masks.cpu(), original_sizes, reshaped_sizes)
     scores = outputs.iou_scores
     return masks, scores