diff --git a/README.md b/README.md index 0a86676..72add75 100644 --- a/README.md +++ b/README.md @@ -56,7 +56,7 @@ removal. | Every production feature | `remove-ai-watermarks[all]` | Lower-level and specialized extras include `pixels`, `heif`, `trustmark`, -`migan`, `lama`, `esrgan`, and `qwen-zimage`. The +`migan`, `lama`, and `qwen-zimage`. The [installation guide](docs/installation.md#feature-extras) documents their exact dependency composition and model requirements. diff --git a/docs/installation.md b/docs/installation.md index 333e588..1c58432 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -98,7 +98,6 @@ application actually uses: | `diffusion` | Diffusion-based invisible watermark removal | `pixels`, Torch, Diffusers | Yes | | `migan` | MI-GAN ONNX fill backend | `visible`, ONNX Runtime | Model download, no Torch | | `lama` | big-LaMa ONNX fill backend | `visible`, ONNX Runtime | Model download, no Torch | -| `esrgan` | Real-ESRGAN upscaling before diffusion | `pixels`, spandrel | Yes | | `qwen-zimage` | CUDA-only Qwen Image plus Z-Image pipeline | `diffusion`, DiffSynth | Yes | | `all` | Every production feature | All rows above | Yes | | `dev` | Tests, linting, typing, and upstream parity checks | `visible`, `detect`, upstream invisible-watermark | Yes, for parity tests | @@ -113,7 +112,6 @@ flowchart LR diffusion --> pixels migan --> visible lama --> visible - esrgan --> pixels qwen["qwen-zimage"] --> diffusion heif trustmark diff --git a/docs/known-limitations.md b/docs/known-limitations.md index b06d4a6..a312c9d 100644 --- a/docs/known-limitations.md +++ b/docs/known-limitations.md @@ -150,13 +150,14 @@ deprecated, emits a warning, and changes nothing. ## Resolution and memory -### Small images are enlarged before SDXL based diffusion +### Small images are processed at their native size -The SDXL, ControlNet, and base Qwen paths use a default minimum long side of -`1024`. Smaller inputs are enlarged before diffusion and restored to their -original dimensions afterward. Set `--min-resolution 0` to disable the floor. +There is no minimum-resolution floor. It existed to enlarge small inputs toward +SDXL's ~1024 training resolution and was removed with the SDXL profiles, which +never applied it anyway. Both surviving profiles run at native geometry, so a +small input is neither enlarged before diffusion nor restored afterward. -`qwen-zimage` does not apply this SDXL minimum resolution floor. +`--max-resolution` still caps very large inputs, and only ever scales down. ### Large images stay at native resolution unless capped diff --git a/docs/module-internals.md b/docs/module-internals.md index ff2acce..ad5b53f 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -492,7 +492,7 @@ stage. router. [`invisible_engine.py`](../src/remove_ai_watermarks/invisible_engine.py) handles -image sizing, optional pre-upscaling, postprocessing, and the public engine +image sizing, postprocessing, and the public engine interface. It delegates model execution to [`_internal/watermark_remover.py`](../src/remove_ai_watermarks/_internal/watermark_remover.py). @@ -725,18 +725,20 @@ Regression coverage: - [`test_tiling.py`](../tests/test_tiling.py) -### Upscaling and postprocessing - -[`upscaler.py`](../src/remove_ai_watermarks/upscaler.py) is the optional -Real-ESRGAN path used only when enlarging a small image to the minimum -resolution floor. Failure or an absent extra falls back to Lanczos. +### Postprocessing [`humanizer.py`](../src/remove_ai_watermarks/humanizer.py) contains explicit grain, unsharp masking, and adaptive polish helpers. +`upscaler.py` held an optional Real-ESRGAN path, reachable only when enlarging a +small image to the minimum-resolution floor. That floor existed to lift small +inputs toward SDXL's ~1024 training size; when the SDXL profiles were removed it +was forced to 0 on every path, so the module, the `--min-resolution` and +`--upscaler` options and the `esrgan` extra were all unreachable and went with +it. Only the `max_resolution` cap can move geometry now, and it only scales down. + Regression coverage: -- [`test_upscaler.py`](../tests/test_upscaler.py) - [`test_humanizer.py`](../tests/test_humanizer.py) ## Image input and output diff --git a/docs/python-api.md b/docs/python-api.md index 7d94872..526b226 100644 --- a/docs/python-api.md +++ b/docs/python-api.md @@ -422,7 +422,7 @@ engine = InvisibleEngine(pipeline="sdxl-zimage") The `qwen-zimage` extra must be installed for that profile. The full `remove_watermark` signature includes strength, steps, guidance, -seeding, tiling, resolution, upscaling, and postprocessing controls. Read the +seeding, tiling, resolution, and postprocessing controls. Read the method signature in [`invisible_engine.py`](../src/remove_ai_watermarks/invisible_engine.py) or use the CLI guide for the concepts. diff --git a/docs/verification-plan.md b/docs/verification-plan.md index a850587..ac14e0d 100644 --- a/docs/verification-plan.md +++ b/docs/verification-plan.md @@ -72,8 +72,10 @@ Run detection, removal, and re-detection over a representative local set. Confir ### A5. Contract sweep across every parameter choice -`scripts/smoke_matrix.py` (exists, 68 rows, 0 skipped with `--diffusion`) covers every -choice-valued flag on fixtures. Extend from fixtures to a stratified corpus slice +`scripts/smoke_matrix.py` covers every choice-valued flag on fixtures. Its knob rows +were rewritten when the CPU/MPS profiles and the ESRGAN chain were removed: most now +assert a knob is REJECTED, and the accepted-knob rows skip without a CUDA device, so +the row count is host-dependent rather than the fixed 68 recorded here before. Extend from fixtures to a stratified corpus slice (~500 images spanning format x provenance x aspect ratio), asserting exit-code semantics rather than just absence of crash. @@ -85,7 +87,7 @@ sweep asserts on stderr, or the codes get split -- the latter is the better fix. 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, -the ESRGAN upscaler, and the ffmpeg audio/video strip. The gap to find is not only +and the ffmpeg audio/video strip. 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/pyproject.toml b/pyproject.toml index 72ffde1..8fe1a21 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -150,19 +150,6 @@ migan = [ "onnxruntime>=1.16.0; python_version >= '3.11'", "huggingface-hub>=0.20.0", ] -# Optional pre-diffusion super-resolution for small inputs (Real-ESRGAN). Loaded via -# spandrel (MIT) -- a pure model-loader with NO basicsr dependency (it pulls only -# torch / torchvision / safetensors / numpy / einops). -# The Real-ESRGAN weights (BSD-3-Clause) download on first use and are cached; they -# are never bundled. CPU works but is slow on large inputs -- it is meant for the -# pre-diffusion upscale of SMALL inputs (and the GPU worker). Guarded by -# upscaler.is_available(); the default upscaler stays Lanczos (cv2, no deps). The -# weights are fetched with torch.hub (bundled with spandrel's torch), so no extra -# download dependency is needed. -esrgan = [ - "remove-ai-watermarks[pixels]", - "spandrel>=0.3.0", -] dev = [ "remove-ai-watermarks[video]", "remove-ai-watermarks[detect]", @@ -178,7 +165,7 @@ dev = [ "uv-outdated>=0.1.0; python_version >= '3.12'", "uv-secure>=0.12.0; python_version >= '3.12'", ] -all = ["remove-ai-watermarks[video,heif,detect,trustmark,diffusion,qwen-zimage,lama,migan,esrgan]"] +all = ["remove-ai-watermarks[video,heif,detect,trustmark,diffusion,qwen-zimage,lama,migan]"] # PyTorch Intel-GPU (XPU) wheel index. ``explicit = true`` keeps it inert for # the default CPU/CUDA install: uv consults it only when a torch install diff --git a/scripts/smoke_matrix.py b/scripts/smoke_matrix.py index 3548268..1c2eb60 100644 --- a/scripts/smoke_matrix.py +++ b/scripts/smoke_matrix.py @@ -354,18 +354,16 @@ def main() -> None: def _knob_rows(r: Runner, tmp: Path, img: Path) -> None: """Every diffusion knob the matrix never touched. - Deliberately cheap (`--steps 4`, `--max-resolution 384`): these rows answer "is the - knob accepted and does the run complete", NOT "is the output good". Quality per knob - needs a per-knob oracle and most of them have none (`--humanize` has no oracle at - all), so claiming more here would be dishonest. + Deliberately cheap (`--max-resolution 384`): these rows answer "is the knob accepted + and does the run complete", NOT "is the output good". Quality per knob needs a + per-knob oracle and most of them have none (`--humanize` has no oracle at all), so + claiming more here would be dishonest. """ - from remove_ai_watermarks import upscaler - - # --steps 20 is the floor that WORKS, not an arbitrary choice: effective timesteps - # are int(steps * strength), so at the default strength 0.15 anything below - # --steps 7 rounds to ZERO and the pipeline dies inside torch. The first version of - # these rows used --steps 4 and every single one failed with the same reshape error. - fast = ["--max-resolution", "384", "--min-resolution", "0", "--steps", "20", "--force", "--seed", "0"] + # Both surviving profiles are CUDA-only and pin a distilled four-step schedule at + # CFG 1.0, so most rows here now assert a knob is REJECTED rather than accepted. + # That is the coverage worth having: a knob the CLI takes and the library refuses + # several layers down is exactly what this matrix exists to catch. + fast = ["--max-resolution", "384", "--force", "--seed", "0"] def run(name: str, extra: list[str], *, tag: str, expect: int | None = 0) -> None: r.run( @@ -375,61 +373,52 @@ def _knob_rows(r: Runner, tmp: Path, img: Path) -> None: timeout=2400, ) - run("--pipeline sdxl", ["--pipeline", "sdxl"], tag="sdxl") - run("--pipeline controlnet", ["--pipeline", "controlnet"], tag="cnet") - run("--strength", ["--strength", "0.2"], tag="strength") - run("--guidance-scale", ["--guidance-scale", "5.0"], tag="gs") - run("--controlnet-scale", ["--controlnet-scale", "0.5"], tag="cns") - run("--humanize", ["--humanize", "0.3"], tag="hum") - run("--unsharp", ["--unsharp", "0.5"], tag="uns") - run("--no-adaptive-polish", ["--no-adaptive-polish"], tag="nap") - run("--tile", ["--tile", "--tile-size", "256", "--tile-overlap", "64"], tag="tile") - run("--device mps", ["--device", "mps"], tag="mps") - run("--upscaler lanczos", ["--upscaler", "lanczos"], tag="lanczos") - run("--auto (deprecated no-op)", ["--auto"], tag="auto") + # Click rejects a retired profile at parse time (exit 2) rather than remapping it. + for retired in ("sdxl", "controlnet", "qwen", "default"): + run(f"--pipeline {retired} is rejected", ["--pipeline", retired], tag=f"retired_{retired}", expect=2) - if upscaler.is_available(): - run("--upscaler esrgan", ["--upscaler", "esrgan"], tag="esrgan") - else: - r.skip("--upscaler esrgan", "the `esrgan` extra is not installed") + # Fixed-graph knobs: accepted by Click, refused by the library (exit 1). + run("--steps 20 is rejected", ["--steps", "20"], tag="steps20", expect=1) + run("--guidance-scale 5.0 is rejected", ["--guidance-scale", "5.0"], tag="gs5", expect=1) + run("--model override is rejected", ["--model", "org/custom"], tag="model", expect=1) - # CPU is correctness-relevant (it is the documented MPS-OOM fallback) but slow, so - # it gets the smallest possible run rather than being skipped. - r.run( - "--device cpu", - [ - "invisible", - str(img), - "-o", - str(tmp / "k_cpu.png"), - "--device", - "cpu", - "--max-resolution", - "256", - "--min-resolution", - "0", - "--steps", - "20", - "--force", - "--seed", - "0", - ], - timeout=3600, - ) - - # qwen is CUDA-class by design (bf16 MMDiT, no MPS fallback). On this host the - # honest outcome is a CLEAN failure, not a crash -- assert it does not hang or - # dump a traceback at the user. try: import torch cuda = bool(torch.cuda.is_available()) except Exception: cuda = False - if cuda: - run("--pipeline qwen", ["--pipeline", "qwen"], tag="qwen") - else: - r.skip("--pipeline qwen", "CUDA-class pipeline; no CUDA device on this host") + + if not cuda: + # Without CUDA the honest outcome is a CLEAN refusal naming the reason, not a + # traceback from inside a half-built pipeline. + run("no CUDA fails cleanly", [], tag="nocuda", expect=1) + for name, extra, tag in ( + ("--device cpu fails cleanly", ["--device", "cpu"], "cpu"), + ("--device mps fails cleanly", ["--device", "mps"], "mps"), + ): + run(name, extra, tag=tag, expect=1) + for label in ( + "--strength", + "--controlnet-scale", + "--humanize", + "--unsharp", + "--no-adaptive-polish", + "--tile", + "--auto", + "--pipeline sdxl-zimage", + ): + r.skip(label, "accepted-knob rows need a CUDA device") + return + + run("--pipeline sdxl-zimage", ["--pipeline", "sdxl-zimage"], tag="sdxlz") + run("--strength", ["--strength", "0.2"], tag="strength") + run("--controlnet-scale", ["--controlnet-scale", "0.5"], tag="cns") + run("--humanize", ["--humanize", "0.3"], tag="hum") + run("--unsharp", ["--unsharp", "0.5"], tag="uns") + run("--no-adaptive-polish", ["--no-adaptive-polish"], tag="nap") + run("--tile", ["--tile", "--tile-size", "256", "--tile-overlap", "64"], tag="tile") + run("--auto (deprecated, requests polish only)", ["--auto"], tag="auto") # --model and --hf-token are deliberately not exercised: one would download a second # multi-GB checkpoint, the other needs a real credential. Skipped loudly, not passed. @@ -596,9 +585,11 @@ def _diffusion_rows(r: Runner, tmp: Path, doubao: Path) -> None: src = imread(str(mj)) h, w = src.shape[:2] box = (w // 4, h // 4, w // 4, h // 4) - rem = WatermarkRemover(pipeline="controlnet") + # 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, num_inference_steps=20, seed=0, region=box) + 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") diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index e55d627..d694efb 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -180,28 +180,10 @@ _controlnet_scale_option = click.option( "(structure/text preservation strength). Higher = closer to original structure.", ) -_min_resolution_option = click.option( - "--min-resolution", - type=int, - default=1024, - help="Upscale long side UP to this (px) before diffusion when the input is smaller, so SDXL runs " - "near 1024 (small inputs distort at native); output is restored to the input size. 0 = off. Default 1024.", -) - _unsharp_option = click.option( "--unsharp", type=float, default=0.0, help="Unsharp-mask sharpening strength (0 = off, typical: 0.3-0.8)." ) -_upscaler_option = click.option( - "--upscaler", - type=click.Choice(["lanczos", "esrgan"]), - default="lanczos", - help="How to upscale a small input to the --min-resolution floor: lanczos (default, cv2, no model) or " - "esrgan (Real-ESRGAN via the 'esrgan' extra; better detail, slower on CPU). Best for photo/texture " - "content -- as a generic GAN with no face/glyph prior it can degrade faces (diffusion mitigates) and " - "thin text, so lanczos stays the default. Falls back to lanczos if the extra is absent. Only when upscaling.", -) - _auto_option = click.option( "--auto", is_flag=True, @@ -376,21 +358,6 @@ def _resolve_profile_polish(auto: bool, adaptive_polish: bool, pipeline: str) -> return adaptive_polish -def _warn_if_esrgan_unavailable(upscaler: str) -> None: - """Tell the user once if ``--upscaler esrgan`` will silently fall back to Lanczos. - - The engine downgrades to Lanczos when the ``esrgan`` extra is absent (fail-safe, so - a batch never breaks mid-run) -- but without this notice the user would believe - Real-ESRGAN ran. Surfaced at the CLI layer, once per invocation (not per image). - """ - if upscaler != "esrgan": - return - from remove_ai_watermarks import upscaler as _upscaler - - if not _upscaler.is_available(): - console.print(" Note: --upscaler esrgan needs the 'esrgan' extra; falling back to Lanczos.") - - def _visible_provenance(path: Path | None) -> frozenset[str]: """Vendor keys local metadata confirms, the EVIDENCE that drives ``auto`` sensitivity. Thin wrapper over the public :func:`api.visible_provenance` (one @@ -896,9 +863,7 @@ def cmd_erase( help="Cap long side (px) before diffusion; 0 = native and preserves the most detail. Raise only on GPU/MPS OOM.", ) @_controlnet_scale_option -@_min_resolution_option @_unsharp_option -@_upscaler_option @_model_option @_guidance_scale_option @_auto_option @@ -920,9 +885,7 @@ def cmd_invisible( humanize: float, unsharp: float, max_resolution: int, - min_resolution: int, controlnet_scale: float, - upscaler: str, model: str | None, guidance_scale: float | None, auto: bool, @@ -952,7 +915,6 @@ def cmd_invisible( source = _validate_image(source) steps = resolve_steps(steps) seed = resolve_seed(seed) - _warn_if_esrgan_unavailable(upscaler) adaptive_polish = _resolve_profile_polish(auto, adaptive_polish, pipeline) if output is None: output = source.with_stem(source.stem + "_clean") @@ -998,8 +960,6 @@ def cmd_invisible( unsharp=unsharp, adaptive_polish=adaptive_polish, max_resolution=max_resolution, - min_resolution=min_resolution, - upscaler=upscaler, vendor=vendor, tile=tile, tile_size=tile_size, @@ -1587,9 +1547,7 @@ def cmd_identify(ctx: click.Context, source: Path, no_visible: bool, as_json: bo help="Cap long side (px) before diffusion; 0 = native and preserves the most detail. Raise only on GPU/MPS OOM.", ) @_controlnet_scale_option -@_min_resolution_option @_unsharp_option -@_upscaler_option @_guidance_scale_option @_auto_option @_adaptive_polish_option @@ -1613,9 +1571,7 @@ def cmd_all( humanize: float, unsharp: float, max_resolution: int, - min_resolution: int, controlnet_scale: float, - upscaler: str, guidance_scale: float | None, auto: bool, adaptive_polish: bool, @@ -1638,7 +1594,6 @@ def cmd_all( source = _validate_image(source) steps = resolve_steps(steps) seed = resolve_seed(seed) - _warn_if_esrgan_unavailable(upscaler) adaptive_polish = _resolve_profile_polish(auto, adaptive_polish, pipeline) if output is None: @@ -1744,8 +1699,6 @@ def cmd_all( unsharp=unsharp, adaptive_polish=adaptive_polish, max_resolution=max_resolution, - min_resolution=min_resolution, - upscaler=upscaler, vendor=vendor, tile=tile, tile_size=tile_size, @@ -1838,9 +1791,7 @@ class _BatchOptions: sensitivity: str = "auto" unsharp: float = 0.0 max_resolution: int = 0 - min_resolution: int = 1024 controlnet_scale: float = 1.0 - upscaler: str = "lanczos" model: str | None = None guidance_scale: float | None = None adaptive_polish: bool = False @@ -1895,8 +1846,6 @@ def _run_batch_invisible( unsharp=options.unsharp, adaptive_polish=options.adaptive_polish, max_resolution=options.max_resolution, - min_resolution=options.min_resolution, - upscaler=options.upscaler, tile=options.tile, tile_size=options.tile_size, tile_overlap=options.tile_overlap, @@ -2030,9 +1979,7 @@ def _process_batch_image( default=0, help="Cap long side (px) before diffusion; 0 = native and preserves the most detail. Raise only on GPU/MPS OOM.", ) -@_min_resolution_option @_unsharp_option -@_upscaler_option @_controlnet_scale_option @_model_option @_guidance_scale_option @@ -2058,9 +2005,7 @@ def cmd_batch( humanize: float, unsharp: float, max_resolution: int, - min_resolution: int, controlnet_scale: float, - upscaler: str, model: str | None, guidance_scale: float | None, auto: bool, @@ -2087,8 +2032,6 @@ def cmd_batch( console.print(f" Found {len(images)} images in {directory}") console.print(f" Output -> {output_dir}") console.print(f" Mode: {mode}") - if mode in ("invisible", "all"): - _warn_if_esrgan_unavailable(upscaler) adaptive_polish = _resolve_profile_polish(auto, adaptive_polish, pipeline) steps = resolve_steps(steps) seed = resolve_seed(seed) @@ -2104,9 +2047,7 @@ def cmd_batch( sensitivity=sensitivity, unsharp=unsharp, max_resolution=max_resolution, - min_resolution=min_resolution, controlnet_scale=controlnet_scale, - upscaler=upscaler, model=model, guidance_scale=guidance_scale, adaptive_polish=adaptive_polish, diff --git a/src/remove_ai_watermarks/invisible_engine.py b/src/remove_ai_watermarks/invisible_engine.py index a254150..2fa40d2 100644 --- a/src/remove_ai_watermarks/invisible_engine.py +++ b/src/remove_ai_watermarks/invisible_engine.py @@ -14,7 +14,7 @@ import logging import os import warnings from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING from ._internal.watermark_profiles import ( DEFAULT_MODEL_ID as DEFAULT_SDXL_MODEL_ID, @@ -48,21 +48,19 @@ def is_available() -> bool: return module_available("diffusers", "torch") -def _target_size(width: int, height: int, max_resolution: int, min_resolution: int = 0) -> tuple[int, int] | None: +def _target_size(width: int, height: int, max_resolution: int) -> tuple[int, int] | None: """Compute the (width, height) to process at, or None for native. - Two opposite long-side adjustments, in precedence order: + One long-side adjustment: if it exceeds ``max_resolution``, scale DOWN to it + (integer-truncated, matching the PIL ``resize`` call site). 0/negative = no cap. + Set only to bound GPU/MPS memory on very large inputs (issue #10). - - ``max_resolution`` (cap): if the long side exceeds it, scale DOWN to it - (integer-truncated, matching the PIL ``resize`` call site). 0/negative = no - cap. Set only to bound GPU/MPS memory on very large inputs (issue #10). - - ``min_resolution`` (floor): else if the long side is below it, scale UP to it - (rounded) so SDXL img2img runs near its ~1024 training resolution instead of - degrading on a tiny latent (a 381x512 portrait distorts badly at native). - The output is restored to the original size by the caller, so the floor is a - transparent quality boost. 0 = no floor. Skipped on a ``min > max`` misconfig. + There was also a ``min_resolution`` floor that scaled small inputs UP toward + SDXL's ~1024 training size. It went with the SDXL profiles: both surviving + profiles run at native geometry, so the floor was forced to 0 on every path and + could not fire. - Returns None when neither applies (native resolution). Pure function so the + Returns None when the cap does not apply (native resolution). Pure function so the resolution decision is unit-testable without loading the diffusion model. """ long_side = max(width, height) @@ -71,9 +69,6 @@ def _target_size(width: int, height: int, max_resolution: int, min_resolution: i # Clamp the short side to >=1: extreme aspect ratios (e.g. 5000x3 capped # at 1024) would otherwise truncate it to 0 and crash image.resize(). return (max(1, int(width * ratio)), max(1, int(height * ratio))) - if min_resolution > 0 and long_side < min_resolution and (max_resolution <= 0 or min_resolution <= max_resolution): - ratio = min_resolution / long_side - return (max(1, round(width * ratio)), max(1, round(height * ratio))) return None @@ -144,32 +139,6 @@ class InvisibleEngine: """ self._remover.preload(global_only=global_only) - def _esrgan_upscale(self, image: Any, target: tuple[int, int]) -> Any: - """Upscale a PIL image to ``target`` with Real-ESRGAN, else Lanczos. - - Runs Real-ESRGAN at its native factor (on the remover's device, CPU fallback), - then resizes to the exact ``target`` with Lanczos. Falls back to a plain Lanczos - resize when the ``esrgan`` extra is absent or the model errors. - """ - import cv2 - import numpy as np - from PIL import Image - - from remove_ai_watermarks import upscaler - - if not upscaler.is_available(): - logger.debug("esrgan upscaler requested but the extra is absent; using Lanczos") - return image.resize(target, Image.Resampling.LANCZOS) - try: - bgr = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR) - big = upscaler.upscale(bgr, device=self._remover.device) - if (big.shape[1], big.shape[0]) != target: - big = cv2.resize(big, target, interpolation=cv2.INTER_LANCZOS4) - return Image.fromarray(cv2.cvtColor(big, cv2.COLOR_BGR2RGB)) - except Exception as e: # never let an optional upscaler break removal - logger.warning("Real-ESRGAN upscale failed (%s); using Lanczos", e) - return image.resize(target, Image.Resampling.LANCZOS) - def remove_watermark( self, image_path: Path, @@ -180,11 +149,9 @@ class InvisibleEngine: seed: int | None = None, humanize: float = 0.0, max_resolution: int = 0, - min_resolution: int = 1024, vendor: str | None = None, unsharp: float = 0.0, adaptive_polish: bool = False, - upscaler: str = "lanczos", tile: bool = False, tile_size: int = 1024, tile_overlap: int = 128, @@ -215,17 +182,6 @@ class InvisibleEngine: = no cap. Set a positive value only to bound GPU/MPS memory on very large inputs (it reintroduces a lossy downscale->upscale round-trip). - min_resolution: Upscale the long side UP to this (px) before diffusion - when the input is smaller, so SDXL runs near its ~1024 training - resolution (small inputs degrade/distort badly at native). 1024 - (default) = on; 0 = off. The output is restored to the original - input size, so this is a transparent quality boost; it adds time - and memory on small inputs. Ignored on a min > max misconfig. - upscaler: How to upscale a small input to the ``min_resolution`` floor: - ``"lanczos"`` (default, cv2, no model download) or ``"esrgan"`` (Real-ESRGAN - via the ``esrgan`` extra). Only applies when UPscaling (the floor - case); a ``max_resolution`` downscale always uses Lanczos. Falls back - to Lanczos if the extra is absent. tile: Process the diffusion pass in overlapping tiles instead of one forward pass. This retains the input's native dimensions instead of applying ``max_resolution``, but each tile is still regenerated. @@ -244,10 +200,7 @@ class InvisibleEngine: from PIL import Image, ImageOps # Resolution policy: a max_resolution cap (0 = none) bounds memory on huge - # inputs, and a min_resolution floor (1024 = default) upscales tiny inputs so - # SDXL img2img runs near its ~1024 training size instead of distorting on a - # tiny latent (a 381x512 portrait wrecks at native -- issue #36 follow-up). - # The output is restored to orig_size below, so the floor is transparent. + # inputs. See _target_size for why it is the only lever left. # Register the HEIF/AVIF opener so a .heic/.avif input (now a SUPPORTED_FORMAT) # decodes here too. The --force skip path bypasses image_io.imread, which is # what would otherwise register it, so a bare Image.open would fail on HEIC. @@ -261,34 +214,16 @@ class InvisibleEngine: # reassigned to the resized copy below; PIL resize returns a new object). reference_pil = image - # qwen-zimage operates at the input's native geometry in its reference graph. - # Keep an explicit max cap available for callers, but do not apply the SDXL - # 1024px minimum-resolution floor to this profile. - effective_min_resolution = ( - 0 if getattr(self._remover, "model_profile", None) in {"qwen-zimage", "sdxl-zimage"} else min_resolution - ) - target = _target_size( - image.width, - image.height, - max_resolution, - effective_min_resolution, - ) + # Both profiles run at the input's native geometry, so only the explicit max + # cap can move it, and it can only ever scale down. + target = _target_size(image.width, image.height, max_resolution) if target is not None: - upscaling = max(target) > max(image.width, image.height) if self._progress_callback: - reason = ( - f"min-resolution floor {min_resolution}px" - if upscaling - else f"max-resolution cap {max_resolution}px" + self._progress_callback( + f"Downscaling {image.width}x{image.height} to {target[0]}x{target[1]} " + f"(max-resolution cap {max_resolution}px)..." ) - verb = "Upscaling" if upscaling else "Downscaling" - self._progress_callback(f"{verb} {image.width}x{image.height} to {target[0]}x{target[1]} ({reason})...") - # Real-ESRGAN only helps when UPscaling (the floor case); a downscale cap - # always uses Lanczos. _esrgan_upscale falls back to Lanczos if the extra is absent. - if upscaling and upscaler == "esrgan": - image = self._esrgan_upscale(image, target) - else: - image = image.resize(target, Image.Resampling.LANCZOS) + image = image.resize(target, Image.Resampling.LANCZOS) # Always persist to a temp file, even without downscaling: WatermarkRemover # reloads by path, so the EXIF-transposed pixels must be saved or rotation diff --git a/src/remove_ai_watermarks/upscaler.py b/src/remove_ai_watermarks/upscaler.py deleted file mode 100644 index a545f13..0000000 --- a/src/remove_ai_watermarks/upscaler.py +++ /dev/null @@ -1,126 +0,0 @@ -"""Optional pre-diffusion super-resolution for small inputs (Real-ESRGAN via spandrel). - -Mirrors ``region_eraser``'s optional-backend pattern: ``is_available()`` guards the -``spandrel`` import, a lazy singleton (double-checked lock) holds the loaded model, and -the weights download on first use (cached by ``torch.hub``) -- they are never bundled. - -The DEFAULT upscaler stays Lanczos (cv2, no model download); this is opt-in via the ``esrgan`` -extra and feeds the ``--upscaler esrgan`` path. ``spandrel`` is a pure model-loader -(MIT) with NO basicsr dependency -- it pulls only torch/torchvision/safetensors/numpy/ -einops -- so it sidesteps the basicsr / ``torchvision.transforms.functional_tensor`` -breakage that the retired ``restore`` (GFPGAN) extra had to shim. Real-ESRGAN weights -are BSD-3-Clause. - -CPU works but is slow on large inputs, so this is meant for the pre-diffusion upscale of -SMALL inputs (and the GPU worker). On a memory-constrained host it is a no-op (the extra -is absent), and the caller falls back to Lanczos. -""" - -# torch/spandrel boundary: these libs ship no usable element types; relax the -# unknown-type rules for this file only. -# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false, reportAttributeAccessIssue=false, reportPrivateImportUsage=false -from __future__ import annotations - -import logging -import threading -from pathlib import Path -from typing import TYPE_CHECKING, Any - -if TYPE_CHECKING: - from numpy.typing import NDArray - -logger = logging.getLogger(__name__) - -# Real-ESRGAN x2plus (BSD-3-Clause), official release. x2 is the right native factor for -# the pre-diffusion floor upscale (small inputs ~512 -> ~1024); spandrel infers the -# architecture and scale from the checkpoint, so swapping the URL is enough to change it. -_MODEL_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth" -_MODEL_FILENAME = "RealESRGAN_x2plus.pth" - -_model: Any = None # lazy singleton (spandrel ImageModelDescriptor) -_model_device: str = "cpu" -_lock = threading.Lock() - - -def is_available() -> bool: - """True if the ``esrgan`` extra (spandrel + torch) is importable.""" - from .optional_deps import module_available - - return module_available("spandrel", "torch") - - -def _model_cache_path() -> Path: - """Path the weights are cached at (the torch.hub checkpoints dir).""" - import torch - - cache_dir = Path(torch.hub.get_dir()) / "checkpoints" - cache_dir.mkdir(parents=True, exist_ok=True) - return cache_dir / _MODEL_FILENAME - - -def _get_model(device: str) -> Any: - """Load the Real-ESRGAN model once (downloading the weights on first use).""" - global _model, _model_device - if _model is not None and _model_device == device: - return _model - with _lock: - if _model is None: - import torch - from spandrel import ImageModelDescriptor, ModelLoader - - dst = _model_cache_path() - if not dst.exists(): - logger.info("Downloading Real-ESRGAN weights to %s", dst) - torch.hub.download_url_to_file(_MODEL_URL, str(dst), progress=False) - model = ModelLoader().load_from_file(str(dst)) - if not isinstance(model, ImageModelDescriptor): - raise RuntimeError(f"Unexpected spandrel model type: {type(model).__name__}") - _model = model.eval() - if _model_device != device: - _model.to(device) - _model_device = device - return _model - - -def scale() -> int: - """The model's native upscale factor (e.g. 2 for x2plus). Loads the model if needed.""" - return int(_get_model("cpu").scale) - - -def upscale(image: NDArray[Any], device: str | None = None) -> NDArray[Any]: - """Upscale a BGR uint8 image by the model's native factor with Real-ESRGAN. - - Returns a BGR uint8 array. Falls back to CPU if the requested device errors (an - MPS/CUDA OOM or unsupported-op on the small pre-diffusion input), mirroring the - diffusion engine's MPS->CPU fallback. - - Raises: - RuntimeError: if the ``esrgan`` extra is not installed (guard with - ``is_available()`` first). - """ - if not is_available(): - raise RuntimeError("Real-ESRGAN upscaler needs the 'esrgan' extra (spandrel). Install it or use Lanczos.") - import cv2 - import numpy as np - import torch - - target_device = (device or "cpu").lower() - if target_device not in {"cpu", "mps", "cuda", "xpu"}: - target_device = "cpu" - rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) - tensor = torch.from_numpy(rgb).permute(2, 0, 1).float().div(255.0).unsqueeze(0) - - def _run(dev: str) -> NDArray[Any]: - model = _get_model(dev) - with torch.no_grad(): - out = model(tensor.to(dev)) - arr = out.clamp(0.0, 1.0).squeeze(0).permute(1, 2, 0).cpu().numpy() * 255.0 - return cv2.cvtColor(arr.round().astype(np.uint8), cv2.COLOR_RGB2BGR) - - try: - return _run(target_device) - except Exception as e: # GPU OOM / unsupported op: fall back to CPU - if target_device == "cpu": - raise - logger.warning("Real-ESRGAN on %s failed (%s); retrying on CPU", target_device, e) - return _run("cpu") diff --git a/tests/test_invisible_engine.py b/tests/test_invisible_engine.py index 80562a6..2445010 100644 --- a/tests/test_invisible_engine.py +++ b/tests/test_invisible_engine.py @@ -61,7 +61,7 @@ class TestNativeOutputSize: out = tmp_path / "out.png" Image.new("RGB", (24, 18), (128, 128, 128)).save(src) - engine.remove_watermark(src, out, min_resolution=0, adaptive_polish=False) + engine.remove_watermark(src, out, adaptive_polish=False) assert Image.open(out).size == (24, 18) @@ -107,35 +107,10 @@ class TestTargetSize: assert _target_size(5000, 3, 1024) == (1024, 1) assert _target_size(3, 5000, 1024) == (1, 1024) - # ── min_resolution floor (small inputs upscaled so SDXL runs near 1024) ── - - def test_floor_default_off(self): - # min_resolution defaults to 0 -> no upscale, preserving legacy behavior. + def test_a_small_input_is_left_at_native_size(self): + """No minimum-resolution floor: only the cap can move geometry.""" assert _target_size(381, 512, 0) is None - - def test_floor_upscales_small_input(self): - # 381x512 portrait, floor 1024 -> long side 512 scaled up to 1024 (x2). - assert _target_size(381, 512, 0, 1024) == (762, 1024) - # Landscape: width is the long side. - assert _target_size(512, 381, 0, 1024) == (1024, 762) - - def test_floor_rounds_short_side(self): - # 333x500, floor 1024: ratio 2.048 -> 333*2.048=681.98 rounds to 682. - assert _target_size(333, 500, 0, 1024) == (682, 1024) - - def test_floor_no_op_at_or_above_floor(self): - # Long side already >= floor -> no upscale (and no cap set -> native). - assert _target_size(1024, 768, 0, 1024) is None - assert _target_size(2000, 1000, 0, 1024) is None - - def test_cap_takes_precedence_over_floor(self): - # A huge input with both set: the cap downscales; the floor never fires. - assert _target_size(2000, 1000, 1024, 1024) == (1024, 512) - - def test_floor_skipped_on_min_above_max_misconfig(self): - # min(1024) > max(800) is a misconfig: the floor must not upscale above the - # cap, so it is skipped and the (within-cap) input stays native. - assert _target_size(500, 400, 800, 1024) is None + assert _target_size(381, 512, 4096) is None class TestEngineDoesNotFabricateAModelId: @@ -165,70 +140,3 @@ class TestEngineDoesNotFabricateAModelId: 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" - - -class TestEsrganUpscale: - """Branches of InvisibleEngine._esrgan_upscale (no diffusion model loaded). - - A SimpleNamespace stands in for the engine so we exercise the helper without - constructing a real InvisibleEngine (which would load WatermarkRemover). - """ - - @staticmethod - def _fake_engine(): - from types import SimpleNamespace - - return SimpleNamespace(_remover=SimpleNamespace(device="cpu")) - - @staticmethod - def _pil(w=120, h=80): - import numpy as np - from PIL import Image - - return Image.fromarray(np.full((h, w, 3), 128, dtype=np.uint8)) - - def test_falls_back_to_lanczos_when_extra_absent(self, monkeypatch): - import numpy as np - from PIL import Image - - from remove_ai_watermarks import upscaler - - monkeypatch.setattr(upscaler, "is_available", lambda: False) - img = self._pil() - out = InvisibleEngine._esrgan_upscale(self._fake_engine(), img, (1024, 683)) - assert out.size == (1024, 683) - # Identical to a plain Lanczos resize (the fallback path). - assert np.array_equal(np.asarray(out), np.asarray(img.resize((1024, 683), Image.Resampling.LANCZOS))) - - def test_resizes_esrgan_output_to_exact_target(self, monkeypatch): - import cv2 - - from remove_ai_watermarks import upscaler - - monkeypatch.setattr(upscaler, "is_available", lambda: True) - - # Fake a 2x upscale that does NOT match the requested target; the helper must - # resize it to the exact target. - def _fake_upscale(bgr, device=None): - return cv2.resize(bgr, (bgr.shape[1] * 2, bgr.shape[0] * 2), interpolation=cv2.INTER_NEAREST) - - monkeypatch.setattr(upscaler, "upscale", _fake_upscale) - out = InvisibleEngine._esrgan_upscale(self._fake_engine(), self._pil(), (1024, 683)) - assert out.size == (1024, 683) - - def test_falls_back_to_lanczos_when_upscale_raises(self, monkeypatch): - import numpy as np - from PIL import Image - - from remove_ai_watermarks import upscaler - - monkeypatch.setattr(upscaler, "is_available", lambda: True) - - def _boom(bgr, device=None): - raise RuntimeError("model exploded") - - monkeypatch.setattr(upscaler, "upscale", _boom) - img = self._pil() - 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))) diff --git a/tests/test_qwen_zimage_pipeline.py b/tests/test_qwen_zimage_pipeline.py index cf46079..8e00eb2 100644 --- a/tests/test_qwen_zimage_pipeline.py +++ b/tests/test_qwen_zimage_pipeline.py @@ -767,11 +767,7 @@ def test_invisible_engine_uses_qwen_zimage_step_default(tmp_image_path, tmp_path engine._remover = MagicMock(model_profile="qwen-zimage") engine._remover.remove_watermark.return_value = tmp_path / "clean.png" - engine.remove_watermark( - tmp_image_path, - tmp_path / "clean.png", - min_resolution=0, - ) + engine.remove_watermark(tmp_image_path, tmp_path / "clean.png") assert engine._remover.remove_watermark.call_args.kwargs["num_inference_steps"] == 4 diff --git a/tests/test_upscaler.py b/tests/test_upscaler.py deleted file mode 100644 index e504c44..0000000 --- a/tests/test_upscaler.py +++ /dev/null @@ -1,32 +0,0 @@ -"""Tests for the optional Real-ESRGAN upscaler (no model download). - -The model-running path is exercised manually (it downloads ~67 MB of BSD-3-Clause -weights on first use); these tests cover the availability guard and the no-model -control flow, mirroring the repo convention for ML-adjacent modules. -""" - -from __future__ import annotations - -import numpy as np -import pytest - -from remove_ai_watermarks import upscaler - - -class TestIsAvailable: - def test_returns_bool(self): - assert isinstance(upscaler.is_available(), bool) - - -class TestUpscaleGuard: - def test_raises_without_extra(self, monkeypatch): - monkeypatch.setattr(upscaler, "is_available", lambda: False) - with pytest.raises(RuntimeError, match="esrgan"): - upscaler.upscale(np.full((32, 32, 3), 128, dtype=np.uint8)) - - -class TestModelCachePath: - def test_cache_path_uses_model_filename(self): - if not upscaler.is_available(): - pytest.skip("esrgan extra (torch) not installed") - assert upscaler._model_cache_path().name == upscaler._MODEL_FILENAME diff --git a/uv.lock b/uv.lock index af7caab..c4933c6 100644 --- a/uv.lock +++ b/uv.lock @@ -3357,7 +3357,6 @@ all = [ { name = "pywavelets", version = "1.8.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "pywavelets", version = "1.9.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "safetensors" }, - { name = "spandrel" }, { name = "tokenizers" }, { name = "torch" }, { name = "torchvision" }, @@ -3397,11 +3396,6 @@ diffusion = [ { name = "torch" }, { name = "transformers" }, ] -esrgan = [ - { name = "numpy" }, - { name = "opencv-python-headless" }, - { name = "spandrel" }, -] heif = [ { name = "pillow-heif" }, ] @@ -3481,16 +3475,14 @@ requires-dist = [ { name = "remove-ai-watermarks", extras = ["diffusion"], marker = "extra == 'qwen-zimage'" }, { name = "remove-ai-watermarks", extras = ["pixels"], marker = "extra == 'detect'" }, { name = "remove-ai-watermarks", extras = ["pixels"], marker = "extra == 'diffusion'" }, - { name = "remove-ai-watermarks", extras = ["pixels"], marker = "extra == 'esrgan'" }, { name = "remove-ai-watermarks", extras = ["pixels"], marker = "extra == 'visible'" }, { name = "remove-ai-watermarks", extras = ["video"], marker = "extra == 'dev'" }, - { name = "remove-ai-watermarks", extras = ["video", "heif", "detect", "trustmark", "diffusion", "qwen-zimage", "lama", "migan", "esrgan"], marker = "extra == 'all'" }, + { name = "remove-ai-watermarks", extras = ["video", "heif", "detect", "trustmark", "diffusion", "qwen-zimage", "lama", "migan"], marker = "extra == 'all'" }, { name = "remove-ai-watermarks", extras = ["visible"], marker = "extra == 'lama'" }, { name = "remove-ai-watermarks", extras = ["visible"], marker = "extra == 'migan'" }, { name = "remove-ai-watermarks", extras = ["visible"], marker = "extra == 'video'" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.4.0" }, { name = "safetensors", marker = "extra == 'diffusion'" }, - { name = "spandrel", marker = "extra == 'esrgan'", specifier = ">=0.3.0" }, { name = "tokenizers", marker = "extra == 'diffusion'", specifier = ">=0.22,<0.23" }, { name = "torch", marker = "extra == 'diffusion'", specifier = ">=2.0.0" }, { name = "torchvision", marker = "extra == 'qwen-zimage'", specifier = ">=0.20.0" }, @@ -3499,7 +3491,7 @@ requires-dist = [ { name = "uv-outdated", marker = "python_full_version >= '3.12' and extra == 'dev'", specifier = ">=0.1.0" }, { name = "uv-secure", marker = "python_full_version >= '3.12' and extra == 'dev'", specifier = ">=0.12.0" }, ] -provides-extras = ["pixels", "heif", "visible", "video", "detect", "diffusion", "qwen-zimage", "trustmark", "lama", "migan", "esrgan", "dev", "all"] +provides-extras = ["pixels", "heif", "visible", "video", "detect", "diffusion", "qwen-zimage", "trustmark", "lama", "migan", "dev", "all"] [[package]] name = "requests" @@ -3662,23 +3654,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, ] -[[package]] -name = "spandrel" -version = "0.4.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "einops" }, - { name = "numpy" }, - { name = "safetensors" }, - { name = "torch" }, - { name = "torchvision" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/2a/8f/ab4565c23dd67a036ab72101a830cebd7ca026b2fddf5771bbf6284f6228/spandrel-0.4.2.tar.gz", hash = "sha256:fefa4ea966c6a5b7721dcf24f3e2062a5a96a395c8bedcb570fb55971fdcbccb", size = 247544, upload-time = "2026-02-21T01:52:26.342Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/74/31/411ea965835534c43d4b98d451968354876e0e867ea1fd42669e4cca0732/spandrel-0.4.2-py3-none-any.whl", hash = "sha256:6c93e3ecbeb0e548fd2df45a605472b34c1614287c56b51bb33cdef7ae5235b5", size = 320811, upload-time = "2026-02-21T01:52:25.015Z" }, -] - [[package]] name = "stamina" version = "26.1.0"