move all sam to cpu

app.py

@@ -107,7 +107,6 @@ except Exception as e:
     raise gr.Error(f"Failed to initialize Blurnonymizer. Check logs. Error: {e}")
 
 # --- Core Processing Function ---
-@spaces.GPU(duration=20)  # add this so that the sam segmentation runs on the gpu
 def anonymise_image(input_image_np: np.ndarray, boxes: list[BoundingBox]):
     """Calls the blurnonymizer instance to censor the image."""
     if not blurnonymizer_instance:

blurnonymize.py

@@ -156,8 +156,9 @@ class ImageBlurnonymizer:
         if self.predictor is not None and not force:
             return
 
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large", device=self.device)
+        # self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = "cpu"
+        self.predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2.1-hiera-small", device=self.device)
 
     @staticmethod
     def _smoothen_mask(mask: np.ndarray):
@@ -216,9 +217,9 @@ class ImageBlurnonymizer:
         if self.predictor is None:
             raise Exception("[-] sam has not been initialized")
 
-        if torch.cuda.is_available() and self.device == "cpu":
-            # class instance was wrongly initialized to run on cpu, but gpu is avaiable
-            self.init_sam(force=True)
+        # if torch.cuda.is_available() and self.device == "cpu":
+        #     # class instance was wrongly initialized to run on cpu, but gpu is avaiable
+        #     self.init_sam(force=True)
 
         x_min, y_min, x_max, y_max = bbox
         x_width = x_max - x_min
@@ -242,7 +243,8 @@ class ImageBlurnonymizer:
         # Ensure points are valid coordinates (e.g., non-negative)
         points = [[max(0, p[0]), max(0, p[1])] for p in points]
 
-        with torch.inference_mode(), torch.autocast(self.device, dtype=torch.bfloat16):
+        # with torch.inference_mode(), torch.autocast(self.device, dtype=torch.bfloat16):
+        with torch.inference_mode(), torch.autocast(self.device):
             self.predictor.set_image(image)
             masks, scores, _ = self.predictor.predict(
                 box=np.array(bbox),  # Predictor might expect numpy array