Update streamlit_app.py

streamlit_app.py

@@ -15,12 +15,9 @@
 import uuid
 from datetime import datetime
 from tempfile import NamedTemporaryFile
-import subprocess
 import streamlit as st
-import cv2
-import numpy as np
-from PIL import Image
 import torch
+
 # Import UI components
 from ui import render_ui
 # Import pipeline functions
@@ -30,11 +27,15 @@
     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)
@@ -56,11 +57,18 @@ def setup_logging(level: int = logging.INFO) -> logging.Logger:
     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
+
+# --- Load models ONCE globally (never per session) ---
+sam2_predictor = load_sam2()
+matanyone_processor = load_matanyone()
+
 # --- Session State Initialization ---
 def initialize_session_state():
     defaults = {
@@ -87,6 +95,7 @@ def initialize_session_state():
             st.session_state[k] = v
     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()
@@ -97,6 +106,7 @@ def set_log_level(name: str):
     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):
     run_id = uuid.uuid4().hex[:8]
@@ -118,13 +128,22 @@ def process_video(uploaded_video, bg_image, bg_color, bg_type):
             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)
+        transparent_path, audio_path = stage1_create_transparent_video(
+            tmp_vid_path,
+            sam2_predictor=sam2_predictor,
+            matanyone_processor=matanyone_processor
+        )
         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
-        final_path = stage2_composite_background(transparent_path, audio_path, background, bg_type.lower())
+        final_path = stage2_composite_background(
+            transparent_path,
+            audio_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"[RUN {run_id}] Stage 2 completed: Final path={final_path}")
@@ -144,6 +163,7 @@ def process_video(uploaded_video, bg_image, bg_color, bg_type):
     finally:
         st.session_state.processing = False
         logger.info(f"[RUN {run_id}] Processing finished")
+
 # --- Main App Entry Point ---
 def main():
     try:
@@ -158,6 +178,7 @@ def main():
     except Exception as e:
         logger.error(f"Main app error: {e}", exc_info=True)
         st.error(f"App startup failed: {str(e)}. Check logs for details.")
+
 if __name__ == "__main__":
     setup_t4_environment()
-    main() 
+    main()