Update app.py

app.py

@@ -1,10 +1,8 @@
 #!/usr/bin/env python3
 # ========================= PRE-IMPORT ENV GUARDS =========================
-# Must run BEFORE importing numpy/cv2/torch to avoid libgomp errors.
 import os
-
-# Remove invalid OMP setting or tame thread counts
-os.environ.pop("OMP_NUM_THREADS", None)           # or set to e.g. "1"
+# Remove invalid OMP setting or tame thread counts BEFORE importing numpy/cv2/torch
+os.environ.pop("OMP_NUM_THREADS", None)           # or set "1"
 os.environ.setdefault("MKL_NUM_THREADS", "1")
 os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
 os.environ.setdefault("NUMEXPR_NUM_THREADS", "1")
@@ -16,7 +14,7 @@
 """
 High-Quality Video Background Replacement - MAIN APPLICATION
 Upload video → Choose professional background → Replace with cinema quality
-Features: SAM2 + MatAnyone with multi-fallback loading, professional backgrounds,
+Features: SAM2 + MatAnyOne with multi-fallback loading, professional backgrounds,
 cinema-quality processing, lazy loading, and enhanced stability
 """
 
@@ -39,16 +37,13 @@
 from typing import Optional, Tuple, Dict, Any
 import logging
 import warnings
+import subprocess
+import importlib
 
-# Import all utilities (must provide: PROFESSIONAL_BACKGROUNDS,
-# create_professional_background, create_procedural_background,
-# segment_person_hq, refine_mask_hq, replace_background_hq, get_model_status)
-from utilities import *
+# Import your utilities
+from utilities import *  # must provide required helpers & PROFESSIONAL_BACKGROUNDS
 
-# Suppress warnings
 warnings.filterwarnings("ignore")
-
-# Setup logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -58,14 +53,12 @@
 try:
     import gradio_client.utils as gc_utils
     original_get_type = gc_utils.get_type
-
     def patched_get_type(schema):
         if not isinstance(schema, dict):
             if isinstance(schema, bool):
                 return "boolean"
             return "string"
         return original_get_type(schema)
-
     gc_utils.get_type = patched_get_type
     logger.info("✅ Applied Gradio schema validation monkey patch.")
 except (ImportError, AttributeError) as e:
@@ -75,10 +68,7 @@ def patched_get_type(schema):
 # SAM2 LOADER (Hydra search path; pass STRING config name to build_sam2)
 # ============================================================================ #
 def load_sam2_predictor(device: str = "cuda", progress: Optional[gr.Progress] = None):
-    """
-    Loads the SAM2 model and returns a SAM2ImagePredictor instance.
-    Uses Hydra only to set the config search path; passes STRING config name to build_sam2.
-    """
+    """Loads SAM2 and returns SAM2ImagePredictor. Uses STRING config name for build_sam2."""
     import hydra
 
     sam_logger = logging.getLogger("SAM2Loader")
