Release 0.22.0 with composable feature extras

src/remove_ai_watermarks/__init__.py

@@ -25,7 +25,7 @@ _os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
 _warnings.filterwarnings("ignore", message=r".*ImageProcessorFast.*")
 
 
-__version__ = "0.21.2"
+__version__ = "0.22.0"
 
 __all__ = ["__version__", "remove_visible", "visible_provenance"]
 

src/remove_ai_watermarks/cli.py

@@ -183,7 +183,7 @@ _upscaler_option = click.option(
     "--upscaler",
     type=click.Choice(["lanczos", "esrgan"]),
     default="lanczos",
-    help="How to upscale a small input to the --min-resolution floor: lanczos (default, cv2, no deps) or "
+    help="How to upscale a small input to the --min-resolution floor: lanczos (default, cv2, no model) or "
     "esrgan (Real-ESRGAN via the 'esrgan' extra; better detail, slower on CPU). Best for photo/texture "
     "content -- as a generic GAN with no face/glyph prior it can degrade faces (diffusion mitigates) and "
     "thin text, so lanczos stays the default. Falls back to lanczos if the extra is absent. Only when upscaling.",
@@ -331,7 +331,7 @@ _visible_backend_option = click.option(
     default="auto",
     help="Fill backend for visible-mark removal (localize -> fill). auto: best available, "
     "LaMa > MI-GAN > cv2 (a learned backend needs the 'lama' or 'migan' extra; else cv2, "
-    "with a warning). cv2: classical inpaint (no deps, smears texture). migan: MI-GAN ONNX "
+    "with a warning). cv2: classical inpaint (no model download, smears texture). migan: MI-GAN ONNX "
     "(light, ~1 GB, the memory-tight pick). lama: big-LaMa ONNX (best quality, ~4.7 GB).",
 )
 
@@ -800,7 +800,7 @@ def _parse_region(spec: str) -> tuple[int, int, int, int]:
     "--backend",
     type=click.Choice(["cv2", "migan", "lama"]),
     default="cv2",
-    help="Inpaint backend. cv2: instant, no deps. migan: light ONNX MI-GAN, ~1 GB RAM, "
+    help="Inpaint backend. cv2: instant, no model download. migan: light ONNX MI-GAN, ~1 GB RAM, "
     "near-LaMa quality (extra 'migan'). lama: big-LaMa, best quality but ~4.7 GB RAM (extra 'lama').",
 )
 @click.option("--inpaint-method", type=click.Choice(["telea", "ns"]), default="telea", help="cv2 inpaint method.")
