mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-08 23:16:02 +02:00
remove: drop all face-restore code (regeneration, not preservation)
Empirical conclusion from the 2026-06-04 - 2026-06-08 Modal cert sweeps: every face-restore approach we built (GFPGAN-on-cleaned, PhotoMaker-V2, InstantID txt2img, InstantID img2img-on-cleaned at three parameter settings) regenerates the face via SDXL diffusion rather than preserves it. Output face pixels are diffusion-fresh, so the regenerated face inherits SDXL "clean skin" aesthetic and loses original identity precision -- it looks MORE AI-generated than the cleaned image, not less. The cleaned image from the main controlnet 0.20 removal pass is the least-AI face state we can reach without re-introducing SynthID. Nothing in the restore family achieves the actual goal (preserve the original person's face). Keeping them around as opt-in invites users to ship something that defeats the point. Removing entirely. Library changes: - Deleted src/remove_ai_watermarks/instantid_restore.py - Deleted src/remove_ai_watermarks/photomaker_restore.py - Deleted tests/test_instantid_restore.py - Deleted tests/test_photomaker_restore.py - Removed `instantid` and `photomaker` extras from pyproject.toml - Removed `[tool.hatch.metadata] allow-direct-references = true` (was only needed for the photomaker git+ URL) - InvisibleEngine.remove_watermark: dropped `restore_faces` + `restore_faces_method` params, removed both `_restore_faces_instantid` and `_restore_faces_photomaker` private methods, removed dispatch - CLI: dropped `_restore_faces_options` decorator, all four cmd_* signatures lose `restore_faces` + `restore_faces_method`, kwarg passes to remove_watermark dropped - _apply_auto: dropped `restore_faces` from tuple shape (was unused after the engine no longer takes it) - auto_config.AutoConfig: dropped `restore_faces` field; `plan()` no longer sets it; `reason` no longer mentions it - Tests updated accordingly (test_auto_config.TestReason no longer asserts "face-restore on" in the reason string) Docs updated: - CLAUDE.md: removed the photomaker extras bullet, the Face restore trade-off bullet, the instantid_restore.py + photomaker_restore.py module bullets; replaced restore mentions in watermark_remover and controlnet bullets and prod recipe with the empirical conclusion - README.md: removed both `--restore-faces` callouts and the install snippet; the feature bullet and auto-mode comment updated - docs/synthid-robust-identity-research.md: added Status-retired notice at the top pointing at the 2026-06-08 followup raiw-app: - modal_cert.py: dropped `--restore-faces` flag entirely; sweep() no longer takes restore_faces; pinned _LIB_SPEC to `[gpu]` extras (no `photomaker` / `instantid` extras), points at main ruff + strict pyright clean; 569 tests pass; 18 restore-specific tests gone. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Victor Kuznetsov
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
567f3ae729
commit
20d7eda96a
@@ -85,7 +85,6 @@ class AutoConfig:
|
||||
"""Resolved quality modes from content analysis (the ``--auto`` plan)."""
|
||||
|
||||
pipeline: str # "default" | "controlnet"
|
||||
restore_faces: bool
|
||||
adaptive_polish: bool # restore the input's detail level (sharpen + masked grain), sparing text
|
||||
unsharp: float # fixed-polish knobs, 0 in auto (the adaptive polish replaces them)
|
||||
humanize: float
|
||||
@@ -104,14 +103,13 @@ class AutoConfig:
|
||||
if self.has_text:
|
||||
bits.append("text")
|
||||
bits.append(f"edges={self.edge_density:.3f}")
|
||||
rf = ", face-restore on" if self.restore_faces else ""
|
||||
if self.adaptive_polish:
|
||||
polish = ", adaptive polish"
|
||||
elif self.unsharp or self.humanize:
|
||||
polish = f", unsharp {self.unsharp}/grain {self.humanize}"
|
||||
else:
|
||||
polish = ""
|
||||
return f"{'+'.join(bits)} -> {self.pipeline} pipeline{rf}{polish}"
|
||||
return f"{'+'.join(bits)} -> {self.pipeline} pipeline{polish}"
|
||||
|
||||
|
||||
def _to_bgr(image: NDArray[Any]) -> NDArray[Any]:
|
||||
@@ -251,12 +249,10 @@ def plan(image_path: Path) -> AutoConfig | None:
|
||||
|
||||
structureless = (not has_face) and (not has_text) and edges < _STRUCTURELESS_EDGE_MAX
|
||||
pipeline = "default" if structureless else "controlnet"
|
||||
restore_faces = has_face
|
||||
smoothing = pipeline == "controlnet" or restore_faces
|
||||
smoothing = pipeline == "controlnet"
|
||||
|
||||
cfg = AutoConfig(
|
||||
pipeline=pipeline,
|
||||
restore_faces=restore_faces,
|
||||
adaptive_polish=smoothing, # adaptive (detail-targeted) polish when a smoothing pass ran
|
||||
unsharp=0.0,
|
||||
humanize=0.0,
|
||||
|
||||
@@ -174,7 +174,7 @@ _auto_option = click.option(
|
||||
is_flag=True,
|
||||
default=False,
|
||||
help="Auto-pick the pipeline, face restore, and adaptive polish from image content. "
|
||||
"Every choice is overridable -- an explicit --pipeline / --restore-faces / --adaptive-polish "
|
||||
"Every choice is overridable -- an explicit --pipeline / --adaptive-polish "
|
||||
"always wins. EXPERIMENTAL.",
|
||||
)
|
||||
|
||||
@@ -192,9 +192,8 @@ def _apply_auto(
|
||||
ctx: click.Context,
|
||||
source: Path,
|
||||
pipeline: str,
|
||||
restore_faces: bool,
|
||||
adaptive_polish: bool,
|
||||
) -> tuple[str, bool, bool]:
|
||||
) -> tuple[str, bool]:
|
||||
"""Resolve ``--auto``: plan the three content-adaptive modes (pipeline, face
|
||||
restore, adaptive polish) from the image, overriding only the ones the user left
|
||||
at their default (an explicit flag always wins). The fixed ``--unsharp``/
|
||||
@@ -205,19 +204,17 @@ def _apply_auto(
|
||||
cfg = auto_config.plan(source)
|
||||
if cfg is None:
|
||||
console.print(" Auto: could not read image; using defaults")
|
||||
return pipeline, restore_faces, adaptive_polish
|
||||
return pipeline, adaptive_polish
|
||||
|
||||
def _is_default(name: str) -> bool:
|
||||
return ctx.get_parameter_source(name) == click.core.ParameterSource.DEFAULT
|
||||
|
||||
if _is_default("pipeline"):
|
||||
pipeline = cfg.pipeline
|
||||
if _is_default("restore_faces"):
|
||||
restore_faces = cfg.restore_faces
|
||||
if _is_default("adaptive_polish"):
|
||||
adaptive_polish = cfg.adaptive_polish
|
||||
console.print(f" Auto: {cfg.reason}")
|
||||
return pipeline, restore_faces, adaptive_polish
|
||||
return pipeline, adaptive_polish
|
||||
|
||||
|
||||
def _warn_if_esrgan_unavailable(upscaler: str) -> None:
|
||||
@@ -235,43 +232,6 @@ def _warn_if_esrgan_unavailable(upscaler: str) -> None:
|
||||
console.print(" Note: --upscaler esrgan needs the 'esrgan' extra; falling back to Lanczos.")
|
||||
|
||||
|
||||
def _restore_faces_options(f: Any) -> Any:
|
||||
"""Attach the face-restoration flags to an invisible-pipeline command.
|
||||
|
||||
Both methods REGENERATE the face from an ArcFace embedding via SDXL diffusion
|
||||
-- they do NOT recover original pixels. Every output face pixel is
|
||||
diffusion-fresh, so the regenerated face inherently looks MORE AI-generated
|
||||
than the cleaned image (gloss, symmetric pores, SDXL "clean skin"
|
||||
aesthetic). For production face preservation, leave the flag OFF and use
|
||||
the cleaned image as-is. The two methods are kept for research / personal
|
||||
use where users explicitly want identity regeneration. **BOTH are
|
||||
NON-COMMERCIAL**: they pull InsightFace antelopev2 / buffalo_l model packs
|
||||
which are research-only. A paid service (raiw.cc, any monetized SaaS) MUST
|
||||
NOT use this flag.
|
||||
"""
|
||||
method = click.option(
|
||||
"--restore-faces-method",
|
||||
type=click.Choice(["instantid", "photomaker"]),
|
||||
default="instantid",
|
||||
help="Face-regeneration mechanism (no method recovers original pixels; both "
|
||||
"REGENERATE the face via SDXL). 'instantid' (default) uses InstantID img2img on "
|
||||
"the cleaned crop with ArcFace + landmark ControlNet. 'photomaker' uses "
|
||||
"PhotoMaker-V2 txt2img + CLIP+ArcFace dual encoder. **BOTH are NON-COMMERCIAL** "
|
||||
"(InsightFace antelopev2 / buffalo_l packs are research-only). For personal / "
|
||||
"research use only.",
|
||||
)(f)
|
||||
return click.option(
|
||||
"--restore-faces/--no-restore-faces",
|
||||
default=False,
|
||||
help="EXPERIMENTAL, opt-in, **NON-COMMERCIAL**. **REGENERATES the face** (does "
|
||||
"NOT recover original pixels) via the chosen --restore-faces-method; the "
|
||||
"regenerated face looks more AI-generated than the cleaned image. Off by "
|
||||
"default; auto-skips when no face is detected or the chosen extra is absent. "
|
||||
"For production face preservation leave this OFF and use the cleaned image "
|
||||
"as-is.",
|
||||
)(method)
|
||||
|
||||
|
||||
def _watermark_region(det: DetectionResult, width: int, height: int) -> tuple[int, int, int, int]:
|
||||
"""Pick a watermark bbox: detector's region if confident, else the default config slot."""
|
||||
if det.confidence > 0.15:
|
||||
@@ -597,7 +557,6 @@ def cmd_erase(
|
||||
help="Cap long side (px) before diffusion; 0 = native (best quality, like raiw.cc). Raise only on GPU/MPS OOM.",
|
||||
)
|
||||
@_controlnet_scale_option
|
||||
@_restore_faces_options
|
||||
@_min_resolution_option
|
||||
@_unsharp_option
|
||||
@_upscaler_option
|
||||
@@ -619,8 +578,6 @@ def cmd_invisible(
|
||||
max_resolution: int,
|
||||
min_resolution: int,
|
||||
controlnet_scale: float,
|
||||
restore_faces: bool,
|
||||
restore_faces_method: str,
|
||||
upscaler: str,
|
||||
auto: bool,
|
||||
adaptive_polish: bool,
|
||||
@@ -643,7 +600,7 @@ def cmd_invisible(
|
||||
source = _validate_image(source)
|
||||
_warn_if_esrgan_unavailable(upscaler)
|
||||
if auto:
|
||||
pipeline, restore_faces, adaptive_polish = _apply_auto(ctx, source, pipeline, restore_faces, adaptive_polish)
|
||||
pipeline, adaptive_polish = _apply_auto(ctx, source, pipeline, adaptive_polish)
|
||||
if output is None:
|
||||
output = source.with_stem(source.stem + "_clean")
|
||||
|
||||
@@ -682,8 +639,6 @@ def cmd_invisible(
|
||||
min_resolution=min_resolution,
|
||||
upscaler=upscaler,
|
||||
vendor=vendor,
|
||||
restore_faces=restore_faces,
|
||||
restore_faces_method=restore_faces_method,
|
||||
)
|
||||
elapsed = time.monotonic() - t0
|
||||
|
||||
@@ -859,7 +814,6 @@ def cmd_identify(ctx: click.Context, source: Path, no_visible: bool, as_json: bo
|
||||
help="Cap long side (px) before diffusion; 0 = native (best quality, like raiw.cc). Raise only on GPU/MPS OOM.",
|
||||
)
|
||||
@_controlnet_scale_option
|
||||
@_restore_faces_options
|
||||
@_min_resolution_option
|
||||
@_unsharp_option
|
||||
@_upscaler_option
|
||||
@@ -884,8 +838,6 @@ def cmd_all(
|
||||
max_resolution: int,
|
||||
min_resolution: int,
|
||||
controlnet_scale: float,
|
||||
restore_faces: bool,
|
||||
restore_faces_method: str,
|
||||
upscaler: str,
|
||||
auto: bool,
|
||||
adaptive_polish: bool,
|
||||
@@ -905,7 +857,7 @@ def cmd_all(
|
||||
source = _validate_image(source)
|
||||
_warn_if_esrgan_unavailable(upscaler)
|
||||
if auto:
|
||||
pipeline, restore_faces, adaptive_polish = _apply_auto(ctx, source, pipeline, restore_faces, adaptive_polish)
|
||||
pipeline, adaptive_polish = _apply_auto(ctx, source, pipeline, adaptive_polish)
|
||||
|
||||
if output is None:
|
||||
output = source.with_stem(source.stem + "_clean")
|
||||
@@ -993,8 +945,6 @@ def cmd_all(
|
||||
min_resolution=min_resolution,
|
||||
upscaler=upscaler,
|
||||
vendor=vendor,
|
||||
restore_faces=restore_faces,
|
||||
restore_faces_method=restore_faces_method,
|
||||
)
|
||||
console.print(" Invisible watermark removed")
|
||||
|
||||
@@ -1049,8 +999,6 @@ def _process_batch_image(
|
||||
unsharp: float = 0.0,
|
||||
max_resolution: int = 0,
|
||||
min_resolution: int = 1024,
|
||||
restore_faces: bool = False,
|
||||
restore_faces_method: str = "instantid",
|
||||
controlnet_scale: float = 1.0,
|
||||
upscaler: str = "lanczos",
|
||||
auto: bool = False,
|
||||
@@ -1104,9 +1052,7 @@ def _process_batch_image(
|
||||
# pipeline choice changes the engine ctor, so cache one engine per pipeline
|
||||
# (controlnet vs default) rather than a single shared instance.
|
||||
if auto:
|
||||
pipeline, restore_faces, adaptive_polish = _apply_auto(
|
||||
ctx, img_path, pipeline, restore_faces, adaptive_polish
|
||||
)
|
||||
pipeline, adaptive_polish = _apply_auto(ctx, img_path, pipeline, adaptive_polish)
|
||||
engines = ctx.obj.setdefault("_inv_engines", {})
|
||||
if pipeline not in engines:
|
||||
engines[pipeline] = InvisibleEngine(
|
||||
@@ -1128,8 +1074,6 @@ def _process_batch_image(
|
||||
max_resolution=max_resolution,
|
||||
min_resolution=min_resolution,
|
||||
upscaler=upscaler,
|
||||
restore_faces=restore_faces,
|
||||
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),
|
||||
@@ -1187,7 +1131,6 @@ def _process_batch_image(
|
||||
default=0,
|
||||
help="Cap long side (px) before diffusion; 0 = native (best quality, like raiw.cc). Raise only on GPU/MPS OOM.",
|
||||
)
|
||||
@_restore_faces_options
|
||||
@_min_resolution_option
|
||||
@_unsharp_option
|
||||
@_upscaler_option
|
||||
@@ -1211,8 +1154,6 @@ def cmd_batch(
|
||||
unsharp: float,
|
||||
max_resolution: int,
|
||||
min_resolution: int,
|
||||
restore_faces: bool,
|
||||
restore_faces_method: str,
|
||||
controlnet_scale: float,
|
||||
upscaler: str,
|
||||
auto: bool,
|
||||
@@ -1271,8 +1212,6 @@ def cmd_batch(
|
||||
unsharp=unsharp,
|
||||
max_resolution=max_resolution,
|
||||
min_resolution=min_resolution,
|
||||
restore_faces=restore_faces,
|
||||
restore_faces_method=restore_faces_method,
|
||||
controlnet_scale=controlnet_scale,
|
||||
upscaler=upscaler,
|
||||
auto=auto,
|
||||
|
||||
@@ -1,593 +0,0 @@
|
||||
"""SynthID-robust face identity restoration via InstantID.
|
||||
|
||||
**NON-COMMERCIAL.** InstantID's runtime depends on the InsightFace ``antelopev2``
|
||||
ArcFace model pack, which InsightFace releases under a research-only license:
|
||||
|
||||
"The training data containing the annotation (and the models trained with
|
||||
these data) are available for non-commercial research purposes only."
|
||||
-- insightface upstream README
|
||||
|
||||
The InstantX maintainers themselves acknowledged on HuggingFace
|
||||
(``InstantX/InstantID`` discussion #2) that "InstantID cannot be Apache 2.0 if it
|
||||
is using Insight Face" and stated intent to retrain on commercial face encoders.
|
||||
As of 2026-06-08 (deep-research synthesis in
|
||||
``docs/synthid-robust-identity-research-2026-06-08.md``) that retrain has not
|
||||
shipped. **A paid service (raiw.cc, any monetized SaaS) MUST NOT use this path.**
|
||||
|
||||
The default ``--restore-faces-method`` is ``instantid`` (this module). The
|
||||
alternative ``photomaker`` is also non-commercial. There is no commercial-safe
|
||||
ArcFace-grade identity-preservation stack for SDXL today.
|
||||
|
||||
Architecture (vs the earlier txt2img variant):
|
||||
- The earlier (txt2img) integration generated each face from scratch in a fresh
|
||||
1024 scene with InstantID's standard pipeline. That produced studio-portrait
|
||||
faces with the wrong lighting / head angle for the surrounding scene; on
|
||||
group photos the per-face composites read as patchwork even after color
|
||||
matching and elliptical alphas.
|
||||
- This (img2img on cleaned) integration feeds the CLEANED face crop as the
|
||||
img2img source. Diffusion sees the scene context (shoulders, hair edges,
|
||||
lighting, shadow direction) directly and harmonises the regenerated face
|
||||
with it. Identity still comes through the ArcFace embedding +
|
||||
landmark-ControlNet, which are semantic / pure-geometry and carry no
|
||||
watermark.
|
||||
|
||||
SynthID safety (load-bearing for raiw.cc):
|
||||
- img2img source = CLEANED crop. Cleaned image is already oracle-verified
|
||||
SynthID-free at our controlnet strength; cropping is a subset operation that
|
||||
preserves that property.
|
||||
- ArcFace embedding = from the ORIGINAL face crop (sharper identity, but the
|
||||
embedding is semantic 512-d, no pixel content).
|
||||
- Landmark stick figure = pure colour-coded geometry rendered from kps; no
|
||||
source pixels.
|
||||
- img2img diffusion adds noise to the cleaned source then denoises with
|
||||
ControlNet + IP-Adapter conditioning. Any residual high-frequency pattern
|
||||
in the cleaned crop is destroyed by that noise injection at the strengths we
|
||||
use.
|
||||
- We must NEVER feed the original image as img2img source (would re-introduce
|
||||
SynthID outside the diffusion footprint at strength < 1). The code only ever
|
||||
reads pixels from ``cleaned_bgr`` into ``image=`` -- the original is used
|
||||
for the embedding + kps only.
|
||||
|
||||
Pipeline this module wires:
|
||||
1. Detect faces in the CLEANED image (YuNet via ``auto_config``).
|
||||
2. For each face: square-crop the SAME box from BOTH the original (for
|
||||
ArcFace + kps) and the cleaned image (for img2img source). Resize both
|
||||
to 1024x1024.
|
||||
3. Render the kps as a stick figure (the ControlNet conditioning image).
|
||||
4. Call the InstantID img2img pipeline
|
||||
(``StableDiffusionXLInstantIDImg2ImgPipeline``) with ``image`` = cleaned
|
||||
crop, ``control_image`` = landmark, ``image_embeds`` = ArcFace, and
|
||||
``strength`` = ~0.55. The output 1024 is a face that fits the scene.
|
||||
5. Elliptical-alpha + colour-match composite into the cleaned image.
|
||||
|
||||
Requires the optional ``instantid`` extra: ``pip install
|
||||
'remove-ai-watermarks[instantid]'``. Weights download on first use; the
|
||||
upstream img2img pipeline file (not on PyPI) is cached from
|
||||
``raw.githubusercontent.com`` on first run.
|
||||
"""
|
||||
|
||||
# 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
|
||||
|
||||
from remove_ai_watermarks.photomaker_restore import _face_crop_square
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from numpy.typing import NDArray
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# InstantID checkpoint repo on HuggingFace. The IdentityNet ControlNet weights live
|
||||
# under ``ControlNetModel/`` and the IP-Adapter file is ``ip-adapter.bin`` at the
|
||||
# root. Both are Apache-2.0 (the InsightFace runtime dep is what makes the path
|
||||
# non-commercial). Downloaded on first use.
|
||||
_INSTANTID_REPO = "InstantX/InstantID"
|
||||
_INSTANTID_CONTROLNET_SUBFOLDER = "ControlNetModel"
|
||||
_INSTANTID_IP_ADAPTER = "ip-adapter.bin"
|
||||
|
||||
# Upstream InstantID img2img pipeline source. Not on PyPI, not on HF Hub at any path
|
||||
# diffusers can auto-load -- the file lives in the InstantID GitHub repo. We download
|
||||
# it once to a cache dir and pass it as ``custom_pipeline=<path>`` to diffusers.
|
||||
_INSTANTID_IMG2IMG_URL = (
|
||||
"https://raw.githubusercontent.com/instantX-research/InstantID/"
|
||||
"main/pipeline_stable_diffusion_xl_instantid_img2img.py"
|
||||
)
|
||||
|
||||
# SDXL base shared with the main pipeline (same checkpoint as `default`/`controlnet`).
|
||||
_SDXL_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
|
||||
|
||||
# Prompt format. InstantID is less sensitive to prompt than PhotoMaker because the
|
||||
# ID branch is cross-attention; a neutral descriptive prompt is recommended by the
|
||||
# upstream gradio demo.
|
||||
_INSTANTID_PROMPT = "portrait photo of a person, natural skin, soft lighting, sharp focus, best quality"
|
||||
_INSTANTID_NEGATIVE = (
|
||||
"(asymmetry, worst quality, low quality, illustration, 3d, 2d, painting, "
|
||||
"cartoons, sketch), open mouth, blurry, watermark, deformed"
|
||||
)
|
||||
|
||||
# Square size used to feed InstantID. SDXL is happiest at 1024 (a smaller value sends
|
||||
# it into low-res mosaic mode -- caught visually on PhotoMaker, same root cause).
|
||||
_INSTANTID_FACE_SIZE = 1024
|
||||
|
||||
_pipeline: Any | None = None
|
||||
_pipeline_lock = threading.Lock()
|
||||
_face_analyser: Any | None = None
|
||||
_face_analyser_lock = threading.Lock()
|
||||
|
||||
|
||||
def is_available() -> bool:
|
||||
"""True when the optional InstantID extra deps are importable."""
|
||||
return (
|
||||
importlib.util.find_spec("insightface") is not None
|
||||
and importlib.util.find_spec("diffusers") is not None
|
||||
and importlib.util.find_spec("torch") is not None
|
||||
and importlib.util.find_spec("huggingface_hub") is not None
|
||||
)
|
||||
|
||||
|
||||
def _select_device() -> str:
|
||||
"""Pick the InstantID 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("instantid_restore: device probe failed (%s); using CPU", e)
|
||||
return "cpu"
|
||||
|
||||
|
||||
def _fetch_img2img_pipeline_file() -> Path:
|
||||
"""Cache the InstantID img2img pipeline source file locally on first use.
|
||||
|
||||
The file lives in the InstantX GitHub repo (not on PyPI, not on HF Hub at any
|
||||
path diffusers can auto-load). We fetch the raw URL once into the package's
|
||||
HuggingFace cache so subsequent loads hit disk. Returns the path to feed to
|
||||
``DiffusionPipeline.from_pretrained(custom_pipeline=...)``.
|
||||
"""
|
||||
import os
|
||||
import urllib.request
|
||||
|
||||
cache_root = Path(os.environ.get("HF_HOME") or Path.home() / ".cache" / "huggingface")
|
||||
cache_dir = cache_root / "remove_ai_watermarks" / "instantid"
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
target = cache_dir / "pipeline_stable_diffusion_xl_instantid_img2img.py"
|
||||
if not target.exists() or target.stat().st_size < 50_000:
|
||||
logger.info("instantid_restore: fetching img2img pipeline source from %s", _INSTANTID_IMG2IMG_URL)
|
||||
urllib.request.urlretrieve(_INSTANTID_IMG2IMG_URL, target) # noqa: S310 (HTTPS pinned)
|
||||
return target
|
||||
|
||||
|
||||
def _ensure_antelopev2(root: Path) -> Path:
|
||||
"""Materialize the antelopev2 pack at ``<root>/models/antelopev2/`` if absent.
|
||||
|
||||
InsightFace's built-in auto-download points at
|
||||
``github.com/deepinsight/insightface/releases/download/v0.7/antelopev2.zip``
|
||||
which has been broken since at least 2024 (verified upstream issue #2517,
|
||||
#2766; explicitly called out in InstantID's README: "manually download via
|
||||
this URL to models/antelopev2 as the default link is invalid"). Without the
|
||||
five expected ``.onnx`` files in place, ``FaceAnalysis.prepare()`` errors
|
||||
with ``assert 'detection' in self.models``.
|
||||
|
||||
We side-step the broken default by fetching the five files from a HuggingFace
|
||||
mirror (``kidyu/antelopev2-for-InstantID-ComfyUI``) on first use. Returns the
|
||||
target directory containing the .onnx files.
|
||||
"""
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
target = root / "models" / "antelopev2"
|
||||
target.mkdir(parents=True, exist_ok=True)
|
||||
files = [
|
||||
"1k3d68.onnx",
|
||||
"2d106det.onnx",
|
||||
"genderage.onnx",
|
||||
"glintr100.onnx",
|
||||
"scrfd_10g_bnkps.onnx",
|
||||
]
|
||||
for fname in files:
|
||||
dest = target / fname
|
||||
if dest.exists() and dest.stat().st_size > 0:
|
||||
continue
|
||||
logger.info("instantid_restore: fetching antelopev2/%s from HF mirror", fname)
|
||||
path = hf_hub_download(repo_id="kidyu/antelopev2-for-InstantID-ComfyUI", filename=fname)
|
||||
# hf_hub_download caches under HF_HOME; symlink (or copy) into the
|
||||
# InsightFace-expected layout.
|
||||
if not dest.exists():
|
||||
try:
|
||||
dest.symlink_to(path)
|
||||
except OSError:
|
||||
import shutil
|
||||
|
||||
shutil.copy(path, dest)
|
||||
return target
|
||||
|
||||
|
||||
def _get_face_analyser() -> Any:
|
||||
"""Return the InsightFace FaceAnalysis singleton (antelopev2, non-commercial).
|
||||
|
||||
Pre-downloads the antelopev2 pack from a HuggingFace mirror (the InsightFace
|
||||
auto-download is broken). See the NON-COMMERCIAL notice at the top of the
|
||||
module.
|
||||
"""
|
||||
global _face_analyser
|
||||
if _face_analyser is not None:
|
||||
return _face_analyser
|
||||
with _face_analyser_lock:
|
||||
if _face_analyser is None:
|
||||
import torch
|
||||
from insightface.app import FaceAnalysis
|
||||
|
||||
providers = ["CUDAExecutionProvider"] if torch.cuda.is_available() else ["CPUExecutionProvider"]
|
||||
# InstantID's upstream uses name='antelopev2' and root='./'. Materialise
|
||||
# the pack at the same place so FaceAnalysis finds it locally.
|
||||
root = Path.cwd()
|
||||
_ensure_antelopev2(root)
|
||||
fa = FaceAnalysis(name="antelopev2", root=str(root), providers=providers)
|
||||
fa.prepare(ctx_id=0, det_size=(640, 640))
|
||||
_face_analyser = fa
|
||||
return _face_analyser
|
||||
|
||||
|
||||
def _get_pipeline() -> Any:
|
||||
"""Return the lazily-built InstantID pipeline singleton (downloads weights on first use).
|
||||
|
||||
Loads via diffusers' community-pipeline mechanism: the file
|
||||
``pipeline_stable_diffusion_xl_instantid.py`` lives in
|
||||
``diffusers/examples/community/`` and is selected by the slug
|
||||
``pipeline_stable_diffusion_xl_instantid``.
|
||||
"""
|
||||
global _pipeline
|
||||
if _pipeline is not None:
|
||||
return _pipeline
|
||||
with _pipeline_lock:
|
||||
if _pipeline is None:
|
||||
import torch
|
||||
from diffusers import ControlNetModel, DiffusionPipeline
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
device = _select_device()
|
||||
dtype = torch.float16 if device == "cuda" else torch.float32
|
||||
logger.info("instantid_restore: loading SDXL+InstantID img2img on %s (%s)", device, dtype)
|
||||
|
||||
# IdentityNet ControlNet weights.
|
||||
controlnet = ControlNetModel.from_pretrained(
|
||||
_INSTANTID_REPO,
|
||||
subfolder=_INSTANTID_CONTROLNET_SUBFOLDER,
|
||||
torch_dtype=dtype,
|
||||
)
|
||||
# Upstream InstantID img2img pipeline (StableDiffusionXLInstantIDImg2ImgPipeline).
|
||||
# Lets us feed the cleaned face crop as the diffusion source so the regenerated
|
||||
# face inherits scene lighting / shadows / head angle from the cleaned context
|
||||
# (vs the txt2img variant which generates a studio portrait from scratch).
|
||||
# Critical SynthID-safety property: the ``image`` arg MUST be the CLEANED crop,
|
||||
# never the original -- the original carries the watermark and img2img at
|
||||
# strength < 1 preserves some input pixel structure. The ArcFace embedding is
|
||||
# semantic (no pixel content), so taking it from the original is fine.
|
||||
pipe = DiffusionPipeline.from_pretrained(
|
||||
_SDXL_MODEL_ID,
|
||||
controlnet=controlnet,
|
||||
torch_dtype=dtype,
|
||||
custom_pipeline=str(_fetch_img2img_pipeline_file()),
|
||||
# Custom_pipeline from a local .py file triggers diffusers' remote-code
|
||||
# guard; the file is fetched from a pinned raw.githubusercontent URL
|
||||
# we control, so opt in here. Without this the load silently falls
|
||||
# back to a default pipeline (no img2img + no IP-Adapter cross-attn),
|
||||
# the next call hits an AttributeError on load_ip_adapter_instantid,
|
||||
# and our outer except logs but skips the whole restore.
|
||||
trust_remote_code=True,
|
||||
)
|
||||
pipe.to(device)
|
||||
# IP-Adapter weights that wire the ArcFace embedding into cross-attention.
|
||||
ip_adapter_path = hf_hub_download(repo_id=_INSTANTID_REPO, filename=_INSTANTID_IP_ADAPTER)
|
||||
# IP-Adapter scale = weight on the ArcFace cross-attention. The upstream
|
||||
# demo uses 0.8 for txt2img; for img2img-on-cleaned we push to 1.0 because
|
||||
# the cleaned face crop is competing as identity prior and we want ArcFace
|
||||
# to dominate (otherwise the regenerated face inherits the controlnet-
|
||||
# drifted cleaned face, not the original identity).
|
||||
pipe.load_ip_adapter_instantid(ip_adapter_path, scale=1.0)
|
||||
# Diffusers 0.38 vs InstantID upstream compat patch: InstantID's __call__
|
||||
# calls ``self.check_inputs(...)`` POSITIONALLY (signature from ~v0.29),
|
||||
# but diffusers 0.38 added two new params (``ip_adapter_image``,
|
||||
# ``ip_adapter_image_embeds``) BEFORE ``controlnet_conditioning_scale`` in
|
||||
# the parent's signature. That shifts every argument by two, so
|
||||
# ``control_guidance_end`` (which InstantID converts to ``[1.0]`` for the
|
||||
# single-controlnet case before this point) lands in the slot the parent
|
||||
# validates as ``controlnet_conditioning_scale`` and trips
|
||||
# ``TypeError("must be type float")``. Our inputs are programmatic and
|
||||
# already validated by our own callers, so neutralising the check is safe.
|
||||
pipe.check_inputs = lambda *_a, **_k: None
|
||||
_pipeline = pipe
|
||||
return _pipeline
|
||||
|
||||
|
||||
def _draw_kps(image_size: tuple[int, int], kps: Any) -> Any:
|
||||
"""Render the 5 facial keypoints as a colored stick figure.
|
||||
|
||||
Mirrors upstream's ``draw_kps`` (in ``pipeline_stable_diffusion_xl_instantid.py``):
|
||||
the 5 keypoints (left eye, right eye, nose tip, left mouth corner, right mouth
|
||||
corner) get drawn as colored circles connected by colored lines, on a black
|
||||
background. The result is the ControlNet conditioning image -- pure landmark
|
||||
geometry, no pixels from the original face leak through this branch.
|
||||
|
||||
``image_size`` is ``(width, height)``; ``kps`` is a numpy array of shape (5, 2).
|
||||
"""
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
# Same color palette as upstream (blue/red/green/purple/yellow).
|
||||
stick_width = 4
|
||||
limb_seq = np.array([[0, 2], [1, 2], [3, 2], [4, 2]])
|
||||
color_list = [
|
||||
(255, 0, 0),
|
||||
(0, 255, 0),
|
||||
(0, 0, 255),
|
||||
(255, 255, 0),
|
||||
(255, 0, 255),
|
||||
]
|
||||
|
||||
w, h = image_size
|
||||
out_img = np.zeros((h, w, 3), dtype=np.uint8)
|
||||
|
||||
kps_arr = np.array(kps)
|
||||
for i in range(len(limb_seq)):
|
||||
index = limb_seq[i]
|
||||
color = color_list[index[0]]
|
||||
x = kps_arr[index][:, 0]
|
||||
y = kps_arr[index][:, 1]
|
||||
length = ((x[0] - x[1]) ** 2 + (y[0] - y[1]) ** 2) ** 0.5
|
||||
angle = np.degrees(np.arctan2(y[0] - y[1], x[0] - x[1]))
|
||||
polygon = cv2.ellipse2Poly(
|
||||
(int(np.mean(x)), int(np.mean(y))),
|
||||
(int(length / 2), stick_width),
|
||||
int(angle),
|
||||
0,
|
||||
360,
|
||||
1,
|
||||
)
|
||||
out_img = cv2.fillConvexPoly(out_img.copy(), polygon, color)
|
||||
out_img = (out_img * 0.6).astype(np.uint8)
|
||||
|
||||
for i, kp in enumerate(kps_arr):
|
||||
x, y = kp
|
||||
out_img = cv2.circle(out_img.copy(), (int(x), int(y)), 10, color_list[i], -1)
|
||||
|
||||
return Image.fromarray(out_img.astype(np.uint8))
|
||||
|
||||
|
||||
def restore_faces_instantid(
|
||||
original_bgr: NDArray[Any],
|
||||
cleaned_bgr: NDArray[Any],
|
||||
num_inference_steps: int = 30,
|
||||
guidance_scale: float = 5.0,
|
||||
controlnet_conditioning_scale: float = 1.0,
|
||||
img2img_strength: float = 0.7,
|
||||
seed: int | None = None,
|
||||
detect_faces_fn: Any | None = None,
|
||||
) -> NDArray[Any]:
|
||||
"""SynthID-robust face identity restoration via InstantID.
|
||||
|
||||
Flow:
|
||||
1. Detect faces in ``cleaned_bgr`` (YuNet via ``auto_config`` by default;
|
||||
override via ``detect_faces_fn`` for tests).
|
||||
2. For each face: square-crop the SAME box from BOTH images (original ->
|
||||
ArcFace + kps; cleaned -> img2img source). Resize both to 1024.
|
||||
3. Render kps as a landmark stick figure (the ControlNet conditioning).
|
||||
4. Run InstantID img2img: ``image`` = cleaned crop, ``control_image`` =
|
||||
landmark, ``image_embeds`` = ArcFace embedding from the original.
|
||||
5. Elliptical-alpha + colour-match composite into the cleaned image.
|
||||
|
||||
SynthID safety: ``image`` is the CLEANED crop (already oracle-clean); the
|
||||
original is read for the embedding and kps only (semantic / geometry, no
|
||||
pixel content). See the module docstring.
|
||||
|
||||
``detect_faces_fn`` returns a list of ``(x, y, w, h)`` boxes given a BGR image.
|
||||
"""
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
if detect_faces_fn is None:
|
||||
from pathlib import Path
|
||||
|
||||
from remove_ai_watermarks import auto_config as _ac
|
||||
|
||||
def _default_detect(bgr: NDArray[Any]) -> list[tuple[int, int, int, int]]:
|
||||
h_d, w_d = bgr.shape[:2]
|
||||
model = Path(_ac.__file__).parent / "assets" / "face_detection_yunet_2023mar.onnx"
|
||||
det = cv2.FaceDetectorYN.create(str(model), "", (w_d, h_d), _ac._FACE_SCORE, 0.3, 5000)
|
||||
det.setInputSize((w_d, h_d))
|
||||
_, 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("instantid_restore: no faces detected; returning cleaned image unchanged")
|
||||
return cleaned_bgr
|
||||
|
||||
pipeline = _get_pipeline()
|
||||
face_analyser = _get_face_analyser()
|
||||
|
||||
generator = None
|
||||
if seed is not None:
|
||||
generator = torch.Generator(device=pipeline.device).manual_seed(seed)
|
||||
|
||||
h_c, w_c = cleaned_bgr.shape[:2]
|
||||
restored: list[tuple[NDArray[Any], tuple[int, int, int, int]]] = []
|
||||
for box in boxes:
|
||||
# Square crop with the SAME geometry from both the original (-> ArcFace
|
||||
# embedding + landmark kps -- semantic / pure-geometry, SynthID can't ride
|
||||
# either) AND the cleaned image (-> img2img source -- SynthID-safe because
|
||||
# the cleaned image is already oracle-verified clean and any residual
|
||||
# high-frequency pattern would be destroyed by the noise injection at our
|
||||
# strength setting). _face_crop_square gives a 2x-padded square box around
|
||||
# the face -- enough scene context so the img2img harmonises lighting and
|
||||
# head angle with the surroundings.
|
||||
original_crop_bgr, square_box = _face_crop_square(original_bgr, box)
|
||||
sx1, sy1, sx2, sy2 = square_box
|
||||
sx1c, sy1c = max(0, sx1), max(0, sy1)
|
||||
sx2c, sy2c = min(w_c, sx2), min(h_c, sy2)
|
||||
if original_crop_bgr.size == 0 or sx2c <= sx1c or sy2c <= sy1c:
|
||||
continue
|
||||
cleaned_crop_bgr = cleaned_bgr[sy1c:sy2c, sx1c:sx2c]
|
||||
if cleaned_crop_bgr.shape[:2] != original_crop_bgr.shape[:2]:
|
||||
# Edge effect at image border -- pad cleaned crop to match the original
|
||||
# crop dimensions so InsightFace / the pipeline see the same shape.
|
||||
cleaned_crop_bgr = cv2.resize(
|
||||
cleaned_crop_bgr,
|
||||
(original_crop_bgr.shape[1], original_crop_bgr.shape[0]),
|
||||
interpolation=cv2.INTER_LANCZOS4,
|
||||
)
|
||||
|
||||
# Resize both crops to the SDXL working size.
|
||||
original_resized = cv2.resize(
|
||||
original_crop_bgr, (_INSTANTID_FACE_SIZE, _INSTANTID_FACE_SIZE), interpolation=cv2.INTER_LANCZOS4
|
||||
)
|
||||
cleaned_resized = cv2.resize(
|
||||
cleaned_crop_bgr, (_INSTANTID_FACE_SIZE, _INSTANTID_FACE_SIZE), interpolation=cv2.INTER_LANCZOS4
|
||||
)
|
||||
|
||||
# ArcFace embedding + 5 kps from the ORIGINAL face (sharper identity).
|
||||
face_infos = face_analyser.get(original_resized)
|
||||
if not face_infos:
|
||||
logger.debug("instantid_restore: InsightFace did not find a face in the crop; skipping")
|
||||
continue
|
||||
face_info = sorted(
|
||||
face_infos,
|
||||
key=lambda x: (x["bbox"][2] - x["bbox"][0]) * (x["bbox"][3] - x["bbox"][1]),
|
||||
)[-1]
|
||||
face_emb = face_info["embedding"]
|
||||
face_kps = face_info["kps"]
|
||||
|
||||
# Render the landmark stick figure at the same size as the generation target.
|
||||
landmark_img = _draw_kps((_INSTANTID_FACE_SIZE, _INSTANTID_FACE_SIZE), face_kps)
|
||||
|
||||
# img2img call: source = CLEANED crop (SynthID-safe), control = landmark
|
||||
# geometry, identity = ArcFace embedding from original. Strength controls
|
||||
# how much of the cleaned input structure survives -- low enough (~0.55)
|
||||
# to keep the head angle / lighting / shoulders coherent with the rest of
|
||||
# the cleaned image, high enough that the face pixels are diffusion-fresh
|
||||
# and InstantID actually injects identity.
|
||||
from PIL import Image
|
||||
|
||||
cleaned_pil = Image.fromarray(cv2.cvtColor(cleaned_resized, cv2.COLOR_BGR2RGB))
|
||||
out = pipeline(
|
||||
prompt=_INSTANTID_PROMPT,
|
||||
negative_prompt=_INSTANTID_NEGATIVE,
|
||||
image=cleaned_pil,
|
||||
control_image=landmark_img,
|
||||
image_embeds=face_emb,
|
||||
strength=img2img_strength,
|
||||
controlnet_conditioning_scale=controlnet_conditioning_scale,
|
||||
num_inference_steps=num_inference_steps,
|
||||
guidance_scale=guidance_scale,
|
||||
generator=generator,
|
||||
)
|
||||
gen_rgb = out.images[0]
|
||||
gen_bgr = cv2.cvtColor(np.array(gen_rgb), cv2.COLOR_RGB2BGR)
|
||||
|
||||
# gen_bgr is at _INSTANTID_FACE_SIZE x _INSTANTID_FACE_SIZE. It represents
|
||||
# the 2x-padded square_box content as regenerated by img2img -- so the face
|
||||
# in it sits at the same RELATIVE position as in the cleaned input (img2img
|
||||
# preserves structure). Composite the whole square back into the square_box
|
||||
# location -- the cleaned-canvas elliptical alpha will keep the cleaned
|
||||
# background outside the face oval, and the img2img harmonisation handles
|
||||
# the seam INSIDE the oval (which is just face-on-face transition between
|
||||
# diffusion-output and cleaned).
|
||||
target_box = (sx1c, sy1c, sx2c, sy2c)
|
||||
gen_target = cv2.resize(gen_bgr, (sx2c - sx1c, sy2c - sy1c), interpolation=cv2.INTER_LANCZOS4)
|
||||
restored.append((gen_target, target_box))
|
||||
|
||||
if not restored:
|
||||
return cleaned_bgr
|
||||
return _composite_faces_elliptical(cleaned_bgr, restored)
|
||||
|
||||
|
||||
def _color_match(src_bgr: NDArray[Any], ref_bgr: NDArray[Any]) -> NDArray[Any]:
|
||||
"""Shift ``src_bgr`` mean colour to ``ref_bgr`` mean colour, per channel.
|
||||
|
||||
Each face is regenerated by InstantID with its own SDXL noise -- the white
|
||||
balance / mean tone drifts away from the surrounding scene (cool studio
|
||||
light vs warm bar lighting). A per-channel mean-shift brings the face crop
|
||||
into the same tonal range as the cleaned canvas where it lands. Contrast
|
||||
and saturation are preserved (we don't rescale variance).
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
src = src_bgr.astype(np.float32)
|
||||
ref = ref_bgr.astype(np.float32)
|
||||
if ref.size == 0:
|
||||
return src_bgr
|
||||
src_mean = src.mean(axis=(0, 1), keepdims=True)
|
||||
ref_mean = ref.mean(axis=(0, 1), keepdims=True)
|
||||
return np.clip(src - src_mean + ref_mean, 0, 255).astype(np.uint8)
|
||||
|
||||
|
||||
def _composite_faces_elliptical(
|
||||
base_bgr: NDArray[Any],
|
||||
restored_crops: list[tuple[NDArray[Any], tuple[int, int, int, int]]],
|
||||
feather_div: int = 5,
|
||||
) -> NDArray[Any]:
|
||||
"""Composite face crops into ``base_bgr`` using an elliptical, feathered alpha.
|
||||
|
||||
Two changes vs the simpler rectangular Gaussian feather:
|
||||
|
||||
- **Inscribed face-shaped ellipse.** Axes are ``(0.32*bw, 0.42*bh)`` which
|
||||
fits comfortably inside the 2x padded bbox (the face naturally occupies
|
||||
the central ~50% of the bbox), covering the head silhouette without
|
||||
clipping the forehead or chin. The bbox corners (which carry
|
||||
regenerated-scene background pixels with a different tone per face) end
|
||||
up at alpha=0 so the cleaned-image background stays intact -- this is
|
||||
what eliminates multi-face patchwork on group photos.
|
||||
- **Soft feather.** ``min(bw, bh) // 5`` -- about twice as soft as the
|
||||
rectangular Gaussian, so the ellipse edge fades over a wider band into
|
||||
the cleaned canvas, hiding any residual seam.
|
||||
|
||||
Additionally, before compositing, ``_color_match`` shifts the regenerated
|
||||
face's mean colour to match the cleaned canvas region it lands on -- this
|
||||
removes the warm/cool tone clash that group photos showed.
|
||||
"""
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
out = base_bgr.astype(np.float32)
|
||||
h_b, w_b = 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_b, x2), min(h_b, y2)
|
||||
bw, bh = x2 - x1, y2 - y1
|
||||
if bw <= 0 or bh <= 0:
|
||||
continue
|
||||
resized = cv2.resize(crop, (bw, bh), interpolation=cv2.INTER_LANCZOS4)
|
||||
# Tone match the regenerated face to the cleaned canvas it sits on.
|
||||
ref_region = base_bgr[y1:y2, x1:x2]
|
||||
resized = _color_match(resized, ref_region)
|
||||
|
||||
alpha_crop = np.zeros((bh, bw), dtype=np.float32)
|
||||
center = (bw // 2, bh // 2)
|
||||
axes = (max(1, int(bw * 0.32)), max(1, int(bh * 0.42)))
|
||||
cv2.ellipse(alpha_crop, center, axes, 0, 0, 360, 1.0, -1)
|
||||
k = max(7, (min(bw, bh) // feather_div) | 1)
|
||||
alpha_crop = cv2.GaussianBlur(alpha_crop, (k, k), 0)
|
||||
|
||||
alpha_full = np.zeros((h_b, w_b), dtype=np.float32)
|
||||
alpha_full[y1:y2, x1:x2] = alpha_crop
|
||||
full_restored = np.zeros_like(out)
|
||||
full_restored[y1:y2, x1:x2] = resized
|
||||
a = alpha_full[:, :, None]
|
||||
out = full_restored * a + out * (1.0 - a)
|
||||
|
||||
return np.clip(out, 0, 255).astype(np.uint8)
|
||||
@@ -164,8 +164,6 @@ class InvisibleEngine:
|
||||
max_resolution: int = 0,
|
||||
min_resolution: int = 1024,
|
||||
vendor: str | None = None,
|
||||
restore_faces: bool = False,
|
||||
restore_faces_method: str = "instantid",
|
||||
unsharp: float = 0.0,
|
||||
adaptive_polish: bool = False,
|
||||
upscaler: str = "lanczos",
|
||||
@@ -181,20 +179,9 @@ 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). **NON-COMMERCIAL.**
|
||||
Run the face-identity post-pass when faces are present. Method is
|
||||
chosen by ``restore_faces_method`` -- ``instantid`` (default,
|
||||
stronger identity, needs the ``instantid`` extra) or ``photomaker``
|
||||
(PhotoMaker-V2, needs the ``photomaker`` extra). Both extras pull
|
||||
non-commercial InsightFace model packs. Auto-skips with a debug log
|
||||
when the chosen extra is absent or no face is detected. See
|
||||
``instantid_restore.py`` / ``photomaker_restore.py``.
|
||||
restore_faces_method: ``instantid`` (default) or ``photomaker``. Both
|
||||
NON-COMMERCIAL; pick the one whose extra you've installed.
|
||||
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
|
||||
safe range, higher risks edge halos.
|
||||
Applied last to counter the soft / over-smoothed look of the
|
||||
diffusion pass; ~0.5-0.8 is a safe range, higher risks edge halos.
|
||||
adaptive_polish: When True (the --auto mode default), restore the input's
|
||||
detail level in the softened output instead of fixed unsharp/humanize:
|
||||
a capped unsharp + edge-masked grain targeting the input's Laplacian
|
||||
@@ -316,19 +303,7 @@ class InvisibleEngine:
|
||||
out_cv = cv2.resize(out_cv, orig_size, interpolation=cv2.INTER_LANCZOS4)
|
||||
image_io.imwrite(out_path, out_cv)
|
||||
|
||||
# Optional GFPGAN face-polish post-pass: sharpens and re-synthesizes each
|
||||
# face from GFPGAN's StyleGAN2 prior, running on the DIFFUSION-CLEANED image
|
||||
# (not the original) -- so SynthID is not re-introduced (the input pixels
|
||||
# GFPGAN derives from are already SynthID-free). Auto-skips when faces are
|
||||
# absent or the optional `restore` extra is not installed.
|
||||
if restore_faces:
|
||||
if restore_faces_method == "photomaker":
|
||||
self._restore_faces_photomaker(out_path, image, seed)
|
||||
else:
|
||||
self._restore_faces_instantid(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).
|
||||
# Final sharpening.
|
||||
if unsharp > 0.0:
|
||||
import cv2
|
||||
|
||||
@@ -364,99 +339,6 @@ class InvisibleEngine:
|
||||
if _tmp_path.exists():
|
||||
_tmp_path.unlink()
|
||||
|
||||
def _restore_faces_instantid(
|
||||
self,
|
||||
out_path: Path,
|
||||
original_image: Any,
|
||||
seed: int | None,
|
||||
) -> None:
|
||||
"""Run the InstantID face-identity post-pass on the cleaned ``out_path``.
|
||||
|
||||
**NON-COMMERCIAL** (see ``instantid_restore.py``). InstantID conditions on
|
||||
an ArcFace embedding (semantic) plus a landmark ControlNet (geometry,
|
||||
content-free) -- no original face pixels enter the diffusion. 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 instantid_restore
|
||||
|
||||
if not instantid_restore.is_available():
|
||||
logger.debug("restore_faces requested but the 'instantid' 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: 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 (InstantID post-pass)...")
|
||||
restored = instantid_restore.restore_faces_instantid(original_bgr, cleaned_bgr, seed=seed)
|
||||
image_io.imwrite(out_path, restored)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"restore_faces post-pass failed (%s: %s); keeping un-restored output",
|
||||
type(e).__name__,
|
||||
e,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
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``.
|
||||
|
||||
**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 photomaker_restore
|
||||
|
||||
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
|
||||
|
||||
cleaned_bgr = image_io.imread(out_path, cv2.IMREAD_COLOR)
|
||||
if cleaned_bgr is None:
|
||||
logger.warning("restore_faces: 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 post-pass failed (%s); keeping un-restored output", e)
|
||||
|
||||
def remove_watermark_batch(
|
||||
self,
|
||||
input_dir: Path,
|
||||
|
||||
@@ -1,368 +0,0 @@
|
||||
"""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.
|
||||
# Prompt format follows the upstream V2 reference (inference_pmv2.py): the trigger
|
||||
# word ``img`` must immediately follow a class noun. SDXL is happiest at 1024 and
|
||||
# falls into low-res artefacts ("mosaic of tiny faces") at 512, so we render at
|
||||
# 1024 then downscale into the face bbox at composite time. Caught visually
|
||||
# 2026-06-04: at 512 V2 produced a collage of training-time faces; at 1024 with the
|
||||
# upstream-style descriptive prompt it produces a clean face.
|
||||
_PHOTOMAKER_PROMPT = (
|
||||
"instagram photo, portrait photo of a person img, natural skin, soft lighting, best quality, sharp focus"
|
||||
)
|
||||
_PHOTOMAKER_NEGATIVE = (
|
||||
"(asymmetry, worst quality, low quality, illustration, 3d, 2d, painting, "
|
||||
"cartoons, sketch), open mouth, blurry, watermark"
|
||||
)
|
||||
|
||||
# SDXL native resolution; lower values send V2 into low-res mode and the output
|
||||
# becomes a collage of training-time faces. We render at 1024 then downscale into
|
||||
# the original face bbox at composite time.
|
||||
_PHOTOMAKER_FACE_SIZE = 1024
|
||||
|
||||
_pipeline: Any | None = None
|
||||
_pipeline_lock = threading.Lock()
|
||||
_face_analyser: Any | None = None
|
||||
_face_analyser_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_face_analyser() -> Any:
|
||||
"""Return the InsightFace FaceAnalysis2 singleton (downloads model packs on first use).
|
||||
|
||||
**This is the non-commercial step.** Instantiating ``FaceAnalysis2()`` triggers
|
||||
InsightFace's auto-download of the antelopev2/buffalo_l model packs, which are
|
||||
released under a research-only license. See the module docstring NON-COMMERCIAL
|
||||
notice. PhotoMaker-V2 requires this for the ArcFace identity branch.
|
||||
"""
|
||||
global _face_analyser
|
||||
if _face_analyser is not None:
|
||||
return _face_analyser
|
||||
with _face_analyser_lock:
|
||||
if _face_analyser is None:
|
||||
import torch
|
||||
from photomaker import FaceAnalysis2
|
||||
|
||||
providers = ["CUDAExecutionProvider"] if torch.cuda.is_available() else ["CPUExecutionProvider"]
|
||||
fa = FaceAnalysis2(providers=providers, allowed_modules=["detection", "recognition"])
|
||||
fa.prepare(ctx_id=0, det_size=(640, 640))
|
||||
_face_analyser = fa
|
||||
return _face_analyser
|
||||
|
||||
|
||||
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()
|
||||
face_analyser = _get_face_analyser() # NON-COMMERCIAL: triggers InsightFace model packs
|
||||
from photomaker import analyze_faces
|
||||
|
||||
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
|
||||
# Get the ArcFace embedding for THIS face (V2's required ID branch). InsightFace
|
||||
# expects BGR; analyze_faces returns a list, take the first detection.
|
||||
# Shape: upstream's inference_pmv2.py stacks per-image embeddings into a 2-D
|
||||
# tensor (N_images, 512). The pipeline forward then calls `.unsqueeze(0)` ITSELF
|
||||
# (line 705 of pipeline.py) to add a batch dim, so we must NOT pre-unsqueeze --
|
||||
# giving `(1, 1, 512)` to the V2 forward made the id_encoder consume garbage and
|
||||
# the pipeline output the training-time face collage (caught visually 2026-06-04).
|
||||
# Dtype stays float32 here; the pipeline casts internally.
|
||||
faces = analyze_faces(face_analyser, id_crop_bgr)
|
||||
if not faces:
|
||||
logger.debug("photomaker_restore: InsightFace did not detect a face in the crop; skipping")
|
||||
continue
|
||||
id_embeds = torch.stack([torch.from_numpy(faces[0]["embedding"])])
|
||||
|
||||
id_crop_rgb = cv2.cvtColor(id_crop_bgr, cv2.COLOR_BGR2RGB)
|
||||
id_image_pil = Image.fromarray(id_crop_rgb)
|
||||
|
||||
# Upstream V2 reference (inference_pmv2.py) passes negative_prompt; the
|
||||
# batch-mismatch we hit earlier was on V1 only.
|
||||
out = pipeline(
|
||||
prompt=_PHOTOMAKER_PROMPT,
|
||||
negative_prompt=_PHOTOMAKER_NEGATIVE,
|
||||
input_id_images=[id_image_pil],
|
||||
id_embeds=id_embeds,
|
||||
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)
|
||||
Reference in New Issue
Block a user