Update app.py

app.py

@@ -1,4 +1,18 @@
 #!/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"
+os.environ.setdefault("MKL_NUM_THREADS", "1")
+os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
+os.environ.setdefault("NUMEXPR_NUM_THREADS", "1")
+# Optional CUDA allocator tuning
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:1024")
+os.environ.setdefault("CUDA_LAUNCH_BLOCKING", "0")
+# ========================================================================
+
 """
 High-Quality Video Background Replacement - MAIN APPLICATION
 Upload video → Choose professional background → Replace with cinema quality
@@ -6,7 +20,6 @@
 cinema-quality processing, lazy loading, and enhanced stability
 """
 
-import os
 import sys
 import tempfile
 import cv2
@@ -25,24 +38,17 @@
 import queue
 from typing import Optional, Tuple, Dict, Any
 import logging
+import warnings
 
-# Import all utilities
+# 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 *
 
-# Fix OpenMP threads issue - remove problematic environment variable
-try:
-    if 'OMP_NUM_THREADS' in os.environ:
-        del os.environ['OMP_NUM_THREADS']
-except:
-    pass
-
-# Suppress warnings and optimize for quality
-import warnings
+# Suppress warnings
 warnings.filterwarnings("ignore")
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:1024'
-os.environ['CUDA_LAUNCH_BLOCKING'] = '0'
 
-# Setup logging for debugging
+# Setup logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -66,97 +72,163 @@ def patched_get_type(schema):
     logger.warning(f"⚠️ Could not apply Gradio monkey patch: {e}")
 
 # ============================================================================ #
-
-def load_sam2_predictor(device="cuda", progress=gr.Progress()):
+# SAM2 LOADER (Hydra cfg, correct build)
+# ============================================================================ #
+def load_sam2_predictor(device: str = "cuda", progress: Optional[gr.Progress] = None):
     """
     Loads the SAM2 model and returns a SAM2ImagePredictor instance.
-    - Tries to load 'sam2_hiera_large' first.
-    - Falls back to 'sam2_hiera_tiny' if large cannot be loaded.  # <--- OBS: Denne linje kan du også opdatere!
-    - Assumes YAML configs are in ./Configs/ (capital C), as required by upstream.
+    - Tries to load 'sam2_hiera_large' model config from ./Configs
+    - Downloads checkpoint via HF if not present
     """
     import hydra
     from omegaconf import OmegaConf
-    import torch
-    import logging
 
-    logger = logging.getLogger("SAM2Loader")
+    sam_logger = logging.getLogger("SAM2Loader")
     configs_dir = os.path.abspath("Configs")
-    logger.info(f"Looking for SAM2 configs in absolute path: {configs_dir}")
+    sam_logger.info(f"Looking for SAM2 configs in absolute path: {configs_dir}")
 
     if not os.path.isdir(configs_dir):
-        logger.error(f"FATAL: Configs directory not found at '{configs_dir}'. Please ensure the 'Configs' folder is at the root of your repository.")
-        raise gr.Error(f"FATAL: SAM2 Configs directory not found. Check repository structure.")
+        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 = []
 
-    def try_load(config_name, checkpoint_name):
+    def _maybe_progress(pct: float, desc: str):
+        if progress is not None:
+            try:
+                progress(pct, desc=desc)
+            except Exception:
+                pass
+
+    def try_load(config_name_with_yaml: str, checkpoint_name: str):
         try:
             checkpoint_path = os.path.join("./checkpoints", checkpoint_name)
-            logger.info(f"Attempting to use checkpoint: {checkpoint_path}")
+            sam_logger.info(f"Attempting to use checkpoint: {checkpoint_path}")
 
             if not os.path.exists(checkpoint_path):
-                logger.info(f"Downloading {checkpoint_name} from Hugging Face Hub...")
-                progress(0.1, desc=f"Downloading {checkpoint_name}...")
+                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
                 checkpoint_path = hf_hub_download(
-                    repo_id=f"facebook/{config_name.replace('.yaml', '')}",  # e.g. facebook/sam2_hiera_large
+                    repo_id=repo,
                     filename=checkpoint_name,
                     cache_dir="./checkpoints",
                     local_dir_use_symlinks=False
                 )
