From f7e0a18e6c8d66b73b276fcd2b7f8b0949bc8113 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Sun, 30 Aug 2026 20:02:14 -0700 Subject: [PATCH] Add --pipeline auto: per-cohort engine router from provenance Routes OpenAI/Microsoft provenance to chroma-zimage (lower floors, better matched-strength fidelity) and Google/Meta/unknown to qwen-zimage, from the 2026-08-29/30 four-cohort calibration. The resolution happens in remove_watermark after the vendor is known and before the strength resolution, so each engine's own floors apply. --- .../_internal/watermark_profiles.py | 26 ++++++-- .../_internal/watermark_remover.py | 19 ++++-- src/remove_ai_watermarks/cli.py | 3 +- tests/test_qwen_zimage_pipeline.py | 59 +++++++++++++++++++ 4 files changed, 98 insertions(+), 9 deletions(-) diff --git a/src/remove_ai_watermarks/_internal/watermark_profiles.py b/src/remove_ai_watermarks/_internal/watermark_profiles.py index 99dd1e5..7f1cb41 100644 --- a/src/remove_ai_watermarks/_internal/watermark_profiles.py +++ b/src/remove_ai_watermarks/_internal/watermark_profiles.py @@ -25,8 +25,9 @@ CONTROLNET_CANNY_MODEL = "xinsir/controlnet-canny-sdxl-1.0" QWEN_ZIMAGE_PROFILE = "qwen-zimage" SDXL_ZIMAGE_PROFILE = "sdxl-zimage" CHROMA_ZIMAGE_PROFILE = "chroma-zimage" +AUTO_PROFILE = "auto" DEFAULT_PROFILE = QWEN_ZIMAGE_PROFILE -PROFILE_CHOICES = (QWEN_ZIMAGE_PROFILE, SDXL_ZIMAGE_PROFILE, CHROMA_ZIMAGE_PROFILE) +PROFILE_CHOICES = (QWEN_ZIMAGE_PROFILE, SDXL_ZIMAGE_PROFILE, CHROMA_ZIMAGE_PROFILE, AUTO_PROFILE) # The modules a real removal run needs, and the extra that installs them. Both live # here, in the only profile module that imports nothing heavy, because the CLI's @@ -47,10 +48,11 @@ INVISIBLE_EXTRA = "'remove-ai-watermarks[qwen-zimage]'" PROFILE_ADAPTIVE_POLISH = { QWEN_ZIMAGE_PROFILE: False, SDXL_ZIMAGE_PROFILE: True, - # Chroma1's output matches the input's detail level, like qwen-zimage: the - # matched-strength calibration measured LPIPS/SSIM parity or better, so - # polishing would add grain to an already-faithful output. CHROMA_ZIMAGE_PROFILE: False, + # Auto resolves to qwen-zimage or chroma-zimage, both of which keep the + # input's detail level. If it ever resolves to an engine that needs polish, + # the resolved profile's own default applies. + AUTO_PROFILE: False, } SDXL_LIGHTNING_MODEL_ID = "ByteDance/SDXL-Lightning" @@ -176,8 +178,24 @@ _ALIASES = { "qwen_zimage": QWEN_ZIMAGE_PROFILE, "sdxl_zimage": SDXL_ZIMAGE_PROFILE, "chroma_zimage": CHROMA_ZIMAGE_PROFILE, + "auto": AUTO_PROFILE, } +# The per-cohort engine router: which profile wins on which vendor, from the +# 2026-08-29/30 four-cohort calibration (docs/chroma1-engine-research.md). +# chroma-zimage has lower floors AND better matched-strength fidelity on OpenAI +# and Microsoft; qwen-zimage wins on Google and Meta. Unknown stays on qwen +# (the shipped default, conservative). +_ENGINE_BY_VENDOR: dict[str, str] = { + "openai": CHROMA_ZIMAGE_PROFILE, + "microsoft": CHROMA_ZIMAGE_PROFILE, +} + + +def resolve_auto_profile(vendor: str | None) -> str: + """Pick the engine for --pipeline auto from the provenance vendor.""" + return _ENGINE_BY_VENDOR.get((vendor or "").casefold(), QWEN_ZIMAGE_PROFILE) + def normalize_profile(profile: str) -> str: """Normalize spelling and resolve the underscore spellings.""" diff --git a/src/remove_ai_watermarks/_internal/watermark_remover.py b/src/remove_ai_watermarks/_internal/watermark_remover.py index d6e6e41..c1097bf 100644 --- a/src/remove_ai_watermarks/_internal/watermark_remover.py +++ b/src/remove_ai_watermarks/_internal/watermark_remover.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any from PIL import Image from remove_ai_watermarks._internal.watermark_profiles import ( + AUTO_PROFILE, CHROMA_ZIMAGE_PROFILE, DEFAULT_PROFILE, INVISIBLE_EXTRA, @@ -18,6 +19,7 @@ from remove_ai_watermarks._internal.watermark_profiles import ( REMOVAL_MODULES, SDXL_ZIMAGE_PROFILE, normalize_profile, + resolve_auto_profile, resolve_seed, resolve_strength, ) @@ -110,6 +112,10 @@ class WatermarkRemover: self.model_profile = normalize_profile(pipeline) if self.model_profile not in PROFILE_CHOICES: raise ValueError(f"Unsupported pipeline '{pipeline}'. Use one of: {', '.join(PROFILE_CHOICES)}.") + # The auto profile resolves to a concrete engine per-image in + # remove_watermark, once the provenance vendor is known. It never + # resolves to sdxl-zimage, so the dtype below is correct either way. + self._auto = self.model_profile == AUTO_PROFILE # There is no ``model_id`` parameter and no ``model_id`` attribute: each # profile pins a fixed model stack, and the dtype below is bound to that # stack's weights. Both used to be constructor overrides that existed only to @@ -132,11 +138,8 @@ class WatermarkRemover: # SDXL ships fp16 weights and an fp16-safe VAE; bf16 would give up the # variant without buying anything on this architecture. self.torch_dtype = torch.float16 # type: ignore[union-attr] - elif self.model_profile == CHROMA_ZIMAGE_PROFILE: - # Chroma1-HD is calibrated bf16 end to end (the floors were measured - # with bf16 weights); the inherited face stage is bf16 regardless. - self.torch_dtype = torch.bfloat16 # type: ignore[union-attr] else: + # qwen-zimage, chroma-zimage, and auto all resolve to bf16 engines. self.torch_dtype = torch.bfloat16 # type: ignore[union-attr] self.cpu_offload = cpu_offload @@ -233,6 +236,14 @@ class WatermarkRemover: with Image.open(image_path) as opened: source = opened.convert("RGB") + # The auto router resolves per-image once the provenance vendor is + # known, BEFORE the strength resolution so the right engine's floors + # are used. If the vendor changed the engine, reset the cached pipeline. + if self._auto: + resolved = resolve_auto_profile(vendor) + if resolved != self.model_profile: + self.model_profile = resolved + self._qwen_zimage_pipeline = None resolved_strength = self._resolve_chroma_strength(strength, vendor, source) if not 0.0 <= resolved_strength <= 1.0: raise ValueError(f"Strength must be between 0.0 and 1.0, got {resolved_strength}") diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index 776cf3f..671ad27 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -256,7 +256,8 @@ _PIPELINE_HELP = ( "followed by SAM-masked Z-Image face repair; sdxl-zimage = the same recipe and the " "same face stage on an SDXL global pass, which needs more denoise; chroma-zimage = " "the same face stage on an Apache-2.0 Chroma1 global pass with lower OpenAI and " - "Microsoft floors and higher Google/Meta floors (docs/chroma1-engine-research.md). " + "Microsoft floors and higher Google/Meta floors; auto = pick the engine from the " + "provenance (chroma-zimage for OpenAI/Microsoft, qwen-zimage for Google/Meta/unknown). " "All are CUDA-ONLY -- install the qwen-zimage extra. There is no CPU or MPS profile for " "invisible-watermark removal." ) diff --git a/tests/test_qwen_zimage_pipeline.py b/tests/test_qwen_zimage_pipeline.py index 2e9ffe8..88508b1 100644 --- a/tests/test_qwen_zimage_pipeline.py +++ b/tests/test_qwen_zimage_pipeline.py @@ -1251,3 +1251,62 @@ def test_chroma_google_face_content_gets_the_lower_adaptive_floor(): assert resolve_strength(0.20, "google", "chroma-zimage", face_count=17) == pytest.approx(0.20) # Other profiles are untouched by the face_count parameter. assert resolve_strength(None, "google", "qwen-zimage", size=(2000, 1850), face_count=17) == pytest.approx(0.27) + + +def test_auto_profile_routes_to_the_measured_engine_per_vendor(): + """--pipeline auto picks the engine the four-cohort calibration supports: + chroma-zimage for OpenAI/Microsoft, qwen-zimage for Google/Meta/unknown.""" + from remove_ai_watermarks._internal.watermark_profiles import ( + CHROMA_ZIMAGE_PROFILE, + QWEN_ZIMAGE_PROFILE, + normalize_profile, + resolve_auto_profile, + ) + + assert normalize_profile("auto") == "auto" + assert resolve_auto_profile("openai") == CHROMA_ZIMAGE_PROFILE + assert resolve_auto_profile("microsoft") == CHROMA_ZIMAGE_PROFILE + assert resolve_auto_profile("google") == QWEN_ZIMAGE_PROFILE + assert resolve_auto_profile("meta") == QWEN_ZIMAGE_PROFILE + assert resolve_auto_profile(None) == QWEN_ZIMAGE_PROFILE + assert resolve_auto_profile("unknown") == QWEN_ZIMAGE_PROFILE + + +def test_auto_profile_strength_uses_the_resolved_engine_floors(tmp_path, monkeypatch): + """After auto resolves to chroma-zimage for an OpenAI file, the strength + must be chroma's 0.09 floor, not qwen's 0.07675.""" + from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover + + _mock_watermark_runtime_deps(monkeypatch) + source = tmp_path / "source.png" + Image.new("RGB", (64, 48)).save(source) + + remover = WatermarkRemover(device="cuda", pipeline="auto") + assert remover.model_profile == "auto" + + runtime = MagicMock() + runtime.run.return_value = Image.new("RGB", (64, 48), (50, 60, 70)) + monkeypatch.setattr(remover, "_load_qwen_zimage_pipeline", lambda: runtime) + + remover.remove_watermark(source, vendor="openai") + + # Auto resolved to chroma-zimage for the OpenAI vendor. + assert remover.model_profile == "chroma-zimage" + _, kwargs = runtime.run.call_args + assert kwargs["strength"] == pytest.approx(0.09) + + +def test_auto_profile_text_manifest_stays_on_qwen(tmp_path, monkeypatch): + """Auto routes OpenAI to chroma-zimage, which rejects text manifests; + the rejection must reach the caller at the boundary.""" + from remove_ai_watermarks._internal.text_restoration import VerifiedTextLine, VerifiedTextManifest + from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover + + _mock_watermark_runtime_deps(monkeypatch) + source = tmp_path / "source.png" + Image.new("RGB", (96, 80)).save(source) + manifest = VerifiedTextManifest("0" * 64, 96, 80, (VerifiedTextLine((4, 4, 20, 16), "x", "alphabetic"),)) + + remover = WatermarkRemover(device="cuda", pipeline="auto") + with pytest.raises(ValueError, match="only by the qwen-zimage profile"): + remover.remove_watermark(source, vendor="openai", text_manifest=manifest)