Rename cache_cleaner.py to utils/cache/cache_cleaner.py

cache_cleaner.py → utils/cache/cache_cleaner.py

@@ -1,6 +1,7 @@
-# ============================================================================ #
-# HARD CACHE CLEANER + WORKING SAM2 LOADER FOR HUGGINGFACE SPACES
-# ============================================================================ #
+"""
+Cache Management and SAM2 Loading Utilities
+Comprehensive cache cleaning system to resolve model loading issues on HF Spaces
+"""
 
 import os
 import gc
@@ -25,7 +26,7 @@ def clean_all_caches(verbose: bool = True):
         """Clean all caches that might interfere with SAM2 loading"""
         
         if verbose:
-            logger.info("🧹 Starting comprehensive cache cleanup...")
+            logger.info("Starting comprehensive cache cleanup...")
         
         # 1. Clean Python module cache
         HardCacheCleaner._clean_python_cache(verbose)
@@ -46,7 +47,7 @@ def clean_all_caches(verbose: bool = True):
         HardCacheCleaner._force_gc_cleanup(verbose)
         
         if verbose:
-            logger.info("✅ Cache cleanup completed")
+            logger.info("Cache cleanup completed")
     
     @staticmethod
     def _clean_python_cache(verbose: bool = True):
@@ -56,7 +57,7 @@ def _clean_python_cache(verbose: bool = True):
             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}")
+                    logger.info(f"Removing cached module: {module}")
                 del sys.modules[module]
             
             # Clear __pycache__ directories
@@ -65,7 +66,7 @@ def _clean_python_cache(verbose: bool = True):
                     if dir_name == "__pycache__":
                         cache_path = os.path.join(root, dir_name)
                         if verbose:
-                            logger.info(f"🗑️  Removing __pycache__: {cache_path}")
+                            logger.info(f"Removing __pycache__: {cache_path}")
                         shutil.rmtree(cache_path, ignore_errors=True)
                         dirs.remove(dir_name)
             
@@ -76,9 +77,14 @@ def _clean_python_cache(verbose: bool = True):
     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/",
             ]
@@ -86,7 +92,7 @@ def _clean_huggingface_cache(verbose: bool = True):
             for cache_path in cache_paths:
                 if os.path.exists(cache_path):
                     if verbose:
-                        logger.info(f"🗑️  Cleaning cache directory: {cache_path}")
+                        logger.info(f"Cleaning cache directory: {cache_path}")
                     
                     # Remove SAM2 specific files
                     for root, dirs, files in os.walk(cache_path):
@@ -96,7 +102,7 @@ def _clean_huggingface_cache(verbose: bool = True):
                                 try:
                                     os.remove(file_path)
                                     if verbose:
-                                        logger.info(f"🗑️  Removed cached file: {file_path}")
+                                        logger.info(f"Removed cached file: {file_path}")
                                 except:
                                     pass
                         
@@ -106,7 +112,7 @@ def _clean_huggingface_cache(verbose: bool = True):
                                 try:
                                     shutil.rmtree(dir_path, ignore_errors=True)
                                     if verbose:
-                                        logger.info(f"🗑️  Removed cached directory: {dir_path}")
+                                        logger.info(f"Removed cached directory: {dir_path}")
                                     dirs.remove(dir_name)
                                 except:
                                     pass
@@ -122,7 +128,7 @@ def _clean_pytorch_cache(verbose: bool = True):
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
                 if verbose:
-                    logger.info("🗑️  Cleared PyTorch CUDA cache")
+                    logger.info("Cleared PyTorch CUDA cache")
         except Exception as e:
             logger.warning(f"PyTorch cache cleanup failed: {e}")
     
@@ -130,7 +136,16 @@ def _clean_pytorch_cache(verbose: bool = True):
     def _clean_temp_directories(verbose: bool = True):
         """Clean temporary directories"""
         try:
-            temp_dirs = [tempfile.gettempdir(), "/tmp", "./tmp", "./temp"]
+            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):
@@ -143,7 +158,7 @@ def _clean_temp_directories(verbose: bool = True):
                                 elif os.path.isdir(item_path):
                                     shutil.rmtree(item_path, ignore_errors=True)
                                 if verbose:
-                                    logger.info(f"🗑️  Removed temp item: {item_path}")
+                                    logger.info(f"Removed temp item: {item_path}")
                             except:
                                 pass
                                 
@@ -160,7 +175,7 @@ def _clear_import_cache(verbose: bool = True):
             importlib.invalidate_caches()
             
             if verbose:
-                logger.info("🗑️  Cleared Python import cache")
+                logger.info("Cleared Python import cache")
                 
         except Exception as e:
             logger.warning(f"Import cache cleanup failed: {e}")
@@ -171,7 +186,7 @@ def _force_gc_cleanup(verbose: bool = True):
         try:
             collected = gc.collect()
             if verbose:
-                logger.info(f"🗑️  Garbage collection freed {collected} objects")
+                logger.info(f"Garbage collection freed {collected} objects")
         except Exception as e:
             logger.warning(f"Garbage collection failed: {e}")
 
@@ -189,7 +204,7 @@ def load_sam2_transformers_approach(device: str = "cuda", model_size: str = "lar
         This method works reliably on HuggingFace Spaces
         """
         try:
-            logger.info("🤖 Loading SAM2 via HuggingFace Transformers...")
+            logger.info("Loading SAM2 via HuggingFace Transformers...")
             
             # Model size mapping
             model_map = {
@@ -212,7 +227,7 @@ def load_sam2_transformers_approach(device: str = "cuda", model_size: str = "lar
                     device=0 if device == "cuda" else -1
                 )
                 
-                logger.info("✅ SAM2 loaded successfully via Transformers pipeline")
+                logger.info("SAM2 loaded successfully via Transformers pipeline")
                 return sam2_pipeline
                 
             except Exception as e:
@@ -225,7 +240,7 @@ def load_sam2_transformers_approach(device: str = "cuda", model_size: str = "lar
                 processor = Sam2Processor.from_pretrained(model_id)
                 model = Sam2Model.from_pretrained(model_id).to(device)
                 
-                logger.info("✅ SAM2 loaded successfully via Transformers classes")
+                logger.info("SAM2 loaded successfully via Transformers classes")
                 return {"model": model, "processor": processor}
                 
             except Exception as e:
@@ -237,7 +252,7 @@ def load_sam2_transformers_approach(device: str = "cuda", model_size: str = "lar
                 
                 predictor = SAM2ImagePredictor.from_pretrained(model_id)
                 
-                logger.info("✅ SAM2 loaded successfully via official from_pretrained")
+                logger.info("SAM2 loaded successfully via official from_pretrained")
                 return predictor
                 
             except Exception as e:
@@ -255,7 +270,7 @@ def load_sam2_fallback_approach(device: str = "cuda") -> Optional[Any]:
         Fallback approach using direct model loading
         """
         try:
-            logger.info("🔄 Trying fallback SAM2 loading approach...")
+            logger.info("Trying fallback SAM2 loading approach...")
             
             # Try the simplest possible approach
             from huggingface_hub import hf_hub_download
@@ -286,10 +301,6 @@ def load_sam2_fallback_approach(device: str = "cuda") -> Optional[Any]:
             return None
 
 
-# ============================================================================ #
-# INTEGRATED MODEL LOADER WITH CACHE CLEANING
-# ============================================================================ #
-
 def load_sam2_with_cache_cleanup(
     device: str = "cuda", 
     model_size: str = "large",
@@ -308,57 +319,32 @@ def load_sam2_with_cache_cleanup(
     try:
         # Step 1: Clean caches if requested
         if force_cache_clean:
-            status_messages.append("🧹 Cleaning caches...")
+            status_messages.append("Cleaning caches...")
             HardCacheCleaner.clean_all_caches(verbose=verbose)
-            status_messages.append("✅ Cache cleanup completed")
+            status_messages.append("Cache cleanup completed")
         
         # Step 2: Try primary loading method
-        status_messages.append("🤖 Loading SAM2 (primary 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!")
+            status_messages.append("SAM2 loaded successfully!")
             return model, "\n".join(status_messages)
         
         # Step 3: Try fallback method
-        status_messages.append("🔄 Trying fallback loading 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)!")
+            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")
+        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}"
+        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")
+        return None, "\n".join(status_messages)