remove-ai-watermarks/src/remove_ai_watermarks/_internal/watermark_remover.py

"""Project-native orchestration for diffusion-based pixel regeneration."""

# 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, reportConstantRedefinition=false, reportUnnecessaryComparison=false
from __future__ import annotations

import logging
import os
import subprocess
from typing import TYPE_CHECKING, Any

from PIL import Image

from remove_ai_watermarks._internal.watermark_profiles import (
    DEFAULT_PROFILE,
    INVISIBLE_EXTRA,
    PROFILE_CHOICES,
    REMOVAL_MODULES,
    SDXL_ZIMAGE_PROFILE,
    normalize_profile,
    resolve_seed,
    resolve_strength,
)
from remove_ai_watermarks.optional_deps import module_available

if TYPE_CHECKING:
    from collections.abc import Callable
    from pathlib import Path

logger = logging.getLogger(__name__)

try:
    import torch

    _HAS_TORCH = True
except ImportError:
    torch = None  # type: ignore[assignment]
    _HAS_TORCH = False

# Probed once at import. ``torch`` is imported above rather than probed because this
# module needs the object, not just the answer.
_HAS_REMOVAL_MODULES = module_available(*(name for name in REMOVAL_MODULES if name != "torch"))


def is_watermark_removal_available() -> bool:
    """Return whether the full removal runtime can be imported."""
    return _HAS_TORCH and _HAS_REMOVAL_MODULES


def _ensure_watermark_deps() -> None:
    if not is_watermark_removal_available():
        raise ImportError(
            f"Invisible watermark regeneration requires the 'qwen-zimage' extra: pip install {INVISIBLE_EXTRA}."
        )


