Update models/sam2_loader.py

models/sam2_loader.py

@@ -1,7 +1,8 @@
 #!/usr/bin/env python3
 """
-SAM2 Loader with T4-optimized predictor wrapper
+SAM2 Loader with Hugging Face Hub integration
 Provides SAM2Predictor class with memory management and optimization features
+Updated to use Hugging Face Hub models instead of direct downloads
 """
 
 import os
@@ -27,109 +28,100 @@ def __init__(self, device: torch.device, model_size: str = "small"):
         self._load_predictor()
         
     def _load_predictor(self):
-        """Load SAM2 predictor with optimizations"""
+        """Load SAM2 predictor with Hugging Face Hub integration"""
         try:
             from sam2.build_sam import build_sam2_video_predictor
             
-            # Download checkpoint if needed
-            checkpoint_path = f"./checkpoints/sam2_hiera_{self.model_size}.pt"
-            if not self._ensure_checkpoint(checkpoint_path):
-                raise RuntimeError(f"Failed to get SAM2 {self.model_size} checkpoint")
+            # Get checkpoint from Hugging Face Hub
+            checkpoint_path = self._get_hf_checkpoint()
+            if not checkpoint_path:
+                raise RuntimeError(f"Failed to get SAM2 {self.model_size} checkpoint from HF Hub")
+            
+            # Get model config
+            model_cfg = self._get_model_config()
             
             # Build predictor
-            model_cfg = f"sam2_hiera_{self.model_size[0]}.yaml"  # small -> s, base -> b, large -> l
             self.predictor = build_sam2_video_predictor(model_cfg, checkpoint_path, device=self.device)
             
             # Apply T4 optimizations
             self._optimize_for_t4()
             
-            logger.info(f"SAM2 {self.model_size} predictor loaded successfully")
+            logger.info(f"SAM2 {self.model_size} predictor loaded successfully from HF Hub")
             
         except ImportError as e:
             logger.error(f"SAM2 import failed: {e}")
-            raise RuntimeError("SAM2 not available - check third_party/sam2 installation")
+            raise RuntimeError("SAM2 not available - check sam2 installation")
         except Exception as e:
             logger.error(f"SAM2 loading failed: {e}")
             raise
     
-    def _ensure_checkpoint(self, checkpoint_path: str) -> bool:
-        """Ensure checkpoint exists, download if needed"""
-        checkpoint_file = Path(checkpoint_path)
-        
-        if checkpoint_file.exists():
-            file_size = checkpoint_file.stat().st_size / (1024**2)
-            if file_size > 50:  # At least 50MB
-                logger.info(f"SAM2 checkpoint exists: {file_size:.1f}MB")
-                return True
-            else:
-                logger.warning(f"Checkpoint too small ({file_size:.1f}MB), re-downloading")
-                checkpoint_file.unlink()
-        
-        return self._download_checkpoint(checkpoint_path)
-    
-    def _download_checkpoint(self, checkpoint_path: str, timeout_seconds: int = 600) -> bool:
-        """Download SAM2 checkpoint"""
+    def _get_hf_checkpoint(self) -> Optional[str]:
+        """Download checkpoint from Hugging Face Hub"""
         try:
-            logger.info(f"Downloading SAM2 {self.model_size} checkpoint...")
-            
-            checkpoint_file = Path(checkpoint_path)
-            checkpoint_file.parent.mkdir(parents=True, exist_ok=True)
+            from huggingface_hub import hf_hub_download
             
