Update models/model_loaders.py

models/model_loaders.py

@@ -11,21 +11,18 @@
 import torch
 import psutil
 from contextlib import contextmanager
-
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
-
 @contextmanager
 def torch_memory_manager():
     try:
-        logger.info("[torch_memory_manager] Enter")  # [LOG+SAFETY PATCH]
+        logger.info("[torch_memory_manager] Enter")
         yield
     finally:
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
         gc.collect()
-        logger.info("[torch_memory_manager] Exit, cleaned up")  # [LOG+SAFETY PATCH]
-
+        logger.info("[torch_memory_manager] Exit, cleaned up")
 def get_memory_usage():
     memory_info = {}
     if torch.cuda.is_available():
@@ -35,22 +32,20 @@ def get_memory_usage():
                                   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}")  # [LOG+SAFETY PATCH]
+    logger.info(f"[get_memory_usage] {memory_info}")
     return memory_info
-
 def clear_model_cache():
-    logger.info("[clear_model_cache] Clearing all model caches...")  # [LOG+SAFETY PATCH]
+    logger.info("[clear_model_cache] Clearing all model caches...")
     if hasattr(st, 'cache_resource'):
         st.cache_resource.clear()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
     gc.collect()
-    logger.info("[clear_model_cache] Model cache cleared")  # [LOG+SAFETY PATCH]
-
+    logger.info("[clear_model_cache] Model cache cleared")
 @st.cache_resource(show_spinner=False)
 def load_sam2_predictor():
     try:
-        logger.info("[load_sam2_predictor] Loading SAM2 image predictor...")  # [LOG+SAFETY PATCH]
+        logger.info("[load_sam2_predictor] Loading SAM2 image predictor...")
         from sam2.build_sam import build_sam2
         from sam2.sam2_image_predictor import SAM2ImagePredictor
         device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -94,19 +89,21 @@ def load_sam2_predictor():
         import traceback
         traceback.print_exc()
         return None, None
-
 def load_sam2():
     predictor, device = load_sam2_predictor()
     return predictor
-
 @st.cache_resource(show_spinner=False)
 def load_matanyone_processor():
     try:
-        logger.info("[load_matanyone_processor] Loading MatAnyone processor...")  # [LOG+SAFETY PATCH]
+        logger.info("[load_matanyone_processor] Loading MatAnyone processor...")
         from matanyone import InferenceCore
         device = "cuda" if torch.cuda.is_available() else "cpu"
         logger.info(f"[load_matanyone_processor] MatAnyone using device: {device}")
-        processor = InferenceCore("PeiqingYang/MatAnyone", device=device)
+        try:
+            processor = InferenceCore("PeiqingYang/MatAnyone", device=device)
+        except Exception as e:
+            logger.warning(f"[load_matanyone_processor] Path warning caught: {e}")
+            processor = InferenceCore("PeiqingYang/MatAnyone", device=device)  # Retry
         if hasattr(processor, 'model'):
             processor.model.to(device)
             processor.model.eval()
@@ -121,11 +118,9 @@ def load_matanyone_processor():
         import traceback
         traceback.print_exc()
         return None, None
-
 def load_matanyone():
     processor, device = load_matanyone_processor()
     return processor
-
 def test_models():
     results = {
         'sam2': {'loaded': False, 'error': None, 'device': None},
@@ -151,9 +146,8 @@ def test_models():
     except Exception as e:
         results['matanyone']['error'] = str(e)
         logger.error(f"[test_models] MatAnyone error: {e}", exc_info=True)
-    logger.info(f"[test_models] Results: {results}")  # [LOG+SAFETY PATCH]
+    logger.info(f"[test_models] Results: {results}")
     return results
-
 def log_memory_usage(stage=""):
     memory_info = get_memory_usage()
     log_msg = f"Memory usage"
@@ -166,17 +160,15 @@ 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
     memory_info = get_memory_usage()
     free_gb = memory_info.get('gpu_free', 0)
-    logger.info(f"[check_memory_available] free_gb={free_gb}, required={required_gb}")  # [LOG+SAFETY PATCH]
+    logger.info(f"[check_memory_available] free_gb={free_gb}, required={required_gb}")
     return free_gb >= required_gb, free_gb
-
 def free_memory_aggressive():
-    logger.info("[free_memory_aggressive] Performing aggressive memory cleanup...")  # [LOG+SAFETY PATCH]
+    logger.info("[free_memory_aggressive] Performing aggressive memory cleanup...")
     print("Performing aggressive memory cleanup...", flush=True)
     clear_model_cache()
     if torch.cuda.is_available():
@@ -189,4 +181,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")