Significantly reduce rendering time with a separate thread for writing frames to stream (ManimCommunity#3888)

* Add separate thread for writing frames to stream

* [pre-commit.ci] pre-commit autoupdate (ManimCommunity#3889)

updates:
- [github.com/astral-sh/ruff-pre-commit: v0.5.4 → v0.5.5](astral-sh/ruff-pre-commit@v0.5.4...v0.5.5)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* Replace the TypeError message code in the _typecheck_input method in … (ManimCommunity#3890)

* Replace the TypeError message code in the _typecheck_input method in the DrawBorderThenFill class.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* Remove print statements used for debugging

* Remove writing process termination

- This is probably leftover from back when manim used subprocess to
  write frames to FFmpeg via stdin

* Add type hints to modified methods & instance vars

* Fix inline code in docstring & type hint for queue

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Irvanal Haq <[email protected]>

Author: 3 people

manim/file_writer/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
@@ -20,7 +22,13 @@
 from manim._config.logger_utils import set_file_logger
 from manim.file_writer.protocols import FileWriterProtocol
 from manim.file_writer.sections import DefaultSectionType, Section
-from manim.utils.file_ops import (
+
+from manim.typing import PixelArray
+
+from .. import config, logger
+from .._config.logger_utils import set_file_logger
+from ..constants import RendererType
+from ..utils.file_ops import (
     add_extension_if_not_present,
     add_version_before_extension,
     guarantee_existence,
@@ -361,7 +369,36 @@ def end_animation(self, allow_write: bool = False) -> None:
             self.close_partial_movie_stream()
             self.num_plays += 1
 
-    def write_frame(self, frame: PixelArray, num_frames: int = 1) -> None:
+    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
+    ):
         """
         Used internally by Manim to write a frame to
         the FFMPEG input buffer.
@@ -374,16 +411,14 @@ def write_frame(self, frame: PixelArray, num_frames: int = 1) -> None:
             The number of times to write frame.
         """
         if write_to_movie():
-            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)
+            frame: np.ndarray = (
+                frame_or_renderer.get_frame()
+                if config.renderer == RendererType.OPENGL
+                else frame_or_renderer
+            )
+
+            msg = (num_frames, frame)
+            self.queue.put(msg)
 
         if is_png_format() and not config.dry_run:
             image = Image.fromarray(frame)
@@ -444,7 +479,7 @@ def finish(self) -> None:
         if self.subcaptions:
             self.write_subcaption_file()
 
-    def open_partial_movie_stream(self, file_path: str | None = None) -> 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
@@ -488,13 +523,20 @@ def open_partial_movie_stream(self, file_path: str | None = None) -> None:
             self.video_container = video_container
             self.video_stream = stream
 
+            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)