From 91aa49b03a170396b6130b555f7cd4210fe784cc Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Wed, 19 Aug 2026 21:33:22 -0700 Subject: [PATCH] Allow verified-text restoration on tiled images Tile the Qwen VAE donor the same way as the global pass. Glyph restore already runs on the blended full frame, so the old tile ban was an artificial gate. --- docs/python-api.md | 5 ++- .../_internal/qwen_zimage_pipeline.py | 19 ++++++++- src/remove_ai_watermarks/invisible_engine.py | 8 ++-- tests/test_invisible_engine.py | 5 ++- tests/test_qwen_zimage_pipeline.py | 41 +++++++++++++++++++ 5 files changed, 70 insertions(+), 8 deletions(-) diff --git a/docs/python-api.md b/docs/python-api.md index 01edd9b..81351a5 100644 --- a/docs/python-api.md +++ b/docs/python-api.md @@ -609,8 +609,9 @@ engine.remove_watermark( Install `remove-ai-watermarks[text-restoration]`. The manifest schema and safety constraints are documented in the CLI guide. The engine verifies its decoded RGB -hash before loading the diffusion models and rejects SDXL, tiling, downscaling, and -postprocessing combinations that were not evaluated. `InvisibleOptions` exposes the +hash before loading the diffusion models and rejects SDXL, downscaling, and +postprocessing combinations that were not evaluated. Tiling is allowed: the VAE +donor uses the same overlapping tiles as the global pass. `InvisibleOptions` exposes the same field for `remove_all`; after a visible-stage edit, the manifest must be built against the staged pixels rather than the pristine source. diff --git a/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py b/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py index b441552..9cda34e 100644 --- a/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py +++ b/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py @@ -1079,7 +1079,24 @@ class QwenZImagePipeline: donor = None if text_manifest is not None: self._progress("Reconstructing the verified text donor with the Qwen VAE...") - donor = self._qwen_vae_roundtrip(image) + if tile and max(image.size) > tile_size: + from remove_ai_watermarks._internal.tiling import run_tiled + + donor = run_tiled( + self._qwen_vae_roundtrip, + image, + tile_size, + tile_overlap, + lambda message: self._progress( + message.replace( + "Tiled diffusion", + "Reconstructing the verified text donor", + 1, + ) + ), + ) + else: + donor = self._qwen_vae_roundtrip(image) global_strength = ( resolution_adaptive_denoise(image.width, image.height) if strength is None else float(strength) ) diff --git a/src/remove_ai_watermarks/invisible_engine.py b/src/remove_ai_watermarks/invisible_engine.py index 93c75bc..cdf9c1e 100644 --- a/src/remove_ai_watermarks/invisible_engine.py +++ b/src/remove_ai_watermarks/invisible_engine.py @@ -186,8 +186,10 @@ class InvisibleEngine: text_manifest: Operator-verified text lines bound to the decoded source pixels. Enables the experimental Qwen-VAE ``vae-glyphs`` post-pass. Requires the ``text-restoration`` extra and the ``qwen-zimage`` - profile. Incompatible with tiling, downscaling, humanize, unsharp, - and adaptive polish because those combinations are not calibrated. + profile. Incompatible with downscaling, humanize, unsharp, and + adaptive polish. Tiling is supported: the VAE donor uses the same + overlapping tiles as the global pass, then glyph restore runs on + the blended full frame. fidelity_anchor: Blend 15% of the Qwen-VAE donor across the whole frame before glyph restoration. OFF by default since 0.27.1: that global blend was measured to return detector-visible OpenAI SynthID on @@ -211,8 +213,6 @@ class InvisibleEngine: raise ValueError("--text-manifest is supported only by the qwen-zimage profile") if max_resolution != 0: raise ValueError("--text-manifest requires --max-resolution 0") - if tile: - raise ValueError("--text-manifest is not calibrated with --tile") if humanize > 0.0 or unsharp > 0.0 or adaptive_polish: raise ValueError("--text-manifest requires humanize=0, unsharp=0, and adaptive polish disabled") from remove_ai_watermarks import region_eraser diff --git a/tests/test_invisible_engine.py b/tests/test_invisible_engine.py index 320fa9e..b7fc46b 100644 --- a/tests/test_invisible_engine.py +++ b/tests/test_invisible_engine.py @@ -61,7 +61,6 @@ class TestVerifiedTextMode: cases = ( ("sdxl-zimage", {}, "qwen-zimage"), ("qwen-zimage", {"max_resolution": 1024}, "max-resolution 0"), - ("qwen-zimage", {"tile": True}, "not calibrated"), ("qwen-zimage", {"humanize": 1.0}, "humanize=0"), ("qwen-zimage", {"adaptive_polish": True}, "polish disabled"), ) @@ -129,6 +128,10 @@ class TestVerifiedTextMode: assert seen["fidelity_anchor"] is True + engine.remove_watermark(source, output, text_manifest=manifest, tile=True) + + assert seen["tile"] is True + class TestNativeOutputSize: """Model-side latent-grid rounding must not change the public output size.""" diff --git a/tests/test_qwen_zimage_pipeline.py b/tests/test_qwen_zimage_pipeline.py index 46a8cf4..3e4775c 100644 --- a/tests/test_qwen_zimage_pipeline.py +++ b/tests/test_qwen_zimage_pipeline.py @@ -659,6 +659,47 @@ def test_no_face_path_still_runs_verified_text_restoration(monkeypatch): restore.assert_called_with(source, anchor, donor, manifest.lines) +def test_tiled_verified_text_runs_vae_donor_per_tile(monkeypatch): + from remove_ai_watermarks._internal import qwen_zimage_pipeline, text_restoration + from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline + from remove_ai_watermarks._internal.text_restoration import VerifiedTextLine, VerifiedTextManifest + + pipeline = object.__new__(QwenZImagePipeline) + pipeline.device = "cuda" + pipeline.progress_callback = None + source = Image.new("RGB", (96, 80), (10, 20, 30)) + restored = Image.new("RGB", (96, 80), (130, 140, 150)) + pipeline._qwen_vae_roundtrip = MagicMock(side_effect=lambda tile: tile) + pipeline._run_global = MagicMock(side_effect=lambda tile, _strength, _seed: tile) + monkeypatch.setattr(qwen_zimage_pipeline, "detect_faces", lambda _image: []) + restore = MagicMock(return_value=restored) + monkeypatch.setattr(text_restoration, "restore_verified_text", restore) + manifest = VerifiedTextManifest( + "0" * 64, + 96, + 80, + (VerifiedTextLine((4, 4, 20, 16), "Exact", "alphabetic"),), + ) + + result = pipeline.run( + source, + strength=0.1, + seed=0, + tile=True, + tile_size=64, + tile_overlap=16, + text_manifest=manifest, + ) + + assert result is restored + assert pipeline._qwen_vae_roundtrip.call_count > 1 + assert pipeline._run_global.call_count > 1 + restore.assert_called_once() + assert restore.call_args.args[0].size == (96, 80) + assert restore.call_args.args[1].size == (96, 80) + assert restore.call_args.args[2].size == (96, 80) + + def test_watermark_remover_dispatches_to_full_pipeline(tmp_path, monkeypatch): from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover