Update streamlit_app.py

streamlit_app.py

@@ -1,4 +1,4 @@
-# streamlit_ui.py - Using video_pipeline.py with anti-shake fixes
+# streamlit_ui.py - Simplified single-button workflow
 
 import streamlit as st
 import os
@@ -49,7 +49,6 @@ def check_gpu():
     if cuda_available:
         logger.info(f"CUDA Version: {torch.version.cuda}")
         logger.info(f"Device Count: {torch.cuda.device_count()}")
-        logger.info(f"Current Device: {torch.cuda.current_device()}")
         logger.info(f"Device Name: {torch.cuda.get_device_name(0)}")
         
         torch.cuda.set_device(0)
@@ -100,14 +99,12 @@ def initialize_session_state():
         'bg_color': "#00FF00",
         'cached_color': None,
         'color_display_cache': None,
-        'transparent_video_path': None,  # Stage 1 output (reusable)
-        'processed_video_path': None,     # Stage 2 output
+        'processed_video_bytes': None,
         'processing': False,
         'gpu_available': None,
         'last_video_id': None,
         'last_bg_image_id': None,
-        'last_error': None,
-        'current_stage': None
+        'last_error': None
     }
     for key, value in defaults.items():
         if key not in st.session_state:
@@ -129,21 +126,6 @@ def main():
             st.success(f"GPU: {torch.cuda.get_device_name(0)}")
         else:
             st.error("GPU: Not Available")
-        
-        st.markdown("---")
-        st.subheader("Pipeline Status")
-        
-        if st.session_state.transparent_video_path:
-            st.success("Stage 1: Transparent video ready")
-        else:
-            st.info("Stage 1: Not started")
-        
-        if st.session_state.processed_video_path:
-            st.success("Stage 2: Final video ready")
-        elif st.session_state.transparent_video_path:
-            st.info("Stage 2: Ready to composite")
-        else:
-            st.info("Stage 2: Waiting for Stage 1")
     
     # Display last error if exists
     if st.session_state.last_error:
@@ -171,8 +153,7 @@ def main():
             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.transparent_video_path = None  # Reset pipeline
-            st.session_state.processed_video_path = None
+            st.session_state.processed_video_bytes = None
             st.session_state.last_error = None
         
         # Video preview with anti-shake placeholder
@@ -263,109 +244,81 @@ def main():
         
         st.header("3. Process Video")
         
-        # Stage 1: Create Transparent Video
-        can_process_stage1 = (
+        can_process = (
             st.session_state.uploaded_video is not None and 
-            not st.session_state.processing and
-            st.session_state.transparent_video_path is None
-        )
-        
-        if st.button("Stage 1: Create Transparent Video", disabled=not can_process_stage1, use_container_width=True):
-            logger.info("=" * 60)
-            logger.info("STAGE 1 STARTED")
-            logger.info("=" * 60)
-            
-            st.session_state.processing = True
-            st.session_state.current_stage = 1
-            st.session_state.last_error = None
-            
-            try:
-                transparent_path = stage1_create_transparent_video(st.session_state.uploaded_video)
-                
-                if transparent_path and os.path.exists(transparent_path):
-                    st.session_state.transparent_video_path = transparent_path
-                    logger.info(f"Stage 1 SUCCESS: {transparent_path}")
-                    st.success("Stage 1 complete! Transparent video created.")
-                else:
-                    error_msg = "Stage 1 failed to create transparent video"
-                    logger.error(error_msg)
-                    st.session_state.last_error = error_msg
-                    st.error(error_msg)
-                    
-            except Exception as e:
-                error_msg = f"Stage 1 Error: {str(e)}"
-                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
-                st.session_state.current_stage = None
-        
-        # Stage 2: Composite with Background
-        can_process_stage2 = (
-            st.session_state.transparent_video_path is not None and 
             not st.session_state.processing
         )
         
-        if st.button("Stage 2: Composite with Background", disabled=not can_process_stage2, use_container_width=True):
+        if st.button("Process Video", disabled=not can_process, use_container_width=True):
             logger.info("=" * 60)
-            logger.info("STAGE 2 STARTED")
+            logger.info("VIDEO PROCESSING STARTED")
             logger.info("=" * 60)
             
             st.session_state.processing = True
-            st.session_state.current_stage = 2
+            st.session_state.processed_video_bytes = None
             st.session_state.last_error = None
             
             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":
                     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)
+                
+                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(
-                    st.session_state.transparent_video_path,
+                    transparent_path,
                     background,
                     bg_type.lower()
                 )
                 
-                if final_path and os.path.exists(final_path):
-                    st.session_state.processed_video_path = final_path
-                    logger.info(f"Stage 2 SUCCESS: {final_path}")
-                    st.success("Stage 2 complete! Final video ready.")
-                else:
-                    error_msg = "Stage 2 failed to create final video"
-                    logger.error(error_msg)
-                    st.session_state.last_error = error_msg
-                    st.error(error_msg)
+                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
+                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")
+                st.success("Video processing complete!")
+                st.session_state.last_error = None
                     
             except Exception as e:
-                error_msg = f"Stage 2 Error: {str(e)}"
-                logger.error(error_msg)
+                error_msg = f"Processing Error: {str(e)}\n\nCheck logs for details."
+                logger.error(f"Exception: {str(e)}")
                 logger.error(traceback.format_exc())
                 st.session_state.last_error = error_msg
                 st.error(error_msg)
                 
             finally:
                 st.session_state.processing = False
-                st.session_state.current_stage = None
         
-        # Show results
-        if st.session_state.processed_video_path and os.path.exists(st.session_state.processed_video_path):
+        # Show processed video if available
+        if st.session_state.processed_video_bytes is not None:
             st.markdown("---")
-            st.markdown("### Final Video")
+            st.markdown("### Processed Video")
             
             try:
-                with open(st.session_state.processed_video_path, 'rb') as f:
-                    video_bytes = f.read()
-                
-                st.video(video_bytes)
+                st.video(st.session_state.processed_video_bytes)
                 st.download_button(
-                    label="Download Final Video",
-                    data=video_bytes,
-                    file_name="final_video.mp4",
+                    label="Download Processed Video",
+                    data=st.session_state.processed_video_bytes,
+                    file_name="processed_video.mp4",
                     mime="video/mp4",
                     use_container_width=True
                 )