-            import requests
-            
-            # Checkpoint URLs
-            urls = {
-                "small": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_small.pt",
-                "base": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_base_plus.pt", 
-                "large": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt"
+            # Repository mapping for different model sizes
+            repo_mapping = {
+                "small": "facebook/sam2-hiera-small",
+                "base": "facebook/sam2-hiera-base-plus", 
+                "large": "facebook/sam2-hiera-large"
             }
             
-            if self.model_size not in urls:
-                raise ValueError(f"Unknown model size: {self.model_size}")
-            
-            checkpoint_url = urls[self.model_size]
-            
-            import time
-            start_time = time.time()
-            response = requests.get(checkpoint_url, stream=True, timeout=30)
-            response.raise_for_status()
-            
-            total_size = int(response.headers.get('content-length', 0))
+            filename_mapping = {
+                "small": "sam2_hiera_small.pt",
+                "base": "sam2_hiera_base_plus.pt",
+                "large": "sam2_hiera_large.pt"
+            }
             
-            temp_path = checkpoint_file.with_suffix('.download')
-            downloaded = 0
-            last_log = start_time
+            if self.model_size not in repo_mapping:
+                logger.error(f"Unknown model size: {self.model_size}")
+                return None
             
-            with open(temp_path, 'wb') as f:
-                for chunk in response.iter_content(chunk_size=1024*1024):
-                    if chunk:
-                        f.write(chunk)
-                        downloaded += len(chunk)
-                        
-                        current_time = time.time()
-                        if current_time - start_time > timeout_seconds:
-                            raise TimeoutError(f"Download timeout after {timeout_seconds}s")
-                        
-                        # Progress logging every 15 seconds
-                        if current_time - last_log > 15:
-                            progress = (downloaded / total_size * 100) if total_size > 0 else 0
-                            speed = downloaded / (current_time - start_time) / (1024**2)
-                            logger.info(f"Download: {progress:.1f}% ({speed:.1f}MB/s)")
-                            last_log = current_time
+            repo_id = repo_mapping[self.model_size]
+            filename = filename_mapping[self.model_size]
             
-            temp_path.rename(checkpoint_file)
+            logger.info(f"Downloading SAM2 {self.model_size} from HF Hub: {repo_id}")
             
-            download_time = time.time() - start_time
-            speed = downloaded / download_time / (1024**2)
-            logger.info(f"Download complete: {downloaded/(1024**2):.1f}MB in {download_time:.1f}s ({speed:.1f}MB/s)")
+            # Download from Hugging Face Hub
+            checkpoint_path = hf_hub_download(
+                repo_id=repo_id,
+                filename=filename,
+                cache_dir=None,  # Use default cache
+                force_download=False,  # Use cached version if available
+                token=None  # No auth token needed for public models
+            )
             
-            return True
+            logger.info(f"SAM2 checkpoint downloaded to: {checkpoint_path}")
+            return checkpoint_path
             
         except Exception as e:
-            logger.error(f"Checkpoint download failed: {e}")
+            logger.error(f"HF Hub download failed: {e}")
+            # Fallback to local checkpoint if HF download fails
+            return self._fallback_local_checkpoint()
+    
+    def _fallback_local_checkpoint(self) -> Optional[str]:
+        """Fallback to local checkpoint files"""
+        try:
+            checkpoint_path = f"./checkpoints/sam2_hiera_{self.model_size}.pt"
             if Path(checkpoint_path).exists():
-                Path(checkpoint_path).unlink()
-            return False
+                logger.info(f"Using local checkpoint: {checkpoint_path}")
+                return checkpoint_path
+            else:
+                logger.error(f"Local checkpoint not found: {checkpoint_path}")
+                return None
+        except Exception as e:
+            logger.error(f"Local checkpoint fallback failed: {e}")
+            return None
+    
+    def _get_model_config(self) -> str:
+        """Get the appropriate model config file"""
+        config_mapping = {
+            "small": "sam2_hiera_s.yaml",
+            "base": "sam2_hiera_b+.yaml", 
+            "large": "sam2_hiera_l.yaml"
+        }
+        
+        return config_mapping.get(self.model_size, "sam2_hiera_s.yaml")
     
     def _optimize_for_t4(self):
         """Apply T4-specific optimizations"""
@@ -175,6 +167,36 @@ def add_new_points(self, inference_state, frame_idx: int, obj_id: int,
             logger.error(f"Failed to add new points: {e}")
             raise
     
+    def add_new_points_or_box(self, inference_state, frame_idx: int, obj_id: int, 
+                             points: np.ndarray, labels: np.ndarray, clear_old_points: bool = True):
+        """Add new points or box for tracking (newer SAM2 API)"""
+        if self.predictor is None:
+            raise RuntimeError("Predictor not loaded")
+        
+        try:
+            # Try the newer API first
+            if hasattr(self.predictor, 'add_new_points_or_box'):
+                return self.predictor.add_new_points_or_box(
+                    inference_state=inference_state,
+                    frame_idx=frame_idx,
+                    obj_id=obj_id,
+                    points=points,
+                    labels=labels,
+                    clear_old_points=clear_old_points
+                )
+            else:
+                # Fallback to older API
+                return self.predictor.add_new_points(
+                    inference_state=inference_state,
+                    frame_idx=frame_idx,
+                    obj_id=obj_id,
+                    points=points,
+                    labels=labels
+                )
+        except Exception as e:
+            logger.error(f"Failed to add new points or box: {e}")
+            raise
+    
     def propagate_in_video(self, inference_state, scale: float = 1.0, **kwargs):
         """Propagate through video with optional scaling"""
         if self.predictor is None: