Update model_loaders.py

model_loaders.py

@@ -64,37 +64,58 @@ def load_sam2_predictor():
     Returns None if loading fails.
     """
     try:
-        logger.info("Loading SAM2 image predictor...")
+        print("Loading SAM2 image predictor...", flush=True)
         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: {device}", flush=True)
+        
         # Try local checkpoints first
         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):
-            logger.warning("Local checkpoints not found, using Hugging Face...")
-            predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large")
+            print("Local checkpoints not found, using Hugging Face...", flush=True)
+            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 gpu_free < 4.0:
-                logger.warning(f"Limited GPU memory ({gpu_free:.1f}GB), using smaller SAM2 model...")
+            if device == "cuda" and gpu_free < 4.0:
+                print(f"Limited GPU memory ({gpu_free:.1f}GB), using smaller SAM2 model...", flush=True)
                 try:
-                    predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-tiny")
+                    predictor = SAM2ImagePredictor.from_pretrained(
+                        "facebook/sam2-hiera-tiny",
+                        device=device
+                    )
                 except Exception:
-                    predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-small")
+                    predictor = SAM2ImagePredictor.from_pretrained(
+                        "facebook/sam2-hiera-small",
+                        device=device
+                    )
             else:
                 # Use local large model
-                predictor = SAM2ImagePredictor(build_sam2(model_cfg, checkpoint_path))
+                sam2_model = build_sam2(model_cfg, checkpoint_path, device=device)
+                predictor = SAM2ImagePredictor(sam2_model)
+        
+        # Verify model is on correct device
+        if hasattr(predictor, 'model'):
+            predictor.model.to(device)
+            print(f"Model moved to {device}", flush=True)
         
-        logger.info("✅ SAM2 image predictor loaded successfully!")
+        print("✅ SAM2 image predictor loaded successfully!", flush=True)
         return predictor
         
     except Exception as e:
-        logger.error(f"Failed to load SAM2 predictor: {e}")
+        print(f"Failed to load SAM2 predictor: {e}", flush=True)
+        import traceback
+        traceback.print_exc()
         return None
 
 # Alias for new app.py
@@ -113,16 +134,22 @@ def load_matanyone_processor():
     Returns None if loading fails.
     """
     try:
-        logger.info("Loading MatAnyone processor...")
+        print("Loading MatAnyone processor...", flush=True)
         from matanyone import InferenceCore
         
-        processor = InferenceCore("PeiqingYang/MatAnyone")
+        # Determine device
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"MatAnyone using device: {device}", flush=True)
+        
+        processor = InferenceCore("PeiqingYang/MatAnyone", device=device)
         
-        logger.info("✅ MatAnyone processor loaded successfully!")
+        print("✅ MatAnyone processor loaded successfully!", flush=True)
         return processor
         
     except Exception as e:
-        logger.error(f"Failed to load MatAnyone: {e}")
+        print(f"Failed to load MatAnyone: {e}", flush=True)
+        import traceback
+        traceback.print_exc()
         return None
 
 # Alias for new app.py
@@ -142,7 +169,7 @@ def load_matanyone():
     enable_segmentation=True,
     min_detection_confidence=0.5
 )
-logger.info("✅ MediaPipe Pose initialized")
+print("✅ MediaPipe Pose initialized", flush=True)
 
 # ============================================================================
 # Model Health Check