drop ZERO GPU

app.py

@@ -2,7 +2,6 @@ from typing import Optional
 
 import gradio as gr
 import numpy as np
-import spaces
 import supervision as sv
 import torch
 from PIL import Image
@@ -18,14 +17,15 @@ video by treating images as single-frame videos. Its design, a simple transforme
 architecture with streaming memory, enables real-time video processing.
 """
 
-DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+# DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+DEVICE = torch.device('cuda')
+
 MASK_ANNOTATOR = sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX)
 CHECKPOINT = "checkpoints/sam2_hiera_large.pt"
 CONFIG = "sam2_hiera_l.yaml"
 
 sam2_model = build_sam2(CONFIG, CHECKPOINT, device=DEVICE, apply_postprocessing=False)
 
-@spaces.GPU
 def process(image_input) -> Optional[Image.Image]:
     mask_generator = SAM2AutomaticMaskGenerator(sam2_model)
     image = np.array(image_input.convert("RGB"))

utils/models.py

@@ -0,0 +1,6 @@
+CHECKPOINTS = {
+    "tiny": ["sam2_hiera_t.yaml", "checkpoints/sam2_hiera_tiny.pt"],
+    "small": ["sam2_hiera_s.yaml", "checkpoints/sam2_hiera_small.pt"],
+    "base_plus": ["sam2_hiera_b+.yaml", "checkpoints/sam2_hiera_base_plus.pt"],
+    "large": ["sam2_hiera_l.yaml", "checkpoints/sam2_hiera_large.pt"],
+}