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:
Victor Kuznetsov
2026-08-03 13:24:53 -07:00
co-authored by Claude Opus 5
parent 691f5d88c2
commit 1a77e24b99
2 changed files with 33 additions and 3 deletions
+4 -3
View File
@@ -121,10 +121,11 @@ class InvisibleEngine:
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(
model_id=effective_model,
model_id=model_id,
device=device,
progress_callback=progress_callback,
hf_token=hf_token,
+29
View File
@@ -138,6 +138,35 @@ class TestTargetSize:
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:
"""Branches of InvisibleEngine._esrgan_upscale (no diffusion model loaded).