Update model_loader.py

model_loader.py

@@ -9,7 +9,9 @@
 
 import os
 import gc
+import sys
 import time
+import shutil
 import logging
 import tempfile
 import traceback
@@ -27,6 +29,171 @@
 
 logger = logging.getLogger(__name__)
 
+# ============================================================================ #
+# HARD CACHE CLEANER
+# ============================================================================ #
+
+class HardCacheCleaner:
+    """
+    Comprehensive cache cleaning system to resolve SAM2 loading issues
+    Clears Python module cache, HuggingFace cache, and temp files
+    """
+    
+    @staticmethod
+    def clean_all_caches(verbose: bool = True):
+        """Clean all caches that might interfere with SAM2 loading"""
+        
+        if verbose:
+            logger.info("Starting comprehensive cache cleanup...")
+        
+        # 1. Clean Python module cache
+        HardCacheCleaner._clean_python_cache(verbose)
+        
+        # 2. Clean HuggingFace cache
+        HardCacheCleaner._clean_huggingface_cache(verbose)
+        
+        # 3. Clean PyTorch cache
+        HardCacheCleaner._clean_pytorch_cache(verbose)
+        
+        # 4. Clean temp directories
+        HardCacheCleaner._clean_temp_directories(verbose)
+        
+        # 5. Clear import cache
+        HardCacheCleaner._clear_import_cache(verbose)
+        
+        # 6. Force garbage collection
+        HardCacheCleaner._force_gc_cleanup(verbose)
+        
+        if verbose:
+            logger.info("Cache cleanup completed")
+    
+    @staticmethod
+    def _clean_python_cache(verbose: bool = True):
+        """Clean Python bytecode cache"""
+        try:
+            # Clear sys.modules cache for SAM2 related modules
+            sam2_modules = [key for key in sys.modules.keys() if 'sam2' in key.lower()]
+            for module in sam2_modules:
+                if verbose:
+                    logger.info(f"Removing cached module: {module}")
+                del sys.modules[module]
+            
+            # Clear __pycache__ directories
+            for root, dirs, files in os.walk("."):
+                for dir_name in dirs[:]:  # Use slice to modify list during iteration
+                    if dir_name == "__pycache__":
+                        cache_path = os.path.join(root, dir_name)
+                        if verbose:
+                            logger.info(f"Removing __pycache__: {cache_path}")
+                        shutil.rmtree(cache_path, ignore_errors=True)
+                        dirs.remove(dir_name)
+            
+        except Exception as e:
+            logger.warning(f"Python cache cleanup failed: {e}")
+    
+    @staticmethod
+    def _clean_huggingface_cache(verbose: bool = True):
+        """Clean HuggingFace model cache"""
+        try:
+            cache_paths = [
+                os.path.expanduser("~/.cache/huggingface/"),
+                os.path.expanduser("~/.cache/torch/"),
+                "./checkpoints/",
+                "./.cache/",
+            ]
+            
+            for cache_path in cache_paths:
+                if os.path.exists(cache_path):
+                    if verbose:
+                        logger.info(f"Cleaning cache directory: {cache_path}")
+                    
+                    # Remove SAM2 specific files
+                    for root, dirs, files in os.walk(cache_path):
+                        for file in files:
+                            if any(pattern in file.lower() for pattern in ['sam2', 'segment-anything-2']):
+                                file_path = os.path.join(root, file)
+                                try:
+                                    os.remove(file_path)
+                                    if verbose:
+                                        logger.info(f"Removed cached file: {file_path}")
+                                except:
+                                    pass
+                        
+                        for dir_name in dirs[:]:
+                            if any(pattern in dir_name.lower() for pattern in ['sam2', 'segment-anything-2']):
+                                dir_path = os.path.join(root, dir_name)
+                                try:
+                                    shutil.rmtree(dir_path, ignore_errors=True)
+                                    if verbose:
+                                        logger.info(f"Removed cached directory: {dir_path}")
+                                    dirs.remove(dir_name)
+                                except:
+                                    pass
+                                    
+        except Exception as e:
+            logger.warning(f"HuggingFace cache cleanup failed: {e}")
+    
+    @staticmethod
+    def _clean_pytorch_cache(verbose: bool = True):
+        """Clean PyTorch cache"""
+        try:
+            import torch
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+                if verbose:
+                    logger.info("Cleared PyTorch CUDA cache")
+        except Exception as e:
+            logger.warning(f"PyTorch cache cleanup failed: {e}")
+    
+    @staticmethod
+    def _clean_temp_directories(verbose: bool = True):
+        """Clean temporary directories"""
+        try:
+            temp_dirs = [tempfile.gettempdir(), "/tmp", "./tmp", "./temp"]
+            
+            for temp_dir in temp_dirs:
+                if os.path.exists(temp_dir):
+                    for item in os.listdir(temp_dir):
+                        if 'sam2' in item.lower() or 'segment' in item.lower():
+                            item_path = os.path.join(temp_dir, item)
+                            try:
+                                if os.path.isfile(item_path):
+                                    os.remove(item_path)
+                                elif os.path.isdir(item_path):
+                                    shutil.rmtree(item_path, ignore_errors=True)
+                                if verbose:
+                                    logger.info(f"Removed temp item: {item_path}")
+                            except:
+                                pass
+                                
+        except Exception as e:
+            logger.warning(f"Temp directory cleanup failed: {e}")
+    
+    @staticmethod
+    def _clear_import_cache(verbose: bool = True):
+        """Clear Python import cache"""
+        try:
+            import importlib
+            
+            # Invalidate import caches
+            importlib.invalidate_caches()
+            
+            if verbose:
+                logger.info("Cleared Python import cache")
+                
+        except Exception as e:
+            logger.warning(f"Import cache cleanup failed: {e}")
+    
+    @staticmethod
+    def _force_gc_cleanup(verbose: bool = True):
+        """Force garbage collection"""
+        try:
+            collected = gc.collect()
+            if verbose:
+                logger.info(f"Garbage collection freed {collected} objects")
+        except Exception as e:
+            logger.warning(f"Garbage collection failed: {e}")
+
 # ============================================================================ #
 # MODEL LOADER CLASS - MAIN INTERFACE
 # ============================================================================ #
