Update app.py

app.py

@@ -10,7 +10,7 @@ from PIL import Image
 import moviepy.editor as mp
 from infer import lotus  # Import the depth model inference function
 import logging
-import io  # Import io module
+import io
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -61,6 +61,8 @@ def process_frame(image, seed=0):
 # 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."""
+    # Create a persistent temporary directory
+    temp_dir = tempfile.mkdtemp()
     try:
         start_time = time.time()
         
@@ -76,80 +78,76 @@ def process_video(video_path, fps=0, seed=0, batch_size=16):
         
         logger.info(f"Processing {total_frames} frames at {fps} FPS...")
         
-        # 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)
+        # 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)
+        processed_frames = []
+        
+        # Process frames in batches
+        for i in range(0, total_frames, batch_size):
+            frames_batch = frames[i:i+batch_size]
+            depth_maps = []
                 
-                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)
+            # 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)
+                    # 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}")
-            
-            logger.info("Creating output files...")
+                    # 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}")
+                
+        logger.info("Creating output files...")
             
-            # 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)
+        # 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)
             
-            # Create MP4 video
-            video_filename = f"depth_video_{int(time.time())}.mp4"
-            output_video_path = os.path.join(temp_dir, video_filename)
+        # 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!")
+        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
+        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"
+        total_time = time.time() - start_time
+        logger.info("Processing complete!")
+        
+        # Yield the file paths
+        yield None, zip_path, output_video_path, f"Processing complete! Total time: {total_time:.2f} seconds"
             
     except Exception as e:
         logger.error(f"Error: {e}")
         yield None, None, None, f"Error processing video: {e}"
+    # Do not delete temp_dir here; we need the files to persist
+    # Cleanup can be handled elsewhere if necessary
 
 # Wrapper function with error handling
 def process_wrapper(video, fps=0, seed=0, batch_size=16):
@@ -165,7 +163,7 @@ def process_wrapper(video, fps=0, seed=0, batch_size=16):
     except Exception as e:
         raise gr.Error(f"Error processing video: {str(e)}")
 
-# Custom CSS for styling
+# Custom CSS for styling (unchanged)
 custom_css = """
     /* Your existing custom CSS */
 """