Update app.py

app.py

@@ -13,6 +13,7 @@ import tempfile
 import uuid
 import time
 import threading
+from concurrent.futures import ThreadPoolExecutor, as_completed
 
 torch.set_float32_matmul_precision("medium")
 
@@ -60,6 +61,33 @@ def cleanup_temp_files():
 cleanup_thread = threading.Thread(target=cleanup_temp_files, daemon=True)
 cleanup_thread.start()
 
+def process(image, bg, fast_mode=False):
+    image_size = image.size
+    input_images = transform_image(image).unsqueeze(0).to(device)
+
+    # Select the model based on fast_mode
+    model = birefnet_lite if fast_mode else birefnet
+
+    # Prediction
+    with torch.no_grad():
+        preds = model(input_images)[-1].sigmoid().cpu()
+    pred = preds[0].squeeze()
+    pred_pil = transforms.ToPILImage()(pred)
+    mask = pred_pil.resize(image_size)
+
+    if isinstance(bg, str) and bg.startswith("#"):
+        color_rgb = tuple(int(bg[i:i+2], 16) for i in (1, 3, 5))
+        background = Image.new("RGBA", image_size, color_rgb + (255,))
+    elif isinstance(bg, Image.Image):
+        background = bg.convert("RGBA").resize(image_size)
+    else:
+        background = Image.open(bg).convert("RGBA").resize(image_size)
+
+    # Composite the image onto the background using the mask
+    image = Image.composite(image, background, mask)
+
+    return 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=True):
@@ -77,11 +105,11 @@ 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
+        # Process frames in parallel
         processed_frames = []
-        yield gr.update(visible=True), gr.update(visible=False), f"Processing started... Elapsed time: 0 seconds"
+        yield gr.update(visible=True), gr.update(visible=False), "Processing started... Elapsed time: 0 seconds"
 
         if bg_type == "Video":
             background_video = mp.VideoFileClip(bg_video)
@@ -96,7 +124,8 @@ 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):
+        # Define a helper function for processing a single frame
+        def process_single_frame(i, frame):
             pil_image = Image.fromarray(frame)
             if bg_type == "Color":
                 processed_image = process(pil_image, color, fast_mode)
@@ -105,23 +134,41 @@ def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=
             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)
+                nonlocal bg_frame_index
+                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"
+            return i, np.array(processed_image)
+
+        with ThreadPoolExecutor(max_workers=4) as executor:
+            # Submit all frame processing tasks
+            future_to_index = {executor.submit(process_single_frame, i, frame): i for i, frame in enumerate(frames)}
+            
+            # As each future completes, process the result
+            for future in as_completed(future_to_index):
+                i, processed_image = future.result()
+                processed_frames.append((i, processed_image))
+                
+                # Update elapsed time
+                elapsed_time = time.time() - start_time
+                # Sort the processed_frames based on index to maintain order
+                processed_frames_sorted = sorted(processed_frames, key=lambda x: x[0])
+                
+                # Yield the first processed image if it's available
+                if len(processed_frames_sorted) == 1:
+                    first_image = Image.fromarray(processed_frames_sorted[0][1])
+                    yield first_image, None, f"Processing frame {processed_frames_sorted[0][0]+1}... Elapsed time: {elapsed_time:.2f} seconds"
+        
+        # Sort all processed frames
+        processed_frames_sorted = sorted(processed_frames, key=lambda x: x[0])
+        final_frames = [frame for i, frame in processed_frames_sorted]
 
         # Create a new video from the processed frames
-        processed_video = mp.ImageSequenceClip(processed_frames, fps=fps)
+        processed_video = mp.ImageSequenceClip(final_frames, fps=fps)
 
         # Add the original audio back to the processed video
         processed_video = processed_video.set_audio(audio)
@@ -136,42 +183,12 @@ 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 None, temp_filepath, f"Processing complete! Elapsed time: {elapsed_time:.2f} seconds"
 
     except Exception as e:
         print(f"Error: {e}")
         elapsed_time = time.time() - start_time
         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")
-
-    # Select the model based on fast_mode
-    model = birefnet_lite if fast_mode else birefnet
-
-    # Prediction
-    with torch.no_grad():
-        preds = model(input_images)[-1].sigmoid().cpu()
-    pred = preds[0].squeeze()
-    pred_pil = transforms.ToPILImage()(pred)
-    mask = pred_pil.resize(image_size)
-
-    if isinstance(bg, str) and bg.startswith("#"):
-        color_rgb = tuple(int(bg[i:i+2], 16) for i in (1, 3, 5))
-        background = Image.new("RGBA", image_size, color_rgb + (255,))
-    elif isinstance(bg, Image.Image):
-        background = bg.convert("RGBA").resize(image_size)
-    else:
-        background = Image.open(bg).convert("RGBA").resize(image_size)
-
-    # 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.")