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

feat: controlnet pipeline for text/face-structure preservation

Browse Source 
Add `--pipeline controlnet` (SDXL base + xinsir canny ControlNet via
StableDiffusionXLControlNetImg2ImgPipeline): the canny edge map conditions the
img2img regeneration so text and face STRUCTURE stay sharp, while the watermark
is still removed by the regeneration (`strength`) -- no original pixels are
copied or frozen, so SynthID does not survive. Oracle-verified clean on OpenAI
with better text/structure fidelity than plain img2img at equal strength.
`--controlnet-scale` tunes structure preservation; fp32 on mps/cpu (fp16-fixed
VAE on cuda/xpu). Shares the img2img runner (live progress + MPS->CPU fallback)
and the fp16-VAE-fix / device-move helpers with the default pipeline.

Remove the superseded subsystems -- ctrlregen (SD1.5 clean-noise),
text-protection (differential / region-hires) and face-protection: they either
destroyed real content or shielded the watermark by re-using original pixels.
controlnet replaces them by regenerating everything under edge conditioning.

Canny preserves face structure but not identity; face IDENTITY is a separate
face-restoration post-pass (CodeFormer/GFPGAN), researched + prototyped but not
yet shipped. An IP-Adapter FaceID attempt was built and removed (footgun: needs
high strength, corrupts faces at removal strength).

