mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-09 23:50:40 +02:00
`InvisibleOptions` promises in its docstring that every default mirrors
`InvisibleEngine`. Two fields made that promise cost something to keep: `force` is
not an engine parameter at all, and `controlnet_scale` was a third spelling of the
engine's `controlnet_conditioning_scale`. The mirror test carried an exception
table for each. This removes both, so the comparison needs no exceptions -- a field
that needs one is a field that belongs somewhere else.
`force` decides WHETHER the engine runs, which is settled before it is built, so it
joins `backend` and `sensitivity` as a parameter of `remove_all` and `remove_batch`
and is threaded to `_run_invisible` as its own argument. `controlnet_scale` takes
the engine's own name; the click option stays `--controlnet-scale` and is now
translated exactly once instead of at three forwarding sites.
Safe to do today: both symbols landed after 0.25.0 and have never been published.
The forwarding turned out to be the weaker half. A defaults comparison cannot see a
hardcoded literal at the seam, and `_run_invisible` passed the entire suite with
`controlnet_conditioning_scale` pinned to a constant. Each of the two knobs also
reaches the engine through TWO paths -- `remove_all` versus `remove_batch(mode="all")`
for `force`, `_run_invisible` versus `_batch_engine` for the scale -- and guarding one
left the other free to hardcode with a green suite. So:
* `test_every_field_arrives_at_the_engine_with_the_caller_s_value` drives the real
seam with all 13 fields set off their defaults; mutating any one of them to its
default fails it.
* `test_force_reaches_the_scrub_gate_in_every_scrubbing_mode` and
`test_batch_controlnet_scale_flows_to_the_cached_engine` are parametrized over
both modes, so neither path can be pinned alone.
Also fixes an order-dependent test surfaced by the added tests reshuffling the xdist
shards. `test_visible_path_decodes_file_once` counted every `image_io.imread` in the
process, but the Gemini engine loads its own bundled capture assets on first
construction, so the count was 3 on a cold engine and 1 on a warm one and the test
passed only when an earlier test happened to build the engine first. It now counts
decodes of the SOURCE, which is the invariant it exists for, and still fails when the
shared decode is broken. The production path was never wrong: the source bitmap is
decoded exactly once.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
680 lines
28 KiB
Python
680 lines
28 KiB
Python
"""High-level convenience API: clean an image (or array) in one call.
|
|
|
|
The low-level building blocks live in ``watermark_registry`` (localize -> fill) and
|
|
``image_io`` (Unicode-safe, alpha-preserving IO). This module ties them into the two
|
|
calls a caller usually wants, so a library user does not have to decode images, wire
|
|
up metadata provenance, or preserve the alpha channel by hand:
|
|
|
|
import remove_ai_watermarks as raiw
|
|
raiw.remove_visible("in.png", "out.png") # path -> file, provenance auto
|
|
result, removed = raiw.remove_visible(bgr_array) # array -> array
|
|
raiw.remove_visible("shot.png", "out.png", sensitivity="strict")
|
|
raiw.visible_provenance("in.png") # -> frozenset({"gemini"})
|
|
|
|
Imports stay lazy (inside the functions), so ``import remove_ai_watermarks`` is cheap.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import suppress
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any, Literal
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable
|
|
|
|
from numpy.typing import NDArray
|
|
|
|
from remove_ai_watermarks.watermark_registry import Backend, Sensitivity
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _VisibleInput:
|
|
"""Normalized visible-removal input with its file-only context."""
|
|
|
|
bgr: NDArray[Any]
|
|
alpha: NDArray[Any] | None = None
|
|
path: Path | None = None
|
|
provenance: frozenset[str] = frozenset()
|
|
|
|
|
|
def visible_provenance(source: str | Path) -> frozenset[str]:
|
|
"""Vendor keys that the file's local metadata confirms, the evidence that drives
|
|
the ``auto`` sensitivity (relaxing a corroborated mark's detection trust gate).
|
|
|
|
Mapping: a Google/Gemini C2PA issuer -> ``"gemini"``; a ``samsung_genai`` marker ->
|
|
``"samsung"``; a China-AIGC (TC260) label -> the vendor its ``ContentProducer``
|
|
names (``KnownMark.tc260_producer_codes``), falling back to ByteDance's two products
|
|
when the producer is absent or unmapped.
|
|
|
|
The TC260 label itself says only "this is AI", so it used to relax Doubao and
|
|
Jimeng on EVERY China-AIGC image -- including one carrying a Qwen or Kling mark,
|
|
where relaxing the wrong pair is pure false-positive risk and the mark actually
|
|
present never reached the relaxed gate its own ``provenance_ncc_factor`` was
|
|
calibrated for. The producer code identifies the signing entity, so it can.
|
|
|
|
Best-effort: any read error yields an empty set (no relaxation). Metadata-only, so
|
|
it never loads cv2/torch.
|
|
"""
|
|
import contextlib
|
|
|
|
path = Path(source)
|
|
with contextlib.suppress(Exception):
|
|
from remove_ai_watermarks import identify
|
|
|
|
rep = identify.identify(path, check_visible=False, check_invisible=False)
|
|
return _provenance_from_report(rep, path)
|
|
return frozenset()
|
|
|
|
|
|
def _tc260_vendors(path: Path) -> frozenset[str]:
|
|
"""Vendor keys a TC260 label confirms, from its ``ContentProducer`` identity.
|
|
|
|
An absent, unreadable or unmapped producer falls back to the historical pair rather
|
|
than to nothing: the caller has already established that the AIGC signal fired, so
|
|
the image IS China-AIGC labelled, and dropping to no relaxation would lose the
|
|
detections the fallback recovers today. The re-read is deliberately isolated -- a
|
|
failure here must narrow the answer, never discard the rest of the provenance.
|
|
"""
|
|
import contextlib
|
|
|
|
from remove_ai_watermarks._internal.constants import TC260_FALLBACK_VENDORS
|
|
|
|
with contextlib.suppress(Exception):
|
|
from remove_ai_watermarks.metadata import aigc_label, uscc_of
|
|
from remove_ai_watermarks.watermark_registry import tc260_producer_vendors
|
|
|
|
producer = (aigc_label(path) or {}).get("ContentProducer", "")
|
|
if producer and (vendor := tc260_producer_vendors().get(uscc_of(producer))):
|
|
return frozenset({vendor})
|
|
return TC260_FALLBACK_VENDORS
|
|
|
|
|
|
def _load_visible_input(source: str | Path | NDArray[Any]) -> _VisibleInput:
|
|
"""Normalize a path/array source without making the public operation stateful."""
|
|
if not isinstance(source, (str, Path)):
|
|
return _VisibleInput(source)
|
|
|
|
from remove_ai_watermarks import image_io
|
|
|
|
path = Path(source)
|
|
bgr, alpha = image_io.read_bgr_and_alpha(path)
|
|
if bgr is None:
|
|
raise ValueError(f"Could not read image: {source}")
|
|
return _VisibleInput(bgr=bgr, alpha=alpha, path=path, provenance=visible_provenance(path))
|
|
|
|
|
|
def _write_visible_result(
|
|
loaded: _VisibleInput,
|
|
result: NDArray[Any],
|
|
removed: list[str],
|
|
output: str | Path,
|
|
*,
|
|
strip_metadata: bool,
|
|
write_noop: bool,
|
|
) -> None:
|
|
"""Write one visible-removal result while preserving a true no-op losslessly."""
|
|
if not removed and not write_noop:
|
|
return
|
|
|
|
from remove_ai_watermarks import image_io
|
|
|
|
out_path = Path(output)
|
|
out_path.parent.mkdir(parents=True, exist_ok=True)
|
|
source_path = loaded.path
|
|
if not removed and source_path is not None and source_path.suffix.lower() == out_path.suffix.lower():
|
|
# Copy the ORIGINAL bytes instead of lossily re-encoding a no-op. An in-place
|
|
# call needs no copy and would otherwise raise shutil.SameFileError.
|
|
if source_path.resolve() != out_path.resolve():
|
|
import shutil
|
|
|
|
shutil.copyfile(source_path, out_path)
|
|
else:
|
|
# imwrite is contractually NON-RAISING, so this bool is the only signal the file
|
|
# was not created. Unchecked, the metadata strip below ran on a nonexistent path
|
|
# and surfaced as a confusing "cannot read image <INPUT>" naming the OUTPUT path
|
|
# (Tier E, 2026-07-20). Raise here so a library caller and the CLI both get an
|
|
# accurate message about the write.
|
|
if not image_io.write_bgr_with_alpha(out_path, result, loaded.alpha):
|
|
raise OSError(f"failed to write output (is the destination writable?): {out_path}")
|
|
|
|
if strip_metadata:
|
|
from remove_ai_watermarks import metadata
|
|
|
|
metadata.remove_ai_metadata(out_path, out_path)
|
|
|
|
|
|
def remove_visible(
|
|
source: str | Path | NDArray[Any],
|
|
output: str | Path | None = None,
|
|
*,
|
|
sensitivity: Sensitivity = "auto",
|
|
backend: Backend = "auto",
|
|
strip_metadata: bool = True,
|
|
write_noop: bool = True,
|
|
) -> tuple[NDArray[Any], list[str]]:
|
|
"""Remove every detected known visible AI mark through localize then fill.
|
|
|
|
The registry currently covers the Gemini sparkle; Doubao, Jimeng, Qwen, Kling,
|
|
Yuanbao, Samsung, RunningHub, Baidu, and LibLibAI text marks; and the Jimeng
|
|
pill. Returns ``(result_bgr, [labels removed])``.
|
|
|
|
``source`` is a file path OR a BGR ndarray. For a PATH, metadata provenance is read
|
|
automatically (so ``sensitivity="auto"`` recovers a moved/faint mark whenever the
|
|
file still carries its provenance) and the alpha channel is preserved on write; for
|
|
an ARRAY there is no metadata to read and no separate alpha plane. When ``output``
|
|
is given the cleaned image is written there (alpha rejoined for a path source); the
|
|
array is always returned as well, so an empty ``removed`` list tells a caller nothing
|
|
known was found (e.g. route to the diffusion ``all`` path or ``erase``).
|
|
|
|
``sensitivity`` (``auto``/``strict``) and ``backend``
|
|
(``auto``/``cv2``/``migan``/``lama``) are the same controls as the CLI.
|
|
|
|
``strip_metadata`` (default True, matching the CLI ``visible --strip-metadata``)
|
|
also strips AI provenance metadata (C2PA/EXIF/XMP/IPTC) from the written output via
|
|
the lossless :func:`metadata.remove_ai_metadata`, so a library call does exactly
|
|
what the CLI does. Only applies when ``output`` is given.
|
|
|
|
``write_noop`` (default True) controls whether ``output`` is written when NOTHING was
|
|
removed: True writes a clean passthrough copy (an idempotent clean); False leaves the
|
|
output path untouched, so a caller that treats "no mark" as "produce nothing" (the CLI
|
|
``visible`` no-mark contract) does not clobber a pre-existing file at that path.
|
|
"""
|
|
from remove_ai_watermarks import watermark_registry
|
|
|
|
# Reject a removed sensitivity loudly; `Sensitivity` is a Literal and not enforced
|
|
# at runtime, so a 0.15 caller would otherwise get `auto` behaviour in silence.
|
|
watermark_registry.validate_sensitivity(sensitivity)
|
|
loaded = _load_visible_input(source)
|
|
result, removed = watermark_registry.remove_auto_marks(
|
|
loaded.bgr,
|
|
sensitivity=sensitivity,
|
|
provenance=loaded.provenance,
|
|
backend=backend,
|
|
)
|
|
if output is not None:
|
|
_write_visible_result(
|
|
loaded,
|
|
result,
|
|
removed,
|
|
output,
|
|
strip_metadata=strip_metadata,
|
|
write_noop=write_noop,
|
|
)
|
|
return result, removed
|
|
|
|
|
|
# ── The three-stage image pipeline (visible -> invisible -> metadata) ──
|
|
# This is the `all` / `batch` pipeline. It lived only in cli.py, written twice with
|
|
# divergences, so a library caller could not run the flagship path at all. The CLI is
|
|
# now a thin wrapper: it builds the options, prints the stage lines through
|
|
# `progress`, and turns the outcome into console text and an exit code.
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class InvisibleOptions:
|
|
"""The invisible stage's knobs, as one value instead of a dozen parameters.
|
|
|
|
ENGINE KNOBS ONLY, under the engine's own names and defaults, so a bare
|
|
``InvisibleOptions()`` behaves exactly like calling the engine with no arguments.
|
|
The engine takes them across TWO callables -- ``__init__`` for the ones that shape
|
|
the loaded stack, ``remove_watermark`` for the per-image ones -- so this is not a
|
|
splat-through bag; ``_run_invisible`` forwards each field to the right one.
|
|
``TestInvisibleOptionsMirrorTheEngine`` compares the signatures field by field with
|
|
no exception table to maintain: a decision made before the engine runs, like
|
|
``force``, is a parameter of ``remove_all`` next to ``backend`` and ``sensitivity``,
|
|
not a knob smuggled in here.
|
|
|
|
Immutable so a batch can build it once and reuse it across every image while the
|
|
engine itself is cached separately.
|
|
"""
|
|
|
|
strength: float | None = None
|
|
pipeline: str = "qwen-zimage"
|
|
seed: int | None = None
|
|
hf_token: str | None = None
|
|
humanize: float = 0.0
|
|
unsharp: float = 0.0
|
|
adaptive_polish: bool | None = None
|
|
max_resolution: int = 0
|
|
controlnet_conditioning_scale: float = 1.0
|
|
cpu_offload: bool = False
|
|
tile: bool = False
|
|
tile_size: int = 1024
|
|
tile_overlap: int = 128
|
|
|
|
|
|
# What the invisible stage did. "unavailable" is the one outcome the caller must
|
|
# surface loudly: the file looks processed but still carries the watermark.
|
|
InvisibleOutcome = Literal["removed", "no-signal", "unavailable"]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RemoveAllResult:
|
|
"""Outcome of :func:`remove_all` -- what each stage actually did."""
|
|
|
|
output: Path
|
|
visible_label: str | None # the removed mark(s), or None when nothing fired
|
|
invisible: InvisibleOutcome
|
|
|
|
|
|
class MetadataStripIncomplete(RuntimeError):
|
|
"""AI metadata survived the strip, so no output was produced.
|
|
|
|
Raised BEFORE the final write, deliberately: the contract is that an incomplete
|
|
strip leaves nothing on disk rather than an AI-readable file the caller might ship.
|
|
"""
|
|
|
|
def __init__(self, surviving: set[str]) -> None:
|
|
self.surviving = surviving
|
|
super().__init__(f"AI metadata survived the strip: {', '.join(sorted(surviving))}")
|
|
|
|
|
|
class _SourceEvidence:
|
|
"""One metadata extraction per source file, serving every stage of one call.
|
|
|
|
``remove_all`` asks the same file two provenance questions -- which vendor the
|
|
metadata confirms (for the visible pass) and whether an invisible target exists
|
|
(for the scrub gate). Both start from the same file-backed extraction, and running
|
|
them independently paid for it twice.
|
|
|
|
Per-CALL only, never module-level: the pipeline REWRITES its input in place on some
|
|
paths (``batch`` with the output directory equal to the input), so a holder that
|
|
outlived one call would answer from pre-write evidence. The individual metadata
|
|
probes are memoized on content, which is what makes that safe -- this holder only
|
|
removes the remaining assembly work.
|
|
|
|
Every accessor fails safe exactly as the function it replaces does: an extraction or
|
|
verdict error yields no provenance (no relaxation) and an invisible target of True
|
|
(scrub rather than skip).
|
|
"""
|
|
|
|
__slots__ = ("_evidence", "_extracted", "_path")
|
|
|
|
def __init__(self, path: Path) -> None:
|
|
self._path = path
|
|
self._extracted = False
|
|
self._evidence: Any | None = None
|
|
|
|
def _extract(self) -> Any | None:
|
|
if not self._extracted:
|
|
self._extracted = True
|
|
with suppress(Exception):
|
|
from remove_ai_watermarks.identify import extract_provenance_evidence
|
|
|
|
self._evidence = extract_provenance_evidence(self._path)
|
|
return self._evidence
|
|
|
|
def visible_provenance(self) -> frozenset[str]:
|
|
"""Vendor keys the metadata confirms; empty on any failure (no relaxation)."""
|
|
evidence = self._extract()
|
|
if evidence is None:
|
|
return frozenset()
|
|
# The suppress spans the VERDICT and the mapping, not just the extraction --
|
|
# `visible_provenance` fails safe as a whole, and a raise here would escape as a
|
|
# traceback where the old code returned an empty set.
|
|
with suppress(Exception):
|
|
from remove_ai_watermarks.identify import identify_from_evidence
|
|
|
|
report = identify_from_evidence(evidence)
|
|
return _provenance_from_report(report, self._path)
|
|
return frozenset()
|
|
|
|
def has_invisible_target(self) -> bool:
|
|
"""Whether a diffusion target exists. True on any failure -- see the fail-safe
|
|
note on ``identify.has_invisible_target``: leaving a watermark on a paid removal
|
|
is worse than over-regenerating a clean image."""
|
|
evidence = self._extract()
|
|
if evidence is None:
|
|
return True
|
|
try:
|
|
from remove_ai_watermarks.identify import identify_from_evidence
|
|
|
|
report = identify_from_evidence(evidence, image_path=self._path, check_invisible=True)
|
|
except Exception:
|
|
return True
|
|
return bool(report.ai_from_metadata)
|
|
|
|
|
|
def _provenance_from_report(report: Any, path: Path) -> frozenset[str]:
|
|
"""Map an already-built report to the vendor keys it confirms.
|
|
|
|
Shared by :func:`visible_provenance` and the evidence holder so the two cannot
|
|
drift; which evidence confirms which mark is data on the registry row.
|
|
"""
|
|
from remove_ai_watermarks.watermark_registry import known_marks
|
|
|
|
signal_names = {signal.name for signal in report.signals}
|
|
platform = (report.platform or "").lower()
|
|
keys: set[str] = set()
|
|
aigc = False
|
|
for mark in known_marks():
|
|
if any(token in platform for token in mark.provenance_platform_tokens):
|
|
keys.add(mark.key)
|
|
if "aigc" in mark.provenance_signals and "aigc" in signal_names:
|
|
aigc = True
|
|
elif any(name in signal_names for name in mark.provenance_signals):
|
|
keys.add(mark.key)
|
|
if aigc:
|
|
# The TC260 label is vendor-agnostic: which of its marks it confirms comes from
|
|
# the producer identity, not from the signal firing.
|
|
keys |= _tc260_vendors(path)
|
|
return frozenset(keys)
|
|
|
|
|
|
def remove_all(
|
|
source: str | Path,
|
|
output: str | Path,
|
|
*,
|
|
backend: Backend = "auto",
|
|
sensitivity: Sensitivity = "auto",
|
|
invisible: InvisibleOptions | None = None,
|
|
force: bool = False,
|
|
engine: Any | None = None,
|
|
progress: Callable[[str, str], None] | None = None,
|
|
) -> RemoveAllResult:
|
|
"""Remove visible marks, the invisible watermark, and AI metadata, in that order.
|
|
|
|
Stages are chained through a file in the SYSTEM temp dir, not next to ``output``:
|
|
the point of staging is that the user never sees a partial output file during a long
|
|
model download, and writing the partial next to the final defeats that.
|
|
|
|
``force`` scrubs even when no invisible watermark is locally detectable. It sits
|
|
here rather than in ``InvisibleOptions`` because it decides WHETHER the engine runs,
|
|
which is settled before the engine is built; the options carry only what the engine
|
|
itself takes.
|
|
|
|
``engine`` accepts an already-constructed ``InvisibleEngine`` so a batch can build
|
|
the model once; leave it None to construct one per call.
|
|
|
|
``progress`` receives ``(stage, detail)`` per step -- ``stage`` is one of
|
|
``visible`` / ``invisible`` / ``metadata`` and ``detail`` is a stable token, not
|
|
prose: the caller owns the wording. The invisible stage reports its
|
|
:data:`InvisibleOutcome`, plus a ``strength=<value>`` line before it runs.
|
|
|
|
Raises :class:`MetadataStripIncomplete` before writing anything when AI metadata
|
|
survives, and ``OSError`` when the output cannot be written.
|
|
"""
|
|
import os
|
|
import tempfile
|
|
|
|
from remove_ai_watermarks import image_io, watermark_registry
|
|
from remove_ai_watermarks.metadata import strip_and_verify
|
|
|
|
def say(stage: str, detail: str) -> None:
|
|
if progress is not None:
|
|
progress(stage, detail)
|
|
|
|
opts = invisible if invisible is not None else InvisibleOptions()
|
|
src, out = Path(source), Path(output)
|
|
watermark_registry.validate_sensitivity(sensitivity)
|
|
|
|
image, alpha = image_io.read_bgr_and_alpha(src)
|
|
if image is None:
|
|
raise ValueError(f"Could not read image: {src}")
|
|
|
|
# One metadata extraction for the whole pipeline: the visible pass asks which
|
|
# vendor is confirmed and the scrub gate asks whether a target exists, and both
|
|
# start from the same evidence.
|
|
evidence = _SourceEvidence(src)
|
|
|
|
tmp_fd, tmp_name = tempfile.mkstemp(suffix=src.suffix)
|
|
os.close(tmp_fd)
|
|
staged = Path(tmp_name)
|
|
try:
|
|
# ── 1. Visible marks ──
|
|
result, removed = watermark_registry.remove_auto_marks(
|
|
image,
|
|
sensitivity=sensitivity,
|
|
provenance=evidence.visible_provenance(),
|
|
backend=backend,
|
|
)
|
|
visible_label = ", ".join(removed) if removed else None
|
|
say("visible", visible_label or "")
|
|
if not image_io.write_bgr_with_alpha(staged, result, alpha):
|
|
raise OSError(f"failed to write the staged intermediate: {staged}")
|
|
|
|
# ── 2. Invisible watermark ──
|
|
outcome = _run_invisible(src, staged, staged, opts, engine, say, evidence, force)
|
|
|
|
# ── 3. AI metadata ──
|
|
# Read the pristine ORIGINAL for provenance above and the STAGED file here:
|
|
# the visible pass has already dropped this file's C2PA.
|
|
_, leftover = strip_and_verify(staged, staged)
|
|
if leftover:
|
|
# Before the write, on purpose -- see MetadataStripIncomplete.
|
|
raise MetadataStripIncomplete(set(leftover))
|
|
say("metadata", "stripped")
|
|
|
|
# The invisible stage (and the cv2.IMREAD_COLOR paths under it) drops alpha, so
|
|
# re-attach the ORIGINAL alpha plane unchanged for transparent formats.
|
|
final_bgr, _ = image_io.read_bgr_and_alpha(staged)
|
|
if final_bgr is None:
|
|
raise OSError(f"failed to read back the staged intermediate: {staged}")
|
|
out.parent.mkdir(parents=True, exist_ok=True)
|
|
if not image_io.write_bgr_with_alpha(out, final_bgr, alpha):
|
|
raise OSError(f"failed to write output (is the destination writable?): {out}")
|
|
finally:
|
|
if staged.exists():
|
|
staged.unlink()
|
|
return RemoveAllResult(out, visible_label, outcome)
|
|
|
|
|
|
def _run_invisible(
|
|
vendor_source: Path,
|
|
in_path: Path,
|
|
out_path: Path,
|
|
opts: InvisibleOptions,
|
|
engine: Any | None,
|
|
say: Callable[[str, str], None],
|
|
evidence: _SourceEvidence,
|
|
force: bool,
|
|
) -> InvisibleOutcome:
|
|
"""Run, or deliberately skip, the diffusion scrub.
|
|
|
|
``vendor_source`` is the PRISTINE original: the staged/output file has already lost
|
|
its C2PA to the visible pass, so gating or resolving the vendor from it would always
|
|
read as "no signal" and "unknown vendor". ``in_path``/``out_path`` are what the
|
|
engine reads and writes (the same staged file for ``remove_all``, input->output for
|
|
an invisible-only batch).
|
|
"""
|
|
from remove_ai_watermarks.invisible_engine import is_available
|
|
|
|
if not is_available():
|
|
say("invisible", "unavailable")
|
|
return "unavailable"
|
|
if not (force or evidence.has_invisible_target()):
|
|
say("invisible", "no-signal")
|
|
return "no-signal"
|
|
|
|
from remove_ai_watermarks._internal.watermark_profiles import resolve_strength, vendor_for_strength
|
|
|
|
vendor = vendor_for_strength(vendor_source)
|
|
# Report the strength the engine will actually execute, resolved the same way it
|
|
# resolves it, so the reported value cannot drift from the executed one.
|
|
with suppress(Exception):
|
|
from PIL import Image
|
|
|
|
with Image.open(vendor_source) as probe:
|
|
say("invisible", f"strength={resolve_strength(opts.strength, vendor, opts.pipeline, size=probe.size)}")
|
|
|
|
if engine is None:
|
|
from remove_ai_watermarks.invisible_engine import InvisibleEngine
|
|
|
|
engine = InvisibleEngine(
|
|
pipeline=opts.pipeline,
|
|
hf_token=opts.hf_token,
|
|
progress_callback=lambda message: say("invisible", message),
|
|
controlnet_conditioning_scale=opts.controlnet_conditioning_scale,
|
|
cpu_offload=opts.cpu_offload,
|
|
)
|
|
engine.remove_watermark(
|
|
image_path=in_path,
|
|
output_path=out_path,
|
|
strength=opts.strength,
|
|
seed=opts.seed,
|
|
humanize=opts.humanize,
|
|
unsharp=opts.unsharp,
|
|
adaptive_polish=opts.adaptive_polish,
|
|
max_resolution=opts.max_resolution,
|
|
vendor=vendor,
|
|
tile=opts.tile,
|
|
tile_size=opts.tile_size,
|
|
tile_overlap=opts.tile_overlap,
|
|
)
|
|
say("invisible", "removed")
|
|
return "removed"
|
|
|
|
|
|
BatchMode = Literal["all", "visible", "metadata", "invisible"]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BatchSummary:
|
|
"""Per-directory outcome of :func:`remove_batch`."""
|
|
|
|
processed: int
|
|
failed: int
|
|
# Files whose invisible watermark was left in place because the GPU extra is
|
|
# missing. Non-empty means the outputs LOOK processed but still carry it.
|
|
invisible_unavailable: list[Path]
|
|
errors: list[tuple[Path, str]]
|
|
|
|
|
|
def remove_batch(
|
|
directory: str | Path,
|
|
output_dir: str | Path,
|
|
*,
|
|
mode: BatchMode = "all",
|
|
backend: Backend = "auto",
|
|
sensitivity: Sensitivity = "auto",
|
|
invisible: InvisibleOptions | None = None,
|
|
force: bool = False,
|
|
engine: Any | None = None,
|
|
progress: Callable[[Path, str, str], None] | None = None,
|
|
) -> BatchSummary:
|
|
"""Run one removal ``mode`` over every supported image in ``directory``.
|
|
|
|
Never raises for a single bad image: a per-file failure is counted and recorded in
|
|
``BatchSummary.errors`` so one unreadable file cannot abandon the rest of the
|
|
directory. ``force`` and ``engine`` are threaded straight through, so a caller that
|
|
passes a constructed ``InvisibleEngine`` loads the model once for the whole run.
|
|
|
|
``progress`` receives ``(path, stage, detail)``. Every image ends with exactly one
|
|
terminal stage -- ``done`` or ``failed`` -- whatever the mode does in between, so a
|
|
caller driving a progress bar can advance on that alone.
|
|
"""
|
|
from remove_ai_watermarks._internal.utils import is_supported_format
|
|
|
|
src_dir, out_dir = Path(directory), Path(output_dir)
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
def say(path: Path, stage: str, detail: str) -> None:
|
|
if progress is not None:
|
|
progress(path, stage, detail)
|
|
|
|
processed = failed = 0
|
|
unavailable: list[Path] = []
|
|
errors: list[tuple[Path, str]] = []
|
|
for img_path in sorted(p for p in src_dir.iterdir() if is_supported_format(p)):
|
|
out_path = out_dir / img_path.name
|
|
try:
|
|
outcome = _run_batch_one(img_path, out_path, mode, backend, sensitivity, invisible, force, engine, say)
|
|
except Exception as exc:
|
|
failed += 1
|
|
errors.append((img_path, str(exc)))
|
|
say(img_path, "failed", str(exc))
|
|
continue
|
|
processed += 1
|
|
if outcome == "unavailable":
|
|
unavailable.append(img_path)
|
|
# Exactly one terminal event per image, in every mode: a caller's progress bar
|
|
# advances on this and nothing else. Keying it off a mode-specific stage line
|
|
# left `visible` and `metadata` runs sitting at 0% for the whole batch.
|
|
say(img_path, "done", outcome or "")
|
|
return BatchSummary(processed, failed, unavailable, errors)
|
|
|
|
|
|
def _run_batch_one(
|
|
img_path: Path,
|
|
out_path: Path,
|
|
mode: BatchMode,
|
|
backend: Backend,
|
|
sensitivity: Sensitivity,
|
|
invisible: InvisibleOptions | None,
|
|
force: bool,
|
|
engine: Any | None,
|
|
say: Callable[[Path, str, str], None],
|
|
) -> InvisibleOutcome | None:
|
|
"""One image, one mode. Returns the invisible outcome when that stage ran."""
|
|
from remove_ai_watermarks import image_io
|
|
from remove_ai_watermarks.metadata import strip_and_verify
|
|
|
|
if mode == "all":
|
|
result = remove_all(
|
|
img_path,
|
|
out_path,
|
|
backend=backend,
|
|
sensitivity=sensitivity,
|
|
invisible=invisible,
|
|
force=force,
|
|
engine=engine,
|
|
progress=lambda stage, detail: say(img_path, stage, detail),
|
|
)
|
|
return result.invisible
|
|
|
|
# NOTE on holders in batch: `out_path` may EQUAL `img_path` (nothing forbids
|
|
# `-o <the input directory>`), and the visible stage rewrites it. So each stage that
|
|
# asks the file a provenance question builds its OWN holder, after any write that
|
|
# precedes it. One holder spanning the write would answer the invisible gate from
|
|
# pre-write evidence and scrub a file that today is correctly skipped.
|
|
if mode == "visible":
|
|
# Deliberately NOT `remove_visible`: its no-op branch copies the original bytes
|
|
# through when nothing was removed, which is right for a single lossless call
|
|
# and wrong here. A batch must produce every output through the one writer, so a
|
|
# failed write RAISES and the run is counted and exits non-zero -- a read-only
|
|
# output directory once produced zero files and still exited 0 (Tier E).
|
|
# Always read the ORIGINAL: a stale out_path from a previous run must not be
|
|
# re-processed as if it were the input.
|
|
from remove_ai_watermarks import watermark_registry
|
|
|
|
image, alpha = image_io.read_bgr_and_alpha(img_path)
|
|
if image is None:
|
|
raise ValueError(f"Could not read image: {img_path}")
|
|
result, _ = watermark_registry.remove_auto_marks(
|
|
image,
|
|
sensitivity=sensitivity,
|
|
provenance=visible_provenance(img_path),
|
|
backend=backend,
|
|
)
|
|
if not image_io.write_bgr_with_alpha(out_path, result, alpha):
|
|
raise OSError(f"failed to write output (is the destination writable?): {out_path}")
|
|
return None
|
|
|
|
if mode == "metadata":
|
|
_, leftover = strip_and_verify(img_path, out_path)
|
|
if leftover:
|
|
raise MetadataStripIncomplete(set(leftover))
|
|
return None
|
|
|
|
# invisible-only: no preceding visible pass, so out_path does not exist yet, and
|
|
# the input IS the pristine original for both the gate and the vendor probe.
|
|
outcome = _run_invisible(
|
|
img_path,
|
|
img_path,
|
|
out_path,
|
|
invisible if invisible is not None else InvisibleOptions(),
|
|
engine,
|
|
lambda stage, detail: say(img_path, stage, detail),
|
|
_SourceEvidence(img_path),
|
|
force,
|
|
)
|
|
if not out_path.exists():
|
|
# Keep the output directory COMPLETE even when the pixels are deliberately
|
|
# left alone; a hole the caller cannot see is worse than an unchanged copy.
|
|
src_bgr, src_alpha = image_io.read_bgr_and_alpha(img_path)
|
|
if src_bgr is None or not image_io.write_bgr_with_alpha(out_path, src_bgr, src_alpha):
|
|
raise OSError(f"failed to copy input through to output: {out_path}")
|
|
return outcome
|