from __future__ import annotations import subprocess import sys import re from contextlib import contextmanager from collections.abc import Iterator from dataclasses import dataclass from pathlib import Path from threading import Thread from typing import Callable from tools.console import PipelineProgressView, ProgressView from tools.video_formatting import chroma_family ENCODER_PRESETS = ( "ultrafast", "superfast", "veryfast", "faster", "fast", "medium", "slow", "slower", "veryslow", "placebo", ) @dataclass(frozen=True) class VideoCodecOptions: codec: str crf: int preset: str extension: str = ".mp4" profile: str | None = None level: str | None = None threads: int | None = None @dataclass(frozen=True) class AudioEncodeOptions: mode: str bitrate: str = "192k" preserve_pitch: bool = True def default_audio_bitrate(mode: str) -> str: return {"aac": "192k", "opus": "128k"}.get(mode, "") @dataclass(frozen=True) class VideoEncodeResult: script: Path output: Path renderer_returncode: int ffmpeg_returncode: int encoder_log: str = "" @dataclass(frozen=True) class Y4MEncodeRun: renderer_returncode: int encoder_returncode: int encoder_log: str renderer_log: str @dataclass class RendererStderrReader: process: subprocess.Popen progress: PipelineProgressView | None lines: list[str] thread: Thread def finish(self) -> str: self.thread.join() if self.progress is not None: self.progress.finish() return "".join(self.lines) def choose_video_pixel_format( ffmpeg: Path, codec: str, bit_depth: int, chroma: str | None = None, ) -> str: family = chroma_family(chroma) supported = supported_pixel_formats(ffmpeg, codec, family) preserving_formats = [ (depth, pix_fmt) for depth, pix_fmt in supported if depth >= bit_depth ] if preserving_formats: return min(preserving_formats, key=lambda item: item[0])[1] raise RuntimeError( f"{codec} does not report a yuv{family} pixel format that preserves " f"the {bit_depth}-bit Avisynth output" ) def pixel_format_bit_depth(pixel_format: str) -> int: for depth in (16, 14, 12, 10): if f"p{depth}" in pixel_format: return depth return 8 def supported_pixel_formats(ffmpeg: Path, codec: str, family: str) -> list[tuple[int, str]]: completed = subprocess.run( [str(ffmpeg), "-hide_banner", "-h", f"encoder={codec}"], capture_output=True, text=True, ) if completed.returncode != 0: return [(8, "yuv420p")] supported: list[tuple[int, str]] = [] for line in completed.stdout.splitlines() + completed.stderr.splitlines(): if "Supported pixel formats:" not in line: continue formats = line.split(":", 1)[1].split() base = f"yuv{family}p" if base in formats: supported.append((8, base)) for depth in (10, 12, 14, 16): pix_fmt = f"{base}{depth}le" if pix_fmt in formats: supported.append((depth, pix_fmt)) break if supported: return supported return [(8, f"yuv{family}p")] if family == "420" else [] def ffmpeg_colorspace_from_matrix(matrix: int | None) -> str | None: if matrix == 1: return "bt709" if matrix in {5, 6}: return "smpte170m" if matrix in {9, 10}: return "bt2020nc" return None def render_audio_wav( *, avs_runner: Path, script: Path, output: Path, ) -> None: output.parent.mkdir(parents=True, exist_ok=True) with output.open("wb") as wav_file: completed = subprocess.run( [str(avs_runner), "--wav", str(script)], stdout=wav_file, stderr=subprocess.PIPE, text=False, ) if completed.returncode != 0: raise RuntimeError( f"Avisynth runner failed to render audio for {script} with exit " f"{completed.returncode}:\n" + completed.stderr.decode(errors="replace") ) def encode_video_only_y4m( *, y4m_command: list[str], ffmpeg: Path, script: Path, audio_source: Path | None, output: Path, options: VideoCodecOptions | None = None, audio_options: AudioEncodeOptions | None = None, audio_start: float = 0.0, audio_duration: float | None = None, audio_segments: list[tuple[float, float]] | None = None, audio_tempo: float = 1.0, audio_sample_rate: int | None = None, source_has_audio: bool = False, allow_audio_copy: bool = False, pixel_format: str = "yuv420p", colorspace: str | None = None, frame_count: int | None = None, progress_label: str | None = None, ) -> VideoEncodeResult: if options is None: options = VideoCodecOptions(codec="libx264", crf=16, preset="slow") if audio_options is None: audio_options = AudioEncodeOptions(mode="none") output.parent.mkdir(parents=True, exist_ok=True) command = ffmpeg_y4m_command( ffmpeg=ffmpeg, audio_source=audio_source, output=output, options=options, audio_options=audio_options, audio_start=audio_start, audio_duration=audio_duration, audio_segments=audio_segments, audio_tempo=audio_tempo, audio_sample_rate=audio_sample_rate, source_has_audio=source_has_audio, allow_audio_copy=allow_audio_copy, pixel_format=pixel_format, colorspace=colorspace, ) run = run_y4m_encoder( y4m_command=y4m_command, script=script, encoder_command=command, frame_count=frame_count, progress_label=progress_label or output.name, read_progress=read_ffmpeg_progress, ) if run.encoder_returncode != 0: raise RuntimeError( f"FFmpeg failed for {script} with exit {run.encoder_returncode}:\n" + run.encoder_log + ( "\nAvisynth runner also exited with " f"{run.renderer_returncode}:\n" + run.renderer_log if run.renderer_returncode != 0 else "" ) ) if run.renderer_returncode != 0: raise RuntimeError( f"Avisynth runner failed for {script} with exit {run.renderer_returncode}:\n" + run.renderer_log ) return VideoEncodeResult( script=script, output=output, renderer_returncode=run.renderer_returncode, ffmpeg_returncode=run.encoder_returncode, encoder_log=run.encoder_log, ) def ffmpeg_y4m_command( *, ffmpeg: Path, audio_source: Path | None, output: Path, options: VideoCodecOptions, audio_options: AudioEncodeOptions, audio_start: float, audio_duration: float | None, audio_segments: list[tuple[float, float]] | None, audio_tempo: float, audio_sample_rate: int | None, source_has_audio: bool, allow_audio_copy: bool, pixel_format: str, colorspace: str | None, ) -> list[str]: command = [ str(ffmpeg), "-y", "-hide_banner", "-loglevel", "info", "-nostats", "-progress", "pipe:2", "-f", "yuv4mpegpipe", "-i", "pipe:0", ] if audio_options.mode != "none" and source_has_audio: if audio_source is None: raise RuntimeError("An audio source is required when audio encoding is enabled") add_audio_input_options( command, audio_options=audio_options, audio_start=audio_start, audio_duration=audio_duration, allow_audio_copy=allow_audio_copy, ) command.extend(["-i", str(audio_source)]) command.extend([ "-map", "0:v:0", "-c:v", options.codec, "-preset", options.preset, "-crf", str(options.crf), "-pix_fmt", pixel_format, ]) add_video_profile_level(command, options) add_video_threads(command, options) if colorspace is not None: command.extend(["-colorspace", colorspace]) add_audio_options( command, audio_options=audio_options, audio_start=audio_start, audio_duration=audio_duration, audio_segments=audio_segments, audio_tempo=audio_tempo, audio_sample_rate=audio_sample_rate, source_has_audio=source_has_audio, allow_audio_copy=allow_audio_copy, ) command.append(str(output)) return command def run_y4m_encoder( *, y4m_command: list[str], script: Path, encoder_command: list[str], frame_count: int | None, progress_label: str, read_progress: Callable[..., str], ) -> Y4MEncodeRun: with video_only_script(script) as render_script: renderer_process = subprocess.Popen( [*y4m_command, str(render_script)], stdout=subprocess.PIPE, stderr=subprocess.PIPE ) assert renderer_process.stdout is not None renderer_reader = start_renderer_stderr_reader( renderer_process, rendering_label=f"Rendering {progress_label}", encoding_label=f"Encoding {progress_label}", encoding_total=frame_count, ) encoder_process = subprocess.Popen( encoder_command, stdin=renderer_process.stdout, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True, ) renderer_process.stdout.close() encoder_log = read_progress( encoder_process, total_frames=frame_count, label=f"Encoding {progress_label}", pipeline_progress=renderer_reader.progress, ) renderer_log = renderer_reader.finish() renderer_returncode = renderer_process.wait() return Y4MEncodeRun( renderer_returncode=renderer_returncode, encoder_returncode=encoder_process.returncode, encoder_log=encoder_log, renderer_log=renderer_log, ) @contextmanager def video_only_script(script: Path) -> Iterator[Path]: wrapper = script.with_name(f".{script.name}.video-only.avs") wrapper.write_text( 'mbt_video_only = true\n' f'Import("{_avs_path(script.name)}")\n' "last\n", encoding="utf-8", ) try: yield wrapper finally: wrapper.unlink(missing_ok=True) def _avs_path(path: str) -> str: return path.replace("\\", "/").replace('"', '\\"') def add_audio_input_options( command: list[str], *, audio_options: AudioEncodeOptions, audio_start: float, audio_duration: float | None, allow_audio_copy: bool, ) -> None: if audio_options.mode != "copy": return if not allow_audio_copy: return if audio_start > 0.001: command.extend(["-ss", f"{audio_start:.6f}"]) if audio_duration is not None: command.extend(["-t", f"{audio_duration:.6f}"]) def add_audio_options( command: list[str], *, audio_options: AudioEncodeOptions, audio_start: float, audio_duration: float | None, audio_segments: list[tuple[float, float]] | None = None, audio_tempo: float, audio_sample_rate: int | None, source_has_audio: bool, allow_audio_copy: bool, ) -> None: if audio_options.mode == "none" or not source_has_audio: command.append("-an") elif audio_options.mode in {"aac", "opus", "flac"}: if audio_segments is not None: add_segmented_encoded_audio( command, audio_segments=audio_segments, audio_options=audio_options, audio_tempo=audio_tempo, audio_sample_rate=audio_sample_rate, ) return if audio_duration is None: raise RuntimeError("audio_duration is required for AAC audio") audio_filters = [ f"atrim=start={audio_start:.6f}:duration={audio_duration:.6f}", "asetpts=PTS-STARTPTS", ] audio_filters.extend( speed_change_filters(audio_tempo, audio_options, audio_sample_rate) ) command.extend(["-map", "1:a:0", "-af", ",".join(audio_filters)]) add_audio_encoder_options(command, audio_options) elif audio_options.mode == "copy": if not allow_audio_copy: raise RuntimeError("Audio copy is only supported for whole-file or simple beginning/end trims.") command.extend(["-map", "1:a:0", "-c:a", "copy"]) else: raise RuntimeError(f"Unsupported audio mode: {audio_options.mode}") def add_video_profile_level(command: list[str], options: VideoCodecOptions) -> None: if options.profile: command.extend(["-profile:v", options.profile]) if options.level: command.extend(["-level:v", options.level]) def add_video_threads(command: list[str], options: VideoCodecOptions) -> None: if options.threads is None: return if options.codec == "libx265": command.extend(["-x265-params", f"pools={options.threads}"]) else: command.extend(["-threads", str(options.threads)]) def print_encoder_statistics(encoder_log: str) -> None: statistics = encoder_statistics_lines(encoder_log) if not statistics: return print(" encoder stats:") for line in statistics: print(f" {line}") def encoder_statistics_lines(encoder_log: str) -> list[str]: pattern = re.compile( r"(frame [IPB]:.*|consecutive B-frames:.*|ref [PB] L[01]:.*)", re.IGNORECASE, ) statistics: list[str] = [] for line in encoder_log.splitlines(): match = pattern.search(line) if match: statistics.append(match.group(1)) return statistics def start_renderer_stderr_reader( process: subprocess.Popen, *, rendering_label: str, encoding_label: str, encoding_total: int | None, ) -> RendererStderrReader: assert process.stderr is not None progress = ( PipelineProgressView( rendering_label=rendering_label, encoding_label=encoding_label, encoding_total=encoding_total, ) if encoding_total and sys.stdout.isatty() else None ) lines: list[str] = [] def consume() -> None: for raw_line in process.stderr: line = raw_line.decode(errors="replace") if progress is not None and line.startswith("mbt_render,"): if _update_renderer_progress(progress, line): continue lines.append(line) thread = Thread(target=consume, daemon=True) thread.start() return RendererStderrReader(process, progress, lines, thread) def _update_renderer_progress(progress: PipelineProgressView, line: str) -> bool: parts = line.strip().split(",") if len(parts) != 6: return False try: processed = int(parts[1]) total = int(parts[2]) elapsed = float(parts[4]) eta = float(parts[5]) except ValueError: return False progress.update_rendering(processed, total, elapsed=elapsed, eta=eta) return True def add_segmented_encoded_audio( command: list[str], *, audio_segments: list[tuple[float, float]], audio_options: AudioEncodeOptions, audio_tempo: float, audio_sample_rate: int | None, ) -> None: if not audio_segments: raise RuntimeError("audio_segments must not be empty") filters: list[str] = [] labels: list[str] = [] for index, (start, duration) in enumerate(audio_segments): label = f"a{index}" filters.append( f"[1:a]atrim=start={start:.6f}:duration={duration:.6f}," f"asetpts=PTS-STARTPTS[{label}]" ) labels.append(f"[{label}]") if len(labels) == 1: current_label = labels[0] else: filters.append( "".join(labels) + f"concat=n={len(labels)}:v=0:a=1[acut]" ) current_label = "[acut]" speed_filters = speed_change_filters(audio_tempo, audio_options, audio_sample_rate) if speed_filters: filters.append(f"{current_label}{','.join(speed_filters)}[aout]") elif current_label != "[aout]": filters.append(f"{current_label}anull[aout]") command.extend(["-filter_complex", ";".join(filters), "-map", "[aout]"]) add_audio_encoder_options(command, audio_options) def add_audio_encoder_options( command: list[str], audio_options: AudioEncodeOptions, ) -> None: if audio_options.mode == "aac": command.extend(["-c:a", "aac", "-b:a", audio_options.bitrate]) elif audio_options.mode == "opus": command.extend(["-c:a", "libopus", "-b:a", audio_options.bitrate]) elif audio_options.mode == "flac": command.extend(["-c:a", "flac"]) else: raise RuntimeError(f"Unsupported encoded audio mode: {audio_options.mode}") def atempo_filters(tempo: float) -> list[str]: if tempo <= 0: raise RuntimeError(f"Invalid audio tempo factor: {tempo}") if abs(tempo - 1.0) < 0.0001: return [] filters: list[str] = [] remaining = tempo while remaining < 0.5: filters.append("atempo=0.5") remaining /= 0.5 while remaining > 2.0: filters.append("atempo=2.0") remaining /= 2.0 filters.append(f"atempo={remaining:.8f}") return filters def speed_change_filters( tempo: float, audio_options: AudioEncodeOptions, audio_sample_rate: int | None, ) -> list[str]: if abs(tempo - 1.0) < 0.0001: return [] if audio_options.preserve_pitch: return atempo_filters(tempo) if audio_sample_rate is None: raise RuntimeError( "Audio sample rate is required for speed changes without pitch preservation" ) changed_rate = max(1, round(audio_sample_rate * tempo)) return [f"asetrate={changed_rate}", f"aresample={audio_sample_rate}"] def read_ffmpeg_progress( process: subprocess.Popen, *, total_frames: int | None, label: str, pipeline_progress: PipelineProgressView | None = None, ) -> str: assert process.stderr is not None errors: list[str] = [] progress = ( ProgressView( total_frames or 0, label, embedded_percent=True, show_rate=True, ) if total_frames and sys.stdout.isatty() and pipeline_progress is None else None ) current_frame = 0 if progress is not None: progress.update(0) for line in process.stderr: text = line.rstrip("\r\n") if text.startswith("frame="): try: current_frame = int(text.split("=", 1)[1].strip()) except ValueError: continue if progress is not None: progress.update(current_frame) elif pipeline_progress is not None: pipeline_progress.update_encoding(current_frame) elif text.startswith(("fps=", "stream_", "bitrate=", "total_size=", "out_time_", "dup_frames=", "drop_frames=", "speed=", "progress=")): continue elif text: errors.append(text) process.wait() if progress is not None: progress.finish(keep=True) return "\n".join(errors)