Update app.py

app.py

@@ -13,7 +13,6 @@ import tempfile
 import uuid
 import time
 import threading
-from queue import Queue
 
 torch.set_float32_matmul_precision("medium")
 
@@ -62,16 +61,10 @@ cleanup_thread = threading.Thread(target=cleanup_temp_files, daemon=True)
 cleanup_thread.start()
 
 
-# Function to process frames in a separate thread
-def process_frame(image, bg, fast_mode, result_queue):
- processed_image = process(image, bg, fast_mode)
- result_queue.put(processed_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):
  try:
- start_time = time.time()
+ start_time = time.time() # Start the timer
 
  # Load the video using moviepy
  video = mp.VideoFileClip(vid)
@@ -103,50 +96,29 @@ def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=
 
  bg_frame_index = 0 # Initialize background frame index
 
- threads = []
- result_queue = Queue()
- frame_batch_size = 4 # Process 4 frames at a time
-
  for i, frame in enumerate(frames):
  pil_image = Image.fromarray(frame)
  if bg_type == "Color":
- background_image = color # Use color directly for threads
+ processed_image = process(pil_image, color, fast_mode)
  elif bg_type == "Image":
- background_image = bg_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:
- background_image = None # Default to no background
-
- # Start a new thread to process the frame
- thread = threading.Thread(target=process_frame, args=(pil_image, background_image, fast_mode, result_queue))
- threads.append(thread)
- thread.start()
-
- # If we have enough threads running or it's the last frame, wait for results
- if len(threads) == frame_batch_size or i == len(list(frames)) - 1:
- for thread in threads:
- thread.join()
- while not result_queue.empty():
- processed_frames.append(np.array(result_queue.get()))
- threads = [] # Reset the threads list
+ 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 the first processed image from the current batch if available
- if processed_frames:
- index = -len(threads) if threads else -min(frame_batch_size, len(processed_frames))
- yield processed_frames[index], None, f"Processing frame {i+1}... Elapsed time: {elapsed_time:.2f} seconds"
- else:
- yield None, None, f"Processing frame {i+1}... Elapsed time: {elapsed_time:.2f} seconds"
-
+ yield processed_image, None, f"Processing frame {i+1}... Elapsed time: {elapsed_time:.2f} seconds"
 
  # Create a new video from the processed frames
  processed_video = mp.ImageSequenceClip(processed_frames, fps=fps)
@@ -164,7 +136,7 @@ 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_frames[-1], temp_filepath, f"Processing complete! Elapsed time: {elapsed_time:.2f} seconds"
+ yield processed_image, temp_filepath, f"Processing complete! Elapsed time: {elapsed_time:.2f} seconds"
 
  except Exception as e:
  print(f"Error: {e}")