def _has_nvidia_gpu() -> bool:
    try:
        subprocess.run(
            ["nvidia-smi"],
            check=True,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
    except (FileNotFoundError, subprocess.CalledProcessError):
        return False
    return True


def _cuda_works() -> bool:
    try:
        probe = torch.tensor([1.0], device="cuda")  # type: ignore[union-attr]
        _ = probe + probe
    except (AssertionError, RuntimeError):
        return False
    return True


def get_device() -> str:
    """Return ``"cuda"`` when a usable CUDA backend is present, else ``"cpu"``.

    Deliberately binary. Both profiles are CUDA-only, so an XPU or MPS answer would
    only travel one frame further to the same refusal in :class:`WatermarkRemover`,
    while costing a probe on each. ``"cpu"`` here means "no CUDA", which is exactly
    what that refusal reports.
    """
    if not _HAS_TORCH:
        return "cpu"
    if torch.cuda.is_available() and _cuda_works():  # type: ignore[union-attr]
        return "cuda"
    if _has_nvidia_gpu():
        logger.warning("NVIDIA GPU detected, but the installed PyTorch build has no working CUDA backend")
    return "cpu"


class WatermarkRemover:
    """Load one regeneration profile and write a metadata-clean raster output."""

    def __init__(
        self,
        device: str | None = None,
        progress_callback: Callable[[str], None] | None = None,
        hf_token: str | None = None,
        pipeline: str = DEFAULT_PROFILE,
        controlnet_conditioning_scale: float = 1.0,
        cpu_offload: bool = False,
    ) -> None:
        self.model_profile = normalize_profile(pipeline)
        if self.model_profile not in PROFILE_CHOICES:
            raise ValueError(f"Unsupported pipeline '{pipeline}'. Use one of: {', '.join(PROFILE_CHOICES)}.")
        # There is no ``model_id`` parameter and no ``model_id`` attribute: each
        # profile pins a fixed model stack, and the dtype below is bound to that
        # stack's weights. Both used to be constructor overrides that existed only to
        # be rejected or to break the run, and the attribute only existed to echo the
        # rejected value back.
        _ensure_watermark_deps()
        selected_device = (device or get_device()).casefold()
        self.device = get_device() if selected_device == "auto" else selected_device
        # CUDA is a precondition of the object, not of the run. Both profiles raise on
        # any other device, so accepting one here only defers a guaranteed failure to
        # model-load time, several layers down and under the wrong profile's name.
        if self.device != "cuda":
            raise ValueError(
                f"Invisible-watermark removal is CUDA-only, so '{self.device}' cannot run it. "
                "Both remaining profiles need an NVIDIA GPU. Visible-mark removal and "
                "every identify command still run on CPU."
            )

        if self.model_profile == SDXL_ZIMAGE_PROFILE:
            # SDXL ships fp16 weights and an fp16-safe VAE; bf16 would give up the
            # variant without buying anything on this architecture.
            self.torch_dtype = torch.float16  # type: ignore[union-attr]
        else:
            self.torch_dtype = torch.bfloat16  # type: ignore[union-attr]

        self.cpu_offload = cpu_offload
        self.controlnet_conditioning_scale = controlnet_conditioning_scale
        self.hf_token = hf_token or os.environ.get("HF_TOKEN")
        self._progress_callback = progress_callback
        self._qwen_zimage_pipeline: Any = None

    def preload(self, *, global_only: bool = False) -> None:
        """Materialize the selected model stack before the first request."""
        self._load_qwen_zimage_pipeline().preload(global_only=global_only)

    def _load_qwen_zimage_pipeline(self) -> Any:
        if self._qwen_zimage_pipeline is None:
            if self.model_profile == SDXL_ZIMAGE_PROFILE:
                from remove_ai_watermarks._internal.sdxl_zimage_pipeline import (
                    SdxlZImagePipeline as _Pipeline,
                )
            else:
                from remove_ai_watermarks._internal.qwen_zimage_pipeline import (
                    QwenZImagePipeline as _Pipeline,
                )

            self._qwen_zimage_pipeline = _Pipeline(
                device=self.device,
                torch_dtype=self.torch_dtype,
                hf_token=self.hf_token,
                progress_callback=self._progress_callback,
                controlnet_conditioning_scale=self.controlnet_conditioning_scale,
                keep_face_models_on_device=False if self.cpu_offload else None,
                keep_global_models_on_device=False if self.cpu_offload else None,
            )
        return self._qwen_zimage_pipeline

    def _write_output(self, image: Image.Image, output_path: Path) -> None:
        import numpy as np

        from remove_ai_watermarks import image_io

        output_path.parent.mkdir(parents=True, exist_ok=True)
        bgr = np.ascontiguousarray(np.asarray(image.convert("RGB"))[:, :, ::-1])
        if not image_io.imwrite(str(output_path), bgr):
            image.save(output_path)
        from remove_ai_watermarks.metadata import remove_ai_metadata

        remove_ai_metadata(output_path, output_path, keep_standard=True)

    def remove_watermark(
        self,
        image_path: Path,
        output_path: Path | None = None,
        strength: float | None = None,
        seed: int | None = None,
        vendor: str | None = None,
        tile: bool = False,
        tile_size: int = 1024,
        tile_overlap: int = 128,
    ) -> Path:
        """Regenerate image pixels and write the result without AI metadata.

        Step count and CFG are not parameters. Each stage of both profiles is a
        distilled schedule that owns its own, so the only thing a caller-supplied
        value could do is break the run or be rejected.
        """
        if not image_path.exists():
            raise FileNotFoundError(f"Image not found: {image_path}")
        destination = output_path or image_path
        with Image.open(image_path) as opened:
            source = opened.convert("RGB")

        resolved_strength = resolve_strength(strength, vendor, self.model_profile, size=source.size)
        if not 0.0 <= resolved_strength <= 1.0:
            raise ValueError(f"Strength must be between 0.0 and 1.0, got {resolved_strength}")

        result = self._load_qwen_zimage_pipeline().run(
            source,
            strength=resolved_strength,
            seed=resolve_seed(seed),
            tile=tile,
            tile_size=tile_size,
            tile_overlap=tile_overlap,
        )
        self._write_output(result, destination)
        return destination