diff --git a/docs/code-provenance.md b/docs/code-provenance.md new file mode 100644 index 0000000..0d7e98a --- /dev/null +++ b/docs/code-provenance.md @@ -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`. diff --git a/docs/controlnet-removal-pipeline-research.md b/docs/controlnet-removal-pipeline-research.md index 65db754..b9f51d7 100644 --- a/docs/controlnet-removal-pipeline-research.md +++ b/docs/controlnet-removal-pipeline-research.md @@ -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 diff --git a/docs/index.md b/docs/index.md index f096e33..ba3810f 100644 --- a/docs/index.md +++ b/docs/index.md @@ -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. | diff --git a/docs/known-limitations.md b/docs/known-limitations.md index c272dcc..b06d4a6 100644 --- a/docs/known-limitations.md +++ b/docs/known-limitations.md @@ -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 diff --git a/docs/module-internals.md b/docs/module-internals.md index 9c401e7..7caa042 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -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. diff --git a/docs/qwen-improvement-research.md b/docs/qwen-improvement-research.md index 90e2f5e..9bb3c8e 100644 --- a/docs/qwen-improvement-research.md +++ b/docs/qwen-improvement-research.md @@ -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 diff --git a/docs/research-doubao-distillation.md b/docs/research-doubao-distillation.md index 2b668c8..1aaf499 100644 --- a/docs/research-doubao-distillation.md +++ b/docs/research-doubao-distillation.md @@ -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.** diff --git a/docs/synthid.md b/docs/synthid.md index 40d9cc3..4bc9d0c 100644 --- a/docs/synthid.md +++ b/docs/synthid.md @@ -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. diff --git a/pyproject.toml b/pyproject.toml index 7d7811f..5e95ddd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/scripts/metadata_removal_audit.py b/scripts/metadata_removal_audit.py index af7d732..01cb47f 100644 --- a/scripts/metadata_removal_audit.py +++ b/scripts/metadata_removal_audit.py @@ -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__) diff --git a/scripts/smoke_matrix.py b/scripts/smoke_matrix.py index a4ee0b1..3548268 100644 --- a/scripts/smoke_matrix.py +++ b/scripts/smoke_matrix.py @@ -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] diff --git a/scripts/synthid_corpus.py b/scripts/synthid_corpus.py index 81beeba..6e007f4 100644 --- a/scripts/synthid_corpus.py +++ b/scripts/synthid_corpus.py @@ -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() diff --git a/scripts/visible_positives.py b/scripts/visible_positives.py index 4154198..8cba624 100644 --- a/scripts/visible_positives.py +++ b/scripts/visible_positives.py @@ -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" diff --git a/src/remove_ai_watermarks/noai/__init__.py b/src/remove_ai_watermarks/_internal/__init__.py similarity index 74% rename from src/remove_ai_watermarks/noai/__init__.py rename to src/remove_ai_watermarks/_internal/__init__.py index 399ebef..e393103 100644 --- a/src/remove_ai_watermarks/noai/__init__.py +++ b/src/remove_ai_watermarks/_internal/__init__.py @@ -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}") diff --git a/src/remove_ai_watermarks/_internal/c2pa.py b/src/remove_ai_watermarks/_internal/c2pa.py new file mode 100644 index 0000000..f432c90 --- /dev/null +++ b/src/remove_ai_watermarks/_internal/c2pa.py @@ -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 diff --git a/src/remove_ai_watermarks/_internal/constants.py b/src/remove_ai_watermarks/_internal/constants.py new file mode 100644 index 0000000..ef05a8c --- /dev/null +++ b/src/remove_ai_watermarks/_internal/constants.py @@ -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} diff --git a/src/remove_ai_watermarks/noai/ebml.py b/src/remove_ai_watermarks/_internal/ebml.py similarity index 100% rename from src/remove_ai_watermarks/noai/ebml.py rename to src/remove_ai_watermarks/_internal/ebml.py diff --git a/src/remove_ai_watermarks/_internal/extractor.py b/src/remove_ai_watermarks/_internal/extractor.py new file mode 100644 index 0000000..df3b706 --- /dev/null +++ b/src/remove_ai_watermarks/_internal/extractor.py @@ -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"" + 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) diff --git a/src/remove_ai_watermarks/noai/flv.py b/src/remove_ai_watermarks/_internal/flv.py similarity index 100% rename from src/remove_ai_watermarks/noai/flv.py rename to src/remove_ai_watermarks/_internal/flv.py diff --git a/src/remove_ai_watermarks/_internal/img2img_runner.py b/src/remove_ai_watermarks/_internal/img2img_runner.py new file mode 100644 index 0000000..2fcd37c --- /dev/null +++ b/src/remove_ai_watermarks/_internal/img2img_runner.py @@ -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() diff --git a/src/remove_ai_watermarks/noai/isobmff.py b/src/remove_ai_watermarks/_internal/isobmff.py similarity index 100% rename from src/remove_ai_watermarks/noai/isobmff.py rename to src/remove_ai_watermarks/_internal/isobmff.py diff --git a/src/remove_ai_watermarks/_internal/progress.py b/src/remove_ai_watermarks/_internal/progress.py new file mode 100644 index 0000000..2fb8f05 --- /dev/null +++ b/src/remove_ai_watermarks/_internal/progress.py @@ -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() diff --git a/src/remove_ai_watermarks/noai/qwen_zimage_pipeline.py b/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py similarity index 90% rename from src/remove_ai_watermarks/noai/qwen_zimage_pipeline.py rename to src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py index 25df2f0..ff9969e 100644 --- a/src/remove_ai_watermarks/noai/qwen_zimage_pipeline.py +++ b/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py @@ -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), diff --git a/src/remove_ai_watermarks/noai/riff.py b/src/remove_ai_watermarks/_internal/riff.py similarity index 100% rename from src/remove_ai_watermarks/noai/riff.py rename to src/remove_ai_watermarks/_internal/riff.py diff --git a/src/remove_ai_watermarks/noai/tiling.py b/src/remove_ai_watermarks/_internal/tiling.py similarity index 100% rename from src/remove_ai_watermarks/noai/tiling.py rename to src/remove_ai_watermarks/_internal/tiling.py diff --git a/src/remove_ai_watermarks/_internal/utils.py b/src/remove_ai_watermarks/_internal/utils.py new file mode 100644 index 0000000..78221bc --- /dev/null +++ b/src/remove_ai_watermarks/_internal/utils.py @@ -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") diff --git a/src/remove_ai_watermarks/_internal/watermark_profiles.py b/src/remove_ai_watermarks/_internal/watermark_profiles.py new file mode 100644 index 0000000..22e99f8 --- /dev/null +++ b/src/remove_ai_watermarks/_internal/watermark_profiles.py @@ -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 diff --git a/src/remove_ai_watermarks/_internal/watermark_remover.py b/src/remove_ai_watermarks/_internal/watermark_remover.py new file mode 100644 index 0000000..891d904 --- /dev/null +++ b/src/remove_ai_watermarks/_internal/watermark_remover.py @@ -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, + ) diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index 3eb95c0..0979657 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -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": diff --git a/src/remove_ai_watermarks/gemini_engine.py b/src/remove_ai_watermarks/gemini_engine.py index fa3dd73..d44f585 100644 --- a/src/remove_ai_watermarks/gemini_engine.py +++ b/src/remove_ai_watermarks/gemini_engine.py @@ -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) diff --git a/src/remove_ai_watermarks/humanizer.py b/src/remove_ai_watermarks/humanizer.py index 5ba484e..0df352e 100644 --- a/src/remove_ai_watermarks/humanizer.py +++ b/src/remove_ai_watermarks/humanizer.py @@ -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) diff --git a/src/remove_ai_watermarks/identify.py b/src/remove_ai_watermarks/identify.py index 9c92c72..3b3e9b6 100644 --- a/src/remove_ai_watermarks/identify.py +++ b/src/remove_ai_watermarks/identify.py @@ -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") diff --git a/src/remove_ai_watermarks/invisible_engine.py b/src/remove_ai_watermarks/invisible_engine.py index 0e38d9e..317bcac 100644 --- a/src/remove_ai_watermarks/invisible_engine.py +++ b/src/remove_ai_watermarks/invisible_engine.py @@ -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 diff --git a/src/remove_ai_watermarks/metadata.py b/src/remove_ai_watermarks/metadata.py index 6d9252e..4e0c07d 100644 --- a/src/remove_ai_watermarks/metadata.py +++ b/src/remove_ai_watermarks/metadata.py @@ -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, diff --git a/src/remove_ai_watermarks/noai/c2pa.py b/src/remove_ai_watermarks/noai/c2pa.py deleted file mode 100644 index 81da1de..0000000 --- a/src/remove_ai_watermarks/noai/c2pa.py +++ /dev/null @@ -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 diff --git a/src/remove_ai_watermarks/noai/constants.py b/src/remove_ai_watermarks/noai/constants.py deleted file mode 100644 index 541a77f..0000000 --- a/src/remove_ai_watermarks/noai/constants.py +++ /dev/null @@ -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 ", 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" diff --git a/src/remove_ai_watermarks/noai/extractor.py b/src/remove_ai_watermarks/noai/extractor.py deleted file mode 100644 index 0c4909b..0000000 --- a/src/remove_ai_watermarks/noai/extractor.py +++ /dev/null @@ -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}: ") - else: - lines.append(f"{key}: {value}") - - return "\n".join(lines) diff --git a/src/remove_ai_watermarks/noai/img2img_runner.py b/src/remove_ai_watermarks/noai/img2img_runner.py deleted file mode 100644 index a5aed90..0000000 --- a/src/remove_ai_watermarks/noai/img2img_runner.py +++ /dev/null @@ -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..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] diff --git a/src/remove_ai_watermarks/noai/progress.py b/src/remove_ai_watermarks/noai/progress.py deleted file mode 100644 index 262a306..0000000 --- a/src/remove_ai_watermarks/noai/progress.py +++ /dev/null @@ -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() diff --git a/src/remove_ai_watermarks/noai/utils.py b/src/remove_ai_watermarks/noai/utils.py deleted file mode 100644 index fdce354..0000000 --- a/src/remove_ai_watermarks/noai/utils.py +++ /dev/null @@ -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" diff --git a/src/remove_ai_watermarks/noai/watermark_profiles.py b/src/remove_ai_watermarks/noai/watermark_profiles.py deleted file mode 100644 index f540c84..0000000 --- a/src/remove_ai_watermarks/noai/watermark_profiles.py +++ /dev/null @@ -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 diff --git a/src/remove_ai_watermarks/noai/watermark_remover.py b/src/remove_ai_watermarks/noai/watermark_remover.py deleted file mode 100644 index b5c3031..0000000 --- a/src/remove_ai_watermarks/noai/watermark_remover.py +++ /dev/null @@ -1,1117 +0,0 @@ -"""Watermark removal using diffusion model regeneration attack. - -Four pipelines (selected by the explicit ``pipeline`` ctor arg): - -0. ``qwen-zimage`` -- Qwen-Image-2512 Lightning + DiffSynth Canny regenerates the - frame, then SAM-masked Z-Image Turbo regenerates original face crops and feathers - them into the global result. CUDA-only and installed through its own optional extra. -1. ``qwen`` -- Qwen-Image (20B MMDiT, Apache-2.0) img2img. The scrub still comes from - the img2img ``strength``; Qwen preserves text (incl. CJK) and structure markedly - better than SDXL at the scrub floor, so it over-regenerates real photos far less. - CUDA/cloud-class (does not fit MPS). See ``watermark_profiles`` for the certified - oracle floors. (Near-threshold scrub is seed-non-deterministic, but the prod path - pins ONE fixed seed, so a certified floor reproduces run-to-run -- pinning the seed - is the release gate, not sweeping seeds.) - -Two SDXL pipelines: -2. ``controlnet`` (DEFAULT) -- SDXL img2img with a canny ControlNet. The watermark - REMOVAL still comes from the img2img regeneration (``strength``); the ControlNet - only PRESERVES structure (text/faces) by conditioning on the edge map. No original - pixels are ever copied or frozen. Because the edge map keeps the regeneration - closer to the original, it needs a higher ``strength`` floor than ``default`` to - destroy SynthID (the certified controlnet ladder; see ``watermark_profiles``). - ``controlnet_conditioning_scale`` is the preservation knob. -3. ``default`` -- plain SDXL img2img. Partial-noise regeneration scrubs the - invisible watermark; ``strength`` controls how much is regenerated. Lighter (no - ControlNet weights), but at the low default strength it leaves SynthID on - flat-graphic content -- use it for inputs without text/faces. -""" - -# torch/diffusers/cv2 boundary: these libs ship no usable types for the tensor and -# array ops below; relax the unknown-type rules for this file only. -# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalCall=false, reportOptionalSubscript=false, reportOptionalOperand=false, reportAttributeAccessIssue=false, reportPrivateImportUsage=false, reportPrivateUsage=false, reportInvalidTypeForm=false, reportConstantRedefinition=false, reportUnnecessaryComparison=false -from __future__ import annotations - -import contextlib -import logging -import os -import sys -import time -from typing import TYPE_CHECKING, Any - -if TYPE_CHECKING: - from collections.abc import Callable - from pathlib import Path - -from PIL import Image - -from remove_ai_watermarks.noai.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, -) - -logger = logging.getLogger(__name__) - -# Check for optional dependencies -_HAS_TORCH = False -_HAS_DIFFUSERS = False - -try: - import torch - - _HAS_TORCH = True -except ImportError: - torch = None # type: ignore - -try: - from diffusers import AutoPipelineForImage2Image as AutoImg2ImgPipeline - - _HAS_DIFFUSERS = True -except ImportError: - AutoImg2ImgPipeline = None # type: ignore - - -def is_watermark_removal_available() -> bool: - """Check if watermark removal dependencies are installed.""" - return _HAS_TORCH and _HAS_DIFFUSERS - - -# Drop-in fp16-safe replacement for the SDXL VAE. The stock SDXL VAE overflows -# to NaN in fp16 and decodes to an all-black image (issue #29, reproduced on a -# CUDA fp16 backend). This community VAE is numerically rescaled to -# stay in fp16 range. SDXL-architecture only. -_SDXL_FP16_VAE_ID = "madebyollin/sdxl-vae-fp16-fix" - - -def _needs_fp16_vae_fix(model_id: str, default_model_id: str, is_fp16: bool) -> bool: - """Whether the plain img2img pipeline must swap in the fp16-fixed SDXL VAE. - - Gated to the default SDXL checkpoint running in fp16: cpu/mps run fp32 (the - stock VAE is fine there) and the differential pipeline upcasts the VAE on its - own, so only this path on a fp16 GPU (CUDA/XPU) hits the NaN/black decode. - A custom non-SDXL ``model_id`` keeps its own VAE (the fix is SDXL-specific). - """ - return is_fp16 and model_id == default_model_id - - -# An fp16 VAE/UNet overflow decodes to NaN, which diffusers' postprocess casts to 0 -# -> a uniform all-black frame (issues #29, #41). The VAE swap above prevents it for -# the default checkpoint, but a custom model_id, a stale install, or a fal/custom -# loader can still bypass it. Detecting a degenerate output and retrying in fp32 (the -# path verified clean) is the model-agnostic safety net: never hand back a black image. -# One threshold serves both guards: a NaN->0 collapse drives mean and variance to ~0. -_DEGENERATE_THRESHOLD = 1.0 - - -def _is_degenerate_image(image: Image.Image) -> bool: - """True if a generated image collapsed to an all-black/NaN frame (#29/#41). - - A NaN fp16 decode casts to 0, so the output is a uniform near-zero image: an - extremely low mean AND near-zero variance. The variance guard keeps a - legitimately dark-but-textured photo (low mean, real detail) from being flagged. - """ - import numpy as np - - arr = np.asarray(image.convert("RGB"), dtype=np.float32) - return float(arr.mean()) < _DEGENERATE_THRESHOLD and float(arr.std()) < _DEGENERATE_THRESHOLD - - -_CUDA_FIX_ENV_KEY = "NOAI_CUDA_FIXED" - - -def _auto_install(packages: list[str], index_url: str | None = None) -> bool: - """Attempt to install missing packages via pip. Returns True on success.""" - import subprocess - - cmd = [sys.executable, "-m", "pip", "install", "-q", *packages] - if index_url: - cmd.extend(["--index-url", index_url]) - try: - subprocess.check_call(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) - return True - except (subprocess.CalledProcessError, FileNotFoundError): - return False - - -def _has_nvidia_gpu() -> bool: - """Check if an NVIDIA GPU is present via nvidia-smi.""" - import subprocess - - try: - subprocess.check_call( - ["nvidia-smi"], - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - return True - except (subprocess.CalledProcessError, FileNotFoundError): - return False - - -def _detect_cuda_index_url() -> str: - """Detect the appropriate PyTorch CUDA index URL from nvidia-smi output.""" - import subprocess - - try: - out = subprocess.check_output( - ["nvidia-smi"], - stderr=subprocess.DEVNULL, - text=True, - ) - for line in out.splitlines(): - if "CUDA Version" in line: - version_str = line.split("CUDA Version:")[-1].strip().rstrip("|").strip() - major, minor = version_str.split(".")[:2] - cuda_tag = f"cu{major}{minor}" - return f"https://download.pytorch.org/whl/{cuda_tag}" - except Exception: # noqa: S110 - pass - return "https://download.pytorch.org/whl/cu121" - - -def _reinstall_torch_cuda_and_restart() -> None: - """Reinstall torch with CUDA support showing live progress, then restart.""" - import re - import subprocess - - from remove_ai_watermarks.noai.progress import run_with_progress - - index_url = _detect_cuda_index_url() - progress_state: dict[str, str] = {"message": "NVIDIA GPU detected — installing CUDA-enabled PyTorch..."} - - pct_re = re.compile(r"(\d+)%") - pkg_re = re.compile(r"(?:Collecting|Downloading|Installing)\s+(\S+)") - - def _run_pip() -> bool: - cmd = [ - sys.executable, - "-m", - "pip", - "install", - "--force-reinstall", - "torch", - "--index-url", - index_url, - ] - proc = subprocess.Popen( - cmd, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - ) - for line in iter(proc.stdout.readline, ""): # type: ignore[union-attr] - stripped = line.strip() - if not stripped: - continue - pkg_m = pkg_re.search(stripped) - pct_m = pct_re.search(stripped) - if pct_m and pkg_m: - progress_state["message"] = f"Downloading {pkg_m.group(1)} ({pct_m.group(1)}%)" - elif pct_m: - progress_state["message"] = f"Downloading CUDA packages ({pct_m.group(1)}%)" - elif pkg_m: - action = "Installing" if stripped.startswith("Installing") else "Downloading" - progress_state["message"] = f"{action} {pkg_m.group(1)}" - elif "Successfully installed" in stripped: - progress_state["message"] = "CUDA-enabled PyTorch installed successfully" - proc.wait() - return proc.returncode == 0 - - try: - success = run_with_progress(_run_pip, progress_state) - except Exception: - success = False - - if not success: - print( - f"\n Failed to install CUDA-enabled PyTorch.\n" - f" Install manually:\n" - f" pip install torch --index-url {index_url}\n", - file=sys.stderr, - ) - return - - os.environ[_CUDA_FIX_ENV_KEY] = "1" - # Re-exec via ``-m`` rather than building a ``-c`` string from repr(sys.argv). - # ``-m`` makes Python set argv[0] to the module path, so forward only the - # actual args (sys.argv[1:]); passing the full argv would re-inject the - # program name as a spurious first argument to Click. - os.execv(sys.executable, [sys.executable, "-m", "remove_ai_watermarks.cli", *sys.argv[1:]]) - - -def _ensure_watermark_deps() -> None: - """Auto-install and re-import missing watermark removal dependencies.""" - global _HAS_TORCH, _HAS_DIFFUSERS, torch, AutoImg2ImgPipeline - missing_pkgs: list[str] = [] - if not _HAS_TORCH: - missing_pkgs.append("torch") - if not _HAS_DIFFUSERS: - missing_pkgs.extend(["diffusers", "transformers", "accelerate"]) - logger.info("Auto-installing missing dependencies: %s", missing_pkgs) - if not _auto_install(missing_pkgs): - raise ImportError( - f"Failed to auto-install missing dependencies: {', '.join(missing_pkgs)}. " - "Try manually: pip install --force-reinstall noai-watermark" - ) - import torch as _torch - - torch = _torch - _HAS_TORCH = True - from diffusers import AutoPipelineForImage2Image - - AutoImg2ImgPipeline = AutoPipelineForImage2Image - _HAS_DIFFUSERS = True - - -def get_device() -> str: - """Get the best available device for inference.""" - if not _HAS_TORCH: - return "cpu" - if torch.cuda.is_available(): # type: ignore - try: - t = torch.tensor([1.0], device="cuda") - _ = t + t - del t - return "cuda" - except (AssertionError, RuntimeError): - pass - # Intel GPU (Arc / Data Center) via the torch XPU backend. The torch.xpu - # namespace exists in stock wheels, but is_available() is only True on an - # XPU-enabled build (download.pytorch.org/whl/xpu), so this is inert on the - # default CPU/CUDA install. Checked before the nvidia-smi path so an Intel - # box never triggers the CUDA reinstaller. - if hasattr(torch, "xpu") and torch.xpu.is_available(): # type: ignore - try: - t = torch.tensor([1.0], device="xpu") - _ = t + t - del t - return "xpu" - except (AssertionError, RuntimeError): - pass - if _has_nvidia_gpu() and not os.environ.get(_CUDA_FIX_ENV_KEY): - _reinstall_torch_cuda_and_restart() - if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - return "mps" - return "cpu" - - -def _make_seed_generator(device: str, seed: int) -> Any: - """Build a seeded ``torch.Generator``, falling back to a CPU generator. - - Some backends have no device-side RNG (notably certain torch-xpu builds), - so ``torch.Generator(device="xpu")`` can raise. A CPU generator is - backend-agnostic and still seeds the pipeline reproducibly, so fall back to - it rather than failing the run when ``--seed`` is used on such a device. - """ - try: - return torch.Generator(device=device).manual_seed(seed) # type: ignore - except (RuntimeError, TypeError): - return torch.Generator().manual_seed(seed) # type: ignore - - -# Canny edge thresholds for the ControlNet control image (xinsir canny recipe: -# cv2.Canny(gray, 100, 200) -> a 3-channel edge map). -_CANNY_LOW = 100 -_CANNY_HIGH = 200 - -# A neutral quality prompt: the goal is faithful regeneration, not creative edits. -_CONTROLNET_PROMPT = "best quality, high quality, sharp, detailed, photographic" -_CONTROLNET_NEGATIVE = "blurry, lowres, deformed, distorted text, garbled text, watermark, jpeg artifacts" - -# Neutral prompts for the Qwen-Image img2img pass (faithful regeneration, not an edit). -_QWEN_PROMPT = "high quality, sharp, detailed, faithful to the original" -_QWEN_NEGATIVE = "blurry, lowres, distorted text, garbled text, artifacts" - - -def _qwen_target_size(width: int, height: int) -> tuple[int, int]: - """Floor (width, height) to a multiple of 16 for Qwen's VAE/patchifier (>= 16). - - Pure; unit-tested. Without explicit dims the img2img pipeline defaults to a 1024x1024 - SQUARE and silently distorts any non-square input. - """ - return max(16, (width // 16) * 16), max(16, (height // 16) * 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 QwenImageImg2ImgPipeline call kwargs (pure; unit-tested without torch). - - Qwen-Image uses ``true_cfg_scale`` (not SDXL's ``guidance_scale``) and takes an - explicit ``negative_prompt``; the scrub still comes from the img2img ``strength``. - Passes an explicit ``height``/``width`` derived from the input (floored to /16): the - pipeline otherwise defaults to a 1024x1024 SQUARE, squishing any non-square input - (the abba mixed-seam test: a 2816x1536 poster came back 1024x1024, distorting the - scene and garbling text). So qwen regenerates at the input's own geometry. - """ - qw, qh = _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, - "height": qh, - "width": qw, - } - - -class WatermarkRemover: - """Remove watermarks from images using diffusion model regeneration. - - Attributes: - model_id: HuggingFace model ID for the diffusion model. - device: Device to run inference on (cuda, xpu, mps, or cpu). - """ - - DEFAULT_MODEL_ID = DEFAULT_MODEL_ID - DEFAULT_STRENGTH = DEFAULT_STRENGTH - CONTROLNET_CANNY_MODEL = CONTROLNET_CANNY_MODEL - - 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: - self.model_id = model_id or self.DEFAULT_MODEL_ID - # Diffusers offloads whole model components between CUDA calls. The custom - # qwen-zimage runtime uses the same flag to keep its face stack off VRAM. - self.cpu_offload = cpu_offload - # The pipeline profile is threaded explicitly (not inferred from model_id): - # both "sdxl" and "controlnet" use the same SDXL base checkpoint. Normalize so - # the legacy "default" alias resolves to "sdxl". - self.model_profile = normalize_profile(pipeline) - self.controlnet_conditioning_scale = controlnet_conditioning_scale - if self.model_profile == QWEN_ZIMAGE_PROFILE and self.model_id != self.DEFAULT_MODEL_ID: - raise ValueError( - "The qwen-zimage pipeline uses a fixed Qwen-Image-2512 + Z-Image model stack; " - "--model is not supported for this profile." - ) - if self.model_profile == QWEN_ZIMAGE_PROFILE: - self.model_id = "Qwen/Qwen-Image-2512 + Tongyi-MAI/Z-Image-Turbo" - - if not is_watermark_removal_available(): - _ensure_watermark_deps() - self.device = (device or get_device()).lower() - if self.device == "auto": - self.device = get_device() - if self.device not in {"cpu", "mps", "cuda", "xpu"}: - raise ValueError(f"Unsupported device '{device}'. Use one of: auto, cpu, mps, cuda, xpu.") - if torch_dtype is None: - if self.device == "cpu" or self.device == "mps": - self.torch_dtype = torch.float32 # type: ignore - elif self.model_profile in {"qwen", QWEN_ZIMAGE_PROFILE}: - # Qwen-Image and Z-Image are published in bf16; fp16 risks overflow. - # cuda/xpu-only by construction: the cpu/mps guard above already forced - # fp32, and the 20B model does not fit MPS anyway. - self.torch_dtype = torch.bfloat16 # type: ignore - else: - self.torch_dtype = torch.float16 # type: ignore - else: - self.torch_dtype = torch_dtype - - self._pipeline: AutoImg2ImgPipeline | None = None - self._controlnet_pipeline: Any = None - self._qwen_pipeline: Any = None - self._qwen_zimage_pipeline: Any = None - self._progress_callback = progress_callback - self.hf_token: str | None = hf_token or os.environ.get("HF_TOKEN") - - def _set_progress(self, message: str) -> None: - """Send a progress update through callback when available.""" - if self._progress_callback is None: - return - with contextlib.suppress(Exception): - self._progress_callback(message) - - # ── Preload ────────────────────────────────────────────────────── - - def preload(self, *, global_only: bool = False) -> None: - """Eagerly load the pipeline so download progress bars are visible. - - ``global_only`` applies to qwen-zimage, whose face stage is optional. - """ - 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() - - # ── Pipeline loading ───────────────────────────────────────────── - - def _maybe_add_fp16_vae(self, load_kwargs: dict[str, Any]) -> None: - """Swap in the fp16-fixed SDXL VAE for the default checkpoint on a fp16 GPU. - - The stock SDXL VAE overflows to NaN in fp16 and decodes to an all-black - image (issue #29). Shared by both pipeline loaders; a no-op on fp32 (cpu/mps) - or a non-SDXL checkpoint. - """ - if _needs_fp16_vae_fix(self.model_id, self.DEFAULT_MODEL_ID, self.torch_dtype == torch.float16): - from diffusers import AutoencoderKL - - self._set_progress("Loading fp16-fixed SDXL VAE (avoids black output)...") - load_kwargs["vae"] = AutoencoderKL.from_pretrained(_SDXL_FP16_VAE_ID, torch_dtype=torch.float16) - - @staticmethod - def _disable_sdxl_watermarker(load_kwargs: dict[str, Any]) -> None: - """Turn off the diffusers default invisible watermarker on an SDXL pipeline. - - diffusers embeds an open "Stable Diffusion XL" DWT-DCT invisible watermark on - EVERY SDXL output whenever ``invisible-watermark`` is installed (kept as a - development parity dependency). A watermark REMOVER must not re-stamp a - detectable AI watermark, or the - cleaned output re-reads as AI (``identify`` -> "Open invisible watermark: Stable - Diffusion XL"). Shared by both SDXL loaders; the ``ControlNetModel`` sub-model - and the Qwen loader never call it (only the pipeline accepts the kwarg). - """ - load_kwargs["add_watermarker"] = False - - def _move_to_device_and_optimize(self, pipeline: Any) -> Any: - """Move a freshly-loaded pipeline to ``self.device`` + enable memory opts. - - Shared by both loaders. On a CUDA move failure (missing CUDA torch build), - trigger the torch-CUDA reinstall+restart. Returns the moved pipeline. - """ - self._set_progress(f"Moving model to device: {self.device}") - if self.cpu_offload and self.device == "cuda": - enable_cpu_offload = getattr(pipeline, "enable_model_cpu_offload", None) - if not callable(enable_cpu_offload): - raise RuntimeError( - "CPU offload was requested, but this pipeline does not support enable_model_cpu_offload()." - ) - self._set_progress("Enabling CUDA model CPU offload (low-VRAM mode)...") - try: - enable_cpu_offload(device=self.device) - except (RuntimeError, AssertionError) as exc: - raise RuntimeError(f"Failed to enable model CPU offload ({exc}).") from exc - else: - try: - pipeline = pipeline.to(self.device) - except (RuntimeError, AssertionError) as exc: - if self.device == "cuda" and not os.environ.get(_CUDA_FIX_ENV_KEY): - self._set_progress("CUDA failed. Reinstalling torch with CUDA support...") - _reinstall_torch_cuda_and_restart() - raise RuntimeError( - f"Failed to move model to {self.device} ({exc}). " - "Install CUDA-enabled PyTorch manually:\n" - f" pip install torch --index-url {_detect_cuda_index_url()}" - ) from exc - - if hasattr(pipeline, "enable_xformers_memory_efficient_attention"): - with contextlib.suppress(Exception): - self._set_progress("Enabling memory optimizations...") - pipeline.enable_xformers_memory_efficient_attention() - - # Mac Float32 memory slicing - if self.device == "mps" and hasattr(pipeline, "enable_attention_slicing"): - with contextlib.suppress(Exception): - pipeline.enable_attention_slicing("max") - - return pipeline - - def _base_load_kwargs(self) -> dict[str, Any]: - """The ``from_pretrained`` kwargs shared by all three loaders (dtype + token). - - Each loader adds its own extras (SDXL safety_checker + fp16 VAE, the ControlNet - model, etc.). Centralizing the dtype/token pair avoids the drift trap of three - copies (a token forgotten on one loader silently breaks gated downloads there). - """ - load_kwargs: dict[str, Any] = {"torch_dtype": self.torch_dtype} - if self.hf_token: - load_kwargs["token"] = self.hf_token - return load_kwargs - - def _load_from_pretrained(self, cls: Any, model_id: str, **load_kwargs: Any) -> Any: - """Call ``cls.from_pretrained`` reading the fp16 weight VARIANT when on fp16. - - When ``torch_dtype`` is float16 (the CUDA/XPU SDXL default), pass - ``variant="fp16"`` so diffusers fetches/reads the half-precision weight files - (~half the bytes of the fp32 defaults) instead of reading the full fp32 files - and downcasting in memory. On a warm weights cache this roughly halves the - cold-start weight read + host->device transfer, which a phase-timed Modal run - measured as ~half of the ~25s cold start. Not every checkpoint publishes an - fp16 variant (a custom ``--model``, or the canny ControlNet if it ships only - default files), so fall back to the default weights if the variant is missing - - worst case is the prior behavior (read fp32, downcast). fp32 (cpu/mps) and bf16 - (qwen) never request the variant. - """ - if self.torch_dtype == torch.float16: - try: - return cls.from_pretrained(model_id, variant="fp16", **load_kwargs) - except Exception as exc: - logger.info("No fp16 weight variant for %s (%s); loading default weights", model_id, exc) - return cls.from_pretrained(model_id, **load_kwargs) - - def _load_pipeline(self) -> AutoImg2ImgPipeline: - """Load the plain SDXL img2img pipeline lazily.""" - if self._pipeline is None: - logger.info("Loading model %s on %s...", self.model_id, self.device) - self._set_progress(f"Loading model weights: {self.model_id}") - - load_kwargs = self._base_load_kwargs() - load_kwargs["safety_checker"] = None - load_kwargs["requires_safety_checker"] = False - self._disable_sdxl_watermarker(load_kwargs) - self._maybe_add_fp16_vae(load_kwargs) - - pipeline = self._load_from_pretrained(AutoImg2ImgPipeline, self.model_id, **load_kwargs) # type: ignore - self._pipeline = self._move_to_device_and_optimize(pipeline) - - logger.info("Model loaded successfully") - self._set_progress("Model initialized. Preparing input image...") - - return self._pipeline # type: ignore - - def _load_controlnet_pipeline(self) -> Any: - """Load the SDXL + canny-ControlNet img2img pipeline lazily. - - Mirrors ``_load_pipeline`` (same fp16-fix VAE, device move, attention - slicing via the shared helpers) but loads the canny ControlNet on top of - the SDXL base. The ControlNet only preserves structure via the edge map; - removal still comes from the img2img regeneration (``strength``). - """ - if self._controlnet_pipeline is None: - from diffusers import ControlNetModel, StableDiffusionXLControlNetImg2ImgPipeline - - logger.info("Loading SDXL + ControlNet (%s) on %s...", CONTROLNET_CANNY_MODEL, self.device) - self._set_progress(f"Loading ControlNet: {CONTROLNET_CANNY_MODEL}") - controlnet = self._load_from_pretrained( - ControlNetModel, CONTROLNET_CANNY_MODEL, torch_dtype=self.torch_dtype - ) - - load_kwargs = self._base_load_kwargs() - load_kwargs["controlnet"] = controlnet - self._disable_sdxl_watermarker(load_kwargs) - self._maybe_add_fp16_vae(load_kwargs) - - self._set_progress(f"Loading model weights: {self.model_id}") - pipeline = self._load_from_pretrained( - StableDiffusionXLControlNetImg2ImgPipeline, self.model_id, **load_kwargs - ) - pipeline = self._move_to_device_and_optimize(pipeline) - with contextlib.suppress(Exception): - pipeline.set_progress_bar_config(disable=True) - - logger.info("ControlNet model loaded successfully") - self._controlnet_pipeline = pipeline - - return self._controlnet_pipeline - - def _load_qwen_pipeline(self) -> Any: - """Load the Qwen-Image img2img pipeline lazily. - - Qwen-Image is its OWN base model (not an SDXL add-on), so it loads - ``QWEN_MODEL_ID`` unless the caller passed a custom ``--model``. Needs a - diffusers build that ships ``QwenImageImg2ImgPipeline``; raises a clear error - otherwise. CUDA/cloud-class (the 20B MMDiT does not fit MPS). - """ - if self._qwen_pipeline is None: - try: - from diffusers import QwenImageImg2ImgPipeline - except ImportError as exc: - raise ImportError( - "The 'qwen' pipeline needs a diffusers version that ships " - "QwenImageImg2ImgPipeline. Upgrade: pip install -U diffusers" - ) from exc - - # Use the Qwen base unless the user explicitly overrode --model. - model = self.model_id if self.model_id != self.DEFAULT_MODEL_ID else QWEN_MODEL_ID - logger.info("Loading Qwen-Image (%s) on %s...", model, self.device) - self._set_progress(f"Loading model weights: {model}") - pipeline = QwenImageImg2ImgPipeline.from_pretrained(model, **self._base_load_kwargs()) - pipeline = self._move_to_device_and_optimize(pipeline) - with contextlib.suppress(Exception): - pipeline.set_progress_bar_config(disable=True) - - logger.info("Qwen-Image model loaded successfully") - self._qwen_pipeline = pipeline - - return self._qwen_pipeline - - def _load_qwen_zimage_pipeline(self) -> Any: - """Load the two-stage Qwen-Image-2512 + Z-Image runtime lazily.""" - if self._qwen_zimage_pipeline is None: - from remove_ai_watermarks.noai.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 - - # ── Core removal ───────────────────────────────────────────────── - - 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: - """Remove watermark from an image using regeneration attack. - - Args: - image_path: Path to the watermarked image. - output_path: Path for the cleaned image. If None, modifies in place. - strength: Denoising strength (0.0-1.0). None -> the vendor-adaptive - default (see ``vendor``). - num_inference_steps: Number of denoising steps. - guidance_scale: Classifier-free guidance scale. - seed: Random seed for reproducibility. None resolves to 0 for - qwen-zimage and stays random for the other profiles. - vendor: SynthID vendor (``"openai"`` / ``"google"`` / None) used to pick the - default strength when ``strength`` is None. Detect it from the ORIGINAL - input with ``watermark_profiles.vendor_for_strength`` before processing - strips the metadata; the caller passes it down so display and execution - agree. - tile: Process the image in overlapping tiles instead of one forward pass. - This keeps the input's native dimensions instead of applying a - ``--max-resolution`` downscale, but every tile is still regenerated. - Only engages when the long side exceeds ``tile_size``; smaller images - run a single pass unchanged. - tile_size: Tile dimension in px (default 1024). - tile_overlap: Overlap between adjacent tiles in px (default 128), feather- - blended so there is no visible seam. - region: Restrict the regeneration to the AI-composited box ``(x, y, w, h)`` - and feather-composite it back over the ORIGINAL pixels everywhere else. - For AI-ENHANCED composites (digitalSourceType - ``compositeWithTrainedAlgorithmicMedia``, surfaced as - ``identify.ProvenanceReport.ai_source_kind == "enhanced"``): the real - photo outside the box is preserved exactly, only the AI region is - scrubbed. The box is supplied by the caller (a C2PA composite manifest - does not carry a reliable machine-readable region). None -> whole frame. - region_feather: Seam taper in px for ``region`` compositing (default 64). - - Returns: - Path to the cleaned image. - - Raises: - FileNotFoundError: If input image doesn't exist. - ValueError: If strength is not in valid range. - """ - if not image_path.exists(): - raise FileNotFoundError(f"Image not found: {image_path}") - - if output_path is None: - output_path = image_path - - self._set_progress("Loading and preprocessing input image...") - init_image = Image.open(image_path).convert("RGB") - w, h = init_image.size - self._set_progress(f"Image loaded: {w}x{h}px | Model: {self.model_id}") - - if self.model_profile == QWEN_ZIMAGE_PROFILE: - from remove_ai_watermarks.noai.qwen_zimage_pipeline import resolution_adaptive_denoise - - strength = strength if strength is not None else resolution_adaptive_denoise(w, h) - else: - strength = resolve_strength(strength, vendor, self.model_profile) - seed = resolve_seed(seed, self.model_profile) - if not 0.0 <= strength <= 1.0: - raise ValueError(f"Strength must be between 0.0 and 1.0, got {strength}") - - num_inference_steps = resolve_steps(num_inference_steps, self.model_profile) - if guidance_scale is None: - guidance_scale = 1.0 if self.model_profile == QWEN_ZIMAGE_PROFILE else 7.5 - elif self.model_profile == QWEN_ZIMAGE_PROFILE and guidance_scale != 1.0: - raise ValueError("The qwen-zimage profile fixes both diffusion stages at CFG 1.0.") - - generator = None - if seed is not None and _HAS_TORCH: - self._set_progress(f"Setting reproducible seed: {seed}") - generator = _make_seed_generator(self.device, seed) - - # A step count whose product with strength rounds to zero kills the pipeline - # inside attention with an opaque reshape error, so raise it to the minimum that - # denoises. Must be applied to the value HANDED TO THE PIPELINE -- the old - # max(1, ...) below only clamped the number in the log line. - if self.model_profile != QWEN_ZIMAGE_PROFILE: - adjusted = viable_steps(num_inference_steps, strength) - if adjusted != num_inference_steps: - logger.warning( - "steps=%s at strength=%s denoises 0 steps and would crash; using steps=%s (1 effective)", - num_inference_steps, - strength, - adjusted, - ) - num_inference_steps = adjusted - elif num_inference_steps != 4: - raise ValueError("The qwen-zimage profile uses the 4-step Lightning LoRA, so --steps must be 4.") - - # DiffSynth keeps all timesteps and compresses their sigma range for a low - # denoise value. Diffusers instead truncates the schedule by strength. - effective_steps = ( - num_inference_steps - if self.model_profile == QWEN_ZIMAGE_PROFILE - else max(1, int(num_inference_steps * strength)) - ) - self._set_progress( - f"Config: strength={strength}, steps={num_inference_steps} " - f"(~{effective_steps} effective), guidance={guidance_scale}, device={self.device}" - ) - - _total_start = time.monotonic() - - def _generate_one(img: Image.Image) -> Image.Image: - if self.model_profile == "qwen": - return self._run_qwen(img, strength, num_inference_steps, guidance_scale, generator) - if self.model_profile == "controlnet": - return self._run_controlnet(img, strength, num_inference_steps, guidance_scale, generator) - return self._run_img2img(img, strength, num_inference_steps, guidance_scale, generator) - - def _generate() -> Image.Image: - # qwen-zimage owns its global-only tiling because its face stage must run - # once after the tiles are blended. Other profiles tile their whole pass. - if self.model_profile == QWEN_ZIMAGE_PROFILE: - return self._run_qwen_zimage( - init_image, - strength, - seed, - tile=tile, - tile_size=tile_size, - tile_overlap=tile_overlap, - ) - # Tile only when asked AND the image is larger than one tile; otherwise a - # single full-image pass (tiling a sub-tile image is pure overhead). - if tile and max(init_image.size) > tile_size: - from remove_ai_watermarks.noai.tiling import run_tiled - - return run_tiled(_generate_one, init_image, tile_size, tile_overlap, self._set_progress) - return _generate_one(init_image) - - cleaned_image = _generate() - - # Safety net for the fp16 all-black/NaN decode (#29/#41): if an fp16 run - # produced a degenerate (uniform black) frame -- the VAE swap did not engage - # for this model/version -- retry once in fp32 on the same device (verified - # clean) so the user never gets a black image. Skipped when an MPS->CPU - # fallback already moved us to fp32. - if self.torch_dtype == torch.float16 and _is_degenerate_image(cleaned_image): - logger.warning("fp16 output was degenerate (all-black/NaN, #29/#41); retrying in fp32 on %s.", self.device) - self._set_progress("Output was black (fp16 overflow); retrying in fp32...") - self.torch_dtype = torch.float32 - self._pipeline = None - self._controlnet_pipeline = None - self._qwen_zimage_pipeline = None - cleaned_image = _generate() - - # Region-targeted regeneration for AI-enhanced composites: keep the real photo - # outside the AI box pixel-exact, blend only the regenerated AI region back in. - if region is not None: - import numpy as np - - from remove_ai_watermarks.noai.tiling import feather_region_composite - - gen = cleaned_image.convert("RGB") - if gen.size != init_image.size: # a downscaled/tiled pass can resize - gen = gen.resize(init_image.size) - cleaned_image = gen - base_rgb = np.asarray(init_image) # original RGB, untouched outside the box - merged = feather_region_composite(base_rgb, np.asarray(gen), region, feather=region_feather) - cleaned_image = Image.fromarray(merged) - self._set_progress(f"Region-targeted regeneration: AI box {region}, real photo preserved") - - self._set_progress(f"Regeneration complete · Output: {w}x{h}px {cleaned_image.mode}") - - output_path.parent.mkdir(parents=True, exist_ok=True) - fmt = output_path.suffix.lower() - if fmt in (".jpg", ".jpeg"): - self._set_progress(f"Encoding as JPEG → {output_path.name}...") - else: - self._set_progress(f"Encoding as PNG → {output_path.name}...") - # Encode through image_io so the regenerated image gets the same quality- - # preserving write as the visible path (JPEG q100 / 4:4:4, HEIC/AVIF via Pillow) - # instead of PIL's default JPEG quality 75, which would crush the diffusion - # output further for no reason. - import numpy as np - - from remove_ai_watermarks import image_io - - rgb = np.asarray(cleaned_image.convert("RGB"))[:, :, ::-1] # PIL RGB -> cv2 BGR - if not image_io.imwrite(str(output_path), np.ascontiguousarray(rgb)): - cleaned_image.save(output_path) # fallback for a container image_io/cv2 cannot encode - - if output_path.exists(): - self._set_progress("Stripping AI metadata from output...") - try: - # The single, robust stripper (byte-level: lossless for JPEG, handles - # every container) -- not the legacy PIL-re-encoding noai.cleaner one. - from remove_ai_watermarks.metadata import remove_ai_metadata - - remove_ai_metadata(output_path, output_path, keep_standard=True) - except Exception: - logger.debug("AI metadata stripping skipped", exc_info=True) - - total_time = time.monotonic() - _total_start - - size_str = "" - try: - file_size = output_path.stat().st_size - if file_size < 1024 * 1024: - size_str = f" ({file_size / 1024:.0f}KB)" - else: - size_str = f" ({file_size / (1024 * 1024):.1f}MB)" - except OSError: - pass - - logger.info("Cleaned image saved to %s", output_path) - self._set_progress(f"✓ Saved {output_path.name}{size_str} · {w}x{h}px · {total_time:.0f}s total") - - return output_path - - # ── Img2img runner ─────────────────────────────────────────────── - - def _run_img2img( - self, - init_image: Image.Image, - strength: float, - num_inference_steps: int, - guidance_scale: float, - generator: Any, - ) -> Image.Image: - """Execute the img2img pipeline with progress and MPS fallback.""" - from remove_ai_watermarks.noai.img2img_runner import run_img2img_with_mps_fallback - - result_image, final_device = run_img2img_with_mps_fallback( - load_pipeline=self._load_pipeline, - image=init_image, - strength=strength, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - generator=generator, - device=self.device, - set_progress=self._set_progress, - reload_on_cpu=self._reload_pipeline_on_cpu, - ) - - if final_device != self.device: - self.device = final_device - self.torch_dtype = torch.float32 # type: ignore[assignment] - - return result_image - - def _reload_pipeline_on_cpu(self) -> Any: - """Reload pipeline on CPU after MPS failure.""" - self.device = "cpu" - self.torch_dtype = torch.float32 # type: ignore[assignment] - self._pipeline = None - return self._load_pipeline() - - # ── ControlNet runner ──────────────────────────────────────────── - - def _build_canny_control_image(self, init_image: Image.Image) -> Image.Image: - """Build the canny ControlNet conditioning image (xinsir recipe). - - cv2.Canny on the RGB->gray array, stacked to 3 channels, wrapped as a PIL - image. The edge map only PRESERVES structure; it never copies pixels. - ``init_image`` is already RGB (``remove_watermark`` converts on load). - """ - import cv2 - import numpy as np - - rgb = np.array(init_image) - gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY) - edges = cv2.Canny(gray, _CANNY_LOW, _CANNY_HIGH) - edges_rgb = np.stack([edges, edges, edges], axis=-1) - return Image.fromarray(edges_rgb) - - def _run_controlnet( - self, - init_image: Image.Image, - strength: float, - num_inference_steps: int, - guidance_scale: float, - generator: Any, - ) -> Image.Image: - """Run the SDXL + canny-ControlNet img2img pass. - - Removal still comes from the img2img regeneration (``strength``); the canny - ControlNet only PRESERVES text and face STRUCTURE via the edge map. No - original pixels are copied/frozen, so SynthID does not survive (canny holds - structure, not face identity). ``controlnet_conditioning_scale`` is the - structure-preservation knob. Shares the img2img runner (live progress + - MPS->CPU fallback) with ``_run_img2img``; the only delta is the extra - ControlNet kwargs (canny control image + conditioning scale + a non-empty - prompt) overlaid via ``extra_kwargs``. - """ - from remove_ai_watermarks.noai.img2img_runner import run_img2img_with_mps_fallback - - extra_kwargs = { - "prompt": _CONTROLNET_PROMPT, - "negative_prompt": _CONTROLNET_NEGATIVE, - "control_image": self._build_canny_control_image(init_image), - "controlnet_conditioning_scale": float(self.controlnet_conditioning_scale), - } - result_image, final_device = run_img2img_with_mps_fallback( - load_pipeline=self._load_controlnet_pipeline, - image=init_image, - strength=strength, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - generator=generator, - device=self.device, - set_progress=self._set_progress, - reload_on_cpu=self._reload_controlnet_on_cpu, - extra_kwargs=extra_kwargs, - ) - - if final_device != self.device: - self.device = final_device - self.torch_dtype = torch.float32 # type: ignore[assignment] - - return result_image - - def _reload_controlnet_on_cpu(self) -> Any: - """Reload the controlnet pipeline on CPU after an MPS failure.""" - self.device = "cpu" - self.torch_dtype = torch.float32 # type: ignore[assignment] - self._controlnet_pipeline = None - return self._load_controlnet_pipeline() - - # ── Qwen runner ────────────────────────────────────────────────── - - def _run_qwen( - self, - init_image: Image.Image, - strength: float, - num_inference_steps: int, - guidance_scale: float, - generator: Any, - ) -> Image.Image: - """Run the Qwen-Image img2img pass. - - Removal comes from the img2img ``strength`` (same lever as the SDXL paths); - Qwen-Image preserves text/structure markedly better at the scrub floor. The - CLI ``guidance_scale`` maps to Qwen's ``true_cfg_scale`` (~4.0 is typical; - the SDXL default of 7.5 is high for Qwen). No MPS->CPU fallback: the 20B MMDiT - is CUDA/cloud-class and does not run on MPS, so an error here propagates. - """ - pipeline = self._load_qwen_pipeline() - self._set_progress(f"Running Qwen-Image img2img (strength={strength}, true_cfg={guidance_scale})...") - kwargs = _build_qwen_kwargs(init_image, strength, num_inference_steps, guidance_scale, generator) - result = pipeline(**kwargs) - return result.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: - """Run the Qwen 2512 Canny pass and masked Z-Image face repair.""" - pipeline = self._load_qwen_zimage_pipeline() - return pipeline.run( - init_image, - strength=strength, - seed=seed, - tile=tile, - tile_size=tile_size, - tile_overlap=tile_overlap, - ) - - # ── Batch ──────────────────────────────────────────────────────── - - 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]: - """Remove watermarks from all images in a directory.""" - if not input_dir.exists(): - raise FileNotFoundError(f"Input directory not found: {input_dir}") - - output_dir.mkdir(parents=True, exist_ok=True) - cleaned_paths: list[Path] = [] - - # Lazy import keeps this module torch-optional; frees device cache per image. - from remove_ai_watermarks.noai.img2img_runner import try_empty_device_cache - - for ext in extensions: - for image_path in input_dir.glob(f"*{ext}"): - output_path = output_dir / image_path.name - try: - result_path = self.remove_watermark( - image_path=image_path, - output_path=output_path, - strength=strength, - num_inference_steps=num_inference_steps, - ) - cleaned_paths.append(result_path) - except Exception as e: - logger.error("Failed to process %s: %s", image_path, e) - try_empty_device_cache(self.device) - - return cleaned_paths - - -# ── Convenience function ───────────────────────────────────────────── - - -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 function to remove watermark from an image. - - ``strength=None`` lets the profile pick its vendor-adaptive default - (0.10 OpenAI / 0.15 Google / 0.15 unknown, from the C2PA SynthID proxy on the - input; same ladder for the controlnet and sdxl pipelines -- the single source of - truth is ``watermark_profiles.py``). Pass a value to override. - - ``region=(x, y, w, h)`` restricts the regeneration to that box and preserves the - real photo elsewhere -- for AI-enhanced composites (see - ``WatermarkRemover.remove_watermark``). - """ - from remove_ai_watermarks.noai.watermark_profiles import vendor_for_strength - - remover = WatermarkRemover(model_id=model_id, device=device, hf_token=hf_token) - return remover.remove_watermark( - image_path=image_path, - output_path=output_path, - strength=strength, - vendor=vendor_for_strength(image_path), - region=region, - ) diff --git a/src/remove_ai_watermarks/video.py b/src/remove_ai_watermarks/video.py index 59fff6b..6f5bd35 100644 --- a/src/remove_ai_watermarks/video.py +++ b/src/remove_ai_watermarks/video.py @@ -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: diff --git a/tests/test_cpu_offload.py b/tests/test_cpu_offload.py index f2c6ffd..8a50b49 100644 --- a/tests/test_cpu_offload.py +++ b/tests/test_cpu_offload.py @@ -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: diff --git a/tests/test_identify.py b/tests/test_identify.py index c32494d..92ce88d 100644 --- a/tests/test_identify.py +++ b/tests/test_identify.py @@ -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): diff --git a/tests/test_img2img_runner.py b/tests/test_img2img_runner.py index 1f964b7..7adeea2 100644 --- a/tests/test_img2img_runner.py +++ b/tests/test_img2img_runner.py @@ -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, ) diff --git a/tests/test_invisible_engine.py b/tests/test_invisible_engine.py index b44f8ce..077bcbe 100644 --- a/tests/test_invisible_engine.py +++ b/tests/test_invisible_engine.py @@ -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) diff --git a/tests/test_metadata.py b/tests/test_metadata.py index 2f6fccc..4c863c6 100644 --- a/tests/test_metadata.py +++ b/tests/test_metadata.py @@ -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'{"Label":"1"}' @@ -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"(c) me" diff --git a/tests/test_noai.py b/tests/test_metadata_internals.py similarity index 93% rename from tests/test_noai.py rename to tests/test_metadata_internals.py index a99a93c..c1900fa 100644 --- a/tests/test_noai.py +++ b/tests/test_metadata_internals.py @@ -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'' + self.AIMARK + b'' 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'plain copyright' 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 -> the packet regex cannot match, so it is left unchanged. data = b'' + self.AIMARK + b"" @@ -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)" diff --git a/tests/test_platform.py b/tests/test_platform.py index 2464649..694f03b 100644 --- a/tests/test_platform.py +++ b/tests/test_platform.py @@ -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)) diff --git a/tests/test_qwen_zimage_pipeline.py b/tests/test_qwen_zimage_pipeline.py index cb3d520..2cf3a18 100644 --- a/tests/test_qwen_zimage_pipeline.py +++ b/tests/test_qwen_zimage_pipeline.py @@ -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) diff --git a/tests/test_security_clamp.py b/tests/test_security_clamp.py index 9e3c194..ffdaae5 100644 --- a/tests/test_security_clamp.py +++ b/tests/test_security_clamp.py @@ -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 diff --git a/tests/test_tiling.py b/tests/test_tiling.py index ab6df75..fd2cccd 100644 --- a/tests/test_tiling.py +++ b/tests/test_tiling.py @@ -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, diff --git a/tests/test_watermark_profiles.py b/tests/test_watermark_profiles.py index 3459b58..85964f1 100644 --- a/tests/test_watermark_profiles.py +++ b/tests/test_watermark_profiles.py @@ -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: