Update core/app.py

core/app.py

@@ -38,6 +38,10 @@
 if str(ROOT) not in sys.path:
     sys.path.insert(0, str(ROOT))
 
+# Create loader directories if they don't exist
+loaders_dir = ROOT / "models" / "loaders"
+loaders_dir.mkdir(parents=True, exist_ok=True)
+
 # ── Gradio schema patch (HF quirk) ───────────────────────────────────────────
 try:
     import gradio_client.utils as gc_utils
@@ -59,28 +63,50 @@ def _patched_get_type(schema):
 from core.exceptions import ModelLoadingError, VideoProcessingError
 from utils.hardware.device_manager import DeviceManager
 from utils.system.memory_manager import MemoryManager
-from models.loaders.model_loader import ModelLoader
+
+# Try to import the new split loaders first, fall back to old if needed
+try:
+    from models.loaders.model_loader import ModelLoader
+    logger.info("Using split loader architecture")
+except ImportError:
+    logger.warning("Split loaders not found, using legacy loader")
+    # Fall back to old loader if split architecture isn't available yet
+    from models.model_loader import ModelLoader  # type: ignore
+
 from processing.video.video_processor import CoreVideoProcessor
 from processing.audio.audio_processor import AudioProcessor
 from utils.monitoring.progress_tracker import ProgressTracker
 from utils.cv_processing import validate_video_file
 
-# ── Optional Two-Stage import (exact module path) ────────────────────────────
+# ── Optional Two-Stage import ────────────────────────────────────────────────
 TWO_STAGE_AVAILABLE = False
 TWO_STAGE_IMPORT_ORIGIN = ""
 TWO_STAGE_IMPORT_ERROR = ""
 CHROMA_PRESETS: Dict[str, Dict[str, Any]] = {"standard": {}}
 TwoStageProcessor = None  # type: ignore
 
-try:
-    from processing.two_stage.two_stage_processor import TwoStageProcessor, CHROMA_PRESETS  # type: ignore
-    TWO_STAGE_AVAILABLE = True
-    TWO_STAGE_IMPORT_ORIGIN = "processing.two_stage.two_stage_processor"
-    logger.info("Two-stage import OK (%s)", TWO_STAGE_IMPORT_ORIGIN)
-except Exception as e:
-    TWO_STAGE_AVAILABLE = False
-    TWO_STAGE_IMPORT_ERROR = f"{repr(e)}\n{traceback.format_exc()}"
-    logger.warning("Two-stage import FAILED: %s", TWO_STAGE_IMPORT_ERROR.strip().splitlines()[-1])
+# Try multiple import paths for two-stage processor
+two_stage_paths = [
+    "processors.two_stage",  # Your fixed version
+    "processing.two_stage.two_stage_processor",
+    "processing.two_stage",
+]
+
+for import_path in two_stage_paths:
+    try:
+        if "processors" in import_path:
+            from processors.two_stage import TwoStageProcessor, CHROMA_PRESETS  # type: ignore
+        else:
+            from processing.two_stage.two_stage_processor import TwoStageProcessor, CHROMA_PRESETS  # type: ignore
+        TWO_STAGE_AVAILABLE = True
+        TWO_STAGE_IMPORT_ORIGIN = import_path
+        logger.info(f"Two-stage import OK ({import_path})")
+        break
+    except Exception as e:
+        continue
+
+if not TWO_STAGE_AVAILABLE:
+    logger.warning("Two-stage import FAILED from all paths")
 
 # ╔══════════════════════════════════════════════════════════════════════════╗
 # ║                           VideoProcessor class                           ║
@@ -99,13 +125,13 @@ def __init__(self):
         self.model_loader = ModelLoader(self.device_manager, self.memory_manager)
 
         self.audio_processor = AudioProcessor()
-        self.core_processor: CoreVideoProcessor | None = None
-        self.two_stage_processor: Any | None = None
+        self.core_processor: Optional[CoreVideoProcessor] = None
+        self.two_stage_processor: Optional[Any] = None
 
         self.models_loaded = False
         self.loading_lock = threading.Lock()
         self.cancel_event = threading.Event()
-        self.progress_tracker: ProgressTracker | None = None
+        self.progress_tracker: Optional[ProgressTracker] = None
 
         logger.info(f"VideoProcessor on device: {self.device_manager.get_optimal_device()}")
 
@@ -124,6 +150,9 @@ def _patch_config_defaults(cfg: Any) -> None:
             "max_model_size_bytes": 0,
             # housekeeping
             "output_dir": str((Path(__file__).resolve().parent.parent) / "outputs"),
+            # MatAnyone settings to ensure it's enabled
+            "matanyone_enabled": True,
+            "use_matanyone": True,
         }
         for k, v in defaults.items():
             if not hasattr(cfg, k):
@@ -161,28 +190,39 @@ def load_models(self, progress_callback: Optional[Callable] = None) -> str:
                 if self.cancel_event.is_set():
                     return "Model loading cancelled"
 
-                # unwrap
-                sam2_predictor = getattr(sam2_loaded, "model", None) if sam2_loaded else None
-                mat_model = getattr(mat_loaded, "model", None) if mat_loaded else None
+                # Get the actual models
+                sam2_predictor = sam2_loaded.model if sam2_loaded else None
+                mat_model = mat_loaded.model if mat_loaded else None
 
-                # single-stage
+                # Initialize core processor
                 self.core_processor = CoreVideoProcessor(config=self.config, models=self.model_loader)
 
-                # two-stage (optional)
+                # Initialize two-stage processor if available
                 self.two_stage_processor = None
-                if TWO_STAGE_AVAILABLE and (TwoStageProcessor is not None) and (sam2_predictor or mat_model):
+                if TWO_STAGE_AVAILABLE and TwoStageProcessor and (sam2_predictor or mat_model):
                     try:
                         self.two_stage_processor = TwoStageProcessor(
                             sam2_predictor=sam2_predictor, matanyone_model=mat_model
                         )
                         logger.info("Two-stage processor initialised")
                     except Exception as e:
-                        logger.warning("Two-stage init failed: %r", e)
+                        logger.warning(f"Two-stage init failed: {e}")
                         self.two_stage_processor = None
 
                 self.models_loaded = True
                 msg = self.model_loader.get_load_summary()
-                msg += ("\n✅ Two-stage processor ready" if self.two_stage_processor else "\n⚠️ Two-stage processor not available")
+                
+                # Add status about processors
+                if self.two_stage_processor:
+                    msg += "\n✅ Two-stage processor ready"
+                else:
+                    msg += "\n⚠️ Two-stage processor not available"
+                    
+                if mat_model:
+                    msg += "\n✅ MatAnyone refinement active"
+                else:
+                    msg += "\n⚠️ MatAnyone not loaded (edges may be rough)"
+                    
                 logger.info(msg)
                 return msg
 
@@ -194,8 +234,7 @@ def load_models(self, progress_callback: Optional[Callable] = None) -> str:
             except Exception as e:
                 self.models_loaded = False
                 err = f"Unexpected error during model loading: {e}"
-                logger.error(err)
-                logger.debug("Traceback:\n%s", traceback.format_exc())
+                logger.error(f"{err}\n{traceback.format_exc()}")
                 return err
 
     # ── Public entry – process video ─────────────────────────────────────────
@@ -212,7 +251,7 @@ def process_video(
         preview_greenscreen: bool = False,
     ) -> Tuple[Optional[str], str]:
         if not self.models_loaded or not self.core_processor:
-            return None, "Models not loaded. Please click “Load Models” first."
+            return None, "Models not loaded. Please click 'Load Models' first."
         if self.cancel_event.is_set():
             return None, "Processing cancelled"
 
@@ -223,9 +262,14 @@ def process_video(
             return None, f"Invalid video: {why}"
 
         try:
+            # Log which mode we're using
+            mode = "two-stage" if use_two_stage else "single-stage"
+            matanyone_status = "enabled" if self.model_loader.get_matanyone() else "disabled"
+            logger.info(f"Processing video in {mode} mode, MatAnyone: {matanyone_status}")
+            
             if use_two_stage:
                 if not TWO_STAGE_AVAILABLE or self.two_stage_processor is None:
-                    return None, "Two-stage processing not available on this build"
+                    return None, "Two-stage processing not available"
                 return self._process_two_stage(
                     video_path,
                     background_choice,
@@ -248,8 +292,7 @@ def process_video(
             logger.error(f"Processing failed: {e}")
             return None, f"Processing failed: {e}"
         except Exception as e:
-            logger.error(f"Unexpected processing error: {e}")
-            logger.debug("Traceback:\n%s", traceback.format_exc())
+            logger.error(f"Unexpected processing error: {e}\n{traceback.format_exc()}")
             return None, f"Unexpected error: {e}"
 
     # ── Private – single-stage ───────────────────────────────────────────────
@@ -268,6 +311,7 @@ def _process_single_stage(
         out_dir.mkdir(parents=True, exist_ok=True)
         out_path = str(out_dir / f"processed_{ts}.mp4")
 
+        # Process video
         result = self.core_processor.process_video(
             input_path=video_path,
             output_path=out_path,
@@ -275,7 +319,9 @@ def _process_single_stage(
                 "background_choice": background_choice,
                 "custom_path": custom_background_path,
             },
+            progress_callback=progress_callback,
         )
+        
         if not result:
             return None, "Video processing failed"
 
@@ -286,16 +332,19 @@ def _process_single_stage(
                     original_video=video_path, processed_video=out_path
                 )
             except Exception as e:
-                logger.warning("Audio mux failed, returning video without audio: %r", e)
+                logger.warning(f"Audio mux failed, returning video without audio: {e}")
                 final_path = out_path
         else:
             final_path = out_path
 
+        # Build status message
+        matanyone_status = "✓" if self.model_loader.get_matanyone() else "✗"
         msg = (
             "Processing completed.\n"
             f"Frames: {result.get('frames', 'unknown')}\n"
             f"Background: {background_choice}\n"
             f"Mode: Single-stage\n"
+            f"MatAnyone: {matanyone_status}\n"
             f"Device: {self.device_manager.get_optimal_device()}"
         )
         return final_path, msg
@@ -321,13 +370,13 @@ def _process_two_stage(
         h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) or 720
         cap.release()
 
-        # Prepare background via core processor
+        # Prepare background
         try:
             background = self.core_processor.prepare_background(
                 background_choice, custom_background_path, w, h
             )
         except Exception as e:
-            logger.error("Background preparation failed: %r", e)
+            logger.error(f"Background preparation failed: {e}")
             return None, f"Failed to prepare background: {e}"
         if background is None:
             return None, "Failed to prepare background"
@@ -338,7 +387,7 @@ def _process_two_stage(
         final_out = str(out_dir / f"final_{ts}.mp4")
 
         chroma_cfg = CHROMA_PRESETS.get(chroma_preset, CHROMA_PRESETS.get("standard", {}))
-        logger.info("Two-stage with preset: %s | key_color_mode=%s", chroma_preset, key_color_mode)
+        logger.info(f"Two-stage with preset: {chroma_preset} | key_color: {key_color_mode}")
 
         result, message = self.two_stage_processor.process_full_pipeline(
             video_path,
@@ -351,19 +400,21 @@ def _process_two_stage(
         if result is None:
             return None, message
 
-        # Mux audio from original (same logic as single-stage)
+        # Mux audio
         try:
             final_path = self.audio_processor.add_audio_to_video(
                 original_video=video_path, processed_video=result
             )
         except Exception as e:
-            logger.warning("Audio mux failed for two-stage; returning video without audio: %r", e)
+            logger.warning(f"Audio mux failed: {e}")
             final_path = result
 
+        matanyone_status = "✓" if self.model_loader.get_matanyone() else "✗"
         msg = (
             "Two-stage processing completed.\n"
             f"Background: {background_choice}\n"
             f"Chroma Preset: {chroma_preset}\n"
+            f"MatAnyone: {matanyone_status}\n"
             f"Device: {self.device_manager.get_optimal_device()}"
         )
         return final_path, msg
@@ -372,23 +423,25 @@ def _process_two_stage(
     def get_status(self) -> Dict[str, Any]:
         status = {
             "models_loaded": self.models_loaded,
-            "two_stage_available": bool(TWO_STAGE_AVAILABLE and (self.two_stage_processor is not None)),
+            "two_stage_available": bool(TWO_STAGE_AVAILABLE and self.two_stage_processor),
             "two_stage_origin": TWO_STAGE_IMPORT_ORIGIN or "",
-            "two_stage_error": TWO_STAGE_IMPORT_ERROR[:2000] if TWO_STAGE_IMPORT_ERROR else "",
             "device": str(self.device_manager.get_optimal_device()),
             "core_processor_loaded": self.core_processor is not None,
             "config": self._safe_config_dict(),
             "memory_usage": self._safe_memory_usage(),
         }
+        
         try:
             status["sam2_loaded"] = self.model_loader.get_sam2() is not None
             status["matanyone_loaded"] = self.model_loader.get_matanyone() is not None
+            status["model_info"] = self.model_loader.get_model_info()
         except Exception:
             status["sam2_loaded"] = False
             status["matanyone_loaded"] = False
 
         if self.progress_tracker:
             status["progress"] = self.progress_tracker.get_all_progress()
+            
         return status
 
     def _safe_config_dict(self) -> Dict[str, Any]:
@@ -396,7 +449,8 @@ def _safe_config_dict(self) -> Dict[str, Any]:
             return self.config.to_dict()
         except Exception:
             keys = ["use_nvenc", "prefer_mp4", "video_codec", "audio_copy",
-                    "ffmpeg_path", "max_model_size", "max_model_size_bytes", "output_dir"]
+                    "ffmpeg_path", "max_model_size", "max_model_size_bytes", 
+                    "output_dir", "matanyone_enabled"]
             return {k: getattr(self.config, k, None) for k in keys}
 
     def _safe_memory_usage(self) -> Dict[str, Any]:
@@ -464,7 +518,14 @@ def main():
     try:
         logger.info("Starting BackgroundFX Pro")
         logger.info(f"Device: {processor.device_manager.get_optimal_device()}")
-        logger.info("Two-stage available (import): %s", TWO_STAGE_AVAILABLE)
+        logger.info(f"Two-stage available: {TWO_STAGE_AVAILABLE}")
+        
+        # Log model loader type
+        try:
+            from models.loaders.model_loader import ModelLoader
+            logger.info("Using split loader architecture")
+        except:
+            logger.info("Using legacy loader")
 
         from ui.ui_components import create_interface
         demo = create_interface()
@@ -479,4 +540,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()