final fix

models/sam2_loader.py

@@ -3,16 +3,14 @@
 SAM2 Loader — Robust loading and mask generation for SAM2
 ========================================================
 - Loads SAM2 model with Hydra config resolution
+- Auto-downloads missing checkpoint files
 - Generates seed masks for MatAnyone
 - Aligned with torch==2.3.1+cu121 and SAM2 commit 3c76f73c1a7e7b4a2e8a0a9a3e5b92f7e6e3f2f5
 
-Changes (2025-09-16):
-- Aligned with torch==2.3.1+cu121 and SAM2 commit
-- Added GPU memory logging for Tesla T4
-- Added SAM2 version logging via importlib.metadata
-- Simplified config resolution to match __init__.py
-- Fixed missing sys and inspect imports
-- Added SAM2Model wrapper class for app_hf.py compatibility
+Changes (2025-09-17):
+- Added automatic checkpoint download functionality
+- Enhanced error handling and logging
+- Fixed missing checkpoint issue that was causing fallback mask generation
 """
 
 from __future__ import annotations
@@ -22,6 +20,8 @@
 import inspect
 import logging
 import importlib.metadata
+import urllib.request
+import urllib.error
 from pathlib import Path
 from typing import Optional, Tuple, Dict, Any
 
@@ -55,6 +55,60 @@ def _add_sys_path(p: Path) -> None:
 
 _add_sys_path(TP_SAM2)
 
+# --------------------------------------------------------------------------------------
+# Checkpoint Download Functionality
+# --------------------------------------------------------------------------------------
+SAM2_CHECKPOINT_URLS = {
+    "sam2_hiera_large.pt": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt",
+    "sam2_hiera_base_plus.pt": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_base_plus.pt", 
+    "sam2_hiera_small.pt": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_small.pt",
+    "sam2_hiera_tiny.pt": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_tiny.pt"
+}
+
+def _download_checkpoint(checkpoint_path: str, checkpoint_name: str) -> bool:
+    """Download SAM2 checkpoint if it doesn't exist."""
+    if os.path.exists(checkpoint_path):
+        logger.info(f"Checkpoint already exists: {checkpoint_path}")
+        return True
+    
+    if checkpoint_name not in SAM2_CHECKPOINT_URLS:
+        logger.error(f"Unknown checkpoint: {checkpoint_name}. Available: {list(SAM2_CHECKPOINT_URLS.keys())}")
+        return False
+    
+    url = SAM2_CHECKPOINT_URLS[checkpoint_name]
+    logger.info(f"Downloading SAM2 checkpoint: {checkpoint_name}")
+    logger.info(f"URL: {url}")
+    logger.info(f"Destination: {checkpoint_path}")
+    
+    try:
+        # Create directory if it doesn't exist
+        os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True)
+        
+        # Download with progress
+        def _progress_hook(block_num, block_size, total_size):
+            if total_size > 0:
+                percent = min(100, block_num * block_size * 100 // total_size)
+                if percent % 10 == 0:  # Log every 10%
+                    logger.info(f"Download progress: {percent}%")
+        
+        urllib.request.urlretrieve(url, checkpoint_path, reporthook=_progress_hook)
+        
+        # Verify the file was downloaded
+        if os.path.exists(checkpoint_path) and os.path.getsize(checkpoint_path) > 0:
+            size_mb = os.path.getsize(checkpoint_path) / (1024 * 1024)
+            logger.info(f"Successfully downloaded {checkpoint_name} ({size_mb:.1f} MB)")
+            return True
+        else:
+            logger.error(f"Download failed: {checkpoint_path} does not exist or is empty")
+            return False
+            
+    except urllib.error.URLError as e:
+        logger.error(f"URL error downloading checkpoint: {e}")
+        return False
+    except Exception as e:
+        logger.error(f"Error downloading checkpoint: {e}")
+        return False
+
 # --------------------------------------------------------------------------------------
 # Safe Torch accessors
 # --------------------------------------------------------------------------------------
@@ -80,7 +134,7 @@ def _pick_device(env_key: str) -> str:
     requested = os.environ.get(env_key, "").strip().lower()
     has_cuda = _has_cuda()
     
-    logger.info(f"CUDA environment variables: {{'SAM2_DEVICE': '{os.environ.get('SAM2_DEVICE', '')}'}}")
+    logger.info(f"CUDA environment variables: {dict((k, v) for k, v in os.environ.items() if 'CUDA' in k or 'SAM2' in k)}")
     logger.info(f"_pick_device({env_key}): requested='{requested}', has_cuda={has_cuda}")
     
     if has_cuda and requested not in {"cpu"}:
@@ -100,12 +154,13 @@ def _pick_device(env_key: str) -> str:
 def _resolve_sam2_cfg(cfg_str: str) -> str:
     """Resolve SAM2 config path - return relative path for Hydra compatibility."""
     logger.info(f"_resolve_sam2_cfg called with cfg_str={cfg_str}")
+    logger.info(f"TP_SAM2 = {TP_SAM2}")
     
-    candidate = os.path.join(TP_SAM2, cfg_str)
+    candidate = TP_SAM2 / cfg_str
     logger.info(f"Candidate path: {candidate}")
-    logger.info(f"Candidate exists: {os.path.exists(candidate)}")
+    logger.info(f"Candidate exists: {candidate.exists()}")
     
-    if os.path.exists(candidate):
+    if candidate.exists():
         if cfg_str.startswith("sam2/configs/"):
             relative_path = cfg_str.replace("sam2/configs/", "configs/")
         else:
@@ -114,15 +169,15 @@ def _resolve_sam2_cfg(cfg_str: str) -> str:
         return relative_path
     
     fallbacks = [
-        os.path.join(TP_SAM2, "sam2", cfg_str),
-        os.path.join(TP_SAM2, "configs", cfg_str),
+        TP_SAM2 / "sam2" / cfg_str,
+        TP_SAM2 / "configs" / cfg_str,
     ]
     
     for fallback in fallbacks:
         logger.info(f"Trying fallback: {fallback}")
-        if os.path.exists(fallback):
-            if "configs/" in fallback:
-                relative_path = "configs/" + fallback.split("configs/")[-1]
+        if fallback.exists():
+            if "configs" in str(fallback):
+                relative_path = "configs/" + str(fallback).split("configs/")[-1]
                 logger.info(f"Returning fallback relative path: {relative_path}")
                 return relative_path
     
@@ -156,7 +211,7 @@ def _find_hiera_config_if_hieradet(cfg_path: str) -> Optional[str]:
     return None
 
 def load_sam2() -> Tuple[Optional[object], bool, Dict[str, Any]]:
-    """Robust SAM2 loader with config resolution and error handling."""
+    """Robust SAM2 loader with config resolution and checkpoint auto-download."""
     meta = {"sam2_import_ok": False, "sam2_init_ok": False}
     try:
         from sam2.build_sam import build_sam2
@@ -174,14 +229,26 @@ def load_sam2() -> Tuple[Optional[object], bool, Dict[str, Any]]:
         logger.info("[SAM2] SAM2 version unknown")
 
     # Check GPU memory before loading
-    if torch and torch.cuda.is_available():
+    if torch.cuda.is_available():
         mem_before = torch.cuda.memory_allocated() / 1024**3
-        logger.info(f"GPU memory before SAM2 load: {mem_before:.2f}GB")
+        logger.info(f"🔍 GPU memory before SAM2 load: {mem_before:.2f}GB")
 
     device = _pick_device("SAM2_DEVICE")
     cfg_env = os.environ.get("SAM2_MODEL_CFG", "sam2/configs/sam2/sam2_hiera_l.yaml")
     cfg = _resolve_sam2_cfg(cfg_env)
-    ckpt = os.environ.get("SAM2_CHECKPOINT", "")
+    
+    # Handle checkpoint with auto-download
+    ckpt = os.environ.get("SAM2_CHECKPOINT", "/home/user/app/checkpoints/sam2_hiera_large.pt")
+    checkpoint_name = os.path.basename(ckpt)
+    
+    # Auto-download checkpoint if missing
+    if not os.path.exists(ckpt):
+        logger.info(f"SAM2 checkpoint not found: {ckpt}")
+        if not _download_checkpoint(ckpt, checkpoint_name):
+            logger.error(f"Failed to download SAM2 checkpoint: {checkpoint_name}")
+            return None, False, meta
+    else:
+        logger.info(f"Using existing SAM2 checkpoint: {ckpt}")
 
     def _try_build(cfg_path: str):
         logger.info(f"_try_build called with cfg_path: {cfg_path}")
@@ -238,12 +305,16 @@ def _try_build(cfg_path: str):
             predictor = SAM2ImagePredictor(sam)
             meta["sam2_init_ok"] = True
             meta["sam2_device"] = device
+            logger.info("✅ SAM2 loaded successfully with auto-downloaded checkpoint")
             return predictor, True, meta
         else:
+            logger.error("❌ SAM2 initialization returned None")
             return None, False, meta
             
     except Exception as e:
-        logger.error(f"SAM2 loading failed: {e}")
+        logger.error(f"❌ SAM2 loading failed: {e}")
+        import traceback
+        logger.error(f"SAM2 loading traceback: {traceback.format_exc()}")
         return None, False, meta
 
 def run_sam2_mask(predictor: object,