CalvinBackup/remove-ai-watermarks
1
0
Fork 0
mirror of https://github.com/wiltodelta/remove-ai-watermarks.git synced 2026-08-07 14:38:35 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat(invisible): add Qwen-Image img2img pipeline (--pipeline qwen)

Browse Source 
A third diffusion pipeline alongside sdxl/controlnet: Qwen-Image (20B MMDiT,
Apache-2.0 code AND weights) img2img. The scrub still comes from the img2img
strength; Qwen preserves text (incl. CJK) and structure markedly better than
SDXL at the scrub floor, so it over-regenerates real photos far less (directly
targets the controlnet over-regeneration that degrades real uploads).

- watermark_profiles: QWEN_MODEL_ID, normalize_profile accepts "qwen".
- WatermarkRemover: _load_qwen_pipeline (bf16, loads Qwen base unless --model
  overridden, clear ImportError if diffusers lacks the class), _run_qwen (no
  MPS fallback -- 20B is CUDA/cloud-class), dispatch in _generate_one/preload,
  pure _build_qwen_kwargs (true_cfg_scale, not guidance_scale).
- Shared _base_load_kwargs() across all three loaders (dtype + token).
- CLI --pipeline gains "qwen"; invisible_engine threads it through.
- scripts/qwen_scrub_prototype.py: standalone PEP 723 GPU experiment.

Prototype oracle floors (Modal A100-80GB, single seed, controls SynthID-positive,
PENDING seed-repeat cert): OpenAI clears at strength ~0.10, Gemini at ~0.30 (0.20
still detected), with CJK text + faces faithful where controlnet plasticizes. The
Gemini floor is higher than the shared default ladder, so pass an explicit
--strength for Gemini on this pipeline until a Qwen-specific ladder is certified.

