Update app.py

app.py

@@ -7,157 +7,151 @@ import time
 import ffmpeg
 import numpy as np
 from PIL import Image
-from concurrent.futures import ThreadPoolExecutor
 import moviepy.editor as mp
 from infer import lotus  # Import the depth model inference function
+import logging
 
-# Set device to use the L40s GPU
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-# Add the preprocess_video function to limit video resolution and frame rate
+# Set device to use the L40s GPU explicitly
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+# Preprocess the video to adjust resolution and frame rate
 def preprocess_video(video_path, target_fps=24, max_resolution=(1920, 1080)):
-    """Preprocess the video to resize and reduce its frame rate."""
+    """Preprocess the video to resize and adjust its frame rate."""
     video = mp.VideoFileClip(video_path)
     
     # Resize video if it's larger than the target resolution
     if video.size[0] > max_resolution[0] or video.size[1] > max_resolution[1]:
         video = video.resize(newsize=max_resolution)
     
-    # Limit FPS
-    video = video.set_fps(target_fps)
+    # Adjust FPS if target_fps is specified
+    if target_fps > 0:
+        video = video.set_fps(target_fps)
     
     return video
 
-def process_frame(frame, seed=0):
+# Process a single frame through the depth model
+def process_frame(image, seed=0):
     """Process a single frame through the depth model and return depth map."""
     try:
-        # Convert frame to PIL Image
-        image = Image.fromarray(frame)
+        # Set seeds for reproducibility
+        torch.manual_seed(seed)
+        np.random.seed(seed)
+        
+        # Process through the depth model (assuming lotus accepts image data)
+        _, output_d = lotus(image, 'depth', seed, device)
         
-        # Save temporary image (lotus requires a file path)
-        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
-            image.save(tmp.name)
-            
-            # Process through the depth model (lotus)
-            _, output_d = lotus(tmp.name, 'depth', seed, device)
-            
-            # Clean up temp file
-            os.unlink(tmp.name)
-            
         # Convert depth output to numpy array
         depth_array = np.array(output_d)
         return depth_array
         
     except Exception as e:
-        print(f"Error processing frame: {e}")
+        logger.error(f"Error processing frame: {e}")
         return None
 
-@spaces.GPU
-def process_video(video_path, fps=0, seed=0, max_workers=32):
+# 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."""
-    temp_dir = None
     try:
         start_time = time.time()
         
         # Preprocess the video
-        video = preprocess_video(video_path)
+        video = preprocess_video(video_path, target_fps=fps)
         
         # Use original video FPS if not specified
         if fps == 0:
             fps = video.fps
                 
-        frames = list(video.iter_frames(fps=fps))
+        frames = list(video.iter_frames(fps=video.fps))
         total_frames = len(frames)
         
-        print(f"Processing {total_frames} frames at {fps} FPS...")
-        
-        # Create temporary directory for frame sequence
-        temp_dir = tempfile.mkdtemp()
-        frames_dir = os.path.join(temp_dir, "frames")
-        os.makedirs(frames_dir, exist_ok=True)
-        
-        # Process frames in larger batches (based on GPU VRAM)
-        batch_size = 50  # Increased batch size to fully utilize the GPU's capabilities
-        processed_frames = []
+        logger.info(f"Processing {total_frames} frames at {fps} FPS...")
         
-        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        # Create temporary directory for frame sequence and outputs
+        with tempfile.TemporaryDirectory() as temp_dir:
+            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):
-                futures = [executor.submit(process_frame, frames[j], seed) for j in range(i, min(i + batch_size, total_frames))]
-                for j, future in enumerate(futures):
-                    try:
-                        result = future.result()
-                        if result is not None:
-                            # Save frame
-                            frame_path = os.path.join(frames_dir, f"frame_{i+j:06d}.png")
-                            Image.fromarray(result).save(frame_path)
-                            
-                            # Collect processed frame for preview
-                            processed_frames.append(result)
-                            
-                            # Update preview (only showing every 10th frame to avoid clutter)
-                            if (i + j + 1) % 10 == 0:
-                                elapsed_time = time.time() - start_time
-                                yield processed_frames[-1], None, None, f"Processed {i+j+1}/{total_frames} frames... Elapsed: {elapsed_time:.2f}s"
-                    except Exception as e:
-                        print(f"Error processing frame {i + j + 1}: {e}")
-        
-        print("Creating output files...")
-        # Create output directory
-        output_dir = os.path.join(os.path.dirname(video_path), "output")
-        os.makedirs(output_dir, exist_ok=True)
-        
-        # Create ZIP of frame sequence
-        zip_filename = f"depth_frames_{int(time.time())}.zip"
-        zip_path = os.path.join(output_dir, zip_filename)
-        shutil.make_archive(zip_path[:-4], 'zip', frames_dir)
-        
-        # Create MP4 video
-        video_filename = f"depth_video_{int(time.time())}.mp4"
-        video_path = os.path.join(output_dir, video_filename)
-        
-        try:
-            # FFmpeg settings for high-quality MP4
-            stream = ffmpeg.input(
-                os.path.join(frames_dir, 'frame_%06d.png'),
-                pattern_type='sequence',
-                framerate=fps
-            )
+                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):
+                    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)
+                        
+                        # Update 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 processed_frames[-1], 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}")
             
-            stream = ffmpeg.output(
-                stream,
-                video_path,
-                vcodec='libx264',
-                pix_fmt='yuv420p',
-                crf=17,  # High quality
-                threads=max_workers
-            )
+            logger.info("Creating output files...")
             
-            ffmpeg.run(stream, overwrite_output=True, capture_stdout=True, capture_stderr=True)
-            print("MP4 video created successfully!")
+            # Create ZIP of frame sequence
+            zip_filename = f"depth_frames_{int(time.time())}.zip"
+            zip_path = os.path.join(temp_dir, zip_filename)
+            shutil.make_archive(zip_path[:-4], 'zip', frames_dir)
             
-        except ffmpeg.Error as e:
-            print(f"Error creating video: {e.stderr.decode() if e.stderr else str(e)}")
-            video_path = None
-    
-        print("Processing complete!")
-        yield None, zip_path, video_path, f"Processing complete! Total time: {time.time() - start_time:.2f} seconds"
+            # Create MP4 video
+            video_filename = f"depth_video_{int(time.time())}.mp4"
+            output_video_path = os.path.join(temp_dir, video_filename)
+            
+            try:
+                # FFmpeg settings for high-quality MP4
+                (
+                    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)
+                )
+                logger.info("MP4 video created successfully!")
+                
+            except ffmpeg.Error as e:
+                logger.error(f"Error creating video: {e.stderr.decode() if e.stderr else str(e)}")
+                output_video_path = None
         
+            total_time = time.time() - start_time
+            logger.info("Processing complete!")
+            
+            # Read output files to return as bytes
+            with open(zip_path, 'rb') as f:
+                zip_data = f.read()
+            with open(output_video_path, 'rb') as f:
+                video_data = f.read()
+            
+            yield None, (zip_filename, zip_data), (video_filename, video_data), f"Processing complete! Total time: {total_time:.2f} seconds"
+            
     except Exception as e:
-        print(f"Error: {e}")
+        logger.error(f"Error: {e}")
         yield None, None, None, f"Error processing video: {e}"
-    finally:
-        if temp_dir and os.path.exists(temp_dir):
-            try:
-                shutil.rmtree(temp_dir)
-            except Exception as e:
-                print(f"Error cleaning up temp directory: {e}")
-                
-def process_wrapper(video, fps=0, seed=0, max_workers=32):
+
+# Wrapper function with error handling
+def process_wrapper(video, fps=0, seed=0, batch_size=16):
     if video is None:
         raise gr.Error("Please upload a video.")
     try:
         outputs = []
-        for output in process_video(video, fps, seed, max_workers):
+        for output in process_video(video.name, fps, seed, batch_size):
             outputs.append(output)
             yield output
         return outputs[-1]
@@ -205,22 +199,23 @@ with gr.Blocks(css=custom_css) as demo:
     
     with gr.Row():
         with gr.Column():
-            video_input = gr.Video(label="Upload Video", interactive=True, show_label=True)
-            fps_slider = gr.Slider(minimum=0, maximum=60, step=1, value=0, label="Output FPS")
-            seed_slider = gr.Slider(minimum=0, maximum=999999999, step=1, value=0, label="Seed")
-            max_workers_slider = gr.Slider(minimum=1, maximum=32, step=1, value=32, label="Max Workers")
-            btn = gr.Button("Process Video", elem_id="submit-button")
+            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")
+            btn = gr.Button("Process Video")
         
         with gr.Column():
-            preview_image = gr.Image(label="Live Preview", show_label=True)
+            preview_image = gr.Image(label="Live Preview")
             output_frames_zip = gr.File(label="Download Frame Sequence (ZIP)")
             output_video = gr.File(label="Download Video (MP4)")
             time_textbox = gr.Textbox(label="Status", interactive=False)
             
-    btn.click(fn=process_wrapper
-
-, inputs=[video_input, fps_slider, seed_slider, max_workers_slider],
-              outputs=[preview_image, output_frames_zip, output_video, time_textbox])
+    btn.click(
+        fn=process_wrapper,
+        inputs=[video_input, fps_slider, seed_slider, batch_size_slider],
+        outputs=[preview_image, output_frames_zip, output_video, time_textbox]
+    )
 
     demo.queue()