Update streamlit_app.py

streamlit_app.py

@@ -1,13 +1,15 @@
 # ==================================================================================
-# STREAMLIT VIDEO BACKGROUND REPLACER - MAIN APPLICATION
+# streamlit_ui.py
+# STREAMLIT VIDEO BACKGROUND REPLACER - MAIN APPLICATION (HF-Ready Logging)
 # ==================================================================================
 # Single-button workflow: Upload video + background → Process → Download result
 # Uses SAM2 + MatAnyone pipeline with temporal smoothing
+# Adds robust logging (stdout + rotating file) and in-app log tail viewer
 # ==================================================================================
 
-# ==================================================================================
+# =========================================
 # CHAPTER 1: IMPORTS AND SETUP
-# ==================================================================================
+# =========================================
 
 import streamlit as st
 import os
@@ -18,9 +20,14 @@
 import numpy as np
 from PIL import Image
 import logging
+import logging.handlers
 import io
 import torch
 import traceback
+import uuid
+from datetime import datetime
+from tempfile import NamedTemporaryFile
+import subprocess
 
 # Project Setup
 sys.path.append(str(Path(__file__).parent.absolute()))
@@ -28,22 +35,134 @@
 # Import pipeline functions
 from pipeline.video_pipeline import stage1_create_transparent_video, stage2_composite_background
 
-# Logging Configuration
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
+APP_NAME = "Advanced Video Background Replacer"
+LOG_FILE = "/tmp/app.log"  # HF Spaces: writable, survives session
+LOG_MAX_BYTES = 5 * 1024 * 1024
+LOG_BACKUPS = 5
+
+# =========================================
+# CHAPTER 2: LOGGING
+# =========================================
+
+def setup_logging(level: int = logging.INFO) -> logging.Logger:
+    logger = logging.getLogger(APP_NAME)
+    logger.setLevel(level)
+    logger.propagate = False  # avoid double logs in Streamlit
+
+    # Clear previous handlers on rerun
+    for h in list(logger.handlers):
+        logger.removeHandler(h)
+
+    # Console handler (stdout → visible in HF logs)
+    ch = logging.StreamHandler(sys.stdout)
+    ch.setLevel(level)
+    ch.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
+
+    # Rotating file handler (local tailing inside the app)
+    fh = logging.handlers.RotatingFileHandler(
+        LOG_FILE, maxBytes=LOG_MAX_BYTES, backupCount=LOG_BACKUPS, encoding="utf-8"
+    )
+    fh.setLevel(level)
+    fh.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
+
+    logger.addHandler(ch)
+    logger.addHandler(fh)
+    return logger
+
+logger = setup_logging()
 
 # Global Exception Hook
 def custom_excepthook(type, value, tb):
     logger.error(f"Unhandled: {type.__name__}: {value}\n{''.join(traceback.format_tb(tb))}", exc_info=True)
 sys.excepthook = custom_excepthook
 