@@ -941,13 +941,14 @@ def cmd_invisible(
     """Remove invisible AI watermarks (SynthID, StableSignature, TreeRing).
 
     Uses diffusion-based regeneration. Requires GPU for reasonable speed.
-    Requires the [gpu] extra: pip install 'remove-ai-watermarks[gpu]'
+    Requires the [diffusion] extra: pip install 'remove-ai-watermarks[diffusion]'
     """
     from remove_ai_watermarks.invisible_engine import is_available as invisible_available
 
     if not invisible_available():
         console.print(
-            "Error: GPU dependencies not installed.\n  Install them with: pip install 'remove-ai-watermarks[gpu]'"
+            "Error: Diffusion dependencies not installed.\n"
+            "  Install them with: pip install 'remove-ai-watermarks[diffusion]'"
         )
         raise SystemExit(1)
 
@@ -1298,7 +1299,7 @@ def cmd_all(
             synthid_skipped = True
             console.print(
                 "    Warning: Skipped - GPU dependencies not installed.\n"
-                "    Install them with: pip install 'remove-ai-watermarks[gpu]'"
+                "    Install them with: pip install 'remove-ai-watermarks[diffusion]'"
             )
         elif _should_skip_invisible_scrub(force, source):
             # No locally-detectable invisible watermark -> skip the destructive
@@ -1404,7 +1405,7 @@ def cmd_all(
             "  visible mark and metadata were stripped.\n"
             "\n"
             "  Install the extra and rerun to remove it:\n"
-            "    pip install 'remove-ai-watermarks[gpu]'\n"
+            "    pip install 'remove-ai-watermarks[diffusion]'\n"
             "  ====================================================================="
         )
         raise SystemExit(1)
@@ -1766,7 +1767,7 @@ def cmd_batch(
             f"\n  WARNING: the invisible (SynthID) watermark was NOT removed on "
             f"{synthid_skipped_count} image(s) -- the GPU dependencies are not installed, "
             f"so those outputs still carry the invisible watermark.\n"
-            f"  Install the extra and rerun: pip install 'remove-ai-watermarks[gpu]'"
+            f"  Install the extra and rerun: pip install 'remove-ai-watermarks[diffusion]'"
         )
 
     # Non-zero exit so a wrapping service detects an incomplete/failed run (batch used

src/remove_ai_watermarks/dwt_dct.py

@@ -0,0 +1,95 @@
+"""DWT-DCT decoder compatible with invisible-watermark's ``dwtDct`` path.
+
+Derived from ShieldMnt/invisible-watermark ``imwatermark/maxDct.py`` (MIT),
+trimmed to the matrix path used by Stable Diffusion, SDXL, and FLUX.
+
+Copyright (c) 2021 ShieldMnt
+
+The complete upstream license is distributed in
+``licenses/invisible-watermark-MIT.txt``.
+"""
+
+# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportMissingTypeStubs=false
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import cv2
+import numpy as np
+import pywt
+
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
+_DEFAULT_SCALES = (0, 36, 36)
+_DEFAULT_BLOCK = 4
+
+
+class _DecodeMaxDct:
+    """Extract frequency-domain bits using the upstream matrix algorithm."""
+
+    def __init__(
+        self,
+        wm_lengths: tuple[int, ...],
+        scales: tuple[int, int, int] = _DEFAULT_SCALES,
+        block: int = _DEFAULT_BLOCK,
+    ) -> None:
+        self._wm_lengths = wm_lengths
+        self._scales = scales
+        self._block = block
+
+    def decode(self, bgr: NDArray[Any]) -> dict[int, NDArray[Any]]:
+        row, col, _channels = bgr.shape
+        yuv = cv2.cvtColor(bgr, cv2.COLOR_BGR2YUV)
+
+        scores_by_length = {wm_len: [[] for _ in range(wm_len)] for wm_len in self._wm_lengths}
+        for channel in range(2):
+            if self._scales[channel] <= 0:
+                continue
+            ca1, _detail = pywt.dwt2(yuv[: row // 4 * 4, : col // 4 * 4, channel], "haar")
+            self._decode_frame(ca1, self._scales[channel], scores_by_length)
+
+        return {
+            wm_len: np.asarray([float(np.asarray(score).mean()) if score else 0.0 for score in scores]) * 255 > 127
+            for wm_len, scores in scores_by_length.items()
+        }
+
+    def _decode_frame(
+        self,
+        frame: NDArray[Any],
+        scale: int,
+        scores_by_length: dict[int, list[list[int]]],
+    ) -> None:
+        row, col = frame.shape
+        bit_index = 0
+        for i in range(row // self._block):
+            for j in range(col // self._block):
+                block = frame[
+                    i * self._block : i * self._block + self._block,
+                    j * self._block : j * self._block + self._block,
+                ]
+                inferred = self._infer_bit(block, scale)
+                for wm_len, scores in scores_by_length.items():
+                    scores[bit_index % wm_len].append(inferred)
+                bit_index += 1
+
+    def _infer_bit(self, block: NDArray[Any], scale: int) -> int:
+        position = int(np.argmax(np.abs(block.flatten()[1:]))) + 1
+        i, j = position // self._block, position % self._block
+        value = abs(float(block[i][j]))
+        return int((value % scale) > 0.5 * scale)
+
+
+def decode_dwt_dct(bgr: NDArray[Any], wm_len: int) -> NDArray[Any]:
+    """Extract ``wm_len`` watermark bits from a BGR image."""
+    return decode_dwt_dct_lengths(bgr, (wm_len,))[wm_len]
+
+
+def decode_dwt_dct_lengths(bgr: NDArray[Any], wm_lengths: tuple[int, ...]) -> dict[int, NDArray[Any]]:
+    """Extract several watermark lengths with one DWT and block scan."""
+    if bgr.size == 0 or min(bgr.shape[:2]) * max(bgr.shape[:2]) < 256 * 256:
+        raise RuntimeError("image too small, should be larger than 256x256")
+    if not wm_lengths or any(wm_len <= 0 for wm_len in wm_lengths):
+        raise ValueError("watermark lengths must be positive")
+    return _DecodeMaxDct(wm_lengths=tuple(dict.fromkeys(wm_lengths))).decode(bgr)

src/remove_ai_watermarks/identify.py

@@ -722,8 +722,8 @@ def _visible_text_marks(image_path: Path, *, image: NDArray[Any] | None = None)
 def _invisible_watermark(image_path: Path) -> str | None:
     """Open invisible-watermark scheme name (SD/SDXL/FLUX) or None.
 
-    Optional: needs the imwatermark decoder (extra ``detect``). Returns None if
-    it is not installed or no known watermark decodes.
+    Optional: needs the torch-free DWT-DCT decoder (extra ``detect``). Returns
+    None if it is not installed or no known watermark decodes.
     """
     from remove_ai_watermarks.invisible_watermark import detect_invisible_watermark
 
@@ -761,6 +761,9 @@ def _collect_visible_signals(
         image = imread(image_path)
     except Exception as exc:  # cv2 missing - detectors fall back / no-op
         logger.debug("visible-mark decode unavailable: %s", exc)
+        return platform
+    if image is None:
+        return platform
 
     sparkle_conf = _visible_sparkle(image_path, image=image)
     if sparkle_conf is not None and sparkle_conf >= _SPARKLE_THRESHOLD:
@@ -1087,8 +1090,8 @@ def identify(
         image_path: Path to the image (PNG, JPEG, WebP, or ISOBMFF container).
         check_visible: Also run the registered visible-mark detectors through cv2.
             Set False for a metadata-only, dependency-light scan.
-        check_invisible: Also decode open invisible watermarks (SD/SDXL/FLUX) via
-            the optional imwatermark library. No-op when it is not installed.
+        check_invisible: Also decode optional open invisible watermarks
+            (SD/SDXL/FLUX). No-op when the decoder extra is not installed.
 
     File-backed metadata extraction runs first. The extracted evidence is then
     evaluated independently, followed by the optional pixel-backed visible and

src/remove_ai_watermarks/invisible_engine.py

@@ -4,7 +4,7 @@ Wraps the vendored noai-watermark code for removing invisible AI watermarks
 (SynthID, StableSignature, TreeRing) via diffusion-based regeneration.
 
 This module requires the 'gpu' extra dependencies:
-    uv pip install 'remove-ai-watermarks[gpu]'
+    uv pip install 'remove-ai-watermarks[diffusion]'
 """
 
 # cv2/torch boundary: this engine wraps cv2 (resize/imwrite/cvtColor) and the
@@ -226,7 +226,7 @@ class InvisibleEngine:
                 input size, so this is a transparent quality boost; it adds time
                 and memory on small inputs. Ignored on a min > max misconfig.
             upscaler: How to upscale a small input to the ``min_resolution`` floor:
-                ``"lanczos"`` (default, cv2, no deps) or ``"esrgan"`` (Real-ESRGAN
+                ``"lanczos"`` (default, cv2, no model download) or ``"esrgan"`` (Real-ESRGAN
                 via the ``esrgan`` extra). Only applies when UPscaling (the floor
                 case); a ``max_resolution`` downscale always uses Lanczos. Falls back
                 to Lanczos if the extra is absent.

src/remove_ai_watermarks/invisible_watermark.py

@@ -14,21 +14,20 @@ source:
 
 The watermark is fragile: it does NOT survive JPEG re-encoding or resizing
 (verified -- gone after JPEG q90), so detection works only on pristine PNG
-originals. Absence is never proof. Requires the optional ``invisible-watermark``
-package (extra: ``detect``); ``detect_invisible_watermark`` returns None when it
-is not installed.
+originals. Absence is never proof. Requires the optional ``detect`` extra;
+``detect_invisible_watermark`` returns None when it is not installed.
 """
 
-# imwatermark ships no type stubs (like cv2); its decoder returns are Unknown.
-# Relax the untyped-library diagnostics for this thin wrapper module only.
+# The optional numeric libraries do not provide complete types for this path.
 # pyright: reportMissingTypeStubs=false, reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false
 
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+    from collections.abc import Iterable
     from pathlib import Path
 
 logger = logging.getLogger(__name__)
@@ -49,10 +48,10 @@ _MATCH_SD1_FRAC = 0.92  # fraction of the 136 string bits that must match
 
 
 def is_available() -> bool:
-    """True if the optional imwatermark decoder is installed."""
+    """True when all dependencies for the optional DWT-DCT decoder exist."""
     from .optional_deps import module_available
 
-    return module_available("imwatermark")
+    return module_available("cv2", "numpy", "pywt")
 
 
 def _bits_match(value: int, ref: int, width: int = 48) -> int:
@@ -68,6 +67,20 @@ def _bytes_match_frac(a: bytes, b: bytes) -> float:
     return 1.0 - diff / (8 * len(b))
 
 
+def _bits_to_int(bits: Iterable[object]) -> int:
+    value = 0
+    for bit in bits:
+        value = (value << 1) | int(bool(bit))
+    return value
+
+
+def _bits_to_bytes(bits: Iterable[object], nbytes: int) -> bytes:
+    import numpy as np
+
+    packed = np.packbits([int(bool(bit)) for bit in bits])
+    return bytes(int(value) for value in packed[:nbytes])
+
+
 def detect_invisible_watermark(image_path: Path) -> str | None:
     """Return the embedding scheme name if a known open watermark is decoded.
 
@@ -78,32 +91,26 @@ def detect_invisible_watermark(image_path: Path) -> str | None:
     """
     if not is_available():
         return None
-    from imwatermark import WatermarkDecoder
-
     from remove_ai_watermarks import image_io
+    from remove_ai_watermarks.dwt_dct import decode_dwt_dct_lengths
 
     img = image_io.imread(image_path)
     if img is None:
         return None
 
-    # 48-bit fixed-message watermarks (SDXL, FLUX.2).
     try:
-        bits = WatermarkDecoder("bits", 48).decode(img, "dwtDct")
-        value = 0
-        for bit in bits:
-            value = (value << 1) | (1 if bit else 0)
-        for name, ref in _BITS_48.items():
-            if _bits_match(value, ref) >= _MATCH_48:
-                return name
+        decoded = decode_dwt_dct_lengths(img, (48, 8 * len(_SD1_STRING)))
     except Exception as exc:  # decode can fail on tiny images
-        logger.debug("48-bit watermark decode failed for %s: %s", image_path, exc)
+        logger.debug("watermark decode failed for %s: %s", image_path, exc)
+        return None
 
-    # 136-bit default string watermark (SD 1.x / 2.x).
-    try:
-        raw = cast("bytes", WatermarkDecoder("bytes", 8 * len(_SD1_STRING)).decode(img, "dwtDct"))
-        if _bytes_match_frac(raw, _SD1_STRING) >= _MATCH_SD1_FRAC:
-            return "Stable Diffusion 1.x / 2.x"
-    except Exception as exc:
-        logger.debug("string watermark decode failed for %s: %s", image_path, exc)
+    value = _bits_to_int(decoded[48])
+    for name, ref in _BITS_48.items():
+        if _bits_match(value, ref) >= _MATCH_48:
+            return name
+
+    raw = _bits_to_bytes(decoded[8 * len(_SD1_STRING)], len(_SD1_STRING))
+    if _bytes_match_frac(raw, _SD1_STRING) >= _MATCH_SD1_FRAC:
+        return "Stable Diffusion 1.x / 2.x"
 
     return None

src/remove_ai_watermarks/licenses/invisible-watermark-MIT.txt

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 ShieldMnt
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

src/remove_ai_watermarks/noai/__init__.py

@@ -7,7 +7,7 @@ is exposed **lazily** via PEP 562 ``__getattr__``: importing a light submodule
 (e.g. ``noai.c2pa`` / ``noai.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 ``gpu``/``detect`` extras are present --
+torch) even in a full install where the ``diffusion`` extra is present --
 otherwise the mere presence of torch in the env inflated identify to ~420 MB and
 risked OOM on a 512 MB host.
 """

src/remove_ai_watermarks/noai/c2pa.py

@@ -43,7 +43,7 @@ from remove_ai_watermarks.noai.constants import (
 
 logger = logging.getLogger(__name__)
 
-# Official C2PA reader (c2pa-python, a core dependency). It is the primary,
+# 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

src/remove_ai_watermarks/noai/constants.py

@@ -7,7 +7,7 @@ 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 core
+# 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).

src/remove_ai_watermarks/noai/watermark_remover.py

@@ -476,8 +476,9 @@ class WatermarkRemover:
         """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 (the ``detect``
-        extra). A watermark REMOVER must not re-stamp a detectable AI watermark, or the
+        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).

src/remove_ai_watermarks/upscaler.py

@@ -4,7 +4,7 @@ Mirrors ``region_eraser``'s optional-backend pattern: ``is_available()`` guards
 ``spandrel`` import, a lazy singleton (double-checked lock) holds the loaded model, and
 the weights download on first use (cached by ``torch.hub``) -- they are never bundled.
 
-The DEFAULT upscaler stays Lanczos (cv2, no deps); this is opt-in via the ``esrgan``
+The DEFAULT upscaler stays Lanczos (cv2, no model download); this is opt-in via the ``esrgan``
 extra and feeds the ``--upscaler esrgan`` path. ``spandrel`` is a pure model-loader
 (MIT) with NO basicsr dependency -- it pulls only torch/torchvision/safetensors/numpy/
 einops -- so it sidesteps the basicsr / ``torchvision.transforms.functional_tensor``