mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-09 23:50:40 +02:00
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:
co-authored by
Claude Opus 5
parent
bf4bfc1ab7
commit
52b2c115e8
@@ -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.
|
||||||
|
|||||||
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
|
||||||
|
|||||||
@@ -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
@@ -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.
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
|||||||
@@ -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
@@ -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"
|
||||||
|
|||||||
@@ -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
@@ -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.
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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
@@ -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
|
||||||
|
|||||||
@@ -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,
|
|
||||||
)
|
|
||||||
|
|||||||
@@ -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
@@ -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:
|
||||||
|
|||||||
@@ -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
@@ -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.
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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))
|
|
||||||
|
|||||||
@@ -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'" },
|
||||||
|
|||||||
Reference in New Issue
Block a user