""" Cache Management and SAM2 Loading Utilities Comprehensive cache cleaning system to resolve model loading issues on HF 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: # Get config for cache directories from config.app_config import get_config config = get_config() cache_paths = [ os.path.expanduser("~/.cache/huggingface/"), os.path.expanduser("~/.cache/torch/"), config.model_cache_dir, "./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: from config.app_config import get_config config = get_config() temp_dirs = [ config.temp_dir, 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 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)