remove-ai-watermarks/src/remove_ai_watermarks/video_invisible.py

"""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,
    VIDEO_SYNTHID_VAE_SCALING_FACTOR,
)
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
    # The gate that validates this factor and the encode/decode calls that apply it
    # must read one value, not three independent reads of the same attribute.
    scaling_factor: float


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()
    scaling_factor = float(vae.config.scaling_factor)
    log.info("Latent scaling factor %.5f", scaling_factor)
    if model != DEFAULT_VIDEO_SYNTHID_VAE:
        log.warning("No oracle-certified profile exists for %s; the shipped noise_std is not calibrated for it", model)
    elif scaling_factor != VIDEO_SYNTHID_VAE_SCALING_FACTOR:
        raise RuntimeError(
            f"{model} loaded with latent scaling factor {scaling_factor}, but the certified "
            f"profile is defined against {VIDEO_SYNTHID_VAE_SCALING_FACTOR}. The perturbation "
            "would be rescaled and the output would no longer match any certified row."
        )
    return VideoVaeRuntime(
        model=model,
        requested_device=device,
        resolved_device=resolved_device,
        vae=vae,
        scaling_factor=scaling_factor,
    )


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,
    scaling_factor: float,
) -> list[Any]:
    """Encode source frames once so every candidate can reuse identical latents."""
    import torch

    latent_batches: list[Any] = []
    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,
    scaling_factor: float,
) -> list[np.ndarray]:
    """Decode cached latents with one perturbation shared across time."""
    import torch

    output: list[np.ndarray] = []
    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,
                    scaling_factor=runtime.scaling_factor,
                )
                latents = latent_batches[0]
                if shared_noise is None:
                    # Removal strength is the ratio of the perturbation to this spread,
                    # not noise_std alone: it is the only local quantity that makes two
                    # models' doses comparable.
                    log.info(
                        "First latent batch spread %.4f against noise_std %.4f",
                        float(latents.float().std()),
                        noise_std,
                    )
                    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,
                    scaling_factor=runtime.scaling_factor,
                )
                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),
        )