Update app.py

app.py

@@ -6,10 +6,13 @@
 
 import early_env  # <<< must be FIRST
 
-import os, time, tempfile
-from pathlib import Path
+import os, time, math
 from typing import Optional, Dict, Any, Callable, Tuple
 
+# Prefer a writable cache on HF/Spaces
+os.environ.setdefault("HF_HOME", "/tmp/hf")
+os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
+
 # 1) CSP-safe Gradio env
 os.environ['GRADIO_ALLOW_FLAGGING'] = 'never'
 os.environ['GRADIO_ANALYTICS_ENABLED'] = 'False'
@@ -19,15 +22,15 @@
 # 2) Gradio schema patch
 try:
     import gradio_client.utils as gc_utils
-    orig_get_type = gc_utils.get_type
-    def patched_get_type(schema):
+    _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
+        return _orig_get_type(schema)
+    gc_utils.get_type = _patched_get_type
 except Exception:
     pass
 
@@ -48,7 +51,7 @@ def patched_get_type(schema):
 
 # 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)
+# Gradient helper (add to utils; fallback here for preview only if missing)
 try:
     from utils import create_gradient_background
 except Exception:
@@ -61,20 +64,18 @@ def _to_rgb(c):
                 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)
+            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)
+            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)
 
@@ -110,20 +111,25 @@ def process(self, image, mask, **kwargs):
 import numpy as np
 import cv2
 from PIL import Image
+from typing import Tuple
 
 PREVIEW_W, PREVIEW_H = 640, 360  # 16:9
 
-def _hex_to_rgb(x: str) -> Tuple[int,int,int]:
-    x = x.strip()
+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)
+        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)
+        arr = arr.clip(0, 255).astype(np.uint8)
     return Image.fromarray(arr)
 
+def _div8(n: int) -> int:
+    # Ensure sizes are multiples of 8 for SD/VAEs
+    return int(math.floor(max(64, n) / 8.0) * 8)
+
 # ---------- main app ----------
 class VideoBackgroundApp:
     def __init__(self):
@@ -134,9 +140,12 @@ def __init__(self):
         self.audio_proc = AudioProcessor()
         self.models_loaded = False
         self.core_processor: Optional[CoreVideoProcessor] = None
+        # Text-to-Image pipeline cache
+        self.t2i_pipe = None
+        self.t2i_model_id = None
         logger.info("VideoBackgroundApp initialized (device=%s)", self.device_mgr.get_optimal_device())
 
-    def load_models(self, progress_callback: Optional[Callable]=None) -> str:
+    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)
@@ -176,7 +185,8 @@ def preview_preset(self, preset_key: str) -> Image.Image:
         return _np_to_pil(bg)
 
     def preview_upload(self, file) -> Optional[Image.Image]:
-        if file is None: return None
+        if file is None:
+            return None
         try:
             img = Image.open(file.name).convert("RGB")
             img = img.resize((PREVIEW_W, PREVIEW_H), Image.LANCZOS)
@@ -187,35 +197,173 @@ def preview_upload(self, file) -> Optional[Image.Image]:
 
     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),
+            "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]:
+    # ---- AI BG: lazy-load + reuse pipe ----
+    def _ensure_t2i(self):
         """
-        Try generating a background with diffusers; save to /tmp and return (img, path, status).
+        Choose and load a text-to-image pipeline once, with memory-efficient settings.
+        Returns (pipe, model_id, msg)
         """
+        if self.t2i_pipe is not None:
+            return self.t2i_pipe, self.t2i_model_id, "AI generator ready"
+
         try:
-            from diffusers import StableDiffusionPipeline
             import torch
+            from diffusers import StableDiffusionPipeline, AutoPipelineForText2Image
+        except Exception as e:
+            return None, None, f"AI generation unavailable (missing deps): {e}"
+
+        # Heuristic: prefer fast/light models when VRAM is small
+        token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
+        device = "cuda" if getattr(torch, "cuda", None) and torch.cuda.is_available() else "cpu"
+
+        vram_gb = None
+        try:
+            vram_gb = self.device_mgr.get_device_memory_gb()
+        except Exception:
+            pass
+
+        # Prefer SD-Turbo if GPU and small VRAM; SDXL-Turbo if large VRAM; fallback to SD 2.1 on CPU
+        if device == "cuda":
+            if vram_gb and vram_gb >= 12:
+                model_id = os.environ.get("BGFX_T2I_MODEL", "stabilityai/sdxl-turbo")
+            else:
+                model_id = os.environ.get("BGFX_T2I_MODEL", "stabilityai/sd-turbo")
+        else:
+            # CPU-friendly (still heavy): classic SD 2.1
             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)
