Update app.py

app.py

@@ -1,8 +1,8 @@
 #!/usr/bin/env python3
 """
-BackgroundFX - Video Background Replacement with Green Screen Workflow
-Fixed for Hugging Face Space - Handles video preview issues
-FIXED: Video display issue by properly handling file stream
+BackgroundFX - Fast Video Background Replacement
+Optimized with Rembg for immediate deployment on HF Space with T4 GPU
+Ready to run in 30 minutes!
 """
 
 import streamlit as st
@@ -14,50 +14,80 @@
 import requests
 from io import BytesIO
 import logging
-import base64
+import torch
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Try to import SAM2 and MatAnyone
+# Check GPU
+CUDA_AVAILABLE = torch.cuda.is_available()
+if CUDA_AVAILABLE:
+    logger.info(f"✅ GPU: {torch.cuda.get_device_name(0)}")
+else:
+    logger.warning("⚠️ Running on CPU")
+
+# Import rembg - the main workhorse
 try:
-    from sam2.build_sam import build_sam2_video_predictor
-    from sam2.sam2_image_predictor import SAM2ImagePredictor
-    SAM2_AVAILABLE = True
-    logger.info("✅ SAM2 loaded successfully")
-except ImportError as e:
-    SAM2_AVAILABLE = False
-    logger.warning(f"⚠️ SAM2 not available: {e}")
+    from rembg import remove, new_session
+    REMBG_AVAILABLE = True
+    logger.info("✅ Rembg loaded")
+except ImportError:
+    REMBG_AVAILABLE = False
+    st.error("❌ Please install rembg: pip install rembg")
 
+# Import MatAnyone for better matting
 try:
-    import matanyone
+    from matanyone import MatAnyone
     MATANYONE_AVAILABLE = True
-    logger.info("✅ MatAnyone loaded successfully")
-except ImportError as e:
+    logger.info("✅ MatAnyone loaded for edge refinement")
+except ImportError:
     MATANYONE_AVAILABLE = False
-    logger.warning(f"⚠️ MatAnyone not available: {e}")
+    logger.info("ℹ️ MatAnyone not available - using standard matting")
+
+# Global session cache
+@st.cache_resource
+def load_rembg_model():
+    """Load and cache the Rembg model"""
+    if REMBG_AVAILABLE:
+        # u2net_human_seg is specifically for people
+        session = new_session('u2net_human_seg')
+        logger.info("✅ U2NET Human Segmentation model loaded")
+        return session
+    return None
+
+@st.cache_resource
+def load_matanyone_model():
+    """Load and cache MatAnyone model for edge refinement"""
+    if MATANYONE_AVAILABLE:
+        try:
+            model = MatAnyone()
+            if CUDA_AVAILABLE:
+                model = model.cuda()
+            logger.info("✅ MatAnyone model loaded for edge refinement")
+            return model
+        except Exception as e:
+            logger.warning(f"Failed to load MatAnyone: {e}")
+            return None
+    return None
 
 def get_video_info(video_path):
-    """Get video information and first frame"""
+    """Get video information"""
     try:
         cap = cv2.VideoCapture(video_path)
         if not cap.isOpened():
             return None, None
         
-        # Get video properties
         fps = int(cap.get(cv2.CAP_PROP_FPS))
         width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         duration = total_frames / fps if fps > 0 else 0
         
-        # Get first frame for thumbnail
         ret, first_frame = cap.read()
         cap.release()
         
         if ret:
-            # Convert BGR to RGB
             first_frame_rgb = cv2.cvtColor(first_frame, cv2.COLOR_BGR2RGB)
             return {
                 'fps': fps,
@@ -68,192 +98,174 @@ def get_video_info(video_path):
             }, first_frame_rgb
         
         return None, None
-        
     except Exception as e:
-        logger.error(f"Error getting video info: {e}")
+        logger.error(f"Error: {e}")
         return None, None
 