The model-running path is CUDA-only (untestable locally); unit tests cover the
pure call-shape (_build_qwen_kwargs) and profile normalization without torch.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Victor Kuznetsov
2026-06-19 20:44:36 -07:00
co-authored by Claude Opus 4.8
parent 0c0c6c6b03
commit 76e3d4154c
10 changed files with 309 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/remove_ai_watermarks/cli.py
+7 -6
View File
@@ -253,15 +253,16 @@ def _normalize_pipeline(ctx: click.Context, param: click.Parameter, value: str |
return normalized
# ``controlnet`` (the default-SELECTED value) and ``sdxl`` (plain SDXL img2img) are the
# two current profiles; ``default`` is an OUTDATED back-compat alias for ``sdxl``
# (warned + normalized away by _normalize_pipeline).
_PIPELINE_CHOICES = ["sdxl", "controlnet", "default"]
# ``controlnet`` (the default-SELECTED value), ``sdxl`` (plain SDXL img2img) and
# ``qwen`` (Qwen-Image, CUDA/cloud-class) are the current profiles; ``default`` is an
# OUTDATED back-compat alias for ``sdxl`` (warned + normalized away by _normalize_pipeline).
_PIPELINE_CHOICES = ["sdxl", "controlnet", "qwen", "default"]
_PIPELINE_HELP = (
"Pipeline profile. controlnet (DEFAULT) = SDXL + canny ControlNet that preserves "
"text/faces via edge conditioning while removing SynthID; sdxl = plain SDXL img2img "
"(lighter, no extra model download, but leaves SynthID on flat-graphic content). "
"('default' is an OUTDATED alias for 'sdxl' -- use sdxl or controlnet.)"
"(lighter, no extra model download, but leaves SynthID on flat-graphic content); "
"qwen = Qwen-Image (20B, Apache-2.0) img2img, best text/structure preservation but "
"CUDA/cloud-class (does not fit MPS). ('default' is an OUTDATED alias for 'sdxl'.)"
)
# Shared --pipeline / --strength decorators so the three diffusion commands
src/remove_ai_watermarks/invisible_engine.py
+3 -2
View File
@@ -103,8 +103,9 @@ class InvisibleEngine:
device: Device for inference (auto/cpu/mps/cuda/xpu). None = auto.
pipeline: Pipeline profile. "controlnet" (DEFAULT; SDXL + canny ControlNet
that preserves text/face structure via edge conditioning while removing
SynthID) or "sdxl" (plain SDXL img2img, lighter but leaves SynthID on
flat-graphic content). "default" is a back-compat alias for "sdxl".
SynthID), "sdxl" (plain SDXL img2img, lighter but leaves SynthID on
flat-graphic content), or "qwen" (Qwen-Image 20B img2img, best text/
structure preservation but CUDA/cloud-class). "default" aliases "sdxl".
hf_token: HuggingFace API token.
progress_callback: Optional callback for progress messages.
controlnet_conditioning_scale: ControlNet structure-preservation
src/remove_ai_watermarks/noai/watermark_profiles.py
+11
View File
@@ -12,6 +12,17 @@ if TYPE_CHECKING:
DEFAULT_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
# Qwen-Image (20B MMDiT, Apache-2.0 code AND weights) base for the ``qwen`` pipeline:
# an img2img alternative to SDXL with native text rendering (incl. CJK). Loaded only
# when ``--pipeline qwen`` is selected; CUDA/cloud-class (does not fit MPS). Prototype
# oracle floors (single-seed, 2026-06-19, pending seed-repeat cert): OpenAI clears at
# strength ~0.10, Google/Gemini at ~0.30 (0.20 still detected) -- the latter is HIGHER
# than the certified controlnet Google floor (0.15), so pass an explicit ``--strength``
# for Gemini content on this pipeline until a Qwen-specific ladder is certified.
# (Dispatch uses the bare "qwen" literal, matching the sdxl/controlnet sites, so there
# is no QWEN_PROFILE constant -- only the model id is referenced from code.)
QWEN_MODEL_ID = "Qwen/Qwen-Image"
# Canonical pipeline-profile names + the back-compat alias. The plain SDXL img2img
# profile is ``sdxl``; ``default`` is kept as an accepted alias (it was the profile's
# name before ``controlnet`` became the default-selected pipeline, 2026-06-09).
src/remove_ai_watermarks/noai/watermark_remover.py
+116 -12
View File
@@ -1,6 +1,14 @@
"""Watermark removal using diffusion model regeneration attack.
Two pipelines:
Three pipelines (selected by the explicit ``pipeline`` ctor arg):
0. ``qwen`` -- Qwen-Image (20B MMDiT, Apache-2.0) img2img. The scrub still comes from
the img2img ``strength``; Qwen preserves text (incl. CJK) and structure markedly
better than SDXL at the scrub floor, so it over-regenerates real photos far less.
CUDA/cloud-class (does not fit MPS). See ``watermark_profiles`` for the prototype
oracle floors (pending seed-repeat cert).
Two SDXL pipelines:
1. ``controlnet`` (DEFAULT) -- SDXL img2img with a canny ControlNet. The watermark
REMOVAL still comes from the img2img regeneration (``strength``); the ControlNet
only PRESERVES structure (text/faces) by conditioning on the edge map. No original
@@ -36,6 +44,7 @@ from remove_ai_watermarks.noai.watermark_profiles import (
CONTROLNET_CANNY_MODEL,
DEFAULT_MODEL_ID,
DEFAULT_STRENGTH,
QWEN_MODEL_ID,
normalize_profile,
resolve_strength,
)
@@ -308,6 +317,29 @@ _CANNY_HIGH = 200
_CONTROLNET_PROMPT = "best quality, high quality, sharp, detailed, photographic"
_CONTROLNET_NEGATIVE = "blurry, lowres, deformed, distorted text, garbled text, watermark, jpeg artifacts"
# Neutral prompts for the Qwen-Image img2img pass (faithful regeneration, not an edit).
_QWEN_PROMPT = "high quality, sharp, detailed, faithful to the original"
_QWEN_NEGATIVE = "blurry, lowres, distorted text, garbled text, artifacts"
def _build_qwen_kwargs(
image: Image.Image, strength: float, num_inference_steps: int, true_cfg_scale: float, generator: Any
) -> dict[str, Any]:
"""Build the QwenImageImg2ImgPipeline call kwargs (pure; unit-tested without torch).
Qwen-Image uses ``true_cfg_scale`` (not SDXL's ``guidance_scale``) and takes an
explicit ``negative_prompt``; the scrub still comes from the img2img ``strength``.
"""
return {
"prompt": _QWEN_PROMPT,
"negative_prompt": _QWEN_NEGATIVE,
"image": image,
"strength": strength,
"num_inference_steps": num_inference_steps,
"true_cfg_scale": true_cfg_scale,
"generator": generator,
}
class WatermarkRemover:
"""Remove watermarks from images using diffusion model regeneration.
@@ -348,6 +380,11 @@ class WatermarkRemover:
if torch_dtype is None:
if self.device == "cpu" or self.device == "mps":
self.torch_dtype = torch.float32 # type: ignore
elif self.model_profile == "qwen":
# Qwen-Image is published in bf16; fp16 risks overflow on the 20B MMDiT.
# cuda/xpu-only by construction: the cpu/mps guard above already forced
# fp32, and the 20B model does not fit MPS anyway.
self.torch_dtype = torch.bfloat16 # type: ignore
else:
self.torch_dtype = torch.float16 # type: ignore
else:
@@ -355,6 +392,7 @@ class WatermarkRemover:
self._pipeline: AutoImg2ImgPipeline | None = None
self._controlnet_pipeline: Any = None
self._qwen_pipeline: Any = None
self._progress_callback = progress_callback
self.hf_token: str | None = hf_token or os.environ.get("HF_TOKEN")
@@ -369,7 +407,9 @@ class WatermarkRemover:
def preload(self) -> None:
"""Eagerly load the pipeline so download progress bars are visible."""
if self.model_profile == "controlnet":
if self.model_profile == "qwen":
self._load_qwen_pipeline()
elif self.model_profile == "controlnet":
self._load_controlnet_pipeline()
else:
self._load_pipeline()
@@ -420,19 +460,27 @@ class WatermarkRemover:
return pipeline
def _base_load_kwargs(self) -> dict[str, Any]:
"""The ``from_pretrained`` kwargs shared by all three loaders (dtype + token).
Each loader adds its own extras (SDXL safety_checker + fp16 VAE, the ControlNet
model, etc.). Centralizing the dtype/token pair avoids the drift trap of three
copies (a token forgotten on one loader silently breaks gated downloads there).
"""
load_kwargs: dict[str, Any] = {"torch_dtype": self.torch_dtype}
if self.hf_token:
load_kwargs["token"] = self.hf_token
return load_kwargs
def _load_pipeline(self) -> AutoImg2ImgPipeline:
"""Load the plain SDXL img2img pipeline lazily."""
if self._pipeline is None:
logger.info("Loading model %s on %s...", self.model_id, self.device)
self._set_progress(f"Loading model weights: {self.model_id}")
load_kwargs: dict[str, Any] = {
"torch_dtype": self.torch_dtype,
"safety_checker": None,
"requires_safety_checker": False,
}
if self.hf_token:
load_kwargs["token"] = self.hf_token
load_kwargs = self._base_load_kwargs()
load_kwargs["safety_checker"] = None
load_kwargs["requires_safety_checker"] = False
self._maybe_add_fp16_vae(load_kwargs)
pipeline = AutoImg2ImgPipeline.from_pretrained(self.model_id, **load_kwargs) # type: ignore
@@ -458,9 +506,8 @@ class WatermarkRemover:
self._set_progress(f"Loading ControlNet: {CONTROLNET_CANNY_MODEL}")
controlnet = ControlNetModel.from_pretrained(CONTROLNET_CANNY_MODEL, torch_dtype=self.torch_dtype)
load_kwargs: dict[str, Any] = {"controlnet": controlnet, "torch_dtype": self.torch_dtype}
if self.hf_token:
load_kwargs["token"] = self.hf_token
load_kwargs = self._base_load_kwargs()
load_kwargs["controlnet"] = controlnet
self._maybe_add_fp16_vae(load_kwargs)
self._set_progress(f"Loading model weights: {self.model_id}")
@@ -474,6 +521,37 @@ class WatermarkRemover:
return self._controlnet_pipeline
def _load_qwen_pipeline(self) -> Any:
"""Load the Qwen-Image img2img pipeline lazily.
Qwen-Image is its OWN base model (not an SDXL add-on), so it loads
``QWEN_MODEL_ID`` unless the caller passed a custom ``--model``. Needs a
diffusers build that ships ``QwenImageImg2ImgPipeline``; raises a clear error
otherwise. CUDA/cloud-class (the 20B MMDiT does not fit MPS).
"""
if self._qwen_pipeline is None:
try:
from diffusers import QwenImageImg2ImgPipeline
except ImportError as exc:
raise ImportError(
"The 'qwen' pipeline needs a diffusers version that ships "
"QwenImageImg2ImgPipeline. Upgrade: pip install -U diffusers"
) from exc
# Use the Qwen base unless the user explicitly overrode --model.
model = self.model_id if self.model_id != self.DEFAULT_MODEL_ID else QWEN_MODEL_ID
logger.info("Loading Qwen-Image (%s) on %s...", model, self.device)
self._set_progress(f"Loading model weights: {model}")
pipeline = QwenImageImg2ImgPipeline.from_pretrained(model, **self._base_load_kwargs())
pipeline = self._move_to_device_and_optimize(pipeline)
with contextlib.suppress(Exception):
pipeline.set_progress_bar_config(disable=True)
logger.info("Qwen-Image model loaded successfully")
self._qwen_pipeline = pipeline
return self._qwen_pipeline
# ── Core removal ─────────────────────────────────────────────────
def remove_watermark(
@@ -552,6 +630,8 @@ class WatermarkRemover:
_total_start = time.monotonic()
def _generate_one(img: Image.Image) -> Image.Image:
if self.model_profile == "qwen":
return self._run_qwen(img, strength, num_inference_steps, guidance_scale, generator)
if self.model_profile == "controlnet":
return self._run_controlnet(img, strength, num_inference_steps, guidance_scale, generator)
return self._run_img2img(img, strength, num_inference_steps, guidance_scale, generator)
@@ -725,6 +805,30 @@ class WatermarkRemover:
self._controlnet_pipeline = None
return self._load_controlnet_pipeline()
# ── Qwen runner ──────────────────────────────────────────────────
def _run_qwen(
self,
init_image: Image.Image,
strength: float,
num_inference_steps: int,
guidance_scale: float,
generator: Any,
) -> Image.Image:
"""Run the Qwen-Image img2img pass.
Removal comes from the img2img ``strength`` (same lever as the SDXL paths);
Qwen-Image preserves text/structure markedly better at the scrub floor. The
CLI ``guidance_scale`` maps to Qwen's ``true_cfg_scale`` (~4.0 is typical;
the SDXL default of 7.5 is high for Qwen). No MPS->CPU fallback: the 20B MMDiT
is CUDA/cloud-class and does not run on MPS, so an error here propagates.
"""
pipeline = self._load_qwen_pipeline()
self._set_progress(f"Running Qwen-Image img2img (strength={strength}, true_cfg={guidance_scale})...")
kwargs = _build_qwen_kwargs(init_image, strength, num_inference_steps, guidance_scale, generator)
result = pipeline(**kwargs)
return result.images[0]
# ── Batch ────────────────────────────────────────────────────────
def remove_watermark_batch(