Update core/app.py

core/app.py

@@ -3,58 +3,30 @@
 BackgroundFX Pro – Main Application Entry Point
 Refactored modular architecture – orchestrates specialised components
 """
-
 from __future__ import annotations
-
 # ── Early env/threading hygiene (safe default to silence libgomp) ────────────
 import os
 if not os.environ.get("OMP_NUM_THREADS", "").isdigit():
     os.environ["OMP_NUM_THREADS"] = "2"
-
 # If you use early_env in your project, keep this import (harmless if absent)
 try:
-    import early_env  # sets OMP/MKL/OPENBLAS + torch threads safely
+    import early_env # sets OMP/MKL/OPENBLAS + torch threads safely
 except Exception:
     pass
-
 import logging
 import threading
 import traceback
 import sys
 from pathlib import Path
-from typing import Optional, Tuple, Dict, Any, Callable
-
-# Mitigate CUDA fragmentation (must be set before importing torch)
-if "PYTORCH_CUDA_ALLOC_CONF" not in os.environ:
-    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True,max_split_size_mb:128"
-
-
-# ── Logging ──────────────────────────────────────────────────────────────────
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-)
-logger = logging.getLogger("core.app")
-
-# ── Ensure project root importable ───────────────────────────────────────────
-PROJECT_FILE = Path(__file__).resolve()
-CORE_DIR = PROJECT_FILE.parent
-ROOT = CORE_DIR.parent
-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)
-
+from utils.cv_processing import validate_video_file
 # ── Gradio schema patch (HF quirk) ───────────────────────────────────────────
 try:
     import gradio_client.utils as gc_utils
     _orig_get_type = gc_utils.get_type
     def _patched_get_type(schema):
         if not isinstance(schema, dict):
-            if isinstance(schema, bool):  return "boolean"
-            if isinstance(schema, str):   return "string"
+            if isinstance(schema, bool): return "boolean"
+            if isinstance(schema, str): return "string"
             if isinstance(schema, (int, float)): return "number"
             return "string"
         return _orig_get_type(schema)
@@ -62,13 +34,9 @@ def _patched_get_type(schema):
     logger.info("Gradio schema patch applied")
 except Exception as e:
     logger.warning(f"Gradio patch failed: {e}")
-
 # ── Core config + components ─────────────────────────────────────────────────
-from config.app_config import get_config
-from core.exceptions import ModelLoadingError, VideoProcessingError
 from utils.hardware.device_manager import DeviceManager
 from utils.system.memory_manager import MemoryManager
-
 # Try to import the new split loaders first, fall back to old if needed
 try:
     from models.loaders.model_loader import ModelLoader
@@ -76,77 +44,63 @@ def _patched_get_type(schema):
 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 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 ────────────────────────────────────────────────
 TWO_STAGE_AVAILABLE = False
 TWO_STAGE_IMPORT_ORIGIN = ""
 TWO_STAGE_IMPORT_ERROR = ""
 CHROMA_PRESETS: Dict[str, Dict[str, Any]] = {"standard": {}}
-TwoStageProcessor = None  # type: ignore
-
+TwoStageProcessor = None # type: ignore
 # Try multiple import paths for two-stage processor
 two_stage_paths = [
-    "processors.two_stage",  # Your fixed version
+    "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
+            from processors.two_stage import TwoStageProcessor, CHROMA_PRESETS # type: ignore
         else:
-            from processing.two_stage.two_stage_processor import TwoStageProcessor, CHROMA_PRESETS  # type: ignore
+            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:
         continue
-
 if not TWO_STAGE_AVAILABLE:
     logger.warning("Two-stage import FAILED from all paths")
-
 # ── Quiet startup self-check (async by default) ──────────────────────────────
 # Place the helper in tools/startup_selfcheck.py (with tools/__init__.py present)
 try:
     from tools.startup_selfcheck import schedule_startup_selfcheck
 except Exception:
-    schedule_startup_selfcheck = None  # graceful if the helper isn't shipped
-
+    schedule_startup_selfcheck = None # graceful if the helper isn't shipped
 # ╔══════════════════════════════════════════════════════════════════════════╗
-# ║                           VideoProcessor class                           ║
+# ║ VideoProcessor class ║
 # ╚══════════════════════════════════════════════════════════════════════════╝
 class VideoProcessor:
     """
     Main orchestrator – coordinates all specialised components.
     """
-
     def __init__(self):
         self.config = get_config()
-        self._patch_config_defaults(self.config)  # avoid AttributeError on older configs
-
+        self._patch_config_defaults(self.config) # avoid AttributeError on older configs
         self.device_manager = DeviceManager()
         self.memory_manager = MemoryManager(self.device_manager.get_optimal_device())
         self.model_loader = ModelLoader(self.device_manager, self.memory_manager)
-
         self.audio_processor = AudioProcessor()
         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: Optional[ProgressTracker] = None
-
         logger.info(f"VideoProcessor on device: {self.device_manager.get_optimal_device()}")
-
     # ── Config hardening: add missing fields safely ───────────────────────────
     @staticmethod
     def _patch_config_defaults(cfg: Any) -> None:
@@ -170,7 +124,6 @@ def _patch_config_defaults(cfg: Any) -> None:
             if not hasattr(cfg, k):
                 setattr(cfg, k, v)
         Path(cfg.output_dir).mkdir(parents=True, exist_ok=True)
-
     # ── Progress helper ───────────────────────────────────────────────────────
     def _init_progress(self, video_path: str, cb: Optional[Callable] = None):
         try:
@@ -184,31 +137,25 @@ def _init_progress(self, video_path: str, cb: Optional[Callable] = None):
         except Exception as e:
             logger.warning(f"Progress init failed: {e}")
             self.progress_tracker = ProgressTracker(100, cb)
-
     # ── Model loading ─────────────────────────────────────────────────────────
     def load_models(self, progress_callback: Optional[Callable] = None) -> str:
         with self.loading_lock:
             if self.models_loaded:
                 return "Models already loaded and validated"
-
             try:
                 self.cancel_event.clear()
                 if progress_callback:
                     progress_callback(0.0, f"Loading on {self.device_manager.get_optimal_device()}")
-
                 sam2_loaded, mat_loaded = self.model_loader.load_all_models(
                     progress_callback=progress_callback, cancel_event=self.cancel_event
                 )
                 if self.cancel_event.is_set():
                     return "Model loading cancelled"
-
                 # Get the actual models
                 sam2_predictor = sam2_loaded.model if sam2_loaded else None
-                mat_model = mat_loaded.model if mat_loaded else None  # NOTE: in our MatAnyone loader this is a stateful adapter (callable)
-
+                mat_model = mat_loaded.model if mat_loaded else None # NOTE: in our MatAnyone loader this is a stateful adapter (callable)
                 # Initialize core processor
                 self.core_processor = CoreVideoProcessor(config=self.config, models=self.model_loader)
-
                 # Initialize two-stage processor if available
                 self.two_stage_processor = None
                 if TWO_STAGE_AVAILABLE and TwoStageProcessor and (sam2_predictor or mat_model):
@@ -220,24 +167,22 @@ def load_models(self, progress_callback: Optional[Callable] = None) -> str:
                     except Exception as 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()
-                
+               
                 # 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
-
             except (AttributeError, ModelLoadingError) as e:
                 self.models_loaded = False
                 err = f"Model loading failed: {e}"
@@ -248,7 +193,6 @@ def load_models(self, progress_callback: Optional[Callable] = None) -> str:
                 err = f"Unexpected error during model loading: {e}"
                 logger.error(f"{err}\n{traceback.format_exc()}")
                 return err
-
     # ── Public entry – process video ─────────────────────────────────────────
     def process_video(
         self,
@@ -261,30 +205,25 @@ def process_video(
         key_color_mode: str = "auto",
         preview_mask: bool = False,
         preview_greenscreen: bool = False,
-    ) -> Tuple[Optional[str], str]:
+    ) -> Tuple[Optional[str], 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, None, "Models not loaded. Please click 'Load Models' first."
         if self.cancel_event.is_set():
-            return None, "Processing cancelled"
-
+            return None, None, "Processing cancelled"
         self._init_progress(video_path, progress_callback)
-
         ok, why = validate_video_file(video_path)
         if not ok:
-            return None, f"Invalid video: {why}"
-
+            return None, 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}")
-
             # IMPORTANT: start each video with a clean MatAnyone memory
             self._reset_matanyone_session()
-
             if use_two_stage:
                 if not TWO_STAGE_AVAILABLE or self.two_stage_processor is None:
-                    return None, "Two-stage processing not available"
+                    return None, None, "Two-stage processing not available"
                 return self._process_two_stage(
                     video_path,
                     background_choice,
@@ -302,14 +241,12 @@ def process_video(
                     preview_mask,
                     preview_greenscreen,
                 )
-
         except VideoProcessingError as e:
             logger.error(f"Processing failed: {e}")
-            return None, f"Processing failed: {e}"
+            return None, None, f"Processing failed: {e}"
         except Exception as e:
             logger.error(f"Unexpected processing error: {e}\n{traceback.format_exc()}")
-            return None, f"Unexpected error: {e}"
-
+            return None, None, f"Unexpected error: {e}"
     # ── Private – per-video MatAnyone reset ──────────────────────────────────
     def _reset_matanyone_session(self):
         """