@@ -86,20 +76,12 @@ def load_sam2_predictor(device: str = "cuda", progress: Optional[gr.Progress] =
     sam_logger.info(f"Looking for SAM2 configs in absolute path: {configs_dir}")
 
     if not os.path.isdir(configs_dir):
-        sam_logger.error(
-            f"FATAL: Configs directory not found at '{configs_dir}'. "
-            f"Please ensure the 'Configs' folder exists at repository root."
-        )
-        raise gr.Error("FATAL: SAM2 Configs directory not found.")
-
-    tried = []
+        raise gr.Error(f"FATAL: SAM2 Configs directory not found at '{configs_dir}'")
 
     def _maybe_progress(pct: float, desc: str):
         if progress is not None:
-            try:
-                progress(pct, desc=desc)
-            except Exception:
-                pass
+            try: progress(pct, desc=desc)
+            except Exception: pass
 
     def try_load(config_name_with_yaml: str, checkpoint_name: str):
         try:
@@ -110,7 +92,7 @@ def try_load(config_name_with_yaml: str, checkpoint_name: str):
                 sam_logger.info(f"Downloading {checkpoint_name} from Hugging Face Hub...")
                 _maybe_progress(0.1, f"Downloading {checkpoint_name}...")
                 from huggingface_hub import hf_hub_download
-                repo = f"facebook/{config_name_with_yaml.replace('.yaml','')}"  # e.g. facebook/sam2_hiera_large
+                repo = f"facebook/{config_name_with_yaml.replace('.yaml','')}"
                 checkpoint_path = hf_hub_download(
                     repo_id=repo,
                     filename=checkpoint_name,
@@ -128,7 +110,7 @@ def try_load(config_name_with_yaml: str, checkpoint_name: str):
                 job_name=f"sam2_load_{int(time.time())}"
             )
 
-            # IMPORTANT: pass STRING config name (no ".yaml") to build_sam2
+            # Pass STRING config name to build_sam2
             config_name = config_name_with_yaml.replace(".yaml", "")
 
             from sam2.build_sam import build_sam2
@@ -142,42 +124,65 @@ def try_load(config_name_with_yaml: str, checkpoint_name: str):
             predictor = SAM2ImagePredictor(sam2_model)
             sam_logger.info(f"✅ Loaded {config_name_with_yaml} successfully on {device}")
             return predictor
+
         except Exception as e:
-            error_msg = f"Failed to load {config_name_with_yaml}: {e}\nTraceback: {traceback.format_exc()}"
-            tried.append(error_msg)
-            sam_logger.warning(error_msg)
+            err = f"Failed to load {config_name_with_yaml}: {e}\nTraceback: {traceback.format_exc()}"
+            sam_logger.warning(err)
             return None
 
     predictor = try_load("sam2_hiera_large.yaml", "sam2_hiera_large.pt")
-
     if predictor is None:
-        error_message = "SAM2 loading failed for large model. Reasons:\n" + "\n".join(tried)
-        sam_logger.error(f"❌ {error_message}")
-        raise gr.Error(error_message)
-
+        raise gr.Error("SAM2 loading failed for large model. Check configs/checkpoint.")
     return predictor
 
 # ============================================================================ #
-# MatAnyOne LOADER (robust cfg + net + core)
+# MatAnyOne LOADER (hard requirement; auto-clone if needed; load weights)
 # ============================================================================ #
+def _git_clone(url: str, dest: str):
+    os.makedirs(os.path.dirname(dest), exist_ok=True)
+    if not os.path.exists(dest):
+        logger.info(f"📥 Cloning {url} → {dest}")
+        subprocess.check_call(["git", "clone", "--depth", "1", url, dest])
+
+def _ensure_repo_on_path(repo_dir: str):
+    if repo_dir not in sys.path:
+        sys.path.insert(0, repo_dir)
+
+def _download_weight_candidates(repo_id: str, filenames: list, cache_dir="./checkpoints") -> Optional[str]:
+    from huggingface_hub import hf_hub_download
+    last_err = None
+    for fn in filenames:
+        try:
+            path = hf_hub_download(repo_id=repo_id, filename=fn, cache_dir=cache_dir, local_dir_use_symlinks=False)
+            return path
+        except Exception as e:
+            last_err = e
+    if last_err:
+        raise RuntimeError(f"Could not download MatAnyOne weights from {repo_id} (tried: {filenames}). Last error: {last_err}")
+    return None
+
 def load_matanyone(device: str):
     """
-    Robust MatAnyOne loader:
-    - Loads an OmegaConf cfg from common paths or creates a minimal default
-    - Builds the network and wraps with InferenceCore
-    - Tries multiple import layouts to accommodate different repos
+    Load MatAnyOne (hard requirement):
+    - Try installed package
+    - Else auto-clone official repo to ./third_party/MatAnyOne and import
+    - Load cfg (or minimal), download weights, build Net + InferenceCore
     """
     from omegaconf import OmegaConf
-
     ma_logger = logging.getLogger("MatAnyOneLoader")
 
-    # Try to locate a config file; otherwise minimal default
-    cfg_path_candidates = [
-        "Configs/matanyone.yaml",
-        "configs/matanyone.yaml",
+    # 1) Try to import from installed package first (two common layouts)
+    import_attempts = [
+        ("matanyone.model.matanyone", "MatAnyOne", "matanyone.inference.inference_core", "InferenceCore"),
+        ("matanyone", "MatAnyOne", "matanyone", "InferenceCore"),
     ]
+
+    # 2) Prepare local clone as backup
+    repo_dir = os.path.abspath("./third_party/MatAnyOne")
+    official_git = "https://github.com/PeiqingYang/MatAnyOne"
+
     cfg = None
-    for p in cfg_path_candidates:
+    for p in ("Configs/matanyone.yaml", "configs/matanyone.yaml"):
         if os.path.exists(p):
             ma_logger.info(f"Loading MatAnyOne cfg: {p}")
             cfg = OmegaConf.load(p)
@@ -190,34 +195,78 @@ def load_matanyone(device: str):
             "device": device,
         })
 
-    last_err = None
-
-    # Layout A (common in forks): separate model + inference modules
-    try:
-        from matanyone.model.matanyone import MatAnyOne
-        from matanyone.inference.inference_core import InferenceCore
-        net = MatAnyOne(cfg)
+    # helper to finalize core from given modules
+    def _build_from_modules(net_mod, core_mod):
+        Net = getattr(net_mod, "MatAnyOne")
+        Core = getattr(core_mod, "InferenceCore")
+        net = Net(cfg)
         net.to(device)
-        core = InferenceCore(net, cfg)
-        ma_logger.info("✅ MatAnyOne loaded (layout A)")
+
+        # Try to get weights if the net exposes a load function or needs explicit weights
+        # Try common filenames
+        try:
+            weight_path = _download_weight_candidates(
+                repo_id="PeiqingYang/MatAnyOne-v1.0",
+                filenames=[
+                    "MatAnyOne_swinB.pth",
+                    "MatAnyOne_SwinB.pth",
+                    "matanyone_swinB.pth",
+                    "weights_swinB.pth",
+                ],
+                cache_dir="./checkpoints"
+            )
+            if weight_path and hasattr(net, "load_state_dict"):
+                state = torch.load(weight_path, map_location=device)
+                # Some repos wrap state under 'state_dict' or similar
+                if isinstance(state, dict) and "state_dict" in state:
+                    state = state["state_dict"]
+                net.load_state_dict(state, strict=False)
+                ma_logger.info(f"✅ Loaded MatAnyOne weights: {os.path.basename(weight_path)}")
+        except Exception as e:
+            ma_logger.warning(f"Could not load MatAnyOne weights automatically (continuing): {e}")
+
+        core = Core(net, cfg)
         return core
-    except Exception as e:
-        last_err = e
-        ma_logger.warning(f"Layout A failed: {e}")
 
-    # Layout B (single package exposing classes)
+    # A) Installed package attempts
+    last_err = None
+    for mod_net, _, mod_core, _ in import_attempts:
+        try:
+            net_mod = importlib.import_module(mod_net)
+            core_mod = importlib.import_module(mod_core)
+            core = _build_from_modules(net_mod, core_mod)
+            ma_logger.info(f"✅ MatAnyOne loaded from installed package ({mod_net} / {mod_core})")
+            return core
+        except Exception as e:
+            last_err = e
+            ma_logger.warning(f"MatAnyOne import attempt failed ({mod_net} / {mod_core}): {e}")
+
+    # B) Local clone fallback
     try:
-        from matanyone import MatAnyOne, InferenceCore
-        net = MatAnyOne(cfg)
-        net.to(device)
-        core = InferenceCore(net, cfg)
-        ma_logger.info("✅ MatAnyOne loaded (layout B)")
-        return core
+        _git_clone(official_git, repo_dir)
+        _ensure_repo_on_path(repo_dir)
+        # Try common in-repo layouts
+        clone_attempts = [
+            ("matanyone.model.matanyone", "matanyone.inference.inference_core"),
+            ("model.matanyone", "inference.inference_core"),
+            ("src.matanyone.model.matanyone", "src.matanyone.inference.inference_core"),
+        ]
+        for mod_net, mod_core in clone_attempts:
+            try:
+                net_mod = importlib.import_module(mod_net)
+                core_mod = importlib.import_module(mod_core)
+                core = _build_from_modules(net_mod, core_mod)
+                ma_logger.info(f"✅ MatAnyOne loaded from cloned repo ({mod_net} / {mod_core})")
+                return core
+            except Exception as e:
+                last_err = e
+                ma_logger.warning(f"MatAnyOne clone import failed ({mod_net} / {mod_core}): {e}")
     except Exception as e:
         last_err = e
-        ma_logger.warning(f"Layout B failed: {e}")
+        ma_logger.error(f"Cloning MatAnyOne failed: {e}")
 
-    raise RuntimeError(f"Failed to initialize MatAnyOne/InferenceCore. Last error: {last_err}")
+    # If we reach here, MatAnyOne is not usable — HARD FAIL by request
+    raise RuntimeError(f"MatAnyOne required but failed to initialize. Last error: {last_err}")
 
 # ============================================================================ #
 # GLOBALS & MODEL SETUP
@@ -228,34 +277,27 @@ def load_matanyone(device: str):
 loading_lock = threading.Lock()
 
 def download_and_setup_models(progress: Optional[gr.Progress] = None):
-    """
-    Download and setup models (SAM2 and MatAnyOne), robust to HF Spaces and local dev.
-    """
+    """Download and setup models. BOTH SAM2 and MatAnyOne are REQUIRED."""
     global sam2_predictor, matanyone_model, models_loaded
 
     with loading_lock:
         if models_loaded:
-            return "✅ High-quality models already loaded"
-        try:
-            logger.info("🔄 Starting ENHANCED model loading with fallback...")
+            return "✅ SAM2 + MatAnyOne already loaded"
 
+        try:
+            logger.info("🔄 Starting ENHANCED model loading...")
             device = "cuda" if torch.cuda.is_available() else "cpu"
 
-            # --- Load SAM2 ---
+            # --- Load SAM2 (required) ---
             local_sam2 = load_sam2_predictor(device=device, progress=progress)
             sam2_predictor = local_sam2
 
-            # --- Load MatAnyOne ---
-            try:
-                local_matanyone = load_matanyone(device)
-                matanyone_model = local_matanyone
-                logger.info("✅ MatAnyOne loaded")
-            except Exception as e:
-                logger.warning(f"❌ MatAnyOne load failed: {e}")
-                raise RuntimeError("MatAnyone model could not be loaded.")
+            # --- Load MatAnyOne (required) ---
+            local_matanyone = load_matanyone(device)
+            matanyone_model = local_matanyone
 
             models_loaded = True
-            logger.info("--- ✅ All models loaded successfully ---")
+            logger.info("--- ✅ All models loaded successfully (SAM2 + MatAnyOne) ---")
             return "✅ SAM2 + MatAnyOne loaded successfully!"
         except Exception as e:
             logger.error(f"❌ Enhanced loading failed: {str(e)}")
@@ -263,32 +305,23 @@ def download_and_setup_models(progress: Optional[gr.Progress] = None):
             return f"❌ Enhanced loading failed: {str(e)}"
 
 # ============================================================================ #
-# TWO-STAGE PROCESSING PIPELINE
+# TWO-STAGE PROCESSING PIPELINE (uses your utilities’ segmentation/compositing)
 # ============================================================================ #
