Update utils/__init__.py

utils/__init__.py

@@ -1,437 +1,431 @@
+#!/usr/bin/env python3
 """
-Complete utils/__init__.py with all required functions
-Device-safe, SAM2↔MatAnyOne interop, and compositing helpers.
+BackgroundFX Pro - CSP-Safe Application Entry Point
+Now with: live background preview + sources: Preset / Upload / Gradient / AI Generate
 """
 
-from __future__ import annotations
-
-import os
-import logging
-import tempfile
-from typing import Optional, Tuple, Dict, Any, List, Iterable, Callable
-
-import cv2
+import early_env  # <<< must be FIRST
+
+import os, time
+from typing import Optional, Dict, Any, Callable, Tuple
+
+# 1) CSP-safe Gradio env
+os.environ['GRADIO_ALLOW_FLAGGING'] = 'never'
+os.environ['GRADIO_ANALYTICS_ENABLED'] = 'False'
+os.environ['GRADIO_SERVER_NAME'] = '0.0.0.0'
+os.environ['GRADIO_SERVER_PORT'] = '7860'
+
+# 2) Gradio schema patch
+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
+except Exception:
+    pass
+
+# 3) Logging early
+from utils.logging_setup import setup_logging, make_logger
+setup_logging(app_name="backgroundfx")
+logger = make_logger("entrypoint")
+logger.info("Entrypoint starting…")
+
+# 4) Imports
+from core.exceptions import ModelLoadingError, VideoProcessingError
+from config.app_config import get_config
+from utils.hardware.device_manager import DeviceManager
+from utils.system.memory_manager import MemoryManager
+from models.loaders.model_loader import ModelLoader
+from processing.video.video_processor import CoreVideoProcessor, ProcessorConfig
+from processing.audio.audio_processor import AudioProcessor
+
+# Background helpers
+from utils import PROFESSIONAL_BACKGROUNDS, validate_video_file, create_professional_background
+# Gradient helper (add to utils; fallback here for preview only if missing)
+try:
+    from utils import create_gradient_background
+except Exception:
+    def create_gradient_background(spec: Dict[str, Any], width: int, height: int):
+        # Lightweight fallback preview (linear only)
+        import numpy as np
+        import cv2
+        def _to_rgb(c):
+            if isinstance(c, (list, tuple)) and len(c) == 3:
+                return tuple(int(x) for x in c)
+            if isinstance(c, str) and c.startswith("#") and len(c) == 7:
+                return tuple(int(c[i:i+2], 16) for i in (1,3,5))
+            return (255, 255, 255)
+        start = _to_rgb(spec.get("start", "#222222"))
+        end   = _to_rgb(spec.get("end", "#888888"))
+        angle = float(spec.get("angle_deg", 0))
+        bg = np.zeros((height, width, 3), np.uint8)
+        for y in range(height):
+            t = y / max(1, height - 1)
+            r = int(start[0] * (1 - t) + end[0] * t)
+            g = int(start[1] * (1 - t) + end[1] * t)
+            b = int(start[2] * (1 - t) + end[2] * t)
+            bg[y, :] = (r, g, b)
+        center = (width / 2, height / 2)
+        rot = cv2.getRotationMatrix2D(center, angle, 1.0)
+        return cv2.warpAffine(bg, rot, (width, height), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT_101)
+
+# 5) CSP-safe fallbacks for models
+class CSPSafeSAM2:
+    def set_image(self, image):
+        self.shape = getattr(image, 'shape', (512, 512, 3))
+    def predict(self, point_coords=None, point_labels=None, box=None, multimask_output=True, **kwargs):
+        import numpy as np
+        h, w = self.shape[:2] if hasattr(self, 'shape') else (512, 512)
+        n = 3 if multimask_output else 1
+        return np.ones((n, h, w), dtype=bool), np.array([0.9, 0.8, 0.7][:n]), np.ones((n, h, w), dtype=np.float32)
+
+class CSPSafeMatAnyone:
+    def step(self, image_tensor, mask_tensor=None, objects=None, first_frame_pred=False, **kwargs):
+        import torch
+        if hasattr(image_tensor, "shape"):
+            if len(image_tensor.shape) == 3:
+                _, H, W = image_tensor.shape
+            elif len(image_tensor.shape) == 4:
+                _, _, H, W = image_tensor.shape
+            else:
+                H, W = 256, 256
+        else:
+            H, W = 256, 256
+        return torch.ones((1, 1, H, W))
+    def output_prob_to_mask(self, output_prob):
+        return (output_prob > 0.5).float()
+    def process(self, image, mask, **kwargs):
+        return mask
+
+# ---------- helpers for UI ----------
 import numpy as np
+import cv2
 from PIL import Image
-import torch
-
-# NEW: interop + bridge imports (add these files from the previous steps)
-from utils.interop import ensure_image_nchw, ensure_mask_for_matanyone, log_shape
-from utils.mask_bridge import sam2_to_matanyone_mask
-
-logger = logging.getLogger(__name__)
-
-# Professional backgrounds configuration
-PROFESSIONAL_BACKGROUNDS = {
-    "office": {"color": (240, 248, 255), "gradient": True},
-    "studio": {"color": (32, 32, 32), "gradient": False},
-    "nature": {"color": (34, 139, 34), "gradient": True},
-    "abstract": {"color": (75, 0, 130), "gradient": True},
-    "white": {"color": (255, 255, 255), "gradient": False},
-    "black": {"color": (0, 0, 0), "gradient": False},
-}
-
-# -------------------------------
-# Utility: device
-# -------------------------------
-def _default_device() -> str:
-    return "cuda" if torch.cuda.is_available() else "cpu"
-
-
-# -------------------------------
-# Video validation
-# -------------------------------
-def validate_video_file(video_path: str) -> bool:
-    """Validate if video file is readable"""
-    try:
-        if not os.path.exists(video_path):
-            return False
-        cap = cv2.VideoCapture(video_path)
-        if not cap.isOpened():
-            return False
-        ret, frame = cap.read()
-        cap.release()
-        return ret and frame is not None
-    except Exception as e:
-        logger.error(f"Video validation failed: {e}")
-        return False
-
-
-# -------------------------------
-# SAM2 person segmentation (first-frame bootstrapping)
-# -------------------------------
-def segment_person_hq(
-    frame_rgb: np.ndarray,
-    *,
-    use_sam2: bool = True,
-    sam2_predictor: Any = None,  # prefer injecting a ready predictor (from your ModelLoader)
-) -> Optional[np.ndarray]:
-    """
-    High-quality person segmentation for a single RGB frame.
-    Returns a float mask HxW in [0,1], or None on failure.
-
-    Preferred path: pass a ready-made SAM2 predictor (e.g., SAM2ImagePredictor).
-    Fallback path: simple color-based segmentation.
-    """
-    try:
-        if use_sam2 and sam2_predictor is not None:
-            try:
-                # SAM2 official predictors accept RGB np.uint8; set + predict.
-                # We use a simple center-point prompt; adapt to your UX if needed.
-                if hasattr(sam2_predictor, "set_image"):
-                    sam2_predictor.set_image(frame_rgb)
-
-                h, w = frame_rgb.shape[:2]
-                center_point = np.array([[w // 2, h // 2]])
-                center_label = np.array([1])
-
-                # Try the SAM2 "predict" API (Meta’s predictor style)
-                if hasattr(sam2_predictor, "predict"):
-                    out = sam2_predictor.predict(
-                        point_coords=center_point,
-                        point_labels=center_label,
-                        multimask_output=True,
-                    )
-                    # Known Meta API returns (masks, scores, logits) as numpy
-                    if isinstance(out, (list, tuple)) and len(out) >= 1:
-                        masks = out[0]
-                        if masks is None or len(masks) == 0:
-                            return None
-                        # masks: (M,H,W); pick best by area
-                        areas = masks.reshape(masks.shape[0], -1).sum(axis=1)
-                        best = int(np.argmax(areas))
-                        m = masks[best].astype(np.float32)
-                        m = (m >= 0.5).astype(np.float32)
-                        return m
-
-                # Some wrappers expose processor/post_process; if you use that, call separately
-                logger.warning("SAM2 predictor provided but unknown API; falling back to simple segmentation")
-            except Exception as e:
-                logger.warning(f"SAM2 segmentation failed: {e}; falling back to simple method")
-
-        # Fallback: color-based person segmentation
-        return _simple_person_segmentation(frame_rgb)
-    except Exception as e:
-        logger.error(f"Person segmentation failed: {e}")
-        return None
-
-
-def _simple_person_segmentation(frame_rgb: np.ndarray) -> np.ndarray:
-    """Simple person segmentation using color-based methods"""
-    hsv = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2HSV)
-    # Green screen detection
-    lower_green = np.array([40, 40, 40])
-    upper_green = np.array([80, 255, 255])
-    green_mask = cv2.inRange(hsv, lower_green, upper_green)
-    # White background detection
-    lower_white = np.array([0, 0, 200])
-    upper_white = np.array([180, 30, 255])
-    white_mask = cv2.inRange(hsv, lower_white, upper_white)
-    # Combine + invert to person
-    bg_mask = cv2.bitwise_or(green_mask, white_mask)
-    person_mask = cv2.bitwise_not(bg_mask)
-    # Morph clean
-    kernel = np.ones((5, 5), np.uint8)
-    person_mask = cv2.morphologyEx(person_mask, cv2.MORPH_CLOSE, kernel)
-    person_mask = cv2.morphologyEx(person_mask, cv2.MORPH_OPEN, kernel)
-    return (person_mask.astype(np.float32) / 255.0)
-
-
-# -------------------------------
-# MatAnyOne integration (first-frame + per-frame)
-# -------------------------------
-def refine_mask_hq(
-    mask_hw_float01: np.ndarray,
-    frame_rgb: np.ndarray,
-    *,
-    use_matanyone: bool = True,
-    mat_core: Any = None,           # prefer injecting a ready InferenceCore from ModelLoader
-    first_frame: bool = True,
-    device: str | None = None,
-) -> np.ndarray:
-    """
-    High-quality mask refinement for a single frame + mask pair using MatAnyOne.
-    Returns refined mask HxW float in [0,1]. If use_matanyone=False or mat_core is None,
-    falls back to simple refinement.
-
-    NOTE: For videos, prefer using seed/refine helpers below that keep temporal memory.
-    """
-    try:
-        if not use_matanyone or mat_core is None:
-            return _simple_mask_refinement(mask_hw_float01, frame_rgb)
-
-        device = device or _default_device()
-
-        # Image → (1,3,H,W)
-        img_nchw = ensure_image_nchw(torch.from_numpy(frame_rgb).to(device), device=device, want_batched=True)
-        log_shape("refine.image_nchw", img_nchw)
-
-        # Mask → (1,H,W)
-        mask_t = torch.from_numpy(mask_hw_float01).to(device)
-        mask_c_hw = ensure_mask_for_matanyone(mask_t, idx_mask=False, threshold=0.5, keep_soft=False, device=device)
-        log_shape("refine.mask_c_hw", mask_c_hw)
-
-        # MatAnyOne step (we let the global guard in ModelLoader do additional checks)
-        pred = mat_core.step(
-            image=img_nchw[0],              # CHW
-            mask=mask_c_hw if first_frame else None,
-            idx_mask=False,
-            matting=True,
-            first_frame_pred=bool(first_frame),
-        )
 
-        # Try to decode output into an alpha HxW float mask
-        refined = _coerce_pred_to_mask(pred, device=device)
-        if refined is None:
-            # If the core doesn’t return alpha directly, fall back
-            return _simple_mask_refinement(mask_hw_float01, frame_rgb)
-
-        return refined
-    except Exception as e:
-        logger.warning(f"MatAnyOne refinement failed: {e}; using simple refinement")
-        return _simple_mask_refinement(mask_hw_float01, frame_rgb)
-
-
-def _coerce_pred_to_mask(pred: Any, device: str = "cuda") -> Optional[np.ndarray]:
-    """
-    Best-effort: extract HxW float mask from MatAnyOne output variants.
-    Supports torch.Tensor, numpy, PIL, or dict with common keys.
-    """
-    try:
-        # Dict-like: look for common keys
-        if isinstance(pred, dict):
-            for k in ("alpha", "mask", "matte", "mattes"):
-                if k in pred:
-                    v = pred[k]
-                    return _coerce_pred_to_mask(v, device=device)
-
-        # Torch tensor
-        if torch.is_tensor(pred):
-            t = pred.detach()
-            # possible shapes: (H,W), (1,H,W), (N,1,H,W)
-            if t.ndim == 4 and t.shape[1] == 1:
-                t = t[0, 0]
-            elif t.ndim == 3 and t.shape[0] == 1:
-                t = t[0]
-            t = t.float().clamp(0, 1).to("cpu").numpy()
-            if t.ndim == 2:
-                return t.astype(np.float32)
-
-        # Numpy
-        if isinstance(pred, np.ndarray):
-            a = pred
-            if a.ndim == 3 and a.shape[0] == 1:
-                a = a[0]
-            if a.ndim == 2:
-                a = a.astype(np.float32)
-                if a.max() > 1.0:
-                    a = a / 255.0
-                return np.clip(a, 0.0, 1.0)
-
-        # PIL Image
-        if isinstance(pred, Image.Image):
-            a = np.array(pred).astype(np.float32)
-            if a.ndim == 3 and a.shape[2] == 1:
-                a = a[:, :, 0]
-            if a.ndim == 2:
-                if a.max() > 1.0:
-                    a = a / 255.0
-                return np.clip(a, 0.0, 1.0)
-
-    except Exception as e:
-        logger.debug(f"_coerce_pred_to_mask fallback due to: {e}")
-    return None
-
-
-def _simple_mask_refinement(mask: np.ndarray, frame_rgb: np.ndarray) -> np.ndarray:
-    """Simple mask refinement using OpenCV operations"""
-    mask_uint8 = (np.clip(mask, 0.0, 1.0) * 255).astype(np.uint8)
-    mask_blurred = cv2.GaussianBlur(mask_uint8, (5, 5), 0)
-    mask_refined = cv2.bilateralFilter(mask_blurred, 9, 75, 75)
-    return (mask_refined.astype(np.float32) / 255.0)
-
-
-# -------------------------------
-# Two-stage video helpers (seed + propagate)
-# -------------------------------
-@torch.inference_mode()
-def seed_with_sam2_post_masks(
-    core: Any,
-    frame0_rgb: np.ndarray,                 # HxWx3 uint8 RGB
-    sam2_post_masks: torch.Tensor,          # (1,M,H,W)
-    iou_scores: Optional[torch.Tensor] = None,
-    *,
-    device: str | None = None,
-    idx_mask: bool = False,
-    threshold: float = 0.5,
-    keep_soft: bool = False,
-) -> Any:
-    """
-    Seed MatAnyOne on the first frame using SAM2 post-processed masks (preferred).
-    """
-    device = device or _default_device()
-    img0 = ensure_image_nchw(torch.from_numpy(frame0_rgb).to(device), device=device, want_batched=True)
-    log_shape("seed.image_nchw", img0)
-
-    if idx_mask:
-        m_c_hw = sam2_to_matanyone_mask(sam2_post_masks.to(device), iou_scores, threshold, "single", keep_soft=False)
-        idx_hw = ensure_mask_for_matanyone(m_c_hw, idx_mask=True, device=device, threshold=threshold)
-        log_shape("seed.idx_hw", idx_hw)
-        return core.step(
-            image=img0[0],
-            mask=idx_hw,
-            idx_mask=True,
-            matting=True,
-            first_frame_pred=True,
+PREVIEW_W, PREVIEW_H = 640, 360  # 16:9
+
+from typing import Tuple
+def _hex_to_rgb(x: str) -> Tuple[int, int, int]:
+    x = (x or "").strip()
+    if x.startswith("#") and len(x) == 7:
+        return tuple(int(x[i:i+2], 16) for i in (1, 3, 5))
+    return (255, 255, 255)
+
+def _np_to_pil(arr: np.ndarray) -> Image.Image:
+    if arr.dtype != np.uint8:
+        arr = arr.clip(0, 255).astype(np.uint8)
+    return Image.fromarray(arr)
+
+# ---------- main app ----------
+class VideoBackgroundApp:
+    def __init__(self):
+        self.config = get_config()
+        self.device_mgr = DeviceManager()
+        self.memory_mgr = MemoryManager(self.device_mgr.get_optimal_device())
+        self.model_loader = ModelLoader(self.device_mgr, self.memory_mgr)
+        self.audio_proc = AudioProcessor()
+        self.models_loaded = False
+        self.core_processor: Optional[CoreVideoProcessor] = None
+        logger.info("VideoBackgroundApp initialized (device=%s)", self.device_mgr.get_optimal_device())
+
+    def load_models(self, progress_callback: Optional[Callable] = None) -> str:
+        logger.info("Loading models (CSP-safe)…")
+        try:
+            sam2, matanyone = self.model_loader.load_all_models(progress_callback=progress_callback)
+        except Exception as e:
+            logger.warning("Model loading failed (%s) - Using CSP-safe fallbacks", e)
+            sam2, matanyone = None, None
+
+        sam2_model = getattr(sam2, "model", sam2) if sam2 else CSPSafeSAM2()
+        matanyone_model = getattr(matanyone, "model", matanyone) if matanyone else CSPSafeMatAnyone()
+
+        cfg = ProcessorConfig(
+            background_preset="office",
+            write_fps=None,
+            max_model_size=1280,
+            use_nvenc=True,
+            nvenc_codec="h264",
+            nvenc_preset="p5",
+            nvenc_cq=18,
+            nvenc_tune_hq=True,
+            nvenc_pix_fmt="yuv420p",
         )
-    else:
-        m_c_hw = sam2_to_matanyone_mask(sam2_post_masks.to(device), iou_scores, threshold, "single", keep_soft=keep_soft)
-        log_shape("seed.mask_c_hw", m_c_hw)
-        return core.step(
-            image=img0[0],
-            mask=m_c_hw,
-            idx_mask=False,
-            matting=True,
-            first_frame_pred=True,
+        self.core_processor = CoreVideoProcessor(config=cfg, models=None)
+        self.core_processor.models = type('FakeModelManager', (), {
+            'get_sam2': lambda self_: sam2_model,
+            'get_matanyone': lambda self_: matanyone_model
+        })()
+
+        self.models_loaded = True
+        logger.info("Models ready (SAM2=%s, MatAnyOne=%s)",
+                    type(sam2_model).__name__, type(matanyone_model).__name__)
+        return "Models loaded (CSP-safe; fallbacks in use if actual AI models failed)."
+
+    # ---- PREVIEWS ----
+    def preview_preset(self, preset_key: str) -> Image.Image:
+        key = preset_key if preset_key in PROFESSIONAL_BACKGROUNDS else "office"
+        bg = create_professional_background(key, PREVIEW_W, PREVIEW_H)  # RGB
+        return _np_to_pil(bg)
+
+    def preview_upload(self, file) -> Optional[Image.Image]:
+        if file is None:
+            return None
+        try:
+            img = Image.open(file.name).convert("RGB")
+            img = img.resize((PREVIEW_W, PREVIEW_H), Image.LANCZOS)
+            return img
+        except Exception as e:
+            logger.warning("Upload preview failed: %s", e)
+            return None
+
+    def preview_gradient(self, gtype: str, color1: str, color2: str, angle: int) -> Image.Image:
+        spec = {
+            "type": (gtype or "linear").lower(),  # "linear" or "radial" (linear in fallback)
+            "start": _hex_to_rgb(color1 or "#222222"),
+            "end": _hex_to_rgb(color2 or "#888888"),
+            "angle_deg": float(angle or 0),
+        }
+        bg = create_gradient_background(spec, PREVIEW_W, PREVIEW_H)
+        return _np_to_pil(bg)
+
+    def ai_generate_background(self, prompt: str, seed: int, width: int, height: int) -> Tuple[Optional[Image.Image], Optional[str], str]:
+        """
+        Try generating a background with diffusers; save to /tmp and return (img, path, status).
+        """
+        try:
+            from diffusers import StableDiffusionPipeline
+            import torch
+            model_id = os.environ.get("BGFX_T2I_MODEL", "stabilityai/stable-diffusion-2-1")
+            dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype).to(device)
+            g = torch.Generator(device=device).manual_seed(int(seed)) if seed is not None else None
+            if device == "cuda":
+                with torch.autocast("cuda"):
+                    img = pipe(prompt, height=height, width=width, guidance_scale=7.0, num_inference_steps=25, generator=g).images[0]
+            else:
+                img = pipe(prompt, height=height, width=width, guidance_scale=7.0, num_inference_steps=25, generator=g).images[0]
+            tmp_path = f"/tmp/ai_bg_{int(time.time())}.png"
+            img.save(tmp_path)
+            return img.resize((PREVIEW_W, PREVIEW_H), Image.LANCZOS), tmp_path, f"AI background generated ✓ ({os.path.basename(tmp_path)})"
+        except Exception as e:
+            logger.warning("AI generation unavailable: %s", e)
+            return None, None, f"AI generation unavailable: {e}"
+
+    # ---- PROCESS VIDEO ----
+    def process_video(
+        self,
+        video: str,
+        bg_source: str,
+        preset_key: str,
+        custom_bg_file,
+        grad_type: str,
+        grad_color1: str,
+        grad_color2: str,
+        grad_angle: int,
+        ai_bg_path: Optional[str],
+    ):
+        if not self.models_loaded:
+            return None, "Models not loaded yet"
+
+        logger.info("process_video called (video=%s, source=%s, preset=%s, file=%s, grad=%s, ai=%s)",
+                    video, bg_source, preset_key, getattr(custom_bg_file, "name", None) if custom_bg_file else None,
+                    {"type": grad_type, "c1": grad_color1, "c2": grad_color2, "angle": grad_angle},
+                    ai_bg_path)
+
+        output_path = f"/tmp/output_{int(time.time())}.mp4"
+
+        # Validate input video
+        ok = validate_video_file(video)
+        if not ok:
+            logger.warning("Invalid/unreadable video: %s", video)
+            return None, "Invalid or unreadable video file"
+
+        # Build bg_config based on source
+        src = (bg_source or "Preset").lower()
+        if src == "upload" and custom_bg_file is not None:
+            bg_cfg: Dict[str, Any] = {"custom_path": custom_bg_file.name}
+        elif src == "gradient":
+            bg_cfg = {
+                "gradient": {
+                    "type": (grad_type or "linear").lower(),
+                    "start": _hex_to_rgb(grad_color1 or "#222222"),
+                    "end": _hex_to_rgb(grad_color2 or "#888888"),
+                    "angle_deg": float(grad_angle or 0),
+                }
+            }
+        elif src == "ai generate" and ai_bg_path:
+            bg_cfg = {"custom_path": ai_bg_path}
+        else:
+            key = preset_key if preset_key in PROFESSIONAL_BACKGROUNDS else "office"
+            bg_cfg = {"background_choice": key}
+
+        try:
+            result = self.core_processor.process_video(
+                input_path=video,
+                output_path=output_path,
+                bg_config=bg_cfg
+            )
+            logger.info("Core processing done → %s", output_path)
+
+            output_with_audio = self.audio_proc.add_audio_to_video(video, output_path)
+            logger.info("Audio merged → %s", output_with_audio)
+
+            frames = (result.get('frames') if isinstance(result, dict) else None) or "n/a"
+            return output_with_audio, f"Processing complete ({frames} frames, background={bg_source})"
+
+        except Exception as e:
+            logger.exception("Processing failed")
+            return None, f"Processing failed: {e}"
+
+# 7) Gradio UI
+def create_csp_safe_gradio():
+    import gradio as gr
+    app = VideoBackgroundApp()
+
+    with gr.Blocks(
+        title="BackgroundFX Pro - CSP Safe",
+        analytics_enabled=False,
+        css="""
+        .gradio-container { max-width: 1100px; margin: auto; }
+        """
+    ) as demo:
+        gr.Markdown("# 🎬 BackgroundFX Pro (CSP-Safe)")
+        gr.Markdown("Replace your video background with cinema-quality AI matting. Now with live background preview.")
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                video = gr.Video(label="Upload Video")
+                bg_source = gr.Radio(
+                    ["Preset", "Upload", "Gradient", "AI Generate"],
+                    value="Preset",
+                    label="Background Source",
+                    interactive=True,
+                )
+
+                # PRESET
+                preset_choices = list(PROFESSIONAL_BACKGROUNDS.keys())
+                default_preset = "office" if "office" in preset_choices else (preset_choices[0] if preset_choices else "office")
+                preset_key = gr.Dropdown(choices=preset_choices, value=default_preset, label="Preset")
+
+                # UPLOAD
+                custom_bg = gr.File(label="Custom Background (Image)", file_types=["image"], visible=False)
+
+                # GRADIENT
+                grad_type = gr.Dropdown(choices=["Linear", "Radial"], value="Linear", label="Gradient Type", visible=False)
+                grad_color1 = gr.ColorPicker(value="#222222", label="Start Color", visible=False)
+                grad_color2 = gr.ColorPicker(value="#888888", label="End Color", visible=False)
+                grad_angle = gr.Slider(0, 360, value=0, step=1, label="Angle (degrees)", visible=False)
+
+                # AI
+                ai_prompt = gr.Textbox(label="AI Prompt", placeholder="e.g., sunlit modern office, soft bokeh, neutral palette", visible=False)
+                ai_seed = gr.Slider(0, 2**31-1, step=1, value=42, label="Seed", visible=False)
+                ai_size = gr.Dropdown(choices=["640x360","960x540","1280x720"], value="640x360", label="AI Image Size", visible=False)
+                ai_go = gr.Button("✨ Generate Background", visible=False, variant="secondary")
+                ai_status = gr.Markdown(visible=False)
+                ai_bg_path_state = gr.State(value=None)  # store /tmp path
+
+                btn_load = gr.Button("🔄 Load Models", variant="secondary")
+                btn_run = gr.Button("🎬 Process Video", variant="primary")
+
+            with gr.Column(scale=1):
+                status = gr.Textbox(label="Status", lines=4)
+                bg_preview = gr.Image(label="Background Preview", width=PREVIEW_W, height=PREVIEW_H, interactive=False)
+                out_video = gr.Video(label="Processed Video")
+
+        # ---------- UI wiring ----------
+
+        # background source → show/hide controls
+        def on_source_toggle(src):
+            src = (src or "Preset").lower()
+            return (
+                gr.update(visible=(src == "preset")),
+                gr.update(visible=(src == "upload")),
+                gr.update(visible=(src == "gradient")),
+                gr.update(visible=(src == "gradient")),
+                gr.update(visible=(src == "gradient")),
+                gr.update(visible=(src == "gradient")),
+                gr.update(visible=(src == "ai generate")),
+                gr.update(visible=(src == "ai generate")),
+                gr.update(visible=(src == "ai generate")),
+                gr.update(visible=(src == "ai generate")),
+                gr.update(visible=(src == "ai generate")),
+            )
+        bg_source.change(
+            fn=on_source_toggle,
+            inputs=[bg_source],
+            outputs=[preset_key, custom_bg, grad_type, grad_color1, grad_color2, grad_angle, ai_prompt, ai_seed, ai_size, ai_go, ai_status],
         )
 
+        # When source changes, also refresh preview based on visible controls
+        def on_source_preview(src, pkey, gt, c1, c2, ang):
+            src_l = (src or "Preset").lower()
+            if src_l == "preset":
+                return app.preview_preset(pkey)
+            elif src_l == "gradient":
+                return app.preview_gradient(gt, c1, c2, ang)
+            # For upload/AI we keep whatever the component change handler sets (don’t overwrite)
+            return gr.update()  # no-op
+        bg_source.change(
+            fn=on_source_preview,
+            inputs=[bg_source, preset_key, grad_type, grad_color1, grad_color2, grad_angle],
+            outputs=[bg_preview]
+        )
 
-@torch.inference_mode()
-def refine_next_frame(core: Any, frame_rgb: np.ndarray, *, device: str | None = None) -> Any:
-    """Step MatAnyOne forward on a subsequent frame (no mask; uses memory)."""
-    device = device or _default_device()
-    img = ensure_image_nchw(torch.from_numpy(frame_rgb).to(device), device=device, want_batched=True)
-    log_shape("refine.image_nchw", img)
-    return core.step(
-        image=img[0],
-        mask=None,
-        idx_mask=False,
-        matting=True,
-        first_frame_pred=False,
-    )
-
+        # live previews
+        preset_key.change(fn=lambda k: app.preview_preset(k), inputs=[preset_key], outputs=[bg_preview])
+        custom_bg.change(fn=lambda f: app.preview_upload(f), inputs=[custom_bg], outputs=[bg_preview])
+        for comp in (grad_type, grad_color1, grad_color2, grad_angle):
+            comp.change(
+                fn=lambda gt, c1, c2, ang: app.preview_gradient(gt, c1, c2, ang),
+                inputs=[grad_type, grad_color1, grad_color2, grad_angle],
+                outputs=[bg_preview],
+            )
+
+        # AI generate
+        def ai_generate(prompt, seed, size):
+            try:
+                w, h = map(int, size.split("x"))
+            except Exception:
+                w, h = PREVIEW_W, PREVIEW_H
+            img, path, msg = app.ai_generate_background(
+                prompt or "professional modern office background, neutral colors, depth of field",
+                int(seed), w, h
+            )
+            return img, (path or None), msg
+        ai_go.click(fn=ai_generate, inputs=[ai_prompt, ai_seed, ai_size], outputs=[bg_preview, ai_bg_path_state, ai_status])
+
+        # model load / run
+        def safe_load():
+            msg = app.load_models()
+            logger.info("UI: models loaded")
+            return msg, app.preview_preset(preset_key.value if hasattr(preset_key, "value") else "office")
+        btn_load.click(fn=safe_load, outputs=[status, bg_preview])
+
+        def safe_process(vid, src, pkey, file, gtype, c1, c2, ang, ai_path):
+            return app.process_video(vid, src, pkey, file, gtype, c1, c2, ang, ai_path)
+        btn_run.click(
+            fn=safe_process,
+            inputs=[video, bg_source, preset_key, custom_bg, grad_type, grad_color1, grad_color2, grad_angle, ai_bg_path_state],
+            outputs=[out_video, status]
+        )
 
-@torch.inference_mode()
-def run_two_stage_matting(
-    core: Any,
-    frames_rgb_iter: Iterable[np.ndarray],     # iterable of HxWx3 uint8 RGB
-    sam2_post_masks: torch.Tensor,             # (1,M,H,W) for the first frame
-    iou_scores: Optional[torch.Tensor] = None,
-    *,
-    device: str | None = None,
-    on_pred: Optional[Callable[[int, Any], None]] = None,
-    progress: Optional[Callable[[int, Optional[int]], None]] = None,
-    total_frames: Optional[int] = None,
-    idx_mask: bool = False,
-    threshold: float = 0.5,
-    keep_soft: bool = False,
-) -> None:
-    """
-    Convenience runner for videos:
-      - Seeds on the first frame using SAM2 post-process outputs
-      - Propagates across the rest (one frame per step)
-    """
-    device = device or _default_device()
-    it = iter(frames_rgb_iter)
-    try:
-        f0 = next(it)
-    except StopIteration:
-        return
-
-    pred0 = seed_with_sam2_post_masks(
-        core, f0, sam2_post_masks, iou_scores,
-        device=device, idx_mask=idx_mask, threshold=threshold, keep_soft=keep_soft
+    return demo
+
+# 8) Launch
+if __name__ == "__main__":
+    logger.info("Launching CSP-safe Gradio interface for Hugging Face Spaces")
+    demo = create_csp_safe_gradio()
+    demo.queue().launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True,
+        debug=False,
+        inbrowser=False
     )
-    if on_pred: on_pred(0, pred0)
-    if progress: progress(1, total_frames)
-
-    t = 1
-    for frgb in it:
-        pred = refine_next_frame(core, frgb, device=device)
-        if on_pred: on_pred(t, pred)
-        t += 1
-        if progress: progress(t, total_frames)
-
-
-# -------------------------------
-# Background replacement
-# -------------------------------
-def replace_background_hq(frame_rgb: np.ndarray, mask_hw_float01: np.ndarray, background_rgb: np.ndarray) -> np.ndarray:
-    """High-quality background replacement with proper compositing"""
-    try:
-        h, w = frame_rgb.shape[:2]
-        background_resized = cv2.resize(background_rgb, (w, h))
-
-        # Ensure mask is HxW float in [0,1]
-        if mask_hw_float01.ndim == 3:
-            mask_hw_float01 = mask_hw_float01[..., 0]
-        m = np.clip(mask_hw_float01.astype(np.float32), 0.0, 1.0)
-
-        # Feather edges lightly
-        m_uint8 = (m * 255).astype(np.uint8)
-        m_feather = cv2.GaussianBlur(m_uint8, (7, 7), 0).astype(np.float32) / 255.0
-        m3 = np.stack([m_feather] * 3, axis=-1)
-
-        result = frame_rgb.astype(np.float32) * m3 + background_resized.astype(np.float32) * (1.0 - m3)
-        return np.clip(result, 0, 255).astype(np.uint8)
-    except Exception as e:
-        logger.error(f"Background replacement failed: {e}")
-        return frame_rgb
-
-
-# -------------------------------
-# Background generators
-# -------------------------------
-def create_professional_background(bg_type: str, width: int, height: int) -> np.ndarray:
-    """Create professional background of specified type and size"""
-    try:
-        if bg_type not in PROFESSIONAL_BACKGROUNDS:
-            bg_type = "office"  # Default fallback
-
-        config = PROFESSIONAL_BACKGROUNDS[bg_type]
-        color = config["color"]
-        use_gradient = config["gradient"]
-
-        if use_gradient:
-            background = _create_gradient_background(color, width, height)
-        else:
-            background = np.full((height, width, 3), color, dtype=np.uint8)
-
-        return background
-    except Exception as e:
-        logger.error(f"Background creation failed: {e}")
-        return np.full((height, width, 3), (255, 255, 255), dtype=np.uint8)
-
-
-def _create_gradient_background(base_color: Tuple[int, int, int], width: int, height: int) -> np.ndarray:
-    """Create a vertical gradient background from base color"""
-    r, g, b = base_color
-    dark = (int(r * 0.7), int(g * 0.7), int(b * 0.7))
-    bg = np.zeros((height, width, 3), dtype=np.uint8)
-    for y in range(height):
-        t = y / max(height, 1)
-        bg[y, :] = [
-            int(dark[0] * (1 - t) + r * t),
-            int(dark[1] * (1 - t) + g * t),
-            int(dark[2] * (1 - t) + b * t),
-        ]
-    return bg
-
-
-# -------------------------------
-# Exports
-# -------------------------------
-__all__ = [
-    # segment / refine (single-frame)
-    "segment_person_hq",
-    "refine_mask_hq",
-    # video runner + steps
-    "seed_with_sam2_post_masks",
-    "refine_next_frame",
-    "run_two_stage_matting",
-    # backgrounds & utils
-    "replace_background_hq",
-    "create_professional_background",
-    "PROFESSIONAL_BACKGROUNDs" if False else "PROFESSIONAL_BACKGROUNDS",
-    "validate_video_file",
-]