Update models/loaders/sam2_loader.py

models/loaders/sam2_loader.py

@@ -107,6 +107,7 @@ 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,10 +116,78 @@ def _load_official(self) -> Optional[Any]:
             predictor.model.eval()
             
         # Set device attribute for the predictor
-        if hasattr(predictor, "device"):
-            predictor.device = self.device
-            
-        return 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)
     
     def _load_transformers(self) -> Optional[Any]:
         """Load using transformers library"""