Update app.py

app.py

@@ -86,7 +86,7 @@ except Exception as e:
 class CustomSAMWrapper:
     def __init__(self, model_path: str, device: str = DEVICE):
         # try:            
-        self.device = "cpu" #torch.device(device)
+        self.device = torch.device(device)
         sam_model = build_sam2("configs/sam2.1/sam2.1_hiera_l.yaml", model_path, self.device)
         sam_model = sam_model.to(self.device)
         self.predictor = SAM2ImagePredictor(sam_model)
@@ -231,7 +231,8 @@ def visualize_masks_on_image(
     blended = cv2.addWeighted(image_np, 1 - alpha, color_mask, alpha, 0)
     return PILImage.fromarray(blended)
 
-
+@spaces.GPU
+@torch.no_grad()
 def run_pipeline(image: PILImage.Image, prompt: str):
     if not model or not processor:
         return "Models not loaded. Please check logs.", None