+
+        logger.info(f"Loading text-to-image model: {model_id} (device={device}, vram={vram_gb} GB)")
+
+        dtype = torch.float16 if device == "cuda" else torch.float32
+
+        pipe = None
+        err = None
+        try:
+            # Newer unified API handles sd-turbo and sdxl-turbo too
+            pipe = AutoPipelineForText2Image.from_pretrained(
+                model_id,
+                torch_dtype=dtype,
+                use_safetensors=True,
+                token=token
+            )
+        except Exception as e1:
+            err = e1
+            try:
+                # Fallback to classic pipeline (works for sd/stable-diffusion-2-1)
+                pipe = StableDiffusionPipeline.from_pretrained(
+                    model_id,
+                    torch_dtype=dtype,
+                    use_safetensors=True,
+                    safety_checker=None,   # disable to avoid false positives for office backgrounds
+                    feature_extractor=None,
+                    use_auth_token=token  # legacy name
+                )
+            except Exception as e2:
+                return None, None, f"AI model load failed: {e1} / {e2}"
+
+        # Memory/perf knobs
+        try:
+            pipe.set_progress_bar_config(disable=True)
+        except Exception:
+            pass
+        try:
+            pipe.enable_attention_slicing()
+        except Exception:
+            pass
+        try:
+            pipe.enable_vae_slicing()
+        except Exception:
+            pass
+        if device == "cuda":
+            try:
+                pipe.enable_xformers_memory_efficient_attention()
+            except Exception:
+                pass
             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]
+        else:
+            # If accelerate is present, offload module-wise to save RAM
+            try:
+                pipe.enable_sequential_cpu_offload()
+            except Exception:
+                pass
+
+        self.t2i_pipe = pipe
+        self.t2i_model_id = model_id
+        return pipe, model_id, f"AI model loaded: {model_id}"
+
+    def ai_generate_background(self, prompt: str, seed: int, width: int, height: int) -> Tuple[Optional[Image.Image], Optional[str], str]:
+        """
+        Generate a background and save to /tmp. Returns (preview_img, path, status).
+        """
+        pipe, model_id, msg = self._ensure_t2i()
+        if pipe is None:
+            logger.warning(msg)
+            return None, None, msg
+
+        # Ensure sane, divisible-by-8 sizes
+        w = _div8(int(width)) if width else PREVIEW_W
+        h = _div8(int(height)) if height else PREVIEW_H
+        w = max(256, min(w, 1536))
+        h = max(256, min(h, 1536))
+
+        # Reasonable defaults for office-like backgrounds
+        prompt = (prompt or "professional modern office background, neutral colors, soft depth of field, clean, minimal, photorealistic")
+        negative = "text, watermark, logo, people, person, artifact, noisy, blurry"
+
+        # Seed & inference
+        try:
+            import torch
+            g = None
+            device = "cuda" if getattr(torch, "cuda", None) and torch.cuda.is_available() else "cpu"
+            try:
+                g = torch.Generator(device=device).manual_seed(int(seed)) if seed is not None else None
+            except Exception:
+                g = None
+
+            # steps: turbo likes very low steps; classic SD needs more
+            steps = 4 if ("turbo" in (model_id or "").lower()) else 25
+            guidance = 1.0 if ("turbo" in (model_id or "").lower()) else 7.0
+
+            with torch.inference_mode():
+                if device == "cuda":
+                    # autocast for fp16
+                    with torch.autocast("cuda"):
+                        out = pipe(
+                            prompt=prompt,
+                            negative_prompt=negative,
+                            height=h,
+                            width=w,
+                            guidance_scale=guidance,
+                            num_inference_steps=steps,
+                            generator=g
+                        )
+                else:
+                    out = pipe(
+                        prompt=prompt,
+                        negative_prompt=negative,
+                        height=h,
+                        width=w,
+                        guidance_scale=guidance,
+                        num_inference_steps=steps,
+                        generator=g
+                    )
+            img = out.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)})"
+            # Return preview-sized display to keep UI snappy
+            return img.resize((PREVIEW_W, PREVIEW_H), Image.LANCZOS), tmp_path, f"{msg} • Generated {w}x{h}"
         except Exception as e:
-            logger.warning("AI generation unavailable: %s", e)
-            return None, None, f"AI generation unavailable: {e}"
+            logger.exception("AI generation error")
+            return None, None, f"AI generation failed: {e}"
 
     # ---- PROCESS VIDEO ----
     def process_video(
@@ -233,6 +381,9 @@ def process_video(
         if not self.models_loaded:
             return None, "Models not loaded yet"
 
+        if not video:
+            return None, "Please upload a video first."
+
         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},
@@ -247,10 +398,9 @@ def process_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}
+            bg_cfg: Dict[str, Any] = {"custom_path": custom_bg_file.name}
         elif src == "gradient":
             bg_cfg = {
                 "gradient": {
@@ -311,9 +461,12 @@ def create_csp_safe_gradio():
 
                 # 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")
+                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)
@@ -339,48 +492,67 @@ def create_csp_safe_gradio():
         # ---------- UI wiring ----------
 
         # background source → show/hide controls
-        def on_source_change(src):
+        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")),
+                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,
+            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],
         )
 
-        # 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])
+        # ✅ Clear any previous AI image path when switching source (avoids stale AI background)
+        def _clear_ai_state(_):
+            return None
+        bg_source.change(fn=_clear_ai_state, inputs=[bg_source], outputs=[ai_bg_path_state])
+
+        # 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]
+        )
 
-        def preview_from_gradient(gt, c1, c2, ang):
-            return app.preview_gradient(gt, c1, c2, ang)
+        # 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=preview_from_gradient, inputs=[grad_type, grad_color1, grad_color2, grad_angle], outputs=[bg_preview])
+            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"))
+                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)
+                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])
 
@@ -388,7 +560,6 @@ def ai_generate(prompt, seed, size):
         def safe_load():
             msg = app.load_models()
             logger.info("UI: models loaded")
-            # 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])