CalvinBackup/remove-ai-watermarks
1
0
Fork 0
mirror of https://github.com/wiltodelta/remove-ai-watermarks.git synced 2026-07-25 17:00:57 +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 15:55:00 -07:00
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
scripts/controlnet_sweep.py
+348
View File
@@ -0,0 +1,348 @@
"""ControlNet-as-removal-pipeline prototype sweep (issue #35 / Jacob).
Research prototype, NOT a shipped pipeline. It tests whether a full-image
SDXL-native ControlNet-conditioned img2img can REPLACE plain SDXL img2img as the
watermark remover: a single structure-guided regeneration that scrubs an invisible
robust watermark (SynthID) everywhere while keeping fine detail and small/CJK text
legible. See docs/controlnet-removal-pipeline-research.md for the full rationale.
The make-or-break tension (from the watermark-removal-attack literature): the
denoise strength high enough to scrub the watermark deforms text, while the
conditioning strong enough to keep text may spare the watermark. There is no local
SynthID detector, so this script CANNOT decide removal on its own -- it produces
one output per (control, strength, conditioning-scale) cell plus an index, and YOU
verify each output by hand in the Gemini app ("Verify with SynthID") and judge text
legibility visually. Fill the verdict columns in the emitted index, then read off
the Pareto cell (oracle clean AND text legible).
Pipeline: stabilityai/stable-diffusion-xl-base-1.0 +
- canny: xinsir/controlnet-canny-sdxl-1.0 (control = cv2.Canny(gray, 100, 200))
- tile: xinsir/controlnet-tile-sdxl-1.0 (control = the resized original, no preproc)
StableDiffusionXLControlNetImg2ImgPipeline (image=init, control_image=control).
Needs the gpu extra (torch + diffusers) and cv2. Runs locally on 32 GB MPS in
fp32 (MPS fp16 decodes to all-black NaN -- issue #29 -- so fp32 is the default on
mps/cpu, fp16 only on cuda/xpu); a dedicated GPU is not required for 1024 px. Run:
uv run python scripts/controlnet_sweep.py path/to/watermarked.png -o sweep_out
uv run python scripts/controlnet_sweep.py img.png --control canny tile \\
--strength 0.3 0.5 0.7 1.0 --scale 0.6 1.0 --size 1024
"""
from __future__ import annotations
# torch/diffusers/cv2 ship no usable types; relax the unknown-type + private-import
# rules for this boundary script (mirrors scripts/visible_alpha_solve.py and the
# cv2/torch engine modules). Pure-logic helpers here stay correct regardless.
# 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, reportPrivateImportUsage=false
import argparse
import contextlib
import csv
import importlib.util
import logging
import sys
from pathlib import Path
from typing import TYPE_CHECKING, Any
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
from _plain_console import Console, Table
if TYPE_CHECKING:
from PIL import Image
logging.basicConfig(level=logging.INFO, format="%(message)s")
log = logging.getLogger(__name__)
console = Console()
BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
FP16_VAE = "madebyollin/sdxl-vae-fp16-fix"
CONTROLNETS = {
"canny": "xinsir/controlnet-canny-sdxl-1.0",
"tile": "xinsir/controlnet-tile-sdxl-1.0",
}
# A neutral quality prompt: the goal is faithful regeneration, not creative edits.
PROMPT = "best quality, high quality, sharp, detailed, photographic"
NEGATIVE_PROMPT = "blurry, lowres, deformed, distorted text, garbled text, watermark, jpeg artifacts"
def pick_device(requested: str) -> str:
"""Resolve the inference device without the CUDA-reinstaller side effect.
Deliberately does NOT call the package ``get_device`` (which can trigger a
torch-CUDA reinstall+restart). A research script should never do that.
"""
import torch
if requested != "auto":
return requested
if torch.cuda.is_available():
return "cuda"
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
return "mps"
return "cpu"
def resolve_dtype(device: str, requested: str) -> Any:
"""fp16 only on cuda/xpu; fp32 on cpu AND mps, unless overridden.
MPS fp16 produces all-black NaN output here (the SDXL UNet/VAE overflows on
the Metal backend -- issue #29; even the fp16-fix VAE does not save it), so the
production pipeline runs fp32 on MPS and so do we. fp32 SDXL + an SDXL ControlNet
at 1024 fits 32 GB unified memory with vae-tiling + attention-slicing.
"""
import torch
if requested == "fp16":
return torch.float16
if requested == "fp32":
return torch.float32
return torch.float16 if device in {"cuda", "xpu"} else torch.float32
def fit_size(image: Image.Image, long_side: int) -> Image.Image:
"""Resize so the long side is ``long_side``, each dim a multiple of 8 (SDXL)."""
from PIL import Image as PILImage
w, h = image.size
scale = long_side / max(w, h)
nw = max(8, round(w * scale) // 8 * 8)
nh = max(8, round(h * scale) // 8 * 8)
if (nw, nh) == (w, h):
return image
return image.resize((nw, nh), PILImage.Resampling.LANCZOS)
def make_control_image(init: Image.Image, control: str) -> Image.Image:
"""Build the ControlNet conditioning image for the given control type.
canny: cv2.Canny(gray, 100, 200) -> 3-channel edge map (xinsir canny recipe).
tile: the init image itself, no preprocessing (xinsir tile recipe).
"""
import cv2
import numpy as np
from PIL import Image as PILImage
if control == "tile":
return init
rgb = np.array(init.convert("RGB"))
gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
edges = cv2.Canny(gray, 100, 200)
edges_rgb = np.stack([edges, edges, edges], axis=-1)
return PILImage.fromarray(edges_rgb)
def psnr(a: Image.Image, b: Image.Image) -> float:
"""Coarse global fidelity proxy vs the original (NOT a text or watermark metric)."""
import numpy as np
x = np.asarray(a.convert("RGB"), dtype=np.float64)
y = np.asarray(b.convert("RGB").resize(a.size), dtype=np.float64)
mse = float(np.mean((x - y) ** 2))
if mse == 0.0:
return 99.0
return float(10.0 * np.log10((255.0**2) / mse))
def load_pipeline(control: str, device: str, dtype: Any) -> Any:
"""Load SDXL base + the chosen SDXL ControlNet as an img2img pipeline."""
import torch
from diffusers import (
AutoencoderKL,
ControlNetModel,
StableDiffusionXLControlNetImg2ImgPipeline,
)
console.print(f"Loading {CONTROLNETS[control]} ({control}) ...")
controlnet = ControlNetModel.from_pretrained(CONTROLNETS[control], torch_dtype=dtype)
load_kwargs: dict[str, Any] = {"controlnet": controlnet, "torch_dtype": dtype}
if dtype == torch.float16:
# The stock SDXL VAE decodes to NaN/black in fp16; the fp16-fix VAE is the
# same swap the production pipeline uses (_SDXL_FP16_VAE_ID).
load_kwargs["vae"] = AutoencoderKL.from_pretrained(FP16_VAE, torch_dtype=dtype)
pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(BASE_MODEL, **load_kwargs)
pipe = pipe.to(device)
pipe.set_progress_bar_config(disable=True)
if device != "cpu":
# Keep the 1024 px + extra-ControlNet peak inside 32 GB unified memory.
with contextlib.suppress(Exception):
pipe.enable_vae_tiling()
with contextlib.suppress(Exception):
pipe.enable_attention_slicing()
return pipe
def run_cell(
pipe: Any,
init: Image.Image,
control_image: Image.Image,
strength: float,
scale: float,
steps: int,
guidance: float,
seed: int,
) -> Image.Image:
"""Run one ControlNet img2img cell and return the regenerated image.
The generator is created on CPU intentionally: a CPU generator is portable
across the mps/cuda/cpu backends (diffusers rejects a device-mismatched one),
matching the production runner's fallback behavior.
"""
import torch
generator = torch.Generator(device="cpu").manual_seed(seed)
result = pipe(
prompt=PROMPT,
negative_prompt=NEGATIVE_PROMPT,
image=init,
control_image=control_image,
controlnet_conditioning_scale=float(scale),
strength=float(strength),
num_inference_steps=steps,
guidance_scale=guidance,
generator=generator,
)
return result.images[0]
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser(description="ControlNet-as-removal-pipeline prototype sweep.")
p.add_argument("image", type=Path, help="Watermarked input image.")
p.add_argument("-o", "--out", type=Path, default=Path("controlnet_sweep_out"), help="Output directory.")
p.add_argument("--control", nargs="+", choices=list(CONTROLNETS), default=["canny", "tile"])
p.add_argument("--strength", nargs="+", type=float, default=[0.3, 0.5, 0.7, 1.0])
p.add_argument("--scale", nargs="+", type=float, default=[0.6, 1.0], help="controlnet_conditioning_scale values.")
p.add_argument("--size", type=int, default=1024, help="Long-side resolution (multiple of 8).")
p.add_argument("--steps", type=int, default=30)
p.add_argument("--guidance", type=float, default=7.5)
p.add_argument("--seed", type=int, default=0)
p.add_argument("--device", default="auto", choices=["auto", "mps", "cuda", "cpu"])
p.add_argument("--dtype", default="auto", choices=["auto", "fp16", "fp32"])
return p.parse_args()
def main() -> int:
args = parse_args()
if not args.image.exists():
log.error("Input image not found: %s", args.image)
return 1
try:
from PIL import Image as PILImage
except ImportError:
log.error("Pillow is required. Install the gpu extra: uv sync --extra gpu --extra dev")
return 1
if importlib.util.find_spec("diffusers") is None or importlib.util.find_spec("torch") is None:
log.error("diffusers/torch are required. Install: uv sync --extra gpu --extra dev")
return 1
device = pick_device(args.device)
dtype = resolve_dtype(device, args.dtype)
console.print(f"Device: {device} | dtype: {str(dtype).split('.')[-1]}")
init_full = PILImage.open(args.image).convert("RGB")
init = fit_size(init_full, args.size)
console.print(f"Input: {args.image.name} {init_full.size[0]}x{init_full.size[1]} -> {init.size[0]}x{init.size[1]}")
args.out.mkdir(parents=True, exist_ok=True)
stem = args.image.stem
init_path = args.out / f"{stem}__INPUT.png"
init.save(init_path)
rows: list[dict[str, Any]] = []
table = Table(title="ControlNet sweep")
for col in ("control", "strength", "scale", "psnr_vs_input", "file"):
table.add_column(col)
# Group by control so SDXL + the ControlNet load once per control type.
for control in args.control:
pipe = load_pipeline(control, device, dtype)
control_image = make_control_image(init, control)
if control == "canny":
control_image.save(args.out / f"{stem}__canny_edges.png")
for strength in args.strength:
for scale in args.scale:
tag = f"{control}_s{strength:g}_c{scale:g}"
console.print(f"Running {tag} ...")
try:
out = run_cell(
pipe,
init,
control_image,
strength,
scale,
args.steps,
args.guidance,
args.seed,
)
except Exception as exc:
log.warning("Cell %s failed: %s", tag, exc)
continue
fname = f"{stem}__{tag}.png"
out.save(args.out / fname)
quality = psnr(init, out)
rows.append(
{
"control": control,
"strength": strength,
"scale": scale,
"psnr_vs_input": round(quality, 2),
"file": fname,
"synthid_oracle": "", # fill: clean / present
"text_legible": "", # fill: yes / no / partial
}
)
table.add_row(control, f"{strength:g}", f"{scale:g}", f"{quality:.2f}", fname)
del pipe
_free_memory(device)
_write_index(args.out, stem, rows, init_path.name)
console.print(table)
console.print(f"\nWrote {len(rows)} cells to {args.out}/")
console.print(f"Next: open {args.out}/sweep_index.csv, run each PNG through the Gemini SynthID oracle,")
console.print("fill synthid_oracle (clean/present) + text_legible (yes/no/partial), find the Pareto cell.")
return 0
def _free_memory(device: str) -> None:
import gc
gc.collect()
with contextlib.suppress(Exception):
import torch
if device == "cuda":
torch.cuda.empty_cache()
elif device == "mps" and hasattr(torch, "mps"):
torch.mps.empty_cache()
def _write_index(out: Path, stem: str, rows: list[dict[str, Any]], input_name: str) -> None:
"""Write the CSV index (with empty verdict columns) and a README."""
fields = ["control", "strength", "scale", "psnr_vs_input", "file", "synthid_oracle", "text_legible"]
with (out / "sweep_index.csv").open("w", newline="", encoding="utf-8") as fh:
writer = csv.DictWriter(fh, fieldnames=fields)
writer.writeheader()
writer.writerows(rows)
readme = (
f"# ControlNet sweep for {stem}\n\n"
f"Input (resized): {input_name}\n\n"
"Each row in sweep_index.csv is one (control, strength, scale) cell. PSNR vs the\n"
"resized input is a COARSE global-fidelity proxy only -- it does NOT measure text\n"
"legibility or watermark presence. Decide those two by hand:\n\n"
"1. synthid_oracle: open the PNG in the Gemini app, 'Verify with SynthID'. Mark\n"
" 'clean' if no SynthID is detected, 'present' if it still is. (No local detector\n"
" exists; this manual check is the only valid SynthID oracle.)\n"
"2. text_legible: eyeball the small/CJK text. Mark yes / partial / no.\n\n"
"The Pareto cell is the one where synthid_oracle=clean AND text_legible=yes at the\n"
"lowest strength. If no cell satisfies both, the canny/tile-ControlNet middle path\n"
"is dead for text and a glyph re-render is required (see\n"
"docs/text-protection-research.md). Record the outcome in\n"
"docs/controlnet-removal-pipeline-research.md.\n"
)
(out / "sweep_README.md").write_text(readme, encoding="utf-8")
if __name__ == "__main__":
raise SystemExit(main())
scripts/text_detection_benchmark.py
-150
View File
@@ -1,150 +0,0 @@
"""Multilingual recall benchmark for the text-protection detector.
Measures the core lever of text protection (`text_protector.TextProtector`): if
the PP-OCRv3 DB detector misses a text region, that text is NOT preserved during
the SDXL watermark-removal pass and gets deformed (issue #14). This renders short
text in several scripts at several font sizes on two canvas sizes, runs detection,
and reports the fraction of each known text bbox the detector covers.
Findings (2026-05-29):
- Detection is script-agnostic: DB segments text *regions*, not characters, so
Latin / Cyrillic / CJK / Hangul / Arabic / digits score identically. Language
was never the lever.
- The only lever is resolution. A fixed small detector input downscaled large
canvases so far that small text was missed. Detecting at the native long side
(capped, see ``text_protector._DET_MAX_LONG_SIDE``) lifts overall hit-rate
from 0.91 to 1.00 and the worst cell (~16 px text on a 2048 canvas) from
0.06 to 1.00.
This needs the detector model (downloaded on first use) and a font that covers
all the scripts (macOS "Arial Unicode"; on Linux install a Noto super-font).
No GPU. Run:
uv run python scripts/text_detection_benchmark.py
"""
from __future__ import annotations
import logging
import sys
from collections import defaultdict
from pathlib import Path
from typing import Any
import cv2
import numpy as np
from _plain_console import Console
from PIL import Image, ImageDraw, ImageFont
from remove_ai_watermarks import text_protector as tp
log = logging.getLogger(__name__)
console = Console()
# A single font covering every tested script isolates "language" from "font".
_FONT_CANDIDATES = [
"/System/Library/Fonts/Supplemental/Arial Unicode.ttf",
"/Library/Fonts/Arial Unicode.ttf",
"/usr/share/fonts/truetype/noto/NotoSansCJK-Regular.ttc",
"/usr/share/fonts/truetype/unifont/unifont.ttf",
]
SCRIPTS = {
"Latin": "Generated by AI",
"Cyrillic": "Сгенерировано ИИ",
"CJK": "豆包AI生成内容",
"Hangul": "AI로 생성됨",
"Arabic": "أنشئ بالذكاء",
"Digits": "0123456789",
}
FONT_SIZES = [16, 24, 32, 48, 64]
CANVASES = [1024, 2048]
PLACEMENTS = [(0.08, 0.15), (0.30, 0.55), (0.10, 0.82)]
BG_COLORS = [(35, 40, 60), (210, 205, 200)]
def _find_font() -> str:
for path in _FONT_CANDIDATES:
if Path(path).exists():
return path
raise SystemExit(
"No multi-script font found. Install one (macOS ships 'Arial Unicode'; "
"on Linux: a Noto CJK/super font) and add its path to _FONT_CANDIDATES."
)
def _render(
font_path: str,
canvas: int,
text: str,
font_size: int,
place: tuple[float, float],
bg: tuple[int, int, int],
) -> tuple[Image.Image, tuple[int, int, int, int]]:
img = Image.new("RGB", (canvas, canvas), bg)
draw = ImageDraw.Draw(img)
font = ImageFont.truetype(font_path, font_size)
x, y = int(place[0] * canvas), int(place[1] * canvas)
fg = (245, 245, 245) if sum(bg) < 360 else (20, 20, 20)
draw.text((x, y), text, font=font, fill=fg)
return img, draw.textbbox((x, y), text, font=font)
def _coverage(boxes: list[Any], bbox: tuple[int, int, int, int], h: int, w: int) -> float | None:
gt = np.zeros((h, w), np.uint8)
cv2.rectangle(gt, (bbox[0], bbox[1]), (bbox[2], bbox[3]), 1, -1)
area = int(gt.sum())
if area == 0:
return None
det = np.zeros((h, w), np.uint8)
if boxes:
cv2.fillPoly(det, [np.asarray(b, np.int32) for b in boxes], 1)
return int((gt & det).sum()) / area
def _hitrate(values: list[float], thr: float = 0.5) -> float:
return sum(c >= thr for c in values) / len(values) if values else float("nan")
def main() -> int:
logging.basicConfig(level=logging.WARNING)
if not tp.is_available():
raise SystemExit("text detector unavailable (need opencv with cv2.dnn.TextDetectionModel_DB)")
font_path = _find_font()
detector = tp.TextProtector()
by_script_size: dict[tuple[str, int], list[float]] = defaultdict(list)
by_size_canvas: dict[tuple[int, int], list[float]] = defaultdict(list)
by_script: dict[str, list[float]] = defaultdict(list)
for canvas in CANVASES:
for script, text in SCRIPTS.items():
for font_size in FONT_SIZES:
for idx, place in enumerate(PLACEMENTS):
img, bbox = _render(font_path, canvas, text, font_size, place, BG_COLORS[idx % len(BG_COLORS)])
bgr = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
cov = _coverage(detector.detect_text_boxes(bgr), bbox, canvas, canvas)
if cov is None:
continue
by_script_size[(script, font_size)].append(cov)
by_size_canvas[(font_size, canvas)].append(cov)
by_script[script].append(cov)
console.print("=== hit-rate (coverage>=0.5) by script x font-size ===")
console.print("script".ljust(10) + "".join(f"{fs:>7}" for fs in FONT_SIZES))
for script in SCRIPTS:
console.print(
script.ljust(10) + "".join(f"{_hitrate(by_script_size[(script, fs)]):>7.2f}" for fs in FONT_SIZES)
)
console.print("\n=== hit-rate by font-size x canvas (the downscale effect) ===")
console.print("size".ljust(8) + "".join(f"{c:>8}" for c in CANVASES))
for fs in FONT_SIZES:
console.print(str(fs).ljust(8) + "".join(f"{_hitrate(by_size_canvas[(fs, c)]):>8.2f}" for c in CANVASES))
overall = _hitrate([c for vals in by_script.values() for c in vals])
console.print(f"\nOVERALL hit-rate: {overall:.2f} (detector max long side = {tp._DET_MAX_LONG_SIDE})")
return 0
if __name__ == "__main__":
sys.exit(main())