Update models/loaders/sam2_loader.py

models/loaders/sam2_loader.py

@@ -3,7 +3,7 @@
 
 """
 SAM2 Loader + Guarded Predictor Adapter (VRAM-friendly, shape-safe, thread-safe, PyTorch 2.x)
-- Official HF load: SAM2ImagePredictor.from_pretrained(...)
+- Uses traditional build_sam2 method with HF hub downloads for SAM 2.1 weights
 - Never assigns predictor.device (read-only) — moves .model to device instead
 - Accepts RGB/BGR, float/uint8; strips alpha; optional BGR→RGB via env
 - Downscale ladder on set_image(); upsample masks back to original H,W
@@ -244,6 +244,12 @@ def __init__(self, device: str = "cuda", cache_dir: str = "./checkpoints/sam2_ca
         self.load_time = 0.0
 
     def _determine_optimal_size(self) -> str:
+        # Check environment variable first
+        env_size = os.environ.get("USE_SAM2", "").lower()
+        if env_size in ["tiny", "small", "base", "large"]:
+            logger.info(f"Using SAM2 size from environment: {env_size}")
+            return env_size
+            
         try:
             if torch.cuda.is_available():
                 props = torch.cuda.get_device_properties(0)
@@ -260,11 +266,12 @@ def load(self, model_size: str = "auto") -> Optional[_SAM2Adapter]:
         if model_size == "auto":
             model_size = self._determine_optimal_size()
 
+        # Use original SAM2 model names (without .1) for compatibility
         model_map = {
-            "tiny":  "facebook/sam2.1-hiera-tiny",
-            "small": "facebook/sam2.1-hiera-small",
-            "base":  "facebook/sam2.1-hiera-base-plus",
-            "large": "facebook/sam2.1-hiera-large",
+            "tiny":  "facebook/sam2-hiera-tiny",
+            "small": "facebook/sam2-hiera-small",
+            "base":  "facebook/sam2-hiera-base-plus",
+            "large": "facebook/sam2-hiera-large",
         }
         self.model_id = model_map.get(model_size, model_map["tiny"])
         logger.info(f"Loading SAM2 model: {self.model_id} (device={self.device})")
@@ -288,17 +295,61 @@ def load(self, model_size: str = "auto") -> Optional[_SAM2Adapter]:
         return None
 
     def _load_official(self):
-        from sam2.sam2_image_predictor import SAM2ImagePredictor
-        predictor = SAM2ImagePredictor.from_pretrained(
-            self.model_id,
-            cache_dir=self.cache_dir,
-            local_files_only=False,
-            trust_remote_code=True,
-        )
-        if hasattr(predictor, "model"):
-            predictor.model = predictor.model.to(self.device)
-            predictor.model.eval()
-        return predictor
+        try:
+            from huggingface_hub import hf_hub_download
+            from sam2.build_sam import build_sam2
+            from sam2.sam2_image_predictor import SAM2ImagePredictor
+        except ImportError as e:
+            logger.error(f"Failed to import SAM2 components: {e}")
+            return None
+        
+        # Map model IDs to config files and checkpoint names
+        config_map = {
+            "facebook/sam2-hiera-tiny": ("sam2_hiera_t.yaml", "sam2_hiera_tiny.pt"),
+            "facebook/sam2-hiera-small": ("sam2_hiera_s.yaml", "sam2_hiera_small.pt"),
+            "facebook/sam2-hiera-base-plus": ("sam2_hiera_b+.yaml", "sam2_hiera_base_plus.pt"),
+            "facebook/sam2-hiera-large": ("sam2_hiera_l.yaml", "sam2_hiera_large.pt"),
+        }
+        
+        config_file, checkpoint_file = config_map.get(self.model_id, (None, None))
+        if not config_file:
+            raise ValueError(f"Unknown model: {self.model_id}")
+        
+        try:
+            # Download the checkpoint from HuggingFace
+            logger.info(f"Downloading checkpoint: {checkpoint_file}")
+            checkpoint_path = hf_hub_download(
+                repo_id=self.model_id,
+                filename=checkpoint_file,
+                cache_dir=self.cache_dir,
+                local_files_only=False
+            )
+            logger.info(f"Checkpoint downloaded to: {checkpoint_path}")
+            
+            # Also download the config file if needed
+            config_path = hf_hub_download(
+                repo_id=self.model_id,
+                filename=config_file,
+                cache_dir=self.cache_dir,
+                local_files_only=False
+            )
+            logger.info(f"Config downloaded to: {config_path}")
+            
+            # Build the model using the traditional method
+            sam2_model = build_sam2(config_path, checkpoint_path, device=self.device)
+            predictor = SAM2ImagePredictor(sam2_model)
+            
+            # Ensure model is on the correct device and in eval mode
+            if hasattr(predictor, "model"):
+                predictor.model = predictor.model.to(self.device)
+                predictor.model.eval()
+            
+            return predictor
+            
+        except Exception as e:
+            logger.error(f"Error loading SAM2 model: {e}")
+            logger.debug(traceback.format_exc())
+            return None
 
     def _load_fallback(self):
         class FallbackSAM2:
@@ -360,4 +411,4 @@ def get_info(self) -> Dict[str, Any]:
     m = out["masks"]
     print("Masks:", m.shape, m.dtype, m.min(), m.max())
     cv2.imwrite("sam2_mask0.png", (np.clip(m[0], 0, 1) * 255).astype(np.uint8))
-    print("Wrote sam2_mask0.png")
+    print("Wrote sam2_mask0.png")