@@ -169,86 +336,123 @@ def load_all_models(self, progress_callback: Optional[callable] = None, cancel_e
             return None, None
 
 # ============================================================================ #
-# SAM2 MODEL LOADING - AUTOMATIC CONFIG DETECTION
+# SAM2 MODEL LOADING - HUGGINGFACE TRANSFORMERS APPROACH
 # ============================================================================ #
 
     def _load_sam2_predictor(self, progress_callback: Optional[callable] = None):
         """
-        Load SAM2 predictor with automatic config detection - no manual config files needed
-        Uses build_sam2_video_predictor for automatic configuration based on checkpoint filename
+        Load SAM2 using HuggingFace Transformers integration with cache cleanup
+        This method works reliably on HuggingFace Spaces without config file issues
         
         Args:
             progress_callback: Progress update callback
             
         Returns:
-            SAM2VideoPredictor or None
+            SAM2 model or None
         """
-        def try_load_sam2_auto(repo_id: str, filename: str, model_name: str):
-            """Attempt to load SAM2 with automatic config detection"""
-            try:
-                checkpoint_path = os.path.join(self.checkpoints_dir, filename)
-                logger.info(f"Attempting SAM2 checkpoint: {checkpoint_path}")
-                
-                # Download checkpoint if needed
-                if not os.path.exists(checkpoint_path):
-                    logger.info(f"Downloading {filename} from Hugging Face Hub...")
-                    if progress_callback:
-                        progress_callback(0.2, f"Downloading {filename}...")
-                    
-                    try:
-                        from huggingface_hub import hf_hub_download
-                        checkpoint_path = hf_hub_download(
-                            repo_id=repo_id,
-                            filename=filename,
-                            cache_dir=self.checkpoints_dir,
-                            local_dir_use_symlinks=False
-                        )
-                        logger.info(f"Download complete: {checkpoint_path}")
-                    except Exception as download_error:
-                        logger.warning(f"Failed to download {filename}: {download_error}")
-                        return None
-                
-                if progress_callback:
-                    progress_callback(0.4, f"Building SAM2 {model_name}...")
-                
-                # Use automatic config detection - NO manual config needed!
-                from sam2.build_sam import build_sam2_video_predictor
-                
-                predictor = build_sam2_video_predictor(checkpoint_path, device=self.device)
-                
-                logger.info(f"SAM2 {model_name} loaded successfully on {self.device}")
-                return predictor
-                
-            except Exception as e:
-                error_msg = f"Failed to load SAM2 {model_name}: {e}"
-                logger.warning(error_msg)
-                return None
-        
-        # Try different SAM2 models with automatic config detection
-        model_attempts = [
-            ("facebook/sam2-hiera-large", "sam2_hiera_large.pt", "hiera_large"),
-            ("facebook/sam2-hiera-base-plus", "sam2_hiera_base_plus.pt", "hiera_base_plus"),
-            ("facebook/sam2-hiera-small", "sam2_hiera_small.pt", "hiera_small"),
-            ("facebook/sam2-hiera-tiny", "sam2_hiera_tiny.pt", "hiera_tiny")
-        ]
+        logger.info("=== USING NEW HF TRANSFORMERS SAM2 LOADER ===")
+        
+        # Step 1: Clean caches before loading
+        if progress_callback:
+            progress_callback(0.15, "Cleaning caches...")
         
-        # Prioritize model size based on device memory
+        HardCacheCleaner.clean_all_caches(verbose=True)
+        
+        # Step 2: Determine model size based on device memory
+        model_size = "large"  # default
         if hasattr(self.device_manager, 'get_device_memory_gb'):
             try:
                 memory_gb = self.device_manager.get_device_memory_gb()
                 if memory_gb < 4:
-                    model_attempts = model_attempts[2:]  # Only tiny and small
+                    model_size = "tiny"
                 elif memory_gb < 8:
-                    model_attempts = model_attempts[1:]  # Skip large
+                    model_size = "base"
+                logger.info(f"Selected SAM2 {model_size} based on {memory_gb}GB memory")
             except Exception as e:
                 logger.warning(f"Could not determine device memory: {e}")
         
-        for repo_id, filename, model_name in model_attempts:
-            predictor = try_load_sam2_auto(repo_id, filename, model_name)
-            if predictor is not None:
-                return predictor
+        # Step 3: Try multiple HuggingFace approaches
+        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"
+        }
+        
+        model_id = model_map.get(model_size, model_map["large"])
+        
+        if progress_callback:
+            progress_callback(0.3, f"Loading SAM2 {model_size}...")
+        
+        # Method 1: HuggingFace Transformers Pipeline (most reliable)
+        try:
+            logger.info("Trying Transformers pipeline approach...")
+            from transformers import pipeline
+            
+            sam2_pipeline = pipeline(
+                "mask-generation",
+                model=model_id,
+                device=0 if str(self.device) == "cuda" else -1
+            )
+            
+            logger.info("SAM2 loaded successfully via Transformers pipeline")
+            return sam2_pipeline
+            
+        except Exception as e:
+            logger.warning(f"Pipeline approach failed: {e}")
+        
+        # Method 2: Direct Transformers classes
+        try:
+            logger.info("Trying direct Transformers classes...")
+            from transformers import Sam2Processor, Sam2Model
+            
+            processor = Sam2Processor.from_pretrained(model_id)
+            model = Sam2Model.from_pretrained(model_id).to(self.device)
+            
+            logger.info("SAM2 loaded successfully via Transformers classes")
+            return {"model": model, "processor": processor}
+            
+        except Exception as e:
+            logger.warning(f"Direct class approach failed: {e}")
+        
+        # Method 3: Official SAM2 with from_pretrained
+        try:
+            logger.info("Trying official SAM2 from_pretrained...")
+            from sam2.sam2_image_predictor import SAM2ImagePredictor
+            
+            predictor = SAM2ImagePredictor.from_pretrained(model_id)
+            
+            logger.info("SAM2 loaded successfully via official from_pretrained")
+            return predictor
+            
+        except Exception as e:
+            logger.warning(f"Official from_pretrained approach failed: {e}")
+        
+        # Method 4: Fallback to direct checkpoint download
+        try:
+            logger.info("Trying fallback checkpoint approach...")
+            from huggingface_hub import hf_hub_download
+            from transformers import Sam2Model
+            
+            # Download checkpoint directly
+            checkpoint_path = hf_hub_download(
+                repo_id=model_id,
+                filename="model.safetensors" if "sam2.1" in model_id else "pytorch_model.bin"
+            )
+            
+            logger.info(f"Downloaded checkpoint to: {checkpoint_path}")
+            
+            # Load with minimal approach
+            model = Sam2Model.from_pretrained(model_id)
+            model = model.to(self.device)
+            
+            logger.info("SAM2 loaded successfully via fallback approach")
+            return model
+            
+        except Exception as e:
+            logger.warning(f"Fallback approach failed: {e}")
         
-        logger.error("All SAM2 model loading attempts failed")
+        logger.error("All SAM2 loading methods failed")
         return None
 
 # ============================================================================ #
@@ -411,7 +615,7 @@ def get_model_info(self) -> Dict[str, Any]:
         
         if self.sam2_predictor is not None:
             try:
-                info['sam2_model_type'] = type(self.sam2_predictor.model).__name__
+                info['sam2_model_type'] = type(self.sam2_predictor).__name__
             except:
                 info['sam2_model_type'] = "Unknown"