Add Avisynth video encoding workflow

This commit is contained in:
ajp_anton
2026-07-21 00:41:00 +00:00
parent f50f465465
commit 49296d6f45
37 changed files with 9487 additions and 76 deletions
+641
View File
@@ -0,0 +1,641 @@
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)