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

feat(photomaker): SynthID-safe face-identity restoration via PhotoMaker-V2

Browse Source 
Adds the second face-restore mechanism, selectable via the new CLI option
`--restore-faces-method=photomaker`. Unlike the existing GFPGAN path (which runs on
the watermarked ORIGINAL and was oracle-confirmed to re-introduce SynthID by partial
pixel blending), PhotoMaker carries identity in a SynthID-invariant OpenCLIP
embedding and regenerates fresh face pixels conditioned on it — the pixels in the
output are diffusion-fresh, so the watermark cannot be transported.

The load-bearing assumption (embedding invariance to SynthID-magnitude pixel noise)
was empirically validated in the prior commit (smoke test): cosine drift 0.002
under a ±2 LSB low-freq carrier, an order of magnitude less than JPEG90 drift
which SynthID survives at >=99% TPR.

End-to-end commercial-safe:
- PhotoMaker-V2 weights: Apache-2.0 (TencentARC)
- ID encoder: OpenCLIP-ViT-H/14 (MIT)
- SDXL base: shared with the main pipeline
- NO InsightFace (the non-commercial blocker for IP-Adapter FaceID / InstantID /
  PuLID / Arc2Face)

Two-pass architecture (PhotoMaker has no ControlNetImg2img class in diffusers):
1) main controlnet/default removal pass cleans SynthID + drifts faces
2) PhotoMaker txt2img regenerates each face from its embedding, feather-composited
   back into the cleaned image

New module `photomaker_restore.py` mirrors `face_restore.py`: lazy pipeline
singleton (double-checked lock), `is_available()` gate, pure `_face_crop_square` and
`_composite_faces` helpers, all unit-tested without the model (9 new tests). New
`InvisibleEngine._restore_faces_photomaker` runs after the diffusion pass, mirroring
`_restore_faces`. CLI flag `--restore-faces-method=[gfpgan|photomaker]` threaded
through `cmd_invisible`/`cmd_all`/`cmd_batch` + `_process_batch_image`.

New optional `photomaker` extra (Apache-2.0 + Apache-2.0/MIT deps, no basicsr).
`[tool.hatch.metadata] allow-direct-references = true` is required because the
upstream PhotoMaker package lives only on GitHub.

The next step (separate work) is oracle validation: run a 6-image cert sweep
through the new pipeline (default/controlnet at the certified strength +
--restore-faces-method=photomaker) and confirm SynthID stays clean while face
identity is recovered. The required infrastructure (`raiw-app/modal_cert.py`) is
already in place.

