Update app.py

app.py

@@ -4,12 +4,11 @@
 # =============================================================================
 """
 Enhanced Video Background Replacement (SAM2 + MatAnyone + AI Backgrounds)
-- Robust memory management & cleanup
-- SAM2 person mask (CUDA)
-- MatAnyone matting w/ **probability** mask on the first frame (no idx_mask)
-- Cleaned tensor shapes (image: 3xHxW, prob: 1xHxW), consistent device/dtype
-- Optional SDXL / Playground / OpenAI background generation
-- Gradio UI with “chapters” in code for quick edits
+- Strict tensor shapes for MatAnyone (image: 3xHxW, first-frame prob mask: 1xHxW)
+- First frame uses PROB path (no idx_mask / objects) to avoid assertion
+- Memory management & cleanup
+- SDXL / Playground / OpenAI backgrounds
+- Gradio UI with “CHAPTER” dividers
 """
 
 # =============================================================================
@@ -22,9 +21,7 @@
 import psutil
 import time
 import json
-import math
 import base64
-import queue
 import random
 import shutil
 import logging
@@ -45,7 +42,7 @@
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger("bgx")
 
-# Environment tuning (safe defaults; do not overwrite if already set)
+# Environment tuning (safe defaults)
 os.environ.setdefault("CUDA_MODULE_LOADING", "LAZY")
 os.environ.setdefault("TORCH_CUDNN_V8_API_ENABLED", "1")
 os.environ.setdefault("PYTHONUNBUFFERED", "1")
@@ -351,7 +348,7 @@ def initialize(self) -> bool:
                 model = build_sam2("sam2.1/sam2.1_hiera_l.yaml", str(ckpt), device="cuda")
                 self.predictor = SAM2ImagePredictor(model)
 
