CalvinBackup/remove-ai-watermarks
1
0
Fork 0
mirror of https://github.com/wiltodelta/remove-ai-watermarks.git synced 2026-08-09 23:50:40 +02:00
Code Issues Packages Projects Releases Wiki Activity

Keep only the two-stage profiles and make CUDA a precondition

Browse Source 
qwen-zimage becomes the default and sdxl-zimage the only alternative. The
controlnet, sdxl, qwen and default profiles are gone, and with them the CPU and
MPS paths for invisible-watermark removal: neither matched the two-stage
recipe's face preservation, so keeping them advertised a quality this library no
longer delivers. Visible-mark removal and every identify command still run
anywhere.

Retired names are rejected rather than remapped. Silently routing --pipeline
sdxl onward would run an old script at a different strength, on a different
model, at a different quality, and report success.

CUDA is now checked when the remover is constructed instead of when the model
loads. Auto-detection cheerfully returned mps on a Mac, so the failure arrived
several layers down, after the dependency check and the pipeline import, in a
message naming whichever internal pipeline happened to raise. _DEVICES collapses
to {"cuda"} and the cpu/mps float32 branch goes with it.

resolve_strength stays total. It briefly returned None for qwen-zimage, meaning
"ask the resolution curve", which pushed a branch onto both callers and left one
of the two strength policies outside the strength module; the CLI copy had
already grown an `or 0.0` guarding a path its own comment called unreachable. It
now takes the image size and answers for both profiles, so the displayed value
cannot drift from the executed one.

Deletion fallout removed with it: img2img_runner and progress.py (the MPS
recovery path and its progress monitor had no callers left), viable_steps, the
fp16 degenerate-output retry, the fp16 VAE fix, and the Qwen img2img call
builders. try_empty_device_cache moved into watermark_remover rather than
leaving a module whose docstring outlived its code. _HAS_DIFFUSERS routes
through optional_deps.module_available, which is what the rest of the library
uses and what correctly rejects a pruned namespace remnant.

--steps, --guidance-scale and --model now have exactly one legal value each and
are still accepted at parse time, then rejected in remove(). Their help text
says so, but validating them beside the option would be better.

