Update core/app.py

core/app.py

@@ -125,30 +125,53 @@ def load_models(self, progress_callback: Optional[Callable] = None) -> str:
                     progress_callback(0.0, f"Starting model loading on {self.device_manager.get_optimal_device()}")
                 
                 # Load models using the specialized loader
-                sam2_predictor, matanyone_model = self.model_loader.load_all_models(
-                    progress_callback=progress_callback,
-                    cancel_event=self.cancel_event
-                )
+                # Note: load_models returns a boolean, not the models themselves
+                success = self.model_loader.load_models()
                 
                 if self.cancel_event.is_set():
                     return "Model loading cancelled"
                 
+                if not success:
+                    return "Model loading failed - check logs for details"
+                
+                # Get the actual model instances using the getter methods
+                sam2_predictor = self.model_loader.get_sam2()
+                matanyone_model = self.model_loader.get_matanyone()
+                
                 # Initialize core processor with correct parameters
                 self.core_processor = CoreVideoProcessor(
                     config=self.config,
                     models=self.model_loader
                 )
                 
-                # Initialize two-stage processor if available
-                if TWO_STAGE_AVAILABLE and sam2_predictor and matanyone_model:
-                    try:
-                        self.two_stage_processor = TwoStageProcessor(sam2_predictor, matanyone_model)
-                        logger.info("Two-stage processor initialized")
-                    except Exception as e:
-                        logger.warning(f"Two-stage processor init failed: {e}")
+                # Initialize two-stage processor if available and models loaded
+                if TWO_STAGE_AVAILABLE:
+                    if sam2_predictor is not None and matanyone_model is not None:
+                        try:
+                            self.two_stage_processor = TwoStageProcessor(
+                                sam2_predictor=sam2_predictor,
+                                matanyone_model=matanyone_model
+                            )
+                            logger.info("✅ Two-stage processor initialized with AI models")
+                        except Exception as e:
+                            logger.warning(f"Two-stage processor init failed: {e}")
+                            self.two_stage_processor = None
+                    else:
+                        logger.warning("Two-stage processor not initialized - models not available")
+                        if sam2_predictor is None:
+                            logger.warning("  - SAM2 predictor is None")
+                        if matanyone_model is None:
+                            logger.warning("  - MatAnyone model is None")
                 
                 self.models_loaded = True
                 message = self.model_loader.get_load_summary()
+                
+                # Add two-stage status to message
+                if self.two_stage_processor is not None:
+                    message += "\n✅ Two-stage processor ready"
+                else:
+                    message += "\n⚠️ Two-stage processor not available"
+                
                 logger.info(message)
                 return message
                 
@@ -192,12 +215,20 @@ def process_video(
         
         try:
             # Route to appropriate processing method
-            if use_two_stage and TWO_STAGE_AVAILABLE and self.two_stage_processor:
+            if use_two_stage:
+                if not TWO_STAGE_AVAILABLE:
+                    return None, "Two-stage processing not available - module not found"
+                    
+                if self.two_stage_processor is None:
+                    return None, "Two-stage processor not initialized - models may not be loaded properly"
+                
+                logger.info("Using two-stage processing pipeline with AI models")
                 return self._process_two_stage(
                     video_path, background_choice, custom_background_path,
                     progress_callback, chroma_preset
                 )
             else:
+                logger.info("Using single-stage processing pipeline")
                 return self._process_single_stage(
                     video_path, background_choice, custom_background_path,
                     progress_callback, preview_mask, preview_greenscreen
@@ -267,6 +298,9 @@ def _process_two_stage(
     ) -> Tuple[Optional[str], str]:
         """Process video using two-stage pipeline"""
         
+        if self.two_stage_processor is None:
+            return None, "Two-stage processor not available"
+        
         # Get video dimensions for background preparation
         import cv2
         cap = cv2.VideoCapture(video_path)
@@ -290,6 +324,7 @@ def _process_two_stage(
         
         chroma_settings = CHROMA_PRESETS.get(chroma_preset, CHROMA_PRESETS['standard'])
         
+        logger.info(f"Starting two-stage processing with chroma preset: {chroma_preset}")
         result, message = self.two_stage_processor.process_full_pipeline(
             video_path,
             background,
@@ -302,10 +337,10 @@ def _process_two_stage(
             return None, message
         
         success_msg = (
-            f"Two-stage success!\n"
+            f"Two-stage processing completed!\n"
             f"Background: {background_choice}\n"
-            f"Preset: {chroma_preset}\n"
-            f"Quality: Cinema-grade\n"
+            f"Chroma Preset: {chroma_preset}\n"
+            f"Quality: Cinema-grade with AI models\n"
             f"Device: {self.device_manager.get_optimal_device()}"
         )
         
@@ -324,11 +359,19 @@ def get_status(self) -> Dict[str, Any]:
         # Add model-specific status if available
         if self.model_loader:
             base_status['model_loader_available'] = True
+            base_status['sam2_loaded'] = self.model_loader.get_sam2() is not None
+            base_status['matanyone_loaded'] = self.model_loader.get_matanyone() is not None
         
         # Add processing status if available
         if self.core_processor:
             base_status['core_processor_loaded'] = True
         
+        # Add two-stage processor status
+        if self.two_stage_processor:
+            base_status['two_stage_processor_ready'] = True
+        else:
+            base_status['two_stage_processor_ready'] = False
+        
         # Add progress tracking if available
         if self.progress_tracker:
             base_status['progress'] = self.progress_tracker.get_all_progress()
@@ -388,7 +431,7 @@ def main():
     try:
         logger.info("Starting Video Background Replacement application")
         logger.info(f"Device: {processor.device_manager.get_optimal_device()}")
-        logger.info(f"Two-stage available: {TWO_STAGE_AVAILABLE}")
+        logger.info(f"Two-stage module available: {TWO_STAGE_AVAILABLE}")
         logger.info("Modular architecture loaded successfully")
         
         # Import and create UI