Files
remove-ai-watermarks/tests/test_cli.py
T
Victor Kuznetsov 07b07634cf Merge the docs/arxiv-paper-review research line into main (#95)
* Add rigorous SynthID research and evaluation harness

* Record SynthID cross-epoch validation results

* Reset SynthID detector hypotheses to paired training

* Add paired wavelet and spectral SynthID probes

* Record closed-corpus SynthID detector ablations

* Add registered SynthID phase probing

* Add periodic SynthID tile probing

* Add calibrated SynthID pixel detector

* Harden scale-registered SynthID detection

* Record SynthID phase and channel rescue findings

* Add production SynthID routing and OpenAI verification

* Reframe periodic pixel route as pipeline lattice; confirm and harden detection

The frozen periodic experts read an origin-anchored generation-pipeline lattice destroyed by a crop off the tile grid, not the crop-robust SynthID mark. Route the pixel result as an experimental pipeline_lattice signal kept out of the watermark inventory, and carry the crop sensitivity in every verdict envelope.

Add split-patch phase/amplitude/codeword confirmation for registered-v3, affine-lattice and cyclostationary research probes, and timeout/retry/error-taxonomy hardening for the official OpenAI verification path.

* Record pipeline-lattice reframing and registered-v3 phase lock

Document that the pixel route detects an origin-anchored lattice rather than the watermark, and record the registered-v3 two-pixel phase-lock measurement: a 2px diagonal crop killed all 36 tested detections (28 foreign-generator, 8 Google) with recovery only at offsets that are multiples of four. Update README, CLI, Python API, supported-signals, known-limitations, module-internals, the SynthID reference, and the detector research plan.

* Split SynthID campaign docs by detector, classifiers, and removal

Keep Model 1 as a research AI-versus-camera result outside the library.
Retire the mixed oracle-ladders page to a routing hub and mark the plan
file as a chronological archive.

* Keep 2026-08-22 SynthID campaign results in research docs only

Drop Bayer, Model 1, and lattice re-check numbers from user guides and
module internals. The library surface is unchanged.

* Unplug lattice detector from package, expand research docs

* Remove lattice detector sources from the public package

* Document SynthID GitHub sweep findings

* Document general AI classifier sweep

* Document modern-negative taxonomy audit and classifier doc split

Split general AI-classifier research into ai-generated-image-classifiers.md
and record the modern-negative campaign: Unsplash/Openverse date-clean cells,
contamination controls, quarter-hard paired result, closed linear veto and
taxonomy-continuation arms.

pre-commit: 1) maintain.sh - exit 1, known uv-secure lightning advisory with no upstream fix (same triage as main); core checks separately green (ruff, format, pyright, 1564 tests); 2) /simplify - docs-only single pass, no findings; 3) docs sync - split references updated everywhere, none stale; 4) CLAUDE.md - routing line already updated, compliant

* Record wild-AI audit, Meta provider class, and stock-negative expansion

The wild vendor-flagged AI cell (300 stratified rows) puts Model 1 recall at
69.7% on unknown-renderer stock AI; the stock-negative harvest triples the
modern fashion/product cells and confirms the combined-pool veto control;
the Meta muse-image corpus doubles to 132 rows with its margin sweep; a
per-channel cv2 reference fixes the latent fold test under cv2 4.10.0.

pre-commit: 1) maintain.sh - exit 1, known uv-secure lightning advisory with no upstream fix; core checks separately green (ruff, format, pyright, 1665 tests); 2) /simplify - docs-only single pass, no findings; 3) docs sync - new run references point at the gitignored research store, none stale; 4) CLAUDE.md - compact, no changes needed

* Close representation-change campaign: both candidates falsified

Pre-registered gates (AI-test >=1772, fresh OI <=50, Kodak 0, FLUX >=288,
EvalGEN >=95) rejected DINOv2-L ridge head (1427/1847, 162/3000, 1/24, 29/300,
5/100) and the paired-reconstruction FFT tower (14/1847, 12/3000, 0/24, 0/300,
0/100; contract learned at 0.949 pairwise separation but no transfer to
production generators). Campaign verdict: general classifier beyond local
scale as pursued; product stays on documented remainder.

pre-commit: 1) maintain.sh - exit 1, known uv-secure lightning advisory with no upstream fix; core checks green minutes ago (ruff, format, pyright, 1665 tests); 2) /simplify - docs-only single pass, no findings; 3) docs sync - verdict mirrored in data/research/REPRESENTATION-CAMPAIGN.md, no stale refs; 4) CLAUDE.md - no changes

* Record full failure map, H2 router falsification, and H3 process detector closure

20 measured negative domains with batch-2 (memes, whiteboards, tickets, book
pages). H2 domain router falsified: 97.9% of AI-test looks structured, every
family 94-100%. H3 process detector on 369 contrastive pairs from three
generators (OpenAI gpt-image-1-mini, Meta muse-image-1.0, Gemini flash-image)
achieves 0.909 pairwise AUC on training pairs but transfers at AUC 0.55-0.60
to production AI; human receipts score higher than any AI family. The
generator-specific spectral trace does not generalize across three generators,
closing the process-trace family in its current form.

pre-commit: 1) maintain.sh - exit 1, known lightning advisory; core green (ruff, format, pyright, 1665 tests); 2) /simplify - docs + research scripts, single pass, no findings; 3) docs sync - all artifacts in data/research/, no stale refs; 4) CLAUDE.md - no changes

* Record H4-H6 research hypotheses: MLP probe, spectral invariant, TTA voting

H4 (MLP on frozen CLIP): 98.6% AI-test recall but 15.8% photo FPR — same
trade-off curve as quarter-hard. H5 (spectral invariant from CVPR 2025):
2.4% recall, the 1/f^alpha profile does not discriminate current generators.
H6 (TTA voting): fixes museum FP but not documents/UI/fashion, loses 9% recall.
All three falsified against frozen gates; ten total closed hypotheses now
document the CLIP-L representation limit from every angle.

Also records the provider cascade across all 20 domains (69.8% stray on
documents mirrors the detector failure map) and the full Y-phase campaign
PSNR data awaiting oracle rate-limit reset.

pre-commit: 1) maintain.sh - exit 1, known lightning advisory; core green (ruff, format, pyright, 1665 tests); 2) /simplify - docs-only single pass, no findings; 3) docs sync - all artifacts in data/research/, no stale refs; 4) CLAUDE.md - no changes

* Fix Y-phase frequency convention, record H4-H6 and campaign results

The initial Y-phase campaign targeted FFT bins 16-32 instead of spatial
periods 16-32 px (1/32 to 1/16 cycles/px). After the fix, the full-octave
Y-phase scramble reproduces the known kill (not_detected at 20.5 dB on s1
at native 1122x1402). Partial sub-bands are visually invisible (42.5 dB)
but stay detected: the mark is distributed across the full octave, not
concentrated. Point intervention is not viable within this frequency band.

H4 MLP probe: 98.6% recall, 15.8% FPR (same trade-off wall).
H5 spectral invariant: 2.4% recall (1/f^alpha does not discriminate).
H6 TTA voting: fixes museum FP but not documents/UI/fashion.

pre-commit: 1) maintain.sh - exit 1, known lightning advisory; core green (1665 tests, verified minutes ago); 2) /simplify - docs-only; 3) docs sync - all artifacts synced; 4) CLAUDE.md - no changes

* Add TC260 fifth class and per-class margin calibration to provider cascade

Five-head cascade (openai/google/tc260/meta/no_ai): openai 68.8%, google
76.6%, tc260 63.8%, meta 71.8% mean recall at uniform margin 0.50.
Per-class margins calibrated against 1.67% photo budget: openai 0.31,
google 0.29, tc260 0.39, meta 0.55. TC260 test recall 71.4% at calibrated
margins on 196 held-out images. Chinese-AIGC is a viable pixel class.

pre-commit: 1) maintain.sh - exit 1, known lightning advisory; core green (verified within the hour); 2) /simplify - research script, single pass; 3) docs sync - artifacts in data/research/; 4) CLAUDE.md - no changes

* Record universal detector v3-v8: diverse training shifts the curve

Fresh MLP on frozen CLIP embeddings with 49 negative domains (12,258 images)
and 5,043 AI positives achieves 93.0% AI-test recall with 39/49 domains at
zero false positives. Digital art FP dropped from 99.3% (Model 1) to 24.8%,
fashion from 62-71% to 12%, UI from 99.3% to 12%. Photo FPR 0.7%. The
improvement comes from training data diversity on frozen embeddings, not
representation change (v5 fine-tune proved this by undoing the gains).
Also adds .local-eval to ruff exclude so the gate covers only tracked files.

pre-commit: 1) maintain.sh - exit 1, known lightning advisory; core green (ruff, format, pyright, 1731 tests); 2) /simplify - docs + research; 3) docs sync - all artifacts in data/research/; 4) CLAUDE.md - no changes

* Add three-tier ensemble detector: DEFINITELY/POSSIBLY/LIKELY_HUMAN

Combines Model 1 and v8 into calibrated confidence tiers. DEFINITELY_AI
(both models agree) achieves 79.6% recall at 0.50% photo FPR with 46/49
domains at zero FP. The union reaches 93.3% recall at 1.80% photo FPR
with Kodak 0/24. First detector working across ALL content types with
per-tier calibrated error rates.

pre-commit: 1) maintain.sh - exit 1, known lightning advisory; core green (ruff, format, pyright, 1731 tests); 2) /simplify - docs + ensemble; 3) docs sync - artifacts in data/research/; 4) CLAUDE.md - no changes

* Add provider v2: diversity training eliminates document leakage

Same diversity-training insight applied to the provider cascade. Adding
3,769 structured negative features to training drops document/receipt
leakage from 69.8% to 5.1%, UI from 10% to 4.3%, fashion from 5-15% to
0-3%. AI recall improves: OpenAI 84% (was 75%), Google 88.8% (was 79%).
Photo leakage 1.4%. Worst regression: memes 47.9% (was 27%).

pre-commit: 1) maintain.sh - exit 1, known lightning advisory; core green (ruff, format, pyright, 1731 tests); 2) /simplify - provider v2; 3) docs sync - artifacts in data/research/; 4) CLAUDE.md - no changes

* Add universal v11: full taxonomy, 136 domains, 124 at zero FP

Master data plan mapped 148 visual domains. Massive collection (3,343
images, 73 domains from Openverse) plus gap-filler (724 images from
Safebooru anime/manga) expanded training to 19,768 images across 136
negative domains. Universal v11 (MLP-512 + Model 1 ensemble):
DEFINITELY 90.3% recall at 0.50% photo FPR, ANY 95.6% recall,
124/134 domains at exactly 0% FP, worst FP 13.8% (digital art).

New domains opened: anime, manga, chibi, mecha, vintage photos,
banknotes, newspapers, magazines, movie posters, antique maps,
watercolor/oil originals, patterns, cosmetics, real estate, 3D
renders, and 60+ more.

pre-commit: 1) maintain.sh - exit 1, known lightning advisory; core green (ruff, format, pyright, 1731 tests); 2) /simplify - docs + research; 3) docs sync - artifacts in data/research/; 4) CLAUDE.md - no changes

* Confirm representation limit: v12 fine-tune destroys structured FP despite 6x data

Fine-tuning CLIP tower on 19,518 images achieves 95.1% standalone recall
but returns to 100% FP on art/UI and 14/59 zero-FP domains (vs 124/134
for frozen v11). This is a representation property, not a data issue.
Frozen CLIP + MLP (v11) is the only working architecture. Also records
provider v3 trained on all negative domains.

pre-commit: 1) maintain.sh - exit 1, known lightning advisory; core green (ruff, format, pyright, 1731 tests); 2) /simplify - docs + research; 3) docs sync - artifacts in data/research/; 4) CLAUDE.md - no changes

* Document v13 plateau: 128/136 zero FP, frozen-CLIP ceiling confirmed

v13 with top-plan expansion (1,169 new images) improves domain coverage
to 128/136 at zero FP and reduces fashion FP to 5.3%, but DEFINITELY
recall plateaus at 89.6% and art FP at 15.9%. The frozen CLIP + MLP
architecture has reached its ceiling.

pre-commit: 1) maintain.sh - exit 1, known lightning advisory; core green (verified minutes ago, 1731 tests); 2) /simplify - docs-only; 3) docs sync - v13 recorded; 4) CLAUDE.md - no changes

* Document the 2026-08-31 photo-classifier freeze.

Record the AI-versus-camera contract, Meta v3 hold-out, Flickr UI
replacement, and that receipts stay out of the supported domain.

* Scrub Spaces catalog details from classifier research docs

* Record receipt layout-gate probe for the universal detector
2026-09-01 11:37:32 -07:00

1269 lines
58 KiB
Python

"""Tests for the CLI entry point."""
from __future__ import annotations
import json
from typing import TYPE_CHECKING
from unittest.mock import MagicMock, patch
if TYPE_CHECKING:
from pathlib import Path
import cv2
import numpy as np
import pytest
from click.testing import CliRunner
from PIL import Image
from PIL.PngImagePlugin import PngInfo
from remove_ai_watermarks.cli import main
@pytest.fixture
def runner():
return CliRunner()
@pytest.fixture
def sample_png(tmp_path: Path) -> Path:
"""Create a sample PNG for CLI testing."""
# Seeded: an unseeded random corner can occasionally trip the Doubao
# visible-mark detector, making `visible --mark auto` flaky.
img = np.random.default_rng(0).integers(0, 255, (200, 200, 3), dtype=np.uint8)
path = tmp_path / "input.png"
cv2.imwrite(str(path), img)
return path
def _make_batch_dir(tmp_path: Path, count: int = 3) -> Path:
"""Create a directory with test images for batch testing."""
input_dir = tmp_path / "input"
input_dir.mkdir()
rng = np.random.default_rng(0)
for i in range(count):
img = rng.integers(0, 255, (100, 100, 3), dtype=np.uint8)
cv2.imwrite(str(input_dir / f"img_{i}.png"), img)
return input_dir
def _make_batch_dir_with_metadata(tmp_path: Path, count: int = 3) -> Path:
"""Create a directory with PNG images containing AI metadata."""
input_dir = tmp_path / "input"
input_dir.mkdir()
for i in range(count):
img = Image.new("RGB", (64, 64), color=(100 + i, 150, 200))
pnginfo = PngInfo()
pnginfo.add_text("parameters", f"Steps: 20, Sampler: Euler, img_{i}")
pnginfo.add_text("prompt", "a test landscape")
img.save(input_dir / f"img_{i}.png", pnginfo=pnginfo)
return input_dir
def _mock_invisible_engine():
"""Create a mock InvisibleEngine that writes a copy of the input image."""
def _mock_remove_watermark(image_path, output_path=None, **kwargs):
out = output_path or image_path.with_stem(image_path.stem + "_clean")
out.parent.mkdir(parents=True, exist_ok=True)
img = Image.open(image_path)
img.save(out)
return out
mock_engine = MagicMock()
mock_engine.remove_watermark.side_effect = _mock_remove_watermark
mock_cls = MagicMock(return_value=mock_engine)
return mock_cls, mock_engine
def _mock_invisible_engine_drops_alpha():
"""Mock InvisibleEngine that mimics the real engine's BGR-only output path.
The real diffusion-based engine reads with cv2.IMREAD_COLOR and writes a
3-channel result. This mock simulates that so we can regression-test alpha
preservation across the ``all`` pipeline.
"""
def _mock_remove_watermark(image_path, output_path=None, **kwargs):
out = output_path or image_path.with_stem(image_path.stem + "_clean")
out.parent.mkdir(parents=True, exist_ok=True)
bgr = cv2.imread(str(image_path), cv2.IMREAD_COLOR)
cv2.imwrite(str(out), bgr)
return out
mock_engine = MagicMock()
mock_engine.remove_watermark.side_effect = _mock_remove_watermark
mock_cls = MagicMock(return_value=mock_engine)
return mock_cls, mock_engine
class TestMainGroup:
"""Tests for the top-level CLI group."""
def test_help(self, runner):
result = runner.invoke(main, ["--help"])
assert result.exit_code == 0
assert "Remove visible and invisible" in result.output
def test_version(self, runner):
result = runner.invoke(main, ["--version"])
assert result.exit_code == 0
assert "remove-ai-watermarks" in result.output
assert "version" in result.output
def test_no_command_shows_banner(self, runner):
result = runner.invoke(main, [])
assert result.exit_code == 0
assert "Remove-AI-Watermarks" in result.output
class TestVisibleCommand:
"""Tests for the 'visible' subcommand."""
def test_visible_help(self, runner):
result = runner.invoke(main, ["visible", "--help"])
assert result.exit_code == 0
assert "visible AI watermark" in result.output
assert "--mark" in result.output
def test_visible_auto_no_mark_exits_two_with_eraser_hint(self, runner, sample_png, tmp_path):
# No known visible mark and no AI provenance signal: the command must not
# re-serve the input as a finished result. It exits EXIT_NO_VISIBLE_MARK
# (2) -- distinct from success (0) and a hard error (1) -- writes no
# output file, and points the user at the region eraser.
output = tmp_path / "clean.png"
result = runner.invoke(main, ["visible", str(sample_png), "-o", str(output)])
assert result.exit_code == 2, result.output
assert not output.exists()
assert "erase" in result.output
# The "no signal" branch must NOT imply the image is clean: a missing
# metadata proxy is not proof an invisible pixel watermark (SynthID) is
# absent, so the message preserves that uncertainty and routes to 'all'.
assert "SynthID" in result.output
assert "all" in result.output
def test_visible_auto_no_mark_routes_to_all_when_metadata(self, runner, tmp_path):
# An image whose only signal is an invisible/metadata watermark (here SD
# generation parameters) has no visible mark to remove; the command must
# exit 2 and upsell the full 'all' pipeline rather than the eraser.
img = Image.fromarray(np.random.default_rng(0).integers(0, 255, (200, 200, 3), dtype=np.uint8))
pnginfo = PngInfo()
pnginfo.add_text("parameters", "Steps: 20, Sampler: Euler, a test landscape")
src = tmp_path / "ai.png"
img.save(src, pnginfo=pnginfo)
output = tmp_path / "clean.png"
result = runner.invoke(main, ["visible", str(src), "-o", str(output)])
assert result.exit_code == 2, result.output
assert not output.exists()
assert "all" in result.output
def test_visible_basic(self, runner, sample_png, tmp_path):
output = tmp_path / "clean.png"
result = runner.invoke(
main,
["visible", str(sample_png), "-o", str(output), "--no-detect"],
)
assert result.exit_code == 0
assert output.exists()
assert "Saved" in result.output
def test_visible_default_output_name(self, runner, sample_png):
result = runner.invoke(main, ["visible", str(sample_png), "--no-detect"])
assert result.exit_code == 0
expected = sample_png.with_stem(sample_png.stem + "_clean")
assert expected.exists()
def test_visible_backend_cv2(self, runner, sample_png, tmp_path):
# The old --inpaint/--no-inpaint flags are gone; the fill backend is picked
# with --backend (localize -> fill). cv2 needs no ONNX model.
output = tmp_path / "clean.png"
result = runner.invoke(
main,
[
"visible",
str(sample_png),
"-o",
str(output),
"--backend",
"cv2",
"--no-detect",
],
)
assert result.exit_code == 0
assert output.exists()
def test_visible_no_detect(self, runner, sample_png, tmp_path):
output = tmp_path / "clean.png"
result = runner.invoke(
main,
["visible", str(sample_png), "-o", str(output), "--no-detect"],
)
assert result.exit_code == 0
def test_visible_nonexistent_file(self, runner):
result = runner.invoke(main, ["visible", "/nonexistent/file.png"])
assert result.exit_code != 0
def test_visible_preserves_rgba_transparency(self, runner, tmp_path):
"""Visible removal on an RGBA PNG must keep the alpha channel,
not silently flatten the image onto an opaque background.
"""
rgba = np.zeros((200, 200, 4), dtype=np.uint8)
rgba[:, :, :3] = 200 # light grey foreground
rgba[50:150, 50:150, 3] = 255 # opaque square in the middle, rest transparent
src = tmp_path / "rgba_in.png"
cv2.imwrite(str(src), rgba)
output = tmp_path / "rgba_out.png"
result = runner.invoke(
main,
["visible", str(src), "-o", str(output), "--no-detect"],
)
assert result.exit_code == 0, result.output
assert output.exists()
out = cv2.imread(str(output), cv2.IMREAD_UNCHANGED)
assert out.ndim == 3, f"output is not 3D: shape={out.shape}"
assert out.shape[2] == 4, f"output is not RGBA: shape={out.shape}"
# The transparent corners must remain transparent.
assert out[0, 0, 3] == 0
assert out[199, 199, 3] == 0
# The opaque center remains opaque (the watermark region default is bottom-right,
# which doesn't overlap the center square at 200x200).
assert out[100, 100, 3] == 255
def test_visible_keeps_alpha_opaque_in_watermark_region(self, runner, tmp_path):
"""Regression for issue #30 (white box): on an opaque RGBA image, the
watermark region must stay OPAQUE. Reverse-alpha recovers real pixels
there, so zeroing alpha would punch a transparent hole that renders as a
solid white box on any non-transparent viewer.
"""
rgba = np.full((200, 200, 4), 255, dtype=np.uint8) # fully opaque white
src = tmp_path / "rgba_full.png"
cv2.imwrite(str(src), rgba)
output = tmp_path / "rgba_kept.png"
result = runner.invoke(
main,
["visible", str(src), "-o", str(output), "--no-detect"],
)
assert result.exit_code == 0, result.output
out = cv2.imread(str(output), cv2.IMREAD_UNCHANGED)
assert out.shape[2] == 4
# Default sparkle position is in the bottom-right; alpha there must stay 255.
from remove_ai_watermarks.gemini_engine import get_watermark_config
cfg = get_watermark_config(200, 200)
px, py = cfg.get_position(200, 200)
size = cfg.logo_size
assert out[py + size // 2, px + size // 2, 3] == 255, "watermark region alpha was zeroed (white-box regression)"
# No pixel anywhere should have been forced transparent.
assert int((out[:, :, 3] == 0).sum()) == 0, "spurious transparent pixels introduced"
def test_visible_rgb_input_stays_rgb(self, runner, sample_png, tmp_path):
"""Regression: a plain RGB PNG must NOT gain a spurious alpha channel."""
output = tmp_path / "rgb_out.png"
result = runner.invoke(
main,
["visible", str(sample_png), "-o", str(output), "--no-detect"],
)
assert result.exit_code == 0, result.output
out = cv2.imread(str(output), cv2.IMREAD_UNCHANGED)
assert out.ndim == 3, f"output is not 3D: shape={out.shape}"
assert out.shape[2] == 3, f"RGB input produced non-RGB output: shape={out.shape}"
class TestInvisibleCommand:
"""Tests for the 'invisible' subcommand."""
def test_invisible_help(self, runner):
result = runner.invoke(main, ["invisible", "--help"])
assert result.exit_code == 0
assert "invisible" in result.output.lower()
def test_invisible_basic(self, runner, sample_png, tmp_path):
mock_cls, mock_engine = _mock_invisible_engine()
output = tmp_path / "clean.png"
with (
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
):
result = runner.invoke(
main,
["invisible", str(sample_png), "-o", str(output), "--force"],
)
assert result.exit_code == 0, result.output
assert output.exists()
mock_engine.remove_watermark.assert_called_once()
def test_invisible_cpu_offload_flows_to_engine(self, runner, sample_png, tmp_path):
mock_cls, _mock_engine = _mock_invisible_engine()
output = tmp_path / "clean.png"
with (
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
):
result = runner.invoke(
main,
["invisible", str(sample_png), "-o", str(output), "--cpu-offload", "--force"],
)
assert result.exit_code == 0, result.output
assert mock_cls.call_args.kwargs["cpu_offload"] is True
def test_invisible_default_output(self, runner, sample_png):
mock_cls, _mock_engine = _mock_invisible_engine()
with (
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
):
result = runner.invoke(main, ["invisible", str(sample_png), "--force"])
assert result.exit_code == 0, result.output
expected = sample_png.with_stem(sample_png.stem + "_clean")
assert expected.exists()
def test_invisible_leaves_the_polish_default_to_the_library(self, runner, sample_png):
"""An untyped --adaptive-polish reaches the engine as None, not as a value.
The per-profile default lives in watermark_profiles, so the CLI must pass the
user's non-choice through rather than resolving it here. Resolving in the CLI
is how the library and the CLI came to disagree on the same profile.
"""
mock_cls, mock_engine = _mock_invisible_engine()
with (
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
):
result = runner.invoke(main, ["invisible", str(sample_png), "--force"])
assert result.exit_code == 0, result.output
assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is None
def test_invisible_no_adaptive_polish_disables(self, runner, sample_png):
mock_cls, mock_engine = _mock_invisible_engine()
with (
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
):
result = runner.invoke(main, ["invisible", str(sample_png), "--no-adaptive-polish", "--force"])
assert result.exit_code == 0, result.output
assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is False
def test_knobs_the_fixed_stack_cannot_honor_are_not_offered(self, runner, sample_png):
"""--model/--steps/--guidance-scale/--device/--auto are gone, not rejected.
Each pinned a value the profiles fix (model stack, per-stage schedule, CFG 1.0,
CUDA), so accepting one only produced an error several layers down -- a flag
that advertises a capability the library does not have. Click now refuses the
option itself, which is the honest answer and the one a caller can act on.
"""
retired = (
["--model", "org/custom-sdxl"],
["--steps", "20"],
["--guidance-scale", "5.5"],
["--device", "cpu"],
["--auto"],
)
# All THREE diffusion commands, not just `invisible`. Each used to declare these
# inline, so removing them from one and not the others is a live possibility.
for command in ("invisible", "all", "batch"):
target = str(sample_png.parent) if command == "batch" else str(sample_png)
for args in retired:
result = runner.invoke(main, [command, target, *args, "--force"])
assert result.exit_code == 2, f"{command} {args[0]}: {result.output}"
assert "No such option" in result.output, f"{command} {args[0]}: {result.output}"
def test_retired_pipeline_names_are_rejected_not_silently_remapped(self, runner, sample_png):
"""default/sdxl/controlnet/qwen were removed with their CPU code paths.
Click rejects them at parse time. Mapping them onward would run a profile the
caller never chose, at a different strength and a different quality.
"""
for retired in ("default", "sdxl", "controlnet", "qwen"):
result = runner.invoke(main, ["invisible", str(sample_png), "--pipeline", retired, "--force"])
assert result.exit_code == 2, result.output
assert "is not one of" in result.output
def test_invisible_nonexistent_file(self, runner):
result = runner.invoke(main, ["invisible", "/nonexistent/file.png"])
assert result.exit_code != 0
def test_invisible_no_signal_skips_and_exits_two(self, runner, sample_png, tmp_path):
"""P0#5: when no invisible AI watermark is locally detectable, the diffusion
scrub must NOT run (it would only degrade a clean image). Mirrors the visible
no-mark contract: write no output, exit 2, and DO NOT imply the image is
clean (stripped SynthID provenance is not proof of absence)."""
mock_cls, mock_engine = _mock_invisible_engine()
output = tmp_path / "clean.png"
with (
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
):
result = runner.invoke(main, ["invisible", str(sample_png), "-o", str(output)])
assert result.exit_code == 2, result.output
assert not output.exists()
mock_engine.remove_watermark.assert_not_called()
assert "--force" in result.output
assert "SynthID" in result.output # the message must preserve removal uncertainty
def test_invisible_force_runs_scrub_on_no_signal(self, runner, sample_png, tmp_path):
"""--force overrides the no-signal skip: the scrub runs regardless."""
mock_cls, mock_engine = _mock_invisible_engine()
output = tmp_path / "clean.png"
with (
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
):
result = runner.invoke(main, ["invisible", str(sample_png), "-o", str(output), "--force"])
assert result.exit_code == 0, result.output
mock_engine.remove_watermark.assert_called_once()
def test_invisible_explicit_vendor_implies_force_and_sets_cohort(self, runner, sample_png, tmp_path):
"""--vendor meta names a cohort the file cannot prove (Content Seal carries
no C2PA), so it must both bypass the no-signal skip and arrive at the engine
as the vendor, where it resolves to the measured Meta floor (not the
area-curve value the same size would otherwise get)."""
mock_cls, mock_engine = _mock_invisible_engine()
output = tmp_path / "clean.png"
with (
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
):
result = runner.invoke(main, ["invisible", str(sample_png), "-o", str(output), "--vendor", "meta"])
assert result.exit_code == 0, result.output
kwargs = mock_engine.remove_watermark.call_args.kwargs
assert kwargs["vendor"] == "meta"
assert "0.1" in result.output # the resolved Meta floor, printed
assert "0.094" not in result.output # not the sample's area-curve value (200x200 -> ~0.0944)
def test_invisible_vendor_auto_keeps_detection_semantics(self, runner, sample_png, tmp_path):
"""--vendor auto is the default spelled out: detection still runs and a
no-signal file still skips."""
mock_cls, mock_engine = _mock_invisible_engine()
output = tmp_path / "clean.png"
with (
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
):
result = runner.invoke(main, ["invisible", str(sample_png), "-o", str(output), "--vendor", "auto"])
assert result.exit_code == 2, result.output
mock_engine.remove_watermark.assert_not_called()
def test_invisible_runs_without_force_when_signal_present(self, runner, tmp_path):
"""An image carrying an AI metadata signal IS a scrub target, so the run
proceeds with no --force needed."""
img = Image.fromarray(np.random.default_rng(0).integers(0, 255, (200, 200, 3), dtype=np.uint8))
pnginfo = PngInfo()
pnginfo.add_text("parameters", "Steps: 20, Sampler: Euler, a test landscape")
src = tmp_path / "ai.png"
img.save(src, pnginfo=pnginfo)
output = tmp_path / "clean.png"
mock_cls, mock_engine = _mock_invisible_engine()
with (
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
):
result = runner.invoke(main, ["invisible", str(src), "-o", str(output)])
assert result.exit_code == 0, result.output
mock_engine.remove_watermark.assert_called_once()
class TestAllCommand:
"""Tests for the 'all' subcommand (full pipeline)."""
def test_all_help(self, runner):
result = runner.invoke(main, ["all", "--help"])
assert result.exit_code == 0
assert "visible" in result.output.lower()
def test_all_basic(self, runner, sample_png, tmp_path):
mock_cls, _mock_engine = _mock_invisible_engine()
output = tmp_path / "clean.png"
with (
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
):
result = runner.invoke(
main,
["all", str(sample_png), "-o", str(output), "--force"],
)
assert result.exit_code == 0, result.output
assert output.exists()
def test_all_cpu_offload_flows_to_engine(self, runner, sample_png, tmp_path):
mock_cls, _mock_engine = _mock_invisible_engine()
output = tmp_path / "clean.png"
with (
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
):
result = runner.invoke(
main,
["all", str(sample_png), "-o", str(output), "--cpu-offload", "--force"],
)
assert result.exit_code == 0, result.output
assert mock_cls.call_args.kwargs["cpu_offload"] is True
def test_all_controlnet_scale_flows_to_engine(self, runner, sample_png, tmp_path):
"""The click option is `--controlnet-scale`, the engine parameter is
`controlnet_conditioning_scale`, and `InvisibleOptions` now uses the engine's
spelling so the translation happens once. Pin the value, not just the name:
hardcoding the constant in `_run_invisible` passed the whole suite before."""
mock_cls, _mock_engine = _mock_invisible_engine()
output = tmp_path / "clean.png"
with (
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
):
result = runner.invoke(
main,
["all", str(sample_png), "-o", str(output), "--controlnet-scale", "0.65", "--force"],
)
assert result.exit_code == 0, result.output
assert mock_cls.call_args.kwargs["controlnet_conditioning_scale"] == pytest.approx(0.65)
def test_all_nonexistent_file(self, runner):
result = runner.invoke(main, ["all", "/nonexistent/file.png"])
assert result.exit_code != 0
def test_all_visible_step_uses_registry(self, runner, sample_png, tmp_path):
"""Regression (#1): the `all` visible step must route through the registry
(remove_auto_marks), so Doubao/Jimeng/Samsung/pill marks are handled -- not
just the Gemini sparkle via a hardcoded GeminiEngine."""
mock_cls, _mock_engine = _mock_invisible_engine()
output = tmp_path / "clean.png"
def _fake_remove_auto(image, **kwargs):
return image, []
with (
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
patch(
"remove_ai_watermarks.watermark_registry.remove_auto_marks", side_effect=_fake_remove_auto
) as mock_auto,
):
result = runner.invoke(main, ["all", str(sample_png), "-o", str(output), "--force"])
assert result.exit_code == 0, result.output
mock_auto.assert_called() # the registry auto-detector drove the visible pass
def test_all_skips_invisible_on_no_signal_but_succeeds(self, runner, sample_png, tmp_path):
"""P0#5: with no detectable invisible watermark and no --force, `all` skips
the destructive step 2 (pixels left intact) but STILL succeeds (exit 0) --
visible removal + metadata strip ran and a file is written. Distinct from the
GPU-missing skip, which is a non-zero failure."""
mock_cls, mock_engine = _mock_invisible_engine()
output = tmp_path / "clean.png"
with (
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
):
result = runner.invoke(main, ["all", str(sample_png), "-o", str(output)])
assert result.exit_code == 0, result.output
assert output.exists()
mock_engine.remove_watermark.assert_not_called()
assert "Skipped (no invisible" in result.output
def test_all_loud_warning_and_nonzero_exit_when_gpu_missing(self, runner, sample_png, tmp_path):
"""Regression (#14/#47): when the GPU extra is absent the invisible step is
skipped, but the output still looks processed -- the run must fail loudly
(prominent banner + non-zero exit) so a skipped SynthID pass is not mistaken
for a clean result. The output file is still written (visible + metadata)."""
output = tmp_path / "clean.png"
with patch("remove_ai_watermarks.invisible_engine.is_available", return_value=False):
result = runner.invoke(main, ["all", str(sample_png), "-o", str(output)])
assert result.exit_code != 0, result.output
assert "NOT removed" in result.output
assert "remove-ai-watermarks[qwen-zimage]" in result.output
assert output.exists() # visible + metadata still produced a file
def test_all_reports_metadata_that_survived_stripping(self, runner, sample_png, tmp_path):
"""The full pipeline must verify the metadata result before reporting success.
``remove_ai_metadata`` is deliberately fail-safe and may copy an undecodable
input through unchanged. Calling it directly let ``all`` print a successful
strip even when a marker survived.
"""
output = tmp_path / "clean.png"
with (
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
# The no-signal gate lives in the library now; patch it there, not on the
# CLI module, or the diffusion stage runs for real.
patch("remove_ai_watermarks.api._SourceEvidence.has_invisible_target", return_value=False),
patch(
"remove_ai_watermarks.metadata.strip_and_verify",
return_value=(tmp_path / "intermediate.png", {"c2pa": True}),
),
):
result = runner.invoke(main, ["all", str(sample_png), "-o", str(output)])
assert result.exit_code != 0
assert "metadata" in result.output.lower()
assert "survived" in result.output.lower()
assert "AI metadata stripped" not in result.output
# An incomplete strip must leave NOTHING on disk: an AI-readable output plus a
# non-zero exit is the failure mode the pre-write raise exists to prevent.
assert not output.exists()
def test_all_preserves_rgba_across_invisible_step(self, runner, tmp_path):
"""Regression: ``all`` must keep transparency even when the invisible
step writes a 3-channel result (as the real diffusion engine does).
"""
rgba = np.zeros((200, 200, 4), dtype=np.uint8)
rgba[:, :, :3] = 200
rgba[50:150, 50:150, 3] = 255 # opaque square; corners transparent
src = tmp_path / "rgba_in.png"
cv2.imwrite(str(src), rgba)
output = tmp_path / "rgba_out.png"
mock_cls, _engine = _mock_invisible_engine_drops_alpha()
with (
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
patch("remove_ai_watermarks.cli.invisible_available", return_value=True, create=True),
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
):
result = runner.invoke(main, ["all", str(src), "-o", str(output), "--force"])
assert result.exit_code == 0, result.output
out = cv2.imread(str(output), cv2.IMREAD_UNCHANGED)
assert out.ndim == 3, f"output not 3D: shape={out.shape}"
assert out.shape[2] == 4, f"output is not RGBA: shape={out.shape}"
assert out[0, 0, 3] == 0
assert out[100, 100, 3] == 255
class TestMetadataCommand:
"""Tests for the 'metadata' subcommand."""
def test_metadata_help(self, runner):
result = runner.invoke(main, ["metadata", "--help"])
assert result.exit_code == 0
def test_metadata_check_clean(self, runner, tmp_clean_png):
result = runner.invoke(main, ["metadata", str(tmp_clean_png), "--check"])
assert result.exit_code == 0
assert "No AI metadata" in result.output
assert "not the same as 'clean'" in result.output
def test_metadata_check_ai(self, runner, tmp_png_with_ai_metadata):
result = runner.invoke(main, ["metadata", str(tmp_png_with_ai_metadata), "--check"])
assert result.exit_code == 0
assert "AI metadata detected" in result.output
assert "not the same as 'clean'" not in result.output
def test_metadata_remove(self, runner, tmp_png_with_ai_metadata, tmp_path):
output = tmp_path / "stripped.png"
result = runner.invoke(
main,
[
"metadata",
str(tmp_png_with_ai_metadata),
"--remove",
"-o",
str(output),
],
)
assert result.exit_code == 0
assert "stripped" in result.output
assert "not the same as 'clean'" in result.output
def test_metadata_remove_reports_failure_when_the_strip_was_a_no_op(self, runner, tmp_path):
"""A file PIL cannot decode is copied through UNCHANGED by the fail-safe.
That is correct (never crash a worker on a partial upload) but the command used
to print "AI metadata stripped ->" and exit 0 for it, so a caller could not tell
a real strip from a no-op and the output still read as AI. A Samsung C2PA
compatibility case exposed the defect.
"""
# PNG signature + a C2PA (caBX) chunk, then garbage: the byte scanner sees the
# marker, PIL cannot decode it.
src = tmp_path / "undecodable.png"
payload = b"c2pa" + b"\x00" * 32
chunk = len(payload).to_bytes(4, "big") + b"caBX" + payload + b"\x00\x00\x00\x00"
src.write_bytes(b"\x89PNG\r\n\x1a\n" + chunk + b"NOTAPNG" * 8)
out = tmp_path / "cleaned.png"
result = runner.invoke(main, ["metadata", str(src), "--remove", "-o", str(out)])
from remove_ai_watermarks.metadata import get_ai_metadata
if not get_ai_metadata(src):
pytest.skip("fixture does not register as an AI-metadata carrier")
assert result.exit_code != 0, "a no-op strip must not report success"
assert "stripped ->" not in result.output
def test_metadata_remove_in_place(self, runner, tmp_png_with_ai_metadata):
"""With ``-o`` omitted, the strip overwrites the source in place (default
output_path=None). Previously every test passed an explicit ``-o``."""
from remove_ai_watermarks.metadata import has_ai_metadata
assert has_ai_metadata(tmp_png_with_ai_metadata) # precondition
result = runner.invoke(main, ["metadata", str(tmp_png_with_ai_metadata), "--remove"])
assert result.exit_code == 0, result.output
assert not has_ai_metadata(tmp_png_with_ai_metadata) # source overwritten, AI metadata gone
class TestIdentifyCommand:
"""Tests for the 'identify' subcommand."""
def test_identify_help(self, runner):
result = runner.invoke(main, ["identify", "--help"])
assert result.exit_code == 0
def test_identify_clean_png(self, runner, tmp_clean_png):
result = runner.invoke(main, ["identify", str(tmp_clean_png), "--no-visible"])
assert result.exit_code == 0
assert "unknown" in result.output
def test_identify_unknown_explains_why(self, runner, tmp_clean_png):
# An unknown verdict must explain itself inline (issue #22: users read a bare
# "unknown" as the tool being broken) rather than only in the caveats section.
result = runner.invoke(main, ["identify", str(tmp_clean_png), "--no-visible"])
assert result.exit_code == 0
assert "No locally-readable AI signal found" in result.output
assert "not the same as 'clean'" in result.output
def test_identify_ai_png_reports_platform(self, runner, tmp_png_with_ai_metadata):
result = runner.invoke(main, ["identify", str(tmp_png_with_ai_metadata), "--no-visible"])
assert result.exit_code == 0
assert "AI-generated" in result.output
assert "Stable Diffusion" in result.output
def test_identify_reports_generated_source_kind(self, runner):
"""The C2PA trainedAlgorithmicMedia source type sharpens the verdict to
'AI-generated (fully synthetic)' at the CLI (the ai_source_kind branch)."""
from pathlib import Path
sample = Path(__file__).resolve().parent.parent / "data" / "fixtures" / "provenance" / "chatgpt-1.png"
if not sample.exists():
pytest.skip("chatgpt sample not present")
result = runner.invoke(main, ["identify", str(sample), "--no-visible"])
assert result.exit_code == 0
assert "AI-generated (fully synthetic)" in result.output
assert "C2PA validation: integrity=valid, signature=valid" in result.output
assert "signer trust=untrusted, signer validity=expired" in result.output
def test_identify_json_is_valid(self, runner, tmp_png_with_ai_metadata):
result = runner.invoke(main, ["identify", str(tmp_png_with_ai_metadata), "--no-visible", "--json"])
assert result.exit_code == 0
payload = json.loads(result.output)
assert payload["is_ai_generated"] is True
assert payload["confidence"] == "high"
def test_identify_nonexistent_file(self, runner):
result = runner.invoke(main, ["identify", "/nonexistent/file.png"])
assert result.exit_code != 0
class TestDetectSynthIDCommandRemoved:
def test_command_is_not_registered(self, runner):
result = runner.invoke(main, ["detect-synthid", "--help"])
assert result.exit_code != 0
assert "No such command" in result.output
class TestVerifyOpenAISynthIDCommand:
def test_help_names_upload_and_pixel_independence(self, runner):
result = runner.invoke(main, ["verify-openai-synthid", "--help"])
assert result.exit_code == 0
assert "--acknowledge-upload" in result.output
assert "independently of C2PA" in result.output
def test_upload_requires_explicit_acknowledgement(self, runner, tmp_clean_png, monkeypatch):
from remove_ai_watermarks import openai_provenance
called = False
def verify(_source, *, acknowledge_upload):
nonlocal called
assert acknowledge_upload is True
called = True
monkeypatch.setattr(openai_provenance, "verify_openai_synthid", verify)
result = runner.invoke(main, ["verify-openai-synthid", str(tmp_clean_png)])
assert result.exit_code != 0
assert "pass --acknowledge-upload" in result.output
assert called is False
def test_json_result_is_machine_readable(self, runner, tmp_clean_png, monkeypatch):
from remove_ai_watermarks import openai_provenance
expected = openai_provenance.OpenAISynthIDDetection(
status="not_detected",
model=None,
generated_at=None,
api_created_at=1_778_000_000,
)
monkeypatch.setattr(
openai_provenance,
"verify_openai_synthid",
lambda _source, *, acknowledge_upload: expected if acknowledge_upload else None,
)
result = runner.invoke(
main,
["verify-openai-synthid", str(tmp_clean_png), "--acknowledge-upload", "--json"],
)
assert result.exit_code == 0, result.output
payload = json.loads(result.output)
assert payload == expected.to_dict()
assert "c2pa" not in payload
def test_text_result_preserves_negative_scope(self, runner, tmp_clean_png, monkeypatch):
from remove_ai_watermarks import openai_provenance
expected = openai_provenance.OpenAISynthIDDetection(
status="not_detected",
model=None,
generated_at=None,
api_created_at=None,
)
monkeypatch.setattr(
openai_provenance,
"verify_openai_synthid",
lambda _source, *, acknowledge_upload: expected if acknowledge_upload else None,
)
result = runner.invoke(
main,
["verify-openai-synthid", str(tmp_clean_png), "--acknowledge-upload"],
)
assert result.exit_code == 0, result.output
assert "AI provenance metadata was stripped" in result.output
assert "not proof" in result.output
class TestBatchCommand:
"""Tests for the 'batch' subcommand."""
def test_batch_help(self, runner):
result = runner.invoke(main, ["batch", "--help"])
assert result.exit_code == 0
def test_batch_empty_dir(self, runner, tmp_path):
empty_dir = tmp_path / "empty"
empty_dir.mkdir()
result = runner.invoke(main, ["batch", str(empty_dir)])
assert result.exit_code == 0
assert "No supported images" in result.output
def test_batch_visible_mode(self, runner, tmp_path):
input_dir = _make_batch_dir(tmp_path)
output_dir = tmp_path / "output"
result = runner.invoke(
main,
["batch", str(input_dir), "-o", str(output_dir), "--mode", "visible"],
)
assert result.exit_code == 0
assert "3 processed" in result.output
assert output_dir.exists()
assert len(list(output_dir.glob("*.png"))) == 3
def test_batch_metadata_mode(self, runner, tmp_path):
input_dir = _make_batch_dir_with_metadata(tmp_path)
output_dir = tmp_path / "output"
result = runner.invoke(
main,
["batch", str(input_dir), "-o", str(output_dir), "--mode", "metadata"],
)
assert result.exit_code == 0
assert "3 processed" in result.output
assert output_dir.exists()
assert len(list(output_dir.glob("*.png"))) == 3
# Verify AI metadata was stripped
for out_img in output_dir.glob("*.png"):
with Image.open(out_img) as img:
assert "parameters" not in img.info
def test_batch_invisible_mode(self, runner, tmp_path):
input_dir = _make_batch_dir(tmp_path)
output_dir = tmp_path / "output"
mock_cls, _mock_engine = _mock_invisible_engine()
with (
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
patch("remove_ai_watermarks.cli.invisible_available", return_value=True, create=True),
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
):
result = runner.invoke(
main,
["batch", str(input_dir), "-o", str(output_dir), "--mode", "invisible", "--force"],
)
assert result.exit_code == 0, result.output
assert "3 processed" in result.output
def test_batch_cpu_offload_flows_to_cached_engine(self, runner, tmp_path):
input_dir = _make_batch_dir(tmp_path)
output_dir = tmp_path / "output"
mock_cls, _mock_engine = _mock_invisible_engine()
with (
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
patch("remove_ai_watermarks.cli.invisible_available", return_value=True, create=True),
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
):
result = runner.invoke(
main,
[
"batch",
str(input_dir),
"-o",
str(output_dir),
"--mode",
"invisible",
"--cpu-offload",
"--force",
],
)
assert result.exit_code == 0, result.output
assert mock_cls.call_args.kwargs["cpu_offload"] is True
@pytest.mark.parametrize("mode", ["invisible", "all"])
def test_batch_controlnet_scale_flows_to_the_cached_engine(self, runner, tmp_path, mode):
"""`_batch_engine` is a SECOND engine-construction site: `all` builds its engine
inside `_run_invisible`, batch prebuilds one for the whole directory. Pinning the
value on the `all` path alone leaves this one free to hardcode a constant."""
input_dir = _make_batch_dir(tmp_path)
output_dir = tmp_path / "output"
mock_cls, _mock_engine = _mock_invisible_engine()
args = ["batch", str(input_dir), "-o", str(output_dir), "--mode", mode]
with (
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
patch("remove_ai_watermarks.cli.invisible_available", return_value=True, create=True),
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
):
result = runner.invoke(main, [*args, "--controlnet-scale", "0.65", "--force"])
assert result.exit_code == 0, result.output
assert mock_cls.call_args.kwargs["controlnet_conditioning_scale"] == pytest.approx(0.65)
def test_batch_invisible_skips_no_signal_and_copies_through(self, runner, tmp_path):
"""P0#5: batch invisible mode skips the scrub on signal-less images (no
--force) and copies the input through, so the output dir is complete with the
pixels left intact and the engine never called."""
input_dir = _make_batch_dir(tmp_path)
output_dir = tmp_path / "output"
mock_cls, mock_engine = _mock_invisible_engine()
with (
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
patch("remove_ai_watermarks.cli.invisible_available", return_value=True, create=True),
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
):
result = runner.invoke(
main,
["batch", str(input_dir), "-o", str(output_dir), "--mode", "invisible"],
)
assert result.exit_code == 0, result.output
assert "3 processed" in result.output
assert len(list(output_dir.glob("*.png"))) == 3 # inputs copied through
mock_engine.remove_watermark.assert_not_called()
def test_batch_all_mode(self, runner, tmp_path):
input_dir = _make_batch_dir(tmp_path)
output_dir = tmp_path / "output"
mock_cls, _mock_engine = _mock_invisible_engine()
with (
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
patch("remove_ai_watermarks.cli.invisible_available", return_value=True, create=True),
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
):
result = runner.invoke(
main,
["batch", str(input_dir), "-o", str(output_dir), "--mode", "all", "--force"],
)
assert result.exit_code == 0, result.output
assert "3 processed" in result.output
def test_batch_all_mode_preserves_rgba(self, runner, tmp_path):
"""Regression: batch ``all`` must keep transparency across the
alpha-dropping invisible step (mirrors test_all_preserves_rgba_...).
"""
input_dir = tmp_path / "input"
input_dir.mkdir()
rgba = np.zeros((200, 200, 4), dtype=np.uint8)
rgba[:, :, :3] = 200
rgba[50:150, 50:150, 3] = 255
cv2.imwrite(str(input_dir / "rgba.png"), rgba)
output_dir = tmp_path / "output"
mock_cls, _engine = _mock_invisible_engine_drops_alpha()
with (
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
patch("remove_ai_watermarks.cli.invisible_available", return_value=True, create=True),
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
):
result = runner.invoke(
main,
["batch", str(input_dir), "-o", str(output_dir), "--mode", "all", "--force"],
)
assert result.exit_code == 0, result.output
out = cv2.imread(str(output_dir / "rgba.png"), cv2.IMREAD_UNCHANGED)
assert out.ndim == 3, f"output not 3D: shape={out.shape}"
assert out.shape[2] == 4, f"output is not RGBA: shape={out.shape}"
assert out[0, 0, 3] == 0
assert out[100, 100, 3] == 255
def test_batch_forwards_an_explicit_adaptive_polish(self, runner, tmp_path):
"""A typed --adaptive-polish must reach the engine as True, not as None.
This mocks the engine, so it covers the CLI's forwarding only; the per-profile
resolution of an UNSET flag happens inside the real engine and is guarded by
test_invisible_engine.py::TestEngineResolvesThePolishPerProfile.
"""
input_dir = _make_batch_dir(tmp_path, count=2)
output_dir = tmp_path / "output"
mock_cls, mock_engine = _mock_invisible_engine()
with (
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
patch("remove_ai_watermarks.cli.invisible_available", return_value=True, create=True),
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
):
result = runner.invoke(
main,
[
"batch",
str(input_dir),
"-o",
str(output_dir),
"--mode",
"invisible",
"--adaptive-polish",
"--force",
],
)
assert result.exit_code == 0, result.output
assert "2 processed" in result.output
assert mock_cls.call_args.kwargs["pipeline"] == "qwen-zimage"
assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is True
def test_batch_default_output_dir(self, runner, tmp_path):
input_dir = _make_batch_dir(tmp_path)
result = runner.invoke(
main,
["batch", str(input_dir), "--mode", "visible"],
)
assert result.exit_code == 0
expected_dir = tmp_path / "input_clean"
assert expected_dir.exists()
def test_batch_errors_exit_nonzero(self, runner, tmp_path):
"""Regression: batch used to always exit 0 even when every image errored,
hiding failure from a wrapping service. A corrupt image must yield a non-zero
exit and an error count."""
input_dir = tmp_path / "input"
input_dir.mkdir()
(input_dir / "corrupt.png").write_bytes(b"this is not a PNG at all" * 50)
result = runner.invoke(main, ["batch", str(input_dir), "--mode", "visible"])
assert result.exit_code != 0, result.output
assert "error" in result.output.lower()
def test_batch_invisible_gpu_missing_writes_output_and_exits_nonzero(self, runner, tmp_path):
"""Regression: batch --mode invisible with a signal-bearing image but no GPU
deps used to write NO output for that image and still exit 0, silently dropping
the files that most needed processing. It must now copy the input through (so the
output dir is complete), warn about the retained SynthID watermark, and exit
non-zero -- mirroring the single ``all`` command."""
input_dir = _make_batch_dir_with_metadata(tmp_path, count=3) # SD params = invisible signal
output_dir = tmp_path / "output"
with patch("remove_ai_watermarks.invisible_engine.is_available", return_value=False):
result = runner.invoke(
main,
["batch", str(input_dir), "-o", str(output_dir), "--mode", "invisible"],
)
assert result.exit_code != 0, result.output
assert "NOT removed" in result.output
assert len(list(output_dir.glob("*.png"))) == 3 # every input copied through, none dropped
class TestGpuHintMarkup:
"""The install hint must name the extra that actually makes a removal run.
It must also survive to the user with its ``[...]`` token intact (plain output
prints it verbatim, with no markup parsing). It used to say ``[diffusion]``,
which installs torch and diffusers but not the DiffSynth face stage both
profiles run -- so following the advice produced a second, different failure.
"""
def test_the_hint_is_a_command_the_user_can_actually_paste(self, runner, sample_png):
"""The extras bracket must be shell-quoted.
A bulk replace once folded these hints into one constant and dropped the quotes
the originals had. `pip install remove-ai-watermarks[qwen-zimage]` dies with
"zsh: no matches found" on the macOS default shell before pip ever runs -- an
install hint that does not install, which is the failure this hint was fixed to
stop producing in the first place.
"""
with patch("remove_ai_watermarks.invisible_engine.is_available", return_value=False):
result = runner.invoke(main, ["invisible", str(sample_png)])
assert "pip install 'remove-ai-watermarks[qwen-zimage]'" in result.output
def test_invisible_install_hint_names_the_working_extra(self, runner, sample_png):
with patch("remove_ai_watermarks.invisible_engine.is_available", return_value=False):
result = runner.invoke(main, ["invisible", str(sample_png)])
assert result.exit_code != 0
assert "remove-ai-watermarks[qwen-zimage]" in result.output
def test_all_install_hint_names_the_working_extra(self, runner, sample_png):
# The `all` pipeline skips the invisible step with a warning that carries
# the same hint; it must name the same extra.
with patch("remove_ai_watermarks.invisible_engine.is_available", return_value=False):
result = runner.invoke(main, ["all", str(sample_png)])
assert "remove-ai-watermarks[qwen-zimage]" in result.output
class TestEraseCommand:
"""Tests for the 'erase' universal region eraser subcommand."""
def test_erase_help(self, runner):
result = runner.invoke(main, ["erase", "--help"])
assert result.exit_code == 0
assert "--region" in result.output
assert "--backend" in result.output
def test_erase_single_region(self, runner, sample_png, tmp_path):
output = tmp_path / "erased.png"
result = runner.invoke(
main,
["erase", str(sample_png), "--region", "10,10,40,40", "-o", str(output)],
)
assert result.exit_code == 0, result.output
assert output.exists()
def test_erase_two_regions(self, runner, sample_png, tmp_path):
output = tmp_path / "erased2.png"
result = runner.invoke(
main,
[
"erase",
str(sample_png),
"--region",
"10,10,30,30",
"--region",
"120,120,30,30",
"-o",
str(output),
],
)
assert result.exit_code == 0, result.output
assert output.exists()
# The banner reports the region count it processed.
assert "2 region(s)" in result.output
def test_erase_default_output_name(self, runner, sample_png):
result = runner.invoke(main, ["erase", str(sample_png), "--region", "10,10,40,40"])
assert result.exit_code == 0, result.output
assert sample_png.with_stem(sample_png.stem + "_clean").exists()
def test_erase_malformed_region_exits_nonzero(self, runner, sample_png, tmp_path):
output = tmp_path / "x.png"
# Only three values: click.BadParameter -> non-zero exit, no output file.
result = runner.invoke(
main,
["erase", str(sample_png), "--region", "1,2,3", "-o", str(output)],
)
assert result.exit_code != 0
assert not output.exists()
def test_erase_nonexistent_file(self, runner):
result = runner.invoke(main, ["erase", "/nonexistent/file.png", "--region", "0,0,10,10"])
assert result.exit_code != 0
def test_erase_lama_backend_without_onnxruntime(self, runner, sample_png, tmp_path):
# The LaMa backend needs onnxruntime; without it the CLI must surface a
# clear error and exit non-zero rather than crash. When onnxruntime IS
# installed there is no missing-dep path to exercise, so skip.
from remove_ai_watermarks.region_eraser import lama_available
if lama_available():
pytest.skip("onnxruntime installed; missing-dep error path not reachable")
output = tmp_path / "y.png"
result = runner.invoke(
main,
["erase", str(sample_png), "--region", "10,10,40,40", "--backend", "lama", "-o", str(output)],
)
assert result.exit_code != 0
assert "onnxruntime" in result.output.lower()
assert not output.exists()
def test_visible_backend_runtime_error_exits_cleanly(runner, tmp_path, monkeypatch):
# A backend whose extra is missing raises RuntimeError in region_eraser.erase;
# the visible --mark auto path must surface it cleanly, not as a raw traceback (#9).
from pathlib import Path
from remove_ai_watermarks import region_eraser
doubao = Path(__file__).resolve().parent.parent / "data" / "fixtures" / "provenance" / "doubao-1.png"
if not doubao.exists():
pytest.skip("doubao sample not present")
def boom(*_a, **_k):
raise RuntimeError("MI-GAN backend requires onnxruntime. Install the extra: ...")
monkeypatch.setattr(region_eraser, "erase", boom)
out = tmp_path / "out.png"
result = runner.invoke(main, ["visible", str(doubao), "-o", str(out), "--backend", "migan"])
assert result.exit_code == 1
assert not isinstance(result.exception, RuntimeError), "RuntimeError leaked as a traceback"
@pytest.mark.parametrize(
("name", "content"),
[
("empty.png", b""),
("notimage.jpg", b"plain text, not an image at all " * 20),
("truncated.png", b"\x89PNG\r\n\x1a\n" + b"\x00" * 40),
],
)
@pytest.mark.parametrize("cmd", [["metadata", "--remove"], ["visible", "--backend", "cv2"]])
def test_unreadable_input_exits_cleanly(runner, tmp_path, name, content, cmd):
"""Regression: a corrupt / empty / non-image file (real prod uploads include ~0.2%
truncated files) must NEVER leak a raw PIL/OSError/ValueError traceback. `metadata
--remove` is fail-safe -- an undecodable file is copied through unchanged (exit 0),
a strip that cannot parse the file is a no-op, not a crash; `visible` must decode to
remove a mark, so it is a clean error (exit 1). Found by the runtime mode fuzz."""
bad = tmp_path / name
bad.write_bytes(content)
out = tmp_path / "out.png"
result = runner.invoke(main, [cmd[0], str(bad), "-o", str(out), *cmd[1:]])
assert result.exception is None or isinstance(result.exception, SystemExit), (
f"leaked a raw traceback: {result.exception!r}"
)
if cmd[0] == "metadata":
assert result.exit_code == 0, result.output # fail-safe copy-through
assert out.read_bytes() == content # input passed through unchanged
else:
assert result.exit_code == 1, result.output # cannot remove a mark from unreadable input
assert "Error" in result.output