Update streamlit_app.py

streamlit_app.py

@@ -53,17 +53,19 @@ def check_gpu():
         logger.info(f"Device Name: {torch.cuda.get_device_name(0)}")
         logger.info(f"Device Capability: {torch.cuda.get_device_capability(0)}")
         
+        # Set as default device
+        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 Tensor Created Successfully on: {test_tensor.device}")
+            logger.info(f"GPU Test Tensor Created 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(f"PyTorch Version: {torch.__version__}")
-        logger.info(f"CUDA Built Version: {torch.version.cuda}")
     
     logger.info("=" * 60)
     return cuda_available
@@ -85,12 +87,6 @@ def check_gpu():
     .stButton>button:hover {
         background-color: #45a049;
     }
-    .stProgress > div > div > div > div {
-        background-color: #4CAF50;
-    }
-    .stAlert {
-        border-radius: 10px;
-    }
     .video-container {
         border: 2px dashed #4CAF50;
         border-radius: 10px;
@@ -106,36 +102,35 @@ def initialize_session_state():
     defaults = {
         'uploaded_video': None,
         'video_bytes_cache': None,
-        'video_preview_placeholder': None,
         'bg_image_cache': None,
-        'bg_preview_placeholder': None,
         'bg_color': "#00FF00",
-        'cached_color': None,
         'color_display_cache': None,
         'processed_video_bytes': None,
         'processing': False,
-        'process_complete': False,
-        'last_video_id': None,
-        'last_bg_image_id': None,
-        'gpu_available': None
+        'gpu_available': None,
+        'files_loaded': False
     }
     for key, value in defaults.items():
         if key not in st.session_state:
             st.session_state[key] = value
     
-    # Run GPU check on first init
     if st.session_state.gpu_available is None:
         st.session_state.gpu_available = check_gpu()
 
 # --- Video Processing ---
 def process_video(input_file, background, bg_type="image"):
-    """Process video with the selected background using SAM2 and MatAnyone pipeline."""
+    """Process video with GPU acceleration"""
     logger.info("=" * 60)
     logger.info("STARTING VIDEO PROCESSING")
     logger.info("=" * 60)
     
     try:
-        # Create a persistent temporary directory
+        # Ensure GPU is default device
+        if torch.cuda.is_available():
+            torch.cuda.set_device(0)
+            logger.info("GPU set as default device for processing")
+        
+        # Create temp directory
         temp_base = Path(tempfile.gettempdir()) / "video_processing"
         temp_base.mkdir(exist_ok=True)
         temp_dir = temp_base / f"session_{int(time.time())}"
@@ -143,19 +138,17 @@ def process_video(input_file, background, bg_type="image"):
         
         logger.info(f"Temp directory: {temp_dir}")
 
-        # Save the uploaded video
+        # Save video
         input_path = str(temp_dir / "input.mp4")
-        logger.info(f"Writing video to {input_path}")
+        logger.info(f"Saving video to {input_path}")
         
         input_file.seek(0)
+        video_data = input_file.read()
         with open(input_path, "wb") as f:
-            written = f.write(input_file.read())
-            logger.info(f"Wrote {written/1e6:.2f}MB")
+            f.write(video_data)
         
-        if not os.path.exists(input_path) or os.path.getsize(input_path) == 0:
-            raise FileNotFoundError(f"Input video not saved properly")
-        
-        logger.info(f"Video file verified: {os.path.getsize(input_path)} bytes")
+        file_size = os.path.getsize(input_path)
+        logger.info(f"Video saved: {file_size} bytes ({file_size/1e6:.2f}MB)")
 
         # Prepare background
         bg_path = None
@@ -164,7 +157,7 @@ def process_video(input_file, background, bg_type="image"):
             bg_cv = cv2.cvtColor(np.array(background), cv2.COLOR_RGB2BGR)
             bg_path = str(temp_dir / "background.jpg")
             cv2.imwrite(bg_path, bg_cv)
-            logger.info(f"Background image written")
+            logger.info(f"Background saved")
             
         elif bg_type == "color":
             logger.info(f"Processing background COLOR: {st.session_state.bg_color}")
@@ -172,91 +165,83 @@ def process_video(input_file, background, bg_type="image"):
             color_rgb = tuple(int(color_hex[i:i+2], 16) for i in (0, 2, 4))
             bg_path = str(temp_dir / "background.jpg")
             cv2.imwrite(bg_path, np.ones((100, 100, 3), dtype=np.uint8) * color_rgb[::-1])
-            logger.info(f"Background color image written")
+            logger.info(f"Background color saved")
 
         # Progress tracking
-        progress_placeholder = st.empty()
-        status_placeholder = st.empty()
+        progress_bar = st.progress(0)
+        status_text = st.empty()
         
         def progress_callback(progress, message):
             progress = max(0, min(1, float(progress)))
             logger.info(f"Progress: {progress*100:.1f}% - {message}")
-            progress_placeholder.progress(progress)
-            status_placeholder.text(f"Status: {message}")
+            progress_bar.progress(progress)
+            status_text.text(f"{message} ({progress*100:.0f}%)")
 
-        # GPU check
+        # GPU status before processing
         if torch.cuda.is_available():
-            device = torch.cuda.get_device_name(0)
-            mem_before = torch.cuda.memory_allocated()/1e9
-            logger.info(f"CUDA Device: {device}")
-            logger.info(f"GPU Memory Before: {mem_before:.2f}GB")
-        else:
-            logger.warning("CUDA NOT AVAILABLE - Processing on CPU")
+            logger.info(f"GPU Memory Before: {torch.cuda.memory_allocated()/1e9:.2f}GB")
 
-        # Process the video
+        # Process video
         output_path = str(temp_dir / "output.mp4")
         click_points = [[0.5, 0.5]]
         
         logger.info("Importing TwoStageProcessor...")
         from pipeline.integrated_pipeline import TwoStageProcessor
 
-        @st.cache_resource
-        def load_processor(temp_dir_str):
-            logger.info(f"Loading processor")
-            return TwoStageProcessor(temp_dir=temp_dir_str)
-        
-        processor = load_processor(str(temp_dir))
-        logger.info("Processor loaded")
+        # Don't cache processor - causes GPU issues
+        logger.info("Creating processor instance...")
+        processor = TwoStageProcessor(temp_dir=str(temp_dir))
+        logger.info("Processor created")
 
-        try:
-            logger.info("Calling process_video...")
-            
-            success = processor.process_video(
-                input_video=input_path,
-                background_video=bg_path if bg_type == "image" else "",
-                click_points=click_points,
-                output_path=output_path,
-                use_matanyone=True,
-                progress_callback=progress_callback
-            )
-            
-            logger.info(f"Processing returned: {success}")
-            
-        except Exception as e:
-            logger.error(f"Pipeline processing failed: {traceback.format_exc()}")
-            raise
+        logger.info("Starting process_video call...")
+        logger.info(f"  Input: {input_path}")
+        logger.info(f"  Background: {bg_path}")
+        logger.info(f"  Output: {output_path}")
+        
+        success = processor.process_video(
+            input_video=input_path,
+            background_video=bg_path if bg_type == "image" else "",
+            click_points=click_points,
+            output_path=output_path,
+            use_matanyone=True,
+            progress_callback=progress_callback
+        )
+        
+        logger.info(f"process_video returned: {success}")
 
         if torch.cuda.is_available():
-            mem_after = torch.cuda.memory_allocated()/1e9
-            logger.info(f"GPU Memory After: {mem_after:.2f}GB")
+            logger.info(f"GPU Memory After: {torch.cuda.memory_allocated()/1e9:.2f}GB")
             torch.cuda.empty_cache()
 
         if not success:
-            raise RuntimeError("Video processing returned False")
+            raise RuntimeError("Processing returned False")
 
-        # Verify output
+        # Check output
         if not os.path.exists(output_path):
-            raise FileNotFoundError(f"Output video not created")
+            raise FileNotFoundError("Output video not created")
         
         output_size = os.path.getsize(output_path)
         logger.info(f"Output: {output_size} bytes ({output_size/1e6:.2f}MB)")
 
-        # Read output into memory
-        logger.info("Reading output into memory...")
+        # Read output
+        logger.info("Reading output...")
         with open(output_path, 'rb') as f:
             video_bytes = f.read()
         
-        logger.info(f"Read {len(video_bytes)/1e6:.2f}MB")
+        logger.info(f"Output loaded: {len(video_bytes)/1e6:.2f}MB")
 
         # Cleanup
         try:
             shutil.rmtree(temp_dir)
-            logger.info("Temp directory cleaned")
+            logger.info("Temp cleaned")
         except Exception as e:
-            logger.warning(f"Could not clean temp: {e}")
+            logger.warning(f"Cleanup warning: {e}")
 
+        progress_bar.empty()
+        status_text.empty()
+        
         logger.info("=" * 60)
-        logger.info("VIDEO PROCESSING COMPLETED")
+        logger.info("PROCESSING COMPLETED SUCCESSFULLY")
         logger.info("=" * 60)
         
         return video_bytes
@@ -264,7 +249,7 @@ def load_processor(temp_dir_str):
     except Exception as e:
         logger.error(f"ERROR: {str(e)}")
         logger.error(traceback.format_exc())
-        st.error(f"An error occurred: {str(e)}")
+        st.error(f"Processing error: {str(e)}")
         return None
 
 # --- Main Application ---
@@ -272,10 +257,9 @@ def main():
     st.title("Advanced Video Background Replacer")
     st.markdown("---")
     
-    # Initialize session state
     initialize_session_state()
     
-    # GPU Status in sidebar
+    # GPU Status
     with st.sidebar:
         st.subheader("System Status")
         if st.session_state.gpu_available:
@@ -285,7 +269,6 @@ def main():
     
     logger.info(f"Rerun - Processing: {st.session_state.processing}")
     
-    # Main layout
     col1, col2 = st.columns([1, 1], gap="large")
     
     with col1:
@@ -297,32 +280,34 @@ def main():
             key="video_uploader"
         )
         
-        # Check if video actually changed using id()
-        current_video_id = id(uploaded)
-        if current_video_id != st.session_state.last_video_id:
-            logger.info(f"New video: {uploaded.name if uploaded else 'None'}")
-            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.process_complete = False
+        # Handle file upload
+        if uploaded is not None:
+            # Always refresh if file exists (handles Streamlit restore bug)
+            if (st.session_state.uploaded_video is None or 
+                not hasattr(st.session_state.uploaded_video, 'name') or
+                st.session_state.uploaded_video.name != uploaded.name):
+                logger.info(f"Loading video: {uploaded.name} ({uploaded.size} bytes)")
+                st.session_state.uploaded_video = uploaded
+                st.session_state.video_bytes_cache = None
+                st.session_state.processed_video_bytes = None
+                st.session_state.files_loaded = True
         
-        # Video preview with placeholder
+        # Video preview
         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:
-            if st.session_state.video_bytes_cache is None:
-                logger.info("Caching video...")
-                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")
-            
-            with st.session_state.video_preview_placeholder.container():
+            try:
+                if st.session_state.video_bytes_cache is None:
+                    logger.info("Caching video...")
+                    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")
+                
                 st.video(st.session_state.video_bytes_cache)
+            except Exception as e:
+                logger.error(f"Video preview error: {e}")
+                st.error(f"Cannot display video: {e}")
         else:
-            st.session_state.video_preview_placeholder.empty()
+            st.info("No video uploaded")
     
     with col2:
         st.header("2. Background Settings")
@@ -341,25 +326,18 @@ def main():
                 key="bg_image_uploader"
             )
             
-            # Check if image actually changed using id()
-            current_bg_id = id(bg_image)
-            if current_bg_id != st.session_state.last_bg_image_id:
-                logger.info(f"New background: {bg_image.name if bg_image else 'None'}")
-                st.session_state.last_bg_image_id = current_bg_id
-                if bg_image is not None:
-                    st.session_state.bg_image_cache = Image.open(bg_image)
-                else:
-                    st.session_state.bg_image_cache = None
-            
-            # Background preview with placeholder
-            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)
+            if bg_image is not None:
+                if (st.session_state.bg_image_cache is None or
+                    not hasattr(st.session_state.bg_image_cache, 'filename') or
+                    getattr(st.session_state.bg_image_cache, 'filename', None) != bg_image.name):
+                    logger.info(f"Loading background: {bg_image.name}")
+                    img = Image.open(bg_image)
+                    img.filename = bg_image.name
+                    st.session_state.bg_image_cache = img
+                
+                st.image(st.session_state.bg_image_cache, caption="Selected Background", use_container_width=True)
             else:
-                st.session_state.bg_preview_placeholder.empty()
+                st.info("No background image")
         
         elif bg_type == "Color":
             selected_color = st.color_picker(
@@ -368,44 +346,38 @@ def main():
                 key="color_picker"
             )
             
