Update app.py

app.py

@@ -29,7 +29,6 @@ from sam2.build_sam import build_sam2_video_predictor
 
 from moviepy.editor import ImageSequenceClip
 from tqdm import tqdm
-import re
 
 def get_video_fps(video_path):
     # Open the video file
@@ -334,34 +333,18 @@ def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_
         fps = original_fps  # Frames per second
         total_frames = len(jpeg_images)
         clip = ImageSequenceClip(jpeg_images, fps=fps)
-        
-        # Create a custom logger class to use tqdm for progress display
-        class TqdmLogger:
-            def __init__(self, total):
-                self.pbar = tqdm(total=total)
-        
-            def __call__(self, message):
-                # Extract progress percentage from the message
-                progress_match = re.search(r"frame\s+(\d+)", message)
-                if progress_match:
-                    current_frame = int(progress_match.group(1))
-                    self.pbar.n = current_frame
-                    self.pbar.refresh()
-        
-            def close(self):
-                self.pbar.close()
-        
-        # Initialize the custom logger with total frames
-        logger = TqdmLogger(total_frames)
-
         # Write the result to a file
         final_vid_output_path = "output_video.mp4"
+        # Use tqdm progress bar while writing the video file
+        with tqdm(total=total_frames, desc="Rendering video") as pbar:
+            def update_progress(current_frame):
+                pbar.update(current_frame - pbar.n)
+            
+            clip.write_videofile("output_video.mp4", codec='libx264', progress_bar=False, verbose=False, logger=None, progress_callback=update_progress)
+
         
-        # Write the result to a file using the custom logger for progress
-        clip.write_videofile(final_vid_output_path, codec='libx264', logger=logger)
         
-        # Close the progress bar after the process is finished
-        logger.close()
+       
         return gr.update(value=None), gr.update(value=final_vid_output_path)
 
 def update_ui(vis_frame_type):