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
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