mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-09 23:50:40 +02:00
Stop the engine from supplying a model id the remover must reject
InvisibleEngine substituted DEFAULT_MODEL_ID whenever model_id was None. When
b0ca205 tightened the remover's fixed-stack check from `not in {None,
DEFAULT_MODEL_ID}` to `is not None`, that substitution turned every single
InvisibleEngine construction into a ValueError - including the deployed Modal
worker's setup(), which is how it was found.
The library suite missed it because these two are tested from opposite sides:
every remover test builds WatermarkRemover directly with model_id unset, and
every engine test mocks the remover away. Nothing exercised the seam between
them. TestEngineDoesNotFabricateAModelId now does, without a GPU.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
691f5d88c2
commit
1a77e24b99
@@ -121,10 +121,11 @@ class InvisibleEngine:
|
|||||||
|
|
||||||
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
|
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
|
||||||
|
|
||||||
effective_model = model_id or self.DEFAULT_MODEL_ID
|
# Pass model_id through untouched. Substituting DEFAULT_MODEL_ID for None here
|
||||||
|
# meant the engine always supplied a model the remover is required to reject,
|
||||||
|
# so EVERY construction raised once that check tightened to "is not None".
|
||||||
self._remover = WatermarkRemover(
|
self._remover = WatermarkRemover(
|
||||||
model_id=effective_model,
|
model_id=model_id,
|
||||||
device=device,
|
device=device,
|
||||||
progress_callback=progress_callback,
|
progress_callback=progress_callback,
|
||||||
hf_token=hf_token,
|
hf_token=hf_token,
|
||||||
|
|||||||
@@ -138,6 +138,35 @@ class TestTargetSize:
|
|||||||
assert _target_size(500, 400, 800, 1024) is None
|
assert _target_size(500, 400, 800, 1024) is None
|
||||||
|
|
||||||
|
|
||||||
|
class TestEngineDoesNotFabricateAModelId:
|
||||||
|
"""The engine must forward model_id untouched, including None.
|
||||||
|
|
||||||
|
It used to substitute DEFAULT_MODEL_ID for None. Once the remover tightened its
|
||||||
|
"you may not override the fixed stack" check from `not in {None, DEFAULT_MODEL_ID}`
|
||||||
|
to `is not None`, that substitution made EVERY InvisibleEngine construction raise -
|
||||||
|
and no test saw it, because the library tests build WatermarkRemover directly while
|
||||||
|
the engine tests mock it. A deployed Modal worker caught it instead.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def test_none_stays_none(self):
|
||||||
|
from unittest.mock import patch
|
||||||
|
|
||||||
|
import remove_ai_watermarks.invisible_engine as engine_module
|
||||||
|
|
||||||
|
with patch("remove_ai_watermarks._internal.watermark_remover.WatermarkRemover") as remover:
|
||||||
|
engine_module.InvisibleEngine(pipeline="qwen-zimage")
|
||||||
|
assert remover.call_args.kwargs["model_id"] is None
|
||||||
|
|
||||||
|
def test_an_explicit_model_id_still_reaches_the_remover_to_be_rejected(self):
|
||||||
|
from unittest.mock import patch
|
||||||
|
|
||||||
|
import remove_ai_watermarks.invisible_engine as engine_module
|
||||||
|
|
||||||
|
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:
|
class TestEsrganUpscale:
|
||||||
"""Branches of InvisibleEngine._esrgan_upscale (no diffusion model loaded).
|
"""Branches of InvisibleEngine._esrgan_upscale (no diffusion model loaded).
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user