Update app.py

app.py

@@ -1,47 +1,43 @@
 #!/usr/bin/env python3
 """
 BackgroundFX Pro - CSP-Safe Application Entry Point
-Built for Hugging Face Spaces (strict Content Security Policy)
-No inline JavaScript, CSP-compliant Gradio, safe environment vars, fallback AI models
+Now with: live background preview + sources: Preset / Upload / Gradient / AI Generate
 """
 
-import early_env  # <<< must be FIRST to sanitize threading/env before anything else
+import early_env  # <<< must be FIRST
 
-import os
+import os, time, tempfile
 from pathlib import Path
-from typing import Optional, Dict, Any, Callable
+from typing import Optional, Dict, Any, Callable, Tuple
 
-# 1️⃣ Set CSP-safe environment variables BEFORE any imports (Gradio will see these)
+# 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️⃣ Patch Gradio schema early for Hugging Face bug compatibility
+# 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"
+            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  # No fatal error if Gradio patch fails
+    pass
 
-# 3️⃣ Initialize logging EARLY (before importing modules that emit logs)
+# 3) Logging early
 from utils.logging_setup import setup_logging, make_logger
-setup_logging(app_name="backgroundfx")  # LOG_LEVEL env var respected; defaults to DEBUG in dev
+setup_logging(app_name="backgroundfx")
 logger = make_logger("entrypoint")
 logger.info("Entrypoint starting…")
 
-# 4️⃣ Import your modular code (assuming your project structure)
+# 4) Imports
 from core.exceptions import ModelLoadingError, VideoProcessingError
 from config.app_config import get_config
 from utils.hardware.device_manager import DeviceManager
@@ -50,10 +46,39 @@ def patched_get_type(schema):
 from processing.video.video_processor import CoreVideoProcessor, ProcessorConfig
 from processing.audio.audio_processor import AudioProcessor
 
-# NOTE: We import background presets (for the UI) and validator from utils
-from utils import PROFESSIONAL_BACKGROUNDS, validate_video_file
-
-# 5️⃣ CSP-safe fallback model stubs
+# Background helpers
+from utils import PROFESSIONAL_BACKGROUNDS, validate_video_file, create_professional_background
+# Gradient helper (add this 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))
+        # build vertical then rotate
+        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)
+        # rotate to angle
+        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))
@@ -66,7 +91,6 @@ def predict(self, point_coords=None, point_labels=None, box=None, multimask_outp
 class CSPSafeMatAnyone:
     def step(self, image_tensor, mask_tensor=None, objects=None, first_frame_pred=False, **kwargs):
         import torch
-        # image_tensor can be CHW or NCHW; our model guard normalizes it upstream
         if hasattr(image_tensor, "shape"):
             if len(image_tensor.shape) == 3:
                 _, H, W = image_tensor.shape
@@ -82,7 +106,25 @@ def output_prob_to_mask(self, output_prob):
     def process(self, image, mask, **kwargs):
         return mask
 
-# 6️⃣ Application main processor object
+# ---------- helpers for UI ----------
+import numpy as np
+import cv2
+from PIL import Image
+
+PREVIEW_W, PREVIEW_H = 640, 360  # 16:9
+
+def _hex_to_rgb(x: str) -> Tuple[int,int,int]:
+    x = x.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()
@@ -105,21 +147,18 @@ def load_models(self, progress_callback: Optional[Callable]=None) -> str:
         sam2_model = getattr(sam2, "model", sam2) if sam2 else CSPSafeSAM2()
         matanyone_model = getattr(matanyone, "model", matanyone) if matanyone else CSPSafeMatAnyone()
 
-        # ⬇️ NEW: fast-but-safe defaults (NVENC + model-only downscale)
         cfg = ProcessorConfig(
-            background_preset="office",   # valid preset key
-            write_fps=None,               # keep source FPS
-            max_model_size=1280,          # model-only downscale; output remains full-res
-            use_nvenc=True,               # try GPU encoder if available
-            nvenc_codec="h264",           # browser-safe preview
-            nvenc_preset="p5",            # HQ preset
-            nvenc_cq=18,                  # lower = higher quality
-            nvenc_tune_hq=True,           # high-quality tuning
-            nvenc_pix_fmt="yuv420p",      # web-compatible pixel format
+            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",
         )
         self.core_processor = CoreVideoProcessor(config=cfg, models=None)
-
-        # Minimal adapter the processor expects
         self.core_processor.models = type('FakeModelManager', (), {
             'get_sam2': lambda self_: sam2_model,
             'get_matanyone': lambda self_: matanyone_model
@@ -130,31 +169,103 @@ def load_models(self, progress_callback: Optional[Callable]=None) -> str:
                     type(sam2_model).__name__, type(matanyone_model).__name__)
         return "Models loaded (CSP-safe; fallbacks in use if actual AI models failed)."
 
-    def process_video(self, video, bg_style, custom_bg_file):
+    # ---- 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.lower(),  # "linear" or "radial" (linear in fallback)
+            "start": _hex_to_rgb(color1),
+            "end": _hex_to_rgb(color2),
+            "angle_deg": float(angle),
+        }
+        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)
+            pipe = pipe.to(device)
+            g = torch.Generator(device=device).manual_seed(int(seed)) if seed is not None else None
+            with torch.autocast(device if device=="cuda" else "cpu"):
+                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, bg_style=%s, custom_bg=%s)",
-                    video, bg_style, getattr(custom_bg_file, "name", None) if custom_bg_file else None)
+        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)
 
-        import time
         output_path = f"/tmp/output_{int(time.time())}.mp4"
 
-        # Background config passed to the processor:
-        # - custom image via {"custom_path": "..."}
-        # - preset via {"background_choice": "<key>"}
-        if custom_bg_file:
-            bg_cfg = {"custom_path": custom_bg_file.name}
-        else:
-            style = bg_style if (bg_style in PROFESSIONAL_BACKGROUNDS) else "office"
-            bg_cfg = {"background_choice": style}
-
-        # Validate input video (utils.validate_video_file returns bool)
+        # 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
+        bg_cfg: Dict[str, Any]
+        src = (bg_source or "Preset").lower()
+        if src == "upload" and custom_bg_file is not None:
+            bg_cfg = {"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,
@@ -167,13 +278,13 @@ def process_video(self, video, bg_style, custom_bg_file):
             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, style={bg_style})"
+            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 interface CSP-safe
+# 7) Gradio UI
 def create_csp_safe_gradio():
     import gradio as gr
     app = VideoBackgroundApp()
@@ -182,46 +293,116 @@ def create_csp_safe_gradio():
         title="BackgroundFX Pro - CSP Safe",
         analytics_enabled=False,
         css="""
-        .gradio-container { max-width: 1000px; margin: auto; }
+        .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. Built for Hugging Face Spaces CSP.")
+        gr.Markdown("Replace your video background with cinema-quality AI matting. Now with live background preview.")
 
         with gr.Row():
-            with gr.Column():
+            with gr.Column(scale=1):
                 video = gr.Video(label="Upload Video")
-                # Ensure default choice exists in PROFESSIONAL_BACKGROUNDS (use 'office')
-                choices = list(PROFESSIONAL_BACKGROUNDS.keys())
-                default_choice = "office" if "office" in choices else (choices[0] if choices else "office")
-                bg_style = gr.Dropdown(
-                    choices=choices,
-                    value=default_choice,
-                    label="Background Style"
+                bg_source = gr.Radio(
+                    ["Preset", "Upload", "Gradient", "AI Generate"],
+                    value="Preset",
+                    label="Background Source",
+                    interactive=True,
                 )
-                custom_bg = gr.File(label="Custom Background (Optional)", file_types=["image"])
+
+                # PRESET
+                preset_choices = list(PROFESSIONAL_BACKGROUNDS.keys())
+                preset_key = gr.Dropdown(choices=preset_choices, value=("office" if "office" in preset_choices else preset_choices[0]), 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():
+            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_change(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_change,
+            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],
+        )
+
+        # live previews
+        def preview_from_preset(key):
+            return app.preview_preset(key)
+        preset_key.change(fn=preview_from_preset, inputs=[preset_key], outputs=[bg_preview])
+
+        def preview_from_upload(file):
+            return app.preview_upload(file)
+        custom_bg.change(fn=preview_from_upload, inputs=[custom_bg], outputs=[bg_preview])
+
+        def preview_from_gradient(gt, c1, c2, ang):
+            return app.preview_gradient(gt, c1, c2, ang)
+        for comp in (grad_type, grad_color1, grad_color2, grad_angle):
+            comp.change(fn=preview_from_gradient, 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
-
-        def safe_process(vid, style, custom_bg_file):
-            logger.info("UI: starting processing (style=%s)", style)
-            return app.process_video(vid, style, custom_bg_file)
-
-        btn_load.click(fn=safe_load, outputs=[status])
-        btn_run.click(fn=safe_process, inputs=[video, bg_style, custom_bg], outputs=[out_video, status])
+            # set initial preview (preset default)
+            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]
+        )
 
     return demo
 
-# 8️⃣ Main entry point (for Hugging Face Spaces)
+# 8) Launch
 if __name__ == "__main__":
     logger.info("Launching CSP-safe Gradio interface for Hugging Face Spaces")
     demo = create_csp_safe_gradio()