cache_cleaner.py

# ============================================================================ #
# HARD CACHE CLEANER + WORKING SAM2 LOADER FOR HUGGINGFACE SPACES
# ============================================================================ #

import os
import gc
import sys
import shutil
import tempfile
import logging
import traceback
from pathlib import Path
from typing import Optional, Dict, Any, Tuple

logger = logging.getLogger(__name__)

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}")


class WorkingSAM2Loader:
    """
    SAM2 loader using HuggingFace Transformers integration - proven to work on HF Spaces
    This avoids all the config file and CUDA compilation issues
    """
    
    @staticmethod
    def load_sam2_transformers_approach(device: str = "cuda", model_size: str = "large") -> Optional[Any]:
        """
        Load SAM2 using HuggingFace Transformers integration
        This method works reliably on HuggingFace Spaces
        """
        try:
            logger.info("🤖 Loading SAM2 via HuggingFace Transformers...")
            
            # Model size mapping
            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"])
            logger.info(f"Using model: {model_id}")
            
            # Method 1: Using Transformers pipeline (most reliable for HF Spaces)
            try:
                from transformers import pipeline
                
                sam2_pipeline = pipeline(
                    "mask-generation",
                    model=model_id,
                    device=0 if 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: Using SAM2 classes directly via Transformers
            try:
                from transformers import Sam2Processor, Sam2Model
                
                processor = Sam2Processor.from_pretrained(model_id)
                model = Sam2Model.from_pretrained(model_id).to(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: Using official SAM2 with .from_pretrained()
            try:
                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}")
            
            return None
            
        except Exception as e:
            logger.error(f"All SAM2 loading methods failed: {e}")
            return None
    
    @staticmethod
    def load_sam2_fallback_approach(device: str = "cuda") -> Optional[Any]:
        """
        Fallback approach using direct model loading
        """
        try:
            logger.info("🔄 Trying fallback SAM2 loading approach...")
            
            # Try the simplest possible approach
            from huggingface_hub import hf_hub_download
            import torch
            
            # Download checkpoint directly
            checkpoint_path = hf_hub_download(
                repo_id="facebook/sam2.1-hiera-large",
                filename="sam2_hiera_large.pt"
            )
            
            logger.info(f"Downloaded checkpoint to: {checkpoint_path}")
            
            # Try to load with minimal dependencies
            try:
                # Method A: Try the working transformers integration
                from transformers import Sam2Model
                model = Sam2Model.from_pretrained("facebook/sam2.1-hiera-large")
                return model.to(device)
                
            except Exception as e:
                logger.warning(f"Transformers fallback failed: {e}")
            
            return None
            
        except Exception as e:
            logger.error(f"Fallback loading failed: {e}")
            return None


# ============================================================================ #
# INTEGRATED MODEL LOADER WITH CACHE CLEANING
# ============================================================================ #

def load_sam2_with_cache_cleanup(
    device: str = "cuda", 
    model_size: str = "large",
    force_cache_clean: bool = True,
    verbose: bool = True
) -> Tuple[Optional[Any], str]:
    """
    Load SAM2 with comprehensive cache cleanup
    
    Returns:
        Tuple of (model, status_message)
    """
    
    status_messages = []
    
    try:
        # Step 1: Clean caches if requested
        if force_cache_clean:
            status_messages.append("🧹 Cleaning caches...")
            HardCacheCleaner.clean_all_caches(verbose=verbose)
            status_messages.append("✅ Cache cleanup completed")
        
        # Step 2: Try primary loading method
        status_messages.append("🤖 Loading SAM2 (primary method)...")
        model = WorkingSAM2Loader.load_sam2_transformers_approach(device, model_size)
        
        if model is not None:
            status_messages.append("✅ SAM2 loaded successfully!")
            return model, "\n".join(status_messages)
        
        # Step 3: Try fallback method
        status_messages.append("🔄 Trying fallback loading method...")
        model = WorkingSAM2Loader.load_sam2_fallback_approach(device)
        
        if model is not None:
            status_messages.append("✅ SAM2 loaded successfully (fallback)!")
            return model, "\n".join(status_messages)
        
        # Step 4: All methods failed
        status_messages.append("❌ All SAM2 loading methods failed")
        return None, "\n".join(status_messages)
        
    except Exception as e:
        error_msg = f"❌ Critical error in SAM2 loading: {e}"
        logger.error(f"{error_msg}\n{traceback.format_exc()}")
        status_messages.append(error_msg)
        return None, "\n".join(status_messages)


# ============================================================================ #
# USAGE EXAMPLE
# ============================================================================ #

if __name__ == "__main__":
    # Clean example usage
    print("Testing SAM2 loader with cache cleanup...")
    
    # Load SAM2 with full cache cleanup
    model, status = load_sam2_with_cache_cleanup(
        device="cuda",
        model_size="large", 
        force_cache_clean=True,
        verbose=True
    )
    
    print("Status:", status)
    
    if model is not None:
        print("SAM2 loaded successfully!")
        print("Model type:", type(model))
    else:
        print("SAM2 loading failed completely")