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

Rewrite internal watermark pipeline and preserve behavior

Browse Source
This commit is contained in:
Victor Kuznetsov
2026-07-31 16:53:41 -07:00
parent 7c922e5133
commit a6c0c1c6f0
54 changed files with 2266 additions and 3771 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/code-provenance.md
+9
View File
@@ -0,0 +1,9 @@
# Code provenance
This page records notices required by source dependencies and licensed derivative work.
## Licensed derivative work
- The DWT-DCT implementation derives from ShieldMnt's
[`invisible-watermark`](https://github.com/ShieldMnt/invisible-watermark), licensed
under MIT. Its notice ships in `src/remove_ai_watermarks/licenses/invisible-watermark-MIT.txt`.
docs/controlnet-removal-pipeline-research.md
+1 -1
View File
@@ -211,7 +211,7 @@ regeneration so strokes exceed the VAE's ~8 px latent floor), but ctrlregen runs
at LOW res, the opposite. CtrlRegen's paper gives no resolution/tiling spec to contradict this.
**Sources.** the former internal
`src/remove_ai_watermarks/noai/ctrlregen/engine.py` (removed after this study);
`src/remove_ai_watermarks/_internal/ctrlregen/engine.py` (removed after this study);
resolution-omission
confirmed against https://arxiv.org/html/2410.05470v1
docs/index.md
+1
View File
@@ -19,6 +19,7 @@ to run the tool. Use the maintainer references only when changing the code.
| Page | Purpose |
| --- | --- |
| [Module internals](module-internals.md) | Current architecture, invariants, and regression guards by module. |
| [Code provenance](code-provenance.md) | Required notices for licensed derivative work. |
| [Verification plan](verification-plan.md) | Verification methods, completed measurements, and remaining validation gaps. |
| [Release and distribution](release-and-distribution.md) | PyPI, Homebrew, Hugging Face Space, and release workflow. |
| [Watermarking landscape](watermarking-landscape.md) | Vendor signals and detection approaches. |
docs/known-limitations.md
+1 -1
View File
@@ -133,7 +133,7 @@ or a different random seed may change the verifier result.
The base Qwen and `qwen-zimage` profiles have profile specific strength
behavior. Consult `remove-ai-watermarks invisible --help` and the source of
[`watermark_profiles.py`](../src/remove_ai_watermarks/noai/watermark_profiles.py)
[`watermark_profiles.py`](../src/remove_ai_watermarks/_internal/watermark_profiles.py)
for the current resolver.
### Pipelines have different quality tradeoffs
docs/module-internals.md
+24 -19
View File
@@ -119,7 +119,7 @@ reported per file without retrying the same multi-GB initialization.
Native MP4/MOV TC260 labels follow TC260-PG-20257A:
`moov.udta.meta.keys` maps an `AIGC` key to a raw JSON value in `ilst`.
[`noai/isobmff.py`](../src/remove_ai_watermarks/noai/isobmff.py) walks those
[`_internal/isobmff.py`](../src/remove_ai_watermarks/_internal/isobmff.py) walks those
nested boxes by seeking, so detection reaches a tail `moov` without reading the
preceding `mdat`. The MP4/MOV/M4V/M4A removal path first validates the top-level
box walk, then copies the source to a sibling temporary file in bounded chunks.
@@ -130,14 +130,14 @@ validated JSON value with same-length spaces. This preserves every box size,
Publication is atomic, and a malformed top-level walk is copied unchanged. A
generic `AIGC` key whose value has no TC260 field is ignored.
[`noai/ebml.py`](../src/remove_ai_watermarks/noai/ebml.py) provides the
[`_internal/ebml.py`](../src/remove_ai_watermarks/_internal/ebml.py) provides the
corresponding bounded Matroska/WebM reader. It seeks over clusters and accepts
only a `Segment.Tags.Tag.SimpleTag` pairing `TagName=AIGC` with a JSON
`TagString` carrying a TC260 field. The existing ffmpeg stream-copy path removes
those container tags without transcoding the encoded streams.
[`noai/riff.py`](../src/remove_ai_watermarks/noai/riff.py) and
[`noai/flv.py`](../src/remove_ai_watermarks/noai/flv.py) implement the remaining
[`_internal/riff.py`](../src/remove_ai_watermarks/_internal/riff.py) and
[`_internal/flv.py`](../src/remove_ai_watermarks/_internal/flv.py) implement the remaining
normative TC260 video placements. The RIFF walker reads only AVI
`LIST/INFO/AIGC` children. The FLV walker skips media tags and parses the AMF0
`script.onMetaData.AIGC` string. Both require a recognized TC260 JSON field and
@@ -284,12 +284,12 @@ Regression coverage:
### C2PA
[`noai/c2pa.py`](../src/remove_ai_watermarks/noai/c2pa.py) reads C2PA with the
[`_internal/c2pa.py`](../src/remove_ai_watermarks/_internal/c2pa.py) reads C2PA with the
official `c2pa-python` reader first. Its byte-level PNG parser remains a fallback
for partial and synthetic fixtures that the official reader rejects.
Vendor attribution comes from the registry in
[`noai/constants.py`](../src/remove_ai_watermarks/noai/constants.py). Derived
[`_internal/constants.py`](../src/remove_ai_watermarks/_internal/constants.py). Derived
issuer and platform maps should not be maintained separately.
### Metadata scanning and stripping
@@ -303,7 +303,7 @@ Key contracts:
- JPEG stripping walks metadata segments and preserves the entropy-coded image
scan.
- ISOBMFF containers use
[`noai/isobmff.py`](../src/remove_ai_watermarks/noai/isobmff.py).
[`_internal/isobmff.py`](../src/remove_ai_watermarks/_internal/isobmff.py).
- Native MP4/MOV TC260 `AIGC` entries are read from
`moov.udta.meta.keys/ilst` and blanked without changing box sizes.
- Native MKV/WebM TC260 `AIGC` entries are read from
@@ -327,7 +327,7 @@ test proves that it no longer appears in the output.
Regression coverage:
- [`test_metadata.py`](../tests/test_metadata.py)
- [`test_noai.py`](../tests/test_noai.py)
- [`test_metadata_internals.py`](../tests/test_metadata_internals.py)
- [`test_security_clamp.py`](../tests/test_security_clamp.py)
### Provenance report
@@ -474,7 +474,7 @@ Regression coverage:
### Profiles and strength
[`noai/watermark_profiles.py`](../src/remove_ai_watermarks/noai/watermark_profiles.py)
[`_internal/watermark_profiles.py`](../src/remove_ai_watermarks/_internal/watermark_profiles.py)
is the source of truth for:
- profile aliases;
@@ -494,12 +494,19 @@ router.
[`invisible_engine.py`](../src/remove_ai_watermarks/invisible_engine.py) handles
image sizing, optional pre-upscaling, postprocessing, and the public engine
interface. It delegates model execution to
[`noai/watermark_remover.py`](../src/remove_ai_watermarks/noai/watermark_remover.py).
[`_internal/watermark_remover.py`](../src/remove_ai_watermarks/_internal/watermark_remover.py).
The Python engine and CLI do not have identical defaults for every optional
postprocessing argument. Integrations that require reproducibility should pass
the relevant values explicitly.
The standard Qwen and ControlNet prompts are calibrated model inputs, and the
ControlNet edge map uses fixed Canny thresholds of 100 and 200. Treat those
values as behavioral compatibility contracts: a refactor must preserve them,
and any deliberate change requires image-quality evaluation rather than only a
unit-test pass. Exact prompt and edge-map regression guards live in
`test_platform.py` and `test_invisible_engine.py`.
Regression coverage:
- [`test_watermark_profiles.py`](../tests/test_watermark_profiles.py)
@@ -519,7 +526,7 @@ Regression coverage:
### Qwen plus Z-Image
[`noai/qwen_zimage_pipeline.py`](../src/remove_ai_watermarks/noai/qwen_zimage_pipeline.py)
[`_internal/qwen_zimage_pipeline.py`](../src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py)
implements the fixed CUDA-only two-stage profile:
1. Qwen Image with Canny conditioning regenerates the frame.
@@ -530,13 +537,11 @@ The profile rejects a custom model identifier. Its global and face model stack
is fixed by the implementation. When tiling is enabled, only the global stage
is tiled; the face stage runs once after the tiles are blended.
The resolution and largest-face adaptive formulas remain exact ports of the
reference workflow. The face stage applies half the reference result because
this port uses a different sampler and composites regenerated SAM pixels rather
than using the reference latent inpaint mask and noise feather. Paired face
evaluations favored this scale on identity, perceptual distance, and full-image
similarity, and the exact OpenAI and Gemini candidates both passed their
matching provider oracle. The global stage stays unchanged.
The maintained implementation preserves the previously oracle-tested strength,
conditioning, crop, and sampler parameters as compatibility contracts. Its Python
orchestration, YuNet integration, SAM selection, masks, sizing helpers, and pixel
compositing are implemented for this runtime. Changing a calibrated model input
requires the same provider-oracle and identity evaluation as a model change.
Regression coverage:
@@ -545,7 +550,7 @@ Regression coverage:
### Tiling
[`noai/tiling.py`](../src/remove_ai_watermarks/noai/tiling.py) contains pure
[`_internal/tiling.py`](../src/remove_ai_watermarks/_internal/tiling.py) contains pure
tile planning, feather weights, tile orchestration, and region compositing.
Tiling engages only when requested and the long side exceeds the tile size.
docs/qwen-improvement-research.md
+8 -8
View File
@@ -79,13 +79,13 @@ Measured on `gemini_3` (18 faces) at the Gemini scrub floor 0.25 vs base-Qwen 0.
a Qwen face fix. The next distinct architecture was Z-Image-Turbo on original masked face
crops, not another Qwen geometry conditioner.
**Implementation follow-up (2026-07-24):** that distinct architecture now exists as the
manual `qwen-zimage` profile. It ports the upstream Synthid-Bypass v2 graph: Qwen-Image-2512
Lightning + DiffSynth Canny for the full frame, then SAM-masked Z-Image Turbo regeneration
from original face crops. The upstream result supplied by the user was Gemini-oracle negative.
The active upstream face path is YOLO + SAM, not the unconnected MediaPipe node. The port
matches its center-point + box prompts, IoU-0.93 proposal selection, detector-box intersection,
crop factor, and paste feather; YuNet is the intentional detector substitution.
**Implementation follow-up (2026-07-24, revised 2026-07-31):** an early experimental
`qwen-zimage` prototype reproduced a broad two-stage shape demonstrated by a public
experiment: structure-guided Qwen regeneration followed by masked Z-Image face
refinement. The maintained profile was subsequently
rewritten with project-owned prompts, adaptive strength and sizing policies, YuNet face
detection, SAM selection, masks, and compositing. It does not include the upstream workflow
JSON or source code. The upstream result supplied by the user was Gemini-oracle negative.
The first exact-path Modal run completed without SAM fallback or face seams. On one crowded
18-face `gemini_3` fixture, ArcFace identity improved materially over controlnet
@@ -108,7 +108,7 @@ checked all six current outputs in the provider-separated
oracles and confirmed that none retained SynthID or the provider generation signal. The
checked bytes used the complete `visible -> qwen-zimage -> metadata` route, the calibrated
YuNet 0.5 gate, and the shipped prompt-cache/model-residency optimizations. This supersedes
the earlier first-port batch check as the release-candidate result, but it is not a
the earlier prototype batch check as the release-candidate result, but it is not a
certification across seeds, resolutions, and content classes. YuNet's threshold was
calibrated independently from upstream YOLO: 0.5 retained the visible faces in the
comparison fixtures while removing the false and duplicate boxes admitted by the copied
docs/research-doubao-distillation.md
+1 -1
View File
@@ -26,7 +26,7 @@ color still leaves a persistent ghost outline.
Diagnosed why, empirically (cached stacks, `/tmp/doubao_distill`): (1) the mark is a clean white overlay with **no dark halo** -- over glyph pixels ~54% are brighter than the clean bg, only ~4% darker -- so the white-logo model `I=(1-α)O+α·255` is correct; (2) but content backgrounds are almost never dark *under* the mark (median darkest available bg over glyph pixels = **58/255**; only ~13% of mark pixels are ever observed on a bg < 40), so on bright backgrounds the equation is ill-conditioned and `α` is unidentifiable; (3) LaMa's `O` is a plausible **hallucination**, not the true pre-mark background, which compounds the error, and per-pixel regression on ~15 obs overfits into color noise.
**Why Gemini's engine is clean (verified in GeminiWatermarkTool `src/core/watermark_engine.cpp`): its alpha map is the watermark stamped on a PURE-BLACK background**, where `watermarked = α·255 + (1-α)·0 = α·255`, so `alpha = capture/255` exactly -- no estimation. (`gemini_bg_*.png` is literally the sparkle in gray on black.) So the real Doubao unlock is the same controlled capture, **not more content images**. The retained black and gray outputs live in `data/calibration/doubao/`; local solid-color seeds are regenerable and are not committed.
**Why Gemini's engine is clean: its alpha map is the watermark stamped on a PURE-BLACK background**, where `watermarked = α·255 + (1-α)·0 = α·255`, so `alpha = capture/255` exactly -- no estimation. (`gemini_bg_*.png` is literally the sparkle in gray on black.) So the real Doubao unlock is the same controlled capture, **not more content images**. The retained black and gray outputs live in `data/calibration/doubao/`; local solid-color seeds are regenerable and are not committed.
**Until black captures arrive, the shipped direction is precise canonical glyph mask + inpaint (cv2 default, lama optional), NOT reverse-alpha.**
docs/synthid.md
+1 -1
View File
@@ -623,7 +623,7 @@ it; (3) **historical engineering conclusion:** this dated run argued for a
higher ControlNet strength than the then-current default. That proposal was
later superseded. The current resolver intentionally shares the 0.10/0.15
ladder between SDXL and ControlNet and uses a separate Qwen ladder; see
`noai/watermark_profiles.py`.
`_internal/watermark_profiles.py`.
Source images are private (faces / product shots), not committed; reproduce on any
photoreal + flat-graphic gpt-image pair, varying the seed, and re-checking the
oracle.
pyproject.toml
+3 -3
View File
@@ -51,7 +51,7 @@ dependencies = [
"python-dotenv>=1.0.0",
# Official C2PA reader (Content Authenticity Initiative, MIT/Apache-2.0). The
# primary, spec-tracking manifest parser for the identify/metadata path; the
# hand-rolled caBX/CBOR scanner in noai/c2pa.py is kept only as a fallback for
# hand-rolled caBX/CBOR scanner in _internal/c2pa.py is kept only as a fallback for
# synthetic/partial blobs the validator rejects. Binary wheel (Rust), but the
# import is light (no torch/numpy) so it fits the dependency-light identify
# host. Prebuilt wheels cover the full CI matrix (linux/macos/windows).
@@ -235,8 +235,8 @@ ignore = [
[tool.ruff.lint.per-file-ignores]
"scripts/*.py" = ["G004", "S108", "S310", "T20"]
"tests/*.py" = ["ANN", "S101", "S105", "S106", "S108"]
"src/remove_ai_watermarks/noai/watermark_remover.py" = ["S603", "S606", "S607", "T201"] # subprocess calls for auto-install/CUDA fix
"src/remove_ai_watermarks/noai/c2pa.py" = ["S110"] # try-except-pass for corrupt file handling
"src/remove_ai_watermarks/_internal/watermark_remover.py" = ["S603", "S606", "S607"] # nvidia-smi capability probe
"src/remove_ai_watermarks/_internal/c2pa.py" = ["S110"] # try-except-pass for corrupt file handling
[tool.ruff.format]
quote-style = "double"
scripts/metadata_removal_audit.py
+1 -1
View File
@@ -32,7 +32,7 @@ from pathlib import Path
import click
import numpy as np
from remove_ai_watermarks.noai.constants import SUPPORTED_FORMATS
from remove_ai_watermarks._internal.constants import SUPPORTED_FORMATS
log = logging.getLogger(__name__)
scripts/smoke_matrix.py
+1 -1
View File
@@ -590,8 +590,8 @@ def _diffusion_rows(r: Runner, tmp: Path, doubao: Path) -> None:
try:
import numpy as np
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
from remove_ai_watermarks.image_io import imread
from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover
src = imread(str(mj))
h, w = src.shape[:2]
scripts/synthid_corpus.py
+1 -1
View File
@@ -29,7 +29,7 @@ import click
from _plain_console import Console, Table
from PIL import Image
from remove_ai_watermarks.noai.c2pa import extract_c2pa_info
from remove_ai_watermarks._internal.c2pa import extract_c2pa_info
log = logging.getLogger(__name__)
console = Console()
scripts/visible_positives.py
+1 -1
View File
@@ -38,7 +38,7 @@ sys.path.insert(0, str(Path(__file__).parent.parent))
# The package's own format set. An inlined copy here silently skipped .heif, which
# CLAUDE.md documents as supported.
from remove_ai_watermarks.noai.constants import SUPPORTED_FORMATS as _EXTS
from remove_ai_watermarks._internal.constants import SUPPORTED_FORMATS as _EXTS
REPO = Path(__file__).resolve().parents[1]
CORPUS = REPO / ".local-eval" / "originals"
src/remove_ai_watermarks/noai/__init__.py → src/remove_ai_watermarks/_internal/__init__.py
+5 -7
View File
@@ -1,10 +1,8 @@
"""Vendored noai-watermark code for invisible watermark removal.
Original: https://github.com/mertizci/noai-watermark (MIT License)
"""Compatibility namespace for metadata and regeneration helpers.
The public API (``WatermarkRemover`` / ``remove_watermark`` / ``remove_ai_metadata``)
is exposed **lazily** via PEP 562 ``__getattr__``: importing a light submodule
(e.g. ``noai.c2pa`` / ``noai.constants`` from ``identify``) must NOT eagerly pull
(e.g. ``_internal.c2pa`` / ``_internal.constants`` from ``identify``) must NOT eagerly pull
``watermark_remover``, which imports torch + diffusers at module top. Keeping this
lazy is what lets ``import remove_ai_watermarks.identify`` stay cheap (~36 MB, no
torch) even in a full install where the ``diffusion`` extra is present --
@@ -17,8 +15,8 @@ from __future__ import annotations
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover, remove_watermark
from remove_ai_watermarks.metadata import remove_ai_metadata
from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover, remove_watermark
__all__ = ["WatermarkRemover", "remove_ai_metadata", "remove_watermark"]
@@ -27,12 +25,12 @@ def __getattr__(name: str) -> object:
"""Resolve the public API on first access (PEP 562), not at package import."""
if name == "remove_ai_metadata":
# Re-export the single, robust stripper (byte-level, lossless-for-JPEG, all
# containers); the old noai.cleaner implementation is retired.
# containers); the old legacy metadata helper implementation is retired.
from remove_ai_watermarks.metadata import remove_ai_metadata
return remove_ai_metadata
if name in ("WatermarkRemover", "remove_watermark"):
from remove_ai_watermarks.noai import watermark_remover
from remove_ai_watermarks._internal import watermark_remover
return getattr(watermark_remover, name)
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
src/remove_ai_watermarks/_internal/c2pa.py
+363
View File
@@ -0,0 +1,363 @@
"""C2PA inspection through the official reader with a bounded PNG fallback."""
from __future__ import annotations
import contextlib
import functools
import json
import logging
import re
import struct
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any, cast
from remove_ai_watermarks._internal.constants import (
C2PA_ACTIONS,
C2PA_AI_TOOLS,
C2PA_CHUNK_TYPE,
C2PA_ISSUERS,
C2PA_SIGNATURES,
C2PA_SOFT_BINDINGS,
PNG_SIGNATURE,
SYNTHID_C2PA_ISSUERS,
)
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from typing import BinaryIO
_C2paReader: Any = None
with contextlib.suppress(Exception):
from c2pa import Reader as _C2paReader # pyright: ignore[reportMissingTypeStubs]
_C2PA_READER_AVAILABLE = _C2paReader is not None
_PNG_HEADER = struct.Struct(">I4s")
@dataclass(frozen=True)
class _PngChunk:
payload: bytes
serialized: bytes
def reader_available() -> bool:
"""Return whether the official C2PA reader loaded successfully."""
return _C2PA_READER_AVAILABLE
def _manifest_json_uncached(path: str) -> str | None:
try:
reader = _C2paReader.try_create(path)
except Exception as error:
logger.debug("C2PA reader rejected %s: %s", path, error)
return None
if reader is None:
return None
try:
with reader:
return cast("str", reader.json())
except Exception as error:
logger.debug("C2PA reader could not serialize %s: %s", path, error)
return None
@functools.lru_cache(maxsize=8)
def _manifest_json_cached(path: str, _mtime_ns: int) -> str | None:
return _manifest_json_uncached(path)
def read_manifest_store_json(image_path: Path) -> str | None:
"""Read the complete manifest-store JSON, caching it until the file changes."""
if not reader_available():
return None
path = str(image_path)
try:
return _manifest_json_cached(path, image_path.stat().st_mtime_ns)
except OSError:
return _manifest_json_uncached(path)
def _find_c2pa_chunk(path: Path) -> _PngChunk | None:
"""Return the first recognizable C2PA chunk without loading the whole PNG."""
try:
stream = path.open("rb")
except OSError:
return None
with stream:
if stream.read(len(PNG_SIGNATURE)) != PNG_SIGNATURE:
return None
file_size = stream.seek(0, 2)
stream.seek(len(PNG_SIGNATURE))
while True:
header = stream.read(_PNG_HEADER.size)
if len(header) != _PNG_HEADER.size:
return None
length, kind = _PNG_HEADER.unpack(header)
if length + 4 > file_size - stream.tell():
return None
if kind == C2PA_CHUNK_TYPE:
payload = stream.read(length)
crc = stream.read(4)
if _looks_like_c2pa(payload):
return _PngChunk(payload, header + payload + crc)
else:
stream.seek(length + 4, 1)
if kind == b"IEND":
return None
def _is_well_formed_png(path: Path) -> bool:
"""Validate PNG chunk bounds with seeks rather than payload allocations."""
try:
stream = path.open("rb")
except OSError:
return False
with stream:
if stream.read(len(PNG_SIGNATURE)) != PNG_SIGNATURE:
return False
file_size = stream.seek(0, 2)
stream.seek(len(PNG_SIGNATURE))
while True:
header = stream.read(_PNG_HEADER.size)
if len(header) != _PNG_HEADER.size:
return False
length, kind = _PNG_HEADER.unpack(header)
if length + 4 > file_size - stream.tell():
return False
stream.seek(length + 4, 1)
if kind == b"IEND":
return True
def _copy_bytes(source: BinaryIO, target: BinaryIO, byte_count: int) -> None:
"""Copy exactly one bounded chunk without allocating its complete payload."""
remaining = byte_count
while remaining:
block = source.read(min(remaining, 1024 * 1024))
if not block:
raise OSError("PNG changed while it was being copied")
target.write(block)
remaining -= len(block)
def _looks_like_c2pa(payload: bytes) -> bool:
lowered = payload.lower()
return any(signature in payload for signature in C2PA_SIGNATURES) or b"c2pa" in lowered or b"jumb" in lowered
def extract_c2pa_chunk(image_path: Path) -> bytes | None:
"""Return the first complete C2PA PNG chunk, including header and CRC."""
if image_path.suffix.casefold() != ".png":
return None
chunk = _find_c2pa_chunk(image_path)
return None if chunk is None else chunk.serialized
def has_c2pa_metadata(image_path: Path) -> bool:
"""Return whether a validly bounded PNG contains a recognizable C2PA chunk."""
return extract_c2pa_chunk(Path(image_path)) is not None
def _active_manifest(store: dict[str, Any]) -> dict[str, Any]:
manifests = store.get("manifests")
if not isinstance(manifests, dict):
return {}
typed_manifests = cast("dict[object, object]", manifests)
active = typed_manifests.get(store.get("active_manifest"))
return cast("dict[str, Any]", active) if isinstance(active, dict) else {}
def _claim_generator_from_store(store: dict[str, Any]) -> str | None:
active = _active_manifest(store)
direct = active.get("claim_generator")
if isinstance(direct, str) and direct.isprintable() and direct:
return direct
candidates = active.get("claim_generator_info")
if isinstance(candidates, list) and candidates and isinstance(candidates[0], dict):
candidate = cast("dict[object, object]", candidates[0])
name = candidate.get("name")
if isinstance(name, str) and name.isprintable() and name:
return name
return None
def synthid_verdict(vendors: str) -> str:
"""Describe why metadata implies a likely pixel-level SynthID watermark."""
return f"likely present ({vendors} embeds SynthID with C2PA)"
def _names_present(buffer: bytes, registry: dict[bytes, str]) -> list[str]:
return sorted({label for token, label in registry.items() if token in buffer})
def synthid_vendors_in(buffer: bytes) -> list[str]:
"""List matching C2PA issuers known to pair their manifests with SynthID."""
registry = {token: label for token, label in C2PA_ISSUERS.items() if token in SYNTHID_C2PA_ISSUERS}
return _names_present(buffer, registry)
def soft_binding_vendors_in(buffer: bytes) -> list[str]:
"""List the soft-binding algorithms named in manifest bytes."""
return _names_present(buffer, C2PA_SOFT_BINDINGS)
def _ordered_matches(buffer: bytes, registry: dict[bytes, str]) -> list[str]:
return list(dict.fromkeys(label for token, label in registry.items() if token in buffer))
def _populate_registry_fields(buffer: bytes, info: dict[str, Any]) -> bool:
issuers = _ordered_matches(buffer, C2PA_ISSUERS)
tools = _ordered_matches(buffer, C2PA_AI_TOOLS)
actions = _ordered_matches(buffer, C2PA_ACTIONS)
if issuers:
info["issuer"] = ", ".join(issuers)
if tools:
info["ai_tool"] = ", ".join(tools)
if actions:
info["actions"] = ", ".join(actions)
ai_source = False
if b"trainedAlgorithmicMedia" in buffer:
info.update(source_type="trainedAlgorithmicMedia (AI-generated)", ai_source_kind="generated")
ai_source = True
elif b"compositeWithTrainedAlgorithmicMedia" in buffer:
info.update(source_type="compositeWithTrainedAlgorithmicMedia (AI-enhanced)", ai_source_kind="enhanced")
ai_source = True
elif b"algorithmicMedia" in buffer:
info["source_type"] = "algorithmicMedia"
synthid = synthid_vendors_in(buffer)
if ai_source and synthid:
info["synthid_vendors"] = synthid
info["synthid_watermark"] = synthid_verdict(", ".join(synthid))
soft_bindings = soft_binding_vendors_in(buffer)
if soft_bindings:
info["soft_binding_vendors"] = soft_bindings
info["soft_binding"] = ", ".join(soft_bindings)
return ai_source
def _base_info(byte_count: int, *, fallback: bool = False) -> dict[str, Any]:
container = "C2PA manifest" if fallback else "C2PA manifest store"
return {
"has_c2pa": True,
"type": "C2PA (Coalition for Content Provenance and Authenticity)",
"c2pa_manifest": f"{container} ({byte_count} bytes)",
}
def _info_from_store(store: dict[str, Any], encoded: bytes) -> dict[str, Any]:
info = _base_info(len(encoded))
_populate_registry_fields(encoded, info)
generator = _claim_generator_from_store(store)
if generator is not None:
info["claim_generator"] = generator
signature_value = _active_manifest(store).get("signature_info")
if isinstance(signature_value, dict):
signature = cast("dict[object, object]", signature_value)
timestamp = signature.get("time")
if timestamp:
info["timestamp"] = str(timestamp)
return info
def c2pa_info_from_manifest_store(store: str | dict[str, Any]) -> dict[str, Any]:
"""Normalize a manifest store supplied as JSON text or a decoded object."""
try:
raw_decoded: object = store if isinstance(store, dict) else json.loads(store)
decoded = cast("dict[str, Any]", raw_decoded) if isinstance(raw_decoded, dict) else None
if not isinstance(decoded, dict) or not decoded or decoded.get("error"):
return {}
encoded = json.dumps(decoded, ensure_ascii=False).encode() if isinstance(store, dict) else store.encode()
except (TypeError, ValueError):
return {}
return _info_from_store(decoded, encoded)
def cbor_text_after(payload: bytes, key: bytes) -> str | None:
"""Decode a definite-length CBOR text value immediately following ``key``."""
key_end = payload.find(key)
if key_end < 0:
return None
cursor = key_end + len(key)
if cursor >= len(payload):
return None
initial = payload[cursor]
if 0x60 <= initial <= 0x77:
length, cursor = initial & 0x1F, cursor + 1
elif initial == 0x78 and cursor + 1 < len(payload):
length, cursor = payload[cursor + 1], cursor + 2
elif initial == 0x79 and cursor + 2 < len(payload):
length = int.from_bytes(payload[cursor + 1 : cursor + 3], "big")
cursor += 3
else:
return None
raw = payload[cursor : cursor + length]
if len(raw) != length:
return None
try:
return raw.decode()
except UnicodeDecodeError:
return raw.decode("latin1", errors="replace")
def _parse_c2pa_chunk(payload: bytes, info: dict[str, Any]) -> None:
info.update(_base_info(len(payload), fallback=True))
_populate_registry_fields(payload, info)
for key, output_key in ((b"name", "claim_generator"), (b"specVersion", "c2pa_spec")):
value = cbor_text_after(payload, key)
if value and value.isprintable():
info[output_key] = value
timestamps = [item.decode() for item in re.findall(rb"\d{14}Z", payload)]
if timestamps:
info["timestamp"] = timestamps[0]
if len(timestamps) > 1:
info["timestamps"] = timestamps[:3]
def _extract_c2pa_info_png(image_path: Path) -> dict[str, Any]:
if image_path.suffix.casefold() != ".png":
return {}
chunk = _find_c2pa_chunk(image_path)
if chunk is None:
return {}
info: dict[str, Any] = {}
_parse_c2pa_chunk(chunk.payload, info)
return info
def extract_c2pa_info(image_path: Path) -> dict[str, Any]:
"""Return normalized C2PA evidence from the official reader or PNG fallback."""
store = read_manifest_store_json(Path(image_path))
if store is not None:
return c2pa_info_from_manifest_store(store)
return _extract_c2pa_info_png(Path(image_path))
def inject_c2pa_chunk(target_path: Path, output_path: Path, c2pa_chunk: bytes) -> None:
"""Replace any C2PA chunks in a PNG and insert ``c2pa_chunk`` before IDAT."""
if target_path.suffix.casefold() != ".png" or output_path.suffix.casefold() != ".png":
raise ValueError("C2PA chunk injection is only supported for PNG files")
if not _is_well_formed_png(target_path):
raise ValueError("Target is not a well-formed PNG file")
output_path.parent.mkdir(parents=True, exist_ok=True)
with target_path.open("rb") as source, output_path.open("wb") as target:
target.write(source.read(len(PNG_SIGNATURE)))
inserted = False
while True:
header = source.read(_PNG_HEADER.size)
length, kind = _PNG_HEADER.unpack(header)
if kind == b"IDAT" and not inserted:
target.write(c2pa_chunk)
inserted = True
if kind == C2PA_CHUNK_TYPE:
source.seek(length + 4, 1)
else:
target.write(header)
_copy_bytes(source, target, length + 4)
if kind == b"IEND":
break
src/remove_ai_watermarks/_internal/constants.py
+152
View File
@@ -0,0 +1,152 @@
"""Registries shared by metadata extraction and provenance classification."""
from __future__ import annotations
from dataclasses import dataclass
def _tokens(value: str) -> tuple[str, ...]:
return tuple(value.split("|"))
SUPPORTED_FORMATS = frozenset(_tokens(".png|.jpg|.jpeg|.webp|.heic|.heif|.avif"))
AI_METADATA_KEYS = _tokens(
"parameters|postprocessing|extras|workflow|prompt|Dream|SD:mode|StableDiffusionVersion|"
"generation_time|Model|Model hash|Seed"
)
PNG_METADATA_KEYS = _tokens(
"Author|Title|Description|Copyright|Creation Time|Software|Disclaimer|Warning|Source|Comment"
)
AI_KEYWORDS = _tokens(
"prompt|negative_prompt|sampler|cfg_scale|lora|diffusion|comfy|midjourney|dall-e|dalle|imagen|firefly|c2pa|chatgpt|gpt-4|sora|openai|truepic|stable_diffusion|invokeai"
)
PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
C2PA_CHUNK_TYPE = b"caBX"
C2PA_SIGNATURES = tuple(
token.encode() for token in _tokens("c2pa|C2PA|jumb|jumd|JUMBF|jumbf|cbor|contentcreds|digid|assertions|manifest")
)
@dataclass(frozen=True, slots=True)
class C2paAiVendor:
"""One issuer signature and its normalized product attribution."""
issuer: bytes
org: str
platform: str | None
needle: str | None
synthid: bool = False
asserts_ai: bool = False
def _vendor(
issuer: bytes | str,
org: str,
platform: str | None,
needle: str | None,
*,
synthid: bool = False,
asserts_ai: bool = False,
) -> C2paAiVendor:
token = issuer.encode() if isinstance(issuer, str) else issuer
return C2paAiVendor(token, org, platform, needle, synthid, asserts_ai)
# Order is product priority when a manifest mentions more than one organization.
C2PA_AI_VENDORS: tuple[C2paAiVendor, ...] = (
_vendor(b"Microsoft", "Microsoft", "Microsoft (Bing Image Creator / Designer)", "Microsoft"),
_vendor(b"Adobe", "Adobe", "Adobe Firefly", "Adobe"),
_vendor(b"OpenAI", "OpenAI", "OpenAI (ChatGPT / gpt-image / DALL-E / Sora)", "OpenAI", synthid=True),
_vendor(b"Google", "Google LLC", "Google (Gemini / Imagen)", "Google", synthid=True),
_vendor(b"Stability AI", "Stability AI", "Stability AI (Stable Image / DreamStudio)", "Stability AI"),
_vendor(b"Black Forest Labs", "Black Forest Labs", "Black Forest Labs (FLUX)", "Black Forest Labs"),
_vendor(b"volcengine", "ByteDance (Volcano Engine)", "ByteDance (Doubao / Jimeng / Volcano Engine)", "ByteDance"),
_vendor(
"北京火山引擎科技有限公司",
"ByteDance (Volcano Engine)",
"ByteDance (Doubao / Jimeng / Volcano Engine)",
"ByteDance",
),
_vendor(b"Byteplus", "BytePlus (ByteDance)", "ByteDance (Doubao / Jimeng / Volcano Engine)", "ByteDance"),
_vendor(
b"Dreamina",
"ByteDance (Dreamina)",
"ByteDance (Doubao / Jimeng / Volcano Engine)",
"ByteDance",
asserts_ai=True,
),
_vendor(b"Canva", "Canva", "Canva (Magic Media)", "Canva"),
_vendor(b"Eleven Labs", "ElevenLabs", "ElevenLabs", "ElevenLabs"),
_vendor(b"fal-ai", "fal.ai", "fal.ai", "fal.ai", asserts_ai=True),
_vendor(b"Bria", "Bria Artificial Intelligence", "Bria AI", "Bria", asserts_ai=True),
_vendor(b"Truepic", "Truepic", None, None),
)
C2PA_ISSUERS = {vendor.issuer: vendor.org for vendor in C2PA_AI_VENDORS}
C2PA_IDENTITY_AI_ORGS = frozenset(vendor.org for vendor in C2PA_AI_VENDORS if vendor.asserts_ai)
SYNTHID_C2PA_ISSUERS = frozenset(vendor.issuer for vendor in C2PA_AI_VENDORS if vendor.synthid)
C2PA_AI_TOOLS = {
token.encode(): label
for token, label in (
("GPT-4o", "GPT-4o"),
("ChatGPT", "ChatGPT"),
("Sora", "Sora"),
("DALL-E", "DALL-E"),
("DALL", "DALL-E"),
("Imagen", "Imagen"),
("Firefly", "Firefly"),
)
}
C2PA_SOFT_BINDINGS = {
b"com.adobe.trustmark": "Adobe TrustMark",
b"com.adobe.icn": "Adobe (content fingerprint)",
b"com.digimarc": "Digimarc",
b"com.imatag.lamark": "Imatag (Lamark)",
b"ai.steg": "Steg.AI",
b"com.microsoft.invismark": "Microsoft InvisMark",
b"com.microsoft.wavmark": "Microsoft WavMark",
b"com.verimatrix": "Verimatrix",
b"com.nagra.nexguard": "NAGRA NexGuard",
b"com.aiwatermark": "AIWatermark (Meta PixelSeal)",
b"ai.trufo": "Trufo",
b"app.overlai": "Overlai",
b"com.markany": "MarkAny",
b"com.mentaport": "Mentaport",
b"es.lumatrace": "LumaTrace",
b"ai.verda": "VerdaAI",
b"ai.contentlens": "ContentLens",
b"io.iscc": "ISCC (content code)",
}
AI_GENERATOR_TOKENS = frozenset(
{
"firefly",
"dall-e",
"dalle",
"midjourney",
"stable diffusion",
"stable-diffusion",
"stablediffusion",
"comfyui",
"automatic1111",
"invokeai",
"imagen",
"gpt-image",
"nightcafe",
"ideogram",
"leonardo",
"flux",
"dreamstudio",
"novelai",
"reve.com",
"aphrodite ai",
"apple photos clean up",
"fal-ai",
}
)
_C2PA_ACTION_NAMES = _tokens("created|converted|edited|filtered|cropped|resized|opened|placed")
C2PA_ACTIONS = {f"c2pa.{action}".encode(): action for action in _C2PA_ACTION_NAMES}
src/remove_ai_watermarks/noai/ebml.py → src/remove_ai_watermarks/_internal/ebml.py
View File
src/remove_ai_watermarks/_internal/extractor.py
+98
View File
@@ -0,0 +1,98 @@
"""Read image metadata without changing the source container."""
from __future__ import annotations
from typing import TYPE_CHECKING, Any, cast
import piexif
from PIL import Image
from remove_ai_watermarks._internal.c2pa import extract_c2pa_chunk, extract_c2pa_info, has_c2pa_metadata
from remove_ai_watermarks._internal.constants import AI_KEYWORDS, AI_METADATA_KEYS
if TYPE_CHECKING:
from pathlib import Path
_EXIF_KEY = "exif"
_AI_KEYS_CASEFOLD = frozenset(key.casefold() for key in AI_METADATA_KEYS)
def _read_pillow_info(source_path: Path) -> dict[str, Any]:
with Image.open(source_path) as image:
return {key: value for key, value in image.info.items() if isinstance(key, str)}
def _decode_exif(raw: object) -> tuple[str, object]:
if not isinstance(raw, bytes):
return "exif_raw", raw
try:
return _EXIF_KEY, piexif.load(raw)
except Exception:
return "exif_raw", raw
def _is_ai_field(key: str) -> bool:
folded = key.casefold()
return folded in _AI_KEYS_CASEFOLD or any(token in folded for token in AI_KEYWORDS)
def _attach_c2pa(source_path: Path, metadata: dict[str, Any]) -> None:
if not has_c2pa_metadata(source_path):
return
metadata["c2pa"] = extract_c2pa_info(source_path)
payload = extract_c2pa_chunk(source_path)
if payload is not None:
metadata["c2pa_chunk"] = payload
def extract_metadata(source_path: Path) -> dict[str, Any]:
"""Return every Pillow-visible field plus decoded EXIF and C2PA data."""
raw_info = _read_pillow_info(source_path)
metadata = dict(raw_info)
if _EXIF_KEY in raw_info:
metadata.pop(_EXIF_KEY, None)
decoded_key, decoded_value = _decode_exif(raw_info[_EXIF_KEY])
metadata[decoded_key] = decoded_value
_attach_c2pa(source_path, metadata)
return metadata
def extract_ai_metadata(source_path: Path) -> dict[str, Any]:
"""Return only metadata keys recognized as AI provenance or generation data."""
metadata = {key: value for key, value in _read_pillow_info(source_path).items() if _is_ai_field(key)}
_attach_c2pa(source_path, metadata)
return metadata
def has_ai_metadata(image_path: Path) -> bool:
"""Return whether a supported metadata signal is present."""
if any(_is_ai_field(key) for key in _read_pillow_info(image_path)):
return True
return has_c2pa_metadata(image_path)
def _summary_value(value: object) -> str:
if isinstance(value, bytes):
return f"<binary data ({len(value)} bytes)>"
text = str(value)
return text if len(text) <= 100 else f"{text[:100]}..."
def get_ai_metadata_summary(source_path: Path) -> str:
"""Format the AI-only metadata view for the command-line report."""
metadata = extract_ai_metadata(source_path)
if not metadata:
return "No AI metadata found."
lines = ["AI Image Metadata:", "-" * 40]
for key, value in metadata.items():
if key == "c2pa_chunk":
continue
if key == "c2pa" and isinstance(value, dict):
lines.append("C2PA Metadata:")
c2pa_fields = cast("dict[str, object]", value)
lines.extend(f" {name}: {_summary_value(item)}" for name, item in c2pa_fields.items())
continue
lines.append(f"{key}: {_summary_value(value)}")
return "\n".join(lines)
src/remove_ai_watermarks/noai/flv.py → src/remove_ai_watermarks/_internal/flv.py
View File
src/remove_ai_watermarks/_internal/img2img_runner.py
+136
View File
@@ -0,0 +1,136 @@
"""Execute Diffusers img2img calls and recover from an MPS runtime failure."""
from __future__ import annotations
import contextlib
import logging
from typing import TYPE_CHECKING, Any
from remove_ai_watermarks._internal.progress import is_mps_error, make_pipeline_progress
if TYPE_CHECKING:
from collections.abc import Callable
from PIL import Image
logger = logging.getLogger(__name__)
def _pipeline_arguments(
image: Image.Image,
strength: float,
num_inference_steps: int,
guidance_scale: float,
generator: Any,
step_callback: Any,
overrides: dict[str, Any] | None,
) -> dict[str, Any]:
arguments: dict[str, Any] = {
"prompt": "",
"image": image,
"strength": strength,
"num_inference_steps": num_inference_steps,
"guidance_scale": guidance_scale,
"generator": generator,
}
arguments.update(overrides or {})
if step_callback is not None:
arguments.update(callback=step_callback, callback_steps=1)
return arguments
def _invoke(pipeline: Any, arguments: dict[str, Any]) -> Image.Image:
response = pipeline(**arguments)
return response.images[0]
def run_img2img(
pipeline: Any,
image: Image.Image,
strength: float,
num_inference_steps: int,
guidance_scale: float,
generator: Any,
device: str,
set_progress: Callable[[str], None],
extra_kwargs: dict[str, Any] | None = None,
) -> Image.Image:
"""Run one img2img request and report denoising progress when supported."""
callback, started, finished, launch_monitor = make_pipeline_progress(
max(1, int(num_inference_steps * strength)), device, set_progress
)
launch_monitor()
arguments = _pipeline_arguments(
image, strength, num_inference_steps, guidance_scale, generator, callback, extra_kwargs
)
try:
try:
return _invoke(pipeline, arguments)
except TypeError as error:
if "callback" not in str(error):
raise
started.set()
arguments.pop("callback", None)
arguments.pop("callback_steps", None)
return _invoke(pipeline, arguments)
finally:
started.set()
finished.set()
def run_img2img_with_mps_fallback(
load_pipeline: Callable[[], Any],
image: Image.Image,
strength: float,
num_inference_steps: int,
guidance_scale: float,
generator: Any,
device: str,
set_progress: Callable[[str], None],
*,
reload_on_cpu: Callable[[], Any],
extra_kwargs: dict[str, Any] | None = None,
) -> tuple[Image.Image, str]:
"""Retry an MPS-specific failure once with a freshly loaded CPU pipeline."""
try:
output = run_img2img(
load_pipeline(),
image,
strength,
num_inference_steps,
guidance_scale,
generator,
device,
set_progress,
extra_kwargs,
)
return output, device
except RuntimeError as error:
if device != "mps" or not is_mps_error(error):
raise
logger.warning("MPS execution failed (%s); retrying on CPU", error)
set_progress("MPS execution failed; retrying on CPU...")
try_empty_device_cache("mps")
output = run_img2img(
reload_on_cpu(),
image,
strength,
num_inference_steps,
guidance_scale,
None,
"cpu",
set_progress,
extra_kwargs,
)
return output, "cpu"
def try_empty_device_cache(device: str) -> None:
"""Ask Torch to release cached accelerator memory when the backend supports it."""
with contextlib.suppress(Exception):
import torch
backend = getattr(torch, device, None)
empty_cache = getattr(backend, "empty_cache", None)
if callable(empty_cache):
empty_cache()
src/remove_ai_watermarks/noai/isobmff.py → src/remove_ai_watermarks/_internal/isobmff.py
View File
src/remove_ai_watermarks/_internal/progress.py
+212
View File
@@ -0,0 +1,212 @@
"""Progress reporting utilities for long-running optional model operations."""
from __future__ import annotations
import contextlib
import io
import os
import sys
import threading
import time
import warnings
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from collections.abc import Callable
_BAR_WIDTH = 28
_SPINNER = ("|", "/", "-", "\\")
def _truncate(text: str, max_len: int = 72) -> str:
if len(text) <= max_len:
return text
return f"{text[: max(0, max_len - 3)]}..."
def _build_bar(step: int) -> str:
position = step % (2 * _BAR_WIDTH - 2)
if position >= _BAR_WIDTH:
position = 2 * _BAR_WIDTH - 2 - position
cells = ["-"] * _BAR_WIDTH
cells[position] = "="
return "".join(cells)
@dataclass
class _TaskResult:
value: Any = None
error: BaseException | None = None
complete: threading.Event = field(default_factory=threading.Event)
def run_with_progress(task: Callable[[], Any], progress_state: dict[str, str] | None = None) -> Any:
"""Run ``task`` on a worker thread and render a compact terminal heartbeat."""
outcome = _TaskResult()
def invoke() -> None:
try:
outcome.value = task()
except BaseException as error: # re-raised on the caller thread
outcome.error = error
finally:
outcome.complete.set()
worker = threading.Thread(target=invoke, name="raiw-progress-task", daemon=True)
worker.start()
started_at = time.monotonic()
frame = 0
terminal = sys.__stderr__
while not outcome.complete.wait(0.1):
message = _truncate((progress_state or {}).get("message", "Processing..."))
elapsed = int(time.monotonic() - started_at)
if terminal is not None:
terminal.write(
f"\r\033[2K {_SPINNER[frame % len(_SPINNER)]} [{_build_bar(frame)}] {elapsed:>3}s {message}"
)
terminal.flush()
frame += 1
worker.join()
elapsed = int(time.monotonic() - started_at)
message = _truncate((progress_state or {}).get("message", "Processing..."))
if terminal is not None:
terminal.write(f"\r\033[2K Completed in {elapsed}s {message}\n")
terminal.flush()
if outcome.error is not None:
raise outcome.error
return outcome.value
def _silence_diffusers() -> None:
from diffusers.utils import logging as diffusers_logging
diffusers_logging.set_verbosity_error()
disable = getattr(diffusers_logging, "disable_progress_bar", None)
if callable(disable):
disable()
def _configure_quiet_libraries() -> None:
os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
operations = (
lambda: __import__("transformers").logging.set_verbosity_error(),
_silence_diffusers,
lambda: __import__("huggingface_hub").logging.set_verbosity_error(),
)
for operation in operations:
with contextlib.suppress(Exception):
operation()
def silence_library_output(
run_func: Callable[[], Any],
set_progress: Callable[[str], None] | None = None,
) -> Callable[[], Any]:
"""Wrap a model call so third-party progress bars do not corrupt our CLI UI."""
def quiet_call() -> Any:
if set_progress is not None:
set_progress("Preparing model runtime...")
_configure_quiet_libraries()
with (
warnings.catch_warnings(),
contextlib.redirect_stdout(io.StringIO()),
contextlib.redirect_stderr(io.StringIO()),
):
warnings.simplefilter("ignore")
if set_progress is not None:
set_progress("Running watermark regeneration...")
return run_func()
return quiet_call
@dataclass
class _PipelineMonitor:
total_steps: int
device: str
update: Callable[[str], None]
bar_len: int
label: str
pre_phases: list[tuple[int, str]]
post_phases: list[tuple[int, str]]
first_step: threading.Event = field(default_factory=threading.Event)
done: threading.Event = field(default_factory=threading.Event)
started_at: float = field(default_factory=time.monotonic)
last_step_at: float = field(default_factory=time.monotonic)
def callback(self, step: int, _timestep: int, _latents: Any) -> None:
self.first_step.set()
now = time.monotonic()
self.last_step_at = now
current = min(self.total_steps, step + 1)
filled = round(self.bar_len * current / self.total_steps)
elapsed = now - self.started_at
eta = elapsed * max(0, self.total_steps - current) / max(1, current)
bar = "#" * filled + "." * (self.bar_len - filled)
self.update(
f"{self.label} [{bar}] {current}/{self.total_steps}, "
f"{elapsed:.0f}s elapsed, ~{eta:.0f}s left, {self.device}"
)
def _phase_message(self, phases: list[tuple[int, str]], elapsed: float) -> str:
message = phases[0][1]
for threshold, candidate in phases:
if elapsed < threshold:
break
message = candidate
return message
def monitor(self) -> None:
while not self.first_step.wait(0.4):
elapsed = time.monotonic() - self.started_at
self.update(self._phase_message(self.pre_phases, elapsed))
decode_started: float | None = None
while not self.done.wait(0.4):
if time.monotonic() - self.last_step_at < 1.5:
decode_started = None
continue
decode_started = decode_started or time.monotonic()
self.update(self._phase_message(self.post_phases, time.monotonic() - decode_started))
def start(self) -> threading.Thread:
self.started_at = self.last_step_at = time.monotonic()
self.first_step.clear()
self.done.clear()
thread = threading.Thread(target=self.monitor, name="raiw-pipeline-progress", daemon=True)
thread.start()
return thread
def make_pipeline_progress(
effective_steps: int,
device: str,
set_progress: Callable[[str], None],
*,
bar_len: int = 20,
label: str = "Denoising",
pre_phases: list[tuple[int, str]] | None = None,
post_phases: list[tuple[int, str]] | None = None,
) -> tuple[Callable[..., None], threading.Event, threading.Event, Callable[[], threading.Thread]]:
"""Build a callback and monitor for the legacy Diffusers callback interface."""
def qualify(entries: list[tuple[int, str]]) -> list[tuple[int, str]]:
return [(second, f"{text} on {device}") for second, text in entries]
monitor = _PipelineMonitor(
total_steps=max(1, effective_steps),
device=device,
update=set_progress,
bar_len=bar_len,
label=label,
pre_phases=pre_phases or qualify([(0, "Encoding image"), (8, "Preparing denoiser"), (20, "Starting sampler")]),
post_phases=post_phases or qualify([(0, "Decoding image"), (10, "Finalizing pixels"), (45, "Still decoding")]),
)
return monitor.callback, monitor.first_step, monitor.done, monitor.start
def is_mps_error(error: Exception) -> bool:
"""Return whether an error message identifies Apple's MPS backend."""
return "mps" in str(error).casefold()
src/remove_ai_watermarks/noai/qwen_zimage_pipeline.py → src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py
+27 -55
View File
@@ -1,17 +1,9 @@
"""Qwen 2512 Canny regeneration followed by masked Z-Image face repair.
"""Project-native Qwen regeneration with an optional masked face refinement pass.
This profile ports the two-stage architecture used by cebeuq/Synthid-Bypass:
1. Qwen-Image-2512 img2img with the 4-step Lightning LoRA and the DiffSynth
blockwise Canny ControlNet regenerates the whole image, optionally as
overlapping feather-blended tiles for large inputs.
2. Faces are detected on the original image, refined to masks with SAM, regenerated
from the original face crops with Z-Image Turbo, and feathered into stage 1.
The runtime intentionally uses permissively licensed YuNet instead of the reference
workflow's Ultralytics detector. All diffusion and segmentation models remain the same
model families. The adaptive formulas are direct ports, while the face result is scaled
for this runtime's different sampler and mask-compositing path.
The profile was inspired by public experiments that combine structure-guided global
regeneration with a second face-only pass. Its orchestration, sizing rules, adaptive
strength policy, detector, masks, prompts, and compositing are implemented here for
this library's Pillow and DiffSynth runtime.
"""
# DiffSynth, torch, transformers, and cv2 expose mostly untyped tensor/array APIs.
@@ -33,7 +25,7 @@ from typing import TYPE_CHECKING, Any
import numpy as np
from PIL import Image
from remove_ai_watermarks.noai.watermark_profiles import resolve_seed
from remove_ai_watermarks._internal.watermark_profiles import resolve_seed
if TYPE_CHECKING:
from collections.abc import Callable
@@ -53,10 +45,8 @@ YUNET_MODEL_URL = (
)
YUNET_MODEL_NAME = "face_detection_yunet_2023mar.onnx"
YUNET_MODEL_SHA256 = "8f2383e4dd3cfbb4553ea8718107fc0423210dc964f9f4280604804ed2552fa4"
# The upstream graph's 0.2 threshold belongs to YOLO and does not transfer to
# YuNet's score calibration. At 0.2 YuNet admitted background and decorative
# false positives, multiplying the serial Z-Image face-stage cost. A 0.5 gate
# retained all visible faces in the public and upstream comparison fixtures.
# This threshold retained the faces in the public validation set without admitting
# decorative background regions as faces.
YUNET_SCORE_THRESHOLD = 0.5
GLOBAL_STEPS = 4
@@ -65,19 +55,13 @@ GLOBAL_CFG = 1.0
FACE_CFG = 1.0
GLOBAL_CONTROLNET_SCALE = 1.0
RESIDENT_FACE_MODEL_MIN_VRAM_GIB = 64.0
# The reference face denoise assumes its ComfyUI detailer sampler, latent
# noise-mask feather, and inpaint path. Applying that value unchanged to this
# DiffSynth crop-regeneration port over-processes faces. Paired public-fixture
# measurements and both provider oracles certified half the reference value.
FACE_DENOISE_SCALE = 0.5
# The source graph uses normalized Canny thresholds 0.05 and 0.25. OpenCV takes
# byte thresholds, so round 255*x to the matching integer values.
_CANNY_LOW = 13
_CANNY_HIGH = 64
# These strings intentionally preserve the reference workflow spelling. They are
# model inputs, not user-facing copy, and changing them would change the port.
# These short model inputs are retained as calibrated compatibility parameters.
# Changing them requires the same provider-oracle and identity evaluation as a model change.
_GLOBAL_PROMPT = "ultra clear and smoothe skin, spotless skin"
_GLOBAL_NEGATIVE = "moles, freckes, high detail skin"
_FACE_PROMPT = ""
@@ -169,28 +153,18 @@ def resolution_adaptive_denoise(
denoise_min: float = 0.08,
denoise_max: float = 0.15,
) -> float:
"""Port the reference resolution-based adaptive denoise calculation.
At neutral level 5, 0.30 MP maps to ``denoise_min`` and 3.70 MP maps to
``denoise_max``. Levels above or below 5 add the same asymmetric spread as the
reference custom node.
"""
image_mp = max(1.0, float(width) * float(height)) / 1_000_000.0
normalized = _clamp((image_mp - 0.30) / (3.70 - 0.30), 0.0, 1.0)
minimum = float(denoise_min)
maximum = float(denoise_max)
if maximum < minimum:
minimum, maximum = maximum, minimum
denoise_range = maximum - minimum
base = minimum + denoise_range * normalized
level = int(adaptive_level)
if level >= 5:
offset = ((float(level) - 5.0) / 5.0) * denoise_range * 0.285714
"""Choose the calibrated global strength from image area and operator level."""
low, high = sorted((float(denoise_min), float(denoise_max)))
megapixels = max(1.0, float(width) * float(height)) * 1e-6
area_fraction = _clamp((megapixels - 0.30) / 3.40, 0.0, 1.0)
strength_range = high - low
strength = low + strength_range * area_fraction
level_delta = float(int(adaptive_level)) - 5.0
if level_delta >= 0.0:
strength += level_delta * strength_range * (0.285714 / 5.0)
else:
offset = -((5.0 - float(level)) / 4.0) * denoise_range * 0.257143
return _clamp(base + offset, 0.0001, 1.0)
strength += level_delta * strength_range * (0.257143 / 4.0)
return _clamp(strength, 0.0001, 1.0)
def largest_face_denoise(
@@ -202,7 +176,7 @@ def largest_face_denoise(
denoise_min: float = 0.05,
denoise_max: float = 0.28,
) -> float:
"""Scale face denoise from the largest face area, matching reference mode."""
"""Choose the calibrated face strength from the largest detected face area."""
width, height = image_size
image_area = max(1.0, float(width) * float(height))
largest_ratio = 0.0
@@ -211,7 +185,7 @@ def largest_face_denoise(
largest_ratio = max(largest_ratio, box_area / image_area)
if largest_ratio <= 0.0:
return _clamp(base_denoise, denoise_min, denoise_max)
scaled = float(base_denoise) * (largest_ratio / max(1e-6, float(adaptive_ratio)))
scaled = float(base_denoise) * largest_ratio / max(1e-6, float(adaptive_ratio))
return _clamp(scaled, denoise_min, denoise_max)
@@ -229,7 +203,7 @@ def _resize_to_target(image: Image.Image) -> Image.Image:
def build_canny_control_image(image: Image.Image) -> Image.Image:
"""Build the three-channel Canny conditioning image used by stage 1."""
"""Build the calibrated three-channel Canny conditioning map."""
import cv2
rgb = np.asarray(image.convert("RGB"))
@@ -259,8 +233,6 @@ def build_global_kwargs(
"seed": seed,
"rand_device": "cpu",
"num_inference_steps": GLOBAL_STEPS,
# The source graph applies ModelSamplingAuraFlow with shift=3. DiffSynth
# expresses the same rational sigma shift as exp(mu), so mu=log(3).
"exponential_shift_mu": math.log(3.0),
"blockwise_controlnet_inputs": [controlnet_input],
}
@@ -312,7 +284,7 @@ def _expanded_box(
*,
factor: float = 2.5,
) -> tuple[int, int, int, int]:
"""Expand a face box around its center, matching the reference crop factor."""
"""Expand a face box around its center to include local lighting context."""
x1, y1, x2, y2 = box
image_width, image_height = image_size
center_x = (x1 + x2) / 2.0
@@ -813,7 +785,7 @@ class QwenZImagePipeline:
scale_for_face = 768.0 / max(1, max(face_width, face_height))
scale_for_crop = 1024.0 / max(1, max(crop_width, crop_height))
scale = min(scale_for_face, scale_for_crop)
# Never shrink below the crop's current size unless the 1024 cap requires it.
# Never shrink below the crop's current size unless the cap requires it.
if max(crop_width, crop_height) <= 1024:
scale = max(1.0, scale)
width = max(16, round(crop_width * scale / 16.0) * 16)
@@ -881,7 +853,7 @@ class QwenZImagePipeline:
resolution_adaptive_denoise(image.width, image.height) if strength is None else float(strength)
)
if tile and max(image.size) > tile_size:
from remove_ai_watermarks.noai.tiling import run_tiled
from remove_ai_watermarks._internal.tiling import run_tiled
global_result = run_tiled(
lambda tile_image: self._run_global(tile_image, global_strength, seed),
src/remove_ai_watermarks/noai/riff.py → src/remove_ai_watermarks/_internal/riff.py
View File
src/remove_ai_watermarks/noai/tiling.py → src/remove_ai_watermarks/_internal/tiling.py
View File
src/remove_ai_watermarks/_internal/utils.py
+30
View File
@@ -0,0 +1,30 @@
"""Small path helpers shared by optional image pipelines."""
from __future__ import annotations
from typing import TYPE_CHECKING
from remove_ai_watermarks._internal.constants import SUPPORTED_FORMATS
if TYPE_CHECKING:
from pathlib import Path
_PIL_FORMAT_BY_SUFFIX = {
".jpg": "JPEG",
".jpeg": "JPEG",
".png": "PNG",
}
def is_supported_format(file_path: Path) -> bool:
"""Return whether ``file_path`` has a supported raster suffix."""
return file_path.suffix.casefold() in SUPPORTED_FORMATS
def get_image_format(file_path: Path) -> str:
"""Return the Pillow save format used by the legacy metadata API.
The metadata writer only has specialized PNG and JPEG paths. Other accepted
inputs therefore use its PNG fallback, matching the established API contract.
"""
return _PIL_FORMAT_BY_SUFFIX.get(file_path.suffix.casefold(), "PNG")
src/remove_ai_watermarks/_internal/watermark_profiles.py
+105
View File
@@ -0,0 +1,105 @@
"""Project-owned configuration for invisible-watermark regeneration profiles."""
from __future__ import annotations
import math
from dataclasses import dataclass
from typing import TYPE_CHECKING, Literal
if TYPE_CHECKING:
from pathlib import Path
DEFAULT_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
QWEN_MODEL_ID = "Qwen/Qwen-Image"
CONTROLNET_CANNY_MODEL = "xinsir/controlnet-canny-sdxl-1.0"
SDXL_PROFILE = "sdxl"
QWEN_ZIMAGE_PROFILE = "qwen-zimage"
OPENAI_STRENGTH = 0.10
GEMINI_STRENGTH = 0.15
UNKNOWN_STRENGTH = GEMINI_STRENGTH
DEFAULT_STRENGTH = UNKNOWN_STRENGTH
QWEN_OPENAI_STRENGTH = 0.10
QWEN_GEMINI_STRENGTH = 0.25
QWEN_UNKNOWN_STRENGTH = QWEN_GEMINI_STRENGTH
@dataclass(frozen=True)
class _StrengthPolicy:
unknown: float
by_vendor: dict[str, float]
def choose(self, vendor: str | None) -> float:
return self.by_vendor.get((vendor or "").casefold(), self.unknown)
_STANDARD_POLICY = _StrengthPolicy(
unknown=UNKNOWN_STRENGTH,
by_vendor={"openai": OPENAI_STRENGTH, "google": GEMINI_STRENGTH},
)
_QWEN_POLICY = _StrengthPolicy(
unknown=QWEN_UNKNOWN_STRENGTH,
by_vendor={"openai": QWEN_OPENAI_STRENGTH, "google": QWEN_GEMINI_STRENGTH},
)
_ALIASES = {"default": SDXL_PROFILE, "qwen_zimage": QWEN_ZIMAGE_PROFILE}
def normalize_profile(profile: str) -> str:
"""Normalize spelling and resolve compatibility aliases."""
value = profile.strip().casefold()
return _ALIASES.get(value, value)
def resolve_steps(num_inference_steps: int | None, pipeline: str) -> int:
"""Return an explicit step count or the selected profile's default."""
if num_inference_steps is not None:
return num_inference_steps
return 4 if normalize_profile(pipeline) == QWEN_ZIMAGE_PROFILE else 50
def resolve_seed(seed: int | None, pipeline: str) -> int | None:
"""Keep the fixed Qwen plus Z-Image profile reproducible by default."""
if seed is not None:
return seed
return 0 if normalize_profile(pipeline) == QWEN_ZIMAGE_PROFILE else None
def strength_default_help() -> str:
"""Describe the live default policy without duplicating its values."""
return (
f"vendor-adaptive (OpenAI {OPENAI_STRENGTH} / Google {GEMINI_STRENGTH} / "
f"unknown {UNKNOWN_STRENGTH}, from the C2PA issuer; qwen-zimage instead uses "
"resolution-adaptive denoise)"
)
def resolve_strength(strength: float | None, vendor: str | None = None, pipeline: str | None = None) -> float:
"""Resolve a user override or the calibrated policy for a profile and vendor."""
if strength is not None:
return strength
policy = _QWEN_POLICY if pipeline is not None and normalize_profile(pipeline) == "qwen" else _STANDARD_POLICY
return policy.choose(vendor)
def viable_steps(num_inference_steps: int, strength: float) -> int:
"""Ensure Diffusers receives at least one effective img2img denoising step."""
if strength <= 0 or int(num_inference_steps * strength) >= 1:
return num_inference_steps
return math.ceil(1.0 / strength)
def vendor_for_strength(image_path: Path) -> Literal["openai", "google"] | None:
"""Select the strength cohort using the input's SynthID provenance proxy."""
try:
from remove_ai_watermarks.metadata import synthid_source
evidence = (synthid_source(image_path) or "").casefold()
except Exception:
return None
if "google" in evidence:
return "google"
if "openai" in evidence:
return "openai"
return None
src/remove_ai_watermarks/_internal/watermark_remover.py
+649
View File
@@ -0,0 +1,649 @@
"""Project-native orchestration for diffusion-based pixel regeneration."""
# 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 contextlib
import logging
import os
import subprocess
from typing import TYPE_CHECKING, Any
from PIL import Image
from remove_ai_watermarks._internal.watermark_profiles import (
CONTROLNET_CANNY_MODEL,
DEFAULT_MODEL_ID,
DEFAULT_STRENGTH,
QWEN_MODEL_ID,
QWEN_ZIMAGE_PROFILE,
normalize_profile,
resolve_seed,
resolve_steps,
resolve_strength,
viable_steps,
)
if TYPE_CHECKING:
from collections.abc import Callable
from pathlib import Path
logger = logging.getLogger(__name__)
try:
import torch
_HAS_TORCH = True
except ImportError:
torch = None # type: ignore[assignment]
_HAS_TORCH = False
try:
from diffusers import AutoPipelineForImage2Image as AutoImg2ImgPipeline
_HAS_DIFFUSERS = True
except ImportError:
AutoImg2ImgPipeline = None # type: ignore[assignment,misc]
_HAS_DIFFUSERS = False
_SDXL_FP16_VAE_ID = "madebyollin/sdxl-vae-fp16-fix"
_DEGENERATE_THRESHOLD = 1.0
_CANNY_LOW = 100
_CANNY_HIGH = 200
_CONTROLNET_PROMPT = "best quality, high quality, sharp, detailed, photographic"
_CONTROLNET_NEGATIVE = "blurry, lowres, deformed, distorted text, garbled text, watermark, jpeg artifacts"
_QWEN_PROMPT = "high quality, sharp, detailed, faithful to the original"
_QWEN_NEGATIVE = "blurry, lowres, distorted text, garbled text, artifacts"
def is_watermark_removal_available() -> bool:
"""Return whether the standard diffusion runtime can be imported."""
return _HAS_TORCH and _HAS_DIFFUSERS
def _ensure_watermark_deps() -> None:
if not is_watermark_removal_available():
raise ImportError(
"Invisible watermark regeneration requires the 'diffusion' extra. Install remove-ai-watermarks[diffusion]."
)
def _needs_fp16_vae_fix(model_id: str, default_model_id: str, is_fp16: bool) -> bool:
"""Return whether the default SDXL pipeline needs the overflow-safe VAE."""
return is_fp16 and model_id == default_model_id
def _is_degenerate_image(image: Image.Image) -> bool:
"""Detect the uniform near-black output produced by an fp16 decode collapse."""
import numpy as np
pixels = np.asarray(image.convert("RGB"), dtype=np.float32)
return float(pixels.mean()) < _DEGENERATE_THRESHOLD and float(pixels.std()) < _DEGENERATE_THRESHOLD
def _has_nvidia_gpu() -> bool:
try:
subprocess.run(
["nvidia-smi"],
check=True,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
)
except (FileNotFoundError, subprocess.CalledProcessError):
return False
return True
def _detect_cuda_index_url() -> str:
"""Return a PyTorch wheel index compatible with the reported CUDA runtime."""
try:
report = subprocess.run(
["nvidia-smi"],
check=True,
capture_output=True,
text=True,
).stdout
except (FileNotFoundError, subprocess.CalledProcessError):
return "https://download.pytorch.org/whl/cu121"
import re
match = re.search(r"CUDA Version:\s*(\d+)\.(\d+)", report)
if match is None:
return "https://download.pytorch.org/whl/cu121"
return f"https://download.pytorch.org/whl/cu{match.group(1)}{match.group(2)}"
def _backend_works(device: str) -> bool:
try:
probe = torch.tensor([1.0], device=device) # type: ignore[union-attr]
_ = probe + probe
except (AssertionError, RuntimeError):
return False
return True
def get_device() -> str:
"""Select CUDA, XPU, MPS, or CPU in that order when each backend is usable."""
if not _HAS_TORCH:
return "cpu"
if torch.cuda.is_available() and _backend_works("cuda"): # type: ignore[union-attr]
return "cuda"
xpu = getattr(torch, "xpu", None)
if xpu is not None and xpu.is_available() and _backend_works("xpu"):
return "xpu"
if _has_nvidia_gpu():
logger.warning("NVIDIA GPU detected, but the installed PyTorch build has no working CUDA backend")
mps = getattr(getattr(torch, "backends", None), "mps", None)
if mps is not None and mps.is_available():
return "mps"
return "cpu"
def _make_seed_generator(device: str, seed: int) -> Any:
"""Create a deterministic generator, using CPU when device RNG is unavailable."""
try:
return torch.Generator(device=device).manual_seed(seed) # type: ignore[union-attr]
except (RuntimeError, TypeError):
return torch.Generator().manual_seed(seed) # type: ignore[union-attr]
def _qwen_target_size(width: int, height: int) -> tuple[int, int]:
"""Floor dimensions to Qwen's 16-pixel latent grid."""
return max(16, width - width % 16), max(16, height - height % 16)
def _build_qwen_kwargs(
image: Image.Image,
strength: float,
num_inference_steps: int,
true_cfg_scale: float,
generator: Any,
) -> dict[str, Any]:
"""Build the Qwen img2img call without importing its optional pipeline class."""
width, height = _qwen_target_size(image.width, image.height)
return {
"prompt": _QWEN_PROMPT,
"negative_prompt": _QWEN_NEGATIVE,
"image": image,
"strength": strength,
"num_inference_steps": num_inference_steps,
"true_cfg_scale": true_cfg_scale,
"generator": generator,
"width": width,
"height": height,
}
class WatermarkRemover:
"""Load one regeneration profile and write a metadata-clean raster output."""
DEFAULT_MODEL_ID = DEFAULT_MODEL_ID
DEFAULT_STRENGTH = DEFAULT_STRENGTH
CONTROLNET_CANNY_MODEL = CONTROLNET_CANNY_MODEL
_DEVICES = frozenset({"cpu", "mps", "cuda", "xpu"})
def __init__(
self,
model_id: str | None = None,
device: str | None = None,
torch_dtype: Any = None,
progress_callback: Callable[[str], None] | None = None,
hf_token: str | None = None,
pipeline: str = "controlnet",
controlnet_conditioning_scale: float = 1.0,
cpu_offload: bool = False,
) -> None:
requested_model = model_id or self.DEFAULT_MODEL_ID
self.model_profile = normalize_profile(pipeline)
if self.model_profile == QWEN_ZIMAGE_PROFILE and model_id not in {None, self.DEFAULT_MODEL_ID}:
raise ValueError("The qwen-zimage profile uses a fixed Qwen-Image-2512 and Z-Image model stack.")
self.model_id = (
"Qwen/Qwen-Image-2512 + Tongyi-MAI/Z-Image-Turbo"
if self.model_profile == QWEN_ZIMAGE_PROFILE
else requested_model
)
_ensure_watermark_deps()
selected_device = (device or get_device()).casefold()
self.device = get_device() if selected_device == "auto" else selected_device
if self.device not in self._DEVICES:
raise ValueError(f"Unsupported device '{device}'. Use one of: auto, cpu, mps, cuda, xpu.")
if torch_dtype is not None:
self.torch_dtype = torch_dtype
elif self.device in {"cpu", "mps"}:
self.torch_dtype = torch.float32 # type: ignore[union-attr]
elif self.model_profile in {"qwen", QWEN_ZIMAGE_PROFILE}:
self.torch_dtype = torch.bfloat16 # type: ignore[union-attr]
else:
self.torch_dtype = torch.float16 # type: ignore[union-attr]
self.cpu_offload = cpu_offload
self.controlnet_conditioning_scale = controlnet_conditioning_scale
self.hf_token = hf_token or os.environ.get("HF_TOKEN")
self._progress_callback = progress_callback
self._pipeline: Any = None
self._controlnet_pipeline: Any = None
self._qwen_pipeline: Any = None
self._qwen_zimage_pipeline: Any = None
def _set_progress(self, message: str) -> None:
if self._progress_callback is not None:
with contextlib.suppress(Exception):
self._progress_callback(message)
def preload(self, *, global_only: bool = False) -> None:
"""Materialize the selected model stack before the first request."""
if self.model_profile == QWEN_ZIMAGE_PROFILE:
self._load_qwen_zimage_pipeline().preload(global_only=global_only)
elif self.model_profile == "qwen":
self._load_qwen_pipeline()
elif self.model_profile == "controlnet":
self._load_controlnet_pipeline()
else:
self._load_pipeline()
def _base_load_kwargs(self) -> dict[str, Any]:
options: dict[str, Any] = {"torch_dtype": self.torch_dtype}
if self.hf_token:
options["token"] = self.hf_token
return options
def _load_from_pretrained(self, cls: Any, model_id: str, **kwargs: Any) -> Any:
if self.torch_dtype == torch.float16: # type: ignore[union-attr]
try:
return cls.from_pretrained(model_id, variant="fp16", **kwargs)
except Exception as error:
logger.info("Model %s has no usable fp16 variant (%s); using default weights", model_id, error)
return cls.from_pretrained(model_id, **kwargs)
def _maybe_add_fp16_vae(self, options: dict[str, Any]) -> None:
if not _needs_fp16_vae_fix(
self.model_id,
self.DEFAULT_MODEL_ID,
self.torch_dtype == torch.float16, # type: ignore[union-attr]
):
return
from diffusers import AutoencoderKL
options["vae"] = AutoencoderKL.from_pretrained(_SDXL_FP16_VAE_ID, torch_dtype=torch.float16)
@staticmethod
def _disable_sdxl_watermarker(options: dict[str, Any]) -> None:
options["add_watermarker"] = False
def _move_to_device_and_optimize(self, pipeline: Any) -> Any:
if self.cpu_offload and self.device == "cuda":
offload = getattr(pipeline, "enable_model_cpu_offload", None)
if not callable(offload):
raise RuntimeError("CPU offload was requested, but this pipeline does not support it.")
offload(device="cuda")
else:
try:
pipeline = pipeline.to(self.device)
except (RuntimeError, AssertionError) as error:
if self.device == "cuda":
raise RuntimeError(
f"Failed to move model to CUDA ({error}). Install a compatible PyTorch wheel from "
f"{_detect_cuda_index_url()}."
) from error
raise
optimize = getattr(pipeline, "enable_xformers_memory_efficient_attention", None)
if callable(optimize):
with contextlib.suppress(Exception):
optimize()
if self.device == "mps":
slice_attention = getattr(pipeline, "enable_attention_slicing", None)
if callable(slice_attention):
with contextlib.suppress(Exception):
slice_attention("max")
return pipeline
def _sdxl_options(self) -> dict[str, Any]:
options = self._base_load_kwargs()
self._disable_sdxl_watermarker(options)
self._maybe_add_fp16_vae(options)
return options
def _load_pipeline(self) -> Any:
if self._pipeline is None:
options = self._sdxl_options()
options.update(safety_checker=None, requires_safety_checker=False)
loaded = self._load_from_pretrained(AutoImg2ImgPipeline, self.model_id, **options)
self._pipeline = self._move_to_device_and_optimize(loaded)
return self._pipeline
def _load_controlnet_pipeline(self) -> Any:
if self._controlnet_pipeline is None:
from diffusers import ControlNetModel, StableDiffusionXLControlNetImg2ImgPipeline
controlnet = self._load_from_pretrained(
ControlNetModel,
CONTROLNET_CANNY_MODEL,
torch_dtype=self.torch_dtype,
)
options = self._sdxl_options()
options["controlnet"] = controlnet
loaded = self._load_from_pretrained(
StableDiffusionXLControlNetImg2ImgPipeline,
self.model_id,
**options,
)
self._controlnet_pipeline = self._move_to_device_and_optimize(loaded)
return self._controlnet_pipeline
def _load_qwen_pipeline(self) -> Any:
if self._qwen_pipeline is None:
try:
from diffusers import QwenImageImg2ImgPipeline
except ImportError as error:
raise ImportError("The qwen profile requires Diffusers with QwenImageImg2ImgPipeline.") from error
model_id = QWEN_MODEL_ID if self.model_id == self.DEFAULT_MODEL_ID else self.model_id
loaded = QwenImageImg2ImgPipeline.from_pretrained(model_id, **self._base_load_kwargs())
self._qwen_pipeline = self._move_to_device_and_optimize(loaded)
return self._qwen_pipeline
def _load_qwen_zimage_pipeline(self) -> Any:
if self._qwen_zimage_pipeline is None:
from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline
self._qwen_zimage_pipeline = QwenZImagePipeline(
device=self.device,
torch_dtype=self.torch_dtype,
hf_token=self.hf_token,
progress_callback=self._progress_callback,
controlnet_conditioning_scale=self.controlnet_conditioning_scale,
keep_face_models_on_device=False if self.cpu_offload else None,
)
return self._qwen_zimage_pipeline
def _reload_on_cpu(self, cache_name: str, loader: Callable[[], Any]) -> Any:
self.device = "cpu"
self.torch_dtype = torch.float32 # type: ignore[union-attr]
setattr(self, cache_name, None)
return loader()
def _run_img2img(
self,
init_image: Image.Image,
strength: float,
num_inference_steps: int,
guidance_scale: float,
generator: Any,
) -> Image.Image:
from remove_ai_watermarks._internal.img2img_runner import run_img2img_with_mps_fallback
output, device = run_img2img_with_mps_fallback(
self._load_pipeline,
init_image,
strength,
num_inference_steps,
guidance_scale,
generator,
self.device,
self._set_progress,
reload_on_cpu=lambda: self._reload_on_cpu("_pipeline", self._load_pipeline),
)
self.device = device
return output
def _build_canny_control_image(self, init_image: Image.Image) -> Image.Image:
import cv2
import numpy as np
gray = cv2.cvtColor(np.asarray(init_image.convert("RGB")), cv2.COLOR_RGB2GRAY)
edges = cv2.Canny(gray, _CANNY_LOW, _CANNY_HIGH)
return Image.fromarray(np.repeat(edges[:, :, None], 3, axis=2))
def _run_controlnet(
self,
init_image: Image.Image,
strength: float,
num_inference_steps: int,
guidance_scale: float,
generator: Any,
) -> Image.Image:
from remove_ai_watermarks._internal.img2img_runner import run_img2img_with_mps_fallback
extras = {
"prompt": _CONTROLNET_PROMPT,
"negative_prompt": _CONTROLNET_NEGATIVE,
"control_image": self._build_canny_control_image(init_image),
"controlnet_conditioning_scale": float(self.controlnet_conditioning_scale),
}
output, device = run_img2img_with_mps_fallback(
self._load_controlnet_pipeline,
init_image,
strength,
num_inference_steps,
guidance_scale,
generator,
self.device,
self._set_progress,
reload_on_cpu=lambda: self._reload_on_cpu("_controlnet_pipeline", self._load_controlnet_pipeline),
extra_kwargs=extras,
)
self.device = device
return output
def _run_qwen(
self,
init_image: Image.Image,
strength: float,
num_inference_steps: int,
guidance_scale: float,
generator: Any,
) -> Image.Image:
response = self._load_qwen_pipeline()(
**_build_qwen_kwargs(init_image, strength, num_inference_steps, guidance_scale, generator)
)
return response.images[0]
def _run_qwen_zimage(
self,
init_image: Image.Image,
strength: float,
seed: int | None,
*,
tile: bool = False,
tile_size: int = 1024,
tile_overlap: int = 128,
) -> Image.Image:
return self._load_qwen_zimage_pipeline().run(
init_image,
strength=strength,
seed=seed,
tile=tile,
tile_size=tile_size,
tile_overlap=tile_overlap,
)
def _generate(
self,
image: Image.Image,
strength: float,
steps: int,
guidance: float,
generator: Any,
seed: int | None,
*,
tile: bool,
tile_size: int,
tile_overlap: int,
) -> Image.Image:
if self.model_profile == QWEN_ZIMAGE_PROFILE:
return self._run_qwen_zimage(
image,
strength,
seed,
tile=tile,
tile_size=tile_size,
tile_overlap=tile_overlap,
)
runner = {
"qwen": self._run_qwen,
"controlnet": self._run_controlnet,
}.get(self.model_profile, self._run_img2img)
if tile and max(image.size) > tile_size:
from remove_ai_watermarks._internal.tiling import run_tiled
return run_tiled(
lambda crop: runner(crop, strength, steps, guidance, generator),
image,
tile_size,
tile_overlap,
self._set_progress,
)
return runner(image, strength, steps, guidance, generator)
def _write_output(self, image: Image.Image, output_path: Path) -> None:
import numpy as np
from remove_ai_watermarks import image_io
output_path.parent.mkdir(parents=True, exist_ok=True)
bgr = np.ascontiguousarray(np.asarray(image.convert("RGB"))[:, :, ::-1])
if not image_io.imwrite(str(output_path), bgr):
image.save(output_path)
from remove_ai_watermarks.metadata import remove_ai_metadata
remove_ai_metadata(output_path, output_path, keep_standard=True)
def remove_watermark(
self,
image_path: Path,
output_path: Path | None = None,
strength: float | None = None,
num_inference_steps: int | None = None,
guidance_scale: float | None = None,
seed: int | None = None,
vendor: str | None = None,
tile: bool = False,
tile_size: int = 1024,
tile_overlap: int = 128,
region: tuple[int, int, int, int] | None = None,
region_feather: int = 64,
) -> Path:
"""Regenerate image pixels and write the result without AI metadata."""
if not image_path.exists():
raise FileNotFoundError(f"Image not found: {image_path}")
destination = output_path or image_path
with Image.open(image_path) as opened:
source = opened.convert("RGB")
if self.model_profile == QWEN_ZIMAGE_PROFILE:
from remove_ai_watermarks._internal.qwen_zimage_pipeline import resolution_adaptive_denoise
resolved_strength = strength if strength is not None else resolution_adaptive_denoise(*source.size)
else:
resolved_strength = resolve_strength(strength, vendor, self.model_profile)
if not 0.0 <= resolved_strength <= 1.0:
raise ValueError(f"Strength must be between 0.0 and 1.0, got {resolved_strength}")
resolved_seed = resolve_seed(seed, self.model_profile)
steps = resolve_steps(num_inference_steps, self.model_profile)
guidance = (
1.0 if guidance_scale is None and self.model_profile == QWEN_ZIMAGE_PROFILE else guidance_scale or 7.5
)
if self.model_profile == QWEN_ZIMAGE_PROFILE:
if steps != 4:
raise ValueError("The qwen-zimage profile requires 4 steps.")
if guidance != 1.0:
raise ValueError("The qwen-zimage profile requires CFG 1.0.")
else:
steps = viable_steps(steps, resolved_strength)
generator = None
if resolved_seed is not None and _HAS_TORCH:
generator = _make_seed_generator(self.device, resolved_seed)
result = self._generate(
source,
resolved_strength,
steps,
guidance,
generator,
resolved_seed,
tile=tile,
tile_size=tile_size,
tile_overlap=tile_overlap,
)
if self.torch_dtype == torch.float16 and _is_degenerate_image(result): # type: ignore[union-attr]
self.torch_dtype = torch.float32 # type: ignore[union-attr]
self._pipeline = self._controlnet_pipeline = self._qwen_pipeline = self._qwen_zimage_pipeline = None
result = self._generate(
source,
resolved_strength,
steps,
guidance,
generator,
resolved_seed,
tile=tile,
tile_size=tile_size,
tile_overlap=tile_overlap,
)
if region is not None:
import numpy as np
from remove_ai_watermarks._internal.tiling import feather_region_composite
if result.size != source.size:
result = result.resize(source.size, Image.Resampling.LANCZOS)
merged = feather_region_composite(
np.asarray(source),
np.asarray(result.convert("RGB")),
region,
feather=region_feather,
)
result = Image.fromarray(merged)
self._write_output(result, destination)
return destination
def remove_watermark_batch(
self,
input_dir: Path,
output_dir: Path,
strength: float | None = None,
num_inference_steps: int | None = None,
extensions: tuple[str, ...] = (".png", ".jpg", ".jpeg", ".webp"),
) -> list[Path]:
"""Process matching files in a directory, logging and continuing on failures."""
if not input_dir.exists():
raise FileNotFoundError(f"Input directory not found: {input_dir}")
output_dir.mkdir(parents=True, exist_ok=True)
from remove_ai_watermarks._internal.img2img_runner import try_empty_device_cache
outputs: list[Path] = []
candidates = sorted(path for path in input_dir.iterdir() if path.suffix.casefold() in extensions)
for source in candidates:
try:
outputs.append(self.remove_watermark(source, output_dir / source.name, strength, num_inference_steps))
except Exception as error:
logger.error("Failed to process %s: %s", source, error)
finally:
try_empty_device_cache(self.device)
return outputs
def remove_watermark(
image_path: Path,
output_path: Path | None = None,
strength: float | None = None,
model_id: str | None = None,
device: str | None = None,
hf_token: str | None = None,
region: tuple[int, int, int, int] | None = None,
) -> Path:
"""Convenience wrapper using the default ControlNet profile."""
from remove_ai_watermarks._internal.watermark_profiles import vendor_for_strength
remover = WatermarkRemover(model_id=model_id, device=device, hf_token=hf_token)
return remover.remove_watermark(
image_path,
output_path,
strength,
vendor=vendor_for_strength(image_path),
region=region,
)
src/remove_ai_watermarks/cli.py
+4 -4
View File
@@ -21,8 +21,8 @@ from typing import TYPE_CHECKING, Any, Literal, NoReturn
import click
from remove_ai_watermarks import __version__, image_io, watermark_registry
from remove_ai_watermarks.noai.constants import SUPPORTED_FORMATS
from remove_ai_watermarks.noai.watermark_profiles import (
from remove_ai_watermarks._internal.constants import SUPPORTED_FORMATS
from remove_ai_watermarks._internal.watermark_profiles import (
resolve_seed,
resolve_steps,
resolve_strength,
@@ -158,7 +158,7 @@ def _resolved_strength_for_display(
if pipeline == "qwen-zimage" and strength is None:
from PIL import Image
from remove_ai_watermarks.noai.qwen_zimage_pipeline import resolution_adaptive_denoise
from remove_ai_watermarks._internal.qwen_zimage_pipeline import resolution_adaptive_denoise
with Image.open(source) as image:
return resolution_adaptive_denoise(image.width, image.height)
@@ -269,7 +269,7 @@ def _normalize_pipeline(ctx: click.Context, param: click.Parameter, value: str |
"""
if value is None:
return None
from remove_ai_watermarks.noai.watermark_profiles import normalize_profile
from remove_ai_watermarks._internal.watermark_profiles import normalize_profile
normalized = normalize_profile(value)
if value.strip().lower() == "default":
src/remove_ai_watermarks/gemini_engine.py
+251 -562
View File
@@ -1,22 +1,6 @@
"""Gemini visible-sparkle detector and localizer (cv2/numpy, no GPU).
"""Locate the visible Gemini sparkle and build a mask for shared inpainting."""
Locates the Google Gemini / Nano Banana sparkle so the shared fill (region_eraser)
can inpaint it. Detection is a multi-scale NCC search against the captured sparkle
alpha template (ported from GeminiWatermarkTool's Snap Engine; original author
Allen Kuo (allenk), https://github.com/allenk/GeminiWatermarkTool), scored by a
spatial + gradient + variance fusion with a false-positive gate. ``footprint_mask``
returns the sparkle footprint (captured alpha thresholded low to include the halo,
then dilated) as a full-frame mask for the fill.
The captured alpha maps are background captures of the sparkle on pure-black
backgrounds (48x48 for small images, 96x96 for large). NB: they are used here only
to DETECT and to shape the removal mask -- the old reverse-alpha pixel recovery
(``original = (watermarked - a*logo)/(1-a)``) is gone; removal is localize -> fill.
"""
# cv2/numpy boundary: cv2 and numpy ship no usable type info for the array ops
# below, so strict pyright cannot know their element types. Relax the unknown-type
# rules for this file only; the public signatures are still annotated with NDArray[Any].
# OpenCV and NumPy expose incomplete types at this array-processing boundary.
# 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
@@ -41,263 +25,172 @@ logger = logging.getLogger(__name__)
class WatermarkSize(Enum):
"""Watermark size mode based on image dimensions."""
"""Provider size tier selected from the source dimensions."""
SMALL = "small" # 48x48, for images <= 1024x1024
LARGE = "large" # 96x96, for images > 1024x1024
SMALL = "small"
LARGE = "large"
@dataclass
class DetectionResult:
"""Result of watermark detection."""
"""Detection decision and its component scores."""
detected: bool = False
confidence: float = 0.0
region: tuple[int, int, int, int] = (0, 0, 0, 0) # x, y, w, h
region: tuple[int, int, int, int] = (0, 0, 0, 0)
size: WatermarkSize = WatermarkSize.SMALL
# stage scores
spatial_score: float = 0.0
gradient_score: float = 0.0
variance_score: float = 0.0
@dataclass
@dataclass(frozen=True, slots=True)
class WatermarkPosition:
"""Watermark position configuration."""
"""Expected provider margins and logo size."""
margin_right: int
margin_bottom: int
logo_size: int
def get_position(self, image_width: int, image_height: int) -> tuple[int, int]:
"""Get top-left position for a given image size."""
x = image_width - self.margin_right - self.logo_size
y = image_height - self.margin_bottom - self.logo_size
return (x, y)
return image_width - self.margin_right - self.logo_size, image_height - self.margin_bottom - self.logo_size
def get_watermark_config(width: int, height: int) -> WatermarkPosition:
"""Get the appropriate watermark configuration based on image size.
@dataclass(frozen=True, slots=True)
class _Candidate:
scale: int
x: int
y: int
spatial: float
gradient: float = 0.0
variance: float = 0.0
Rules discovered from Gemini:
- W > 1024 AND H > 1024: 96x96 logo at (W-64-96, H-64-96)
- Otherwise: 48x48 logo at (W-32-48, H-32-48)
"""
if width > 1024 and height > 1024:
return WatermarkPosition(margin_right=64, margin_bottom=64, logo_size=96)
return WatermarkPosition(margin_right=32, margin_bottom=32, logo_size=48)
@property
def fused(self) -> float:
if self.spatial < 0.25:
return max(0.0, self.spatial * 0.5)
return self.spatial * 0.50 + self.gradient * 0.30 + self.variance * 0.20
def get_watermark_size(width: int, height: int) -> WatermarkSize:
"""Determine watermark size mode from image dimensions."""
if width > 1024 and height > 1024:
return WatermarkSize.LARGE
return WatermarkSize.SMALL
"""Return the provider's large tier only when both axes exceed 1024."""
return WatermarkSize.LARGE if width > 1024 and height > 1024 else WatermarkSize.SMALL
def _calculate_alpha_map(bg_capture: NDArray[Any]) -> NDArray[Any]:
"""Calculate alpha map from a background capture.
def get_watermark_config(width: int, height: int) -> WatermarkPosition:
"""Return the observed standard placement for the selected size tier."""
if get_watermark_size(width, height) is WatermarkSize.LARGE:
return WatermarkPosition(64, 64, 96)
return WatermarkPosition(32, 32, 48)
The alpha map represents how much the watermark affects each pixel.
alpha = max(R, G, B) / 255.0
"""
if len(bg_capture.shape) == 2:
gray = bg_capture.astype(np.float32)
elif bg_capture.shape[2] >= 3:
# Use max of channels for brightness
gray = np.max(bg_capture[:, :, :3], axis=2).astype(np.float32)
def _calculate_alpha_map(background_capture: NDArray[Any]) -> NDArray[Any]:
"""Convert a black-background sparkle capture to a normalized opacity map."""
if background_capture.ndim == 2:
intensity = background_capture
elif background_capture.shape[2] >= 3:
intensity = background_capture[:, :, :3].max(axis=2)
else:
gray = bg_capture[:, :, 0].astype(np.float32)
return gray / 255.0
intensity = background_capture[:, :, 0]
return intensity.astype(np.float32) / 255.0
def _load_embedded_asset(name: str) -> NDArray[Any]:
"""Load an embedded PNG asset and decode it with OpenCV."""
asset_path = Path(__file__).parent / "assets" / name
if not asset_path.exists():
raise FileNotFoundError(f"Embedded asset not found: {asset_path}")
data = asset_path.read_bytes()
buf = np.frombuffer(data, dtype=np.uint8)
img = cv2.imdecode(buf, cv2.IMREAD_COLOR)
if img is None:
raise RuntimeError(f"Failed to decode embedded asset: {name}")
return img
def _load_capture(filename: str, expected_side: int) -> NDArray[Any]:
capture = image_io.imread(Path(__file__).parent / "assets" / filename, cv2.IMREAD_COLOR)
if capture is None:
raise RuntimeError(f"Failed to decode embedded asset: {filename}")
if capture.shape[:2] != (expected_side, expected_side):
capture = cv2.resize(capture, (expected_side, expected_side), interpolation=cv2.INTER_AREA)
return capture
# Single source of truth for the multi-scale template ladder (aggressively downscaled to
# slightly upscaled): the precomputed `_tmpl_cache` and the `_scan_scales` loop must use
# the SAME scales or a scan scale would miss the cache and KeyError.
_TEMPLATE_SCALES: tuple[int, ...] = tuple(range(16, 120, 2))
def _gray_float(image: NDArray[Any]) -> NDArray[Any]:
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if image.ndim == 3 and image.shape[2] >= 3 else image
return gray.astype(np.float32) / 255.0
def _overlaps(candidate: _Candidate, selected: _Candidate) -> bool:
radius = 0.5 * max(candidate.scale, selected.scale)
return abs(candidate.x - selected.x) < radius and abs(candidate.y - selected.y) < radius
_TEMPLATE_SCALES = tuple(range(16, 120, 2))
class GeminiEngine:
"""Detects and localizes the visible Gemini sparkle for the shared fill removal.
"""Project-native detector and mask builder for the white Gemini sparkle."""
The multi-scale NCC detection is a Python port of the GeminiWatermarkTool C++
Snap Engine; ``footprint_mask`` turns a detection into a removal mask.
"""
# Body pixels at >= this fraction of the peak captured alpha define the sparkle
# "core", sampled by the detection FP-gate's core-vs-ring brightness margin
# (:meth:`_core_and_bg`).
_CORE_ALPHA_FRAC = 0.8
# Sparkle false-positive gate. A real Gemini sparkle is a bright WHITE overlay,
# so its core sits above the local background; a shape-only NCC match on ornate
# or flat content (text, banners, hatching) can score >0.5 without that lift.
# Demote a detection that is BOTH low-confidence AND low core-ring brightness
# margin -- the joint signature of a content false positive (verified on the
# detector calibration: demoted examples were visual false positives or a
# near-invisible white-on-white sparkle whose AI verdict is held by metadata
# anyway). Real sparkles escape via EITHER high confidence
# (white-bg sparkles score >=0.79 despite a low margin) OR high margin (dark/mid
# backgrounds, incl. the #36 faint-corner case, lift well clear), so both must
# fail to demote.
_SPARKLE_FP_CONF = 0.65
_SPARKLE_FP_MARGIN = 5.0
# Bright-background content false positives (2026-06-26 landing-page FPs: a snow+sky
# photo and a white-background product render both scored ~0.51). The margin gate
# above cannot catch them -- a bright background gives the "core" a HIGH core-ring
# margin (it is genuinely brighter than its surroundings), so the brightness check
# reads it as a real overlay. The discriminating signature is the GRADIENT NCC: a
# real white sparkle is a crisp star silhouette (grad ~0.97-1.0 on the synthetic
# composites, ~0.96 on the real #36 corner sparkle), while a smooth luminance blob
# that shape-NCC-matches the rough outline has low gradient fidelity (the two FPs
# measured 0.105 and 0.463). So ALSO demote a low-confidence match whose gradient
# NCC is below this floor, regardless of margin -- 0.55 sits well above the worst FP
# (0.463) and far below every real sparkle (>=0.8). This only ADDS demotions on
# bright backgrounds (a real bright-bg sparkle keeps grad ~0.97), so it cannot
# regress a dark/mid sparkle (already kept by margin) or a white-bg one (kept by
# confidence >= 0.65, above the gate).
_SPARKLE_FP_GRAD = 0.55
# White-core rescue for the gate above. A real but FAINT sparkle -- a soft white
# star on a bright/textured background -- has a high core-ring margin but low
# gradient fidelity, the SAME signature the grad gate uses to demote the smooth
# colored-corner FP, so faint real sparkles get demoted with it. The separator the
# grad gate discards is the CORE COLOR: a real Gemini sparkle core is near-WHITE
# (low saturation), while a clean bright corner that shape-matches (sky, sun, a warm
# light) is COLORED. So do NOT demote a low-grad match that already clears the trust
# confidence (_SPARKLE_KEEP_CONF -- the registry's 0.5 sparkle gate plus a small
# margin so the ~0.51 bright-background FPs the grad gate was added for stay demoted)
# AND has a bright (margin) near-neutral core (_core_saturation <= _SPARKLE_WHITE_SAT).
# Calibrated on metadata-stripped faint sparkles to recover low-gradient marks
# without materially increasing clean false fires.
_SPARKLE_KEEP_CONF = 0.52
_SPARKLE_WHITE_SAT = 0.20
# Corner promotion (issue #36): the size weight that suppresses tiny-patch
# false positives also buries a small, near-perfect sparkle when a larger,
# mediocre match sits elsewhere (e.g. a bright collar in a portrait). A small
# faint sparkle on a busy background therefore loses the global argmax and the
# image reads as clean -- the regression osachub reported when the search
# window widened 256px -> 512px (v0.7.2's tighter window still found it).
# Remedy: if the bottom-right corner holds a very-high-fidelity raw-NCC match,
# trust it regardless of size, without reverting the wider window (which is
# needed for variant margins). The threshold sits midway between the worst
# real-photo corner match (~0.78 across native + downscaled real photos) and a
# genuine faint sparkle (~0.93), so it adds true detections without adding
# false ones; it only ever overrides a lower-fidelity global pick, so it cannot
# weaken an existing detection.
_CORNER_PROMOTE_NCC = 0.85
# Bottom-right corner side for the promotion search, as a fraction of the
# image's short side, clamped to an absolute pixel band. Relative so the corner
# stays a true corner at every scale: a fixed 256 px is a genuine corner on a
# large image but covers ~70% of a small portrait, where a busy real photo can
# then raw-match the star template at ~0.81 (only 0.04 below the promote gate).
# Scaling the side down on small images drops that worst case to ~0.69, while
# the upper clamp stops it ballooning on huge images (more corner area = more
# random texture to false-match -- a real photo reached ~0.83 at 512 px). The
# Gemini sparkle sits ~60-160 px from the corner (fixed margins, not
# proportional), and the [96, 384] band covers that at every measured size.
_CORNER_PROMOTE_FRAC = 0.20
_CORNER_PROMOTE_MIN = 96
_CORNER_PROMOTE_MAX = 384
# Number of top size-weighted spatial candidates scored by full fusion before one
# is selected. The single size-weighted argmax can bury a genuine mid-size sparkle
# under a LARGER, lower-fidelity shape match (the 256->512 search-widening
# regression: a real corner sparkle at raw ~0.77 lost to a decoy at raw ~0.63).
# Scoring the top-K by gradient-bearing fusion rescues it. Top-K (NOT the raw-NCC
# argmax) keeps the tiny-patch suppression intact: a coincidental 16 px match never
# ranks in the size-weighted top-K, so widening selection cannot add a false
# positive on non-Gemini content (verified on the doubao/jimeng visible corpora).
_SELECT_TOPK = 3
_MASK_ALPHA = 0.04
_MASK_DILATE_FRAC = 0.18
def __init__(self, logo_value: float = 255.0) -> None:
"""Initialize the engine with embedded alpha maps.
Args:
logo_value: The logo brightness value (default 255.0 = white).
"""
self.logo_value = logo_value
# Load embedded background captures
bg_small = _load_embedded_asset("gemini_bg_48.png")
bg_large = _load_embedded_asset("gemini_bg_96.png")
# Ensure correct sizes
if bg_small.shape[:2] != (48, 48):
bg_small = cv2.resize(bg_small, (48, 48), interpolation=cv2.INTER_AREA)
if bg_large.shape[:2] != (96, 96):
bg_large = cv2.resize(bg_large, (96, 96), interpolation=cv2.INTER_AREA)
# Calculate alpha maps
self._alpha_small = _calculate_alpha_map(bg_small)
self._alpha_large = _calculate_alpha_map(bg_large)
# Per-scale resized templates are constant (``_alpha_large`` never changes),
# so precompute the whole fixed 16..118 ladder once: ``_scan_scales`` runs it on
# every image (twice -- global + corner), and re-``resize``-ing the 96x96 source
# each time is pure repeated work. Prebuilt (not lazy) so the dict is read-only
# after construction and safe to share across threads via the module singleton.
self._alpha_small = _calculate_alpha_map(_load_capture("gemini_bg_48.png", 48))
self._alpha_large = _calculate_alpha_map(_load_capture("gemini_bg_96.png", 96))
self._tmpl_cache: dict[int, NDArray[Any]] = {
scale: cv2.resize(self._alpha_large, (scale, scale), interpolation=cv2.INTER_AREA)
for scale in _TEMPLATE_SCALES
side: cv2.resize(self._alpha_large, (side, side), interpolation=cv2.INTER_AREA) for side in _TEMPLATE_SCALES
}
logger.debug(
"Alpha maps loaded: small=%s, large=%s",
self._alpha_small.shape,
self._alpha_large.shape,
)
def get_alpha_map(self, size: WatermarkSize) -> NDArray[Any]:
"""Get the base alpha map for a specific standard size."""
if size == WatermarkSize.SMALL:
return self._alpha_small
return self._alpha_large
return self._alpha_small if size is WatermarkSize.SMALL else self._alpha_large
def get_interpolated_alpha(self, size_px: int) -> NDArray[Any]:
"""Create an interpolated alpha map dynamically scaled from the high-res 96x96 base."""
source = self._alpha_large
if size_px == source.shape[1]:
return source.copy()
interp = cv2.INTER_LINEAR if size_px > source.shape[1] else cv2.INTER_AREA
return cv2.resize(source, (size_px, size_px), interpolation=interp)
# ── Detection ────────────────────────────────────────────────────
if size_px == self._alpha_large.shape[1]:
return self._alpha_large.copy()
method = cv2.INTER_LINEAR if size_px > self._alpha_large.shape[1] else cv2.INTER_AREA
return cv2.resize(self._alpha_large, (size_px, size_px), interpolation=method)
def _scan_scales(self, gray: NDArray[Any]) -> Iterator[tuple[int, float, tuple[int, int]]]:
"""Yield ``(scale, max_ncc, max_loc)`` for the alpha template matched at each scale.
Shared multi-scale ``TM_CCOEFF_NORMED`` primitive over a normalized [0, 1]
grayscale region, used by both the size-weighted global search in
``detect_watermark`` and the raw-NCC corner pass in ``_corner_promote`` --
each applies its own scoring/argmax to the yielded values. The 96x96
``_alpha_large`` is the high-quality source downscaled per scale; the range
covers aggressively downscaled to slightly upscaled logos.
"""
for scale in _TEMPLATE_SCALES:
if scale > gray.shape[0] or scale > gray.shape[1]:
"""Yield the strongest normalized template match at every usable scale."""
height, width = gray.shape[:2]
for side, template in self._tmpl_cache.items():
if side > height or side > width:
continue
match_res = cv2.matchTemplate(gray, self._tmpl_cache[scale], cv2.TM_CCOEFF_NORMED)
_, max_val, _, max_loc = cv2.minMaxLoc(match_res)
yield scale, float(max_val), max_loc
response = cv2.matchTemplate(gray, template, cv2.TM_CCOEFF_NORMED)
_minimum, maximum, _min_location, max_location = cv2.minMaxLoc(response)
yield side, float(maximum), max_location
def _global_candidates(self, image: NDArray[Any]) -> list[_Candidate]:
height, width = image.shape[:2]
search_side = min(height, width, 512)
origin_x, origin_y = width - search_side, height - search_side
gray = _gray_float(image[origin_y:height, origin_x:width])
ranked = sorted(
(
(
score * min(1.0, (side / 96.0) ** 0.5),
_Candidate(side, origin_x + location[0], origin_y + location[1], score),
)
for side, score, location in self._scan_scales(gray)
),
key=lambda item: (item[0], item[1].scale, item[1].spatial, item[1].x, item[1].y),
reverse=True,
)
selected: list[_Candidate] = []
for _weighted, candidate in ranked:
if any(_overlaps(candidate, prior) for prior in selected):
continue
selected.append(candidate)
if len(selected) == self._SELECT_TOPK:
break
return selected
def _score_candidate(self, image: NDArray[Any], candidate: _Candidate) -> _Candidate:
if candidate.spatial < 0.25:
return candidate
gradient, variance = self._grad_var_scores(image, candidate.scale, candidate.x, candidate.y)
return _Candidate(candidate.scale, candidate.x, candidate.y, candidate.spatial, gradient, variance)
def detect_watermark(
self,
@@ -306,241 +199,94 @@ class GeminiEngine:
*,
trust_provenance: bool = False,
) -> DetectionResult:
"""Detect Gemini watermark using multi-scale Snap Engine logic (ported from C++ vendor algorithm).
``trust_provenance`` signals that external metadata already proves this is a
Google generation (C2PA issuer "Google"/"Gemini"). The false-positive gate
exists only to reject content that shape-matches the sparkle on NON-Google
images (Doubao text, ornate corners); when provenance confirms Google, that
gate would demote a genuine sparkle the vendor moved/re-rendered (bigger,
lighter, shifted), so it is skipped. The caller (registry) still applies the
relaxed provenance trust gate to the returned confidence."""
"""Return the strongest sparkle-shaped bottom-right candidate."""
result = DetectionResult()
if image is None or image.size == 0:
return result
# Normalize to 3-channel BGR: the multi-scale search tolerates grayscale, but
# the FP-gate / alpha-gain helpers (_core_and_bg) reduce over axis=2 and would
# crash on a 2D/BGRA input reaching this public entry point (e.g. via the
# registry detect adapter or the library API).
image = image_io.to_bgr(image)
h, w = image.shape[:2]
base_size = force_size or get_watermark_size(w, h)
result.size = base_size
# Dynamically search bottom-right corner. 512 covers up to 512px from the
# corner -- enough for known Gemini margin variations (standard: 64+96=160px;
# observed variants up to ~300px). 256 was too tight and caused misses.
search_size = int(min(min(w, h), 512))
sx1 = max(0, w - search_size)
sy1 = max(0, h - search_size)
search_region = image[sy1:h, sx1:w]
if len(search_region.shape) == 3 and search_region.shape[2] >= 3:
gray_sr = cv2.cvtColor(search_region, cv2.COLOR_BGR2GRAY)
else:
gray_sr = search_region.copy()
gray_sr_f = gray_sr.astype(np.float32) / 255.0
# Phase 1 & 2: multi-scale spatial NCC search. The size weight (mimicking the
# C++ vendor weight) overcomes the NCC bias toward tiny patches, but its single
# argmax can bury a genuine mid-size sparkle under a LARGER, lower-fidelity
# shape match (the 256->512 search-widening regression). So score the top-K
# size-weighted candidates by the FULL fusion and keep the highest -- the
# gradient term separates a true white sparkle from a shape-only decoy. See
# _SELECT_TOPK for why top-K (not the raw-NCC argmax) preserves tiny-patch
# suppression and so cannot add a false positive on non-Gemini content.
scored: list[tuple[float, int, int, int, float]] = [] # (adj, scale, raw, x, y)
for scale, max_val, max_loc in self._scan_scales(gray_sr_f):
adj_val = max_val * min(1.0, (scale / 96.0) ** 0.5)
scored.append((adj_val, scale, max_val, sx1 + max_loc[0], sy1 + max_loc[1]))
scored.sort(reverse=True)
# Top-K candidates at distinct locations (NMS: drop a lower-ranked match that
# overlaps an already-kept one -- the same sparkle matches at adjacent scales).
candidates: list[tuple[int, int, int, float]] = []
for _adj, scale, raw, x, y in scored:
if any(
abs(x - px) < 0.5 * max(scale, ps) and abs(y - py) < 0.5 * max(scale, ps)
for ps, px, py, _ in candidates
):
continue
candidates.append((scale, x, y, raw))
if len(candidates) >= self._SELECT_TOPK:
break
# Corner promotion: a near-perfect small bottom-right sparkle the size weight
# buries even below the top-K (see _CORNER_PROMOTE_NCC) -- add it as a candidate.
promoted = self._corner_promote(image, candidates[0][3] if candidates else -1.0)
source = image_io.to_bgr(image)
height, width = source.shape[:2]
result.size = force_size or get_watermark_size(width, height)
candidates = self._global_candidates(source)
promoted = self._corner_promote(source, candidates[0].spatial if candidates else -1.0)
if promoted is not None:
candidates.append(promoted)
# No candidate at any scale: the search region is smaller than the 16px template
# floor (an image whose short side is < 16px), so nothing is detectable. Return
# the empty (detected=False) result rather than dereferencing candidates[0].
candidates.append(_Candidate(promoted[0], promoted[1], promoted[2], promoted[3]))
if not candidates:
return result
# Select the candidate with the highest full-fusion confidence (pre-FP-gate).
best_scale, pos_x, pos_y, best_raw_ncc = candidates[0]
grad_score, var_score, best_fused = 0.0, 0.0, -1.0
for c_scale, c_x, c_y, c_raw in candidates:
if c_raw < 0.25:
c_grad, c_var, c_fused = 0.0, 0.0, max(0.0, c_raw * 0.5)
else:
c_grad, c_var = self._grad_var_scores(image, c_scale, c_x, c_y)
c_fused = c_raw * 0.50 + c_grad * 0.30 + c_var * 0.20
if c_fused > best_fused:
best_fused = c_fused
best_scale, pos_x, pos_y = c_scale, c_x, c_y
best_raw_ncc, grad_score, var_score = c_raw, c_grad, c_var
best = max((self._score_candidate(source, candidate) for candidate in candidates), key=lambda item: item.fused)
result.region = (best.x, best.y, best.scale, best.scale)
result.spatial_score = float(best.spatial)
result.gradient_score = float(best.gradient)
result.variance_score = float(best.variance)
result.region = (pos_x, pos_y, best_scale, best_scale)
result.spatial_score = float(best_raw_ncc)
result.gradient_score = float(grad_score)
result.variance_score = float(var_score)
if result.spatial_score < 0.25:
result.confidence = float(max(0.0, result.spatial_score * 0.5))
return result
# ── Fusion ───────────────────────────────────────────────────
# best_fused is the selected candidate's spatial*0.5 + grad*0.3 + var*0.2.
confidence = best_fused
# False-positive gate: a low-confidence match that shows NEITHER real-sparkle
# signature is a content false positive, not a white sparkle overlay. A real
# sparkle proves itself by a bright core (high core-ring margin, on dark/mid
# backgrounds) OR a crisp star silhouette (high gradient NCC, on any background
# incl. bright). Demote when both are weak -- this catches the dark/mid no-core
# FP (low margin) AND the bright-background smooth-blob FP (high margin but low
# gradient), which the margin check alone misses. See _SPARKLE_FP_GRAD.
if confidence < self._SPARKLE_FP_CONF and not trust_provenance:
alpha = self.get_interpolated_alpha(best_scale)
pos = (pos_x, pos_y)
margin = self._core_ring_margin(image, alpha, pos)
low_margin = margin is not None and margin < self._SPARKLE_FP_MARGIN
low_grad = grad_score < self._SPARKLE_FP_GRAD
if low_margin or low_grad:
# White-core rescue: a real faint sparkle clears the trust confidence,
# has a bright core (not low_margin), and a near-WHITE core -- unlike the
# colored-corner FP the low-grad demotion targets. See _SPARKLE_WHITE_SAT.
core_sat = self._core_saturation(image, alpha, pos)
white_core = not low_margin and core_sat is not None and core_sat <= self._SPARKLE_WHITE_SAT
if not (confidence >= self._SPARKLE_KEEP_CONF and white_core):
logger.debug(
"Sparkle FP gate: conf=%.3f, margin=%s, grad=%.3f, core_sat=%s; demoting.",
confidence,
f"{margin:.1f}" if margin is not None else "n/a",
grad_score,
f"{core_sat:.2f}" if core_sat is not None else "n/a",
)
confidence = min(confidence, 0.30)
result.confidence = float(max(0.0, min(1.0, confidence)))
confidence = best.fused
if best.spatial >= 0.25 and confidence < self._SPARKLE_FP_CONF and not trust_provenance:
confidence = self._apply_false_positive_gate(source, best, confidence)
result.confidence = float(np.clip(confidence, 0.0, 1.0))
result.detected = result.confidence >= 0.35
logger.debug(
"Detection: spatial=%.3f, grad=%.3f, var=%.3f → conf=%.3f (%s)",
result.spatial_score,
result.gradient_score,
var_score,
result.confidence,
"DETECTED" if result.detected else "not detected",
)
return result
def _grad_var_scores(
self,
image: NDArray[Any],
scale: int,
pos_x: int,
pos_y: int,
) -> tuple[float, float]:
"""Return ``(gradient_score, variance_score)`` for a candidate sparkle.
Factored out of ``detect_watermark`` so each top-K candidate can be scored by
the full fusion before one is selected. The gradient NCC correlates
Sobel-magnitude maps (shape fidelity, contrast-robust); the variance score
rewards a flat overlay region against the row band above it.
"""
h, w = image.shape[:2]
x1, y1 = pos_x, pos_y
x2, y2 = min(w, x1 + scale), min(h, y1 + scale)
region = image[y1:y2, x1:x2]
gray_region = cv2.cvtColor(region, cv2.COLOR_BGR2GRAY) if region.ndim == 3 and region.shape[2] >= 3 else region
gray_f = gray_region.astype(np.float32) / 255.0
alpha_region = self.get_interpolated_alpha(scale)[: y2 - y1, : x2 - x1]
# ── Gradient NCC ──
img_gmag = cv2.magnitude(
cv2.Sobel(gray_f, cv2.CV_32F, 1, 0, ksize=3), cv2.Sobel(gray_f, cv2.CV_32F, 0, 1, ksize=3)
def _apply_false_positive_gate(self, image: NDArray[Any], candidate: _Candidate, confidence: float) -> float:
alpha = self.get_interpolated_alpha(candidate.scale)
position = (candidate.x, candidate.y)
margin = self._core_ring_margin(image, alpha, position)
low_margin = margin is not None and margin < self._SPARKLE_FP_MARGIN
low_gradient = candidate.gradient < self._SPARKLE_FP_GRAD
if not low_margin and not low_gradient:
return confidence
saturation = self._core_saturation(image, alpha, position)
neutral_core = not low_margin and saturation is not None and saturation <= self._SPARKLE_WHITE_SAT
if confidence >= self._SPARKLE_KEEP_CONF and neutral_core:
return confidence
logger.debug(
"Sparkle candidate demoted: confidence=%.3f, margin=%s, gradient=%.3f, saturation=%s",
confidence,
margin,
candidate.gradient,
saturation,
)
alpha_gmag = cv2.magnitude(
cv2.Sobel(alpha_region, cv2.CV_32F, 1, 0, ksize=3), cv2.Sobel(alpha_region, cv2.CV_32F, 0, 1, ksize=3)
return min(confidence, 0.30)
def _grad_var_scores(self, image: NDArray[Any], scale: int, pos_x: int, pos_y: int) -> tuple[float, float]:
height, width = image.shape[:2]
x2, y2 = min(width, pos_x + scale), min(height, pos_y + scale)
region = image[pos_y:y2, pos_x:x2]
gray = _gray_float(region)
alpha = self.get_interpolated_alpha(scale)[: y2 - pos_y, : x2 - pos_x]
image_edges = cv2.magnitude(
cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=3),
cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=3),
)
_, grad_score, _, _ = cv2.minMaxLoc(cv2.matchTemplate(img_gmag, alpha_gmag, cv2.TM_CCOEFF_NORMED))
# ── Variance ──
var_score = 0.0
ref_h = min(y1, scale)
if ref_h > 8:
ref_region = image[y1 - ref_h : y1, x1:x2]
gray_ref = cv2.cvtColor(ref_region, cv2.COLOR_BGR2GRAY) if ref_region.ndim == 3 else ref_region
_, s_wm = cv2.meanStdDev(gray_region)
_, s_ref = cv2.meanStdDev(gray_ref)
if s_ref[0][0] > 5.0:
var_score = max(0.0, min(1.0, 1.0 - (s_wm[0][0] / s_ref[0][0])))
return float(grad_score), float(var_score)
def _corner_promote(
self,
image: NDArray[Any],
current_raw_ncc: float,
) -> tuple[int, int, int, float] | None:
"""Search the bottom-right corner for a very-high-fidelity sparkle match.
Returns ``(scale, x, y, raw_ncc)`` when the corner holds a match with raw
NCC >= ``_CORNER_PROMOTE_NCC`` that beats the global pick's ``current_raw_ncc``,
else None. Used to rescue a small sparkle that the size weight buried under
a larger, lower-fidelity match elsewhere. See ``_CORNER_PROMOTE_NCC`` and
``_CORNER_PROMOTE_FRAC`` for the corner sizing.
"""
h, w = image.shape[:2]
side = max(
self._CORNER_PROMOTE_MIN, min(self._CORNER_PROMOTE_MAX, round(min(w, h) * self._CORNER_PROMOTE_FRAC))
alpha_edges = cv2.magnitude(
cv2.Sobel(alpha, cv2.CV_32F, 1, 0, ksize=3),
cv2.Sobel(alpha, cv2.CV_32F, 0, 1, ksize=3),
)
cs = int(min(min(w, h), side))
cx1, cy1 = max(0, w - cs), max(0, h - cs)
corner = image[cy1:h, cx1:w]
gray = cv2.cvtColor(corner, cv2.COLOR_BGR2GRAY) if corner.ndim == 3 and corner.shape[2] >= 3 else corner
gray = gray.astype(np.float32) / 255.0
response = cv2.matchTemplate(image_edges, alpha_edges, cv2.TM_CCOEFF_NORMED)
_minimum, gradient, _min_location, _max_location = cv2.minMaxLoc(response)
best_raw = -1.0
best_scale = 0
best_loc = (0, 0)
for scale, max_val, max_loc in self._scan_scales(gray):
if max_val > best_raw:
best_raw = max_val
best_scale = scale
best_loc = max_loc
variance = 0.0
reference_height = min(pos_y, scale)
if reference_height > 8:
reference = image[pos_y - reference_height : pos_y, pos_x:x2]
reference_gray = cv2.cvtColor(reference, cv2.COLOR_BGR2GRAY) if reference.ndim == 3 else reference
_mean, region_std = cv2.meanStdDev((gray * 255.0).astype(np.uint8))
_reference_mean, reference_std = cv2.meanStdDev(reference_gray)
if reference_std[0][0] > 5.0:
variance = float(np.clip(1.0 - region_std[0][0] / reference_std[0][0], 0.0, 1.0))
return float(gradient), variance
if best_raw >= self._CORNER_PROMOTE_NCC and best_raw > current_raw_ncc:
return best_scale, cx1 + best_loc[0], cy1 + best_loc[1], float(best_raw)
return None
# ── Removal ──────────────────────────────────────────────────────
# Footprint mask for the localize -> fill removal path. The mask must cover the
# WHOLE sparkle including its faint semi-transparent halo, not just the bright
# core, or the fill leaves a visible ring. Threshold the captured alpha low
# (>_MASK_ALPHA catches the halo the core-only 0.10 misses) then dilate by a
# sparkle-relative margin so alignment slop and the outermost halo are absorbed.
_MASK_ALPHA = 0.04
_MASK_DILATE_FRAC = 0.18 # dilation radius as a fraction of the sparkle scale
def _corner_promote(self, image: NDArray[Any], current_raw_ncc: float) -> tuple[int, int, int, float] | None:
height, width = image.shape[:2]
desired = round(min(width, height) * self._CORNER_PROMOTE_FRAC)
side = min(min(width, height), max(self._CORNER_PROMOTE_MIN, min(self._CORNER_PROMOTE_MAX, desired)))
origin_x, origin_y = width - side, height - side
matches = self._scan_scales(_gray_float(image[origin_y:height, origin_x:width]))
best = max(matches, key=lambda item: item[1], default=None)
if best is None or best[1] < self._CORNER_PROMOTE_NCC or best[1] <= current_raw_ncc:
return None
return best[0], origin_x + best[2][0], origin_y + best[2][1], float(best[1])
def footprint_mask(
self,
@@ -550,50 +296,37 @@ class GeminiEngine:
dilate: int | None = None,
region: tuple[int, int, int, int] | None = None,
) -> NDArray[Any] | None:
"""Full-frame uint8 mask (255 = sparkle) of the sparkle footprint, for the
shared fill removal path (cv2 / MI-GAN / LaMa), or None.
The footprint is the interpolated captured alpha at the detected scale,
thresholded LOW so the faint halo is included, then dilated by a
sparkle-relative margin. When ``force`` and nothing is detected, falls back to
the default sparkle slot for the image size (the ``--no-detect`` path).
``region`` is the already-resolved ``(x, y, scale)`` from the caller's detection
(the registry passes the decision's provenance-aware region). When given, the
mask is built from it directly WITHOUT a second internal detect -- otherwise a
provenance/assume-relaxed sparkle would be re-demoted by the strict re-detect and
yield no mask (reported-removed-but-unchanged). Absent ``region``, direct callers
keep the detect-then-force behavior.
"""
"""Build a full-frame mask from a resolved or newly detected sparkle."""
if image is None or image.size == 0:
return None # guard before to_bgr (cvtColor raises on an empty Mat); mirror detect_watermark
image = image_io.to_bgr(image)
h, w = image.shape[:2]
return None
source = image_io.to_bgr(image)
height, width = source.shape[:2]
if region is not None:
x, y, scale = region[0], region[1], region[2]
x, y, scale = region[:3]
else:
det = self.detect_watermark(image)
if det.detected:
x, y, scale = det.region[0], det.region[1], det.region[2]
detection = self.detect_watermark(source)
if detection.detected:
x, y, scale = detection.region[:3]
elif force:
cfg = get_watermark_config(w, h)
x, y = cfg.get_position(w, h)
scale = cfg.logo_size
config = get_watermark_config(width, height)
x, y = config.get_position(width, height)
scale = config.logo_size
else:
return None
alpha = self.get_interpolated_alpha(scale)
fp = self._footprint_indices(alpha, (x, y), image.shape)
if fp is None:
placed = self._footprint_indices(self.get_interpolated_alpha(scale), (x, y), source.shape)
if placed is None:
return None
aroi, (y1, y2, x1, x2) = fp
sil = (aroi > self._MASK_ALPHA).astype(np.uint8) * 255
if int((sil > 0).sum()) == 0:
alpha, (y1, y2, x1, x2) = placed
silhouette = (alpha > self._MASK_ALPHA).astype(np.uint8) * 255
if not silhouette.any():
return None
mask = np.zeros((h, w), np.uint8)
mask[y1:y2, x1:x2] = sil
d = dilate if dilate is not None else max(13, int(scale * self._MASK_DILATE_FRAC))
if d > 0:
mask = cv2.dilate(mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * d + 1, 2 * d + 1)))
mask = np.zeros((height, width), dtype=np.uint8)
mask[y1:y2, x1:x2] = silhouette
radius = dilate if dilate is not None else max(13, int(scale * self._MASK_DILATE_FRAC))
if radius > 0:
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * radius + 1, 2 * radius + 1))
mask = cv2.dilate(mask, kernel)
return mask
def _footprint_indices(
@@ -602,21 +335,35 @@ class GeminiEngine:
position: tuple[int, int],
image_shape: tuple[int, ...],
) -> tuple[NDArray[Any], tuple[int, int, int, int]] | None:
"""Return (alpha_roi, (y1, y2, x1, x2)) for the placed footprint, or None.
Shared by the over-subtraction test and the inpaint mask so both operate on
exactly the same clipped, in-bounds region.
"""
x, y = position
ah, aw = alpha_map.shape[:2]
ih, iw = image_shape[:2]
alpha_height, alpha_width = alpha_map.shape[:2]
image_height, image_width = image_shape[:2]
x1, y1 = max(0, x), max(0, y)
x2, y2 = min(iw, x + aw), min(ih, y + ah)
x2, y2 = min(image_width, x + alpha_width), min(image_height, y + alpha_height)
if x1 >= x2 or y1 >= y2:
return None
ax1, ay1 = x1 - x, y1 - y
alpha_roi = alpha_map[ay1 : ay1 + (y2 - y1), ax1 : ax1 + (x2 - x1)]
return alpha_roi, (y1, y2, x1, x2)
alpha_x, alpha_y = x1 - x, y1 - y
clipped = alpha_map[alpha_y : alpha_y + y2 - y1, alpha_x : alpha_x + x2 - x1]
return clipped, (y1, y2, x1, x2)
def _core_mask_and_box(
self,
image: NDArray[Any],
alpha_map: NDArray[Any],
position: tuple[int, int],
) -> tuple[NDArray[Any], NDArray[Any], tuple[int, int, int, int], float] | None:
placed = self._footprint_indices(alpha_map, position, image.shape)
if placed is None:
return None
alpha, bounds = placed
peak = float(alpha.max())
if peak < 0.2:
return None
core = alpha >= peak * self._CORE_ALPHA_FRAC
if not core.any():
return None
y1, y2, x1, x2 = bounds
return core, image[y1:y2, x1:x2], bounds, peak
def _core_and_bg(
self,
@@ -624,41 +371,22 @@ class GeminiEngine:
alpha_map: NDArray[Any],
position: tuple[int, int],
) -> tuple[float, float, float] | None:
"""Return ``(core_obs, bg, a_cap)`` for the placed sparkle, or None.
``core_obs`` is the bright-core brightness (75th pct over the high-alpha
core), ``bg`` the local background ring median, ``a_cap`` the captured peak
alpha. Shared by the alpha-gain estimate and the false-positive margin gate.
None when the footprint or the background ring cannot be sampled.
"""
placed = self._footprint_indices(alpha_map, position, image.shape)
if placed is None:
sample = self._core_mask_and_box(image, alpha_map, position)
if sample is None:
return None
alpha_roi, (y1, y2, x1, x2) = placed
a_cap = float(alpha_roi.max())
if a_cap < 0.2:
return None
core = alpha_roi >= a_cap * self._CORE_ALPHA_FRAC
if not bool(core.any()):
return None
# Convert only the footprint+ring crop to gray, not the whole image: every
# sample below lives inside the ring box, so a full-image mean is wasted work
# that scales with resolution (~70 ms on a 12 MP image, recomputed for both
# the alpha-gain estimate and the over-subtraction gate). The crop is sized by
# the footprint, so this is O(footprint^2) regardless of image size.
ih, iw = image.shape[:2]
pad = int((x2 - x1) * 0.7)
ry1, ry2 = max(0, y1 - pad), min(ih, y2 + pad)
rx1, rx2 = max(0, x1 - pad), min(iw, x2 + pad)
ring = image[ry1:ry2, rx1:rx2].astype(np.float32).mean(axis=2)
# Footprint box expressed in ring-crop coordinates.
core, _box, (y1, y2, x1, x2), peak = sample
height, width = image.shape[:2]
padding = int((x2 - x1) * 0.7)
ry1, ry2 = max(0, y1 - padding), min(height, y2 + padding)
rx1, rx2 = max(0, x1 - padding), min(width, x2 + padding)
luminance = image[ry1:ry2, rx1:rx2].astype(np.float32).mean(axis=2)
fy1, fy2, fx1, fx2 = y1 - ry1, y2 - ry1, x1 - rx1, x2 - rx1
core_obs = float(np.percentile(ring[fy1:fy2, fx1:fx2][core], 75))
ring_mask = np.ones(ring.shape, dtype=bool)
ring_mask[fy1:fy2, fx1:fx2] = False
if int(ring_mask.sum()) < 10:
core_value = float(np.percentile(luminance[fy1:fy2, fx1:fx2][core], 75))
background = np.ones(luminance.shape, dtype=bool)
background[fy1:fy2, fx1:fx2] = False
if background.sum() < 10:
return None
return core_obs, float(np.median(ring[ring_mask])), a_cap
return core_value, float(np.median(luminance[background])), peak
def _core_ring_margin(
self,
@@ -666,14 +394,8 @@ class GeminiEngine:
alpha_map: NDArray[Any],
position: tuple[int, int],
) -> float | None:
"""Bright-core brightness minus the local background ring (gray levels).
A real white sparkle overlay lifts its core above the surroundings; a
shape-only NCC false positive on ornate/flat content does not. None when the
background ring cannot be sampled.
"""
cb = self._core_and_bg(image, alpha_map, position)
return None if cb is None else cb[0] - cb[1]
sample = self._core_and_bg(image, alpha_map, position)
return None if sample is None else sample[0] - sample[1]
def _core_saturation(
self,
@@ -681,57 +403,24 @@ class GeminiEngine:
alpha_map: NDArray[Any],
position: tuple[int, int],
) -> float | None:
"""Median color saturation of the sparkle core (0 = white/neutral, higher =
colored). A real Gemini sparkle is a white star, so its core is near-neutral;
a clean bright corner that shape-matches (sky, sun, a warm light) is colored,
so a high core saturation flags the false positive the brightness/gradient
gates miss. Samples the same high-alpha core pixels as :meth:`_core_and_bg`.
None when the footprint cannot be placed or the core is empty.
"""
placed = self._footprint_indices(alpha_map, position, image.shape)
if placed is None:
sample = self._core_mask_and_box(image, alpha_map, position)
if sample is None:
return None
alpha_roi, (y1, y2, x1, x2) = placed
a_cap = float(alpha_roi.max())
if a_cap < 0.2:
return None
core = alpha_roi >= a_cap * self._CORE_ALPHA_FRAC
box = image[y1:y2, x1:x2]
if box.shape[:2] != core.shape or not bool(core.any()):
return None
px = box[core].astype(np.float32) # (N, 3) BGR core pixels
hi = px.max(axis=1)
lo = px.min(axis=1)
return float(np.median((hi - lo) / (hi + 1.0)))
core, box, _bounds, _peak = sample
pixels = box[core].astype(np.float32)
brightest = pixels.max(axis=1)
darkest = pixels.min(axis=1)
return float(np.median((brightest - darkest) / (brightest + 1.0)))
@functools.lru_cache(maxsize=1)
def _shared_engine() -> GeminiEngine:
"""Process-wide default ``GeminiEngine`` singleton.
The engine holds only constant assets (embedded captures, alpha maps, the
precomputed template ladder) and takes the image as a method argument, so one
instance is reused across every ``detect_sparkle_confidence`` call instead of
reloading assets + recomputing alpha maps + rebuilding the template cache on
each of the ~34k images an ``identify`` batch scans. Output is identical."""
return GeminiEngine()
def detect_sparkle_confidence(image_path: Path, *, image: NDArray[Any] | None = None) -> float | None:
"""Visible-sparkle detection confidence for a file, for provenance use.
Loads the image with cv2 and runs :meth:`GeminiEngine.detect_watermark`.
Returns the NCC confidence in [0, 1], or None if the image cannot be read
(cv2 returns None for unsupported containers such as HEIC). Kept here so the
cv2 dependency stays in this module; callers apply their own threshold.
``image`` lets a caller that has already decoded the file (e.g. ``identify``
running several visible-mark detectors) pass the BGR array to avoid a second
full decode; when None the file is read from ``image_path``.
"""
from remove_ai_watermarks import image_io
img = image if image is not None else image_io.imread(image_path)
if img is None:
"""Return the local sparkle confidence, or None when decoding fails."""
decoded = image if image is not None else image_io.imread(image_path)
if decoded is None:
return None
return float(_shared_engine().detect_watermark(img).confidence)
return float(_shared_engine().detect_watermark(decoded).confidence)
src/remove_ai_watermarks/humanizer.py
+14 -21
View File
@@ -1,7 +1,7 @@
"""Post-processing filters for the cleaned output.
``apply_analog_humanizer`` injects film grain and chromatic aberration to defeat
digital AI-perfection classifiers (ported from NeuralBleach); ``unsharp_mask``
``apply_analog_humanizer`` injects film grain and chromatic aberration to reduce
overly uniform digital surfaces; ``unsharp_mask``
counters the soft, over-smoothed look that the diffusion pass leaves behind
(itself a common "this is AI" tell).
"""
@@ -16,10 +16,7 @@ from numpy.typing import NDArray
def apply_analog_humanizer(image: NDArray, grain_intensity: float = 4.0, chromatic_shift: int = 1) -> NDArray:
"""
Apply Analog Humanizer (film grain and chromatic aberration) to an image.
This simulates analog film imperfections to defeat digital AI perfection classifiers.
Ported from NeuralBleach.
Apply shared-luminance grain and a small lateral color offset.
Args:
image: BGR image as numpy array (uint8).
@@ -33,26 +30,22 @@ def apply_analog_humanizer(image: NDArray, grain_intensity: float = 4.0, chromat
if len(image.shape) != 3 or image.shape[2] != 3:
return image.copy()
# Split channels (OpenCV uses BGR)
# B = 0, G = 1, R = 2
# Translate the outer color channels without circular edge wrapping.
b, g, r = cv2.split(image)
# 1. Chromatic Aberration
# Shift R channel left, B channel right. np.roll is circular, so it wraps
# the opposite edge into a thin colored fringe at the L/R borders; replicate
# the original edge columns there to keep the intended offset interior-only.
# Clamp so the edge-replication slices below always have a source column: a shift
# >= width would leave them empty and crash the broadcast (r[:, -shift:] = (H, 0)).
shift = min(chromatic_shift, image.shape[1] - 1)
shift = min(max(0, chromatic_shift), max(0, image.shape[1] - 1))
if shift > 0:
r = np.roll(r, -shift, axis=1)
r[:, -shift:] = r[:, -shift - 1 : -shift]
b = np.roll(b, shift, axis=1)
b[:, :shift] = b[:, shift : shift + 1]
shifted_b = np.empty_like(b)
shifted_b[:, :shift] = b[:, :1]
shifted_b[:, shift:] = b[:, :-shift]
b = shifted_b
shifted_r = np.empty_like(r)
shifted_r[:, :-shift] = r[:, shift:]
shifted_r[:, -shift:] = r[:, -1:]
r = shifted_r
merged = cv2.merge((b, g, r))
# 2. Film Grain (Gaussian Noise)
if grain_intensity > 0:
img_f = merged.astype(np.float32)
noise = np.random.normal(0, grain_intensity, img_f.shape).astype(np.float32)
src/remove_ai_watermarks/identify.py
+13 -13
View File
@@ -25,6 +25,18 @@ import logging
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, cast
from remove_ai_watermarks._internal.c2pa import (
c2pa_info_from_manifest_store,
cbor_text_after,
extract_c2pa_info,
soft_binding_vendors_in,
)
from remove_ai_watermarks._internal.constants import (
C2PA_AI_TOOLS,
C2PA_AI_VENDORS,
C2PA_IDENTITY_AI_ORGS,
C2PA_ISSUERS,
)
from remove_ai_watermarks.metadata import (
AI_METADATA_KEYS,
AIGC_MARKERS,
@@ -46,18 +58,6 @@ from remove_ai_watermarks.metadata import (
xai_signature,
xai_signature_pair,
)
from remove_ai_watermarks.noai.c2pa import (
c2pa_info_from_manifest_store,
cbor_text_after,
extract_c2pa_info,
soft_binding_vendors_in,
)
from remove_ai_watermarks.noai.constants import (
C2PA_AI_TOOLS,
C2PA_AI_VENDORS,
C2PA_IDENTITY_AI_ORGS,
C2PA_ISSUERS,
)
from remove_ai_watermarks.watermark_registry import GEMINI_SPARKLE_TRUST_CONF
if TYPE_CHECKING:
@@ -822,7 +822,7 @@ def _identify_from_evidence(
issuers = [info["issuer"]] if info.get("issuer") else _issuers_in(head)
# Full AI generation (trainedAlgorithmicMedia) vs an AI-enhanced real photo
# (compositeWithTrainedAlgorithmicMedia). The structured kind is parsed once in
# noai.c2pa._populate_registry_fields (covers PNG + any container the c2pa-python
# _internal.c2pa._populate_registry_fields (covers PNG + any container the c2pa-python
# reader handles); fall back to a raw head scan for the non-PNG raw-blob path
# where extract_c2pa_info returns {}. Full generation wins when both appear.
c2pa_source_kind = info.get("ai_source_kind")
src/remove_ai_watermarks/invisible_engine.py
+4 -10
View File
@@ -1,7 +1,4 @@
"""Invisible watermark removal engine.
Wraps the vendored noai-watermark code for removing invisible AI watermarks
(SynthID, StableSignature, TreeRing) via diffusion-based regeneration.
"""Diffusion engine for regenerating images that carry invisible AI watermarks.
This module requires the 'gpu' extra dependencies:
uv pip install 'remove-ai-watermarks[diffusion]'
@@ -19,10 +16,10 @@ import warnings
from pathlib import Path
from typing import TYPE_CHECKING, Any
from .noai.watermark_profiles import (
from ._internal.watermark_profiles import (
DEFAULT_MODEL_ID as DEFAULT_SDXL_MODEL_ID,
)
from .noai.watermark_profiles import (
from ._internal.watermark_profiles import (
resolve_seed,
)
@@ -81,9 +78,6 @@ def _target_size(width: int, height: int, max_resolution: int, min_resolution: i
class InvisibleEngine:
"""Remove invisible AI watermarks using diffusion model regeneration.
Based on noai-watermark by mertizci:
https://github.com/mertizci/noai-watermark
The approach encodes the image into latent space, injects controlled noise
to break watermark patterns, and reconstructs via reverse diffusion.
"""
@@ -124,7 +118,7 @@ class InvisibleEngine:
residency. CUDA only.
"""
from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
effective_model = model_id or self.DEFAULT_MODEL_ID
src/remove_ai_watermarks/metadata.py
+15 -18
View File
@@ -1,7 +1,4 @@
"""AI metadata detection and removal.
Wraps the noai-watermark metadata handling for stripping AI-generation
metadata (EXIF, PNG text chunks, C2PA provenance) from images.
"""Detect and remove AI provenance metadata from image containers.
For metadata-only operations, the heavy ML dependencies are NOT required.
"""
@@ -104,7 +101,7 @@ IPTC_AI_MARKERS: tuple[bytes, ...] = (
# (Meta / Instagram / MidJourney) use ``trainedAlgorithmicMedia``. Including the bare
# token flagged clean procedural images as AI (is_ai=high + has_invisible_target=True ->
# a diffusion scrub of clean content), contradicting the c2pa layer, which sets
# source_type without ai_source for it (tests/test_noai.py::test_plain_algorithmic_media_not_flagged_ai).
# source_type without ai_source for it (tests/test_metadata_internals.py::test_plain_algorithmic_media_not_flagged_ai).
# It is not a substring of the trained/composite tokens, so its removal does not affect
# their detection.
@@ -218,7 +215,7 @@ def _is_ai_value(value: str) -> bool:
detection: NovelAI stamps a generic ``Title``/``Source`` text chunk (an
AI-shaped value under a non-AI key) that ``_is_ai_key`` alone would keep.
"""
from remove_ai_watermarks.noai.constants import AI_GENERATOR_TOKENS
from remove_ai_watermarks._internal.constants import AI_GENERATOR_TOKENS
value_lower = value.lower()
return any(token in value_lower for token in AI_GENERATOR_TOKENS)
@@ -310,7 +307,7 @@ def _scan_head_impl(image_path: Path, size: int) -> bytes:
with open(image_path, "rb") as f:
head = f.read(size)
# Lazy import: isobmff imports this module's constants at top level.
from remove_ai_watermarks.noai import isobmff
from remove_ai_watermarks._internal import isobmff
if isobmff.is_isobmff(head):
region = isobmff.scan_c2pa_region(image_path)
@@ -348,7 +345,7 @@ def has_ai_metadata(image_path: Path) -> bool:
# Check C2PA — via the official c2pa-python reader first (spec-tracking, every
# container it supports), then a binary scan that also catches AVIF/HEIF/JPEG-XL
# containers and synthetic/partial blobs the validator rejects.
from remove_ai_watermarks.noai.c2pa import read_manifest_store_json
from remove_ai_watermarks._internal.c2pa import read_manifest_store_json
if read_manifest_store_json(image_path) is not None:
return True
@@ -466,7 +463,7 @@ def aigc_label(image_path: Path) -> dict[str, str] | None:
# in ``moov.udta.meta.keys`` and points to a raw JSON value in ``ilst``.
# Read it through the bounded box walker so a tail ``moov`` after a large
# ``mdat`` is found without loading or scanning the media payload.
from remove_ai_watermarks.noai.isobmff import tc260_aigc_payloads
from remove_ai_watermarks._internal.isobmff import tc260_aigc_payloads
isobmff_candidates = tuple(payload.decode("utf-8", "replace") for payload in tc260_aigc_payloads(image_path))
if result := aigc_label_from_metadata(b"", isobmff_candidates):
@@ -475,7 +472,7 @@ def aigc_label(image_path: Path) -> dict[str, str] | None:
# Native MKV/WebM TC260 metadata: ``Segment.Tags.Tag.SimpleTag`` carries
# ``TagName=AIGC`` and the raw JSON in ``TagString``. The EBML walker seeks
# over clusters and reads only bounded metadata values.
from remove_ai_watermarks.noai.ebml import tc260_aigc_payloads as ebml_tc260_aigc_payloads
from remove_ai_watermarks._internal.ebml import tc260_aigc_payloads as ebml_tc260_aigc_payloads
ebml_candidates = tuple(payload.decode("utf-8", "replace") for payload in ebml_tc260_aigc_payloads(image_path))
if result := aigc_label_from_metadata(b"", ebml_candidates):
@@ -486,11 +483,11 @@ def aigc_label(image_path: Path) -> dict[str, str] | None:
# that could collide inside compressed video.
legacy_payloads: tuple[bytes, ...] = ()
if image_path.suffix.lower() == ".avi":
from remove_ai_watermarks.noai.riff import tc260_aigc_payloads as riff_tc260_aigc_payloads
from remove_ai_watermarks._internal.riff import tc260_aigc_payloads as riff_tc260_aigc_payloads
legacy_payloads = riff_tc260_aigc_payloads(image_path)
elif image_path.suffix.lower() == ".flv":
from remove_ai_watermarks.noai.flv import tc260_aigc_payloads as flv_tc260_aigc_payloads
from remove_ai_watermarks._internal.flv import tc260_aigc_payloads as flv_tc260_aigc_payloads
legacy_payloads = flv_tc260_aigc_payloads(image_path)
legacy_candidates = tuple(payload.decode("utf-8", "replace") for payload in legacy_payloads)
@@ -672,7 +669,7 @@ def synthid_source(image_path: Path) -> str | None:
Returns:
Comma-joined vendor name(s) (e.g. ``"OpenAI"``) or None.
"""
from remove_ai_watermarks.noai.c2pa import extract_c2pa_info, synthid_vendors_in
from remove_ai_watermarks._internal.c2pa import extract_c2pa_info, synthid_vendors_in
# PNG: the caBX chunk parser gives a clean, structured issuer.
vendors = extract_c2pa_info(image_path).get("synthid_vendors")
@@ -694,7 +691,7 @@ def synthid_source(image_path: Path) -> str | None:
def generator_from_metadata(candidates: Iterable[str], scan: bytes = b"") -> str | None:
"""Return a known AI generator from collected EXIF, PNG, or XMP values."""
from remove_ai_watermarks.noai.constants import AI_GENERATOR_TOKENS
from remove_ai_watermarks._internal.constants import AI_GENERATOR_TOKENS
creator_tools = (
match.group(1).decode("latin1", "replace")
@@ -844,7 +841,7 @@ def _ai_exif_targets(loaded: dict[str, Any]) -> list[tuple[str, int, bytes, str]
"""
import piexif
from remove_ai_watermarks.noai.constants import AI_GENERATOR_TOKENS
from remove_ai_watermarks._internal.constants import AI_GENERATOR_TOKENS
ifd0: dict[int, Any] = loaded.get("0th") or {}
ifde: dict[int, Any] = loaded.get("Exif") or {}
@@ -903,7 +900,7 @@ def get_ai_metadata(image_path: Path) -> dict[str, str]:
"""
from PIL import Image
from remove_ai_watermarks.noai.c2pa import extract_c2pa_info, soft_binding_vendors_in, synthid_verdict
from remove_ai_watermarks._internal.c2pa import extract_c2pa_info, soft_binding_vendors_in, synthid_verdict
result: dict[str, str] = {}
@@ -923,7 +920,7 @@ def get_ai_metadata(image_path: Path) -> dict[str, str]:
except Exception as exc:
logger.debug("PIL could not open %s for AI-metadata scan: %s", image_path, exc)
# C2PA manifest fields from the single canonical parser (noai/c2pa.py).
# C2PA manifest fields from the single canonical parser (_internal/c2pa.py).
c2pa = extract_c2pa_info(image_path)
for key in (
"c2pa_manifest",
@@ -1215,7 +1212,7 @@ def remove_ai_metadata(
# offset-preserving streaming path; images retain the in-memory item scrub
# needed for XMP/EXIF inside mdat/idat. Route the remaining formats by suffix
# OR by an ``ftyp`` content sniff.
from remove_ai_watermarks.noai.isobmff import (
from remove_ai_watermarks._internal.isobmff import (
blank_ai_exif_tokens,
blank_ai_xmp_packets,
blank_tc260_aigc_tags,
src/remove_ai_watermarks/noai/c2pa.py
-562
View File
@@ -1,562 +0,0 @@
"""C2PA (Coalition for Content Provenance and Authenticity) metadata handling.
Reading goes through the official c2pa-python ``Reader`` first (any container it
supports), via ``extract_c2pa_info`` / ``read_manifest_store_json``. The
hand-rolled PNG ``caBX`` JUMBF-chunk tools below (``has_c2pa_metadata`` /
``extract_c2pa_chunk`` / ``inject_c2pa_chunk`` and the ``_extract_c2pa_info_png``
fallback) cover raw-chunk extraction, re-injection, and the cases the validator
rejects (synthetic/partial blobs, a broken/absent wheel). Known issuers:
- Google Imagen
- Adobe Firefly
- Microsoft Designer
- OpenAI (ChatGPT, GPT-4o, Sora, DALL-E)
- Truepic (signing authority)
The fallback parser uses byte-level scanning — it does not validate JUMBF/CBOR
structure but reliably identifies known signatures, issuers, tools, and actions.
The vendor / source-type / SynthID / soft-binding registry scan
(``_populate_registry_fields``) is shared by both the reader and fallback paths.
"""
from __future__ import annotations
import contextlib
import functools
import json
import logging
import re
import struct
from pathlib import Path
from typing import Any, cast
from remove_ai_watermarks.noai.constants import (
C2PA_ACTIONS,
C2PA_AI_TOOLS,
C2PA_CHUNK_TYPE,
C2PA_ISSUERS,
C2PA_SIGNATURES,
C2PA_SOFT_BINDINGS,
PNG_SIGNATURE,
SYNTHID_C2PA_ISSUERS,
)
logger = logging.getLogger(__name__)
# Official C2PA reader (c2pa-python, a default dependency). It is the primary,
# spec-tracking manifest parser; the hand-rolled caBX/CBOR scanner below stays as
# a fallback for synthetic/partial blobs the validator rejects. The import is
# guarded so a partially-broken install degrades to the byte-scan rather than
# crashing the dependency-light identify path.
_C2paReader: Any = None
with contextlib.suppress(Exception): # broken/absent wheel -> byte-scan fallback
from c2pa import Reader as _C2paReader # pyright: ignore[reportMissingTypeStubs]
_C2PA_READER_AVAILABLE = _C2paReader is not None
def reader_available() -> bool:
"""True when the official c2pa-python Reader imported successfully."""
return _C2PA_READER_AVAILABLE
def read_manifest_store_json(image_path: Path) -> str | None:
"""Return the full C2PA manifest-store JSON for ``image_path``, or None.
Uses the official c2pa-python ``Reader`` (any container it supports: PNG,
JPEG, WebP, AVIF/HEIF, MP4, ...). Returns None when the reader is unavailable,
the file carries no parseable manifest, or parsing fails. The JSON is the
WHOLE store (every manifest plus ingredient manifests), matching the
whole-chunk semantics of the legacy byte scan -- an AI-source marker in a
parent/ingredient manifest (e.g. a ChatGPT edit of a Sora generation) is
still seen.
Memoized per (path, mtime): one identify/get_ai_metadata call invokes the
structured parser ~3 times on the same file, so the cache turns the repeated
crypto-validating reads into one.
"""
if not _C2PA_READER_AVAILABLE:
return None
try:
mtime = image_path.stat().st_mtime_ns
except OSError:
return _read_manifest_store_impl(str(image_path))
return _read_manifest_store_cached(str(image_path), mtime)
@functools.lru_cache(maxsize=8)
def _read_manifest_store_cached(path_str: str, _mtime_ns: int) -> str | None:
"""Cache shim: ``_mtime_ns`` is part of the key only (invalidates on change)."""
return _read_manifest_store_impl(path_str)
def _read_manifest_store_impl(path_str: str) -> str | None:
# try_create returns None when there is no manifest; a default Reader does no
# trust enforcement, so an untrusted signer still yields the manifest content
# (we report what is in the file, we do not gate on certificate trust).
try:
reader = _C2paReader.try_create(path_str)
except Exception as exc: # malformed manifest, unsupported container, etc.
logger.debug("c2pa Reader could not parse %s: %s", path_str, exc)
return None
if reader is None:
return None
try:
with reader:
return reader.json()
except Exception as exc: # pragma: no cover - reader opened but json() failed
logger.debug("c2pa Reader.json() failed on %s: %s", path_str, exc)
return None
def has_c2pa_metadata(image_path: Path) -> bool:
"""
Check if an image contains C2PA metadata.
Args:
image_path: Path to the image file.
Returns:
True if C2PA metadata is detected, False otherwise.
"""
image_path = Path(image_path)
if image_path.suffix.lower() != ".png":
return False
try:
with open(image_path, "rb") as f:
signature = f.read(8)
if signature != PNG_SIGNATURE:
return False
file_size = f.seek(0, 2)
f.seek(8)
while True:
chunk_header = f.read(8)
if len(chunk_header) < 8:
break
length = struct.unpack(">I", chunk_header[:4])[0]
chunk_type = chunk_header[4:8]
# Clamp the attacker-controlled 32-bit length to the bytes that
# actually remain, so a malformed huge length can't allocate GBs.
safe_length = max(0, min(length, file_size - f.tell()))
if chunk_type == C2PA_CHUNK_TYPE:
chunk_data = f.read(safe_length)
# Check for any C2PA signature
for sig in C2PA_SIGNATURES:
if sig in chunk_data:
return True
# Also check if chunk_data itself contains C2PA-like patterns
if b"jumb" in chunk_data.lower() or b"c2pa" in chunk_data.lower():
return True
f.read(4)
else:
f.seek(safe_length + 4, 1)
if chunk_type == b"IEND":
break
except Exception:
pass
return False
def _claim_generator_from_store(store: dict[str, Any]) -> str | None:
"""Structured claim-generator name from the active manifest of a store dict.
Prefers the top-level ``claim_generator`` string (Firefly: "Adobe_Firefly"),
falling back to the first ``claim_generator_info[].name`` (ChatGPT keys it
only there). isprintable() guards against odd binary-ish values.
"""
active = _active_manifest(store)
generator: Any = active.get("claim_generator")
if not (isinstance(generator, str) and generator):
info_list: list[Any] = active.get("claim_generator_info") or []
if info_list and isinstance(first := info_list[0], dict):
generator = cast("dict[str, Any]", first).get("name")
return generator if isinstance(generator, str) and generator and generator.isprintable() else None
def _active_manifest(store: dict[str, Any]) -> dict[str, Any]:
"""The active manifest dict from a manifest-store dict, or {} when absent."""
manifests: Any = store.get("manifests")
if not isinstance(manifests, dict):
return {}
active = cast("dict[str, Any]", manifests).get(store.get("active_manifest", ""))
return cast("dict[str, Any]", active) if isinstance(active, dict) else {}
def _info_from_store(store: dict[str, Any], store_bytes: bytes) -> dict[str, Any]:
"""Build normalized C2PA info from one parsed manifest store."""
c2pa_info: dict[str, Any] = {
"has_c2pa": True,
"type": "C2PA (Coalition for Content Provenance and Authenticity)",
"c2pa_manifest": f"C2PA manifest store ({len(store_bytes)} bytes)",
}
# The whole-store JSON carries every vendor / source-type / SynthID /
# soft-binding signature (across active + ingredient manifests), so the same
# registry scan that runs on the raw caBX chunk applies unchanged here.
_populate_registry_fields(store_bytes, c2pa_info)
if generator := _claim_generator_from_store(store):
c2pa_info["claim_generator"] = generator
sig: Any = _active_manifest(store).get("signature_info")
if isinstance(sig, dict) and (time := cast("dict[str, Any]", sig).get("time")):
c2pa_info["timestamp"] = str(time)
return c2pa_info
def _info_from_store_json(store_json: str) -> dict[str, Any]:
"""Build the C2PA info dict from a c2pa-python manifest-store JSON string."""
store_bytes = store_json.encode("utf-8")
try:
parsed: Any = json.loads(store_json)
except (ValueError, TypeError):
parsed = {}
store = cast("dict[str, Any]", parsed) if isinstance(parsed, dict) else {}
return _info_from_store(store, store_bytes)
def c2pa_info_from_manifest_store(store: str | dict[str, Any]) -> dict[str, Any]:
"""Build normalized C2PA evidence from an externally collected manifest store.
``store`` may be the JSON string returned by ``c2pa.Reader.json()`` or its
decoded dictionary form. This is the non-file-backed counterpart to
:func:`extract_c2pa_info`.
"""
if isinstance(store, dict):
parsed = store
try:
store_json = json.dumps(store, ensure_ascii=False)
except (TypeError, ValueError):
return {}
else:
store_json = store
try:
decoded: Any = json.loads(store_json)
except (TypeError, ValueError):
return {}
if not isinstance(decoded, dict):
return {}
parsed = cast("dict[str, Any]", decoded)
if not store_json or not parsed or parsed.get("error"):
return {}
return _info_from_store(parsed, store_json.encode("utf-8"))
def extract_c2pa_info(image_path: Path) -> dict[str, Any]:
"""
Extract C2PA metadata information from an image.
Uses the official c2pa-python reader first (any supported container), falling
back to the hand-rolled PNG caBX parser when the reader is unavailable or the
file carries no parseable manifest (synthetic/partial blobs).
Args:
image_path: Path to the image file.
Returns:
Dictionary containing C2PA metadata info, or {} when none is found.
"""
image_path = Path(image_path)
if (store_json := read_manifest_store_json(image_path)) is not None:
return _info_from_store_json(store_json)
return _extract_c2pa_info_png(image_path)
def _extract_c2pa_info_png(image_path: Path) -> dict[str, Any]:
"""Fallback PNG caBX parser, used when the c2pa-python reader finds nothing."""
c2pa_info: dict[str, Any] = {}
if not has_c2pa_metadata(image_path):
return c2pa_info
c2pa_info["has_c2pa"] = True
c2pa_info["type"] = "C2PA (Coalition for Content Provenance and Authenticity)"
try:
with open(image_path, "rb") as f:
signature = f.read(8)
if signature != PNG_SIGNATURE:
return c2pa_info
file_size = f.seek(0, 2)
f.seek(8)
while True:
chunk_header = f.read(8)
if len(chunk_header) < 8:
break
length = struct.unpack(">I", chunk_header[:4])[0]
chunk_type = chunk_header[4:8]
# Clamp the attacker-controlled 32-bit length to the bytes that
# actually remain, so a malformed huge length can't allocate GBs.
safe_length = max(0, min(length, file_size - f.tell()))
if chunk_type == C2PA_CHUNK_TYPE:
chunk_data = f.read(safe_length)
_parse_c2pa_chunk(chunk_data, c2pa_info)
f.read(4)
else:
f.seek(safe_length + 4, 1)
if chunk_type == b"IEND":
break
except Exception:
pass
return c2pa_info
def cbor_text_after(payload: bytes, key: bytes) -> str | None:
"""Return the CBOR text-string immediately following ``key`` in ``payload``.
Handles CBOR major-type 3 length prefixes: direct (0x60-0x77), 1-byte
(0x78 NN), and 2-byte (0x79 NN NN). This reads the actual encoded value, so
it avoids the byte-grabbing artifacts a loose regex produces (e.g. the
leading length byte showing up as ``fGPT-4o``).
"""
idx = payload.find(key)
if idx < 0:
return None
p = idx + len(key)
if p >= len(payload):
return None
head = payload[p]
if 0x60 <= head <= 0x77:
length, start = head - 0x60, p + 1
elif head == 0x78 and p + 1 < len(payload):
length, start = payload[p + 1], p + 2
elif head == 0x79 and p + 2 < len(payload):
length, start = (payload[p + 1] << 8) | payload[p + 2], p + 3
else:
return None
raw_str = payload[start : start + length]
try:
return raw_str.decode("utf-8")
except UnicodeDecodeError:
return raw_str.decode("latin1", errors="replace")
def synthid_verdict(vendors: str) -> str:
"""Human-readable SynthID-source verdict, shared by all callers."""
return f"likely present ({vendors} embeds SynthID with C2PA)"
def synthid_vendors_in(buffer: bytes) -> list[str]:
"""Return SynthID-using C2PA issuer names whose signature appears in ``buffer``.
Shared by the PNG caBX parser and the format-agnostic binary scan so both
apply the same SYNTHID_C2PA_ISSUERS rule against their respective bytes.
"""
return sorted({name for sig, name in C2PA_ISSUERS.items() if sig in buffer and sig in SYNTHID_C2PA_ISSUERS})
def soft_binding_vendors_in(buffer: bytes) -> list[str]:
"""Return forensic-watermark vendor names whose C2PA soft-binding ``alg``
identifier appears in ``buffer``.
A ``c2pa.soft-binding`` assertion names the watermark scheme that stamped the
pixels (Adobe TrustMark, Digimarc, Imatag, Steg.AI, ...). Shared by the PNG
caBX parser and the format-agnostic binary scan so both apply the same
C2PA_SOFT_BINDINGS rule against their respective bytes.
"""
return sorted({name for sig, name in C2PA_SOFT_BINDINGS.items() if sig in buffer})
def _populate_registry_fields(buf: bytes, c2pa_info: dict[str, Any]) -> bool:
"""Populate the registry-driven C2PA fields by scanning ``buf``.
Shared by the legacy caBX-chunk parser and the c2pa-python store-JSON path so
both produce an identical dict shape. ``buf`` is the raw manifest bytes for
the former and the manifest-store JSON (UTF-8) for the latter; the vendor /
tool / action / source-type / SynthID / soft-binding signatures appear in
both. Sets ``issuer``, ``ai_tool``, ``actions``, ``source_type``,
``synthid_vendors`` / ``synthid_watermark``, ``soft_binding_vendors`` /
``soft_binding`` when present and returns whether the source type is AI.
"""
if issuers := [name for sig, name in C2PA_ISSUERS.items() if sig in buf]:
c2pa_info["issuer"] = ", ".join(dict.fromkeys(issuers))
if ai_tools := [name for sig, name in C2PA_AI_TOOLS.items() if sig in buf]:
c2pa_info["ai_tool"] = ", ".join(dict.fromkeys(ai_tools))
if actions := [name for sig, name in C2PA_ACTIONS.items() if sig in buf]:
c2pa_info["actions"] = ", ".join(actions)
# Digital source type (matched anywhere in the store, including ingredient
# manifests -- a ChatGPT edit of a Sora generation carries the AI marker on
# the parent, not the active manifest).
# ``ai_source_kind`` is the structured generated-vs-enhanced split the caller
# branches on (full-frame scrub vs region-targeted clean); ``source_type`` is the
# human-readable form. The two byte strings are unambiguous:
# "compositeWithTrainedAlgorithmicMedia" capitalizes the inner "Trained", so a
# lowercase "trainedAlgorithmicMedia" match is standalone full generation, which
# wins when both appear (an edit chain).
ai_source = False
if b"trainedAlgorithmicMedia" in buf:
c2pa_info["source_type"] = "trainedAlgorithmicMedia (AI-generated)"
c2pa_info["ai_source_kind"] = "generated"
ai_source = True
elif b"compositeWithTrainedAlgorithmicMedia" in buf:
# Checked BEFORE bare ``algorithmicMedia``: a manifest can carry both tokens
# (an AI-enhanced composite with a procedural ingredient), and the bare-token
# branch would otherwise fire first and misclassify the AI composite as non-AI.
c2pa_info["source_type"] = "compositeWithTrainedAlgorithmicMedia (AI-enhanced)"
c2pa_info["ai_source_kind"] = "enhanced"
ai_source = True
elif b"algorithmicMedia" in buf:
c2pa_info["source_type"] = "algorithmicMedia"
# SynthID pixel-watermark proxy: a C2PA manifest from a SynthID-using
# vendor (Google/OpenAI) on AI-generated content implies an invisible
# SynthID watermark in the pixels (see SYNTHID_C2PA_ISSUERS).
synthid_vendors = synthid_vendors_in(buf)
if synthid_vendors and ai_source:
c2pa_info["synthid_vendors"] = synthid_vendors
c2pa_info["synthid_watermark"] = synthid_verdict(", ".join(synthid_vendors))
# Soft-binding: a forensic/third-party watermark vendor named in the
# manifest (Adobe TrustMark, Digimarc, ...), independent of the issuer.
soft_binding_vendors = soft_binding_vendors_in(buf)
if soft_binding_vendors:
c2pa_info["soft_binding_vendors"] = soft_binding_vendors
c2pa_info["soft_binding"] = ", ".join(soft_binding_vendors)
return ai_source
def _parse_c2pa_chunk(chunk_data: bytes, c2pa_info: dict[str, Any]) -> None:
"""Parse a raw caBX chunk payload and populate the info dictionary.
The fallback path, used when the official c2pa-python reader is unavailable
or rejects the file (synthetic/partial blobs, broken installs).
"""
c2pa_info["c2pa_manifest"] = f"C2PA manifest ({len(chunk_data)} bytes)"
_populate_registry_fields(chunk_data, c2pa_info)
# Claim generator and spec version: read the CBOR text-string values
# directly (regex byte-grabbing produced artifacts like ``fGPT-4o``).
# Guard with isprintable(): on some manifests (e.g. Microsoft Designer) the
# first ``name`` key precedes a binary field (a hash), not the generator
# string, which would otherwise surface as control-char garbage.
if (generator := cbor_text_after(chunk_data, b"name")) and generator.isprintable():
c2pa_info["claim_generator"] = generator
if (spec := cbor_text_after(chunk_data, b"specVersion")) and spec.isprintable():
c2pa_info["c2pa_spec"] = spec
# Find timestamps
timestamp_matches = re.findall(rb"(\d{14}Z)", chunk_data)
if timestamp_matches:
c2pa_info["timestamp"] = timestamp_matches[0].decode("utf-8")
if len(timestamp_matches) > 1:
c2pa_info["timestamps"] = [t.decode("utf-8") for t in timestamp_matches[:3]]
def extract_c2pa_chunk(image_path: Path) -> bytes | None:
"""
Extract the raw C2PA JUMBF chunk from a PNG file.
Args:
image_path: Path to the source PNG file.
Returns:
Raw bytes of the C2PA chunk or None.
"""
if image_path.suffix.lower() != ".png":
return None
try:
with open(image_path, "rb") as f:
signature = f.read(8)
if signature != PNG_SIGNATURE:
return None
file_size = f.seek(0, 2)
f.seek(8)
while True:
chunk_header = f.read(8)
if len(chunk_header) < 8:
break
length = struct.unpack(">I", chunk_header[:4])[0]
chunk_type = chunk_header[4:8]
# Clamp the attacker-controlled 32-bit length to the bytes that
# actually remain, so a malformed huge length can't allocate GBs.
safe_length = max(0, min(length, file_size - f.tell()))
if chunk_type == C2PA_CHUNK_TYPE:
chunk_data = f.read(safe_length)
crc = f.read(4)
# Check for any C2PA signature
for sig in C2PA_SIGNATURES:
if sig in chunk_data:
return chunk_header + chunk_data + crc
# Also check lowercase variants
if b"jumb" in chunk_data.lower() or b"c2pa" in chunk_data.lower():
return chunk_header + chunk_data + crc
else:
f.seek(safe_length + 4, 1)
if chunk_type == b"IEND":
break
except Exception:
pass
return None
def inject_c2pa_chunk(target_path: Path, output_path: Path, c2pa_chunk: bytes) -> None:
"""
Inject a C2PA JUMBF chunk into a PNG file.
Args:
target_path: Path to the target PNG file.
output_path: Path where the output file will be saved.
c2pa_chunk: Raw bytes of the C2PA chunk to inject.
Raises:
ValueError: If not PNG files.
"""
if target_path.suffix.lower() != ".png" or output_path.suffix.lower() != ".png":
raise ValueError("C2PA chunk injection is only supported for PNG files")
output_path.parent.mkdir(parents=True, exist_ok=True)
with open(target_path, "rb") as f_in, open(output_path, "wb") as f_out:
f_out.write(f_in.read(8))
c2pa_injected = False
while True:
chunk_header = f_in.read(8)
if len(chunk_header) < 8:
break
length = struct.unpack(">I", chunk_header[:4])[0]
chunk_type = chunk_header[4:8]
chunk_data = f_in.read(length)
crc = f_in.read(4)
if chunk_type == b"IDAT" and not c2pa_injected:
f_out.write(c2pa_chunk)
c2pa_injected = True
if chunk_type == C2PA_CHUNK_TYPE:
continue
f_out.write(chunk_header)
f_out.write(chunk_data)
f_out.write(crc)
if chunk_type == b"IEND":
break
src/remove_ai_watermarks/noai/constants.py
-342
View File
@@ -1,342 +0,0 @@
"""Shared constants for AI metadata detection, C2PA parsing, and format support.
All modules reference these constants rather than hard-coding values,
so adding a new AI tool or metadata key requires updating only this file.
"""
from typing import NamedTuple
# Supported image formats for the pixel/removal path (CLI input validation + batch
# discovery). PNG/JPEG/WebP decode+encode via cv2; HEIC/HEIF/AVIF via the optional
# pillow-heif dep (image_io.imread Pillow fallback + imwrite _pil_write), so batch
# now picks them up and the CLI no longer warns on an iPhone HEIC. JPEG-XL is left
# out on purpose -- it is metadata/strip-only (no pixel decoder without pillow-jxl).
SUPPORTED_FORMATS = {".png", ".jpg", ".jpeg", ".webp", ".heic", ".heif", ".avif"}
# AI-generated image metadata keys (Stable Diffusion, ComfyUI, Midjourney, etc.)
AI_METADATA_KEYS = [
"parameters", # Stable Diffusion WebUI (AUTOMATIC1111, Vladmandic)
"postprocessing", # SD WebUI post-processing info
"extras", # SD WebUI extras
"workflow", # ComfyUI workflow JSON
"prompt", # Some AI tools
"Dream", # DreamStudio
"SD:mode", # Stability AI
"StableDiffusionVersion", # SD version info
"generation_time", # Generation time info
"Model", # Model name
"Model hash", # Model hash
"Seed", # Seed value
]
# Standard PNG metadata keys
PNG_METADATA_KEYS = [
"Author",
"Title",
"Description",
"Copyright",
"Creation Time",
"Software",
"Disclaimer",
"Warning",
"Source",
"Comment",
]
# AI-related keywords for detection
AI_KEYWORDS = [
"prompt",
"negative_prompt",
"sampler",
"cfg_scale",
"lora",
"diffusion",
"comfy",
"midjourney",
"dall-e",
"dalle",
"imagen",
"firefly",
"c2pa",
"chatgpt",
"gpt-4",
"sora",
"openai",
"truepic",
"stable_diffusion",
"invokeai",
]
# C2PA (Coalition for Content Provenance and Authenticity) constants
# Used by Google Imagen, Adobe Firefly, Microsoft Designer, OpenAI, etc.
C2PA_CHUNK_TYPE = b"caBX" # JUMBF container chunk type for C2PA
C2PA_SIGNATURES = [
b"c2pa",
b"C2PA",
b"jumb",
b"jumd",
b"JUMBF",
b"jumbf",
b"cbor",
b"contentcreds",
b"digid",
b"assertions",
b"manifest",
]
# Single source of truth for every C2PA-signing vendor. The three per-vendor
# facts that used to live in separate tables -- the issuer byte signature
# (C2PA_ISSUERS), the SynthID pairing (SYNTHID_C2PA_ISSUERS), and the human
# platform label (identify._ISSUER_PLATFORM) -- are all fields here, so adding a
# new C2PA vendor is a single append below; the views derive automatically.
class C2paAiVendor(NamedTuple):
issuer: bytes # distinctive byte signature scanned in the manifest (cert org / signer)
org: str # resolved issuer/cert-org display name (the old C2PA_ISSUERS value)
# Human platform label for identify; None marks a signing authority / non-generator
# (e.g. Truepic), which never names an AI platform on its own.
platform: str | None
# Substring matched against the joined issuer-org names for platform attribution
# (usually a shorter form of org, e.g. "Google" for "Google LLC"); None when platform is.
needle: str | None
synthid: bool = False # vendor pairs an invisible SynthID pixel watermark with its C2PA manifest
# The vendor's mere presence in the manifest asserts AI generation even without
# a digitalSourceType (``trainedAlgorithmicMedia``) assertion. Set ONLY for a
# pure-generator brand whose issuer/generator byte string is unambiguous (e.g.
# "Dreamina"). Do NOT set for common-word issuers (Adobe/Google/OpenAI/Microsoft):
# those appear incidentally in unrelated XMP/trust-chain bytes, so they stay
# source-type-gated in identify._attribute_platform.
asserts_ai: bool = False
# C2PA known vendors, ORDERED for first-match-wins platform attribution: when a
# manifest names several issuers (Microsoft Designer signs as "OpenAI, Microsoft"),
# the earlier entry wins so the product, not the backend engine, is named.
# Used by Google Imagen, Adobe Firefly, Microsoft Designer, OpenAI, etc.
C2PA_AI_VENDORS: tuple[C2paAiVendor, ...] = (
# Microsoft signs both Designer and Bing Image Creator; Bing now runs its own
# MAI-Image model (not DALL-E), so the label stays model-neutral.
C2paAiVendor(b"Microsoft", "Microsoft", "Microsoft (Bing Image Creator / Designer)", "Microsoft"),
C2paAiVendor(b"Adobe", "Adobe", "Adobe Firefly", "Adobe"),
C2paAiVendor(b"OpenAI", "OpenAI", "OpenAI (ChatGPT / gpt-image / DALL-E / Sora)", "OpenAI", synthid=True),
C2paAiVendor(b"Google", "Google LLC", "Google (Gemini / Imagen)", "Google", synthid=True),
# Stability AI signs C2PA as "Stability AI" (cert org "Stability AI Ltd").
# Verified on a live Brand Studio (DreamStudio successor) output, 2026-05-24.
C2paAiVendor(b"Stability AI", "Stability AI", "Stability AI (Stable Image / DreamStudio)", "Stability AI"),
# Black Forest Labs (FLUX) API output: claim_generator_info "Black Forest
# Labs API" + a c2pa.ai_generated_content assertion + trainedAlgorithmicMedia.
# Verified on a real signed FLUX JPEG, 2026-05-29.
C2paAiVendor(b"Black Forest Labs", "Black Forest Labs", "Black Forest Labs (FLUX)", "Black Forest Labs"),
# ByteDance's Volcano Engine (Volcengine) signs its AI image output with a
# cert from certificate_center@volcengine.com -- the platform behind Doubao /
# Jimeng. Verified on two real signed JPEGs, 2026-05-29.
C2paAiVendor(
b"volcengine", "ByteDance (Volcano Engine)", "ByteDance (Doubao / Jimeng / Volcano Engine)", "ByteDance"
),
# Some Volcano Engine certs name the signer with the Chinese legal entity
# "北京火山引擎科技有限公司" (Beijing Volcano Engine Technology Co., Ltd.) rather
# than the latin "volcengine" -- the latin needle misses it entirely. The issuer is the
# UTF-8 of the Chinese name (it appears UTF-8-encoded in the manifest-store
# JSON and the raw caBX bytes alike); it normalizes to the same "ByteDance"
# needle and platform as the volcengine row, so the two collapse together for
# clash detection. Verified against compatible signed samples.
C2paAiVendor(
"北京火山引擎科技有限公司".encode(),
"ByteDance (Volcano Engine)",
"ByteDance (Doubao / Jimeng / Volcano Engine)",
"ByteDance",
),
# ByteDance's international brand (BytePlus / Seedream / Seededit) signs its
# cert as "Byteplus Pte. Ltd." -- the bare ``volcengine`` needle misses it, so
# real BytePlus AI output was mis-attributed (an incidental "Adobe XMP" string
# in the file's XMP made it read "Adobe Firefly"). Adding the issuer means the
# clean manifest issuer matches "BytePlus (ByteDance)" directly. The platform
# string mirrors the volcengine row: both share the "ByteDance" needle, so the
# earlier row's label wins anyway -- they normalize together for clash
# detection. Verified on compatible signed samples.
C2paAiVendor(b"Byteplus", "BytePlus (ByteDance)", "ByteDance (Doubao / Jimeng / Volcano Engine)", "ByteDance"),
# Dreamina (ByteDance's international Jimeng brand) signs C2PA as "Bytedance
# Pte. Ltd." with a "Dreamina/x.y" claim generator and, unlike the Volcano
# Engine output, NO digitalSourceType assertion -- so the generator name is the
# only AI signal. It is registered by that generator token (which the caBX /
# store-JSON byte scan sees across active + ingredient manifests, where the
# active manifest is often a plain c2pa-tool transcode). ``asserts_ai`` lets the
# issuer alone flag AI without trainedAlgorithmicMedia; "Dreamina" is a
# distinctive brand string, so it does not risk the incidental-mention problem
# the common-word issuers have. Verified on compatible signed samples.
# Normalizes to the same "ByteDance" needle/platform as the
# volcengine row (they collapse together for clash detection).
C2paAiVendor(
b"Dreamina",
"ByteDance (Dreamina)",
"ByteDance (Doubao / Jimeng / Volcano Engine)",
"ByteDance",
asserts_ai=True,
),
# Canva Magic Media signs AI-generated images as "Canva" with a generic
# c2pa-rs claim generator + trainedAlgorithmicMedia; without this entry the
# source read AI but no platform was attributed. Verified on compatible signed
# samples. Canva does not use SynthID.
C2paAiVendor(b"Canva", "Canva", "Canva (Magic Media)", "Canva"),
# ElevenLabs is a pure generative-AI company (AI voice / audio, and image /
# video via its API); it signs output as "Eleven Labs Inc.", so the C2PA
# manifest alone marks AI generation. Verified on compatible signed samples.
# ElevenLabs does not use SynthID.
C2paAiVendor(b"Eleven Labs", "ElevenLabs", "ElevenLabs", "ElevenLabs"),
# fal.ai (generative inference platform, issuer "fal - Features & Labels
# Inc." / common name "fal.ai", claim generators like "fal-ai/seedvr",
# "fal-ai/gpt-image-2"). The files carry trainedAlgorithmicMedia, so the
# verdict already fired, but the platform stayed unattributed. fal.ai is
# a pure generative platform, so ``asserts_ai`` also covers its output
# that omits the source-type.
C2paAiVendor(b"fal-ai", "fal.ai", "fal.ai", "fal.ai", asserts_ai=True),
# Bria AI (bria.ai, generative platform) signs as "Bria Artificial
# Intelligence" with a "Bria Ai" claim generator and source type
# ``empty`` (NOT trainedAlgorithmicMedia), so a real signed file was
# completely missed by identify. A pure-AI
# vendor with distinctive strings, so ``asserts_ai`` is safe here.
C2paAiVendor(b"Bria", "Bria Artificial Intelligence", "Bria AI", "Bria", asserts_ai=True),
# Truepic is a C2PA signing authority, not an AI generator: no platform label,
# never asserts is_ai (the verdict comes from the digital-source-type).
C2paAiVendor(b"Truepic", "Truepic", None, None),
)
# Deliberately NOT registered as AI-generation vendors:
# - TikTok Inc.: signs C2PA as a content-provenance / AI-labeling authority on
# uploads, not as an image generator. The is_ai verdict keys off the
# digitalSourceType (trainedAlgorithmicMedia), which is already honored; a
# bare TikTok signer marks distribution provenance, not generation, so adding
# it as a generator needle would mis-label human uploads as AI.
# - PixelBin.io (issuer "Fynd"): an image transformation / optimization / CDN
# service. Its C2PA stamps a transform/upload step, not a generation event.
# Both are excluded to avoid false-positive AI attribution; re-evaluate only
# against a real signed file whose manifest carries a trainedAlgorithmicMedia
# digital-source type produced by the vendor itself.
# Derived view -- add a vendor to C2PA_AI_VENDORS above, not here.
# C2PA issuer signature -> resolved org name, for the manifest byte-scan.
C2PA_ISSUERS: dict[bytes, str] = {v.issuer: v.org for v in C2PA_AI_VENDORS}
# Resolved org names of the vendors whose presence asserts AI generation on its
# own (no digitalSourceType needed) -- see the ``asserts_ai`` field. identify uses
# this to lift the AI verdict for an identity-AI issuer (e.g. Dreamina) that ships
# no trainedAlgorithmicMedia. Derived from the flag -- set it on the vendor, not here.
C2PA_IDENTITY_AI_ORGS: frozenset[str] = frozenset(v.org for v in C2PA_AI_VENDORS if v.asserts_ai)
# C2PA issuers whose signed outputs also carry an invisible SynthID pixel
# watermark -- a metadata proxy for "SynthID is in the pixels":
# - Google (Imagen/Gemini): embeds SynthID, long-standing (DeepMind docs).
# - OpenAI (ChatGPT/Codex/API): pairs SynthID with C2PA since ~2026-05-20.
# Confirmed by OpenAI's Help Center ("C2PA and SynthID in OpenAI-generated
# images", updated 2026-05-21): "Images generated with ChatGPT, Codex, and
# our API include both C2PA metadata and SynthID watermarks." OpenAI also
# notes a signal may be absent if "the image was created before these
# signals were available" -- so OpenAI images from before the rollout can
# carry C2PA without SynthID. For OpenAI the proxy is therefore "likely",
# not certain; the verdict string is hedged accordingly. OpenAI's own oracle
# is openai.com/verify (Google's is the Gemini app "Verify with SynthID").
# The issuer byte ("OpenAI"/"Google") is verified locally against data/fixtures/provenance;
# the SynthID pairing is documented behavior (Google: DeepMind; OpenAI: above).
# Adobe Firefly and Microsoft Designer sign C2PA but do NOT use SynthID, so a
# C2PA manifest alone is not a SynthID signal -- the issuer is. The pixel
# watermark is not locally detectable (proprietary decoder); the C2PA companion
# is the proxy, and only while the manifest is intact.
# Derived from the `synthid` flag on C2PA_AI_VENDORS -- set it there, not here.
SYNTHID_C2PA_ISSUERS: frozenset[bytes] = frozenset(v.issuer for v in C2PA_AI_VENDORS if v.synthid)
# C2PA known AI tools
C2PA_AI_TOOLS = {
b"GPT-4o": "GPT-4o",
b"ChatGPT": "ChatGPT",
b"Sora": "Sora",
b"DALL-E": "DALL-E",
b"DALL": "DALL-E",
b"Imagen": "Imagen",
b"Firefly": "Firefly",
}
# C2PA ``c2pa.soft-binding`` algorithm identifiers -> the forensic-watermark
# vendor that stamped the pixels. The manifest's ``alg`` field names the
# watermark scheme even when the watermark itself cannot be decoded locally, so
# a byte-scan for these (keyed on a distinctive prefix to catch all variants)
# tells us a third-party forensic watermark is present and whose. Verified
# against the official C2PA registry (github.com/c2pa-org/softbinding-algorithm-list).
# Adobe TrustMark is additionally decodable locally (see ``trustmark_detector``);
# the rest (Digimarc, Imatag, Steg.AI, etc.) are proprietary oracle-only decoders.
C2PA_SOFT_BINDINGS = {
b"com.adobe.trustmark": "Adobe TrustMark",
b"com.adobe.icn": "Adobe (content fingerprint)",
b"com.digimarc": "Digimarc",
b"com.imatag.lamark": "Imatag (Lamark)",
b"ai.steg": "Steg.AI",
b"com.microsoft.invismark": "Microsoft InvisMark",
b"com.microsoft.wavmark": "Microsoft WavMark",
b"com.verimatrix": "Verimatrix",
b"com.nagra.nexguard": "NAGRA NexGuard",
b"com.aiwatermark": "AIWatermark (Meta PixelSeal)",
b"ai.trufo": "Trufo",
b"app.overlai": "Overlai",
b"com.markany": "MarkAny",
b"com.mentaport": "Mentaport",
b"es.lumatrace": "LumaTrace",
b"ai.verda": "VerdaAI",
b"ai.contentlens": "ContentLens",
b"io.iscc": "ISCC (content code)",
}
# Lowercased substrings that mark an AI generator when found in an EXIF
# ``Software`` / XMP ``CreatorTool`` value. Conservative on purpose: plain
# editors like "Adobe Photoshop" or "GIMP" must NOT match (no AI token), so only
# generator names land here. Add new generators here, not inline.
AI_GENERATOR_TOKENS: frozenset[str] = frozenset(
{
"firefly",
"dall-e",
"dalle",
"midjourney",
"stable diffusion",
"stable-diffusion",
"stablediffusion",
"comfyui",
"automatic1111",
"invokeai",
"imagen",
"gpt-image",
"nightcafe",
"ideogram",
"leonardo",
"flux",
"dreamstudio",
# Generator stamps without C2PA:
# - NovelAI (anime SD): PNG tEXt Software="NovelAI", Source="NovelAI
# Diffusion V4.5 <hash>", Title="NovelAI generated image".
# - Reve Image (reve.com): EXIF Software / XMP CreatorTool = "reve.com"
# (the bare token "reve" would false-positive on "forever"/"reverie").
# - Aphrodite AI: EXIF Make / Software = "Aphrodite AI[ v1.0]".
"novelai",
"reve.com",
"aphrodite ai",
# Additional verified markers:
# - Apple Photos Clean Up (Apple Intelligence object removal): XMP
# photoshop:Credit / IPTC credit value; composite source-type
# covered detection, this token covers removal parity.
# - fal-ai: generative-platform generator string.
"apple photos clean up",
"fal-ai",
}
)
# C2PA action types
C2PA_ACTIONS = {
b"c2pa.created": "created",
b"c2pa.converted": "converted",
b"c2pa.edited": "edited",
b"c2pa.filtered": "filtered",
b"c2pa.cropped": "cropped",
b"c2pa.resized": "resized",
b"c2pa.opened": "opened",
b"c2pa.placed": "placed",
}
# PNG signature
PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
src/remove_ai_watermarks/noai/extractor.py
-155
View File
@@ -1,155 +0,0 @@
"""Read-only metadata extraction from PNG and JPEG images.
Provides functions to pull all metadata, AI-only metadata, or a
human-readable summary without modifying the source file.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Any, cast
if TYPE_CHECKING:
from pathlib import Path
import piexif
from PIL import Image
from remove_ai_watermarks.noai.c2pa import extract_c2pa_chunk, extract_c2pa_info, has_c2pa_metadata
from remove_ai_watermarks.noai.constants import AI_KEYWORDS, AI_METADATA_KEYS, PNG_METADATA_KEYS
def extract_metadata(source_path: Path) -> dict[str, Any]:
"""
Extract all metadata from a PNG or JPG file.
Args:
source_path: Path to the source image file.
Returns:
Dictionary containing all extracted metadata.
"""
metadata: dict[str, Any] = {}
with Image.open(source_path) as img:
# Extract EXIF data
if "exif" in img.info:
try:
exif_dict = piexif.load(img.info["exif"])
metadata["exif"] = exif_dict
except Exception:
metadata["exif_raw"] = img.info["exif"]
# Extract standard PNG metadata
for key in PNG_METADATA_KEYS:
if key in img.info:
metadata[key] = img.info[key]
# Extract all other metadata including AI-specific
for key, value in img.info.items():
if not isinstance(key, str):
continue
if key not in metadata and key not in ["exif"]:
metadata[key] = value
# Extract DPI and gamma if present
if "dpi" in img.info:
metadata["dpi"] = img.info["dpi"]
if "gamma" in img.info:
metadata["gamma"] = img.info["gamma"]
# Check for C2PA metadata
if has_c2pa_metadata(source_path):
metadata["c2pa"] = extract_c2pa_info(source_path)
c2pa_chunk = extract_c2pa_chunk(source_path)
if c2pa_chunk:
metadata["c2pa_chunk"] = c2pa_chunk
return metadata
def extract_ai_metadata(source_path: Path) -> dict[str, Any]:
"""
Extract only AI-generated metadata from a PNG or JPG file.
Args:
source_path: Path to the source image file.
Returns:
Dictionary containing only AI-related metadata.
"""
ai_metadata: dict[str, Any] = {}
with Image.open(source_path) as img:
for key in AI_METADATA_KEYS:
if key in img.info:
ai_metadata[key] = img.info[key]
for key, value in img.info.items():
if not isinstance(key, str):
continue
key_lower = key.lower()
if key not in ai_metadata and any(kw in key_lower for kw in AI_KEYWORDS):
ai_metadata[key] = value
# Check for C2PA metadata
if has_c2pa_metadata(source_path):
ai_metadata["c2pa"] = extract_c2pa_info(source_path)
c2pa_chunk = extract_c2pa_chunk(source_path)
if c2pa_chunk:
ai_metadata["c2pa_chunk"] = c2pa_chunk
return ai_metadata
def has_ai_metadata(image_path: Path) -> bool:
"""
Check if an image contains AI-generated metadata.
Args:
image_path: Path to the image file.
Returns:
True if AI metadata is detected, False otherwise.
"""
with Image.open(image_path) as img:
for key in AI_METADATA_KEYS:
if key in img.info:
return True
return bool(has_c2pa_metadata(image_path))
def get_ai_metadata_summary(source_path: Path) -> str:
"""
Get a human-readable summary of AI metadata.
Args:
source_path: Path to the source image file.
Returns:
Formatted string with AI metadata summary.
"""
ai_meta = extract_ai_metadata(source_path)
if not ai_meta:
return "No AI metadata found."
lines = ["AI Image Metadata:"]
lines.append("-" * 40)
for key, value in ai_meta.items():
if key == "c2pa_chunk":
continue
if key == "c2pa" and isinstance(value, dict):
lines.append("C2PA Metadata:")
for ck, cv in cast("dict[str, Any]", value).items():
lines.append(f" {ck}: {cv}")
elif isinstance(value, str) and len(value) > 100:
value = value[:100] + "..."
lines.append(f"{key}: {value}")
elif isinstance(value, bytes):
lines.append(f"{key}: <binary data ({len(value)} bytes)>")
else:
lines.append(f"{key}: {value}")
return "\n".join(lines)
src/remove_ai_watermarks/noai/img2img_runner.py
-161
View File
@@ -1,161 +0,0 @@
"""Img2img pipeline execution with progress monitoring and MPS fallback.
Extracted from ``watermark_remover.py`` to keep the ``WatermarkRemover``
class focused on orchestration.
"""
from __future__ import annotations
import contextlib
import logging
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from collections.abc import Callable
from PIL import Image
from remove_ai_watermarks.noai.progress import is_mps_error, make_pipeline_progress
logger = logging.getLogger(__name__)
def run_img2img(
pipeline: Any,
image: Image.Image,
strength: float,
num_inference_steps: int,
guidance_scale: float,
generator: Any,
device: str,
set_progress: Callable[[str], None],
extra_kwargs: dict[str, Any] | None = None,
) -> Image.Image:
"""Execute img2img with live progress and return the generated image.
``extra_kwargs`` overlays additional pipeline arguments (e.g. the ControlNet
``control_image`` / ``controlnet_conditioning_scale`` and a non-empty prompt),
so a ControlNet img2img pass reuses the same progress + fallback machinery.
"""
effective_steps = max(1, int(num_inference_steps * strength))
step_cb, first_step, done_ev, start_updater = make_pipeline_progress(
effective_steps,
device,
set_progress,
)
start_updater()
try:
result = _call_pipeline(
pipeline, image, strength, num_inference_steps, guidance_scale, generator, step_cb, extra_kwargs
)
done_ev.set()
return result.images[0]
except TypeError as exc:
# The only TypeError we retry is the deprecated-callback case: `_call_pipeline`
# passes the legacy `callback`/`callback_steps` kwargs, and a diffusers version
# that removed them raises TypeError("... unexpected keyword argument
# 'callback'"). We then re-run once WITHOUT the progress callback. Any OTHER
# TypeError (e.g. a bad control_image/dtype in the forward pass) is a real error
# -- re-raise it instead of silently re-running the whole diffusion pass and
# masking the cause.
if "callback" not in str(exc):
raise
first_step.set()
result = _call_pipeline(
pipeline, image, strength, num_inference_steps, guidance_scale, generator, None, extra_kwargs
)
done_ev.set()
return result.images[0]
finally:
first_step.set()
done_ev.set()
def run_img2img_with_mps_fallback(
load_pipeline: Callable[[], Any],
image: Image.Image,
strength: float,
num_inference_steps: int,
guidance_scale: float,
generator: Any,
device: str,
set_progress: Callable[[str], None],
*,
reload_on_cpu: Callable[[], Any],
extra_kwargs: dict[str, Any] | None = None,
) -> tuple[Image.Image, str]:
"""Run img2img; on MPS error, fall back to CPU.
``extra_kwargs`` overlays extra pipeline arguments (used by the ControlNet
path). Returns ``(result_image, final_device)`` — device may change to
``"cpu"`` on fallback.
"""
pipeline = load_pipeline()
try:
img = run_img2img(
pipeline,
image,
strength,
num_inference_steps,
guidance_scale,
generator,
device,
set_progress,
extra_kwargs,
)
return img, device
except RuntimeError as error:
if device == "mps" and is_mps_error(error):
logger.warning("MPS error detected: %s. Falling back to CPU.", error)
set_progress("MPS error! Clearing cache and retrying on CPU...")
try_empty_device_cache("mps")
pipeline = reload_on_cpu()
img = run_img2img(
pipeline, image, strength, num_inference_steps, guidance_scale, None, "cpu", set_progress, extra_kwargs
)
return img, "cpu"
raise
def _call_pipeline(
pipeline: Any,
image: Image.Image,
strength: float,
num_inference_steps: int,
guidance_scale: float,
generator: Any,
step_callback: Any,
extra_kwargs: dict[str, Any] | None = None,
) -> Any:
kwargs: dict[str, Any] = {
"prompt": "",
"image": image,
"strength": strength,
"num_inference_steps": num_inference_steps,
"guidance_scale": guidance_scale,
"generator": generator,
}
if extra_kwargs:
kwargs.update(extra_kwargs)
if step_callback is not None:
kwargs["callback"] = step_callback
kwargs["callback_steps"] = 1
return pipeline(**kwargs)
def try_empty_device_cache(device: str) -> None:
"""Best-effort free of cached GPU/MPS/XPU memory for ``device``.
``torch.<device>.empty_cache()`` exists for cuda/mps/xpu but not cpu (the
hasattr guard skips the cpu no-op). Never raises -- callers use it as cleanup
(the MPS->CPU fallback here, and the batch loop in watermark_remover).
"""
with contextlib.suppress(Exception):
import torch
backend = getattr(torch, device, None)
if backend is not None and hasattr(backend, "empty_cache"):
backend.empty_cache() # type: ignore[attr-defined]
src/remove_ai_watermarks/noai/progress.py
-332
View File
@@ -1,332 +0,0 @@
"""Terminal progress animation and library output suppression.
This module provides two main capabilities for the CLI:
1. ``run_with_progress`` — a styled two-line terminal animation that
displays a bouncing highlight bar, a braille spinner, elapsed time,
and a live operation message while a background task executes.
2. ``silence_library_output`` — a wrapper that suppresses noisy log
output produced by third-party ML libraries (transformers, diffusers,
huggingface_hub, tqdm) so the user only sees our own progress messages.
"""
from __future__ import annotations
import contextlib
import io
import os
import sys
import threading
import time
import warnings
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from collections.abc import Callable
# ── ANSI color constants ────────────────────────────────────────────
_CYAN = "\033[36m"
_YELLOW = "\033[33m"
_GREEN = "\033[32m"
_DIM = "\033[2m"
_BOLD = "\033[1m"
_RESET = "\033[0m"
# Bar geometry
_BAR_WIDTH = 32
_HIGHLIGHT_WIDTH = 5
def _no_color() -> bool:
"""Respect the NO_COLOR convention (https://no-color.org/)."""
return bool(os.environ.get("NO_COLOR"))
def _truncate(text: str, max_len: int = 72) -> str:
"""Shorten a string with an ellipsis if it exceeds *max_len*."""
return text if len(text) <= max_len else text[: max_len - 1] + "…"
def _build_bar(step: int) -> str:
"""Build a flowing highlight bar that bounces across the width.
The highlight segment (5 chars wide) travels left→right→left
continuously, giving the user a visual "working" signal.
"""
cycle = _BAR_WIDTH * 2 - 2
pos = step % cycle
if pos >= _BAR_WIDTH:
pos = cycle - pos
hl_start = max(0, pos - _HIGHLIGHT_WIDTH // 2)
hl_end = min(_BAR_WIDTH, pos + _HIGHLIGHT_WIDTH // 2 + 1)
before = "━" * hl_start
highlight = "━" * (hl_end - hl_start)
after = "━" * (_BAR_WIDTH - hl_end)
if _no_color():
return before + highlight + after
return f"{_DIM}{before}{_RESET}{_BOLD}{_YELLOW}{highlight}{_RESET}{_DIM}{after}{_RESET}"
def run_with_progress(
task: Callable[[], Any],
progress_state: dict[str, str] | None = None,
) -> Any:
"""Execute *task* in a background thread while showing a progress animation.
The animation renders two lines to ``sys.__stderr__``:
- **Line 1**: braille spinner + bouncing bar + elapsed seconds
- **Line 2**: current operation message from *progress_state*
When the task finishes, a green "Completed" line replaces the animation.
Args:
task: A zero-argument callable to run in the background.
progress_state: Mutable dict whose ``"message"`` key is read
by the animation loop to display the current operation.
Returns:
Whatever *task* returns.
Raises:
Any exception raised by *task* is re-raised after the animation
is cleaned up.
"""
done = threading.Event()
output_holder: dict[str, Any] = {"result": None, "error": None}
def worker() -> None:
try:
output_holder["result"] = task()
except Exception as error: # pragma: no cover - passthrough
output_holder["error"] = error
finally:
done.set()
thread = threading.Thread(target=worker, daemon=True)
thread.start()
spinner_frames = "⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏"
idx = 0
start_time = time.time()
no_color = _no_color()
def _get_operation() -> str:
if isinstance(progress_state, dict):
return progress_state.get("message", "Processing...")
return "Processing..."
# ── Animation loop ──────────────────────────────────────────────
while not done.is_set():
spinner = spinner_frames[idx % len(spinner_frames)]
elapsed = int(time.time() - start_time)
bar_str = _build_bar(idx)
operation = _truncate(_get_operation())
if no_color:
line1 = f" {spinner} Processing {bar_str} {elapsed:>3}s"
line2 = f" ╰─ {operation}"
else:
line1 = f" {_CYAN}{spinner}{_RESET} Processing {bar_str} {_BOLD}{_YELLOW}{elapsed:>3}s{_RESET}"
line2 = f" {_DIM}╰─ {operation}{_RESET}"
print(
f"\r\033[2K{line1}\n\033[2K{line2}\033[1A\r",
end="",
flush=True,
file=sys.__stderr__,
)
time.sleep(0.08)
idx += 1
# ── Final "done" frame ──────────────────────────────────────────
thread.join()
total = int(time.time() - start_time)
final_operation = _truncate(_get_operation())
done_bar = "━" * _BAR_WIDTH
if no_color:
final_line1 = f" ✓ Completed {done_bar} {total:>3}s"
final_line2 = f" ╰─ {final_operation}"
else:
final_line1 = (
f" {_GREEN}{_BOLD}✓{_RESET} {_GREEN}Completed{_RESET} "
f"{_GREEN}{done_bar}{_RESET} {_BOLD}{_GREEN}{total:>3}s{_RESET}"
)
final_line2 = f" {_DIM}╰─ {final_operation}{_RESET}"
print(
f"\r\033[2K{final_line1}\n\033[2K{final_line2}",
file=sys.__stderr__,
)
if output_holder["error"] is not None:
raise output_holder["error"]
return output_holder["result"]
def silence_library_output(
run_func: Callable[[], Any],
set_progress: Callable[[str], None] | None = None,
) -> Callable[[], Any]:
"""Return a wrapper that silences noisy ML library output.
The wrapper:
1. Disables HuggingFace Hub progress bars via env var.
2. Sets ``transformers``, ``diffusers``, and ``huggingface_hub``
loggers to *error* level.
3. Redirects ``stdout`` and ``stderr`` to ``io.StringIO`` sinks so
that stray ``tqdm`` bars and model-loading chatter are invisible.
4. Suppresses all Python warnings during the call.
Args:
run_func: The callable to execute silently.
set_progress: Optional callback to report phase changes.
Returns:
A zero-argument callable that, when invoked, runs *run_func*
inside the silent context.
"""
def wrapped() -> Any:
if set_progress:
set_progress("Configuring runtime and suppressing noisy logs...")
os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
for _silence in (
lambda: __import__("transformers").logging.set_verbosity_error(),
lambda: _silence_diffusers(),
lambda: __import__("huggingface_hub").logging.set_verbosity_error(),
):
with contextlib.suppress(Exception):
_silence()
with warnings.catch_warnings():
warnings.simplefilter("ignore")
with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()):
if set_progress:
set_progress("Executing watermark removal pipeline...")
return run_func()
return wrapped
def _silence_diffusers() -> None:
"""Silence diffusers logging and progress bars."""
from diffusers.utils import logging as diffusers_logging
diffusers_logging.set_verbosity_error()
if hasattr(diffusers_logging, "disable_progress_bar"):
diffusers_logging.disable_progress_bar()
# ── Shared pipeline progress helpers ─────────────────────────────────
_DEFAULT_PRE_PHASES: list[tuple[int, str]] = [
(0, "Encoding image with VAE encoder"),
(3, "Mapping pixel data → latent space"),
(7, "Injecting noise into latent representation"),
(12, "Building denoiser schedule"),
(18, "Starting reverse diffusion sampler"),
(30, "Running first denoising iteration"),
(50, "Still processing — this can take a while"),
(90, "Pipeline running — may take a few minutes"),
]
_DEFAULT_POST_PHASES: list[tuple[int, str]] = [
(0, "Denoising complete · Running VAE decoder"),
(2, "Decoding latent channels → RGB color space"),
(5, "Reconstructing pixel grid from latents"),
(10, "Applying color space conversion and normalization"),
(18, "Finalizing pixel output"),
(30, "Still decoding — large images take longer"),
(60, "Almost done — large images take longer to decode"),
]
def make_pipeline_progress(
effective_steps: int,
device: str,
set_progress: Callable[[str], None],
*,
bar_len: int = 20,
label: str = "Denoising",
pre_phases: list[tuple[int, str]] | None = None,
post_phases: list[tuple[int, str]] | None = None,
) -> tuple[Callable[..., None], threading.Event, threading.Event, Callable[[], threading.Thread]]:
"""Create step callback and background updater for a diffusion pipeline.
Returns:
(step_callback, first_step_event, pipeline_done_event, start_updater)
where ``start_updater()`` launches and returns the background thread.
"""
pre = pre_phases or [(s, f"{m} on {device}") for s, m in _DEFAULT_PRE_PHASES]
post = post_phases or [(s, f"{m} on {device}") for s, m in _DEFAULT_POST_PHASES]
t0_holder: list[float] = [time.monotonic()]
first_step = threading.Event()
pipeline_done = threading.Event()
last_cb_time: list[float] = [t0_holder[0]]
def _background_updater() -> None:
idx = 0
while not first_step.is_set():
elapsed = time.monotonic() - t0_holder[0]
while idx < len(pre) - 1 and elapsed >= pre[idx + 1][0]:
idx += 1
set_progress(pre[idx][1])
first_step.wait(timeout=0.4)
idx = 0
post_start: float | None = None
while not pipeline_done.is_set():
since_cb = time.monotonic() - last_cb_time[0]
if since_cb >= 1.5:
if post_start is None:
post_start = time.monotonic()
elapsed = time.monotonic() - post_start
while idx < len(post) - 1 and elapsed >= post[idx + 1][0]:
idx += 1
set_progress(post[idx][1])
else:
post_start = None
idx = 0
pipeline_done.wait(timeout=0.4)
def step_callback(step: int, timestep: int, latents: Any) -> None:
first_step.set()
last_cb_time[0] = time.monotonic()
elapsed = time.monotonic() - t0_holder[0]
current = step + 1
per_step = elapsed / max(1, current)
remaining = per_step * max(0, effective_steps - current)
filled = int(bar_len * current / max(1, effective_steps))
bar = "█" * filled + "░" * (bar_len - filled)
set_progress(
f"{label} [{bar}] {current}/{effective_steps} | {elapsed:.0f}s elapsed, ~{remaining:.0f}s left | {device}"
)
def start_updater() -> threading.Thread:
t0_holder[0] = time.monotonic()
last_cb_time[0] = t0_holder[0]
first_step.clear()
pipeline_done.clear()
t = threading.Thread(target=_background_updater, daemon=True)
t.start()
return t
return step_callback, first_step, pipeline_done, start_updater
# ── MPS fallback helper ──────────────────────────────────────────────
def is_mps_error(error: Exception) -> bool:
"""Check whether an exception is an MPS-related runtime error."""
return "mps" in str(error).lower()
src/remove_ai_watermarks/noai/utils.py
-43
View File
@@ -1,43 +0,0 @@
"""Low-level utility helpers used across the metadata pipeline.
Kept deliberately small — only format detection lives here so that
higher-level modules can import without circular dependencies.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from pathlib import Path
from remove_ai_watermarks.noai.constants import SUPPORTED_FORMATS
def is_supported_format(file_path: Path) -> bool:
"""
Check if the file format is supported.
Args:
file_path: Path to the image file.
Returns:
True if the format is supported, False otherwise.
"""
return file_path.suffix.lower() in SUPPORTED_FORMATS
def get_image_format(file_path: Path) -> str:
"""
Get the image format from file path.
Args:
file_path: Path to the image file.
Returns:
Format string (PNG, JPEG, etc.).
"""
suffix = file_path.suffix.lower()
if suffix in {".jpg", ".jpeg"}:
return "JPEG"
return "PNG"
src/remove_ai_watermarks/noai/watermark_profiles.py
-207
View File
@@ -1,207 +0,0 @@
"""Watermark removal model profiles and the default strength.
Pure configuration and lookup functions with no ML dependencies.
"""
from __future__ import annotations
import math
from typing import TYPE_CHECKING, Literal
if TYPE_CHECKING:
from pathlib import Path
DEFAULT_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
# Qwen-Image (20B MMDiT, Apache-2.0 code AND weights) base for the ``qwen`` pipeline:
# an img2img alternative to SDXL with native text rendering (incl. CJK). Loaded only
# when ``--pipeline qwen`` is selected; CUDA/cloud-class (does not fit MPS). CERTIFIED
# oracle floors (2026-06-20): OpenAI **0.10** (seed-robust -- clean on seeds 0-4) and
# Google/Gemini **0.25** (seed 0 verified on 2 images; pin a seed in prod, the Gemini
# oracle rate-limits volume seed-repeat). The Gemini floor (0.25) is HIGHER than the
# certified controlnet Gemini floor (0.15); ``resolve_strength(..., pipeline="qwen")``
# now carries this via ``_QWEN_VENDOR_STRENGTH`` (below), so ``--pipeline qwen`` gets the
# right floor automatically -- the old manual "pass --strength 0.25 for Gemini on qwen"
# workaround is retired.
# (Dispatch uses the bare "qwen" literal, matching the sdxl/controlnet sites, so there
# is no QWEN_PROFILE constant -- only the model id is referenced from code.)
QWEN_MODEL_ID = "Qwen/Qwen-Image"
# Canonical pipeline-profile names + the back-compat alias. The plain SDXL img2img
# profile is ``sdxl``; ``default`` is kept as an accepted alias (it was the profile's
# name before ``controlnet`` became the default-selected pipeline, 2026-06-09).
SDXL_PROFILE = "sdxl"
QWEN_ZIMAGE_PROFILE = "qwen-zimage"
_PROFILE_ALIASES = {
"default": SDXL_PROFILE,
"qwen_zimage": QWEN_ZIMAGE_PROFILE,
}
def normalize_profile(profile: str) -> str:
"""Canonicalize a pipeline-profile name, resolving the ``default`` -> ``sdxl`` alias."""
normalized = profile.strip().lower()
return _PROFILE_ALIASES.get(normalized, normalized)
def resolve_steps(num_inference_steps: int | None, pipeline: str) -> int:
"""Resolve a profile-specific step default while preserving explicit values.
The Lightning LoRA in ``qwen-zimage`` is distilled for four steps. Existing
SDXL and Qwen profiles keep the long-standing 50-step CLI default.
"""
if num_inference_steps is not None:
return num_inference_steps
return 4 if normalize_profile(pipeline) == QWEN_ZIMAGE_PROFILE else 50
def resolve_seed(seed: int | None, pipeline: str) -> int | None:
"""Keep the oracle-verified qwen-zimage profile deterministic by default."""
if seed is not None:
return seed
return 0 if normalize_profile(pipeline) == QWEN_ZIMAGE_PROFILE else None
# The SDXL-native canny ControlNet used by the ``controlnet`` pipeline. The
# ControlNet is an add-on to the SDXL base checkpoint (DEFAULT_MODEL_ID), not a
# separate base model, so both the ``sdxl`` and ``controlnet`` profiles load the
# same base weights and share the same vendor-adaptive strength ladder (see below).
CONTROLNET_CANNY_MODEL = "xinsir/controlnet-canny-sdxl-1.0"
# Vendor-adaptive default denoising strength for the SDXL img2img scrub, overridable
# from the CLI (`--strength`). The right strength depends on which vendor's SynthID is
# present (detected from the C2PA issuer, metadata.synthid_source). The SAME ladder
# applies to BOTH pipelines (`sdxl` plain img2img and `controlnet`) -- see "why one
# ladder" below.
#
# Data basis (see docs/synthid.md sections 2.2 / 5.5): ORACLE-CERTIFIED controlnet floors.
# Oracle re-testing
# LOWERED the ladder back to OpenAI 0.10 / Google 0.15: each output verified on its own
# oracle (openai.com/verify for OpenAI, the Google Gemini app for Google), all clean ->
# - OpenAI 0.10: 2 photoreal images (1402 / 1448 px), SynthID not found on either.
# - Google 0.15: 2 NATIVE-resolution images (both 2816x1536), SynthID not found on
# either -- this directly retires the earlier "native ~2816 likely needs ~0.35+"
# guess, which was speculative and never oracle-checked at that resolution.
# This supersedes the 2026-06-04 cert (OpenAI 0.20 / Google 0.30), whose higher floor a
# pixel-fidelity sweep showed was ~2x the removal floor and over-regenerated for no
# efficacy gain (Google MAE -20% at 0.15 vs 0.30, no SynthID returning). Unknown vendor
# tracks the Google (more robust watermark) value -> 0.15, still safe-by-default and the
# floor that real (no-vendor) photos hit, so it also minimizes damage when there is in
# fact nothing to remove. CAVEAT: the re-test is n=2 per vendor on photoreal / landscape
# content; FLAT-GRAPHIC hard cases (the historical `sdxl` weak spot) were NOT in the
# sample, so if an oracle still reads SynthID on a flat output, raise `--strength`.
#
# Why ONE ladder for both pipelines (2026-06-09): the certification was run on
# controlnet, and it does NOT transfer to `sdxl` by symmetry -- the two pipelines have
# OPPOSITE hard cases (controlnet leaves SynthID on photoreal, `sdxl` leaves it on flat
# graphics; the content-x-pipeline table in docs/synthid.md §5.1). BUT on its OWN hard
# case (flat fills) `sdxl` is the WEAKER remover -- plain img2img at low strength barely
# perturbs a flat region -- so it needs AT LEAST as much strength as controlnet, not
# less. Hence the certified controlnet floor is the right floor for `sdxl` too. The
# higher strength costs little quality where it matters. `controlnet` is now the default
# pipeline and `sdxl` is reached only through an explicit `--pipeline sdxl`. NOTE:
# this is a MARGIN argument for `sdxl`, not a fresh certification -- there is no local
# SynthID detector, so if an oracle still reads SynthID on a flat `sdxl` output, raise
# `--strength`.
OPENAI_STRENGTH = 0.10
GEMINI_STRENGTH = 0.15
UNKNOWN_STRENGTH = 0.15
# Backwards-compatible alias: the vendor-unknown value (what a caller gets without a
# detected vendor). Kept as DEFAULT_STRENGTH for existing references.
DEFAULT_STRENGTH = UNKNOWN_STRENGTH
# Detected-vendor -> default strength. Vendor strings come from `vendor_for_strength`.
_VENDOR_STRENGTH = {"openai": OPENAI_STRENGTH, "google": GEMINI_STRENGTH}
# Qwen has its OWN certified floors (Modal A100-80GB, 2026-06-20), DIFFERENT from the
# SDXL ladder above: OpenAI 0.10 (seed-robust), Gemini 0.25 (HIGHER than controlnet's
# 0.15 -- the 20B MMDiT perturbs less per denoising step, so it needs more strength to
# clear Gemini SynthID). Unknown vendor tracks the higher (Gemini) value, safe-by-default.
# `resolve_strength(..., pipeline="qwen")` uses this table so `--pipeline qwen` carries the
# right floor automatically -- retiring the old manual "pass --strength 0.25 for Gemini on
# qwen" workaround.
QWEN_OPENAI_STRENGTH = 0.10
QWEN_GEMINI_STRENGTH = 0.25
QWEN_UNKNOWN_STRENGTH = 0.25
_QWEN_VENDOR_STRENGTH = {"openai": QWEN_OPENAI_STRENGTH, "google": QWEN_GEMINI_STRENGTH}
def strength_default_help() -> str:
"""One-line description of the vendor-adaptive default, derived from the constants.
Single source of truth for the CLI ``--strength`` help so the numbers can never
drift from the actual ladder (they did once when the per-pipeline split was unified).
"""
return (
f"vendor-adaptive (OpenAI {OPENAI_STRENGTH} / Google {GEMINI_STRENGTH} / "
f"unknown {UNKNOWN_STRENGTH}, from the C2PA issuer; qwen-zimage instead uses "
"resolution-adaptive denoise)"
)
def resolve_strength(strength: float | None, vendor: str | None = None, pipeline: str | None = None) -> float:
"""Resolve the denoising strength, applying the vendor default when unset.
``None`` means "the user did not pass ``--strength``", which resolves
**vendor-adaptively**: ``vendor`` (``"openai"`` / ``"google"`` / None, from
``vendor_for_strength``) selects the per-vendor floor. The ``sdxl`` and ``controlnet``
pipelines share ONE ladder (``OPENAI_STRENGTH`` / ``GEMINI_STRENGTH`` /
``UNKNOWN_STRENGTH`` -- see the module comment for why); ``qwen`` has its OWN higher
ladder (``_QWEN_VENDOR_STRENGTH``, Gemini 0.25 vs controlnet 0.15), selected when
``pipeline`` normalizes to ``"qwen"``. An explicit value always wins (including
``0.0`` -- the check is ``is None``, not falsiness). Shared by the CLI (for display)
and the engine (for execution) so the two never disagree -- both must pass the SAME
``vendor`` and ``pipeline``.
"""
if strength is not None:
return strength
if pipeline is not None and normalize_profile(pipeline) == "qwen":
return _QWEN_VENDOR_STRENGTH.get(vendor or "", QWEN_UNKNOWN_STRENGTH)
return _VENDOR_STRENGTH.get(vendor or "", UNKNOWN_STRENGTH)
def viable_steps(num_inference_steps: int, strength: float) -> int:
"""The smallest step count >= ``num_inference_steps`` that actually denoises.
diffusers derives its img2img timesteps as ``int(steps * strength)``. When that
rounds to ZERO the pipeline builds an empty latent and dies deep inside attention
with ``cannot reshape tensor of 0 elements into shape [0, -1, 1, 512]`` -- an opaque
torch error for what is really "these two options cannot work together".
The combination is reachable with entirely valid CLI arguments: at the default
strength 0.15 every ``--steps`` below 7 crashed, and nothing told the user that
``--steps`` and ``--strength`` interact (found by the release smoke matrix,
2026-07-19). Raising the count to the minimum that denoises keeps the caller's intent
-- they asked for "few steps", not "zero" -- and the engine logs the adjustment.
A non-positive ``strength`` cannot denoise at any step count; return the caller's
value unchanged rather than dividing by zero.
"""
if strength <= 0:
return num_inference_steps
if int(num_inference_steps * strength) >= 1:
return num_inference_steps
return math.ceil(1 / strength)
def vendor_for_strength(image_path: Path) -> Literal["openai", "google"] | None:
"""Detect the SynthID vendor for strength selection: ``"openai"`` / ``"google"`` / None.
Reads the C2PA SynthID proxy (``metadata.synthid_source``) on the ORIGINAL input,
so it must run before any pass that strips metadata. When both issuers appear (a
rare multi-sign anomaly) Google wins -- the more-robust watermark -> safer (higher)
strength. Returns None when metadata is stripped or the issuer is neither vendor,
which maps to ``UNKNOWN_STRENGTH``. Lazy-imports ``metadata`` to keep this module
dependency-light.
"""
try:
from remove_ai_watermarks.metadata import synthid_source
src = (synthid_source(image_path) or "").lower()
except Exception: # metadata unreadable -> treat as unknown vendor
return None
if "google" in src:
return "google"
if "openai" in src:
return "openai"
return None
src/remove_ai_watermarks/noai/watermark_remover.py
-1117
View File
File diff suppressed because it is too large Load Diff
src/remove_ai_watermarks/video.py
+1 -1
View File
@@ -318,7 +318,7 @@ def _select_stable_visible_mark(
def _platform_from_video_metadata(markers: dict[str, str]) -> str | None:
"""Map supported C2PA-derived marker text to its generating platform."""
from remove_ai_watermarks.noai.constants import C2PA_AI_VENDORS
from remove_ai_watermarks._internal.constants import C2PA_AI_VENDORS
marker_text = "\n".join(markers.values()).casefold()
if not marker_text:
tests/test_cpu_offload.py
+1 -1
View File
@@ -12,7 +12,7 @@ from unittest.mock import Mock
import pytest
from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
def _remover(device: str, cpu_offload: bool) -> WatermarkRemover:
tests/test_identify.py
+1 -1
View File
@@ -718,7 +718,7 @@ class TestSparkleDetectRemoveAlignment:
class TestIdentifyImportIsLight:
"""`import identify` must stay torch-free (lazy noai/__init__): the package
"""`import identify` must stay torch-free (lazy _internal/__init__): the package
is deployed on a 512 MB host where eagerly pulling torch/diffusers OOMs."""
def test_import_identify_does_not_pull_torch(self):
tests/test_img2img_runner.py
+2 -2
View File
@@ -13,8 +13,8 @@ from unittest.mock import Mock
import pytest
from remove_ai_watermarks.noai import img2img_runner
from remove_ai_watermarks.noai.img2img_runner import (
from remove_ai_watermarks._internal import img2img_runner
from remove_ai_watermarks._internal.img2img_runner import (
run_img2img,
run_img2img_with_mps_fallback,
)
tests/test_invisible_engine.py
+8 -2
View File
@@ -215,7 +215,7 @@ class TestCannyControlImage:
pytest.skip("diffusion extra (torch/diffusers) not installed")
import numpy as np
from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
rng = np.random.default_rng(0)
img = Image.fromarray(rng.integers(0, 256, (64, 80, 3), dtype=np.uint8))
@@ -224,4 +224,10 @@ class TestCannyControlImage:
arr = np.array(out)
assert out.mode == "RGB"
assert arr.shape == (64, 80, 3)
assert arr.max() <= 255
import cv2
gray = cv2.cvtColor(np.asarray(img.convert("RGB")), cv2.COLOR_RGB2GRAY)
expected = cv2.Canny(gray, 100, 200)
assert np.array_equal(arr[:, :, 0], expected)
assert np.array_equal(arr[:, :, 1], expected)
assert np.array_equal(arr[:, :, 2], expected)
tests/test_metadata.py
+13 -13
View File
@@ -110,8 +110,8 @@ class TestHasAiMetadata:
def test_strip_c2pa_boxes_removes_uuid_box(self, tmp_path: Path):
"""ISOBMFF strip should drop the C2PA uuid box and keep everything else."""
from remove_ai_watermarks._internal.isobmff import strip_c2pa_boxes
from remove_ai_watermarks.metadata import C2PA_UUID
from remove_ai_watermarks.noai.isobmff import strip_c2pa_boxes
ftyp = b"\x00\x00\x00\x18ftypavif\x00\x00\x00\x00avifmif1"
# uuid box: size(4) + 'uuid' + 16-byte UUID + minimal payload (8 bytes -> total 32)
@@ -123,7 +123,7 @@ class TestHasAiMetadata:
def test_strip_c2pa_boxes_passthrough_for_non_isobmff(self):
"""Non-ISOBMFF input must be returned unchanged."""
from remove_ai_watermarks.noai.isobmff import strip_c2pa_boxes
from remove_ai_watermarks._internal.isobmff import strip_c2pa_boxes
data = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR" + b"\x00" * 100
cleaned, stripped = strip_c2pa_boxes(data)
@@ -1375,7 +1375,7 @@ class TestSoftBinding:
"""C2PA soft-binding alg identifier -> forensic-watermark vendor name."""
def test_vendors_in_recognizes_known_algs(self):
from remove_ai_watermarks.noai.c2pa import soft_binding_vendors_in
from remove_ai_watermarks._internal.c2pa import soft_binding_vendors_in
assert soft_binding_vendors_in(b"...alg...com.adobe.trustmark.P...") == ["Adobe TrustMark"]
assert soft_binding_vendors_in(b"com.digimarc.validate.1") == ["Digimarc"]
@@ -1385,7 +1385,7 @@ class TestSoftBinding:
assert soft_binding_vendors_in(b"io.iscc.v0") == ["ISCC (content code)"]
def test_vendors_in_empty_when_absent(self):
from remove_ai_watermarks.noai.c2pa import soft_binding_vendors_in
from remove_ai_watermarks._internal.c2pa import soft_binding_vendors_in
assert soft_binding_vendors_in(b"no soft binding here") == []
@@ -1472,8 +1472,8 @@ class TestLateProvenanceBox:
return p
def test_scan_c2pa_region_finds_late_box(self, tmp_path: Path):
from remove_ai_watermarks._internal.isobmff import scan_c2pa_region
from remove_ai_watermarks.metadata import C2PA_UUID
from remove_ai_watermarks.noai.isobmff import scan_c2pa_region
region = scan_c2pa_region(self._mp4_late_c2pa(tmp_path))
assert C2PA_UUID in region
@@ -1495,7 +1495,7 @@ class TestLateProvenanceBox:
assert has_ai_metadata(self._mp4_late_c2pa(tmp_path)) is True
def test_scan_c2pa_region_non_isobmff_is_empty(self, tmp_path: Path):
from remove_ai_watermarks.noai.isobmff import scan_c2pa_region
from remove_ai_watermarks._internal.isobmff import scan_c2pa_region
p = tmp_path / "not.bin"
p.write_bytes(b"\x89PNG\r\n\x1a\n not an isobmff file")
@@ -1505,7 +1505,7 @@ class TestLateProvenanceBox:
"""A 64-bit largesize (size32 == 1) uuid box must be walked and collected."""
import struct
from remove_ai_watermarks.noai.isobmff import scan_c2pa_region
from remove_ai_watermarks._internal.isobmff import scan_c2pa_region
payload = b"LARGESIZE-C2PA-MANIFEST"
total = 16 + len(payload) # 4 (size32=1) + 4 (type) + 8 (largesize) + payload
@@ -1516,7 +1516,7 @@ class TestLateProvenanceBox:
def test_scan_c2pa_region_caps_at_max_total(self, tmp_path: Path):
"""The collected payload is bounded by ``max_total`` (never unbounded)."""
from remove_ai_watermarks.noai.isobmff import scan_c2pa_region
from remove_ai_watermarks._internal.isobmff import scan_c2pa_region
p = tmp_path / "big.mp4"
p.write_bytes(_MP4_FTYP + _box(b"uuid", b"A" * 5000))
@@ -1550,7 +1550,7 @@ class TestMetaBoxXmpBlanking:
in place (same length -> iloc offsets and image data stay intact)."""
def test_blanks_ai_packet_only(self):
from remove_ai_watermarks.noai.isobmff import blank_ai_xmp_packets
from remove_ai_watermarks._internal.isobmff import blank_ai_xmp_packets
before, after = b"IMG_BEFORE" * 4, b"IMG_AFTER" * 4
data = before + _AI_XMP + after + _PLAIN_XMP
@@ -1563,13 +1563,13 @@ class TestMetaBoxXmpBlanking:
assert b"dc:rights" in out # plain XMP left alone
def test_no_packet_is_noop(self):
from remove_ai_watermarks.noai.isobmff import blank_ai_xmp_packets
from remove_ai_watermarks._internal.isobmff import blank_ai_xmp_packets
data = b"just some mdat bytes, no xmp here"
assert blank_ai_xmp_packets(data) == (data, 0)
def test_plain_xmp_untouched(self):
from remove_ai_watermarks.noai.isobmff import blank_ai_xmp_packets
from remove_ai_watermarks._internal.isobmff import blank_ai_xmp_packets
out, n = blank_ai_xmp_packets(_PLAIN_XMP)
assert n == 0
@@ -1600,7 +1600,7 @@ class TestIsobmffMetadataRemoval:
def test_strips_ai_xmp_uuid_box(self):
# A uuid box carrying a TC260 AIGC label is dropped by content match,
# regardless of the (non-C2PA) XMP UUID's byte order.
from remove_ai_watermarks.noai.isobmff import strip_c2pa_boxes
from remove_ai_watermarks._internal.isobmff import strip_c2pa_boxes
xmp_uuid = bytes(range(16)) # arbitrary, not the C2PA UUID
payload = b'<x:xmpmeta><TC260:AIGC>{"Label":"1"}</TC260:AIGC></x:xmpmeta>'
@@ -1611,7 +1611,7 @@ class TestIsobmffMetadataRemoval:
def test_keeps_plain_non_ai_xmp(self):
# A uuid box with ordinary (non-AI) XMP must be preserved.
from remove_ai_watermarks.noai.isobmff import strip_c2pa_boxes
from remove_ai_watermarks._internal.isobmff import strip_c2pa_boxes
xmp_uuid = bytes(range(16))
payload = b"<x:xmpmeta><dc:rights>(c) me</dc:rights></x:xmpmeta>"
tests/test_noai.py → tests/test_metadata_internals.py
+27 -27
View File
@@ -1,5 +1,5 @@
"""Tests for vendored noai submodules: constants, extractor, c2pa, plus the
consolidated metadata strip (formerly noai.cleaner)."""
"""Tests for metadata compatibility submodules: constants, extractor, C2PA, plus the
consolidated metadata strip (formerly legacy metadata helper)."""
from __future__ import annotations
@@ -8,10 +8,7 @@ from pathlib import Path
import pytest
from remove_ai_watermarks.metadata import (
remove_ai_metadata as noai_remove_ai_metadata,
)
from remove_ai_watermarks.noai.c2pa import (
from remove_ai_watermarks._internal.c2pa import (
_parse_c2pa_chunk,
cbor_text_after,
extract_c2pa_chunk,
@@ -20,24 +17,27 @@ from remove_ai_watermarks.noai.c2pa import (
inject_c2pa_chunk,
synthid_verdict,
)
from remove_ai_watermarks.noai.constants import (
from remove_ai_watermarks._internal.constants import (
AI_KEYWORDS,
AI_METADATA_KEYS,
C2PA_CHUNK_TYPE,
PNG_SIGNATURE,
SUPPORTED_FORMATS,
)
from remove_ai_watermarks.noai.extractor import (
from remove_ai_watermarks._internal.extractor import (
extract_ai_metadata,
extract_metadata,
get_ai_metadata_summary,
has_ai_metadata,
)
from remove_ai_watermarks.noai.isobmff import (
from remove_ai_watermarks._internal.isobmff import (
blank_ai_exif_tokens,
is_isobmff,
strip_c2pa_boxes,
)
from remove_ai_watermarks.metadata import (
remove_ai_metadata as remove_metadata,
)
# ── Constants ───────────────────────────────────────────────────────
@@ -77,7 +77,7 @@ class TestConstants:
class TestExtractor:
"""Tests for noai.extractor functions."""
"""Tests for internal metadata extraction helpers."""
def test_extract_metadata_returns_dict(self, tmp_clean_png):
meta = extract_metadata(tmp_clean_png)
@@ -115,11 +115,11 @@ class TestExtractor:
class TestCleaner:
"""Metadata stripping via the single, consolidated ``metadata.remove_ai_metadata``
(the legacy ``noai.cleaner`` duplicate was retired)."""
(the legacy ``legacy metadata helper`` duplicate was retired)."""
def test_remove_ai_metadata(self, tmp_png_with_ai_metadata, tmp_path):
output = tmp_path / "cleaned.png"
noai_remove_ai_metadata(tmp_png_with_ai_metadata, output)
remove_metadata(tmp_png_with_ai_metadata, output)
assert output.exists()
# Verify AI metadata removed
meta = extract_ai_metadata(output)
@@ -201,7 +201,7 @@ class TestC2PARealSamples:
def test_extract_info_uses_reader_store(self):
"""The c2pa-python reader path: structured (not heuristic) extraction."""
from remove_ai_watermarks.noai import c2pa
from remove_ai_watermarks._internal import c2pa
assert c2pa.reader_available()
info = extract_c2pa_info(SAMPLES_DIR / "chatgpt-1.png")
@@ -213,7 +213,7 @@ class TestC2PARealSamples:
def test_fallback_to_png_parser_when_reader_unavailable(self, monkeypatch):
"""With the reader disabled, the hand-rolled PNG parser still works."""
from remove_ai_watermarks.noai import c2pa
from remove_ai_watermarks._internal import c2pa
monkeypatch.setattr(c2pa, "_C2PA_READER_AVAILABLE", False)
info = extract_c2pa_info(SAMPLES_DIR / "chatgpt-1.png")
@@ -326,7 +326,7 @@ class TestC2PADigitalSourceType:
(AI-enhanced) and a bare procedural ``algorithmicMedia`` token must classify as
AI-enhanced. Before the reorder the bare-token elif fired first and returned
non-AI, dropping the composite AI signal (a false negative)."""
from remove_ai_watermarks.noai.c2pa import _populate_registry_fields
from remove_ai_watermarks._internal.c2pa import _populate_registry_fields
info: dict = {}
_populate_registry_fields(b"x compositeWithTrainedAlgorithmicMedia x algorithmicMedia x", info)
@@ -453,7 +453,7 @@ class TestISOBMFF:
assert out == FTYP + b"\x00\x00\x00\x0cmdat" + b"pixels!!"
def test_streaming_malformed_walk_copies_input_unchanged(self, tmp_path: Path):
from remove_ai_watermarks.noai.isobmff import strip_isobmff_media_file
from remove_ai_watermarks._internal.isobmff import strip_isobmff_media_file
source = tmp_path / "malformed.mp4"
output = tmp_path / "clean.mp4"
@@ -471,7 +471,7 @@ class TestISOBMFF:
monkeypatch: pytest.MonkeyPatch,
):
from remove_ai_watermarks import metadata
from remove_ai_watermarks.noai import isobmff
from remove_ai_watermarks._internal import isobmff
source = tmp_path / "source.mp4"
output = tmp_path / "clean.mp4"
@@ -495,7 +495,7 @@ class TestIterTopLevelBoxes:
"""The box walker's three size encodings and its underflow/overflow guards."""
def test_64bit_largesize(self):
from remove_ai_watermarks.noai.isobmff import _iter_top_level_boxes
from remove_ai_watermarks._internal.isobmff import _iter_top_level_boxes
# size32 == 1 -> a 64-bit largesize follows the type; total box length = 24.
box = struct.pack(">I", 1) + b"uuid" + struct.pack(">Q", 24) + b"payload!"
@@ -505,7 +505,7 @@ class TestIterTopLevelBoxes:
assert (start, end, btype, payload_off) == (0, 24, b"uuid", 16)
def test_size0_runs_to_eof(self):
from remove_ai_watermarks.noai.isobmff import _iter_top_level_boxes
from remove_ai_watermarks._internal.isobmff import _iter_top_level_boxes
box = struct.pack(">I", 0) + b"mdat" + b"tail-to-eof"
boxes = list(_iter_top_level_boxes(box))
@@ -514,13 +514,13 @@ class TestIterTopLevelBoxes:
assert (start, end, btype, payload_off) == (0, len(box), b"mdat", 8)
def test_underflow_size_stops_safely(self):
from remove_ai_watermarks.noai.isobmff import _iter_top_level_boxes
from remove_ai_watermarks._internal.isobmff import _iter_top_level_boxes
# size (4) < the 8-byte header -> the guard returns without yielding a box.
assert list(_iter_top_level_boxes(struct.pack(">I", 4) + b"ftyp" + b"more")) == []
def test_overflow_size_stops_safely(self):
from remove_ai_watermarks.noai.isobmff import _iter_top_level_boxes
from remove_ai_watermarks._internal.isobmff import _iter_top_level_boxes
# size claims 999 but the buffer is far shorter -> guard returns, no partial box.
assert list(_iter_top_level_boxes(struct.pack(">I", 999) + b"uuid" + b"x")) == []
@@ -533,7 +533,7 @@ class TestBlankAiXmpPackets:
AIMARK = b"trainedAlgorithmicMedia"
def test_ai_packet_blanked_same_length(self):
from remove_ai_watermarks.noai.isobmff import blank_ai_xmp_packets
from remove_ai_watermarks._internal.isobmff import blank_ai_xmp_packets
packet = b'<?xpacket begin="x"?><x:xmpmeta>' + self.AIMARK + b'</x:xmpmeta><?xpacket end="w"?>'
data = b"boxhdr" + packet + b"tail"
@@ -545,7 +545,7 @@ class TestBlankAiXmpPackets:
assert b"tail" in out
def test_clean_packet_left_intact(self):
from remove_ai_watermarks.noai.isobmff import blank_ai_xmp_packets
from remove_ai_watermarks._internal.isobmff import blank_ai_xmp_packets
packet = b'<?xpacket begin="x"?><x:xmpmeta>plain copyright</x:xmpmeta><?xpacket end="w"?>'
out, n = blank_ai_xmp_packets(packet)
@@ -553,7 +553,7 @@ class TestBlankAiXmpPackets:
assert out == packet
def test_missing_end_delimiter_not_blanked(self):
from remove_ai_watermarks.noai.isobmff import blank_ai_xmp_packets
from remove_ai_watermarks._internal.isobmff import blank_ai_xmp_packets
# No <?xpacket end?> -> the packet regex cannot match, so it is left unchanged.
data = b'<?xpacket begin="x"?><x:xmpmeta>' + self.AIMARK + b"</x:xmpmeta>"
@@ -568,7 +568,7 @@ class TestC2paBufferScans:
as vendors are added."""
def test_soft_binding_vendors_in(self):
from remove_ai_watermarks.noai.c2pa import C2PA_SOFT_BINDINGS, soft_binding_vendors_in
from remove_ai_watermarks._internal.c2pa import C2PA_SOFT_BINDINGS, soft_binding_vendors_in
sig, name = next(iter(C2PA_SOFT_BINDINGS.items()))
assert name in soft_binding_vendors_in(b"...manifest..." + sig + b"...tail...")
@@ -576,7 +576,7 @@ class TestC2paBufferScans:
assert soft_binding_vendors_in(b"no soft-binding assertion here") == []
def test_synthid_vendors_in_requires_synthid_issuer(self):
from remove_ai_watermarks.noai.c2pa import C2PA_ISSUERS, SYNTHID_C2PA_ISSUERS, synthid_vendors_in
from remove_ai_watermarks._internal.c2pa import C2PA_ISSUERS, SYNTHID_C2PA_ISSUERS, synthid_vendors_in
syn_sig = next(s for s in C2PA_ISSUERS if s in SYNTHID_C2PA_ISSUERS)
non_sig = next(s for s in C2PA_ISSUERS if s not in SYNTHID_C2PA_ISSUERS)
@@ -585,7 +585,7 @@ class TestC2paBufferScans:
assert C2PA_ISSUERS[non_sig] not in synthid_vendors_in(b"x" + non_sig + b"x")
def test_synthid_verdict_format(self):
from remove_ai_watermarks.noai.c2pa import synthid_verdict
from remove_ai_watermarks._internal.c2pa import synthid_verdict
assert synthid_verdict("Google LLC") == "likely present (Google LLC embeds SynthID with C2PA)"
tests/test_platform.py
+29 -30
View File
@@ -13,9 +13,9 @@ import numpy as np
import pytest
from PIL import Image
from remove_ai_watermarks.noai.progress import is_mps_error
from remove_ai_watermarks.noai.utils import get_image_format, is_supported_format
from remove_ai_watermarks.noai.watermark_profiles import (
from remove_ai_watermarks._internal.progress import is_mps_error
from remove_ai_watermarks._internal.utils import get_image_format, is_supported_format
from remove_ai_watermarks._internal.watermark_profiles import (
DEFAULT_STRENGTH,
GEMINI_STRENGTH,
OPENAI_STRENGTH,
@@ -24,7 +24,7 @@ from remove_ai_watermarks.noai.watermark_profiles import (
resolve_strength,
strength_default_help,
)
from remove_ai_watermarks.noai.watermark_remover import get_device, is_watermark_removal_available
from remove_ai_watermarks._internal.watermark_remover import get_device, is_watermark_removal_available
# ── Device detection ────────────────────────────────────────────────
@@ -42,7 +42,7 @@ class TestDeviceDetection:
# Just verify it doesn't crash and returns a valid string
assert isinstance(device, str)
@patch("remove_ai_watermarks.noai.watermark_remover._HAS_TORCH", False)
@patch("remove_ai_watermarks._internal.watermark_remover._HAS_TORCH", False)
def test_no_torch_returns_cpu(self):
assert get_device() == "cpu"
@@ -55,7 +55,7 @@ class TestDeviceDetection:
fake_torch = MagicMock()
fake_torch.cuda.is_available.return_value = False
fake_torch.xpu.is_available.return_value = True
with patch("remove_ai_watermarks.noai.watermark_remover.torch", fake_torch):
with patch("remove_ai_watermarks._internal.watermark_remover.torch", fake_torch):
assert get_device() == "xpu"
fake_torch.tensor.assert_called_with([1.0], device="xpu")
@@ -65,7 +65,7 @@ class TestDeviceDetection:
pytest.skip("torch/diffusers not installed")
import torch
from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
remover = WatermarkRemover(device="xpu")
assert remover.device == "xpu"
@@ -74,7 +74,7 @@ class TestDeviceDetection:
def test_seed_generator_falls_back_to_cpu_when_device_rng_unsupported(self):
"""A device with no RNG backend (e.g. some torch-xpu builds) falls back
to a CPU generator instead of raising when --seed is used."""
from remove_ai_watermarks.noai import watermark_remover as wr
from remove_ai_watermarks._internal import watermark_remover as wr
def fake_generator(device="cpu"):
if device == "xpu":
@@ -138,7 +138,7 @@ class TestFp16WeightVariant:
def _remover(self, dtype: object):
if not is_watermark_removal_available():
pytest.skip("torch/diffusers not installed")
from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
# device="cpu" alone would force fp32; the explicit torch_dtype override lets us
# exercise the fp16 path with no GPU (construction loads no weights).
@@ -192,12 +192,12 @@ class TestNoReembeddedWatermark:
def _remover(self, profile: str):
if not is_watermark_removal_available():
pytest.skip("torch/diffusers not installed")
from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
return WatermarkRemover(device="cpu", pipeline=profile)
def _capture(self, monkeypatch, remover):
from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
calls: list[tuple[str, dict]] = []
@@ -242,7 +242,7 @@ class TestQwenKwargs:
"""
def test_uses_true_cfg_not_guidance_scale(self):
from remove_ai_watermarks.noai.watermark_remover import _build_qwen_kwargs
from remove_ai_watermarks._internal.watermark_remover import _build_qwen_kwargs
gen = object()
img = _StubImage(2816, 1536)
@@ -254,14 +254,13 @@ class TestQwenKwargs:
assert kwargs["strength"] == 0.3
assert kwargs["image"] is img
assert kwargs["generator"] is gen
# Faithful-regeneration prompt + an explicit negative prompt.
assert kwargs["prompt"]
assert kwargs["negative_prompt"]
assert kwargs["prompt"] == "high quality, sharp, detailed, faithful to the original"
assert kwargs["negative_prompt"] == "blurry, lowres, distorted text, garbled text, artifacts"
def test_passes_explicit_aspect_preserving_size(self):
# Without height/width the pipeline defaults to 1024x1024 and squishes non-square
# input (the abba mixed-seam regression). Both already multiples of 16 -> unchanged.
from remove_ai_watermarks.noai.watermark_remover import _build_qwen_kwargs
from remove_ai_watermarks._internal.watermark_remover import _build_qwen_kwargs
kwargs = _build_qwen_kwargs(
_StubImage(2816, 1536), strength=0.25, num_inference_steps=40, true_cfg_scale=4.0, generator=None
@@ -270,14 +269,14 @@ class TestQwenKwargs:
assert kwargs["height"] == 1536
def test_qwen_target_size_floors_to_multiple_of_16(self):
from remove_ai_watermarks.noai.watermark_remover import _qwen_target_size
from remove_ai_watermarks._internal.watermark_remover import _qwen_target_size
assert _qwen_target_size(2816, 1536) == (2816, 1536) # already /16
assert _qwen_target_size(1122, 1402) == (1120, 1392) # floored
assert _qwen_target_size(10, 10) == (16, 16) # min clamp, never 0
def test_qwen_model_id_is_qwen_image(self):
from remove_ai_watermarks.noai.watermark_profiles import QWEN_MODEL_ID
from remove_ai_watermarks._internal.watermark_profiles import QWEN_MODEL_ID
assert QWEN_MODEL_ID == "Qwen/Qwen-Image"
@@ -303,7 +302,7 @@ class TestResolveStrength:
# Qwen's certified Gemini floor (0.25) is HIGHER than controlnet's (0.15); OpenAI
# matches (0.10). Unknown vendor on qwen tracks the higher Gemini value. This retires
# the old manual "pass --strength 0.25 for Gemini on qwen" workaround.
from remove_ai_watermarks.noai.watermark_profiles import QWEN_GEMINI_STRENGTH, QWEN_OPENAI_STRENGTH
from remove_ai_watermarks._internal.watermark_profiles import QWEN_GEMINI_STRENGTH, QWEN_OPENAI_STRENGTH
assert QWEN_GEMINI_STRENGTH == 0.25
assert QWEN_OPENAI_STRENGTH == 0.10
@@ -354,32 +353,32 @@ class TestVendorForStrength:
return patch("remove_ai_watermarks.metadata.synthid_source", return_value=value)
def test_openai(self):
from remove_ai_watermarks.noai.watermark_profiles import vendor_for_strength
from remove_ai_watermarks._internal.watermark_profiles import vendor_for_strength
with self._patch("OpenAI"):
assert vendor_for_strength(Path("x.png")) == "openai"
def test_google(self):
from remove_ai_watermarks.noai.watermark_profiles import vendor_for_strength
from remove_ai_watermarks._internal.watermark_profiles import vendor_for_strength
with self._patch("Google"):
assert vendor_for_strength(Path("x.png")) == "google"
def test_both_issuers_google_wins(self):
# The more-robust watermark wins -> safer (higher) strength.
from remove_ai_watermarks.noai.watermark_profiles import vendor_for_strength
from remove_ai_watermarks._internal.watermark_profiles import vendor_for_strength
with self._patch("OpenAI, Google"):
assert vendor_for_strength(Path("x.png")) == "google"
def test_none_when_no_synthid_source(self):
from remove_ai_watermarks.noai.watermark_profiles import vendor_for_strength
from remove_ai_watermarks._internal.watermark_profiles import vendor_for_strength
with self._patch(None):
assert vendor_for_strength(Path("x.png")) is None
def test_unreadable_metadata_is_none(self):
from remove_ai_watermarks.noai.watermark_profiles import vendor_for_strength
from remove_ai_watermarks._internal.watermark_profiles import vendor_for_strength
with patch("remove_ai_watermarks.metadata.synthid_source", side_effect=OSError):
assert vendor_for_strength(Path("x.png")) is None
@@ -478,19 +477,19 @@ class TestFp16VaeFix:
DEFAULT = "stabilityai/stable-diffusion-xl-base-1.0"
def test_default_sdxl_on_fp16_needs_fix(self):
from remove_ai_watermarks.noai.watermark_remover import _needs_fp16_vae_fix
from remove_ai_watermarks._internal.watermark_remover import _needs_fp16_vae_fix
assert _needs_fp16_vae_fix(self.DEFAULT, self.DEFAULT, is_fp16=True) is True
def test_fp32_does_not_need_fix(self):
"""cpu/mps run fp32, where the stock SDXL VAE is fine."""
from remove_ai_watermarks.noai.watermark_remover import _needs_fp16_vae_fix
from remove_ai_watermarks._internal.watermark_remover import _needs_fp16_vae_fix
assert _needs_fp16_vae_fix(self.DEFAULT, self.DEFAULT, is_fp16=False) is False
def test_non_default_model_keeps_own_vae(self):
"""A custom (non-SDXL) checkpoint must not get the SDXL-specific VAE."""
from remove_ai_watermarks.noai.watermark_remover import _needs_fp16_vae_fix
from remove_ai_watermarks._internal.watermark_remover import _needs_fp16_vae_fix
assert _needs_fp16_vae_fix("runwayml/stable-diffusion-v1-5", self.DEFAULT, is_fp16=True) is False
@@ -500,13 +499,13 @@ class TestDegenerateOutputGuard:
``remove_watermark`` can retry in fp32. Pure image statistics, no model needed."""
def test_all_black_is_degenerate(self):
from remove_ai_watermarks.noai.watermark_remover import _is_degenerate_image
from remove_ai_watermarks._internal.watermark_remover import _is_degenerate_image
black = Image.fromarray(np.zeros((64, 64, 3), np.uint8))
assert _is_degenerate_image(black) is True
def test_normal_image_is_not_degenerate(self):
from remove_ai_watermarks.noai.watermark_remover import _is_degenerate_image
from remove_ai_watermarks._internal.watermark_remover import _is_degenerate_image
rng = np.random.default_rng(0)
normal = Image.fromarray(rng.integers(0, 256, (64, 64, 3), dtype=np.uint8))
@@ -514,7 +513,7 @@ class TestDegenerateOutputGuard:
def test_dark_but_textured_image_is_not_degenerate(self):
"""A legitimately dark photo with real detail must NOT be flagged (variance guard)."""
from remove_ai_watermarks.noai.watermark_remover import _is_degenerate_image
from remove_ai_watermarks._internal.watermark_remover import _is_degenerate_image
rng = np.random.default_rng(1)
dark = Image.fromarray(rng.integers(0, 40, (64, 64, 3), dtype=np.uint8))
tests/test_qwen_zimage_pipeline.py
+50 -44
View File
@@ -13,7 +13,7 @@ from PIL import Image
def _mock_watermark_runtime_deps(monkeypatch):
"""Bypass optional GPU imports while testing Qwen Z-Image routing."""
from remove_ai_watermarks.noai import watermark_remover
from remove_ai_watermarks._internal import watermark_remover
fake_torch = MagicMock()
fake_torch.float16 = object()
@@ -23,26 +23,24 @@ def _mock_watermark_runtime_deps(monkeypatch):
monkeypatch.setattr(watermark_remover, "is_watermark_removal_available", lambda: True)
def test_resolution_adaptive_denoise_matches_reference_formula():
from remove_ai_watermarks.noai.qwen_zimage_pipeline import resolution_adaptive_denoise
def test_resolution_adaptive_denoise_preserves_calibrated_values():
from remove_ai_watermarks._internal.qwen_zimage_pipeline import resolution_adaptive_denoise
# The reference node maps 0.30 MP to the lower bound and 3.70 MP to the
# upper bound. Level 6 adds one fifth of the configured upward spread.
assert resolution_adaptive_denoise(600, 500, adaptive_level=6) == pytest.approx(0.084)
assert resolution_adaptive_denoise(2000, 1850, adaptive_level=6) == pytest.approx(0.154)
def test_largest_face_denoise_matches_reference_formula():
from remove_ai_watermarks.noai.qwen_zimage_pipeline import largest_face_denoise
def test_largest_face_denoise_preserves_calibrated_values():
from remove_ai_watermarks._internal.qwen_zimage_pipeline import largest_face_denoise
image_size = (1000, 1000)
assert largest_face_denoise([(0, 0, 300, 100)], image_size) == 0.10
assert largest_face_denoise([(0, 0, 150, 100)], image_size) == 0.05
assert largest_face_denoise([(0, 0, 900, 900)], image_size) == 0.28
assert largest_face_denoise([(0, 0, 300, 100)], image_size) == pytest.approx(0.10)
assert largest_face_denoise([(0, 0, 150, 100)], image_size) == pytest.approx(0.05)
assert largest_face_denoise([(0, 0, 900, 900)], image_size) == pytest.approx(0.28)
def test_global_kwargs_use_lightning_and_diffsynth_controlnet_shape():
from remove_ai_watermarks.noai.qwen_zimage_pipeline import build_global_kwargs
from remove_ai_watermarks._internal.qwen_zimage_pipeline import build_global_kwargs
image = Image.new("RGB", (1122, 1402))
kwargs = build_global_kwargs(image, strength=0.11, seed=7, controlnet_input="CONTROL")
@@ -56,10 +54,12 @@ def test_global_kwargs_use_lightning_and_diffsynth_controlnet_shape():
assert kwargs["width"] == 1120
assert kwargs["height"] == 1392
assert kwargs["exponential_shift_mu"] == pytest.approx(math.log(3.0))
assert kwargs["prompt"] == "ultra clear and smoothe skin, spotless skin"
assert kwargs["negative_prompt"] == "moles, freckes, high detail skin"
def test_face_kwargs_use_zimage_reference_settings():
from remove_ai_watermarks.noai.qwen_zimage_pipeline import build_face_kwargs
def test_face_kwargs_use_project_zimage_settings():
from remove_ai_watermarks._internal.qwen_zimage_pipeline import build_face_kwargs
crop = Image.new("RGB", (713, 941))
kwargs = build_face_kwargs(crop, strength=0.17, seed=9)
@@ -71,10 +71,12 @@ def test_face_kwargs_use_zimage_reference_settings():
assert kwargs["seed"] == 9
assert kwargs["width"] == 704
assert kwargs["height"] == 928
assert kwargs["prompt"] == ""
assert kwargs["negative_prompt"] == "blurry, ugly, bad quality,"
def test_canny_control_image_matches_reference_thresholds():
from remove_ai_watermarks.noai.qwen_zimage_pipeline import build_canny_control_image
def test_canny_control_image_is_three_channel_and_detects_an_edge():
from remove_ai_watermarks._internal.qwen_zimage_pipeline import build_canny_control_image
source = np.zeros((64, 80, 3), dtype=np.uint8)
source[:, 40:] = 255
@@ -83,11 +85,21 @@ def test_canny_control_image_matches_reference_thresholds():
assert result.shape == (64, 80, 3)
assert np.array_equal(result[:, :, 0], result[:, :, 1])
assert np.array_equal(result[:, :, 1], result[:, :, 2])
assert result.max() == 255
import cv2
expected = cv2.Canny(cv2.cvtColor(source, cv2.COLOR_RGB2GRAY), 13, 64)
assert np.array_equal(result[:, :, 0], expected)
def test_face_crop_geometry_preserves_calibrated_values():
from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline, _expanded_box
assert _expanded_box((100, 100, 200, 200), (500, 500)) == (25, 25, 275, 275)
assert QwenZImagePipeline._detail_size((500, 400), (100, 80)) == (1024, 816)
def test_yunet_download_targets_verified_lfs_artifact():
from remove_ai_watermarks.noai.qwen_zimage_pipeline import (
from remove_ai_watermarks._internal.qwen_zimage_pipeline import (
YUNET_MODEL_SHA256,
YUNET_MODEL_URL,
YUNET_SCORE_THRESHOLD,
@@ -95,14 +107,11 @@ def test_yunet_download_targets_verified_lfs_artifact():
assert YUNET_MODEL_URL.startswith("https://media.githubusercontent.com/media/opencv/opencv_zoo/")
assert YUNET_MODEL_SHA256 == "8f2383e4dd3cfbb4553ea8718107fc0423210dc964f9f4280604804ed2552fa4"
# YuNet scores are not calibrated like the upstream YOLO detector's scores.
# A 0.2 YuNet threshold admitted background and decorative false positives,
# multiplying the serial Z-Image face-stage cost on crowded scenes.
assert pytest.approx(0.5) == YUNET_SCORE_THRESHOLD
def test_resident_face_models_disable_vram_offload():
from remove_ai_watermarks.noai.qwen_zimage_pipeline import (
from remove_ai_watermarks._internal.qwen_zimage_pipeline import (
QwenZImagePipeline,
_pin_vram_managed_models,
resolve_face_model_residency,
@@ -154,7 +163,7 @@ def test_resident_face_models_disable_vram_offload():
def test_static_prompt_cache_reuses_embeddings_without_caching_image_edits():
from remove_ai_watermarks.noai.qwen_zimage_pipeline import _cache_static_prompt_embeddings
from remove_ai_watermarks._internal.qwen_zimage_pipeline import _cache_static_prompt_embeddings
class PromptUnit:
output_params = ("prompt_embeds",)
@@ -186,7 +195,7 @@ def test_static_prompt_cache_reuses_embeddings_without_caching_image_edits():
def test_sam_pixels_match_model_dtype_without_casting_boxes():
import torch
from remove_ai_watermarks.noai.qwen_zimage_pipeline import _prepare_sam_inputs
from remove_ai_watermarks._internal.qwen_zimage_pipeline import _prepare_sam_inputs
class Inputs(dict[str, torch.Tensor]):
def to(self, device: str):
@@ -206,7 +215,7 @@ def test_sam_pixels_match_model_dtype_without_casting_boxes():
def test_sam_prompts_match_impact_center_and_clip_masks_to_boxes():
from remove_ai_watermarks.noai.qwen_zimage_pipeline import (
from remove_ai_watermarks._internal.qwen_zimage_pipeline import (
_clip_sam_masks_to_boxes,
_sam_point_prompts,
)
@@ -226,7 +235,7 @@ def test_sam_prompts_match_impact_center_and_clip_masks_to_boxes():
def test_sam_proposal_selection_matches_impact_sub_threshold():
from remove_ai_watermarks.noai.qwen_zimage_pipeline import _select_sam_masks
from remove_ai_watermarks._internal.qwen_zimage_pipeline import _select_sam_masks
masks = np.zeros((2, 3, 8, 8), dtype=np.float32)
masks[0, 0, 1:3, 1:3] = 1.0
@@ -256,7 +265,7 @@ def test_sam_proposal_selection_matches_impact_sub_threshold():
def test_sam_bfloat16_outputs_convert_to_numpy_float32():
import torch
from remove_ai_watermarks.noai.qwen_zimage_pipeline import _sam_outputs_to_numpy
from remove_ai_watermarks._internal.qwen_zimage_pipeline import _sam_outputs_to_numpy
masks = torch.ones((1, 2, 3, 4, 4), dtype=torch.bfloat16)
scores = torch.tensor([[[0.95, 0.75, 0.50], [0.99, 0.80, 0.60]]], dtype=torch.bfloat16)
@@ -269,7 +278,7 @@ def test_sam_bfloat16_outputs_convert_to_numpy_float32():
def test_face_composite_preserves_every_pixel_outside_mask():
from remove_ai_watermarks.noai.qwen_zimage_pipeline import composite_face
from remove_ai_watermarks._internal.qwen_zimage_pipeline import composite_face
base = np.full((32, 32, 3), 10, dtype=np.uint8)
detail = np.full((32, 32, 3), 240, dtype=np.uint8)
@@ -284,7 +293,7 @@ def test_face_composite_preserves_every_pixel_outside_mask():
def test_profile_defaults_to_four_global_steps():
from remove_ai_watermarks.noai.watermark_profiles import (
from remove_ai_watermarks._internal.watermark_profiles import (
normalize_profile,
resolve_seed,
resolve_steps,
@@ -305,7 +314,7 @@ def test_cli_exposes_qwen_zimage_profile():
assert "qwen-zimage" in _PIPELINE_CHOICES
def test_cli_qwen_zimage_keeps_upstream_postprocess_default(tmp_image_path, monkeypatch):
def test_cli_qwen_zimage_keeps_profile_postprocess_default(tmp_image_path, monkeypatch):
from remove_ai_watermarks import cli
mock_engine = MagicMock()
@@ -338,7 +347,7 @@ def test_cli_qwen_zimage_keeps_upstream_postprocess_default(tmp_image_path, monk
def test_watermark_remover_dispatches_to_full_pipeline(tmp_path, monkeypatch):
from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
_mock_watermark_runtime_deps(monkeypatch)
source = tmp_path / "source.png"
@@ -364,7 +373,7 @@ def test_watermark_remover_dispatches_to_full_pipeline(tmp_path, monkeypatch):
def test_watermark_remover_dispatches_qwen_tiling_to_full_pipeline(tmp_path, monkeypatch):
from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
_mock_watermark_runtime_deps(monkeypatch)
source = tmp_path / "source.png"
@@ -395,11 +404,11 @@ def test_watermark_remover_dispatches_qwen_tiling_to_full_pipeline(tmp_path, mon
def test_qwen_tiling_runs_global_tiles_then_one_full_frame_face_stage(monkeypatch):
from remove_ai_watermarks.noai.qwen_zimage_pipeline import (
from remove_ai_watermarks._internal.qwen_zimage_pipeline import (
QwenZImagePipeline,
resolution_adaptive_denoise,
)
from remove_ai_watermarks.noai.tiling import plan_tiles
from remove_ai_watermarks._internal.tiling import plan_tiles
image = Image.new("RGB", (1500, 1500), (20, 30, 40))
runtime = QwenZImagePipeline(device="cuda", torch_dtype="bf16")
@@ -415,7 +424,7 @@ def test_qwen_tiling_runs_global_tiles_then_one_full_frame_face_stage(monkeypatc
monkeypatch.setattr(runtime, "_run_global", fake_global)
monkeypatch.setattr(runtime, "_run_faces", face_stage)
monkeypatch.setattr(
"remove_ai_watermarks.noai.qwen_zimage_pipeline.detect_faces",
"remove_ai_watermarks._internal.qwen_zimage_pipeline.detect_faces",
lambda _image: [(100, 100, 300, 300)],
)
monkeypatch.setattr(runtime, "_sam_masks", lambda _image, _boxes: [np.ones((1500, 1500), dtype=np.uint8)])
@@ -435,9 +444,6 @@ def test_qwen_tiling_runs_global_tiles_then_one_full_frame_face_stage(monkeypatc
assert all(strength == pytest.approx(resolution_adaptive_denoise(1500, 1500)) for _, strength, _ in global_calls)
assert all(seed == 0 for _, _, seed in global_calls)
face_stage.assert_called_once()
# Keep this literal independent from the runtime helper: the port deliberately
# uses half the upstream face denoise because it lacks the reference latent
# noise-mask feather and uses a different sampler/runtime.
assert face_stage.call_args.kwargs["strength"] == pytest.approx(0.0296296296)
assert face_stage.call_args.args[0] is image
assert face_stage.call_args.args[1].size == image.size
@@ -445,7 +451,7 @@ def test_qwen_tiling_runs_global_tiles_then_one_full_frame_face_stage(monkeypatc
def test_global_only_preload_skips_face_models(monkeypatch):
from remove_ai_watermarks.noai.qwen_zimage_pipeline import QwenZImagePipeline
from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline
runtime = QwenZImagePipeline(device="cuda", torch_dtype="bf16")
qwen = MagicMock()
@@ -456,7 +462,7 @@ def test_global_only_preload_skips_face_models(monkeypatch):
monkeypatch.setattr(runtime, "_load_zimage", zimage)
monkeypatch.setattr(runtime, "_load_sam", sam)
monkeypatch.setattr(
"remove_ai_watermarks.noai.qwen_zimage_pipeline._yunet_model_path",
"remove_ai_watermarks._internal.qwen_zimage_pipeline._yunet_model_path",
yunet,
)
@@ -469,7 +475,7 @@ def test_global_only_preload_skips_face_models(monkeypatch):
def test_full_preload_still_loads_face_models(monkeypatch):
from remove_ai_watermarks.noai.qwen_zimage_pipeline import QwenZImagePipeline
from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline
runtime = QwenZImagePipeline(device="cuda", torch_dtype="bf16")
qwen = MagicMock()
@@ -480,7 +486,7 @@ def test_full_preload_still_loads_face_models(monkeypatch):
monkeypatch.setattr(runtime, "_load_zimage", zimage)
monkeypatch.setattr(runtime, "_load_sam", sam)
monkeypatch.setattr(
"remove_ai_watermarks.noai.qwen_zimage_pipeline._yunet_model_path",
"remove_ai_watermarks._internal.qwen_zimage_pipeline._yunet_model_path",
yunet,
)
@@ -493,7 +499,7 @@ def test_full_preload_still_loads_face_models(monkeypatch):
def test_watermark_remover_forwards_global_only_preload(monkeypatch):
from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
runtime = MagicMock()
remover = WatermarkRemover.__new__(WatermarkRemover)
@@ -506,7 +512,7 @@ def test_watermark_remover_forwards_global_only_preload(monkeypatch):
def test_qwen_zimage_rejects_runtime_knobs_that_change_fixed_graph(tmp_path, monkeypatch):
from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
_mock_watermark_runtime_deps(monkeypatch)
with pytest.raises(ValueError, match="fixed Qwen-Image-2512"):
@@ -517,7 +523,7 @@ def test_qwen_zimage_rejects_runtime_knobs_that_change_fixed_graph(tmp_path, mon
remover = WatermarkRemover(device="cpu", pipeline="qwen-zimage")
with pytest.raises(ValueError, match=r"CFG 1\.0"):
remover.remove_watermark(source, guidance_scale=2.0)
with pytest.raises(ValueError, match="4-step Lightning"):
with pytest.raises(ValueError, match="requires 4 steps"):
remover.remove_watermark(source, num_inference_steps=8)
tests/test_security_clamp.py
+1 -1
View File
@@ -20,7 +20,7 @@ import struct
import tracemalloc
from remove_ai_watermarks import metadata
from remove_ai_watermarks.noai import c2pa, isobmff
from remove_ai_watermarks._internal import c2pa, isobmff
PNG_SIG = b"\x89PNG\r\n\x1a\n"
_HUGE = 0x7FFFFFFF # ~2 GiB declared length on a tiny file
tests/test_tiling.py
+1 -1
View File
@@ -12,7 +12,7 @@ import numpy as np
import pytest
from PIL import Image
from remove_ai_watermarks.noai.tiling import (
from remove_ai_watermarks._internal.tiling import (
Tile,
_axis_positions,
feather_region_composite,
tests/test_watermark_profiles.py
+1 -1
View File
@@ -4,7 +4,7 @@ from __future__ import annotations
import pytest
from remove_ai_watermarks.noai.watermark_profiles import resolve_strength, viable_steps
from remove_ai_watermarks._internal.watermark_profiles import resolve_strength, viable_steps
class TestViableSteps: