Update app.py

app.py

@@ -13,20 +13,17 @@ import tempfile
 import uuid
 import time
 import threading
+from concurrent.futures import ThreadPoolExecutor
 
 torch.set_float32_matmul_precision("medium")
-
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load both BiRefNet models
 birefnet = AutoModelForImageSegmentation.from_pretrained(
-    "ZhengPeng7/BiRefNet", trust_remote_code=True
-)
+    "ZhengPeng7/BiRefNet", trust_remote_code=True)
 birefnet.to(device)
-
 birefnet_lite = AutoModelForImageSegmentation.from_pretrained(
-    "ZhengPeng7/BiRefNet_lite", trust_remote_code=True
-)
+    "ZhengPeng7/BiRefNet_lite", trust_remote_code=True)
 birefnet_lite.to(device)
 
 transform_image = transforms.Compose(
@@ -37,7 +34,6 @@ transform_image = transforms.Compose(
     ]
 )
 
-
 # Function to delete files older than 10 minutes in the temp directory
 def cleanup_temp_files():
     while True:
@@ -55,11 +51,29 @@ def cleanup_temp_files():
                             print(f"Error deleting file {filepath}: {e}")
         time.sleep(60)  # Check every minute
 
-
 # Start the cleanup thread
 cleanup_thread = threading.Thread(target=cleanup_temp_files, daemon=True)
 cleanup_thread.start()
 
+# Function to process a single frame
+def process_frame(frame, bg_type, bg, fast_mode, bg_frame_index, background_frames, color):
+    try:
+        pil_image = Image.fromarray(frame)
+        if bg_type == "Color":
+            processed_image = process(pil_image, color, fast_mode)
+        elif bg_type == "Image":
+            processed_image = process(pil_image, bg, fast_mode)
+        elif bg_type == "Video":
+            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
+        return np.array(processed_image), bg_frame_index
+    except Exception as e:
+        print(f"Error processing frame: {e}")
+        return frame, bg_frame_index
 
 @spaces.GPU
 def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=0, video_handling="slow_down", fast_mode=True):
@@ -77,7 +91,7 @@ def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=
         audio = video.audio
 
         # Extract frames at the specified FPS
-        frames = video.iter_frames(fps=fps)
+        frames = list(video.iter_frames(fps=fps))
 
         # Process each frame for background removal
         processed_frames = []
@@ -96,29 +110,14 @@ def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=
 
         bg_frame_index = 0  # Initialize background frame index
 
-        for i, frame in enumerate(frames):
-            pil_image = Image.fromarray(frame)
-            if bg_type == "Color":
-                processed_image = process(pil_image, color, fast_mode)
-            elif bg_type == "Image":
-                processed_image = process(pil_image, bg_image, fast_mode)
-            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
-
-            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"
+        # Use ThreadPoolExecutor for parallel processing
+        with ThreadPoolExecutor(max_workers=4) as executor:
+            futures = [executor.submit(process_frame, frames[i], bg_type, bg_image, fast_mode, bg_frame_index, background_frames, color) for i in range(len(frames))]
+            for future in futures:
+                result, bg_frame_index = future.result()
+                processed_frames.append(result)
+                elapsed_time = time.time() - start_time
+                yield result, None, f"Processing frame {len(processed_frames)}... Elapsed time: {elapsed_time:.2f} seconds"
 
         # Create a new video from the processed frames
         processed_video = mp.ImageSequenceClip(processed_frames, fps=fps)
@@ -135,8 +134,9 @@ 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}")
@@ -144,7 +144,6 @@ def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=
         yield gr.update(visible=False), gr.update(visible=True), f"Error processing video: {e}. Elapsed time: {elapsed_time:.2f} seconds"
         yield None, f"Error processing video: {e}", f"Error processing video: {e}. Elapsed time: {elapsed_time:.2f} seconds"
 
-
 def process(image, bg, fast_mode=False):
     image_size = image.size
     input_images = transform_image(image).unsqueeze(0).to("cuda")
@@ -169,12 +168,10 @@ def process(image, bg, fast_mode=False):
 
     # Composite the image onto the background using the mask
     image = Image.composite(image, background, mask)
-
     return image
 
-
 with gr.Blocks(theme=gr.themes.Ocean()) as demo:
-    gr.Markdown("# Video Background Remover & Changer\n### You can replace image background with any color, image or video.\nNOTE: As this Space is running on ZERO GPU it has limit. It can handle approx 200frmaes at once. So, if you have big video than use small chunks or Duplicate this space.")
+    gr.Markdown("# Video Background Remover & Changer\n### You can replace image background with any color, image or video.\nNOTE: As this Space is running on ZERO GPU it has limit. It can handle approx 200 frames at once. So, if you have a big video than use small chunks or Duplicate this space.")
     with gr.Row():
         in_video = gr.Video(label="Input Video", interactive=True)
         stream_image = gr.Image(label="Streaming Output", visible=False)
@@ -196,8 +193,7 @@ with gr.Blocks(theme=gr.themes.Ocean()) as demo:
         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=True, interactive=True)
-
-    time_textbox = gr.Textbox(label="Time Elapsed", interactive=False) # Add time textbox
+    time_textbox = gr.Textbox(label="Time Elapsed", interactive=False)  # Add time textbox
 
     def update_visibility(bg_type):
         if bg_type == "Color":