From 2f7225799654309959e232f52a10854c00e2dfa9 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Mon, 3 Aug 2026 16:31:39 -0700 Subject: [PATCH] Fix what the verification pass found in the knob removal An adversarial review of 52b2c11 (five independent audits, each finding put to two skeptics, plus a completeness critic) found four defects that commit introduced and several stale claims it should have caught. The install hint no longer installs -- again. Folding five hints into one INVISIBLE_EXTRA constant dropped the shell quoting the originals had, so the printed remediation was `pip install remove-ai-watermarks[qwen-zimage]`. Bare brackets are a glob in zsh, the macOS default shell: it dies with "no matches found" before pip runs. That is the exact failure 52b2c11 existed to stop producing, reintroduced in a different form by a bulk replace. The constant is quoted now, and a test asserts the quotes rather than the bare substring -- the old assertions passed either way, which is why nothing caught it. Three tests were not guarding what they claimed: - The commit's headline behaviour change, per-profile polish resolution inside the engine, had no test at all. Rebinding resolve_adaptive_polish to the pre-commit `bool(value)` left the full suite green. Now covered by a test that drives the real engine and observes whether humanizer.adaptive_polish ran; that mutation now fails it. - TestAvailability still asserted the pre-commit (torch, diffusers) contract, so in a diffusion-only environment it was simply wrong, and comparing each gate to a tuple copied from itself could never catch the two gates disagreeing -- the drift the shared REMOVAL_MODULES was introduced to prevent. Replaced with a test that simulates each module's absence and requires BOTH gates to close. - Both CUDA-refusal guards skipped in every environment, including CI: they were gated on the diffusion stack, which no CI job installs. The refusal fires before any torch attribute is read, so they now run everywhere; only the dtype assertion keeps its skip. Also: the retired-knob test covered `invisible` but not `all` or `batch`, though all three declared those options separately; and smoke_matrix.py still called remove_watermark(region=...), a parameter 52b2c11 deleted, with the resulting TypeError swallowed into a skip by a broad except. Stale documentation the previous sweep missed: known-limitations still described an MPS out-of-memory fallback and a lighter-pipeline escape that no code can produce; module-internals declared Canny thresholds of 100/200 as a compatibility contract while the code uses 13/64, attributed enable_model_cpu_offload to deleted profiles, and still warned that the engine and CLI defaults differ (this commit's predecessor made them identical); cli.md gated `all` on the `diffusion` extra; python-api claimed "cuda" was the only accepted explicit device when "auto" is too. The claim that `device` is not a parameter was wrong in both module-internals and .claude/rules/development.md -- it is one, deliberately, and now says so. `--cpu-offload` help and the pipeline's CUDA guard both still pointed at MPS. Not fixed here, reported instead -- both are outside this repo: - ComfyUI-remove-ai-watermarks nodes.py:332 passes num_inference_steps and guidance_scale (plus min_resolution/upscaler from bf4bfc1). distribute.yml's comfyui job runs on every release and fails the release if the node sync fails, so 0.25.0 needs that node updated first. - raiw-app modal_app.py:422-425 forwards the same two kwargs into remove_watermark. Latent: it is pinned to 1a77e24 and nothing supplies a value today, so it fires on the next pin bump. pre-commit: 1) maintain.sh - exit 0 (1093 tests, Pyright 0 errors, no vulnerabilities); 2) /simplify - not re-run, this commit is the applied output of a five-dimension adversarial review; 3) docs sync - grepped MPS/mps, the extras names and every symbol touched across README, docs/, scripts/, .claude/; updated 6 docs; 4) CLAUDE.md - corrected the device claim in .claude/rules/development.md and added the shell-quoting rule Verified by execution, not assertion: smoke_matrix --quick 51 pass / 0 fail, _knob_rows driven directly 10 pass / 0 fail / 7 skip (no CUDA), the install hint rendered and round-tripped through zsh, and each new test confirmed to fail under the mutation it is meant to catch. Co-Authored-By: Claude Opus 5 --- .claude/rules/development.md | 6 +- docs/cli.md | 2 +- docs/known-limitations.md | 26 +++--- docs/module-internals.md | 53 +++++++----- docs/python-api.md | 6 +- docs/verification-plan.md | 6 +- scripts/smoke_matrix.py | 46 ++-------- .../_internal/qwen_zimage_pipeline.py | 5 +- .../_internal/watermark_profiles.py | 6 +- src/remove_ai_watermarks/cli.py | 3 +- tests/test_cli.py | 37 ++++++-- tests/test_invisible_engine.py | 61 ++++++++++--- tests/test_platform.py | 86 +++++++++++++------ 13 files changed, 213 insertions(+), 130 deletions(-) diff --git a/.claude/rules/development.md b/.claude/rules/development.md index 6c128fd..094dbe4 100644 --- a/.claude/rules/development.md +++ b/.claude/rules/development.md @@ -11,7 +11,11 @@ Every single-image command declares `source` with `dir_okay=False`; `batch` decl Exit-code and no-signal behavior is a public contract. Read the command-line section of [`../../docs/module-internals.md`](../../docs/module-internals.md) before changing it. -Do not add an option whose only outcome is an error. Model id, step count, CFG and any non-CUDA device are fixed by the profile, so none of them is a parameter of the CLI, `InvisibleEngine`, or `WatermarkRemover` -- they were accepted-then-rejected for a while, which moved the failure several frames below the caller and advertised choices the pinned stack cannot honor. If a value cannot vary, delete the knob rather than validating it. The same rule applies to install hints: name the extra that actually makes the command work (`qwen-zimage`, not `diffusion`). +Do not add an option whose only outcome is an error. Model id, step count and CFG are fixed by the profile, so none of them is a parameter of the CLI, `InvisibleEngine`, or `WatermarkRemover` -- they were accepted-then-rejected for a while, which moved the failure several frames below the caller and advertised choices the pinned stack cannot honor. If a value cannot vary, delete the knob rather than validating it. + +`device` is the deliberate exception and stays a library parameter: `None`/`"auto"` detect, `"cuda"` pins without detecting, and everything else raises at construction. It is not a CLI option, because the only value a user could usefully type is the one auto-detection already finds. + +The same rule applies to install hints: name the extra that actually makes the command work (`qwen-zimage`, not `diffusion`), and keep the printed command shell-quoted -- bare `pkg[extra]` is a glob in zsh. ## Local gate diff --git a/docs/cli.md b/docs/cli.md index bcce973..fbbedf7 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -433,7 +433,7 @@ The command runs: The visible options and diffusion options are also available on `all`. -If diffusion is required but the `diffusion` extra is unavailable, `all` still +If invisible removal is required but the `qwen-zimage` extra is unavailable, `all` still writes the result of the visible and metadata stages, prints a prominent warning, and exits with code 1. This prevents a partial result from being reported as complete. diff --git a/docs/known-limitations.md b/docs/known-limitations.md index eb72104..0934d7c 100644 --- a/docs/known-limitations.md +++ b/docs/known-limitations.md @@ -179,23 +179,25 @@ overlapping tiles. It avoids the explicit downscale, but it is not pixel lossless: each tile is independently regenerated. With `qwen-zimage`, only the global Qwen stage is tiled; the face stage runs after tile blending. -### CPU offload is CUDA only +### CPU offload trades speed for VRAM -`--cpu-offload` moves Diffusers model components between CPU and CUDA instead of -keeping the complete standard pipeline in GPU memory. For `qwen-zimage`, it -forces the face stack to use its offload path. +`--cpu-offload` forces both stacks of the two-stage profile out of automatic +device residency, streaming weights instead of pinning them. It reduces CUDA +memory pressure at the cost of speed. -The option reduces CUDA memory pressure at the cost of speed. It has no effect -on CPU or MPS and fails loudly when a CUDA Diffusers pipeline does not expose -the required offload method. +Residency is otherwise chosen from the card's total VRAM. On a card large enough +to hold a stack, offloading is pure waste: DiffSynth drops weights to the meta +device and re-reads every parameter from disk on each stage transition. -### MPS may fall back to CPU +### There is no CPU or MPS fallback -The SDXL paths include an MPS out-of-memory fallback that reloads on CPU. A run -that appears much slower after an MPS failure may be continuing on CPU. +Invisible-watermark removal refuses any device but CUDA at construction. There is +no MPS out-of-memory fallback that continues on CPU, and no lighter profile to +drop to: both remaining profiles need an NVIDIA GPU. -Memory needs depend on the pipeline, input size, dtype, and machine. Use tiling, -a resolution cap, or a lighter pipeline when necessary. +When a run does not fit, the levers are `--tile` (native geometry, tiled +diffusion), `--max-resolution` (an explicit downscale), and `--cpu-offload`. +Memory needs depend on the profile, input size, dtype, and card. ## Metadata and formats diff --git a/docs/module-internals.md b/docs/module-internals.md index be0c721..be574f3 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -502,15 +502,21 @@ mandatory global stage and YuNet while leaving the optional Z-Image and SAM face stack lazy until a face is detected. The default `preload()` still loads every stage. -**What is deliberately not a parameter.** Model id, step count, CFG and any -non-CUDA device are fixed by the profile, so none of them appears in -`WatermarkRemover.__init__`, `remove_watermark`, `InvisibleEngine`, or the CLI. -They used to be accepted and then rejected several frames down; a signature that -refuses the argument outright fails where the caller can act on it, and stops a -wrapper from threading a value that would silently do nothing. The step count and -CFG live with the stage that runs them (`GLOBAL_STEPS`, `FACE_STEPS`, `GLOBAL_CFG`, -`FACE_CFG` in `qwen_zimage_pipeline.py`). The dtype is likewise profile-owned: see -"Face-stage dtype" for what an override cost the last time one existed. +**What is deliberately not a parameter.** Model id, step count and CFG are fixed +by the profile, so none of them appears in `WatermarkRemover.__init__`, +`remove_watermark`, `InvisibleEngine`, or the CLI. They used to be accepted and +then rejected several frames down; a signature that refuses the argument outright +fails where the caller can act on it, and stops a wrapper from threading a value +that would silently do nothing. The step count and CFG live with the stage that +runs them (`GLOBAL_STEPS`, `FACE_STEPS`, `GLOBAL_CFG`, `FACE_CFG` in +`qwen_zimage_pipeline.py`). The dtype is likewise profile-owned: see "Face-stage +dtype" for what an override cost the last time one existed. + +`device` is the exception and remains a library parameter: `None` or `"auto"` +detect, `"cuda"` pins without detecting (which is what a container that knows its +hardware wants), and any other value raises at construction. It is not a CLI +option, because the only useful value a user could type is the one detection +already returns. [`invisible_engine.py`](../src/remove_ai_watermarks/invisible_engine.py) handles image sizing, postprocessing, and the public engine @@ -523,16 +529,20 @@ and reporting it implied an Apple-silicon or Intel-GPU path that does not exist. The refusal names the *resolved* device, so `device=None` on a CUDA-less host says `'cpu'` rather than `'None'`. -The Python engine and CLI do not have identical defaults for every optional -postprocessing argument. Integrations that require reproducibility should pass -the relevant values explicitly. +The Python engine and the CLI now resolve the same defaults: the CLI forwards an +unset `--adaptive-polish` and `--seed` as `None` and the engine applies the +profile's answer, so a library caller and a CLI caller on one profile produce the +same pixels. They diverged before, in opposite directions, for exactly this knob. -The standard Qwen and ControlNet prompts are calibrated model inputs, and the -ControlNet edge map uses fixed Canny thresholds of 100 and 200. Treat those -values as behavioral compatibility contracts: a refactor must preserve them, -and any deliberate change requires image-quality evaluation rather than only a -unit-test pass. Exact prompt and edge-map regression guards live in -`test_platform.py` and `test_invisible_engine.py`. +The global and face prompts are calibrated model inputs, and the Canny edge map +uses fixed thresholds of `_CANNY_LOW = 13` / `_CANNY_HIGH = 64` +(`qwen_zimage_pipeline.py`). Treat those values as behavioral compatibility +contracts: a refactor must preserve them, and any deliberate change requires +image-quality evaluation rather than only a unit-test pass. The prompt and +edge-map regression guards are +`test_qwen_zimage_pipeline.py::test_global_kwargs_use_lightning_and_diffsynth_controlnet_shape`, +`::test_face_kwargs_use_project_zimage_settings` and +`::test_canny_control_image_is_three_channel_and_detects_an_edge`. Regression coverage: @@ -542,9 +552,10 @@ Regression coverage: ### CPU offload -CPU offload is enabled only when requested on CUDA. The standard Diffusers -profiles call `enable_model_cpu_offload`. The `qwen-zimage` profile uses the -same flag to force **both** its stacks out of automatic device residency. +CPU offload is enabled only when requested. Nothing calls Diffusers' +`enable_model_cpu_offload` any more -- that belonged to the deleted single-stage +profiles. `--cpu-offload` now forces **both** stacks of the two-stage profiles out +of automatic device residency. Residency is otherwise chosen from the card's total VRAM, once per stack: `resolve_global_model_residency` gates the mandatory Qwen stack at diff --git a/docs/python-api.md b/docs/python-api.md index 7458538..8977fc2 100644 --- a/docs/python-api.md +++ b/docs/python-api.md @@ -401,9 +401,9 @@ engine.remove_watermark( ) ``` -`device=None` detects CUDA. The only other accepted value is `"cuda"`; anything -else raises at construction rather than deferring a guaranteed failure to model -load time. +`device=None` and `device="auto"` both run detection. `"cuda"` pins it without +detecting. Every other value raises at construction rather than deferring a +guaranteed failure to model-load time. For limited CUDA memory: diff --git a/docs/verification-plan.md b/docs/verification-plan.md index 2fcf6cc..99e19b5 100644 --- a/docs/verification-plan.md +++ b/docs/verification-plan.md @@ -86,8 +86,10 @@ sweep asserts on stderr, or the codes get split -- the latter is the better fix. #### Coverage method Compare the flags and values exercised by the matrix against the options declared by -the CLI. Include optional backends, batch modes, tiling, region-targeted composition, -and the ffmpeg audio/video strip. The gap to find is not only +the CLI. Include optional backends, batch modes, tiling, and the ffmpeg audio/video +strip. (Region-targeted diffusion composition was on this list until the parameter +was deleted: nothing but a caller-less convenience wrapper ever reached it, and the +user-facing region path is `erase`, which inpaints instead.) The gap to find is not only "logic untested" but "never executed on real data", which is precisely what this campaign is for. diff --git a/scripts/smoke_matrix.py b/scripts/smoke_matrix.py index 26809b8..7326b95 100644 --- a/scripts/smoke_matrix.py +++ b/scripts/smoke_matrix.py @@ -304,7 +304,7 @@ def main() -> None: _diffusion_rows(r, tmp, doubao) else: for name in ("invisible", "all", "batch --mode invisible"): - r.skip(f"{name} (model-running body)", "pass --diffusion to exercise it (needs the SDXL weights)") + r.skip(f"{name} (model-running body)", "pass --diffusion to exercise it (needs CUDA and the model weights)") # ---- real-data formats and shapes --------------------------------------- if not a.quick: @@ -486,7 +486,7 @@ def _sdxl_watermark_bits(img: object) -> float: def _diffusion_rows(r: Runner, tmp: Path, doubao: Path) -> None: - """Exercise the model-running bodies at a reduced resolution (MPS-friendly). + """Exercise the model-running bodies at a reduced resolution (CUDA required). Bounded with `--max-resolution 512` and a fixed seed: the point is that the paths RUN and keep their contracts, not to certify removal strength (that needs the @@ -508,7 +508,7 @@ def _diffusion_rows(r: Runner, tmp: Path, doubao: Path) -> None: # OUTPUT and the row degrades to a skip rather than a false pass if it goes fragile. mj = SAMPLES / "mj-1.png" inv = tmp / "inv_mj.png" - res = r.run("invisible runs (mps, 512px)", ["invisible", str(mj), "-o", str(inv), "--force", *small], timeout=1800) + res = r.run("invisible runs (512px)", ["invisible", str(mj), "-o", str(inv), "--force", *small], timeout=1800) if res.status == "pass" and inv.exists(): src_img, out_img = imread(str(mj)), imread(str(inv)) r.check( @@ -536,7 +536,7 @@ def _diffusion_rows(r: Runner, tmp: Path, doubao: Path) -> None: # `all`: every stage must land -- the visible mark AND the metadata both gone. allout = tmp / "all_out.png" - res = r.run("all runs (mps, 512px)", ["all", str(doubao), "-o", str(allout), *small], timeout=1800) + res = r.run("all runs (512px)", ["all", str(doubao), "-o", str(allout), *small], timeout=1800) if res.status == "pass" and allout.exists(): rep = json.loads(_capture(["identify", str(allout), "--json"])) r.check( @@ -555,7 +555,7 @@ def _diffusion_rows(r: Runner, tmp: Path, doubao: Path) -> None: shutil.copy(f, bd / f.name) bout = tmp / "batch_inv_out" res = r.run( - "batch --mode invisible runs (mps, 512px)", + "batch --mode invisible runs (512px)", ["batch", str(bd), "--mode", "invisible", "-o", str(bout), *small], expect_exit=None, timeout=3600, @@ -567,46 +567,12 @@ def _diffusion_rows(r: Runner, tmp: Path, doubao: Path) -> None: # `batch --mode all` -- the only --mode value the matrix never ran. aout = tmp / "batch_all_out" r.run( - "batch --mode all runs (mps, 512px)", + "batch --mode all runs (512px)", ["batch", str(bd), "--mode", "all", "-o", str(aout), *small], expect_exit=None, timeout=3600, ) - # The AI-enhanced composite path: regenerate ONLY a region and feather it back, - # leaving everything outside the box pixel-exact. Library-level -- the CLI has no - # flag for it, so it would otherwise never be exercised on real data. - try: - import numpy as np - - from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover - from remove_ai_watermarks.image_io import imread - - src = imread(str(mj)) - h, w = src.shape[:2] - box = (w // 4, h // 4, w // 4, h // 4) - # Default profile, default four-step schedule: the remover rejects any other - # step count now, and the retired controlnet profile no longer exists. - rem = WatermarkRemover() - rout = tmp / "region_composite.png" - rem.remove_watermark(mj, rout, strength=0.15, seed=0, region=box) - got = imread(str(rout)) - if got is None or got.shape != src.shape: - r.check("region composite keeps the frame outside the box", False, "shape changed or unreadable") - else: - mask = np.ones(src.shape[:2], dtype=bool) - x, y, bw, bh = box - # Outside the box PLUS the feather margin must be untouched. - pad = 96 - mask[max(0, y - pad) : y + bh + pad, max(0, x - pad) : x + bw + pad] = False - r.check( - "region composite keeps the frame outside the box", - bool(np.array_equal(src[mask], got[mask])), - "pixels changed outside the regenerated region", - ) - except Exception as e: - r.skip("region composite (remove_watermark(region=...))", f"{type(e).__name__}: {e}"[:120]) - if __name__ == "__main__": main() diff --git a/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py b/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py index 90fd41e..bbc0657 100644 --- a/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py +++ b/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py @@ -636,8 +636,9 @@ class QwenZImagePipeline: def _require_cuda(self) -> None: if self.device != "cuda": raise RuntimeError( - "The qwen-zimage pipeline is CUDA-only. Its Qwen-Image-2512 and " - "Z-Image models do not fit the supported MPS path." + "The qwen-zimage pipeline is CUDA-only; there is no CPU or MPS path " + "for it. WatermarkRemover already refuses a non-CUDA device, so " + "reaching this guard means the pipeline was constructed directly." ) def _vram_limit(self) -> float | None: diff --git a/src/remove_ai_watermarks/_internal/watermark_profiles.py b/src/remove_ai_watermarks/_internal/watermark_profiles.py index dee61a9..57e27be 100644 --- a/src/remove_ai_watermarks/_internal/watermark_profiles.py +++ b/src/remove_ai_watermarks/_internal/watermark_profiles.py @@ -34,7 +34,11 @@ PROFILE_CHOICES = (QWEN_ZIMAGE_PROFILE, SDXL_ZIMAGE_PROFILE) # the CLI passed on a torch+diffusers environment and the run then died at the # DiffSynth face stage, telling the user to install an extra that does not contain it. REMOVAL_MODULES = ("torch", "diffusers", "diffsynth") -INVISIBLE_EXTRA = "remove-ai-watermarks[qwen-zimage]" +# Shell-quoted, because this string is printed as a command the user copy-pastes. +# Bare brackets are a glob in zsh (the macOS default shell): an unquoted +# ``pip install remove-ai-watermarks[qwen-zimage]`` dies with "no matches found" +# before pip ever runs -- another install hint that does not install. +INVISIBLE_EXTRA = "'remove-ai-watermarks[qwen-zimage]'" # qwen-zimage's output already matches the input's detail level, so polishing it is a # no-op at best. sdxl-zimage's global pass leaves the softer output the polish exists diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index d2da61a..9afd9c8 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -286,8 +286,7 @@ _cpu_offload_option = click.option( help=( "Offload model components to CPU between CUDA calls instead of keeping the " "whole pipeline in VRAM, at the cost of speed. For qwen-zimage, forces the " - "face stack to offload instead of using automatic residency. No effect on " - "cpu/mps." + "face stack to offload instead of using automatic residency." ), ) diff --git a/tests/test_cli.py b/tests/test_cli.py index afb1e96..92f0610 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -363,16 +363,21 @@ class TestInvisibleCommand: 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 ( + retired = ( ["--model", "org/custom-sdxl"], ["--steps", "20"], ["--guidance-scale", "5.5"], ["--device", "cpu"], ["--auto"], - ): - 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}" + ) + # All THREE diffusion commands, not just `invisible`. Each used to declare these + # inline, so removing them from one and not the others is a live possibility. + for command in ("invisible", "all", "batch"): + target = str(sample_png.parent) if command == "batch" else str(sample_png) + for args in retired: + result = runner.invoke(main, [command, target, *args, "--force"]) + assert result.exit_code == 2, f"{command} {args[0]}: {result.output}" + assert "No such option" in result.output, f"{command} {args[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. @@ -862,11 +867,12 @@ class TestBatchCommand: assert out[0, 0, 3] == 0 assert out[100, 100, 3] == 255 - 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. + def test_batch_forwards_an_explicit_adaptive_polish(self, runner, tmp_path): + """A typed --adaptive-polish must reach the engine as True, not as None. - The off is a parameter-source check, not a changed default, so it must yield to - an explicit --adaptive-polish rather than swallowing it. + This mocks the engine, so it covers the CLI's forwarding only; the per-profile + resolution of an UNSET flag happens inside the real engine and is guarded by + test_invisible_engine.py::TestEngineResolvesThePolishPerProfile. """ input_dir = _make_batch_dir(tmp_path, count=2) output_dir = tmp_path / "output" @@ -943,6 +949,19 @@ class TestGpuHintMarkup: profiles run -- so following the advice produced a second, different failure. """ + def test_the_hint_is_a_command_the_user_can_actually_paste(self, runner, sample_png): + """The extras bracket must be shell-quoted. + + A bulk replace once folded these hints into one constant and dropped the quotes + the originals had. `pip install remove-ai-watermarks[qwen-zimage]` dies with + "zsh: no matches found" on the macOS default shell before pip ever runs -- an + install hint that does not install, which is the failure this hint was fixed to + stop producing in the first place. + """ + with patch("remove_ai_watermarks.invisible_engine.is_available", return_value=False): + result = runner.invoke(main, ["invisible", str(sample_png)]) + assert "pip install 'remove-ai-watermarks[qwen-zimage]'" in result.output + 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)]) diff --git a/tests/test_invisible_engine.py b/tests/test_invisible_engine.py index f78272e..8844c43 100644 --- a/tests/test_invisible_engine.py +++ b/tests/test_invisible_engine.py @@ -16,21 +16,19 @@ class TestIsAvailable: result = is_available() assert isinstance(result, bool) - def test_available_reflects_every_module_a_run_needs(self): - """True iff every module in REMOVAL_MODULES imports, diffsynth included. + def test_the_module_list_includes_the_face_stage_runtime(self): + """diffsynth is part of the answer, not an optional upgrade. - 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. + Both profiles repair faces with the DiffSynth Z-Image stage, so a + torch+diffusers-only environment used to clear this gate and then die there. + The discriminating guard -- that this gate and the remover's precondition + BOTH close when any one module is missing -- lives in + ``test_platform.py::TestAvailability``; comparing ``is_available()`` to a + tuple derived from itself passes on every host and proves nothing. """ - import importlib.util - 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: @@ -144,3 +142,46 @@ class TestEngineConstructsWithoutAModelId: with pytest.raises(TypeError): engine_module.InvisibleEngine(model_id="org/custom", pipeline="qwen-zimage") # type: ignore[call-arg] + + +class TestEngineResolvesThePolishPerProfile: + """The engine, not the CLI, turns an unset adaptive_polish into the profile default. + + This is the change that stopped a library caller and a CLI caller on one profile + from producing different pixels, and it had no test: rebinding + ``resolve_adaptive_polish`` to ``bool(value)`` -- exactly the pre-commit behaviour -- + left the whole suite green. + """ + + @staticmethod + def _engine(profile: str): + from unittest.mock import MagicMock + + engine = object.__new__(InvisibleEngine) + engine._progress_callback = None + engine._remover = MagicMock(model_profile=profile) + return engine + + def _polish_used(self, profile: str, requested, tmp_path, monkeypatch) -> bool: + seen: list[bool] = [] + monkeypatch.setattr( + "remove_ai_watermarks.humanizer.adaptive_polish", + lambda out, ref, seed=None: (seen.append(True), out)[1], + ) + src = tmp_path / f"{profile}_{requested}.png" + Image.new("RGB", (32, 32), (90, 120, 150)).save(src) + engine = self._engine(profile) + engine._remover.remove_watermark.side_effect = lambda **kw: ( + Image.open(kw["image_path"]).save(kw["output_path"]), + kw["output_path"], + )[1] + engine.remove_watermark(src, tmp_path / f"out_{profile}_{requested}.png", adaptive_polish=requested) + return bool(seen) + + def test_unset_follows_the_profile_not_the_signature_default(self, tmp_path, monkeypatch): + assert self._polish_used("qwen-zimage", None, tmp_path, monkeypatch) is False + assert self._polish_used("sdxl-zimage", None, tmp_path, monkeypatch) is True + + def test_an_explicit_value_still_wins_on_both_profiles(self, tmp_path, monkeypatch): + assert self._polish_used("qwen-zimage", True, tmp_path, monkeypatch) is True + assert self._polish_used("sdxl-zimage", False, tmp_path, monkeypatch) is False diff --git a/tests/test_platform.py b/tests/test_platform.py index 2d99765..e4a7159 100644 --- a/tests/test_platform.py +++ b/tests/test_platform.py @@ -14,6 +14,7 @@ import pytest from remove_ai_watermarks._internal.utils import get_image_format, is_supported_format from remove_ai_watermarks._internal.watermark_profiles import ( PROFILE_CHOICES, + REMOVAL_MODULES, SDXL_ZIMAGE_GEMINI_STRENGTH, SDXL_ZIMAGE_OPENAI_STRENGTH, SDXL_ZIMAGE_UNKNOWN_STRENGTH, @@ -67,20 +68,19 @@ class TestDeviceDetection: 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. + + Deliberately NOT gated on the diffusion stack. It used to be, and since no CI + job installs diffusers or diffsynth (the dev extra pulls torch only, via + invisible-watermark) the guard skipped in every environment it ran in -- + including the maintainer's. The refusal fires before any torch attribute is + touched, so faking the dependency probe is enough to reach it. """ - if not is_watermark_removal_available(): - pytest.skip("torch/diffusers not installed") - import torch + from remove_ai_watermarks._internal import watermark_remover as module - from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover - - for device in ("cpu", "mps", "xpu"): - with pytest.raises(ValueError, match="CUDA-only"): - WatermarkRemover(device=device) - - remover = WatermarkRemover(device="cuda") - assert remover.device == "cuda" - assert remover.torch_dtype == torch.bfloat16 + with patch.object(module, "is_watermark_removal_available", return_value=True): + for device in ("cpu", "mps", "xpu"): + with pytest.raises(ValueError, match="CUDA-only"): + module.WatermarkRemover(device=device) 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". @@ -88,16 +88,28 @@ class TestDeviceDetection: 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 with ( + patch.object(module, "is_watermark_removal_available", return_value=True), patch.object(module, "get_device", return_value="cpu"), pytest.raises(ValueError, match="'cpu' cannot run it"), ): module.WatermarkRemover(device=None) + def test_a_cuda_remover_picks_the_profile_dtype(self): + """The dtype half still needs real torch, so it keeps its skip.""" + if not is_watermark_removal_available(): + pytest.skip("the qwen-zimage extra is not installed") + import torch + + from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover + + remover = WatermarkRemover(device="cuda") + assert remover.device == "cuda" + assert remover.torch_dtype == torch.bfloat16 + assert WatermarkRemover(device="cuda", pipeline="sdxl-zimage").torch_dtype == torch.float16 + class TestModelProfiles: """Only the two CUDA-only two-stage profiles remain.""" @@ -155,7 +167,7 @@ class TestNoReembeddedWatermark: def test_sdxl_global_stage_disables_watermarker(self, monkeypatch: pytest.MonkeyPatch): if not is_watermark_removal_available(): - pytest.skip("torch/diffusers not installed") + pytest.skip("the qwen-zimage extra is not installed") import diffusers from remove_ai_watermarks._internal.sdxl_zimage_pipeline import SdxlZImagePipeline @@ -315,25 +327,47 @@ class TestFormatUtils: class TestAvailability: - """Tests for dependency availability checks.""" + """The CLI gate and the remover precondition must answer from the same module list. - def test_watermark_removal_available(self): - # Reflects the actual environment: True iff torch + diffusers (the gpu - # extra) are importable. The default+dev CI env has no diffusers, so this - # must not assume the full stack is present. - import importlib.util + Both used to hardcode (torch, diffusers) while the code moved to REMOVAL_MODULES, + which includes diffsynth. In a torch+diffusers-only environment the assertions were + then simply wrong -- and, worse, comparing each gate against a tuple copied from + itself can never catch the two disagreeing, which is the drift that let the CLI pass + an environment the run then died in. + """ - expected = all(importlib.util.find_spec(m) is not None for m in ("torch", "diffusers")) - assert is_watermark_removal_available() is expected - - def test_invisible_is_available(self): + def test_both_gates_agree_and_read_the_shared_module_list(self): import importlib.util from remove_ai_watermarks.invisible_engine import is_available - expected = all(importlib.util.find_spec(m) is not None for m in ("torch", "diffusers")) + assert "diffsynth" in REMOVAL_MODULES + expected = all(importlib.util.find_spec(m) is not None for m in REMOVAL_MODULES) + assert is_watermark_removal_available() is expected assert is_available() is expected + def test_a_missing_module_closes_both_gates(self, monkeypatch: pytest.MonkeyPatch): + """Discriminating, not vacuous: it must FAIL if either gate stops requiring one. + + Comparing the live answer to a tuple derived from the same constant passes on + any host -- with the full stack (True == True) and with none of it + (False == False). Simulate each module's absence instead. + """ + import remove_ai_watermarks.invisible_engine as engine_module + from remove_ai_watermarks._internal import watermark_remover as remover_module + + for missing in REMOVAL_MODULES: + present = {name: name != missing for name in REMOVAL_MODULES} + monkeypatch.setattr( + "remove_ai_watermarks.optional_deps.module_available", + lambda *names, _p=present: all(_p.get(n, True) for n in names), + ) + # The remover probes at import time, so drive its cached flags directly. + monkeypatch.setattr(remover_module, "_HAS_TORCH", missing != "torch") + monkeypatch.setattr(remover_module, "_HAS_REMOVAL_MODULES", missing == "torch") + assert engine_module.is_available() is False, f"engine gate ignores a missing {missing}" + assert remover_module.is_watermark_removal_available() is False, f"remover gate ignores a missing {missing}" + # ── Platform-specific path handling ─────────────────────────────────