mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-07-05 07:57:50 +02:00
fix(invisible): preserve native output dimensions
This commit is contained in:
@@ -269,7 +269,11 @@ class InvisibleEngine:
|
||||
# each stage independently imread/imwrote the full-res output, so a run
|
||||
# with several stages PNG-decoded+re-encoded the same image 2-4 times.
|
||||
# PNG is lossless, so the single-write output is byte-identical.
|
||||
needs_restore = target is not None # the input was resized before diffusion
|
||||
# Diffusers rounds native dimensions down to the latent grid (multiples
|
||||
# of 8), even when our own resolution policy did not resize the input.
|
||||
# Route those outputs through the same final resize so --no-polish does
|
||||
# not silently change e.g. 1448x1086 into 1448x1080.
|
||||
needs_restore = target is not None or any(dimension % 8 for dimension in orig_size)
|
||||
if humanize > 0.0 or unsharp > 0.0 or adaptive_polish or needs_restore:
|
||||
import cv2
|
||||
|
||||
|
||||
@@ -2,6 +2,10 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from remove_ai_watermarks.invisible_engine import InvisibleEngine, _target_size, is_available
|
||||
|
||||
|
||||
@@ -31,6 +35,29 @@ class TestInvisibleEngineInit:
|
||||
assert InvisibleEngine.DEFAULT_MODEL_ID == "stabilityai/stable-diffusion-xl-base-1.0"
|
||||
|
||||
|
||||
class TestNativeOutputSize:
|
||||
"""Model-side latent-grid rounding must not change the public output size."""
|
||||
|
||||
def test_no_polish_restores_native_non_multiple_of_eight_size(self, tmp_path):
|
||||
engine = object.__new__(InvisibleEngine)
|
||||
|
||||
def _remove_watermark(image_path, output_path=None, **_kwargs):
|
||||
out = output_path or image_path.with_stem(image_path.stem + "_clean")
|
||||
# Model-side latent-grid rounding: 18px becomes 16px.
|
||||
Image.open(image_path).crop((0, 0, 24, 16)).save(out)
|
||||
return out
|
||||
|
||||
engine._remover = SimpleNamespace(remove_watermark=_remove_watermark)
|
||||
engine._progress_callback = None
|
||||
src = tmp_path / "src.png"
|
||||
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)
|
||||
|
||||
assert Image.open(out).size == (24, 18)
|
||||
|
||||
|
||||
class TestTargetSize:
|
||||
"""Regression guard for the native-resolution decision (issues #10 / #15).
|
||||
|
||||
|
||||
Reference in New Issue
Block a user