mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-06-10 04:43:54 +02:00
refactor(face-restore): drop GFPGAN, ship PhotoMaker-V2 as the sole restore (non-commercial)
Visual review of the GFPGAN-on-cleaned output (9-face grid, 1448x1086) showed it only polished the already-drifted face without restoring identity — useless for the "restore who is in the photo" intent. Dropping it. The shipped restore path is now PhotoMaker-V2, which delivers true identity-from- embedding face regeneration via a CLIP+ArcFace dual encoder. The ArcFace branch pulls InsightFace antelopev2/buffalo_l model packs at runtime, which InsightFace releases under a research-only license, so the whole extra is **NON-COMMERCIAL**. raiw.cc and any monetized deployment must NOT install the `photomaker` extra. This is called out at every entry point: CLI flag help, module docstring, pyproject extra block, CLAUDE.md extras bullet, README install snippet. Changes: - Deleted `src/remove_ai_watermarks/face_restore.py` and its tests. - Deleted the `restore` extra (gfpgan/facexlib/basicsr + scipy<1.18 / numba<0.60 pins) and the basicsr setuptools<69 build pin from pyproject.toml. - Restored `src/remove_ai_watermarks/photomaker_restore.py` (V2 this time: `TencentARC/PhotoMaker-V2`, `photomaker-v2.bin`, no `pm_version='v1'` override). - Restored the `photomaker` extra in pyproject with all the upstream-compat pins (einops, peft, onnxruntime, insightface) and the `allow-direct-references` hatch metadata block. - `InvisibleEngine` swapped `_restore_faces` -> `_restore_faces_photomaker`; `--restore-faces-method` removed (only one method, no choice). - CLI flag help, CLAUDE.md, README, docs/synthid.md, and docs/controlnet-removal-pipeline-research.md all updated. - docs/synthid-robust-identity-research.md status notice rewritten to list both abandoned commercial-safe attempts (V1 + GFPGAN-on-cleaned) and the non-commercial trade-off we accepted. ruff + strict pyright(src/) clean; 578 tests pass (the 9 GFPGAN tests are gone, the 11 PhotoMaker tests stay green). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Victor Kuznetsov
@@ -23,7 +23,7 @@ If this tool saves you time, consider [sponsoring its development](https://githu
|
||||
- **AI metadata stripping** — EXIF, PNG text chunks, C2PA provenance manifests (PNG / JPEG / AVIF / HEIF / JPEG-XL, **MP4 / MOV / M4V / M4A** at the container level, and **WebM / MP3 / WAV / FLAC / OGG** losslessly via ffmpeg), XMP DigitalSourceType
|
||||
- **"Made with AI" label removal** — removes the AI-disclosure metadata that platforms read to apply automatic labels (useful for clearing a false-positive label from a human-edited photograph)
|
||||
- **Analog Humanizer** — optional film grain and chromatic aberration post-processing
|
||||
- **Text and face preservation (experimental)** — optional `--pipeline controlnet` adds a canny ControlNet that keeps text and face structure sharp through the removal pass (without copying original pixels, so SynthID is still removed). Canny preserves face *structure*, not *identity* (the regenerated face drifts in likeness); face detail is polished by the `--restore-faces` GFPGAN post-pass on the cleaned image (opt-in, SynthID-safe). Both are experimental and off by default.
|
||||
- **Text and face preservation (experimental)** — optional `--pipeline controlnet` adds a canny ControlNet that keeps text and face structure sharp through the removal pass (without copying original pixels, so SynthID is still removed). Canny preserves face *structure*, not *identity* (the regenerated face drifts in likeness); identity is regenerated by the `--restore-faces` PhotoMaker-V2 post-pass (opt-in, **NON-COMMERCIAL** — pulls non-commercial InsightFace model packs). Both are experimental and off by default.
|
||||
- **Batch processing** — process entire directories
|
||||
- **Detection** — three-stage NCC watermark detection with confidence scoring
|
||||
- **Provenance detection (`identify`)** — aggregate C2PA issuer, the C2PA soft-binding forensic-watermark vendor (Adobe TrustMark, Digimarc, Imatag, ...), IPTC "Made with AI" plus the IPTC 2025.1 `AISystemUsed` field, embedded SD/ComfyUI params, EXIF/XMP generator tags, the xAI/Grok EXIF signature, the China TC260 AIGC label (XMP, PNG chunk, or EXIF), the HuggingFace `hf-job-id` job marker, the SynthID metadata proxy, the visible marks (Gemini sparkle plus the Doubao "豆包AI生成" / Jimeng "即梦AI" / Samsung Galaxy AI "Contenuti generati dall'AI" text marks), the open SD/SDXL/FLUX invisible watermark, and (with the `trustmark` extra) the open Adobe TrustMark watermark into one origin-platform + watermark-inventory verdict (`--json` for machine output)
|
||||
@@ -128,7 +128,7 @@ image → encode to latent space (VAE) at native resolution
|
||||
>
|
||||
> **`--pipeline controlnet` preserves text and face structure (experimental, opt-in).** It runs the same SDXL img2img scrub but adds a canny ControlNet that conditions the regeneration on the image's edge map, so text and structure stay sharp at the strengths that remove SynthID. The watermark removal still comes from the img2img regeneration (`--strength`); the ControlNet only preserves structure — no original pixels are copied or frozen, so SynthID does not survive. `--controlnet-scale` tunes the preservation strength (higher = closer to the original structure). Runs fp32 on mps/cpu (fp16 only on cuda/xpu, where the fp16-fixed SDXL VAE is loaded automatically).
|
||||
>
|
||||
> **`--restore-faces` polishes faces on the cleaned image (GFPGAN, experimental, opt-in, SynthID-safe).** Canny preserves where a face is, but not who it is — the regenerated face drifts in likeness. The `--restore-faces` post-pass (experimental, off by default; needs the `restore` extra) runs GFPGAN on the diffusion-CLEANED image (not the original) and feather-composites each polished face into the cleaned result. Because the input pixels GFPGAN derives from are already SynthID-free, the partial pixel-blend at weight 0.5 cannot re-introduce the watermark (this is a fix to the earlier original-source variant that was oracle-confirmed to re-add SynthID). Identity fidelity is limited by GFPGAN's StyleGAN2 prior conditioned on the cleaned face; a true identity-as-embedding stack (PhotoMaker-V1) was researched but blocked by upstream / diffusers-version compatibility issues — see `docs/synthid-robust-identity-research.md`. Commercial-safe: GFPGAN Apache-2.0, RetinaFace MIT; CodeFormer stays non-commercial and is not shipped.
|
||||
> **`--restore-faces` regenerates faces from a CLIP+ArcFace embedding (PhotoMaker-V2, experimental, opt-in, NON-COMMERCIAL).** Canny preserves where a face is, but not who it is — the regenerated face drifts in likeness. The `--restore-faces` post-pass (experimental, off by default; needs the `photomaker` extra) crops each face from the original, feeds it to PhotoMaker-V2 as an identity reference, and regenerates a fresh face from a CLIP+ArcFace embedding which is then feather-composited into the cleaned image. The pixels are diffusion-fresh so SynthID is not re-introduced. **NON-COMMERCIAL:** PhotoMaker-V2's ID encoder pulls InsightFace antelopev2/buffalo_l model packs at runtime, which are released under a research-only license — a paid service must NOT use this flag. (A commercial-safe path was attempted via PhotoMaker-V1 + GFPGAN-on-cleaned but neither was a good fit: V1 hit upstream / diffusers-0.38 compatibility walls, and GFPGAN only polished the already-drifted face without restoring identity. See `docs/synthid-robust-identity-research.md`.)
|
||||
|
||||
SDXL is the default since May 2026: empirically defeats SynthID v2 on Gemini 3 Pro outputs, where the older SD-1.5 pipeline at 768 px did not. The SD-1.5 path was removed once it was verified not to handle v2. Note the scope: this defeats the SynthID *verifier*, which is not the same as being forensically indistinguishable from a real photo. Recent work ([arXiv:2605.09203](https://arxiv.org/abs/2605.09203)) shows watermark-removal pipelines leave detectable traces, so a separate "this image was processed" classifier can still flag the output.
|
||||
|
||||
@@ -136,7 +136,7 @@ SDXL is the default since May 2026: empirically defeats SynthID v2 on Gemini 3 P
|
||||
|
||||
> **Technical deep-dive:** see [`docs/synthid.md`](docs/synthid.md) for a primary-source-cited breakdown of how SynthID works mechanically (post-hoc encoder/decoder, 136-bit payload, pixel-space embedding), what it empirically survives (JPEG, crop, resize: ~99.98% TPR at 0.1% FPR from arXiv:2510.09263), what removes it, and the forensic-stealth tradeoff (all known removal attacks are detectable at >98% TPR@1%FPR per arXiv:2605.09203).
|
||||
|
||||
**Text and face preservation** (experimental, opt-in `--pipeline controlnet`): adds a canny ControlNet so text and face *structure* stay sharp through the removal pass, without copying or freezing any original pixels (so SynthID is still removed). Tune the preservation strength with `--controlnet-scale`. Canny preserves structure but not face *identity* (face detail is polished by the `--restore-faces` GFPGAN post-pass on the cleaned image, experimental and off by default — see the callout above). Both features are experimental.
|
||||
**Text and face preservation** (experimental, opt-in `--pipeline controlnet`): adds a canny ControlNet so text and face *structure* stay sharp through the removal pass, without copying or freezing any original pixels (so SynthID is still removed). Tune the preservation strength with `--controlnet-scale`. Canny preserves structure but not face *identity* (identity is regenerated by the `--restore-faces` PhotoMaker-V2 post-pass, experimental and off by default, **non-commercial** — see the callout above). Both features are experimental.
|
||||
|
||||
**Analog Humanizer**: optional film grain and chromatic aberration injection that mimics a photo of a screen, raising the bar for AI-generated image classifiers. (It frustrates generic classifiers but does not guarantee forensic invisibility — see the [arXiv:2605.09203](https://arxiv.org/abs/2605.09203) note above.)
|
||||
|
||||
@@ -214,13 +214,15 @@ After installation the `remove-ai-watermarks` command is available system-wide.
|
||||
> pip install -e ".[trustmark]" # or: uv pip install -e ".[trustmark]"
|
||||
> ```
|
||||
>
|
||||
> To polish face detail after invisible removal (the `--restore-faces` GFPGAN
|
||||
> post-pass on the cleaned image, experimental and opt-in, SynthID-safe by
|
||||
> construction), install the `restore` extra. The GFPGANv1.4 and RetinaFace weights
|
||||
> download on first use:
|
||||
> To regenerate face identity after invisible removal (the `--restore-faces`
|
||||
> PhotoMaker-V2 post-pass, experimental and opt-in, **NON-COMMERCIAL** because
|
||||
> PhotoMaker-V2 pulls non-commercial InsightFace model packs at runtime), install
|
||||
> the `photomaker` extra. The PhotoMaker-V2 adapter weights and InsightFace face
|
||||
> packs download on first use. Do NOT install this extra in a commercial / paid
|
||||
> service:
|
||||
>
|
||||
> ```bash
|
||||
> pip install -e ".[restore]" # or: uv pip install -e ".[restore]"
|
||||
> pip install -e ".[photomaker]" # or: uv pip install -e ".[photomaker]"
|
||||
> ```
|
||||
>
|
||||
> For sharper upscaling of small inputs before diffusion (`--upscaler esrgan`,
|
||||
|
||||
@@ -124,13 +124,14 @@ Gemini app; the two payloads are vendor-specific and never cross-checked):
|
||||
- **Fix the seed in prod.** The non-determinism is purely `seed=None` (random); a fixed
|
||||
`--seed` makes every run reproduce the certified-clean result, so you ship a
|
||||
deterministic, re-certifiable config (and the seed sweep collapses to one config).
|
||||
- **`--restore-faces` is SynthID-safe by construction now (GFPGAN-on-cleaned, 2026-06-04).**
|
||||
The GFPGAN-on-original path that re-added SynthID was fixed by running GFPGAN on the
|
||||
diffusion-CLEANED image instead — the input pixels GFPGAN derives from are already
|
||||
SynthID-free, so the partial pixel-blend cannot transport the watermark. Needs the
|
||||
`restore` extra. (The PhotoMaker-V1 identity-as-embedding alternative was researched
|
||||
but blocked by upstream / diffusers-version compatibility issues; see
|
||||
`docs/synthid-robust-identity-research.md`.)
|
||||
- **`--restore-faces` is PhotoMaker-V2 (NON-COMMERCIAL).** The GFPGAN-on-cleaned path
|
||||
was tried and rejected: it polished but did not restore identity. PhotoMaker-V2
|
||||
regenerates faces from a CLIP+ArcFace embedding (so pixels are fresh, SynthID is not
|
||||
re-introduced) but pulls InsightFace antelopev2/buffalo_l model packs at runtime,
|
||||
which are research-only. Needs the `photomaker` extra; **a paid service MUST NOT
|
||||
use this flag.** PhotoMaker-V1 was attempted as a commercial-safe alternative but
|
||||
blocked by a CFG batch-dim mismatch in the upstream pipeline (forked from diffusers
|
||||
0.29; we ship 0.38) — see `docs/synthid-robust-identity-research.md`.
|
||||
- **No local SynthID detector exists** → the service can't self-verify; bake in strength
|
||||
margin and periodic oracle spot-checks.
|
||||
- **Lesson:** visual-quality / face-identity recovery does NOT prove removal — only the
|
||||
|
||||
@@ -30,23 +30,39 @@ is the correct commercial-safe target: its `PhotoMakerIDEncoder` (model.py)
|
||||
forward takes only `(id_pixel_values, prompt_embeds, class_tokens_mask)` -- no
|
||||
ArcFace branch -- so identity is CLIP-only.
|
||||
|
||||
**Status notice (2026-06-04, end of session).** Even on V1, the cert sweep hit a
|
||||
cascade of upstream compatibility issues with the diffusers version we ship
|
||||
(0.38): missing `einops` declaration, missing `peft` declaration, default
|
||||
`pm_version='v2'` that mis-loads V1 weights into the V2 encoder, custom
|
||||
`id_encoder` left on CPU after `pipe.to(device)`, and a CFG-batch tensor-shape
|
||||
mismatch in the denoising loop (`Expected size 2 but got size 1`). 7 cascading
|
||||
fixes did not get the pipeline running end-to-end. The PhotoMaker `pipeline.py`
|
||||
header notes it was forked from diffusers v0.29.1; SDXL prompt-encoder handling
|
||||
changed significantly between 0.29 and 0.38, so making this work end-to-end is a
|
||||
proper fork or a diffusers downgrade -- both expensive. **The shipped path is
|
||||
GFPGAN on the diffusion-CLEANED image** (`face_restore.py`, the `restore`
|
||||
extra): a one-line change from the original GFPGAN-on-watermarked design that
|
||||
made the pass SynthID-safe by construction. Identity fidelity is lower than what
|
||||
a working identity-as-embedding stack would deliver, but the pipeline runs, the
|
||||
oracle is satisfied, and the dependency footprint is small. PhotoMaker remains
|
||||
the right north-star for a future identity-fidelity upgrade once the upstream
|
||||
compat work is done (or once a `diffusers ~0.29` forked pipeline is vendored).
|
||||
**Status notice (2026-06-04, end of session).** Two commercial-safe paths were
|
||||
tried and abandoned:
|
||||
|
||||
1. **PhotoMaker-V1** (commercial-safe by license but blocked by upstream compat).
|
||||
The cert sweep hit a cascade of upstream compatibility issues with the diffusers
|
||||
version we ship (0.38): missing `einops` declaration, missing `peft` declaration,
|
||||
default `pm_version='v2'` that mis-loads V1 weights into the V2 encoder, custom
|
||||
`id_encoder` left on CPU after `pipe.to(device)`, and a CFG-batch tensor-shape
|
||||
mismatch in the denoising loop (`Expected size 2 but got size 1`). 7 cascading
|
||||
fixes did not get the pipeline running end-to-end. The PhotoMaker `pipeline.py`
|
||||
header notes it was forked from diffusers v0.29.1; SDXL prompt-encoder handling
|
||||
changed significantly between 0.29 and 0.38.
|
||||
2. **GFPGAN on the diffusion-CLEANED image** (commercial-safe, but no identity
|
||||
recovery). A one-line change made it SynthID-safe (input pixels are already
|
||||
clean, so the partial blend cannot transport the watermark), but visual
|
||||
inspection of the cert output showed it only polished the already-drifted face
|
||||
without actually restoring identity. Trade-off was real and the value too low
|
||||
to keep.
|
||||
|
||||
**The shipped path is PhotoMaker-V2** (`photomaker_restore.py`, the `photomaker`
|
||||
extra). V2 uses a DUAL ID encoder (CLIP image features + ArcFace embedding),
|
||||
which delivers true identity-from-embedding face regeneration. The cost is that
|
||||
the ArcFace embedding comes from InsightFace's `antelopev2`/`buffalo_l` model
|
||||
packs, which are released under a non-commercial / research-only license. **So
|
||||
the shipped restore path is NON-COMMERCIAL.** raiw.cc and any other monetized
|
||||
deployment must NOT install the `photomaker` extra. The CLI flag and module
|
||||
docstring both call this out at every entry point.
|
||||
|
||||
A future commercial-safe path would need either (a) the PhotoMaker upstream to
|
||||
land its diffusers 0.38 compat fix so V1 can run, or (b) an equally good
|
||||
ArcFace-class face-recognition model released under a permissive license that
|
||||
PhotoMaker-V2 can be retargeted to. Neither is on a near-term horizon as of
|
||||
this writing.
|
||||
|
||||
## 1. Why identity-by-embedding (not by pixel) is the only SynthID-robust path
|
||||
|
||||
|
||||
+1
-1
@@ -570,7 +570,7 @@ table.
|
||||
schedule to `resolve_strength`, do not reuse the default ladder; (2) the
|
||||
`--restore-faces` pass is now SynthID-safe by construction (the GFPGAN-on-original
|
||||
path that re-added SynthID was removed 2026-06-04; the shipped restore is
|
||||
GFPGAN-on-cleaned, see `face_restore.py`); (3)
|
||||
PhotoMaker-V2, NON-COMMERCIAL, see `photomaker_restore.py`); (3)
|
||||
removal near threshold is seed-non-deterministic -> FIX the prod seed (kills the
|
||||
coin-flip; ship a deterministic certified config).
|
||||
|
||||
|
||||
+26
-24
@@ -76,22 +76,26 @@ lama = [
|
||||
"onnxruntime>=1.16.0",
|
||||
"huggingface-hub>=0.20.0",
|
||||
]
|
||||
# Optional GFPGAN face-polish post-pass (commercial-safe: GFPGAN Apache-2.0 +
|
||||
# RetinaFace MIT). Polishes face detail in the DIFFUSION-CLEANED image (not the
|
||||
# original) using GFPGAN's StyleGAN2 prior, so SynthID is NOT re-introduced -- the
|
||||
# input pixels GFPGAN derives from are already SynthID-free. This is the shipped path
|
||||
# because the alternative we wanted (PhotoMaker-V1 identity-as-embedding) has
|
||||
# significant upstream / diffusers-version compatibility issues; see
|
||||
# `src/remove_ai_watermarks/face_restore.py` and
|
||||
# `docs/synthid-robust-identity-research.md`. gfpgan/basicsr/facexlib are an OLD
|
||||
# ecosystem and pin numpy<2: scipy<1.18 (>=1.18 uses np.long, gone in numpy 1.24-1.26)
|
||||
# and numba<0.60. Kept OUT of `all` (heavy + model download).
|
||||
restore = [
|
||||
"gfpgan>=1.3.8",
|
||||
"facexlib>=0.3.0",
|
||||
"basicsr>=1.4.2",
|
||||
"scipy<1.18",
|
||||
"numba<0.60",
|
||||
# **NON-COMMERCIAL.** Optional PhotoMaker-V2 SynthID-robust face-identity post-pass.
|
||||
# PhotoMaker-V2's ID encoder pulls an InsightFace ArcFace embedding at runtime, and
|
||||
# the pretrained InsightFace model packs (antelopev2, buffalo_l) are released under a
|
||||
# non-commercial / research-only license. A paid service (raiw.cc, any monetized SaaS)
|
||||
# MUST NOT use this extra. See `src/remove_ai_watermarks/photomaker_restore.py` and
|
||||
# `docs/synthid-robust-identity-research.md`. The PhotoMaker adapter weights
|
||||
# (photomaker-v2.bin) are Apache-2.0 and download on first use; the InsightFace model
|
||||
# packs download on first FaceAnalysis() (only triggered inside PhotoMaker's V2 forward).
|
||||
# Pins beyond the upstream PhotoMaker package itself patch missing declarations that
|
||||
# would otherwise break the load chain (verified empirically via the Modal cert sweep
|
||||
# 2026-06-04): einops (used in forward), peft (required by diffusers.fuse_lora),
|
||||
# onnxruntime (transitive via insightface), and insightface itself (required for the
|
||||
# package's __init__.py to even import).
|
||||
photomaker = [
|
||||
"photomaker @ git+https://github.com/TencentARC/PhotoMaker.git",
|
||||
"huggingface-hub>=0.20.0",
|
||||
"einops>=0.7.0",
|
||||
"insightface>=0.7.3",
|
||||
"onnxruntime>=1.16.0",
|
||||
"peft>=0.10.0",
|
||||
]
|
||||
# Optional pre-diffusion super-resolution for small inputs (Real-ESRGAN). Loaded via
|
||||
# spandrel (MIT) -- a pure model-loader with NO basicsr dependency (it pulls only
|
||||
@@ -121,14 +125,6 @@ all = ["remove-ai-watermarks[gpu,detect,trustmark,lama,dev]"]
|
||||
[tool.uv]
|
||||
prerelease = "allow"
|
||||
|
||||
# basicsr 1.4.2 (pulled by the `restore` GFPGAN extra) ships sdist-only and its
|
||||
# setup.py get_version() reads basicsr/version.py in a way that newer setuptools
|
||||
# (>= 69) breaks with ``KeyError: '__version__'`` under isolated PEP 517 builds.
|
||||
# Pin an old setuptools as its build dependency so the sdist builds; this is
|
||||
# scoped to basicsr and does not affect the rest of the resolution.
|
||||
[tool.uv.extra-build-dependencies]
|
||||
basicsr = ["setuptools<69"]
|
||||
|
||||
# PyTorch Intel-GPU (XPU) wheel index. ``explicit = true`` keeps it inert for
|
||||
# the default CPU/CUDA install: uv consults it only when a torch install
|
||||
# explicitly targets it (see the ``gpu`` extra comment), so it does not alter
|
||||
@@ -159,6 +155,12 @@ Repository = "https://github.com/wiltodelta/remove-ai-watermarks"
|
||||
requires = ["hatchling<1.31"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
# Allow the `photomaker` extra to reference the upstream git URL directly (the
|
||||
# TencentARC/PhotoMaker package is not on PyPI). The extra itself is NON-COMMERCIAL
|
||||
# (see the photomaker block above and `photomaker_restore.py`).
|
||||
[tool.hatch.metadata]
|
||||
allow-direct-references = true
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["src/remove_ai_watermarks"]
|
||||
|
||||
|
||||
@@ -238,16 +238,18 @@ def _warn_if_esrgan_unavailable(upscaler: str) -> None:
|
||||
def _restore_faces_options(f: Any) -> Any:
|
||||
"""Attach the face-restoration flag to an invisible-pipeline command.
|
||||
|
||||
The post-pass runs GFPGAN on the DIFFUSION-CLEANED image (not the original), so
|
||||
SynthID is not re-introduced (the input pixels GFPGAN derives from are already
|
||||
SynthID-free). See ``face_restore.py``.
|
||||
The post-pass uses PhotoMaker-V2 to regenerate each face from a CLIP+ArcFace
|
||||
embedding. **NON-COMMERCIAL** -- PhotoMaker-V2 pulls InsightFace antelopev2/
|
||||
buffalo_l model packs at runtime, which are research-only. A paid service
|
||||
(raiw.cc, any monetized SaaS) MUST NOT use this flag.
|
||||
"""
|
||||
return click.option(
|
||||
"--restore-faces/--no-restore-faces",
|
||||
default=False,
|
||||
help="EXPERIMENTAL, opt-in. Polish face detail with a GFPGAN post-pass on the "
|
||||
"cleaned image 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, **NON-COMMERCIAL** -- needs the 'photomaker' extra "
|
||||
"which pulls non-commercial InsightFace model packs. Restores face identity via "
|
||||
"PhotoMaker-V2 (CLIP+ArcFace embedding -> fresh face); off by default, auto-skips "
|
||||
"when no face is detected or the extra is absent.",
|
||||
)(f)
|
||||
|
||||
|
||||
|
||||
@@ -1,210 +0,0 @@
|
||||
"""Optional GFPGAN face-polish post-pass for the invisible removal pipeline.
|
||||
|
||||
The diffusion removal pass scrubs the watermark everywhere but lets faces drift in
|
||||
likeness (canny holds face *structure*, not *identity*). This module sharpens and
|
||||
re-synthesizes each face from GFPGAN's StyleGAN2 prior, running on the
|
||||
DIFFUSION-CLEANED image -- not on the original.
|
||||
|
||||
**Why "cleaned, not original":** an earlier version of this module ran GFPGAN on the
|
||||
ORIGINAL (watermarked) image and was oracle-confirmed (2026-06-04) to re-introduce
|
||||
SynthID into the face regions, because GFPGAN at fidelity weight 0.5 blends ~half
|
||||
the input pixels with the prior, and SynthID is robust to that partial blend. The
|
||||
fix is to feed GFPGAN the already-clean image -- whatever pixels it preserves are
|
||||
already SynthID-free, so the composited face stays clean. Identity is recovered from
|
||||
the StyleGAN2 prior conditioned on the already-drifted cleaned face (not on the
|
||||
original face), so identity fidelity is somewhat lower than the would-have-been
|
||||
identity-as-embedding stack (PhotoMaker-V1), but the upstream PhotoMaker package has
|
||||
significant compatibility issues with the diffusers version we ship, so this is the
|
||||
shipping path.
|
||||
|
||||
Both GFPGAN (Apache-2.0) and its RetinaFace detector (MIT) are commercial-safe.
|
||||
The GFPGANv1.4 weights and the RetinaFace detector download on first use and are
|
||||
never bundled. Requires the optional ``restore`` extra (gfpgan/facexlib/basicsr).
|
||||
"""
|
||||
|
||||
# cv2/torch/gfpgan boundary: gfpgan/basicsr/facexlib ship no usable type stubs and
|
||||
# this module wraps cv2 (feather composite) and torch; 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
|
||||
|
||||
import logging
|
||||
import sys
|
||||
import threading
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from numpy.typing import NDArray
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# GFPGANv1.4 weights (Apache-2.0). Downloaded on first use, never bundled.
|
||||
_GFPGAN_MODEL_URL = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
|
||||
_GFPGAN_ARCH = "clean"
|
||||
_GFPGAN_CHANNEL_MULTIPLIER = 2
|
||||
|
||||
_restorer: Any | None = None
|
||||
_restorer_lock = threading.Lock()
|
||||
|
||||
|
||||
def is_available() -> bool:
|
||||
"""True when the optional GFPGAN face-restoration deps are importable."""
|
||||
import importlib.util
|
||||
|
||||
return importlib.util.find_spec("gfpgan") is not None and importlib.util.find_spec("facexlib") is not None
|
||||
|
||||
|
||||
def _apply_basicsr_shim() -> None:
|
||||
"""Install the ``torchvision.transforms.functional_tensor`` compatibility shim.
|
||||
|
||||
basicsr (a GFPGAN dependency) imports ``rgb_to_grayscale`` from the
|
||||
``torchvision.transforms.functional_tensor`` module, which newer torchvision
|
||||
removed. Recreate that module pointing at the public functional API. Idempotent:
|
||||
only installed when the real module is missing.
|
||||
"""
|
||||
import importlib.util
|
||||
|
||||
if importlib.util.find_spec("torchvision.transforms.functional_tensor") is not None:
|
||||
return
|
||||
if "torchvision.transforms.functional_tensor" in sys.modules:
|
||||
return
|
||||
|
||||
import types
|
||||
|
||||
import torchvision.transforms.functional as tv_functional
|
||||
|
||||
shim = types.ModuleType("torchvision.transforms.functional_tensor")
|
||||
shim.rgb_to_grayscale = tv_functional.rgb_to_grayscale
|
||||
sys.modules["torchvision.transforms.functional_tensor"] = shim
|
||||
|
||||
|
||||
def _select_device() -> str:
|
||||
"""Pick the GFPGAN device: CUDA when present, else CPU.
|
||||
|
||||
The pip GFPGANer has an MPS device-mismatch bug, and this is a cheap post-pass
|
||||
on a few face crops, so MPS is deliberately avoided -- CPU is the safe default
|
||||
on Apple silicon.
|
||||
"""
|
||||
try:
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
return "cuda"
|
||||
except Exception as e:
|
||||
logger.debug("face_restore: CUDA probe failed (%s); using CPU", e)
|
||||
return "cpu"
|
||||
|
||||
|
||||
def _get_restorer() -> Any:
|
||||
"""Return the lazily-built GFPGANer singleton (downloads weights on first use)."""
|
||||
global _restorer
|
||||
if _restorer is not None:
|
||||
return _restorer
|
||||
with _restorer_lock:
|
||||
if _restorer is None:
|
||||
_apply_basicsr_shim()
|
||||
from gfpgan import GFPGANer
|
||||
|
||||
_restorer = GFPGANer(
|
||||
model_path=_GFPGAN_MODEL_URL,
|
||||
upscale=1,
|
||||
arch=_GFPGAN_ARCH,
|
||||
channel_multiplier=_GFPGAN_CHANNEL_MULTIPLIER,
|
||||
device=_select_device(),
|
||||
)
|
||||
return _restorer
|
||||
|
||||
|
||||
def _composite_faces(
|
||||
base_bgr: NDArray[Any],
|
||||
restored_bgr: NDArray[Any],
|
||||
boxes: list[tuple[float, float, float, float]],
|
||||
pad: int = 14,
|
||||
feather_div: int = 6,
|
||||
) -> NDArray[Any]:
|
||||
"""Feather-composite restored face regions from ``restored_bgr`` into ``base_bgr``.
|
||||
|
||||
Pure cv2/numpy helper (no gfpgan), so it is unit-testable without the model.
|
||||
For each ``(x1, y1, x2, y2)`` box: pad and clip to the image, build a Gaussian-
|
||||
feathered rectangular alpha, and blend ``restored * a + base * (1 - a)``. Boxes
|
||||
that fall fully outside the image (or an empty list) leave ``base_bgr`` unchanged.
|
||||
"""
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
out = base_bgr.astype(np.float32)
|
||||
h, w = base_bgr.shape[:2]
|
||||
|
||||
for box in boxes:
|
||||
x1 = int(box[0]) - pad
|
||||
y1 = int(box[1]) - pad
|
||||
x2 = int(box[2]) + pad
|
||||
y2 = int(box[3]) + pad
|
||||
x1 = max(0, min(x1, w))
|
||||
y1 = max(0, min(y1, h))
|
||||
x2 = max(0, min(x2, w))
|
||||
y2 = max(0, min(y2, h))
|
||||
bw = x2 - x1
|
||||
bh = y2 - y1
|
||||
if bw <= 0 or bh <= 0:
|
||||
continue
|
||||
|
||||
alpha = np.zeros((h, w), dtype=np.float32)
|
||||
alpha[y1:y2, x1:x2] = 1.0
|
||||
k = max(3, (min(bw, bh) // feather_div) | 1) # odd kernel >= 3
|
||||
alpha = cv2.GaussianBlur(alpha, (k, k), 0)
|
||||
alpha = alpha[:, :, None]
|
||||
out = restored_bgr.astype(np.float32) * alpha + out * (1.0 - alpha)
|
||||
|
||||
return np.clip(out, 0, 255).astype(np.uint8)
|
||||
|
||||
|
||||
def restore_faces(
|
||||
original_bgr: NDArray[Any], # legacy positional kept for API stability; unused
|
||||
cleaned_bgr: NDArray[Any],
|
||||
weight: float = 0.5,
|
||||
pad: int = 14,
|
||||
feather_div: int = 6,
|
||||
) -> NDArray[Any]:
|
||||
"""Restore face identity in ``cleaned_bgr`` by running GFPGAN on the CLEANED image.
|
||||
|
||||
GFPGAN is a fidelity-restoration net: it sharpens and re-synthesizes face details
|
||||
from its StyleGAN2 prior conditioned on the INPUT face. **Running it on the
|
||||
diffusion-cleaned image (not the original)** is what makes this pass SynthID-safe:
|
||||
the input pixels GFPGAN derives from are already SynthID-free, so the partial
|
||||
pixel-blend at the default weight 0.5 cannot re-introduce the watermark.
|
||||
|
||||
The earlier version of this module ran GFPGAN on the ORIGINAL (watermarked) image
|
||||
and was oracle-confirmed (2026-06-04) to re-introduce SynthID into the face
|
||||
regions. The fix is the single-line source swap below.
|
||||
|
||||
The ``original_bgr`` argument is kept for positional API stability with the
|
||||
earlier signature but is no longer used; pass it for legacy callers, ignore it
|
||||
in new code.
|
||||
|
||||
Args:
|
||||
original_bgr: UNUSED (legacy; kept for positional API stability).
|
||||
cleaned_bgr: The diffusion-cleaned image as cv2 BGR (faces drifted from the
|
||||
removal pass). GFPGAN runs on THIS, polishing each face without changing
|
||||
the watermark state of the source pixels.
|
||||
weight: GFPGAN fidelity weight (0-1); lower = more StyleGAN2 regeneration of
|
||||
the face from the prior.
|
||||
pad: Pixels to grow each face box before compositing.
|
||||
feather_div: Larger = sharper composite edge (box-min // feather_div kernel).
|
||||
"""
|
||||
restorer = _get_restorer()
|
||||
_, _, restored_img = restorer.enhance(
|
||||
cleaned_bgr,
|
||||
has_aligned=False,
|
||||
only_center_face=False,
|
||||
paste_back=True,
|
||||
weight=weight,
|
||||
)
|
||||
|
||||
det_faces = getattr(restorer.face_helper, "det_faces", None) or []
|
||||
boxes = [(float(b[0]), float(b[1]), float(b[2]), float(b[3])) for b in det_faces]
|
||||
if not boxes:
|
||||
logger.debug("face_restore: no faces detected; returning cleaned image unchanged")
|
||||
return cleaned_bgr
|
||||
|
||||
return _composite_faces(cleaned_bgr, restored_img, boxes, pad=pad, feather_div=feather_div)
|
||||
@@ -180,11 +180,11 @@ class InvisibleEngine:
|
||||
guidance_scale: Classifier-free guidance scale.
|
||||
seed: Random seed for reproducibility.
|
||||
humanize: Intensity of Analog Humanizer film grain (0 = off).
|
||||
restore_faces: EXPERIMENTAL, opt-in (default False). Run the GFPGAN
|
||||
face-polish post-pass when faces are present (needs the ``restore``
|
||||
extra). Runs on the diffusion-CLEANED image (not the original), so
|
||||
SynthID is not re-introduced. Auto-skips with a debug log when the
|
||||
extra is absent or no face is detected.
|
||||
restore_faces: EXPERIMENTAL, opt-in (default False). **NON-COMMERCIAL.**
|
||||
Run the PhotoMaker-V2 face-identity post-pass when faces are present
|
||||
(needs the ``photomaker`` extra, which pulls non-commercial InsightFace
|
||||
model packs). Auto-skips with a debug log when the extra is absent or no
|
||||
face is detected. See ``photomaker_restore.py`` for the legal notice.
|
||||
unsharp: Final unsharp-mask sharpening strength (0 = off, default).
|
||||
Applied last (after face restoration) to counter the soft,
|
||||
over-smoothed look of the diffusion + restoration; ~0.5-0.8 is a
|
||||
@@ -316,7 +316,7 @@ class InvisibleEngine:
|
||||
# GFPGAN derives from are already SynthID-free). Auto-skips when faces are
|
||||
# absent or the optional `restore` extra is not installed.
|
||||
if restore_faces:
|
||||
self._restore_faces(out_path)
|
||||
self._restore_faces_photomaker(out_path, image, seed)
|
||||
|
||||
# Final sharpening, LAST so it crisps the face-restored result too (a
|
||||
# pre-restore sharpen would be smoothed back over by the face pass).
|
||||
@@ -355,24 +355,29 @@ class InvisibleEngine:
|
||||
if _tmp_path.exists():
|
||||
_tmp_path.unlink()
|
||||
|
||||
def _restore_faces(self, out_path: Path) -> None:
|
||||
"""Run the GFPGAN face-polish post-pass on the cleaned ``out_path``.
|
||||
def _restore_faces_photomaker(
|
||||
self,
|
||||
out_path: Path,
|
||||
original_image: Any,
|
||||
seed: int | None,
|
||||
) -> None:
|
||||
"""Run the PhotoMaker-V2 face-identity post-pass on the cleaned ``out_path``.
|
||||
|
||||
SynthID-safe: GFPGAN is run on the diffusion-CLEANED image (not the original),
|
||||
so the partial pixel-blend it does at fidelity weight 0.5 cannot re-introduce
|
||||
the watermark -- the input pixels GFPGAN derives from are already SynthID-free.
|
||||
Best-effort: any failure logs a warning and leaves the un-restored cleaned
|
||||
output in place; a missing ``restore`` extra is logged at debug and skipped
|
||||
(the flag must never error when the extra is absent or no face is present).
|
||||
**NON-COMMERCIAL** (see ``photomaker_restore.py``). PhotoMaker carries identity
|
||||
in a CLIP+ArcFace embedding and regenerates fresh face pixels conditioned on
|
||||
it, so the watermark is not transported. Best-effort: any failure (missing
|
||||
extra, model load, runtime error) logs a warning and leaves the un-restored
|
||||
cleaned output in place.
|
||||
"""
|
||||
from remove_ai_watermarks import face_restore
|
||||
from remove_ai_watermarks import photomaker_restore
|
||||
|
||||
if not face_restore.is_available():
|
||||
logger.debug("restore_faces requested but the 'restore' extra is not installed; skipping")
|
||||
if not photomaker_restore.is_available():
|
||||
logger.debug("restore_faces requested but the 'photomaker' extra is not installed; skipping")
|
||||
return
|
||||
|
||||
try:
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from remove_ai_watermarks import image_io
|
||||
|
||||
@@ -381,13 +386,15 @@ class InvisibleEngine:
|
||||
logger.warning("restore_faces: could not read cleaned output %s; skipping", out_path)
|
||||
return
|
||||
|
||||
if self._progress_callback:
|
||||
self._progress_callback("Polishing face identity (GFPGAN on cleaned image)...")
|
||||
# original_bgr is unused (GFPGAN runs on cleaned_bgr); pass an empty array
|
||||
# for positional API stability with the legacy signature.
|
||||
import numpy as np
|
||||
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)
|
||||
|
||||
restored = face_restore.restore_faces(np.empty((0, 0, 3), dtype=np.uint8), cleaned_bgr)
|
||||
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 post-pass failed (%s); keeping un-restored output", e)
|
||||
|
||||
@@ -0,0 +1,317 @@
|
||||
"""SynthID-robust face identity restoration via PhotoMaker-V2.
|
||||
|
||||
**NON-COMMERCIAL.** This module uses PhotoMaker-V2, whose ID encoder
|
||||
(``PhotoMakerIDEncoder_CLIPInsightfaceExtendtoken``) requires an ArcFace embedding
|
||||
from InsightFace's pretrained ``antelopev2`` / ``buffalo_l`` model packs. Those packs
|
||||
are released by InsightFace under a **non-commercial / research-only license**:
|
||||
|
||||
"The pretrained models we provided with this library are available for
|
||||
non-commercial research purposes only."
|
||||
-- insightface PyPI README
|
||||
|
||||
The PyPI ``insightface`` package itself is MIT-licensed code, but the model weights
|
||||
it downloads on first ``FaceAnalysis()`` are not commercial. **A paid service
|
||||
(raiw.cc, any monetized SaaS, any enterprise deployment) MUST NOT use this path.**
|
||||
The default ``--restore-faces`` method is ``gfpgan`` (commercial-safe, ships with
|
||||
the ``restore`` extra); ``--restore-faces-method photomaker`` is an explicit opt-in
|
||||
for non-commercial use only. See ``docs/synthid-robust-identity-research.md``.
|
||||
|
||||
The diffusion removal pass scrubs the pixel watermark from the WHOLE image, including
|
||||
faces, but lets faces drift in identity. PhotoMaker-V2 carries identity in two
|
||||
semantic streams (an OpenCLIP-ViT-H/14 image embedding AND an ArcFace identity
|
||||
embedding) and uses them to CONDITION a fresh txt2img generation -- the pixels are
|
||||
new, so the watermark cannot be transported.
|
||||
|
||||
That embeddings do not carry an invisible pixel watermark like SynthID is the
|
||||
load-bearing assumption of the whole approach; the OpenCLIP smoke test (cosine
|
||||
0.9977 invariance to SynthID-magnitude pixel noise) supports it for the CLIP
|
||||
stream, and ArcFace is even more invariant to small perceptual changes by design.
|
||||
|
||||
Architecture: PhotoMaker-V2 is a fine-tuned OpenCLIP-ViT-H/14 + InsightFace dual ID
|
||||
encoder plus LoRA on the SDXL UNet attention layers. It ships as a single
|
||||
``photomaker-v2.bin`` checkpoint loaded into a ``PhotoMakerStableDiffusionXLPipeline``
|
||||
(txt2img). 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 the dual 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.
|
||||
|
||||
Requires the optional ``photomaker`` extra: ``pip install
|
||||
'remove-ai-watermarks[photomaker]'`` -- this pulls the upstream PhotoMaker package
|
||||
(Apache-2.0), ``insightface`` (MIT code), ``einops``, ``peft``, ``onnxruntime``,
|
||||
and ``huggingface-hub``. Weights and InsightFace model packs 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 adapter; ID encoder pulls non-commercial
|
||||
# InsightFace model packs at runtime -- see the NON-COMMERCIAL notice in the module
|
||||
# docstring). Downloaded on first use; never bundled.
|
||||
_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)
|
||||
# Move SDXL submodules to the device BEFORE loading the PhotoMaker adapter:
|
||||
# ``load_photomaker_adapter`` reads ``self.device`` / ``self.unet.dtype`` to
|
||||
# place the new ID encoder. If we ``.to(device)`` after, the SDXL submodules
|
||||
# move but the id_encoder stays where it was (custom attribute, not in the
|
||||
# auto-managed module tree), and inference errors with
|
||||
# "Input type (torch.cuda.HalfTensor) and weight type (torch.HalfTensor)
|
||||
# should be the same" (caught empirically 2026-06-04).
|
||||
pipe.to(device)
|
||||
# Default ``pm_version`` is "v2"; we load the V2 weights (photomaker-v2.bin)
|
||||
# into the V2 encoder (PhotoMakerIDEncoder_CLIPInsightfaceExtendtoken). The V2
|
||||
# encoder takes BOTH the CLIP image features AND an InsightFace ArcFace
|
||||
# embedding -- the latter is what makes this path non-commercial.
|
||||
pipe.load_photomaker_adapter(
|
||||
str(Path(adapter_path).parent),
|
||||
subfolder="",
|
||||
weight_name=_PHOTOMAKER_FILE,
|
||||
trigger_word="img",
|
||||
)
|
||||
pipe.fuse_lora()
|
||||
# Belt: also explicitly cast the loaded id_encoder, because some
|
||||
# diffusers/torch combinations leave the encoder buffers untouched even
|
||||
# though ``pipe.to(device)`` ran first.
|
||||
if hasattr(pipe, "id_encoder") and pipe.id_encoder is not None:
|
||||
pipe.id_encoder = pipe.id_encoder.to(device=device, dtype=dtype)
|
||||
_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)
|
||||
|
||||
# Don't pass negative_prompt: the PhotoMaker pipeline manages its own CFG by
|
||||
# concatenating [negative_prompt_embeds, prompt_embeds]; if we pass a custom
|
||||
# negative the upstream code splits text_only vs id-injected branches and
|
||||
# the resulting embed batch dims can mismatch (we saw
|
||||
# "Sizes of tensors must match except in dimension 1. Expected size 2 but got
|
||||
# size 1" on a real run). The default empty negative is what the upstream
|
||||
# gradio demo uses.
|
||||
out = pipeline(
|
||||
prompt=_PHOTOMAKER_PROMPT,
|
||||
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)
|
||||
@@ -1,85 +0,0 @@
|
||||
"""Tests for the GFPGAN face-restoration post-pass.
|
||||
|
||||
The pure feather-composite helper is unit-tested without the model; the
|
||||
model-running paths are gated behind ``is_available()`` (a multi-hundred-MB
|
||||
download), matching the discipline used for the other ML-adjacent modules.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from remove_ai_watermarks import face_restore
|
||||
|
||||
|
||||
class TestIsAvailable:
|
||||
def test_returns_bool(self):
|
||||
assert isinstance(face_restore.is_available(), bool)
|
||||
|
||||
def test_reflects_dependencies(self):
|
||||
import importlib.util
|
||||
|
||||
expected = all(importlib.util.find_spec(m) is not None for m in ("gfpgan", "facexlib"))
|
||||
assert face_restore.is_available() is expected
|
||||
|
||||
|
||||
class TestCompositeFaces:
|
||||
"""Unit tests for the pure ``_composite_faces`` helper (cv2/numpy only)."""
|
||||
|
||||
def _base_and_restored(self, h: int = 100, w: int = 120):
|
||||
base = np.zeros((h, w, 3), dtype=np.uint8) # black
|
||||
restored = np.full((h, w, 3), 255, dtype=np.uint8) # white
|
||||
return base, restored
|
||||
|
||||
def test_output_shape_and_dtype(self):
|
||||
base, restored = self._base_and_restored()
|
||||
out = face_restore._composite_faces(base, restored, [(40.0, 30.0, 80.0, 70.0)])
|
||||
assert out.shape == base.shape
|
||||
assert out.dtype == np.uint8
|
||||
|
||||
def test_box_region_pulls_toward_restored(self):
|
||||
base, restored = self._base_and_restored()
|
||||
out = face_restore._composite_faces(base, restored, [(40.0, 30.0, 80.0, 70.0)])
|
||||
# Center of the box should be near the restored (white) value.
|
||||
cy, cx = 50, 60
|
||||
assert out[cy, cx].mean() > 200
|
||||
|
||||
def test_far_from_box_stays_base(self):
|
||||
base, restored = self._base_and_restored()
|
||||
out = face_restore._composite_faces(base, restored, [(40.0, 30.0, 80.0, 70.0)], pad=2)
|
||||
# Top-left corner is far from the box and feather, so it stays black.
|
||||
assert out[0, 0].mean() < 5
|
||||
|
||||
def test_empty_boxes_returns_base_unchanged(self):
|
||||
base, restored = self._base_and_restored()
|
||||
out = face_restore._composite_faces(base, restored, [])
|
||||
assert np.array_equal(out, base)
|
||||
|
||||
def test_box_fully_outside_is_skipped(self):
|
||||
base, restored = self._base_and_restored(h=100, w=120)
|
||||
# Box entirely beyond the right/bottom edge -> clipped to empty -> no-op.
|
||||
out = face_restore._composite_faces(base, restored, [(200.0, 200.0, 260.0, 260.0)], pad=0)
|
||||
assert np.array_equal(out, base)
|
||||
|
||||
def test_near_edge_box_clips_without_error(self):
|
||||
base, restored = self._base_and_restored(h=100, w=120)
|
||||
# Box reaching past the bottom-right corner must clip, not raise.
|
||||
out = face_restore._composite_faces(base, restored, [(100.0, 80.0, 130.0, 110.0)], pad=10)
|
||||
assert out.shape == base.shape
|
||||
# The clipped in-bounds region still pulls toward white.
|
||||
assert out[95, 115].mean() > 100
|
||||
|
||||
|
||||
@pytest.mark.skipif(not face_restore.is_available(), reason="requires the 'restore' extra (gfpgan/facexlib)")
|
||||
class TestRestoreFacesModel:
|
||||
"""Model-running smoke test, gated behind the optional extra."""
|
||||
|
||||
def test_no_faces_returns_cleaned_unchanged(self):
|
||||
# A flat gray image has no faces; restore_faces must return the cleaned
|
||||
# input unchanged (the no-op path).
|
||||
cleaned = np.full((128, 128, 3), 127, dtype=np.uint8)
|
||||
original = np.full((128, 128, 3), 127, dtype=np.uint8)
|
||||
out = face_restore.restore_faces(original, cleaned)
|
||||
assert out.shape == cleaned.shape
|
||||
assert np.array_equal(out, cleaned)
|
||||
@@ -0,0 +1,138 @@
|
||||
"""Tests for the PhotoMaker-V2 face identity restoration helper.
|
||||
|
||||
These tests cover the pure-Python parts (face crop math, composite, the no-faces
|
||||
no-op, the is_available guard) WITHOUT loading PhotoMaker or SDXL -- the model-loading
|
||||
path is gated behind ``is_available()`` and exercised manually via the Modal cert
|
||||
sweep, mirroring the convention used for ``face_restore`` and ``upscaler``.
|
||||
|
||||
The end-to-end PhotoMaker run is monkey-patched: we replace ``_get_pipeline`` with a
|
||||
fake pipeline whose ``__call__`` returns a known constant-color face, so we can verify
|
||||
that the right boxes get the right pixels composited back.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from remove_ai_watermarks import photomaker_restore
|
||||
|
||||
|
||||
class TestIsAvailable:
|
||||
def test_returns_bool(self):
|
||||
assert isinstance(photomaker_restore.is_available(), bool)
|
||||
|
||||
|
||||
class TestV2WeightPins:
|
||||
"""Pin the V2 repo + weights so a maintainer change is visible in a code review."""
|
||||
|
||||
def test_repo_is_v2(self):
|
||||
assert photomaker_restore._PHOTOMAKER_REPO == "TencentARC/PhotoMaker-V2"
|
||||
|
||||
def test_weight_filename_is_v2(self):
|
||||
assert photomaker_restore._PHOTOMAKER_FILE == "photomaker-v2.bin"
|
||||
|
||||
|
||||
class TestFaceCropSquare:
|
||||
def test_centers_on_face_box(self):
|
||||
img = np.full((400, 400, 3), 128, dtype=np.uint8)
|
||||
crop, box = photomaker_restore._face_crop_square(img, (100, 150, 80, 80))
|
||||
x1, y1, x2, y2 = box
|
||||
# The crop covers the requested box (with padding)
|
||||
assert x1 <= 100
|
||||
assert y1 <= 150
|
||||
assert x2 >= 180
|
||||
assert y2 >= 230
|
||||
assert crop.shape[0] == y2 - y1
|
||||
assert crop.shape[1] == x2 - x1
|
||||
|
||||
def test_clips_at_image_edges(self):
|
||||
img = np.full((200, 200, 3), 128, dtype=np.uint8)
|
||||
crop, (x1, y1, x2, y2) = photomaker_restore._face_crop_square(img, (180, 180, 30, 30))
|
||||
# Box must be clipped within the image
|
||||
assert x1 >= 0
|
||||
assert y1 >= 0
|
||||
assert x2 <= 200
|
||||
assert y2 <= 200
|
||||
assert crop.shape[0] == y2 - y1
|
||||
assert crop.shape[1] == x2 - x1
|
||||
|
||||
def test_pad_widens_the_crop(self):
|
||||
img = np.full((400, 400, 3), 128, dtype=np.uint8)
|
||||
_, no_pad = photomaker_restore._face_crop_square(img, (150, 150, 50, 50), pad=0.0)
|
||||
_, with_pad = photomaker_restore._face_crop_square(img, (150, 150, 50, 50), pad=0.5)
|
||||
assert (with_pad[2] - with_pad[0]) > (no_pad[2] - no_pad[0])
|
||||
|
||||
|
||||
class TestCompositeFaces:
|
||||
def test_empty_list_returns_base_unchanged(self):
|
||||
base = np.full((100, 100, 3), 64, dtype=np.uint8)
|
||||
out = photomaker_restore._composite_faces(base, [])
|
||||
assert np.array_equal(out, base)
|
||||
|
||||
def test_box_outside_image_is_skipped(self):
|
||||
base = np.full((100, 100, 3), 64, dtype=np.uint8)
|
||||
crop = np.full((40, 40, 3), 200, dtype=np.uint8)
|
||||
out = photomaker_restore._composite_faces(base, [(crop, (200, 200, 240, 240))])
|
||||
assert np.array_equal(out, base)
|
||||
|
||||
def test_composited_box_pulls_pixel_value_toward_crop(self):
|
||||
base = np.full((200, 200, 3), 40, dtype=np.uint8)
|
||||
crop = np.full((50, 50, 3), 220, dtype=np.uint8)
|
||||
# Place the crop fully inside the image at (60, 60)..(110, 110)
|
||||
out = photomaker_restore._composite_faces(base, [(crop, (60, 60, 110, 110))])
|
||||
# The box center should be heavily biased toward the crop color (>120) ...
|
||||
assert out[85, 85, 0] > 120
|
||||
# ... and corners (well outside the feathered region) stay close to base
|
||||
assert int(out[0, 0, 0]) - int(base[0, 0, 0]) <= 1
|
||||
|
||||
|
||||
class TestRestoreFacesPhotomakerControlFlow:
|
||||
"""End-to-end control flow with a fake pipeline -- no diffusion model loaded."""
|
||||
|
||||
@staticmethod
|
||||
def _fake_pipeline_class(fill_value: int = 200):
|
||||
"""Class-based fake (no ``__call__`` on a SimpleNamespace, which Python won't dispatch)."""
|
||||
from PIL import Image
|
||||
|
||||
size = photomaker_restore._PHOTOMAKER_FACE_SIZE
|
||||
fake_face = Image.fromarray(np.full((size, size, 3), fill_value, dtype=np.uint8))
|
||||
|
||||
class _FakePipe:
|
||||
device = "cpu"
|
||||
|
||||
def __call__(self, **_kwargs):
|
||||
return SimpleNamespace(images=[fake_face])
|
||||
|
||||
return _FakePipe()
|
||||
|
||||
def test_no_faces_returns_cleaned_unchanged(self, monkeypatch):
|
||||
# Force is_available so we never hit the missing-extra branch
|
||||
monkeypatch.setattr(photomaker_restore, "is_available", lambda: True)
|
||||
monkeypatch.setattr(photomaker_restore, "_get_pipeline", lambda: self._fake_pipeline_class())
|
||||
|
||||
orig = np.full((200, 200, 3), 30, dtype=np.uint8)
|
||||
cleaned = np.full((200, 200, 3), 90, dtype=np.uint8)
|
||||
out = photomaker_restore.restore_faces_photomaker(orig, cleaned, detect_faces_fn=lambda _b: [])
|
||||
assert np.array_equal(out, cleaned)
|
||||
|
||||
def test_one_face_gets_composited_into_cleaned(self, monkeypatch):
|
||||
monkeypatch.setattr(photomaker_restore, "is_available", lambda: True)
|
||||
monkeypatch.setattr(photomaker_restore, "_get_pipeline", lambda: self._fake_pipeline_class(fill_value=210))
|
||||
|
||||
orig = np.full((400, 400, 3), 30, dtype=np.uint8)
|
||||
cleaned = np.full((400, 400, 3), 90, dtype=np.uint8)
|
||||
# Mark the original face region with a distinctive color so we can confirm the
|
||||
# crop reached the pipeline (not strictly tested here, but useful sanity).
|
||||
cv2.rectangle(orig, (150, 150), (250, 250), (200, 100, 50), -1)
|
||||
|
||||
out = photomaker_restore.restore_faces_photomaker(
|
||||
orig, cleaned, detect_faces_fn=lambda _b: [(150, 150, 100, 100)]
|
||||
)
|
||||
# The cleaned image should have shifted toward the fake-face fill (210) inside
|
||||
# the face region.
|
||||
assert out[200, 200, 0] > 150
|
||||
# And the corner pixels (well outside the feather) should still be near the base.
|
||||
assert int(out[0, 0, 0]) - int(cleaned[0, 0, 0]) <= 1
|
||||
Reference in New Issue
Block a user