mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-06 22:18:36 +02:00
Add global-only preload for qwen-zimage
This commit is contained in:
@@ -73,7 +73,7 @@ Compact map. The full per-module detail (design decisions, tuned thresholds, cal
|
||||
- `invisible_watermark.py` — decodes the OPEN DWT-DCT watermarks (SD / SDXL / FLUX) via `imwatermark` (extra `detect`, pulls torch). Fragile two ways: (1) does not survive JPEG re-encode/resize; (2) **carrier-fragile on a broad class of pristine images** -- a clean encode->decode round-trip recovers 48/48 on chatgpt/firefly/random but FAILS (28-39/48, below the `_MATCH_48`=44 gate) on the FLUX fox, doubao, a flat FLUX generation, AND a clean synthetic flat fill with no watermark. The failure does NOT track texture; it goes with a degenerate **all-ones decode that is a CARRIER ARTIFACT, not a watermark** (synthetic clean image reproduces it). So `detect_invisible_watermark` is **positive-only**: trust a hit; a `None` is inconclusive unless a same-carrier positive-control embed first recovers >=44. Verified 2026-06-19; full caveat in `docs/watermarking-landscape.md`.
|
||||
- `trustmark_detector.py` — Adobe TrustMark open decoder (extra `trustmark`). Do NOT remove the JPEG re-encode false-positive gate — a lone TrustMark hit without it is almost always content noise.
|
||||
- `noai/watermark_remover.py` — `WatermarkRemover` with four diffusion pipelines selected by the explicit `pipeline` ctor arg, never inferred from `model_id`: `sdxl` (plain SDXL img2img), `controlnet` (SDXL + canny ControlNet, **the compatibility and cost DEFAULT since 2026-06-09**), `qwen` (Qwen-Image 20B img2img), and `qwen-zimage` (delegates to the fixed two-stage runtime below). Removal comes from img2img strength. Both SDXL loaders pass `add_watermarker=False`; diffusers otherwise re-stamps an open SDXL DWT-DCT watermark. Qwen's certified floors and fidelity results remain as documented below. The base `qwen` profile stays the manual text lane; `qwen-zimage` is the recommended high-quality manual mode, especially for face identity, while remaining experimental rather than an auto-router.
|
||||
- `noai/qwen_zimage_pipeline.py` — CUDA-only Qwen-Image-2512 Lightning + DiffSynth Canny full-frame regeneration, followed by YuNet face boxes, SAM masks, and Z-Image Turbo regeneration from original face crops. Ports both adaptive denoise formulas from Synthid-Bypass v2. The active upstream face path is YOLO + SAM; this port keeps its center-point and box prompts, proposal selection, detector-box intersection, crop factor, and paste feather while replacing YOLO with YuNet to avoid an AGPL runtime. The port is architectural, not bit-identical: it uses full safetensors instead of GGUF, DiffSynth samplers instead of the Comfy sampler pairs, and no latent detailer feather. DiffSynth input pixels, Canny control, and explicit dimensions must share the same /16 grid. SAM pixels follow the model dtype, geometric prompts stay float32, and bfloat16 outputs convert through float32 before NumPy. The YuNet download verifies its SHA-256. Separate `qwen-zimage` extra; fixed four-step global and eight-step face schedules; no custom `--model`; `--tile` applies only to the global stage, followed by one full-frame face stage; CLI adaptive polish defaults off. GPUs with at least 64 GiB VRAM keep the face stack resident, while smaller devices retain CPU offload. Fixed prompt embeddings are cached only when they do not depend on an edit image. The exact seed-0 release candidate passed the corresponding provider-oracle checks; broader seeded text, face, and tiled-output certification remains open.
|
||||
- `noai/qwen_zimage_pipeline.py` — CUDA-only Qwen-Image-2512 Lightning + DiffSynth Canny full-frame regeneration, followed by YuNet face boxes, SAM masks, and Z-Image Turbo regeneration from original face crops. Ports both adaptive denoise formulas from Synthid-Bypass v2. The active upstream face path is YOLO + SAM; this port keeps its center-point and box prompts, proposal selection, detector-box intersection, crop factor, and paste feather while replacing YOLO with YuNet to avoid an AGPL runtime. The port is architectural, not bit-identical: it uses full safetensors instead of GGUF, DiffSynth samplers instead of the Comfy sampler pairs, and no latent detailer feather. DiffSynth input pixels, Canny control, and explicit dimensions must share the same /16 grid. SAM pixels follow the model dtype, geometric prompts stay float32, and bfloat16 outputs convert through float32 before NumPy. The YuNet download verifies its SHA-256. Separate `qwen-zimage` extra; fixed four-step global and eight-step face schedules; no custom `--model`; `--tile` applies only to the global stage, followed by one full-frame face stage; CLI adaptive polish defaults off. `InvisibleEngine.preload(global_only=True)` warms Qwen and YuNet while leaving Z-Image and SAM lazy until a face is detected; the default `preload()` remains a full preload. GPUs with at least 64 GiB VRAM keep the face stack resident, while smaller devices retain CPU offload. Fixed prompt embeddings are cached only when they do not depend on an edit image. The exact seed-0 release candidate passed the corresponding provider-oracle checks; broader seeded text, face, and tiled-output certification remains open.
|
||||
- `noai/tiling.py` — sliding-window tiled diffusion for large inputs (CLI `--tile`). The SDXL, ControlNet, and base Qwen paths branch to `run_tiled` when `tile` is set AND the long side exceeds `tile_size`, refactoring the single-pass `_generate` into a per-tile `_generate_one` (the ControlNet edge map is rebuilt per tile inside it). `qwen-zimage` instead calls `run_tiled` only around its global Qwen stage, blends the tiles, then runs one full-frame face stage. Pure helpers `plan_tiles` (uniform-size tiles, last one flush to the edge) and `feather_weights` (strictly-positive separable taper -> partition-of-unity blend) are unit-tested without the model. Also home to `feather_region_composite(base, regenerated, box, *, feather)` — the pure region-targeted compositor for **AI-enhanced composites** (`ai_source_kind == "enhanced"`): blends the regenerated AI box back over the original with a feathered seam, leaving the real photo OUTSIDE the box pixel-exact. It backs `WatermarkRemover.remove_watermark(region=...)` (regenerate ONLY the AI region, not the whole frame); the no-model lossless region path stays `region_eraser.erase`. New tile/region-blend tuning goes in these pure helpers; do not inline blend math into the runner.
|
||||
- `auto_config.py` + the content-detection layer were REMOVED 2026-06-09; `--auto` is a deprecated no-op (controlnet is the default pipeline and adaptive polish is ON by default for the original profiles, while `qwen-zimage` leaves it off to preserve the upstream two-stage output).
|
||||
- `upscaler.py` — optional Real-ESRGAN pre-diffusion super-resolution for small inputs (extra `esrgan`, spandrel only). Manual opt-in; the default `--upscaler` stays `lanczos` and the engine always falls back to Lanczos on absence/error. ESRGAN can degrade faces and thin text.
|
||||
|
||||
@@ -266,6 +266,10 @@ is the source of truth for:
|
||||
- the minimum viable step calculation.
|
||||
|
||||
The current profiles are `controlnet`, `sdxl`, `qwen`, and `qwen-zimage`.
|
||||
For serverless cold starts, `InvisibleEngine.preload(global_only=True)` loads the
|
||||
mandatory Qwen stage and YuNet while leaving the optional Z-Image and SAM face
|
||||
stack lazy until a face is detected. The default `preload()` still loads every
|
||||
stage.
|
||||
`default` is a legacy alias for `sdxl`. There is no content-dependent automatic
|
||||
router.
|
||||
|
||||
|
||||
@@ -139,9 +139,14 @@ class InvisibleEngine:
|
||||
)
|
||||
self._progress_callback = progress_callback
|
||||
|
||||
def preload(self) -> None:
|
||||
"""Eagerly load the pipeline so download progress is visible."""
|
||||
self._remover.preload()
|
||||
def preload(self, *, global_only: bool = False) -> None:
|
||||
"""Eagerly load the pipeline so download progress is visible.
|
||||
|
||||
For ``qwen-zimage``, ``global_only=True`` loads the mandatory Qwen stage
|
||||
and leaves the optional Z-Image and SAM face stack lazy until a face is
|
||||
detected. Other profiles have no optional stage and ignore the flag.
|
||||
"""
|
||||
self._remover.preload(global_only=global_only)
|
||||
|
||||
def _esrgan_upscale(self, image: Any, target: tuple[int, int]) -> Any:
|
||||
"""Upscale a PIL image to ``target`` with Real-ESRGAN, else Lanczos.
|
||||
|
||||
@@ -767,12 +767,13 @@ class QwenZImagePipeline:
|
||||
log.warning("SAM face-mask refinement failed (%s); using box-derived ellipse masks", exc)
|
||||
return _ellipse_masks(boxes, image.size)
|
||||
|
||||
def preload(self) -> None:
|
||||
"""Eagerly load both diffusion stages and the face segmentation model."""
|
||||
def preload(self, *, global_only: bool = False) -> None:
|
||||
"""Eagerly load the mandatory stage and, by default, the face stack."""
|
||||
self._load_qwen()
|
||||
self._load_zimage()
|
||||
self._load_sam()
|
||||
_yunet_model_path()
|
||||
if not global_only:
|
||||
self._load_zimage()
|
||||
self._load_sam()
|
||||
|
||||
def _run_global(self, image: Image.Image, strength: float, seed: int | None) -> Image.Image:
|
||||
pipe, controlnet_input_cls = self._load_qwen()
|
||||
|
||||
@@ -442,10 +442,13 @@ class WatermarkRemover:
|
||||
|
||||
# ── Preload ──────────────────────────────────────────────────────
|
||||
|
||||
def preload(self) -> None:
|
||||
"""Eagerly load the pipeline so download progress bars are visible."""
|
||||
def preload(self, *, global_only: bool = False) -> None:
|
||||
"""Eagerly load the pipeline so download progress bars are visible.
|
||||
|
||||
``global_only`` applies to qwen-zimage, whose face stage is optional.
|
||||
"""
|
||||
if self.model_profile == QWEN_ZIMAGE_PROFILE:
|
||||
self._load_qwen_zimage_pipeline().preload()
|
||||
self._load_qwen_zimage_pipeline().preload(global_only=global_only)
|
||||
elif self.model_profile == "qwen":
|
||||
self._load_qwen_pipeline()
|
||||
elif self.model_profile == "controlnet":
|
||||
|
||||
@@ -35,6 +35,14 @@ class TestInvisibleEngineInit:
|
||||
# SDXL base became the default in May 2026 (defeats SynthID v2).
|
||||
assert InvisibleEngine.DEFAULT_MODEL_ID == "stabilityai/stable-diffusion-xl-base-1.0"
|
||||
|
||||
def test_preload_forwards_global_only(self):
|
||||
engine = object.__new__(InvisibleEngine)
|
||||
engine._remover = SimpleNamespace(preload=lambda **kwargs: setattr(engine, "_preload_kwargs", kwargs))
|
||||
|
||||
engine.preload(global_only=True)
|
||||
|
||||
assert engine._preload_kwargs == {"global_only": True}
|
||||
|
||||
|
||||
class TestNativeOutputSize:
|
||||
"""Model-side latent-grid rounding must not change the public output size."""
|
||||
|
||||
@@ -440,6 +440,67 @@ def test_qwen_tiling_runs_global_tiles_then_one_full_frame_face_stage(monkeypatc
|
||||
assert result.size == image.size
|
||||
|
||||
|
||||
def test_global_only_preload_skips_face_models(monkeypatch):
|
||||
from remove_ai_watermarks.noai.qwen_zimage_pipeline import QwenZImagePipeline
|
||||
|
||||
runtime = QwenZImagePipeline(device="cuda", torch_dtype="bf16")
|
||||
qwen = MagicMock()
|
||||
zimage = MagicMock()
|
||||
sam = MagicMock()
|
||||
yunet = MagicMock()
|
||||
monkeypatch.setattr(runtime, "_load_qwen", qwen)
|
||||
monkeypatch.setattr(runtime, "_load_zimage", zimage)
|
||||
monkeypatch.setattr(runtime, "_load_sam", sam)
|
||||
monkeypatch.setattr(
|
||||
"remove_ai_watermarks.noai.qwen_zimage_pipeline._yunet_model_path",
|
||||
yunet,
|
||||
)
|
||||
|
||||
runtime.preload(global_only=True)
|
||||
|
||||
qwen.assert_called_once_with()
|
||||
zimage.assert_not_called()
|
||||
sam.assert_not_called()
|
||||
yunet.assert_called_once_with()
|
||||
|
||||
|
||||
def test_full_preload_still_loads_face_models(monkeypatch):
|
||||
from remove_ai_watermarks.noai.qwen_zimage_pipeline import QwenZImagePipeline
|
||||
|
||||
runtime = QwenZImagePipeline(device="cuda", torch_dtype="bf16")
|
||||
qwen = MagicMock()
|
||||
zimage = MagicMock()
|
||||
sam = MagicMock()
|
||||
yunet = MagicMock()
|
||||
monkeypatch.setattr(runtime, "_load_qwen", qwen)
|
||||
monkeypatch.setattr(runtime, "_load_zimage", zimage)
|
||||
monkeypatch.setattr(runtime, "_load_sam", sam)
|
||||
monkeypatch.setattr(
|
||||
"remove_ai_watermarks.noai.qwen_zimage_pipeline._yunet_model_path",
|
||||
yunet,
|
||||
)
|
||||
|
||||
runtime.preload()
|
||||
|
||||
qwen.assert_called_once_with()
|
||||
zimage.assert_called_once_with()
|
||||
sam.assert_called_once_with()
|
||||
yunet.assert_called_once_with()
|
||||
|
||||
|
||||
def test_watermark_remover_forwards_global_only_preload(monkeypatch):
|
||||
from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover
|
||||
|
||||
runtime = MagicMock()
|
||||
remover = WatermarkRemover.__new__(WatermarkRemover)
|
||||
remover.model_profile = "qwen-zimage"
|
||||
monkeypatch.setattr(remover, "_load_qwen_zimage_pipeline", lambda: runtime)
|
||||
|
||||
remover.preload(global_only=True)
|
||||
|
||||
runtime.preload.assert_called_once_with(global_only=True)
|
||||
|
||||
|
||||
def test_qwen_zimage_rejects_runtime_knobs_that_change_fixed_graph(tmp_path, monkeypatch):
|
||||
from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover
|
||||
|
||||
|
||||
Reference in New Issue
Block a user