mem split fix

requirements.txt

@@ -12,5 +12,5 @@ safetensors
 matplotlib
 torchvision
 pydantic==2.10.6
-git+https://github.com/facebookresearch/sam2.git
+sam2
 gradio_image_prompter

sam2_mask.py

@@ -1,5 +1,5 @@
 # K-I-S-S
-
+import spaces
 import gradio as gr
 from gradio_image_prompter import ImagePrompter
 from sam2.sam2_image_predictor import SAM2ImagePredictor
@@ -10,10 +10,12 @@ from PIL import Image as PILImage
 # Initialize SAM2 predictor
 MODEL = "facebook/sam2.1-hiera-large"
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-PREDICTOR = SAM2ImagePredictor.from_pretrained(MODEL, device=DEVICE)
 
+@spaces.GPU()
 def predict_masks(image, points):
     """Predict a single mask from the image based on selected points."""
+    global PREDICTOR
+    PREDICTOR = SAM2ImagePredictor.from_pretrained(MODEL, device=DEVICE)
     image_np = np.array(image)
     points_list = [[point["x"], point["y"]] for point in points]
     input_labels = [1] * len(points_list)