Update app.py

app.py

@@ -1,55 +1,98 @@
 # app.py
-import streamlit as st
-from pathlib import Path
-import tempfile
+import sys
 import os
+import tempfile
+from pathlib import Path
+import streamlit as st
 from PIL import Image
-from pipeline.two_stage_pipeline import process_two_stage
 
+# Force app root and pipeline to be importable (for HF docker)
+sys.path.append(str(Path(__file__).parent.resolve()))
+sys.path.append(str(Path(__file__).parent.resolve() / "pipeline"))
 
-st.set_page_config(page_title="BackgroundFX Pro", layout="wide")
+# === CUDA DIAGNOSTICS (LOGS) ===
+try:
+    import torch
+    print("=== CUDA Info ===")
+    print("CUDA available:", torch.cuda.is_available())
+    print("CUDA device count:", torch.cuda.device_count())
+    print("torch version:", torch.__version__)
+    print("=================")
+except Exception as e:
+    print(f"[WARN] Could not print CUDA info: {e}")
 
-st.title("🎥 BackgroundFX Pro — AI Video Background Replacer")
-st.markdown(
-    "Upload a video and a background image. The system will remove the original background "
-    "using SAM2 & MatAnyone, and composite the subject over your new background image."
+# === Import the two-stage pipeline ===
+try:
+    from pipeline.two_stage_pipeline import process_two_stage
+except ImportError as e:
+    st.error(f"Pipeline not found or misconfigured: {e}")
+    st.stop()
+
+st.set_page_config(
+    page_title="BackgroundFX Pro",
+    layout="wide",
+    page_icon="🎬"
+)
+
+st.title("🎬 BackgroundFX Pro — AI Video Background Replacer")
+st.write(
+    "Upload a video and a background image (optional). The system will remove the original background "
+    "using **SAM2** & **MatAnyone**, and composite the subject over your new background image."
 )
 
 def progress_callback(msg, prog):
-    """Streamlit-friendly progress reporting."""
-    if not hasattr(st.session_state, "progress_bar"):
+    """Streamlit-friendly progress reporting (persistent)."""
+    if "progress_bar" not in st.session_state:
         st.session_state.progress_bar = st.progress(0)
         st.session_state.progress_text = st.empty()
     st.session_state.progress_bar.progress(min(max(prog,0.0),1.0))
     st.session_state.progress_text.text(msg)
 
-# File uploaders
-video_file = st.file_uploader("1️⃣ Upload video file (mp4, mov, avi)", type=["mp4", "mov", "avi"])
-bg_file = st.file_uploader("2️⃣ Upload new background (jpg/png, optional)", type=["jpg", "jpeg", "png"])
+# --- Uploaders ---
+video_file = st.file_uploader(
+    "1️⃣ Upload video file (mp4, mov, avi, mkv)", 
+    type=["mp4", "mov", "avi", "mkv"]
+)
+bg_file = st.file_uploader(
+    "2️⃣ Upload new background (jpg/png, optional)", 
+    type=["jpg", "jpeg", "png"]
+)
+
+video_path, bg_path = None, None
 
-if video_file is not None:
+if video_file:
+    # Save video to temp file
     with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as vf:
         vf.write(video_file.read())
         video_path = vf.name
     st.video(video_path)
 
-    # Use uploaded background, or fallback to solid green
-    if bg_file:
-        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as bf:
-            bf.write(bg_file.read())
-            bg_path = bf.name
-        background_img = Image.open(bg_path).convert("RGB")
-        st.image(background_img, caption="Background Image")
-    else:
-        # Fallback: green background
-        background_img = Image.new("RGB", (1280, 720), (0,255,0))
-        st.info("No background image uploaded. Using default green background.")
+if bg_file:
+    # Save background to temp file
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as bf:
+        bf.write(bg_file.read())
+        bg_path = bf.name
+    background_img = Image.open(bg_path).convert("RGB")
+    st.image(background_img, caption="Background Image")
+elif video_file:
+    # Fallback: green background sized to video or 1280x720
+    try:
+        import cv2
+        cap = cv2.VideoCapture(video_path)
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) or 1280
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) or 720
+        cap.release()
+    except Exception:
+        width, height = 1280, 720
+    background_img = Image.new("RGB", (width, height), (0,255,0))
+    st.info("No background image uploaded. Using default green background.")
 
-    # Button to start processing
+# --- Processing ---
+if video_path and background_img:
     if st.button("🚀 Run Background Replacement"):
         st.session_state.progress_bar = st.progress(0)
         st.session_state.progress_text = st.empty()
-        with st.spinner("Processing video... This may take several minutes for long videos."):
+        with st.spinner("Processing video... (May take several minutes for long videos)"):
             try:
                 out_path = process_two_stage(
                     video_path,
@@ -68,10 +111,15 @@ def progress_callback(msg, prog):
             except Exception as e:
                 st.error(f"❌ Processing failed: {e}")
 
-        # Clean up temp files (optional)
+        # Clean up temp files (best effort)
         try:
-            os.unlink(video_path)
+            if video_path and os.path.exists(video_path):
+                os.unlink(video_path)
         except Exception: pass
-        if bg_file:
-            try: os.unlink(bg_path)
-            except Exception: pass
+        try:
+            if bg_path and os.path.exists(bg_path):
+                os.unlink(bg_path)
+        except Exception: pass
+
+# HF Spaces tip
+st.caption("Built for Hugging Face T4 Spaces • CUDA & GPU acceleration required • v2025.10")