CalvinBackup/remove-ai-watermarks
1
0
Fork 0
mirror of https://github.com/wiltodelta/remove-ai-watermarks.git synced 2026-08-09 23:50:40 +02:00
Code Issues Packages Projects Releases Wiki Activity

Remove the unreachable ESRGAN upscale chain

Browse Source 
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
2026-08-03 14:54:26 -07:00
co-authored by Claude Opus 5
parent 95a6964e04
commit bf4bfc1ab7
15 changed files with 99 additions and 521 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_invisible_engine.py
+4 -96
View File
@@ -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)))
tests/test_qwen_zimage_pipeline.py
+1 -5
View File
@@ -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
tests/test_upscaler.py
-32
View File
@@ -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