mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-06-10 12:53:56 +02:00
Victor Kuznetsov
01fe98bf54
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>
86 lines
3.5 KiB
Python
86 lines
3.5 KiB
Python
"""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)
|