-            if selected_color != st.session_state.cached_color:
-                logger.info(f"Color: {selected_color}")
+            if selected_color != st.session_state.bg_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[::-1]
                 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()
+                st.image(st.session_state.color_display_cache, caption="Selected Color", width=200)
         
         st.header("3. Process & Download")
         
-        if st.button(
-            "Process Video",
-            disabled=st.session_state.uploaded_video is None or st.session_state.processing,
-            use_container_width=True
-        ):
+        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("PROCESS BUTTON CLICKED")
+            logger.info("=" * 60)
+            
             st.session_state.processing = True
-            st.session_state.process_complete = False
             st.session_state.processed_video_bytes = None
             
-            with st.spinner("Processing video..."):
+            with st.spinner("Processing video (this may take several minutes)..."):
                 try:
                     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}")
                     
                     video_bytes = process_video(
                         st.session_state.uploaded_video,
@@ -415,21 +387,21 @@ def main():
                     
                     if video_bytes and len(video_bytes) > 0:
                         st.session_state.processed_video_bytes = video_bytes
-                        st.session_state.process_complete = True
-                        logger.info(f"Complete! {len(video_bytes)/1e6:.2f}MB")
+                        logger.info(f"SUCCESS: {len(video_bytes)/1e6:.2f}MB")
                         st.success("Video processing complete!")
                     else:
                         logger.error("Processing returned empty")
-                        st.error("Failed to process video")
+                        st.error("Processing failed")
                         
                 except Exception as e:
                     logger.error(f"Exception: {str(e)}")
+                    logger.error(traceback.format_exc())
                     st.error(f"Error: {str(e)}")
                 finally:
                     st.session_state.processing = False
         
         # Show processed video
-        if st.session_state.processed_video_bytes is not None and len(st.session_state.processed_video_bytes) > 0:
+        if st.session_state.processed_video_bytes is not None:
             st.markdown("---")
             st.markdown("### Processed Video")
             
@@ -443,8 +415,8 @@ def main():
                     use_container_width=True
                 )
             except Exception as e:
-                logger.error(f"Display error: {str(e)}")
-                st.error(f"Error: {str(e)}")
+                logger.error(f"Display error: {e}")
+                st.error(f"Display error: {e}")
 
 if __name__ == "__main__":
     main()