refactor: unify C2PA vendor registry + code-health fixes + uv publish

Three P2 cleanups from a library-wide review. Detection -- single C2PA_AI_VENDORS registry (noai/constants.py): - C2PA_ISSUERS, SYNTHID_C2PA_ISSUERS, and identify._ISSUER_PLATFORM now derive from one C2paAiVendor table, so adding a C2PA vendor is one entry instead of edits in three places across two files. Behavior-identical (262 detection tests pass; the kept `needle` field is load-bearing -- it differs from `org` for Google and ByteDance, with no mechanical derivation). Code-health: - region_eraser.erase_lama now accepts grayscale/BGRA like erase_cv2 (it crashed on grayscale and silently dropped alpha on BGRA). +2 regression tests. - batch frees the device cache between images via a shared try_empty_device_cache helper (generalized from the MPS-only _try_clear_mps_cache, now reused by both the MPS->CPU fallback and the batch loop). - batch gained --controlnet-scale (parity with invisible/all). CI / packaging: - publish.yml uploads via `uv publish` (PyPI trusted publishing over OIDC), replacing pypa/gh-action-pypi-publish so uploads no longer depend on that action's bundled twine accepting the Metadata-Version. Workflow filename + pypi environment unchanged, so PyPI's trusted-publisher entry still matches. - hatchling pin relaxed <1.28 -> <1.31 (verified against hatch's changelog: 1.30.0 made Metadata 2.5 the default, 1.30.1 reverted to 2.4; 1.27-1.29 were always 2.4). Kept as belt-and-suspenders so the first uv-publish release ships 2.4, isolating the uploader swap from the metadata-version bump. Docs (CLAUDE.md, pyproject) synced; corrected the inaccurate "hatchling 1.28+ emits 2.5" note. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-04 18:18:00 +02:00 · 2026-06-03 20:54:39 -07:00
parent 9bd2c17cc4
commit 5cf68a6a3d
10 changed files with 155 additions and 42 deletions
@@ -30,5 +30,12 @@ jobs:
      - name: Build package
        run: uv build

+      # Publish with `uv publish` (its own uploader, not the twine bundled in the
+      # old pypa action). Trusted publishing is automatic in GitHub Actions: the
+      # `id-token: write` permission + the `pypi` environment supply the OIDC token,
+      # and PyPI's trusted-publisher entry matches on repo + workflow filename +
+      # environment name (all unchanged from the pypa-action setup), so no API token
+      # is needed. uv's uploader also accepts Metadata-Version 2.5 -- the permanent
+      # fix that lets the hatchling pin be relaxed (see pyproject [build-system]).
      - name: Publish to PyPI
-        uses: pypa/gh-action-pypi-publish@release/v1
+        run: uv publish
@@ -132,12 +132,18 @@ remove-ai-watermarks = "remove_ai_watermarks.cli:main"
 Repository = "https://github.com/wiltodelta/remove-ai-watermarks"

 [build-system]
-# Pin hatchling < 1.28: 1.28+ emits Metadata-Version 2.5 (PEP 639), which the twine
-# bundled in pypa/gh-action-pypi-publish@release/v1 rejects ("'2.5' is not a valid
-# Metadata-Version"), failing the PyPI upload (v0.8.3, 2026-06-01). 1.27.x emits 2.4,
-# which uploads fine (0.8.2 shipped on it). Lift this pin once the publish action's
-# twine is upgraded to >= 6.1.0 (2.5-aware) or the workflow moves to `uv publish`.
-requires = ["hatchling<1.28"]
+# Pin hatchling < 1.31. hatchling 1.30.0 made Metadata-Version 2.5 (PEP 794) the
+# default, which the twine bundled in pypa/gh-action-pypi-publish@release/v1 rejects
+# ("'2.5' is not a valid Metadata-Version"), failing the v0.8.3 PyPI upload
+# (2026-06-01) when unpinned requires = ["hatchling"] pulled 1.30.0. hatchling 1.30.1
+# reverted the default to 2.4 ("kept at 2.4 until more tools support 2.5"), and
+# 1.27-1.29 were always 2.4 -- so < 1.31 keeps `uv build` on a 2.4-emitting hatchling
+# (it resolves to the latest allowed, 1.30.1). The publish workflow now uses
+# `uv publish`, whose uploader accepts 2.5, so this pin is belt-and-suspenders, not
+# load-bearing: keeping it makes the first uv-publish release ship 2.4 metadata
+# (isolating the uploader swap from the metadata-version bump). Drop to
+# `requires = ["hatchling"]` once that release confirms the path.
+requires = ["hatchling<1.31"]
 build-backend = "hatchling.build"

 [tool.hatch.build.targets.wheel]
@@ -990,6 +990,7 @@ def _process_batch_image(
    min_resolution: int = 1024,
    restore_faces: bool = False,
    restore_faces_weight: float = 0.5,
+    controlnet_scale: float = 1.0,
 ) -> None:
    """Process a single image for batch mode.

@@ -1040,6 +1041,7 @@ def _process_batch_image(
                    device=None if device == "auto" else device,
                    pipeline=pipeline,
                    hf_token=hf_token,
+                    controlnet_conditioning_scale=controlnet_scale,
                )
            engine_inv = ctx.obj["_inv_engine"]
            engine_inv.remove_watermark(
@@ -1114,6 +1116,7 @@ def _process_batch_image(
@_restore_faces_options
@_min_resolution_option
@_unsharp_option
+@_controlnet_scale_option
@click.pass_context
 def cmd_batch(
    ctx: click.Context,
@@ -1133,6 +1136,7 @@ def cmd_batch(
    min_resolution: int,
    restore_faces: bool,
    restore_faces_weight: float,
+    controlnet_scale: float,
 ) -> None:
    """Process all images in a directory."""
    _banner()
@@ -1187,6 +1191,7 @@ def cmd_batch(
                    min_resolution=min_resolution,
                    restore_faces=restore_faces,
                    restore_faces_weight=restore_faces_weight,
+                    controlnet_scale=controlnet_scale,
                )
                processed += 1

@@ -40,7 +40,7 @@ from remove_ai_watermarks.metadata import (
    xai_signature,
 )
 from remove_ai_watermarks.noai.c2pa import cbor_text_after, extract_c2pa_info, soft_binding_vendors_in
-from remove_ai_watermarks.noai.constants import C2PA_AI_TOOLS, C2PA_ISSUERS
+from remove_ai_watermarks.noai.constants import C2PA_AI_TOOLS, C2PA_AI_VENDORS, C2PA_ISSUERS

 if TYPE_CHECKING:
    from pathlib import Path
@@ -59,19 +59,13 @@ _SCAN_BYTES = 1024 * 1024
 # signal (e.g. an OpenAI image scored 0.37 -- below threshold, correctly dropped).
 _SPARKLE_THRESHOLD = 0.5

-# Issuer (C2PA signer) -> human-readable generating platform. Ordered: when a
-# manifest names several issuers (Microsoft Designer signs as "OpenAI,
-# Microsoft"), the first match wins so the product, not the backend, is named.
-_ISSUER_PLATFORM: tuple[tuple[str, str], ...] = (
-    # 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.
-    ("Microsoft", "Microsoft (Bing Image Creator / Designer)"),
-    ("Adobe", "Adobe Firefly"),
-    ("OpenAI", "OpenAI (ChatGPT / gpt-image / DALL-E / Sora)"),
-    ("Google", "Google (Gemini / Imagen)"),
-    ("Stability AI", "Stability AI (Stable Image / DreamStudio)"),
-    ("Black Forest Labs", "Black Forest Labs (FLUX)"),
-    ("ByteDance", "ByteDance (Doubao / Jimeng / Volcano Engine)"),
+# Issuer (C2PA signer) -> human-readable generating platform, derived from the
+# single C2PA_AI_VENDORS registry. Ordered: when a manifest names several issuers
+# (Microsoft Designer signs as "OpenAI, Microsoft"), the first match wins so the
+# product, not the backend, is named -- the registry order encodes that priority.
+# Signing authorities without an AI platform (e.g. Truepic) are skipped here.
+_ISSUER_PLATFORM: tuple[tuple[str, str], ...] = tuple(
+    (v.needle, v.platform) for v in C2PA_AI_VENDORS if v.platform is not None and v.needle is not None
 )

 # PNG-text / EXIF keys that indicate a local diffusion pipeline (vs. a hosted
@@ -4,6 +4,8 @@ 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
 SUPPORTED_FORMATS = {".png", ".jpg", ".jpeg", ".webp"}

@@ -78,25 +80,56 @@ C2PA_SIGNATURES = [
    b"manifest",
 ]

-# C2PA known issuers
-C2PA_ISSUERS = {
-    b"Google": "Google LLC",
-    b"Adobe": "Adobe",
-    b"Microsoft": "Microsoft",
-    b"OpenAI": "OpenAI",
-    b"Truepic": "Truepic",
+
+# 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
+
+
+# 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.
-    b"Stability AI": "Stability AI",
+    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.
-    b"Black Forest Labs": "Black Forest Labs",
+    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.
-    b"volcengine": "ByteDance (Volcano Engine)",
-}
+    C2paAiVendor(
+        b"volcengine", "ByteDance (Volcano Engine)", "ByteDance (Doubao / Jimeng / Volcano Engine)", "ByteDance"
+    ),
+    # 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),
+)
+
+# 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}

 # C2PA issuers whose signed outputs also carry an invisible SynthID pixel
 # watermark -- a metadata proxy for "SynthID is in the pixels":
@@ -117,7 +150,8 @@ C2PA_ISSUERS = {
 # 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.
-SYNTHID_C2PA_ISSUERS = frozenset({b"Google", b"OpenAI"})
+# 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 = {
@@ -102,7 +102,7 @@ def run_img2img_with_mps_fallback(
        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_clear_mps_cache()
+            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
@@ -137,9 +137,16 @@ def _call_pipeline(
    return pipeline(**kwargs)


-def _try_clear_mps_cache() -> None:
+def try_empty_device_cache(device: str) -> None:
+    """Best-effort free of cached GPU/MPS/XPU memory for ``device``.
+
+    ``torch.<device>.empty_cache()`` exists for cuda/mps/xpu but not cpu (the
+    hasattr guard skips the cpu no-op). Never raises -- callers use it as cleanup
+    (the MPS->CPU fallback here, and the batch loop in watermark_remover).
+    """
    with contextlib.suppress(Exception):
        import torch

-        if hasattr(torch, "mps"):
-            torch.mps.empty_cache()  # type: ignore[attr-defined]
+        backend = getattr(torch, device, None)
+        if backend is not None and hasattr(backend, "empty_cache"):
+            backend.empty_cache()  # type: ignore[attr-defined]
@@ -692,6 +692,9 @@ class WatermarkRemover:
        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
@@ -705,6 +708,7 @@ class WatermarkRemover:
                    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

@@ -108,7 +108,19 @@ def erase_lama(image_bgr: NDArray[Any], mask: NDArray[Any]) -> NDArray[Any]:
    LaMa runs at a fixed square input size. To preserve full-image resolution we
    crop a padded region around the mask, inpaint that crop at the model size,
    and paste only the masked pixels back -- so untouched areas stay pixel-exact.
+
+    Like ``erase_cv2``, accepts 1-channel (grayscale) and 4-channel (BGRA) input:
+    LaMa runs on 3-channel BGR, so grayscale is promoted to BGR (result demoted
+    back) and a BGRA alpha plane is split off and re-attached unchanged. Without
+    this the ``cv2.cvtColor(..., BGR2RGB)`` below would crash on grayscale and
+    silently drop alpha on BGRA.
    """
+    if image_bgr.ndim == 2:
+        bgr = erase_lama(cv2.cvtColor(image_bgr, cv2.COLOR_GRAY2BGR), mask)
+        return cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
+    if image_bgr.ndim == 3 and image_bgr.shape[2] == 4:
+        bgr = erase_lama(np.ascontiguousarray(image_bgr[:, :, :3]), mask)
+        return np.dstack([bgr, image_bgr[:, :, 3]])
    session = _get_lama_session()
    inp = session.get_inputs()  # type: ignore[attr-defined]
    img_name = inp[0].name
@@ -90,3 +90,47 @@ class TestLamaBackend:
            pytest.skip("onnxruntime installed; cannot test the unavailable path")
        with pytest.raises(RuntimeError, match="onnxruntime"):
            erase(img, boxes=[(10, 10, 20, 20)], backend="lama")
+
+
+class TestLamaChannelHandling:
+    """erase_lama must accept grayscale (2D) and BGRA (4-channel) like erase_cv2.
+
+    The real ONNX model is never loaded -- the session is faked to an identity
+    inpaint, so this exercises only the channel promote/split wrapper (the fix for
+    LaMa crashing on grayscale and dropping alpha on BGRA).
+    """
+
+    @pytest.fixture
+    def _fake_lama(self, monkeypatch: pytest.MonkeyPatch):
+        from remove_ai_watermarks import region_eraser
+
+        class _In:
+            def __init__(self, name: str, shape: list[int]):
+                self.name = name
+                self.shape = shape
+
+        class _FakeSession:
+            def get_inputs(self):
+                return [_In("image", [1, 3, 512, 512]), _In("mask", [1, 1, 512, 512])]
+
+            def run(self, _outputs, feeds):
+                # Identity inpaint: echo the image tensor (1,3,size,size) back.
+                return [feeds["image"]]
+
+        monkeypatch.setattr(region_eraser, "lama_available", lambda: True)
+        monkeypatch.setattr(region_eraser, "_get_lama_session", lambda: _FakeSession())
+
+    @pytest.mark.usefixtures("_fake_lama")
+    def test_grayscale_2d_does_not_raise(self):
+        gray = np.full((100, 100), 120, np.uint8)
+        out = erase(gray, boxes=[(40, 40, 20, 20)], backend="lama")
+        assert out.ndim == 2
+        assert out.shape == gray.shape
+
+    @pytest.mark.usefixtures("_fake_lama")
+    def test_bgra_preserves_alpha(self):
+        bgra = np.full((100, 100, 4), 120, np.uint8)
+        bgra[..., 3] = 200  # opaque-ish alpha plane
+        out = erase(bgra, boxes=[(40, 40, 20, 20)], backend="lama")
+        assert out.shape == bgra.shape
+        assert np.array_equal(out[..., 3], bgra[..., 3])  # alpha carried through unchanged