Update app.py

app.py

@@ -1,12 +1,9 @@
-hERE IS MY APP
-
-
 #!/usr/bin/env python3
 """
 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
@@ -21,6 +18,7 @@
 from typing import Optional, Tuple, Dict, Any
 import logging
 from huggingface_hub import hf_hub_download
+
 # Import utilities - CRITICAL: Use these functions, don't duplicate!
 from utilities import (
     segment_person_hq,
@@ -31,25 +29,30 @@
     PROFESSIONAL_BACKGROUNDS,
     validate_video_file
 )
+
 # Import two-stage processor if available
 try:
     from two_stage_processor import TwoStageProcessor, CHROMA_PRESETS
     TWO_STAGE_AVAILABLE = True
 except ImportError:
     TWO_STAGE_AVAILABLE = False
+
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
+
 # ============================================================================ #
 # OPTIMIZATION SETTINGS
 # ============================================================================ #
 KEYFRAME_INTERVAL = 5  # Process MatAnyone every 5th frame
 FRAME_SKIP = 1        # Process every frame (set to 2 for every other frame)
 MEMORY_CLEANUP_INTERVAL = 30  # Clean memory every 30 frames
+
 # ============================================================================ #
 # MODEL CACHING SYSTEM
 # ============================================================================ #
 CACHE_DIR = Path("/tmp/model_cache")
 CACHE_DIR.mkdir(exist_ok=True, parents=True)
+
 # ============================================================================ #
 # GLOBAL MODEL STATE
 # ============================================================================ #
@@ -59,14 +62,33 @@
 loading_lock = threading.Lock()
 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:
             progress_callback(pct, desc)
+
         # Format progress info for display in the UI
         if "Frame" in desc and "|" in desc:
             parts = desc.split("|")
@@ -87,8 +109,10 @@ def _prog(pct: float, desc: str):
                     f.write(display_text)
             except Exception as e:
                 logger.warning(f"Error writing processing info: {e}")
+
     try:
         _prog(0.1, "Initializing SAM2...")
+
         # Download checkpoint with caching
         checkpoint_path = hf_hub_download(
             repo_id="facebook/sam2-hiera-large",
@@ -96,14 +120,18 @@ 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
         from sam2.build_sam import build_sam2
         from sam2.sam2_image_predictor import SAM2ImagePredictor
+
         # Build model with explicit config
         sam2_model = build_sam2("sam2_hiera_l.yaml", checkpoint_path)
         sam2_model.to(device)
         predictor = SAM2ImagePredictor(sam2_model)
+
         # Test the predictor with dummy data
         _prog(0.8, "Testing SAM2 functionality...")
         test_image = np.zeros((256, 256, 3), dtype=np.uint8)
@@ -115,15 +143,19 @@ def _prog(pct: float, desc: str):
             point_labels=test_labels,
             multimask_output=False
         )
+
         if masks is None or len(masks) == 0:
             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:
         logger.error(f"SAM2 loading failed: {str(e)}")
         logger.error(f"Full traceback: {traceback.format_exc()}")
         raise Exception(f"SAM2 loading failed: {str(e)}")
+
 # ============================================================================ #
 # MATANYONE LOADER WITH VALIDATION
 # ============================================================================ #
@@ -132,15 +164,19 @@ def load_matanyone_fixed(progress_callback: Optional[callable] = None) -> Any:
     def _prog(pct: float, desc: str):
         if progress_callback:
             progress_callback(pct, desc)
+
     try:
         _prog(0.2, "Loading MatAnyone...")
+
         from matanyone import InferenceCore
         processor = InferenceCore("PeiqingYang/MatAnyone")
+
         # Test MatAnyone with dummy data
         _prog(0.8, "Testing MatAnyone functionality...")
         test_image = np.zeros((256, 256, 3), dtype=np.uint8)
         test_mask = np.zeros((256, 256), dtype=np.uint8)
         test_mask[64:192, 64:192] = 255
+
         # Test the processor
         try:
             if hasattr(processor, 'process') or hasattr(processor, '__call__'):
@@ -149,13 +185,16 @@ def _prog(pct: float, desc: str):
                 logger.warning("MatAnyone interface unclear, will use fallback refinement")
         except Exception as test_e:
             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:
         logger.error(f"MatAnyone loading failed: {str(e)}")
         logger.error(f"Full traceback: {traceback.format_exc()}")
         raise Exception(f"MatAnyone loading failed: {str(e)}")
+
 # ============================================================================ #
 # MODEL MANAGEMENT FUNCTIONS
 # ============================================================================ #
@@ -165,53 +204,68 @@ 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]:
     """Get current cache status"""
     return {
         "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:
     """Load models with comprehensive validation"""
     global sam2_predictor, matanyone_model, models_loaded, two_stage_processor, PROCESS_CANCELLED
+
     with loading_lock:
         if models_loaded and not PROCESS_CANCELLED:
             return "Models already loaded and validated"
+
         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"
+
             # Load MatAnyone with validation
             matanyone_model = load_matanyone_fixed(progress_callback=progress_callback)
+
             if PROCESS_CANCELLED:
                 return "Model loading cancelled by user"
+
             models_loaded = True
+
             # Initialize two-stage processor if available
             if TWO_STAGE_AVAILABLE:
                 two_stage_processor = TwoStageProcessor(sam2_predictor, matanyone_model)
                 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)
             return message
+
         except Exception as e:
             models_loaded = False
             error_msg = f"Model loading failed: {str(e)}"
             logger.error(error_msg)
             return error_msg
+
 # ============================================================================ #
 # MAIN VIDEO PROCESSING - USING UTILITIES FUNCTIONS
 # ============================================================================ #
@@ -227,21 +281,28 @@ def process_video_fixed(
 ) -> Tuple[Optional[str], str]:
     """Optimized video processing using proper functions from utilities"""
     global PROCESS_CANCELLED
+
     if PROCESS_CANCELLED:
         return None, "Processing cancelled by user"
+
     if not models_loaded:
         return None, "Models not loaded. Call load_models_with_validation() first."
+
     if not video_path or not os.path.exists(video_path):
         return None, f"Video file not found: {video_path}"
+
     # Validate video file
     is_valid, validation_msg = validate_video_file(video_path)
     if not is_valid:
         return None, f"Invalid video: {validation_msg}"
+
     def _prog(pct: float, desc: str):
         if PROCESS_CANCELLED:
             raise Exception("Processing cancelled by user")
+
         if progress_callback:
             progress_callback(pct, desc)
+
         # Update processing info file
         if "Frame" in desc and "|" in desc:
             parts = desc.split("|")
@@ -249,6 +310,7 @@ 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}
@@ -262,24 +324,31 @@ def _prog(pct: float, desc: str):
                     f.write(display_text)
             except Exception as e:
                 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:
             if not TWO_STAGE_AVAILABLE:
                 return None, "Two-stage mode not available. Please add two_stage_processor.py file."
+
             if two_stage_processor is None:
                 return None, "Two-stage processor not initialized. Please reload models."
+
             _prog(0.05, "Starting TWO-STAGE green screen processing...")
+
             # Get video dimensions
             cap = cv2.VideoCapture(video_path)
             frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
             frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
             cap.release()
+
             # Prepare background
             if background_choice == "custom" and custom_background_path:
                 if not os.path.exists(custom_background_path):
                     return None, f"Custom background not found: {custom_background_path}"
+
                 background = cv2.imread(custom_background_path)
                 if background is None:
                     return None, "Could not read custom background image."
@@ -291,11 +360,14 @@ def _prog(pct: float, desc: str):
                     background_name = bg_config["name"]
                 else:
                     return None, f"Invalid background selection: {background_choice}"
+
             # Get chroma settings
             chroma_settings = CHROMA_PRESETS.get(chroma_preset, CHROMA_PRESETS['standard'])
+
             # Run two-stage pipeline
             timestamp = int(time.time())
             final_output = f"/tmp/twostage_final_{timestamp}.mp4"
+
             result, message = two_stage_processor.process_full_pipeline(
                 video_path,
                 background,
@@ -303,13 +375,17 @@ def _prog(pct: float, desc: str):
                 chroma_settings=chroma_settings,
                 progress_callback=_prog
             )
+
             if PROCESS_CANCELLED:
                 return None, "Processing cancelled by user"
+
             if result is None:
                 return None, message
+
             # Add audio back
             _prog(0.9, "Adding audio...")
             final_with_audio = f"/tmp/twostage_audio_{timestamp}.mp4"
+
             try:
                 audio_cmd = (
                     f'ffmpeg -y -i "{final_output}" -i "{video_path}" '
@@ -324,34 +400,46 @@ def _prog(pct: float, desc: str):
             except Exception as e:
                 logger.warning(f"Audio processing error: {e}")
                 final_with_audio = final_output  # Fallback to video without audio
+
             _prog(1.0, "TWO-STAGE processing complete!")
+
             success_message = (
                 f"TWO-STAGE Success!\n"
                 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():
             return None, "Could not open video file."
+
         fps = cap.get(cv2.CAP_PROP_FPS)
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
         if total_frames == 0:
             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
         background_name = ""
+
         if background_choice == "custom" and custom_background_path:
             if not os.path.exists(custom_background_path):
                 return None, f"Custom background not found: {custom_background_path}"
+
             background = cv2.imread(custom_background_path)
             if background is None:
                 return None, "Could not read custom background image."
@@ -363,11 +451,15 @@ def _prog(pct: float, desc: str):
                 background_name = bg_config["name"]
             else:
                 return None, f"Invalid background selection: {background_choice}"
+
         if background is None:
             return None, "Failed to create background."
+
         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:
             temp_output = f"/tmp/preview_{timestamp}.mp4"
@@ -375,13 +467,17 @@ def _prog(pct: float, desc: str):
         else:
             final_path = f"/tmp/output_{timestamp}.mp4"
             final_writer = cv2.VideoWriter(final_path, fourcc, fps, (frame_width, frame_height))
+
         if not final_writer.isOpened():
             return None, "Could not create output video file."
+
         frame_count = 0
         successful_frames = 0
         last_refined_mask = None
+
         # Processing stats
         start_time = time.time()
+
         while True:
             if PROCESS_CANCELLED:
                 cap.release()
@@ -389,13 +485,16 @@ def _prog(pct: float, desc: str):
                 if os.path.exists(final_path):
                     os.remove(final_path)
                 return None, "Processing cancelled by user"
+
             ret, frame = cap.read()
             if not ret:
                 break
+
             # Skip frames if FRAME_SKIP > 1
             if frame_count % FRAME_SKIP != 0:
                 frame_count += 1
                 continue
+
             try:
                 # Update progress with detailed timing info and ETA
                 elapsed_time = time.time() - start_time
@@ -403,13 +502,17 @@ def _prog(pct: float, desc: str):
                 remaining_frames = total_frames - frame_count
                 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)
                 _prog(0.1 + (frame_count / max(1, total_frames)) * 0.8, progress_msg)
+
                 # CRITICAL: Use functions from utilities.py, not local implementations!
                 # SAM2 segmentation using utilities function
                 mask = segment_person_hq(frame, sam2_predictor)
+
                 if preview_mask:
                     # Save mask visualization
                     mask_vis = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
@@ -417,15 +520,17 @@ def _prog(pct: float, desc: str):
                     final_writer.write(mask_vis)
                     frame_count += 1
                     continue
+
                 # MatAnyone refinement on keyframes using utilities function
                 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
                     refined_mask = cv2.addWeighted(mask, alpha, last_refined_mask, 1-alpha, 0)
+
                 if preview_greenscreen:
                     # Create green screen preview
                     green_bg = np.zeros_like(frame)
@@ -437,15 +542,19 @@ def _prog(pct: float, desc: str):
                     final_writer.write(preview_frame.astype(np.uint8))
                     frame_count += 1
                     continue
+
                 # CRITICAL: Use replace_background_hq from utilities which has the transparency fix!
                 result_frame = replace_background_hq(frame, refined_mask, background)
                 final_writer.write(result_frame)
                 successful_frames += 1
+
             except Exception as frame_error:
                 logger.warning(f"Error processing frame {frame_count}: {frame_error}")
                 # Write original frame if processing fails
                 final_writer.write(frame)
+
             frame_count += 1
+
             # Memory management
             if frame_count % MEMORY_CLEANUP_INTERVAL == 0:
                 gc.collect()
@@ -454,25 +563,32 @@ 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()
+
         if PROCESS_CANCELLED:
             if os.path.exists(final_path):
                 os.remove(final_path)
             return None, "Processing cancelled by user"
+
         if successful_frames == 0:
             return None, "No frames were processed successfully with AI."
+
         # Calculate processing stats
         total_time = time.time() - start_time
         avg_fps = frame_count / total_time if total_time > 0 else 0
+
         _prog(0.9, "Finalizing output...")
+
         if preview_mask or preview_greenscreen:
             final_output = temp_output
         else:
             # Add audio back for final output
             _prog(0.9, "Adding audio...")
             final_output = f"/tmp/final_{timestamp}.mp4"
+
             try:
                 audio_cmd = (
                     f'ffmpeg -y -i "{final_path}" -i "{video_path}" '
@@ -486,13 +602,16 @@ def _prog(pct: float, desc: str):
             except Exception as e:
                 logger.warning(f"Audio processing error: {e}")
                 shutil.copy2(final_path, final_output)
+
             # Cleanup
             try:
                 if os.path.exists(final_path):
                     os.remove(final_path)
             except Exception as e:
                 logger.warning(f"Cleanup error: {e}")
+
         _prog(1.0, "Processing complete!")
+
         success_message = (
             f"Success!\n"
             f"Background: {background_name}\n"
@@ -502,12 +621,16 @@ 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
+
     except Exception as e:
         logger.error(f"Processing error: {traceback.format_exc()}")
         return None, f"Processing Error: {str(e)}"
+
 # ============================================================================ #
 # MAIN - IMPORT UI COMPONENTS
 # ============================================================================ #
@@ -517,13 +640,18 @@ 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
         from ui_components import create_interface
+
         os.makedirs("/tmp/MyAvatar/My_Videos/", exist_ok=True)
         CACHE_DIR.mkdir(exist_ok=True, parents=True)
+
         print("Creating interface...")
         demo = create_interface()
+
         print("Launching...")
         demo.launch(
             server_name="0.0.0.0",
@@ -533,10 +661,10 @@ def main():
             debug=True,
             enable_queue=True
         )
+
     except Exception as e:
         logger.error(f"Startup failed: {e}")
         print(f"Startup failed: {e}")
+
 if __name__ == "__main__":
     main()
-
-pLEASE UPDATE