Update models/model_loaders.py

models/model_loaders.py

@@ -2,6 +2,7 @@
 """
 Model Loading and Memory Management
 Handles lazy loading of SAM2 and MatAnyone models with caching
+(Enhanced logging, error handling, and memory safety)
 """
 
 import os
@@ -11,172 +12,133 @@
 import torch
 import psutil
 import mediapipe as mp
+from contextlib import contextmanager
 
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Context manager for CUDA memory cleanup
-from contextlib import contextmanager
-
 @contextmanager
 def torch_memory_manager():
-    """Context manager for CUDA memory cleanup."""
     try:
+        logger.info("[torch_memory_manager] Enter")  # [LOG+SAFETY PATCH]
         yield
     finally:
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
         gc.collect()
+        logger.info("[torch_memory_manager] Exit, cleaned up")  # [LOG+SAFETY PATCH]
 
 def get_memory_usage():
-    """Get current memory usage statistics."""
     memory_info = {}
-    
-    # GPU memory if available
     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
-    
-    # RAM memory
     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]
     return memory_info
 
 def clear_model_cache():
-    """Clear all cached models and free memory."""
+    logger.info("[clear_model_cache] Clearing all model caches...")  # [LOG+SAFETY PATCH]
     if hasattr(st, 'cache_resource'):
         st.cache_resource.clear()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
     gc.collect()
-    logger.info("Model cache cleared")
-
-# ============================================================================
-# SAM2 Model Loading
-# ============================================================================
+    logger.info("[clear_model_cache] Model cache cleared")  # [LOG+SAFETY PATCH]
 
 @st.cache_resource(show_spinner=False)
 def load_sam2_predictor():
-    """
-    Lazy load SAM2 image predictor with fallback strategies.
-    Returns (predictor, device) tuple. Returns (None, None) if loading fails.
-    """
     try:
-        print("Loading SAM2 image predictor...", flush=True)
+        logger.info("[load_sam2_predictor] Loading SAM2 image predictor...")  # [LOG+SAFETY PATCH]
         from sam2.build_sam import build_sam2
         from sam2.sam2_image_predictor import SAM2ImagePredictor
-        
-        # Determine device
+
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        print(f"Using device for SAM2: {device}", flush=True)
-        
-        # Try local checkpoints first
+        logger.info(f"[load_sam2_predictor] Using device: {device}")
+
         checkpoint_path = "/home/user/app/checkpoints/sam2.1_hiera_large.pt"
         model_cfg = "/home/user/app/configs/sam2.1/sam2.1_hiera_l.yaml"
-        
+
         if not os.path.exists(checkpoint_path) or not os.path.exists(model_cfg):
-            print("Local checkpoints not found, using Hugging Face...", flush=True)
+            logger.warning("[load_sam2_predictor] Local checkpoints not found, using Hugging Face.")
             predictor = SAM2ImagePredictor.from_pretrained(
                 "facebook/sam2-hiera-large",
                 device=device
             )
         else:
-            # Check available GPU memory
             memory_info = get_memory_usage()
             gpu_free = memory_info.get('gpu_free', 0)
-            
             if device == "cuda" and gpu_free < 4.0:
-                print(f"Limited GPU memory ({gpu_free:.1f}GB), using smaller SAM2 model...", flush=True)
+                logger.warning(f"[load_sam2_predictor] Limited GPU memory ({gpu_free:.1f}GB), using smaller SAM2 model.")
                 try:
                     predictor = SAM2ImagePredictor.from_pretrained(
                         "facebook/sam2-hiera-tiny",
                         device=device
                     )
-                except Exception:
+                except Exception as e:
+                    logger.warning(f"[load_sam2_predictor] Tiny model failed, trying small. {e}")
                     predictor = SAM2ImagePredictor.from_pretrained(
                         "facebook/sam2-hiera-small",
                         device=device
                     )
             else:
-                # Use local large model
+                logger.info("[load_sam2_predictor] Using local large model")
                 sam2_model = build_sam2(model_cfg, checkpoint_path, device=device)
                 predictor = SAM2ImagePredictor(sam2_model)
-        
-        # CRITICAL: Verify and force model to correct device
+
         if hasattr(predictor, 'model'):
             predictor.model.to(device)
             predictor.model.eval()
-            print(f"SAM2 model moved to {device} and set to eval mode", flush=True)
-        
-        print(f"✅ SAM2 loaded successfully on {device}!", flush=True)
+            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
-        
+
     except Exception as e:
-        print(f"❌ Failed to load SAM2 predictor: {e}", flush=True)
+        logger.error(f"❌ Failed to load SAM2 predictor: {e}", exc_info=True)
         import traceback
         traceback.print_exc()
         return None, None
 
-# Alias for backward compatibility
 def load_sam2():
-    """Alias for load_sam2_predictor() - returns just predictor for compatibility"""
     predictor, device = load_sam2_predictor()
     return predictor
 
-# ============================================================================
-# MatAnyone Model Loading
-# ============================================================================
-
 @st.cache_resource(show_spinner=False)
 def load_matanyone_processor():
-    """
-    Lazy load MatAnyone processor with explicit GPU placement.
-    Returns (processor, device) tuple. Returns (None, None) if loading fails.
-    """
     try:
