Update streamlit_app.py

streamlit_app.py

@@ -1,54 +1,41 @@
 #!/usr/bin/env python3
 """
-Core Application Logic for Video Background Replacer
-- Handles video processing pipeline (SAM2 + MatAnyone + FFmpeg)
-- Integrates with UI (imported from ui.py)
-- Manages session state and logging
+Advanced Video Background Replacer - Streamlit Entrypoint
 """
 import os
 import sys
 import time
 from pathlib import Path
 import logging
-import logging.handlers
 import traceback
 import uuid
-from datetime import datetime
 from tempfile import NamedTemporaryFile
 import streamlit as st
-import torch
 
-# Import UI components
 from ui import render_ui
-# Import pipeline functions
 from pipeline.video_pipeline import (
     stage1_create_transparent_video,
     stage2_composite_background,
     setup_t4_environment,
     check_gpu
 )
-# Import model loaders (the new robust way)
 from models.model_loaders import load_sam2, load_matanyone
 
-# --- Constants ---
 APP_NAME = "Advanced Video Background Replacer"
 LOG_FILE = "/tmp/app.log"
 LOG_MAX_BYTES = 5 * 1024 * 1024
 LOG_BACKUPS = 5
 
-# --- Logging Setup ---
 def setup_logging(level: int = logging.INFO) -> logging.Logger:
     logger = logging.getLogger(APP_NAME)
     logger.setLevel(level)
     logger.propagate = False
-    # Clear previous handlers on rerun
+    # Remove previous handlers (Streamlit reruns)
     for h in list(logger.handlers):
         logger.removeHandler(h)
-    # Console handler
     ch = logging.StreamHandler(sys.stdout)
     ch.setLevel(level)
     ch.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
-    # Rotating file handler
     fh = logging.handlers.RotatingFileHandler(
         LOG_FILE, maxBytes=LOG_MAX_BYTES, backupCount=LOG_BACKUPS, encoding="utf-8"
     )
@@ -60,16 +47,14 @@ def setup_logging(level: int = logging.INFO) -> logging.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
 
-# --- Load models ONCE globally (never per session) ---
+# Only load once
 sam2_predictor = load_sam2()
 matanyone_processor = load_matanyone()
 
-# --- Session State Initialization ---
 def initialize_session_state():
     defaults = {
         'uploaded_video': None,
@@ -88,7 +73,8 @@ def initialize_session_state():
         'last_error': None,
         'log_level_name': 'INFO',
         'auto_refresh_logs': False,
-        'log_tail_lines': 400
+        'log_tail_lines': 400,
+        'generated_bg': None,
     }
     for k, v in defaults.items():
         if k not in st.session_state:
@@ -96,38 +82,23 @@ def initialize_session_state():
     if st.session_state.gpu_available is None:
         st.session_state.gpu_available = check_gpu(logger)
 
-# --- Set Log Level ---
-def set_log_level(name: str):
-    name = (name or "INFO").upper()
-    lvl = getattr(logging, name, logging.INFO)
-    setup_logging(lvl)
-    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}")
-
-# --- Main Processing Function ---
-def process_video(uploaded_video, bg_image, bg_color, bg_type):
+def process_video(uploaded_video, background, bg_type):
     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(f"[RUN {run_id}] VIDEO PROCESSING STARTED at {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())}")
     logger.info(f"[RUN {run_id}] Video size={len(uploaded_video.read()) / 1e6:.2f}MB, BG type={bg_type}")
-    uploaded_video.seek(0)  # Reset for later read
-    logger.info("=" * 80)
+    uploaded_video.seek(0)
     st.session_state.processing = True
     st.session_state.processed_video_bytes = None
     st.session_state.last_error = None
     t0 = time.time()
     try:
-        # Materialize uploaded video to temp file
         suffix = Path(uploaded_video.name).suffix or ".mp4"
         with NamedTemporaryFile(delete=False, suffix=suffix) as tmp_vid:
             uploaded_video.seek(0)
             tmp_vid.write(uploaded_video.read())
             tmp_vid_path = tmp_vid.name
         logger.info(f"[RUN {run_id}] Temporary video path: {tmp_vid_path}")
-        # Stage 1: Create transparent video and extract audio
         transparent_path, audio_path = stage1_create_transparent_video(
             tmp_vid_path,
             sam2_predictor=sam2_predictor,
@@ -136,8 +107,7 @@ def process_video(uploaded_video, bg_image, bg_color, bg_type):
         if not transparent_path or not os.path.exists(transparent_path):
             raise RuntimeError("Stage 1 failed: Transparent video not created")
         logger.info(f"[RUN {run_id}] Stage 1 completed: Transparent path={transparent_path}, Audio path={audio_path}")
-        # Stage 2: Composite with background and restore audio
-        background = bg_image.convert("RGB") if bg_type == "Image" else bg_color
+        # For "color", background is "#RRGGBB", for others it's PIL.Image
         final_path = stage2_composite_background(
             transparent_path,
             audio_path,
@@ -147,7 +117,6 @@ def process_video(uploaded_video, bg_image, bg_color, bg_type):
         if not final_path or not os.path.exists(final_path):
             raise RuntimeError("Stage 2 failed: Final video not created")
         logger.info(f"[RUN {run_id}] Stage 2 completed: Final path={final_path}")
-        # Load final video into memory for download
         with open(final_path, 'rb') as f:
             st.session_state.processed_video_bytes = f.read()
         total = time.time() - t0
@@ -164,7 +133,6 @@ def process_video(uploaded_video, bg_image, bg_color, bg_type):
         st.session_state.processing = False
         logger.info(f"[RUN {run_id}] Processing finished")
 
-# --- Main App Entry Point ---
 def main():
     try:
         st.set_page_config(