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=(512, 512)):
+def preprocess_video(video_path, target_fps=24, max_resolution=(256, 256)):
     """Preprocess the video to resize and adjust its frame rate."""
     video = mp.VideoFileClip(video_path)
 
@@ -38,7 +38,7 @@ 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)):
+def process_frames_batch(frames_batch, seed=0, target_size=(256, 256)):
     """Process a batch of frames and return depth maps."""
     try:
         torch.cuda.empty_cache()  # Clear GPU cache
@@ -64,7 +64,7 @@ def process_video(video_path, fps=0, seed=0, batch_size=4):
         start_time = time.time()
 
         # Preprocess the video
-        video = preprocess_video(video_path, target_fps=fps, max_resolution=(512, 512))
+        video = preprocess_video(video_path, target_fps=fps, max_resolution=(256, 256))
 
         # Use original video FPS if not specified
         if fps == 0:
@@ -84,14 +84,15 @@ def process_video(video_path, fps=0, seed=0, batch_size=4):
         # Process frames in batches
         for i in range(0, total_frames, batch_size):
             current_batch_size = batch_size
-            while current_batch_size > 0:
+            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)
-                    break
+                    success = True
                 except RuntimeError as e:
                     if 'out of memory' in str(e):
-                        current_batch_size = current_batch_size // 2
+                        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:
@@ -147,6 +148,10 @@ def process_video(video_path, fps=0, seed=0, batch_size=4):
         logger.error(f"Error: {e}")
         yield None, None, None, f"Error processing video: {e}"
 
+    finally:
+        # Clean up temporary directory
+        pass  # Remove if you decide to delete temp_dir
+
 # Wrapper function with error handling
 def process_wrapper(video, fps=0, seed=0, batch_size=4):
     if video is None: