Update ui.py

ui.py

@@ -3,7 +3,7 @@
 Modern UI for Video Background Replacer (PRO)
 - Clean, responsive Streamlit layout
 - Shows live stage progress and all errors during processing
-- NO automatic reruns (no st.rerun) to avoid killing long-running jobs
+- Progress bars and visual feedback to keep session alive
 """
 import streamlit as st
 import os
@@ -14,7 +14,7 @@
 import time
 
 logger = logging.getLogger("Advanced Video Background Replacer")
-UI_BUILD = "ui-2025-10-04-10-30Z"  # visible in sidebar so you know this file is active
+UI_BUILD = "ui-2025-10-04-16-00Z"
 
 def tail_file(path: str, lines: int = 400) -> str:
     if not os.path.exists(path):
@@ -36,7 +36,6 @@ def read_file_bytes(path: str) -> bytes:
         return b""
 
 def _render_background_settings():
-    # Available stock images (may not exist in Space)
     stock_images = {
         "Sunset Beach": "stock_images/sunset_beach.jpg",
         "Urban Office": "stock_images/urban_office.jpg",
@@ -94,7 +93,6 @@ def _render_background_settings():
         prompt = st.text_input("Describe the background to generate (AI):", key="ai_bg_prompt")
         ai_ready = False
         if st.button("Generate Background", key="gen_bg_btn") and prompt:
-            # TODO: Plug in real AI image generator
             background = Image.new("RGB", (512, 320), (64, 32, 96))
             st.session_state.generated_bg = background
             st.success("AI-generated background (stub). Replace with your generator!")
@@ -115,7 +113,6 @@ def render_ui(process_video_func):
             st.caption(f"UI build: {UI_BUILD}")
             st.markdown("**Log file:** `/tmp/app.log`")
 
-            # GPU indicator (best-effort)
             if st.session_state.get('gpu_available', False):
                 try:
                     import torch
@@ -126,7 +123,6 @@ def render_ui(process_video_func):
             else:
                 st.error("GPU: Not Available")
 
-            # Log tail controls (NO auto refresh, only manual button)
             st.number_input("Tail last N lines", min_value=50, max_value=5000, step=50, key="log_tail_lines")
             log_bytes = read_file_bytes("/tmp/app.log")
             st.download_button(
@@ -160,7 +156,7 @@ def render_ui(process_video_func):
                 try:
                     uploaded_video.seek(0)
                     video_bytes = uploaded_video.read()
-                    st.session_state.video_bytes_cache = video_bytes  # persist across reruns
+                    st.session_state.video_bytes_cache = video_bytes
                     with video_preview_placeholder.container():
                         st.video(video_bytes)
                 except Exception as e:
@@ -173,23 +169,73 @@ def render_ui(process_video_func):
         with col2:
             background, bg_type = _render_background_settings()
             st.header("3. Process Video")
-            status_box = st.empty()  # progress area
-
+            
+            # Progress tracking UI elements
+            progress_container = st.container()
+            with progress_container:
+                progress_bar = st.progress(0)
+                status_text = st.empty()
+                stage_status = st.empty()
+                
             can_process = (
                 st.session_state.get('video_bytes_cache') is not None
                 and not st.session_state.get('processing', False)
                 and (background is not None)
             )
 
+            # Enhanced progress callback with visual updates
+            progress_state = {"step": 0, "total_steps": 12}
+            
             def progress_callback(msg):
-                status_box.info(msg)
+                status_text.info(f"**Status:** {msg}")
                 logger.info(f"[STAGE] {msg}")
+                
+                # Update progress based on message content
+                if "Stage 1 initiated" in msg:
+                    progress_state["step"] = 1
+                    stage_status.markdown("**Current Stage:** Stage 1 - Transparent Video Creation")
+                elif "GPU engaged" in msg or "SAM2" in msg:
+                    progress_state["step"] = 2
+                elif "SAM2 processing frame" in msg:
+                    progress_state["step"] = 3
+                elif "SAM2 complete" in msg or "clearing GPU" in msg:
+                    progress_state["step"] = 4
+                elif "MatAnyone starting" in msg:
+                    progress_state["step"] = 5
+                    stage_status.markdown("**Current Stage:** Stage 1 - MatAnyone Processing")
+                elif "MatAnyone processing" in msg:
+                    progress_state["step"] = 6
+                elif "MatAnyone complete" in msg:
+                    progress_state["step"] = 7
+                elif "Smoothing" in msg:
+                    progress_state["step"] = 8
+                elif "transparent video" in msg and "Creating" in msg:
+                    progress_state["step"] = 9
+                elif "Stage 1 complete" in msg:
+                    progress_state["step"] = 10
+                elif "Stage 2 begun" in msg:
+                    progress_state["step"] = 10
+                    stage_status.markdown("**Current Stage:** Stage 2 - Background Compositing")
+                elif "Compositing frames" in msg:
+                    progress_state["step"] = 11
+                elif "Stage 2 complete" in msg:
+                    progress_state["step"] = 12
+                    stage_status.markdown("**Current Stage:** Complete!")
+                
+                # Update progress bar
+                progress_pct = min(100, int((progress_state["step"] / progress_state["total_steps"]) * 100))
+                progress_bar.progress(progress_pct)
 
             if st.button("Process Video", disabled=not can_process, use_container_width=True):
                 try:
                     logger.info("Process Video button clicked")
+                    
+                    # Reset progress
+                    progress_state["step"] = 0
+                    progress_bar.progress(0)
+                    status_text.empty()
+                    stage_status.empty()
 
-                    # Rebuild a file-like object from cached bytes so we don't depend on the upload widget object
                     import io
                     class _MemFile:
                         def __init__(self, name, data):
@@ -209,15 +255,16 @@ def seek(self, pos):
 
                     st.session_state.processing = False
                     if success:
-                        status_box.success("✅ Video processing complete!")
+                        progress_bar.progress(100)
+                        status_text.success("**Status:** Processing complete!")
                         st.success("Video processing complete!")
                     else:
-                        status_box.error("❌ Video processing failed. Check logs for details.")
+                        status_text.error("**Status:** Processing failed. Check logs for details.")
                         st.error("Video processing failed. Check logs for details.")
                 except Exception as e:
                     st.session_state.processing = False
                     logger.error(f"[UI] Process video error: {e}", exc_info=True)
-                    status_box.error(f"Processing error: {str(e)}. Check logs for details.")
+                    status_text.error(f"**Status:** Error - {str(e)}")
                     st.error(f"Processing error: {str(e)}. Check logs for details.")
 
         # --- Results Display ---
@@ -239,11 +286,11 @@ def seek(self, pos):
 
         # --- Error Display ---
         if st.session_state.get('last_error'):
-            with st.expander("⚠️ Last Error", expanded=True):
+            with st.expander("Last Error", expanded=True):
                 st.error(st.session_state.last_error)
                 if st.button("Clear Error"):
                     st.session_state.last_error = None
-                    # No st.rerun() — the next interaction will refresh naturally
+                    
     except Exception as e:
         logger.error(f"[UI] Render UI error: {e}", exc_info=True)
-        st.error(f"UI rendering error: {str(e)}. Check logs for details.")
+        st.error(f"UI rendering error: {str(e)}. Check logs for details.")