-                # Quick smoke test
+                # Smoke test
                 test = np.zeros((64, 64, 3), dtype=np.uint8)
                 self.predictor.set_image(test)
                 masks, scores, _ = self.predictor.predict(
@@ -416,47 +413,37 @@ def __init__(self):
         self.core = None
         self.initialized = False
 
-    # ----- small tensor helpers -----
-    def _to_chw_float(self, img01: np.ndarray) -> torch.Tensor:
-        # img01: HxWx3, float32 [0..1]
+    # ----- tensor helpers -----
+    def _to_chw_float(self, img01: np.ndarray) -> "torch.Tensor":
         assert img01.ndim == 3 and img01.shape[2] == 3, f"Expected HxWx3, got {img01.shape}"
         t = torch.from_numpy(img01.transpose(2, 0, 1)).contiguous().float()  # 3xHxW
         return t.to(DEVICE, non_blocking=CUDA_AVAILABLE)
 
-    def _prob_from_mask_u8(self, mask_u8: np.ndarray, w: int, h: int) -> torch.Tensor:
-        # returns 1xHxW float32 [0..1]
+    def _prob_from_mask_u8(self, mask_u8: np.ndarray, w: int, h: int) -> "torch.Tensor":
         if mask_u8.shape[0] != h or mask_u8.shape[1] != w:
             mask_u8 = cv2.resize(mask_u8, (w, h), interpolation=cv2.INTER_NEAREST)
         prob = (mask_u8.astype(np.float32) / 255.0)[None, ...]  # 1xHxW
         t = torch.from_numpy(prob).contiguous().float()
         return t.to(DEVICE, non_blocking=CUDA_AVAILABLE)
 
-    def _alpha_to_u8_hw(self, alpha_like: torch.Tensor) -> np.ndarray:
-        # Accepts tensor with shapes: (1,H,W) or (H,W) or (K,H,W) where K==1
+    def _alpha_to_u8_hw(self, alpha_like) -> np.ndarray:
         if isinstance(alpha_like, (list, tuple)) and len(alpha_like) > 1:
-            # Many MatAnyone step() return (indices, probs)
-            alpha_like = alpha_like[1]
+            alpha_like = alpha_like[1]  # handle (indices, probs)
         if isinstance(alpha_like, torch.Tensor):
             t = alpha_like.detach()
             if t.is_cuda:
                 t = t.cpu()
-            t = t.float().clamp(0, 1)
-            a = t.numpy()
+            a = t.float().clamp(0, 1).numpy()
         else:
             a = np.asarray(alpha_like, dtype=np.float32)
             a = np.clip(a, 0, 1)
-
-        if a.ndim == 3 and a.shape[0] == 1:
-            a = a[0]            # (H,W)
-        elif a.ndim == 3 and a.shape[0] > 1:
-            a = a[0]            # take first object
-        elif a.ndim == 2:
-            pass
-        else:
-            # try to squeeze any trailing singleton dims
-            a = np.squeeze(a)
-            if a.ndim != 2:
-                raise ValueError(f"Alpha map must be HxW; got shape {a.shape}")
+        a = np.squeeze(a)
+        if a.ndim != 2:
+            # handle shapes (1,H,W) or (K,H,W) → pick first
+            if a.ndim == 3 and a.shape[0] >= 1:
+                a = a[0]
+            else:
+                raise ValueError(f"Alpha must be HxW; got {a.shape}")
         return (np.clip(a * 255.0, 0, 255).astype(np.uint8))
 
     def initialize(self) -> bool:
@@ -498,7 +485,7 @@ def initialize(self) -> bool:
                 state.matanyone_error = f"MatAnyone init error: {e}"
                 return False
 
-    # ----- main video matting using PROB mask on first frame -----
+    # ----- video matting using first-frame PROB mask -----
     def process_video(self, input_path: str, mask_path: str, output_path: str) -> str:
         if not self.initialized or self.core is None:
             raise RuntimeError("MatAnyone not initialized")
@@ -526,7 +513,7 @@ def process_video(self, input_path: str, mask_path: str, output_path: str) -> st
 
         frame_idx = 0
 
-        # ---------- First frame (with PROB mask) ----------
+        # First frame (with PROB mask)
         ok, frame_bgr = cap.read()
         if not ok or frame_bgr is None:
             cap.release()
@@ -536,14 +523,13 @@ def process_video(self, input_path: str, mask_path: str, output_path: str) -> st
         prob_chw = self._prob_from_mask_u8(seed_mask, w, h)   # 1xHxW
 
         with torch.no_grad():
-            # Use PROB path (no idx_mask, no objects). Some forks require `matting=True`
             out_prob = self.core.step(img_chw, prob=prob_chw, matting=True)
 
         alpha_u8 = self._alpha_to_u8_hw(out_prob)
         cv2.imwrite(str(tmp_dir / f"{frame_idx:06d}.png"), alpha_u8)
         frame_idx += 1
 
-        # ---------- Remaining frames (no mask) ----------
+        # Remaining frames (no mask)
         while True:
             ok, frame_bgr = cap.read()
             if not ok or frame_bgr is None:
@@ -579,7 +565,7 @@ def process_video(self, input_path: str, mask_path: str, output_path: str) -> st
         return str(alpha_path)
 
 # =============================================================================
-# CHAPTER 7: AI BACKGROUNDS (SDXL / Playground / OpenAI)
+# CHAPTER 7: AI BACKGROUNDS
 # =============================================================================
 def _maybe_enable_xformers(pipe):
     try:
@@ -639,7 +625,7 @@ def generate_sdxl_background(width:int, height:int, prompt:str, steps:int=30, gu
             generator=generator
         ).images[0]
 
-        out = TEMP_DIR / f"sdxl_bg_{int(time.time())}_{seed:08d}.jpg"
+        out = TEMP_DIR / f"sdxl_bg_{int(time.time())}_{seed or 0:08d}.jpg"
         img.save(out, quality=95, optimize=True)
         memory_manager.register_temp_file(str(out))
         del pipe, img
@@ -680,7 +666,7 @@ def generate_playground_v25_background(width:int, height:int, prompt:str, steps:
             generator=generator
         ).images[0]
 
-        out = TEMP_DIR / f"pg25_bg_{int(time.time())}_{seed:08d}.jpg"
+        out = TEMP_DIR / f"pg25_bg_{int(time.time())}_{seed or 0:08d}.jpg"
         img.save(out, quality=95, optimize=True)
         memory_manager.register_temp_file(str(out))
         del pipe, img
@@ -726,7 +712,7 @@ def generate_sd15_background(width:int, height:int, prompt:str, steps:int=25, gu
             generator=generator
         ).images[0]
 
-        out = TEMP_DIR / f"sd15_bg_{int(time.time())}_{seed:08d}.jpg"
+        out = TEMP_DIR / f"sd15_bg_{int(time.time())}_{seed or 0:08d}.jpg"
         img.save(out, quality=95, optimize=True)
         memory_manager.register_temp_file(str(out))
         del pipe, img
@@ -785,7 +771,7 @@ def generate_ai_background_router(width:int, height:int, prompt:str, model:str="
             return str(out)
 
 # =============================================================================
-# CHAPTER 8: CHUNKED PROCESSOR (optional, unchanged)
+# CHAPTER 8: CHUNKED PROCESSOR (optional)
 # =============================================================================
 class ChunkedVideoProcessor:
     def __init__(self, chunk_size_frames: int = 60):
@@ -934,7 +920,7 @@ def process_video_main(
             alpha_clip = VideoFileClip(alpha_video)
 
             if background_path and os.path.exists(background_path):
-                messages.append(f"🖼️ Using background file")
+                messages.append("🖼️ Using background file")
                 bg_bgr = cv2.imread(background_path)
                 bg_bgr = cv2.resize(bg_bgr, (w, h))
                 bg_rgb = cv2.cvtColor(bg_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
@@ -1049,17 +1035,13 @@ def preload(ai_model, openai_key, force_gpu, progress=gr.Progress()):
             progress(0, desc="Preloading...")
             msg = ""
             if ai_model in ("SDXL", "Playground v2.5", "SD 1.5 (fallback)"):
-                # “preload lite”: generate tiny image once (2 steps)
                 try:
                     if ai_model == "SDXL":
-                        _ = generate_sdxl_background(64, 64, "plain background", steps=2, guidance=3.5,
-                                                     seed=42, require_gpu=bool(force_gpu))
+                        _ = generate_sdxl_background(64, 64, "plain", steps=2, guidance=3.5, seed=42, require_gpu=bool(force_gpu))
                     elif ai_model == "Playground v2.5":
-                        _ = generate_playground_v25_background(64, 64, "plain background", steps=2, guidance=3.5,
-                                                               seed=42, require_gpu=bool(force_gpu))
+                        _ = generate_playground_v25_background(64, 64, "plain", steps=2, guidance=3.5, seed=42, require_gpu=bool(force_gpu))
                     else:
-                        _ = generate_sd15_background(64, 64, "plain background", steps=2, guidance=3.5,
-                                                     seed=42, require_gpu=bool(force_gpu))
+                        _ = generate_sd15_background(64, 64, "plain", steps=2, guidance=3.5, seed=42, require_gpu=bool(force_gpu))
                     msg += f"{ai_model} preloaded.\n"
                 except Exception as e:
                     msg += f"{ai_model} preload failed: {e}\n"
@@ -1152,12 +1134,18 @@ def approve_background(bg_path):
                 gr.Markdown("### Background")
                 bg_method = gr.Radio(choices=["Upload Image", "Gradients", "AI Generated"],
                                      value="AI Generated", label="Background Method")
+
+                # Upload group (hidden by default)
                 with gr.Group(visible=False) as upload_group:
                     upload_img = gr.Image(label="Background Image", type="filepath")
-                with gr.Group(visible=True) as gradient_group:
+
+                # Gradient group (hidden by default)
+                with gr.Group(visible=False) as gradient_group:
                     gradient_choice = gr.Dropdown(label="Gradient Style",
                                                   choices=list(GRADIENT_PRESETS.keys()),
                                                   value="Slate")
+
+                # AI group (visible by default)
                 with gr.Group(visible=True) as ai_group:
                     prompt_suggestions = gr.Dropdown(label="💡 Prompt Inspiration",
                                                      choices=AI_PROMPT_SUGGESTIONS,
@@ -1216,6 +1204,7 @@ def approve_background(bg_path):
 
         # --- Wiring ---
         def update_background_visibility(method):
+            # return visibilities for: upload_group, gradient_group, ai_group
             return (
                 gr.update(visible=(method == "Upload Image")),
                 gr.update(visible=(method == "Gradients")),
@@ -1228,7 +1217,7 @@ def update_prompt_from_suggestion(suggestion):
             return gr.update(value=suggestion)
 
         bg_method.change(
-            lambda m: update_background_visibility(m),
+            update_background_visibility,
             inputs=[bg_method],
             outputs=[upload_group, gradient_group, ai_group]
         )
@@ -1255,10 +1244,11 @@ def update_prompt_from_suggestion(suggestion):
         diagnostics_btn.click(diag, outputs=[diagnostics_output])
         cleanup_btn.click(cleanup, outputs=[diagnostics_output])
 
+        # ----- FIXED: use upload_img (Image component), not upload_group (Group) -----
         def process_video(
             video_file,
             bg_method,
-            upload_img,
+            upload_img,             # <-- correct input
             gradient_choice,
             approved_background_path,
             last_generated_bg,
@@ -1321,7 +1311,7 @@ def process_video(
             inputs=[
                 video_input,
                 bg_method,
-                upload_group,  # this group passes the image component value
+                upload_img,            # <-- FIXED here
                 gradient_choice,
                 approved_background_path, last_generated_bg,
                 trim_enabled, trim_seconds, crf_value, audio_enabled,