Update app.py

app.py

@@ -2,8 +2,8 @@
 """
 Final Fixed Video Background Replacement
 Uses proper functions from utilities.py to avoid transparency issues
-NEW: Added GPU detection, model caching, batch processing support,
-     and improved error handling
+NEW: Added automatic device detection for Hugging Face Spaces compatibility,
+     improved error handling, and better resource management
 """
 import sys
 import cv2
@@ -63,10 +63,27 @@
 two_stage_processor = None
 PROCESS_CANCELLED = False
 
+# ============================================================================ #
+# DEVICE DETECTION FOR HUGGING FACE SPACES
+# ============================================================================ #
+def get_device():
+    """Automatically detect the best available device (CPU or GPU)"""
+    if torch.cuda.is_available():
+        # Get the current CUDA device name
+        device_name = torch.cuda.get_device_name(0)
+        logger.info(f"Using GPU: {device_name}")
+        return "cuda"
+    else:
+        logger.info("Using CPU (no GPU available)")
+        return "cpu"
+
+# Set the device globally
+DEVICE = get_device()
+
 # ============================================================================ #
 # SAM2 LOADER WITH VALIDATION
 # ============================================================================ #
-def load_sam2_predictor_fixed(device: str = "cuda", progress_callback: Optional[callable] = None) -> Any:
+def load_sam2_predictor_fixed(device: str = DEVICE, progress_callback: Optional[callable] = None) -> Any:
     """Load SAM2 with proper error handling and validation"""
     def _prog(pct: float, desc: str):
         if progress_callback:
@@ -79,7 +96,6 @@ def _prog(pct: float, desc: str):
             time_info = parts[1].strip() if len(parts) > 1 else ""
             fps_info = parts[2].strip() if len(parts) > 2 else ""
             eta_info = parts[3].strip() if len(parts) > 3 else ""
-
             display_text = f"""📊 PROCESSING STATUS
 ━━━━━━━━━━━━━━━━━━━━━━━━━━
 🎬 {frame_info}
@@ -88,7 +104,6 @@ def _prog(pct: float, desc: str):
 🎯 {eta_info}
 ━━━━━━━━━━━━━━━━━━━━━━━━━━
 📈 Progress: {pct*100:.1f}%"""
-
             try:
                 with open("/tmp/processing_info.txt", 'w') as f:
                     f.write(display_text)
@@ -105,6 +120,7 @@ def _prog(pct: float, desc: str):
             cache_dir=str(CACHE_DIR / "sam2_checkpoint"),
             force_download=False
         )
+
         _prog(0.5, "SAM2 checkpoint downloaded, building model...")
 
         # Import and build
@@ -132,7 +148,7 @@ def _prog(pct: float, desc: str):
             raise Exception("SAM2 predictor test failed - no masks generated")
 
         _prog(1.0, "SAM2 loaded and validated successfully!")
-        logger.info("SAM2 predictor loaded and tested successfully")
+        logger.info(f"SAM2 predictor loaded and tested successfully on {device}")
         return predictor
 
     except Exception as e:
@@ -171,7 +187,7 @@ def _prog(pct: float, desc: str):
             logger.warning(f"MatAnyone test failed: {test_e}, will use enhanced OpenCV")
 
         _prog(1.0, "MatAnyone loaded successfully!")
-        logger.info("MatAnyone processor loaded successfully")
+        logger.info(f"MatAnyone processor loaded successfully on {DEVICE}")
         return processor
 
     except Exception as e:
@@ -188,7 +204,8 @@ def get_model_status() -> Dict[str, str]:
     return {
         'sam2': 'Ready' if sam2_predictor is not None else 'Not loaded',
         'matanyone': 'Ready' if matanyone_model is not None else 'Not loaded',
-        'validated': models_loaded
+        'validated': models_loaded,
+        'device': DEVICE
     }
 
 def get_cache_status() -> Dict[str, Any]:
@@ -197,7 +214,8 @@ def get_cache_status() -> Dict[str, Any]:
         "sam2_loaded": sam2_predictor is not None,
         "matanyone_loaded": matanyone_model is not None,
         "models_validated": models_loaded,
-        "two_stage_available": TWO_STAGE_AVAILABLE
+        "two_stage_available": TWO_STAGE_AVAILABLE,
+        "device": DEVICE
     }
 
 def load_models_with_validation(progress_callback: Optional[callable] = None) -> str:
@@ -211,14 +229,13 @@ def load_models_with_validation(progress_callback: Optional[callable] = None) ->
         try:
             PROCESS_CANCELLED = False
             start_time = time.time()
-            device = "cuda" if torch.cuda.is_available() else "cpu"
-            logger.info(f"Starting model loading on {device}")
+            logger.info(f"Starting model loading on {DEVICE}")
 
             if progress_callback:
-                progress_callback(0.0, "Starting model loading...")
+                progress_callback(0.0, f"Starting model loading on {DEVICE}...")
 
             # Load SAM2 with validation
-            sam2_predictor = load_sam2_predictor_fixed(device=device, progress_callback=progress_callback)
+            sam2_predictor = load_sam2_predictor_fixed(device=DEVICE, progress_callback=progress_callback)
 
             if PROCESS_CANCELLED:
                 return "Model loading cancelled by user"
@@ -237,7 +254,7 @@ def load_models_with_validation(progress_callback: Optional[callable] = None) ->
                 logger.info("Two-stage processor initialized")
 
             load_time = time.time() - start_time
-            message = f"SUCCESS: SAM2 + MatAnyone loaded and validated in {load_time:.1f}s"
+            message = f"SUCCESS: SAM2 + MatAnyone loaded and validated in {load_time:.1f}s on {DEVICE}"
             if TWO_STAGE_AVAILABLE:
                 message += " (Two-stage mode available)"
             logger.info(message)
