Significantly reduce rendering time with a separate thread for writing frames to stream

manim/scene/scene_file_writer.py

@@ -7,6 +7,8 @@
 import json
 import shutil
 from pathlib import Path
+from queue import Queue
+from threading import Thread
 from typing import TYPE_CHECKING, Any
 
 import av
@@ -16,6 +18,7 @@
 from pydub import AudioSegment
 
 from manim import __version__
+from manim.typing import PixelArray
 
 from .. import config, logger
 from .._config.logger_utils import set_file_logger
@@ -359,6 +362,33 @@ def end_animation(self, allow_write: bool = False):
         if write_to_movie() and allow_write:
             self.close_partial_movie_stream()
 
+    def listen_and_write(self):
+        """
+        For internal use only: blocks until new frame is available on the queue.
+        """
+        while True:
+            num_frames, frame_data = self.queue.get()
+            if frame_data is None:
+                break
+
+            self.encode_and_write_frame(frame_data, num_frames)
+
+    def encode_and_write_frame(self, frame: PixelArray, num_frames: int) -> None:
+        """
+        For internal use only: takes a given frame in ``np.ndarray`` format and
+        write it to the stream
+        """
+        for _ in range(num_frames):
+            # Notes: precomputing reusing packets does not work!
+            # I.e., you cannot do `packets = encode(...)`
+            # and reuse it, as it seems that `mux(...)`
+            # consumes the packet.
+            # The same issue applies for `av_frame`,
+            # reusing it renders weird-looking frames.
+            av_frame = av.VideoFrame.from_ndarray(frame, format="rgba")
+            for packet in self.video_stream.encode(av_frame):
+                self.video_container.mux(packet)
+
     def write_frame(
         self, frame_or_renderer: np.ndarray | OpenGLRenderer, num_frames: int = 1
     ):
@@ -379,16 +409,9 @@ def write_frame(
                 if config.renderer == RendererType.OPENGL
                 else frame_or_renderer
             )
-            for _ in range(num_frames):
-                # Notes: precomputing reusing packets does not work!
-                # I.e., you cannot do `packets = encode(...)`
-                # and reuse it, as it seems that `mux(...)`
-                # consumes the packet.
-                # The same issue applies for `av_frame`,
-                # reusing it renders weird-looking frames.
-                av_frame = av.VideoFrame.from_ndarray(frame, format="rgba")
-                for packet in self.video_stream.encode(av_frame):
-                    self.video_container.mux(packet)
+
+            msg = (num_frames, frame)
+            self.queue.put(msg)
 
         if is_png_format() and not config["dry_run"]:
             image: Image = (
@@ -430,7 +453,7 @@ def save_final_image(self, image: np.ndarray):
         image.save(self.image_file_path)
         self.print_file_ready_message(self.image_file_path)
 
-    def finish(self):
+    def finish(self) -> None:
         """
         Finishes writing to the FFMPEG buffer or writing images
         to output directory.
@@ -440,8 +463,6 @@ def finish(self):
         frame in the default image directory.
         """
         if write_to_movie():
-            if hasattr(self, "writing_process"):
-                self.writing_process.terminate()
             self.combine_to_movie()
             if config.save_sections:
                 self.combine_to_section_videos()
@@ -455,7 +476,7 @@ def finish(self):
         if self.subcaptions:
             self.write_subcaption_file()
 
-    def open_partial_movie_stream(self, file_path=None):
+    def open_partial_movie_stream(self, file_path=None) -> None:
         """Open a container holding a video stream.
 
         This is used internally by Manim initialize the container holding
@@ -499,13 +520,20 @@ def open_partial_movie_stream(self, file_path=None):
             self.video_container = video_container
             self.video_stream = stream
 
-    def close_partial_movie_stream(self):
+            self.queue: Queue[tuple[int, PixelArray | None]] = Queue()
+            self.writer_thread = Thread(target=self.listen_and_write, args=())
+            self.writer_thread.start()
+
+    def close_partial_movie_stream(self) -> None:
         """Close the currently opened video container.
 
         Used internally by Manim to first flush the remaining packages
         in the video stream holding a partial file, and then close
         the corresponding container.
         """
+        self.queue.put((-1, None))
+        self.writer_thread.join()
+
         for packet in self.video_stream.encode():
             self.video_container.mux(packet)