Update models/model_loaders.py

models/model_loaders.py

@@ -1,9 +1,10 @@
 #!/usr/bin/env python3
 """
 Model Loading and Memory Management
-Handles lazy loading of SAM2 and MatAnyone models with caching
+Handles lazy loading of SAM2 and MatAnyone models with caching.
 (Enhanced logging, error handling, and memory safety)
 """
+
 import os
 import gc
 import logging
@@ -11,8 +12,10 @@
 import torch
 import psutil
 from contextlib import contextmanager
+
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
+
 @contextmanager
 def torch_memory_manager():
     try:
@@ -23,18 +26,21 @@ def torch_memory_manager():
             torch.cuda.empty_cache()
         gc.collect()
         logger.info("[torch_memory_manager] Exit, cleaned up")
+
 def get_memory_usage():
     memory_info = {}
     if torch.cuda.is_available():
         memory_info['gpu_allocated'] = torch.cuda.memory_allocated() / 1e9
         memory_info['gpu_reserved'] = torch.cuda.memory_reserved() / 1e9
         memory_info['gpu_free'] = (torch.cuda.get_device_properties(0).total_memory -
-                                  torch.cuda.memory_allocated()) / 1e9
+                                   torch.cuda.memory_allocated()) / 1e9
     memory_info['ram_used'] = psutil.virtual_memory().used / 1e9
     memory_info['ram_available'] = psutil.virtual_memory().available / 1e9
     logger.info(f"[get_memory_usage] {memory_info}")
     return memory_info
+
 def clear_model_cache():
+    """Manual/debug only: Clear Streamlit resource cache and free memory."""
     logger.info("[clear_model_cache] Clearing all model caches...")
     if hasattr(st, 'cache_resource'):
         st.cache_resource.clear()
@@ -42,8 +48,10 @@ def clear_model_cache():
         torch.cuda.empty_cache()
     gc.collect()
     logger.info("[clear_model_cache] Model cache cleared")
+
 @st.cache_resource(show_spinner=False)
 def load_sam2_predictor():
+    """Load SAM2 image predictor, choosing model size based on available GPU memory."""
     try:
         logger.info("[load_sam2_predictor] Loading SAM2 image predictor...")
         from sam2.build_sam import build_sam2
@@ -83,17 +91,21 @@ def load_sam2_predictor():
             predictor.model.eval()
             logger.info(f"[load_sam2_predictor] SAM2 model moved to {device} and set to eval mode")
         logger.info(f"✅ SAM2 loaded successfully on {device}!")
-        return predictor, device
+        return predictor
     except Exception as e:
         logger.error(f"❌ Failed to load SAM2 predictor: {e}", exc_info=True)
         import traceback
         traceback.print_exc()
-        return None, None
+        return None
+
 def load_sam2():
-    predictor, device = load_sam2_predictor()
+    """Convenience alias for legacy code: returns only the predictor object."""
+    predictor = load_sam2_predictor()
     return predictor
+
 @st.cache_resource(show_spinner=False)
 def load_matanyone_processor():
+    """Load MatAnyone processor (inference core) on the best available device."""
     try:
         logger.info("[load_matanyone_processor] Loading MatAnyone processor...")
         from matanyone import InferenceCore
@@ -112,35 +124,37 @@ def load_matanyone_processor():
             processor.device = device
             logger.info(f"[load_matanyone_processor] Set processor.device to {device}")
         logger.info(f"✅ MatAnyone loaded successfully on {device}!")
-        return processor, device
+        return processor
     except Exception as e:
         logger.error(f"❌ Failed to load MatAnyone: {e}", exc_info=True)
         import traceback
         traceback.print_exc()
-        return None, None
+        return None
+
 def load_matanyone():
-    processor, device = load_matanyone_processor()
+    """Convenience alias for legacy code: returns only the processor object."""
+    processor = load_matanyone_processor()
     return processor
+
 def test_models():
+    """For admin/diagnosis: attempts to load both models and returns status."""
     results = {
-        'sam2': {'loaded': False, 'error': None, 'device': None},
-        'matanyone': {'loaded': False, 'error': None, 'device': None}
+        'sam2': {'loaded': False, 'error': None},
+        'matanyone': {'loaded': False, 'error': None}
     }
     try:
-        sam2_predictor, sam2_device = load_sam2_predictor()
+        sam2_predictor = load_sam2_predictor()
         if sam2_predictor is not None:
             results['sam2']['loaded'] = True
-            results['sam2']['device'] = sam2_device
         else:
             results['sam2']['error'] = "Predictor returned None"
     except Exception as e:
         results['sam2']['error'] = str(e)
         logger.error(f"[test_models] SAM2 error: {e}", exc_info=True)
     try:
-        matanyone_processor, matanyone_device = load_matanyone_processor()
+        matanyone_processor = load_matanyone_processor()
         if matanyone_processor is not None:
             results['matanyone']['loaded'] = True
-            results['matanyone']['device'] = matanyone_device
         else:
             results['matanyone']['error'] = "Processor returned None"
     except Exception as e:
@@ -148,6 +162,7 @@ def test_models():
         logger.error(f"[test_models] MatAnyone error: {e}", exc_info=True)
     logger.info(f"[test_models] Results: {results}")
     return results
+
 def log_memory_usage(stage=""):
     memory_info = get_memory_usage()
     log_msg = f"Memory usage"
@@ -160,6 +175,7 @@ def log_memory_usage(stage=""):
     print(log_msg, flush=True)
     logger.info(log_msg)
     return memory_info
+
 def check_memory_available(required_gb=2.0):
     if not torch.cuda.is_available():
         return False, 0.0
@@ -167,7 +183,9 @@ def check_memory_available(required_gb=2.0):
     free_gb = memory_info.get('gpu_free', 0)
     logger.info(f"[check_memory_available] free_gb={free_gb}, required={required_gb}")
     return free_gb >= required_gb, free_gb
+
 def free_memory_aggressive():
+    """For emergency/manual use only! Do NOT call after every video or from UI!"""
     logger.info("[free_memory_aggressive] Performing aggressive memory cleanup...")
     print("Performing aggressive memory cleanup...", flush=True)
     clear_model_cache()
@@ -181,4 +199,4 @@ def free_memory_aggressive():
     gc.collect()
     print("Memory cleanup complete", flush=True)
     logger.info("Memory cleanup complete")
-    log_memory_usage("after cleanup")
+    log_memory_usage("after cleanup")