From e911ec30964bfcaa24e97ec7a6c7841602c19308 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=A9rome=20Eertmans?= Date: Tue, 28 Jan 2025 23:05:59 +0000 Subject: [PATCH] feat(lib): smarter files reversing (#439) * feat(lib): smarter files reversing Implement a smarter generation of reversed files by splitting the video into smaller segments. Closes #434 * chore(lib): change default length * chore: use suffix * chore(docs): update * chore(docs): fix docs and add changelog entry * chore(tests): coverage * chore(docs): typo --- CHANGELOG.md | 13 +++++++ manim_slides/slide/base.py | 15 ++++++-- manim_slides/slide/manim.py | 17 +++++++-- manim_slides/utils.py | 70 +++++++++++++++++++++++++++++++++++-- tests/test_slide.py | 20 +++++++++++ uv.lock | 2 +- 6 files changed, 129 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ece27057..e539d745 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 (unreleased)= ## [Unreleased](https://github.com/jeertmans/manim-slides/compare/v5.4.2...HEAD) +(unreleased-added)= +### Added + +- Added `max_duration_before_split_reverse` and `num_processes` class variables. + [#439](https://github.com/jeertmans/manim-slides/pull/439) + +(unreleased-changed)= +### Changed + +- Automatically split large video animations into smaller chunks + for lightweight (and potentially faster) reversed animations generation. + [#439](https://github.com/jeertmans/manim-slides/pull/439) + (v5.4.2)= ## [v5.4.2](https://github.com/jeertmans/manim-slides/compare/v5.4.1...v5.4.2) diff --git a/manim_slides/slide/base.py b/manim_slides/slide/base.py index c648251a..a276c6ff 100644 --- a/manim_slides/slide/base.py +++ b/manim_slides/slide/base.py @@ -36,6 +36,8 @@ class BaseSlide: disable_caching: bool = False flush_cache: bool = False skip_reversing: bool = False + max_duration_before_split_reverse: float | None = 4.0 + num_processes: int | None = None def __init__( self, *args: Any, output_folder: Path = FOLDER_PATH, **kwargs: Any @@ -530,10 +532,11 @@ def _save_slides( for pre_slide_config in tqdm( self._slides, - desc=f"Concatenating animation files to '{scene_files_folder}' and generating reversed animations", + desc=f"Concatenating animations to '{scene_files_folder}' and generating reversed animations", leave=self._leave_progress_bar, ascii=True if platform.system() == "Windows" else None, disable=not self._show_progress_bar, + unit=" slides", ): if pre_slide_config.skip_animations: continue @@ -557,7 +560,15 @@ def _save_slides( if skip_reversing: rev_file = dst_file else: - reverse_video_file(dst_file, rev_file) + reverse_video_file( + dst_file, + rev_file, + max_segment_duration=self.max_duration_before_split_reverse, + num_processes=self.num_processes, + leave=self._leave_progress_bar, + ascii=True if platform.system() == "Windows" else None, + disable=not self._show_progress_bar, + ) slides.append( SlideConfig.from_pre_slide_config_and_files( diff --git a/manim_slides/slide/manim.py b/manim_slides/slide/manim.py index 68842193..fb8a7a2b 100644 --- a/manim_slides/slide/manim.py +++ b/manim_slides/slide/manim.py @@ -11,7 +11,7 @@ class Slide(BaseSlide, Scene): # type: ignore[misc] """ - Inherits from :class:`Scene` and provide necessary tools + Inherits from :class:`Scene` and provides necessary tools for slides rendering. :param args: Positional arguments passed to scene object. @@ -20,15 +20,26 @@ class Slide(BaseSlide, Scene): # type: ignore[misc] :cvar bool disable_caching: :data:`False`: Whether to disable the use of cached animation files. :cvar bool flush_cache: :data:`False`: Whether to flush the cache. - Unlike with Manim, flushing is performed before rendering. :cvar bool skip_reversing: :data:`False`: Whether to generate reversed animations. - If set to :data:`False`, and no cached reversed animation exists (or caching is disabled) for a given slide, then the reversed animation will be simply the same as the original one, i.e., ``rev_file = file``, for the current slide config. + :cvar typing.Optional[float] max_duration_before_split_reverse: :data:`4.0`: Maximum duration + before of a video animation before it is reversed by splitting the file into smaller chunks. + Generating reversed animations can require an important amount of + memory (because the whole video needs to be kept in memory), + and splitting the video into multiple chunks usually speeds + up the process (because it can be done in parallel) while taking + less memory. + Set this to :data:`None` to disable splitting the file into chunks. + :cvar typing.Optional[int] num_processes: :data:`None`: Number of processes + to use for parallelizable operations. + If :data:`None`, defaults to :func:`os.process_cpu_count`. + This is currently used when generating reversed animations, and can + increase memory consumption. """ def __init__(self, *args: Any, **kwargs: Any) -> None: diff --git a/manim_slides/utils.py b/manim_slides/utils.py index dc6e92b9..0703bf1d 100644 --- a/manim_slides/utils.py +++ b/manim_slides/utils.py @@ -2,9 +2,12 @@ import os import tempfile from collections.abc import Iterator +from multiprocessing import Pool from pathlib import Path +from typing import Any, Optional import av +from tqdm import tqdm from .logger import logger @@ -89,8 +92,9 @@ def link_nodes(*nodes: av.filter.context.FilterContext) -> None: c.link_to(n) -def reverse_video_file(src: Path, dest: Path) -> None: +def reverse_video_file_in_one_chunk(src_and_dest: tuple[Path, Path]) -> None: """Reverses a video file, writing the result to `dest`.""" + src, dest = src_and_dest with ( av.open(str(src)) as input_container, av.open(str(dest), mode="w") as output_container, @@ -120,8 +124,70 @@ def reverse_video_file(src: Path, dest: Path) -> None: for _ in range(frames_count): frame = graph.pull() - frame.pict_type = 5 # Otherwise we get a warning saying it is changed + frame.pict_type = "NONE" # Otherwise we get a warning saying it is changed output_container.mux(output_stream.encode(frame)) for packet in output_stream.encode(): output_container.mux(packet) + + +def reverse_video_file( + src: Path, + dest: Path, + max_segment_duration: Optional[float] = 4.0, + num_processes: Optional[int] = None, + **tqdm_kwargs: Any, +) -> None: + """Reverses a video file, writing the result to `dest`.""" + with av.open(str(src)) as input_container: # Fast path if file is short enough + input_stream = input_container.streams.video[0] + if max_segment_duration is None: + return reverse_video_file_in_one_chunk((src, dest)) + elif input_stream.duration: + if ( + float(input_stream.duration * input_stream.time_base) + <= max_segment_duration + ): + return reverse_video_file_in_one_chunk((src, dest)) + else: # pragma: no cover + logger.debug( + f"Could not determine duration of {src}, falling back to segmentation." + ) + + with tempfile.TemporaryDirectory() as tmpdirname: + tmpdir = Path(tmpdirname) + with av.open( + str(tmpdir / f"%04d.{src.suffix}"), + "w", + format="segment", + options={"segment_time": str(max_segment_duration)}, + ) as output_container: + output_stream = output_container.add_stream( + template=input_stream, + ) + + for packet in input_container.demux(input_stream): + if packet.dts is None: + continue + + packet.stream = output_stream + output_container.mux(packet) + + src_files = list(tmpdir.iterdir()) + rev_files = [ + src_file.with_stem("rev_" + src_file.stem) for src_file in src_files + ] + + with Pool(num_processes, maxtasksperchild=1) as pool: + for _ in tqdm( + pool.imap_unordered( + reverse_video_file_in_one_chunk, zip(src_files, rev_files) + ), + desc="Reversing large file by cutting it in segments", + total=len(src_files), + unit=" files", + **tqdm_kwargs, + ): + pass # We just consume the iterator + + concatenate_video_files(rev_files[::-1], dest) diff --git a/tests/test_slide.py b/tests/test_slide.py index c4baeba2..5c038741 100644 --- a/tests/test_slide.py +++ b/tests/test_slide.py @@ -314,6 +314,26 @@ def construct(self) -> None: self.play(dot.animate.move_to(LEFT)) self.play(dot.animate.move_to(DOWN)) + def test_split_reverse(self) -> None: + @assert_renders + class _(CESlide): + max_duration_before_split_reverse = 3.0 + + def construct(self) -> None: + self.wait(2.0) + for _ in range(3): + self.next_slide() + self.wait(10.0) + + @assert_renders + class __(CESlide): + max_duration_before_split_reverse = None + + def construct(self) -> None: + self.wait(5.0) + self.next_slide() + self.wait(5.0) + def test_file_too_long(self) -> None: @assert_renders class _(CESlide): diff --git a/uv.lock b/uv.lock index aaa6df30..2f903402 100644 --- a/uv.lock +++ b/uv.lock @@ -1526,7 +1526,7 @@ wheels = [ [[package]] name = "manim-slides" -version = "5.4.1" +version = "5.4.2" source = { editable = "." } dependencies = [ { name = "av" },