migrating SAM2 space from T4 to ZERO

app.py

@@ -9,6 +9,7 @@ from gradio_image_prompter import ImagePrompter
 
 from utils.models import load_models, CHECKPOINT_NAMES, MODE_NAMES, \
     MASK_GENERATION_MODE, BOX_PROMPT_MODE
+import spaces
 
 MARKDOWN = """
 # Segment Anything Model 2 🔥
@@ -37,11 +38,20 @@ EXAMPLES = [
     ["tiny", MASK_GENERATION_MODE, "https://media.roboflow.com/notebooks/examples/dog-4.jpeg", None],
 ]
 
-DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+DEVICE = torch.device('cuda')
+
+torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
+if torch.cuda.get_device_properties(0).major >= 8:
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
 MASK_ANNOTATOR = sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX)
 IMAGE_PREDICTORS, MASK_GENERATORS = load_models(device=DEVICE)
 
 
+@spaces.GPU
+@torch.inference_mode()
+@torch.autocast(device_type="cuda", dtype=torch.bfloat16)
 def process(
     checkpoint_dropdown,
     mode_dropdown,
@@ -147,4 +157,4 @@ with gr.Blocks() as demo:
         outputs=[image_output_component]
     )
 
-demo.launch(debug=False, show_error=True, max_threads=1)
+demo.launch(debug=False, show_error=True)

requirements.txt

@@ -1,3 +1,4 @@
+spaces
 samv2
 gradio
 supervision

utils/models.py

@@ -6,9 +6,9 @@ from sam2.build_sam import build_sam2
 from sam2.sam2_image_predictor import SAM2ImagePredictor
 
 BOX_PROMPT_MODE = "box prompt"
+POINT_PROMPT_MODE = "point prompt"
 MASK_GENERATION_MODE = "mask generation"
-VIDEO_SEGMENTATION_MODE = "video segmentation"
-MODE_NAMES = [BOX_PROMPT_MODE, MASK_GENERATION_MODE]
+MODE_NAMES = [BOX_PROMPT_MODE, POINT_PROMPT_MODE, MASK_GENERATION_MODE]
 
 CHECKPOINT_NAMES = ["tiny", "small", "base_plus", "large"]
 CHECKPOINTS = {