Update sam2_mask.py

sam2_mask.py

@@ -111,6 +111,7 @@ def process_mask(mask, expand_contract_px, expand, feathering_enabled, feather_s
         mask = feather_mask(mask, feather_size)
     return mask
 
+@spaces.GPU()
 def sam_process(input_image, checkpoint, tracking_points, trackings_input_label, expand_contract_px, expand, feathering_enabled, feather_size):
     image = Image.open(input_image)
     image = np.array(image.convert("RGB"))
@@ -123,7 +124,7 @@ def sam_process(input_image, checkpoint, tracking_points, trackings_input_label,
     # sam2_checkpoint, model_cfg = checkpoint_map[checkpoint]
     # Use CPU for both model and computations
     # sam2_model = build_sam2(model_cfg, sam2_checkpoint, device="cpu")
-    predictor = SAM2ImagePredictor.from_pretrained(sam21_hfmap[checkpoint], device="cpu")
+    predictor = SAM2ImagePredictor.from_pretrained(sam21_hfmap[checkpoint], device="cuda")
 
     # predictor = SAM2ImagePredictor(sam2_model)
     predictor.set_image(image)
@@ -152,8 +153,7 @@ with gr.Blocks() as demo:
     tracking_points = gr.State([])
     trackings_input_label = gr.State([])
     with gr.Column():
-        gr.Markdown("# SAM2 Image Predictor (CPU Version)")
-        gr.Markdown("This version runs entirely on CPU")
+        gr.Markdown("# SAM2 Image Predictor / Masking Assistant")
         with gr.Row():
             with gr.Column():
                 input_image = gr.Image(label="input image", interactive=False, type="filepath", visible=False)