Docs: docs/controlnet-removal-pipeline-research.md, scripts/controlnet_sweep.py.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Victor Kuznetsov
2026-06-03 16:59:28 -07:00
co-authored by Claude Opus 4.8
parent 175609b60a
commit d90d5d886a
28 changed files with 1239 additions and 3541 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/remove_ai_watermarks/cli.py
+26 -39
View File
@@ -136,23 +136,14 @@ def _validate_image(path: Path) -> Path:
_ALPHA_FORMATS = {".png", ".webp"}
# Shared option decorators for commands that run the invisible-watermark pipeline.
# Both cmd_invisible and cmd_all expose these flags; defining them once avoids
# Shared option decorator for commands that run the invisible-watermark pipeline.
# Both cmd_invisible and cmd_all expose this flag; defining it once avoids
# copy-paste drift.
_protect_text_option = click.option(
"--protect-text",
is_flag=True,
default=False,
help=(
"Enable text region protection (experimental: re-scrubs text blocks at high resolution). "
"May prevent SynthID removal in text areas -- verify with oracle before relying on it."
),
)
_protect_faces_option = click.option(
"--protect-faces",
is_flag=True,
default=False,
help="Enable face protection (experimental: YOLO detect + blend original faces back).",
_controlnet_scale_option = click.option(
"--controlnet-scale",
type=float,
default=1.0,
help="ControlNet conditioning scale (structure/text preservation strength), controlnet pipeline only.",
)
@@ -453,14 +444,15 @@ def cmd_erase(
type=float,
default=None,
help="Denoising strength (0.0-1.0). Default: vendor-adaptive (OpenAI 0.10 / Google 0.15 / "
"unknown 0.15, from the C2PA issuer); ctrlregen uses 1.0.",
"unknown 0.15, from the C2PA issuer).",
)
@click.option("--steps", type=int, default=50, help="Number of denoising steps. Default: 50.")
@click.option(
"--pipeline",
type=click.Choice(["default", "ctrlregen"]),
type=click.Choice(["default", "controlnet"]),
default="default",
help="Pipeline profile (default=SDXL; ctrlregen=CtrlRegen, EXPERIMENTAL/destructive at clean-noise).",
help="Pipeline profile (default=SDXL img2img; controlnet=SDXL + canny ControlNet that preserves "
"text/faces via edge conditioning while removing SynthID).",
)
@click.option(
"--device",
@@ -479,8 +471,7 @@ def cmd_erase(
default=0,
help="Cap long side (px) before diffusion; 0 = native (best quality, like raiw.cc). Raise only on GPU/MPS OOM.",
)
@_protect_text_option
@_protect_faces_option
@_controlnet_scale_option
@click.pass_context
def cmd_invisible(
ctx: click.Context,
@@ -494,8 +485,7 @@ def cmd_invisible(
hf_token: str | None,
humanize: float,
max_resolution: int,
protect_text: bool,
protect_faces: bool,
controlnet_scale: float,
) -> None:
"""Remove invisible AI watermarks (SynthID, StableSignature, TreeRing).
@@ -526,6 +516,7 @@ def cmd_invisible(
pipeline=pipeline,
hf_token=hf_token,
progress_callback=progress_cb,
controlnet_conditioning_scale=controlnet_scale,
)
# Detect the SynthID vendor from the ORIGINAL (before processing strips C2PA) so the
@@ -533,7 +524,7 @@ def cmd_invisible(
vendor = vendor_for_strength(source)
console.print(f" Input: {source.name}")
console.print(f" Pipeline: {pipeline}")
console.print(f" Strength: {resolve_strength(strength, pipeline, vendor)} Steps: {steps}")
console.print(f" Strength: {resolve_strength(strength, vendor)} Steps: {steps}")
t0 = time.monotonic()
result_path = engine.remove_watermark(
@@ -544,8 +535,6 @@ def cmd_invisible(
guidance_scale=None,
seed=seed,
humanize=humanize,
protect_text=protect_text,
protect_faces=protect_faces,
max_resolution=max_resolution,
vendor=vendor,
)
@@ -694,15 +683,15 @@ def cmd_identify(ctx: click.Context, source: Path, no_visible: bool, as_json: bo
"--strength",
type=float,
default=None,
help="Invisible watermark denoising strength. Default: vendor-adaptive "
"(OpenAI 0.10 / Google 0.15 / unknown 0.15); ctrlregen uses 1.0.",
help="Invisible watermark denoising strength. Default: vendor-adaptive (OpenAI 0.10 / Google 0.15 / unknown 0.15).",
)
@click.option("--steps", type=int, default=50, help="Number of denoising steps for invisible removal.")
@click.option(
"--pipeline",
type=click.Choice(["default", "ctrlregen"]),
type=click.Choice(["default", "controlnet"]),
default="default",
help="Pipeline profile (default=SDXL; ctrlregen=CtrlRegen, EXPERIMENTAL/destructive at clean-noise).",
help="Pipeline profile (default=SDXL img2img; controlnet=SDXL + canny ControlNet that preserves "
"text/faces via edge conditioning while removing SynthID).",
)
@click.option("--model", type=str, default=None, help="HuggingFace model ID for invisible removal.")
@click.option(
@@ -722,8 +711,7 @@ def cmd_identify(ctx: click.Context, source: Path, no_visible: bool, as_json: bo
default=0,
help="Cap long side (px) before diffusion; 0 = native (best quality, like raiw.cc). Raise only on GPU/MPS OOM.",
)
@_protect_text_option
@_protect_faces_option
@_controlnet_scale_option
@click.pass_context
def cmd_all(
ctx: click.Context,
@@ -740,8 +728,7 @@ def cmd_all(
hf_token: str | None,
humanize: float,
max_resolution: int,
protect_text: bool,
protect_faces: bool,
controlnet_scale: float,
) -> None:
"""Remove ALL watermarks: visible + invisible + metadata.
@@ -822,13 +809,14 @@ def cmd_all(
pipeline=pipeline,
hf_token=hf_token,
progress_callback=progress_cb,
controlnet_conditioning_scale=controlnet_scale,
)
# Detect the vendor from the pristine ORIGINAL (`source`); `tmp_path` has
# already lost its C2PA to the visible-removal pass, so reading it would
# always resolve to the unknown-vendor default.
vendor = vendor_for_strength(source)
console.print(f" Strength: {resolve_strength(strength, pipeline, vendor)} Steps: {steps}")
console.print(f" Strength: {resolve_strength(strength, vendor)} Steps: {steps}")
inv_engine.remove_watermark(
image_path=tmp_path,
output_path=tmp_path,
@@ -836,8 +824,6 @@ def cmd_all(
num_inference_steps=steps,
seed=seed,
humanize=humanize,
protect_text=protect_text,
protect_faces=protect_faces,
max_resolution=max_resolution,
vendor=vendor,
)
@@ -990,9 +976,10 @@ def _process_batch_image(
)
@click.option(
"--pipeline",
type=click.Choice(["default", "ctrlregen"]),
type=click.Choice(["default", "controlnet"]),
default="default",
help="Pipeline profile (default=SDXL; ctrlregen=CtrlRegen, EXPERIMENTAL/destructive at clean-noise).",
help="Pipeline profile (default=SDXL img2img; controlnet=SDXL + canny ControlNet that preserves "
"text/faces via edge conditioning while removing SynthID).",
)
@click.option(
"--device",
src/remove_ai_watermarks/face_protector.py
-150
View File
@@ -1,150 +0,0 @@
"""YOLO-based face detection and soft-blend restoration for diffusion pipelines."""
# cv2/numpy/ultralytics boundary: these libs ship no usable element types; relax
# the unknown-type rules for this file only.
# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalCall=false, reportOptionalSubscript=false, reportOptionalOperand=false, reportAttributeAccessIssue=false, reportPrivateImportUsage=false, reportPrivateUsage=false, reportInvalidTypeForm=false, reportConstantRedefinition=false, reportUnnecessaryComparison=false, reportPossiblyUnboundVariable=false
import logging
from pathlib import Path
import cv2
import numpy as np
try:
from ultralytics import YOLO
HAS_YOLO = True
except ImportError:
HAS_YOLO = False
logger = logging.getLogger(__name__)
class FaceProtector:
"""
Detects faces in an image and provides methods to seamlessly paste them back
onto the an upscaled/processed image to preserve facial details that may have
been destroyed by latent diffusion or other algorithms.
"""
def __init__(self, use_yolo: bool = True, model_name: str = "yolov8n.pt") -> None:
self.use_yolo = use_yolo and HAS_YOLO
self.detector = None
self.haar_cascade = None
if self.use_yolo:
# Fix SSL certificate issues on macOS (fresh Python installs)
self._fix_ssl_certs()
logger.info("Loading YOLO model '%s' for face protection...", model_name)
self.detector = YOLO(model_name)
else:
if use_yolo and not HAS_YOLO:
logger.warning(
"ultralytics YOLO is not installed. Falling back to OpenCV Haar "
"Cascades. Install ultralytics with `pip install ultralytics` "
"for better face detection."
)
logger.info("Loading OpenCV Haar Cascade for face protection...")
cascade_path = Path(cv2.__file__).parent / "data" / "haarcascade_frontalface_default.xml"
if not cascade_path.exists():
cascade_path = "haarcascade_frontalface_default.xml"
self.haar_cascade = cv2.CascadeClassifier(str(cascade_path))
def detect_face_bboxes(self, image: np.ndarray) -> list[tuple[int, int, int, int]]:
"""
Detect faces and return bounding boxes as (x1, y1, x2, y2).
"""
if self.use_yolo and self.detector is not None:
# For standard YOLOv8n, 'person' is class 0. We'll use person bounding boxes
# as a proxy for faces/people to protect them. If using a specific face model, adjust classes.
results = self.detector(image, verbose=False, classes=[0])
bboxes = []
for r in results:
boxes = r.boxes
for box in boxes:
x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
bboxes.append((int(x1), int(y1), int(x2), int(y2)))
return bboxes
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
faces = self.haar_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
bboxes = []
for x, y, w, h in faces:
# Add a 20% margin around the haar cascade face box
margin_x = int(w * 0.2)
margin_y = int(h * 0.2)
x1 = max(0, x - margin_x)
y1 = max(0, y - int(margin_y * 1.5)) # more margin on top for hair
x2 = min(image.shape[1], x + w + margin_x)
y2 = min(image.shape[0], y + h + margin_y)
bboxes.append((x1, y1, x2, y2))
return bboxes
@staticmethod
def _fix_ssl_certs() -> None:
"""Set SSL_CERT_FILE from certifi if not already set (macOS fix)."""
import os
if os.environ.get("SSL_CERT_FILE"):
return
try:
import certifi
os.environ["SSL_CERT_FILE"] = certifi.where()
except ImportError:
pass
def extract_faces(self, image: np.ndarray) -> list[tuple[tuple[int, int, int, int], np.ndarray]]:
"""
Extract faces from the image.
Returns a list of (bbox, face_crop) tuples.
"""
bboxes = self.detect_face_bboxes(image)
faces = []
for bbox in bboxes:
x1, y1, x2, y2 = bbox
faces.append((bbox, image[y1:y2, x1:x2].copy()))
return faces
def restore_faces(
self, processed_image: np.ndarray, original_faces: list[tuple[tuple[int, int, int, int], np.ndarray]]
) -> np.ndarray:
"""
Paste original faces back onto the processed image using seamless cloning
or soft blending so the edges don't show.
"""
if not original_faces:
return processed_image
result = processed_image.copy()
for (x1, y1, x2, y2), face_crop in original_faces:
h, w = face_crop.shape[:2]
# If the processed image was resized, we'd need to resize face_crop, but
# pipeline ensures the output from InvisibleEngine is the same size or we resize it back before this.
if result.shape[:2] != processed_image.shape[:2]:
continue # Safety bypass
try:
# Create a soft alpha mask for the face crop to smoothly blend it
mask = np.zeros((h, w), dtype=np.float32)
# Inner ellipse is pure white
cv2.ellipse(mask, (w // 2, h // 2), (int(w * 0.4), int(h * 0.4)), 0, 0, 360, 1.0, -1)
# Blur the mask heavily for soft edges
blur_size = max(w, h) // 4
if blur_size % 2 == 0:
blur_size += 1
mask = cv2.GaussianBlur(mask, (blur_size, blur_size), 0)
mask = cv2.merge([mask, mask, mask])
# Blend
target_roi = result[y1:y2, x1:x2].astype(np.float32)
src_roi = face_crop.astype(np.float32)
blended = src_roi * mask + target_roi * (1.0 - mask)
result[y1:y2, x1:x2] = blended.astype(np.uint8)
except Exception as e:
logger.warning("Failed to restore face at %d,%d to %d,%d: %s", x1, y1, x2, y2, e)
return result
src/remove_ai_watermarks/invisible_engine.py
+24 -63
View File
@@ -7,9 +7,9 @@ This module requires the 'gpu' extra dependencies:
uv pip install 'remove-ai-watermarks[gpu]'
"""
# cv2/torch boundary: this engine wraps cv2 (resize/imwrite/cvtColor), the YOLO
# face protector, and the humanizer, none of which carry usable element types;
# relax the unknown-type rules for this file only.
# cv2/torch boundary: this engine wraps cv2 (resize/imwrite/cvtColor) and the
# humanizer, none of which carry usable element types; relax the unknown-type
# rules for this file only.
# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalCall=false, reportOptionalSubscript=false, reportOptionalOperand=false, reportAttributeAccessIssue=false, reportPrivateImportUsage=false, reportPrivateUsage=false, reportInvalidTypeForm=false, reportConstantRedefinition=false, reportUnnecessaryComparison=false
from __future__ import annotations
@@ -70,12 +70,9 @@ class InvisibleEngine:
to break watermark patterns, and reconstructs via reverse diffusion.
"""
# SDXL base is the default since May 2026; the current Google SynthID is
# removed at strength ~0.30 / steps=50 / native res (oracle-verified, n=3 fresh
# Gemini -- 0.10/0.15/0.2 still detected). See CLAUDE.md "Known limitations" for
# the strength study and the regression evidence ruling out SD-1.5 pipelines.
# SDXL base is the default since May 2026; the vendor-adaptive strength
# removes the current SynthID (see watermark_profiles + docs/synthid.md).
DEFAULT_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
CTRLREGEN_MODEL_ID = "yepengliu/ctrlregen"
def __init__(
self,
@@ -84,31 +81,33 @@ class InvisibleEngine:
pipeline: str = "default",
hf_token: str | None = None,
progress_callback: Callable[[str], None] | None = None,
controlnet_conditioning_scale: float = 1.0,
) -> None:
"""Initialize the invisible watermark removal engine.
Args:
model_id: HuggingFace model ID. None = use default for pipeline.
model_id: HuggingFace model ID. None = use the SDXL base default.
device: Device for inference (auto/cpu/mps/cuda/xpu). None = auto.
pipeline: Pipeline profile. "default" (SDXL base, defeats SynthID
v2) or "ctrlregen" (CtrlRegen).
pipeline: Pipeline profile. "default" (plain SDXL img2img) or
"controlnet" (SDXL + canny ControlNet that preserves text/face
structure via edge conditioning while removing SynthID).
hf_token: HuggingFace API token.
progress_callback: Optional callback for progress messages.
controlnet_conditioning_scale: ControlNet structure-preservation
strength (controlnet pipeline only).
"""
from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover
effective_model = model_id
if pipeline == "ctrlregen" and model_id is None:
effective_model = self.CTRLREGEN_MODEL_ID
elif model_id is None:
effective_model = self.DEFAULT_MODEL_ID
effective_model = model_id or self.DEFAULT_MODEL_ID
self._remover = WatermarkRemover(
model_id=effective_model,
device=device,
progress_callback=progress_callback,
hf_token=hf_token,
pipeline=pipeline,
controlnet_conditioning_scale=controlnet_conditioning_scale,
)
self._progress_callback = progress_callback
@@ -125,8 +124,6 @@ class InvisibleEngine:
guidance_scale: float | None = None,
seed: int | None = None,
humanize: float = 0.0,
protect_faces: bool = False,
protect_text: bool = False,
max_resolution: int = 0,
vendor: str | None = None,
) -> Path:
@@ -135,16 +132,12 @@ class InvisibleEngine:
Args:
image_path: Path to the watermarked image.
output_path: Output path (None = overwrite source).
strength: Denoising strength (0.0-1.0). None -> profile default
(0.10 for SDXL, 1.0 clean-noise for ctrlregen).
strength: Denoising strength (0.0-1.0). None -> the vendor-adaptive
default.
steps: Number of denoising steps.
guidance_scale: Classifier-free guidance scale.
seed: Random seed for reproducibility.
humanize: Intensity of Analog Humanizer film grain (0 = off).
protect_faces: Boolean to extract and restore faces intact.
protect_text: Detect text regions and preserve them via Differential
Diffusion when any are found, so glyphs (incl. CJK) survive the
removal pass. On by default; the detector decides per image.
max_resolution: Cap the long side (px) before diffusion. 0 (default)
= native resolution, no pre-downscale -- matches the hosted
raiw.cc backend. Set a positive value only to bound GPU/MPS
@@ -189,27 +182,6 @@ class InvisibleEngine:
image_path = _tmp_path
try:
# Optional: Face protection (Phase 1 - Extraction)
original_faces = []
if protect_faces:
try:
import cv2
from remove_ai_watermarks.face_protector import FaceProtector
if self._progress_callback:
self._progress_callback("Detecting and extracting faces (protect-faces)...")
# Convert PIL to CV2 BGR
import numpy as np
cv_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
protector = FaceProtector(use_yolo=True)
original_faces = protector.extract_faces(cv_img)
if self._progress_callback:
self._progress_callback(f"Extracted {len(original_faces)} face(s) for protection.")
except Exception as e:
logger.error("Failed to extract faces: %s", e)
out_path = self._remover.remove_watermark(
image_path=image_path,
output_path=output_path,
@@ -217,14 +189,12 @@ class InvisibleEngine:
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
seed=seed,
protect_text=protect_text,
vendor=vendor,
)
# Optional: Face restoration & Humanizer (Phase 2 - Post-processing)
if protect_faces or humanize > 0.0:
# Post-processing: optional Humanizer, then restore original resolution.
if humanize > 0.0:
import cv2
import numpy as np
from remove_ai_watermarks import image_io
@@ -232,20 +202,11 @@ class InvisibleEngine:
if out_cv is None:
return out_path
if protect_faces and original_faces:
if self._progress_callback:
self._progress_callback("Restoring protected faces with soft blending...")
from remove_ai_watermarks.face_protector import FaceProtector
if self._progress_callback:
self._progress_callback(f"Applying Analog Humanizer (grain: {humanize})...")
from remove_ai_watermarks.humanizer import apply_analog_humanizer
protector = FaceProtector(use_yolo=True)
out_cv = protector.restore_faces(out_cv, original_faces)
if humanize > 0.0:
if self._progress_callback:
self._progress_callback(f"Applying Analog Humanizer (grain: {humanize})...")
from remove_ai_watermarks.humanizer import apply_analog_humanizer
out_cv = apply_analog_humanizer(out_cv, grain_intensity=humanize, chromatic_shift=1)
out_cv = apply_analog_humanizer(out_cv, grain_intensity=humanize, chromatic_shift=1)
# Restore original resolution
if (out_cv.shape[1], out_cv.shape[0]) != orig_size:
@@ -259,7 +220,7 @@ class InvisibleEngine:
image_io.imwrite(out_path, out_cv)
else:
# Even if no protect_faces or humanize, we must restore original size if needed
# No humanize: still restore the original size if it was capped.
import cv2
from remove_ai_watermarks import image_io
src/remove_ai_watermarks/metadata.py
+6 -7
View File
@@ -255,10 +255,9 @@ def has_ai_metadata(image_path: Path) -> bool:
"""
from PIL import Image
# PIL may not handle AVIF/HEIF/JPEG-XL without the optional plugins
# (ultralytics also monkey-patches Image.open in a way that can raise
# ModuleNotFoundError when pi_heif autoload fails), so any open failure
# falls through to the binary scan.
# PIL may not handle AVIF/HEIF/JPEG-XL without the optional plugins, and a
# third-party plugin autoload can raise a non-OSError (e.g. ModuleNotFoundError),
# so any open failure falls through to the binary scan.
try:
with Image.open(image_path) as img:
for key in img.info:
@@ -655,9 +654,9 @@ def get_ai_metadata(image_path: Path) -> dict[str, str]:
result: dict[str, str] = {}
# PIL may not open AVIF/HEIF/JPEG-XL without optional plugins (and
# ultralytics' Image.open patch can raise ModuleNotFoundError); fall through
# to the C2PA/binary path on any open failure. See CLAUDE.md.
# PIL may not open AVIF/HEIF/JPEG-XL without optional plugins (and a
# third-party plugin autoload can raise a non-OSError); fall through to the
# C2PA/binary path on any open failure. See CLAUDE.md.
try:
with Image.open(image_path) as img:
for key, value in img.info.items():
src/remove_ai_watermarks/noai/ctrlregen/__init__.py
-18
View File
@@ -1,18 +0,0 @@
"""CtrlRegen watermark removal via controllable regeneration.
Implements the pipeline from "Image Watermarks Are Removable Using
Controllable Regeneration from Clean Noise" (ICLR 2025) by Liu et al.
This sub-package uses a ControlNet for spatial guidance (canny edges)
and a DINOv2-based IP Adapter for semantic guidance to regenerate
watermarked images from partially noised latents.
Attribution:
Based on https://github.com/yepengliu/CtrlRegen .
"""
from __future__ import annotations
from remove_ai_watermarks.noai.ctrlregen.engine import CtrlRegenEngine, is_ctrlregen_available
__all__ = ["CtrlRegenEngine", "is_ctrlregen_available"]
src/remove_ai_watermarks/noai/ctrlregen/color.py
-40
View File
@@ -1,40 +0,0 @@
"""Color matching post-processing for CtrlRegen output.
After diffusion-based regeneration, the output image may have slight
color shifts. This module uses histogram-based color transfer to
align the regenerated image's color distribution back to the original.
Attribution:
Adapted from https://github.com/yepengliu/CtrlRegen .
"""
from __future__ import annotations
import numpy as np
from color_matcher import ColorMatcher
from color_matcher.normalizer import Normalizer
from PIL import Image
def color_match(reference: Image.Image, source: Image.Image) -> Image.Image:
"""Transfer the color distribution of *reference* onto *source*.
Uses a two-pass histogram matching approach (``hm-mkl-hm``) that
preserves fine-grained color relationships while correcting global
shifts introduced by the regeneration pipeline.
Args:
reference: The original (watermarked) image whose colors should
be preserved.
source: The regenerated image whose colors will be adjusted.
Returns:
A new PIL Image with the structure of *source* but the color
palette of *reference*.
"""
cm = ColorMatcher()
ref_np = Normalizer(np.asarray(reference)).type_norm()
src_np = Normalizer(np.asarray(source)).type_norm()
result = cm.transfer(src=src_np, ref=ref_np, method="hm-mkl-hm")
result = Normalizer(result).uint8_norm()
return Image.fromarray(result)
src/remove_ai_watermarks/noai/ctrlregen/engine.py
-365
View File
@@ -1,365 +0,0 @@
"""CtrlRegen engine — orchestrates the full watermark removal pipeline.
Loads the base SD 1.5 model with a ControlNet (spatial control from
canny edges) and a DINOv2-based IP Adapter (semantic control), then
runs controllable regeneration with optional color matching.
Attribution:
Based on https://github.com/yepengliu/CtrlRegen .
"""
from __future__ import annotations
import contextlib
import logging
import os
import sys
import time
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from collections.abc import Callable
import torch
from PIL import Image
from remove_ai_watermarks.noai.progress import make_pipeline_progress
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Availability checks — these imports are optional.
# ---------------------------------------------------------------------------
_HAS_CONTROLNET_AUX = False
_HAS_COLOR_MATCHER = False
_HAS_DIFFUSERS = False
try:
from diffusers import AutoencoderKL, ControlNetModel, UniPCMultistepScheduler
from remove_ai_watermarks.noai.ctrlregen.pipeline import CustomCtrlRegenPipeline
_HAS_DIFFUSERS = True
except ImportError:
AutoencoderKL = None # type: ignore[assignment,misc]
ControlNetModel = None # type: ignore[assignment,misc]
UniPCMultistepScheduler = None # type: ignore[assignment,misc]
CustomCtrlRegenPipeline = None # type: ignore[assignment,misc]
try:
from controlnet_aux import CannyDetector
_HAS_CONTROLNET_AUX = True
except ImportError:
CannyDetector = None # type: ignore[assignment,misc]
try:
from remove_ai_watermarks.noai.ctrlregen.color import color_match
_HAS_COLOR_MATCHER = True
except ImportError:
color_match = None # type: ignore[assignment]
CTRLREGEN_HF_REPO = "yepengliu/ctrlregen"
SPATIAL_SUBFOLDER = "spatialnet_ckp/spatial_control_ckp_14000"
SEMANTIC_SUBFOLDER = "semanticnet_ckp/models"
SEMANTIC_WEIGHT_NAME = "semantic_control_ckp_435000.bin"
DEFAULT_BASE_MODEL = "SG161222/Realistic_Vision_V4.0_noVAE"
CUSTOM_VAE_ID = "stabilityai/sd-vae-ft-mse"
PROCESS_SIZE = 512
DEFAULT_GUIDANCE_SCALE = 2.0
QUALITY_PROMPT = "best quality, high quality"
NEGATIVE_PROMPT = "monochrome, lowres, bad anatomy, worst quality, low quality"
CANNY_LOW_THRESHOLD = 100
CANNY_HIGH_THRESHOLD = 150
TILE_SIZE = 512
TILE_OVERLAP = 192
def is_ctrlregen_available() -> bool:
"""Return True when all CtrlRegen-specific dependencies are installed."""
return _HAS_DIFFUSERS and _HAS_CONTROLNET_AUX and _HAS_COLOR_MATCHER
class CtrlRegenEngine:
"""End-to-end CtrlRegen watermark removal engine.
Handles model loading, canny edge extraction, controlled denoising,
and color-matched post-processing in a single ``run()`` call.
"""
def __init__(
self,
base_model_id: str | None = None,
device: str = "cpu",
torch_dtype: torch.dtype | None = None,
hf_token: str | None = None,
progress_callback: Callable[[str], None] | None = None,
) -> None:
if not is_ctrlregen_available():
missing: list[str] = []
if not _HAS_DIFFUSERS:
missing.extend(["diffusers", "transformers", "accelerate"])
if not _HAS_CONTROLNET_AUX:
missing.append("controlnet-aux")
if not _HAS_COLOR_MATCHER:
missing.append("color-matcher")
logger.info("Auto-installing missing dependencies: %s", missing)
import subprocess
try:
subprocess.check_call(
[sys.executable, "-m", "pip", "install", *missing],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
)
except (subprocess.CalledProcessError, FileNotFoundError) as exc:
raise ImportError(
"Failed to auto-install missing dependencies: "
+ ", ".join(missing)
+ ". Try manually: pip install --force-reinstall noai-watermark"
) from exc
self.base_model_id = base_model_id or DEFAULT_BASE_MODEL
self.device = device
self.torch_dtype = torch_dtype or (torch.float32 if device in ("cpu", "mps") else torch.float16)
self.hf_token: str | None = hf_token or os.environ.get("HF_TOKEN")
self._progress_callback = progress_callback
self._pipeline: CustomCtrlRegenPipeline | None = None # type: ignore[assignment]
self._canny_detector: CannyDetector | None = None # type: ignore[assignment]
def _set_progress(self, message: str) -> None:
if self._progress_callback is None:
return
with contextlib.suppress(Exception):
self._progress_callback(message)
# ------------------------------------------------------------------
# Loading
# ------------------------------------------------------------------
def load(self) -> None:
"""Download and assemble the full CtrlRegen pipeline."""
if self._pipeline is not None:
return
token_kwargs: dict[str, Any] = {}
if self.hf_token:
token_kwargs["token"] = self.hf_token
self._set_progress(f"Loading CtrlRegen spatial ControlNet from {CTRLREGEN_HF_REPO}...")
logger.info("Loading ControlNet from %s/%s", CTRLREGEN_HF_REPO, SPATIAL_SUBFOLDER)
controlnet = [
ControlNetModel.from_pretrained(
CTRLREGEN_HF_REPO,
subfolder=SPATIAL_SUBFOLDER,
torch_dtype=self.torch_dtype,
**token_kwargs,
)
]
self._set_progress(f"Loading SD base model ({self.base_model_id}) for CtrlRegen pipeline...")
logger.info("Loading base pipeline from %s", self.base_model_id)
pipe = CustomCtrlRegenPipeline.from_pretrained(
self.base_model_id,
controlnet=controlnet,
torch_dtype=self.torch_dtype,
safety_checker=None,
requires_safety_checker=False,
**token_kwargs,
)
self._set_progress(f"Loading CtrlRegen semantic IP-Adapter + DINOv2 from {CTRLREGEN_HF_REPO}...")
logger.info("Loading IP-Adapter from %s/%s", CTRLREGEN_HF_REPO, SEMANTIC_SUBFOLDER)
pipe.load_ctrlregen_ip_adapter(
CTRLREGEN_HF_REPO,
subfolder=SEMANTIC_SUBFOLDER,
weight_name=SEMANTIC_WEIGHT_NAME,
**token_kwargs,
)
from transformers import AutoImageProcessor, AutoModel
pipe.image_encoder = AutoModel.from_pretrained("facebook/dinov2-giant").to(self.device, dtype=self.torch_dtype)
pipe.feature_extractor = AutoImageProcessor.from_pretrained("facebook/dinov2-giant")
self._set_progress(f"Loading custom VAE ({CUSTOM_VAE_ID})...")
logger.info("Loading VAE from %s", CUSTOM_VAE_ID)
pipe.vae = AutoencoderKL.from_pretrained(
CUSTOM_VAE_ID,
torch_dtype=self.torch_dtype,
**token_kwargs,
)
self._set_progress("Configuring UniPC scheduler...")
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe.set_ip_adapter_scale(1.0)
self._set_progress(f"Moving CtrlRegen pipeline to {self.device}...")
pipe = pipe.to(self.device)
if hasattr(pipe, "enable_xformers_memory_efficient_attention"):
with contextlib.suppress(Exception):
pipe.enable_xformers_memory_efficient_attention()
self._pipeline = pipe
self._canny_detector = CannyDetector()
self._set_progress("CtrlRegen pipeline ready.")
logger.info("CtrlRegen pipeline loaded on %s", self.device)
# ------------------------------------------------------------------
# Inference — public entry point
# ------------------------------------------------------------------
def run(
self,
image: Image.Image,
strength: float = 0.5,
num_inference_steps: int = 50,
guidance_scale: float = DEFAULT_GUIDANCE_SCALE,
seed: int | None = None,
) -> Image.Image:
"""Run CtrlRegen watermark removal on a single image.
Images that fit within ``TILE_SIZE`` (512) are processed as a
single pass. Larger images are split into overlapping tiles.
"""
self.load()
assert self._pipeline is not None
assert self._canny_detector is not None
orig_w, orig_h = image.size
orig_image = image
t0 = time.monotonic()
needs_tiling = orig_w > TILE_SIZE or orig_h > TILE_SIZE
if needs_tiling:
from remove_ai_watermarks.noai.ctrlregen.tiling import resize_center_crop, run_tiled
aligned_w = orig_w // 8 * 8
aligned_h = orig_h // 8 * 8
if aligned_w != orig_w or aligned_h != orig_h:
image = image.resize((aligned_w, aligned_h), Image.LANCZOS)
regen_image = run_tiled(
pipeline=self._pipeline,
canny_detector=self._canny_detector,
image=image,
strength=strength,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
seed=seed,
tile_size=TILE_SIZE,
tile_overlap=TILE_OVERLAP,
quality_prompt=QUALITY_PROMPT,
negative_prompt=NEGATIVE_PROMPT,
canny_low=CANNY_LOW_THRESHOLD,
canny_high=CANNY_HIGH_THRESHOLD,
device=self.device,
set_progress=self._set_progress,
ip_adapter_image=orig_image,
)
else:
from remove_ai_watermarks.noai.ctrlregen.tiling import resize_center_crop
proc_image = resize_center_crop(image, PROCESS_SIZE)
self._set_progress(f"Preprocessed {orig_w}x{orig_h}px → {proc_image.size[0]}x{proc_image.size[1]}px")
regen_image = self._run_single(
proc_image,
strength,
num_inference_steps,
guidance_scale,
seed,
)
if regen_image.size != (orig_w, orig_h):
self._set_progress(f"Resizing {regen_image.size[0]}x{regen_image.size[1]}px → {orig_w}x{orig_h}px...")
regen_image = regen_image.resize((orig_w, orig_h), Image.LANCZOS)
self._set_progress(f"Applying color matching at {orig_w}x{orig_h}px...")
output = color_match(reference=orig_image, source=regen_image)
self._set_progress(f"✓ CtrlRegen done · {orig_w}x{orig_h}px · {time.monotonic() - t0:.0f}s total")
return output
# ------------------------------------------------------------------
# Single-image path (image <= 512x512)
# ------------------------------------------------------------------
def _run_single(
self,
image: Image.Image,
strength: float,
num_inference_steps: int,
guidance_scale: float,
seed: int | None,
) -> Image.Image:
"""Process a single 512x512 image through the CtrlRegen pipeline."""
w, h = image.size
effective_steps = max(1, int(num_inference_steps * strength))
self._set_progress(
f"Extracting canny edges ({w}x{h}px, thresholds {CANNY_LOW_THRESHOLD}/{CANNY_HIGH_THRESHOLD})..."
)
control_image = self._canny_detector(
image,
low_threshold=CANNY_LOW_THRESHOLD,
high_threshold=CANNY_HIGH_THRESHOLD,
)
generator = torch.manual_seed(seed if seed is not None else 0)
self._set_progress(
f"Config: strength={strength}, steps={num_inference_steps} "
f"(~{effective_steps} effective), guidance={guidance_scale}"
)
step_cb, first_step, pipeline_done, start_updater = make_pipeline_progress(
effective_steps,
self.device,
self._set_progress,
label="CtrlRegen denoising",
)
start_updater()
try:
result = self._pipeline(
prompt=QUALITY_PROMPT,
negative_prompt=NEGATIVE_PROMPT,
image=[image],
control_image=[control_image],
controlnet_conditioning_scale=1.0,
ip_adapter_image=[image],
strength=strength,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
generator=generator,
control_guidance_start=0.0,
control_guidance_end=1.0,
callback=step_cb,
callback_steps=1,
)
except TypeError:
first_step.set()
result = self._pipeline(
prompt=QUALITY_PROMPT,
negative_prompt=NEGATIVE_PROMPT,
image=[image],
control_image=[control_image],
controlnet_conditioning_scale=1.0,
ip_adapter_image=[image],
strength=strength,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
generator=generator,
)
finally:
first_step.set()
pipeline_done.set()
return result.images[0]
src/remove_ai_watermarks/noai/ctrlregen/ip_adapter.py
-149
View File
@@ -1,149 +0,0 @@
"""Custom IP-Adapter mixin using DINOv2 as the image encoder.
The standard diffusers ``IPAdapterMixin`` uses a CLIP image encoder.
CtrlRegen replaces it with ``facebook/dinov2-giant`` for richer
semantic features. This mixin provides ``load_ctrlregen_ip_adapter``
which handles the custom weight format and encoder swap.
Attribution:
Adapted from https://github.com/yepengliu/CtrlRegen .
"""
from __future__ import annotations
import logging
from typing import Any
import torch
from diffusers.models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT
from diffusers.utils import (
_get_model_file,
is_accelerate_available,
is_torch_version,
)
from diffusers.utils import (
logging as diffusers_logging,
)
from huggingface_hub.utils import validate_hf_hub_args
from safetensors import safe_open
from transformers import AutoImageProcessor, AutoModel
logger = logging.getLogger(__name__)
_diffusers_logger = diffusers_logging.get_logger(__name__)
DINOV2_MODEL_ID = "facebook/dinov2-giant"
class CustomIPAdapterMixin:
"""Mixin that adds ``load_ctrlregen_ip_adapter`` to a diffusers pipeline."""
@validate_hf_hub_args
def load_ctrlregen_ip_adapter(
self,
pretrained_model_name_or_path_or_dict: str | list[str] | dict[str, torch.Tensor],
subfolder: str | list[str],
weight_name: str | list[str],
image_encoder_folder: str | None = "image_encoder",
**kwargs: Any,
) -> None:
"""Load CtrlRegen IP-Adapter weights and DINOv2 image encoder.
Parameters mirror ``IPAdapterMixin.load_ip_adapter`` but the
image encoder is always ``facebook/dinov2-giant`` regardless of
the ``image_encoder_folder`` value in the checkpoint.
"""
if not isinstance(weight_name, list):
weight_name = [weight_name]
if not isinstance(pretrained_model_name_or_path_or_dict, list):
pretrained_model_name_or_path_or_dict = [pretrained_model_name_or_path_or_dict]
if len(pretrained_model_name_or_path_or_dict) == 1:
pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict * len(weight_name)
if not isinstance(subfolder, list):
subfolder = [subfolder]
if len(subfolder) == 1:
subfolder = subfolder * len(weight_name)
if len(weight_name) != len(pretrained_model_name_or_path_or_dict):
raise ValueError("`weight_name` and `pretrained_model_name_or_path_or_dict` must have the same length.")
if len(weight_name) != len(subfolder):
raise ValueError("`weight_name` and `subfolder` must have the same length.")
cache_dir = kwargs.pop("cache_dir", None)
force_download = kwargs.pop("force_download", False)
kwargs.pop("resume_download", None)
proxies = kwargs.pop("proxies", None)
local_files_only = kwargs.pop("local_files_only", None)
token = kwargs.pop("token", None)
revision = kwargs.pop("revision", None)
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT)
if low_cpu_mem_usage and not is_accelerate_available():
low_cpu_mem_usage = False
_diffusers_logger.warning(
"Cannot initialize model with low cpu memory usage because "
"`accelerate` was not found. Defaulting to "
"`low_cpu_mem_usage=False`."
)
if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"):
raise NotImplementedError("Low memory initialization requires torch >= 1.9.0.")
user_agent = {
"file_type": "attn_procs_weights",
"framework": "pytorch",
}
state_dicts: list[dict] = []
for path_or_dict, wn, sf in zip(pretrained_model_name_or_path_or_dict, weight_name, subfolder, strict=False):
if not isinstance(path_or_dict, dict):
model_file = _get_model_file(
path_or_dict,
weights_name=wn,
cache_dir=cache_dir,
force_download=force_download,
proxies=proxies,
local_files_only=local_files_only,
token=token,
revision=revision,
subfolder=sf,
user_agent=user_agent,
)
if wn.endswith(".safetensors"):
state_dict: dict = {"image_proj": {}, "ip_adapter": {}}
with safe_open(model_file, framework="pt", device="cpu") as f:
for key in f.keys(): # noqa: SIM118
if key.startswith("image_proj."):
state_dict["image_proj"][key.replace("image_proj.", "")] = f.get_tensor(key)
elif key.startswith("ip_adapter."):
state_dict["ip_adapter"][key.replace("ip_adapter.", "")] = f.get_tensor(key)
else:
state_dict = torch.load(model_file, map_location="cpu")
else:
state_dict = path_or_dict
keys = list(state_dict.keys())
if keys != ["image_proj", "ip_adapter"]:
raise ValueError("Required keys (`image_proj` and `ip_adapter`) missing from the state dict.")
state_dicts.append(state_dict)
# Always use DINOv2-giant as the image encoder.
has_encoder_attr = hasattr(self, "image_encoder") and getattr(self, "image_encoder", None) is None
if has_encoder_attr and image_encoder_folder is not None:
logger.info("Loading DINOv2-giant image encoder for CtrlRegen")
enc_dtype = getattr(self, "dtype", torch.float32) # type: ignore[attr-defined]
image_encoder = AutoModel.from_pretrained(DINOV2_MODEL_ID).to(
self.device,
dtype=enc_dtype, # type: ignore[attr-defined]
)
self.register_modules(image_encoder=image_encoder) # type: ignore[attr-defined]
if hasattr(self, "feature_extractor") and getattr(self, "feature_extractor", None) is None:
feature_extractor = AutoImageProcessor.from_pretrained(DINOV2_MODEL_ID)
self.register_modules(feature_extractor=feature_extractor) # type: ignore[attr-defined]
unet = (
getattr(self, self.unet_name) # type: ignore[attr-defined]
if not hasattr(self, "unet")
else self.unet # type: ignore[attr-defined]
)
unet._load_ip_adapter_weights(state_dicts, low_cpu_mem_usage=low_cpu_mem_usage)
src/remove_ai_watermarks/noai/ctrlregen/pipeline.py
-35
View File
@@ -1,35 +0,0 @@
"""Custom Stable Diffusion ControlNet Img2Img pipeline for CtrlRegen.
Extends ``StableDiffusionControlNetImg2ImgPipeline`` with the
``load_ctrlregen_ip_adapter`` method (via ``CustomIPAdapterMixin``)
that swaps in DINOv2-giant as the image encoder and loads the
CtrlRegen semantic-control adapter weights.
No ``encode_image`` override is needed — the CtrlRegen checkpoint
creates an ``IPAdapterPlusImageProjection`` which tells diffusers to
call ``encode_image`` with ``output_hidden_states=True``. The
default implementation then uses ``hidden_states[-2]`` from DINOv2,
which is exactly what the projection was trained on.
Attribution:
Adapted from https://github.com/yepengliu/CtrlRegen .
"""
from __future__ import annotations
from diffusers import StableDiffusionControlNetImg2ImgPipeline
from remove_ai_watermarks.noai.ctrlregen.ip_adapter import CustomIPAdapterMixin
class CustomCtrlRegenPipeline(
StableDiffusionControlNetImg2ImgPipeline,
CustomIPAdapterMixin,
):
"""SD ControlNet Img2Img pipeline with DINOv2 IP-Adapter support.
MRO mirrors the original CtrlRegen repository: the base diffusers
pipeline comes first so all standard methods are resolved from it,
while ``CustomIPAdapterMixin`` only adds the
``load_ctrlregen_ip_adapter`` method.
"""
src/remove_ai_watermarks/noai/ctrlregen/tiling.py
-179
View File
@@ -1,179 +0,0 @@
"""Tile-based processing for large images in the CtrlRegen pipeline.
Extracted from ``ctrlregen.engine`` to keep the engine focused on
single-image inference and model orchestration.
"""
from __future__ import annotations
import math
import time
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from collections.abc import Callable
import numpy as np
import torch
from PIL import Image
def tile_positions(total: int, tile: int, overlap: int) -> list[int]:
"""Compute evenly-spaced tile start positions covering *total* pixels."""
if not (0 <= overlap < tile):
raise ValueError(f"overlap must satisfy 0 <= overlap < tile (got overlap={overlap}, tile={tile})")
if total <= tile:
return [0]
n = max(2, math.ceil((total - overlap) / (tile - overlap)))
stride = (total - tile) / (n - 1)
return [round(i * stride) for i in range(n)]
def make_blend_weight(h: int, w: int, overlap: int) -> np.ndarray:
"""2-D weight mask: 1.0 in center, cosine ramp in overlap margins."""
wy = np.ones(h, dtype=np.float64)
wx = np.ones(w, dtype=np.float64)
if overlap > 0:
ramp = 0.5 - 0.5 * np.cos(np.linspace(0, np.pi, overlap))
wy[:overlap] = np.minimum(wy[:overlap], ramp)
wy[-overlap:] = np.minimum(wy[-overlap:], ramp[::-1])
wx[:overlap] = np.minimum(wx[:overlap], ramp)
wx[-overlap:] = np.minimum(wx[-overlap:], ramp[::-1])
return np.outer(wy, wx)
def resize_center_crop(image: Image.Image, size: int = 512) -> Image.Image:
"""Resize shortest edge to *size*, then center-crop to a square.
Matches the ``transforms.Resize(512) + CenterCrop(512)`` pipeline
used in the original CtrlRegen repository.
"""
w, h = image.size
short = min(w, h)
scale = size / short
new_w, new_h = round(w * scale), round(h * scale)
image = image.resize((new_w, new_h), Image.BILINEAR)
left = (new_w - size) // 2
top = (new_h - size) // 2
return image.crop((left, top, left + size, top + size))
def run_tiled(
pipeline: Any,
canny_detector: Any,
image: Image.Image,
strength: float,
num_inference_steps: int,
guidance_scale: float,
seed: int | None,
*,
tile_size: int,
tile_overlap: int,
quality_prompt: str,
negative_prompt: str,
canny_low: int,
canny_high: int,
device: str,
set_progress: Callable[[str], None],
ip_adapter_image: Image.Image | None = None,
) -> Image.Image:
"""Split a large image into overlapping tiles, process each, blend."""
w, h = image.size
xs = tile_positions(w, tile_size, tile_overlap)
ys = tile_positions(h, tile_size, tile_overlap)
n_tiles = len(xs) * len(ys)
grid = f"{len(xs)}x{len(ys)}"
effective_steps = max(1, int(num_inference_steps * strength))
set_progress(f"Tiling {w}x{h}px → {n_tiles} tiles ({grid} grid, {tile_size}px, overlap {tile_overlap}px)")
canvas = np.zeros((h, w, 3), dtype=np.float64)
weight_sum = np.zeros((h, w), dtype=np.float64)
blend_w = make_blend_weight(tile_size, tile_size, tile_overlap)
t0 = time.monotonic()
bar_len = 20
tile_idx = 0
for ty in ys:
for tx in xs:
tile_idx += 1
prefix = f"[Tile {tile_idx}/{n_tiles}]"
tile = image.crop((tx, ty, tx + tile_size, ty + tile_size))
set_progress(f"{prefix} Extracting canny edges...")
control = canny_detector(
tile,
low_threshold=canny_low,
high_threshold=canny_high,
)
gen = None
if seed is not None:
gen = torch.Generator(device=device).manual_seed(seed + tile_idx)
tile_t0 = time.monotonic()
def _make_cb(
_prefix: str = prefix,
_t0: float = tile_t0,
_es: int = effective_steps,
) -> Callable:
def _cb(step: int, timestep: int, latents: Any) -> None:
elapsed = time.monotonic() - _t0
cur = step + 1
per = elapsed / max(1, cur)
rem = per * max(0, _es - cur)
filled = int(bar_len * cur / max(1, _es))
bar = "█" * filled + "░" * (bar_len - filled)
set_progress(f"{_prefix} [{bar}] {cur}/{_es} | {elapsed:.0f}s, ~{rem:.0f}s left")
return _cb
sem_image = ip_adapter_image if ip_adapter_image is not None else tile
try:
result = pipeline(
prompt=quality_prompt,
negative_prompt=negative_prompt,
image=[tile],
control_image=[control],
controlnet_conditioning_scale=1.0,
ip_adapter_image=[sem_image],
strength=strength,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
generator=gen,
control_guidance_start=0.0,
control_guidance_end=1.0,
callback=_make_cb(),
callback_steps=1,
)
except TypeError:
result = pipeline(
prompt=quality_prompt,
negative_prompt=negative_prompt,
image=[tile],
control_image=[control],
controlnet_conditioning_scale=1.0,
ip_adapter_image=[sem_image],
strength=strength,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
generator=gen,
)
proc_arr = np.array(result.images[0], dtype=np.float64)
th, tw = proc_arr.shape[:2]
mask = blend_w[:th, :tw]
canvas[ty : ty + th, tx : tx + tw] += proc_arr * mask[..., None]
weight_sum[ty : ty + th, tx : tx + tw] += mask
tile_time = time.monotonic() - tile_t0
total_elapsed = time.monotonic() - t0
set_progress(f"{prefix} Done ({tile_time:.0f}s) · Total: {total_elapsed:.0f}s")
set_progress(f"Blending {n_tiles} tiles → {w}x{h}px...")
canvas /= np.maximum(weight_sum[..., None], 1e-8)
return Image.fromarray(np.clip(canvas, 0, 255).astype(np.uint8))
src/remove_ai_watermarks/noai/img2img_runner.py
+18 -115
View File
@@ -29,8 +29,14 @@ def run_img2img(
generator: Any,
device: str,
set_progress: Callable[[str], None],
extra_kwargs: dict[str, Any] | None = None,
) -> Image.Image:
"""Execute img2img with live progress and return the generated image."""
"""Execute img2img with live progress and return the generated image.
``extra_kwargs`` overlays additional pipeline arguments (e.g. the ControlNet
``control_image`` / ``controlnet_conditioning_scale`` and a non-empty prompt),
so a ControlNet img2img pass reuses the same progress + fallback machinery.
"""
effective_steps = max(1, int(num_inference_steps * strength))
step_cb, first_step, done_ev, start_updater = make_pipeline_progress(
@@ -42,26 +48,14 @@ def run_img2img(
try:
result = _call_pipeline(
pipeline,
image,
strength,
num_inference_steps,
guidance_scale,
generator,
step_cb,
pipeline, image, strength, num_inference_steps, guidance_scale, generator, step_cb, extra_kwargs
)
done_ev.set()
return result.images[0]
except TypeError:
first_step.set()
result = _call_pipeline(
pipeline,
image,
strength,
num_inference_steps,
guidance_scale,
generator,
None,
pipeline, image, strength, num_inference_steps, guidance_scale, generator, None, extra_kwargs
)
done_ev.set()
return result.images[0]
@@ -81,11 +75,13 @@ def run_img2img_with_mps_fallback(
set_progress: Callable[[str], None],
*,
reload_on_cpu: Callable[[], Any],
extra_kwargs: dict[str, Any] | None = None,
) -> tuple[Image.Image, str]:
"""Run img2img; on MPS error, fall back to CPU.
Returns:
(result_image, final_device) — device may change to ``"cpu"`` on fallback.
``extra_kwargs`` overlays extra pipeline arguments (used by the ControlNet
path). Returns ``(result_image, final_device)`` — device may change to
``"cpu"`` on fallback.
"""
pipeline = load_pipeline()
@@ -99,6 +95,7 @@ def run_img2img_with_mps_fallback(
generator,
device,
set_progress,
extra_kwargs,
)
return img, device
except RuntimeError as error:
@@ -108,104 +105,7 @@ def run_img2img_with_mps_fallback(
_try_clear_mps_cache()
pipeline = reload_on_cpu()
img = run_img2img(
pipeline,
image,
strength,
num_inference_steps,
guidance_scale,
None,
"cpu",
set_progress,
)
return img, "cpu"
raise
def run_differential(
pipeline: Any,
image: Image.Image,
change_map: Any,
strength: float,
num_inference_steps: int,
guidance_scale: float,
generator: Any,
device: str,
set_progress: Callable[[str], None],
) -> Image.Image:
"""Run the SDXL Differential-Diffusion pipeline and return the image.
Unlike standard img2img, the differential pipeline needs pre-processed image
tensors plus a per-pixel change map (HxW float32 in [0, 1]); white preserves
the original pixels, black regenerates them. Runs without a step callback --
the community pipeline's callback signature differs across diffusers
versions, and a protect-text pass is short.
"""
import torch
image_tensor = pipeline.image_processor.preprocess(image).to(device)
map_tensor = torch.from_numpy(change_map)[None].to(device) # pyright: ignore[reportPrivateImportUsage, reportUnknownMemberType]
set_progress(f"Running protected regeneration ({device}, strength={strength})...")
result = pipeline(
prompt="",
image=image_tensor,
original_image=image_tensor,
map=map_tensor,
strength=strength,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
generator=generator,
)
return result.images[0]
def run_differential_with_mps_fallback(
load_pipeline: Callable[[], Any],
image: Image.Image,
change_map: Any,
strength: float,
num_inference_steps: int,
guidance_scale: float,
generator: Any,
device: str,
set_progress: Callable[[str], None],
*,
reload_on_cpu: Callable[[], Any],
) -> tuple[Image.Image, str]:
"""Run differential img2img; on MPS error, fall back to CPU.
Returns:
(result_image, final_device) -- device may change to ``"cpu"`` on fallback.
"""
pipeline = load_pipeline()
try:
img = run_differential(
pipeline,
image,
change_map,
strength,
num_inference_steps,
guidance_scale,
generator,
device,
set_progress,
)
return img, device
except RuntimeError as error:
if device == "mps" and is_mps_error(error):
logger.warning("MPS error detected: %s. Falling back to CPU.", error)
set_progress("MPS error! Clearing cache and retrying on CPU...")
_try_clear_mps_cache()
pipeline = reload_on_cpu()
img = run_differential(
pipeline,
image,
change_map,
strength,
num_inference_steps,
guidance_scale,
None,
"cpu",
set_progress,
pipeline, image, strength, num_inference_steps, guidance_scale, None, "cpu", set_progress, extra_kwargs
)
return img, "cpu"
raise
@@ -219,6 +119,7 @@ def _call_pipeline(
guidance_scale: float,
generator: Any,
step_callback: Any,
extra_kwargs: dict[str, Any] | None = None,
) -> Any:
kwargs: dict[str, Any] = {
"prompt": "",
@@ -228,6 +129,8 @@ def _call_pipeline(
"guidance_scale": guidance_scale,
"generator": generator,
}
if extra_kwargs:
kwargs.update(extra_kwargs)
if step_callback is not None:
kwargs["callback"] = step_callback
kwargs["callback_steps"] = 1
src/remove_ai_watermarks/noai/watermark_profiles.py
+27 -50
View File
@@ -1,4 +1,4 @@
"""Watermark removal model profiles, the default strength, and profile detection.
"""Watermark removal model profiles and the default strength.
Pure configuration and lookup functions with no ML dependencies.
"""
@@ -11,13 +11,18 @@ if TYPE_CHECKING:
from pathlib import Path
DEFAULT_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
CTRLREGEN_MODEL_ID = "yepengliu/ctrlregen"
# The SDXL-native canny ControlNet used by the ``controlnet`` pipeline. The
# ControlNet is an add-on to the SDXL base checkpoint (DEFAULT_MODEL_ID), not a
# separate base model, so both the ``default`` and ``controlnet`` profiles load
# the same base weights and share the same vendor-adaptive strength.
CONTROLNET_CANNY_MODEL = "xinsir/controlnet-canny-sdxl-1.0"
# Vendor-adaptive default denoising strength for the SDXL img2img scrub, overridable
# from the CLI (`--strength`). The right strength depends on which vendor's SynthID is
# present, detected from the C2PA issuer (metadata.synthid_source). Oracle-verified
# controlled study (2026-06-01, clean v0.8.6 with protect_text/faces OFF, per-image
# openai.com/verify or Gemini-app verdict; see docs/synthid.md section 2.2):
# controlled study (2026-06-01, clean v0.8.6, per-image openai.com/verify or Gemini-app
# verdict; see docs/synthid.md section 2.2):
# - OpenAI gpt-image: removed at 0.05 across 1024-1600 (n=4), resolution-independent.
# OPENAI_STRENGTH 0.10 = the 0.05 floor plus a 2x margin (keeps quality high).
# - Google Gemini: removed at 0.15 on the capped-1536 path (n=4); 0.05/0.10 do NOT
@@ -29,8 +34,8 @@ CTRLREGEN_MODEL_ID = "yepengliu/ctrlregen"
# - Unknown vendor (metadata stripped, or non-OpenAI/Google C2PA): UNKNOWN_STRENGTH
# 0.15, the safe middle that clears both vendors at the tested resolutions.
# The dominant factor is VENDOR, not resolution: Google's SynthID is ~3x more robust
# than OpenAI's. The earlier single 0.30 default (and the "resolution dependence" lore)
# came from contaminated tests run with protect_text ON -- see docs/synthid.md 2.2.
# than OpenAI's. The ``controlnet`` pipeline shares these strengths (same SDXL base; the
# canny ControlNet only preserves structure, the strength still drives removal).
OPENAI_STRENGTH = 0.10
GEMINI_STRENGTH = 0.15
UNKNOWN_STRENGTH = 0.15
@@ -41,45 +46,21 @@ DEFAULT_STRENGTH = UNKNOWN_STRENGTH
# Detected-vendor -> default strength. Vendor strings come from `vendor_for_strength`.
_VENDOR_STRENGTH = {"openai": OPENAI_STRENGTH, "google": GEMINI_STRENGTH}
# CtrlRegen removes watermarks by regenerating from (near) clean Gaussian noise,
# NOT by the light-touch partial-noise img2img the SDXL default uses. The research
# is explicit (CtrlRegen, ICLR 2025, arXiv:2410.05470): partial-noise regeneration
# "struggles with high-perturbation watermarks" because a small noise step "retains"
# watermark information that diffuses back into the output; the fix is to start from
# clean noise. With the StableDiffusionControlNetImg2ImgPipeline that maps to a high
# strength (~1.0 = full noise at the first timestep, structure held by the canny
# ControlNet + DINOv2 IP-Adapter, not by the watermarked latent). So the ctrlregen
# profile must NOT inherit the SDXL default (`DEFAULT_STRENGTH`, a partial-noise
# value) -- at that low strength it loads ControlNet + DINOv2-giant and then barely
# changes the image (a no-op for removal). Tunable via
# `--strength`; lower it to trade removal strength for fidelity (the CtrlRegen+ regime).
#
# EXPERIMENTAL -- NOT recommended for production. The same GPU study that set the 0.3
# SDXL threshold tested ctrlregen at its clean-noise strength and found it DESTROYS
# images: smooth/background regions fill with hallucinated micro-text garbage, and it
# is heavy (~8.5 min / ~$0.30 vs ~25 s / ~$0.02 for SDXL on a large image). The pipeline
# is effectively binary -- low strength = no-op, high strength = destroys -- with no
# usable middle, so the literature's "clean-noise is the lever" (arXiv:2410.05470) did
# NOT survive empirical testing on real content. SDXL img2img at ~0.3 is the shippable
# path; ctrlregen stays opt-in and flagged experimental.
CTRLREGEN_DEFAULT_STRENGTH = 1.0
def resolve_strength(strength: float | None, vendor: str | None = None) -> float:
"""Resolve the denoising strength, applying the vendor default when unset.
def resolve_strength(strength: float | None, profile: str, vendor: str | None = None) -> float:
"""Resolve the denoising strength, applying the profile/vendor default when unset.
``None`` means "the user did not pass ``--strength``". ``ctrlregen`` resolves to
``CTRLREGEN_DEFAULT_STRENGTH`` (clean-noise regeneration). The SDXL default profile
resolves **vendor-adaptively**: ``vendor`` (``"openai"`` / ``"google"`` / None, from
``None`` means "the user did not pass ``--strength``", which resolves
**vendor-adaptively**: ``vendor`` (``"openai"`` / ``"google"`` / None, from
``vendor_for_strength``) selects ``OPENAI_STRENGTH`` / ``GEMINI_STRENGTH`` /
``UNKNOWN_STRENGTH``. An explicit value always wins (including ``0.0`` -- the check is
``is None``, not falsiness). Shared by the CLI (for display) and the engine (for
``UNKNOWN_STRENGTH``. An explicit value always wins (including ``0.0`` -- the check
is ``is None``, not falsiness). The ``default`` and ``controlnet`` profiles share
the same SDXL base (the ControlNet only preserves structure), so the default does
NOT depend on the profile. Shared by the CLI (for display) and the engine (for
execution) so the two never disagree -- both must pass the SAME ``vendor``.
"""
if strength is not None:
return strength
if profile == "ctrlregen":
return CTRLREGEN_DEFAULT_STRENGTH
return _VENDOR_STRENGTH.get(vendor or "", UNKNOWN_STRENGTH)
@@ -107,17 +88,13 @@ def vendor_for_strength(image_path: Path) -> Literal["openai", "google"] | None:
def get_model_id_for_profile(profile: str) -> str:
"""Map CLI model profile names to concrete Hugging Face model IDs."""
"""Map CLI model profile names to concrete Hugging Face model IDs.
Both ``default`` and ``controlnet`` use the SDXL base checkpoint -- the canny
ControlNet (``CONTROLNET_CANNY_MODEL``) is an add-on loaded on top of it, not a
separate base model.
"""
normalized = profile.strip().lower()
if normalized == "default":
if normalized in ("default", "controlnet"):
return DEFAULT_MODEL_ID
if normalized == "ctrlregen":
return CTRLREGEN_MODEL_ID
raise ValueError(f"Unknown model profile '{profile}'. Use one of: default, ctrlregen.")
def detect_model_profile(model_id: str) -> str:
"""Infer model profile from model identifier."""
if "ctrlregen" in model_id.lower():
return "ctrlregen"
return "default"
raise ValueError(f"Unknown model profile '{profile}'. Use one of: default, controlnet.")
src/remove_ai_watermarks/noai/watermark_remover.py
+152 -339
View File
@@ -1,13 +1,13 @@
"""Watermark removal using diffusion model regeneration attack.
Based on the paper "Image Watermarks Are Removable Using Controllable
Regeneration from Clean Noise" (ICLR 2025).
This module implements a simple regeneration attack that:
1. Encodes the watermarked image to latent space
2. Adds noise via forward diffusion process
3. Denoises via reverse diffusion process
4. Decodes back to pixel space
Two pipelines:
1. ``default`` -- plain SDXL img2img. Partial-noise regeneration scrubs the
invisible watermark; ``strength`` controls how much is regenerated.
2. ``controlnet`` -- 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
pixels are ever copied or frozen, so SynthID does not survive.
``controlnet_conditioning_scale`` is the preservation knob.
"""
# torch/diffusers/cv2 boundary: these libs ship no usable types for the tensor and
@@ -29,10 +29,9 @@ if TYPE_CHECKING:
from PIL import Image
from remove_ai_watermarks.noai.watermark_profiles import (
CTRLREGEN_MODEL_ID,
CONTROLNET_CANNY_MODEL,
DEFAULT_MODEL_ID,
DEFAULT_STRENGTH,
detect_model_profile,
resolve_strength,
)
@@ -273,21 +272,14 @@ def _make_seed_generator(device: str, seed: int) -> Any:
return torch.Generator().manual_seed(seed) # type: ignore
def _generator_device(generator: Any) -> str:
"""Best-effort device type of a ``torch.Generator`` (e.g. ``"cpu"``, ``"mps"``)."""
device = getattr(generator, "device", None)
return getattr(device, "type", str(device)) if device is not None else "cpu"
# Canny edge thresholds for the ControlNet control image (xinsir canny recipe:
# cv2.Canny(gray, 100, 200) -> a 3-channel edge map).
_CANNY_LOW = 100
_CANNY_HIGH = 200
# Keep legacy name available for backwards compatibility
_detect_model_profile_from_id = detect_model_profile
# SDXL Differential-Diffusion community pipeline, pinned to the installed
# diffusers version so the fetched pipeline code matches the library (see #21).
# Diffusers' dynamic-module loader resolves ``custom_revision`` against the
# package version string (``0.38.0``), NOT the GitHub git tag (``v0.38.0``).
_DIFF_PIPELINE_NAME = "pipeline_stable_diffusion_xl_differential_img2img"
_DIFF_PIPELINE_REVISION = "0.38.0"
# A neutral quality prompt: the goal is faithful regeneration, not creative edits.
_CONTROLNET_PROMPT = "best quality, high quality, sharp, detailed, photographic"
_CONTROLNET_NEGATIVE = "blurry, lowres, deformed, distorted text, garbled text, watermark, jpeg artifacts"
class WatermarkRemover:
@@ -299,8 +291,8 @@ class WatermarkRemover:
"""
DEFAULT_MODEL_ID = DEFAULT_MODEL_ID
CTRLREGEN_MODEL_ID = CTRLREGEN_MODEL_ID
DEFAULT_STRENGTH = DEFAULT_STRENGTH
CONTROLNET_CANNY_MODEL = CONTROLNET_CANNY_MODEL
def __init__(
self,
@@ -309,9 +301,14 @@ class WatermarkRemover:
torch_dtype: Any = None,
progress_callback: Callable[[str], None] | None = None,
hf_token: str | None = None,
pipeline: str = "default",
controlnet_conditioning_scale: float = 1.0,
) -> None:
self.model_id = model_id or self.DEFAULT_MODEL_ID
self.model_profile = detect_model_profile(self.model_id)
# The pipeline profile is threaded explicitly (not inferred from model_id):
# both "default" and "controlnet" use the same SDXL base checkpoint.
self.model_profile = pipeline
self.controlnet_conditioning_scale = controlnet_conditioning_scale
if not is_watermark_removal_available():
_ensure_watermark_deps()
@@ -329,8 +326,7 @@ class WatermarkRemover:
self.torch_dtype = torch_dtype
self._pipeline: AutoImg2ImgPipeline | None = None
self._diff_pipeline: Any = None
self._ctrlregen_engine: Any = None
self._controlnet_pipeline: Any = None
self._progress_callback = progress_callback
self.hf_token: str | None = hf_token or os.environ.get("HF_TOKEN")
@@ -345,44 +341,59 @@ class WatermarkRemover:
def preload(self) -> None:
"""Eagerly load the pipeline so download progress bars are visible."""
if self.model_profile == "ctrlregen":
self._run_ctrlregen_preload()
if self.model_profile == "controlnet":
self._load_controlnet_pipeline()
else:
self._load_pipeline()
def _run_ctrlregen_preload(self) -> None:
"""Ensure the CtrlRegen engine and all its models are loaded."""
from remove_ai_watermarks.noai.ctrlregen import is_ctrlregen_available
if not is_ctrlregen_available():
missing_pkgs = ["controlnet-aux", "color-matcher", "safetensors"]
logger.info("Auto-installing missing CtrlRegen dependencies: %s", missing_pkgs)
if not _auto_install(missing_pkgs):
raise ImportError(
f"Failed to auto-install missing dependencies: {', '.join(missing_pkgs)}. "
"Try manually: pip install --force-reinstall noai-watermark"
)
if self._ctrlregen_engine is None:
self._ctrlregen_engine = self._make_ctrlregen_engine()
self._ctrlregen_engine.load()
def _make_ctrlregen_engine(self) -> Any:
"""Create a new CtrlRegenEngine with current settings."""
from remove_ai_watermarks.noai.ctrlregen import CtrlRegenEngine
base_model = self.model_id if self.model_id != self.CTRLREGEN_MODEL_ID else None
return CtrlRegenEngine(
base_model_id=base_model,
device=self.device,
torch_dtype=self.torch_dtype,
hf_token=self.hf_token,
progress_callback=self._progress_callback,
)
# ── Pipeline loading ─────────────────────────────────────────────
def _maybe_add_fp16_vae(self, load_kwargs: dict[str, Any]) -> None:
"""Swap in the fp16-fixed SDXL VAE for the default checkpoint on a fp16 GPU.
The stock SDXL VAE overflows to NaN in fp16 and decodes to an all-black
image (issue #29). Shared by both pipeline loaders; a no-op on fp32 (cpu/mps)
or a non-SDXL checkpoint.
"""
if _needs_fp16_vae_fix(self.model_id, self.DEFAULT_MODEL_ID, self.torch_dtype == torch.float16):
from diffusers import AutoencoderKL
self._set_progress("Loading fp16-fixed SDXL VAE (avoids black output)...")
load_kwargs["vae"] = AutoencoderKL.from_pretrained(_SDXL_FP16_VAE_ID, torch_dtype=torch.float16)
def _move_to_device_and_optimize(self, pipeline: Any) -> Any:
"""Move a freshly-loaded pipeline to ``self.device`` + enable memory opts.
Shared by both loaders. On a CUDA move failure (missing CUDA torch build),
trigger the torch-CUDA reinstall+restart. Returns the moved pipeline.
"""
self._set_progress(f"Moving model to device: {self.device}")
try:
pipeline = pipeline.to(self.device)
except (RuntimeError, AssertionError) as exc:
if self.device == "cuda" and not os.environ.get(_CUDA_FIX_ENV_KEY):
self._set_progress("CUDA failed. Reinstalling torch with CUDA support...")
_reinstall_torch_cuda_and_restart()
raise RuntimeError(
f"Failed to move model to {self.device} ({exc}). "
"Install CUDA-enabled PyTorch manually:\n"
f" pip install torch --index-url {_detect_cuda_index_url()}"
) from exc
if hasattr(pipeline, "enable_xformers_memory_efficient_attention"):
with contextlib.suppress(Exception):
self._set_progress("Enabling memory optimizations...")
pipeline.enable_xformers_memory_efficient_attention()
# Mac Float32 memory slicing
if self.device == "mps" and hasattr(pipeline, "enable_attention_slicing"):
with contextlib.suppress(Exception):
pipeline.enable_attention_slicing("max")
return pipeline
def _load_pipeline(self) -> AutoImg2ImgPipeline:
"""Load the diffusion pipeline lazily."""
"""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}")
@@ -394,48 +405,47 @@ class WatermarkRemover:
}
if self.hf_token:
load_kwargs["token"] = self.hf_token
self._maybe_add_fp16_vae(load_kwargs)
# Avoid the SDXL fp16 NaN/all-black decode (issue #29) by loading the
# fp16-fixed VAE for the default SDXL checkpoint on a fp16 GPU.
if _needs_fp16_vae_fix(self.model_id, self.DEFAULT_MODEL_ID, self.torch_dtype == torch.float16):
from diffusers import AutoencoderKL
self._set_progress("Loading fp16-fixed SDXL VAE (avoids black output)...")
load_kwargs["vae"] = AutoencoderKL.from_pretrained(_SDXL_FP16_VAE_ID, torch_dtype=torch.float16)
self._pipeline = AutoImg2ImgPipeline.from_pretrained( # type: ignore
self.model_id,
**load_kwargs,
)
self._set_progress(f"Moving model to device: {self.device}")
try:
self._pipeline = self._pipeline.to(self.device) # type: ignore
except (RuntimeError, AssertionError) as exc:
if self.device == "cuda" and not os.environ.get(_CUDA_FIX_ENV_KEY):
self._set_progress("CUDA failed. Reinstalling torch with CUDA support...")
_reinstall_torch_cuda_and_restart()
raise RuntimeError(
f"Failed to move model to {self.device} ({exc}). "
"Install CUDA-enabled PyTorch manually:\n"
f" pip install torch --index-url {_detect_cuda_index_url()}"
) from exc
if hasattr(self._pipeline, "enable_xformers_memory_efficient_attention"):
with contextlib.suppress(Exception):
self._set_progress("Enabling memory optimizations...")
self._pipeline.enable_xformers_memory_efficient_attention() # type: ignore
# Mac Float32 memory slicing
if self.device == "mps" and hasattr(self._pipeline, "enable_attention_slicing"):
with contextlib.suppress(Exception):
self._pipeline.enable_attention_slicing("max")
pipeline = AutoImg2ImgPipeline.from_pretrained(self.model_id, **load_kwargs) # type: ignore
self._pipeline = self._move_to_device_and_optimize(pipeline)
logger.info("Model loaded successfully")
self._set_progress("Model initialized. Preparing input image...")
return self._pipeline # type: ignore
def _load_controlnet_pipeline(self) -> Any:
"""Load the SDXL + canny-ControlNet img2img pipeline lazily.
Mirrors ``_load_pipeline`` (same fp16-fix VAE, device move, attention
slicing via the shared helpers) but loads the canny ControlNet on top of
the SDXL base. The ControlNet only preserves structure via the edge map;
removal still comes from the img2img regeneration (``strength``).
"""
if self._controlnet_pipeline is None:
from diffusers import ControlNetModel, StableDiffusionXLControlNetImg2ImgPipeline
logger.info("Loading SDXL + ControlNet (%s) on %s...", CONTROLNET_CANNY_MODEL, self.device)
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
self._maybe_add_fp16_vae(load_kwargs)
self._set_progress(f"Loading model weights: {self.model_id}")
pipeline = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(self.model_id, **load_kwargs)
pipeline = self._move_to_device_and_optimize(pipeline)
with contextlib.suppress(Exception):
pipeline.set_progress_bar_config(disable=True)
logger.info("ControlNet model loaded successfully")
self._controlnet_pipeline = pipeline
return self._controlnet_pipeline
# ── Core removal ─────────────────────────────────────────────────
def remove_watermark(
@@ -446,7 +456,6 @@ class WatermarkRemover:
num_inference_steps: int = 50,
guidance_scale: float | None = None,
seed: int | None = None,
protect_text: bool = True,
vendor: str | None = None,
) -> Path:
"""Remove watermark from an image using regeneration attack.
@@ -459,10 +468,6 @@ class WatermarkRemover:
num_inference_steps: Number of denoising steps.
guidance_scale: Classifier-free guidance scale.
seed: Random seed for reproducibility.
protect_text: Detect text regions and preserve them via Differential
Diffusion when any are found (SDXL default profile only). On by
default; the detector decides per image, and text-free inputs run
the standard pass at no extra cost.
vendor: SynthID vendor (``"openai"`` / ``"google"`` / None) used to pick the
default strength when ``strength`` is None. Detect it from the ORIGINAL
input with ``watermark_profiles.vendor_for_strength`` before processing
@@ -482,13 +487,13 @@ class WatermarkRemover:
if output_path is None:
output_path = image_path
strength = resolve_strength(strength, self.model_profile, vendor)
strength = resolve_strength(strength, vendor)
if not 0.0 <= strength <= 1.0:
raise ValueError(f"Strength must be between 0.0 and 1.0, got {strength}")
if guidance_scale is None:
guidance_scale = 2.0 if self.model_profile == "ctrlregen" else 7.5
guidance_scale = 7.5
self._set_progress("Loading and preprocessing input image...")
init_image = Image.open(image_path).convert("RGB")
@@ -508,16 +513,8 @@ class WatermarkRemover:
_total_start = time.monotonic()
if self.model_profile == "ctrlregen":
cleaned_image = self._run_ctrlregen(
init_image,
strength,
num_inference_steps,
guidance_scale,
generator,
)
elif protect_text and self._can_protect_text():
cleaned_image = self._run_region_hires(
if self.model_profile == "controlnet":
cleaned_image = self._run_controlnet(
init_image,
strength,
num_inference_steps,
@@ -525,12 +522,6 @@ class WatermarkRemover:
generator,
)
else:
if protect_text:
logger.debug(
"Text protection unavailable "
"(needs the SDXL default model and the cv2 text detector); "
"running standard img2img."
)
cleaned_image = self._run_img2img(
init_image,
strength,
@@ -613,148 +604,25 @@ class WatermarkRemover:
self._pipeline = None
return self._load_pipeline()
# ── Text-protected differential runner ───────────────────────────
# ── ControlNet runner ────────────────────────────────────────────
def _can_protect_text(self) -> bool:
"""True when text protection can run: SDXL default model + cv2 detector."""
from remove_ai_watermarks import text_protector
def _build_canny_control_image(self, init_image: Image.Image) -> Image.Image:
"""Build the canny ControlNet conditioning image (xinsir recipe).
return self.model_id == self.DEFAULT_MODEL_ID and text_protector.is_available()
def _load_differential_pipeline(self) -> Any:
"""Load the SDXL Differential-Diffusion community pipeline lazily."""
if self._diff_pipeline is None:
from diffusers import DiffusionPipeline
self._set_progress("Loading Differential-Diffusion pipeline (protect-text)...")
use_fp16 = self.device in {"mps", "cuda", "xpu"}
load_kwargs: dict[str, Any] = {
"custom_pipeline": _DIFF_PIPELINE_NAME,
"custom_revision": _DIFF_PIPELINE_REVISION,
"torch_dtype": torch.float16 if use_fp16 else torch.float32, # type: ignore[attr-defined]
"use_safetensors": True,
}
if use_fp16:
load_kwargs["variant"] = "fp16"
if self.hf_token:
load_kwargs["token"] = self.hf_token
pipeline = DiffusionPipeline.from_pretrained(self.model_id, **load_kwargs).to(self.device)
# The differential pipeline upcasts the SDXL VAE to fp32 internally
# (the fp16 VAE decodes to NaN/black otherwise), so we add no extra
# VAE handling here. Attention slicing is also left off on MPS: it
# produced NaN latents with this pipeline, and the protect-text pass
# is short enough not to need it.
with contextlib.suppress(Exception):
pipeline.set_progress_bar_config(disable=True)
self._diff_pipeline = pipeline
return self._diff_pipeline
def _reload_differential_on_cpu(self) -> Any:
"""Reload the differential pipeline on CPU after an MPS failure."""
self.device = "cpu"
self.torch_dtype = torch.float32 # type: ignore[assignment]
self._diff_pipeline = None
return self._load_differential_pipeline()
# Region high-res text scrub: defaults tuned so each text block is upscaled
# enough that strokes exceed the VAE's ~8px latent cell, capped so a single
# region never blows past the GPU/MPS memory budget.
_REGION_HIRES_SCALE = 3.0
_REGION_MAX_MEGAPIXELS = 1.3
def _run_region_hires(
self,
init_image: Image.Image,
strength: float,
num_inference_steps: int,
guidance_scale: float,
generator: Any,
) -> Image.Image:
"""Scrub the whole image, then RE-scrub each detected text block at high
resolution and composite it back.
Unlike the Differential-Diffusion path (which freezes text in latent space
and so leaves the watermark intact there), every pixel here is regenerated
-- the watermark is removed everywhere. Small text survives because each
text block is upscaled before its img2img pass, so strokes span more than
one VAE latent cell (the ~8px floor that softens text at native scale);
the scrubbed crop is downscaled and feather-composited back. Falls back to
the plain global scrub when no text is detected.
cv2.Canny on the RGB->gray array, stacked to 3 channels, wrapped as a PIL
image. The edge map only PRESERVES structure; it never copies pixels.
``init_image`` is already RGB (``remove_watermark`` converts on load).
"""
import math
import cv2
import numpy as np
from remove_ai_watermarks import text_protector
rgb = np.array(init_image)
gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
edges = cv2.Canny(gray, _CANNY_LOW, _CANNY_HIGH)
edges_rgb = np.stack([edges, edges, edges], axis=-1)
return Image.fromarray(edges_rgb)
base = self._run_img2img(init_image, strength, num_inference_steps, guidance_scale, generator)
# The base pass may have fallen back from MPS to CPU (it flips
# self.device). The generator was built for the original device, and
# diffusers rejects a device-mismatched generator ("Expected a 'cpu'
# device generator but found 'mps'"), so drop it for the per-region
# passes -- they then seed from the global RNG, which is fine here.
if generator is not None and self.device == "cpu" and _generator_device(generator) != "cpu":
generator = None
bgr = cv2.cvtColor(np.array(init_image), cv2.COLOR_RGB2BGR)
try:
boxes = text_protector.TextProtector().detect_text_boxes(bgr)
except Exception as exc:
logger.warning("Text detection failed (%s); keeping the global scrub.", exc)
return base
if not boxes:
self._set_progress("No text detected; global scrub only.")
return base
width, height = init_image.size
regions = text_protector.merge_text_regions(boxes, height, width)
orig_bgr = cv2.cvtColor(np.array(init_image), cv2.COLOR_RGB2BGR)
out_bgr = cv2.cvtColor(np.array(base), cv2.COLOR_RGB2BGR)
budget = self._REGION_MAX_MEGAPIXELS * 1_000_000
done = 0
for x, y, w, h in regions:
area = max(1, w * h)
# INTEGER scale so the upscale -> scrub -> downscale round-trip is an
# exact dimensional inverse (a fractional factor truncates and shifts
# the composited text ~1-2px, which is invisible but tanks alignment).
scale = int(min(self._REGION_HIRES_SCALE, math.sqrt(budget / area)))
if scale < 2:
# Region too large to even double within the budget: upscaling
# buys nothing here; the global scrub covers it (documented limit
# for very large text areas -- tiling is the future fix).
continue
crop = orig_bgr[y : y + h, x : x + w]
up = cv2.resize(crop, (w * scale, h * scale), interpolation=cv2.INTER_LANCZOS4)
up_pil = Image.fromarray(cv2.cvtColor(up, cv2.COLOR_BGR2RGB))
scrubbed = self._run_img2img(up_pil, strength, num_inference_steps, guidance_scale, generator)
down = cv2.resize(cv2.cvtColor(np.array(scrubbed), cv2.COLOR_RGB2BGR), (w, h), interpolation=cv2.INTER_AREA)
# The up -> scrub -> down round-trip can offset the re-rendered text by
# a pixel or two (the diffusion pipeline rounds dims to a multiple of
# 8, so the inverse resize is not perfectly centered). Phase-correlate
# the patch back to the original crop and translate it so the glyphs
# land exactly where they were -- otherwise a sub-pixel shift garbles
# the composite even though the text is crisp.
cg = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY).astype(np.float32)
dg = cv2.cvtColor(down, cv2.COLOR_BGR2GRAY).astype(np.float32)
(sx, sy), resp = cv2.phaseCorrelate(cg, dg)
# Only correct for the real 1-2px round-trip shift. On a near-flat /
# low-contrast crop phaseCorrelate returns a spurious large offset at
# a tiny response (e.g. (19,19) at resp ~0.005); warping by that
# garbles the composite -- the exact failure this was meant to
# prevent. Gate on both a confident response and a plausible offset.
if resp > 0.3 and abs(sx) < 4 and abs(sy) < 4 and (abs(sx) > 0.1 or abs(sy) > 0.1):
m = np.float32([[1, 0, -sx], [0, 1, -sy]])
down = cv2.warpAffine(down, m, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
out_bgr = text_protector.feather_paste(out_bgr, down, x, y)
done += 1
self._set_progress(f"Re-scrubbed {done}/{len(regions)} text region(s) at high resolution.")
return Image.fromarray(cv2.cvtColor(out_bgr, cv2.COLOR_BGR2RGB))
def _run_differential(
def _run_controlnet(
self,
init_image: Image.Image,
strength: float,
@@ -762,105 +630,50 @@ class WatermarkRemover:
guidance_scale: float,
generator: Any,
) -> Image.Image:
"""Run differential img2img that preserves detected text regions."""
import cv2
import numpy as np
"""Run the SDXL + canny-ControlNet img2img pass.
from remove_ai_watermarks import text_protector
Removal still comes from the img2img regeneration (``strength``); the canny
ControlNet only PRESERVES text and face STRUCTURE via the edge map. No
original pixels are copied/frozen, so SynthID does not survive (canny holds
structure, not face identity). ``controlnet_conditioning_scale`` is the
structure-preservation knob. Shares the img2img runner (live progress +
MPS->CPU fallback) with ``_run_img2img``; the only delta is the extra
ControlNet kwargs (canny control image + conditioning scale + a non-empty
prompt) overlaid via ``extra_kwargs``.
"""
from remove_ai_watermarks.noai.img2img_runner import run_img2img_with_mps_fallback
self._set_progress("Detecting text regions to protect (protect-text)...")
bgr = cv2.cvtColor(np.array(init_image), cv2.COLOR_RGB2BGR)
try:
boxes = text_protector.TextProtector().detect_text_boxes(bgr)
except Exception as exc:
logger.warning("Text detection failed (%s); running standard img2img.", exc)
return self._run_img2img(init_image, strength, num_inference_steps, guidance_scale, generator)
if not boxes:
self._set_progress("No text detected; running standard img2img.")
return self._run_img2img(init_image, strength, num_inference_steps, guidance_scale, generator)
width, height = init_image.size
change_map = text_protector.build_change_map(boxes, height, width)
self._set_progress(f"Protecting {len(boxes)} text region(s) via Differential Diffusion...")
from remove_ai_watermarks.noai.img2img_runner import run_differential_with_mps_fallback
result_image, final_device = run_differential_with_mps_fallback(
load_pipeline=self._load_differential_pipeline,
extra_kwargs = {
"prompt": _CONTROLNET_PROMPT,
"negative_prompt": _CONTROLNET_NEGATIVE,
"control_image": self._build_canny_control_image(init_image),
"controlnet_conditioning_scale": float(self.controlnet_conditioning_scale),
}
result_image, final_device = run_img2img_with_mps_fallback(
load_pipeline=self._load_controlnet_pipeline,
image=init_image,
change_map=change_map,
strength=strength,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
generator=generator,
device=self.device,
set_progress=self._set_progress,
reload_on_cpu=self._reload_differential_on_cpu,
reload_on_cpu=self._reload_controlnet_on_cpu,
extra_kwargs=extra_kwargs,
)
if final_device != self.device:
self.device = final_device
self.torch_dtype = torch.float32 # type: ignore[assignment]
return result_image
# ── CtrlRegen runner ─────────────────────────────────────────────
def _run_ctrlregen(
self,
init_image: Image.Image,
strength: float,
num_inference_steps: int,
guidance_scale: float,
generator: Any,
) -> Image.Image:
"""Run CtrlRegen pipeline with MPS fallback."""
from remove_ai_watermarks.noai.ctrlregen import is_ctrlregen_available
from remove_ai_watermarks.noai.progress import is_mps_error
if not is_ctrlregen_available():
missing_pkgs = ["controlnet-aux", "color-matcher", "safetensors"]
logger.info("Auto-installing missing CtrlRegen dependencies: %s", missing_pkgs)
if not _auto_install(missing_pkgs):
raise ImportError(
f"Failed to auto-install missing dependencies: {', '.join(missing_pkgs)}. "
"Try manually: pip install --force-reinstall noai-watermark"
)
if self._ctrlregen_engine is None:
self._ctrlregen_engine = self._make_ctrlregen_engine()
seed = None
if generator is not None and hasattr(generator, "initial_seed"):
seed = generator.initial_seed()
try:
return self._ctrlregen_engine.run(
image=init_image,
strength=strength,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
seed=seed,
)
except RuntimeError as error:
if self.device == "mps" and is_mps_error(error):
logger.warning("MPS out of memory during CtrlRegen. Falling back to CPU.")
self._set_progress("MPS out of memory! Retrying CtrlRegen on CPU...")
with contextlib.suppress(Exception):
if _HAS_TORCH and hasattr(torch, "mps"):
torch.mps.empty_cache() # type: ignore[attr-defined]
self.device = "cpu"
self.torch_dtype = torch.float32 # type: ignore[assignment]
self._ctrlregen_engine = self._make_ctrlregen_engine()
return self._ctrlregen_engine.run(
image=init_image,
strength=strength,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
seed=seed,
)
raise
def _reload_controlnet_on_cpu(self) -> Any:
"""Reload the controlnet pipeline on CPU after an MPS failure."""
self.device = "cpu"
self.torch_dtype = torch.float32 # type: ignore[assignment]
self._controlnet_pipeline = None
return self._load_controlnet_pipeline()
# ── Batch ────────────────────────────────────────────────────────
@@ -909,9 +722,9 @@ def remove_watermark(
) -> Path:
"""Convenience function to remove watermark from an image.
``strength=None`` lets the profile pick its default: vendor-adaptive for SDXL
``strength=None`` lets the profile pick its vendor-adaptive SDXL default
(0.10 OpenAI / 0.15 Google / 0.15 unknown, from the C2PA SynthID proxy on the
input), clean-noise 1.0 for ctrlregen. Pass a value to override.
input). Pass a value to override.
"""
from remove_ai_watermarks.noai.watermark_profiles import vendor_for_strength
src/remove_ai_watermarks/text_protector.py
-271
View File
@@ -1,271 +0,0 @@
"""Text-region protection for diffusion-based watermark removal.
SDXL img2img (the ``invisible`` pipeline) regenerates every pixel, so small text
and CJK glyphs get deformed at the strengths that defeat SynthID (issue #21).
This module detects text regions and builds a per-pixel "change map" for
Differential Diffusion: the background is regenerated normally while text
regions are largely preserved, so glyphs survive the watermark-removal pass.
Detection uses only OpenCV's DNN module (no torch): the PP-OCRv3 text detector
is a ~2.4 MB ONNX model (Apache-2.0, from opencv_zoo) that is CJK-native and
returns rotated quadrilaterals. The model is downloaded and cached on first use;
it is never bundled in this repo.
Change-map polarity (verified empirically against the differential pipeline):
white (1.0) = PRESERVE the original pixels, black (0.0) = MAXIMUM change. So the
map is black everywhere except the text polygons, which are painted toward
white. ``preserve`` stays below a hard 1.0 freeze by default: SynthID is
designed to survive cropping, so totally freezing text pixels would leave the
watermark intact there. A high-but-partial preserve still scrubs lightly.
"""
# cv2 ships no type stubs; mirror the pragma used by the other cv2-using modules.
# pyright: reportMissingTypeStubs=false, reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false, reportCallIssue=false, reportArgumentType=false, reportReturnType=false
from __future__ import annotations
import logging
import os
import tempfile
import urllib.request
from pathlib import Path
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from numpy.typing import NDArray
logger = logging.getLogger(__name__)
# PP-OCRv3 Chinese text detector (DB head), opencv_zoo, Apache-2.0.
_MODEL_URL = (
"https://github.com/opencv/opencv_zoo/raw/main/models/text_detection_ppocr/text_detection_cn_ppocrv3_2023may.onnx"
)
_MODEL_FILENAME = "text_detection_cn_ppocrv3_2023may.onnx"
# DB detector input: the image is detected at its NATIVE long side, capped at
# this value (rounded to a multiple of 32), never upscaled. A fixed small input
# (the old 736) downscaled large images so far that small text fell below the
# detector's resolution and was missed -- the cause of the "small text still
# distorts" reports (issue #14). Detection is script-agnostic (DB segments text
# *regions*, not characters), so this recall fix applies to every language; the
# only lever that mattered was resolution. 1536 recovers full recall down to
# ~12 px text on a 2048 canvas at ~100 ms on CPU (a fixed 736 missed it); going
# higher buys no measured recall at 2x+ the cost. Benchmarked in
# scripts/text_detection_benchmark.py. Very large canvases with tiny text may
# still need tiling -- a documented limit, not yet built.
_DET_MAX_LONG_SIDE = 1536
# ImageNet mean (x255) and 1/255 scale -- the normalization PP-OCRv3 expects.
_DET_MEAN = (0.485 * 255, 0.456 * 255, 0.406 * 255)
_DET_SCALE = 1 / 255.0
def is_available() -> bool:
"""True when OpenCV's DNN text-detection model is importable."""
try:
import cv2
return hasattr(cv2.dnn, "TextDetectionModel_DB")
except ImportError:
return False
def _cache_dir() -> Path:
"""Local cache directory for the detector model (created on demand)."""
cache = Path.home() / ".cache" / "remove-ai-watermarks"
cache.mkdir(parents=True, exist_ok=True)
return cache
def _model_path() -> Path:
"""Return the cached detector path, downloading it on first use."""
target = _cache_dir() / _MODEL_FILENAME
if target.exists() and target.stat().st_size > 0:
return target
logger.info("Downloading PP-OCRv3 text detector (~2.4 MB) to %s", target)
# Download to a temp file in the same dir, then atomically rename so a
# partial download never leaves a corrupt model cached.
fd, tmp_name = tempfile.mkstemp(dir=str(target.parent), suffix=".onnx.part")
tmp_path = Path(tmp_name)
try:
os.close(fd)
with urllib.request.urlopen(_MODEL_URL) as resp: # noqa: S310 (trusted GitHub URL)
tmp_path.write_bytes(resp.read())
tmp_path.replace(target)
finally:
if tmp_path.exists():
tmp_path.unlink()
return target
def _detection_input_size(height: int, width: int) -> tuple[int, int]:
"""DB-detector input ``(in_w, in_h)`` for an image of the given size.
Detect at the native long side, capped at ``_DET_MAX_LONG_SIDE`` and never
upscaled, each side rounded down to a multiple of 32 (the DB head requires
/32 dims), floored at 32. Pure function so the resolution contract (the
issue #14 small-text recall fix) is unit-testable without the model.
"""
long_side = max(height, width)
scale = min(_DET_MAX_LONG_SIDE, long_side) / long_side
in_w = max((round(width * scale) // 32) * 32, 32)
in_h = max((round(height * scale) // 32) * 32, 32)
return in_w, in_h
def build_change_map(
boxes: list[NDArray[Any]],
height: int,
width: int,
preserve: float = 0.9,
feather: int = 15,
) -> NDArray[Any]:
"""Build a Differential-Diffusion change map from text polygons.
Args:
boxes: Text-region polygons as arrays of (x, y) vertices.
height: Output map height in pixels.
width: Output map width in pixels.
preserve: Map value painted inside text polygons (0..1). White (1.0)
fully preserves the original pixels; the default 0.9 preserves
strongly while still letting a light scrub through.
feather: Gaussian-blur kernel size for soft polygon edges (forced odd).
Returns:
Float32 HxW array in [0, 1]: ~0 in the background (full change),
``preserve`` inside text regions, blended at the edges.
"""
import cv2
import numpy as np
change_map = np.zeros((height, width), np.float32)
if boxes:
polys = [np.asarray(b, np.int32) for b in boxes]
cv2.fillPoly(change_map, polys, float(preserve))
if feather > 0:
if feather % 2 == 0:
feather += 1
change_map = cv2.GaussianBlur(change_map, (feather, feather), 0)
# GaussianBlur can overshoot the painted value by a float epsilon; keep
# the contract that the map stays a valid [0, 1] change map.
np.clip(change_map, 0.0, 1.0, out=change_map)
return change_map
def merge_text_regions(
boxes: list[NDArray[Any]],
height: int,
width: int,
dilate_frac: float = 0.012,
pad_frac: float = 0.02,
max_regions: int = 8,
) -> list[tuple[int, int, int, int]]:
"""Group detected text polygons into a few padded axis-aligned rectangles.
The DB detector returns one box per word/line; the region-high-res text scrub
runs a separate diffusion pass per region, so we coalesce nearby boxes into a
handful of *local* blocks (a light dilation merges within a paragraph but not
across the whole image, so each block stays small enough to upscale within a
memory budget). Returns ``(x, y, w, h)`` rects, largest-area first, clipped to
the image and capped at ``max_regions``.
"""
import cv2
import numpy as np
mask = np.zeros((height, width), np.uint8)
if not boxes:
return []
cv2.fillPoly(mask, [np.asarray(b, np.int32) for b in boxes], 1)
k = max(1, int(min(height, width) * dilate_frac))
mask = cv2.dilate(mask, cv2.getStructuringElement(cv2.MORPH_RECT, (k, k)))
n, _labels, stats, _c = cv2.connectedComponentsWithStats(mask, 8)
pad = int(min(height, width) * pad_frac)
rects: list[tuple[int, int, int, int]] = []
for i in range(1, n):
x, y, w, h = (
int(stats[i, cv2.CC_STAT_LEFT]),
int(stats[i, cv2.CC_STAT_TOP]),
int(stats[i, cv2.CC_STAT_WIDTH]),
int(stats[i, cv2.CC_STAT_HEIGHT]),
)
x0, y0 = max(0, x - pad), max(0, y - pad)
x1, y1 = min(width, x + w + pad), min(height, y + h + pad)
rects.append((x0, y0, x1 - x0, y1 - y0))
rects.sort(key=lambda r: -(r[2] * r[3]))
return rects[:max_regions]
def feather_paste(
base: NDArray[Any],
patch: NDArray[Any],
x: int,
y: int,
feather: int = 8,
) -> NDArray[Any]:
"""Alpha-composite ``patch`` into ``base`` at ``(x, y)`` with a feathered edge.
Used to drop a separately re-scrubbed (high-resolution) text region back into
the globally-scrubbed image without a visible seam. Returns a new array;
``base`` is not modified. ``patch`` is clipped to ``base`` bounds.
"""
import numpy as np
out = base.copy()
bh, bw = base.shape[:2]
ph, pw = patch.shape[:2]
x0, y0 = max(0, x), max(0, y)
x1, y1 = min(bw, x + pw), min(bh, y + ph)
if x1 <= x0 or y1 <= y0:
return out
patch_roi = patch[y0 - y : y1 - y, x0 - x : x1 - x].astype(np.float32)
base_roi = out[y0:y1, x0:x1].astype(np.float32)
rh, rw = base_roi.shape[:2]
alpha = np.ones((rh, rw), np.float32)
f = max(0, min(feather, rh // 2, rw // 2))
if f > 0:
ramp = np.linspace(0.0, 1.0, f, dtype=np.float32)
alpha[:f, :] *= ramp[:, None]
alpha[rh - f :, :] *= ramp[::-1, None]
alpha[:, :f] *= ramp[None, :]
alpha[:, rw - f :] *= ramp[None, ::-1]
a3 = alpha[:, :, None]
out[y0:y1, x0:x1] = (patch_roi * a3 + base_roi * (1.0 - a3)).astype(base.dtype)
return out
class TextProtector:
"""Detect text regions with PP-OCRv3 for diffusion change-map protection."""
def __init__(
self,
binary_threshold: float = 0.3,
polygon_threshold: float = 0.5,
max_candidates: int = 200,
unclip_ratio: float = 2.0,
) -> None:
import cv2
self._detector = cv2.dnn.TextDetectionModel_DB(str(_model_path()))
self._detector.setBinaryThreshold(binary_threshold)
self._detector.setPolygonThreshold(polygon_threshold)
self._detector.setMaxCandidates(max_candidates)
self._detector.setUnclipRatio(unclip_ratio)
def detect_text_boxes(self, bgr_image: NDArray[Any]) -> list[NDArray[Any]]:
"""Detect text regions, returning a list of rotated quad polygons.
Args:
bgr_image: Image as an HxWx3 BGR uint8 array (OpenCV convention).
Returns:
One array of four (x, y) vertices per detected text region.
"""
height, width = bgr_image.shape[:2]
in_w, in_h = _detection_input_size(height, width)
self._detector.setInputParams(
scale=_DET_SCALE,
size=(in_w, in_h),
mean=_DET_MEAN,
swapRB=True,
)
boxes, _confidences = self._detector.detect(bgr_image)
return list(boxes)