Update app.py

app.py

@@ -11,6 +11,7 @@ import moviepy.editor as mp
 from infer import lotus  # Import the depth model inference function
 import logging
 import io
+from multiprocessing import Pool, cpu_count
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -35,13 +36,17 @@ def preprocess_video(video_path, target_fps=24, max_resolution=(1920, 1080)):
     return video
 
 # Process a single frame through the depth model
-def process_frame(image, seed=0):
+def process_frame(args):
     """Process a single frame through the depth model and return depth map."""
+    frame_index, frame_data, seed = args
     try:
         # Set seeds for reproducibility
         torch.manual_seed(seed)
         np.random.seed(seed)
         
+        # Convert frame data to PIL Image
+        image = Image.fromarray(frame_data).convert('RGB')
+        
         # Save image to an in-memory file
         img_bytes = io.BytesIO()
         image.save(img_bytes, format='PNG')
@@ -52,15 +57,15 @@ def process_frame(image, seed=0):
         
         # Convert depth output to numpy array
         depth_array = np.array(output_d)
-        return depth_array
-        
+        return (frame_index, depth_array)
+            
     except Exception as e:
-        logger.error(f"Error processing frame: {e}")
-        return None
+        logger.error(f"Error processing frame {frame_index}: {e}")
+        return (frame_index, None)
 
 # Process video frames and generate depth maps
-def process_video(video_path, fps=0, seed=0, batch_size=16):
-    """Process video, batch frames, and use L40s GPU to generate depth maps."""
+def process_video(video_path, fps=0, seed=0, num_workers=4):
+    """Process video frames in parallel and generate depth maps."""
     # Create a persistent temporary directory
     temp_dir = tempfile.mkdtemp()
     try:
@@ -81,37 +86,28 @@ def process_video(video_path, fps=0, seed=0, batch_size=16):
         # Create directory for frame sequence and outputs
         frames_dir = os.path.join(temp_dir, "frames")
         os.makedirs(frames_dir, exist_ok=True)
-            
-        processed_frames = []
         
-        # Process frames in batches
-        for i in range(0, total_frames, batch_size):
-            frames_batch = frames[i:i+batch_size]
-            depth_maps = []
-                
-            # Process each frame in the batch
-            for frame in frames_batch:
-                depth_map = process_frame(Image.fromarray(frame), seed)
-                depth_maps.append(depth_map)
-                    
-            for j, depth_map in enumerate(depth_maps):
+        # Prepare arguments for multiprocessing
+        args_list = [(i, frames[i], seed) for i in range(total_frames)]
+        
+        # Use multiprocessing Pool to process frames in parallel
+        with Pool(processes=num_workers) as pool:
+            results = []
+            for result in pool.imap_unordered(process_frame, args_list):
+                frame_index, depth_map = result
                 if depth_map is not None:
                     # Save frame
-                    frame_index = i + j
                     frame_path = os.path.join(frames_dir, f"frame_{frame_index:06d}.png")
-                    Image.fromarray(depth_map).save(frame_path)
-                        
-                    # Collect processed frame for preview
-                    processed_frames.append(depth_map)
-                        
+                    Image.fromarray(depth_map.squeeze()).save(frame_path)
+                    
                     # Update preview every 10% progress
-                    if frame_index % max(1, total_frames // 10) == 0:
+                    if (frame_index + 1) % max(1, total_frames // 10) == 0:
                         elapsed_time = time.time() - start_time
-                        progress = (frame_index / total_frames) * 100
-                        yield processed_frames[-1], None, None, f"Processed {frame_index}/{total_frames} frames... ({progress:.2f}%) Elapsed: {elapsed_time:.2f}s"
+                        progress = ((frame_index + 1) / total_frames) * 100
+                        yield depth_map, None, None, f"Processed {frame_index + 1}/{total_frames} frames... ({progress:.2f}%) Elapsed: {elapsed_time:.2f}s"
                 else:
-                    logger.error(f"Error processing frame {frame_index}")
-                
+                    logger.error(f"Frame {frame_index} failed to process.")
+        
         logger.info("Creating output files...")
             
         # Create ZIP of frame sequence
@@ -129,7 +125,7 @@ def process_video(video_path, fps=0, seed=0, batch_size=16):
                 ffmpeg
                 .input(os.path.join(frames_dir, 'frame_%06d.png'), pattern_type='sequence', framerate=fps)
                 .output(output_video_path, vcodec='libx264', pix_fmt='yuv420p', crf=17)
-                .run(overwrite_output=True)
+                .run(overwrite_output=True, quiet=True)
             )
             logger.info("MP4 video created successfully!")
                 
@@ -150,13 +146,13 @@ def process_video(video_path, fps=0, seed=0, batch_size=16):
     # Cleanup can be handled elsewhere if necessary
 
 # Wrapper function with error handling
-def process_wrapper(video, fps=0, seed=0, batch_size=16):
+def process_wrapper(video, fps=0, seed=0, num_workers=4):
     if video is None:
         raise gr.Error("Please upload a video.")
     try:
         outputs = []
         # Use video directly, since it's the file path
-        for output in process_video(video, fps, seed, batch_size):
+        for output in process_video(video, fps, seed, num_workers):
             outputs.append(output)
             yield output
         return outputs[-1]
@@ -181,7 +177,7 @@ with gr.Blocks(css=custom_css) as demo:
             video_input = gr.Video(label="Upload Video", interactive=True)
             fps_slider = gr.Slider(minimum=0, maximum=60, step=1, value=0, label="Output FPS (0 for original)")
             seed_slider = gr.Number(value=0, label="Seed")
-            batch_size_slider = gr.Slider(minimum=1, maximum=64, step=1, value=16, label="Batch Size")
+            num_workers_slider = gr.Slider(minimum=1, maximum=cpu_count(), step=1, value=4, label="Number of Workers")
             btn = gr.Button("Process Video")
         
         with gr.Column():
@@ -192,7 +188,7 @@ with gr.Blocks(css=custom_css) as demo:
             
     btn.click(
         fn=process_wrapper,
-        inputs=[video_input, fps_slider, seed_slider, batch_size_slider],
+        inputs=[video_input, fps_slider, seed_slider, num_workers_slider],
         outputs=[preview_image, output_frames_zip, output_video, time_textbox]
     )