diff --git a/manim/scene/scene_file_writer.py b/manim/scene/scene_file_writer.py index f3e5c08a4a..97a948a993 100644 --- a/manim/scene/scene_file_writer.py +++ b/manim/scene/scene_file_writer.py @@ -5,10 +5,11 @@ __all__ = ["SceneFileWriter"] import json -import os import shutil import subprocess from pathlib import Path +from queue import Queue +from threading import Thread from typing import TYPE_CHECKING, Any import numpy as np @@ -17,7 +18,7 @@ from pydub import AudioSegment from manim import __version__ - +from .section import DefaultSectionType, Section from .. import config, logger from .._config.logger_utils import set_file_logger from ..constants import RendererType @@ -33,7 +34,6 @@ write_to_movie, ) from ..utils.sounds import get_full_sound_file_path -from .section import DefaultSectionType, Section if TYPE_CHECKING: from manim.renderer.opengl_renderer import OpenGLRenderer @@ -78,6 +78,7 @@ def __init__(self, renderer, scene_name, **kwargs): self.partial_movie_files: list[str] = [] self.subcaptions: list[srt.Subtitle] = [] self.sections: list[Section] = [] + self.queue = Queue() # first section gets automatically created for convenience # if you need the first section to be skipped, add a first section by hand, it will replace this one self.next_section( @@ -388,15 +389,14 @@ def write_frame(self, frame_or_renderer: np.ndarray | OpenGLRenderer): elif config.renderer == RendererType.CAIRO: frame = frame_or_renderer if write_to_movie(): - self.writing_process.stdin.write(frame.tobytes()) + self.queue.put(frame.tobytes()) if is_png_format() and not config["dry_run"]: self.output_image_from_array(frame) def write_opengl_frame(self, renderer: OpenGLRenderer): if write_to_movie(): - self.writing_process.stdin.write( - renderer.get_raw_frame_buffer_object_data(), - ) + self.queue.put(renderer.get_raw_frame_buffer_object_data()) + elif is_png_format() and not config["dry_run"]: target_dir = self.image_file_path.parent / self.image_file_path.stem extension = self.image_file_path.suffix @@ -453,6 +453,8 @@ def finish(self): """ if write_to_movie(): if hasattr(self, "writing_process"): + self.queue.put(-1) + self.writer_thread.join() self.writing_process.terminate() self.combine_to_movie() if config.save_sections: @@ -515,12 +517,26 @@ def open_movie_pipe(self, file_path=None): else: command += ["-vcodec", "libx264", "-pix_fmt", "yuv420p"] command += [file_path] + self.writing_process = subprocess.Popen(command, stdin=subprocess.PIPE) + self.writer_thread = Thread(target=self.listen_and_write, args=()) + self.writer_thread.start() + + def listen_and_write(self) -> None: + while True: + buf = self.queue.get() + if buf == -1: + return + + self.writing_process.stdin.write(buf) + # TODO(saveliy): this method needs to close the thread def close_movie_pipe(self): """ Used internally by Manim to gracefully stop writing to FFMPEG's input buffer """ + self.queue.put(-1) + self.writer_thread.join() self.writing_process.stdin.close() self.writing_process.wait() @@ -667,7 +683,8 @@ def combine_to_movie(self): self.print_file_ready_message(str(movie_file_path)) if write_to_movie(): for file_path in partial_movie_files: - # We have to modify the accessed time so if we have to clean the cache we remove the one used the longest. + # We have to modify the accessed time so if we have to clean the cache we remove the one used the + # longest. modify_atime(file_path) def combine_to_section_videos(self) -> None: