Update ui/ui_components.py

ui/ui_components.py

@@ -1,398 +1,597 @@
 #!/usr/bin/env python3
 """
-UI Components for BackgroundFX Pro (forced Two-Stage)
------------------------------------------------------
-* Pure layout (tiny wrapper to set env for quality)
-* All heavy logic stays in ui/callbacks.py
-* Two-stage mode is always active (checkbox removed)
+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
-import gradio as gr
+os.environ["OMP_NUM_THREADS"] = "2"  # Force valid value early
+# 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
+except Exception:
+    pass
 import logging
-
-from ui.callbacks import (
-    cb_load_models,
-    cb_process_video,
-    cb_cancel,
-    cb_status,
-    cb_clear,
-    cb_generate_bg,
-    cb_use_gen_bg,
-    cb_preset_bg_preview,
+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(__name__)
-
-# Typography & UI polish: sharper text + cleaner cards
-CSS = """
-:root {
-  --radius: 16px;
-  --font-sans: 'Inter', system-ui, -apple-system, 'Segoe UI', Roboto,
-               'Helvetica Neue', Arial, sans-serif;
-}
-
-/* Global crisp text */
-html, body, .gradio-container, .gradio-container * {
-  font-family: var(--font-sans) !important;
-  -webkit-font-smoothing: antialiased !important;
-  -moz-osx-font-smoothing: grayscale !important;
-  text-rendering: optimizeLegibility !important;
-  font-synthesis-weight: none;
-}
-
-/* Headings tighter & bolder */
-.gradio-container h1, .gradio-container h2, .gradio-container h3 {
-  letter-spacing: -0.01em;
-  font-weight: 700;
-}
-
-/* Body copy slightly tighter */
-#hero .prose, .gr-markdown, .gr-text {
-  letter-spacing: -0.003em;
-}
-
-/* Card look */
-.card {
-  border-radius: var(--radius);
-  border: 1px solid rgba(0,0,0,.08);
-  padding: 16px;
-  background: linear-gradient(180deg, rgba(255,255,255,.94), rgba(248,250,252,.94));
-  box-shadow: 0 10px 30px rgba(0,0,0,.06);
-}
-
-.footer-note { opacity: 0.7; font-size: 12px; }
-.sm { font-size: 13px; opacity: 0.85; }
-#statusbox { min-height: 120px; }
-.preview-img { border-radius: var(--radius); border: 1px solid rgba(0,0,0,.08); }
-
-/* Buttons get a tiny weight bump for clarity */
-button, .gr-button { font-weight: 600; }
-
-/* Inline Quality select between buttons */
-.inline-quality .wrap-inner { min-width: 170px; }
-"""
-
-# Keep in sync with utils/cv_processing.PROFESSIONAL_BACKGROUNDS
-_BG_CHOICES = [
-    "minimalist",
-    "office_modern",
-    "studio_blue",
-    "studio_green",
-    "warm_gradient",
-    "tech_dark",
+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)
+# ── 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, (int, float)): return "number"
+            return "string"
+        return _orig_get_type(schema)
+    gc_utils.get_type = _patched_get_type
+    logger.info("Gradio schema patch applied")
+except Exception as e:
+    logger.warning(f"Gradio patch failed: {e}")
+# ── Core config + components ─────────────────────────────────────────────────
+try:
+    from config.app_config import get_config
+except ImportError:
+    # Dummy if missing
+    class DummyConfig:
+        def to_dict(self):
+            return {}
+    get_config = lambda: DummyConfig()
+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
+    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 ────────────────────────────────────────────────
+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 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",
 ]
-PRO_IMAGE_CHOICES = ["minimalist", "office_modern", "studio_blue", "studio_green"]
-GRADIENT_COLOR_CHOICES = ["warm_gradient", "tech_dark"]
-
-
-def get_background_status(bg_method, bg_style, custom_bg_path=None, pro_image=None, gradient=None):
-    """Get current background status text for display"""
-    if bg_method == "Upload file" and custom_bg_path:
-        filename = os.path.basename(custom_bg_path) if custom_bg_path else "none"
-        return f"Background: custom({filename})"
-    elif bg_method == "Pre-loaded professional images" and pro_image:
-        return f"Background: {pro_image}"
-    elif bg_method == "Pre-loaded Gradients / Colors" and gradient:
-        return f"Background: {gradient}"
-    elif bg_method == "AI generated background":
-        return f"Background: AI-generated"
-    elif bg_style:
-        return f"Background: {bg_style}"
-    else:
-        return "Background: none"
-
-
-def on_quality_change(quality_value):
-    """Called when quality dropdown changes - sets environment variable"""
-    os.environ["BFX_QUALITY"] = quality_value.lower()
-    logger.info(f"UI: Quality preference set to {quality_value}")
-    return f"Quality mode: {quality_value} (refine_every={'1' if quality_value=='max' else '3' if quality_value=='balanced' else '10'} frames)"
-
-
-def create_interface() -> gr.Blocks:
-    with gr.Blocks(
-        title="🎬 BackgroundFX Pro",
-        css=CSS,
-        analytics_enabled=False,
-        theme=gr.themes.Soft()
-    ) as demo:
-
-        # ------------------------------------------------------------------
-        # HERO
-        # ------------------------------------------------------------------
-        with gr.Row(elem_id="hero"):
-            gr.Markdown(
-                "## 🎬 BackgroundFX Pro (CSP-Safe)\n"
-                "Replace your video background with cinema-quality AI matting. "
-                "Built for Hugging Face Spaces CSP.\n\n"
-                "_Tip: press **Load Models** once after the Space spins up._"
-            )
-
-        # ------------------------------------------------------------------
-        # TAB – Quick Start
-        # ------------------------------------------------------------------
-        with gr.Tab("🏁 Quick Start"):
-
-            with gr.Row():
-                # ── Left column ────────────────────────────────────────────
-                with gr.Column(scale=1):
-                    video = gr.Video(label="Upload Video", interactive=True)
-
-                    # Hidden: effective preset key (still used by callbacks / defaults)
-                    bg_style = gr.Dropdown(
-                        label="Background Style (hidden)",
-                        choices=_BG_CHOICES,
-                        value="minimalist",
-                        visible=False,
-                    )
-
-                    # =======================
-                    # Background Source block
-                    # =======================
-                    gr.Markdown("### 🖼️ Background Source")
-
-                    bg_method = gr.Radio(
-                        label="Choose method",
-                        choices=[
-                            "Upload file",
-                            "Pre-loaded professional images",
-                            "Pre-loaded Gradients / Colors",
-                            "AI generated background",
-                        ],
-                        value="Upload file",
-                    )
-
-                    # a) Upload file option
-                    with gr.Group(visible=True) as grp_upload:
-                        custom_bg = gr.Image(
-                            label="Upload Background Image",
-                            interactive=True,
-                            type="filepath",  # returns file path
-                            elem_classes=["preview-img"]
-                        )
-
-                    # b) Pre-loaded professional images
-                    with gr.Group(visible=False) as grp_pro_images:
-                        pro_image_dd = gr.Dropdown(
-                            label="Professional Images",
-                            choices=PRO_IMAGE_CHOICES,
-                            value=PRO_IMAGE_CHOICES[0],
-                            info="Pre-defined photo-like backgrounds",
-                        )
-                        gr.Markdown(
-                            "<span class='sm'>Selecting a preset updates the preview below.</span>"
-                        )
-
-                    # c) Pre-loaded gradients & full colors
-                    with gr.Group(visible=False) as grp_gradients:
-                        gradient_dd = gr.Dropdown(
-                            label="Gradients & Full Colors",
-                            choices=GRADIENT_COLOR_CHOICES,
-                            value=GRADIENT_COLOR_CHOICES[0],
-                            info="Clean gradients or solid color styles",
-                        )
-                        gr.Markdown(
-                            "<span class='sm'>Selecting a preset updates the preview below.</span>"
-                        )
-
-                    # d) AI-generated background (inline, lightweight)
-                    with gr.Group(visible=False) as grp_ai:
-                        prompt = gr.Textbox(
-                            label="Describe vibe",
-                            value="modern office",
-                            info="e.g. 'soft sunset studio', 'cool tech dark', 'forest ambience'"
-                        )
-                        with gr.Row():
-                            gen_width  = gr.Slider(640, 1920, 1280, step=10, label="Width")
-                            gen_height = gr.Slider(360, 1080, 720,  step=10, label="Height")
-                        with gr.Row():
-                            bokeh    = gr.Slider(0, 30, 8,    step=1,    label="Bokeh Blur")
-                            vignette = gr.Slider(0, 0.6, 0.15, step=0.01, label="Vignette")
-                            contrast = gr.Slider(0.8, 1.4, 1.05, step=0.01, label="Contrast")
-                        with gr.Row():
-                            btn_gen_bg_inline = gr.Button("✨ Generate Background", variant="primary")
-                            use_gen_as_custom_inline = gr.Button("📌 Use as Custom Background", variant="secondary")
-                        gen_preview = gr.Image(
-                            label="Generated Background",
-                            interactive=False,
-                            elem_classes=["preview-img"]
-                        )
-                        gen_path = gr.Textbox(label="Saved Path", interactive=False)
-
-                    # ── Advanced options accordion ───────────────────────
-                    with gr.Accordion("Advanced", open=False):
-                        chroma_preset = gr.Dropdown(
-                            label="Chroma Preset",
-                            choices=["standard"],    # can add 'studio', 'outdoor' later
-                            value="standard"
-                        )
-                        key_color_mode = gr.Dropdown(
-                            label="Key-Colour Mode",
-                            choices=["auto", "green", "blue", "cyan", "magenta"],
-                            value="auto",
-                            info="Auto picks a colour far from your clothes; override if needed."
-                        )
-                        preview_mask = gr.Checkbox(
-                            label="Preview Mask only (mute audio)",
-                            value=False
+for import_path in two_stage_paths:
+    try:
+        exec(f"from {import_path} import TwoStageProcessor, CHROMA_PRESETS")
+        TWO_STAGE_AVAILABLE = True
+        TWO_STAGE_IMPORT_ORIGIN = import_path
+        logger.info(f"Two-stage import OK ({import_path})")
+        break
+    except Exception as e:
+        TWO_STAGE_IMPORT_ERROR = str(e)
+        continue
+if not TWO_STAGE_AVAILABLE:
+    logger.warning(f"Two-stage import FAILED from all paths: {TWO_STAGE_IMPORT_ERROR}")
+# ── 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
+# Dummy exceptions if core.exceptions not available
+class ModelLoadingError(Exception):
+    pass
+class VideoProcessingError(Exception):
+    pass
+# ╔══════════════════════════════════════════════════════════════════════════╗
+# ║ 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.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:
+        defaults = {
+            # video / i/o
+            "use_nvenc": False,
+            "prefer_mp4": True,
+            "video_codec": "mp4v",
+            "audio_copy": True,
+            "ffmpeg_path": "ffmpeg",
+            # model/resource guards
+            "max_model_size": 0,
+            "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):
+                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:
+            import cv2
+            cap = cv2.VideoCapture(video_path)
+            total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+            cap.release()
+            if total <= 0:
+                total = 100
+            self.progress_tracker = ProgressTracker(total, cb)
+        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)
+                # 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):
+                    try:
+                        self.two_stage_processor = TwoStageProcessor(
+                            sam2_predictor=sam2_predictor, matanyone_model=mat_model
                         )
-                        preview_greenscreen = gr.Checkbox(
-                            label="Preview Green-screen only (mute audio)",
-                            value=False
-                        )
-
-                    # ── Controls row: Load → Quality → Process / Cancel ──
-                    with gr.Row():
-                        btn_load = gr.Button("🔄 Load Models", variant="secondary")
-
-                        quality = gr.Dropdown(
-                            label="Quality",
-                            choices=["speed", "balanced", "max"],
-                            value=os.getenv("BFX_QUALITY", "balanced"),
-                            info="Speed = fastest; Max = best edges & spill control.",
-                            elem_classes=["inline-quality"],
-                        )
-
-                        btn_run = gr.Button("🎬 Process Video", variant="primary")
-                        btn_cancel = gr.Button("⏹️ Cancel", variant="secondary")
-
-                # ── Right column ──────────────────────────────────────────
-                with gr.Column(scale=1):
-                    out_video = gr.Video(label="Processed Output", interactive=False)
-                    statusbox = gr.Textbox(label="Status", lines=8, elem_id="statusbox")
-                    quality_status = gr.Textbox(label="Quality Settings", lines=1, value="Quality mode: balanced")
-                    with gr.Row():
-                        btn_refresh = gr.Button("🔍 Refresh Status", variant="secondary")
-                        btn_clear   = gr.Button("🧹 Clear",           variant="secondary")
-
-        # ------------------------------------------------------------------
-        # TAB – Status & settings
-        # ------------------------------------------------------------------
-        with gr.Tab("📈 Status & Settings"):
-            with gr.Row():
-                with gr.Column(scale=1, elem_classes=["card"]):
-                    model_status = gr.JSON(label="Model Status")
-                with gr.Column(scale=1, elem_classes=["card"]):
-                    cache_status = gr.JSON(label="Cache / System Status")
-
-            gr.Markdown(
-                "<div class='footer-note'>If models fail to load, fallbacks keep the UI responsive. "
-                "Check the runtime log for details.</div>"
-            )
-
-        # ------------------------------------------------------------------
-        # Callback wiring
-        # ------------------------------------------------------------------
-
-        # Toggle which background sub-section is visible
-        def _toggle_bg_sections(choice: str):
-            return (
-                gr.update(visible=(choice == "Upload file")),
-                gr.update(visible=(choice == "Pre-loaded professional images")),
-                gr.update(visible=(choice == "Pre-loaded Gradients / Colors")),
-                gr.update(visible=(choice == "AI generated background")),
-            )
-
-        bg_method.change(
-            _toggle_bg_sections,
-            inputs=[bg_method],
-            outputs=[grp_upload, grp_pro_images, grp_gradients, grp_ai],
+                        logger.info("Two-stage processor initialised")
+                    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}"
+                logger.error(err)
+                return err
+            except Exception as e:
+                self.models_loaded = False
+                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,
+        video_path: str,
+        background_choice: str,
+        custom_background_path: Optional[str] = None,
+        progress_callback: Optional[Callable] = None,
+        use_two_stage: bool = False,
+        chroma_preset: str = "standard",
+        key_color_mode: str = "auto",
+        preview_mask: bool = False,
+        preview_greenscreen: bool = False,
+    ) -> Tuple[Optional[str], Optional[str], str]:
+        
+        # ===== BACKGROUND PATH DEBUG & FIX =====
+        logger.info("=" * 60)
+        logger.info("BACKGROUND PATH DEBUGGING")
+        logger.info(f"background_choice: {background_choice}")
+        logger.info(f"custom_background_path type: {type(custom_background_path)}")
+        logger.info(f"custom_background_path value: {custom_background_path}")
+        
+        # Fix 1: Handle if Gradio sends a dict
+        if isinstance(custom_background_path, dict):
+            original = custom_background_path
+            custom_background_path = custom_background_path.get('name') or custom_background_path.get('path')
+            logger.info(f"Extracted path from dict: {original} -> {custom_background_path}")
+        
+        # Fix 2: Handle PIL Image objects
+        try:
+            from PIL import Image
+            if isinstance(custom_background_path, Image.Image):
+                import tempfile
+                with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
+                    custom_background_path.save(tmp.name)
+                    custom_background_path = tmp.name
+                    logger.info(f"Saved PIL Image to: {custom_background_path}")
+        except ImportError:
+            pass
+        
+        # Fix 3: Verify file exists when using custom background
+        if background_choice == "custom" or custom_background_path:
+            if custom_background_path:
+                if Path(custom_background_path).exists():
+                    logger.info(f"✅ Background file exists: {custom_background_path}")
+                else:
+                    logger.warning(f"⚠️ Background file does not exist: {custom_background_path}")
+                    # Try to find it in Gradio temp directories
+                    import glob
+                    patterns = [
+                        "/tmp/gradio*/**/*.jpg",
+                        "/tmp/gradio*/**/*.jpeg", 
+                        "/tmp/gradio*/**/*.png",
+                        "/tmp/**/*.jpg",
+                        "/tmp/**/*.jpeg",
+                        "/tmp/**/*.png",
+                    ]
+                    for pattern in patterns:
+                        files = glob.glob(pattern, recursive=True)
+                        if files:
+                            # Get the most recent file
+                            newest = max(files, key=os.path.getmtime)
+                            logger.info(f"Found potential background: {newest}")
+                            # Only use it if it was created in the last 5 minutes
+                            if (time.time() - os.path.getmtime(newest)) < 300:
+                                custom_background_path = newest
+                                logger.info(f"✅ Using recent temp file: {custom_background_path}")
+                                break
+            else:
+                logger.error("❌ Custom background mode but path is None!")
+        
+        logger.info(f"Final custom_background_path: {custom_background_path}")
+        logger.info("=" * 60)
+        
+        if not self.models_loaded or not self.core_processor:
+            return None, None, "Models not loaded. Please click 'Load Models' first."
+        if self.cancel_event.is_set():
+            return None, None, "Processing cancelled"
+        self._init_progress(video_path, progress_callback)
+        ok, why = validate_video_file(video_path)
+        if not ok:
+            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, None, "Two-stage processing not available"
+                final, green, msg = self._process_two_stage(
+                    video_path,
+                    background_choice,
+                    custom_background_path,
+                    progress_callback,
+                    chroma_preset,
+                    key_color_mode,
+                )
+                return final, green, msg
+            else:
+                final, green, msg = self._process_single_stage(
+                    video_path,
+                    background_choice,
+                    custom_background_path,
+                    progress_callback,
+                    preview_mask,
+                    preview_greenscreen,
+                )
+                return final, green, msg
+        except VideoProcessingError as e:
+            logger.error(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, None, f"Unexpected error: {e}"
+    # ── Private – per-video MatAnyone reset ──────────────────────────────────
+    def _reset_matanyone_session(self):
+        """
+        Ensure a fresh MatAnyone memory per video. The MatAnyone loader we use returns a
+        callable *stateful adapter*. If present, reset() clears its InferenceCore memory.
+        """
+        try:
+            mat = self.model_loader.get_matanyone()
+        except Exception:
+            mat = None
+        if mat is not None and hasattr(mat, "reset") and callable(mat.reset):
+            try:
+                mat.reset()
+                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,
+        video_path: str,
+        background_choice: str,
+        custom_background_path: Optional[str],
+        progress_callback: Optional[Callable],
+        preview_mask: bool,
+        preview_greenscreen: bool,
+    ) -> Tuple[Optional[str], Optional[str], str]:
+        import time
+        
+        # Additional debug logging for single-stage
+        logger.info(f"[Single-stage] background_choice: {background_choice}")
+        logger.info(f"[Single-stage] custom_background_path: {custom_background_path}")
+        
+        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,
+            output_path=out_path,
+            bg_config={
+                "background_choice": background_choice,
+                "custom_path": custom_background_path,
+            },
+            progress_callback=progress_callback,
         )
-
-        # Load models
-        btn_load.click(cb_load_models, outputs=statusbox)
-
-        # Quality dropdown handler - updates environment and status
-        quality.change(
-            on_quality_change,
-            inputs=[quality],
-            outputs=[quality_status]
+       
+        if not result:
+            return None, None, "Video processing failed"
+        # Mux audio unless preview-only
+        if not (preview_mask or preview_greenscreen):
+            try:
+                final_path = self.audio_processor.add_audio_to_video(
+                    original_video=video_path, processed_video=out_path
+                )
+            except Exception as e:
+                logger.warning(f"Audio mux failed, returning video without audio: {e}")
+                final_path = out_path
+        else:
+            final_path = out_path
+        # Build status message
+        try:
+            mat_loaded = bool(self.model_loader.get_matanyone())
+        except Exception:
+            mat_loaded = False
+        matanyone_status = "✓" if mat_loaded 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()}"
         )
-
-        # Tiny wrapper to set env for quality before calling the existing callback
-        def _run_with_quality(video_pth, bg_style_val, custom_bg_pth,
-                              use_two_stage_state, chroma_p, key_mode,
-                              prev_mask, prev_gs, quality_val, bg_method_val, 
-                              pro_image_val, gradient_val):
-            os.environ["BFX_QUALITY"] = (quality_val or "balanced")
-            
-            # Update status with background info
-            bg_status = get_background_status(
-                bg_method_val, bg_style_val, custom_bg_pth, 
-                pro_image_val, gradient_val
-            )
-            
-            result = cb_process_video(
-                video_pth, bg_style_val, custom_bg_pth,
-                use_two_stage_state, chroma_p, key_mode, prev_mask, prev_gs
+        return final_path, None, msg  # No green in single-stage
+    # ── Private – two-stage ─────────────────────────────────────────────────
+    def _process_two_stage(
+        self,
+        video_path: str,
+        background_choice: str,
+        custom_background_path: Optional[str],
+        progress_callback: Optional[Callable],
+        chroma_preset: str,
+        key_color_mode: str,
+    ) -> Tuple[Optional[str], Optional[str], str]:
+        if self.two_stage_processor is None:
+            return None, None, "Two-stage processor not available"
+        
+        # Additional debug logging for two-stage
+        logger.info(f"[Two-stage] background_choice: {background_choice}")
+        logger.info(f"[Two-stage] custom_background_path: {custom_background_path}")
+        
+        import cv2, time
+        cap = cv2.VideoCapture(video_path)
+        if not cap.isOpened():
+            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(
+                background_choice, custom_background_path, w, h
             )
-            
-            # Append background status to the result message
-            if isinstance(result, tuple) and len(result) >= 2:
-                video_out, msg = result[0], result[1]
-                return video_out, f"{msg}\n{bg_status}"
-            return result
-
-        # Always two-stage: pass use_two_stage=True to callback via State
-        btn_run.click(
-            _run_with_quality,
-            inputs=[
-                video,
-                bg_style,
-                custom_bg,
-                gr.State(value=True),  # Always two-stage
-                chroma_preset,
-                key_color_mode,
-                preview_mask,
-                preview_greenscreen,
-                quality,               # <-- the new control
-                bg_method,            # Added for status
-                pro_image_dd,         # Added for status
-                gradient_dd,          # Added for status
-            ],
-            outputs=[out_video, statusbox],
+        except Exception as e:
+            logger.error(f"Background preparation failed: {e}")
+            return None, None, f"Failed to prepare background: {e}"
+        if background is None:
+            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)
+        final_path, green_path, stage2_msg = self.two_stage_processor.process_full_pipeline(
+            video_path,
+            background,
+            final_out,
+            key_color_mode=key_color_mode,
+            chroma_settings=chroma_cfg,
+            progress_callback=progress_callback,
         )
-
-        # Cancel / Status / Clear
-        btn_cancel.click(cb_cancel, outputs=statusbox)
-        btn_refresh.click(cb_status, outputs=[model_status, cache_status])
-
-        btn_clear.click(
-            cb_clear,
-            outputs=[out_video, statusbox, gen_preview, gen_path, custom_bg]
-        )
-
-        # Preloaded presets → update preview (write into custom_bg)
-        pro_image_dd.change(
-            cb_preset_bg_preview,
-            inputs=[pro_image_dd],
-            outputs=[custom_bg],
-        )
-        gradient_dd.change(
-            cb_preset_bg_preview,
-            inputs=[gradient_dd],
-            outputs=[custom_bg],
-        )
-
-        # AI background generation (inline)
-        btn_gen_bg_inline.click(
-            cb_generate_bg,
-            inputs=[prompt, gen_width, gen_height, bokeh, vignette, contrast],
-            outputs=[gen_preview, gen_path],
-        )
-        use_gen_as_custom_inline.click(
-            cb_use_gen_bg,
-            inputs=[gen_path],
-            outputs=[custom_bg],
+        if final_path is None:
+            return None, None, stage2_msg
+        # Mux audio
+        try:
+            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_with_audio = final_path
+        try:
+            mat_loaded = bool(self.model_loader.get_matanyone())
+        except Exception:
+            mat_loaded = False
+        matanyone_status = "✓" if mat_loaded 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()}"
         )
-
-        # Initialize with a default preset preview on load
-        demo.load(
-            cb_preset_bg_preview,
-            inputs=[bg_style],
-            outputs=[custom_bg]
+        return final_with_audio, green_path, msg
+    # ── Status helpers ───────────────────────────────────────────────────────
+    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),
+            "two_stage_origin": TWO_STAGE_IMPORT_ORIGIN or "",
+            "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]:
+        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",
+                    "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()
+        except Exception:
+            pass
+        try:
+            self.model_loader.cleanup()
+        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,
+    custom_background_path: Optional[str],
+    progress_callback: Optional[Callable] = None,
+    use_two_stage: bool = False,
+    chroma_preset: str = "standard",
+    key_color_mode: str = "auto",
+    preview_mask: bool = False,
+    preview_greenscreen: bool = False,
+) -> Tuple[Optional[str], Optional[str], str]:
+    return processor.process_video(
+        video_path,
+        background_choice,
+        custom_background_path,
+        progress_callback,
+        use_two_stage,
+        chroma_preset,
+        key_color_mode,
+        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(
+            server_name="0.0.0.0",
+            server_port=7860,
+            show_error=True,
+            debug=False,
         )
-
-    return demo
+    finally:
+        processor.cleanup_resources()
+if __name__ == "__main__":
+    main()