Update app.py

app.py

@@ -13,17 +13,14 @@ import tempfile
 import uuid
 import time
 from concurrent.futures import ThreadPoolExecutor
-import asyncio
 
 torch.set_float32_matmul_precision("medium")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load both BiRefNet models
-birefnet = AutoModelForImageSegmentation.from_pretrained(
- "ZhengPeng7/BiRefNet", trust_remote_code=True)
+birefnet = AutoModelForImageSegmentation.from_pretrained("ZhengPeng7/BiRefNet", trust_remote_code=True)
 birefnet.to(device)
-birefnet_lite = AutoModelForImageSegmentation.from_pretrained(
- "ZhengPeng7/BiRefNet_lite", trust_remote_code=True)
+birefnet_lite = AutoModelForImageSegmentation.from_pretrained("ZhengPeng7/BiRefNet_lite", trust_remote_code=True)
 birefnet_lite.to(device)
 
 transform_image = transforms.Compose([
@@ -32,74 +29,77 @@ transform_image = transforms.Compose([
  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
+# 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
 
 @spaces.GPU
-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
-
- bg_frame_index = 0
+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
+ 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
+ 
+ bg_frame_index = 0 # Initialize background frame index
 
- # 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)
+ 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"
+ 
+ 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 result, None, f"Processing frame {len(processed_frames)}... Elapsed time: {elapsed_time:.2f} seconds"
+ 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"
  
- 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"
+ 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