Update models/loaders/sam2_loader.py

models/loaders/sam2_loader.py

@@ -99,7 +99,7 @@ def _determine_optimal_size(self) -> str:
         return "tiny"  # Conservative default
     
     def _load_official(self) -> Optional[Any]:
-        """Load using official SAM2 API"""
+        """Load using official SAM2 API - return directly without wrapper"""
         from sam2.sam2_image_predictor import SAM2ImagePredictor
         
         predictor = SAM2ImagePredictor.from_pretrained(
@@ -107,7 +107,6 @@ def _load_official(self) -> Optional[Any]:
             cache_dir=self.cache_dir,
             local_files_only=False,
             trust_remote_code=True,
-            device=self.device,  # Pass device directly
         )
         
         # Move to device and set to eval mode
@@ -115,79 +114,13 @@ def _load_official(self) -> Optional[Any]:
             predictor.model = predictor.model.to(self.device)
             predictor.model.eval()
             
-        # Set device attribute for the predictor
-        predictor.device = self.device
-        
-        # Wrap to ensure proper automatic mask generation
-        class SAM2Wrapper:
-            def __init__(self, predictor, device):
-                self.predictor = predictor
-                self.device = device
-                self._image_set = False
-                
-            def set_image(self, image):
-                """Set image for processing"""
-                self.predictor.set_image(image)
-                self._image_set = True
-                
-            def predict(self, point_coords=None, point_labels=None, box=None, **kwargs):
-                """Generate masks with automatic detection if no prompts given"""
-                if not self._image_set:
-                    # Auto-set image if not already done
-                    logger.warning("Image not set, returning empty mask")
-                    return {
-                        "masks": np.zeros((1, 512, 512), dtype=np.float32),
-                        "scores": np.array([0.0]),
-                        "logits": np.zeros((1, 512, 512), dtype=np.float32),
-                    }
-                
-                # If no prompts, generate automatic mask
-                if point_coords is None and box is None:
-                    # Use center point as default
-                    h, w = 512, 512  # Default size
-                    point_coords = np.array([[w//2, h//2]], dtype=np.float32)
-                    point_labels = np.array([1], dtype=np.int32)
-                
-                return self.predictor.predict(
-                    point_coords=point_coords,
-                    point_labels=point_labels,
-                    box=box,
-                    **kwargs
-                )
-                
-            def generate_automatic_masks(self, image):
-                """Generate masks automatically for the entire image"""
-                self.set_image(image)
-                # Generate with points in a grid
-                h, w = image.shape[:2]
-                points = []
-                labels = []
-                
-                # Create a grid of points
-                for y in range(h//4, h, h//2):
-                    for x in range(w//4, w, w//2):
-                        points.append([x, y])
-                        labels.append(1)
-                
-                if points:
-                    masks, scores, logits = self.predictor.predict(
-                        point_coords=np.array(points, dtype=np.float32),
-                        point_labels=np.array(labels, dtype=np.int32),
-                        multimask_output=True
-                    )
-                    
-                    # Return best mask
-                    if len(scores) > 0:
-                        best_idx = scores.argmax()
-                        return masks[best_idx], scores[best_idx]
-                
-                return np.ones((h, w), dtype=np.float32), 1.0
-                
-            def __getattr__(self, name):
-                """Forward other attributes to predictor"""
-                return getattr(self.predictor, name)
-        
-        return SAM2Wrapper(predictor, self.device)
+        # Set device attribute if it exists
+        if hasattr(predictor, "device"):
+            predictor.device = self.device
+            
+        # Return the predictor directly - no wrapper!
+        # The calling code expects the standard SAM2 interface
+        return predictor
     
     def _load_transformers(self) -> Optional[Any]:
         """Load using transformers library"""