-                logger.info(f"✅ Download complete: {checkpoint_path}")
+                sam_logger.info(f"✅ Download complete: {checkpoint_path}")
 
+            # Reset & init Hydra
             if hydra.core.global_hydra.GlobalHydra.instance().is_initialized():
                 hydra.core.global_hydra.GlobalHydra.instance().clear()
 
             hydra.initialize(config_path=os.path.relpath(configs_dir), job_name=f"sam2_load_{int(time.time())}")
-            cfg = hydra.compose(config_name=config_name)
+            cfg_name = config_name_with_yaml.replace(".yaml", "")
+            cfg = hydra.compose(config_name=cfg_name)
 
-            logger.info(f"Trying to load {config_name} on {device} with checkpoint {checkpoint_path}")
-            progress(0.3, desc=f"Loading {config_name}...")
-          
             from sam2.build_sam import build_sam2
             from sam2.sam2_image_predictor import SAM2ImagePredictor
 
-            sam2_model = build_sam2(config_name, checkpoint_path)
+            sam_logger.info(f"Trying to load {config_name_with_yaml} on {device} with checkpoint {checkpoint_path}")
+            _maybe_progress(0.3, f"Loading {config_name_with_yaml}...")
+
+            # IMPORTANT: pass cfg (not the string)
+            sam2_model = build_sam2(cfg, checkpoint_path)
             sam2_model.to(device)
             predictor = SAM2ImagePredictor(sam2_model)
-            logger.info(f"✅ Loaded {config_name} successfully on {device}")
+            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}: {e}\nTraceback: {traceback.format_exc()}"
+            error_msg = f"Failed to load {config_name_with_yaml}: {e}\nTraceback: {traceback.format_exc()}"
             tried.append(error_msg)
-            logger.warning(error_msg)
+            sam_logger.warning(error_msg)
             return None
 
     predictor = try_load("sam2_hiera_large.yaml", "sam2_hiera_large.pt")
-    # if predictor is None:
-    #     logger.warning("Could not load large model, falling back to tiny model.")
-    #     predictor = try_load("sam2_hiera_tiny.yaml", "sam2_hiera_tiny.pt")
-    #     if predictor:
-    #         logger.warning("⚠️ Using Tiny model as fallback (less accurate, but faster and lighter).")
 
     if predictor is None:
-        error_message = "SAM2 loading failed for large model. Reasons: \n" + "\n".join(tried)
-        logger.error(f"❌ {error_message}")
+        error_message = "SAM2 loading failed for large model. Reasons:\n" + "\n".join(tried)
+        sam_logger.error(f"❌ {error_message}")
         raise gr.Error(error_message)
 
     return predictor
 
-# -------------------------------------------------------------------
+# ============================================================================ #
+# MatAnyOne LOADER (robust cfg + net + core)
+# ============================================================================ #
+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
+    """
+    from omegaconf import OmegaConf
 
-# Global variables for models (lazy loading)
+    ma_logger = logging.getLogger("MatAnyOneLoader")
+
+    # Try to locate a config file; otherwise minimal default
+    cfg_path_candidates = [
+        "Configs/matanyone.yaml",
+        "configs/matanyone.yaml",
+    ]
+    cfg = None
+    for p in cfg_path_candidates:
+        if os.path.exists(p):
+            ma_logger.info(f"Loading MatAnyOne cfg: {p}")
+            cfg = OmegaConf.load(p)
+            break
+    if cfg is None:
+        ma_logger.warning("No MatAnyOne cfg found, using minimal defaults.")
+        cfg = OmegaConf.create({
+            "model": {"backbone": "swinB"},
+            "inference": {"amp": True},
+            "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)
+        net.to(device)
+        core = InferenceCore(net, cfg)
+        ma_logger.info("✅ MatAnyOne loaded (layout A)")
+        return core
+    except Exception as e:
+        last_err = e
+        ma_logger.warning(f"Layout A failed: {e}")
+
+    # Layout B (single package exposing classes)
+    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
+    except Exception as e:
+        last_err = e
+        ma_logger.warning(f"Layout B failed: {e}")
+
+    raise RuntimeError(f"Failed to initialize MatAnyOne/InferenceCore. Last error: {last_err}")
+
+# ============================================================================ #
+# GLOBALS & MODEL SETUP
+# ============================================================================ #
 sam2_predictor = None
 matanyone_model = None
 models_loaded = False
 loading_lock = threading.Lock()
 
-# ------- Robust download_and_setup_models() using above loader --------
-
-def download_and_setup_models(progress=gr.Progress()):
+def download_and_setup_models(progress: Optional[gr.Progress] = None):
     """
