diff --git a/README.md b/README.md index 6074a81..4bf37f3 100644 --- a/README.md +++ b/README.md @@ -126,6 +126,13 @@ uv pip install -e . After installation the `remove-ai-watermarks` command is available system-wide. +> **Note**: The base install covers visible watermark removal and metadata stripping. +> For invisible watermark removal (SynthID etc.), install GPU dependencies: +> +> ```bash +> pip install -e ".[gpu]" # or: uv pip install -e ".[gpu]" +> ``` + #### Invisible watermark removal Invisible removal uses diffusion models and a GPU for reasonable speed. @@ -237,7 +244,7 @@ pip install certifi - [noai-watermark](https://github.com/mertizci/noai-watermark) by mertizci — invisible watermark removal engine - [GeminiWatermarkTool](https://github.com/allenk/GeminiWatermarkTool) by Allen Kuo (MIT) — visible watermark removal algorithm - [CtrlRegen](https://github.com/yepengliu/CtrlRegen) by Liu et al. (ICLR 2025) — controllable regeneration pipeline -- [NeuralBleach](https://github.com/...) (MIT) — analog humanizer technique +- NeuralBleach (MIT) — analog humanizer technique ## ⚠️ Disclaimer diff --git a/pyproject.toml b/pyproject.toml index 6f0b911..d15e869 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "remove-ai-watermarks" -version = "0.3.0" +version = "0.3.1" description = "Remove visible and invisible AI watermarks from images (Gemini / Nano Banana, ChatGPT, Stable Diffusion)" readme = "README.md" requires-python = ">=3.10" @@ -26,7 +26,6 @@ gpu = [ "controlnet-aux>=0.0.9", "safetensors", "ultralytics>=8.0.0", - "opencv-python>=4.8.0", ] dev = [ "pytest>=8.0.0", diff --git a/src/remove_ai_watermarks/__init__.py b/src/remove_ai_watermarks/__init__.py index 168093b..057c792 100644 --- a/src/remove_ai_watermarks/__init__.py +++ b/src/remove_ai_watermarks/__init__.py @@ -1,3 +1,3 @@ """Remove-AI-Watermarks: Unified tool for removing visible and invisible AI watermarks.""" -__version__ = "0.3.0" +__version__ = "0.3.1" diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index c6f7d28..222106b 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -202,7 +202,9 @@ def cmd_visible( @click.option("--device", type=click.Choice(["auto", "cpu", "mps", "cuda"]), default="auto", help="Inference device.") @click.option("--seed", type=int, default=None, help="Random seed for reproducibility.") @click.option("--hf-token", type=str, default=None, help="HuggingFace API token.") -@click.option("--humanize", type=float, default=0.0, help="Humanization strength (0.0–1.0) for invisible removal.") +@click.option( + "--humanize", type=float, default=0.0, help="Analog Humanizer film grain intensity (0 = off, typical: 2.0–6.0)." +) @click.pass_context def cmd_invisible( ctx: click.Context, @@ -219,7 +221,17 @@ def cmd_invisible( """Remove invisible AI watermarks (SynthID, StableSignature, TreeRing). Uses diffusion-based regeneration. Requires GPU for reasonable speed. + Requires the [gpu] extra: pip install 'remove-ai-watermarks[gpu]' """ + from remove_ai_watermarks.invisible_engine import is_available as invisible_available + + if not invisible_available(): + console.print( + "[red]Error:[/] GPU dependencies not installed.\n" + " Install them with: [bold]pip install 'remove-ai-watermarks[gpu]'[/]" + ) + raise SystemExit(1) + from remove_ai_watermarks.invisible_engine import InvisibleEngine source = _validate_image(source) @@ -333,7 +345,9 @@ def cmd_metadata( @click.option("--device", type=click.Choice(["auto", "cpu", "mps", "cuda"]), default="auto", help="Inference device.") @click.option("--seed", type=int, default=None, help="Random seed for reproducibility.") @click.option("--hf-token", type=str, default=None, help="HuggingFace API token.") -@click.option("--humanize", type=float, default=0.0, help="Humanization strength (0.0–1.0) for invisible removal.") +@click.option( + "--humanize", type=float, default=0.0, help="Analog Humanizer film grain intensity (0 = off, typical: 2.0–6.0)." +) @click.pass_context def cmd_all( ctx: click.Context, @@ -415,31 +429,39 @@ def cmd_all( # ── Step 2: Invisible watermark ────────────────────────────── console.print("\n [bold cyan]② Invisible watermark removal[/]") - from remove_ai_watermarks.invisible_engine import InvisibleEngine + from remove_ai_watermarks.invisible_engine import is_available as invisible_available - device_str = None if device == "auto" else device + if not invisible_available(): + console.print( + " [yellow]⚠[/] Skipped — GPU dependencies not installed.\n" + " Install them with: [bold]pip install 'remove-ai-watermarks[gpu]'[/]" + ) + else: + from remove_ai_watermarks.invisible_engine import InvisibleEngine - def progress_cb(msg: str) -> None: - console.print(f" [dim]{msg}[/]") + device_str = None if device == "auto" else device - inv_engine = InvisibleEngine( - model_id=model, - device=device_str, - pipeline=pipeline, - hf_token=hf_token, - progress_callback=progress_cb, - ) + def progress_cb(msg: str) -> None: + console.print(f" [dim]{msg}[/]") - console.print(f" [dim]Strength:[/] {strength} Steps: {steps}") - inv_engine.remove_watermark( - image_path=tmp_path, - output_path=tmp_path, - strength=strength, - num_inference_steps=steps, - seed=seed, - humanize=humanize, - ) - console.print(" [green]✓[/] Invisible watermark removed") + inv_engine = InvisibleEngine( + model_id=model, + device=device_str, + pipeline=pipeline, + hf_token=hf_token, + progress_callback=progress_cb, + ) + + console.print(f" [dim]Strength:[/] {strength} Steps: {steps}") + inv_engine.remove_watermark( + image_path=tmp_path, + output_path=tmp_path, + strength=strength, + num_inference_steps=steps, + seed=seed, + humanize=humanize, + ) + console.print(" [green]✓[/] Invisible watermark removed") # ── Step 3: Metadata ───────────────────────────────────────── console.print("\n [bold cyan]③ AI metadata stripping[/]") @@ -486,7 +508,9 @@ def cmd_all( @click.option("--strength", type=float, default=None, help="Denoising strength (invisible mode).") @click.option("--steps", type=int, default=50, help="Number of denoising steps (invisible mode).") @click.option("--inpaint/--no-inpaint", default=True, help="Apply inpainting (visible mode).") -@click.option("--humanize", type=float, default=0.0, help="Humanization strength (0.0–1.0) for invisible removal.") +@click.option( + "--humanize", type=float, default=0.0, help="Analog Humanizer film grain intensity (0 = off, typical: 2.0–6.0)." +) @click.option("--pipeline", type=click.Choice(["default", "ctrlregen"]), default="default", help="Pipeline profile.") @click.option("--device", type=click.Choice(["auto", "cpu", "mps", "cuda"]), default="auto", help="Inference device.") @click.option("--seed", type=int, default=None, help="Random seed for reproducibility.") @@ -588,7 +612,13 @@ def cmd_batch( if invisible_available(): from remove_ai_watermarks.invisible_engine import InvisibleEngine - engine_inv = InvisibleEngine() + if "_inv_engine" not in ctx.obj: + ctx.obj["_inv_engine"] = InvisibleEngine( + device=None if device == "auto" else device, + pipeline=pipeline, + hf_token=hf_token, + ) + engine_inv = ctx.obj["_inv_engine"] engine_inv.remove_watermark( img_path if mode == "invisible" else out_path, out_path, diff --git a/src/remove_ai_watermarks/invisible_engine.py b/src/remove_ai_watermarks/invisible_engine.py index ac33715..a774311 100644 --- a/src/remove_ai_watermarks/invisible_engine.py +++ b/src/remove_ai_watermarks/invisible_engine.py @@ -3,8 +3,8 @@ Wraps the vendored noai-watermark code for removing invisible AI watermarks (SynthID, StableSignature, TreeRing) via diffusion-based regeneration. -This module requires the 'invisible' extra dependencies: - uv pip install 'remove-ai-watermarks[invisible]' +This module requires the 'gpu' extra dependencies: + uv pip install 'remove-ai-watermarks[gpu]' """ from __future__ import annotations diff --git a/tests/test_cli.py b/tests/test_cli.py index f971961..3bac3c2 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -79,7 +79,7 @@ class TestMainGroup: def test_version(self, runner): result = runner.invoke(main, ["--version"]) assert result.exit_code == 0 - assert "0.3.0" in result.output + assert "0.3.1" in result.output def test_no_command_shows_banner(self, runner): result = runner.invoke(main, []) @@ -151,8 +151,9 @@ class TestInvisibleCommand: 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.cli.InvisibleEngine", mock_cls, create=True), patch( - "remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls + with ( + patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True), + patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls), ): result = runner.invoke( main, @@ -164,8 +165,9 @@ class TestInvisibleCommand: def test_invisible_default_output(self, runner, sample_png): 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 + with ( + 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)]) assert result.exit_code == 0, result.output @@ -188,8 +190,9 @@ class TestAllCommand: 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 + with ( + patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True), + patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls), ): result = runner.invoke( main, @@ -282,10 +285,11 @@ class TestBatchCommand: 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 + 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, @@ -298,10 +302,11 @@ class TestBatchCommand: 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 + 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, @@ -319,4 +324,3 @@ class TestBatchCommand: assert result.exit_code == 0 expected_dir = tmp_path / "input_clean" assert expected_dir.exists() - diff --git a/uv.lock b/uv.lock index 70a4caf..498f9b1 100644 --- a/uv.lock +++ b/uv.lock @@ -1993,64 +1993,73 @@ wheels = [ [[package]] name = "remove-ai-watermarks" -version = "0.3.0" +version = "0.3.1" source = { editable = "." } dependencies = [ - { name = "accelerate" }, { name = "click" }, { name = "color-matcher" }, - { name = "controlnet-aux" }, - { name = "diffusers" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "opencv-python" }, + { name = "opencv-python-headless" }, { name = "piexif" }, { name = "pillow" }, { name = "python-dotenv" }, { name = "requests" }, { name = "rich" }, - { name = "safetensors" }, - { name = "torch" }, - { name = "transformers" }, - { name = "ultralytics" }, ] [package.optional-dependencies] all = [ + { name = "accelerate" }, + { name = "controlnet-aux" }, + { name = "diffusers" }, { name = "pytest" }, { name = "pytest-cov" }, { name = "ruff" }, + { name = "safetensors" }, + { name = "torch" }, + { name = "transformers" }, + { name = "ultralytics" }, ] dev = [ { name = "pytest" }, { name = "pytest-cov" }, { name = "ruff" }, ] +gpu = [ + { name = "accelerate" }, + { name = "controlnet-aux" }, + { name = "diffusers" }, + { name = "safetensors" }, + { name = "torch" }, + { name = "transformers" }, + { name = "ultralytics" }, +] [package.metadata] requires-dist = [ - { name = "accelerate", specifier = ">=0.25.0" }, + { name = "accelerate", marker = "extra == 'gpu'", specifier = ">=0.25.0" }, { name = "click", specifier = ">=8.0.0" }, { name = "color-matcher" }, - { name = "controlnet-aux", specifier = ">=0.0.9" }, - { name = "diffusers", specifier = ">=0.25.0" }, + { name = "controlnet-aux", marker = "extra == 'gpu'", specifier = ">=0.0.9" }, + { name = "diffusers", marker = "extra == 'gpu'", specifier = ">=0.25.0" }, { name = "numpy", specifier = ">=1.24.0" }, - { name = "opencv-python", specifier = ">=4.8.0" }, + { name = "opencv-python-headless", specifier = ">=4.8.0" }, { name = "piexif", specifier = ">=1.1.3" }, { name = "pillow", specifier = ">=10.0.0" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0.0" }, { name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=4.1.0" }, { name = "python-dotenv", specifier = ">=1.0.0" }, - { name = "remove-ai-watermarks", extras = ["dev"], marker = "extra == 'all'" }, + { name = "remove-ai-watermarks", extras = ["gpu", "dev"], marker = "extra == 'all'" }, { name = "requests", specifier = ">=2.33.0" }, { name = "rich", specifier = ">=13.0.0" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.4.0" }, - { name = "safetensors" }, - { name = "torch", specifier = ">=2.0.0" }, - { name = "transformers", specifier = ">=4.35.0" }, - { name = "ultralytics", specifier = ">=8.0.0" }, + { name = "safetensors", marker = "extra == 'gpu'" }, + { name = "torch", marker = "extra == 'gpu'", specifier = ">=2.0.0" }, + { name = "transformers", marker = "extra == 'gpu'", specifier = ">=4.35.0" }, + { name = "ultralytics", marker = "extra == 'gpu'", specifier = ">=8.0.0" }, ] -provides-extras = ["dev", "all"] +provides-extras = ["gpu", "dev", "all"] [[package]] name = "requests"