Update app.py

app.py

@@ -1,4 +1,4 @@
-import gradio as gr
+ import gradio as gr
 from loadimg import load_img
 import spaces
 from transformers import AutoModelForImageSegmentation
@@ -12,8 +12,8 @@ import os
 import tempfile
 import uuid
 import time
-import threading
 from concurrent.futures import ThreadPoolExecutor
+import asyncio
 
 torch.set_float32_matmul_precision("medium")
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -26,115 +26,92 @@ birefnet_lite = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet_lite", trust_remote_code=True)
 birefnet_lite.to(device)
 
-transform_image = transforms.Compose(
-    [
-        transforms.Resize((1024, 1024)),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-    ]
-)
-
-# 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
+transform_image = transforms.Compose([
+    transforms.Resize((1024, 1024)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+])
+
+# Function to process a single frame asynchronously
+async def process_frame_async(frame, bg_type, bg, fast_mode, bg_frame_index, background_frames, color):
+    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
 
 @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, max_workers=6):
-    try:
-        start_time = time.time()  # Start the timer
-
-        # Load the video using moviepy
-        video = mp.VideoFileClip(vid)
-
-        # Load original fps if fps value is equal to 0
-        if fps == 0:
-            fps = video.fps
-
-        # Extract audio from the video
-        audio = video.audio
-
-        # Extract frames at the specified FPS
-        frames = list(video.iter_frames(fps=fps))
-
-        # Process each frame for background removal
-        processed_frames = []
-        yield gr.update(visible=True), gr.update(visible=False), f"Processing started... Elapsed time: 0 seconds"
-
-        if bg_type == "Video":
-            background_video = mp.VideoFileClip(bg_video)
-            if background_video.duration < video.duration:
-                if video_handling == "slow_down":
-                    background_video = background_video.fx(mp.vfx.speedx, factor=video.duration / background_video.duration)
-                else:  # video_handling == "loop"
-                    background_video = mp.concatenate_videoclips([background_video] * int(video.duration / background_video.duration + 1))
-            background_frames = list(background_video.iter_frames(fps=fps))  # Convert to list
-        else:
-            background_frames = None
-
-        bg_frame_index = 0  # Initialize background frame index
-
-        # Use ThreadPoolExecutor for parallel processing with specified max_workers
-        with ThreadPoolExecutor(max_workers=max_workers) 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)
-
-        # Add the original audio back to the processed video
-        processed_video = processed_video.set_audio(audio)
-
-        # Save the processed video to a temporary file using tempfile
-        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_file:
-            temp_filepath = temp_file.name
-            processed_video.write_videofile(temp_filepath, codec="libx264")
-
-        elapsed_time = time.time() - start_time
-        yield gr.update(visible=False), gr.update(visible=True), f"Processing complete! Elapsed time: {elapsed_time:.2f} seconds"
+async def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=0, video_handling="slow_down", fast_mode=True, max_workers=6):
+    start_time = time.time()  # Start the timer
+    
+    video = mp.VideoFileClip(vid)
+    if fps == 0:
+        fps = video.fps
+    
+    audio = video.audio
+    frames = list(video.iter_frames(fps=fps))
+    
+    processed_frames = []
+    yield gr.update(visible=True), gr.update(visible=False), f"Processing started... Elapsed time: 0 seconds"
+    
+    if bg_type == "Video":
+        background_video = mp.VideoFileClip(bg_video)
+        if background_video.duration < video.duration:
+            if video_handling == "slow_down":
+                background_video = background_video.fx(mp.vfx.speedx, factor=video.duration / background_video.duration)
+            else:  # video_handling == "loop"
+                background_video = mp.concatenate_videoclips([background_video] * int(video.duration / background_video.duration + 1))
+        background_frames = list(background_video.iter_frames(fps=fps))
+    else:
+        background_frames = None
 
-        # Return the path to the temporary file
-        yield processed_frames[-1], temp_filepath, f"Processing complete! Elapsed time: {elapsed_time:.2f} seconds"
+    bg_frame_index = 0
 
-    except Exception as e:
-        print(f"Error: {e}")
+    # Use ThreadPoolExecutor for parallel processing with specified max_workers
+    loop = asyncio.get_event_loop()
+    tasks = [
+        loop.run_in_executor(
+            None, process_frame_async, frames[i], bg_type, bg_image, fast_mode, bg_frame_index, background_frames, color
+        )
+        for i in range(len(frames))
+    ]
+    
+    for future in asyncio.as_completed(tasks):
+        result, bg_frame_index = await future
+        processed_frames.append(result)
         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"
+        yield result, None, f"Processing frame {len(processed_frames)}... Elapsed time: {elapsed_time:.2f} seconds"
+    
+    processed_video = mp.ImageSequenceClip(processed_frames, fps=fps)
+    processed_video = processed_video.set_audio(audio)
+    
+    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_file:
+        temp_filepath = temp_file.name
+        processed_video.write_videofile(temp_filepath, codec="libx264")
+    
+    elapsed_time = time.time() - start_time
+    yield gr.update(visible=False), gr.update(visible=True), f"Processing complete! Elapsed time: {elapsed_time:.2f} seconds"
+    yield processed_frames[-1], temp_filepath, f"Processing complete! 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
+    input_images = transform_image(image).unsqueeze(0).to(device)
     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,))
@@ -142,18 +119,20 @@ def process(image, bg, fast_mode=False):
         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 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)
         out_video = gr.Video(label="Final Output Video")
+    
     submit_button = gr.Button("Change Background", interactive=True)
+    
     with gr.Row():
         fps_slider = gr.Slider(
             minimum=0,
@@ -167,12 +146,14 @@ with gr.Blocks(theme=gr.themes.Ocean()) as demo:
         color_picker = gr.ColorPicker(label="Background Color", value="#00FF00", visible=True, interactive=True)
         bg_image = gr.Image(label="Background Image", type="filepath", visible=False, interactive=True)
         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=True, interactive=True)
-        max_workers_slider = gr.Slider( minimum=1, maximum=32, step=1, value=6, label="Max Workers", info="Determines how many Franes to process parallel", interactive=True
-        )
-    time_textbox = gr.Textbox(label="Time Elapsed", interactive=False)  # Add time textbox
+        max_workers_slider = gr.Slider( minimum=1, maximum=32, step=1, value=6, label="Max Workers", info="Determines how many frames to process in parallel", interactive=True )
+
+    time_textbox = gr.Textbox(label="Time Elapsed", interactive=False)
 
     def update_visibility(bg_type):
         if bg_type == "Color":
@@ -184,10 +165,8 @@ with gr.Blocks(theme=gr.themes.Ocean()) as demo:
         else:
             return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
-
     bg_type.change(update_visibility, inputs=bg_type, outputs=[color_picker, bg_image, bg_video, video_handling_options])
 
-
     examples = gr.Examples(
         [
             ["rickroll-2sec.mp4", "Video", None, "background.mp4"],
@@ -201,7 +180,6 @@ with gr.Blocks(theme=gr.themes.Ocean()) as demo:
         cache_mode="eager",
     )
 
-
     submit_button.click(
         fn,
         inputs=[in_video, bg_type, bg_image, bg_video, color_picker, fps_slider, video_handling_radio, fast_mode_checkbox, max_workers_slider],
@@ -209,4 +187,4 @@ with gr.Blocks(theme=gr.themes.Ocean()) as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(show_error=True)
+    demo.launch(show_error=True)