-        print("Loading MatAnyone processor...", flush=True)
+        logger.info("[load_matanyone_processor] Loading MatAnyone processor...")  # [LOG+SAFETY PATCH]
         from matanyone import InferenceCore
-        
-        # Determine device
+
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        print(f"MatAnyone using device: {device}", flush=True)
-        
-        # Load processor with explicit device
+        logger.info(f"[load_matanyone_processor] MatAnyone using device: {device}")
+
         processor = InferenceCore("PeiqingYang/MatAnyone", device=device)
-        
-        # CRITICAL: Verify the processor's model is actually on GPU
         if hasattr(processor, 'model'):
             processor.model.to(device)
             processor.model.eval()
-            print(f"MatAnyone model explicitly moved to {device}", flush=True)
-        
-        # Check if processor has device attribute and set it
+            logger.info(f"[load_matanyone_processor] MatAnyone model explicitly moved to {device}")
+
         if not hasattr(processor, 'device'):
             processor.device = device
-            print(f"Set processor.device to {device}", flush=True)
-        
-        print(f"✅ MatAnyone loaded successfully on {device}!", flush=True)
+            logger.info(f"[load_matanyone_processor] Set processor.device to {device}")
+
+        logger.info(f"✅ MatAnyone loaded successfully on {device}!")
         return processor, device
-        
+
     except Exception as e:
-        print(f"❌ Failed to load MatAnyone: {e}", flush=True)
+        logger.error(f"❌ Failed to load MatAnyone: {e}", exc_info=True)
         import traceback
         traceback.print_exc()
         return None, None
 
-# Alias for backward compatibility
 def load_matanyone():
-    """Alias for load_matanyone_processor() - returns just processor for compatibility"""
     processor, device = load_matanyone_processor()
     return processor
 
-# ============================================================================
-# MediaPipe Pose
-# ============================================================================
-
-# Initialize MediaPipe Pose as a module-level variable
 mp_pose = mp.solutions.pose
 pose = mp_pose.Pose(
     static_image_mode=False,
@@ -184,23 +146,13 @@ def load_matanyone():
     enable_segmentation=True,
     min_detection_confidence=0.5
 )
-print("✅ MediaPipe Pose initialized", flush=True)
-
-# ============================================================================
-# Model Health Check
-# ============================================================================
+logger.info("✅ MediaPipe Pose initialized",)  # [LOG+SAFETY PATCH]
 
 def test_models():
-    """
-    Test if both models can load successfully.
-    Returns dict with test results.
-    """
     results = {
         'sam2': {'loaded': False, 'error': None, 'device': None},
         'matanyone': {'loaded': False, 'error': None, 'device': None}
     }
-    
-    # Test SAM2
     try:
         sam2_predictor, sam2_device = load_sam2_predictor()
         if sam2_predictor is not None:
@@ -210,8 +162,7 @@ def test_models():
             results['sam2']['error'] = "Predictor returned None"
     except Exception as e:
         results['sam2']['error'] = str(e)
-    
-    # Test MatAnyone
+        logger.error(f"[test_models] SAM2 error: {e}", exc_info=True)
     try:
         matanyone_processor, matanyone_device = load_matanyone_processor()
         if matanyone_processor is not None:
@@ -221,53 +172,35 @@ def test_models():
             results['matanyone']['error'] = "Processor returned None"
     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]
     return results
 
-# ============================================================================
-# Memory Monitoring
-# ============================================================================
-
 def log_memory_usage(stage=""):
-    """Log current memory usage with optional stage label."""
     memory_info = get_memory_usage()
-    
     log_msg = f"Memory usage"
     if stage:
         log_msg += f" ({stage})"
     log_msg += ":"
-    
     if 'gpu_allocated' in memory_info:
         log_msg += f" GPU {memory_info['gpu_allocated']:.1f}GB allocated, {memory_info['gpu_free']:.1f}GB free"
-    
     log_msg += f" | RAM {memory_info['ram_used']:.1f}GB used"
-    
     print(log_msg, flush=True)
     logger.info(log_msg)
     return memory_info
 
 def check_memory_available(required_gb=2.0):
-    """
-    Check if enough GPU memory is available.
-    Returns (bool, float) - (is_available, free_gb)
-    """
     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]
     return free_gb >= required_gb, free_gb
 
 def free_memory_aggressive():
-    """Aggressively free GPU and system memory."""
+    logger.info("[free_memory_aggressive] Performing aggressive memory cleanup...")  # [LOG+SAFETY PATCH]
     print("Performing aggressive memory cleanup...", flush=True)
-    logger.info("Performing aggressive memory cleanup...")
-    
-    # Clear model cache
     clear_model_cache()
-    
-    # CUDA cleanup
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
         torch.cuda.synchronize()
@@ -275,10 +208,7 @@ def free_memory_aggressive():
             torch.cuda.ipc_collect()
         except Exception:
             pass
-    
-    # System cleanup
     gc.collect()
-    
     print("Memory cleanup complete", flush=True)
     logger.info("Memory cleanup complete")
-    log_memory_usage("after cleanup")
+    log_memory_usage("after cleanup")