Not addressed, and worth its own decision: invisible_engine forces
min_resolution to 0 for both profiles, so the --min-resolution floor, --upscaler,
_esrgan_upscale, upscaler.py and the esrgan extra are all unreachable.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Victor Kuznetsov
2026-08-03 12:16:14 -07:00
co-authored by Claude Opus 5
parent 3d43bac6a5
commit b0ca2054f6
23 changed files with 331 additions and 1458 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.claude/rules/development.md
-1
View File
@@ -23,7 +23,6 @@ Do not classify an entire module as untestable because its main path downloads a
- target-size selection in `test_invisible_engine.py`;
- unsharp and adaptive-polish helpers in `test_humanizer.py`;
- mocked device fallback in `test_img2img_runner.py`;
- tiling geometry and blending in `test_tiling.py`;
- prompt-embedding cache keying, storage round-trip, and the cross-pipeline reuse
that lets a stack load without its text encoder, in `test_qwen_zimage_pipeline.py`;
docs/cli.md
+9 -4
View File
@@ -363,10 +363,15 @@ remove-ai-watermarks invisible image.png -o clean.png --force
| Pipeline | When to use it |
| --- | --- |
| `controlnet` | Default compatibility profile with structural conditioning |
| `sdxl` | Lighter plain SDXL regeneration |
| `qwen` | Large CUDA oriented Qwen Image profile |
| `qwen-zimage` | CUDA only high fidelity profile with a separate face stage |
| `qwen-zimage` | Default. Qwen-Image-2512 global pass plus a SAM-masked Z-Image face stage |
| `sdxl-zimage` | The same recipe and face stage on an SDXL global pass, at a higher denoise |
**Both are CUDA-only.** There is no CPU or MPS profile for invisible-watermark
removal. The former `controlnet`, `sdxl`, `qwen` and `default` profiles were removed
rather than kept as a CPU path: none of them matched this recipe's face preservation,
so offering them implied a quality the library no longer delivers. Passing a retired
name is rejected at parse time rather than remapped. Visible-mark removal and every
identify path still run anywhere.
Example:
docs/python-api.md
+5 -4
View File
@@ -389,7 +389,7 @@ from pathlib import Path
from remove_ai_watermarks.invisible_engine import InvisibleEngine
engine = InvisibleEngine(
pipeline="controlnet",
pipeline="qwen-zimage", # the default; the only other value is "sdxl-zimage"
device=None,
cpu_offload=False,
)
@@ -407,15 +407,16 @@ For limited CUDA memory:
```python
engine = InvisibleEngine(
pipeline="controlnet",
pipeline="qwen-zimage",
cpu_offload=True,
)
```
For the CUDA only high fidelity profile:
Both profiles are CUDA-only, so `device=None` resolving to CPU or MPS cannot run
invisible-watermark removal at all. For the SDXL global stage instead of Qwen:
```python
engine = InvisibleEngine(pipeline="qwen-zimage")
engine = InvisibleEngine(pipeline="sdxl-zimage")
```
The `qwen-zimage` extra must be installed for that profile.
docs/synthid-robust-identity-research.md
+3 -1
View File
@@ -246,7 +246,9 @@ from the test set + this doc).
- If no -> the assumption is wrong; PhotoMaker would re-introduce the
watermark. Stop and reconsider.
2. **PhotoMaker-V1 prototype** in the existing `controlnet` pipeline:
- Mirror the `_load_controlnet_pipeline` path: add a PhotoMaker variant that
- Mirror the `_load_controlnet_pipeline` path (removed in 0.24.0 with the
controlnet profile; the equivalent seam is now `SdxlZImagePipeline._load_sdxl`):
add a PhotoMaker variant that
loads SDXL + canny ControlNet + PhotoMaker adapter on the same engine.
- Extract the OpenCLIP face embedding from the watermarked face crops (use
OpenCV YuNet, already bundled for `auto`, to find the face boxes).
docs/synthid.md
+10 -3
View File
@@ -514,10 +514,17 @@ study (section 2.2) gives empirical floors:
- **Google native 2816**: 0.15 clears (n=2, deployed controlnet worker, 2026-06-14) --
the same rung as capped 1536, so no resolution penalty was observed.
The default is **vendor-adaptive** (`watermark_profiles.resolve_strength` +
`vendor_for_strength`): the tool reads the C2PA issuer on the original input and picks
> **Superseded in 0.24.0.** The `sdxl`, `controlnet`, `qwen` and `default` profiles
> were removed, and `OPENAI_STRENGTH` / `GEMINI_STRENGTH` / `UNKNOWN_STRENGTH` went
> with them. Everything from here to the end of this section is a record of what was
> measured on those profiles, kept because the oracle verdicts are still the evidence
> base. For the strength policy in force now see `module-internals.md`: `qwen-zimage`
> uses `resolution_adaptive_denoise`, `sdxl-zimage` a flat vendor ladder.
The default was **vendor-adaptive** (`watermark_profiles.resolve_strength` +
`vendor_for_strength`): the tool read the C2PA issuer on the original input and picked
`OPENAI_STRENGTH` 0.10 / `GEMINI_STRENGTH` 0.15 / `UNKNOWN_STRENGTH` 0.15 **(LOWERED
2026-06-14 from the 2026-06-04 cert floors 0.20/0.30/0.30)**. **The SAME ladder applies
2026-06-14 from the 2026-06-04 cert floors 0.20/0.30/0.30)**. **The SAME ladder applied
to both pipelines** (`sdxl` and `controlnet`). The 2026-06-14 re-test on the deployed
Modal controlnet worker (v0.10.0) cleared SynthID on the oracle at OpenAI 0.10 (2
photoreal) and Google 0.15 (2 NATIVE 2816x1536, contradicting the "native >= 0.30" guess
packaging/conda/recipe.yaml
+1 -1
View File
@@ -1,7 +1,7 @@
schema_version: 1
context:
version: "0.23.0"
version: "0.24.0"
python_min: "3.10"
package:
pyproject.toml
+1 -1
View File
@@ -1,6 +1,6 @@
[project]
name = "remove-ai-watermarks"
version = "0.23.0"
version = "0.24.0"
description = "AI watermark remover for visible, invisible, and provenance marks in images and video"
readme = "README.md"
requires-python = ">=3.10.1"
src/remove_ai_watermarks/__init__.py
+1 -1
View File
@@ -32,7 +32,7 @@ _os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
_warnings.filterwarnings("ignore", message=r".*ImageProcessorFast.*")
__version__ = "0.23.0"
__version__ = "0.24.0"
__all__ = [
"__version__",
src/remove_ai_watermarks/_internal/img2img_runner.py
-136
View File
@@ -1,136 +0,0 @@
"""Execute Diffusers img2img calls and recover from an MPS runtime failure."""
from __future__ import annotations
import contextlib
import logging
from typing import TYPE_CHECKING, Any
from remove_ai_watermarks._internal.progress import is_mps_error, make_pipeline_progress
if TYPE_CHECKING:
from collections.abc import Callable
from PIL import Image
logger = logging.getLogger(__name__)
def _pipeline_arguments(
image: Image.Image,
strength: float,
num_inference_steps: int,
guidance_scale: float,
generator: Any,
step_callback: Any,
overrides: dict[str, Any] | None,
) -> dict[str, Any]:
arguments: dict[str, Any] = {
"prompt": "",
"image": image,
"strength": strength,
"num_inference_steps": num_inference_steps,
"guidance_scale": guidance_scale,
"generator": generator,
}
arguments.update(overrides or {})
if step_callback is not None:
arguments.update(callback=step_callback, callback_steps=1)
return arguments
def _invoke(pipeline: Any, arguments: dict[str, Any]) -> Image.Image:
response = pipeline(**arguments)
return response.images[0]
def run_img2img(
pipeline: Any,
image: Image.Image,
strength: float,
num_inference_steps: int,
guidance_scale: float,
generator: Any,
device: str,
set_progress: Callable[[str], None],
extra_kwargs: dict[str, Any] | None = None,
) -> Image.Image:
"""Run one img2img request and report denoising progress when supported."""
callback, started, finished, launch_monitor = make_pipeline_progress(
max(1, int(num_inference_steps * strength)), device, set_progress
)
launch_monitor()
arguments = _pipeline_arguments(
image, strength, num_inference_steps, guidance_scale, generator, callback, extra_kwargs
)
try:
try:
return _invoke(pipeline, arguments)
except TypeError as error:
if "callback" not in str(error):
raise
started.set()
arguments.pop("callback", None)
arguments.pop("callback_steps", None)
return _invoke(pipeline, arguments)
finally:
started.set()
finished.set()
def run_img2img_with_mps_fallback(
load_pipeline: Callable[[], Any],
image: Image.Image,
strength: float,
num_inference_steps: int,
guidance_scale: float,
generator: Any,
device: str,
set_progress: Callable[[str], None],
*,
reload_on_cpu: Callable[[], Any],
extra_kwargs: dict[str, Any] | None = None,
) -> tuple[Image.Image, str]:
"""Retry an MPS-specific failure once with a freshly loaded CPU pipeline."""
try:
output = run_img2img(
load_pipeline(),
image,
strength,
num_inference_steps,
guidance_scale,
generator,
device,
set_progress,
extra_kwargs,
)
return output, device
except RuntimeError as error:
if device != "mps" or not is_mps_error(error):
raise
logger.warning("MPS execution failed (%s); retrying on CPU", error)
set_progress("MPS execution failed; retrying on CPU...")
try_empty_device_cache("mps")
output = run_img2img(
reload_on_cpu(),
image,
strength,
num_inference_steps,
guidance_scale,
None,
"cpu",
set_progress,
extra_kwargs,
)
return output, "cpu"
def try_empty_device_cache(device: str) -> None:
"""Ask Torch to release cached accelerator memory when the backend supports it."""
with contextlib.suppress(Exception):
import torch
backend = getattr(torch, device, None)
empty_cache = getattr(backend, "empty_cache", None)
if callable(empty_cache):
empty_cache()
src/remove_ai_watermarks/_internal/progress.py
-212
View File
@@ -1,212 +0,0 @@
"""Progress reporting utilities for long-running optional model operations."""
from __future__ import annotations
import contextlib
import io
import os
import sys
import threading
import time
import warnings
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from collections.abc import Callable
_BAR_WIDTH = 28
_SPINNER = ("|", "/", "-", "\\")
def _truncate(text: str, max_len: int = 72) -> str:
if len(text) <= max_len:
return text
return f"{text[: max(0, max_len - 3)]}..."
def _build_bar(step: int) -> str:
position = step % (2 * _BAR_WIDTH - 2)
if position >= _BAR_WIDTH:
position = 2 * _BAR_WIDTH - 2 - position
cells = ["-"] * _BAR_WIDTH
cells[position] = "="
return "".join(cells)
@dataclass
class _TaskResult:
value: Any = None
error: BaseException | None = None
complete: threading.Event = field(default_factory=threading.Event)
def run_with_progress(task: Callable[[], Any], progress_state: dict[str, str] | None = None) -> Any:
"""Run ``task`` on a worker thread and render a compact terminal heartbeat."""
outcome = _TaskResult()
def invoke() -> None:
try:
outcome.value = task()
except BaseException as error: # re-raised on the caller thread
outcome.error = error
finally:
outcome.complete.set()
worker = threading.Thread(target=invoke, name="raiw-progress-task", daemon=True)
worker.start()
started_at = time.monotonic()
frame = 0
terminal = sys.__stderr__
while not outcome.complete.wait(0.1):
message = _truncate((progress_state or {}).get("message", "Processing..."))
elapsed = int(time.monotonic() - started_at)
if terminal is not None:
terminal.write(
f"\r\033[2K {_SPINNER[frame % len(_SPINNER)]} [{_build_bar(frame)}] {elapsed:>3}s {message}"
)
terminal.flush()
frame += 1
worker.join()
elapsed = int(time.monotonic() - started_at)
message = _truncate((progress_state or {}).get("message", "Processing..."))
if terminal is not None:
terminal.write(f"\r\033[2K Completed in {elapsed}s {message}\n")
terminal.flush()
if outcome.error is not None:
raise outcome.error
return outcome.value
def _silence_diffusers() -> None:
from diffusers.utils import logging as diffusers_logging
diffusers_logging.set_verbosity_error()
disable = getattr(diffusers_logging, "disable_progress_bar", None)
if callable(disable):
disable()
def _configure_quiet_libraries() -> None:
os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
operations = (
lambda: __import__("transformers").logging.set_verbosity_error(),
_silence_diffusers,
lambda: __import__("huggingface_hub").logging.set_verbosity_error(),
)
for operation in operations:
with contextlib.suppress(Exception):
operation()
def silence_library_output(
run_func: Callable[[], Any],
set_progress: Callable[[str], None] | None = None,
) -> Callable[[], Any]:
"""Wrap a model call so third-party progress bars do not corrupt our CLI UI."""
def quiet_call() -> Any:
if set_progress is not None:
set_progress("Preparing model runtime...")
_configure_quiet_libraries()
with (
warnings.catch_warnings(),
contextlib.redirect_stdout(io.StringIO()),
contextlib.redirect_stderr(io.StringIO()),
):
warnings.simplefilter("ignore")
if set_progress is not None:
set_progress("Running watermark regeneration...")
return run_func()
return quiet_call
@dataclass
class _PipelineMonitor:
total_steps: int
device: str
update: Callable[[str], None]
bar_len: int
label: str
pre_phases: list[tuple[int, str]]
post_phases: list[tuple[int, str]]
first_step: threading.Event = field(default_factory=threading.Event)
done: threading.Event = field(default_factory=threading.Event)
started_at: float = field(default_factory=time.monotonic)
last_step_at: float = field(default_factory=time.monotonic)
def callback(self, step: int, _timestep: int, _latents: Any) -> None:
self.first_step.set()
now = time.monotonic()
self.last_step_at = now
current = min(self.total_steps, step + 1)
filled = round(self.bar_len * current / self.total_steps)
elapsed = now - self.started_at
eta = elapsed * max(0, self.total_steps - current) / max(1, current)
bar = "#" * filled + "." * (self.bar_len - filled)
self.update(
f"{self.label} [{bar}] {current}/{self.total_steps}, "
f"{elapsed:.0f}s elapsed, ~{eta:.0f}s left, {self.device}"
)
def _phase_message(self, phases: list[tuple[int, str]], elapsed: float) -> str:
message = phases[0][1]
for threshold, candidate in phases:
if elapsed < threshold:
break
message = candidate
return message
def monitor(self) -> None:
while not self.first_step.wait(0.4):
elapsed = time.monotonic() - self.started_at
self.update(self._phase_message(self.pre_phases, elapsed))
decode_started: float | None = None
while not self.done.wait(0.4):
if time.monotonic() - self.last_step_at < 1.5:
decode_started = None
continue
decode_started = decode_started or time.monotonic()
self.update(self._phase_message(self.post_phases, time.monotonic() - decode_started))
def start(self) -> threading.Thread:
self.started_at = self.last_step_at = time.monotonic()
self.first_step.clear()
self.done.clear()
thread = threading.Thread(target=self.monitor, name="raiw-pipeline-progress", daemon=True)
thread.start()
return thread
def make_pipeline_progress(
effective_steps: int,
device: str,
set_progress: Callable[[str], None],
*,
bar_len: int = 20,
label: str = "Denoising",
pre_phases: list[tuple[int, str]] | None = None,
post_phases: list[tuple[int, str]] | None = None,
) -> tuple[Callable[..., None], threading.Event, threading.Event, Callable[[], threading.Thread]]:
"""Build a callback and monitor for the legacy Diffusers callback interface."""
def qualify(entries: list[tuple[int, str]]) -> list[tuple[int, str]]:
return [(second, f"{text} on {device}") for second, text in entries]
monitor = _PipelineMonitor(
total_steps=max(1, effective_steps),
device=device,
update=set_progress,
bar_len=bar_len,
label=label,
pre_phases=pre_phases or qualify([(0, "Encoding image"), (8, "Preparing denoiser"), (20, "Starting sampler")]),
post_phases=post_phases or qualify([(0, "Decoding image"), (10, "Finalizing pixels"), (45, "Still decoding")]),
)
return monitor.callback, monitor.first_step, monitor.done, monitor.start
def is_mps_error(error: Exception) -> bool:
"""Return whether an error message identifies Apple's MPS backend."""
return "mps" in str(error).casefold()
src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py
+1 -1
View File
@@ -1039,7 +1039,7 @@ class QwenZImagePipeline:
) -> Image.Image:
"""Execute global regeneration and masked face repair."""
self._require_cuda()
seed = resolve_seed(seed, "qwen-zimage")
seed = resolve_seed(seed)
global_strength = (
resolution_adaptive_denoise(image.width, image.height) if strength is None else float(strength)
)
src/remove_ai_watermarks/_internal/watermark_profiles.py
+48 -52
View File
@@ -1,33 +1,39 @@
"""Project-owned configuration for invisible-watermark regeneration profiles."""
"""Project-owned configuration for invisible-watermark regeneration profiles.
Two profiles remain, and both are CUDA-only: ``qwen-zimage`` (the default) and
``sdxl-zimage``. The older ``controlnet``, ``sdxl``, ``qwen`` and ``default`` profiles
were removed rather than kept as a CPU path, because none of them matched the two-stage
recipe's face preservation and keeping them implied a quality this library no longer
offers. Removing invisible watermarks therefore needs a CUDA device; the visible-mark
registry and every identify path still run anywhere.
"""
from __future__ import annotations
import math
from dataclasses import dataclass
from typing import TYPE_CHECKING, Literal
if TYPE_CHECKING:
from pathlib import Path
# 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.
DEFAULT_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
QWEN_MODEL_ID = "Qwen/Qwen-Image"
CONTROLNET_CANNY_MODEL = "xinsir/controlnet-canny-sdxl-1.0"
SDXL_PROFILE = "sdxl"
QWEN_ZIMAGE_PROFILE = "qwen-zimage"
SDXL_ZIMAGE_PROFILE = "sdxl-zimage"
DEFAULT_PROFILE = QWEN_ZIMAGE_PROFILE
PROFILE_CHOICES = (QWEN_ZIMAGE_PROFILE, SDXL_ZIMAGE_PROFILE)
SDXL_LIGHTNING_MODEL_ID = "ByteDance/SDXL-Lightning"
SDXL_LIGHTNING_PATTERN = "sdxl_lightning_4step_lora.safetensors"
OPENAI_STRENGTH = 0.10
GEMINI_STRENGTH = 0.15
UNKNOWN_STRENGTH = GEMINI_STRENGTH
DEFAULT_STRENGTH = UNKNOWN_STRENGTH
QWEN_OPENAI_STRENGTH = 0.10
QWEN_GEMINI_STRENGTH = 0.25
QWEN_UNKNOWN_STRENGTH = QWEN_GEMINI_STRENGTH
# Both profiles run the same distilled four-step schedule, and both are certified at a
# fixed seed because SynthID removal near the strength floor is seed-dependent.
PROFILE_STEPS = 4
PROFILE_SEED = 0
PROFILE_CFG = 1.0
# 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
@@ -55,74 +61,64 @@ class _StrengthPolicy:
return self.by_vendor.get((vendor or "").casefold(), self.unknown)
_STANDARD_POLICY = _StrengthPolicy(
unknown=UNKNOWN_STRENGTH,
by_vendor={"openai": OPENAI_STRENGTH, "google": GEMINI_STRENGTH},
)
_QWEN_POLICY = _StrengthPolicy(
unknown=QWEN_UNKNOWN_STRENGTH,
by_vendor={"openai": QWEN_OPENAI_STRENGTH, "google": QWEN_GEMINI_STRENGTH},
)
_SDXL_ZIMAGE_POLICY = _StrengthPolicy(
unknown=SDXL_ZIMAGE_UNKNOWN_STRENGTH,
by_vendor={"openai": SDXL_ZIMAGE_OPENAI_STRENGTH, "google": SDXL_ZIMAGE_GEMINI_STRENGTH},
)
_ALIASES = {
"default": SDXL_PROFILE,
"qwen_zimage": QWEN_ZIMAGE_PROFILE,
"sdxl_zimage": SDXL_ZIMAGE_PROFILE,
}
_FOUR_STEP_PROFILES = frozenset({QWEN_ZIMAGE_PROFILE, SDXL_ZIMAGE_PROFILE})
def normalize_profile(profile: str) -> str:
"""Normalize spelling and resolve compatibility aliases."""
"""Normalize spelling and resolve the underscore spellings."""
value = profile.strip().casefold()
return _ALIASES.get(value, value)
def resolve_steps(num_inference_steps: int | None, pipeline: str) -> int:
"""Return an explicit step count or the selected profile's default."""
if num_inference_steps is not None:
return num_inference_steps
return 4 if normalize_profile(pipeline) in _FOUR_STEP_PROFILES else 50
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, pipeline: str) -> int | None:
"""Keep the fixed four-step Z-Image profiles reproducible by default."""
if seed is not None:
return seed
return 0 if normalize_profile(pipeline) in _FOUR_STEP_PROFILES else None
def resolve_seed(seed: int | None) -> int:
"""Keep both profiles reproducible by default."""
return PROFILE_SEED if seed is None else seed
def strength_default_help() -> str:
"""Describe the live default policy without duplicating its values."""
return (
f"vendor-adaptive (OpenAI {OPENAI_STRENGTH} / Google {GEMINI_STRENGTH} / "
f"unknown {UNKNOWN_STRENGTH}, from the C2PA issuer; qwen-zimage instead uses "
"resolution-adaptive denoise)"
"profile-adaptive (qwen-zimage uses resolution-adaptive denoise; sdxl-zimage "
f"uses OpenAI {SDXL_ZIMAGE_OPENAI_STRENGTH} / Google {SDXL_ZIMAGE_GEMINI_STRENGTH} / "
f"unknown {SDXL_ZIMAGE_UNKNOWN_STRENGTH}, from the C2PA issuer)"
)
def resolve_strength(strength: float | None, vendor: str | None = None, pipeline: str | None = None) -> float:
"""Resolve a user override or the calibrated policy for a profile and vendor."""
def resolve_strength(
strength: float | None,
vendor: str | None = None,
pipeline: str | None = None,
*,
size: tuple[int, int] | None = None,
) -> float:
"""Resolve a user override or the calibrated policy for a profile and vendor.
Total by design. qwen-zimage picks its strength from image area rather than from
the vendor, so it needs ``size``; returning ``None`` for it instead would push that
branch onto every caller and move one of the two strength policies outside this
module. ``size`` is required for qwen-zimage without an explicit strength.
"""
if strength is not None:
return strength
profile = normalize_profile(pipeline) if pipeline is not None else ""
if profile == "qwen":
policy = _QWEN_POLICY
elif profile == SDXL_ZIMAGE_PROFILE:
policy = _SDXL_ZIMAGE_POLICY
else:
policy = _STANDARD_POLICY
return policy.choose(vendor)
if normalize_profile(pipeline or "") == SDXL_ZIMAGE_PROFILE:
return _SDXL_ZIMAGE_POLICY.choose(vendor)
if size is None:
raise ValueError("qwen-zimage resolves strength from image area, so size is required")
from remove_ai_watermarks._internal.qwen_zimage_pipeline import resolution_adaptive_denoise
def viable_steps(num_inference_steps: int, strength: float) -> int:
"""Ensure Diffusers receives at least one effective img2img denoising step."""
if strength <= 0 or int(num_inference_steps * strength) >= 1:
return num_inference_steps
return math.ceil(1.0 / strength)
return resolution_adaptive_denoise(*size)
def vendor_for_strength(image_path: Path) -> Literal["openai", "google"] | None:
src/remove_ai_watermarks/_internal/watermark_remover.py
+53 -359
View File
@@ -12,18 +12,19 @@ from typing import TYPE_CHECKING, Any
from PIL import Image
from remove_ai_watermarks._internal.watermark_profiles import (
CONTROLNET_CANNY_MODEL,
DEFAULT_MODEL_ID,
DEFAULT_STRENGTH,
QWEN_MODEL_ID,
DEFAULT_PROFILE,
PROFILE_CFG,
PROFILE_CHOICES,
PROFILE_STEPS,
QWEN_ZIMAGE_PROFILE,
SDXL_ZIMAGE_PROFILE,
normalize_profile,
resolve_seed,
resolve_steps,
resolve_strength,
viable_steps,
)
from remove_ai_watermarks.optional_deps import module_available
if TYPE_CHECKING:
from collections.abc import Callable
@@ -33,7 +34,6 @@ 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_PROFILES = frozenset({QWEN_ZIMAGE_PROFILE, SDXL_ZIMAGE_PROFILE})
_ZIMAGE_STACKS = {
QWEN_ZIMAGE_PROFILE: "Qwen-Image-2512 and Z-Image",
SDXL_ZIMAGE_PROFILE: "SDXL and Z-Image",
@@ -47,22 +47,7 @@ except ImportError:
torch = None # type: ignore[assignment]
_HAS_TORCH = False
try:
from diffusers import AutoPipelineForImage2Image as AutoImg2ImgPipeline
_HAS_DIFFUSERS = True
except ImportError:
AutoImg2ImgPipeline = None # type: ignore[assignment,misc]
_HAS_DIFFUSERS = False
_SDXL_FP16_VAE_ID = "madebyollin/sdxl-vae-fp16-fix"
_DEGENERATE_THRESHOLD = 1.0
_CANNY_LOW = 100
_CANNY_HIGH = 200
_CONTROLNET_PROMPT = "best quality, high quality, sharp, detailed, photographic"
_CONTROLNET_NEGATIVE = "blurry, lowres, deformed, distorted text, garbled text, watermark, jpeg artifacts"
_QWEN_PROMPT = "high quality, sharp, detailed, faithful to the original"
_QWEN_NEGATIVE = "blurry, lowres, distorted text, garbled text, artifacts"
_HAS_DIFFUSERS = module_available("diffusers")
def is_watermark_removal_available() -> bool:
@@ -77,19 +62,6 @@ def _ensure_watermark_deps() -> None:
)
def _needs_fp16_vae_fix(model_id: str, default_model_id: str, is_fp16: bool) -> bool:
"""Return whether the default SDXL pipeline needs the overflow-safe VAE."""
return is_fp16 and model_id == default_model_id
def _is_degenerate_image(image: Image.Image) -> bool:
"""Detect the uniform near-black output produced by an fp16 decode collapse."""
import numpy as np
pixels = np.asarray(image.convert("RGB"), dtype=np.float32)
return float(pixels.mean()) < _DEGENERATE_THRESHOLD and float(pixels.std()) < _DEGENERATE_THRESHOLD
def _has_nvidia_gpu() -> bool:
try:
subprocess.run(
@@ -103,23 +75,21 @@ def _has_nvidia_gpu() -> bool:
return True
def _detect_cuda_index_url() -> str:
"""Return a PyTorch wheel index compatible with the reported CUDA runtime."""
try:
report = subprocess.run(
["nvidia-smi"],
check=True,
capture_output=True,
text=True,
).stdout
except (FileNotFoundError, subprocess.CalledProcessError):
return "https://download.pytorch.org/whl/cu121"
import re
def try_empty_device_cache(device: str) -> None:
"""Ask Torch to release cached accelerator memory when the backend supports it.
match = re.search(r"CUDA Version:\s*(\d+)\.(\d+)", report)
if match is None:
return "https://download.pytorch.org/whl/cu121"
return f"https://download.pytorch.org/whl/cu{match.group(1)}{match.group(2)}"
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:
@@ -148,48 +118,11 @@ def get_device() -> str:
return "cpu"
def _make_seed_generator(device: str, seed: int) -> Any:
"""Create a deterministic generator, using CPU when device RNG is unavailable."""
try:
return torch.Generator(device=device).manual_seed(seed) # type: ignore[union-attr]
except (RuntimeError, TypeError):
return torch.Generator().manual_seed(seed) # type: ignore[union-attr]
def _qwen_target_size(width: int, height: int) -> tuple[int, int]:
"""Floor dimensions to Qwen's 16-pixel latent grid."""
return max(16, width - width % 16), max(16, height - height % 16)
def _build_qwen_kwargs(
image: Image.Image,
strength: float,
num_inference_steps: int,
true_cfg_scale: float,
generator: Any,
) -> dict[str, Any]:
"""Build the Qwen img2img call without importing its optional pipeline class."""
width, height = _qwen_target_size(image.width, image.height)
return {
"prompt": _QWEN_PROMPT,
"negative_prompt": _QWEN_NEGATIVE,
"image": image,
"strength": strength,
"num_inference_steps": num_inference_steps,
"true_cfg_scale": true_cfg_scale,
"generator": generator,
"width": width,
"height": height,
}
class WatermarkRemover:
"""Load one regeneration profile and write a metadata-clean raster output."""
DEFAULT_MODEL_ID = DEFAULT_MODEL_ID
DEFAULT_STRENGTH = DEFAULT_STRENGTH
CONTROLNET_CANNY_MODEL = CONTROLNET_CANNY_MODEL
_DEVICES = frozenset({"cpu", "mps", "cuda", "xpu"})
_DEVICES = frozenset({"cuda"})
def __init__(
self,
@@ -198,47 +131,48 @@ class WatermarkRemover:
torch_dtype: Any = None,
progress_callback: Callable[[str], None] | None = None,
hf_token: str | None = None,
pipeline: str = "controlnet",
pipeline: str = DEFAULT_PROFILE,
controlnet_conditioning_scale: float = 1.0,
cpu_offload: bool = False,
) -> None:
requested_model = model_id or self.DEFAULT_MODEL_ID
self.model_profile = normalize_profile(pipeline)
if self.model_profile in _ZIMAGE_PROFILES and model_id not in {None, self.DEFAULT_MODEL_ID}:
if self.model_profile not in PROFILE_CHOICES:
raise ValueError(f"Unsupported pipeline '{pipeline}'. Use one of: {', '.join(PROFILE_CHOICES)}.")
if model_id is not None:
raise ValueError(
f"The {self.model_profile} profile uses a fixed {_ZIMAGE_STACKS[self.model_profile]} model stack."
)
self.model_id = (
"Qwen/Qwen-Image-2512 + Tongyi-MAI/Z-Image-Turbo"
if self.model_profile == QWEN_ZIMAGE_PROFILE
else requested_model
else f"{DEFAULT_MODEL_ID} + Tongyi-MAI/Z-Image-Turbo"
)
_ensure_watermark_deps()
selected_device = (device or get_device()).casefold()
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
# 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.
if self.device not in self._DEVICES:
raise ValueError(f"Unsupported device '{device}'. Use one of: auto, cpu, mps, cuda, xpu.")
raise ValueError(
f"Invisible-watermark removal is CUDA-only, so '{device}' cannot run it. "
"Both remaining profiles need an NVIDIA GPU. Visible-mark removal and "
"every identify command still run on CPU."
)
if torch_dtype is not None:
self.torch_dtype = torch_dtype
elif self.device in {"cpu", "mps"}:
self.torch_dtype = torch.float32 # type: ignore[union-attr]
elif self.model_profile == SDXL_ZIMAGE_PROFILE:
# SDXL ships fp16 weights and an fp16-safe VAE; bf16 would give up the
# variant without buying anything on this architecture.
self.torch_dtype = torch.float16 # type: ignore[union-attr]
elif self.model_profile in {"qwen", QWEN_ZIMAGE_PROFILE}:
self.torch_dtype = torch.bfloat16 # type: ignore[union-attr]
else:
self.torch_dtype = torch.float16 # type: ignore[union-attr]
self.torch_dtype = torch.bfloat16 # type: ignore[union-attr]
self.cpu_offload = cpu_offload
self.controlnet_conditioning_scale = controlnet_conditioning_scale
self.hf_token = hf_token or os.environ.get("HF_TOKEN")
self._progress_callback = progress_callback
self._pipeline: Any = None
self._controlnet_pipeline: Any = None
self._qwen_pipeline: Any = None
self._qwen_zimage_pipeline: Any = None
def _set_progress(self, message: str) -> None:
@@ -248,114 +182,7 @@ class WatermarkRemover:
def preload(self, *, global_only: bool = False) -> None:
"""Materialize the selected model stack before the first request."""
if self.model_profile in _ZIMAGE_PROFILES:
self._load_qwen_zimage_pipeline().preload(global_only=global_only)
elif self.model_profile == "qwen":
self._load_qwen_pipeline()
elif self.model_profile == "controlnet":
self._load_controlnet_pipeline()
else:
self._load_pipeline()
def _base_load_kwargs(self) -> dict[str, Any]:
options: dict[str, Any] = {"torch_dtype": self.torch_dtype}
if self.hf_token:
options["token"] = self.hf_token
return options
def _load_from_pretrained(self, cls: Any, model_id: str, **kwargs: Any) -> Any:
if self.torch_dtype == torch.float16: # type: ignore[union-attr]
try:
return cls.from_pretrained(model_id, variant="fp16", **kwargs)
except Exception as error:
logger.info("Model %s has no usable fp16 variant (%s); using default weights", model_id, error)
return cls.from_pretrained(model_id, **kwargs)
def _maybe_add_fp16_vae(self, options: dict[str, Any]) -> None:
if not _needs_fp16_vae_fix(
self.model_id,
self.DEFAULT_MODEL_ID,
self.torch_dtype == torch.float16, # type: ignore[union-attr]
):
return
from diffusers import AutoencoderKL
options["vae"] = AutoencoderKL.from_pretrained(_SDXL_FP16_VAE_ID, torch_dtype=torch.float16)
@staticmethod
def _disable_sdxl_watermarker(options: dict[str, Any]) -> None:
options["add_watermarker"] = False
def _move_to_device_and_optimize(self, pipeline: Any) -> Any:
if self.cpu_offload and self.device == "cuda":
offload = getattr(pipeline, "enable_model_cpu_offload", None)
if not callable(offload):
raise RuntimeError("CPU offload was requested, but this pipeline does not support it.")
offload(device="cuda")
else:
try:
pipeline = pipeline.to(self.device)
except (RuntimeError, AssertionError) as error:
if self.device == "cuda":
raise RuntimeError(
f"Failed to move model to CUDA ({error}). Install a compatible PyTorch wheel from "
f"{_detect_cuda_index_url()}."
) from error
raise
optimize = getattr(pipeline, "enable_xformers_memory_efficient_attention", None)
if callable(optimize):
with contextlib.suppress(Exception):
optimize()
if self.device == "mps":
slice_attention = getattr(pipeline, "enable_attention_slicing", None)
if callable(slice_attention):
with contextlib.suppress(Exception):
slice_attention("max")
return pipeline
def _sdxl_options(self) -> dict[str, Any]:
options = self._base_load_kwargs()
self._disable_sdxl_watermarker(options)
self._maybe_add_fp16_vae(options)
return options
def _load_pipeline(self) -> Any:
if self._pipeline is None:
options = self._sdxl_options()
options.update(safety_checker=None, requires_safety_checker=False)
loaded = self._load_from_pretrained(AutoImg2ImgPipeline, self.model_id, **options)
self._pipeline = self._move_to_device_and_optimize(loaded)
return self._pipeline
def _load_controlnet_pipeline(self) -> Any:
if self._controlnet_pipeline is None:
from diffusers import ControlNetModel, StableDiffusionXLControlNetImg2ImgPipeline
controlnet = self._load_from_pretrained(
ControlNetModel,
CONTROLNET_CANNY_MODEL,
torch_dtype=self.torch_dtype,
)
options = self._sdxl_options()
options["controlnet"] = controlnet
loaded = self._load_from_pretrained(
StableDiffusionXLControlNetImg2ImgPipeline,
self.model_id,
**options,
)
self._controlnet_pipeline = self._move_to_device_and_optimize(loaded)
return self._controlnet_pipeline
def _load_qwen_pipeline(self) -> Any:
if self._qwen_pipeline is None:
try:
from diffusers import QwenImageImg2ImgPipeline
except ImportError as error:
raise ImportError("The qwen profile requires Diffusers with QwenImageImg2ImgPipeline.") from error
model_id = QWEN_MODEL_ID if self.model_id == self.DEFAULT_MODEL_ID else self.model_id
loaded = QwenImageImg2ImgPipeline.from_pretrained(model_id, **self._base_load_kwargs())
self._qwen_pipeline = self._move_to_device_and_optimize(loaded)
return self._qwen_pipeline
self._load_qwen_zimage_pipeline().preload(global_only=global_only)
def _load_qwen_zimage_pipeline(self) -> Any:
if self._qwen_zimage_pipeline is None:
@@ -379,88 +206,6 @@ class WatermarkRemover:
)
return self._qwen_zimage_pipeline
def _reload_on_cpu(self, cache_name: str, loader: Callable[[], Any]) -> Any:
self.device = "cpu"
self.torch_dtype = torch.float32 # type: ignore[union-attr]
setattr(self, cache_name, None)
return loader()
def _run_img2img(
self,
init_image: Image.Image,
strength: float,
num_inference_steps: int,
guidance_scale: float,
generator: Any,
) -> Image.Image:
from remove_ai_watermarks._internal.img2img_runner import run_img2img_with_mps_fallback
output, device = run_img2img_with_mps_fallback(
self._load_pipeline,
init_image,
strength,
num_inference_steps,
guidance_scale,
generator,
self.device,
self._set_progress,
reload_on_cpu=lambda: self._reload_on_cpu("_pipeline", self._load_pipeline),
)
self.device = device
return output
def _build_canny_control_image(self, init_image: Image.Image) -> Image.Image:
import cv2
import numpy as np
gray = cv2.cvtColor(np.asarray(init_image.convert("RGB")), cv2.COLOR_RGB2GRAY)
edges = cv2.Canny(gray, _CANNY_LOW, _CANNY_HIGH)
return Image.fromarray(np.repeat(edges[:, :, None], 3, axis=2))
def _run_controlnet(
self,
init_image: Image.Image,
strength: float,
num_inference_steps: int,
guidance_scale: float,
generator: Any,
) -> Image.Image:
from remove_ai_watermarks._internal.img2img_runner import run_img2img_with_mps_fallback
extras = {
"prompt": _CONTROLNET_PROMPT,
"negative_prompt": _CONTROLNET_NEGATIVE,
"control_image": self._build_canny_control_image(init_image),
"controlnet_conditioning_scale": float(self.controlnet_conditioning_scale),
}
output, device = run_img2img_with_mps_fallback(
self._load_controlnet_pipeline,
init_image,
strength,
num_inference_steps,
guidance_scale,
generator,
self.device,
self._set_progress,
reload_on_cpu=lambda: self._reload_on_cpu("_controlnet_pipeline", self._load_controlnet_pipeline),
extra_kwargs=extras,
)
self.device = device
return output
def _run_qwen(
self,
init_image: Image.Image,
strength: float,
num_inference_steps: int,
guidance_scale: float,
generator: Any,
) -> Image.Image:
response = self._load_qwen_pipeline()(
**_build_qwen_kwargs(init_image, strength, num_inference_steps, guidance_scale, generator)
)
return response.images[0]
def _run_qwen_zimage(
self,
init_image: Image.Image,
@@ -484,40 +229,20 @@ class WatermarkRemover:
self,
image: Image.Image,
strength: float,
steps: int,
guidance: float,
generator: Any,
seed: int | None,
*,
tile: bool,
tile_size: int,
tile_overlap: int,
) -> Image.Image:
if self.model_profile in _ZIMAGE_PROFILES:
return self._run_qwen_zimage(
image,
strength,
seed,
tile=tile,
tile_size=tile_size,
tile_overlap=tile_overlap,
)
runner = {
"qwen": self._run_qwen,
"controlnet": self._run_controlnet,
}.get(self.model_profile, self._run_img2img)
if tile and max(image.size) > tile_size:
from remove_ai_watermarks._internal.tiling import run_tiled
return run_tiled(
lambda crop: runner(crop, strength, steps, guidance, generator),
image,
tile_size,
tile_overlap,
self._set_progress,
)
return runner(image, strength, steps, guidance, generator)
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:
import numpy as np
@@ -554,56 +279,27 @@ class WatermarkRemover:
with Image.open(image_path) as opened:
source = opened.convert("RGB")
if self.model_profile == QWEN_ZIMAGE_PROFILE:
from remove_ai_watermarks._internal.qwen_zimage_pipeline import resolution_adaptive_denoise
resolved_strength = strength if strength is not None else resolution_adaptive_denoise(*source.size)
else:
resolved_strength = resolve_strength(strength, vendor, self.model_profile)
resolved_strength = resolve_strength(strength, vendor, self.model_profile, size=source.size)
if not 0.0 <= resolved_strength <= 1.0:
raise ValueError(f"Strength must be between 0.0 and 1.0, got {resolved_strength}")
resolved_seed = resolve_seed(seed, self.model_profile)
steps = resolve_steps(num_inference_steps, self.model_profile)
guidance = 1.0 if guidance_scale is None and self.model_profile in _ZIMAGE_PROFILES else guidance_scale or 7.5
if self.model_profile in _ZIMAGE_PROFILES:
if steps != 4:
raise ValueError(f"The {self.model_profile} profile requires 4 steps.")
if guidance != 1.0:
raise ValueError(f"The {self.model_profile} profile requires CFG 1.0.")
else:
steps = viable_steps(steps, resolved_strength)
# Both profiles are distilled four-step schedules at CFG 1.0. Anything else is
# 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}.")
generator = None
if resolved_seed is not None and _HAS_TORCH:
generator = _make_seed_generator(self.device, resolved_seed)
result = self._generate(
source,
resolved_strength,
steps,
guidance,
generator,
resolved_seed,
resolve_seed(seed),
tile=tile,
tile_size=tile_size,
tile_overlap=tile_overlap,
)
if self.torch_dtype == torch.float16 and _is_degenerate_image(result): # type: ignore[union-attr]
self.torch_dtype = torch.float32 # type: ignore[union-attr]
self._pipeline = self._controlnet_pipeline = self._qwen_pipeline = self._qwen_zimage_pipeline = None
result = self._generate(
source,
resolved_strength,
steps,
guidance,
generator,
resolved_seed,
tile=tile,
tile_size=tile_size,
tile_overlap=tile_overlap,
)
if region is not None:
import numpy as np
@@ -634,8 +330,6 @@ class WatermarkRemover:
if not input_dir.exists():
raise FileNotFoundError(f"Input directory not found: {input_dir}")
output_dir.mkdir(parents=True, exist_ok=True)
from remove_ai_watermarks._internal.img2img_runner import try_empty_device_cache
outputs: list[Path] = []
candidates = sorted(path for path in input_dir.iterdir() if path.suffix.casefold() in extensions)
for source in candidates:
src/remove_ai_watermarks/cli.py
+50 -65
View File
@@ -24,6 +24,9 @@ from remove_ai_watermarks import __version__, image_io, watermark_registry
from remove_ai_watermarks._internal.constants import SUPPORTED_FORMATS
from remove_ai_watermarks._internal.utils import is_supported_format
from remove_ai_watermarks._internal.watermark_profiles import (
DEFAULT_PROFILE,
PROFILE_CHOICES,
QWEN_ZIMAGE_PROFILE,
resolve_seed,
resolve_steps,
resolve_strength,
@@ -155,15 +158,15 @@ def _resolved_strength_for_display(
vendor: str | None,
pipeline: str,
) -> float:
"""Resolve the same profile-specific strength the engine will execute."""
if pipeline == "qwen-zimage" and strength is None:
from PIL import Image
"""Resolve the same profile-specific strength the engine will execute.
from remove_ai_watermarks._internal.qwen_zimage_pipeline import resolution_adaptive_denoise
One call for both profiles, so the printed value cannot drift from the executed
one; the size is what qwen-zimage derives its strength from.
"""
from PIL import Image
with Image.open(source) as image:
return resolution_adaptive_denoise(image.width, image.height)
return resolve_strength(strength, vendor, pipeline)
with Image.open(source) as image:
return resolve_strength(strength, vendor, pipeline, size=image.size)
# Shared option decorator for commands that run the invisible-watermark pipeline.
@@ -173,8 +176,8 @@ _controlnet_scale_option = click.option(
"--controlnet-scale",
type=float,
default=1.0,
help="ControlNet conditioning scale (structure/text preservation strength); "
"applies to the controlnet pipeline (the default). Higher = closer to original structure.",
help="Canny ControlNet conditioning scale on the global stage "
"(structure/text preservation strength). Higher = closer to original structure.",
)
_min_resolution_option = click.option(
@@ -203,9 +206,9 @@ _auto_option = click.option(
"--auto",
is_flag=True,
default=False,
help="DEPRECATED: controlnet and adaptive polish are already the defaults, so "
"--auto only emits a warning and changes nothing. Use --no-adaptive-polish "
"to disable polishing.",
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(
@@ -251,49 +254,29 @@ _model_option = click.option(
"--model",
type=str,
default=None,
help="HuggingFace model ID for the diffusion pipeline. Default: the SDXL base checkpoint.",
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). Default: 7.5, except qwen-zimage "
"fixes CFG at 1.0. Lower = follow the prompt less / stay closer to the input.",
help="Classifier-free guidance scale (CFG). Both profiles are distilled and fix "
"CFG at 1.0, so any other value is rejected.",
)
def _normalize_pipeline(ctx: click.Context, param: click.Parameter, value: str | None) -> str | None:
"""Resolve the legacy ``default`` profile name to ``sdxl`` (click option callback).
Emits a one-line deprecation notice when the user explicitly passes the outdated
``default`` value, pointing at the two current choices (``sdxl`` / ``controlnet``).
"""
if value is None:
return None
from remove_ai_watermarks._internal.watermark_profiles import normalize_profile
normalized = normalize_profile(value)
if value.strip().lower() == "default":
click.echo(
"Warning: --pipeline default is deprecated and maps to 'sdxl'. "
"Use --pipeline sdxl (plain SDXL) or --pipeline controlnet (the default).",
err=True,
)
return normalized
# ``controlnet`` (the default-SELECTED value), ``sdxl`` (plain SDXL img2img) and
# ``qwen`` (Qwen-Image, CUDA/cloud-class) are the current profiles; ``default`` is an
# OUTDATED back-compat alias for ``sdxl`` (warned + normalized away by _normalize_pipeline).
_PIPELINE_CHOICES = ["sdxl", "controlnet", "qwen", "qwen-zimage", "default"]
# 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
# recipe's face preservation, so offering them implied a quality the library no longer
# delivers. BOTH remaining profiles are CUDA-only.
_PIPELINE_CHOICES = list(PROFILE_CHOICES)
_PIPELINE_HELP = (
"Pipeline profile. controlnet (DEFAULT) = SDXL + canny ControlNet that preserves "
"text/faces via edge conditioning while removing SynthID; sdxl = plain SDXL img2img "
"(lighter, no extra model download, but leaves SynthID on flat-graphic content); "
"qwen = Qwen-Image (20B, Apache-2.0) img2img, best text/structure preservation but "
"CUDA/cloud-class; qwen-zimage = Qwen-Image-2512 + Lightning + Canny, followed by "
"SAM-masked Z-Image face repair (CUDA-only; install the qwen-zimage extra). "
"('default' is an OUTDATED alias for 'sdxl'.)"
"Pipeline profile. qwen-zimage (DEFAULT) = Qwen-Image-2512 + Lightning + Canny, "
"followed by SAM-masked Z-Image face repair; sdxl-zimage = the same recipe and the "
"same face stage on an SDXL global pass, which needs more denoise. Both are "
"CUDA-ONLY -- install the qwen-zimage extra. There is no CPU or MPS profile for "
"invisible-watermark removal."
)
# Shared --pipeline / --strength decorators so the three diffusion commands
@@ -302,8 +285,7 @@ _PIPELINE_HELP = (
_pipeline_option = click.option(
"--pipeline",
type=click.Choice(_PIPELINE_CHOICES),
default="controlnet",
callback=_normalize_pipeline,
default=DEFAULT_PROFILE,
help=_PIPELINE_HELP,
)
_strength_option = click.option(
@@ -362,12 +344,10 @@ _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, but the
pipeline is now always controlnet (the default) and the adaptive polish is ON by
default (it self-gates by detail level), so the content detectors were removed and
``--auto`` is now a no-op alias: the polish it used to enable is already the default,
and an explicit ``--no-adaptive-polish`` still wins. So it only emits a deprecation
warning and passes ``adaptive_polish`` through.
``--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(
@@ -379,9 +359,14 @@ def _resolve_auto_polish(auto: bool, adaptive_polish: bool) -> bool:
def _resolve_profile_polish(auto: bool, adaptive_polish: bool, pipeline: str) -> bool:
"""Keep the upstream qwen-zimage output unchanged unless polish was explicit."""
"""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":
if pipeline != QWEN_ZIMAGE_PROFILE or auto:
return adaptive_polish
ctx = click.get_current_context(silent=True)
if ctx is None:
@@ -885,7 +870,7 @@ def cmd_erase(
"--steps",
type=int,
default=None,
help="Number of denoising steps. Default: 4 for qwen-zimage, 50 otherwise.",
help="Number of denoising steps. Both profiles are distilled four-step schedules, so 4 is the only accepted value.",
)
@_pipeline_option
@click.option(
@@ -965,8 +950,8 @@ def cmd_invisible(
from remove_ai_watermarks.invisible_engine import InvisibleEngine
source = _validate_image(source)
steps = resolve_steps(steps, pipeline)
seed = resolve_seed(seed, pipeline)
steps = resolve_steps(steps)
seed = resolve_seed(seed)
_warn_if_esrgan_unavailable(upscaler)
adaptive_polish = _resolve_profile_polish(auto, adaptive_polish, pipeline)
if output is None:
@@ -1575,7 +1560,7 @@ def cmd_identify(ctx: click.Context, source: Path, no_visible: bool, as_json: bo
"--steps",
type=int,
default=None,
help="Number of denoising steps. Default: 4 for qwen-zimage, 50 otherwise.",
help="Number of denoising steps. Both profiles are distilled four-step schedules, so 4 is the only accepted value.",
)
@_pipeline_option
@_model_option
@@ -1651,8 +1636,8 @@ def cmd_all(
"""
_banner()
source = _validate_image(source)
steps = resolve_steps(steps, pipeline)
seed = resolve_seed(seed, pipeline)
steps = resolve_steps(steps)
seed = resolve_seed(seed)
_warn_if_esrgan_unavailable(upscaler)
adaptive_polish = _resolve_profile_polish(auto, adaptive_polish, pipeline)
@@ -2018,7 +2003,7 @@ def _process_batch_image(
"--steps",
type=int,
default=None,
help="Number of denoising steps. Default: 4 for qwen-zimage, 50 otherwise.",
help="Number of denoising steps. Both profiles are distilled four-step schedules, so 4 is the only accepted value.",
)
@_visible_backend_option
@_visible_sensitivity_option
@@ -2105,8 +2090,8 @@ def cmd_batch(
if mode in ("invisible", "all"):
_warn_if_esrgan_unavailable(upscaler)
adaptive_polish = _resolve_profile_polish(auto, adaptive_polish, pipeline)
steps = resolve_steps(steps, pipeline)
seed = resolve_seed(seed, pipeline)
steps = resolve_steps(steps)
seed = resolve_seed(seed)
options = _BatchOptions(
strength=strength,
steps=steps,
src/remove_ai_watermarks/invisible_engine.py
+13 -17
View File
@@ -20,7 +20,9 @@ from ._internal.watermark_profiles import (
DEFAULT_MODEL_ID as DEFAULT_SDXL_MODEL_ID,
)
from ._internal.watermark_profiles import (
DEFAULT_PROFILE,
resolve_seed,
resolve_steps,
)
if TYPE_CHECKING:
@@ -90,7 +92,7 @@ class InvisibleEngine:
self,
model_id: str | None = None,
device: str | None = None,
pipeline: str = "controlnet",
pipeline: str = DEFAULT_PROFILE,
hf_token: str | None = None,
progress_callback: Callable[[str], None] | None = None,
controlnet_conditioning_scale: float = 1.0,
@@ -101,19 +103,16 @@ class InvisibleEngine:
Args:
model_id: HuggingFace model ID. None = use the SDXL base default.
device: Device for inference (auto/cpu/mps/cuda/xpu). None = auto.
pipeline: Pipeline profile. "controlnet" (DEFAULT; SDXL + canny ControlNet
that preserves text/face structure via edge conditioning while removing
SynthID), "sdxl" (plain SDXL img2img, lighter but leaves SynthID on
flat-graphic content), or "qwen" (Qwen-Image 20B img2img, best text/
structure preservation but CUDA/cloud-class), or "qwen-zimage"
(Qwen-Image-2512 Lightning + Canny, then SAM-masked Z-Image face
repair; CUDA-only), or "sdxl-zimage" (the same recipe and the same face
stage on an SDXL global pass, vendor-adaptive strength because an SDXL
global stage needs more of it; CUDA-only). "default" aliases "sdxl".
pipeline: Pipeline profile, one of "qwen-zimage" (DEFAULT;
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
global pass, vendor-adaptive strength because an SDXL global stage
needs more of it). BOTH ARE CUDA-ONLY -- there is no CPU or MPS path
for invisible-watermark removal.
hf_token: HuggingFace API token.
progress_callback: Optional callback for progress messages.
controlnet_conditioning_scale: ControlNet structure-preservation
strength (controlnet pipeline only).
controlnet_conditioning_scale: Canny ControlNet structure-preservation
strength on the global stage of both profiles.
cpu_offload: Offload model components to CPU between CUDA calls instead
of keeping the whole pipeline in VRAM, at the cost of speed. For
qwen-zimage, force the face stack to offload instead of using automatic
@@ -238,11 +237,8 @@ class InvisibleEngine:
"""
import tempfile
if num_inference_steps is None:
profile = getattr(self._remover, "model_profile", None)
num_inference_steps = 4 if profile in {"qwen-zimage", "sdxl-zimage"} else 100
profile = getattr(self._remover, "model_profile", "controlnet")
seed = resolve_seed(seed, profile)
num_inference_steps = resolve_steps(num_inference_steps)
seed = resolve_seed(seed)
from PIL import Image, ImageOps
tests/test_cli.py
+19 -29
View File
@@ -327,7 +327,7 @@ class TestInvisibleCommand:
expected = sample_png.with_stem(sample_png.stem + "_clean")
assert expected.exists()
def test_invisible_adaptive_polish_on_by_default(self, runner, sample_png):
def test_invisible_adaptive_polish_off_by_default_under_qwen_zimage(self, runner, sample_png):
mock_cls, mock_engine = _mock_invisible_engine()
with (
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
@@ -336,8 +336,10 @@ class TestInvisibleCommand:
):
result = runner.invoke(main, ["invisible", str(sample_png), "--force"])
assert result.exit_code == 0, result.output
# adaptive_polish is ON by default (self-gating, so a no-op where not needed).
assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is True
# The default profile is qwen-zimage, and _resolve_profile_polish keeps its
# 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
@@ -368,30 +370,16 @@ class TestInvisibleCommand:
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_pipeline_default_alias_warns_and_maps_to_sdxl(self, runner, sample_png):
mock_cls, _mock_engine = _mock_invisible_engine()
with (
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
):
result = runner.invoke(main, ["invisible", str(sample_png), "--pipeline", "default", "--force"])
assert result.exit_code == 0, result.output
# The legacy value warns and is normalized to "sdxl" before the engine is built.
assert "deprecated" in result.output.lower()
assert mock_cls.call_args.kwargs["pipeline"] == "sdxl"
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.
def test_pipeline_sdxl_does_not_warn(self, runner, sample_png):
mock_cls, _mock_engine = _mock_invisible_engine()
with (
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
):
result = runner.invoke(main, ["invisible", str(sample_png), "--pipeline", "sdxl", "--force"])
assert result.exit_code == 0, result.output
assert "deprecated" not in result.output.lower()
assert mock_cls.call_args.kwargs["pipeline"] == "sdxl"
Click rejects them at parse time. Mapping them onward would run a profile the
caller never chose, at a different strength and a different quality.
"""
for retired in ("default", "sdxl", "controlnet", "qwen"):
result = runner.invoke(main, ["invisible", str(sample_png), "--pipeline", retired, "--force"])
assert result.exit_code == 2, result.output
assert "is not one of" in result.output
def test_invisible_nonexistent_file(self, runner):
result = runner.invoke(main, ["invisible", "/nonexistent/file.png"])
@@ -871,8 +859,10 @@ class TestBatchCommand:
assert out[100, 100, 3] == 255
def test_batch_auto_is_deprecated_and_enables_polish(self, runner, tmp_path):
"""--auto is retired: it warns and just enables the adaptive polish (the
pipeline is always the default controlnet now)."""
"""--auto is retired: it warns and just enables the adaptive polish.
It no longer selects a pipeline: qwen-zimage is the only default there is.
"""
input_dir = _make_batch_dir(tmp_path, count=2)
output_dir = tmp_path / "output"
mock_cls, mock_engine = _mock_invisible_engine()
@@ -890,7 +880,7 @@ class TestBatchCommand:
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"] == "controlnet"
assert mock_cls.call_args.kwargs["pipeline"] == "qwen-zimage"
assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is True
def test_batch_default_output_dir(self, runner, tmp_path):
tests/test_cpu_offload.py
-56
View File
@@ -8,10 +8,6 @@ core CI matrix needs no diffusion dependency, model download, or GPU.
from __future__ import annotations
from unittest.mock import Mock
import pytest
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
@@ -21,55 +17,3 @@ def _remover(device: str, cpu_offload: bool) -> WatermarkRemover:
remover.cpu_offload = cpu_offload
remover._progress_callback = None
return remover
class TestCpuOffloadPlacement:
def test_offload_enabled_on_cuda_streams_instead_of_moving(self):
remover = _remover("cuda", cpu_offload=True)
pipeline = Mock()
returned = remover._move_to_device_and_optimize(pipeline)
pipeline.enable_model_cpu_offload.assert_called_once_with(device="cuda")
pipeline.to.assert_not_called()
# Offload leaves the pipeline object in place (accelerate hooks handle it).
assert returned is pipeline
def test_no_offload_moves_whole_pipeline_to_cuda(self):
remover = _remover("cuda", cpu_offload=False)
pipeline = Mock()
remover._move_to_device_and_optimize(pipeline)
pipeline.to.assert_called_once_with("cuda")
pipeline.enable_model_cpu_offload.assert_not_called()
def test_offload_flag_ignored_off_cuda(self):
# The flag is CUDA-only: on cpu it must still be a plain .to("cpu").
remover = _remover("cpu", cpu_offload=True)
pipeline = Mock()
remover._move_to_device_and_optimize(pipeline)
pipeline.to.assert_called_once_with("cpu")
pipeline.enable_model_cpu_offload.assert_not_called()
def test_offload_fails_loudly_when_pipeline_lacks_support(self):
remover = _remover("cuda", cpu_offload=True)
pipeline = Mock(spec=["to"])
with pytest.raises(RuntimeError, match="does not support"):
remover._move_to_device_and_optimize(pipeline)
pipeline.to.assert_not_called()
def test_qwen_zimage_forces_face_stack_offload(self):
remover = _remover("cuda", cpu_offload=True)
remover.torch_dtype = object()
remover.hf_token = None
remover.controlnet_conditioning_scale = 1.0
remover._qwen_zimage_pipeline = None
runtime = remover._load_qwen_zimage_pipeline()
assert runtime.keep_face_models_on_device is False
tests/test_img2img_runner.py
-133
View File
@@ -1,133 +0,0 @@
"""Unit tests for the MPS->CPU fallback orchestration (no GPU/model required).
``img2img_runner`` has no torch import at module top -- the pipeline is
injected as a plain callable -- so the fallback control flow is fully
mockable. This guards the exact behavior hit in production on Apple Silicon:
a native-resolution SDXL run that OOMs on MPS must transparently retry on CPU,
while any non-MPS error must propagate unchanged.
"""
from __future__ import annotations
from unittest.mock import Mock
import pytest
from remove_ai_watermarks._internal import img2img_runner
from remove_ai_watermarks._internal.img2img_runner import (
run_img2img,
run_img2img_with_mps_fallback,
)
_MPS_OOM = "MPS backend out of memory (MPS allocated: 17.21 GiB, max allowed: 20.13 GiB)"
def _result(image: object) -> Mock:
"""A stand-in for a diffusers pipeline output object (has .images)."""
out = Mock()
out.images = [image]
return out
class TestMpsFallback:
def test_mps_error_reloads_on_cpu_and_retries(self, monkeypatch: pytest.MonkeyPatch):
sentinel = object()
inner = Mock(side_effect=[RuntimeError(_MPS_OOM), sentinel])
monkeypatch.setattr(img2img_runner, "run_img2img", inner)
load_pipeline = Mock(return_value="gpu_pipe")
reload_on_cpu = Mock(return_value="cpu_pipe")
img, device = run_img2img_with_mps_fallback(
load_pipeline, object(), 0.05, 50, 7.5, "gen", "mps", lambda _m: None, reload_on_cpu=reload_on_cpu
)
assert (img, device) == (sentinel, "cpu")
reload_on_cpu.assert_called_once()
assert inner.call_count == 2
# Retry must use the reloaded CPU pipeline, device "cpu", and drop the
# MPS generator (generator=None) so CPU runs deterministically.
retry_args = inner.call_args_list[1].args
assert retry_args[0] == "cpu_pipe"
assert retry_args[5] is None # generator
assert retry_args[6] == "cpu" # device
def test_happy_path_returns_original_device_without_reload(self, monkeypatch: pytest.MonkeyPatch):
sentinel = object()
monkeypatch.setattr(img2img_runner, "run_img2img", Mock(return_value=sentinel))
reload_on_cpu = Mock()
img, device = run_img2img_with_mps_fallback(
Mock(return_value="gpu_pipe"),
object(),
0.05,
50,
7.5,
"gen",
"mps",
lambda _m: None,
reload_on_cpu=reload_on_cpu,
)
assert (img, device) == (sentinel, "mps")
reload_on_cpu.assert_not_called()
def test_non_mps_runtime_error_propagates(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(img2img_runner, "run_img2img", Mock(side_effect=RuntimeError("CUDA out of memory")))
reload_on_cpu = Mock()
with pytest.raises(RuntimeError, match="CUDA"):
run_img2img_with_mps_fallback(
Mock(return_value="gpu_pipe"),
object(),
0.05,
50,
7.5,
"gen",
"mps",
lambda _m: None,
reload_on_cpu=reload_on_cpu,
)
reload_on_cpu.assert_not_called()
def test_mps_error_on_non_mps_device_propagates(self, monkeypatch: pytest.MonkeyPatch):
# An "mps"-worded error while running on cpu must NOT trigger the reload.
monkeypatch.setattr(img2img_runner, "run_img2img", Mock(side_effect=RuntimeError(_MPS_OOM)))
reload_on_cpu = Mock()
with pytest.raises(RuntimeError, match="MPS backend"):
run_img2img_with_mps_fallback(
Mock(return_value="cpu_pipe"),
object(),
0.05,
50,
7.5,
None,
"cpu",
lambda _m: None,
reload_on_cpu=reload_on_cpu,
)
reload_on_cpu.assert_not_called()
class TestRunImg2Img:
def test_returns_first_image_from_pipeline_result(self):
sentinel = object()
pipeline = Mock(return_value=_result(sentinel))
out = run_img2img(pipeline, object(), 0.05, 50, 7.5, None, "cpu", lambda _m: None)
assert out is sentinel
def test_typeerror_on_callback_retries_without_callback(self):
# Older diffusers reject the progress callback kwarg with TypeError;
# run_img2img must retry once without it rather than fail.
sentinel = object()
pipeline = Mock(side_effect=[TypeError("unexpected keyword 'callback'"), _result(sentinel)])
out = run_img2img(pipeline, object(), 0.05, 50, 7.5, None, "cpu", lambda _m: None)
assert out is sentinel
assert pipeline.call_count == 2
# First attempt passes the progress callback; the retry omits it.
assert "callback" in pipeline.call_args_list[0].kwargs
assert "callback" not in pipeline.call_args_list[1].kwargs
tests/test_invisible_engine.py
-28
View File
@@ -4,7 +4,6 @@ from __future__ import annotations
from types import SimpleNamespace
import pytest
from PIL import Image
from remove_ai_watermarks.invisible_engine import InvisibleEngine, _target_size, is_available
@@ -204,30 +203,3 @@ class TestEsrganUpscale:
out = InvisibleEngine._esrgan_upscale(self._fake_engine(), img, (512, 341))
assert out.size == (512, 341)
assert np.array_equal(np.asarray(out), np.asarray(img.resize((512, 341), Image.Resampling.LANCZOS)))
class TestCannyControlImage:
"""The ControlNet canny conditioning image builder (pure cv2/numpy; behind the gpu
extra since it lives on WatermarkRemover). Skips when torch/diffusers are absent."""
def test_edge_map_is_3channel_rgb(self):
if not is_available():
pytest.skip("diffusion extra (torch/diffusers) not installed")
import numpy as np
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
rng = np.random.default_rng(0)
img = Image.fromarray(rng.integers(0, 256, (64, 80, 3), dtype=np.uint8))
# The method uses no instance state, so call it unbound with a dummy self.
out = WatermarkRemover._build_canny_control_image(None, img) # type: ignore[arg-type]
arr = np.array(out)
assert out.mode == "RGB"
assert arr.shape == (64, 80, 3)
import cv2
gray = cv2.cvtColor(np.asarray(img.convert("RGB")), cv2.COLOR_RGB2GRAY)
expected = cv2.Canny(gray, 100, 200)
assert np.array_equal(arr[:, :, 0], expected)
assert np.array_equal(arr[:, :, 1], expected)
assert np.array_equal(arr[:, :, 2], expected)
tests/test_platform.py
+102 -288
View File
@@ -9,17 +9,14 @@ from __future__ import annotations
from pathlib import Path
from unittest.mock import MagicMock, patch
import numpy as np
import pytest
from PIL import Image
from remove_ai_watermarks._internal.progress import is_mps_error
from remove_ai_watermarks._internal.utils import get_image_format, is_supported_format
from remove_ai_watermarks._internal.watermark_profiles import (
DEFAULT_STRENGTH,
GEMINI_STRENGTH,
OPENAI_STRENGTH,
UNKNOWN_STRENGTH,
PROFILE_CHOICES,
SDXL_ZIMAGE_GEMINI_STRENGTH,
SDXL_ZIMAGE_OPENAI_STRENGTH,
SDXL_ZIMAGE_UNKNOWN_STRENGTH,
normalize_profile,
resolve_strength,
strength_default_help,
@@ -59,290 +56,158 @@ class TestDeviceDetection:
assert get_device() == "xpu"
fake_torch.tensor.assert_called_with([1.0], device="xpu")
def test_init_accepts_xpu_and_selects_fp16(self):
"""WatermarkRemover accepts device='xpu' and picks fp16 (not fp32)."""
def test_non_cuda_devices_are_refused_at_construction(self):
"""CUDA is a precondition of the object, not of the run.
Both remaining profiles raise on any other device, so accepting cpu/mps/xpu
here only defers a guaranteed failure to model-load time - several layers down,
after the dependency check and the pipeline import, under a message naming
whichever profile the internal pipeline happens to be.
"""
if not is_watermark_removal_available():
pytest.skip("torch/diffusers not installed")
import torch
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
remover = WatermarkRemover(device="xpu")
assert remover.device == "xpu"
assert remover.torch_dtype == torch.float16
for device in ("cpu", "mps", "xpu"):
with pytest.raises(ValueError, match="CUDA-only"):
WatermarkRemover(device=device)
def test_seed_generator_falls_back_to_cpu_when_device_rng_unsupported(self):
"""A device with no RNG backend (e.g. some torch-xpu builds) falls back
to a CPU generator instead of raising when --seed is used."""
from remove_ai_watermarks._internal import watermark_remover as wr
def fake_generator(device="cpu"):
if device == "xpu":
raise RuntimeError("Device type xpu is not supported for torch.Generator()")
gen = MagicMock()
gen.manual_seed.return_value = f"gen:{device}"
return gen
fake_torch = MagicMock()
fake_torch.Generator.side_effect = fake_generator
with patch.object(wr, "torch", fake_torch):
assert wr._make_seed_generator("xpu", 123) == "gen:cpu"
assert wr._make_seed_generator("cuda", 123) == "gen:cuda"
remover = WatermarkRemover(device="cuda")
assert remover.device == "cuda"
assert remover.torch_dtype == torch.bfloat16
class TestMpsErrorDetection:
"""Tests for MPS error detection helper."""
class TestEmptyDeviceCache:
"""try_empty_device_cache is all that remains of the img2img runner.
def test_detects_mps_error(self):
err = RuntimeError("MPS backend out of memory")
assert is_mps_error(err) is True
Its module lost run_img2img and the MPS fallback along with the CPU/MPS profiles;
both surviving profiles are CUDA-only, so there is no MPS failure left to recover
from. The helper must stay silent on a backend that cannot empty a cache, because
it runs in cleanup paths where a raise would replace the real error.
"""
def test_non_mps_error(self):
err = RuntimeError("CUDA out of memory")
assert is_mps_error(err) is False
def test_unknown_backend_is_a_silent_no_op(self):
from remove_ai_watermarks._internal.watermark_remover import try_empty_device_cache
def test_generic_error(self):
err = RuntimeError("something went wrong")
assert is_mps_error(err) is False
# ── Model profiles ──────────────────────────────────────────────────
try_empty_device_cache("cpu")
try_empty_device_cache("definitely-not-a-backend")
class TestModelProfiles:
"""Tests for watermark_profiles.py profile-name normalization."""
"""Only the two CUDA-only two-stage profiles remain."""
def test_canonical_profiles_unchanged(self):
assert normalize_profile("sdxl") == "sdxl"
assert normalize_profile("controlnet") == "controlnet"
assert normalize_profile("qwen") == "qwen"
assert normalize_profile("qwen-zimage") == "qwen-zimage"
assert normalize_profile("sdxl-zimage") == "sdxl-zimage"
def test_default_alias_resolves_to_sdxl(self):
# "default" is the legacy alias for "sdxl" (back-compat for existing scripts).
assert normalize_profile("default") == "sdxl"
def test_underscore_spellings_resolve(self):
assert normalize_profile("qwen_zimage") == "qwen-zimage"
assert normalize_profile(" SDXL_ZImage ") == "sdxl-zimage"
def test_normalize_is_case_and_whitespace_insensitive(self):
assert normalize_profile(" Default ") == "sdxl"
assert normalize_profile("CONTROLNET") == "controlnet"
def test_retired_names_no_longer_resolve_to_a_profile(self):
"""default/sdxl/controlnet/qwen were removed, not aliased onward.
class TestFp16WeightVariant:
"""_load_from_pretrained reads the fp16 weight variant on fp16, with a fallback.
Loading the fp16 ``variant`` reads the half-precision weight files (~half the bytes)
instead of the fp32 defaults + a downcast, which roughly halves the cold-start weight
read. fp32 (cpu/mps) and bf16 (qwen) must never request the variant; a checkpoint
without fp16 files must fall back to the default weights (prior behavior).
"""
def _remover(self, dtype: object):
if not is_watermark_removal_available():
pytest.skip("torch/diffusers not installed")
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
# device="cpu" alone would force fp32; the explicit torch_dtype override lets us
# exercise the fp16 path with no GPU (construction loads no weights).
return WatermarkRemover(device="cpu", torch_dtype=dtype)
def test_fp16_requests_variant(self):
import torch
remover = self._remover(torch.float16)
cls = MagicMock()
cls.from_pretrained.return_value = "PIPE"
out = remover._load_from_pretrained(cls, "some/model", token="t")
assert out == "PIPE"
cls.from_pretrained.assert_called_once_with("some/model", variant="fp16", token="t")
def test_fp16_falls_back_when_variant_missing(self):
import torch
remover = self._remover(torch.float16)
cls = MagicMock()
cls.from_pretrained.side_effect = [OSError("no fp16 weight files"), "PIPE"]
out = remover._load_from_pretrained(cls, "some/model", token="t")
assert out == "PIPE"
assert cls.from_pretrained.call_count == 2
first, second = cls.from_pretrained.call_args_list
assert first.kwargs.get("variant") == "fp16"
assert "variant" not in second.kwargs # the fallback drops the variant
def test_fp32_never_requests_variant(self):
import torch
remover = self._remover(torch.float32)
cls = MagicMock()
cls.from_pretrained.return_value = "PIPE"
remover._load_from_pretrained(cls, "some/model")
cls.from_pretrained.assert_called_once_with("some/model")
assert "variant" not in cls.from_pretrained.call_args.kwargs
Silently mapping them at the alias layer would route an old script into a
profile it never asked for; the remover raises on the unknown name instead.
"""
for retired in ("default", "sdxl", "controlnet", "qwen"):
assert normalize_profile(retired) not in PROFILE_CHOICES
class TestNoReembeddedWatermark:
"""F2 regression: the SDXL removal pipelines must disable the diffusers default
invisible watermarker (``add_watermarker=False``).
"""F2 regression: the SDXL global stage must disable the diffusers watermarker.
diffusers stamps an open "Stable Diffusion XL" DWT-DCT watermark onto every SDXL
output whenever ``invisible-watermark`` is installed. A watermark REMOVER that left
it on would replace one detectable AI watermark (SynthID) with another -- the cleaned
output re-reads as AI. The ControlNet sub-model load must NOT receive the kwarg (it
is not a pipeline and does not accept it).
it on would replace one detectable AI watermark (SynthID) with another -- the
cleaned output re-reads as AI. The ControlNet sub-model load must NOT receive the
kwarg, since it is not a pipeline and does not accept it.
Only sdxl-zimage carries an SDXL pipeline now; qwen-zimage's global stage is
DiffSynth, which has no such watermarker.
"""
def _remover(self, profile: str):
def test_sdxl_global_stage_disables_watermarker(self, monkeypatch: pytest.MonkeyPatch):
if not is_watermark_removal_available():
pytest.skip("torch/diffusers not installed")
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
import diffusers
return WatermarkRemover(device="cpu", pipeline=profile)
from remove_ai_watermarks._internal.sdxl_zimage_pipeline import SdxlZImagePipeline
def _capture(self, monkeypatch, remover):
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
calls: dict[str, dict] = {}
calls: list[tuple[str, dict]] = []
def record(name):
def fake(*_args, **kwargs):
calls[name] = kwargs
return MagicMock()
def fake_load(self, cls, model_id, **kwargs):
calls.append((getattr(cls, "__name__", str(cls)), kwargs))
return MagicMock()
return fake
monkeypatch.setattr(WatermarkRemover, "_load_from_pretrained", fake_load)
monkeypatch.setattr(WatermarkRemover, "_move_to_device_and_optimize", lambda self, p: p)
return calls
monkeypatch.setattr(diffusers.ControlNetModel, "from_pretrained", record("controlnet"))
monkeypatch.setattr(diffusers.AutoencoderKL, "from_pretrained", record("vae"))
monkeypatch.setattr(diffusers.StableDiffusionXLControlNetImg2ImgPipeline, "from_pretrained", record("pipeline"))
monkeypatch.setattr("huggingface_hub.hf_hub_download", lambda *a, **k: "lora.safetensors")
# from_config would otherwise resolve the mock's config as a repo id.
monkeypatch.setattr(diffusers.EulerDiscreteScheduler, "from_config", lambda *a, **k: MagicMock())
def test_sdxl_pipeline_disables_watermarker(self, monkeypatch: pytest.MonkeyPatch):
remover = self._remover("sdxl")
calls = self._capture(monkeypatch, remover)
remover._load_pipeline()
assert any(kw.get("add_watermarker") is False for _, kw in calls), calls
pipeline = SdxlZImagePipeline(device="cuda", torch_dtype=None)
monkeypatch.setattr(type(pipeline), "_require_cuda", lambda self: None)
pipeline._load_sdxl()
def test_controlnet_pipeline_disables_watermarker(self, monkeypatch: pytest.MonkeyPatch):
remover = self._remover("controlnet")
calls = self._capture(monkeypatch, remover)
remover._load_controlnet_pipeline()
by_cls = dict(calls)
# the SDXL pipeline load disables the watermarker...
assert by_cls["StableDiffusionXLControlNetImg2ImgPipeline"].get("add_watermarker") is False
# ...but the ControlNet sub-model load must not carry the kwarg (it would error).
assert "add_watermarker" not in by_cls["ControlNetModel"]
class _StubImage:
"""Minimal PIL.Image stand-in: just the ``width``/``height`` the pure helper reads."""
def __init__(self, width: int, height: int) -> None:
self.width = width
self.height = height
class TestQwenKwargs:
"""_build_qwen_kwargs is pure (no torch); guards the Qwen-Image call shape.
watermark_remover imports torch under a try/except, so the module (and this pure
helper) imports fine in the default+dev CI env where torch is absent.
"""
def test_uses_true_cfg_not_guidance_scale(self):
from remove_ai_watermarks._internal.watermark_remover import _build_qwen_kwargs
gen = object()
img = _StubImage(2816, 1536)
kwargs = _build_qwen_kwargs(img, strength=0.3, num_inference_steps=40, true_cfg_scale=4.0, generator=gen)
# Qwen uses true_cfg_scale, NOT SDXL's guidance_scale.
assert kwargs["true_cfg_scale"] == 4.0
assert "guidance_scale" not in kwargs
# The scrub still comes from strength; image + generator pass through.
assert kwargs["strength"] == 0.3
assert kwargs["image"] is img
assert kwargs["generator"] is gen
assert kwargs["prompt"] == "high quality, sharp, detailed, faithful to the original"
assert kwargs["negative_prompt"] == "blurry, lowres, distorted text, garbled text, artifacts"
def test_passes_explicit_aspect_preserving_size(self):
# Without height/width the pipeline defaults to 1024x1024 and squishes non-square
# input (the abba mixed-seam regression). Both already multiples of 16 -> unchanged.
from remove_ai_watermarks._internal.watermark_remover import _build_qwen_kwargs
kwargs = _build_qwen_kwargs(
_StubImage(2816, 1536), strength=0.25, num_inference_steps=40, true_cfg_scale=4.0, generator=None
)
assert kwargs["width"] == 2816
assert kwargs["height"] == 1536
def test_qwen_target_size_floors_to_multiple_of_16(self):
from remove_ai_watermarks._internal.watermark_remover import _qwen_target_size
assert _qwen_target_size(2816, 1536) == (2816, 1536) # already /16
assert _qwen_target_size(1122, 1402) == (1120, 1392) # floored
assert _qwen_target_size(10, 10) == (16, 16) # min clamp, never 0
def test_qwen_model_id_is_qwen_image(self):
from remove_ai_watermarks._internal.watermark_profiles import QWEN_MODEL_ID
assert QWEN_MODEL_ID == "Qwen/Qwen-Image"
assert calls["pipeline"].get("add_watermarker") is False
assert "add_watermarker" not in calls["controlnet"]
class TestResolveStrength:
"""resolve_strength applies the vendor default only when strength is unset."""
"""resolve_strength answers for sdxl-zimage and defers for qwen-zimage."""
def test_none_is_vendor_adaptive(self):
# No vendor -> unknown default; OpenAI lower, Google == unknown. The sdxl/controlnet
# pipelines share this ladder (the certified controlnet floors); qwen has its own
# (see test_qwen_pipeline_uses_its_own_higher_ladder).
assert resolve_strength(None) == UNKNOWN_STRENGTH
assert resolve_strength(None, "openai") == OPENAI_STRENGTH
assert resolve_strength(None, "google") == GEMINI_STRENGTH
assert resolve_strength(None, None) == UNKNOWN_STRENGTH
# An unrecognized vendor string falls through to the unknown default.
assert resolve_strength(None, "adobe") == UNKNOWN_STRENGTH
# sdxl/controlnet pipelines (and the "default" alias) use the same shared ladder.
assert resolve_strength(None, "google", "controlnet") == GEMINI_STRENGTH
assert resolve_strength(None, "google", "sdxl") == GEMINI_STRENGTH
def test_qwen_zimage_answers_from_the_resolution_curve(self):
"""The function is total: it owns both policies rather than returning None.
def test_qwen_pipeline_uses_its_own_higher_ladder(self):
# Qwen's certified Gemini floor (0.25) is HIGHER than controlnet's (0.15); OpenAI
# matches (0.10). Unknown vendor on qwen tracks the higher Gemini value. This retires
# the old manual "pass --strength 0.25 for Gemini on qwen" workaround.
from remove_ai_watermarks._internal.watermark_profiles import QWEN_GEMINI_STRENGTH, QWEN_OPENAI_STRENGTH
qwen-zimage picks strength from image area, so it takes the size. Returning
None for it would push that branch onto every caller and leave one of the two
strength policies living outside this module. The vendor is ignored here on
purpose - the curve, not the issuer, is what was calibrated.
"""
assert resolve_strength(None, "google", "qwen-zimage", size=(2000, 1850)) == pytest.approx(0.154)
assert resolve_strength(None, None, "qwen-zimage", size=(600, 500)) == pytest.approx(0.084)
assert QWEN_GEMINI_STRENGTH == 0.25
assert QWEN_OPENAI_STRENGTH == 0.10
assert resolve_strength(None, "google", "qwen") == QWEN_GEMINI_STRENGTH
assert resolve_strength(None, "openai", "qwen") == QWEN_OPENAI_STRENGTH
assert resolve_strength(None, None, "qwen") == QWEN_GEMINI_STRENGTH # unknown -> higher floor
assert resolve_strength(None, "google", "qwen") > resolve_strength(None, "google", "controlnet")
# An explicit strength still wins on qwen.
assert resolve_strength(0.12, "google", "qwen") == 0.12
def test_qwen_zimage_without_a_size_fails_loudly(self):
"""A missing size must not silently fall back to some vendor value."""
with pytest.raises(ValueError, match="size is required"):
resolve_strength(None, "google", "qwen-zimage")
def test_ladder_is_the_certified_controlnet_floors(self):
# The unified ladder == the oracle-certified controlnet floors. Lowered on the
# 2026-06-14 Modal re-test (OpenAI 0.10, Google/unknown 0.15); Google is the
# more-robust watermark, so it is higher.
assert OPENAI_STRENGTH == 0.10
assert GEMINI_STRENGTH == 0.15
assert UNKNOWN_STRENGTH == 0.15
assert OPENAI_STRENGTH < GEMINI_STRENGTH
def test_sdxl_zimage_uses_its_flat_vendor_ladder(self):
def test_default_strength_alias_is_unknown_vendor_value(self):
assert DEFAULT_STRENGTH == UNKNOWN_STRENGTH
assert OPENAI_STRENGTH < UNKNOWN_STRENGTH
assert SDXL_ZIMAGE_OPENAI_STRENGTH == 0.15
assert SDXL_ZIMAGE_GEMINI_STRENGTH == 0.25
assert SDXL_ZIMAGE_UNKNOWN_STRENGTH == SDXL_ZIMAGE_GEMINI_STRENGTH
assert resolve_strength(None, "openai", "sdxl-zimage") == SDXL_ZIMAGE_OPENAI_STRENGTH
assert resolve_strength(None, "google", "sdxl-zimage") == SDXL_ZIMAGE_GEMINI_STRENGTH
# An unrecognised issuer takes the stricter Gemini value, not the OpenAI one.
assert resolve_strength(None, "adobe", "sdxl-zimage") == SDXL_ZIMAGE_UNKNOWN_STRENGTH
assert resolve_strength(None, None, "sdxl-zimage") == SDXL_ZIMAGE_UNKNOWN_STRENGTH
def test_strength_default_help_derives_from_constants(self):
# The CLI --strength help is built from this, so it can never drift from the ladder.
h = strength_default_help()
assert str(OPENAI_STRENGTH) in h
assert str(GEMINI_STRENGTH) in h
assert str(UNKNOWN_STRENGTH) in h
assert str(SDXL_ZIMAGE_OPENAI_STRENGTH) in h
assert str(SDXL_ZIMAGE_GEMINI_STRENGTH) in h
def test_explicit_value_overrides_vendor(self):
assert resolve_strength(0.3) == 0.3
assert resolve_strength(0.3, "openai") == 0.3
assert resolve_strength(0.3, "openai", "sdxl-zimage") == 0.3
assert resolve_strength(0.3, None, "qwen-zimage") == 0.3
def test_explicit_zero_is_respected_not_treated_as_unset(self):
# 0.0 is falsy but explicit -- must not fall through to the vendor default
# 0.0 is falsy but explicit -- it must not fall through to the vendor default
# (the old `strength or DEFAULT` bug would have). Range validation lives in
# remove_watermark, not here.
assert resolve_strength(0.0) == 0.0
assert resolve_strength(0.0, "google") == 0.0
assert resolve_strength(0.0, "google", "sdxl-zimage") == 0.0
assert resolve_strength(0.0, None, "qwen-zimage") == 0.0
class TestVendorForStrength:
@@ -467,54 +332,3 @@ class TestPlatformPaths:
# If we get here without error, asset loading works
assert engine._alpha_small.shape == (48, 48)
assert engine._alpha_large.shape == (96, 96)
class TestFp16VaeFix:
"""The plain SDXL img2img pipeline must swap in the fp16-fixed VAE on fp16
GPUs to avoid the NaN/all-black decode (issue #29). Pure decision logic, no
torch or model download needed."""
DEFAULT = "stabilityai/stable-diffusion-xl-base-1.0"
def test_default_sdxl_on_fp16_needs_fix(self):
from remove_ai_watermarks._internal.watermark_remover import _needs_fp16_vae_fix
assert _needs_fp16_vae_fix(self.DEFAULT, self.DEFAULT, is_fp16=True) is True
def test_fp32_does_not_need_fix(self):
"""cpu/mps run fp32, where the stock SDXL VAE is fine."""
from remove_ai_watermarks._internal.watermark_remover import _needs_fp16_vae_fix
assert _needs_fp16_vae_fix(self.DEFAULT, self.DEFAULT, is_fp16=False) is False
def test_non_default_model_keeps_own_vae(self):
"""A custom (non-SDXL) checkpoint must not get the SDXL-specific VAE."""
from remove_ai_watermarks._internal.watermark_remover import _needs_fp16_vae_fix
assert _needs_fp16_vae_fix("runwayml/stable-diffusion-v1-5", self.DEFAULT, is_fp16=True) is False
class TestDegenerateOutputGuard:
"""The fp16 black-output safety net (#29/#41): detect an all-black/NaN frame so
``remove_watermark`` can retry in fp32. Pure image statistics, no model needed."""
def test_all_black_is_degenerate(self):
from remove_ai_watermarks._internal.watermark_remover import _is_degenerate_image
black = Image.fromarray(np.zeros((64, 64, 3), np.uint8))
assert _is_degenerate_image(black) is True
def test_normal_image_is_not_degenerate(self):
from remove_ai_watermarks._internal.watermark_remover import _is_degenerate_image
rng = np.random.default_rng(0)
normal = Image.fromarray(rng.integers(0, 256, (64, 64, 3), dtype=np.uint8))
assert _is_degenerate_image(normal) is False
def test_dark_but_textured_image_is_not_degenerate(self):
"""A legitimately dark photo with real detail must NOT be flagged (variance guard)."""
from remove_ai_watermarks._internal.watermark_remover import _is_degenerate_image
rng = np.random.default_rng(1)
dark = Image.fromarray(rng.integers(0, 40, (64, 64, 3), dtype=np.uint8))
assert _is_degenerate_image(dark) is False
tests/test_qwen_zimage_pipeline.py
+14 -17
View File
@@ -530,12 +530,10 @@ def test_profile_defaults_to_four_global_steps():
)
assert normalize_profile("qwen-zimage") == "qwen-zimage"
assert resolve_steps(None, "qwen-zimage") == 4
assert resolve_steps(None, "controlnet") == 50
assert resolve_steps(12, "qwen-zimage") == 12
assert resolve_seed(None, "qwen-zimage") == 0
assert resolve_seed(None, "controlnet") is None
assert resolve_seed(17, "qwen-zimage") == 17
assert resolve_steps(None) == 4
assert resolve_steps(12) == 12
assert resolve_seed(None) == 0
assert resolve_seed(17) == 17
def test_cli_exposes_qwen_zimage_profile():
@@ -586,7 +584,7 @@ def test_watermark_remover_dispatches_to_full_pipeline(tmp_path, monkeypatch):
runtime = MagicMock()
runtime.run.return_value = Image.new("RGB", (64, 48), (50, 60, 70))
remover = WatermarkRemover(device="cpu", pipeline="qwen-zimage")
remover = WatermarkRemover(device="cuda", pipeline="qwen-zimage")
monkeypatch.setattr(remover, "_load_qwen_zimage_pipeline", lambda: runtime)
assert remover.model_id == "Qwen/Qwen-Image-2512 + Tongyi-MAI/Z-Image-Turbo"
@@ -612,7 +610,7 @@ def test_watermark_remover_dispatches_qwen_tiling_to_full_pipeline(tmp_path, mon
runtime = MagicMock()
runtime.run.return_value = Image.new("RGB", (96, 80), (50, 60, 70))
remover = WatermarkRemover(device="cpu", pipeline="qwen-zimage")
remover = WatermarkRemover(device="cuda", pipeline="qwen-zimage")
monkeypatch.setattr(remover, "_load_qwen_zimage_pipeline", lambda: runtime)
remover.remove_watermark(
@@ -746,11 +744,11 @@ def test_qwen_zimage_rejects_runtime_knobs_that_change_fixed_graph(tmp_path, mon
_mock_watermark_runtime_deps(monkeypatch)
with pytest.raises(ValueError, match="fixed Qwen-Image-2512"):
WatermarkRemover(model_id="custom/model", device="cpu", pipeline="qwen-zimage")
WatermarkRemover(model_id="custom/model", device="cuda", pipeline="qwen-zimage")
source = tmp_path / "source.png"
Image.new("RGB", (64, 48)).save(source)
remover = WatermarkRemover(device="cpu", pipeline="qwen-zimage")
remover = WatermarkRemover(device="cuda", pipeline="qwen-zimage")
with pytest.raises(ValueError, match=r"CFG 1\.0"):
remover.remove_watermark(source, guidance_scale=2.0)
with pytest.raises(ValueError, match="requires 4 steps"):
@@ -786,10 +784,11 @@ def test_sdxl_zimage_strength_is_vendor_adaptive_and_leaves_other_profiles_alone
assert resolve_strength(None, "google", "sdxl-zimage") == pytest.approx(SDXL_ZIMAGE_GEMINI_STRENGTH)
# Unknown provenance takes the stricter of the two.
assert resolve_strength(None, None, "sdxl-zimage") == pytest.approx(SDXL_ZIMAGE_GEMINI_STRENGTH)
# An explicit value still wins, and the older profiles are untouched.
# An explicit value still wins, and qwen-zimage is untouched by this ladder: it
# defers to its resolution curve rather than to a vendor value.
assert resolve_strength(0.4, "google", "sdxl-zimage") == pytest.approx(0.4)
assert resolve_strength(None, "openai", "controlnet") == pytest.approx(0.10)
assert resolve_strength(None, "google", "controlnet") == pytest.approx(0.15)
assert resolve_strength(None, "openai", "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():
@@ -800,10 +799,8 @@ def test_sdxl_zimage_shares_the_four_step_seed_and_step_contract():
)
assert normalize_profile("sdxl_zimage") == "sdxl-zimage"
assert resolve_steps(None, "sdxl-zimage") == 4
assert resolve_seed(None, "sdxl-zimage") == 0
assert resolve_steps(None, "controlnet") == 50
assert resolve_seed(None, "controlnet") is None
assert resolve_steps(None) == 4
assert resolve_seed(None) == 0
def test_sdxl_requested_steps_compensate_for_the_diffusers_truncation():
tests/test_watermark_profiles.py
-48
View File
@@ -1,48 +0,0 @@
"""Pure tests for the strength/steps profile helpers (no model, no torch needed)."""
from __future__ import annotations
import pytest
from remove_ai_watermarks._internal.watermark_profiles import resolve_strength, viable_steps
class TestViableSteps:
"""Guards the crash found by the release smoke matrix on 2026-07-19.
diffusers derives its img2img timesteps as ``int(steps * strength)``. When that
rounds to zero the pipeline builds an empty tensor and dies deep inside attention
with "cannot reshape tensor of 0 elements into shape [0, -1, 1, 512]". At the
default strength 0.15 that was every ``--steps`` below 7, reachable with entirely
valid CLI arguments and no special flags.
"""
@pytest.mark.parametrize(
("steps", "strength"),
[(1, 0.15), (2, 0.15), (5, 0.15), (6, 0.15), (5, 0.10), (9, 0.10), (1, 0.5)],
)
def test_never_returns_a_count_that_denoises_zero_steps(self, steps: int, strength: float):
assert int(viable_steps(steps, strength) * strength) >= 1
@pytest.mark.parametrize(
("steps", "strength"),
[(50, 0.15), (20, 0.15), (7, 0.15), (10, 0.10), (2, 0.5), (50, 1.0)],
)
def test_leaves_a_workable_count_untouched(self, steps: int, strength: float):
assert viable_steps(steps, strength) == steps
def test_raises_only_to_the_minimum_needed(self):
# strength 0.15 needs 7 (int(7*0.15)==1); it must not jump to some larger default.
assert viable_steps(5, 0.15) == 7
assert viable_steps(1, 0.10) == 10
def test_the_vendor_defaults_all_have_a_reachable_floor(self):
for vendor in (None, "openai", "google"):
strength = resolve_strength(None, vendor)
assert int(viable_steps(1, strength) * strength) >= 1
@pytest.mark.parametrize("strength", [0.0, -0.1])
def test_a_non_positive_strength_cannot_loop_or_divide_by_zero(self, strength: float):
# No denoising is possible at all here; return the caller's value rather than
# dividing by zero or spinning.
assert viable_steps(20, strength) == 20
uv.lock
Generated
+1 -1
View File
@@ -3331,7 +3331,7 @@ wheels = [
[[package]]
name = "remove-ai-watermarks"
version = "0.23.0"
version = "0.24.0"
source = { editable = "." }
dependencies = [
{ name = "c2pa-python" },