Update app.py

app.py

@@ -13,6 +13,7 @@ import tempfile
 import uuid
 import time
 import threading
+from queue import Queue
 
 torch.set_float32_matmul_precision("medium")
 
@@ -61,10 +62,16 @@ cleanup_thread = threading.Thread(target=cleanup_temp_files, daemon=True)
 cleanup_thread.start()
 
 
+# Function to process frames in a separate thread
+def process_frame(image, bg, fast_mode, result_queue):
+    processed_image = process(image, bg, fast_mode)
+    result_queue.put(processed_image)
+
+
 @spaces.GPU
-def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=0, video_handling="slow_down", fast_mode=False):
+def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=0, video_handling="slow_down", fast_mode=True):
     try:
-        start_time = time.time()  # Start the timer
+        start_time = time.time()
 
         # Load the video using moviepy
         video = mp.VideoFileClip(vid)
@@ -96,29 +103,44 @@ def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=
 
         bg_frame_index = 0  # Initialize background frame index
 
+        threads = []
+        result_queue = Queue()
+        frame_batch_size = 4 # Process 4 frames at a time
+
         for i, frame in enumerate(frames):
             pil_image = Image.fromarray(frame)
             if bg_type == "Color":
-                processed_image = process(pil_image, color, fast_mode)
+                background_image = color # Use color directly for threads
             elif bg_type == "Image":
-                processed_image = process(pil_image, bg_image, fast_mode)
+                background_image = bg_image
             elif bg_type == "Video":
                 if video_handling == "slow_down":
                     background_frame = background_frames[bg_frame_index % len(background_frames)]
                     bg_frame_index += 1
                     background_image = Image.fromarray(background_frame)
-                    processed_image = process(pil_image, background_image, fast_mode)
                 else:  # video_handling == "loop"
                     background_frame = background_frames[bg_frame_index % len(background_frames)]
                     bg_frame_index += 1
                     background_image = Image.fromarray(background_frame)
-                    processed_image = process(pil_image, background_image, fast_mode)
             else:
-                processed_image = pil_image  # Default to original image if no background is selected
+                background_image = None  # Default to no background
+
+            # Start a new thread to process the frame
+            thread = threading.Thread(target=process_frame, args=(pil_image, background_image, fast_mode, result_queue))
+            threads.append(thread)
+            thread.start()
+
+            # If we have enough threads running or it's the last frame, wait for results
+            if len(threads) == frame_batch_size or i == len(list(frames)) - 1:
+                for thread in threads:
+                    thread.join()
+                while not result_queue.empty():
+                    processed_frames.append(np.array(result_queue.get()))
+                threads = [] # Reset the threads list
 
-            processed_frames.append(np.array(processed_image))
             elapsed_time = time.time() - start_time
-            yield processed_image, None, f"Processing frame {i+1}... Elapsed time: {elapsed_time:.2f} seconds"
+            # Yield the first processed image from the current batch
+            yield processed_frames[-len(threads) if threads else -frame_batch_size], None, f"Processing frame {i+1}... Elapsed time: {elapsed_time:.2f} seconds"
 
         # Create a new video from the processed frames
         processed_video = mp.ImageSequenceClip(processed_frames, fps=fps)
@@ -136,7 +158,7 @@ def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=
         elapsed_time = time.time() - start_time
         yield gr.update(visible=False), gr.update(visible=True), f"Processing complete! Elapsed time: {elapsed_time:.2f} seconds"
         # Return the path to the temporary file
-        yield processed_image, temp_filepath, f"Processing complete! Elapsed time: {elapsed_time:.2f} seconds"
+        yield processed_frames[-1], temp_filepath, f"Processing complete! Elapsed time: {elapsed_time:.2f} seconds"
 
     except Exception as e:
         print(f"Error: {e}")
@@ -195,7 +217,7 @@ with gr.Blocks(theme=gr.themes.Ocean()) as demo:
         bg_video = gr.Video(label="Background Video", visible=False, interactive=True)
         with gr.Column(visible=False) as video_handling_options:
             video_handling_radio = gr.Radio(["slow_down", "loop"], label="Video Handling", value="slow_down", interactive=True)
-        fast_mode_checkbox = gr.Checkbox(label="Fast Mode (Use BiRefNet_lite)", value=False, interactive=True)
+        fast_mode_checkbox = gr.Checkbox(label="Fast Mode (Use BiRefNet_lite)", value=True, interactive=True)
 
     time_textbox = gr.Textbox(label="Time Elapsed", interactive=False) # Add time textbox