From 735de62f60d98b505a3dcddbf620373e8a057f3c Mon Sep 17 00:00:00 2001 From: Saveliy Yusufov Date: Sun, 4 Aug 2024 17:15:58 -0400 Subject: [PATCH] Significantly reduce rendering time with a separate thread for writing frames to stream (#3888) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add separate thread for writing frames to stream * [pre-commit.ci] pre-commit autoupdate (#3889) updates: - [github.com/astral-sh/ruff-pre-commit: v0.5.4 → v0.5.5](https://github.com/astral-sh/ruff-pre-commit/compare/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 … (#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 <125118413+irvanalhaq9@users.noreply.github.com> --- manim/file_writer/file_writer.py | 68 ++++++++++++++++++++++++++------ 1 file changed, 55 insertions(+), 13 deletions(-) diff --git a/manim/file_writer/file_writer.py b/manim/file_writer/file_writer.py index 2a0675c22f..d0770ec258 100644 --- a/manim/file_writer/file_writer.py +++ b/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,6 +523,10 @@ 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. @@ -495,6 +534,9 @@ def close_partial_movie_stream(self) -> None: 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)