mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-09 23:50:40 +02:00
470 lines
17 KiB
Python
470 lines
17 KiB
Python
"""Oracle-certified VAE regeneration for video SynthID removal.
|
|
|
|
Google does not publish a local video SynthID decoder. The default profile is
|
|
therefore certified against Google's matching content-verification flow and
|
|
also reports local fidelity metrics.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
# torch/diffusers/cv2 expose incomplete types at this boundary. Pure helpers
|
|
# remain annotated while third-party tensor and image calls are relaxed here.
|
|
# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalCall=false, reportOptionalSubscript=false, reportOptionalOperand=false, reportAttributeAccessIssue=false, reportPrivateImportUsage=false, reportPrivateUsage=false, reportInvalidTypeForm=false
|
|
import logging
|
|
import math
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from remove_ai_watermarks.video_encoding import (
|
|
abort_raw_video_encoder,
|
|
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 (
|
|
DEFAULT_VIDEO_SYNTHID_FPS,
|
|
DEFAULT_VIDEO_SYNTHID_LONG_SIDE,
|
|
DEFAULT_VIDEO_SYNTHID_NOISE_STD,
|
|
DEFAULT_VIDEO_SYNTHID_VAE,
|
|
VIDEO_SYNTHID_LATENT_MULTIPLE,
|
|
)
|
|
from remove_ai_watermarks.video_temporal import (
|
|
_backward_map,
|
|
_motion_residual,
|
|
build_temporal_reference,
|
|
temporal_residual_ratio,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Iterable, Sequence
|
|
from pathlib import Path
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
__all__ = ["build_temporal_reference", "temporal_residual_ratio"]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RegenerationMetrics:
|
|
"""Measured properties of one regenerated video."""
|
|
|
|
frames: int
|
|
fps: float
|
|
width: int
|
|
height: int
|
|
psnr_db: float
|
|
temporal_residual_ratio: float
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class VideoVaeRuntime:
|
|
"""Loaded VAE state reusable across multiple video regenerations."""
|
|
|
|
model: str
|
|
requested_device: str
|
|
resolved_device: str
|
|
vae: Any
|
|
|
|
|
|
def is_available() -> bool:
|
|
"""Return whether the optional VAE runtime can be imported."""
|
|
from remove_ai_watermarks.optional_deps import module_available
|
|
|
|
return module_available("torch", "diffusers")
|
|
|
|
|
|
def _fit_size(width: int, height: int, long_side: int) -> tuple[int, int]:
|
|
"""Fit dimensions to ``long_side`` while preserving aspect and VAE alignment."""
|
|
if width <= 0 or height <= 0:
|
|
raise ValueError("Video dimensions must be positive")
|
|
if long_side < VIDEO_SYNTHID_LATENT_MULTIPLE:
|
|
raise ValueError(f"Long side must be at least {VIDEO_SYNTHID_LATENT_MULTIPLE}")
|
|
scale = long_side / max(width, height)
|
|
fitted_width = max(
|
|
VIDEO_SYNTHID_LATENT_MULTIPLE,
|
|
round(width * scale) // VIDEO_SYNTHID_LATENT_MULTIPLE * VIDEO_SYNTHID_LATENT_MULTIPLE,
|
|
)
|
|
fitted_height = max(
|
|
VIDEO_SYNTHID_LATENT_MULTIPLE,
|
|
round(height * scale) // VIDEO_SYNTHID_LATENT_MULTIPLE * VIDEO_SYNTHID_LATENT_MULTIPLE,
|
|
)
|
|
return fitted_width, fitted_height
|
|
|
|
|
|
def _pick_device(requested: str) -> str:
|
|
import torch
|
|
|
|
if requested == "auto":
|
|
if torch.cuda.is_available():
|
|
return "cuda"
|
|
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
|
return "mps"
|
|
return "cpu"
|
|
if requested == "cuda" and not torch.cuda.is_available():
|
|
raise RuntimeError("CUDA was requested but is not available")
|
|
if requested == "mps" and not (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()):
|
|
raise RuntimeError("MPS was requested but is not available")
|
|
return requested
|
|
|
|
|
|
def load_video_vae_runtime(
|
|
*,
|
|
model: str = DEFAULT_VIDEO_SYNTHID_VAE,
|
|
device: str = "auto",
|
|
) -> VideoVaeRuntime:
|
|
"""Load one reusable video VAE runtime."""
|
|
if device not in {"auto", "cuda", "mps", "cpu"}:
|
|
raise ValueError("device must be auto, cuda, mps, or cpu")
|
|
if not is_available():
|
|
raise RuntimeError("Video SynthID regeneration requires the diffusion extra")
|
|
|
|
import torch
|
|
from diffusers import AutoencoderKL
|
|
|
|
resolved_device = _pick_device(device)
|
|
dtype = torch.float16 if resolved_device == "cuda" else torch.float32
|
|
log.info("Loading %s on %s", model, resolved_device)
|
|
vae = AutoencoderKL.from_pretrained(model, torch_dtype=dtype).to(resolved_device)
|
|
vae.eval()
|
|
vae.enable_slicing()
|
|
return VideoVaeRuntime(
|
|
model=model,
|
|
requested_device=device,
|
|
resolved_device=resolved_device,
|
|
vae=vae,
|
|
)
|
|
|
|
|
|
def _shared_latent_noise(
|
|
spatial_shape: Sequence[int],
|
|
*,
|
|
seed: int,
|
|
device: str,
|
|
dtype: Any,
|
|
) -> Any:
|
|
"""Return one deterministic spatial noise field for reuse across time."""
|
|
import torch
|
|
|
|
if len(spatial_shape) != 3 or any(size <= 0 for size in spatial_shape):
|
|
raise ValueError("Expected a positive CHW latent shape")
|
|
generator = torch.Generator(device="cpu").manual_seed(seed)
|
|
noise = torch.randn((1, *spatial_shape), generator=generator, dtype=torch.float32)
|
|
return noise.to(device=device, dtype=dtype)
|
|
|
|
|
|
def paired_psnr(reference: np.ndarray, candidate: np.ndarray) -> float:
|
|
"""Return paired PSNR over uint8 frame stacks."""
|
|
if reference.shape != candidate.shape:
|
|
raise ValueError("PSNR inputs must have matching shapes")
|
|
mse = float(np.mean((reference.astype(np.float32) - candidate.astype(np.float32)) ** 2))
|
|
if mse == 0.0:
|
|
return math.inf
|
|
return 20.0 * math.log10(255.0 / math.sqrt(mse))
|
|
|
|
|
|
def _probe_video(source: Path) -> tuple[int, int, float]:
|
|
capture = cv2.VideoCapture(str(source))
|
|
if not capture.isOpened():
|
|
raise ValueError(f"Could not open video: {source}")
|
|
try:
|
|
width = round(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
height = round(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
source_fps = float(capture.get(cv2.CAP_PROP_FPS))
|
|
finally:
|
|
capture.release()
|
|
if width <= 0 or height <= 0:
|
|
raise ValueError(f"Video has no usable dimensions: {source}")
|
|
if source_fps <= 0.0:
|
|
raise ValueError(f"Video has no usable frame rate: {source}")
|
|
return width, height, source_fps
|
|
|
|
|
|
def read_sampled_frames(
|
|
source: Path,
|
|
*,
|
|
duration: float | None,
|
|
output_fps: float,
|
|
size: tuple[int, int],
|
|
) -> tuple[list[np.ndarray], float]:
|
|
"""Read uniformly sampled frames and resize them to the VAE geometry."""
|
|
_width, _height, source_fps = _probe_video(source)
|
|
effective_fps = min(output_fps, source_fps)
|
|
frames = list(
|
|
_iter_sampled_frames(
|
|
source,
|
|
source_fps=source_fps,
|
|
duration=duration,
|
|
effective_fps=effective_fps,
|
|
size=size,
|
|
)
|
|
)
|
|
if len(frames) < 2:
|
|
raise ValueError("The selected clip produced fewer than two frames")
|
|
return frames, effective_fps
|
|
|
|
|
|
def _iter_sampled_frames(
|
|
source: Path,
|
|
*,
|
|
source_fps: float,
|
|
duration: float | None,
|
|
effective_fps: float,
|
|
size: tuple[int, int],
|
|
) -> Iterable[np.ndarray]:
|
|
"""Yield uniformly sampled frames without retaining the full video."""
|
|
capture = cv2.VideoCapture(str(source))
|
|
if not capture.isOpened():
|
|
raise ValueError(f"Could not open video: {source}")
|
|
sample_period = 1.0 / effective_fps
|
|
next_sample_time = 0.0
|
|
frame_index = 0
|
|
try:
|
|
while True:
|
|
ok, frame = capture.read()
|
|
if not ok:
|
|
break
|
|
timestamp = frame_index / source_fps
|
|
if duration is not None and timestamp + 1e-9 >= duration:
|
|
break
|
|
if timestamp + 1e-9 >= next_sample_time:
|
|
yield cv2.resize(frame, size, interpolation=cv2.INTER_LANCZOS4)
|
|
next_sample_time += sample_period
|
|
frame_index += 1
|
|
finally:
|
|
capture.release()
|
|
|
|
|
|
def _frame_batches(frames: Sequence[np.ndarray], batch_size: int) -> Iterable[Sequence[np.ndarray]]:
|
|
for start in range(0, len(frames), batch_size):
|
|
yield frames[start : start + batch_size]
|
|
|
|
|
|
def _stream_batches(frames: Iterable[np.ndarray], batch_size: int) -> Iterable[list[np.ndarray]]:
|
|
batch: list[np.ndarray] = []
|
|
for frame in frames:
|
|
batch.append(frame)
|
|
if len(batch) == batch_size:
|
|
yield batch
|
|
batch = []
|
|
if batch:
|
|
yield batch
|
|
|
|
|
|
def _encode_frame_latents(
|
|
frames: Sequence[np.ndarray],
|
|
*,
|
|
vae: Any,
|
|
device: str,
|
|
batch_size: int,
|
|
) -> list[Any]:
|
|
"""Encode source frames once so every candidate can reuse identical latents."""
|
|
import torch
|
|
|
|
latent_batches: list[Any] = []
|
|
scaling_factor = float(vae.config.scaling_factor)
|
|
with torch.inference_mode():
|
|
for batch in _frame_batches(frames, batch_size):
|
|
rgb = np.stack([frame[:, :, ::-1] for frame in batch])
|
|
tensor = torch.from_numpy(np.ascontiguousarray(rgb)).permute(0, 3, 1, 2)
|
|
tensor = tensor.to(device=device, dtype=vae.dtype) / 127.5 - 1.0
|
|
latents = vae.encode(tensor).latent_dist.mode() * scaling_factor
|
|
latent_batches.append(latents)
|
|
return latent_batches
|
|
|
|
|
|
def _decode_frame_latents(
|
|
latent_batches: Sequence[Any],
|
|
*,
|
|
vae: Any,
|
|
noise_std: float,
|
|
shared_noise: Any,
|
|
) -> list[np.ndarray]:
|
|
"""Decode cached latents with one perturbation shared across time."""
|
|
import torch
|
|
|
|
output: list[np.ndarray] = []
|
|
scaling_factor = float(vae.config.scaling_factor)
|
|
with torch.inference_mode():
|
|
for latents in latent_batches:
|
|
perturbed = latents + noise_std * shared_noise.expand(latents.shape[0], -1, -1, -1)
|
|
decoded = vae.decode(perturbed / scaling_factor).sample
|
|
decoded = ((decoded / 2.0 + 0.5).clamp(0.0, 1.0) * 255.0).round().to(torch.uint8)
|
|
decoded = decoded.permute(0, 2, 3, 1).cpu().numpy()
|
|
output.extend(np.ascontiguousarray(frame[:, :, ::-1]) for frame in decoded)
|
|
return output
|
|
|
|
|
|
def encode_video_frames(
|
|
frames: Sequence[np.ndarray],
|
|
source: Path,
|
|
output: Path,
|
|
*,
|
|
fps: float,
|
|
) -> None:
|
|
"""Encode regenerated frames, copy audio, and omit source metadata."""
|
|
if not frames:
|
|
raise ValueError("At least one frame is required for video encoding")
|
|
height, width = frames[0].shape[:2]
|
|
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, 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(
|
|
source: Path,
|
|
output: Path,
|
|
*,
|
|
noise_std: float = DEFAULT_VIDEO_SYNTHID_NOISE_STD,
|
|
long_side: int = DEFAULT_VIDEO_SYNTHID_LONG_SIDE,
|
|
fps: float = DEFAULT_VIDEO_SYNTHID_FPS,
|
|
batch_size: int = 4,
|
|
seed: int = 0,
|
|
model: str = DEFAULT_VIDEO_SYNTHID_VAE,
|
|
device: str = "auto",
|
|
duration: float | None = None,
|
|
runtime: VideoVaeRuntime | None = None,
|
|
) -> RegenerationMetrics:
|
|
"""Regenerate video pixels and return fidelity metrics.
|
|
|
|
The default profile is certified against the provider oracle. This function
|
|
does not perform a per-file SynthID decode because Google exposes no local
|
|
decoder.
|
|
"""
|
|
if not 0.0 <= noise_std <= 1.0:
|
|
raise ValueError("noise_std must be between 0 and 1")
|
|
if fps < 1.0:
|
|
raise ValueError("fps must be at least 1")
|
|
if batch_size < 1:
|
|
raise ValueError("batch_size must be at least 1")
|
|
if duration is not None and duration <= 0:
|
|
raise ValueError("duration must be positive")
|
|
if device not in {"auto", "cuda", "mps", "cpu"}:
|
|
raise ValueError("device must be auto, cuda, mps, or cpu")
|
|
width, height, source_fps = _probe_video(source)
|
|
size = _fit_size(width, height, long_side)
|
|
effective_fps = min(fps, source_fps)
|
|
|
|
if runtime is None:
|
|
runtime = load_video_vae_runtime(model=model, device=device)
|
|
elif runtime.model != model or runtime.requested_device != device:
|
|
raise ValueError("The supplied video VAE runtime does not match the requested model and device")
|
|
resolved_device = runtime.resolved_device
|
|
vae = runtime.vae
|
|
|
|
with staged_video_output(output) as (encoded_video, temporary_output):
|
|
process = start_raw_video_encoder(
|
|
raw_video_command(
|
|
encoded_video,
|
|
width=size[0],
|
|
height=size[1],
|
|
fps=effective_fps,
|
|
crf=18,
|
|
profile=probe_video_encode_profile(source),
|
|
)
|
|
)
|
|
frame_pipe = process.stdin
|
|
frame_count = 0
|
|
squared_error = 0.0
|
|
pixel_count = 0
|
|
temporal_baseline = 0.0
|
|
temporal_candidate = 0.0
|
|
previous_gray: np.ndarray | None = None
|
|
previous_reference_f32: np.ndarray | None = None
|
|
previous_candidate_f32: np.ndarray | None = None
|
|
shared_noise: Any | None = None
|
|
try:
|
|
sampled_frames = _iter_sampled_frames(
|
|
source,
|
|
source_fps=source_fps,
|
|
duration=duration,
|
|
effective_fps=effective_fps,
|
|
size=size,
|
|
)
|
|
for frames in _stream_batches(sampled_frames, batch_size):
|
|
latent_batches = _encode_frame_latents(
|
|
frames,
|
|
vae=vae,
|
|
device=resolved_device,
|
|
batch_size=batch_size,
|
|
)
|
|
latents = latent_batches[0]
|
|
if shared_noise is None:
|
|
shared_noise = _shared_latent_noise(
|
|
latents.shape[1:],
|
|
seed=seed,
|
|
device=resolved_device,
|
|
dtype=latents.dtype,
|
|
)
|
|
regenerated = _decode_frame_latents(
|
|
latent_batches,
|
|
vae=vae,
|
|
noise_std=noise_std,
|
|
shared_noise=shared_noise,
|
|
)
|
|
for reference, candidate in zip(frames, regenerated, strict=True):
|
|
frame_pipe.write(candidate.tobytes())
|
|
reference_f32 = reference.astype(np.float32)
|
|
candidate_f32 = candidate.astype(np.float32)
|
|
difference = reference_f32 - candidate_f32
|
|
squared_error += float(np.sum(difference * difference, dtype=np.float64))
|
|
pixel_count += reference.size
|
|
current_gray = cv2.cvtColor(reference, cv2.COLOR_BGR2GRAY)
|
|
if (
|
|
previous_gray is not None
|
|
and previous_reference_f32 is not None
|
|
and previous_candidate_f32 is not None
|
|
):
|
|
frame_maps = _backward_map(current_gray, previous_gray)
|
|
temporal_baseline += _motion_residual(reference_f32, previous_reference_f32, frame_maps)
|
|
temporal_candidate += _motion_residual(candidate_f32, previous_candidate_f32, frame_maps)
|
|
previous_gray = current_gray
|
|
previous_reference_f32 = reference_f32
|
|
previous_candidate_f32 = candidate_f32
|
|
frame_count += 1
|
|
if frame_count < 2:
|
|
raise ValueError("The selected clip produced fewer than two frames")
|
|
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
|
|
|
|
mse = squared_error / pixel_count
|
|
return RegenerationMetrics(
|
|
frames=frame_count,
|
|
fps=effective_fps,
|
|
width=size[0],
|
|
height=size[1],
|
|
psnr_db=math.inf if mse == 0.0 else 20.0 * math.log10(255.0 / math.sqrt(mse)),
|
|
temporal_residual_ratio=temporal_candidate / max(temporal_baseline, 1e-6),
|
|
)
|