Skip to content

Commit

Permalink
Significantly reduce rendering time with a separate thread for writin…
Browse files Browse the repository at this point in the history
…g frames to stream (#3888)

* 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](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 … (#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>
  • Loading branch information
3 people authored Aug 4, 2024
1 parent 17e5a77 commit 11ea416
Showing 1 changed file with 43 additions and 15 deletions.
58 changes: 43 additions & 15 deletions manim/scene/scene_file_writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
):
Expand All @@ -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 = (
Expand Down Expand Up @@ -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.
Expand All @@ -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()
Expand All @@ -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
Expand Down Expand Up @@ -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)

Expand Down

0 comments on commit 11ea416

Please sign in to comment.