-    Download and setup models (SAM2 and MatAnyone), robust to Hugging Face Spaces and local dev.
-    Uses local YAML config, falls back to Tiny if Large can't be loaded.
+    Download and setup models (SAM2 and MatAnyOne), robust to HF Spaces and local dev.
     """
     global sam2_predictor, matanyone_model, models_loaded
 
@@ -166,77 +238,79 @@ def download_and_setup_models(progress=gr.Progress()):
         try:
             logger.info("🔄 Starting ENHANCED model loading with fallback...")
 
-            # --- Load SAM2 ---
             device = "cuda" if torch.cuda.is_available() else "cpu"
-            sam2_predictor_local = load_sam2_predictor(device, progress)
-            sam2_predictor = sam2_predictor_local
 
-            # --- Load MatAnyone (your original robust loader logic) ---
-            matanyone_model_local = None
-            matanyone_loaded = False
+            # --- Load SAM2 ---
+            local_sam2 = load_sam2_predictor(device=device, progress=progress)
+            # keep global
+            sam2_predictor = local_sam2
+
+            # --- Load MatAnyOne ---
             try:
-                from huggingface_hub import hf_hub_download
-                from matanyone import InferenceCore
-                matanyone_model_local = InferenceCore(
-                    "PeiqingYang/MatAnyone-v1.0",
-                    device=device,
-                    cfg={}
-                )
-                matanyone_loaded = True
-                logger.info("✅ MatAnyone loaded via HuggingFace Hub")
+                local_matanyone = load_matanyone(device)
+                matanyone_model = local_matanyone
+                logger.info("✅ MatAnyOne loaded")
             except Exception as e:
-                logger.warning(f"❌ MatAnyone load failed: {e}")
-
-            if not matanyone_loaded:
+                logger.warning(f"❌ MatAnyOne load failed: {e}")
                 raise RuntimeError("MatAnyone model could not be loaded.")
 
-            matanyone_model = matanyone_model_local
-
             models_loaded = True
             logger.info("--- ✅ All models loaded successfully ---")
-            return "✅ SAM2 + MatAnyone loaded successfully!"
+            return "✅ SAM2 + MatAnyOne loaded successfully!"
         except Exception as e:
             logger.error(f"❌ Enhanced loading failed: {str(e)}")
             logger.error(f"Full traceback: {traceback.format_exc()}")
             return f"❌ Enhanced loading failed: {str(e)}"
-# =======================================================================
-# [START REST OF YOUR MAIN APP, UNCHANGED]
-# =======================================================================
 
-def process_video_hq(video_path, background_choice, custom_background_path, progress=gr.Progress()):
+# ============================================================================ #
+# TWO-STAGE PROCESSING PIPELINE
+# ============================================================================ #
+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(0, desc="🎬 Initializing TWO-STAGE processing...")
-        
+        _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}"
-        
+
         cap = cv2.VideoCapture(video_path)
         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
         background = None
         background_name = ""
-        
+
         if background_choice == "custom" and custom_background_path:
             try:
                 background = cv2.imread(custom_background_path)
@@ -258,125 +332,125 @@ def process_video_hq(video_path, background_choice, custom_background_path, prog
                     return None, f"❌ Error creating background: {str(e)}"
             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)
-        progress(0.1, desc="🟢 STAGE 1: Creating green screen version...")
+        _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
-                progress(progress_pct, desc=f"🟢 Green screen frame {frame_count + 1}/{total_frames}")
-                
+                _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)
-                
+
                 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
-        
+
         greenscreen_writer.release()
         cap.release()
-        
+
         # STAGE 2: Replace green screen with final background (Green Screen → Final)
-        progress(0.5, desc=f"🎨 STAGE 2: Replacing green screen with {background_name}...")
-        
+        _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
-                progress(progress_pct, desc=f"🎬 Final compositing frame {frame_count + 1}/{total_frames}")
-                
+                _prog(progress_pct, f"🎬 Final compositing frame {frame_count + 1}/{total_frames}")
+
                 # Detect green screen with wider detection range
                 hsv = cv2.cvtColor(green_frame, cv2.COLOR_BGR2HSV)
-                lower_green = np.array([25, 30, 30])  # More sensitive (default: [35,40,40])
-                upper_green = np.array([100, 255, 255])  # Wider range (default: [85,255,255])
+                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)
+                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)
                 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
-        
+
         greenscreen_cap.release()
         final_writer.release()
-        
+
         # Cleanup intermediate green screen file
         try:
             os.remove(greenscreen_path)
-        except:
+        except Exception:
             pass
-        
+
         if frame_count == 0:
             return None, "❌ No frames were processed successfully."
-        
-        progress(0.9, desc="🎵 Adding high-quality audio...")
-        
+
+        _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}" '
@@ -384,13 +458,12 @@ def process_video_hq(video_path, background_choice, custom_background_path, prog
                 f'-c:a aac -b:a 192k -ac 2 -ar 48000 '
                 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:
@@ -398,50 +471,53 @@ def process_video_hq(video_path, background_choice, custom_background_path, prog
             except Exception as e2:
                 logger.error(f"Failed to copy video file: {e2}")
                 return None, f"❌ Failed to finalize video: {str(e2)}"
-        
+
         # Save to MyAvatar/My Videos directory
         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)
-        except:
+        except Exception:
             pass
-        
-        progress(1.0, desc="✅ TWO-STAGE processing complete!")
-        
+
+        _prog(1.0, "✅ TWO-STAGE processing complete!")
+
         success_message = (
             f"✅ TWO-STAGE Success!\n"
             f"🟢 Stage 1: Original → Green Screen\n"
             f"🎬 Stage 2: Green Screen → {background_name}\n"
-            f"📊 Processed: {frame_count}/{total_frames} frames\n"
+            f"📊 Processed: {frame_count} frames\n"
             f"📁 Saved: MyAvatar/My Videos/{saved_filename}\n"
-            f"🎯 Quality: Cinema-grade with SAM2 + MatAnyone\n"
+            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
 
+# ============================================================================ #
+# 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:
@@ -449,36 +525,32 @@ def extract_video_path(v):
         return v
 
     with gr.Blocks(
-        title="ENHANCED High-Quality Video Background Replacement", 
+        title="ENHANCED High-Quality Video Background Replacement",
         theme=gr.themes.Soft(),
         css="""
