From 4a896cd4b59b20571d1975eaf751d63e3776db11 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Tue, 4 Aug 2026 19:33:54 -0700 Subject: [PATCH] Make InvisibleOptions engine knobs only and pin the forwarding `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 --- .claude/rules/development.md | 20 +++++- docs/module-internals.md | 22 +++++-- docs/python-api.md | 8 ++- src/remove_ai_watermarks/api.py | 44 +++++++++---- src/remove_ai_watermarks/cli.py | 10 +-- tests/test_api.py | 112 +++++++++++++++++++++++++++++--- tests/test_cli.py | 40 ++++++++++++ tests/test_identify.py | 10 ++- 8 files changed, 227 insertions(+), 39 deletions(-) diff --git a/.claude/rules/development.md b/.claude/rules/development.md index 0522e96..c58230b 100644 --- a/.claude/rules/development.md +++ b/.claude/rules/development.md @@ -41,7 +41,25 @@ Do not classify an entire module as untestable because its main path downloads a mirror another, compare them field by field rather than pinning the values you happen to know about, so the next field added on one side and not the other fails at the seam. Two of these defaults drifted in practice and neither needed a GPU to catch; - the incident is recorded in `docs/module-internals.md`. + the incident is recorded in `docs/module-internals.md`. Keep the comparison free of + an exception table: a field that needs one is a field that belongs elsewhere, which + is what `force` turned out to be. + +A defaults comparison is not a forwarding test, and the two fail differently. Pin the +VALUE at the seam, not just the name -- `_run_invisible` passed the whole suite with +`controlnet_conditioning_scale` hardcoded, because nothing asserted the caller's value +arrived. `test_every_field_arrives_at_the_engine_with_the_caller_s_value` drives the +real seam with every field set off its default, so one test covers the whole bag +instead of one assertion per knob. + +Count the seams before believing a knob is covered. Each of `force` and +`controlnet_conditioning_scale` reaches the engine through TWO paths -- `remove_all` +versus `remove_batch(mode="all")` for the first, `_run_invisible` versus `_batch_engine` +for the second -- and in both cases guarding one path left the other free to hardcode a +constant with a green suite. The mode-parametrized guards in +`TestRemoveBatchLibrary::test_force_reaches_the_scrub_gate_in_every_scrubbing_mode` and +`TestBatchCommand::test_batch_controlnet_scale_flows_to_the_cached_engine` exist because +that is what actually happened. Use availability checks only for paths that actually load large models. diff --git a/docs/module-internals.md b/docs/module-internals.md index 7807757..ed4a425 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -87,13 +87,21 @@ over: - `remove_all`, returning a `RemoveAllResult` after the visible, invisible, and metadata stages - `remove_batch`, returning a `BatchSummary` for one directory and one mode -- `InvisibleOptions`, the invisible stage's knobs as one immutable value. Every - default mirrors `InvisibleEngine`, so a bare `InvisibleOptions()` behaves - exactly like calling the engine with no arguments. Two silently stopped: - `max_resolution=None` reached `_target_size`'s `max_resolution > 0` and raised - `TypeError` on every library call, and `cpu_offload=True` made a library run - slower than the identical CLI run. `TestInvisibleOptionsMirrorTheEngine` - compares the two signatures field by field +- `InvisibleOptions`, the invisible stage's knobs as one immutable value. 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 what shapes the + loaded stack and `remove_watermark` for the per-image ones, so `_run_invisible` + forwards each field to the right one rather than splatting the whole bag. Two + defaults silently stopped + mirroring: `max_resolution=None` reached `_target_size`'s `max_resolution > 0` + and raised `TypeError` on every library call, and `cpu_offload=True` made a + library run slower than the identical CLI run. + `TestInvisibleOptionsMirrorTheEngine` compares the two signatures field by + field, and deliberately keeps no exception table: a field needing one is a + field that belongs elsewhere. `force` was such a field, and it decides whether + the engine runs rather than how, so it is a parameter of `remove_all` and + `remove_batch` next to `backend` and `sensitivity` - `MetadataStripIncomplete`, raised before any write when AI metadata survives `remove_all` reports progress as `(stage, detail)` pairs of stable tokens, not diff --git a/docs/python-api.md b/docs/python-api.md index 7ac1941..a28353e 100644 --- a/docs/python-api.md +++ b/docs/python-api.md @@ -106,11 +106,17 @@ from remove_ai_watermarks import InvisibleOptions raiw.remove_all( "input.png", "clean.png", - invisible=InvisibleOptions(strength=0.35, force=True), + invisible=InvisibleOptions(strength=0.35), + force=True, progress=print, ) ``` +`InvisibleOptions` carries only what `InvisibleEngine` itself takes, and uses the +engine's own parameter names and defaults. `force`, which decides whether the +engine runs at all, is a parameter of `remove_all` and `remove_batch` alongside +`backend` and `sensitivity`. + If AI metadata survives the strip, `remove_all` raises `MetadataStripIncomplete` **before** writing anything: an AI-readable output is worse than no output. diff --git a/src/remove_ai_watermarks/api.py b/src/remove_ai_watermarks/api.py index eae2fc6..8d27eaf 100644 --- a/src/remove_ai_watermarks/api.py +++ b/src/remove_ai_watermarks/api.py @@ -216,11 +216,18 @@ def remove_visible( class InvisibleOptions: """The invisible stage's knobs, as one value instead of a dozen parameters. - Every default MIRRORS ``InvisibleEngine``, so a bare ``InvisibleOptions()`` behaves - exactly like calling the engine with no arguments; two once drifted and broke - silently, so ``TestInvisibleOptionsMirrorTheEngine`` compares the two signatures - field by field. Immutable so a batch can build it once and reuse it across every - image while the engine itself is cached separately. + 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 @@ -231,13 +238,11 @@ class InvisibleOptions: unsharp: float = 0.0 adaptive_polish: bool | None = None max_resolution: int = 0 - controlnet_scale: float = 1.0 + controlnet_conditioning_scale: float = 1.0 cpu_offload: bool = False tile: bool = False tile_size: int = 1024 tile_overlap: int = 128 - # Scrub even when no invisible watermark is locally detectable. - force: bool = False # What the invisible stage did. "unavailable" is the one outcome the caller must @@ -365,6 +370,7 @@ def remove_all( 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: @@ -374,6 +380,11 @@ def remove_all( 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. @@ -425,7 +436,7 @@ def remove_all( raise OSError(f"failed to write the staged intermediate: {staged}") # ── 2. Invisible watermark ── - outcome = _run_invisible(src, staged, staged, opts, engine, say, evidence) + 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: @@ -458,6 +469,7 @@ def _run_invisible( engine: Any | None, say: Callable[[str, str], None], evidence: _SourceEvidence, + force: bool, ) -> InvisibleOutcome: """Run, or deliberately skip, the diffusion scrub. @@ -472,7 +484,7 @@ def _run_invisible( if not is_available(): say("invisible", "unavailable") return "unavailable" - if not (opts.force or evidence.has_invisible_target()): + if not (force or evidence.has_invisible_target()): say("invisible", "no-signal") return "no-signal" @@ -494,7 +506,7 @@ def _run_invisible( pipeline=opts.pipeline, hf_token=opts.hf_token, progress_callback=lambda message: say("invisible", message), - controlnet_conditioning_scale=opts.controlnet_scale, + controlnet_conditioning_scale=opts.controlnet_conditioning_scale, cpu_offload=opts.cpu_offload, ) engine.remove_watermark( @@ -538,6 +550,7 @@ def remove_batch( 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: @@ -545,8 +558,8 @@ def remove_batch( 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. ``engine`` is threaded straight through, so a caller that passes a - constructed ``InvisibleEngine`` loads the model once for the whole run. + 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 @@ -567,7 +580,7 @@ def remove_batch( 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, engine, say) + 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))) @@ -590,6 +603,7 @@ def _run_batch_one( backend: Backend, sensitivity: Sensitivity, invisible: InvisibleOptions | None, + force: bool, engine: Any | None, say: Callable[[Path, str, str], None], ) -> InvisibleOutcome | None: @@ -604,6 +618,7 @@ def _run_batch_one( backend=backend, sensitivity=sensitivity, invisible=invisible, + force=force, engine=engine, progress=lambda stage, detail: say(img_path, stage, detail), ) @@ -653,6 +668,7 @@ def _run_batch_one( 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 diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index 49331f9..cda3f1f 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -1496,13 +1496,13 @@ def cmd_all( unsharp=unsharp, adaptive_polish=adaptive_polish, max_resolution=max_resolution, - controlnet_scale=controlnet_scale, + controlnet_conditioning_scale=controlnet_scale, cpu_offload=cpu_offload, tile=tile, tile_size=tile_size, tile_overlap=tile_overlap, - force=force, ), + force=force, progress=progress, ) except MetadataStripIncomplete as e: @@ -1560,7 +1560,7 @@ def _batch_engine(mode: str, options: InvisibleOptions) -> object | None: return InvisibleEngine( pipeline=options.pipeline, hf_token=options.hf_token, - controlnet_conditioning_scale=options.controlnet_scale, + controlnet_conditioning_scale=options.controlnet_conditioning_scale, cpu_offload=options.cpu_offload, ) @@ -1642,12 +1642,11 @@ def cmd_batch( unsharp=unsharp, adaptive_polish=adaptive_polish, max_resolution=max_resolution, - controlnet_scale=controlnet_scale, + controlnet_conditioning_scale=controlnet_scale, cpu_offload=cpu_offload, tile=tile, tile_size=tile_size, tile_overlap=tile_overlap, - force=force, ) with Progress( @@ -1678,6 +1677,7 @@ def cmd_batch( backend=backend, # type: ignore[arg-type] sensitivity=_parse_sensitivity(sensitivity), invisible=invisible_options, + force=force, engine=_batch_engine(mode, invisible_options), progress=on_progress, ) diff --git a/tests/test_api.py b/tests/test_api.py index 67c7fc9..2b0f35f 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -196,10 +196,12 @@ class TestRemoveVisibleOutputPath: class TestInvisibleOptionsMirrorTheEngine: - """``InvisibleOptions`` forwards to ``InvisibleEngine`` and promises every default - mirrors it. Compare the signatures field by field rather than pinning the values we - happen to know about, so the next field added on one side and not the other fails - here. See the incident record in ``docs/module-internals.md``.""" + """``InvisibleOptions`` forwards to ``InvisibleEngine`` and promises every NAME and + default mirrors it. Compare the signatures field by field rather than pinning the + values we happen to know about, so the next field added on one side and not the + other fails here. The comparison deliberately has no exception table: a field that + needs one is a field that belongs somewhere else, which is why ``force`` is a + parameter of ``remove_all``. See the incident record in ``docs/module-internals.md``.""" def test_every_default_matches_the_engine(self): import dataclasses @@ -213,12 +215,67 @@ class TestInvisibleOptionsMirrorTheEngine: for name, p in inspect.signature(method).parameters.items() if p.default is not inspect.Parameter.empty } - # ``force`` is a pipeline decision made before the engine runs, so it has no - # engine counterpart; ``remove_watermark`` spells the controlnet knob out in full. - renamed = {"controlnet_scale": "controlnet_conditioning_scale"} - options = {f.name: f.default for f in dataclasses.fields(api.InvisibleOptions) if f.name != "force"} + options = {f.name: f.default for f in dataclasses.fields(api.InvisibleOptions)} - assert options == {name: engine.get(renamed.get(name, name), "") for name in options} + assert options == {name: engine.get(name, "") for name in options} + + @pytest.mark.skipif(not CHATGPT.exists(), reason="sample image not present") + def test_every_field_arrives_at_the_engine_with_the_caller_s_value(self, monkeypatch, tmp_path): + """A defaults comparison is not a forwarding test. `_run_invisible` hands each + field to one of TWO engine callables by hand, and a hardcoded literal there is + invisible to the mirror above -- `controlnet_conditioning_scale` shipped that way + and the whole suite stayed green. Drive the real seam with every field set OFF + its default and assert the caller's value arrives, whichever callable takes it.""" + import dataclasses + + from remove_ai_watermarks import invisible_engine + + # Every value differs from the default, so a hardcoded default cannot pass. + opts = api.InvisibleOptions( + strength=0.42, + pipeline="sdxl-zimage", + seed=7, + hf_token="token", + humanize=0.3, + unsharp=0.2, + adaptive_polish=True, + max_resolution=1536, + controlnet_conditioning_scale=0.65, + cpu_offload=True, + tile=True, + tile_size=768, + tile_overlap=64, + ) + seen: dict[str, object] = {} + + class FakeEngine: + def __init__(self, **kwargs): + seen.update(kwargs) + + def remove_watermark(self, **kwargs): + seen.update(kwargs) + + monkeypatch.setattr(invisible_engine, "is_available", lambda: True) + monkeypatch.setattr(invisible_engine, "InvisibleEngine", FakeEngine) + api._run_invisible( + CHATGPT, + CHATGPT, + tmp_path / "out.png", + opts, + None, + lambda _stage, _detail: None, + api._SourceEvidence(CHATGPT), + True, + ) + + missing = {f.name: getattr(opts, f.name) for f in dataclasses.fields(opts) if f.name not in seen} + assert not missing, f"never forwarded to the engine: {missing}" + wrong = { + f.name: (getattr(opts, f.name), seen[f.name]) + for f in dataclasses.fields(opts) + if seen[f.name] != getattr(opts, f.name) + } + assert not wrong, f"forwarded a value the caller did not pass (want, got): {wrong}" class TestRemoveAllLibrary: @@ -381,6 +438,43 @@ class TestRemoveBatchLibrary: assert summary.processed == 0 assert summary.failed == 1 + @pytest.mark.parametrize("mode", ["all", "invisible"]) + @pytest.mark.parametrize(("force", "expected"), [(True, "removed"), (False, "no-signal")]) + def test_force_reaches_the_scrub_gate_in_every_scrubbing_mode(self, monkeypatch, tmp_path, mode, force, expected): + """``force`` reaches the gate through a DIFFERENT seam per mode: ``all`` re-enters + ``remove_all``, ``invisible`` calls ``_run_invisible`` directly. Only the second + was guarded, and pinning the first to False passed the whole suite while every + output kept its watermark and the run still reported the files as processed.""" + from remove_ai_watermarks import api, image_io, invisible_engine + + src = tmp_path / "in" + src.mkdir() + for i in range(2): + image_io.imwrite(src / f"img{i}.png", np.full((64, 64, 3), 120, np.uint8)) + scrubbed: list[Path] = [] + + class FakeEngine: + def remove_watermark(self, **kwargs): + scrubbed.append(kwargs["image_path"]) + image_io.imwrite(kwargs["output_path"], np.full((64, 64, 3), 120, np.uint8)) + + monkeypatch.setattr(invisible_engine, "is_available", lambda: True) + events: list[tuple[str, str]] = [] + + summary = api.remove_batch( + src, + tmp_path / "out", + mode=mode, + backend="cv2", + force=force, + engine=FakeEngine(), + progress=lambda _p, stage, detail: events.append((stage, detail)), + ) + + assert summary.processed == 2 + assert ("invisible", expected) in events + assert len(scrubbed) == (2 if force else 0) + class TestSourceEvidenceHolder: """One metadata extraction per source file, per call. diff --git a/tests/test_cli.py b/tests/test_cli.py index cb744f3..e6184d2 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -485,6 +485,26 @@ class TestAllCommand: 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 @@ -802,6 +822,26 @@ class TestBatchCommand: 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 diff --git a/tests/test_identify.py b/tests/test_identify.py index ca039a8..ee68391 100644 --- a/tests/test_identify.py +++ b/tests/test_identify.py @@ -782,13 +782,19 @@ class TestIdentifyVisibleTextMarks: """The web path identify(check_visible=True, check_invisible=False) must decode the image exactly once and share the array across the sparkle + text-mark detectors. Two decodes of the same bitmap spiked memory on the - small web worker (the OOM the decode-once refactor addresses).""" + small web worker (the OOM the decode-once refactor addresses). + + Count decodes OF THE SOURCE, not every ``imread`` in the process: the Gemini + engine loads its own bundled capture assets on first construction, so a + process-wide count was 3 on a cold engine and 1 on a warm one, and the test + passed only when some earlier test happened to build the engine first.""" import remove_ai_watermarks.image_io as image_io real_imread = image_io.imread with patch.object(image_io, "imread", side_effect=real_imread) as mock_imread: identify(tmp_clean_png, check_visible=True, check_invisible=False) - assert mock_imread.call_count == 1 + source_decodes = [c for c in mock_imread.call_args_list if Path(c.args[0]) == Path(tmp_clean_png)] + assert len(source_decodes) == 1 def test_missing_pixel_extra_preserves_metadata_verdict(self, tmp_png_with_ai_metadata: Path): import remove_ai_watermarks.image_io as image_io