diff --git a/docs/module-internals.md b/docs/module-internals.md index 1f4c892..2c94312 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -68,10 +68,14 @@ Important contracts: - `visible` and `erase` strip and verify AI metadata on a staged sibling and exit `1` with nothing published when a marker survives; they used to swallow a failed strip unless `--verbose` was set. The library `remove_visible` does the same and - raises `MetadataStripIncomplete`, as `remove_all` already did, and the diffusion - writer behind `invisible` raises when its own strip leaves a marker. + raises `MetadataStripIncomplete`, as `remove_all` already did. The diffusion + writer behind `invisible` stages the same way and raises with nothing published + (fixed 2026-10-19; it used to leave the marked file at the output path). `_internal.utils.staged_sibling` is the shared stage: it never pre-creates the - name, so the published file keeps the default permissions. + name, so the published file keeps the default permissions. These writers strip + through `metadata.strip_written_output`, which skips the readability probe: + the file is the package's own encode, and the probe's full decode doubled the + cost of the strip (decided 2026-10-19). - A quiet `metadata --check` and a successful `metadata --remove` (same for `video metadata`) end by repeating the `identify` limit: the pixel channel is untouched and a watermark such as SynthID has no local decoder once its diff --git a/src/remove_ai_watermarks/_internal/watermark_remover.py b/src/remove_ai_watermarks/_internal/watermark_remover.py index 06a5ef2..6c5dabc 100644 --- a/src/remove_ai_watermarks/_internal/watermark_remover.py +++ b/src/remove_ai_watermarks/_internal/watermark_remover.py @@ -212,17 +212,18 @@ class WatermarkRemover: import numpy as np from remove_ai_watermarks import image_io + from remove_ai_watermarks._internal.utils import staged_sibling + from remove_ai_watermarks.metadata import strip_written_output output_path.parent.mkdir(parents=True, exist_ok=True) bgr = np.ascontiguousarray(np.asarray(image.convert("RGB"))[:, :, ::-1]) - if not image_io.imwrite(str(output_path), bgr, display_tags_from=display_tags_from): - image.save(output_path) - from remove_ai_watermarks.metadata import strip_and_verify - - # Verified, so a marker that survives fails the run instead of shipping with it. - _, leftover = strip_and_verify(output_path, output_path, keep_standard=True) - if leftover: - raise RuntimeError(f"AI metadata survived the strip: {', '.join(sorted(leftover))}") + # Staged and verified, so a marker that survives fails the run with nothing + # published at output_path. + with staged_sibling(output_path) as staged: + if not image_io.imwrite(str(staged), bgr, display_tags_from=display_tags_from): + image.save(staged) + if leftover := strip_written_output(staged): + raise RuntimeError(f"AI metadata survived the strip: {', '.join(sorted(leftover))}") def remove_watermark( self, diff --git a/src/remove_ai_watermarks/api.py b/src/remove_ai_watermarks/api.py index 39f7ec4..d1ef1ee 100644 --- a/src/remove_ai_watermarks/api.py +++ b/src/remove_ai_watermarks/api.py @@ -157,10 +157,9 @@ def _write_visible_result( raise OSError(f"failed to write output (is the destination writable?): {out_path}") if strip_metadata: - from remove_ai_watermarks.metadata import strip_and_verify + from remove_ai_watermarks.metadata import strip_written_output - _, leftover = strip_and_verify(staged, staged) - if leftover: + if leftover := strip_written_output(staged): raise MetadataStripIncomplete(set(leftover)) @@ -464,7 +463,7 @@ def remove_all( import tempfile from remove_ai_watermarks import image_io, watermark_registry - from remove_ai_watermarks.metadata import strip_and_verify + from remove_ai_watermarks.metadata import strip_written_output def say(stage: str, detail: str) -> None: if progress is not None: @@ -519,8 +518,7 @@ def remove_all( # ── 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: + if leftover := strip_written_output(staged): # Before the write, on purpose -- see MetadataStripIncomplete. raise MetadataStripIncomplete(set(leftover)) say("metadata", "stripped") diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index b1af397..d40f8c7 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -467,7 +467,7 @@ def _write_output_or_exit( published, so a marker that survives exits 1 with nothing written. Both commands using this used to swallow a failed strip unless ``--verbose`` was set.""" from remove_ai_watermarks._internal.utils import staged_sibling - from remove_ai_watermarks.metadata import strip_and_verify + from remove_ai_watermarks.metadata import strip_written_output output.parent.mkdir(parents=True, exist_ok=True) with staged_sibling(output) as staged: @@ -477,7 +477,7 @@ def _write_output_or_exit( if not strip_metadata: return try: - _, leftover = strip_and_verify(staged, staged) + leftover = strip_written_output(staged) except (OSError, RuntimeError, ValueError) as e: console.print(f" Error: failed to strip AI metadata from {output}: {e}") raise SystemExit(1) from e diff --git a/src/remove_ai_watermarks/metadata.py b/src/remove_ai_watermarks/metadata.py index c846c4f..f8cd9ac 100644 --- a/src/remove_ai_watermarks/metadata.py +++ b/src/remove_ai_watermarks/metadata.py @@ -1968,6 +1968,25 @@ def strip_and_verify( """ if (reason := unreadable_reason(source_path)) is not None: raise ValueError(f"cannot read {source_path.name}: {reason}") + return _strip_and_rescan(source_path, output_path, keep_standard=keep_standard) + + +def strip_written_output(path: Path, *, keep_standard: bool = True) -> dict[str, str]: + """Strip ``path`` in place and return the markers that survived. + + For a file whose pixels this run has already decoded or written itself, so the + readability probe that :func:`strip_and_verify` runs first would only decode + them a second time. + """ + return _strip_and_rescan(path, path, keep_standard=keep_standard)[1] + + +def _strip_and_rescan( + source_path: Path, + output_path: Path | None, + *, + keep_standard: bool, +) -> tuple[Path, dict[str, str]]: out = remove_ai_metadata(source_path, output_path, keep_standard=keep_standard) remaining = get_ai_metadata(out) if not remaining: diff --git a/tests/test_api.py b/tests/test_api.py index 5a18988..44017c3 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -430,7 +430,7 @@ class TestRemoveAllLibrary: from remove_ai_watermarks import api, invisible_engine, metadata monkeypatch.setattr(invisible_engine, "is_available", lambda: False) - monkeypatch.setattr(metadata, "strip_and_verify", lambda src, dst: (dst, {"c2pa"})) + monkeypatch.setattr(metadata, "strip_written_output", lambda _path: {"c2pa": "survived"}) out = tmp_path / "clean.png" with pytest.raises(api.MetadataStripIncomplete, match="c2pa"): diff --git a/tests/test_cli.py b/tests/test_cli.py index be0e12d..a3d608a 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -675,10 +675,7 @@ class TestAllCommand: # 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}), - ), + patch("remove_ai_watermarks.metadata.strip_written_output", return_value={"c2pa": "survived"}), ): result = runner.invoke(main, ["all", str(sample_png), "-o", str(output)]) diff --git a/tests/test_fail_closed.py b/tests/test_fail_closed.py index f04fdfa..a203f23 100644 --- a/tests/test_fail_closed.py +++ b/tests/test_fail_closed.py @@ -239,6 +239,29 @@ class TestSkippedStepsAreFailures: remover = object.__new__(WatermarkRemover) # the writer needs no model state with pytest.raises(RuntimeError, match="survived"): remover._write_output(Image.new("RGB", (16, 16)), tmp_path / "out.png") + assert list(tmp_path.iterdir()) == [], "an output that kept its AI metadata was published" + + @pytest.mark.parametrize("command", ["erase", "remove_visible", "invisible_writer"]) + def test_an_output_the_package_just_wrote_is_not_probed_again(self, tmp_path, clean_png, monkeypatch, command): + """The readability probe decodes every pixel; on our own encode it is pure waste.""" + from PIL import Image + + from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover + from remove_ai_watermarks.api import remove_visible + + probed: list[Path] = [] + real = metadata.unreadable_reason + monkeypatch.setattr(metadata, "unreadable_reason", lambda path: probed.append(path) or real(path)) + out = tmp_path / "out.png" + if command == "erase": + result = _invoke(["erase", "--region", "2,2,10,10", str(clean_png), "-o", str(out)]) + assert result.exit_code == 0, result.output + elif command == "remove_visible": + remove_visible(clean_png, out, write_noop=True) + else: + object.__new__(WatermarkRemover)._write_output(Image.new("RGB", (16, 16)), out) + assert out.exists() + assert [p for p in probed if p.name != clean_png.name] == [] def test_metadata_remove_reports_a_missing_ffmpeg_without_a_traceback(self, tmp_path, monkeypatch): monkeypatch.setattr(shutil, "which", lambda _name: None)