-        .gradio-container {
-            max-width: 1200px !important;
-        }
-        .progress-bar {
-            background: linear-gradient(90deg, #3498db, #2ecc71) !important;
-        }
+        .gradio-container { max-width: 1200px !important; }
+        .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("*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", 
+                    label="🎥 Drop your video here",
                     height=300
                 )
 
-                # ================= ADD VIDEO PREVIEW BLOCK =================
+                # Video preview
                 video_preview = gr.Video(
                     label="📺 Preview of Uploaded Video",
                     height=200,
@@ -489,23 +561,24 @@ def extract_video_path(v):
                     inputs=video_input,
                     outputs=video_preview
                 )
-                # ================= END VIDEO PREVIEW BLOCK =================
 
                 gr.Markdown("### 🎨 Step 2: Choose Background Method")
                 gr.Markdown("*Select your preferred background creation method*")
-                
-                # 4 Background Methods
+
+                # FIXED Radio (flat choices)
                 background_method = gr.Radio(
-                    choices=[
-                        ("A) 📷 Upload Image", "upload"),
-                        ("B) 🎨 Professional Presets", "professional"), 
-                        ("C) 🌈 Colors/Gradients", "colors"),
-                        ("D) 🤖 AI Generated", "ai")
-                    ],
+                    choices=["upload", "professional", "colors", "ai"],
                     value="professional",
                     label="Background Method"
                 )
-                
+                # Labels hint
+                gr.Markdown(
+                    "- **upload** = 📷 Upload Image  \n"
+                    "- **professional** = 🎨 Professional Presets  \n"
+                    "- **colors** = 🌈 Colors/Gradients  \n"
+                    "- **ai** = 🤖 AI Generated"
+                )
+
                 # Method A: Upload Image
                 with gr.Group(visible=False) as upload_group:
                     gr.Markdown("**📷 Upload Your Background Image**")
@@ -513,7 +586,7 @@ def extract_video_path(v):
                         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**")
@@ -522,114 +595,110 @@ def extract_video_path(v):
                         value="office_modern",
                         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 == "professional")), # professional_group
                         gr.update(visible=(method == "colors")),       # colors_group
                         gr.update(visible=(method == "ai"))            # ai_group
                     )
