Update model_loaders.py

model_loaders.py

@@ -61,7 +61,7 @@ def clear_model_cache():
 def load_sam2_predictor():
     """
     Lazy load SAM2 image predictor with fallback strategies.
-    Returns None if loading fails.
+    Returns (predictor, device) tuple. Returns (None, None) if loading fails.
     """
     try:
         print("Loading SAM2 image predictor...", flush=True)
@@ -70,7 +70,7 @@ def load_sam2_predictor():
         
         # Determine device
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        print(f"Using device: {device}", flush=True)
+        print(f"Using device for SAM2: {device}", flush=True)
         
         # Try local checkpoints first
         checkpoint_path = "/home/user/app/checkpoints/sam2.1_hiera_large.pt"
@@ -104,24 +104,26 @@ def load_sam2_predictor():
                 sam2_model = build_sam2(model_cfg, checkpoint_path, device=device)
                 predictor = SAM2ImagePredictor(sam2_model)
         
-        # Verify model is on correct device
+        # CRITICAL: Verify and force model to correct device
         if hasattr(predictor, 'model'):
             predictor.model.to(device)
-            print(f"Model moved to {device}", flush=True)
+            predictor.model.eval()
+            print(f"SAM2 model moved to {device} and set to eval mode", flush=True)
         
-        print("✅ SAM2 image predictor loaded successfully!", flush=True)
-        return predictor
+        print(f"✅ SAM2 loaded successfully on {device}!", flush=True)
+        return predictor, device
         
     except Exception as e:
-        print(f"Failed to load SAM2 predictor: {e}", flush=True)
+        print(f"❌ Failed to load SAM2 predictor: {e}", flush=True)
         import traceback
         traceback.print_exc()
-        return None
+        return None, None
 
-# Alias for new app.py
+# Alias for backward compatibility
 def load_sam2():
-    """Alias for load_sam2_predictor() - for compatibility with app.py"""
-    return load_sam2_predictor()
+    """Alias for load_sam2_predictor() - returns just predictor for compatibility"""
+    predictor, device = load_sam2_predictor()
+    return predictor
 
 # ============================================================================
 # MatAnyone Model Loading
@@ -130,8 +132,8 @@ def load_sam2():
 @st.cache_resource(show_spinner=False)
 def load_matanyone_processor():
     """
-    Lazy load MatAnyone processor.
-    Returns None if loading fails.
+    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)
@@ -141,21 +143,34 @@ def load_matanyone_processor():
         device = "cuda" if torch.cuda.is_available() else "cpu"
         print(f"MatAnyone using device: {device}", flush=True)
         
+        # Load processor with explicit device
         processor = InferenceCore("PeiqingYang/MatAnyone", device=device)
         
-        print("✅ MatAnyone processor loaded successfully!", flush=True)
-        return processor
+        # 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
+        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)
+        return processor, device
         
     except Exception as e:
-        print(f"Failed to load MatAnyone: {e}", flush=True)
+        print(f"❌ Failed to load MatAnyone: {e}", flush=True)
         import traceback
         traceback.print_exc()
-        return None
+        return None, None
 
-# Alias for new app.py
+# Alias for backward compatibility
 def load_matanyone():
-    """Alias for load_matanyone_processor() - for compatibility with app.py"""
-    return load_matanyone_processor()
+    """Alias for load_matanyone_processor() - returns just processor for compatibility"""
+    processor, device = load_matanyone_processor()
+    return processor
 
 # ============================================================================
 # MediaPipe Pose
@@ -181,15 +196,16 @@ def test_models():
     Returns dict with test results.
     """
     results = {
-        'sam2': {'loaded': False, 'error': None},
-        'matanyone': {'loaded': False, 'error': None}
+        'sam2': {'loaded': False, 'error': None, 'device': None},
+        'matanyone': {'loaded': False, 'error': None, 'device': None}
     }
     
     # Test SAM2
     try:
-        sam2_predictor = load_sam2_predictor()
+        sam2_predictor, sam2_device = 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:
@@ -197,9 +213,10 @@ def test_models():
     
     # Test MatAnyone
     try:
-        matanyone_processor = load_matanyone_processor()
+        matanyone_processor, matanyone_device = 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:
@@ -225,6 +242,7 @@ def log_memory_usage(stage=""):
     
     log_msg += f" | RAM {memory_info['ram_used']:.1f}GB used"
     
+    print(log_msg, flush=True)
     logger.info(log_msg)
     return memory_info
 
@@ -243,6 +261,7 @@ def check_memory_available(required_gb=2.0):
 
 def free_memory_aggressive():
     """Aggressively free GPU and system memory."""
+    print("Performing aggressive memory cleanup...", flush=True)
     logger.info("Performing aggressive memory cleanup...")
     
     # Clear model cache
@@ -260,5 +279,6 @@ def free_memory_aggressive():
     # System cleanup
     gc.collect()
     
+    print("Memory cleanup complete", flush=True)
     logger.info("Memory cleanup complete")
     log_memory_usage("after cleanup")