Delete every knob the fixed profiles cannot honor

The CLI still advertised --model, --steps, --guidance-scale, --device and a
deprecated --auto. Each pinned a value the two surviving profiles fix -- the
model stack, the per-stage distilled schedule, CFG 1.0, CUDA -- so the only
outcome any of them had was an error raised several frames below the caller,
under a message naming an internal profile. A flag whose sole result is a
refusal is worse than no flag: it advertises a capability that does not exist,
and it lets a wrapper thread a value that will silently do nothing. They are
gone from the parser, from InvisibleEngine, and from WatermarkRemover, so the
failure is now a TypeError or a Click "No such option" at the point the caller
can act on.

The install hint was wrong in the same way. is_available() checked torch and
diffusers, then told the user to install [diffusion] -- which contains neither
DiffSynth nor the Z-Image face stage both profiles run. Following the advice
produced a second, different failure. The module list and the extra name now
live once in watermark_profiles (REMOVAL_MODULES, INVISIBLE_EXTRA) and are read
by both the CLI gate and the remover's precondition, which cannot drift apart
because they are the same tuple.

The adaptive-polish default moved out of the argument parser. It was resolved by
reading Click's parameter source, which put per-profile data in the CLI layer,
left the engine declaring the opposite default (False vs True) so a library
caller and a CLI caller on one profile got different output, and lost the polish
entirely for anything that supplies the flag non-interactively. The flag is now
tri-state (default=None) and resolve_adaptive_polish owns the per-profile
answer. The seed follows the same rule: the CLI stopped pre-resolving it.

Dead code removed with it: six scan_*_video wrappers and the _scan_video helper
none of them had a caller for, PNG_METADATA_KEYS, feather_region_composite and
the remover region path that was only reachable from a no-caller convenience
wrapper, remove_watermark_batch on both layers, try_empty_device_cache, the
_generate/_run_qwen_zimage pass-through pair, self.model_id, and the _internal
PEP 562 shim that no caller ever went through. get_device now answers cuda or
cpu only: mps and xpu travelled one frame to the same CUDA-only refusal while
costing a device probe each, and that refusal now names the resolved device, so
device=None on a CUDA-less host says 'cpu' rather than 'None'. The XPU wheel
index went with them.

Docs: README, cli, installation, python-api, supported-signals,
known-limitations and module-internals all still described the removed profiles,
the CPU/MPS/XPU ladder, a `default`->`sdxl` alias, and the wrong extra.
known-limitations still listed the retired SDXL strength ladder as current.
scripts/smoke_matrix.py and real_examples_e2e.py drove --device mps.

Next release is 0.25.0, not a patch: this removes public parameters and
narrows a published extra on top of the released 0.24.0.

pre-commit: 1) maintain.sh - exit 0 (1091 tests, Pyright 0 errors, no
vulnerabilities); 2) /simplify - 4 agents, 11 findings applied, 2 skipped
(dropping the `device` parameter entirely, which raiw-app pins; folding
diffsynth into the `diffusion` extra, which video-only callers do not need);
3) docs sync - grepped every removed identifier across README, docs/, scripts/,
.claude/; updated 9 docs; 4) CLAUDE.md - added the no-error-only-knobs rule to
.claude/rules/development.md

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Victor Kuznetsov
2026-08-03 15:38:40 -07:00
co-authored by Claude Opus 5
parent bf4bfc1ab7
commit 52b2c115e8
28 changed files with 543 additions and 955 deletions
+2
View File
@@ -11,6 +11,8 @@ Every single-image command declares `source` with `dir_okay=False`; `batch` decl
Exit-code and no-signal behavior is a public contract. Read the command-line section of [`../../docs/module-internals.md`](../../docs/module-internals.md) before changing it. Exit-code and no-signal behavior is a public contract. Read the command-line section of [`../../docs/module-internals.md`](../../docs/module-internals.md) before changing it.
Do not add an option whose only outcome is an error. Model id, step count, CFG and any non-CUDA device are fixed by the profile, so none of them is a parameter of the CLI, `InvisibleEngine`, or `WatermarkRemover` -- they were accepted-then-rejected for a while, which moved the failure several frames below the caller and advertised choices the pinned stack cannot honor. If a value cannot vary, delete the knob rather than validating it. The same rule applies to install hints: name the extra that actually makes the command work (`qwen-zimage`, not `diffusion`).
## Local gate ## Local gate
Run `bash maintain.sh` from the repository root. The authoritative type gate is scoped to `src/`; full-project Pyright can exhaust Node memory on the ML dependency graph. Run `bash maintain.sh` from the repository root. The authoritative type gate is scoped to `src/`; full-project Pyright can exhaust Node memory on the ML dependency graph.
+14 -12
View File
@@ -52,11 +52,11 @@ removal.
| Visible video processing | `remove-ai-watermarks[video]` | | Visible video processing | `remove-ai-watermarks[video]` |
| Video SynthID removal | `remove-ai-watermarks[video,diffusion]` | | Video SynthID removal | `remove-ai-watermarks[video,diffusion]` |
| Torch-free DWT-DCT detection | `remove-ai-watermarks[detect]` | | Torch-free DWT-DCT detection | `remove-ai-watermarks[detect]` |
| Diffusion removal | `remove-ai-watermarks[diffusion]` | | Invisible image removal (needs CUDA) | `remove-ai-watermarks[qwen-zimage]` |
| Every production feature | `remove-ai-watermarks[all]` | | Every production feature | `remove-ai-watermarks[all]` |
Lower-level and specialized extras include `pixels`, `heif`, `trustmark`, Lower-level and specialized extras include `pixels`, `heif`, `trustmark`,
`migan`, `lama`, and `qwen-zimage`. The `migan`, `lama`, and `diffusion`. The
[installation guide](docs/installation.md#feature-extras) documents their exact [installation guide](docs/installation.md#feature-extras) documents their exact
dependency composition and model requirements. dependency composition and model requirements.
@@ -180,10 +180,11 @@ check. Google does not publish a local decoder, so a fresh provider check
remains useful for unusually important files or after provider changes, but it remains useful for unusually important files or after provider changes, but it
is not a product result state. is not a product result state.
For invisible watermark removal, install the diffusion dependencies: For invisible watermark removal, install the `qwen-zimage` extra. **An NVIDIA GPU
is required**: both profiles are CUDA-only, and there is no CPU or MPS fallback.
```bash ```bash
uv tool install --force "remove-ai-watermarks[diffusion]" uv tool install --force "remove-ai-watermarks[qwen-zimage]"
remove-ai-watermarks invisible image.png -o clean.png remove-ai-watermarks invisible image.png -o clean.png
``` ```
@@ -207,14 +208,14 @@ features, and development setup.
### High quality invisible removal ### High quality invisible removal
The `qwen-zimage` profile is the highest fidelity option for face heavy images. `qwen-zimage` is the default profile: a Qwen-Image-2512 Lightning pass under Canny
It is CUDA only and uses a much larger model stack than the default ControlNet ControlNet, followed by SAM-masked Z-Image repair of any detected face. The
profile. alternative, `sdxl-zimage`, swaps the global stage for SDXL and keeps the same face
stage. Both are CUDA only.
```bash ```bash
uv tool install --force "remove-ai-watermarks[qwen-zimage]" uv tool install --force "remove-ai-watermarks[qwen-zimage]"
remove-ai-watermarks invisible image.png -o clean.png \ remove-ai-watermarks invisible image.png -o clean.png --force
--pipeline qwen-zimage --force
``` ```
| OpenAI example before | OpenAI example after | | OpenAI example before | OpenAI example after |
@@ -273,7 +274,7 @@ remove-ai-watermarks invisible image.png -o clean.png \
``` ```
CPU offload lowers CUDA memory pressure by moving model components between CPU CPU offload lowers CUDA memory pressure by moving model components between CPU
and GPU. It is slower and has no effect on CPU or MPS. and GPU, at the cost of speed.
### Process a directory ### Process a directory
@@ -378,8 +379,9 @@ invisible removal.
The shipped profile is oracle-certified, but no public local decoder can The shipped profile is oracle-certified, but no public local decoder can
certify an arbitrary output at runtime. Recheck unusually important outputs certify an arbitrary output at runtime. Recheck unusually important outputs
after provider changes. after provider changes.
- `qwen-zimage` requires CUDA. The other diffusion profiles also support the - Invisible-watermark removal requires CUDA. Both profiles refuse any other
devices listed by `remove-ai-watermarks invisible --help`. device at construction rather than falling back to one that cannot run them.
Visible removal, metadata stripping and `identify` still run anywhere.
- Provider watermark systems can change. Validate important outputs with the - Provider watermark systems can change. Validate important outputs with the
provider's own verifier when one is available. provider's own verifier when one is available.
+10 -7
View File
@@ -20,8 +20,7 @@ defaults. This page focuses on choosing the right command.
| `visible` and `erase` with OpenCV | `remove-ai-watermarks[visible]` (`pixels` is the minimal runtime) | | `visible` and `erase` with OpenCV | `remove-ai-watermarks[visible]` (`pixels` is the minimal runtime) |
| `visible` or `erase` with MI-GAN | `remove-ai-watermarks[migan]` | | `visible` or `erase` with MI-GAN | `remove-ai-watermarks[migan]` |
| `visible` or `erase` with big-LaMa | `remove-ai-watermarks[lama]` | | `visible` or `erase` with big-LaMa | `remove-ai-watermarks[lama]` |
| `invisible` | `remove-ai-watermarks[diffusion]` | | `invisible` and `all` (needs CUDA) | `remove-ai-watermarks[qwen-zimage]` |
| `invisible --pipeline qwen-zimage` | `remove-ai-watermarks[qwen-zimage]` |
| `video metadata` and `video identify --no-visible` | Default package | | `video metadata` and `video identify --no-visible` | Default package |
| `video identify`, `video visible`, and visible/all batch modes | `remove-ai-watermarks[video]` | | `video identify`, `video visible`, and visible/all batch modes | `remove-ai-watermarks[video]` |
| `video invisible` and `video all --invisible` | `remove-ai-watermarks[video,diffusion]` | | `video invisible` and `video all --invisible` | `remove-ai-watermarks[video,diffusion]` |
@@ -340,10 +339,11 @@ failed encode does not overwrite an existing result.
## Remove invisible watermarks ## Remove invisible watermarks
Install the diffusion dependencies first: Install the removal dependencies first. Both profiles are CUDA-only and both
run the DiffSynth Z-Image face stage, so this is the extra either one needs:
```bash ```bash
uv tool install --force "remove-ai-watermarks[diffusion]" uv tool install --force "remove-ai-watermarks[qwen-zimage]"
``` ```
Then run: Then run:
@@ -380,8 +380,11 @@ remove-ai-watermarks invisible image.png -o clean.png \
--pipeline qwen-zimage --force --pipeline qwen-zimage --force
``` ```
The legacy `default` value is an alias for `sdxl`. The `--auto` option is There is no `--model`, `--steps`, `--guidance-scale` or `--device` option, and the
deprecated, emits a warning, and changes nothing. deprecated `--auto` is gone. Each profile pins its model stack, its per-stage
schedule, CFG 1.0 and CUDA, so every one of those flags existed only to be refused
several layers down. They are not parsed at all now, which fails at the point the
user can act on rather than after a model load.
### Work with limited memory ### Work with limited memory
@@ -415,7 +418,7 @@ It is a memory strategy, not a guarantee of better quality.
The `all` command and the `all` installation extra are separate concepts. The The `all` command and the `all` installation extra are separate concepts. The
command runs every applicable stage. Installing `remove-ai-watermarks[all]` command runs every applicable stage. Installing `remove-ai-watermarks[all]`
makes every production backend available; a smaller installation such as makes every production backend available; a smaller installation such as
`remove-ai-watermarks[visible,diffusion]` can also run the command with fewer `remove-ai-watermarks[visible,qwen-zimage]` can also run the command with fewer
optional backends. optional backends.
```bash ```bash
+19 -15
View File
@@ -65,22 +65,24 @@ uv tool install --force "remove-ai-watermarks[video,diffusion]"
## Invisible watermark removal ## Invisible watermark removal
Diffusion based removal needs the `diffusion` extra: Install the `qwen-zimage` extra:
```bash
uv tool install --force "remove-ai-watermarks[diffusion]"
```
The code supports CUDA, XPU, MPS, and CPU devices. A GPU is recommended because
CPU inference is slow.
For the CUDA only Qwen Image plus Z-Image profile:
```bash ```bash
uv tool install --force "remove-ai-watermarks[qwen-zimage]" uv tool install --force "remove-ai-watermarks[qwen-zimage]"
``` ```
The `qwen-zimage` extra includes the normal `diffusion` dependencies. Both remaining profiles run a Z-Image face stage on the DiffSynth runtime, so
both need this extra. It includes the `diffusion` dependencies; `diffusion` on its
own covers the torch and diffusers imports but not the face stage, so it is not
enough to run a removal.
**An NVIDIA GPU is required.** `qwen-zimage` and `sdxl-zimage` are CUDA-only, and
construction refuses any other device rather than falling back to a slow or broken
one. There is no CPU, MPS or XPU path for invisible-watermark removal. Visible-mark
removal, metadata stripping and every `identify` command still run anywhere.
Video SynthID regeneration is a separate VAE path and does still run on CPU or MPS;
it needs the `diffusion` extra, not this one.
## Feature extras ## Feature extras
@@ -95,10 +97,10 @@ application actually uses:
| `video` | Visible video identification/removal and timestamp preservation | `visible`, PyAV | No | | `video` | Visible video identification/removal and timestamp preservation | `visible`, PyAV | No |
| `detect` | Open DWT-DCT detection for Stable Diffusion, SDXL, and FLUX | `pixels`, PyWavelets | No | | `detect` | Open DWT-DCT detection for Stable Diffusion, SDXL, and FLUX | `pixels`, PyWavelets | No |
| `trustmark` | Adobe TrustMark detection | trustmark | Yes | | `trustmark` | Adobe TrustMark detection | trustmark | Yes |
| `diffusion` | Diffusion-based invisible watermark removal | `pixels`, Torch, Diffusers | Yes | | `diffusion` | Torch and Diffusers runtime; video SynthID regeneration | `pixels`, Torch, Diffusers | Yes |
| `migan` | MI-GAN ONNX fill backend | `visible`, ONNX Runtime | Model download, no Torch | | `migan` | MI-GAN ONNX fill backend | `visible`, ONNX Runtime | Model download, no Torch |
| `lama` | big-LaMa ONNX fill backend | `visible`, ONNX Runtime | Model download, no Torch | | `lama` | big-LaMa ONNX fill backend | `visible`, ONNX Runtime | Model download, no Torch |
| `qwen-zimage` | CUDA-only Qwen Image plus Z-Image pipeline | `diffusion`, DiffSynth | Yes | | `qwen-zimage` | Invisible image-watermark removal, both CUDA-only profiles | `diffusion`, DiffSynth | Yes |
| `all` | Every production feature | All rows above | Yes | | `all` | Every production feature | All rows above | Yes |
| `dev` | Tests, linting, typing, and upstream parity checks | `visible`, `detect`, upstream invisible-watermark | Yes, for parity tests | | `dev` | Tests, linting, typing, and upstream parity checks | `visible`, `detect`, upstream invisible-watermark | Yes, for parity tests |
@@ -214,5 +216,7 @@ The normal behavior is to skip diffusion when no supported local signal is
found. A missing signal does not prove that the image is clean. If you know the found. A missing signal does not prove that the image is clean. If you know the
image came from a relevant generator, use `--force`. image came from a relevant generator, use `--force`.
If the CLI reports that diffusion dependencies are unavailable, install the If the CLI reports that the removal dependencies are unavailable, install the
`diffusion` extra. Video SynthID removal needs both `video` and `diffusion`. `qwen-zimage` extra. `diffusion` alone covers Torch and Diffusers but not the
DiffSynth face stage that both profiles run. Video SynthID removal is a separate
path and needs `video` and `diffusion`.
+26 -17
View File
@@ -120,33 +120,42 @@ prefix so it can reuse identical latents across candidate strengths.
### Strength is content and seed dependent ### Strength is content and seed dependent
For SDXL and ControlNet, the CLI resolves an unset strength from the detected The two profiles resolve an unset strength differently, because different things
vendor: were measured for each.
- OpenAI: `0.10`; `qwen-zimage` reads it from image area, through the resolution-adaptive denoise
- Google: `0.15`; curve. The vendor is deliberately ignored: the curve, not the issuer, is what was
- unknown: `0.15`. calibrated.
An explicit `--strength` overrides these defaults. The defaults are operating `sdxl-zimage` reads it from the C2PA issuer, on a flat ladder:
points, not universal guarantees. Near a removal threshold, different content
or a different random seed may change the verifier result.
The base Qwen and `qwen-zimage` profiles have profile specific strength - OpenAI: `0.15`;
behavior. Consult `remove-ai-watermarks invisible --help` and the source of - Google: `0.25`;
[`watermark_profiles.py`](../src/remove_ai_watermarks/_internal/watermark_profiles.py) - unknown: `0.25`, following the stricter of the two.
for the current resolver.
An SDXL global pass needs more denoise than Qwen at the same fidelity, and the
values are flat rather than a curve because flat values are what was measured: each
verdict came from a fixed strength at one size, and no size dependence has been
established for that stage.
An explicit `--strength` overrides both. The defaults are operating points, not
universal guarantees. Near a removal threshold, different content or a different
random seed may change the verifier result, which is why both profiles are
certified at a fixed seed. The live resolver is
[`watermark_profiles.py`](../src/remove_ai_watermarks/_internal/watermark_profiles.py).
### Pipelines have different quality tradeoffs ### Pipelines have different quality tradeoffs
| Pipeline | Main limit | | Pipeline | Main limit |
| --- | --- | | --- | --- |
| `controlnet` | Edge conditioning can preserve a watermark carrying region too closely, and faces may drift. |
| `sdxl` | Flat graphics and precise structure may receive too little or unhelpful change. |
| `qwen` | Large CUDA oriented model; face smoothing can still be significant. |
| `qwen-zimage` | CUDA only, large model stack, and limited broad certification across seeds and content. | | `qwen-zimage` | CUDA only, large model stack, and limited broad certification across seeds and content. |
| `sdxl-zimage` | CUDA only. Its strength ladder is flat per vendor, not a resolution curve, because flat values are what was measured. |
The legacy `default` profile name maps to `sdxl`. The `--auto` flag is The `controlnet`, `sdxl`, `qwen` and `default` profiles were removed, not aliased
deprecated, emits a warning, and changes nothing. onward: a retired name is rejected at parse time rather than routed into a profile
the caller never chose. There is no `--model`, `--steps`, `--guidance-scale`,
`--device` or `--auto` option either; each profile pins its model stack, its
per-stage schedule, CFG 1.0 and CUDA.
## Resolution and memory ## Resolution and memory
+45 -17
View File
@@ -53,9 +53,19 @@ The decorators for diffusion options are shared by `invisible`, `all`, and
`batch`. The runtime help generated by Click is the source of truth for option `batch`. The runtime help generated by Click is the source of truth for option
names and defaults. names and defaults.
The deprecated `--auto` option does not select a pipeline or change adaptive `--adaptive-polish` is tri-state: it declares `default=None`, so "the user did not
polishing. [`_resolve_auto_polish`](../src/remove_ai_watermarks/cli.py) emits a choose" is a value the CLI passes through rather than a default it has to invent.
warning and returns the explicit polish value unchanged. `resolve_adaptive_polish` in `watermark_profiles.py` turns that `None` into the
profile's answer (off for `qwen-zimage`, whose output already matches the input's
detail level; on for `sdxl-zimage`). The same call runs inside
`InvisibleEngine.remove_watermark`, so a library caller and a CLI caller on one
profile get the same output.
It used to read Click's parameter source in the CLI instead. That put per-profile
data in the argument-parsing layer, left the engine declaring the opposite default,
and silently lost the polish for anything supplying the flag non-interactively (an
envvar default or a wrapper calling `main()` with a defaulted list is classified
`DEFAULT`). The seed follows the same rule: the CLI does not pre-resolve it either.
Regression coverage: Regression coverage:
@@ -477,25 +487,42 @@ Regression coverage:
[`_internal/watermark_profiles.py`](../src/remove_ai_watermarks/_internal/watermark_profiles.py) [`_internal/watermark_profiles.py`](../src/remove_ai_watermarks/_internal/watermark_profiles.py)
is the source of truth for: is the source of truth for:
- profile aliases; - profile names and their underscore spellings;
- default model identifiers; - the fixed seed;
- default steps and seeds; - the SDXL global-stage checkpoint id (`SDXL_MODEL_ID`) and the Canny ControlNet id;
- vendor-adaptive strength resolution; - strength resolution for both profiles.
- the minimum viable step calculation.
The current profiles are `qwen-zimage` (the default) and `sdxl-zimage`, and both
are CUDA-only. `controlnet`, `sdxl`, `qwen` and `default` were removed rather than
kept as a CPU path, and are rejected rather than aliased onward. There is no
content-dependent automatic router.
The current profiles are `controlnet`, `sdxl`, `qwen`, and `qwen-zimage`.
For serverless cold starts, `InvisibleEngine.preload(global_only=True)` loads the For serverless cold starts, `InvisibleEngine.preload(global_only=True)` loads the
mandatory Qwen stage and YuNet while leaving the optional Z-Image and SAM face mandatory global stage and YuNet while leaving the optional Z-Image and SAM face
stack lazy until a face is detected. The default `preload()` still loads every stack lazy until a face is detected. The default `preload()` still loads every
stage. stage.
`default` is a legacy alias for `sdxl`. There is no content-dependent automatic
router. **What is deliberately not a parameter.** Model id, step count, CFG and any
non-CUDA device are fixed by the profile, so none of them appears in
`WatermarkRemover.__init__`, `remove_watermark`, `InvisibleEngine`, or the CLI.
They used to be accepted and then rejected several frames down; a signature that
refuses the argument outright fails where the caller can act on it, and stops a
wrapper from threading a value that would silently do nothing. The step count and
CFG live with the stage that runs them (`GLOBAL_STEPS`, `FACE_STEPS`, `GLOBAL_CFG`,
`FACE_CFG` in `qwen_zimage_pipeline.py`). The dtype is likewise profile-owned: see
"Face-stage dtype" for what an override cost the last time one existed.
[`invisible_engine.py`](../src/remove_ai_watermarks/invisible_engine.py) handles [`invisible_engine.py`](../src/remove_ai_watermarks/invisible_engine.py) handles
image sizing, postprocessing, and the public engine image sizing, postprocessing, and the public engine
interface. It delegates model execution to interface. It delegates model execution to
[`_internal/watermark_remover.py`](../src/remove_ai_watermarks/_internal/watermark_remover.py). [`_internal/watermark_remover.py`](../src/remove_ai_watermarks/_internal/watermark_remover.py).
`get_device` in that module answers only `cuda` or `cpu`. An `mps` or `xpu` answer
would travel one frame to the same CUDA-only refusal while costing a device probe,
and reporting it implied an Apple-silicon or Intel-GPU path that does not exist.
The refusal names the *resolved* device, so `device=None` on a CUDA-less host says
`'cpu'` rather than `'None'`.
The Python engine and CLI do not have identical defaults for every optional The Python engine and CLI do not have identical defaults for every optional
postprocessing argument. Integrations that require reproducibility should pass postprocessing argument. Integrations that require reproducibility should pass
the relevant values explicitly. the relevant values explicitly.
@@ -509,9 +536,8 @@ unit-test pass. Exact prompt and edge-map regression guards live in
Regression coverage: Regression coverage:
- [`test_watermark_profiles.py`](../tests/test_watermark_profiles.py)
- [`test_invisible_engine.py`](../tests/test_invisible_engine.py) - [`test_invisible_engine.py`](../tests/test_invisible_engine.py)
- [`test_img2img_runner.py`](../tests/test_img2img_runner.py) - [`test_qwen_zimage_pipeline.py`](../tests/test_qwen_zimage_pipeline.py)
- [`test_platform.py`](../tests/test_platform.py) - [`test_platform.py`](../tests/test_platform.py)
### CPU offload ### CPU offload
@@ -712,14 +738,16 @@ Regression coverage:
### Tiling ### Tiling
[`_internal/tiling.py`](../src/remove_ai_watermarks/_internal/tiling.py) contains pure [`_internal/tiling.py`](../src/remove_ai_watermarks/_internal/tiling.py) contains pure
tile planning, feather weights, tile orchestration, and region compositing. tile planning, feather weights, and tile orchestration.
Tiling engages only when requested and the long side exceeds the tile size. Tiling engages only when requested and the long side exceeds the tile size.
It avoids an explicit full-image downscale but does not make diffusion It avoids an explicit full-image downscale but does not make diffusion
pixel-preserving. Each tile is still regenerated. pixel-preserving. Each tile is still regenerated.
`feather_region_composite` changes only the requested box and leaves pixels It also held a `feather_region_composite` for AI-*enhanced* composites, where only
outside it unchanged. the edited region should change. Nothing ever reached it: the `erase` command
inpaints through `region_eraser`, and the remover's `region` argument was only
reachable from a module-level convenience wrapper with no callers. Both went.
Regression coverage: Regression coverage:
+18 -12
View File
@@ -6,8 +6,9 @@ and pipeline modules are intended for maintainers and specialized workflows.
Dependency groups are identical for the CLI and Python API. The default install Dependency groups are identical for the CLI and Python API. The default install
covers metadata extraction, normalization, verdict logic, and stripping. covers metadata extraction, normalization, verdict logic, and stripping.
Array/pixel APIs use `pixels`; visible removal uses `visible`; DWT-DCT detection Array/pixel APIs use `pixels`; visible removal uses `visible`; DWT-DCT detection
uses `detect`; diffusion removal uses `diffusion`; and visible video processing uses `detect`; invisible image removal uses `qwen-zimage` and an NVIDIA GPU; and
uses `video`. Video SynthID removal combines `video` and `diffusion`. Add `heif` visible video processing uses `video`. Video SynthID removal is a separate VAE
path that still runs on CPU and combines `video` and `diffusion`. Add `heif`
independently when path-based pixel APIs must decode HEIC, HEIF, or AVIF. See independently when path-based pixel APIs must decode HEIC, HEIF, or AVIF. See
the complete [feature-extra matrix](installation.md#feature-extras). the complete [feature-extra matrix](installation.md#feature-extras).
@@ -380,8 +381,8 @@ to 8-bit SDR.
## Remove invisible watermarks ## Remove invisible watermarks
Install `remove-ai-watermarks[diffusion]` for the standard pipelines or Install `remove-ai-watermarks[qwen-zimage]`. Both profiles need it, and both
`remove-ai-watermarks[qwen-zimage]` for the CUDA-only high-fidelity profile. need an NVIDIA GPU.
```python ```python
from pathlib import Path from pathlib import Path
@@ -400,8 +401,9 @@ engine.remove_watermark(
) )
``` ```
`device=None` selects the device automatically. Supported explicit values are `device=None` detects CUDA. The only other accepted value is `"cuda"`; anything
defined by the CLI and runtime device resolver. else raises at construction rather than deferring a guaranteed failure to model
load time.
For limited CUDA memory: For limited CUDA memory:
@@ -412,18 +414,22 @@ engine = InvisibleEngine(
) )
``` ```
Both profiles are CUDA-only, so `device=None` resolving to CPU or MPS cannot run Both profiles are CUDA-only, so on a machine without an NVIDIA GPU `device=None`
invisible-watermark removal at all. For the SDXL global stage instead of Qwen: resolves to `cpu` and construction raises. For the SDXL global stage instead of
Qwen:
```python ```python
engine = InvisibleEngine(pipeline="sdxl-zimage") engine = InvisibleEngine(pipeline="sdxl-zimage")
``` ```
The `qwen-zimage` extra must be installed for that profile. The `qwen-zimage` extra is required for both profiles: each runs the same
DiffSynth Z-Image face stage.
The full `remove_watermark` signature includes strength, steps, guidance, `remove_watermark` takes strength, seed, tiling, resolution, and postprocessing
seeding, tiling, resolution, and postprocessing controls. Read the controls. It takes no model id, step count or guidance scale, and neither does the
method signature in constructor: each profile pins its model stack, its per-stage schedule and CFG
1.0, so passing one raises `TypeError` at the call rather than being accepted and
refused several layers down. Read the method signature in
[`invisible_engine.py`](../src/remove_ai_watermarks/invisible_engine.py) or use [`invisible_engine.py`](../src/remove_ai_watermarks/invisible_engine.py) or use
the CLI guide for the concepts. the CLI guide for the concepts.
Defaults can differ between the Python method and CLI profile resolution, so Defaults can differ between the Python method and CLI profile resolution, so
+6 -6
View File
@@ -114,13 +114,13 @@ The `invisible` command uses diffusion regeneration. It targets watermark
patterns by changing the image rather than decoding and deleting a known patterns by changing the image rather than decoding and deleting a known
payload. payload.
Current pipeline values: Current pipeline values, both CUDA-only:
- `controlnet`; - `qwen-zimage`, the default;
- `sdxl`; - `sdxl-zimage`, the same recipe and the same face stage on an SDXL global pass.
- `qwen`;
- `qwen-zimage`; The `controlnet`, `sdxl`, `qwen` and `default` values were removed. A retired name
- legacy alias `default`, which resolves to `sdxl`. is rejected at parse time rather than remapped onto a surviving profile.
SynthID does not have a public local pixel decoder in this project. The tool can SynthID does not have a public local pixel decoder in this project. The tool can
infer likely presence from supported provenance metadata, but after that infer likely presence from supported provenance metadata, but after that
+5 -4
View File
@@ -714,10 +714,11 @@ end has simply never been through the Gemini oracle on any pipeline. Do not reas
a resolution trend here; measure it. a resolution trend here; measure it.
**Current implication:** the old floor table remains evidence about the dated **Current implication:** the old floor table remains evidence about the dated
test set, not the current resolver. The shipped SDXL and ControlNet defaults are test set, not the current resolver. The SDXL and ControlNet profiles it measured
defined in `watermark_profiles.py`, and face restoration is available only no longer exist; the shipped defaults are defined in `watermark_profiles.py`, and
through the separate `qwen-zimage` profile. Removal near a threshold remains both surviving profiles run face repair as a built-in second stage rather than as
seed dependent, so reproducible verification requires a fixed seed. an optional restore. Removal near a threshold remains seed dependent, so
reproducible verification requires a fixed seed.
--- ---
+14 -8
View File
@@ -91,20 +91,26 @@ and the ffmpeg audio/video strip. The gap to find is not only
"logic untested" but "logic untested" but
"never executed on real data", which is precisely what this campaign is for. "never executed on real data", which is precisely what this campaign is for.
#### Bug found by the extension: `--steps` below ~7 crashes inside torch #### Bug found by the extension: `--steps` below ~7 crashed inside torch
Effective timesteps are `int(steps * strength)`. At the vendor-adaptive default strength **Fixed by deletion.** `--steps` no longer exists, on the CLI or in the Python API,
(0.15, or 0.10 for OpenAI) any `--steps` under 7 rounds to **zero**, and the pipeline dies so this class of failure is unreachable. Kept as a record of why.
with a raw traceback:
Effective timesteps were `int(steps * strength)`. At the vendor-adaptive default
strength (0.15, or 0.10 for OpenAI) any `--steps` under 7 rounded to **zero**, and the
pipeline died with a raw traceback:
``` ```
$ remove-ai-watermarks invisible img.png --steps 5
RuntimeError: cannot reshape tensor of 0 elements into shape [0, -1, 1, 512] RuntimeError: cannot reshape tensor of 0 elements into shape [0, -1, 1, 512]
``` ```
Fully valid CLI arguments, no special flags, no `--force`. The value is accepted, the Fully valid CLI arguments, no special flags, no `--force`. The value was accepted, the
crash is a torch internal, and nothing tells the user that steps and strength interact. crash was a torch internal, and nothing told the user that steps and strength interact.
Fix is either a clamp to >=1 effective step or an up-front validation naming both values. The considered fixes were a clamp to >=1 effective step or an up-front validation
naming both values; what shipped instead is that each stage owns its own distilled
schedule and no caller can set it. The general lesson stands: a knob whose valid range
depends on another knob's value needs the interaction validated where both are known,
or it needs to not be a knob.
Method note: the first run of the knob rows failed 12 times with this identical error, Method note: the first run of the knob rows failed 12 times with this identical error,
which read like twelve broken features. It was one bad harness parameter (`--steps 4`) which read like twelve broken features. It was one bad harness parameter (`--steps 4`)
+7 -19
View File
@@ -86,16 +86,11 @@ detect = [
] ]
diffusion = [ diffusion = [
"remove-ai-watermarks[pixels]", "remove-ai-watermarks[pixels]",
# A CUDA-enabled torch build is required: invisible-watermark removal has no
# CPU, MPS or XPU path. The default PyPI wheel carries CUDA on Linux/Windows;
# on macOS there is no CUDA build and this extra installs only for the
# non-diffusion imports it shares.
"torch>=2.0.0", "torch>=2.0.0",
# The default PyPI torch wheel is a CPU/CUDA build. To drive an Intel GPU
# (Arc / Data Center) via ``--device xpu`` you need an XPU-enabled torch
# from PyTorch's XPU wheel index (Linux/Windows only -- there is no macOS
# XPU build). Install that build first, then this extra (torch is then
# already satisfied and won't be re-pulled):
# pip install torch --index-url https://download.pytorch.org/whl/xpu
# pip install 'remove-ai-watermarks[diffusion]'
# uv users can target the ``pytorch-xpu`` index declared under [tool.uv]:
# uv pip install torch --index-url https://download.pytorch.org/whl/xpu
"diffusers>=0.38.0", "diffusers>=0.38.0",
# diffusers 0.38's auto-pipeline registry imports ``Qwen3VLForConditional # diffusers 0.38's auto-pipeline registry imports ``Qwen3VLForConditional
# Generation`` (its ``nucleusmoe_image`` pipeline), which only exists in # Generation`` (its ``nucleusmoe_image`` pipeline), which only exists in
@@ -165,16 +160,9 @@ dev = [
"uv-outdated>=0.1.0; python_version >= '3.12'", "uv-outdated>=0.1.0; python_version >= '3.12'",
"uv-secure>=0.12.0; python_version >= '3.12'", "uv-secure>=0.12.0; python_version >= '3.12'",
] ]
all = ["remove-ai-watermarks[video,heif,detect,trustmark,diffusion,qwen-zimage,lama,migan]"] # ``qwen-zimage`` already pulls ``diffusion``; naming both would suggest diffusion is
# independently sufficient for a removal, which it is not.
# PyTorch Intel-GPU (XPU) wheel index. ``explicit = true`` keeps it inert for all = ["remove-ai-watermarks[video,heif,detect,trustmark,qwen-zimage,lama,migan]"]
# the default CPU/CUDA install: uv consults it only when a torch install
# explicitly targets it (see the ``diffusion`` extra comment), so it does not alter
# the locked CPU/CUDA resolution. Linux/Windows only -- no macOS XPU build.
[[tool.uv.index]]
name = "pytorch-xpu"
url = "https://download.pytorch.org/whl/xpu"
explicit = true
[project.scripts] [project.scripts]
remove-ai-watermarks = "remove_ai_watermarks.cli:main" remove-ai-watermarks = "remove_ai_watermarks.cli:main"
+14 -10
View File
@@ -13,20 +13,20 @@ WHAT IT COVERS
metadata real AI-metadata files -> --check detects, --remove strip-and-verifies clean metadata real AI-metadata files -> --check detects, --remove strip-and-verifies clean
visible real marked images per mark -> the mark is gone on re-detect, output written visible real marked images per mark -> the mark is gone on re-detect, output written
erase a real image, each fill backend (cv2 / migan / lama) -> output written erase a real image, each fill backend (cv2 / migan / lama) -> output written
invisible a real SynthID image on MPS at reduced resolution -> a CHANGED image is written invisible a real SynthID image at reduced resolution -> a CHANGED image is written
all a real marked image through the full pipeline -> output written all a real marked image through the full pipeline -> output written
batch a real directory -> every input produces an output batch a real directory -> every input produces an output
invisible/all run the diffusion model, so they are gated behind --diffusion and run at a invisible/all run the diffusion model, so they are gated behind --diffusion and run at a
small --max-resolution on MPS (the user's "reduced size on MPS" path). Everything else is small --max-resolution. Both profiles are CUDA-only, so that section needs an NVIDIA
cv2/numpy and fast. GPU; everything else is cv2/numpy, fast, and runs anywhere.
DATA SAFETY DATA SAFETY
Treat input datasets as sensitive and read-only. Output stays in a gitignored temp Treat input datasets as sensitive and read-only. Output stays in a gitignored temp
dir. Records example uids and pass/fail, never image content. dir. Records example uids and pass/fail, never image content.
uv run python scripts/real_examples_e2e.py # fast surface (no diffusion) uv run python scripts/real_examples_e2e.py # fast surface (no diffusion)
uv run python scripts/real_examples_e2e.py --diffusion # + invisible/all on MPS uv run python scripts/real_examples_e2e.py --diffusion # + invisible/all (needs CUDA)
""" """
from __future__ import annotations from __future__ import annotations
@@ -239,12 +239,16 @@ def check_erase(res: Results, tmp: Path) -> None:
def check_diffusion(res: Results, tmp: Path, gemini_src: str, openai_src: str) -> None: def check_diffusion(res: Results, tmp: Path, gemini_src: str, openai_src: str) -> None:
"""The GPU path: invisible + all on MPS at reduced resolution -> a CHANGED image.""" """The GPU path: invisible + all at reduced resolution -> a CHANGED image.
CUDA-only. On a machine without an NVIDIA GPU every row here fails with the
library's clean refusal; run this section on a GPU box.
"""
import numpy as np import numpy as np
from remove_ai_watermarks.image_io import imread from remove_ai_watermarks.image_io import imread
print("\ninvisible / all -- real SynthID image on MPS, reduced resolution") print("\ninvisible / all -- real SynthID image, reduced resolution (CUDA required)")
for label, src in (("invisible/gemini", gemini_src), ("invisible/openai", openai_src)): for label, src in (("invisible/gemini", gemini_src), ("invisible/openai", openai_src)):
if not src: if not src:
res.add(label, "-", True, "no real positive found (skipped)") res.add(label, "-", True, "no real positive found (skipped)")
@@ -252,7 +256,7 @@ def check_diffusion(res: Results, tmp: Path, gemini_src: str, openai_src: str) -
sp = Path(src) sp = Path(src)
outp = tmp / f"inv_{sp.stem}.png" outp = tmp / f"inv_{sp.stem}.png"
code, out = run( code, out = run(
["invisible", src, "-o", str(outp), "--device", "mps", "--max-resolution", "512", "--seed", "0"], ["invisible", src, "-o", str(outp), "--max-resolution", "512", "--seed", "0"],
timeout=1200, timeout=1200,
) )
if not outp.exists(): if not outp.exists():
@@ -276,7 +280,7 @@ def check_diffusion(res: Results, tmp: Path, gemini_src: str, openai_src: str) -
sp = Path(gemini_src) sp = Path(gemini_src)
outp = tmp / f"all_{sp.stem}.png" outp = tmp / f"all_{sp.stem}.png"
code, out = run( code, out = run(
["all", gemini_src, "-o", str(outp), "--device", "mps", "--max-resolution", "512", "--seed", "0"], ["all", gemini_src, "-o", str(outp), "--max-resolution", "512", "--seed", "0"],
timeout=1200, timeout=1200,
) )
ok = outp.exists() and outp.stat().st_size > 0 ok = outp.exists() and outp.stat().st_size > 0
@@ -304,7 +308,7 @@ def check_batch(res: Results, tmp: Path) -> None:
def main() -> None: def main() -> None:
ap = argparse.ArgumentParser() ap = argparse.ArgumentParser()
ap.add_argument("--diffusion", action="store_true", help="also run invisible/all on MPS (slow)") ap.add_argument("--diffusion", action="store_true", help="also run invisible/all (slow, needs CUDA)")
ap.add_argument("--gemini", default="") ap.add_argument("--gemini", default="")
ap.add_argument("--openai", default="") ap.add_argument("--openai", default="")
a = ap.parse_args() a = ap.parse_args()
@@ -320,7 +324,7 @@ def main() -> None:
if a.diffusion: if a.diffusion:
check_diffusion(res, tmp, a.gemini, a.openai) check_diffusion(res, tmp, a.gemini, a.openai)
else: else:
print("\n(diffusion skipped -- pass --diffusion to run invisible/all on MPS)") print("\n(diffusion skipped -- pass --diffusion to run invisible/all)")
raise SystemExit(res.report()) raise SystemExit(res.report())
+23 -23
View File
@@ -17,9 +17,10 @@ WHAT IT COVERS AND WHY THAT SHAPE
work" report); a no-op must be byte-identical; `metadata --remove` must actually work" report); a no-op must be byte-identical; `metadata --remove` must actually
strip; a JPEG strip must not touch the pixels. Note the pixel-lossless contract is strip; a JPEG strip must not touch the pixels. Note the pixel-lossless contract is
the DEFAULT path's -- `--remove-all` deliberately re-encodes (see metadata.py). the DEFAULT path's -- `--remove-all` deliberately re-encodes (see metadata.py).
* The diffusion bodies under `--diffusion`, at `--max-resolution 512` so they fit MPS * The diffusion bodies under `--diffusion`, at a small `--max-resolution`. Both
(~1 min/image on 32 GB unified memory). Not just exit codes: `invisible` must profiles are CUDA-only, so those rows need an NVIDIA GPU and are reported as
restore the input resolution and must NOT re-stamp SDXL's own open watermark. skips elsewhere. Not just exit codes: `invisible` must restore the input
resolution and must NOT re-stamp SDXL's own open watermark.
WHAT IT DOES NOT COVER, DELIBERATELY AND LOUDLY WHAT IT DOES NOT COVER, DELIBERATELY AND LOUDLY
* Without `--diffusion`, the model-running bodies are reported as SKIPPED with a * Without `--diffusion`, the model-running bodies are reported as SKIPPED with a
@@ -34,7 +35,7 @@ WHAT IT DOES NOT COVER, DELIBERATELY AND LOUDLY
uv run python scripts/smoke_matrix.py # corpus + fixtures uv run python scripts/smoke_matrix.py # corpus + fixtures
uv run python scripts/smoke_matrix.py --quick # fixtures only, no corpus uv run python scripts/smoke_matrix.py --quick # fixtures only, no corpus
uv run python scripts/smoke_matrix.py --diffusion # + the SDXL model paths uv run python scripts/smoke_matrix.py --diffusion # + the model paths (needs CUDA)
""" """
from __future__ import annotations from __future__ import annotations
@@ -113,7 +114,7 @@ def main() -> None:
ap = argparse.ArgumentParser() ap = argparse.ArgumentParser()
ap.add_argument("--quick", action="store_true", help="fixtures only; skip the corpus rows") ap.add_argument("--quick", action="store_true", help="fixtures only; skip the corpus rows")
ap.add_argument( ap.add_argument(
"--diffusion", action="store_true", help="also run the model-running paths (SDXL weights, ~1 min/image)" "--diffusion", action="store_true", help="also run the model-running paths (needs CUDA, ~1 min/image)"
) )
a = ap.parse_args() a = ap.parse_args()
@@ -359,10 +360,12 @@ def _knob_rows(r: Runner, tmp: Path, img: Path) -> None:
per-knob oracle and most of them have none (`--humanize` has no oracle at all), so per-knob oracle and most of them have none (`--humanize` has no oracle at all), so
claiming more here would be dishonest. claiming more here would be dishonest.
""" """
# Both surviving profiles are CUDA-only and pin a distilled four-step schedule at # Both surviving profiles are CUDA-only, pin a fixed model stack and let each stage
# CFG 1.0, so most rows here now assert a knob is REJECTED rather than accepted. # own its schedule and CFG. The knobs that used to contradict that (--model, --steps,
# That is the coverage worth having: a knob the CLI takes and the library refuses # --guidance-scale, --device, --auto) are gone from the parser, so the rows below
# several layers down is exactly what this matrix exists to catch. # assert Click itself refuses them. That is the coverage worth having: an option the
# CLI accepts and the library refuses several layers down is what this matrix exists
# to catch, and the cheapest way to keep it caught is for the option not to exist.
fast = ["--max-resolution", "384", "--force", "--seed", "0"] fast = ["--max-resolution", "384", "--force", "--seed", "0"]
def run(name: str, extra: list[str], *, tag: str, expect: int | None = 0) -> None: def run(name: str, extra: list[str], *, tag: str, expect: int | None = 0) -> None:
@@ -377,10 +380,15 @@ def _knob_rows(r: Runner, tmp: Path, img: Path) -> None:
for retired in ("sdxl", "controlnet", "qwen", "default"): for retired in ("sdxl", "controlnet", "qwen", "default"):
run(f"--pipeline {retired} is rejected", ["--pipeline", retired], tag=f"retired_{retired}", expect=2) run(f"--pipeline {retired} is rejected", ["--pipeline", retired], tag=f"retired_{retired}", expect=2)
# Fixed-graph knobs: accepted by Click, refused by the library (exit 1). # Retired options: no longer parsed at all (exit 2, "No such option").
run("--steps 20 is rejected", ["--steps", "20"], tag="steps20", expect=1) for args, tag in (
run("--guidance-scale 5.0 is rejected", ["--guidance-scale", "5.0"], tag="gs5", expect=1) (["--steps", "20"], "steps20"),
run("--model override is rejected", ["--model", "org/custom"], tag="model", expect=1) (["--guidance-scale", "5.0"], "gs5"),
(["--model", "org/custom"], "model"),
(["--device", "cpu"], "devcpu"),
(["--auto"], "auto"),
):
run(f"{args[0]} is no longer an option", args, tag=tag, expect=2)
try: try:
import torch import torch
@@ -393,11 +401,6 @@ def _knob_rows(r: Runner, tmp: Path, img: Path) -> None:
# Without CUDA the honest outcome is a CLEAN refusal naming the reason, not a # Without CUDA the honest outcome is a CLEAN refusal naming the reason, not a
# traceback from inside a half-built pipeline. # traceback from inside a half-built pipeline.
run("no CUDA fails cleanly", [], tag="nocuda", expect=1) run("no CUDA fails cleanly", [], tag="nocuda", expect=1)
for name, extra, tag in (
("--device cpu fails cleanly", ["--device", "cpu"], "cpu"),
("--device mps fails cleanly", ["--device", "mps"], "mps"),
):
run(name, extra, tag=tag, expect=1)
for label in ( for label in (
"--strength", "--strength",
"--controlnet-scale", "--controlnet-scale",
@@ -405,7 +408,6 @@ def _knob_rows(r: Runner, tmp: Path, img: Path) -> None:
"--unsharp", "--unsharp",
"--no-adaptive-polish", "--no-adaptive-polish",
"--tile", "--tile",
"--auto",
"--pipeline sdxl-zimage", "--pipeline sdxl-zimage",
): ):
r.skip(label, "accepted-knob rows need a CUDA device") r.skip(label, "accepted-knob rows need a CUDA device")
@@ -418,11 +420,9 @@ def _knob_rows(r: Runner, tmp: Path, img: Path) -> None:
run("--unsharp", ["--unsharp", "0.5"], tag="uns") run("--unsharp", ["--unsharp", "0.5"], tag="uns")
run("--no-adaptive-polish", ["--no-adaptive-polish"], tag="nap") run("--no-adaptive-polish", ["--no-adaptive-polish"], tag="nap")
run("--tile", ["--tile", "--tile-size", "256", "--tile-overlap", "64"], tag="tile") run("--tile", ["--tile", "--tile-size", "256", "--tile-overlap", "64"], tag="tile")
run("--auto (deprecated, requests polish only)", ["--auto"], tag="auto")
# --model and --hf-token are deliberately not exercised: one would download a second # --hf-token needs a real credential, so it cannot be exercised meaningfully here.
# multi-GB checkpoint, the other needs a real credential. Skipped loudly, not passed. # Skipped loudly rather than silently passed over.
r.skip("--model", "would download a second multi-GB checkpoint")
r.skip("--hf-token", "needs a real credential; cannot be exercised meaningfully here") r.skip("--hf-token", "needs a real credential; cannot be exercised meaningfully here")
# CONTRACT, not just execution: the same seed must reproduce the same pixels. # CONTRACT, not just execution: the same seed must reproduce the same pixels.
+13 -34
View File
@@ -1,36 +1,15 @@
"""Compatibility namespace for metadata and regeneration helpers. """Private namespace for metadata parsing and regeneration internals.
The public API (``WatermarkRemover`` / ``remove_watermark`` / ``remove_ai_metadata``) Deliberately empty. It carried a PEP 562 ``__getattr__`` re-exporting
is exposed **lazily** via PEP 562 ``__getattr__``: importing a light submodule ``WatermarkRemover`` and ``remove_ai_metadata`` as a "compatibility namespace",
(e.g. ``_internal.c2pa`` / ``_internal.constants`` from ``identify``) must NOT eagerly pull but nothing ever reached for either through this package -- every caller imports
``watermark_remover``, which imports torch + diffusers at module top. Keeping this the submodule directly. The laziness it defended is real and still enforced, just
lazy is what lets ``import remove_ai_watermarks.identify`` stay cheap (~36 MB, no elsewhere: importing a light submodule (``_internal.c2pa`` / ``_internal.constants``
torch) even in a full install where the ``diffusion`` extra is present -- from ``identify``) must not pull ``watermark_remover``, which imports torch at
otherwise the mere presence of torch in the env inflated identify to ~420 MB and module top. That property comes from those direct submodule imports, not from a
risked OOM on a 512 MB host. shim here; a re-export in this file would be the one thing that could break it.
Keep this module free of imports. ``import remove_ai_watermarks.identify`` stays
around 36 MB even in a full install where torch is present; routing anything heavy
through here inflated it to roughly 420 MB and risked OOM on a 512 MB host.
""" """
from __future__ import annotations
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover, remove_watermark
from remove_ai_watermarks.metadata import remove_ai_metadata
__all__ = ["WatermarkRemover", "remove_ai_metadata", "remove_watermark"]
def __getattr__(name: str) -> object:
"""Resolve the public API on first access (PEP 562), not at package import."""
if name == "remove_ai_metadata":
# Re-export the single, robust stripper (byte-level, lossless-for-JPEG, all
# containers); the old legacy metadata helper implementation is retired.
from remove_ai_watermarks.metadata import remove_ai_metadata
return remove_ai_metadata
if name in ("WatermarkRemover", "remove_watermark"):
from remove_ai_watermarks._internal import watermark_remover
return getattr(watermark_remover, name)
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
@@ -14,9 +14,6 @@ AI_METADATA_KEYS = _tokens(
"parameters|postprocessing|extras|workflow|prompt|Dream|SD:mode|StableDiffusionVersion|" "parameters|postprocessing|extras|workflow|prompt|Dream|SD:mode|StableDiffusionVersion|"
"generation_time|Model|Model hash|Seed" "generation_time|Model|Model hash|Seed"
) )
PNG_METADATA_KEYS = _tokens(
"Author|Title|Description|Copyright|Creation Time|Software|Disclaimer|Warning|Source|Comment"
)
AI_KEYWORDS = _tokens( AI_KEYWORDS = _tokens(
"prompt|negative_prompt|sampler|cfg_scale|lora|diffusion|comfy|midjourney|dall-e|dalle|imagen|firefly|c2pa|chatgpt|gpt-4|sora|openai|truepic|stable_diffusion|invokeai" "prompt|negative_prompt|sampler|cfg_scale|lora|diffusion|comfy|midjourney|dall-e|dalle|imagen|firefly|c2pa|chatgpt|gpt-4|sora|openai|truepic|stable_diffusion|invokeai"
) )
@@ -30,9 +30,9 @@ from remove_ai_watermarks._internal.qwen_zimage_pipeline import (
) )
from remove_ai_watermarks._internal.watermark_profiles import ( from remove_ai_watermarks._internal.watermark_profiles import (
CONTROLNET_CANNY_MODEL, CONTROLNET_CANNY_MODEL,
DEFAULT_MODEL_ID,
SDXL_LIGHTNING_MODEL_ID, SDXL_LIGHTNING_MODEL_ID,
SDXL_LIGHTNING_PATTERN, SDXL_LIGHTNING_PATTERN,
SDXL_MODEL_ID,
) )
log = logging.getLogger(__name__) log = logging.getLogger(__name__)
@@ -88,7 +88,7 @@ class SdxlZImagePipeline(QwenZImagePipeline):
controlnet = ControlNetModel.from_pretrained(CONTROLNET_CANNY_MODEL, torch_dtype=torch.float16, **token) controlnet = ControlNetModel.from_pretrained(CONTROLNET_CANNY_MODEL, torch_dtype=torch.float16, **token)
vae = AutoencoderKL.from_pretrained(SDXL_VAE_MODEL_ID, torch_dtype=torch.float16, **token) vae = AutoencoderKL.from_pretrained(SDXL_VAE_MODEL_ID, torch_dtype=torch.float16, **token)
pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained( pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
DEFAULT_MODEL_ID, SDXL_MODEL_ID,
controlnet=controlnet, controlnet=controlnet,
vae=vae, vae=vae,
torch_dtype=torch.float16, torch_dtype=torch.float16,
+6 -59
View File
@@ -1,11 +1,11 @@
"""Sliding-window tiled diffusion for large images. """Sliding-window tiled diffusion for large images.
The img2img / ControlNet pipeline denoises the WHOLE image in one forward pass, The global stage denoises the WHOLE image in one forward pass, so it OOMs on a
so it OOMs on MPS/GPU above ~2K (issue #10). Tiling splits the image into GPU above ~2K (issue #10). Tiling splits the image into overlapping tiles -- each
overlapping tiles -- each kept near SDXL's ~1024 training size -- regenerates kept near the ~1024 training size -- regenerates each tile independently, and
each tile independently, and feather-blends the overlaps. The result retains the feather-blends the overlaps. The result retains the input's native dimensions
input's native dimensions without an explicit ``--max-resolution`` downscale, but without an explicit ``--max-resolution`` downscale, but it is not pixel-lossless
it is not pixel-lossless because every tile is regenerated. because every tile is regenerated.
The geometry (``plan_tiles``) and the blend weighting (``feather_weights``) are The geometry (``plan_tiles``) and the blend weighting (``feather_weights``) are
pure functions, unit-tested without the diffusion model. ``run_tiled`` is the pure functions, unit-tested without the diffusion model. ``run_tiled`` is the
@@ -100,59 +100,6 @@ def feather_weights(width: int, height: int, overlap: int) -> NDArray[Any]:
return weights return weights
def feather_region_composite(
base: NDArray[Any],
regenerated: NDArray[Any],
box: tuple[int, int, int, int],
*,
feather: int = 64,
) -> NDArray[Any]:
"""Composite ``regenerated`` over ``base`` inside ``box`` only, feathering the seam.
For AI-ENHANCED composites (digitalSourceType ``compositeWithTrainedAlgorithmicMedia``):
the diffusion remover regenerates the whole frame, but only the AI-composited
REGION should change -- the rest is a real photo that must be preserved. This
blends the regenerated pixels in over ``box = (x, y, w, h)`` with a separable
linear taper of ``feather`` px at the box edges, so the result equals ``base``
EXACTLY outside the box and ramps smoothly (no hard seam) at the boundary.
Pure and model-free (unit-tested): ``base`` and ``regenerated`` must be the same
shape (H x W, or H x W x C). The output preserves ``base``'s dtype. ``feather`` is
clamped to half the box on each axis, so a small region still tapers symmetrically;
``feather=0`` is a hard-edged paste.
"""
import numpy as np
if base.shape != regenerated.shape:
raise ValueError(f"shape mismatch: base {base.shape} vs regenerated {regenerated.shape}")
h, w = base.shape[:2]
x, y, bw, bh = box
x0, y0 = max(0, x), max(0, y)
x1, y1 = min(w, x + bw), min(h, y + bh)
out = base.copy()
if x1 <= x0 or y1 <= y0:
return out # empty / off-image box -> nothing regenerated
def taper(n: int) -> NDArray[Any]:
win = np.ones(n, dtype=np.float32)
f = min(max(feather, 0), n // 2)
if f > 0:
ramp = (np.arange(f, dtype=np.float32) + 1.0) / (f + 1.0) # in (0, 1), 0 at the edge
win[:f] = ramp
win[n - f :] = ramp[::-1]
return win
rh, rw = y1 - y0, x1 - x0
wmap = np.outer(taper(rh), taper(rw)) # ~0 at the box edge, 1 in the interior
if base.ndim == 3:
wmap = wmap[:, :, None]
roi_base = base[y0:y1, x0:x1].astype(np.float32)
roi_gen = regenerated[y0:y1, x0:x1].astype(np.float32)
blended = roi_base * (1.0 - wmap) + roi_gen * wmap
out[y0:y1, x0:x1] = np.clip(blended, 0, 255).astype(base.dtype)
return out
def run_tiled( def run_tiled(
generate_tile: Callable[[PILImage.Image], PILImage.Image], generate_tile: Callable[[PILImage.Image], PILImage.Image],
image: PILImage.Image, image: PILImage.Image,
@@ -17,8 +17,10 @@ if TYPE_CHECKING:
from pathlib import Path from pathlib import Path
# SDXL base is no longer a profile of its own, but it is still the global stage of # SDXL base is no longer a profile of its own, but it is still the global stage of
# sdxl-zimage, so the checkpoint id stays. # sdxl-zimage, so the checkpoint id stays. Named for what it is rather than
DEFAULT_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0" # ``DEFAULT_MODEL_ID``: there is no user-selectable model any more, so "default"
# implied an override that both profiles reject.
SDXL_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
CONTROLNET_CANNY_MODEL = "xinsir/controlnet-canny-sdxl-1.0" CONTROLNET_CANNY_MODEL = "xinsir/controlnet-canny-sdxl-1.0"
QWEN_ZIMAGE_PROFILE = "qwen-zimage" QWEN_ZIMAGE_PROFILE = "qwen-zimage"
@@ -26,14 +28,27 @@ SDXL_ZIMAGE_PROFILE = "sdxl-zimage"
DEFAULT_PROFILE = QWEN_ZIMAGE_PROFILE DEFAULT_PROFILE = QWEN_ZIMAGE_PROFILE
PROFILE_CHOICES = (QWEN_ZIMAGE_PROFILE, SDXL_ZIMAGE_PROFILE) PROFILE_CHOICES = (QWEN_ZIMAGE_PROFILE, SDXL_ZIMAGE_PROFILE)
# The modules a real removal run needs, and the extra that installs them. Both live
# here, in the only profile module that imports nothing heavy, because the CLI's
# availability gate and the remover's own precondition must agree: when they drifted,
# the CLI passed on a torch+diffusers environment and the run then died at the
# DiffSynth face stage, telling the user to install an extra that does not contain it.
REMOVAL_MODULES = ("torch", "diffusers", "diffsynth")
INVISIBLE_EXTRA = "remove-ai-watermarks[qwen-zimage]"
# qwen-zimage's output already matches the input's detail level, so polishing it is a
# no-op at best. sdxl-zimage's global pass leaves the softer output the polish exists
# for. This is per-profile data, not a CLI concern: the flag defaults to None so that
# "the user did not choose" stays a value rather than an inference from Click state.
PROFILE_ADAPTIVE_POLISH = {QWEN_ZIMAGE_PROFILE: False, SDXL_ZIMAGE_PROFILE: True}
SDXL_LIGHTNING_MODEL_ID = "ByteDance/SDXL-Lightning" SDXL_LIGHTNING_MODEL_ID = "ByteDance/SDXL-Lightning"
SDXL_LIGHTNING_PATTERN = "sdxl_lightning_4step_lora.safetensors" SDXL_LIGHTNING_PATTERN = "sdxl_lightning_4step_lora.safetensors"
# Both profiles run the same distilled four-step schedule, and both are certified at a # Both profiles are certified at a fixed seed because SynthID removal near the
# fixed seed because SynthID removal near the strength floor is seed-dependent. # strength floor is seed-dependent. The step count and CFG are not settable at all --
PROFILE_STEPS = 4 # each stage owns them (``GLOBAL_STEPS`` / ``FACE_STEPS`` in qwen_zimage_pipeline).
PROFILE_SEED = 0 PROFILE_SEED = 0
PROFILE_CFG = 1.0
# sdxl-zimage runs the qwen-zimage recipe on an SDXL global stage, and strength is # sdxl-zimage runs the qwen-zimage recipe on an SDXL global stage, and strength is
# architecture-bound: at Qwen's 0.154 an SDXL global pass leaves SynthID on a native # architecture-bound: at Qwen's 0.154 an SDXL global pass leaves SynthID on a native
@@ -77,16 +92,18 @@ def normalize_profile(profile: str) -> str:
return _ALIASES.get(value, value) return _ALIASES.get(value, value)
def resolve_steps(num_inference_steps: int | None) -> int:
"""Return an explicit step count or the distilled four-step default."""
return PROFILE_STEPS if num_inference_steps is None else num_inference_steps
def resolve_seed(seed: int | None) -> int: def resolve_seed(seed: int | None) -> int:
"""Keep both profiles reproducible by default.""" """Keep both profiles reproducible by default."""
return PROFILE_SEED if seed is None else seed return PROFILE_SEED if seed is None else seed
def resolve_adaptive_polish(adaptive_polish: bool | None, pipeline: str) -> bool:
"""Return an explicit polish choice, or the profile's calibrated default."""
if adaptive_polish is not None:
return adaptive_polish
return PROFILE_ADAPTIVE_POLISH.get(normalize_profile(pipeline), True)
def strength_default_help() -> str: def strength_default_help() -> str:
"""Describe the live default policy without duplicating its values.""" """Describe the live default policy without duplicating its values."""
return ( return (
@@ -3,7 +3,6 @@
# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalCall=false, reportOptionalSubscript=false, reportOptionalOperand=false, reportAttributeAccessIssue=false, reportPrivateImportUsage=false, reportPrivateUsage=false, reportInvalidTypeForm=false, reportConstantRedefinition=false, reportUnnecessaryComparison=false # pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalCall=false, reportOptionalSubscript=false, reportOptionalOperand=false, reportAttributeAccessIssue=false, reportPrivateImportUsage=false, reportPrivateUsage=false, reportInvalidTypeForm=false, reportConstantRedefinition=false, reportUnnecessaryComparison=false
from __future__ import annotations from __future__ import annotations
import contextlib
import logging import logging
import os import os
import subprocess import subprocess
@@ -12,16 +11,13 @@ from typing import TYPE_CHECKING, Any
from PIL import Image from PIL import Image
from remove_ai_watermarks._internal.watermark_profiles import ( from remove_ai_watermarks._internal.watermark_profiles import (
DEFAULT_MODEL_ID,
DEFAULT_PROFILE, DEFAULT_PROFILE,
PROFILE_CFG, INVISIBLE_EXTRA,
PROFILE_CHOICES, PROFILE_CHOICES,
PROFILE_STEPS, REMOVAL_MODULES,
QWEN_ZIMAGE_PROFILE,
SDXL_ZIMAGE_PROFILE, SDXL_ZIMAGE_PROFILE,
normalize_profile, normalize_profile,
resolve_seed, resolve_seed,
resolve_steps,
resolve_strength, resolve_strength,
) )
from remove_ai_watermarks.optional_deps import module_available from remove_ai_watermarks.optional_deps import module_available
@@ -32,13 +28,6 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
# Both two-stage profiles share the face stage, the four-step schedule, CFG 1.0, the
# fixed model stack and the native-resolution contract; only the global model differs.
_ZIMAGE_STACKS = {
QWEN_ZIMAGE_PROFILE: "Qwen-Image-2512 and Z-Image",
SDXL_ZIMAGE_PROFILE: "SDXL and Z-Image",
}
try: try:
import torch import torch
@@ -47,18 +36,20 @@ except ImportError:
torch = None # type: ignore[assignment] torch = None # type: ignore[assignment]
_HAS_TORCH = False _HAS_TORCH = False
_HAS_DIFFUSERS = module_available("diffusers") # Probed once at import. ``torch`` is imported above rather than probed because this
# module needs the object, not just the answer.
_HAS_REMOVAL_MODULES = module_available(*(name for name in REMOVAL_MODULES if name != "torch"))
def is_watermark_removal_available() -> bool: def is_watermark_removal_available() -> bool:
"""Return whether the standard diffusion runtime can be imported.""" """Return whether the full removal runtime can be imported."""
return _HAS_TORCH and _HAS_DIFFUSERS return _HAS_TORCH and _HAS_REMOVAL_MODULES
def _ensure_watermark_deps() -> None: def _ensure_watermark_deps() -> None:
if not is_watermark_removal_available(): if not is_watermark_removal_available():
raise ImportError( raise ImportError(
"Invisible watermark regeneration requires the 'diffusion' extra. Install remove-ai-watermarks[diffusion]." f"Invisible watermark regeneration requires the 'qwen-zimage' extra: pip install {INVISIBLE_EXTRA}."
) )
@@ -75,26 +66,9 @@ def _has_nvidia_gpu() -> bool:
return True return True
def try_empty_device_cache(device: str) -> None: def _cuda_works() -> bool:
"""Ask Torch to release cached accelerator memory when the backend supports it.
Moved here when ``img2img_runner`` was deleted: the runner and its MPS recovery
path went with the CPU/MPS profiles, leaving this as that module's only content.
Silent by design -- it runs in cleanup paths where a raise would replace the real
error.
"""
if not _HAS_TORCH:
return
backend = getattr(torch, device, None) # type: ignore[union-attr]
empty_cache = getattr(backend, "empty_cache", None)
if callable(empty_cache):
with contextlib.suppress(Exception):
empty_cache()
def _backend_works(device: str) -> bool:
try: try:
probe = torch.tensor([1.0], device=device) # type: ignore[union-attr] probe = torch.tensor([1.0], device="cuda") # type: ignore[union-attr]
_ = probe + probe _ = probe + probe
except (AssertionError, RuntimeError): except (AssertionError, RuntimeError):
return False return False
@@ -102,33 +76,28 @@ def _backend_works(device: str) -> bool:
def get_device() -> str: def get_device() -> str:
"""Select CUDA, XPU, MPS, or CPU in that order when each backend is usable.""" """Return ``"cuda"`` when a usable CUDA backend is present, else ``"cpu"``.
Deliberately binary. Both profiles are CUDA-only, so an XPU or MPS answer would
only travel one frame further to the same refusal in :class:`WatermarkRemover`,
while costing a probe on each. ``"cpu"`` here means "no CUDA", which is exactly
what that refusal reports.
"""
if not _HAS_TORCH: if not _HAS_TORCH:
return "cpu" return "cpu"
if torch.cuda.is_available() and _backend_works("cuda"): # type: ignore[union-attr] if torch.cuda.is_available() and _cuda_works(): # type: ignore[union-attr]
return "cuda" return "cuda"
xpu = getattr(torch, "xpu", None)
if xpu is not None and xpu.is_available() and _backend_works("xpu"):
return "xpu"
if _has_nvidia_gpu(): if _has_nvidia_gpu():
logger.warning("NVIDIA GPU detected, but the installed PyTorch build has no working CUDA backend") logger.warning("NVIDIA GPU detected, but the installed PyTorch build has no working CUDA backend")
mps = getattr(getattr(torch, "backends", None), "mps", None)
if mps is not None and mps.is_available():
return "mps"
return "cpu" return "cpu"
class WatermarkRemover: class WatermarkRemover:
"""Load one regeneration profile and write a metadata-clean raster output.""" """Load one regeneration profile and write a metadata-clean raster output."""
DEFAULT_MODEL_ID = DEFAULT_MODEL_ID
_DEVICES = frozenset({"cuda"})
def __init__( def __init__(
self, self,
model_id: str | None = None,
device: str | None = None, device: str | None = None,
torch_dtype: Any = None,
progress_callback: Callable[[str], None] | None = None, progress_callback: Callable[[str], None] | None = None,
hf_token: str | None = None, hf_token: str | None = None,
pipeline: str = DEFAULT_PROFILE, pipeline: str = DEFAULT_PROFILE,
@@ -138,31 +107,25 @@ class WatermarkRemover:
self.model_profile = normalize_profile(pipeline) self.model_profile = normalize_profile(pipeline)
if self.model_profile not in PROFILE_CHOICES: if self.model_profile not in PROFILE_CHOICES:
raise ValueError(f"Unsupported pipeline '{pipeline}'. Use one of: {', '.join(PROFILE_CHOICES)}.") raise ValueError(f"Unsupported pipeline '{pipeline}'. Use one of: {', '.join(PROFILE_CHOICES)}.")
if model_id is not None: # There is no ``model_id`` parameter and no ``model_id`` attribute: each
raise ValueError( # profile pins a fixed model stack, and the dtype below is bound to that
f"The {self.model_profile} profile uses a fixed {_ZIMAGE_STACKS[self.model_profile]} model stack." # stack's weights. Both used to be constructor overrides that existed only to
) # be rejected or to break the run, and the attribute only existed to echo the
self.model_id = ( # rejected value back.
"Qwen/Qwen-Image-2512 + Tongyi-MAI/Z-Image-Turbo"
if self.model_profile == QWEN_ZIMAGE_PROFILE
else f"{DEFAULT_MODEL_ID} + Tongyi-MAI/Z-Image-Turbo"
)
_ensure_watermark_deps() _ensure_watermark_deps()
selected_device = (device or get_device()).casefold() selected_device = (device or get_device()).casefold()
self.device = get_device() if selected_device == "auto" else selected_device self.device = get_device() if selected_device == "auto" else selected_device
# CUDA is a precondition of the object, not of the run. Both profiles raise on # CUDA is a precondition of the object, not of the run. Both profiles raise on
# any other device, so accepting one here only defers a guaranteed failure to # any other device, so accepting one here only defers a guaranteed failure to
# model-load time, several layers down and under the wrong profile's name. # model-load time, several layers down and under the wrong profile's name.
if self.device not in self._DEVICES: if self.device != "cuda":
raise ValueError( raise ValueError(
f"Invisible-watermark removal is CUDA-only, so '{device}' cannot run it. " f"Invisible-watermark removal is CUDA-only, so '{self.device}' cannot run it. "
"Both remaining profiles need an NVIDIA GPU. Visible-mark removal and " "Both remaining profiles need an NVIDIA GPU. Visible-mark removal and "
"every identify command still run on CPU." "every identify command still run on CPU."
) )
if torch_dtype is not None: if self.model_profile == SDXL_ZIMAGE_PROFILE:
self.torch_dtype = torch_dtype
elif self.model_profile == SDXL_ZIMAGE_PROFILE:
# SDXL ships fp16 weights and an fp16-safe VAE; bf16 would give up the # SDXL ships fp16 weights and an fp16-safe VAE; bf16 would give up the
# variant without buying anything on this architecture. # variant without buying anything on this architecture.
self.torch_dtype = torch.float16 # type: ignore[union-attr] self.torch_dtype = torch.float16 # type: ignore[union-attr]
@@ -175,18 +138,13 @@ class WatermarkRemover:
self._progress_callback = progress_callback self._progress_callback = progress_callback
self._qwen_zimage_pipeline: Any = None self._qwen_zimage_pipeline: Any = None
def _set_progress(self, message: str) -> None:
if self._progress_callback is not None:
with contextlib.suppress(Exception):
self._progress_callback(message)
def preload(self, *, global_only: bool = False) -> None: def preload(self, *, global_only: bool = False) -> None:
"""Materialize the selected model stack before the first request.""" """Materialize the selected model stack before the first request."""
self._load_qwen_zimage_pipeline().preload(global_only=global_only) self._load_qwen_zimage_pipeline().preload(global_only=global_only)
def _load_qwen_zimage_pipeline(self) -> Any: def _load_qwen_zimage_pipeline(self) -> Any:
if self._qwen_zimage_pipeline is None: if self._qwen_zimage_pipeline is None:
if getattr(self, "model_profile", QWEN_ZIMAGE_PROFILE) == SDXL_ZIMAGE_PROFILE: if self.model_profile == SDXL_ZIMAGE_PROFILE:
from remove_ai_watermarks._internal.sdxl_zimage_pipeline import ( from remove_ai_watermarks._internal.sdxl_zimage_pipeline import (
SdxlZImagePipeline as _Pipeline, SdxlZImagePipeline as _Pipeline,
) )
@@ -206,44 +164,6 @@ class WatermarkRemover:
) )
return self._qwen_zimage_pipeline return self._qwen_zimage_pipeline
def _run_qwen_zimage(
self,
init_image: Image.Image,
strength: float,
seed: int | None,
*,
tile: bool = False,
tile_size: int = 1024,
tile_overlap: int = 128,
) -> Image.Image:
return self._load_qwen_zimage_pipeline().run(
init_image,
strength=strength,
seed=seed,
tile=tile,
tile_size=tile_size,
tile_overlap=tile_overlap,
)
def _generate(
self,
image: Image.Image,
strength: float,
seed: int | None,
*,
tile: bool,
tile_size: int,
tile_overlap: int,
) -> Image.Image:
return self._run_qwen_zimage(
image,
strength,
seed,
tile=tile,
tile_size=tile_size,
tile_overlap=tile_overlap,
)
def _write_output(self, image: Image.Image, output_path: Path) -> None: def _write_output(self, image: Image.Image, output_path: Path) -> None:
import numpy as np import numpy as np
@@ -262,17 +182,18 @@ class WatermarkRemover:
image_path: Path, image_path: Path,
output_path: Path | None = None, output_path: Path | None = None,
strength: float | None = None, strength: float | None = None,
num_inference_steps: int | None = None,
guidance_scale: float | None = None,
seed: int | None = None, seed: int | None = None,
vendor: str | None = None, vendor: str | None = None,
tile: bool = False, tile: bool = False,
tile_size: int = 1024, tile_size: int = 1024,
tile_overlap: int = 128, tile_overlap: int = 128,
region: tuple[int, int, int, int] | None = None,
region_feather: int = 64,
) -> Path: ) -> Path:
"""Regenerate image pixels and write the result without AI metadata.""" """Regenerate image pixels and write the result without AI metadata.
Step count and CFG are not parameters. Each stage of both profiles is a
distilled schedule that owns its own, so the only thing a caller-supplied
value could do is break the run or be rejected.
"""
if not image_path.exists(): if not image_path.exists():
raise FileNotFoundError(f"Image not found: {image_path}") raise FileNotFoundError(f"Image not found: {image_path}")
destination = output_path or image_path destination = output_path or image_path
@@ -283,82 +204,13 @@ class WatermarkRemover:
if not 0.0 <= resolved_strength <= 1.0: if not 0.0 <= resolved_strength <= 1.0:
raise ValueError(f"Strength must be between 0.0 and 1.0, got {resolved_strength}") raise ValueError(f"Strength must be between 0.0 and 1.0, got {resolved_strength}")
# Both profiles are distilled four-step schedules at CFG 1.0. Anything else is result = self._load_qwen_zimage_pipeline().run(
# a caller error rather than a knob, so it is rejected instead of coerced.
steps = resolve_steps(num_inference_steps)
if steps != PROFILE_STEPS:
raise ValueError(f"The {self.model_profile} profile requires {PROFILE_STEPS} steps.")
if guidance_scale is not None and guidance_scale != PROFILE_CFG:
raise ValueError(f"The {self.model_profile} profile requires CFG {PROFILE_CFG}.")
result = self._generate(
source, source,
resolved_strength, strength=resolved_strength,
resolve_seed(seed), seed=resolve_seed(seed),
tile=tile, tile=tile,
tile_size=tile_size, tile_size=tile_size,
tile_overlap=tile_overlap, tile_overlap=tile_overlap,
) )
if region is not None:
import numpy as np
from remove_ai_watermarks._internal.tiling import feather_region_composite
if result.size != source.size:
result = result.resize(source.size, Image.Resampling.LANCZOS)
merged = feather_region_composite(
np.asarray(source),
np.asarray(result.convert("RGB")),
region,
feather=region_feather,
)
result = Image.fromarray(merged)
self._write_output(result, destination) self._write_output(result, destination)
return destination return destination
def remove_watermark_batch(
self,
input_dir: Path,
output_dir: Path,
strength: float | None = None,
num_inference_steps: int | None = None,
extensions: tuple[str, ...] = (".png", ".jpg", ".jpeg", ".webp"),
) -> list[Path]:
"""Process matching files in a directory, logging and continuing on failures."""
if not input_dir.exists():
raise FileNotFoundError(f"Input directory not found: {input_dir}")
output_dir.mkdir(parents=True, exist_ok=True)
outputs: list[Path] = []
candidates = sorted(path for path in input_dir.iterdir() if path.suffix.casefold() in extensions)
for source in candidates:
try:
outputs.append(self.remove_watermark(source, output_dir / source.name, strength, num_inference_steps))
except Exception as error:
logger.error("Failed to process %s: %s", source, error)
finally:
try_empty_device_cache(self.device)
return outputs
def remove_watermark(
image_path: Path,
output_path: Path | None = None,
strength: float | None = None,
model_id: str | None = None,
device: str | None = None,
hf_token: str | None = None,
region: tuple[int, int, int, int] | None = None,
) -> Path:
"""Convenience wrapper using the default ControlNet profile."""
from remove_ai_watermarks._internal.watermark_profiles import vendor_for_strength
remover = WatermarkRemover(model_id=model_id, device=device, hf_token=hf_token)
return remover.remove_watermark(
image_path,
output_path,
strength,
vendor=vendor_for_strength(image_path),
region=region,
)
+55 -226
View File
@@ -25,10 +25,8 @@ from remove_ai_watermarks._internal.constants import SUPPORTED_FORMATS
from remove_ai_watermarks._internal.utils import is_supported_format from remove_ai_watermarks._internal.utils import is_supported_format
from remove_ai_watermarks._internal.watermark_profiles import ( from remove_ai_watermarks._internal.watermark_profiles import (
DEFAULT_PROFILE, DEFAULT_PROFILE,
INVISIBLE_EXTRA,
PROFILE_CHOICES, PROFILE_CHOICES,
QWEN_ZIMAGE_PROFILE,
resolve_seed,
resolve_steps,
resolve_strength, resolve_strength,
strength_default_help, strength_default_help,
vendor_for_strength, vendor_for_strength,
@@ -184,23 +182,14 @@ _unsharp_option = click.option(
"--unsharp", type=float, default=0.0, help="Unsharp-mask sharpening strength (0 = off, typical: 0.3-0.8)." "--unsharp", type=float, default=0.0, help="Unsharp-mask sharpening strength (0 = off, typical: 0.3-0.8)."
) )
_auto_option = click.option(
"--auto",
is_flag=True,
default=False,
help="DEPRECATED: it no longer selects a pipeline. It now only requests the "
"adaptive polish, which the two-stage profiles otherwise leave off to keep their "
"output untouched. Prefer --adaptive-polish.",
)
_adaptive_polish_option = click.option( _adaptive_polish_option = click.option(
"--adaptive-polish/--no-adaptive-polish", "--adaptive-polish/--no-adaptive-polish",
default=True, default=None,
help="Restore the input's detail level after removal (capped unsharp + edge-masked grain " help="Restore the input's detail level after removal (capped unsharp + edge-masked grain "
"targeting the input's sharpness, sparing text), countering the over-smoothed look. ON by " "targeting the input's sharpness, sparing text), countering the over-smoothed look. "
"default except for qwen-zimage, whose upstream-matching output is left unchanged; it " "Unset follows the profile: ON for sdxl-zimage, OFF for qwen-zimage, whose "
"self-limits where there is no detail deficit (text/flat graphics). Pass --adaptive-polish " "upstream-matching output is left unchanged. It self-limits where there is no detail "
"or --no-adaptive-polish to override. Independent of --unsharp/--humanize.", "deficit (text/flat graphics). Independent of --unsharp/--humanize.",
) )
@@ -230,23 +219,11 @@ def _tile_options(f: Any) -> Any:
)(f) )(f)
# HuggingFace model + CFG knobs, shared by the diffusion commands (invisible/all/batch) # There is deliberately no --model, --steps, --guidance-scale or --device option.
# so the surface stays identical across them. # Each profile pins a fixed model stack, a distilled per-stage schedule, CFG 1.0 and
_model_option = click.option( # CUDA; every one of those knobs existed only so the library could reject it several
"--model", # layers down. A flag whose sole outcome is an error is worse than no flag at all --
type=str, # it advertises a capability that does not exist.
default=None,
help="HuggingFace model ID. Both profiles pin a fixed model stack, so anything "
"other than the default is rejected rather than silently ignored.",
)
_guidance_scale_option = click.option(
"--guidance-scale",
type=float,
default=None,
help="Classifier-free guidance scale (CFG). Both profiles are distilled and fix "
"CFG at 1.0, so any other value is rejected.",
)
# The two-stage profiles are the only ones left. The former controlnet, sdxl, qwen and # The two-stage profiles are the only ones left. The former controlnet, sdxl, qwen and
# default profiles were removed rather than kept as a CPU path: none matched this # default profiles were removed rather than kept as a CPU path: none matched this
@@ -276,6 +253,23 @@ _strength_option = click.option(
default=None, default=None,
help=f"Denoising strength (0.0-1.0). Default: {strength_default_help()}.", help=f"Denoising strength (0.0-1.0). Default: {strength_default_help()}.",
) )
_seed_option = click.option(
"--seed",
type=int,
default=None,
help="Random seed for reproducibility. Default 0: both profiles are certified "
"at a fixed seed, because SynthID removal near the strength floor is seed-dependent.",
)
_hf_token_option = click.option("--hf-token", type=str, default=None, help="HuggingFace API token.")
_humanize_option = click.option(
"--humanize", type=float, default=0.0, help="Analog Humanizer film grain intensity (0 = off, typical: 2.0-6.0)."
)
_max_resolution_option = click.option(
"--max-resolution",
type=int,
default=0,
help="Cap long side (px) before diffusion; 0 = native and preserves the most detail. Raise only on GPU OOM.",
)
_force_option = click.option( _force_option = click.option(
"--force/--no-force", "--force/--no-force",
default=False, default=False,
@@ -323,41 +317,6 @@ _visible_sensitivity_option = click.option(
) )
def _resolve_auto_polish(auto: bool, adaptive_polish: bool) -> bool:
"""Warn on the retired ``--auto`` flag, returning ``adaptive_polish`` unchanged.
``--auto`` used to plan the pipeline + polish from content detection. There is now
only one default pipeline, and the content detectors were removed, so the flag
survives purely as a polish request: it emits a deprecation warning and passes
``adaptive_polish`` through, with an explicit ``--no-adaptive-polish`` still winning.
"""
if auto:
click.echo(
"Warning: --auto is deprecated and now does nothing (the adaptive polish it "
"enabled is ON by default). Use --no-adaptive-polish to turn the polish off.",
err=True,
)
return adaptive_polish
def _resolve_profile_polish(auto: bool, adaptive_polish: bool, pipeline: str) -> bool:
"""Keep the upstream qwen-zimage output unchanged unless polish was explicit.
``--auto`` counts as explicit. It is deprecated, but it is still a request for the
polish, and once qwen-zimage became the DEFAULT pipeline the source check below
would otherwise have silently turned that flag into a no-op for every caller.
"""
adaptive_polish = _resolve_auto_polish(auto, adaptive_polish)
if pipeline != QWEN_ZIMAGE_PROFILE or auto:
return adaptive_polish
ctx = click.get_current_context(silent=True)
if ctx is None:
return adaptive_polish
if ctx.get_parameter_source("adaptive_polish") == click.core.ParameterSource.DEFAULT:
return False
return adaptive_polish
def _visible_provenance(path: Path | None) -> frozenset[str]: def _visible_provenance(path: Path | None) -> frozenset[str]:
"""Vendor keys local metadata confirms, the EVIDENCE that drives ``auto`` """Vendor keys local metadata confirms, the EVIDENCE that drives ``auto``
sensitivity. Thin wrapper over the public :func:`api.visible_provenance` (one sensitivity. Thin wrapper over the public :func:`api.visible_provenance` (one
@@ -833,40 +792,13 @@ def cmd_erase(
"-o", "--output", type=click.Path(path_type=Path), default=None, help="Output path (default: <source>_clean.<ext>)." "-o", "--output", type=click.Path(path_type=Path), default=None, help="Output path (default: <source>_clean.<ext>)."
) )
@_strength_option @_strength_option
@click.option(
"--steps",
type=int,
default=None,
help="Number of denoising steps. Both profiles are distilled four-step schedules, so 4 is the only accepted value.",
)
@_pipeline_option @_pipeline_option
@click.option( @_seed_option
"--device", @_hf_token_option
type=click.Choice(["auto", "cpu", "mps", "cuda", "xpu"]), @_humanize_option
default="auto", @_max_resolution_option
help="Inference device.",
)
@click.option(
"--seed",
type=int,
default=None,
help="Random seed for reproducibility. Default: 0 for qwen-zimage, random otherwise.",
)
@click.option("--hf-token", type=str, default=None, help="HuggingFace API token.")
@click.option(
"--humanize", type=float, default=0.0, help="Analog Humanizer film grain intensity (0 = off, typical: 2.0-6.0)."
)
@click.option(
"--max-resolution",
type=int,
default=0,
help="Cap long side (px) before diffusion; 0 = native and preserves the most detail. Raise only on GPU/MPS OOM.",
)
@_controlnet_scale_option @_controlnet_scale_option
@_unsharp_option @_unsharp_option
@_model_option
@_guidance_scale_option
@_auto_option
@_adaptive_polish_option @_adaptive_polish_option
@_tile_options @_tile_options
@_force_option @_force_option
@@ -877,19 +809,14 @@ def cmd_invisible(
source: Path, source: Path,
output: Path | None, output: Path | None,
strength: float | None, strength: float | None,
steps: int | None,
pipeline: str, pipeline: str,
device: str,
seed: int | None, seed: int | None,
hf_token: str | None, hf_token: str | None,
humanize: float, humanize: float,
unsharp: float, unsharp: float,
max_resolution: int, max_resolution: int,
controlnet_scale: float, controlnet_scale: float,
model: str | None, adaptive_polish: bool | None,
guidance_scale: float | None,
auto: bool,
adaptive_polish: bool,
tile: bool, tile: bool,
tile_size: int, tile_size: int,
tile_overlap: int, tile_overlap: int,
@@ -898,29 +825,24 @@ def cmd_invisible(
) -> None: ) -> None:
"""Remove invisible AI watermarks (SynthID, StableSignature, TreeRing). """Remove invisible AI watermarks (SynthID, StableSignature, TreeRing).
Uses diffusion-based regeneration. Requires GPU for reasonable speed. Regenerates the pixels with the two-stage diffusion profile. CUDA-only:
Requires the [diffusion] extra: pip install 'remove-ai-watermarks[diffusion]' pip install 'remove-ai-watermarks[qwen-zimage]'
""" """
from remove_ai_watermarks.invisible_engine import is_available as invisible_available from remove_ai_watermarks.invisible_engine import is_available as invisible_available
if not invisible_available(): if not invisible_available():
console.print( console.print(
"Error: Diffusion dependencies not installed.\n" "Error: the invisible-removal dependencies are not installed.\n"
" Install them with: pip install 'remove-ai-watermarks[diffusion]'" f" Install them with: pip install {INVISIBLE_EXTRA}"
) )
raise SystemExit(1) raise SystemExit(1)
from remove_ai_watermarks.invisible_engine import InvisibleEngine from remove_ai_watermarks.invisible_engine import InvisibleEngine
source = _validate_image(source) source = _validate_image(source)
steps = resolve_steps(steps)
seed = resolve_seed(seed)
adaptive_polish = _resolve_profile_polish(auto, adaptive_polish, pipeline)
if output is None: if output is None:
output = source.with_stem(source.stem + "_clean") output = source.with_stem(source.stem + "_clean")
device_str = None if device == "auto" else device
# Gate BEFORE building the engine: skip the destructive regeneration when no # Gate BEFORE building the engine: skip the destructive regeneration when no
# invisible AI watermark is locally detectable (it would only degrade a clean # invisible AI watermark is locally detectable (it would only degrade a clean
# image -- dominant paid score-0 cause), so the common skip path pays nothing for # image -- dominant paid score-0 cause), so the common skip path pays nothing for
@@ -932,8 +854,6 @@ def cmd_invisible(
console.print(f" {msg}") console.print(f" {msg}")
engine = InvisibleEngine( engine = InvisibleEngine(
model_id=model,
device=device_str,
pipeline=pipeline, pipeline=pipeline,
hf_token=hf_token, hf_token=hf_token,
progress_callback=progress_cb, progress_callback=progress_cb,
@@ -946,15 +866,13 @@ def cmd_invisible(
vendor = vendor_for_strength(source) vendor = vendor_for_strength(source)
console.print(f" Input: {source.name}") console.print(f" Input: {source.name}")
console.print(f" Pipeline: {pipeline}") console.print(f" Pipeline: {pipeline}")
console.print(f" Strength: {_resolved_strength_for_display(source, strength, vendor, pipeline)} Steps: {steps}") console.print(f" Strength: {_resolved_strength_for_display(source, strength, vendor, pipeline)}")
t0 = time.monotonic() t0 = time.monotonic()
result_path = engine.remove_watermark( result_path = engine.remove_watermark(
image_path=source, image_path=source,
output_path=output, output_path=output,
strength=strength, strength=strength,
num_inference_steps=steps,
guidance_scale=guidance_scale,
seed=seed, seed=seed,
humanize=humanize, humanize=humanize,
unsharp=unsharp, unsharp=unsharp,
@@ -1516,40 +1434,13 @@ def cmd_identify(ctx: click.Context, source: Path, no_visible: bool, as_json: bo
@_visible_backend_option @_visible_backend_option
@_visible_sensitivity_option @_visible_sensitivity_option
@_strength_option @_strength_option
@click.option(
"--steps",
type=int,
default=None,
help="Number of denoising steps. Both profiles are distilled four-step schedules, so 4 is the only accepted value.",
)
@_pipeline_option @_pipeline_option
@_model_option @_seed_option
@click.option( @_hf_token_option
"--device", @_humanize_option
type=click.Choice(["auto", "cpu", "mps", "cuda", "xpu"]), @_max_resolution_option
default="auto",
help="Inference device.",
)
@click.option(
"--seed",
type=int,
default=None,
help="Random seed for reproducibility. Default: 0 for qwen-zimage, random otherwise.",
)
@click.option("--hf-token", type=str, default=None, help="HuggingFace API token.")
@click.option(
"--humanize", type=float, default=0.0, help="Analog Humanizer film grain intensity (0 = off, typical: 2.0-6.0)."
)
@click.option(
"--max-resolution",
type=int,
default=0,
help="Cap long side (px) before diffusion; 0 = native and preserves the most detail. Raise only on GPU/MPS OOM.",
)
@_controlnet_scale_option @_controlnet_scale_option
@_unsharp_option @_unsharp_option
@_guidance_scale_option
@_auto_option
@_adaptive_polish_option @_adaptive_polish_option
@_tile_options @_tile_options
@_force_option @_force_option
@@ -1562,19 +1453,14 @@ def cmd_all(
backend: str, backend: str,
sensitivity: str, sensitivity: str,
strength: float | None, strength: float | None,
steps: int | None,
pipeline: str, pipeline: str,
model: str | None,
device: str,
seed: int | None, seed: int | None,
hf_token: str | None, hf_token: str | None,
humanize: float, humanize: float,
unsharp: float, unsharp: float,
max_resolution: int, max_resolution: int,
controlnet_scale: float, controlnet_scale: float,
guidance_scale: float | None, adaptive_polish: bool | None,
auto: bool,
adaptive_polish: bool,
tile: bool, tile: bool,
tile_size: int, tile_size: int,
tile_overlap: int, tile_overlap: int,
@@ -1592,9 +1478,6 @@ def cmd_all(
""" """
_banner() _banner()
source = _validate_image(source) source = _validate_image(source)
steps = resolve_steps(steps)
seed = resolve_seed(seed)
adaptive_polish = _resolve_profile_polish(auto, adaptive_polish, pipeline)
if output is None: if output is None:
output = source.with_stem(source.stem + "_clean") output = source.with_stem(source.stem + "_clean")
@@ -1649,7 +1532,7 @@ def cmd_all(
synthid_skipped = True synthid_skipped = True
console.print( console.print(
" Warning: Skipped - GPU dependencies not installed.\n" " Warning: Skipped - GPU dependencies not installed.\n"
" Install them with: pip install 'remove-ai-watermarks[diffusion]'" f" Install them with: pip install {INVISIBLE_EXTRA}"
) )
elif _should_skip_invisible_scrub(force, source): elif _should_skip_invisible_scrub(force, source):
# No locally-detectable invisible watermark -> skip the destructive # No locally-detectable invisible watermark -> skip the destructive
@@ -1666,14 +1549,10 @@ def cmd_all(
else: else:
from remove_ai_watermarks.invisible_engine import InvisibleEngine from remove_ai_watermarks.invisible_engine import InvisibleEngine
device_str = None if device == "auto" else device
def progress_cb(msg: str) -> None: def progress_cb(msg: str) -> None:
console.print(f" {msg}") console.print(f" {msg}")
inv_engine = InvisibleEngine( inv_engine = InvisibleEngine(
model_id=model,
device=device_str,
pipeline=pipeline, pipeline=pipeline,
hf_token=hf_token, hf_token=hf_token,
progress_callback=progress_cb, progress_callback=progress_cb,
@@ -1685,15 +1564,11 @@ def cmd_all(
# already lost its C2PA to the visible-removal pass, so reading it would # already lost its C2PA to the visible-removal pass, so reading it would
# always resolve to the unknown-vendor default. # always resolve to the unknown-vendor default.
vendor = vendor_for_strength(source) vendor = vendor_for_strength(source)
console.print( console.print(f" Strength: {_resolved_strength_for_display(source, strength, vendor, pipeline)}")
f" Strength: {_resolved_strength_for_display(source, strength, vendor, pipeline)} Steps: {steps}"
)
inv_engine.remove_watermark( inv_engine.remove_watermark(
image_path=tmp_path, image_path=tmp_path,
output_path=tmp_path, output_path=tmp_path,
strength=strength, strength=strength,
num_inference_steps=steps,
guidance_scale=guidance_scale,
seed=seed, seed=seed,
humanize=humanize, humanize=humanize,
unsharp=unsharp, unsharp=unsharp,
@@ -1753,7 +1628,7 @@ def cmd_all(
" visible mark and metadata were stripped.\n" " visible mark and metadata were stripped.\n"
"\n" "\n"
" Install the extra and rerun to remove it:\n" " Install the extra and rerun to remove it:\n"
" pip install 'remove-ai-watermarks[diffusion]'\n" f" pip install {INVISIBLE_EXTRA}\n"
" =====================================================================" " ====================================================================="
) )
raise SystemExit(1) raise SystemExit(1)
@@ -1775,15 +1650,13 @@ class _BatchOptions:
"""Validated processing options shared by every image in one batch. """Validated processing options shared by every image in one batch.
Click necessarily exposes these as individual command parameters, but the Click necessarily exposes these as individual command parameters, but the
processing core should receive one coherent value instead of a 21-argument processing core should receive one coherent value instead of a long positional
call. Keeping the object immutable also makes it safe to reuse while the call. Keeping the object immutable also makes it safe to reuse while the
batch caches model instances in ``ctx.obj``. batch caches model instances in ``ctx.obj``.
""" """
strength: float | None strength: float | None
steps: int
pipeline: str pipeline: str
device: str
seed: int | None seed: int | None
hf_token: str | None hf_token: str | None
humanize: float humanize: float
@@ -1792,9 +1665,8 @@ class _BatchOptions:
unsharp: float = 0.0 unsharp: float = 0.0
max_resolution: int = 0 max_resolution: int = 0
controlnet_scale: float = 1.0 controlnet_scale: float = 1.0
model: str | None = None # None means "the user did not choose"; the library resolves it per profile.
guidance_scale: float | None = None adaptive_polish: bool | None = None
adaptive_polish: bool = False
tile: bool = False tile: bool = False
tile_size: int = 1024 tile_size: int = 1024
tile_overlap: int = 128 tile_overlap: int = 128
@@ -1828,8 +1700,6 @@ def _run_batch_invisible(
engines = ctx.obj.setdefault("_inv_engines", {}) engines = ctx.obj.setdefault("_inv_engines", {})
if options.pipeline not in engines: if options.pipeline not in engines:
engines[options.pipeline] = InvisibleEngine( engines[options.pipeline] = InvisibleEngine(
model_id=options.model,
device=None if options.device == "auto" else options.device,
pipeline=options.pipeline, pipeline=options.pipeline,
hf_token=options.hf_token, hf_token=options.hf_token,
controlnet_conditioning_scale=options.controlnet_scale, controlnet_conditioning_scale=options.controlnet_scale,
@@ -1839,8 +1709,6 @@ def _run_batch_invisible(
img_path if mode == "invisible" else out_path, img_path if mode == "invisible" else out_path,
out_path, out_path,
strength=options.strength, strength=options.strength,
num_inference_steps=options.steps,
guidance_scale=options.guidance_scale,
seed=options.seed, seed=options.seed,
humanize=options.humanize, humanize=options.humanize,
unsharp=options.unsharp, unsharp=options.unsharp,
@@ -1948,42 +1816,15 @@ def _process_batch_image(
"--mode", type=click.Choice(["visible", "invisible", "metadata", "all"]), default="visible", help="Processing mode." "--mode", type=click.Choice(["visible", "invisible", "metadata", "all"]), default="visible", help="Processing mode."
) )
@_strength_option @_strength_option
@click.option(
"--steps",
type=int,
default=None,
help="Number of denoising steps. Both profiles are distilled four-step schedules, so 4 is the only accepted value.",
)
@_visible_backend_option @_visible_backend_option
@_visible_sensitivity_option @_visible_sensitivity_option
@click.option( @_humanize_option
"--humanize", type=float, default=0.0, help="Analog Humanizer film grain intensity (0 = off, typical: 2.0-6.0)."
)
@_pipeline_option @_pipeline_option
@click.option( @_seed_option
"--device", @_hf_token_option
type=click.Choice(["auto", "cpu", "mps", "cuda", "xpu"]), @_max_resolution_option
default="auto",
help="Inference device.",
)
@click.option(
"--seed",
type=int,
default=None,
help="Random seed for reproducibility. Default: 0 for qwen-zimage, random otherwise.",
)
@click.option("--hf-token", type=str, default=None, help="HuggingFace API token.")
@click.option(
"--max-resolution",
type=int,
default=0,
help="Cap long side (px) before diffusion; 0 = native and preserves the most detail. Raise only on GPU/MPS OOM.",
)
@_unsharp_option @_unsharp_option
@_controlnet_scale_option @_controlnet_scale_option
@_model_option
@_guidance_scale_option
@_auto_option
@_adaptive_polish_option @_adaptive_polish_option
@_tile_options @_tile_options
@_force_option @_force_option
@@ -1995,9 +1836,7 @@ def cmd_batch(
mode: str, mode: str,
output_dir: Path | None, output_dir: Path | None,
strength: float | None, strength: float | None,
steps: int | None,
pipeline: str, pipeline: str,
device: str,
seed: int | None, seed: int | None,
hf_token: str | None, hf_token: str | None,
backend: str, backend: str,
@@ -2006,10 +1845,7 @@ def cmd_batch(
unsharp: float, unsharp: float,
max_resolution: int, max_resolution: int,
controlnet_scale: float, controlnet_scale: float,
model: str | None, adaptive_polish: bool | None,
guidance_scale: float | None,
auto: bool,
adaptive_polish: bool,
tile: bool, tile: bool,
tile_size: int, tile_size: int,
tile_overlap: int, tile_overlap: int,
@@ -2032,14 +1868,9 @@ def cmd_batch(
console.print(f" Found {len(images)} images in {directory}") console.print(f" Found {len(images)} images in {directory}")
console.print(f" Output -> {output_dir}") console.print(f" Output -> {output_dir}")
console.print(f" Mode: {mode}") console.print(f" Mode: {mode}")
adaptive_polish = _resolve_profile_polish(auto, adaptive_polish, pipeline)
steps = resolve_steps(steps)
seed = resolve_seed(seed)
options = _BatchOptions( options = _BatchOptions(
strength=strength, strength=strength,
steps=steps,
pipeline=pipeline, pipeline=pipeline,
device=device,
seed=seed, seed=seed,
hf_token=hf_token, hf_token=hf_token,
humanize=humanize, humanize=humanize,
@@ -2048,8 +1879,6 @@ def cmd_batch(
unsharp=unsharp, unsharp=unsharp,
max_resolution=max_resolution, max_resolution=max_resolution,
controlnet_scale=controlnet_scale, controlnet_scale=controlnet_scale,
model=model,
guidance_scale=guidance_scale,
adaptive_polish=adaptive_polish, adaptive_polish=adaptive_polish,
tile=tile, tile=tile,
tile_size=tile_size, tile_size=tile_size,
@@ -2103,7 +1932,7 @@ def cmd_batch(
f"\n WARNING: the invisible (SynthID) watermark was NOT removed on " f"\n WARNING: the invisible (SynthID) watermark was NOT removed on "
f"{synthid_skipped_count} image(s) -- the GPU dependencies are not installed, " f"{synthid_skipped_count} image(s) -- the GPU dependencies are not installed, "
f"so those outputs still carry the invisible watermark.\n" f"so those outputs still carry the invisible watermark.\n"
f" Install the extra and rerun: pip install 'remove-ai-watermarks[diffusion]'" f" Install the extra and rerun: pip install {INVISIBLE_EXTRA}"
) )
# Non-zero exit so a wrapping service detects an incomplete/failed run (batch used # Non-zero exit so a wrapping service detects an incomplete/failed run (batch used
+31 -59
View File
@@ -1,7 +1,7 @@
"""Diffusion engine for regenerating images that carry invisible AI watermarks. """Diffusion engine for regenerating images that carry invisible AI watermarks.
This module requires the 'gpu' extra dependencies: Requires the 'qwen-zimage' extra and a CUDA device:
uv pip install 'remove-ai-watermarks[diffusion]' uv pip install 'remove-ai-watermarks[qwen-zimage]'
""" """
# cv2/torch boundary: this engine wraps cv2 (resize/imwrite/cvtColor) and the # cv2/torch boundary: this engine wraps cv2 (resize/imwrite/cvtColor) and the
@@ -16,13 +16,11 @@ import warnings
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from ._internal.watermark_profiles import (
DEFAULT_MODEL_ID as DEFAULT_SDXL_MODEL_ID,
)
from ._internal.watermark_profiles import ( from ._internal.watermark_profiles import (
DEFAULT_PROFILE, DEFAULT_PROFILE,
REMOVAL_MODULES,
resolve_adaptive_polish,
resolve_seed, resolve_seed,
resolve_steps,
) )
if TYPE_CHECKING: if TYPE_CHECKING:
@@ -42,10 +40,15 @@ logger = logging.getLogger(__name__)
def is_available() -> bool: def is_available() -> bool:
"""Check if invisible watermark removal dependencies are installed.""" """Whether the dependencies for a real removal run are installed.
Shares :data:`REMOVAL_MODULES` with the remover's own precondition so the two
cannot drift. When they did, a torch+diffusers-only environment passed this gate
and then died at the DiffSynth face stage.
"""
from .optional_deps import module_available from .optional_deps import module_available
return module_available("diffusers", "torch") return module_available(*REMOVAL_MODULES)
def _target_size(width: int, height: int, max_resolution: int) -> tuple[int, int] | None: def _target_size(width: int, height: int, max_resolution: int) -> tuple[int, int] | None:
@@ -79,13 +82,8 @@ class InvisibleEngine:
to break watermark patterns, and reconstructs via reverse diffusion. to break watermark patterns, and reconstructs via reverse diffusion.
""" """
# SDXL base is the default since May 2026; the vendor-adaptive strength
# removes the current SynthID (see watermark_profiles + docs/synthid.md).
DEFAULT_MODEL_ID = DEFAULT_SDXL_MODEL_ID
def __init__( def __init__(
self, self,
model_id: str | None = None,
device: str | None = None, device: str | None = None,
pipeline: str = DEFAULT_PROFILE, pipeline: str = DEFAULT_PROFILE,
hf_token: str | None = None, hf_token: str | None = None,
@@ -96,8 +94,9 @@ class InvisibleEngine:
"""Initialize the invisible watermark removal engine. """Initialize the invisible watermark removal engine.
Args: Args:
model_id: HuggingFace model ID. None = use the SDXL base default. device: Device for inference. Both profiles are CUDA-only, so the
device: Device for inference (auto/cpu/mps/cuda/xpu). None = auto. usable values are "cuda" and None/"auto" (which detects it);
anything else raises rather than falling back.
pipeline: Pipeline profile, one of "qwen-zimage" (DEFAULT; pipeline: Pipeline profile, one of "qwen-zimage" (DEFAULT;
Qwen-Image-2512 Lightning + Canny, then SAM-masked Z-Image face repair) Qwen-Image-2512 Lightning + Canny, then SAM-masked Z-Image face repair)
or "sdxl-zimage" (the same recipe and the same face stage on an SDXL or "sdxl-zimage" (the same recipe and the same face stage on an SDXL
@@ -116,11 +115,7 @@ class InvisibleEngine:
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
# Pass model_id through untouched. Substituting DEFAULT_MODEL_ID for None here
# meant the engine always supplied a model the remover is required to reject,
# so EVERY construction raised once that check tightened to "is not None".
self._remover = WatermarkRemover( self._remover = WatermarkRemover(
model_id=model_id,
device=device, device=device,
progress_callback=progress_callback, progress_callback=progress_callback,
hf_token=hf_token, hf_token=hf_token,
@@ -144,14 +139,12 @@ class InvisibleEngine:
image_path: Path, image_path: Path,
output_path: Path | None = None, output_path: Path | None = None,
strength: float | None = None, strength: float | None = None,
num_inference_steps: int | None = None,
guidance_scale: float | None = None,
seed: int | None = None, seed: int | None = None,
humanize: float = 0.0, humanize: float = 0.0,
max_resolution: int = 0, max_resolution: int = 0,
vendor: str | None = None, vendor: str | None = None,
unsharp: float = 0.0, unsharp: float = 0.0,
adaptive_polish: bool = False, adaptive_polish: bool | None = None,
tile: bool = False, tile: bool = False,
tile_size: int = 1024, tile_size: int = 1024,
tile_overlap: int = 128, tile_overlap: int = 128,
@@ -161,27 +154,26 @@ class InvisibleEngine:
Args: Args:
image_path: Path to the watermarked image. image_path: Path to the watermarked image.
output_path: Output path (None = overwrite source). output_path: Output path (None = overwrite source).
strength: Denoising strength (0.0-1.0). None -> the vendor-adaptive strength: Denoising strength (0.0-1.0). None -> the profile's calibrated
default. default (resolution-adaptive for qwen-zimage, vendor-adaptive for
num_inference_steps: Number of denoising steps. None keeps the existing sdxl-zimage).
100-step library default, except qwen-zimage uses its required seed: Random seed for reproducibility. None resolves to 0, because both
four-step Lightning schedule. profiles are certified at a fixed seed.
guidance_scale: Classifier-free guidance scale.
seed: Random seed for reproducibility. None resolves to 0 for
qwen-zimage and stays random for the other profiles.
humanize: Intensity of Analog Humanizer film grain (0 = off). humanize: Intensity of Analog Humanizer film grain (0 = off).
unsharp: Final unsharp-mask sharpening strength (0 = off, default). unsharp: Final unsharp-mask sharpening strength (0 = off, default).
Applied last to counter the soft / over-smoothed look of the Applied last to counter the soft / over-smoothed look of the
diffusion pass; ~0.5-0.8 is a safe range, higher risks edge halos. diffusion pass; ~0.5-0.8 is a safe range, higher risks edge halos.
adaptive_polish: When True (the CLI default), restore the input's detail adaptive_polish: Restore the input's detail level in the softened
level in the softened output: a capped unsharp + edge-masked grain output: a capped unsharp + edge-masked grain targeting the input's
targeting the input's Laplacian variance. Self-limiting -- a no-op when Laplacian variance. Self-limiting -- a no-op when the output already
the output already meets the input's detail level (text/flat graphics), meets the input's detail level (text/flat graphics), so it only acts on
so it only acts on over-smoothed photo/face texture. Runs LAST. over-smoothed photo/face texture. Runs LAST. None (the default) follows
the profile: off for qwen-zimage, on for sdxl-zimage. This resolves
through the same ``resolve_adaptive_polish`` the CLI uses, so a library
caller and a CLI caller on one profile get the same output.
max_resolution: Cap the long side (px) before diffusion. 0 (default) max_resolution: Cap the long side (px) before diffusion. 0 (default)
= no cap. Set a positive value only to bound GPU/MPS memory on = no cap. Set a positive value only to bound GPU memory on very large
very large inputs (it reintroduces a lossy downscale->upscale inputs (it reintroduces a lossy downscale->upscale round-trip).
round-trip).
tile: Process the diffusion pass in overlapping tiles instead of one tile: Process the diffusion pass in overlapping tiles instead of one
forward pass. This retains the input's native dimensions instead forward pass. This retains the input's native dimensions instead
of applying ``max_resolution``, but each tile is still regenerated. of applying ``max_resolution``, but each tile is still regenerated.
@@ -194,8 +186,8 @@ class InvisibleEngine:
""" """
import tempfile import tempfile
num_inference_steps = resolve_steps(num_inference_steps)
seed = resolve_seed(seed) seed = resolve_seed(seed)
adaptive_polish = resolve_adaptive_polish(adaptive_polish, self._remover.model_profile)
from PIL import Image, ImageOps from PIL import Image, ImageOps
@@ -243,8 +235,6 @@ class InvisibleEngine:
image_path=image_path, image_path=image_path,
output_path=output_path, output_path=output_path,
strength=strength, strength=strength,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
seed=seed, seed=seed,
vendor=vendor, vendor=vendor,
tile=tile, tile=tile,
@@ -315,21 +305,3 @@ class InvisibleEngine:
# _tmp_path is always set above (we persist the image unconditionally). # _tmp_path is always set above (we persist the image unconditionally).
if _tmp_path.exists(): if _tmp_path.exists():
_tmp_path.unlink() _tmp_path.unlink()
def remove_watermark_batch(
self,
input_dir: Path,
output_dir: Path,
strength: float | None = None,
steps: int | None = None,
) -> list[Path]:
"""Remove invisible watermarks from all images in a directory."""
if steps is None:
profile = getattr(self._remover, "model_profile", None)
steps = 4 if profile in {"qwen-zimage", "sdxl-zimage"} else 50
return self._remover.remove_watermark_batch(
input_dir=input_dir,
output_dir=output_dir,
strength=strength,
num_inference_steps=steps,
)
-38
View File
@@ -1138,14 +1138,6 @@ def _scan_video_detectors(
} }
def _scan_video(
source: Path,
detector: Any,
) -> VideoScan:
"""Decode a video once and collect one untrusted candidate per frame."""
return _scan_video_detectors(source, {"selected": detector})["selected"]
def scan_video_marks( def scan_video_marks(
source: Path, source: Path,
marks: tuple[str, ...] = VIDEO_VISIBLE_MARKS, marks: tuple[str, ...] = VIDEO_VISIBLE_MARKS,
@@ -1181,36 +1173,6 @@ def scan_video_marks(
) )
def scan_sora_video(source: Path) -> VideoScan:
"""Decode a video once and collect one untrusted Sora candidate per frame."""
return _scan_video(source, detect_sora_frame)
def scan_veo_video(source: Path) -> VideoScan:
"""Decode a video once and collect one untrusted Veo candidate per frame."""
return _scan_video(source, detect_veo_frame)
def scan_seedance_video(source: Path) -> VideoScan:
"""Decode a video once and collect one untrusted Seedance candidate per frame."""
return _scan_video(source, detect_seedance_frame)
def scan_dola_video(source: Path) -> VideoScan:
"""Decode a video once and collect one untrusted Dola candidate per frame."""
return _scan_video(source, detect_dola_frame)
def scan_hailuo_video(source: Path) -> VideoScan:
"""Decode a video once and collect one untrusted Hailuo candidate per frame."""
return _scan_video(source, detect_hailuo_frame)
def scan_kling_video(source: Path) -> VideoScan:
"""Decode a video once and collect one untrusted Kling candidate per frame."""
return _scan_video(source, detect_kling_frame)
def _mask_for_region( def _mask_for_region(
frame_bgr: NDArray[Any], frame_bgr: NDArray[Any],
region: Region, region: Region,
+52 -35
View File
@@ -327,7 +327,13 @@ class TestInvisibleCommand:
expected = sample_png.with_stem(sample_png.stem + "_clean") expected = sample_png.with_stem(sample_png.stem + "_clean")
assert expected.exists() assert expected.exists()
def test_invisible_adaptive_polish_off_by_default_under_qwen_zimage(self, runner, sample_png): 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() mock_cls, mock_engine = _mock_invisible_engine()
with ( with (
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True), patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
@@ -336,13 +342,7 @@ class TestInvisibleCommand:
): ):
result = runner.invoke(main, ["invisible", str(sample_png), "--force"]) result = runner.invoke(main, ["invisible", str(sample_png), "--force"])
assert result.exit_code == 0, result.output assert result.exit_code == 0, result.output
# The default profile is qwen-zimage, and _resolve_profile_polish keeps its assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is None
# output untouched unless polish was asked for explicitly. It stays available:
# passing --adaptive-polish still turns it on (covered separately).
assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is False
# Default model is None (the SDXL base) and CFG is None (the library's 7.5).
assert mock_cls.call_args.kwargs["model_id"] is None
assert mock_engine.remove_watermark.call_args.kwargs["guidance_scale"] is None
def test_invisible_no_adaptive_polish_disables(self, runner, sample_png): def test_invisible_no_adaptive_polish_disables(self, runner, sample_png):
mock_cls, mock_engine = _mock_invisible_engine() mock_cls, mock_engine = _mock_invisible_engine()
@@ -355,20 +355,24 @@ class TestInvisibleCommand:
assert result.exit_code == 0, result.output assert result.exit_code == 0, result.output
assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is False assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is False
def test_invisible_model_and_guidance_scale_flow_to_engine(self, runner, sample_png): def test_knobs_the_fixed_stack_cannot_honor_are_not_offered(self, runner, sample_png):
mock_cls, mock_engine = _mock_invisible_engine() """--model/--steps/--guidance-scale/--device/--auto are gone, not rejected.
with (
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True), Each pinned a value the profiles fix (model stack, per-stage schedule, CFG 1.0,
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True), CUDA), so accepting one only produced an error several layers down -- a flag
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls), 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.
"""
for retired in (
["--model", "org/custom-sdxl"],
["--steps", "20"],
["--guidance-scale", "5.5"],
["--device", "cpu"],
["--auto"],
): ):
result = runner.invoke( result = runner.invoke(main, ["invisible", str(sample_png), *retired, "--force"])
main, assert result.exit_code == 2, f"{retired[0]}: {result.output}"
["invisible", str(sample_png), "--model", "org/custom-sdxl", "--guidance-scale", "5.5", "--force"], assert "No such option" in result.output, f"{retired[0]}: {result.output}"
)
assert result.exit_code == 0, result.output
assert mock_cls.call_args.kwargs["model_id"] == "org/custom-sdxl"
assert mock_engine.remove_watermark.call_args.kwargs["guidance_scale"] == 5.5
def test_retired_pipeline_names_are_rejected_not_silently_remapped(self, runner, sample_png): 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. """default/sdxl/controlnet/qwen were removed with their CPU code paths.
@@ -530,7 +534,7 @@ class TestAllCommand:
result = runner.invoke(main, ["all", str(sample_png), "-o", str(output)]) result = runner.invoke(main, ["all", str(sample_png), "-o", str(output)])
assert result.exit_code != 0, result.output assert result.exit_code != 0, result.output
assert "NOT removed" in result.output assert "NOT removed" in result.output
assert "remove-ai-watermarks[diffusion]" in result.output assert "remove-ai-watermarks[qwen-zimage]" in result.output
assert output.exists() # visible + metadata still produced a file assert output.exists() # visible + metadata still produced a file
def test_all_reports_metadata_that_survived_stripping(self, runner, sample_png, tmp_path): def test_all_reports_metadata_that_survived_stripping(self, runner, sample_png, tmp_path):
@@ -858,10 +862,11 @@ class TestBatchCommand:
assert out[0, 0, 3] == 0 assert out[0, 0, 3] == 0
assert out[100, 100, 3] == 255 assert out[100, 100, 3] == 255
def test_batch_auto_is_deprecated_and_enables_polish(self, runner, tmp_path): def test_batch_explicit_adaptive_polish_overrides_the_qwen_zimage_off(self, runner, tmp_path):
"""--auto is retired: it warns and just enables the adaptive polish. """qwen-zimage leaves the polish off by default; a typed flag still turns it on.
It no longer selects a pipeline: qwen-zimage is the only default there is. The off is a parameter-source check, not a changed default, so it must yield to
an explicit --adaptive-polish rather than swallowing it.
""" """
input_dir = _make_batch_dir(tmp_path, count=2) input_dir = _make_batch_dir(tmp_path, count=2)
output_dir = tmp_path / "output" output_dir = tmp_path / "output"
@@ -874,12 +879,19 @@ class TestBatchCommand:
): ):
result = runner.invoke( result = runner.invoke(
main, main,
["batch", str(input_dir), "-o", str(output_dir), "--mode", "invisible", "--auto", "--force"], [
"batch",
str(input_dir),
"-o",
str(output_dir),
"--mode",
"invisible",
"--adaptive-polish",
"--force",
],
) )
assert result.exit_code == 0, result.output assert result.exit_code == 0, result.output
assert "2 processed" in result.output assert "2 processed" in result.output
assert "deprecated" in result.output.lower()
# Pipeline stays the default controlnet; --auto only turned the polish on.
assert mock_cls.call_args.kwargs["pipeline"] == "qwen-zimage" assert mock_cls.call_args.kwargs["pipeline"] == "qwen-zimage"
assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is True assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is True
@@ -923,21 +935,26 @@ class TestBatchCommand:
class TestGpuHintMarkup: class TestGpuHintMarkup:
"""The diffusion install hint must reach the user with the ``[diffusion]`` token """The install hint must name the extra that actually makes a removal run.
intact (plain output prints it verbatim, with no markup parsing)."""
def test_invisible_install_hint_keeps_gpu_extra(self, runner, sample_png): 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_invisible_install_hint_names_the_working_extra(self, runner, sample_png):
with patch("remove_ai_watermarks.invisible_engine.is_available", return_value=False): with patch("remove_ai_watermarks.invisible_engine.is_available", return_value=False):
result = runner.invoke(main, ["invisible", str(sample_png)]) result = runner.invoke(main, ["invisible", str(sample_png)])
assert result.exit_code != 0 assert result.exit_code != 0
assert "remove-ai-watermarks[diffusion]" in result.output assert "remove-ai-watermarks[qwen-zimage]" in result.output
def test_all_install_hint_keeps_gpu_extra(self, runner, sample_png): 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 `all` pipeline skips the invisible step with a warning that carries
# the same hint; it must keep the [diffusion] extra too. # the same hint; it must name the same extra.
with patch("remove_ai_watermarks.invisible_engine.is_available", return_value=False): with patch("remove_ai_watermarks.invisible_engine.is_available", return_value=False):
result = runner.invoke(main, ["all", str(sample_png)]) result = runner.invoke(main, ["all", str(sample_png)])
assert "remove-ai-watermarks[diffusion]" in result.output assert "remove-ai-watermarks[qwen-zimage]" in result.output
class TestEraseCommand: class TestEraseCommand:
+27 -23
View File
@@ -16,24 +16,26 @@ class TestIsAvailable:
result = is_available() result = is_available()
assert isinstance(result, bool) assert isinstance(result, bool)
def test_available_reflects_dependencies(self): def test_available_reflects_every_module_a_run_needs(self):
"""is_available() is True iff torch + diffusers (the diffusion extra) import. """True iff every module in REMOVAL_MODULES imports, diffsynth included.
Must not assume the full stack: the default+dev CI env has no diffusers. Derived from the same tuple the remover's precondition uses, so this cannot
pass while the two disagree -- the drift that let a torch+diffusers-only
environment clear the CLI gate and then die at the DiffSynth face stage.
Must not assume the full stack: the default+dev CI env has none of it.
""" """
import importlib.util import importlib.util
expected = all(importlib.util.find_spec(m) is not None for m in ("torch", "diffusers")) from remove_ai_watermarks._internal.watermark_profiles import REMOVAL_MODULES
assert "diffsynth" in REMOVAL_MODULES
expected = all(importlib.util.find_spec(m) is not None for m in REMOVAL_MODULES)
assert is_available() is expected assert is_available() is expected
class TestInvisibleEngineInit: class TestInvisibleEngineInit:
"""Tests for InvisibleEngine construction (no GPU required).""" """Tests for InvisibleEngine construction (no GPU required)."""
def test_default_model_id(self):
# SDXL base became the default in May 2026 (defeats SynthID v2).
assert InvisibleEngine.DEFAULT_MODEL_ID == "stabilityai/stable-diffusion-xl-base-1.0"
def test_preload_forwards_global_only(self): def test_preload_forwards_global_only(self):
engine = object.__new__(InvisibleEngine) engine = object.__new__(InvisibleEngine)
engine._remover = SimpleNamespace(preload=lambda **kwargs: setattr(engine, "_preload_kwargs", kwargs)) engine._remover = SimpleNamespace(preload=lambda **kwargs: setattr(engine, "_preload_kwargs", kwargs))
@@ -55,7 +57,7 @@ class TestNativeOutputSize:
Image.open(image_path).crop((0, 0, 24, 16)).save(out) Image.open(image_path).crop((0, 0, 24, 16)).save(out)
return out return out
engine._remover = SimpleNamespace(remove_watermark=_remove_watermark) engine._remover = SimpleNamespace(remove_watermark=_remove_watermark, model_profile="qwen-zimage")
engine._progress_callback = None engine._progress_callback = None
src = tmp_path / "src.png" src = tmp_path / "src.png"
out = tmp_path / "out.png" out = tmp_path / "out.png"
@@ -113,30 +115,32 @@ class TestTargetSize:
assert _target_size(381, 512, 4096) is None assert _target_size(381, 512, 4096) is None
class TestEngineDoesNotFabricateAModelId: class TestEngineConstructsWithoutAModelId:
"""The engine must forward model_id untouched, including None. """Plain construction must reach the remover, and must not name a model.
It used to substitute DEFAULT_MODEL_ID for None. Once the remover tightened its The engine used to take a ``model_id`` and substitute the SDXL default for None.
"you may not override the fixed stack" check from `not in {None, DEFAULT_MODEL_ID}` Once the remover tightened its "you may not override the fixed stack" check to
to `is not None`, that substitution made EVERY InvisibleEngine construction raise - ``is not None``, that substitution made EVERY construction raise -- and no test
and no test saw it, because the library tests build WatermarkRemover directly while saw it, because the library tests build WatermarkRemover directly while the engine
the engine tests mock it. A deployed Modal worker caught it instead. tests mock it. A deployed Modal worker caught it instead. The parameter is gone on
both sides now, so guard the property that broke: a default construction reaches
the remover, carrying no model at all.
""" """
def test_none_stays_none(self): def test_default_construction_names_no_model(self):
from unittest.mock import patch from unittest.mock import patch
import remove_ai_watermarks.invisible_engine as engine_module import remove_ai_watermarks.invisible_engine as engine_module
with patch("remove_ai_watermarks._internal.watermark_remover.WatermarkRemover") as remover: with patch("remove_ai_watermarks._internal.watermark_remover.WatermarkRemover") as remover:
engine_module.InvisibleEngine(pipeline="qwen-zimage") engine_module.InvisibleEngine(pipeline="qwen-zimage")
assert remover.call_args.kwargs["model_id"] is None assert remover.call_count == 1
assert "model_id" not in remover.call_args.kwargs
def test_an_explicit_model_id_still_reaches_the_remover_to_be_rejected(self): def test_a_model_id_is_not_accepted(self):
from unittest.mock import patch import pytest
import remove_ai_watermarks.invisible_engine as engine_module import remove_ai_watermarks.invisible_engine as engine_module
with patch("remove_ai_watermarks._internal.watermark_remover.WatermarkRemover") as remover: with pytest.raises(TypeError):
engine_module.InvisibleEngine(model_id="org/custom", pipeline="qwen-zimage") engine_module.InvisibleEngine(model_id="org/custom", pipeline="qwen-zimage") # type: ignore[call-arg]
assert remover.call_args.kwargs["model_id"] == "org/custom"
+57 -33
View File
@@ -1,7 +1,7 @@
"""Tests for cross-platform and cross-device compatibility. """Tests for device detection, profile resolution, and platform-specific paths.
Verifies that device detection, MPS fallback, and platform-specific Invisible-watermark removal is CUDA-only, so the device tests here assert a binary
code paths work correctly on CPU, MPS (macOS), and CUDA (Linux/Windows). answer and a clean refusal rather than a fallback ladder.
""" """
from __future__ import annotations from __future__ import annotations
@@ -27,34 +27,38 @@ from remove_ai_watermarks._internal.watermark_remover import get_device, is_wate
class TestDeviceDetection: class TestDeviceDetection:
"""Tests for get_device() across platforms.""" """get_device() is binary: CUDA, or the "cpu" that names its absence."""
def test_returns_valid_device(self): def test_answer_is_cuda_or_cpu(self):
device = get_device() """No mps/xpu answer exists. Both would travel one frame to the same refusal.
assert device in ("cpu", "mps", "cuda", "xpu")
def test_cpu_fallback_when_no_gpu(self): Reporting them anyway cost a device probe each and let a caller believe the
"""On CI / machines without GPU, should fall back to cpu or mps.""" library had an Apple-silicon or Intel-GPU path that it does not.
device = get_device() """
# Just verify it doesn't crash and returns a valid string assert get_device() in ("cpu", "cuda")
assert isinstance(device, str)
@patch("remove_ai_watermarks._internal.watermark_remover._HAS_TORCH", False) @patch("remove_ai_watermarks._internal.watermark_remover._HAS_TORCH", False)
def test_no_torch_returns_cpu(self): def test_no_torch_returns_cpu(self):
assert get_device() == "cpu" assert get_device() == "cpu"
def test_xpu_selected_when_available(self): def test_working_cuda_is_selected_and_probed(self):
"""An XPU-enabled torch (no CUDA) routes to the Intel GPU backend. """A reported CUDA device is smoke-tested before it is returned.
The whole torch module is mocked so the smoke-test ops succeed without torch.cuda.is_available() can be True on a build whose CUDA backend then
any real device; cuda must read False so the cuda branch is skipped. raises on the first real op; without the probe that surfaced much later.
""" """
fake_torch = MagicMock() fake_torch = MagicMock()
fake_torch.cuda.is_available.return_value = False fake_torch.cuda.is_available.return_value = True
fake_torch.xpu.is_available.return_value = True
with patch("remove_ai_watermarks._internal.watermark_remover.torch", fake_torch): with patch("remove_ai_watermarks._internal.watermark_remover.torch", fake_torch):
assert get_device() == "xpu" assert get_device() == "cuda"
fake_torch.tensor.assert_called_with([1.0], device="xpu") fake_torch.tensor.assert_called_with([1.0], device="cuda")
def test_broken_cuda_backend_falls_back_to_cpu(self):
fake_torch = MagicMock()
fake_torch.cuda.is_available.return_value = True
fake_torch.tensor.side_effect = RuntimeError("no kernel image")
with patch("remove_ai_watermarks._internal.watermark_remover.torch", fake_torch):
assert get_device() == "cpu"
def test_non_cuda_devices_are_refused_at_construction(self): def test_non_cuda_devices_are_refused_at_construction(self):
"""CUDA is a precondition of the object, not of the run. """CUDA is a precondition of the object, not of the run.
@@ -78,21 +82,21 @@ class TestDeviceDetection:
assert remover.device == "cuda" assert remover.device == "cuda"
assert remover.torch_dtype == torch.bfloat16 assert remover.torch_dtype == torch.bfloat16
def test_the_refusal_names_the_resolved_device_not_a_bare_none(self):
"""``device=None`` on a CUDA-less host must report "cpu", not "None".
class TestEmptyDeviceCache: The message used to interpolate the raw argument, so the common auto-detect
"""try_empty_device_cache is all that remains of the img2img runner. path told the user that ``'None'`` cannot run the removal.
"""
if not is_watermark_removal_available():
pytest.skip("torch/diffusers not installed")
from remove_ai_watermarks._internal import watermark_remover as module
Its module lost run_img2img and the MPS fallback along with the CPU/MPS profiles; with (
both surviving profiles are CUDA-only, so there is no MPS failure left to recover patch.object(module, "get_device", return_value="cpu"),
from. The helper must stay silent on a backend that cannot empty a cache, because pytest.raises(ValueError, match="'cpu' cannot run it"),
it runs in cleanup paths where a raise would replace the real error. ):
""" module.WatermarkRemover(device=None)
def test_unknown_backend_is_a_silent_no_op(self):
from remove_ai_watermarks._internal.watermark_remover import try_empty_device_cache
try_empty_device_cache("cpu")
try_empty_device_cache("definitely-not-a-backend")
class TestModelProfiles: class TestModelProfiles:
@@ -116,6 +120,26 @@ class TestModelProfiles:
assert normalize_profile(retired) not in PROFILE_CHOICES assert normalize_profile(retired) not in PROFILE_CHOICES
class TestResolveAdaptivePolish:
"""The polish default is per-profile data, not a CLI parameter-source inference."""
def test_unset_follows_the_profile(self):
from remove_ai_watermarks._internal.watermark_profiles import resolve_adaptive_polish
# qwen-zimage already matches the input's detail level, so polishing it only
# moves the output away from upstream. An SDXL global pass leaves the softer
# result the polish exists for.
assert resolve_adaptive_polish(None, "qwen-zimage") is False
assert resolve_adaptive_polish(None, "sdxl-zimage") is True
assert resolve_adaptive_polish(None, "qwen_zimage") is False
def test_an_explicit_choice_always_wins(self):
from remove_ai_watermarks._internal.watermark_profiles import resolve_adaptive_polish
assert resolve_adaptive_polish(True, "qwen-zimage") is True
assert resolve_adaptive_polish(False, "sdxl-zimage") is False
class TestNoReembeddedWatermark: class TestNoReembeddedWatermark:
"""F2 regression: the SDXL global stage must disable the diffusers watermarker. """F2 regression: the SDXL global stage must disable the diffusers watermarker.
+32 -27
View File
@@ -175,6 +175,7 @@ def test_cpu_offload_forces_both_stacks_to_stream(monkeypatch, cpu_offload, expe
monkeypatch.setattr(pipeline_module, "QwenZImagePipeline", Recorder) monkeypatch.setattr(pipeline_module, "QwenZImagePipeline", Recorder)
remover = module.WatermarkRemover.__new__(module.WatermarkRemover) remover = module.WatermarkRemover.__new__(module.WatermarkRemover)
remover.model_profile = "qwen-zimage"
remover.device = "cuda" remover.device = "cuda"
remover.torch_dtype = None remover.torch_dtype = None
remover.hf_token = None remover.hf_token = None
@@ -526,16 +527,13 @@ def test_face_composite_preserves_every_pixel_outside_mask():
assert np.all(result[12:20, 12:20] == 240) assert np.all(result[12:20, 12:20] == 240)
def test_profile_defaults_to_four_global_steps(): def test_profile_defaults_to_four_global_steps_and_a_fixed_seed():
from remove_ai_watermarks._internal.watermark_profiles import ( """The step count belongs to the stage, not to a caller-settable profile knob."""
normalize_profile, from remove_ai_watermarks._internal.qwen_zimage_pipeline import GLOBAL_STEPS
resolve_seed, from remove_ai_watermarks._internal.watermark_profiles import normalize_profile, resolve_seed
resolve_steps,
)
assert normalize_profile("qwen-zimage") == "qwen-zimage" assert normalize_profile("qwen-zimage") == "qwen-zimage"
assert resolve_steps(None) == 4 assert GLOBAL_STEPS == 4
assert resolve_steps(12) == 12
assert resolve_seed(None) == 0 assert resolve_seed(None) == 0
assert resolve_seed(17) == 17 assert resolve_seed(17) == 17
@@ -560,8 +558,11 @@ def test_cli_qwen_zimage_keeps_profile_postprocess_default(tmp_image_path, monke
) )
assert result.exit_code == 0, result.output assert result.exit_code == 0, result.output
assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is False # Both defaults are the profile's, resolved once by the library rather than
assert mock_engine.remove_watermark.call_args.kwargs["seed"] == 0 # pre-resolved here: the CLI passes them through unset so a library caller on the
# same profile gets the same answer.
assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is None
assert mock_engine.remove_watermark.call_args.kwargs["seed"] is None
result = CliRunner().invoke( result = CliRunner().invoke(
cli.main, cli.main,
@@ -590,7 +591,6 @@ def test_watermark_remover_dispatches_to_full_pipeline(tmp_path, monkeypatch):
runtime.run.return_value = Image.new("RGB", (64, 48), (50, 60, 70)) runtime.run.return_value = Image.new("RGB", (64, 48), (50, 60, 70))
remover = WatermarkRemover(device="cuda", pipeline="qwen-zimage") remover = WatermarkRemover(device="cuda", pipeline="qwen-zimage")
monkeypatch.setattr(remover, "_load_qwen_zimage_pipeline", lambda: runtime) monkeypatch.setattr(remover, "_load_qwen_zimage_pipeline", lambda: runtime)
assert remover.model_id == "Qwen/Qwen-Image-2512 + Tongyi-MAI/Z-Image-Turbo"
remover.remove_watermark( remover.remove_watermark(
source, source,
@@ -743,23 +743,30 @@ def test_watermark_remover_forwards_global_only_preload(monkeypatch):
runtime.preload.assert_called_once_with(global_only=True) runtime.preload.assert_called_once_with(global_only=True)
def test_qwen_zimage_rejects_runtime_knobs_that_change_fixed_graph(tmp_path, monkeypatch): def test_the_fixed_graph_offers_no_runtime_knob_to_reject(tmp_path, monkeypatch):
"""model_id, steps and CFG are not parameters at any layer.
They used to be accepted and then rejected, which put the failure several frames
below the caller and made the surface advertise choices the pinned stack cannot
honor. TypeError from the signature is the earlier, clearer answer -- and it is
what keeps a wrapper from threading a value that would silently do nothing.
"""
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
_mock_watermark_runtime_deps(monkeypatch) _mock_watermark_runtime_deps(monkeypatch)
with pytest.raises(ValueError, match="fixed Qwen-Image-2512"): with pytest.raises(TypeError):
WatermarkRemover(model_id="custom/model", device="cuda", pipeline="qwen-zimage") WatermarkRemover(model_id="custom/model", device="cuda", pipeline="qwen-zimage") # type: ignore[call-arg]
source = tmp_path / "source.png" source = tmp_path / "source.png"
Image.new("RGB", (64, 48)).save(source) Image.new("RGB", (64, 48)).save(source)
remover = WatermarkRemover(device="cuda", pipeline="qwen-zimage") remover = WatermarkRemover(device="cuda", pipeline="qwen-zimage")
with pytest.raises(ValueError, match=r"CFG 1\.0"): with pytest.raises(TypeError):
remover.remove_watermark(source, guidance_scale=2.0) remover.remove_watermark(source, guidance_scale=2.0) # type: ignore[call-arg]
with pytest.raises(ValueError, match="requires 4 steps"): with pytest.raises(TypeError):
remover.remove_watermark(source, num_inference_steps=8) remover.remove_watermark(source, num_inference_steps=8) # type: ignore[call-arg]
def test_invisible_engine_uses_qwen_zimage_step_default(tmp_image_path, tmp_path): def test_invisible_engine_passes_the_seed_but_never_a_step_count(tmp_image_path, tmp_path):
from remove_ai_watermarks.invisible_engine import InvisibleEngine from remove_ai_watermarks.invisible_engine import InvisibleEngine
engine = InvisibleEngine.__new__(InvisibleEngine) engine = InvisibleEngine.__new__(InvisibleEngine)
@@ -769,7 +776,10 @@ def test_invisible_engine_uses_qwen_zimage_step_default(tmp_image_path, tmp_path
engine.remove_watermark(tmp_image_path, tmp_path / "clean.png") engine.remove_watermark(tmp_image_path, tmp_path / "clean.png")
assert engine._remover.remove_watermark.call_args.kwargs["num_inference_steps"] == 4 kwargs = engine._remover.remove_watermark.call_args.kwargs
assert kwargs["seed"] == 0
assert "num_inference_steps" not in kwargs
assert "guidance_scale" not in kwargs
def test_sdxl_zimage_strength_is_vendor_adaptive_and_leaves_other_profiles_alone(): def test_sdxl_zimage_strength_is_vendor_adaptive_and_leaves_other_profiles_alone():
@@ -791,15 +801,10 @@ def test_sdxl_zimage_strength_is_vendor_adaptive_and_leaves_other_profiles_alone
assert resolve_strength(None, "google", "qwen-zimage", size=(2000, 1850)) == pytest.approx(0.154) assert resolve_strength(None, "google", "qwen-zimage", size=(2000, 1850)) == pytest.approx(0.154)
def test_sdxl_zimage_shares_the_four_step_seed_and_step_contract(): def test_sdxl_zimage_shares_the_fixed_seed_contract():
from remove_ai_watermarks._internal.watermark_profiles import ( from remove_ai_watermarks._internal.watermark_profiles import normalize_profile, resolve_seed
normalize_profile,
resolve_seed,
resolve_steps,
)
assert normalize_profile("sdxl_zimage") == "sdxl-zimage" assert normalize_profile("sdxl_zimage") == "sdxl-zimage"
assert resolve_steps(None) == 4
assert resolve_seed(None) == 0 assert resolve_seed(None) == 0
-70
View File
@@ -15,7 +15,6 @@ from PIL import Image
from remove_ai_watermarks._internal.tiling import ( from remove_ai_watermarks._internal.tiling import (
Tile, Tile,
_axis_positions, _axis_positions,
feather_region_composite,
feather_weights, feather_weights,
plan_tiles, plan_tiles,
run_tiled, run_tiled,
@@ -139,72 +138,3 @@ class TestRunTiled:
image = Image.new("RGB", (1500, 1100), (200, 100, 50)) image = Image.new("RGB", (1500, 1100), (200, 100, 50))
out = run_tiled(generate, image, tile_size=1024, overlap=128) out = run_tiled(generate, image, tile_size=1024, overlap=128)
assert out.size == (1500, 1100) assert out.size == (1500, 1100)
class TestFeatherRegionComposite:
"""Region-targeted compositing for AI-enhanced composites: only the AI box is
regenerated, the real photo outside it stays pixel-exact (roadmap P1#8)."""
@staticmethod
def _frames(h=200, w=300):
base = np.full((h, w, 3), 80, np.uint8)
regenerated = np.full((h, w, 3), 200, np.uint8)
return base, regenerated
def test_outside_box_is_pixel_exact(self):
base, regen = self._frames()
out = feather_region_composite(base, regen, (100, 60, 80, 50), feather=8)
# Far corners are well outside the box -> identical to base.
assert np.array_equal(out[:50, :80], base[:50, :80])
assert np.array_equal(out[150:, 220:], base[150:, 220:])
def test_interior_equals_regenerated(self):
base, regen = self._frames()
out = feather_region_composite(base, regen, (100, 60, 80, 50), feather=8)
# Deep interior of the box (past the feather ramp) is fully regenerated.
assert np.array_equal(out[80:90, 130:150], regen[80:90, 130:150])
def test_hard_paste_when_no_feather(self):
base, regen = self._frames()
out = feather_region_composite(base, regen, (100, 60, 80, 50), feather=0)
assert np.array_equal(out[60:110, 100:180], regen[60:110, 100:180])
assert np.array_equal(out[:60], base[:60])
def test_seam_is_monotonic_ramp(self):
base, regen = self._frames()
out = feather_region_composite(base, regen, (100, 60, 80, 50), feather=10).astype(np.float32)
# Along a horizontal line crossing the left edge, values rise from base(80)
# toward regenerated(200) monotonically through the feather band.
row = out[85, 100:115, 0]
assert row[0] < row[-1]
assert np.all(np.diff(row) >= -1e-3)
def test_dtype_preserved(self):
base, regen = self._frames()
out = feather_region_composite(base, regen, (50, 50, 40, 40), feather=4)
assert out.dtype == base.dtype
def test_grayscale_2d_supported(self):
base = np.full((100, 120), 30, np.uint8)
regen = np.full((100, 120), 220, np.uint8)
out = feather_region_composite(base, regen, (40, 30, 30, 30), feather=4)
assert out.shape == base.shape
assert np.array_equal(out[:30], base[:30])
def test_empty_or_offimage_box_returns_base(self):
base, regen = self._frames()
assert np.array_equal(feather_region_composite(base, regen, (0, 0, 0, 0)), base)
assert np.array_equal(feather_region_composite(base, regen, (500, 500, 40, 40)), base)
def test_box_clamped_to_image_bounds(self):
base, regen = self._frames()
# Box overhangs the bottom-right; only the in-image part is composited.
out = feather_region_composite(base, regen, (280, 180, 60, 60), feather=0)
assert np.array_equal(out[180:, 280:], regen[180:, 280:])
assert out.shape == base.shape
def test_shape_mismatch_raises(self):
base, _ = self._frames(200, 300)
bad = np.full((100, 100, 3), 200, np.uint8)
with pytest.raises(ValueError, match="shape mismatch"):
feather_region_composite(base, bad, (10, 10, 20, 20))
Generated
+1 -1
View File
@@ -3477,7 +3477,7 @@ requires-dist = [
{ name = "remove-ai-watermarks", extras = ["pixels"], marker = "extra == 'diffusion'" }, { name = "remove-ai-watermarks", extras = ["pixels"], marker = "extra == 'diffusion'" },
{ name = "remove-ai-watermarks", extras = ["pixels"], marker = "extra == 'visible'" }, { name = "remove-ai-watermarks", extras = ["pixels"], marker = "extra == 'visible'" },
{ name = "remove-ai-watermarks", extras = ["video"], marker = "extra == 'dev'" }, { name = "remove-ai-watermarks", extras = ["video"], marker = "extra == 'dev'" },
{ name = "remove-ai-watermarks", extras = ["video", "heif", "detect", "trustmark", "diffusion", "qwen-zimage", "lama", "migan"], marker = "extra == 'all'" }, { name = "remove-ai-watermarks", extras = ["video", "heif", "detect", "trustmark", "qwen-zimage", "lama", "migan"], marker = "extra == 'all'" },
{ name = "remove-ai-watermarks", extras = ["visible"], marker = "extra == 'lama'" }, { name = "remove-ai-watermarks", extras = ["visible"], marker = "extra == 'lama'" },
{ name = "remove-ai-watermarks", extras = ["visible"], marker = "extra == 'migan'" }, { name = "remove-ai-watermarks", extras = ["visible"], marker = "extra == 'migan'" },
{ name = "remove-ai-watermarks", extras = ["visible"], marker = "extra == 'video'" }, { name = "remove-ai-watermarks", extras = ["visible"], marker = "extra == 'video'" },