-# ==================================================================================
-# CHAPTER 2: STREAMLIT CONFIGURATION
-# ==================================================================================
+# =========================================
+# CHAPTER 3: DIAGNOSTICS HELPERS
+# =========================================
+
+def try_run(cmd, timeout=5):
+    """Run a shell command safely and return (ok, stdout|stderr)."""
+    try:
+        out = subprocess.check_output(cmd, stderr=subprocess.STDOUT, timeout=timeout, text=True)
+        return True, out.strip()
+    except Exception as e:
+        return False, str(e)
+
+def tail_file(path: str, lines: int = 400) -> str:
+    """Read last N lines from a text file efficiently."""
+    if not os.path.exists(path):
+        return "(log file not found)"
+    try:
+        # A simple, safe tail (files are small due to rotation)
+        with open(path, "r", encoding="utf-8", errors="replace") as f:
+            content = f.readlines()
+        return "".join(content[-lines:])
+    except Exception as e:
+        return f"(failed to read log: {e})"
+
+def check_gpu(logger: logging.Logger):
+    """Check GPU availability and log details."""
+    logger.info("=" * 60)
+    logger.info("GPU DIAGNOSTIC")
+    logger.info("=" * 60)
+
+    logger.info(f"Python: {sys.version.split()[0]}")
+    logger.info(f"Torch: {torch.__version__}")
+    logger.info(f"CUDA compiled version in torch: {torch.version.cuda}")
+
+    cuda_available = torch.cuda.is_available()
+    logger.info(f"torch.cuda.is_available(): {cuda_available}")
+
+    # Try to probe nvidia-smi if present
+    ok, out = try_run(["bash", "-lc", "command -v nvidia-smi && nvidia-smi -L && nvidia-smi -q -d MEMORY | head -n 40"], timeout=6)
+    if ok:
+        logger.info("nvidia-smi probe:\n" + out)
+    else:
+        logger.info(f"nvidia-smi probe not available: {out}")
+
+    if cuda_available:
+        try:
+            count = torch.cuda.device_count()
+            logger.info(f"CUDA Device Count: {count}")
+            for i in range(count):
+                logger.info(f"Device {i}: {torch.cuda.get_device_name(i)}")
+            torch.cuda.set_device(0)
+            logger.info("Set CUDA device 0 as default")
+            test_tensor = torch.randn(64, 64).cuda()
+            logger.info(f"GPU Test: SUCCESS on {test_tensor.device}")
+            del test_tensor
+            torch.cuda.empty_cache()
+        except Exception as e:
+            logger.error(f"GPU test failed: {e}", exc_info=True)
+    else:
+        logger.warning("CUDA NOT AVAILABLE")
+
+    logger.info("=" * 60)
+    return cuda_available
+
+def dump_env(logger: logging.Logger):
+    keys_of_interest = [
+        "HF_HOME", "HF_TOKEN", "PYTORCH_CUDA_ALLOC_CONF", "CUDA_VISIBLE_DEVICES",
+        "TORCH_ALLOW_TF32_CUBLAS_OVERRIDE", "OMP_NUM_THREADS", "NUMEXPR_MAX_THREADS",
+        "PYTHONPATH", "PATH"
+    ]
+    logger.info("=== ENV VARS (sanitized) ===")
+    for k in keys_of_interest:
+        v = os.environ.get(k)
+        if not v:
+            continue
+        # Mask tokens/secrets
+        if "TOKEN" in k or "KEY" in k or "SECRET" in k or "PASSWORD" in k:
+            v = "[REDACTED]"
+        logger.info(f"{k}={v}")
+    logger.info("============================")
+
+# =========================================
+# CHAPTER 4: STREAMLIT CONFIGURATION
+# =========================================
 
