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

refactor(face-restore): rollback PhotoMaker, restore GFPGAN on the CLEANED image

Browse Source 
After 7 cascading upstream-compat fixes (insightface dep, peft dep, pm_version,
device, etc.), the PhotoMaker V1 cert sweep still hit a CFG batch-dim mismatch
inside the denoising loop. The upstream PhotoMaker `pipeline.py` is forked from
diffusers v0.29.1 and our env runs 0.38; SDXL prompt-encoder handling changed
significantly between those versions, so making PhotoMaker work end-to-end
needs a proper fork or a diffusers downgrade — both expensive. Not worth
shipping today.

Pivot: restore `face_restore.py` (GFPGAN) with a single-line fix that makes it
SynthID-safe by construction. The previous design ran GFPGAN.enhance on the
ORIGINAL watermarked image and was oracle-confirmed to re-add SynthID via the
weight-0.5 pixel blend. The fix is to run GFPGAN on the diffusion-CLEANED
image — whatever pixels GFPGAN derives from are already SynthID-free, so the
partial blend cannot transport the watermark. Identity fidelity is lower than
a true identity-as-embedding stack would deliver, but it ships and works.

Changes:
- `src/remove_ai_watermarks/face_restore.py` restored from pre-wipe state with
  one line changed: `restorer.enhance(cleaned_bgr, ...)` instead of
  `restorer.enhance(original_bgr, ...)`. `original_bgr` is kept as an unused
  positional argument for API stability.
- `src/remove_ai_watermarks/photomaker_restore.py` and its tests REMOVED. The
  research note (`docs/synthid-robust-identity-research.md`) keeps a "status
  notice" documenting why PhotoMaker is parked for now and what the path back
  in would look like.
- `pyproject.toml` `restore` extra restored (gfpgan/facexlib/basicsr +
  scipy<1.18 + numba<0.60 pins + the basicsr setuptools<69 build pin), plus
  `photomaker` extra (with its einops/insightface/peft pile) and the
  `[tool.hatch.metadata] allow-direct-references = true` block REMOVED.
- `InvisibleEngine._restore_faces_photomaker` removed; `_restore_faces`
  restored. The `--restore-faces` CLI flag and its plumbing through cmd_*
  signatures are unchanged.
- CLAUDE.md, README.md, docs/synthid.md, docs/controlnet-removal-pipeline-
  research.md updated to describe the shipped GFPGAN-on-cleaned design and to
  reference PhotoMaker only as the parked alternative.

ruff + strict pyright(src/) clean; 578 tests pass.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Victor Kuznetsov
2026-06-08 16:55:45 -07:00
parent d1b85ee6a8
commit 01fe98bf54
13 changed files with 1273 additions and 851 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_face_restore.py
+85
View File
@@ -0,0 +1,85 @@
"""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)
tests/test_photomaker_restore.py
-150
View File
@@ -1,150 +0,0 @@
"""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 TestV1OnlyCommercialSafetyGuard:
"""The module must lock to PhotoMaker-V1 (Apache + CLIP-only encoder).
V2 pulls InsightFace antelopev2/buffalo_l face packs which are non-commercial.
A maintainer touching ``_PHOTOMAKER_FILE`` for any reason must trip this guard.
"""
def test_repo_is_v1(self):
assert photomaker_restore._PHOTOMAKER_REPO == "TencentARC/PhotoMaker"
def test_weight_filename_is_v1(self):
assert photomaker_restore._PHOTOMAKER_FILE == "photomaker-v1.bin"
def test_module_source_does_not_call_face_analysis(self):
"""We may IMPORT `insightface` (transitive) but must never instantiate FaceAnalysis."""
import inspect
src = inspect.getsource(photomaker_restore)
assert "FaceAnalysis" not in src
assert "insightface.app" not in src
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