Update app.py

app.py

@@ -23,7 +23,7 @@ task_name = 'depth'
 pipe_g, pipe_d = load_models(task_name, device)
 
 # Preprocess the video to adjust resolution and frame rate
-def preprocess_video(video_path, target_fps=24, max_resolution=(1920, 1080)):
+def preprocess_video(video_path, target_fps=24, max_resolution=(512, 512)):
     """Preprocess the video to resize and adjust its frame rate."""
     video = mp.VideoFileClip(video_path)
 
@@ -38,11 +38,13 @@ def preprocess_video(video_path, target_fps=24, max_resolution=(1920, 1080)):
     return video
 
 # Process a batch of frames through the depth model
-def process_frames_batch(frames_batch, seed=0):
+def process_frames_batch(frames_batch, seed=0, target_size=(512, 512)):
     """Process a batch of frames and return depth maps."""
     try:
-        # Convert frames to PIL Images
-        images_batch = [Image.fromarray(frame).convert('RGB') for frame in frames_batch]
+        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]
 
         # Run batch inference
         depth_maps = lotus(images_batch, 'depth', seed, device, pipe_g, pipe_d)
@@ -54,7 +56,7 @@ def process_frames_batch(frames_batch, seed=0):
         return [None] * len(frames_batch)
 
 # Process video frames and generate depth maps
-def process_video(video_path, fps=0, seed=0, batch_size=16):
+def process_video(video_path, fps=0, seed=0, batch_size=4):
     """Process video frames in batches and generate depth maps."""
     # Create a persistent temporary directory
     temp_dir = tempfile.mkdtemp()
@@ -62,7 +64,7 @@ def process_video(video_path, fps=0, seed=0, batch_size=16):
         start_time = time.time()
 
         # Preprocess the video
-        video = preprocess_video(video_path, target_fps=fps)
+        video = preprocess_video(video_path, target_fps=fps, max_resolution=(512, 512))
 
         # Use original video FPS if not specified
         if fps == 0:
@@ -81,8 +83,19 @@ def process_video(video_path, fps=0, seed=0, batch_size=16):
 
         # Process frames in batches
         for i in range(0, total_frames, batch_size):
-            frames_batch = frames[i:i+batch_size]
-            depth_maps = process_frames_batch(frames_batch, seed)
+            current_batch_size = batch_size
+            while current_batch_size > 0:
+                try:
+                    frames_batch = frames[i:i+current_batch_size]
+                    depth_maps = process_frames_batch(frames_batch, seed)
+                    break
+                except RuntimeError as e:
+                    if 'out of memory' in str(e):
+                        current_batch_size = 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
@@ -135,7 +148,7 @@ def process_video(video_path, fps=0, seed=0, batch_size=16):
         yield None, None, None, f"Error processing video: {e}"
 
 # Wrapper function with error handling
-def process_wrapper(video, fps=0, seed=0, batch_size=16):
+def process_wrapper(video, fps=0, seed=0, batch_size=4):
     if video is None:
         raise gr.Error("Please upload a video.")
     try:
@@ -192,7 +205,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")
+            batch_size_slider = gr.Slider(minimum=1, maximum=16, step=1, value=4, label="Batch Size")
             btn = gr.Button("Process Video")
 
         with gr.Column():