mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-06-05 10:38:00 +02:00
Victor Kuznetsov
d7e4fe8835
Two quality knobs for the SDXL invisible pass: - min_resolution floor (default 1024, --min-resolution): small inputs are upscaled to a 1024px long-side floor before diffusion, since SDXL img2img distorts on a tiny latent (a 381x512 portrait wrecks at native). The output is restored to the original input size, so it is a transparent quality boost; it adds time/memory on small inputs. 0 disables. Extends the pure _target_size helper (now cap-or-floor-or-native, min skipped on a min>max misconfig), unit-tested without a model. - unsharp post-filter (humanizer.unsharp_mask, --unsharp, opt-in default 0): applied LAST, after the GFPGAN face pass (a pre-GFPGAN sharpen would be smoothed back over), to counter the soft/over-smoothed look that diffusion + restoration leave behind (an AI tell). Pairs with --humanize (grain). Both threaded through invisible/all/batch + the module-level helper. Verified end-to-end on a 381x512 portrait: upscaled to 1024, sharpened, restored to 381x512. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
104 lines
4.4 KiB
Python
104 lines
4.4 KiB
Python
"""Tests for the invisible watermark engine (unit tests, no GPU required)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from remove_ai_watermarks.invisible_engine import InvisibleEngine, _target_size, is_available
|
|
|
|
|
|
class TestIsAvailable:
|
|
"""Tests for dependency checking."""
|
|
|
|
def test_returns_bool(self):
|
|
result = is_available()
|
|
assert isinstance(result, bool)
|
|
|
|
def test_available_reflects_dependencies(self):
|
|
"""is_available() is True iff torch + diffusers (the gpu extra) import.
|
|
|
|
Must not assume the full stack: the core+dev CI env has no diffusers.
|
|
"""
|
|
import importlib.util
|
|
|
|
expected = all(importlib.util.find_spec(m) is not None for m in ("torch", "diffusers"))
|
|
assert is_available() is expected
|
|
|
|
|
|
class TestInvisibleEngineInit:
|
|
"""Tests for InvisibleEngine construction (no GPU required)."""
|
|
|
|
def test_default_model_id(self):
|
|
# SDXL base became the default in May 2026 (defeats SynthID v2).
|
|
assert InvisibleEngine.DEFAULT_MODEL_ID == "stabilityai/stable-diffusion-xl-base-1.0"
|
|
|
|
|
|
class TestTargetSize:
|
|
"""Regression guard for the native-resolution decision (issues #10 / #15).
|
|
|
|
max_resolution=0 must NOT downscale -- the forced downscale->upscale
|
|
round-trip was the quality loss in #10, and downscaling at all let SynthID
|
|
survive in #15 (the native SDXL pass at strength ~0.05 is what defeats it).
|
|
"""
|
|
|
|
def test_native_default_no_downscale(self):
|
|
# The default (0) means native resolution: no resize, regardless of size.
|
|
assert _target_size(4096, 4096, 0) is None
|
|
assert _target_size(123, 456, 0) is None
|
|
|
|
def test_negative_cap_treated_as_native(self):
|
|
assert _target_size(4096, 4096, -1) is None
|
|
|
|
def test_cap_below_long_side_downscales(self):
|
|
# 2000x1000, cap 1024 -> long side scaled to 1024, aspect preserved.
|
|
assert _target_size(2000, 1000, 1024) == (1024, 512)
|
|
|
|
def test_cap_uses_long_side_for_portrait(self):
|
|
# Portrait: height is the long side, so it drives the ratio.
|
|
assert _target_size(1000, 2000, 1024) == (512, 1024)
|
|
|
|
def test_cap_at_or_above_long_side_no_downscale(self):
|
|
# Already within the cap (and exactly equal) -> no resize.
|
|
assert _target_size(800, 600, 1024) is None
|
|
assert _target_size(1024, 768, 1024) is None
|
|
|
|
def test_integer_truncation_matches_pil_call_site(self):
|
|
# 1254x1254 (the gpt-image sample) capped at 1000: int(1254*1000/1254)=1000.
|
|
assert _target_size(1254, 1254, 1000) == (1000, 1000)
|
|
# Non-divisible ratio truncates toward zero like int() at the call site.
|
|
assert _target_size(1000, 333, 500) == (500, 166)
|
|
|
|
def test_extreme_aspect_ratio_clamps_short_side_to_one(self):
|
|
# 5000x3 capped at 1024: int(3 * 1024/5000) = 0 would crash resize();
|
|
# the short side must clamp to 1, never 0.
|
|
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.
|
|
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
|