Update app.py

app.py

@@ -11,6 +11,7 @@ import numpy as np
 import os
 import tempfile
 import uuid
+import threading
 
 torch.set_float32_matmul_precision("highest")
 
@@ -28,6 +29,31 @@ transform_image = transforms.Compose(
 
 BATCH_SIZE = 3
 
+def process_batch(frame_batch, bg_type, bg_image, bg_video, color, fps, video_handling, bg_frame_index, background_frames):
+ pil_images = [Image.fromarray(f) for f in frame_batch]
+ processed_images = []
+
+ 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":
+ for _ in range(len(frame_batch)):
+ if video_handling == "slow_down":
+ background_frame = background_frames[int(bg_frame_index)]
+ bg_frame_index += len(background_frames) / (len(frame_batch) * (len(background_frames) / (fps*mp.VideoFileClip(bg_video).duration)))
+ 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)
+
+ processed_images.append(process(pil_images[_], background_image))
+ else:
+ processed_images = pil_images
+
+ return processed_images, 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"):
  try:
@@ -35,7 +61,7 @@ def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=
  if fps == 0:
  fps = video.fps
  audio = video.audio
- frames = video.iter_frames(fps=fps)
+ frames = list(video.iter_frames(fps=fps))
  processed_frames = []
  yield gr.update(visible=True), gr.update(visible=False)
 
@@ -52,40 +78,22 @@ def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=
 
  bg_frame_index = 0
  frame_batch = []
+ threads = []
+
 
  for i, frame in enumerate(frames):
  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 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)
- 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_images.append(process(pil_images[_],background_image))
-
-
- else:
- processed_images = pil_images
+ if len(frame_batch) == BATCH_SIZE or i == len(frames) - 1: # Process batch or last frames
+ thread = threading.Thread(target=lambda : processed_frames.extend(process_batch(frame_batch, bg_type, bg_image, bg_video, color, fps, video_handling, bg_frame_index, background_frames)[0]))
+ threads.append(thread)
+ thread.start()
+ frame_batch = []
 
- for processed_image in processed_images:
- processed_frames.append(np.array(processed_image))
- yield processed_image, None
- frame_batch = [] # Clear the batch
+ for thread in threads:
+ thread.join()
 
+ for processed_image in processed_frames:
+ yield processed_image, None
 
  processed_video = mp.ImageSequenceClip(processed_frames, fps=fps)
  processed_video = processed_video.set_audio(audio)