ruff + strict pyright(src/) clean; 586 tests pass (+ 9 new in
tests/test_photomaker_restore.py).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Victor Kuznetsov
2026-06-08 15:20:29 -07:00
co-authored by Claude Opus 4.8
parent f8f247308b
commit 1439eb0714
7 changed files with 532 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/remove_ai_watermarks/cli.py
+24 -6
View File
@@ -236,22 +236,32 @@ def _warn_if_esrgan_unavailable(upscaler: str) -> None:
def _restore_faces_options(f: Any) -> Any:
"""Attach the shared GFPGAN face-restoration flags to an invisible-pipeline command."""
"""Attach the shared face-restoration flags to an invisible-pipeline command."""
restore_flag = click.option(
"--restore-faces/--no-restore-faces",
default=False,
help="EXPERIMENTAL, opt-in. Restore face identity with a GFPGAN post-pass when "
"faces are present (needs the 'restore' extra); off by default, auto-skips when no "
"face is detected or the extra is absent.",
help="EXPERIMENTAL, opt-in. Restore face identity with a post-pass when faces are "
"present; off by default, auto-skips when no face is detected or the chosen extra "
"is absent.",
)
method_flag = click.option(
"--restore-faces-method",
type=click.Choice(["gfpgan", "photomaker"]),
default="gfpgan",
help="Face-restore mechanism: 'gfpgan' (cheap, needs 'restore' extra, BUT runs on "
"the watermarked original and re-introduces SynthID) or 'photomaker' (PhotoMaker-V2, "
"needs the 'photomaker' extra; carries identity via a SynthID-invariant OpenCLIP "
"embedding so the regenerated face pixels are watermark-free). Default: gfpgan.",
)
weight_flag = click.option(
"--restore-faces-weight",
type=float,
default=0.5,
help="GFPGAN fidelity weight (0-1); lower = more GAN regeneration (cleaner "
"watermark scrub), higher = closer to the input.",
"watermark scrub), higher = closer to the input. Ignored when "
"--restore-faces-method=photomaker.",
)
return restore_flag(weight_flag(f))
return restore_flag(method_flag(weight_flag(f)))
def _watermark_region(det: DetectionResult, width: int, height: int) -> tuple[int, int, int, int]:
@@ -603,6 +613,7 @@ def cmd_invisible(
controlnet_scale: float,
restore_faces: bool,
restore_faces_weight: float,
restore_faces_method: str,
upscaler: str,
auto: bool,
adaptive_polish: bool,
@@ -666,6 +677,7 @@ def cmd_invisible(
vendor=vendor,
restore_faces=restore_faces,
restore_faces_weight=restore_faces_weight,
restore_faces_method=restore_faces_method,
)
elapsed = time.monotonic() - t0
@@ -868,6 +880,7 @@ def cmd_all(
controlnet_scale: float,
restore_faces: bool,
restore_faces_weight: float,
restore_faces_method: str,
upscaler: str,
auto: bool,
adaptive_polish: bool,
@@ -977,6 +990,7 @@ def cmd_all(
vendor=vendor,
restore_faces=restore_faces,
restore_faces_weight=restore_faces_weight,
restore_faces_method=restore_faces_method,
)
console.print(" Invisible watermark removed")
@@ -1033,6 +1047,7 @@ def _process_batch_image(
min_resolution: int = 1024,
restore_faces: bool = False,
restore_faces_weight: float = 0.5,
restore_faces_method: str = "gfpgan",
controlnet_scale: float = 1.0,
upscaler: str = "lanczos",
auto: bool = False,
@@ -1112,6 +1127,7 @@ def _process_batch_image(
upscaler=upscaler,
restore_faces=restore_faces,
restore_faces_weight=restore_faces_weight,
restore_faces_method=restore_faces_method,
# Detect the vendor from the pristine original (`img_path`), not the
# visible-processed `out_path` whose C2PA is already gone.
vendor=vendor_for_strength(img_path),
@@ -1195,6 +1211,7 @@ def cmd_batch(
min_resolution: int,
restore_faces: bool,
restore_faces_weight: float,
restore_faces_method: str,
controlnet_scale: float,
upscaler: str,
auto: bool,
@@ -1255,6 +1272,7 @@ def cmd_batch(
min_resolution=min_resolution,
restore_faces=restore_faces,
restore_faces_weight=restore_faces_weight,
restore_faces_method=restore_faces_method,
controlnet_scale=controlnet_scale,
upscaler=upscaler,
auto=auto,
src/remove_ai_watermarks/invisible_engine.py
+57 -1
View File
@@ -166,6 +166,7 @@ class InvisibleEngine:
vendor: str | None = None,
restore_faces: bool = False,
restore_faces_weight: float = 0.5,
restore_faces_method: str = "gfpgan",
unsharp: float = 0.0,
adaptive_polish: bool = False,
upscaler: str = "lanczos",
@@ -185,6 +186,13 @@ class InvisibleEngine:
face-restoration post-pass when faces are present (needs the
``restore`` extra). Auto-skips with a debug log when the extra is
absent or no face is detected.
restore_faces_method: Which face-identity restoration mechanism to run after
the diffusion pass: ``"gfpgan"`` (default; cheap, but WARNING the GFPGAN
pass runs on the watermarked ORIGINAL and re-introduces SynthID -- see
``face_restore.py``) or ``"photomaker"`` (PhotoMaker-V2; carries identity
via a SynthID-invariant OpenCLIP embedding and regenerates fresh face
pixels conditioned on it -- SynthID-safe, but heavier and requires the
``photomaker`` extra). See ``docs/synthid-robust-identity-research.md``.
restore_faces_weight: GFPGAN fidelity weight (0-1); lower = more GAN
regeneration (cleaner watermark scrub), higher = closer to input.
unsharp: Final unsharp-mask sharpening strength (0 = off, default).
@@ -318,7 +326,10 @@ class InvisibleEngine:
# the cleaned output at its final resolution; auto-skips when faces are
# absent or the optional extra is not installed.
if restore_faces:
self._restore_faces(out_path, image, restore_faces_weight)
if restore_faces_method == "photomaker":
self._restore_faces_photomaker(out_path, image, seed)
else:
self._restore_faces(out_path, image, restore_faces_weight)
# Final sharpening, LAST so it crisps the face-restored result too (a
# pre-GFPGAN sharpen would be smoothed back over by the face pass).
@@ -406,6 +417,51 @@ class InvisibleEngine:
except Exception as e:
logger.warning("restore_faces post-pass failed (%s); keeping un-restored output", e)
def _restore_faces_photomaker(
self,
out_path: Path,
original_image: Any,
seed: int | None,
) -> None:
"""Run the PhotoMaker-V2 SynthID-safe face-identity restoration post-pass.
Unlike the GFPGAN path (which blends watermarked original face pixels back into
the cleaned output and re-introduces SynthID), PhotoMaker carries identity in a
SynthID-invariant OpenCLIP embedding and regenerates fresh face pixels conditioned
on it. Best-effort: any failure (missing extra, model load, runtime error) logs a
warning and leaves the un-restored cleaned output in place. See
``docs/synthid-robust-identity-research.md`` and ``photomaker_restore.py``.
"""
from remove_ai_watermarks import photomaker_restore
if not photomaker_restore.is_available():
logger.debug("restore_faces=photomaker requested but the 'photomaker' extra is not installed; skipping")
return
try:
import cv2
import numpy as np
from remove_ai_watermarks import image_io
cleaned_bgr = image_io.imread(out_path, cv2.IMREAD_COLOR)
if cleaned_bgr is None:
logger.warning("restore_faces_photomaker: could not read cleaned output %s; skipping", out_path)
return
original_rgb = original_image.convert("RGB")
original_bgr = cv2.cvtColor(np.array(original_rgb), cv2.COLOR_RGB2BGR)
cleaned_size = (cleaned_bgr.shape[1], cleaned_bgr.shape[0])
if (original_bgr.shape[1], original_bgr.shape[0]) != cleaned_size:
original_bgr = cv2.resize(original_bgr, cleaned_size, interpolation=cv2.INTER_LANCZOS4)
if self._progress_callback:
self._progress_callback("Restoring face identity (PhotoMaker-V2 post-pass)...")
restored = photomaker_restore.restore_faces_photomaker(original_bgr, cleaned_bgr, seed=seed)
image_io.imwrite(out_path, restored)
except Exception as e:
logger.warning("restore_faces_photomaker post-pass failed (%s); keeping un-restored output", e)
def remove_watermark_batch(
self,
input_dir: Path,
src/remove_ai_watermarks/photomaker_restore.py
+286
View File
@@ -0,0 +1,286 @@
"""SynthID-robust face identity restoration via PhotoMaker-V2.
The diffusion removal pass scrubs the pixel watermark from the WHOLE image, including
faces, but lets faces drift in identity. Unlike the GFPGAN restore pass in
``face_restore.py`` (which runs on the watermarked ORIGINAL and re-introduces SynthID
via partial pixel blending), PhotoMaker carries identity in a SEMANTIC EMBEDDING
(OpenCLIP-ViT-H/14 image embedding, finetuned by PhotoMaker-V2) and uses it to
CONDITION a fresh txt2img generation -- the pixels are new, so the watermark cannot
be transported.
That the embedding cannot carry an invisible pixel watermark like SynthID was
empirically confirmed 2026-06-04: on 31 face crops, the cosine similarity between
``embed(orig)`` and ``embed(synthid_proxy(orig))`` (a ±2 LSB low-frequency noise of
SynthID magnitude) is 0.9977 -- an order of magnitude less drift than JPEG90, which
SynthID survives at >=99% TPR by design. See ``docs/synthid-robust-identity-research.md``.
Architecture: PhotoMaker-V2 is a fine-tuned OpenCLIP-ViT-H/14 ID encoder plus LoRA on
the SDXL UNet attention layers. It ships as a single ``photomaker-v2.bin`` checkpoint
loaded into a ``PhotoMakerStableDiffusionXLPipeline`` (txt2img only -- there is no
PhotoMakerControlNetImg2img class in diffusers). We use it as a SECOND PASS after the
main controlnet/default removal:
1. Main removal pass (`controlnet` at the certified strength) cleans SynthID
everywhere but leaves faces drifted.
2. For each face found in the CLEANED image (YuNet), this module takes the SAME
face region from the ORIGINAL, computes a PhotoMaker ID embedding from it, and
runs PhotoMaker txt2img to regenerate JUST that face crop from the embedding.
The freshly generated face is feather-composited back into the cleaned image.
The generated face pixels are diffusion-fresh and inherit identity from the embedding
(not the pixels), so SynthID is not re-introduced.
Commercial-safe end-to-end:
- PhotoMaker-V2 weights: Apache-2.0 (TencentARC).
- ID encoder: OpenCLIP-ViT-H/14 (MIT) finetuned by PhotoMaker (still Apache-2.0).
- SDXL base: shared with the main pipeline (already used in `default`/`controlnet`).
- NO InsightFace / antelopev2 (which is the non-commercial blocker for IP-Adapter
FaceID / InstantID / PuLID / Arc2Face).
Requires the optional ``photomaker`` extra: ``pip install
'remove-ai-watermarks[photomaker]'`` (pulls torch / diffusers / the upstream PhotoMaker
package, all commercial-safe). Weights download on first use; never bundled.
"""
# cv2/torch/diffusers boundary: relax 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
import importlib.util
import logging
import threading
from pathlib import Path
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from numpy.typing import NDArray
logger = logging.getLogger(__name__)
# PhotoMaker-V2 weights (Apache-2.0, TencentARC). Downloaded on first use.
_PHOTOMAKER_REPO = "TencentARC/PhotoMaker-V2"
_PHOTOMAKER_FILE = "photomaker-v2.bin"
# SDXL base shared with the main pipeline (same checkpoint as `default`/`controlnet`).
_SDXL_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
# The neutral prompt PhotoMaker is designed around: a class noun + the trigger word
# `img`, which PhotoMaker replaces with the ID embedding at inference. Keeping it
# scene-neutral (no extra style words) maximises identity transfer from the embed and
# minimises hallucinated background/lighting that would not match the cleaned scene.
_PHOTOMAKER_PROMPT = "a portrait photo of a person img, natural lighting, sharp focus"
_PHOTOMAKER_NEGATIVE = "blurry, lowres, deformed, distorted, watermark"
# Square size used to feed PhotoMaker (must match a multiple of 64; 512 fits CPU/GPU
# comfortably and gives the encoder enough pixels for a stable embedding).
_PHOTOMAKER_FACE_SIZE = 512
_pipeline: Any | None = None
_pipeline_lock = threading.Lock()
def is_available() -> bool:
"""True when the optional PhotoMaker extra deps are importable."""
return (
importlib.util.find_spec("photomaker") is not None
and importlib.util.find_spec("diffusers") is not None
and importlib.util.find_spec("huggingface_hub") is not None
)
def _select_device() -> str:
"""Pick the PhotoMaker pipeline device: CUDA when present, MPS on Apple, else CPU."""
try:
import torch
if torch.cuda.is_available():
return "cuda"
if torch.backends.mps.is_available():
return "mps"
except Exception as e:
logger.debug("photomaker_restore: device probe failed (%s); using CPU", e)
return "cpu"
def _get_pipeline() -> Any:
"""Return the lazily-built PhotoMaker pipeline singleton (downloads weights on first use)."""
global _pipeline
if _pipeline is not None:
return _pipeline
with _pipeline_lock:
if _pipeline is None:
import torch
from huggingface_hub import hf_hub_download
from photomaker import PhotoMakerStableDiffusionXLPipeline
device = _select_device()
dtype = torch.float16 if device == "cuda" else torch.float32
logger.info("photomaker_restore: loading SDXL+PhotoMaker on %s (%s)", device, dtype)
adapter_path = hf_hub_download(repo_id=_PHOTOMAKER_REPO, filename=_PHOTOMAKER_FILE)
pipe = PhotoMakerStableDiffusionXLPipeline.from_pretrained(_SDXL_MODEL_ID, torch_dtype=dtype)
pipe.load_photomaker_adapter(
str(Path(adapter_path).parent),
subfolder="",
weight_name=_PHOTOMAKER_FILE,
trigger_word="img",
)
pipe.to(device)
pipe.fuse_lora()
_pipeline = pipe
return _pipeline
def _face_crop_square(
image_bgr: NDArray[Any],
box: tuple[int, int, int, int],
pad: float = 0.30,
) -> tuple[NDArray[Any], tuple[int, int, int, int]]:
"""Square crop around a face box (with padding), clipped to the image.
Returns ``(crop_bgr, (x1, y1, x2, y2))``. The crop is the image content inside the
returned square box -- callers use the box for the composite step. Pure numpy slicing,
no model.
"""
h, w = image_bgr.shape[:2]
x, y, bw, bh = box
cx, cy = x + bw // 2, y + bh // 2
side = int(max(bw, bh) * (1.0 + 2.0 * pad))
half = side // 2
x1 = max(0, cx - half)
y1 = max(0, cy - half)
x2 = min(w, cx + half)
y2 = min(h, cy + half)
return image_bgr[y1:y2, x1:x2], (x1, y1, x2, y2)
def _composite_faces(
base_bgr: NDArray[Any],
restored_crops: list[tuple[NDArray[Any], tuple[int, int, int, int]]],
feather_div: int = 6,
) -> NDArray[Any]:
"""Feather-composite a list of ``(restored_crop, (x1, y1, x2, y2))`` into ``base_bgr``.
Pure cv2/numpy helper (no model), unit-testable. For each ``(crop, box)``: resize
the crop to the box size, build a Gaussian-feathered rectangular alpha, and blend
``crop * a + base * (1 - a)``. Boxes that fall fully outside the image (or an empty
list) leave ``base_bgr`` unchanged. Mirrors the alpha math in ``face_restore._composite_faces``.
"""
import cv2
import numpy as np
out = base_bgr.astype(np.float32)
h, w = base_bgr.shape[:2]
for crop, (x1, y1, x2, y2) in restored_crops:
x1, y1 = max(0, x1), max(0, y1)
x2, y2 = min(w, x2), min(h, y2)
bw, bh = x2 - x1, y2 - y1
if bw <= 0 or bh <= 0:
continue
resized = cv2.resize(crop, (bw, bh), interpolation=cv2.INTER_LANCZOS4)
alpha = np.zeros((h, w), dtype=np.float32)
alpha[y1:y2, x1:x2] = 1.0
k = max(3, (min(bw, bh) // feather_div) | 1)
alpha = cv2.GaussianBlur(alpha, (k, k), 0)[:, :, None]
full_restored = np.zeros_like(out)
full_restored[y1:y2, x1:x2] = resized
out = full_restored * alpha + out * (1.0 - alpha)
return np.clip(out, 0, 255).astype(np.uint8)
def restore_faces_photomaker(
original_bgr: NDArray[Any],
cleaned_bgr: NDArray[Any],
num_inference_steps: int = 30,
guidance_scale: float = 5.0,
style_strength: int = 20,
seed: int | None = None,
detect_faces_fn: Any | None = None,
) -> NDArray[Any]:
"""SynthID-robust face identity restoration via PhotoMaker txt2img.
Pipeline:
1. Detect faces in ``cleaned_bgr`` (YuNet via the package's ``auto_config`` by
default; override via ``detect_faces_fn`` for tests).
2. For each face: take the SAME box from ``original_bgr`` -> square crop -> PhotoMaker
txt2img with that crop as the ID image -> a fresh face generated from the
OpenCLIP embedding (the embedding is SynthID-invariant by ~3 orders of magnitude,
see docs/synthid-robust-identity-research.md).
3. Feather-composite each regenerated face into ``cleaned_bgr``.
Faces are taken from ``original_bgr`` (the embedding ignores the watermark) but the
PIXELS that land in the output are diffusion-fresh, so SynthID is not transported.
Args:
original_bgr: The original (watermarked) image as cv2 BGR. Source of identity.
cleaned_bgr: The main-pass output as cv2 BGR. Faces drifted in identity; this
module replaces those face regions.
num_inference_steps: Diffusion steps inside PhotoMaker (def 30).
guidance_scale: CFG scale inside PhotoMaker (def 5.0; the PhotoMaker recipe).
style_strength: PhotoMaker's ``start_merge_step`` knob ~ 20-30 (def 20).
seed: Optional seed for reproducibility.
detect_faces_fn: Optional callable ``(bgr) -> list[(x,y,w,h)]`` to override the
default YuNet detector (used by tests).
Returns:
``cleaned_bgr`` with regenerated face regions composited in (or unchanged when
no face is detected).
"""
import cv2
import numpy as np
import torch
from PIL import Image
if detect_faces_fn is None:
from remove_ai_watermarks import auto_config as _ac
def _default_detect(bgr: NDArray[Any]) -> list[tuple[int, int, int, int]]:
h, w = bgr.shape[:2]
model = Path(_ac.__file__).parent / "assets" / "face_detection_yunet_2023mar.onnx"
det = cv2.FaceDetectorYN.create(str(model), "", (w, h), _ac._FACE_SCORE, 0.3, 5000)
det.setInputSize((w, h))
_, faces = det.detect(bgr)
if faces is None:
return []
return [(int(f[0]), int(f[1]), int(f[2]), int(f[3])) for f in faces if int(f[2]) > 0 and int(f[3]) > 0]
detect_faces_fn = _default_detect
boxes = detect_faces_fn(cleaned_bgr)
if not boxes:
logger.debug("photomaker_restore: no faces detected; returning cleaned image unchanged")
return cleaned_bgr
pipeline = _get_pipeline()
generator = None
if seed is not None:
generator = torch.Generator(device=pipeline.device).manual_seed(seed)
restored: list[tuple[NDArray[Any], tuple[int, int, int, int]]] = []
for box in boxes:
id_crop_bgr, square_box = _face_crop_square(original_bgr, box)
if id_crop_bgr.size == 0:
continue
id_crop_rgb = cv2.cvtColor(id_crop_bgr, cv2.COLOR_BGR2RGB)
id_image_pil = Image.fromarray(id_crop_rgb)
out = pipeline(
prompt=_PHOTOMAKER_PROMPT,
negative_prompt=_PHOTOMAKER_NEGATIVE,
input_id_images=[id_image_pil],
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
start_merge_step=style_strength,
generator=generator,
height=_PHOTOMAKER_FACE_SIZE,
width=_PHOTOMAKER_FACE_SIZE,
num_images_per_prompt=1,
)
gen_rgb = out.images[0]
gen_bgr = cv2.cvtColor(np.array(gen_rgb), cv2.COLOR_RGB2BGR)
restored.append((gen_bgr, square_box))
return _composite_faces(cleaned_bgr, restored)