-def load_background_image(background_url):
-    """Load background image from URL"""
-    try:
-        response = requests.get(background_url)
-        response.raise_for_status()
-        image = Image.open(BytesIO(response.content))
-        return np.array(image.convert('RGB'))
-    except Exception as e:
-        logger.error(f"Failed to load background image: {e}")
-        return create_default_background()
-
-def create_default_background():
-    """Create a default brick wall background"""
-    height, width = 720, 1280
-    background = np.ones((height, width, 3), dtype=np.uint8) * 150
-    
-    # Add brick pattern
-    brick_height, brick_width = 40, 80
-    for y in range(0, height, brick_height):
-        for x in range(0, width, brick_width):
-            offset = brick_width // 2 if (y // brick_height) % 2 else 0
-            x_pos = (x + offset) % width
-            
-            cv2.rectangle(background, 
-                         (x_pos, y), 
-                         (min(x_pos + brick_width - 2, width), min(y + brick_height - 2, height)), 
-                         (180, 120, 80), -1)
-            cv2.rectangle(background, 
-                         (x_pos, y), 
-                         (min(x_pos + brick_width - 2, width), min(y + brick_height - 2, height)), 
-                         (120, 80, 40), 2)
-    
+def create_gradient_background(width=1280, height=720, color1=(70, 130, 180), color2=(255, 140, 90)):
+    """Create a nice gradient background"""
+    background = np.zeros((height, width, 3), dtype=np.uint8)
+    for y in range(height):
+        blend = y / height
+        color = [
+            int(color1[0] * (1 - blend) + color2[0] * blend),
+            int(color1[1] * (1 - blend) + color2[1] * blend),
+            int(color1[2] * (1 - blend) + color2[2] * blend)
+        ]
+        background[y, :] = color
     return background
 
-def segment_person_sam2(frame):
-    """Segment person using SAM2"""
+def process_frame_rembg(frame, session, background, use_green_screen=False, matanyone_model=None):
+    """Process a single frame with Rembg and optional MatAnyone refinement"""
     try:
-        predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large")
-        predictor.set_image(frame)
+        # Get RGBA output from rembg
+        frame_pil = Image.fromarray(frame)
+        output = remove(frame_pil, session=session, alpha_matting=True)  # Enable alpha matting for better edges
+        output_np = np.array(output)
+        
+        # Extract alpha channel
+        if output_np.shape[2] == 4:
+            alpha = output_np[:, :, 3].astype(float) / 255.0
+            person_rgb = output_np[:, :, :3]
+        else:
+            alpha = np.ones(frame.shape[:2])
+            person_rgb = output_np
+        
+        # Apply MatAnyone for edge refinement if available
+        if matanyone_model is not None and MATANYONE_AVAILABLE:
+            try:
+                # MatAnyone expects torch tensors
+                import torch
+                
+                # Convert to tensor
+                frame_tensor = torch.from_numpy(frame).float().permute(2, 0, 1).unsqueeze(0) / 255.0
+                alpha_tensor = torch.from_numpy(alpha).float().unsqueeze(0).unsqueeze(0)
+                
+                if CUDA_AVAILABLE:
+                    frame_tensor = frame_tensor.cuda()
+                    alpha_tensor = alpha_tensor.cuda()
+                
+                # Refine alpha with MatAnyone
+                with torch.no_grad():
+                    refined_alpha = matanyone_model.refine_alpha(frame_tensor, alpha_tensor)
+                
+                # Convert back to numpy
+                alpha = refined_alpha.squeeze().cpu().numpy()
+                
+                logger.info("✅ Applied MatAnyone edge refinement")
+            except Exception as e:
+                logger.warning(f"MatAnyone refinement failed, using original alpha: {e}")
         
+        # Resize background
         h, w = frame.shape[:2]
