Update app.py

app.py

@@ -37,27 +37,27 @@ def preprocess_video(video_path, target_fps=24, max_resolution=(512, 512)):
 
     return video
 
-# Process a batch of frames through the depth model
-def process_frames_batch(frames_batch, seed=0, target_size=(512, 512)):
-    """Process a batch of frames and return depth maps."""
+# Process a single frame through the depth model
+def process_frame(frame, seed=0, target_size=(512, 512)):
+    """Process a single frame and return depth map."""
     try:
         torch.cuda.empty_cache()  # Clear GPU cache
 
-        # Resize frames to the target size
-        images_batch = [Image.fromarray(frame).convert('RGB').resize(target_size, Image.BILINEAR) for frame in frames_batch]
+        # Resize frame to the target size
+        image = Image.fromarray(frame).convert('RGB').resize(target_size, Image.BILINEAR)
 
-        # Run batch inference
-        depth_maps = lotus(images_batch, 'depth', seed, device, pipe_g, pipe_d)
+        # Run inference
+        depth_map = lotus(image, 'depth', seed, device, pipe_g, pipe_d)
 
-        return depth_maps
+        return depth_map
 
     except Exception as e:
-        logger.error(f"Error processing batch: {e}")
-        return [None] * len(frames_batch)
+        logger.error(f"Error processing frame: {e}")
+        return None
 
 # Process video frames and generate depth maps
-def process_video(video_path, fps=0, seed=0, batch_size=4):
-    """Process video frames in batches and generate depth maps."""
+def process_video(video_path, fps=0, seed=0):
+    """Process video frames individually and generate depth maps."""
     # Create a persistent temporary directory
     temp_dir = tempfile.mkdtemp()
     try:
@@ -79,39 +79,22 @@ def process_video(video_path, fps=0, seed=0, batch_size=4):
         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):
-            current_batch_size = batch_size
-            success = False
-            while current_batch_size > 0 and not success:
-                try:
-                    frames_batch = frames[i:i+current_batch_size]
-                    depth_maps = process_frames_batch(frames_batch, seed)
-                    success = True
-                except RuntimeError as e:
-                    if 'out of memory' in str(e):
-                        current_batch_size = max(1, current_batch_size // 2)
-                        logger.warning(f"Reducing batch size to {current_batch_size} due to out of memory error.")
-                        torch.cuda.empty_cache()
-                    else:
-                        raise e
-
-            for j, depth_map in enumerate(depth_maps):
-                frame_index = i + j
-                if depth_map is not None:
-                    # Save frame
-                    frame_path = os.path.join(frames_dir, f"frame_{frame_index:06d}.png")
-                    depth_map.save(frame_path)
-
-                    # Update live preview every 10% progress
-                    if frame_index % max(1, total_frames // 10) == 0:
-                        elapsed_time = time.time() - start_time
-                        progress = (frame_index / total_frames) * 100
-                        yield depth_map, None, None, f"Processed {frame_index}/{total_frames} frames... ({progress:.2f}%) Elapsed: {elapsed_time:.2f}s"
-                else:
-                    logger.error(f"Error processing frame {frame_index}")
+        # Process frames individually
+        for i, frame in enumerate(frames):
+            depth_map = process_frame(frame, seed)
+
+            if depth_map is not None:
+                # Save frame
+                frame_path = os.path.join(frames_dir, f"frame_{i:06d}.png")
+                depth_map.save(frame_path)
+
+                # Update live preview every 10% progress
+                if i % max(1, total_frames // 10) == 0:
+                    elapsed_time = time.time() - start_time
+                    progress = (i / total_frames) * 100
+                    yield depth_map, None, None, f"Processed {i}/{total_frames} frames... ({progress:.2f}%) Elapsed: {elapsed_time:.2f}s"
+            else:
+                logger.error(f"Error processing frame {i}")
 
         logger.info("Creating output files...")
 
@@ -153,13 +136,13 @@ def process_video(video_path, fps=0, seed=0, batch_size=4):
         pass
 
 # Wrapper function with error handling
-def process_wrapper(video, fps=0, seed=0, batch_size=4):
+def process_wrapper(video, fps=0, seed=0):
     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):
             outputs.append(output)
             yield output
         return outputs[-1]
@@ -210,7 +193,6 @@ 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=16, step=1, value=4, label="Batch Size")
             btn = gr.Button("Process Video")
 
         with gr.Column():
@@ -221,7 +203,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],
         outputs=[preview_image, output_frames_zip, output_video, time_textbox]
     )