mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-09 23:50:40 +02:00
Remove the unreachable ESRGAN upscale chain
The min-resolution floor lifted small inputs toward SDXL's ~1024 training size, and Real-ESRGAN was an optional way to do that lifting. Both surviving profiles run at native geometry, so the engine forced the floor to 0 on every path; the floor never fired, `upscaling` was never true, and nothing downstream of it could execute. Gone: upscaler.py, _esrgan_upscale, the min_resolution and upscaler parameters, _target_size's floor branch, --min-resolution, --upscaler, _warn_if_esrgan_unavailable and the `esrgan` extra. max_resolution stays and is now the only lever on geometry; it can only scale down. scripts/smoke_matrix.py was the one live consumer and neither gate saw it - Pyright is scoped to src/ and Ruff cannot resolve its function-local import - so `--diffusion` would have died at import. Its knob rows were written for the removed profiles besides (--pipeline sdxl, --steps 20, --guidance-scale 5.0, --device mps), so they are rewritten rather than patched: most now assert a knob is REJECTED, which is the coverage worth having when the CLI accepts a value the library refuses several layers down. Accepted-knob rows skip without CUDA, so the row count is host-dependent and verification-plan.md no longer claims a fixed 68. This removes a public module, a CLI option and a published extra, so the next release is 0.25.0, not a patch. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Victor Kuznetsov
co-authored by
Claude Opus 5
Claude Opus 5
parent
95a6964e04
commit
bf4bfc1ab7
@@ -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.
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
+1
-1
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
+1
-14
@@ -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
|
||||
|
||||
+51
-60
@@ -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")
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
@@ -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)))
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user