-                
+
                 background_method.change(
                     fn=switch_background_method,
                     inputs=background_method,
                     outputs=[upload_group, professional_group, colors_group, ai_group]
                 )
-                
+
                 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", 
-                        size="lg"
+                        "🚀 Step 1: Load AI Models",
+                        variant="secondary",
                     )
                     process_btn = gr.Button(
-                        "✨ Step 2: Process Video", 
-                        variant="primary", 
-                        size="lg"
+                        "✨ Step 2: Process Video",
+                        variant="primary",
                     )
-                
+
                 # System status
                 status_text = gr.Textbox(
-                    label="🔧 System Status", 
-                    value=get_model_status(), 
+                    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", 
+                    label="🎬 Your Processed Video",
                     height=400
                 )
-                
+
                 result_text = gr.Textbox(
-                    label="📊 Processing Results", 
+                    label="📊 Processing Results",
                     interactive=False,
                     lines=6,
                     placeholder="Processing status and results will appear here..."
                 )
-                
+
                 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]
-                    
                     bg_preview_html += f"""
                     <div style='
                         padding: 12px 8px; 
@@ -648,23 +717,20 @@ def switch_background_method(method):
                         </div>
                     </div>
                     """
-                
                 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:
-                    # Save generated image
-                    import tempfile
                     with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
                         cv2.imwrite(tmp.name, bg_image)
                         return tmp.name, f"✅ Background generated: {prompt[:50]}..."
@@ -673,90 +739,89 @@ def generate_ai_background(prompt, style):
             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, 
-                                 progress=gr.Progress()):
+        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,
+            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:
-                progress(0, desc="🎬 Preparing background...")
-                
-                # Determine which background to use based on method
                 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."
-                
+
                 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}"
-                
+
                 elif bg_method == "colors":
-                    # Create custom gradient as temporary image
                     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,
+                            "colors": colors if grad_type != "solid" else [colors[0]],
                             "direction": grad_type if grad_type != "solid" else "vertical"
                         }
-                        
-                        if grad_type == "solid":
-                            bg_config["colors"] = [colors[0]]
-                        
-                        # Create temporary image for gradient
+
                         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."
-                
+
                 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)}"
-        
-        # Connect all the functions
+
+        # 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]
         )
