Split video encoding from audio muxing

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

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)

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

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):

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,

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,