@@ -326,7 +263,6 @@ def _reset_matanyone_session(self):
                 logger.info("MatAnyone session reset for new video")
             except Exception as e:
                 logger.warning(f"MatAnyone session reset failed (continuing): {e}")
-
     # ── Private – single-stage ───────────────────────────────────────────────
     def _process_single_stage(
         self,
@@ -336,13 +272,12 @@ def _process_single_stage(
         progress_callback: Optional[Callable],
         preview_mask: bool,
         preview_greenscreen: bool,
-    ) -> Tuple[Optional[str], str]:
+    ) -> Tuple[Optional[str], Optional[str], str]:
         import time
         ts = int(time.time())
         out_dir = Path(self.config.output_dir) / "single_stage"
         out_dir.mkdir(parents=True, exist_ok=True)
         out_path = str(out_dir / f"processed_{ts}.mp4")
-
         # Process video via your CoreVideoProcessor
         result = self.core_processor.process_video(
             input_path=video_path,
@@ -353,10 +288,9 @@ def _process_single_stage(
             },
             progress_callback=progress_callback,
         )
-        
+       
         if not result:
-            return None, "Video processing failed"
-
+            return None, None, "Video processing failed"
         # Mux audio unless preview-only
         if not (preview_mask or preview_greenscreen):
             try:
@@ -368,7 +302,6 @@ def _process_single_stage(
                 final_path = out_path
         else:
             final_path = out_path
-
         # Build status message
         try:
             mat_loaded = bool(self.model_loader.get_matanyone())
@@ -383,8 +316,7 @@ def _process_single_stage(
             f"MatAnyone: {matanyone_status}\n"
             f"Device: {self.device_manager.get_optimal_device()}"
         )
-        return final_path, msg
-
+        return final_path, None, msg  # No green in single-stage
     # ── Private – two-stage ─────────────────────────────────────────────────
     def _process_two_stage(
         self,
@@ -394,18 +326,16 @@ def _process_two_stage(
         progress_callback: Optional[Callable],
         chroma_preset: str,
         key_color_mode: str,
-    ) -> Tuple[Optional[str], str]:
+    ) -> Tuple[Optional[str], Optional[str], str]:
         if self.two_stage_processor is None:
-            return None, "Two-stage processor not available"
-
+            return None, None, "Two-stage processor not available"
         import cv2, time
         cap = cv2.VideoCapture(video_path)
         if not cap.isOpened():
-            return None, "Could not open input video"
+            return None, None, "Could not open input video"
         w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) or 1280
         h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) or 720
         cap.release()
-
         # Prepare background
         try:
             background = self.core_processor.prepare_background(
@@ -413,20 +343,17 @@ def _process_two_stage(
             )
         except Exception as e:
             logger.error(f"Background preparation failed: {e}")
-            return None, f"Failed to prepare background: {e}"
+            return None, None, f"Failed to prepare background: {e}"
         if background is None:
-            return None, "Failed to prepare background"
-
+            return None, None, "Failed to prepare background"
         ts = int(time.time())
         out_dir = Path(self.config.output_dir) / "two_stage"
         out_dir.mkdir(parents=True, exist_ok=True)
         final_out = str(out_dir / f"final_{ts}.mp4")
-
         chroma_cfg = CHROMA_PRESETS.get(chroma_preset, CHROMA_PRESETS.get("standard", {}))
         logger.info(f"Two-stage with preset: {chroma_preset} | key_color: {key_color_mode}")
-
         # (Per-video reset already called in process_video)
-        result, message = self.two_stage_processor.process_full_pipeline(
+        final_path, green_path, stage2_msg = self.two_stage_processor.process_full_pipeline(
             video_path,
             background,
             final_out,
@@ -434,18 +361,16 @@ def _process_two_stage(
             chroma_settings=chroma_cfg,
             progress_callback=progress_callback,
         )
-        if result is None:
-            return None, message
-
+        if final_path is None:
+            return None, None, stage2_msg
         # Mux audio
         try:
-            final_path = self.audio_processor.add_audio_to_video(
-                original_video=video_path, processed_video=result
+            final_with_audio = self.audio_processor.add_audio_to_video(
+                original_video=video_path, processed_video=final_path
             )
         except Exception as e:
             logger.warning(f"Audio mux failed: {e}")
-            final_path = result
-
+            final_with_audio = final_path
         try:
             mat_loaded = bool(self.model_loader.get_matanyone())
         except Exception:
@@ -458,8 +383,7 @@ def _process_two_stage(
             f"MatAnyone: {matanyone_status}\n"
             f"Device: {self.device_manager.get_optimal_device()}"
         )
-        return final_path, msg
-
+        return final_with_audio, green_path, msg
     # ── Status helpers ───────────────────────────────────────────────────────
     def get_status(self) -> Dict[str, Any]:
         status = {
@@ -471,7 +395,7 @@ def get_status(self) -> Dict[str, Any]:
             "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
@@ -479,31 +403,26 @@ def get_status(self) -> Dict[str, Any]:
         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]:
         try:
             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", 
+                    "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]:
         try:
             return self.memory_manager.get_memory_usage()
         except Exception:
             return {}
-
     def cancel_processing(self):
         self.cancel_event.set()
         logger.info("Cancellation requested")
-
     def cleanup_resources(self):
         try:
             self.memory_manager.cleanup_aggressive()
@@ -514,14 +433,10 @@ def cleanup_resources(self):
         except Exception:
             pass
         logger.info("Resources cleaned up")
-
-
 # ── Singleton + thin wrappers (used by UI callbacks) ────────────────────────
 processor = VideoProcessor()
-
 def load_models_with_validation(progress_callback: Optional[Callable] = None) -> str:
     return processor.load_models(progress_callback)
-
 def process_video_fixed(
     video_path: str,
     background_choice: str,
@@ -532,7 +447,7 @@ def process_video_fixed(
     key_color_mode: str = "auto",
     preview_mask: bool = False,
     preview_greenscreen: bool = False,
-) -> Tuple[Optional[str], str]:
+) -> Tuple[Optional[str], Optional[str], str]:
     return processor.process_video(
         video_path,
         background_choice,
@@ -544,37 +459,29 @@ def process_video_fixed(
         preview_mask,
         preview_greenscreen,
     )
-
 def get_model_status() -> Dict[str, Any]:
     return processor.get_status()
-
 def get_cache_status() -> Dict[str, Any]:
     return processor.get_status()
-
 PROCESS_CANCELLED = processor.cancel_event
-
-
 # ── CLI entrypoint (must exist; app.py imports main) ─────────────────────────
 def main():
     try:
         logger.info("Starting BackgroundFX Pro")
         logger.info(f"Device: {processor.device_manager.get_optimal_device()}")
         logger.info(f"Two-stage available: {TWO_STAGE_AVAILABLE}")
-
         # 🔹 Quiet model self-check (defaults to async; set SELF_CHECK_MODE=sync to block)
         if schedule_startup_selfcheck is not None:
             try:
                 schedule_startup_selfcheck(mode=os.getenv("SELF_CHECK_MODE", "async"))
             except Exception as e:
                 logger.error(f"Startup self-check skipped: {e}", exc_info=True)
-
         # Log model loader type
         try:
             from models.loaders.model_loader import ModelLoader
             logger.info("Using split loader architecture")
         except Exception:
             logger.info("Using legacy loader")
-
         from ui.ui_components import create_interface
         demo = create_interface()
         demo.queue().launch(
@@ -585,7 +492,5 @@ def main():
         )
     finally:
         processor.cleanup_resources()
-
-
 if __name__ == "__main__":
-    main()
+    main()