-def process_video_hq(
-    video_path,
-    background_choice,
-    custom_background_path,
-    progress: Optional[gr.Progress] = None
-):
+def process_video_hq(video_path, background_choice, custom_background_path, progress: Optional[gr.Progress] = None):
     """TWO-STAGE High-quality video processing: Original → Green Screen → Final Background"""
     if not models_loaded:
         return None, "❌ Models not loaded. Click 'Load Models' first."
-
     if not video_path:
         return None, "❌ No video file provided."
 
     def _prog(pct: float, desc: str):
         if progress is not None:
-            try:
-                progress(pct, desc=desc)
-            except Exception:
-                pass
+            try: progress(pct, desc=desc)
+            except Exception: pass
 
     try:
         _prog(0.0, "🎬 Initializing TWO-STAGE processing...")
 
-        # Validate and read video
         if not os.path.exists(video_path):
             return None, f"❌ Video file not found: {video_path}"
 
@@ -296,161 +329,119 @@ def _prog(pct: float, desc: str):
         if not cap.isOpened():
             return None, "❌ Could not open video file. Please check the format."
 
-        # Get video properties
         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))
-
         logger.info(f"Video properties: {frame_width}x{frame_height}, {fps}fps, {total_frames} frames")
 
         if total_frames == 0:
             return None, "❌ Video appears to be empty or corrupted."
 
-        # Prepare final background for Stage 2
+        # Prepare final background (Stage 2)
         background = None
         background_name = ""
 
         if background_choice == "custom" and custom_background_path:
-            try:
-                background = cv2.imread(custom_background_path)
-                if background is None:
-                    return None, "❌ Could not read custom background image. Please check the file format."
-                background_name = "Custom Image"
-                logger.info("Using custom background image")
-            except Exception as e:
-                return None, f"❌ Error loading custom background: {str(e)}"
+            background = cv2.imread(custom_background_path)
+            if background is None:
+                return None, "❌ Could not read custom background image. Please check the file format."
+            background_name = "Custom Image"
+            logger.info("Using custom background image")
         else:
             if background_choice in PROFESSIONAL_BACKGROUNDS:
-                try:
-                    bg_config = PROFESSIONAL_BACKGROUNDS[background_choice]
-                    background = create_professional_background(bg_config, frame_width, frame_height)
-                    background_name = bg_config["name"]
-                    logger.info(f"Using professional background: {background_name}")
-                except Exception as e:
-                    logger.error(f"Error creating professional background: {e}")
-                    return None, f"❌ Error creating background: {str(e)}"
+                bg_config = PROFESSIONAL_BACKGROUNDS[background_choice]
+                background = create_professional_background(bg_config, frame_width, frame_height)
+                background_name = bg_config["name"]
+                logger.info(f"Using professional background: {background_name}")
             else:
                 return None, f"❌ Invalid background selection: {background_choice}"
 
         if background is None:
             return None, "❌ Failed to create background."
 
-        # Setup codec and timestamp
         timestamp = int(time.time())
         fourcc = cv2.VideoWriter_fourcc(*'mp4v')
 
-        # STAGE 1: Create green screen video (Original → Green Screen)
+        # STAGE 1: Original → Green Screen
         _prog(0.1, "🟢 STAGE 1: Creating green screen version...")
         greenscreen_path = f"/tmp/greenscreen_{timestamp}.mp4"
         greenscreen_writer = cv2.VideoWriter(greenscreen_path, fourcc, fps, (frame_width, frame_height))
-
         if not greenscreen_writer.isOpened():
             return None, "❌ Could not create green screen video file."
 
         frame_count = 0
-
-        # Process original video to green screen
         while True:
             ret, frame = cap.read()
             if not ret:
                 break
-
             try:
-                progress_pct = 0.1 + (frame_count / total_frames) * 0.4
-                _prog(progress_pct, f"🟢 Green screen frame {frame_count + 1}/{total_frames}")
-
-                # Segment person and create green screen frame
-                mask = segment_person_hq(frame)
-                refined_mask = refine_mask_hq(frame, mask)
-                green_screen = create_green_screen_background(frame)
-                green_screen_frame = replace_background_hq(frame, refined_mask, green_screen)
-
+                _prog(0.1 + (frame_count / max(1, total_frames)) * 0.4, f"🟢 Green screen frame {frame_count + 1}/{total_frames}")
+                mask = segment_person_hq(frame)            # from utilities
+                refined_mask = refine_mask_hq(frame, mask) # from utilities
+                green_screen = create_green_screen_background(frame)  # from utilities
+                green_screen_frame = replace_background_hq(frame, refined_mask, green_screen)  # from utilities
                 greenscreen_writer.write(green_screen_frame)
-                frame_count += 1
-
-                if frame_count % 100 == 0:
-                    gc.collect()
-                    if torch.cuda.is_available():
-                        torch.cuda.empty_cache()
-
             except Exception as e:
                 logger.warning(f"Error in Stage 1 frame {frame_count}: {e}")
                 greenscreen_writer.write(frame)
-                frame_count += 1
+            frame_count += 1
+            if frame_count % 100 == 0:
+                gc.collect()
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
 
         greenscreen_writer.release()
         cap.release()
 
-        # STAGE 2: Replace green screen with final background (Green Screen → Final)
+        # STAGE 2: Green Screen → Final Background
         _prog(0.5, f"🎨 STAGE 2: Replacing green screen with {background_name}...")
-
         final_path = f"/tmp/final_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 final output video file."
 
-        # Open green screen video
         greenscreen_cap = cv2.VideoCapture(greenscreen_path)
         if not greenscreen_cap.isOpened():
             return None, "❌ Could not open green screen video."
 
         frame_count = 0
-
-        # Process green screen video to final background with enhanced green detection
         while True:
             ret, green_frame = greenscreen_cap.read()
             if not ret:
                 break
-
             try:
-                progress_pct = 0.5 + (frame_count / total_frames) * 0.4
-                _prog(progress_pct, f"🎬 Final compositing frame {frame_count + 1}/{total_frames}")
-
-                # Detect green screen with wider detection range
+                _prog(0.5 + (frame_count / max(1, total_frames)) * 0.4, f"🎬 Final compositing frame {frame_count + 1}/{total_frames}")
                 hsv = cv2.cvtColor(green_frame, cv2.COLOR_BGR2HSV)
                 lower_green = np.array([25, 30, 30])
                 upper_green = np.array([100, 255, 255])
                 green_mask = cv2.inRange(hsv, lower_green, upper_green)
-
-                # Additional mask processing for cleaner edges
                 kernel = np.ones((3, 3), np.uint8)
                 green_mask = cv2.morphologyEx(green_mask, cv2.MORPH_OPEN, kernel)
                 green_mask = cv2.morphologyEx(green_mask, cv2.MORPH_CLOSE, kernel)
-                green_mask = 255 - green_mask  # Invert mask (person = white, green = black)
-
-                result_frame = replace_background_hq(green_frame, green_mask, background)
+                green_mask = 255 - green_mask
+                result_frame = replace_background_hq(green_frame, green_mask, background)  # from utilities
                 final_writer.write(result_frame)
-                frame_count += 1
-
-                if frame_count % 100 == 0:
-                    gc.collect()
-                    if torch.cuda.is_available():
-                        torch.cuda.empty_cache()
-
             except Exception as e:
                 logger.warning(f"Error in Stage 2 frame {frame_count}: {e}")
                 final_writer.write(green_frame)
-                frame_count += 1
+            frame_count += 1
+            if frame_count % 100 == 0:
+                gc.collect()
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
 
         greenscreen_cap.release()
         final_writer.release()
 
-        # Cleanup intermediate green screen file
-        try:
-            os.remove(greenscreen_path)
-        except Exception:
-            pass
+        try: os.remove(greenscreen_path)
+        except Exception: pass
 
         if frame_count == 0:
             return None, "❌ No frames were processed successfully."
 
         _prog(0.9, "🎵 Adding high-quality audio...")
-
-        # Add audio back with high quality settings
         final_output = f"/tmp/final_output_hq_{timestamp}.mp4"
-
         try:
             audio_cmd = (
                 f'ffmpeg -y -i "{final_path}" -i "{video_path}" '
@@ -459,15 +450,12 @@ def _prog(pct: float, desc: str):
                 f'-map 0:v:0 -map 1:a:0? -shortest "{final_output}"'
             )
             result = os.system(audio_cmd)
-
             if result != 0 or not os.path.exists(final_output):
                 logger.warning("Audio merging failed, using video without audio")
                 shutil.copy2(final_path, final_output)
-
         except Exception as e:
             logger.warning(f"Audio processing error: {e}, using video without audio")
-            try:
-                shutil.copy2(final_path, final_output)
+            try: shutil.copy2(final_path, final_output)
             except Exception as e2:
                 logger.error(f"Failed to copy video file: {e2}")
                 return None, f"❌ Failed to finalize video: {str(e2)}"
@@ -476,17 +464,14 @@ def _prog(pct: float, desc: str):
         try:
             myavatar_path = "/tmp/MyAvatar/My_Videos/"
             os.makedirs(myavatar_path, exist_ok=True)
-
             saved_filename = f"two_stage_bg_replaced_{timestamp}.mp4"
             saved_path = os.path.join(myavatar_path, saved_filename)
             shutil.copy2(final_output, saved_path)
-
             logger.info(f"Video saved to: {saved_path}")
         except Exception as e:
             logger.warning(f"Could not save to MyAvatar directory: {e}")
             saved_filename = os.path.basename(final_output)
 
-        # Cleanup temporary files
         try:
             if os.path.exists(final_path):
                 os.remove(final_path)
@@ -494,7 +479,6 @@ def _prog(pct: float, desc: str):
             pass
 
         _prog(1.0, "✅ TWO-STAGE processing complete!")
-
         success_message = (
             f"✅ TWO-STAGE Success!\n"
             f"🟢 Stage 1: Original → Green Screen\n"
@@ -504,22 +488,17 @@ def _prog(pct: float, desc: str):
             f"🎯 Quality: Cinema-grade with SAM2 + MatAnyOne\n"
             f"🚀 Method: Professional two-stage compositing"
         )