-        center_point = np.array([[w//2, h//2]])
-        center_label = np.array([1])
-        
-        masks, scores, _ = predictor.predict(
-            point_coords=center_point,
-            point_labels=center_label,
-            multimask_output=False
-        )
-        
-        return masks[0] if len(masks) > 0 else None
-        
-    except Exception as e:
-        logger.error(f"SAM2 segmentation failed: {e}")
-        return None
-
-def segment_person_fallback(frame):
-    """Fallback person segmentation using color-based method"""
-    try:
-        hsv = cv2.cvtColor(frame, cv2.COLOR_RGB2HSV)
-        
-        lower_skin = np.array([0, 20, 70])
-        upper_skin = np.array([20, 255, 255])
-        
-        skin_mask = cv2.inRange(hsv, lower_skin, upper_skin)
-        
-        kernel = np.ones((5, 5), np.uint8)
-        skin_mask = cv2.morphologyEx(skin_mask, cv2.MORPH_CLOSE, kernel)
-        skin_mask = cv2.morphologyEx(skin_mask, cv2.MORPH_OPEN, kernel)
-        
-        contours, _ = cv2.findContours(skin_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        bg_resized = cv2.resize(background, (w, h))
         
-        if contours:
-            largest_contour = max(contours, key=cv2.contourArea)
-            mask = np.zeros(frame.shape[:2], dtype=np.uint8)
-            cv2.fillPoly(mask, [largest_contour], 255)
-            kernel = np.ones((20, 20), np.uint8)
-            mask = cv2.dilate(mask, kernel, iterations=2)
-            return mask.astype(bool)
-        
-        return None
-        
-    except Exception as e:
-        logger.error(f"Fallback segmentation failed: {e}")
-        return None
-
-def insert_green_screen(frame, person_mask):
-    """Insert green screen background while preserving person"""
-    try:
-        green_background = np.zeros_like(frame)
-        green_background[:, :] = [0, 255, 0]  # Pure green
+        if use_green_screen:
+            # Green screen workflow with refined alpha
+            green = np.zeros_like(frame)
+            green[:, :] = [0, 255, 0]
+            
+            # Composite person on green with smooth alpha
+            green_composite = np.zeros_like(frame, dtype=np.float32)
+            for c in range(3):
+                green_composite[:, :, c] = alpha * person_rgb[:, :, c] + (1 - alpha) * green[:, :, c]
+            
+            green_composite = green_composite.astype(np.uint8)
+            
+            # Replace green with background
+            lower_green = np.array([0, 200, 0])
+            upper_green = np.array([100, 255, 100])
+            green_mask = cv2.inRange(green_composite, lower_green, upper_green)
+            green_mask_inv = cv2.bitwise_not(green_mask)
+            
+            result = cv2.bitwise_and(green_composite, green_composite, mask=green_mask_inv)
+            bg_part = cv2.bitwise_and(bg_resized, bg_resized, mask=green_mask)
+            result = cv2.add(result, bg_part)
+        else:
+            # Direct composite with refined alpha (faster and usually better)
+            # Apply slight edge smoothing to alpha
+            alpha_smooth = cv2.GaussianBlur(alpha.astype(np.float32), (3, 3), 0)
+            
+            result = np.zeros_like(frame, dtype=np.float32)
+            for c in range(3):
+                result[:, :, c] = alpha_smooth * person_rgb[:, :, c] + (1 - alpha_smooth) * bg_resized[:, :, c]
+            
+            result = result.astype(np.uint8)
         
-        result = np.where(person_mask[..., None], frame, green_background)
         return result
         
     except Exception as e:
-        logger.error(f"Green screen insertion failed: {e}")
+        logger.error(f"Frame processing error: {e}")
         return frame
 
-def replace_background(frame, person_mask, background):
-    """Replace background with new image"""
+def process_video_fast(video_path, output_path, background, progress_callback=None, 
+                       skip_frames=1, use_green_screen=False, use_matanyone=True):
+    """Fast video processing with Rembg and optional MatAnyone"""
     try:
-        # Resize background to match frame dimensions
-        h, w = frame.shape[:2]
-        background_resized = cv2.resize(background, (w, h))
-        
-        # Create smooth edges
-        mask_float = person_mask.astype(np.float32)
-        kernel = np.ones((5, 5), np.float32) / 25
-        mask_float = cv2.filter2D(mask_float, -1, kernel)
-        
-        # Composite the images
-        result = np.zeros_like(frame)
-        for c in range(3):
-            result[:, :, c] = (mask_float * frame[:, :, c] + 
-                              (1 - mask_float) * background_resized[:, :, c])
-        
-        return result.astype(np.uint8)
+        # Load models once
+        session = load_rembg_model()
+        if session is None:
+            st.error("Failed to load Rembg model")
+            return False
+        
+        # Load MatAnyone if requested and available
+        matanyone_model = None
+        if use_matanyone and MATANYONE_AVAILABLE:
+            matanyone_model = load_matanyone_model()
+            if matanyone_model:
+                st.info("✨ Using MatAnyone for edge refinement")
         
-    except Exception as e:
-        logger.error(f"Background replacement failed: {e}")
-        return frame
-
-def process_video(video_path, output_path, background_image, use_green_screen=True, progress_callback=None):
-    """Process video with background replacement"""
-    try:
         cap = cv2.VideoCapture(video_path)
         
-        # Get video properties
         fps = int(cap.get(cv2.CAP_PROP_FPS))
         width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         
-        # Setup video writer
+        # Adjust FPS if skipping frames
+        output_fps = max(fps // skip_frames, 15)  # Minimum 15 fps
+        
+        # Use MP4V codec for compatibility
         fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+        out = cv2.VideoWriter(output_path, fourcc, output_fps, (width, height))
         
         frame_count = 0
+        processed_count = 0
         
         while True:
             ret, frame = cap.read()
             if not ret:
                 break
             
-            # Convert BGR to RGB for processing
-            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frame_count += 1
             
-            # Segment person
-            if SAM2_AVAILABLE:
-                person_mask = segment_person_sam2(frame_rgb)
-            else:
-                person_mask = segment_person_fallback(frame_rgb)
+            # Skip frames for speed
+            if skip_frames > 1 and (frame_count - 1) % skip_frames != 0:
+                continue
             
-            if person_mask is not None:
-                if use_green_screen:
-                    # First create green screen
-                    frame_rgb = insert_green_screen(frame_rgb, person_mask)
-                    # Then replace green with background
-                    # Create green mask
-                    lower_green = np.array([0, 200, 0])
-                    upper_green = np.array([100, 255, 100])
-                    green_mask = cv2.inRange(frame_rgb, lower_green, upper_green)
-                    green_mask = green_mask.astype(bool)
-                    # Replace green areas with background
-                    frame_rgb = replace_background(frame_rgb, ~green_mask, background_image)
-                else:
-                    # Direct background replacement
-                    frame_rgb = replace_background(frame_rgb, person_mask, background_image)
+            # Process frame with MatAnyone if available
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            processed_frame = process_frame_rembg(
+                frame_rgb, session, background, use_green_screen, matanyone_model
+            )
             
-            # Convert back to BGR for writing
-            frame_bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)
+            # Write frame
+            frame_bgr = cv2.cvtColor(processed_frame, cv2.COLOR_RGB2BGR)
             out.write(frame_bgr)
             
-            frame_count += 1
-            if progress_callback:
-                progress_callback(frame_count / total_frames)
+            processed_count += 1
+            
+            if progress_callback and frame_count % 5 == 0:  # Update progress every 5 frames
+                progress = frame_count / total_frames
+                progress_callback(progress)
         
         cap.release()
         out.release()
         
+        # Clear GPU memory
+        if CUDA_AVAILABLE:
+            torch.cuda.empty_cache()
+        
         return True
         
     except Exception as e:
@@ -263,185 +275,226 @@ def process_video(video_path, output_path, background_image, use_green_screen=Tr
 # Streamlit UI
 def main():
     st.set_page_config(
-        page_title="BackgroundFX",
-        page_icon="🎬",
+        page_title="BackgroundFX - Fast",
+        page_icon="🚀",
         layout="wide"
     )
     
-    st.title("🎬 BackgroundFX - Video Background Replacement")
-    st.markdown("Replace video backgrounds with AI-powered segmentation")
+    st.title("🚀 BackgroundFX - Fast Background Replacement")
+    st.markdown("**Optimized for speed** - Using Rembg U2NET for human segmentation")
     
-    # Check dependencies
-    if not SAM2_AVAILABLE:
-        st.warning("⚠️ SAM2 not available - using fallback segmentation method")
-    if not MATANYONE_AVAILABLE:
-        st.info("ℹ️ MatAnyone not available - using standard matting")
+    # Quick status check
+    cols = st.columns(4)
+    with cols[0]:
+        if CUDA_AVAILABLE:
+            st.success(f"✅ GPU Active")
+        else:
+            st.warning("⚠️ CPU Mode")
+    with cols[1]:
+        if REMBG_AVAILABLE:
+            st.success("✅ Rembg Ready")
+        else:
+            st.error("❌ Install rembg")
+    with cols[2]:
+        if MATANYONE_AVAILABLE:
+            st.success("✅ MatAnyone")
+        else:
+            st.info("ℹ️ Basic Matting")
+    with cols[3]:
+        st.info("⚡ Fast Mode")
     
-    # Sidebar for settings
-    with st.sidebar:
-        st.header("⚙️ Settings")
-        
-        use_green_screen = st.checkbox(
-            "Use Green Screen Workflow",
-            value=True,
-            help="First create green screen, then replace with background"
+    # Two columns layout
+    col1, col2 = st.columns(2)
+    
+    with col1:
+        st.header("📹 Upload Video")
+        
+        # Speed preset at the top for visibility
+        speed_mode = st.select_slider(
+            "⚡ Speed Mode",
+            options=["Quality", "Balanced", "Fast", "Ultra Fast"],
+            value="Fast"
         )
         
-        st.subheader("📸 Background Options")
-        bg_option = st.radio(
-            "Choose background source:",
-            ["Default Brick Wall", "Custom URL", "Upload Image"]
+        # MatAnyone edge refinement option
+        use_matanyone = st.checkbox(
+            "✨ Use MatAnyone Edge Refinement",
+            value=MATANYONE_AVAILABLE,
+            disabled=not MATANYONE_AVAILABLE,
+            help="Improves edges around hair and clothing (adds ~20% processing time)"
         )
         
-        background_image = None
+        # Set parameters based on mode
+        if speed_mode == "Ultra Fast":
+            skip_frames = 3
+            use_green = False
+            st.caption("⚡ Every 3rd frame, direct compositing")
+        elif speed_mode == "Fast":
+            skip_frames = 2
+            use_green = False
+            st.caption("⚡ Every 2nd frame, direct compositing")
+        elif speed_mode == "Balanced":
+            skip_frames = 1
+            use_green = False
+            st.caption("⚡ All frames, direct compositing")
+        else:  # Quality
+            skip_frames = 1
+            use_green = True
+            st.caption("⚡ All frames, green screen + edge refinement")
         
-        if bg_option == "Custom URL":
-            bg_url = st.text_input(
-                "Background Image URL",
-                value="https://images.unsplash.com/photo-1557683316-973673baf926",
-                help="Enter a direct image URL"
-            )
-            if bg_url:
-                with st.spinner("Loading background..."):
-                    background_image = load_background_image(bg_url)
-                    if background_image is not None:
-                        st.success("✅ Background loaded")
-                        st.image(background_image, caption="Background Preview", use_column_width=True)
-        
-        elif bg_option == "Upload Image":
-            uploaded_bg = st.file_uploader(
-                "Upload Background Image",
-                type=['jpg', 'jpeg', 'png'],
-                help="Upload your own background image"
-            )
-            if uploaded_bg is not None:
-                background_image = np.array(Image.open(uploaded_bg).convert('RGB'))
-                st.success("✅ Background uploaded")
-                st.image(background_image, caption="Background Preview", use_column_width=True)
-        
-        else:  # Default Brick Wall
-            background_image = create_default_background()
-            st.info("Using default brick wall background")
-            st.image(background_image, caption="Default Background", use_column_width=True)
-    
-    # Main content area
-    col1, col2 = st.columns(2)
-    
-    with col1:
-        st.header("📹 Input Video")
         uploaded_video = st.file_uploader(
-            "Upload your video",
+            "Choose video file",
             type=['mp4', 'avi', 'mov', 'mkv'],
-            help="Upload a video to process"
+            help="For best results, use videos under 30 seconds"
         )
         
         if uploaded_video is not None:
-            # FIXED: Read bytes once and reuse
+            # Save video
             video_bytes = uploaded_video.read()
             
-            # Save to temporary file for processing
             with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tmp_file:
                 tmp_file.write(video_bytes)
-                temp_video_path = tmp_file.name
+                temp_path = tmp_file.name
             
-            # Display video using the original bytes (not consuming the stream)
             st.video(video_bytes)
             
-            # Get video info using the temp file
-            video_info, first_frame = get_video_info(temp_video_path)
-            
-            if video_info and first_frame is not None:
-                st.success(f"✅ Video loaded: {video_info['width']}x{video_info['height']}, "
-                          f"{video_info['fps']} fps, {video_info['duration']:.1f}s")
+            # Get info
+            info, first_frame = get_video_info(temp_path)
+            if info:
+                st.success(f"✅ Ready: {info['duration']:.1f}s @ {info['fps']}fps")
                 
-                # Show first frame as thumbnail
-                st.image(first_frame, caption="First Frame Preview", use_column_width=True)
+                # Time estimate
+                process_time = (info['duration'] / skip_frames) * 0.5  # Rough estimate
+                st.info(f"⏱️ Estimated time: {process_time:.0f} seconds")
                 
-                # Store paths in session state
-                if 'temp_video_path' not in st.session_state:
-                    st.session_state.temp_video_path = temp_video_path
+                st.session_state.video_path = temp_path
+                st.session_state.video_info = info
+        
+        # Background selection
+        st.subheader("🎨 Background")
+        bg_type = st.radio("Choose:", ["Gradient", "Color", "Image URL", "Upload"])
+        
+        background = None
+        
+        if bg_type == "Gradient":
+            col_a, col_b = st.columns(2)
+            with col_a:
+                color1 = st.color_picker("Top", "#4682B4")
+            with col_b:
+                color2 = st.color_picker("Bottom", "#FF8C5A")
+            
+            # Convert hex to RGB
+            c1 = tuple(int(color1[i:i+2], 16) for i in (1, 3, 5))
+            c2 = tuple(int(color2[i:i+2], 16) for i in (1, 3, 5))
+            
+            if 'video_info' in st.session_state:
+                w = st.session_state.video_info['width']
+                h = st.session_state.video_info['height']
+                background = create_gradient_background(w, h, c1, c2)
             else:
-                st.error("Failed to read video information")
+                background = create_gradient_background(1280, 720, c1, c2)
+        
+        elif bg_type == "Color":
+            color = st.color_picker("Pick color", "#00FF00")
+            rgb = tuple(int(color[i:i+2], 16) for i in (1, 3, 5))
+            background = np.full((720, 1280, 3), rgb, dtype=np.uint8)
+        
+        elif bg_type == "Image URL":
+            url = st.text_input("Image URL", "https://images.unsplash.com/photo-1557683316-973673baf926")
+            if url:
+                try:
+                    response = requests.get(url)
+                    img = Image.open(BytesIO(response.content))
+                    background = np.array(img.convert('RGB'))
+                    st.image(background, caption="Background", use_column_width=True)
+                except:
+                    st.error("Failed to load image")
+        
+        else:  # Upload
+            uploaded_bg = st.file_uploader("Upload image", type=['jpg', 'jpeg', 'png'])
+            if uploaded_bg:
+                img = Image.open(uploaded_bg)
+                background = np.array(img.convert('RGB'))
+                st.image(background, caption="Background", use_column_width=True)
     
     with col2:
-        st.header("🎯 Output")
+        st.header("🎬 Result")
         
-        if uploaded_video is not None and background_image is not None:
-            if st.button("🚀 Process Video", type="primary"):
+        if uploaded_video and background is not None:
+            if st.button("🚀 Process Video", type="primary", use_container_width=True):
+                
+                # Check if rembg is available
+                if not REMBG_AVAILABLE:
+                    st.error("Please install rembg first!")
+                    st.code("pip install rembg", language="bash")
+                    return
+                
                 try:
-                    # Create output path
                     output_path = tempfile.mktemp(suffix='.mp4')
                     
-                    # Process video with progress bar
-                    progress_bar = st.progress(0)
-                    status_text = st.empty()
+                    progress = st.progress(0)
+                    status = st.empty()
                     
-                    def update_progress(progress):
-                        progress_bar.progress(progress)
-                        status_text.text(f"Processing: {int(progress * 100)}%")
+                    def update_progress(value):
+                        progress.progress(value)
+                        status.text(f"Processing: {int(value * 100)}%")
                     
-                    status_text.text("Starting video processing...")
+                    status.text("🔄 Starting processing...")
                     
-                    # Use the temp file path we saved
-                    success = process_video(
-                        st.session_state.temp_video_path,
+                    success = process_video_fast(
+                        st.session_state.video_path,
                         output_path,
-                        background_image,
-                        use_green_screen,
-                        update_progress
+                        background,
+                        update_progress,
+                        skip_frames,
+                        use_green_screen=use_green,
+                        use_matanyone=use_matanyone
                     )
                     
                     if success and os.path.exists(output_path):
-                        status_text.text("✅ Processing complete!")
+                        status.text("✅ Done!")
                         
-                        # Read the processed video
                         with open(output_path, 'rb') as f:
-                            processed_video = f.read()
+                            result_video = f.read()
                         
-                        # Display processed video
-                        st.video(processed_video)
+                        st.video(result_video)
                         
-                        # Download button
                         st.download_button(
-                            label="📥 Download Processed Video",
-                            data=processed_video,
-                            file_name="backgroundfx_output.mp4",
-                            mime="video/mp4"
+                            "💾 Download Result",
+                            data=result_video,
+                            file_name=f"backgroundfx_{speed_mode.lower()}.mp4",
+                            mime="video/mp4",
+                            use_container_width=True
                         )
                         
-                        # Cleanup
+                        # Show stats
+                        size_mb = len(result_video) / (1024 * 1024)
+                        st.success(f"✅ Output size: {size_mb:.1f} MB")
+                        
                         os.unlink(output_path)
                     else:
-                        st.error("❌ Video processing failed")
-                    
+                        st.error("Processing failed!")
+                
                 except Exception as e:
-                    st.error(f"Error during processing: {str(e)}")
+                    st.error(f"Error: {str(e)}")
                     logger.error(f"Processing error: {e}")
-        
-        elif uploaded_video is None:
-            st.info("👈 Please upload a video to begin")
-        elif background_image is None:
-            st.info("👈 Please select or upload a background image")
-    
-    # Cleanup temporary files on session end
-    if 'temp_video_path' in st.session_state and os.path.exists(st.session_state.temp_video_path):
-        try:
-            os.unlink(st.session_state.temp_video_path)
-            del st.session_state.temp_video_path
-        except:
-            pass
+        else:
+            if not uploaded_video:
+                st.info("👈 Upload a video to start")
+            else:
+                st.info("👈 Select a background")
     
-    # Footer
+    # Footer with tips
     st.markdown("---")
-    st.markdown(
-        """
-        <div style='text-align: center'>
-            <p>BackgroundFX v1.0 | AI-Powered Video Background Replacement</p>
-            <p>Using SAM2 for person segmentation and green screen technology</p>
-        </div>
-        """,
-        unsafe_allow_html=True
-    )
+    with st.expander("💡 Quick Tips"):
+        st.markdown("""
+        - **Ultra Fast**: Best for quick previews (3x faster)
+        - **Fast**: Good balance of speed and quality (2x faster)
+        - **Balanced**: Full quality, still fast
+        - **Quality**: Best edges with green screen workflow
+        - Videos under 30 seconds process fastest
+        - Gradient backgrounds render instantly
+        """)
 
 if __name__ == "__main__":
     main()