Update app.py

app.py

@@ -17,6 +17,7 @@ torch.set_float32_matmul_precision("highest")
 birefnet = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet", trust_remote_code=True
 ).to("cuda")
+
 transform_image = transforms.Compose(
     [
         transforms.Resize((1024, 1024)),
@@ -25,6 +26,8 @@ transform_image = transforms.Compose(
     ]
 )
 
+BATCH_SIZE = 3
+
 @spaces.GPU
 def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=0, video_handling="slow_down"):
     try:
@@ -44,48 +47,44 @@ def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=
                 else:
                     background_video = mp.concatenate_videoclips([background_video] * int(video.duration / background_video.duration + 1))
             background_frames = list(background_video.iter_frames(fps=fps))
-        elif bg_type in ["Color", "Image"]:
-             # Prepare background once if it's a static image or color
-            if bg_type == "Color":
-                color_rgb = tuple(int(color[i:i+2], 16) for i in (1, 3, 5))
-                background_pil = Image.new("RGBA", (1024, 1024), color_rgb + (255,))
-            else: # bg_type == "Image":
-                background_pil = Image.open(bg_image).convert("RGBA").resize((1024, 1024))
-            background_tensor = transforms.ToTensor(background_pil).to("cuda")
         else:
-            background_tensor = None
-
+            background_frames = None
 
         bg_frame_index = 0
         frame_batch = []
+
         for i, frame in enumerate(frames):
-            frame = Image.fromarray(frame)
-            frame = transforms.ToTensor(frame).to('cuda')
             frame_batch.append(frame)
+            if len(frame_batch) == BATCH_SIZE or i == video.fps * video.duration -1:  # Process batch or last frames
+                pil_images = [Image.fromarray(f) for f in frame_batch]
+
 
-            if len(frame_batch) >= 3 or i == int(video.fps * video.duration) - 1 :
-                input_images = torch.stack(frame_batch).to("cuda")
-                with torch.no_grad():
-                    preds = birefnet(input_images)[-1].sigmoid()
-                for j, pred in enumerate(preds):
-                    if bg_type == "Video":
+                if bg_type == "Color":
+                    processed_images = [process(img, color) for img in pil_images]
+                elif bg_type == "Image":
+                    processed_images = [process(img, bg_image) for img in pil_images]
+                elif bg_type == "Video":
+                    processed_images = []
+                    for _ in range(len(frame_batch)):
                         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).resize((1024, 1024))
-                            background_tensor = transforms.ToTensor(background_image).to("cuda")
+                            background_image = Image.fromarray(background_frame)
                         else:  # video_handling == "loop"
                             background_frame = background_frames[bg_frame_index % len(background_frames)]
                             bg_frame_index += 1
-                            background_image = Image.fromarray(background_frame).resize((1024, 1024))
-                            background_tensor = transforms.ToTensor(background_image).to("cuda")
-                    mask = transforms.ToPILImage()(pred.cpu().squeeze())
-                    processed_image = Image.composite(transforms.ToPILImage()(frame_batch[j].cpu()), transforms.ToPILImage()(background_tensor.cpu()), mask).resize(video.size)
+                            background_image = Image.fromarray(background_frame)
+
+                        processed_images.append(process(pil_images[_],background_image))
 
+
+                else:
+                    processed_images = pil_images
+
+                for processed_image in processed_images:
                     processed_frames.append(np.array(processed_image))
                     yield processed_image, None
-
-                frame_batch = []
+                frame_batch = []  # Clear the batch
 
 
         processed_video = mp.ImageSequenceClip(processed_frames, fps=fps)
@@ -107,6 +106,30 @@ def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=
         yield None, f"Error processing video: {e}"
 
 
+def process(image, bg):
+    image_size = image.size
+    input_images = transform_image(image).unsqueeze(0).to("cuda")
+    # Prediction
+    with torch.no_grad():
+        preds = birefnet(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:
     with gr.Row():
         in_video = gr.Video(label="Input Video", interactive=True)