@@ -302,7 +319,6 @@ def _prog(pct: float, desc: str):
 🎯 {eta_info}
 ━━━━━━━━━━━━━━━━━━━━━━━━━━
 📈 Progress: {pct*100:.1f}%"""
-
             try:
                 with open("/tmp/processing_info.txt", 'w') as f:
                     f.write(display_text)
@@ -310,7 +326,7 @@ def _prog(pct: float, desc: str):
                 logger.warning(f"Error writing processing info: {e}")
 
     try:
-        _prog(0.0, f"Starting {'TWO-STAGE' if use_two_stage else 'SINGLE-STAGE'} processing...")
+        _prog(0.0, f"Starting {'TWO-STAGE' if use_two_stage else 'SINGLE-STAGE'} processing on {DEVICE}...")
 
         # Check if two-stage mode is requested
         if use_two_stage:
@@ -337,7 +353,6 @@ def _prog(pct: float, desc: str):
                 if background is None:
                     return None, "Could not read custom background image."
                 background_name = "Custom Image"
-
             else:
                 if background_choice in PROFESSIONAL_BACKGROUNDS:
                     bg_config = PROFESSIONAL_BACKGROUNDS[background_choice]
@@ -393,13 +408,14 @@ def _prog(pct: float, desc: str):
                 f"Background: {background_name}\n"
                 f"Method: Green Screen Chroma Key\n"
                 f"Preset: {chroma_preset}\n"
-                f"Quality: Professional cinema-grade"
+                f"Quality: Professional cinema-grade\n"
+                f"Device: {DEVICE}"
             )
 
             return final_output, success_message
 
         # Single-stage processing
-        _prog(0.05, "Starting SINGLE-STAGE processing...")
+        _prog(0.05, f"Starting SINGLE-STAGE processing on {DEVICE}...")
 
         cap = cv2.VideoCapture(video_path)
         if not cap.isOpened():
@@ -414,7 +430,7 @@ def _prog(pct: float, desc: str):
             return None, "Video appears to be empty."
 
         # Log video info
-        logger.info(f"Video info: {frame_width}x{frame_height}, {fps}fps, {total_frames} frames")
+        logger.info(f"Video info: {frame_width}x{frame_height}, {fps}fps, {total_frames} frames, processing on {DEVICE}")
 
         # Prepare background
         background = None
@@ -428,7 +444,6 @@ def _prog(pct: float, desc: str):
             if background is None:
                 return None, "Could not read custom background image."
             background_name = "Custom Image"
-
         else:
             if background_choice in PROFESSIONAL_BACKGROUNDS:
                 bg_config = PROFESSIONAL_BACKGROUNDS[background_choice]
@@ -443,7 +458,7 @@ def _prog(pct: float, desc: str):
         timestamp = int(time.time())
         fourcc = cv2.VideoWriter_fourcc(*'avc1')  # H.264 for better compatibility
 
-        _prog(0.1, f"Processing {total_frames} frames with {'TWO-STAGE' if use_two_stage else 'SINGLE-STAGE'} processing...")
+        _prog(0.1, f"Processing {total_frames} frames with {'TWO-STAGE' if use_two_stage else 'SINGLE-STAGE'} processing on {DEVICE}...")
 
         # Create temporary output for preview if needed
         if preview_mask or preview_greenscreen:
@@ -488,7 +503,7 @@ def _prog(pct: float, desc: str):
                 eta_seconds = remaining_frames / current_fps if current_fps > 0 else 0
                 eta_display = f"{int(eta_seconds//60)}m {int(eta_seconds%60)}s" if eta_seconds > 60 else f"{int(eta_seconds)}s"
 
-                progress_msg = f"Frame {frame_count + 1}/{total_frames} | {elapsed_time:.1f}s | {current_fps:.1f} fps | ETA: {eta_display}"
+                progress_msg = f"Frame {frame_count + 1}/{total_frames} | {elapsed_time:.1f}s | {current_fps:.1f} fps | ETA: {eta_display} | Device: {DEVICE}"
 
                 # Log and display progress
                 logger.info(progress_msg)
@@ -510,7 +525,7 @@ def _prog(pct: float, desc: str):
                 if (frame_count % KEYFRAME_INTERVAL == 0) or (last_refined_mask is None):
                     refined_mask = refine_mask_hq(frame, mask, matanyone_model)
                     last_refined_mask = refined_mask.copy()
-                    logger.info(f"Keyframe refinement at frame {frame_count}")
+                    logger.info(f"Keyframe refinement at frame {frame_count} on {DEVICE}")
                 else:
                     # Blend SAM2 mask with last refined mask for temporal smoothness
                     alpha = 0.7
@@ -548,7 +563,7 @@ def _prog(pct: float, desc: str):
                 elapsed = time.time() - start_time
                 fps_actual = frame_count / elapsed
                 eta = (total_frames - frame_count) / fps_actual if fps_actual > 0 else 0
-                logger.info(f"Progress: {frame_count}/{total_frames}, FPS: {fps_actual:.1f}, ETA: {eta:.0f}s")
+                logger.info(f"Progress: {frame_count}/{total_frames}, FPS: {fps_actual:.1f}, ETA: {eta:.0f}s, Device: {DEVICE}")
 
         cap.release()
         final_writer.release()
@@ -606,7 +621,8 @@ def _prog(pct: float, desc: str):
             f"Processing time: {total_time:.1f}s\n"
             f"Average FPS: {avg_fps:.1f}\n"
             f"Keyframe interval: {KEYFRAME_INTERVAL}\n"
-            f"Mode: {'TWO-STAGE' if use_two_stage else 'SINGLE-STAGE'}"
+            f"Mode: {'TWO-STAGE' if use_two_stage else 'SINGLE-STAGE'}\n"
+            f"Device: {DEVICE}"
         )
 
         return final_output, success_message
@@ -624,6 +640,7 @@ def main():
         print(f"Keyframe interval: {KEYFRAME_INTERVAL} frames")
         print(f"Frame skip: {FRAME_SKIP} (1=all frames, 2=every other)")
         print(f"Two-stage mode: {'AVAILABLE' if TWO_STAGE_AVAILABLE else 'NOT AVAILABLE'}")
+        print(f"Device: {DEVICE}")
         print("Loading UI components...")
 
         # Import UI components