-        
+
         process_btn.click(
             fn=process_video_enhanced,
             inputs=[
                 video_input,           # video_path
-                background_method,     # bg_method  
+                background_method,     # bg_method
                 custom_background,     # custom_img
                 professional_choice,   # prof_choice
                 gradient_type,         # grad_type
@@ -765,92 +830,58 @@ def process_video_enhanced(video_path, bg_method, custom_img, prof_choice, grad_
             ],
             outputs=[video_output, result_text]
         )
-        
-        # Comprehensive info section
+
+        # Info
         with gr.Accordion("ℹ️ ENHANCED Quality & Features", open=False):
             gr.Markdown("""
             ### 🏆 TWO-STAGE Cinema-Quality Features:
-            
-            **🎬 Two-Stage Processing:**
-            - **Stage 1**: Original Video → Green Screen Video (SAM2 + MatAnyone segmentation)
-            - **Stage 2**: Green Screen Video → Final Background (Professional chroma key replacement)
-            - **Why Two-Stage?**: Better edge quality, cleaner separation, professional results
-            
-            **🤖 Advanced AI Models:**
-            - **SAM2**: State-of-the-art segmentation (Large/Tiny auto-selection)
-            - **MatAnyone**: CVPR 2025 professional matting technology
-            - **Multi-Fallback Loading**: 4+ methods each for maximum reliability
-            - **OpenCV Fallbacks**: Enhanced backup systems for compatibility
-            
-            **🎨 4 Background Methods:**
-            - **A) Upload Image**: Use any custom image as background
-            - **B) Professional Presets**: 15+ high-quality professional backgrounds
-            - **C) Colors/Gradients**: Custom color combinations with 6 gradient types
-            - **D) AI Generated**: Procedural backgrounds from text prompts
-            
-            **🎬 Professional Quality:**
-            - **✨ Edge Feathering**: Smooth, natural transitions
-            - **🎬 Gamma Correction**: Professional color compositing
-            - **🔍 Multi-Point Segmentation**: 7-point strategic person detection
-            - **🧹 Morphological Processing**: Advanced mask cleanup
-            - **🟢 Green Screen Intermediate**: Professional chroma key workflow
-            
-            **🎵 Audio & Video:**
-            - **High-Quality Audio**: 192kbps AAC preservation
-            - **📺 H.264 Codec**: CRF 18 for broadcast quality
-            - **🎞️ Frame Processing**: Advanced error handling
-            - **💾 Smart Caching**: Optimized memory management
-            
-            ### 💡 Usage Tips:
-            - Upload videos in common formats (MP4, MOV, AVI)
-            - For best results, ensure good lighting in original video
-            - Custom backgrounds work best with high resolution images
-            - AI prompts: Try "modern office", "sunset mountain", "abstract tech"
-            - GPU processing is faster but CPU fallback always available
-            - Two-stage processing gives cinema-quality results
+            **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.
             """)
-        
-        # Footer
+
         gr.Markdown("---")
-        gr.Markdown(
-            "*🎬 Cinema-Quality Video Background Replacement - "
-            "Enhanced with TWO-STAGE processing and 4-method background system*"
-        )
-    
+        gr.Markdown("*🎬 Cinema-Quality Video Background Replacement — TWO-STAGE pipeline*")
+
     return demo
 
+# ============================================================================ #
+# MAIN
+# ============================================================================ #
 def main():
     """Main application entry point"""
     try:
         print("🎬 Cinema-Quality Video Background Replacement")
         print("=" * 50)
-        
-        # Initialize application
+
+        # Initialize application paths
         os.makedirs("/tmp/MyAvatar/My_Videos/", exist_ok=True)
         os.makedirs(os.path.expanduser("~/.cache/sam2"), exist_ok=True)
-        
+
         print("🚀 Features:")
-        print("  • SAM2 + MatAnyone AI models")
+        print("  • SAM2 + MatAnyOne AI models")
         print("  • TWO-STAGE processing (Original → Green Screen → Final)")
         print("  • 4 background methods (Upload/Professional/Colors/AI)")
         print("  • Multi-fallback loading system")
         print("  • Cinema-quality processing")
         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
         )
-        
+
     except KeyboardInterrupt:
         logger.info("🛑 Application stopped by user")
         print("\n🛑 Application stopped by user")