Update app.py

app.py

@@ -4,6 +4,7 @@
 - Fixed output path parameter (now uses proper .mp4 extension)
 - Fixed parameter order for TwoStageProcessor.process_video()
 - Fixed background handling (pass background, not preset to processor)
+- Added robust vertical gradient + background size guard to prevent broadcasting errors
 - All other functionality maintained from working version
 """
 
@@ -1686,16 +1687,28 @@ def _ensure_file_exists(path_str: str) -> bool:
     except Exception:
         return False
 
-def _make_vertical_gradient(w: int, h: int, c1, c2) -> np.ndarray:
-    """Make a vertical BGR gradient image (uint8)."""
-    top = np.array(c1, dtype=np.float32)
-    bot = np.array(c2, dtype=np.float32)
-    grad = np.zeros((h, w, 3), dtype=np.uint8)
-    # vectorized lerp
-    t = np.linspace(0.0, 1.0, num=h, dtype=np.float32)[:, None]
-    line = (top * (1.0 - t) + bot * t).astype(np.uint8)
-    grad[:] = line
-    return grad
+def _ensure_size_bgr(img: np.ndarray, width: int, height: int) -> np.ndarray:
+    """Guarantee image is (height,width,3) BGR uint8."""
+    if img is None or getattr(img, "size", 0) == 0:
+        raise ValueError("Background image is empty")
+    if img.ndim == 2:
+        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+    if img.shape[2] != 3:
+        img = img[:, :, :3]
+    if (img.shape[1], img.shape[0]) != (width, height):
+        img = cv2.resize(img, (width, height), interpolation=cv2.INTER_AREA)
+    return img.astype(np.uint8)
+
+def _make_vertical_gradient(width: int, height: int, c1, c2) -> np.ndarray:
+    """
+    Make a vertical BGR gradient image (uint8).
+    Returns shape = (height, width, 3). Broadcast-safe via explicit repeat.
+    """
+    top = np.array(c1, dtype=np.float32)   # (3,)
+    bot = np.array(c2, dtype=np.float32)   # (3,)
+    rows = np.linspace(top, bot, num=max(1, height), dtype=np.float32)   # (height, 3)
+    grad = np.repeat(rows[:, None, :], repeats=max(1, width), axis=1)    # (height, width, 3)
+    return np.clip(grad, 0, 255).astype(np.uint8)
 
 def resolve_background(
     video_path: str,
@@ -1725,7 +1738,7 @@ def resolve_background(
         img = cv2.imread(upload_img_path)
         if img is None:
             raise RuntimeError("Failed to read uploaded background image.")
-        img = cv2.resize(img, (w, h), interpolation=cv2.INTER_AREA)
+        img = _ensure_size_bgr(img, w, h)
         cv2.imwrite(out_path, img)
         return out_path
 
@@ -1743,7 +1756,7 @@ def resolve_background(
         img = cv2.imread(path)
         if img is None:
             raise RuntimeError(f"Failed to load preset image: {preset_name}")
-        img = cv2.resize(img, (w, h), interpolation=cv2.INTER_AREA)
+        img = _ensure_size_bgr(img, w, h)
         cv2.imwrite(out_path, img)
         return out_path
 
@@ -1752,6 +1765,7 @@ def resolve_background(
             raise RuntimeError("Please select a gradient preset.")
         c1, c2 = GRADIENT_PRESETS[gradient_name]
         grad = _make_vertical_gradient(w, h, c1, c2)
+        grad = _ensure_size_bgr(grad, w, h)
         cv2.imwrite(out_path, grad)
         return out_path
 
@@ -1760,6 +1774,7 @@ def resolve_background(
         # If you wire up a diffusion pipeline later, render to (w,h) here and save to out_path.
         logger.info("AI Background requested: %s (using placeholder gradient for now)", ai_prompt or "(no prompt)")
         grad = _make_vertical_gradient(w, h, (80, 72, 96), (180, 180, 210))
+        grad = _ensure_size_bgr(grad, w, h)
         cv2.imwrite(out_path, grad)
         return out_path
 
@@ -1917,6 +1932,10 @@ def _toggle_bg_fields(src):
                 gr.update(visible=(src == "AI Background")),         # ai_prompt
             )
 
+        background_source.change(
+            _toggle_bg_fields,
+            inputs=[background_source],
+           
         background_source.change(
             _toggle_bg_fields,
             inputs=[background_source],
@@ -1944,7 +1963,7 @@ def _toggle_bg_fields(src):
         gr.Markdown("""
         **Background Notes**
         - *Professional Presets* expect files under `assets/backgrounds/professional/`.
-        - *Gradients* render on-the-fly (fastest).
+        - *Gradients* render on-the-fly (fastest) and are auto-sized to your video to avoid broadcasting issues.
         - *AI Background* is currently a placeholder gradient — wire a diffusion
           pipeline here later to render images from prompts, then save and return the path.
         """)
@@ -1957,7 +1976,6 @@ def _toggle_bg_fields(src):
 
 def main():
     logger.info("Starting Video Background Replacement - SAM2 LARGE (REVERTED TO WORKING VERSION)")
-    
     interface = create_interface()
     interface.launch(
         server_name="0.0.0.0",
@@ -1968,4 +1986,4 @@ def main():
     )
 
 if __name__ == "__main__":
-    main()
+    main()