-# Streamlit Page Config
 st.set_page_config(
-    page_title="Advanced Video Background Replacer",
+    page_title=APP_NAME,
     page_icon="🎥",
     layout="wide",
     initial_sidebar_state="expanded"
@@ -52,10 +171,7 @@ def custom_excepthook(type, value, tb):
 # Custom CSS
 st.markdown("""
 <style>
-    .main .block-container {
-        padding-top: 2rem;
-        padding-bottom: 2rem;
-    }
+    .main .block-container { padding-top: 1.0rem; padding-bottom: 2rem; }
     .stButton>button {
         width: 100%;
         background-color: #4CAF50;
@@ -63,52 +179,16 @@ def custom_excepthook(type, value, tb):
         font-weight: bold;
         transition: all 0.3s;
     }
-    .stButton>button:hover {
-        background-color: #45a049;
-    }
+    .stButton>button:hover { background-color: #45a049; }
+    .small-mono { font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, "Liberation Mono", "Courier New", monospace; font-size: 12px; white-space: pre-wrap; }
 </style>
 """, unsafe_allow_html=True)
 
-# ==================================================================================
-# CHAPTER 3: GPU DIAGNOSTICS
-# ==================================================================================
-
-def check_gpu():
-    """Check GPU availability and log details"""
-    logger.info("=" * 60)
-    logger.info("GPU DIAGNOSTIC")
-    logger.info("=" * 60)
-    
-    cuda_available = torch.cuda.is_available()
-    logger.info(f"torch.cuda.is_available(): {cuda_available}")
-    
-    if cuda_available:
-        logger.info(f"CUDA Version: {torch.version.cuda}")
-        logger.info(f"Device Count: {torch.cuda.device_count()}")
-        logger.info(f"Device Name: {torch.cuda.get_device_name(0)}")
-        
-        torch.cuda.set_device(0)
-        logger.info("Set CUDA device 0 as default")
-        
-        try:
-            test_tensor = torch.randn(100, 100).cuda()
-            logger.info(f"GPU Test: SUCCESS on {test_tensor.device}")
-            del test_tensor
-            torch.cuda.empty_cache()
-        except Exception as e:
-            logger.error(f"GPU Test Failed: {e}")
-    else:
-        logger.warning("CUDA NOT AVAILABLE")
-    
-    logger.info("=" * 60)
-    return cuda_available
-
-# ==================================================================================
-# CHAPTER 4: SESSION STATE INITIALIZATION
-# ==================================================================================
+# =========================================
+# CHAPTER 5: SESSION STATE
+# =========================================
 
 def initialize_session_state():
-    """Initialize all session state variables"""
     defaults = {
         'uploaded_video': None,
         'video_bytes_cache': None,
@@ -123,245 +203,276 @@ def initialize_session_state():
         'gpu_available': None,
         'last_video_id': None,
         'last_bg_image_id': None,
-        'last_error': None
+        'last_error': None,
+        'log_level_name': 'INFO',
+        'auto_refresh_logs': False,
+        'log_tail_lines': 400
     }
-    for key, value in defaults.items():
-        if key not in st.session_state:
-            st.session_state[key] = value
-    
+    for k, v in defaults.items():
+        if k not in st.session_state:
+            st.session_state[k] = v
+
     if st.session_state.gpu_available is None:
-        st.session_state.gpu_available = check_gpu()
+        # On first load, run diagnostics once
+        dump_env(logger)
+        st.session_state.gpu_available = check_gpu(logger)
 
-# ==================================================================================
-# CHAPTER 5: MAIN APPLICATION
-# ==================================================================================
+def set_log_level(name: str):
+    name = (name or "INFO").upper()
+    lvl = getattr(logging, name, logging.INFO)
+    setup_logging(lvl)  # rebuild handlers at new level
+    global logger
+    logger = logging.getLogger(APP_NAME)
+    logger.setLevel(lvl)
+    st.session_state.log_level_name = name
+    logger.info(f"Log level set to {name}")
+
+# =========================================
+# CHAPTER 6: MAIN APP
+# =========================================
 
 def main():
-    st.title("Advanced Video Background Replacer")
+    st.title(APP_NAME)
+    st.caption("HF-ready with robust logs & live tail")
     st.markdown("---")
-    
+
     initialize_session_state()
-    
-    # ==================================================================================
-    # SIDEBAR: System Status
-    # ==================================================================================
-    
+
+    # ------------- Sidebar: System & Logs -----------------
     with st.sidebar:
         st.subheader("System Status")
+
+        # Log level control
+        log_level = st.selectbox(
+            "Log level",
+            ["DEBUG", "INFO", "WARNING", "ERROR"],
+            index=["DEBUG", "INFO", "WARNING", "ERROR"].index(st.session_state.log_level_name),
+            key="log_level_select"
+        )
+        if log_level != st.session_state.log_level_name:
+            set_log_level(log_level)
+
         if st.session_state.gpu_available:
-            st.success(f"GPU: {torch.cuda.get_device_name(0)}")
+            try:
+                dev = torch.cuda.get_device_name(0)
+            except Exception:
+                dev = "Detected (name unavailable)"
+            st.success(f"GPU: {dev}")
         else:
             st.error("GPU: Not Available")
-    
-    # ==================================================================================
-    # ERROR DISPLAY
-    # ==================================================================================
-    
+
+        st.markdown("**Log file:** `/tmp/app.log`")
+        st.checkbox("Auto-refresh log tail (every 2s)", key="auto_refresh_logs")
+
+        st.number_input("Tail last N lines", min_value=50, max_value=5000, step=50, key="log_tail_lines")
+        if st.button("Refresh Logs"):
+            st.session_state._force_log_refresh = True  # trigger rerun
+
+        with st.expander("View Log Tail", expanded=True):
+            # Auto-refresh: trigger reruns
+            if st.session_state.auto_refresh_logs:
+                st.experimental_rerun  # no-op reference to indicate awareness
+                time.sleep(2)
+                st.rerun()
+            log_text = tail_file(LOG_FILE, st.session_state.log_tail_lines)
+            st.code(log_text, language="text")
+
+    # ------------- Error Display -----------------
     if st.session_state.last_error:
         with st.expander("⚠️ Last Error", expanded=False):
             st.error(st.session_state.last_error)
             if st.button("Clear Error"):
                 st.session_state.last_error = None
                 st.rerun()
-    
+
     col1, col2 = st.columns([1, 1], gap="large")
-    
-    # ==================================================================================
-    # COLUMN 1: VIDEO UPLOAD
-    # ==================================================================================
-    
+
+    # ------------- Column 1: Video Upload -----------------
     with col1:
         st.header("1. Upload Video")
-        
         uploaded = st.file_uploader(
             "Upload Video",
-            type=["mp4", "mov", "avi"],
+            type=["mp4", "mov", "avi", "mkv", "webm"],
             key="video_uploader"
         )
-        
-        # FIX: Only update session state when uploaded is not None
+
         current_video_id = id(uploaded)
         if uploaded is not None and current_video_id != st.session_state.last_video_id:
-            logger.info(f"New video: {uploaded.name}")
+            logger.info(f"[UI] New video selected: name={uploaded.name}, size≈{getattr(uploaded, 'size', 'n/a')} bytes")
             st.session_state.uploaded_video = uploaded
             st.session_state.last_video_id = current_video_id
             st.session_state.video_bytes_cache = None
             st.session_state.processed_video_bytes = None
             st.session_state.last_error = None
-        
-        # Video preview with anti-shake placeholder
+
         st.markdown("### Video Preview")
         if st.session_state.video_preview_placeholder is None:
             st.session_state.video_preview_placeholder = st.empty()
-        
+
         if st.session_state.uploaded_video is not None:
             try:
                 if st.session_state.video_bytes_cache is None:
-                    logger.info("Caching video...")
+                    logger.info("[UI] Caching uploaded video for preview...")
                     st.session_state.uploaded_video.seek(0)
                     st.session_state.video_bytes_cache = st.session_state.uploaded_video.read()
-                    logger.info(f"Cached {len(st.session_state.video_bytes_cache)/1e6:.2f}MB")
-                
-                # ALWAYS show cached video (even after errors)
+                    logger.info(f"[UI] Cached {len(st.session_state.video_bytes_cache)/1e6:.2f} MB for preview")
+
                 with st.session_state.video_preview_placeholder.container():
                     st.video(st.session_state.video_bytes_cache)
-                    
+
             except Exception as e:
-                logger.error(f"Video preview error: {e}")
+                logger.error(f"[UI] Video preview error: {e}", exc_info=True)
                 st.session_state.video_preview_placeholder.error(f"Cannot display video: {e}")
         else:
             st.session_state.video_preview_placeholder.empty()
-    
-    # ==================================================================================
-    # COLUMN 2: BACKGROUND SETTINGS & PROCESSING
-    # ==================================================================================
-    
+
+    # ------------- Column 2: Background + Processing -----------------
     with col2:
         st.header("2. Background Settings")
-        
+
         bg_type = st.radio(
             "Select Background Type:",
             ["Image", "Color"],
             horizontal=True,
             key="bg_type_radio"
         )
-        
-        # ==================================================================================
-        # BACKGROUND IMAGE UPLOAD
-        # ==================================================================================
-        
+
+        # Background Image
         if bg_type == "Image":
             bg_image = st.file_uploader(
                 "Upload Background Image",
                 type=["jpg", "png", "jpeg"],
                 key="bg_image_uploader"
             )
-            
-            # FIX: Only update session state when bg_image is not None
             current_bg_id = id(bg_image)
             if bg_image is not None and current_bg_id != st.session_state.last_bg_image_id:
-                logger.info(f"New background: {bg_image.name}")
+                logger.info(f"[UI] New background image: {bg_image.name}")
                 st.session_state.last_bg_image_id = current_bg_id
-                st.session_state.bg_image_cache = Image.open(bg_image)
-            
-            # Background preview with anti-shake placeholder
+                try:
+                    st.session_state.bg_image_cache = Image.open(bg_image)
+                    # Lazy convert to display; real convert is just before Stage 2
+                except Exception as e:
+                    st.session_state.bg_image_cache = None
+                    logger.error(f"[UI] Failed to load background image: {e}", exc_info=True)
+
             if st.session_state.bg_preview_placeholder is None:
                 st.session_state.bg_preview_placeholder = st.empty()
-            
+
             if st.session_state.bg_image_cache is not None:
                 with st.session_state.bg_preview_placeholder.container():
                     st.image(st.session_state.bg_image_cache, caption="Selected Background", use_container_width=True)
             else:
                 st.session_state.bg_preview_placeholder.empty()
-        
-        # ==================================================================================
-        # BACKGROUND COLOR PICKER
-        # ==================================================================================
-        
-        elif bg_type == "Color":
+
+        # Background Color
+        else:
             selected_color = st.color_picker(
                 "Choose Background Color",
                 st.session_state.bg_color,
                 key="color_picker"
             )
-            
             if selected_color != st.session_state.cached_color:
-                logger.info(f"Color: {selected_color}")
+                logger.info(f"[UI] New background color: {selected_color}")
                 st.session_state.bg_color = selected_color
                 st.session_state.cached_color = selected_color
-                
+
                 color_rgb = tuple(int(selected_color.lstrip('#')[i:i+2], 16) for i in (0, 2, 4))
                 color_display = np.zeros((100, 100, 3), dtype=np.uint8)
                 color_display[:, :] = color_rgb
                 st.session_state.color_display_cache = color_display
-            
+
             if st.session_state.bg_preview_placeholder is None:
                 st.session_state.bg_preview_placeholder = st.empty()
-            
+
             if st.session_state.color_display_cache is not None:
                 with st.session_state.bg_preview_placeholder.container():
                     st.image(st.session_state.color_display_cache, caption="Selected Color", width=200)
             else:
                 st.session_state.bg_preview_placeholder.empty()
-        
-        # ==================================================================================
-        # VIDEO PROCESSING SECTION
-        # ==================================================================================
-        
+
+        # ------------- Processing -------------
         st.header("3. Process Video")
-        
-        can_process = (
-            uploaded is not None and 
-            not st.session_state.processing
-        )
-        
+        can_process = (st.session_state.uploaded_video is not None and not st.session_state.processing)
+
         if st.button("Process Video", disabled=not can_process, use_container_width=True):
-            logger.info("=" * 60)
-            logger.info("VIDEO PROCESSING STARTED")
-            logger.info("=" * 60)
-            
+            run_id = uuid.uuid4().hex[:8]
+            logger.info("=" * 80)
+            logger.info(f"[RUN {run_id}] VIDEO PROCESSING STARTED at {datetime.utcnow().isoformat()}Z")
+            logger.info("=" * 80)
+
             st.session_state.processing = True
             st.session_state.processed_video_bytes = None
             st.session_state.last_error = None
-            
+
+            t0 = time.time()
             try:
-                # Determine background
-                background = None
-                if bg_type == "Image" and st.session_state.bg_image_cache is not None:
-                    background = st.session_state.bg_image_cache
-                    logger.info("Using IMAGE background")
-                elif bg_type == "Color":
+                # Validate background
+                if bg_type == "Image":
+                    if st.session_state.bg_image_cache is None:
+                        raise RuntimeError("Background type is Image, but no image was provided.")
+                    background = st.session_state.bg_image_cache.convert("RGB")
+                    logger.info(f"[RUN {run_id}] Using IMAGE background (RGB)")
+                else:
+                    if not st.session_state.bg_color:
+                        raise RuntimeError("Background type is Color, but no color was selected.")
                     background = st.session_state.bg_color
-                    logger.info(f"Using COLOR background: {background}")
-                
-                # Stage 1: Create transparent video (0-70%)
-                logger.info("Starting Stage 1: Transparent video creation...")
-                transparent_path = stage1_create_transparent_video(st.session_state.uploaded_video)
-                
+                    logger.info(f"[RUN {run_id}] Using COLOR background: {background}")
+
+                # Materialize uploaded video to temp file
+                if st.session_state.video_bytes_cache is None:
+                    logger.info(f"[RUN {run_id}] Reading uploaded video into cache for processing...")
+                    st.session_state.uploaded_video.seek(0)
+                    st.session_state.video_bytes_cache = st.session_state.uploaded_video.read()
+                    logger.info(f"[RUN {run_id}] Cached {len(st.session_state.video_bytes_cache)/1e6:.2f} MB for processing")
+
+                suffix = Path(st.session_state.uploaded_video.name).suffix or ".mp4"
+                with NamedTemporaryFile(delete=False, suffix=suffix) as tmp_vid:
+                    tmp_vid.write(st.session_state.video_bytes_cache)
+                    tmp_vid_path = tmp_vid.name
+
+                logger.info(f"[RUN {run_id}] Temp video path: {tmp_vid_path}")
+
+                # Stage 1
+                t1 = time.time()
+                logger.info(f"[RUN {run_id}] Stage 1: Transparent video creation START")
+                transparent_path = stage1_create_transparent_video(tmp_vid_path)
                 if not transparent_path or not os.path.exists(transparent_path):
                     raise RuntimeError("Stage 1 failed: Transparent video not created")
-                
-                logger.info(f"Stage 1 complete: {transparent_path}")
-                
-                # Stage 2: Composite with background (70-100%)
-                logger.info("Starting Stage 2: Compositing with background...")
-                final_path = stage2_composite_background(
-                    transparent_path,
-                    background,
-                    bg_type.lower()
-                )
-                
+                logger.info(f"[RUN {run_id}] Stage 1: DONE → {transparent_path} (Δ {time.time()-t1:.2f}s)")
+
+                # Stage 2
+                t2 = time.time()
+                logger.info(f"[RUN {run_id}] Stage 2: Compositing START (bg_type={bg_type.lower()})")
+                final_path = stage2_composite_background(transparent_path, background, bg_type.lower())
                 if not final_path or not os.path.exists(final_path):
                     raise RuntimeError("Stage 2 failed: Final video not created")
-                
-                logger.info(f"Stage 2 complete: {final_path}")
-                
-                # Load final video into session state
+                logger.info(f"[RUN {run_id}] Stage 2: DONE → {final_path} (Δ {time.time()-t2:.2f}s)")
+
+                # Load final into memory (Streamlit download)
                 with open(final_path, 'rb') as f:
                     st.session_state.processed_video_bytes = f.read()
-                
-                logger.info(f"Processing SUCCESS: {len(st.session_state.processed_video_bytes)/1e6:.2f}MB")
+
+                total = time.time() - t0
+                logger.info(f"[RUN {run_id}] SUCCESS size={len(st.session_state.processed_video_bytes)/1e6:.2f}MB, total Δ={total:.2f}s")
                 st.success("Video processing complete!")
                 st.session_state.last_error = None
-                    
+
             except Exception as e:
-                error_msg = f"Processing Error: {str(e)}\n\nCheck logs for details."
-                logger.error(f"Exception: {str(e)}")
+                total = time.time() - t0
+                error_msg = f"[RUN {run_id}] Processing Error: {str(e)} (Δ {total:.2f}s)\n\nCheck logs for details."
+                logger.error(error_msg)
                 logger.error(traceback.format_exc())
                 st.session_state.last_error = error_msg
                 st.error(error_msg)
-                
+
             finally:
                 st.session_state.processing = False
-        
-        # ==================================================================================
-        # RESULTS DISPLAY
-        # ==================================================================================
-        
+                logger.info(f"[RUN {run_id}] Processing finished (processing flag cleared)")
+
+        # Results
         if st.session_state.processed_video_bytes is not None:
             st.markdown("---")
             st.markdown("### Processed Video")
-            
             try:
                 st.video(st.session_state.processed_video_bytes)
                 st.download_button(
@@ -372,12 +483,12 @@ def main():
                     use_container_width=True
                 )
             except Exception as e:
-                logger.error(f"Display error: {e}")
+                logger.error(f"[UI] Display error: {e}", exc_info=True)
                 st.error(f"Display error: {e}")
 
-# ==================================================================================
-# CHAPTER 6: APPLICATION ENTRY POINT
-# ==================================================================================
+# =========================================
+# CHAPTER 7: ENTRY POINT
+# =========================================
 
 if __name__ == "__main__":
-    main() 
+    main()