CalvinBackup/remove-ai-watermarks
1
0
Fork 0
mirror of https://github.com/wiltodelta/remove-ai-watermarks.git synced 2026-08-10 16:10:33 +02:00
Code Issues Packages Projects Releases Wiki Activity

Split video encoding from audio muxing

Browse Source
This commit is contained in:
Victor Kuznetsov
2026-07-31 15:22:41 -07:00
parent 9b9f9c6ced
commit 7e63856a90
6 changed files with 261 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/remove_ai_watermarks/video_encoding.py
+108 -37
View File
@@ -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
+27 -28
View File
@@ -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
+12 -6
View File
@@ -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):