-
         return final_output, success_message
 
     except Exception as e:
-        error_msg = f"❌ TWO-STAGE Processing Error: {str(e)}"
         logger.error(f"Video processing error: {traceback.format_exc()}")
-        return None, error_msg
+        return None, f"❌ TWO-STAGE Processing Error: {str(e)}"
 
 # ============================================================================ #
 # GRADIO UI
 # ============================================================================ #
 def create_interface():
-    """Create enhanced Gradio interface with comprehensive features and 4-method background system"""
-
     def extract_video_path(v):
-        # Robustly extract file path from input (tuple, list, or string)
         if isinstance(v, (tuple, list)) and len(v) > 0:
             return v[0]
         return v
@@ -532,46 +511,27 @@ def extract_video_path(v):
         .progress-bar { background: linear-gradient(90deg, #3498db, #2ecc71) !important; }
         """
     ) as demo:
-
-        # Header
         gr.Markdown("# 🎬 Cinema-Quality Video Background Replacement")
         gr.Markdown("**Upload a video → Choose a background → Get professional results with AI**")
         gr.Markdown("*Powered by SAM2 + MatAnyOne with multi-fallback loading for maximum reliability*")
         gr.Markdown("---")
 
         with gr.Row():
-            # Left column - Input and controls
             with gr.Column(scale=1):
                 gr.Markdown("### 📥 Step 1: Upload Your Video")
                 gr.Markdown("*Supports MP4, MOV, AVI, and other common formats*")
+                video_input = gr.Video(label="🎥 Drop your video here", height=300)
 
-                video_input = gr.Video(
-                    label="🎥 Drop your video here",
-                    height=300
-                )
-
-                # Video preview
-                video_preview = gr.Video(
-                    label="📺 Preview of Uploaded Video",
-                    height=200,
-                    interactive=False
-                )
-                video_input.change(
-                    fn=extract_video_path,
-                    inputs=video_input,
-                    outputs=video_preview
-                )
+                video_preview = gr.Video(label="📺 Preview of Uploaded Video", height=200, interactive=False)
+                video_input.change(fn=extract_video_path, inputs=video_input, outputs=video_preview)
 
                 gr.Markdown("### 🎨 Step 2: Choose Background Method")
                 gr.Markdown("*Select your preferred background creation method*")
-
-                # FIXED Radio (flat choices)
                 background_method = gr.Radio(
                     choices=["upload", "professional", "colors", "ai"],
                     value="professional",
                     label="Background Method"
                 )
-                # Labels hint
                 gr.Markdown(
                     "- **upload** = 📷 Upload Image  \n"
                     "- **professional** = 🎨 Professional Presets  \n"
@@ -579,15 +539,10 @@ def extract_video_path(v):
                     "- **ai** = 🤖 AI Generated"
                 )
 
-                # Method A: Upload Image
                 with gr.Group(visible=False) as upload_group:
                     gr.Markdown("**📷 Upload Your Background Image**")
-                    custom_background = gr.Image(
-                        label="Drop your background image here",
-                        type="filepath"
-                    )
+                    custom_background = gr.Image(label="Drop your background image here", type="filepath")
 
-                # Method B: Professional Presets
                 with gr.Group(visible=True) as professional_group:
                     gr.Markdown("**🎨 Professional Background Presets**")
                     professional_choice = gr.Dropdown(
@@ -596,53 +551,43 @@ def extract_video_path(v):
                         label="Select Professional Background"
                     )
 
-                # Method C: Colors/Gradients
                 with gr.Group(visible=False) as colors_group:
                     gr.Markdown("**🌈 Custom Colors & Gradients**")
-
                     gradient_type = gr.Dropdown(
                         choices=["solid", "vertical", "horizontal", "diagonal", "radial", "soft_radial"],
                         value="vertical",
                         label="Gradient Type"
                     )
-
                     with gr.Row():
                         color1 = gr.ColorPicker(label="🎨 Color 1", value="#3498db")
                         color2 = gr.ColorPicker(label="🎨 Color 2", value="#2ecc71")
-
                     with gr.Row():
                         color3 = gr.ColorPicker(label="🎨 Color 3", value="#e74c3c")
                         use_third_color = gr.Checkbox(label="Use 3rd color", value=False)
 
-                # Method D: AI Generated
                 with gr.Group(visible=False) as ai_group:
                     gr.Markdown("**🤖 AI Generated Background**")
-
                     ai_prompt = gr.Textbox(
                         label="Describe your background",
                         placeholder="e.g., 'modern office with plants', 'sunset over mountains', 'abstract tech pattern'",
                         lines=2
                     )
-
                     ai_style = gr.Dropdown(
                         choices=["photorealistic", "artistic", "abstract", "minimalist", "corporate", "nature"],
                         value="photorealistic",
                         label="Style"
                     )
-
                     with gr.Row():
                         generate_ai_btn = gr.Button("🎨 Generate Background", variant="secondary")
                         ai_generated_image = gr.Image(label="Generated Background", type="filepath", visible=False)
 
-                # Background method switching function
                 def switch_background_method(method):
                     return (
-                        gr.update(visible=(method == "upload")),       # upload_group
-                        gr.update(visible=(method == "professional")), # professional_group
-                        gr.update(visible=(method == "colors")),       # colors_group
-                        gr.update(visible=(method == "ai"))            # ai_group
+                        gr.update(visible=(method == "upload")),
+                        gr.update(visible=(method == "professional")),
+                        gr.update(visible=(method == "colors")),
+                        gr.update(visible=(method == "ai"))
                     )
-
                 background_method.change(
                     fn=switch_background_method,
                     inputs=background_method,
@@ -651,35 +596,16 @@ def switch_background_method(method):
 
                 gr.Markdown("### 🎬 Processing Controls")
                 gr.Markdown("*First load the AI models, then process your video*")
-
                 with gr.Row():
-                    load_models_btn = gr.Button(
-                        "🚀 Step 1: Load AI Models",
-                        variant="secondary",
-                    )
-                    process_btn = gr.Button(
-                        "✨ Step 2: Process Video",
-                        variant="primary",
-                    )
+                    load_models_btn = gr.Button("🚀 Step 1: Load AI Models", variant="secondary")
+                    process_btn = gr.Button("✨ Step 2: Process Video", variant="primary")
 
-                # System status
-                status_text = gr.Textbox(
-                    label="🔧 System Status",
-                    value=get_model_status(),
-                    interactive=False,
-                    lines=3
-                )
+                status_text = gr.Textbox(label="🔧 System Status", value=get_model_status(), interactive=False, lines=3)
 
-            # Right column - Results and preview
             with gr.Column(scale=1):
                 gr.Markdown("### 📤 Your Results")
                 gr.Markdown("*Processed video will appear here after Step 2*")
-
-                video_output = gr.Video(
-                    label="🎬 Your Processed Video",
-                    height=400
-                )
-
+                video_output = gr.Video(label="🎬 Your Processed Video", height=400)
                 result_text = gr.Textbox(
                     label="📊 Processing Results",
                     interactive=False,
@@ -688,29 +614,15 @@ def switch_background_method(method):
                 )
 
                 gr.Markdown("### 🎨 Professional Backgrounds Available")
-
-                # Create background preview grid
                 bg_preview_html = """
                 <div style='display: grid; grid-template-columns: repeat(3, 1fr); gap: 8px; padding: 10px; max-height: 400px; overflow-y: auto; border: 1px solid #ddd; border-radius: 8px;'>
                 """
                 for key, config in PROFESSIONAL_BACKGROUNDS.items():
                     colors = config["colors"]
-                    if len(colors) >= 2:
-                        gradient = f"linear-gradient(45deg, {colors[0]}, {colors[-1]})"
-                    else:
-                        gradient = colors[0]
+                    gradient = f"linear-gradient(45deg, {colors[0]}, {colors[-1]})" if len(colors) >= 2 else colors[0]
                     bg_preview_html += f"""
-                    <div style='
-                        padding: 12px 8px; 
-                        border: 1px solid #ddd; 
-                        border-radius: 6px; 
-                        text-align: center; 
-                        background: {gradient};
-                        min-height: 60px;
-                        display: flex;
-                        align-items: center;
-                        justify-content: center;
-                    '>
+                    <div style='padding: 12px 8px; border: 1px solid #ddd; border-radius: 6px; text-align: center; background: {gradient};
+                                min-height: 60px; display: flex; align-items: center; justify-content: center;'>
                         <div>
                             <strong style='color: white; text-shadow: 1px 1px 2px rgba(0,0,0,0.8); font-size: 12px; display: block;'>{config["name"]}</strong>
                             <small style='color: rgba(255,255,255,0.9); text-shadow: 1px 1px 1px rgba(0,0,0,0.6); font-size: 10px;'>{config.get("description", "")[:30]}...</small>
@@ -720,126 +632,79 @@ def switch_background_method(method):
                 bg_preview_html += "</div>"
                 gr.HTML(bg_preview_html)
 
-        # AI Background Generation Function
         def generate_ai_background(prompt, style):
-            """Generate AI background using procedural methods"""
             if not prompt or not prompt.strip():
                 return None, "❌ Please enter a prompt"
-
             try:
-                # Create procedural background based on prompt
                 bg_image = create_procedural_background(prompt, style, 1920, 1080)
-
                 if bg_image is not None:
                     with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
                         cv2.imwrite(tmp.name, bg_image)
                         return tmp.name, f"✅ Background generated: {prompt[:50]}..."
-                else:
-                    return None, "❌ Generation failed, try different prompt"
+                return None, "❌ Generation failed, try different prompt"
             except Exception as e:
                 logger.error(f"AI generation error: {e}")
                 return None, f"❌ Generation error: {str(e)}"
 
-        # Enhanced video processing function that handles all 4 methods
         def process_video_enhanced(
-            video_path,
-            bg_method,
-            custom_img,
-            prof_choice,
-            grad_type,
-            color1, color2, color3, use_third,
-            ai_prompt, ai_style, ai_img,
+            video_path, bg_method, custom_img, prof_choice, grad_type,
+            color1, color2, color3, use_third, ai_prompt, ai_style, ai_img,
             progress: Optional[gr.Progress] = None
         ):
-            """Process video with any of the 4 background methods using TWO-STAGE approach"""
-
             if not models_loaded:
                 return None, "❌ Models not loaded. Click 'Load Models' first."
-
             if not video_path:
                 return None, "❌ No video file provided."
-
             try:
                 if bg_method == "upload":
                     if custom_img and os.path.exists(custom_img):
                         return process_video_hq(video_path, "custom", custom_img, progress)
-                    else:
-                        return None, "❌ No image uploaded. Please upload a background image."
-
+                    return None, "❌ No image uploaded. Please upload a background image."
                 elif bg_method == "professional":
                     if prof_choice and prof_choice in PROFESSIONAL_BACKGROUNDS:
                         return process_video_hq(video_path, prof_choice, None, progress)
-                    else:
-                        return None, f"❌ Invalid professional background: {prof_choice}"
-
+                    return None, f"❌ Invalid professional background: {prof_choice}"
                 elif bg_method == "colors":
                     try:
                         colors = [color1 or "#3498db", color2 or "#2ecc71"]
                         if use_third and color3:
                             colors.append(color3)
-
                         bg_config = {
                             "type": "gradient" if grad_type != "solid" else "color",
                             "colors": colors if grad_type != "solid" else [colors[0]],
                             "direction": grad_type if grad_type != "solid" else "vertical"
                         }
-
                         gradient_bg = create_professional_background(bg_config, 1920, 1080)
                         temp_path = f"/tmp/gradient_{int(time.time())}.png"
                         cv2.imwrite(temp_path, gradient_bg)
-
                         return process_video_hq(video_path, "custom", temp_path, progress)
                     except Exception as e:
                         return None, f"❌ Error creating gradient: {str(e)}"
-
                 elif bg_method == "ai":
                     if ai_img and os.path.exists(ai_img):
                         return process_video_hq(video_path, "custom", ai_img, progress)
-                    else:
-                        return None, "❌ No AI background generated. Click 'Generate Background' first."
-
+                    return None, "❌ No AI background generated. Click 'Generate Background' first."
                 else:
                     return None, f"❌ Unknown background method: {bg_method}"
-
             except Exception as e:
                 logger.error(f"Enhanced processing error: {e}")
                 return None, f"❌ Processing error: {str(e)}"
 
-        # Wire up callbacks
-        load_models_btn.click(
-            fn=download_and_setup_models,
-            outputs=status_text
-        )
-
-        generate_ai_btn.click(
-            fn=generate_ai_background,
-            inputs=[ai_prompt, ai_style],
-            outputs=[ai_generated_image, status_text]
-        )
-
+        load_models_btn.click(fn=download_and_setup_models, outputs=status_text)
+        generate_ai_btn.click(fn=generate_ai_background, inputs=[ai_prompt, ai_style], outputs=[ai_generated_image, status_text])
         process_btn.click(
             fn=process_video_enhanced,
-            inputs=[
-                video_input,           # video_path
-                background_method,     # bg_method
-                custom_background,     # custom_img
-                professional_choice,   # prof_choice
-                gradient_type,         # grad_type
-                color1, color2, color3, use_third_color,  # colors
-                ai_prompt, ai_style, ai_generated_image   # AI
-            ],
+            inputs=[video_input, background_method, custom_background, professional_choice,
+                    gradient_type, color1, color2, color3, use_third_color,
+                    ai_prompt, ai_style, ai_generated_image],
             outputs=[video_output, result_text]
         )
 
-        # Info
         with gr.Accordion("ℹ️ ENHANCED Quality & Features", open=False):
             gr.Markdown("""
             ### 🏆 TWO-STAGE Cinema-Quality Features:
             **Stage 1**: Original → Green Screen (SAM2 + MatAnyOne)
             **Stage 2**: Green Screen → Final Background (professional chroma key)
-
-            **Background Methods**: Upload image / Professional presets / Gradients / AI generated
-
             **Quality**: Edge feathering, gamma correction, mask cleanup, H.264 CRF 18, AAC 192kbps.
             """)
 
@@ -852,12 +717,10 @@ def process_video_enhanced(
 # MAIN
 # ============================================================================ #
 def main():
-    """Main application entry point"""
     try:
+        print(f"===== Application Startup at {time.strftime('%Y-%m-%d %H:%M:%S')} =====\n")
         print("🎬 Cinema-Quality Video Background Replacement")
         print("=" * 50)
-
-        # Initialize application paths
         os.makedirs("/tmp/MyAvatar/My_Videos/", exist_ok=True)
         os.makedirs(os.path.expanduser("~/.cache/sam2"), exist_ok=True)
 
@@ -870,17 +733,11 @@ def main():
         print("  • Enhanced stability & error handling")
         print("=" * 50)
 
-        # Create and launch interface
         logger.info("🌐 Creating Gradio interface...")
         demo = create_interface()
 
         logger.info("🚀 Launching application...")
-        demo.launch(
-            server_name="0.0.0.0",
-            server_port=7860,
-            share=True,
-            show_error=True
-        )
+        demo.launch(server_name="0.0.0.0", server_port=7860, share=True, show_error=True)
 
     except KeyboardInterrupt:
         logger.info("🛑 Application stopped by user")