From 7e63856a90b432c0d446ae8189586a52d792f32b Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Fri, 31 Jul 2026 15:22:41 -0700 Subject: [PATCH] Split video encoding from audio muxing --- docs/module-internals.md | 20 ++- src/remove_ai_watermarks/video_encoding.py | 145 +++++++++++++++----- src/remove_ai_watermarks/video_invisible.py | 55 ++++---- src/remove_ai_watermarks/video_visible.py | 18 ++- tests/test_video.py | 43 +++++- tests/test_video_invisible.py | 72 ++++++++-- 6 files changed, 261 insertions(+), 92 deletions(-) diff --git a/docs/module-internals.md b/docs/module-internals.md index 284f6fd..b640e00 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -158,10 +158,15 @@ tail when diagnostics are unusually large. Aborts release it even when ffmpeg has already exited. A real subprocess regression writes diagnostics beyond pipe capacity while streaming frames, checks bounded failure reporting, and the Linux full-clip CI job guards the complete path. -The finite source file is opened before the frame pipe, so ffmpeg can initialize -the copied audio stream before producer backpressure is possible. Stream and -metadata mappings are source-indexed accordingly; regressions assert the input -order and both map targets. +Frame encoding and source-audio copying run as two ffmpeg processes in sequence. +The streaming encoder has only the frame pipe as input, so input probing or +demux queues cannot deadlock the producer against a second input. After that +pipe reaches EOF, a finite stream-copy mux combines the encoded video with the +source audio and applies the requested metadata/chapter policy. Both stages use +sibling temporary files, and only the completed mux is published atomically. +The mux also redirects diagnostics to disk and reports only a bounded head and +tail. Command regressions assert the single-input encoder and final map targets; +failure regressions cover bounded mux diagnostics and atomic cleanup. `probe_video_encode_profile` reads the first source video stream with ffprobe and preserves the supported properties that survive the 8-bit BGR boundary: `yuv420p`/`yuv422p`/`yuv444p` chroma sampling, recognized color tags, encoder @@ -196,9 +201,10 @@ space, applies one seeded spatial-noise field across the entire sequence, and decodes fresh pixels. Reusing a single noise field avoids independent frame-to-frame noise. The shipped path retains only one configured frame batch, updates PSNR and temporal residuals incrementally, and streams BGR frames -directly to ffmpeg. ffmpeg encodes H.264 video, maps optional source audio, and -drops all source metadata. The result is written through a same-directory -temporary file and atomically replaced only after a successful encode. +directly to the video-only ffmpeg encoder. A separate stream-copy mux then adds +optional source audio and drops all source metadata. The result is written +through same-directory temporary files and atomically replaced only after both +stages succeed. The engine returns PSNR and a motion-compensated temporal-residual ratio as quality measurements. Neither is a watermark detector. The high-level result diff --git a/src/remove_ai_watermarks/video_encoding.py b/src/remove_ai_watermarks/video_encoding.py index fe5f5f4..36ce507 100644 --- a/src/remove_ai_watermarks/video_encoding.py +++ b/src/remove_ai_watermarks/video_encoding.py @@ -106,21 +106,7 @@ class _RawVideoEncoder: if self._stderr_buffer.closed: raise RuntimeError("ffmpeg diagnostics have already been collected") try: - self._stderr_buffer.seek(0, os.SEEK_END) - size = self._stderr_buffer.tell() - if size <= _FFMPEG_STDERR_LIMIT: - self._stderr_buffer.seek(0) - raw_stderr = self._stderr_buffer.read() - else: - payload_limit = _FFMPEG_STDERR_LIMIT - len(_FFMPEG_STDERR_TRUNCATION) - head_size = payload_limit // 2 - tail_size = payload_limit - head_size - self._stderr_buffer.seek(0) - head = self._stderr_buffer.read(head_size) - self._stderr_buffer.seek(-tail_size, os.SEEK_END) - tail = self._stderr_buffer.read(tail_size) - raw_stderr = head + _FFMPEG_STDERR_TRUNCATION + tail - return raw_stderr.decode("utf-8", errors="replace") + return _read_bounded_stderr(self._stderr_buffer) finally: self._stderr_buffer.close() @@ -131,6 +117,25 @@ class _RawVideoEncoder: self._stderr_buffer.close() +def _read_bounded_stderr(buffer: BinaryIO) -> str: + """Read bounded head-and-tail diagnostics from a seekable binary stream.""" + buffer.seek(0, os.SEEK_END) + size = buffer.tell() + if size <= _FFMPEG_STDERR_LIMIT: + buffer.seek(0) + raw_stderr = buffer.read() + else: + payload_limit = _FFMPEG_STDERR_LIMIT - len(_FFMPEG_STDERR_TRUNCATION) + head_size = payload_limit // 2 + tail_size = payload_limit - head_size + buffer.seek(0) + head = buffer.read(head_size) + buffer.seek(-tail_size, os.SEEK_END) + tail = buffer.read(tail_size) + raw_stderr = head + _FFMPEG_STDERR_TRUNCATION + tail + return raw_stderr.decode("utf-8", errors="replace") + + def _known_value(value: object, allowed: frozenset[str]) -> str | None: """Return a supported ffmpeg enum value, otherwise omit it.""" return value if isinstance(value, str) and value in allowed else None @@ -268,11 +273,11 @@ def probe_video_timestamps(source: Path) -> tuple[float, ...]: @contextmanager -def atomic_video_output(output: Path) -> Generator[Path]: - """Yield a sibling temporary path and publish it only after success.""" +def _temporary_video_path(output: Path, *, prefix: str) -> Generator[Path]: + """Yield one sibling temporary path and remove it on exit.""" output.parent.mkdir(parents=True, exist_ok=True) with tempfile.NamedTemporaryFile( - prefix=f".{output.stem}-", + prefix=prefix, suffix=output.suffix, dir=output.parent, delete=False, @@ -280,11 +285,18 @@ def atomic_video_output(output: Path) -> Generator[Path]: temporary_output = Path(stream.name) try: yield temporary_output - os.replace(temporary_output, output) finally: temporary_output.unlink(missing_ok=True) +@contextmanager +def atomic_video_output(output: Path) -> Generator[Path]: + """Yield a sibling temporary path and publish it only after success.""" + with _temporary_video_path(output, prefix=f".{output.stem}-") as temporary_output: + yield temporary_output + os.replace(temporary_output, output) + + def _video_codec_args(suffix: str, *, crf: int, profile: VideoEncodeProfile) -> list[str]: if suffix == ".webm": return ["-c:v", "libvpx-vp9", "-crf", str(crf), "-b:v", "0"] @@ -320,19 +332,17 @@ def _profile_args(profile: VideoEncodeProfile) -> list[str]: def raw_video_command( - source: Path, output: Path, *, width: int, height: int, fps: float, - strip_metadata: bool, crf: int, profile: VideoEncodeProfile, timestamped_input: bool = False, copy_input_timestamps: bool = False, ) -> list[str]: - """Build a source-aware ffmpeg command for BGR frames on standard input.""" + """Build an ffmpeg command for BGR frames on standard input.""" ffmpeg = shutil.which("ffmpeg") if ffmpeg is None: raise RuntimeError("Video processing requires ffmpeg on PATH") @@ -358,36 +368,97 @@ def raw_video_command( "-loglevel", "error", *(["-copyts"] if copy_input_timestamps else []), - "-i", - str(source), *frame_input, "-map", - "1:v:0", - "-map", - "0:a?", + "0:v:0", *_video_codec_args(output.suffix.lower(), crf=crf, profile=profile), *_profile_args(profile), - "-c:a", - "copy", "-map_metadata", - "-1" if strip_metadata else "0", + "-1", "-map_chapters", - "-1" if strip_metadata else "0", + "-1", ] if timestamped_input: command.extend(["-fps_mode", "passthrough"]) if copy_input_timestamps: command.extend(["-avoid_negative_ts", "disabled"]) - if output.suffix.lower() in {".mp4", ".mov", ".m4v"}: - if profile.time_base is not None: - numerator, denominator = (int(part) for part in profile.time_base.split("/", 1)) - if numerator == 1: - command.extend(["-video_track_timescale", str(denominator)]) - command.extend(["-movflags", "+faststart"]) + if output.suffix.lower() in {".mp4", ".mov", ".m4v"} and profile.time_base is not None: + numerator, denominator = (int(part) for part in profile.time_base.split("/", 1)) + if numerator == 1: + command.extend(["-video_track_timescale", str(denominator)]) command.append(str(output)) return command +def mux_encoded_video( + encoded_video: Path, + source: Path, + output: Path, + *, + strip_metadata: bool, + copy_input_timestamps: bool = False, +) -> None: + """Copy encoded video and source audio into the final container.""" + ffmpeg = shutil.which("ffmpeg") + if ffmpeg is None: + raise RuntimeError("Video processing requires ffmpeg on PATH") + command = [ + ffmpeg, + "-y", + "-loglevel", + "error", + *(["-copyts"] if copy_input_timestamps else []), + "-i", + str(encoded_video), + "-i", + str(source), + "-map", + "0:v:0", + "-map", + "1:a?", + "-c", + "copy", + "-map_metadata", + "-1" if strip_metadata else "1", + "-map_chapters", + "-1" if strip_metadata else "1", + ] + if copy_input_timestamps: + command.extend(["-avoid_negative_ts", "disabled"]) + if output.suffix.lower() in {".mp4", ".mov", ".m4v"}: + command.extend(["-movflags", "+faststart"]) + command.append(str(output)) + with tempfile.TemporaryFile(mode="w+b") as stderr_buffer: + result = subprocess.run( # noqa: S603 + command, + stdout=subprocess.DEVNULL, + stderr=stderr_buffer, + check=False, + ) + stderr = _read_bounded_stderr(stderr_buffer) + log.info( + "ffmpeg video mux: command=%s status=%s stderr=%s", + command, + result.returncode, + stderr, + ) + if result.returncode != 0: + raise RuntimeError(f"ffmpeg failed to mux {output}: {stderr.strip()[:500]}") + + +@contextmanager +def staged_video_output(output: Path) -> Generator[tuple[Path, Path]]: + """Yield video-only and final temporary paths, then publish atomically.""" + with ( + atomic_video_output(output) as temporary_output, + _temporary_video_path( + output, + prefix=f".{output.stem}-video-", + ) as encoded_video, + ): + yield encoded_video, temporary_output + + def start_raw_video_encoder(command: list[str]) -> _RawVideoEncoder: """Start ffmpeg and validate its raw-frame pipes.""" log.info("Starting ffmpeg video encode: command=%s", command) diff --git a/src/remove_ai_watermarks/video_invisible.py b/src/remove_ai_watermarks/video_invisible.py index 832b95d..45da702 100644 --- a/src/remove_ai_watermarks/video_invisible.py +++ b/src/remove_ai_watermarks/video_invisible.py @@ -21,10 +21,11 @@ import numpy as np from remove_ai_watermarks.video_encoding import ( abort_raw_video_encoder, - atomic_video_output, finish_raw_video_encoder, + mux_encoded_video, probe_video_encode_profile, raw_video_command, + staged_video_output, start_raw_video_encoder, ) from remove_ai_watermarks.video_synthid import ( @@ -309,29 +310,28 @@ def encode_video_frames( if not frames: raise ValueError("At least one frame is required for video encoding") height, width = frames[0].shape[:2] - output.parent.mkdir(parents=True, exist_ok=True) - process = start_raw_video_encoder( - raw_video_command( - source, - output, - width=width, - height=height, - fps=fps, - strip_metadata=True, - crf=18, - profile=probe_video_encode_profile(source), + with staged_video_output(output) as (encoded_video, temporary_output): + process = start_raw_video_encoder( + raw_video_command( + encoded_video, + width=width, + height=height, + fps=fps, + crf=18, + profile=probe_video_encode_profile(source), + ) ) - ) - frame_pipe = process.stdin - try: - for frame in frames: - if frame.shape[:2] != (height, width): - raise ValueError("Video frames must have matching dimensions") - frame_pipe.write(frame.tobytes()) - finish_raw_video_encoder(process, output, operation="SynthID removal encode") - except Exception: - abort_raw_video_encoder(process) - raise + frame_pipe = process.stdin + try: + for frame in frames: + if frame.shape[:2] != (height, width): + raise ValueError("Video frames must have matching dimensions") + frame_pipe.write(frame.tobytes()) + finish_raw_video_encoder(process, encoded_video, operation="SynthID removal encode") + mux_encoded_video(encoded_video, source, temporary_output, strip_metadata=True) + except Exception: + abort_raw_video_encoder(process) + raise def regenerate_video_candidate( @@ -375,15 +375,13 @@ def regenerate_video_candidate( resolved_device = runtime.resolved_device vae = runtime.vae - with atomic_video_output(output) as temporary_output: + with staged_video_output(output) as (encoded_video, temporary_output): process = start_raw_video_encoder( raw_video_command( - source, - temporary_output, + encoded_video, width=size[0], height=size[1], fps=effective_fps, - strip_metadata=True, crf=18, profile=probe_video_encode_profile(source), ) @@ -451,9 +449,10 @@ def regenerate_video_candidate( raise ValueError("The selected clip produced fewer than two frames") finish_raw_video_encoder( process, - temporary_output, + encoded_video, operation="SynthID removal encode", ) + mux_encoded_video(encoded_video, source, temporary_output, strip_metadata=True) except Exception: abort_raw_video_encoder(process) raise diff --git a/src/remove_ai_watermarks/video_visible.py b/src/remove_ai_watermarks/video_visible.py index a6328d5..649f514 100644 --- a/src/remove_ai_watermarks/video_visible.py +++ b/src/remove_ai_watermarks/video_visible.py @@ -39,11 +39,12 @@ from PIL import Image, ImageDraw, ImageFont from remove_ai_watermarks.video import VIDEO_VISIBLE_MARKS from remove_ai_watermarks.video_encoding import ( abort_raw_video_encoder, - atomic_video_output, finish_raw_video_encoder, + mux_encoded_video, probe_video_encode_profile, probe_video_timestamps, raw_video_command, + staged_video_output, start_raw_video_encoder, ) from remove_ai_watermarks.video_temporal import stabilize_filled_frame @@ -1331,7 +1332,7 @@ def encode_clean_video( if len(regions) != len(scan.detections): raise ValueError("Temporal localization count does not match the scanned frame count") - with atomic_video_output(output) as temporary_output: + with staged_video_output(output) as (encoded_video, temporary_output): profile = probe_video_encode_profile(source) if (profile.component_depth or 0) > 8 or profile.color_transfer in _HDR_TRANSFERS: source_format = profile.source_pixel_format or "unknown pixel format" @@ -1350,12 +1351,10 @@ def encode_clean_video( ) process = start_raw_video_encoder( raw_video_command( - source, - temporary_output, + encoded_video, width=scan.width, height=scan.height, fps=scan.fps, - strip_metadata=strip_metadata, crf=14, profile=profile, timestamped_input=timestamped_input, @@ -1426,9 +1425,16 @@ def encode_clean_video( timestamped_writer = None finish_raw_video_encoder( process, - temporary_output, + encoded_video, operation="visible-watermark encode", ) + mux_encoded_video( + encoded_video, + source, + temporary_output, + strip_metadata=strip_metadata, + copy_input_timestamps=preserve_start_offset, + ) except Exception: if timestamped_writer is not None: with suppress(Exception): diff --git a/tests/test_video.py b/tests/test_video.py index 5a2b894..927808f 100644 --- a/tests/test_video.py +++ b/tests/test_video.py @@ -1992,7 +1992,7 @@ class TestVideoVisibleEncoding: targets: list[Path] = [] - def fake_command(_source: Path, target: Path, **_kwargs: object) -> list[str]: + def fake_command(target: Path, **_kwargs: object) -> list[str]: targets.append(target) return ["ffmpeg", str(target)] @@ -2001,11 +2001,15 @@ class TestVideoVisibleEncoding: if fail: raise RuntimeError("synthetic encode failure") + def fake_mux(encoded: Path, _source: Path, target: Path, **_kwargs: object) -> None: + target.write_bytes(encoded.read_bytes()) + process = FakeProcess() monkeypatch.setattr(video_visible.cv2, "VideoCapture", lambda _path: FakeCapture()) monkeypatch.setattr(video_visible, "raw_video_command", fake_command) monkeypatch.setattr(video_visible, "start_raw_video_encoder", lambda _command: process) monkeypatch.setattr(video_visible, "finish_raw_video_encoder", fake_finish) + monkeypatch.setattr(video_visible, "mux_encoded_video", fake_mux) monkeypatch.setattr(watermark_registry, "resolve_backend", lambda _backend: "cv2") return process, targets @@ -2071,6 +2075,43 @@ class TestVideoVisibleEncoding: assert output.read_bytes() == b"previous" assert not targets[0].exists() + def test_failed_mux_preserves_existing_output( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks import video_visible + from remove_ai_watermarks.video_visible import FrameLocalization, VideoScan, encode_clean_video + + source = tmp_path / "source.mp4" + source.write_bytes(b"source") + output = tmp_path / "clean.mp4" + output.write_bytes(b"previous") + scan = VideoScan(8, 8, 24.0, (FrameLocalization(0, 0.0, None),)) + _process, targets = self._patch_single_frame_encode( + monkeypatch, + encoded_bytes=b"complete video stream", + fail=False, + ) + + def fail_mux(*_args: object, **_kwargs: object) -> None: + raise RuntimeError("synthetic mux failure") + + monkeypatch.setattr(video_visible, "mux_encoded_video", fail_mux) + + with pytest.raises(RuntimeError, match="synthetic mux failure"): + encode_clean_video( + source, + output, + scan, + [None], + backend="cv2", + strip_metadata=True, + ) + + assert output.read_bytes() == b"previous" + assert not targets[0].exists() + def test_rejects_high_bit_depth_before_silent_downconversion( self, tmp_path: Path, diff --git a/tests/test_video_invisible.py b/tests/test_video_invisible.py index 6df1cd2..abc030d 100644 --- a/tests/test_video_invisible.py +++ b/tests/test_video_invisible.py @@ -5,7 +5,7 @@ from __future__ import annotations import sys import threading from types import SimpleNamespace -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast import pytest @@ -14,6 +14,7 @@ from remove_ai_watermarks.video_synthid import DEFAULT_VIDEO_SYNTHID_NOISE_STD if TYPE_CHECKING: from pathlib import Path + from typing import BinaryIO def test_encoder_redirects_large_stderr_while_frames_are_streaming( @@ -84,7 +85,7 @@ def test_regeneration_rejects_noise_outside_unit_interval(tmp_path: Path) -> Non ) -def test_encoder_command_discards_metadata_and_copies_audio( +def test_encoder_and_mux_commands_separate_streaming_frames_from_source_audio( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -102,20 +103,16 @@ def test_encoder_command_discards_metadata_and_copies_audio( ) command = video_encoding.raw_video_command( - source, output, width=8, height=8, fps=2.0, - strip_metadata=True, crf=18, profile=profile, ) metadata_index = command.index("-map_metadata") assert command[metadata_index + 1] == "-1" - audio_codec_index = command.index("-c:a") - assert command[audio_codec_index + 1] == "copy" output_pixel_format_index = command.index("-pix_fmt", command.index("-c:v")) assert command[output_pixel_format_index + 1] == "yuv420p" assert command[command.index("-color_range") + 1] == "tv" @@ -128,28 +125,77 @@ def test_encoder_command_discards_metadata_and_copies_audio( "colorprim=bt709:transfer=bt709:colormatrix=bt709:range=limited" ) assert "pipe:0" in command - assert command.index(str(source)) < command.index("pipe:0") - assert command[command.index("-map") + 1] == "1:v:0" - second_map = command.index("-map", command.index("-map") + 1) - assert command[second_map + 1] == "0:a?" + assert str(source) not in command + assert command[command.index("-map") + 1] == "0:v:0" assert "-shortest" not in command + calls: list[list[str]] = [] + + def fake_run(mux_command: list[str], **_kwargs: object) -> SimpleNamespace: + calls.append(mux_command) + return SimpleNamespace(returncode=0, stdout="", stderr="") + + monkeypatch.setattr(video_encoding.subprocess, "run", fake_run) + encoded_video = tmp_path / "encoded.mp4" + video_encoding.mux_encoded_video( + encoded_video, + source, + output, + strip_metadata=True, + ) + + mux_command = calls[0] + assert mux_command.index(str(encoded_video)) < mux_command.index(str(source)) + assert mux_command[mux_command.index("-map") + 1] == "0:v:0" + second_map = mux_command.index("-map", mux_command.index("-map") + 1) + assert mux_command[second_map + 1] == "1:a?" + assert mux_command[mux_command.index("-c") + 1] == "copy" + assert mux_command[mux_command.index("-map_metadata") + 1] == "-1" + assert "-shortest" not in mux_command + + +def test_mux_reports_bounded_disk_backed_stderr( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + diagnostic = b"synthetic mux diagnostic" + tail_diagnostic = b"synthetic mux diagnostic tail" + caplog.set_level("INFO", logger=video_encoding.__name__) + monkeypatch.setattr(video_encoding.shutil, "which", lambda _name: "/usr/bin/ffmpeg") + + def fake_run(_command: list[str], **kwargs: object) -> SimpleNamespace: + stderr = cast("BinaryIO", kwargs["stderr"]) + stderr.write(diagnostic + b"x" * 262144 + tail_diagnostic) + stderr.flush() + return SimpleNamespace(returncode=7) + + monkeypatch.setattr(video_encoding.subprocess, "run", fake_run) + + with pytest.raises(RuntimeError, match=diagnostic.decode()): + video_encoding.mux_encoded_video( + tmp_path / "encoded.mp4", + tmp_path / "source.mp4", + tmp_path / "output.mp4", + strip_metadata=True, + ) + + assert "ffmpeg stderr truncated" in caplog.text + assert tail_diagnostic.decode() in caplog.text + def test_timestamped_encoder_reads_nut_and_passes_pts_through( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: - source = tmp_path / "source.mp4" output = tmp_path / "candidate.mp4" monkeypatch.setattr(video_encoding.shutil, "which", lambda _name: "/usr/bin/ffmpeg") command = video_encoding.raw_video_command( - source, output, width=8, height=8, fps=24.0, - strip_metadata=True, crf=18, profile=video_encoding.VideoEncodeProfile(time_base="1/90000"), timestamped_input=True,