mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-09 23:50:40 +02:00
Delete every knob the fixed profiles cannot honor
The CLI still advertised --model, --steps, --guidance-scale, --device and a deprecated --auto. Each pinned a value the two surviving profiles fix -- the model stack, the per-stage distilled schedule, CFG 1.0, CUDA -- so the only outcome any of them had was an error raised several frames below the caller, under a message naming an internal profile. A flag whose sole result is a refusal is worse than no flag: it advertises a capability that does not exist, and it lets a wrapper thread a value that will silently do nothing. They are gone from the parser, from InvisibleEngine, and from WatermarkRemover, so the failure is now a TypeError or a Click "No such option" at the point the caller can act on. The install hint was wrong in the same way. is_available() checked torch and diffusers, then told the user to install [diffusion] -- which contains neither DiffSynth nor the Z-Image face stage both profiles run. Following the advice produced a second, different failure. The module list and the extra name now live once in watermark_profiles (REMOVAL_MODULES, INVISIBLE_EXTRA) and are read by both the CLI gate and the remover's precondition, which cannot drift apart because they are the same tuple. The adaptive-polish default moved out of the argument parser. It was resolved by reading Click's parameter source, which put per-profile data in the CLI layer, left the engine declaring the opposite default (False vs True) so a library caller and a CLI caller on one profile got different output, and lost the polish entirely for anything that supplies the flag non-interactively. The flag is now tri-state (default=None) and resolve_adaptive_polish owns the per-profile answer. The seed follows the same rule: the CLI stopped pre-resolving it. Dead code removed with it: six scan_*_video wrappers and the _scan_video helper none of them had a caller for, PNG_METADATA_KEYS, feather_region_composite and the remover region path that was only reachable from a no-caller convenience wrapper, remove_watermark_batch on both layers, try_empty_device_cache, the _generate/_run_qwen_zimage pass-through pair, self.model_id, and the _internal PEP 562 shim that no caller ever went through. get_device now answers cuda or cpu only: mps and xpu travelled one frame to the same CUDA-only refusal while costing a device probe each, and that refusal now names the resolved device, so device=None on a CUDA-less host says 'cpu' rather than 'None'. The XPU wheel index went with them. Docs: README, cli, installation, python-api, supported-signals, known-limitations and module-internals all still described the removed profiles, the CPU/MPS/XPU ladder, a `default`->`sdxl` alias, and the wrong extra. known-limitations still listed the retired SDXL strength ladder as current. scripts/smoke_matrix.py and real_examples_e2e.py drove --device mps. Next release is 0.25.0, not a patch: this removes public parameters and narrows a published extra on top of the released 0.24.0. pre-commit: 1) maintain.sh - exit 0 (1091 tests, Pyright 0 errors, no vulnerabilities); 2) /simplify - 4 agents, 11 findings applied, 2 skipped (dropping the `device` parameter entirely, which raiw-app pins; folding diffsynth into the `diffusion` extra, which video-only callers do not need); 3) docs sync - grepped every removed identifier across README, docs/, scripts/, .claude/; updated 9 docs; 4) CLAUDE.md - added the no-error-only-knobs rule to .claude/rules/development.md Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
bf4bfc1ab7
commit
52b2c115e8
+52
-35
@@ -327,7 +327,13 @@ class TestInvisibleCommand:
|
||||
expected = sample_png.with_stem(sample_png.stem + "_clean")
|
||||
assert expected.exists()
|
||||
|
||||
def test_invisible_adaptive_polish_off_by_default_under_qwen_zimage(self, runner, sample_png):
|
||||
def test_invisible_leaves_the_polish_default_to_the_library(self, runner, sample_png):
|
||||
"""An untyped --adaptive-polish reaches the engine as None, not as a value.
|
||||
|
||||
The per-profile default lives in watermark_profiles, so the CLI must pass the
|
||||
user's non-choice through rather than resolving it here. Resolving in the CLI
|
||||
is how the library and the CLI came to disagree on the same profile.
|
||||
"""
|
||||
mock_cls, mock_engine = _mock_invisible_engine()
|
||||
with (
|
||||
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
|
||||
@@ -336,13 +342,7 @@ class TestInvisibleCommand:
|
||||
):
|
||||
result = runner.invoke(main, ["invisible", str(sample_png), "--force"])
|
||||
assert result.exit_code == 0, result.output
|
||||
# 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
|
||||
assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is None
|
||||
|
||||
def test_invisible_no_adaptive_polish_disables(self, runner, sample_png):
|
||||
mock_cls, mock_engine = _mock_invisible_engine()
|
||||
@@ -355,20 +355,24 @@ class TestInvisibleCommand:
|
||||
assert result.exit_code == 0, result.output
|
||||
assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is False
|
||||
|
||||
def test_invisible_model_and_guidance_scale_flow_to_engine(self, runner, sample_png):
|
||||
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),
|
||||
def test_knobs_the_fixed_stack_cannot_honor_are_not_offered(self, runner, sample_png):
|
||||
"""--model/--steps/--guidance-scale/--device/--auto are gone, not rejected.
|
||||
|
||||
Each pinned a value the profiles fix (model stack, per-stage schedule, CFG 1.0,
|
||||
CUDA), so accepting one only produced an error several layers down -- a flag
|
||||
that advertises a capability the library does not have. Click now refuses the
|
||||
option itself, which is the honest answer and the one a caller can act on.
|
||||
"""
|
||||
for retired in (
|
||||
["--model", "org/custom-sdxl"],
|
||||
["--steps", "20"],
|
||||
["--guidance-scale", "5.5"],
|
||||
["--device", "cpu"],
|
||||
["--auto"],
|
||||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["invisible", str(sample_png), "--model", "org/custom-sdxl", "--guidance-scale", "5.5", "--force"],
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert mock_cls.call_args.kwargs["model_id"] == "org/custom-sdxl"
|
||||
assert mock_engine.remove_watermark.call_args.kwargs["guidance_scale"] == 5.5
|
||||
result = runner.invoke(main, ["invisible", str(sample_png), *retired, "--force"])
|
||||
assert result.exit_code == 2, f"{retired[0]}: {result.output}"
|
||||
assert "No such option" in result.output, f"{retired[0]}: {result.output}"
|
||||
|
||||
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.
|
||||
@@ -530,7 +534,7 @@ class TestAllCommand:
|
||||
result = runner.invoke(main, ["all", str(sample_png), "-o", str(output)])
|
||||
assert result.exit_code != 0, result.output
|
||||
assert "NOT removed" in result.output
|
||||
assert "remove-ai-watermarks[diffusion]" in result.output
|
||||
assert "remove-ai-watermarks[qwen-zimage]" in result.output
|
||||
assert output.exists() # visible + metadata still produced a file
|
||||
|
||||
def test_all_reports_metadata_that_survived_stripping(self, runner, sample_png, tmp_path):
|
||||
@@ -858,10 +862,11 @@ class TestBatchCommand:
|
||||
assert out[0, 0, 3] == 0
|
||||
assert out[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.
|
||||
def test_batch_explicit_adaptive_polish_overrides_the_qwen_zimage_off(self, runner, tmp_path):
|
||||
"""qwen-zimage leaves the polish off by default; a typed flag still turns it on.
|
||||
|
||||
It no longer selects a pipeline: qwen-zimage is the only default there is.
|
||||
The off is a parameter-source check, not a changed default, so it must yield to
|
||||
an explicit --adaptive-polish rather than swallowing it.
|
||||
"""
|
||||
input_dir = _make_batch_dir(tmp_path, count=2)
|
||||
output_dir = tmp_path / "output"
|
||||
@@ -874,12 +879,19 @@ class TestBatchCommand:
|
||||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["batch", str(input_dir), "-o", str(output_dir), "--mode", "invisible", "--auto", "--force"],
|
||||
[
|
||||
"batch",
|
||||
str(input_dir),
|
||||
"-o",
|
||||
str(output_dir),
|
||||
"--mode",
|
||||
"invisible",
|
||||
"--adaptive-polish",
|
||||
"--force",
|
||||
],
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "2 processed" in result.output
|
||||
assert "deprecated" in result.output.lower()
|
||||
# Pipeline stays the default controlnet; --auto only turned the polish on.
|
||||
assert mock_cls.call_args.kwargs["pipeline"] == "qwen-zimage"
|
||||
assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is True
|
||||
|
||||
@@ -923,21 +935,26 @@ class TestBatchCommand:
|
||||
|
||||
|
||||
class TestGpuHintMarkup:
|
||||
"""The diffusion install hint must reach the user with the ``[diffusion]`` token
|
||||
intact (plain output prints it verbatim, with no markup parsing)."""
|
||||
"""The install hint must name the extra that actually makes a removal run.
|
||||
|
||||
def test_invisible_install_hint_keeps_gpu_extra(self, runner, sample_png):
|
||||
It must also survive to the user with its ``[...]`` token intact (plain output
|
||||
prints it verbatim, with no markup parsing). It used to say ``[diffusion]``,
|
||||
which installs torch and diffusers but not the DiffSynth face stage both
|
||||
profiles run -- so following the advice produced a second, different failure.
|
||||
"""
|
||||
|
||||
def test_invisible_install_hint_names_the_working_extra(self, runner, sample_png):
|
||||
with patch("remove_ai_watermarks.invisible_engine.is_available", return_value=False):
|
||||
result = runner.invoke(main, ["invisible", str(sample_png)])
|
||||
assert result.exit_code != 0
|
||||
assert "remove-ai-watermarks[diffusion]" in result.output
|
||||
assert "remove-ai-watermarks[qwen-zimage]" in result.output
|
||||
|
||||
def test_all_install_hint_keeps_gpu_extra(self, runner, sample_png):
|
||||
def test_all_install_hint_names_the_working_extra(self, runner, sample_png):
|
||||
# The `all` pipeline skips the invisible step with a warning that carries
|
||||
# the same hint; it must keep the [diffusion] extra too.
|
||||
# the same hint; it must name the same extra.
|
||||
with patch("remove_ai_watermarks.invisible_engine.is_available", return_value=False):
|
||||
result = runner.invoke(main, ["all", str(sample_png)])
|
||||
assert "remove-ai-watermarks[diffusion]" in result.output
|
||||
assert "remove-ai-watermarks[qwen-zimage]" in result.output
|
||||
|
||||
|
||||
class TestEraseCommand:
|
||||
|
||||
@@ -16,24 +16,26 @@ class TestIsAvailable:
|
||||
result = is_available()
|
||||
assert isinstance(result, bool)
|
||||
|
||||
def test_available_reflects_dependencies(self):
|
||||
"""is_available() is True iff torch + diffusers (the diffusion extra) import.
|
||||
def test_available_reflects_every_module_a_run_needs(self):
|
||||
"""True iff every module in REMOVAL_MODULES imports, diffsynth included.
|
||||
|
||||
Must not assume the full stack: the default+dev CI env has no diffusers.
|
||||
Derived from the same tuple the remover's precondition uses, so this cannot
|
||||
pass while the two disagree -- the drift that let a torch+diffusers-only
|
||||
environment clear the CLI gate and then die at the DiffSynth face stage.
|
||||
Must not assume the full stack: the default+dev CI env has none of it.
|
||||
"""
|
||||
import importlib.util
|
||||
|
||||
expected = all(importlib.util.find_spec(m) is not None for m in ("torch", "diffusers"))
|
||||
from remove_ai_watermarks._internal.watermark_profiles import REMOVAL_MODULES
|
||||
|
||||
assert "diffsynth" in REMOVAL_MODULES
|
||||
expected = all(importlib.util.find_spec(m) is not None for m in REMOVAL_MODULES)
|
||||
assert is_available() is expected
|
||||
|
||||
|
||||
class TestInvisibleEngineInit:
|
||||
"""Tests for InvisibleEngine construction (no GPU required)."""
|
||||
|
||||
def test_default_model_id(self):
|
||||
# SDXL base became the default in May 2026 (defeats SynthID v2).
|
||||
assert InvisibleEngine.DEFAULT_MODEL_ID == "stabilityai/stable-diffusion-xl-base-1.0"
|
||||
|
||||
def test_preload_forwards_global_only(self):
|
||||
engine = object.__new__(InvisibleEngine)
|
||||
engine._remover = SimpleNamespace(preload=lambda **kwargs: setattr(engine, "_preload_kwargs", kwargs))
|
||||
@@ -55,7 +57,7 @@ class TestNativeOutputSize:
|
||||
Image.open(image_path).crop((0, 0, 24, 16)).save(out)
|
||||
return out
|
||||
|
||||
engine._remover = SimpleNamespace(remove_watermark=_remove_watermark)
|
||||
engine._remover = SimpleNamespace(remove_watermark=_remove_watermark, model_profile="qwen-zimage")
|
||||
engine._progress_callback = None
|
||||
src = tmp_path / "src.png"
|
||||
out = tmp_path / "out.png"
|
||||
@@ -113,30 +115,32 @@ class TestTargetSize:
|
||||
assert _target_size(381, 512, 4096) is None
|
||||
|
||||
|
||||
class TestEngineDoesNotFabricateAModelId:
|
||||
"""The engine must forward model_id untouched, including None.
|
||||
class TestEngineConstructsWithoutAModelId:
|
||||
"""Plain construction must reach the remover, and must not name a model.
|
||||
|
||||
It used to substitute DEFAULT_MODEL_ID for None. Once the remover tightened its
|
||||
"you may not override the fixed stack" check from `not in {None, DEFAULT_MODEL_ID}`
|
||||
to `is not None`, that substitution made EVERY InvisibleEngine construction raise -
|
||||
and no test saw it, because the library tests build WatermarkRemover directly while
|
||||
the engine tests mock it. A deployed Modal worker caught it instead.
|
||||
The engine used to take a ``model_id`` and substitute the SDXL default for None.
|
||||
Once the remover tightened its "you may not override the fixed stack" check to
|
||||
``is not None``, that substitution made EVERY construction raise -- and no test
|
||||
saw it, because the library tests build WatermarkRemover directly while the engine
|
||||
tests mock it. A deployed Modal worker caught it instead. The parameter is gone on
|
||||
both sides now, so guard the property that broke: a default construction reaches
|
||||
the remover, carrying no model at all.
|
||||
"""
|
||||
|
||||
def test_none_stays_none(self):
|
||||
def test_default_construction_names_no_model(self):
|
||||
from unittest.mock import patch
|
||||
|
||||
import remove_ai_watermarks.invisible_engine as engine_module
|
||||
|
||||
with patch("remove_ai_watermarks._internal.watermark_remover.WatermarkRemover") as remover:
|
||||
engine_module.InvisibleEngine(pipeline="qwen-zimage")
|
||||
assert remover.call_args.kwargs["model_id"] is None
|
||||
assert remover.call_count == 1
|
||||
assert "model_id" not in remover.call_args.kwargs
|
||||
|
||||
def test_an_explicit_model_id_still_reaches_the_remover_to_be_rejected(self):
|
||||
from unittest.mock import patch
|
||||
def test_a_model_id_is_not_accepted(self):
|
||||
import pytest
|
||||
|
||||
import remove_ai_watermarks.invisible_engine as engine_module
|
||||
|
||||
with patch("remove_ai_watermarks._internal.watermark_remover.WatermarkRemover") as remover:
|
||||
engine_module.InvisibleEngine(model_id="org/custom", pipeline="qwen-zimage")
|
||||
assert remover.call_args.kwargs["model_id"] == "org/custom"
|
||||
with pytest.raises(TypeError):
|
||||
engine_module.InvisibleEngine(model_id="org/custom", pipeline="qwen-zimage") # type: ignore[call-arg]
|
||||
|
||||
+57
-33
@@ -1,7 +1,7 @@
|
||||
"""Tests for cross-platform and cross-device compatibility.
|
||||
"""Tests for device detection, profile resolution, and platform-specific paths.
|
||||
|
||||
Verifies that device detection, MPS fallback, and platform-specific
|
||||
code paths work correctly on CPU, MPS (macOS), and CUDA (Linux/Windows).
|
||||
Invisible-watermark removal is CUDA-only, so the device tests here assert a binary
|
||||
answer and a clean refusal rather than a fallback ladder.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -27,34 +27,38 @@ from remove_ai_watermarks._internal.watermark_remover import get_device, is_wate
|
||||
|
||||
|
||||
class TestDeviceDetection:
|
||||
"""Tests for get_device() across platforms."""
|
||||
"""get_device() is binary: CUDA, or the "cpu" that names its absence."""
|
||||
|
||||
def test_returns_valid_device(self):
|
||||
device = get_device()
|
||||
assert device in ("cpu", "mps", "cuda", "xpu")
|
||||
def test_answer_is_cuda_or_cpu(self):
|
||||
"""No mps/xpu answer exists. Both would travel one frame to the same refusal.
|
||||
|
||||
def test_cpu_fallback_when_no_gpu(self):
|
||||
"""On CI / machines without GPU, should fall back to cpu or mps."""
|
||||
device = get_device()
|
||||
# Just verify it doesn't crash and returns a valid string
|
||||
assert isinstance(device, str)
|
||||
Reporting them anyway cost a device probe each and let a caller believe the
|
||||
library had an Apple-silicon or Intel-GPU path that it does not.
|
||||
"""
|
||||
assert get_device() in ("cpu", "cuda")
|
||||
|
||||
@patch("remove_ai_watermarks._internal.watermark_remover._HAS_TORCH", False)
|
||||
def test_no_torch_returns_cpu(self):
|
||||
assert get_device() == "cpu"
|
||||
|
||||
def test_xpu_selected_when_available(self):
|
||||
"""An XPU-enabled torch (no CUDA) routes to the Intel GPU backend.
|
||||
def test_working_cuda_is_selected_and_probed(self):
|
||||
"""A reported CUDA device is smoke-tested before it is returned.
|
||||
|
||||
The whole torch module is mocked so the smoke-test ops succeed without
|
||||
any real device; cuda must read False so the cuda branch is skipped.
|
||||
torch.cuda.is_available() can be True on a build whose CUDA backend then
|
||||
raises on the first real op; without the probe that surfaced much later.
|
||||
"""
|
||||
fake_torch = MagicMock()
|
||||
fake_torch.cuda.is_available.return_value = False
|
||||
fake_torch.xpu.is_available.return_value = True
|
||||
fake_torch.cuda.is_available.return_value = True
|
||||
with patch("remove_ai_watermarks._internal.watermark_remover.torch", fake_torch):
|
||||
assert get_device() == "xpu"
|
||||
fake_torch.tensor.assert_called_with([1.0], device="xpu")
|
||||
assert get_device() == "cuda"
|
||||
fake_torch.tensor.assert_called_with([1.0], device="cuda")
|
||||
|
||||
def test_broken_cuda_backend_falls_back_to_cpu(self):
|
||||
fake_torch = MagicMock()
|
||||
fake_torch.cuda.is_available.return_value = True
|
||||
fake_torch.tensor.side_effect = RuntimeError("no kernel image")
|
||||
with patch("remove_ai_watermarks._internal.watermark_remover.torch", fake_torch):
|
||||
assert get_device() == "cpu"
|
||||
|
||||
def test_non_cuda_devices_are_refused_at_construction(self):
|
||||
"""CUDA is a precondition of the object, not of the run.
|
||||
@@ -78,21 +82,21 @@ class TestDeviceDetection:
|
||||
assert remover.device == "cuda"
|
||||
assert remover.torch_dtype == torch.bfloat16
|
||||
|
||||
def test_the_refusal_names_the_resolved_device_not_a_bare_none(self):
|
||||
"""``device=None`` on a CUDA-less host must report "cpu", not "None".
|
||||
|
||||
class TestEmptyDeviceCache:
|
||||
"""try_empty_device_cache is all that remains of the img2img runner.
|
||||
The message used to interpolate the raw argument, so the common auto-detect
|
||||
path told the user that ``'None'`` cannot run the removal.
|
||||
"""
|
||||
if not is_watermark_removal_available():
|
||||
pytest.skip("torch/diffusers not installed")
|
||||
from remove_ai_watermarks._internal import watermark_remover as module
|
||||
|
||||
Its module lost run_img2img and the MPS fallback along with the CPU/MPS profiles;
|
||||
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_unknown_backend_is_a_silent_no_op(self):
|
||||
from remove_ai_watermarks._internal.watermark_remover import try_empty_device_cache
|
||||
|
||||
try_empty_device_cache("cpu")
|
||||
try_empty_device_cache("definitely-not-a-backend")
|
||||
with (
|
||||
patch.object(module, "get_device", return_value="cpu"),
|
||||
pytest.raises(ValueError, match="'cpu' cannot run it"),
|
||||
):
|
||||
module.WatermarkRemover(device=None)
|
||||
|
||||
|
||||
class TestModelProfiles:
|
||||
@@ -116,6 +120,26 @@ class TestModelProfiles:
|
||||
assert normalize_profile(retired) not in PROFILE_CHOICES
|
||||
|
||||
|
||||
class TestResolveAdaptivePolish:
|
||||
"""The polish default is per-profile data, not a CLI parameter-source inference."""
|
||||
|
||||
def test_unset_follows_the_profile(self):
|
||||
from remove_ai_watermarks._internal.watermark_profiles import resolve_adaptive_polish
|
||||
|
||||
# qwen-zimage already matches the input's detail level, so polishing it only
|
||||
# moves the output away from upstream. An SDXL global pass leaves the softer
|
||||
# result the polish exists for.
|
||||
assert resolve_adaptive_polish(None, "qwen-zimage") is False
|
||||
assert resolve_adaptive_polish(None, "sdxl-zimage") is True
|
||||
assert resolve_adaptive_polish(None, "qwen_zimage") is False
|
||||
|
||||
def test_an_explicit_choice_always_wins(self):
|
||||
from remove_ai_watermarks._internal.watermark_profiles import resolve_adaptive_polish
|
||||
|
||||
assert resolve_adaptive_polish(True, "qwen-zimage") is True
|
||||
assert resolve_adaptive_polish(False, "sdxl-zimage") is False
|
||||
|
||||
|
||||
class TestNoReembeddedWatermark:
|
||||
"""F2 regression: the SDXL global stage must disable the diffusers watermarker.
|
||||
|
||||
|
||||
@@ -175,6 +175,7 @@ def test_cpu_offload_forces_both_stacks_to_stream(monkeypatch, cpu_offload, expe
|
||||
monkeypatch.setattr(pipeline_module, "QwenZImagePipeline", Recorder)
|
||||
|
||||
remover = module.WatermarkRemover.__new__(module.WatermarkRemover)
|
||||
remover.model_profile = "qwen-zimage"
|
||||
remover.device = "cuda"
|
||||
remover.torch_dtype = None
|
||||
remover.hf_token = None
|
||||
@@ -526,16 +527,13 @@ def test_face_composite_preserves_every_pixel_outside_mask():
|
||||
assert np.all(result[12:20, 12:20] == 240)
|
||||
|
||||
|
||||
def test_profile_defaults_to_four_global_steps():
|
||||
from remove_ai_watermarks._internal.watermark_profiles import (
|
||||
normalize_profile,
|
||||
resolve_seed,
|
||||
resolve_steps,
|
||||
)
|
||||
def test_profile_defaults_to_four_global_steps_and_a_fixed_seed():
|
||||
"""The step count belongs to the stage, not to a caller-settable profile knob."""
|
||||
from remove_ai_watermarks._internal.qwen_zimage_pipeline import GLOBAL_STEPS
|
||||
from remove_ai_watermarks._internal.watermark_profiles import normalize_profile, resolve_seed
|
||||
|
||||
assert normalize_profile("qwen-zimage") == "qwen-zimage"
|
||||
assert resolve_steps(None) == 4
|
||||
assert resolve_steps(12) == 12
|
||||
assert GLOBAL_STEPS == 4
|
||||
assert resolve_seed(None) == 0
|
||||
assert resolve_seed(17) == 17
|
||||
|
||||
@@ -560,8 +558,11 @@ def test_cli_qwen_zimage_keeps_profile_postprocess_default(tmp_image_path, monke
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is False
|
||||
assert mock_engine.remove_watermark.call_args.kwargs["seed"] == 0
|
||||
# Both defaults are the profile's, resolved once by the library rather than
|
||||
# pre-resolved here: the CLI passes them through unset so a library caller on the
|
||||
# same profile gets the same answer.
|
||||
assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is None
|
||||
assert mock_engine.remove_watermark.call_args.kwargs["seed"] is None
|
||||
|
||||
result = CliRunner().invoke(
|
||||
cli.main,
|
||||
@@ -590,7 +591,6 @@ def test_watermark_remover_dispatches_to_full_pipeline(tmp_path, monkeypatch):
|
||||
runtime.run.return_value = Image.new("RGB", (64, 48), (50, 60, 70))
|
||||
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"
|
||||
|
||||
remover.remove_watermark(
|
||||
source,
|
||||
@@ -743,23 +743,30 @@ def test_watermark_remover_forwards_global_only_preload(monkeypatch):
|
||||
runtime.preload.assert_called_once_with(global_only=True)
|
||||
|
||||
|
||||
def test_qwen_zimage_rejects_runtime_knobs_that_change_fixed_graph(tmp_path, monkeypatch):
|
||||
def test_the_fixed_graph_offers_no_runtime_knob_to_reject(tmp_path, monkeypatch):
|
||||
"""model_id, steps and CFG are not parameters at any layer.
|
||||
|
||||
They used to be accepted and then rejected, which put the failure several frames
|
||||
below the caller and made the surface advertise choices the pinned stack cannot
|
||||
honor. TypeError from the signature is the earlier, clearer answer -- and it is
|
||||
what keeps a wrapper from threading a value that would silently do nothing.
|
||||
"""
|
||||
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
|
||||
|
||||
_mock_watermark_runtime_deps(monkeypatch)
|
||||
with pytest.raises(ValueError, match="fixed Qwen-Image-2512"):
|
||||
WatermarkRemover(model_id="custom/model", device="cuda", pipeline="qwen-zimage")
|
||||
with pytest.raises(TypeError):
|
||||
WatermarkRemover(model_id="custom/model", device="cuda", pipeline="qwen-zimage") # type: ignore[call-arg]
|
||||
|
||||
source = tmp_path / "source.png"
|
||||
Image.new("RGB", (64, 48)).save(source)
|
||||
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"):
|
||||
remover.remove_watermark(source, num_inference_steps=8)
|
||||
with pytest.raises(TypeError):
|
||||
remover.remove_watermark(source, guidance_scale=2.0) # type: ignore[call-arg]
|
||||
with pytest.raises(TypeError):
|
||||
remover.remove_watermark(source, num_inference_steps=8) # type: ignore[call-arg]
|
||||
|
||||
|
||||
def test_invisible_engine_uses_qwen_zimage_step_default(tmp_image_path, tmp_path):
|
||||
def test_invisible_engine_passes_the_seed_but_never_a_step_count(tmp_image_path, tmp_path):
|
||||
from remove_ai_watermarks.invisible_engine import InvisibleEngine
|
||||
|
||||
engine = InvisibleEngine.__new__(InvisibleEngine)
|
||||
@@ -769,7 +776,10 @@ def test_invisible_engine_uses_qwen_zimage_step_default(tmp_image_path, tmp_path
|
||||
|
||||
engine.remove_watermark(tmp_image_path, tmp_path / "clean.png")
|
||||
|
||||
assert engine._remover.remove_watermark.call_args.kwargs["num_inference_steps"] == 4
|
||||
kwargs = engine._remover.remove_watermark.call_args.kwargs
|
||||
assert kwargs["seed"] == 0
|
||||
assert "num_inference_steps" not in kwargs
|
||||
assert "guidance_scale" not in kwargs
|
||||
|
||||
|
||||
def test_sdxl_zimage_strength_is_vendor_adaptive_and_leaves_other_profiles_alone():
|
||||
@@ -791,15 +801,10 @@ def test_sdxl_zimage_strength_is_vendor_adaptive_and_leaves_other_profiles_alone
|
||||
assert resolve_strength(None, "google", "qwen-zimage", size=(2000, 1850)) == pytest.approx(0.154)
|
||||
|
||||
|
||||
def test_sdxl_zimage_shares_the_four_step_seed_and_step_contract():
|
||||
from remove_ai_watermarks._internal.watermark_profiles import (
|
||||
normalize_profile,
|
||||
resolve_seed,
|
||||
resolve_steps,
|
||||
)
|
||||
def test_sdxl_zimage_shares_the_fixed_seed_contract():
|
||||
from remove_ai_watermarks._internal.watermark_profiles import normalize_profile, resolve_seed
|
||||
|
||||
assert normalize_profile("sdxl_zimage") == "sdxl-zimage"
|
||||
assert resolve_steps(None) == 4
|
||||
assert resolve_seed(None) == 0
|
||||
|
||||
|
||||
|
||||
@@ -15,7 +15,6 @@ from PIL import Image
|
||||
from remove_ai_watermarks._internal.tiling import (
|
||||
Tile,
|
||||
_axis_positions,
|
||||
feather_region_composite,
|
||||
feather_weights,
|
||||
plan_tiles,
|
||||
run_tiled,
|
||||
@@ -139,72 +138,3 @@ class TestRunTiled:
|
||||
image = Image.new("RGB", (1500, 1100), (200, 100, 50))
|
||||
out = run_tiled(generate, image, tile_size=1024, overlap=128)
|
||||
assert out.size == (1500, 1100)
|
||||
|
||||
|
||||
class TestFeatherRegionComposite:
|
||||
"""Region-targeted compositing for AI-enhanced composites: only the AI box is
|
||||
regenerated, the real photo outside it stays pixel-exact (roadmap P1#8)."""
|
||||
|
||||
@staticmethod
|
||||
def _frames(h=200, w=300):
|
||||
base = np.full((h, w, 3), 80, np.uint8)
|
||||
regenerated = np.full((h, w, 3), 200, np.uint8)
|
||||
return base, regenerated
|
||||
|
||||
def test_outside_box_is_pixel_exact(self):
|
||||
base, regen = self._frames()
|
||||
out = feather_region_composite(base, regen, (100, 60, 80, 50), feather=8)
|
||||
# Far corners are well outside the box -> identical to base.
|
||||
assert np.array_equal(out[:50, :80], base[:50, :80])
|
||||
assert np.array_equal(out[150:, 220:], base[150:, 220:])
|
||||
|
||||
def test_interior_equals_regenerated(self):
|
||||
base, regen = self._frames()
|
||||
out = feather_region_composite(base, regen, (100, 60, 80, 50), feather=8)
|
||||
# Deep interior of the box (past the feather ramp) is fully regenerated.
|
||||
assert np.array_equal(out[80:90, 130:150], regen[80:90, 130:150])
|
||||
|
||||
def test_hard_paste_when_no_feather(self):
|
||||
base, regen = self._frames()
|
||||
out = feather_region_composite(base, regen, (100, 60, 80, 50), feather=0)
|
||||
assert np.array_equal(out[60:110, 100:180], regen[60:110, 100:180])
|
||||
assert np.array_equal(out[:60], base[:60])
|
||||
|
||||
def test_seam_is_monotonic_ramp(self):
|
||||
base, regen = self._frames()
|
||||
out = feather_region_composite(base, regen, (100, 60, 80, 50), feather=10).astype(np.float32)
|
||||
# Along a horizontal line crossing the left edge, values rise from base(80)
|
||||
# toward regenerated(200) monotonically through the feather band.
|
||||
row = out[85, 100:115, 0]
|
||||
assert row[0] < row[-1]
|
||||
assert np.all(np.diff(row) >= -1e-3)
|
||||
|
||||
def test_dtype_preserved(self):
|
||||
base, regen = self._frames()
|
||||
out = feather_region_composite(base, regen, (50, 50, 40, 40), feather=4)
|
||||
assert out.dtype == base.dtype
|
||||
|
||||
def test_grayscale_2d_supported(self):
|
||||
base = np.full((100, 120), 30, np.uint8)
|
||||
regen = np.full((100, 120), 220, np.uint8)
|
||||
out = feather_region_composite(base, regen, (40, 30, 30, 30), feather=4)
|
||||
assert out.shape == base.shape
|
||||
assert np.array_equal(out[:30], base[:30])
|
||||
|
||||
def test_empty_or_offimage_box_returns_base(self):
|
||||
base, regen = self._frames()
|
||||
assert np.array_equal(feather_region_composite(base, regen, (0, 0, 0, 0)), base)
|
||||
assert np.array_equal(feather_region_composite(base, regen, (500, 500, 40, 40)), base)
|
||||
|
||||
def test_box_clamped_to_image_bounds(self):
|
||||
base, regen = self._frames()
|
||||
# Box overhangs the bottom-right; only the in-image part is composited.
|
||||
out = feather_region_composite(base, regen, (280, 180, 60, 60), feather=0)
|
||||
assert np.array_equal(out[180:, 280:], regen[180:, 280:])
|
||||
assert out.shape == base.shape
|
||||
|
||||
def test_shape_mismatch_raises(self):
|
||||
base, _ = self._frames(200, 300)
|
||||
bad = np.full((100, 100, 3), 200, np.uint8)
|
||||
with pytest.raises(ValueError, match="shape mismatch"):
|
||||
feather_region_composite(base, bad, (10, 10, 20, 20))
|
||||
|
||||
Reference in New Issue
Block a user