diff --git a/.github/workflows/publish-photo-classify-hf.yml b/.github/workflows/publish-photo-classify-hf.yml index 5a71729..a1deda0 100644 --- a/.github/workflows/publish-photo-classify-hf.yml +++ b/.github/workflows/publish-photo-classify-hf.yml @@ -6,8 +6,9 @@ name: Publish photo classifier to Hugging Face # Hub model: wiltodelta/raiw-photo-classify # Secret: HF_TOKEN, write-role Hugging Face token (same secret Space sync uses). # -# Weights are not in git. Full mode downloads them from a GitHub Release -# of this repository, then uploads with the card from docs/photo-classify-hf/. +# Weights are not in git. Full mode downloads the frozen PyTorch files from a +# GitHub Release, exports the vision-only ONNX graph, then uploads everything +# with the card from docs/photo-classify-hf/. on: workflow_dispatch: @@ -52,6 +53,14 @@ jobs: gh release download "$TAG" --repo "${{ github.repository }}" --dir weights ls -l weights + - name: Export the vision-only ONNX graph + if: ${{ github.event.inputs.mode == 'full' }} + run: | + set -euo pipefail + uv run --extra classify --with onnx python scripts/export_photo_classify_onnx.py \ + --checkpoint weights/clip-l-ft.pt \ + --out weights/clip-l-ft-vision-fp32.onnx + - name: Publish to Hugging Face env: HF_TOKEN: ${{ secrets.HF_TOKEN }} diff --git a/docs/development.md b/docs/development.md index 3302f74..4eda755 100644 --- a/docs/development.md +++ b/docs/development.md @@ -29,26 +29,22 @@ Do not add advisory ignores or suppress the scanner's exit status. Run and report Ruff, Pyright scoped to `src/`, and tests separately when security blocks the gate. -Rechecked on 2026-09-08 against PyPI release metadata and the advisory sources: +Rechecked on 2026-09-10 against the locked versions and the advisory sources: -- `lightning`, pulled by the optional `trustmark` extra, remains affected by - [PYSEC-2026-3624](https://github.com/pypa/advisory-database/blob/main/vulns/lightning/PYSEC-2026-3624.yaml) - and its alias [GHSA-qqmf-gpg7-g8gw](https://github.com/advisories/GHSA-qqmf-gpg7-g8gw). - These describe one checkpoint-loading vulnerability. The latest PyPI release - is still affected; the advisory's `2022.6.15` fix entry is not a newer release - on the supported 2.x line. The adapter selects TrustMark's supplied weights, - rather than accepting a caller-supplied Lightning checkpoint. +- `lightning` and `pytorch-lightning`, pulled by the optional `trustmark` + extra, are pinned to 2.6.6. That release clears PYSEC-2026-3624 and + PYSEC-2026-3967; do not reintroduce an older transitive resolution. - `accelerate`, required by the diffusion stack, is affected by - [GHSA-4j2p-28q2-5m79](https://github.com/advisories/GHSA-4j2p-28q2-5m79). - The sharded-checkpoint loader accepts unsafe `weight_map` paths; the advisory - lists no patched release. The project has no direct call to + [GHSA-4j2p-28q2-5m79](https://github.com/advisories/GHSA-4j2p-28q2-5m79) + and PYSEC-2026-3804. The sharded-checkpoint loader accepts unsafe + `weight_map` paths; neither advisory lists a patched release. The project has no direct call to `load_checkpoint_in_model` or `load_checkpoint_and_dispatch`, but this does not prove indirect model-loading paths are unaffected. -Before carrying either block forward, rerun `uvx uv-secure uv.lock` and check +Before carrying the remaining block forward, rerun `uvx uv-secure uv.lock` and check PyPI for a fixed release. Upgrade a fixed package with -`uv lock --upgrade-package ` and rerun the gate. Removing either -current dependency would remove a supported optional feature, so that is not a +`uv lock --upgrade-package ` and rerun the gate. Removing the current +dependency would remove a supported optional feature, so that is not a maintenance-only repair. Minor-only lock refreshes must retain each package's current major and lower diff --git a/docs/known-limitations.md b/docs/known-limitations.md index 777204b..1d2423e 100644 --- a/docs/known-limitations.md +++ b/docs/known-limitations.md @@ -516,11 +516,9 @@ remain identical because isolated decoder hits can otherwise be content noise. It accepts Variant P schemas 0-2. Variant Q requires a different model, and schema 3 is rejected at the measured precision threshold. Its NumPy 1.x runtime limits the extra to Python 3.11-3.12; the rest of the package remains supported -through Python 3.14. The current TrustMark dependency line also resolves lightning 2.6.5 -(PYSEC-2026-3624, no fixed release yet); the vulnerable `load_from_checkpoint` -path is unreachable here because TrustMark loads its checksummed checkpoints -with plain `torch.load`. Bump lightning and cut a patch release when a fix -ships. The calibration history is in +through Python 3.14. The current TrustMark dependency line resolves +`lightning` and `pytorch-lightning` 2.6.6, which contain the published fixes +for their checkpoint-loading advisories. The calibration history is in [module internals](module-internals.md#metadata-and-provenance). A generic metadata-free AI-generated-image classifier is not shipped. It is a diff --git a/docs/module-internals.md b/docs/module-internals.md index f65d94d..a280eb5 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -1085,9 +1085,9 @@ module. Weights stay out of git. The Hub snapshot is `wiltodelta/raiw-photo-classify`. `RAIW_CLASSIFY_WEIGHTS` overrides it. The extra is `classify`. The runtime -resolves the snapshot through `WEIGHTS_ALLOW_PATTERNS`, exported so an -offline deploy's pre-cache installs the exact runtime request as one name -(`tests/test_classify.py` pins the seam). User guide: +narrows its snapshot download to the selected backend's model and shared heads. +`WEIGHTS_ALLOW_PATTERNS` exports the union so an offline deploy can pre-cache +both backends explicitly (`tests/test_classify.py` pins the seam). User guide: [photo-classify.md](photo-classify.md). Hub card: [photo-classify-hf/README.md](photo-classify-hf/README.md). @@ -1103,17 +1103,32 @@ after. `tests/test_clip_l_ft.py` rejects a loader that invokes random Linear initialization and verifies that checkpoint parameters and constructor-created buffers are fully materialized. -A same-day vision-only FP32 ONNX probe was 1,216,102,506 bytes versus the -1,710,697,163-byte full PyTorch checkpoint, loaded an ONNX Runtime CPU session -in 0.419 seconds, ran the probe in 0.136 seconds, and preserved all 42 tracked -fixture verdicts. It is not a shipped path: the full checkpoint contains -494,601,216 bytes of unused text-tower weights, but replacing it requires a new -published vision artifact and a full labeled-corpus certification. Dynamic INT8 -reduced the ONNX file to 305,462,280 bytes but changed 3 of the 42 fixture -verdicts, in both directions across the public decision boundary, so it is -rejected without retraining and a new operating point. Splitting Model 2 does -not address cold start: its detector and provider files total about 2 MB and -loaded in roughly 3 ms in the same stage profile; CLIP-L Model 1 dominates. +The optional CPU-only `backend="onnx"` path reads a static batch-one, +vision-only FP32 graph produced by `scripts/export_photo_classify_onnx.py`. +PyTorch remains the default until that artifact is published in the pinned Hub +snapshot. The graph is 1,216,102,506 bytes versus the 1,710,697,163-byte full +PyTorch checkpoint; the removed text tower alone is 494,601,216 bytes. A +2026-09-10 stress comparison covered 500 locally available catalog photographs +closest to either Model 1 threshold plus 500 deterministic family-stratified +photographs. The minimum embedding cosine was 0.999999979, mean absolute delta +was 1.54e-7, and no detector or receipt-gate decision changed. All 42 tracked +fixtures also kept byte-identical public records. + +Three alternating fresh-process CPU pairs on the same synthetic 512x384 PNG +gave PyTorch 7.037/5.235/5.003 seconds and ONNX 6.058/4.956/4.528 seconds under +the then-current machine load: ONNX won all three pairs, but its 4.956-second +median was only 5% below PyTorch's 5.235 seconds. Median peak RSS fell from +3.78 GB to 1.92 GB. These are process-cold, OS-cache-warm local results, not a +hosted-container estimate; the memory reduction is the stronger result. ONNX +Runtime graph optimization is disabled because rebuilding the already +constant-folded 1.2 GB graph increased session startup. CoreML is not selected: +its static session took 18.8 seconds to construct and its dynamic-batch probe +failed. Dynamic INT8 reduced the graph to 305,462,280 bytes but changed 3 of the +42 fixture verdicts in both directions across the public decision boundary, so +it remains rejected without retraining and a new operating point. Splitting +Model 2 does not address cold start: its detector and provider files total about +2 MB and loaded in roughly 3 ms in the same stage profile; CLIP-L Model 1 +dominates. ### Pixel forensics diff --git a/docs/photo-classify-hf/README.md b/docs/photo-classify-hf/README.md index ef24e1f..a16893d 100644 --- a/docs/photo-classify-hf/README.md +++ b/docs/photo-classify-hf/README.md @@ -258,6 +258,7 @@ silently change the heads. | File | Role | | --- | --- | | `clip-l-ft.pt` | HeadedCLIP state dict (CLIP-L plus unused linear head) | +| `clip-l-ft-vision-fp32.onnx` | Static batch-one vision tower and projection for the optional CPU backend | | `probe-weights-clip-l-ft.npz` | Ridge mean, scale, weights, `thr_oi_1pct` | | `detector.pt` | Freeze MLP | | `provider.pt` | Focal heads keyed `openai`, `google`, `bytedance`, `tc260`, `meta_muse_image`, `no_ai` | diff --git a/docs/photo-classify-hf/operating-point.json b/docs/photo-classify-hf/operating-point.json index 91f5f4c..78a9278 100644 --- a/docs/photo-classify-hf/operating-point.json +++ b/docs/photo-classify-hf/operating-point.json @@ -52,6 +52,7 @@ }, "files": [ "clip-l-ft.pt", + "clip-l-ft-vision-fp32.onnx", "probe-weights-clip-l-ft.npz", "detector.pt", "provider.pt" diff --git a/docs/photo-classify.md b/docs/photo-classify.md index 608c468..7ccee97 100644 --- a/docs/photo-classify.md +++ b/docs/photo-classify.md @@ -38,7 +38,11 @@ print(result.label, result.detector, result.provider) Missing extra raises `RuntimeError` with `pip install 'remove-ai-watermarks[classify]'`. `device` is a library -parameter (`None` / `"auto"` / `"cpu"` / `"cuda"`), not a CLI option. +parameter (`None` / `"auto"` / `"cpu"` / `"cuda"`), not a CLI option. The +explicit `backend="onnx"` path is CPU-only and requires +`clip-l-ft-vision-fp32.onnx` in `RAIW_CLASSIFY_WEIGHTS`; PyTorch remains the +default until that artifact is published in the pinned Hub snapshot. Install +the experimental runtime with `pip install 'remove-ai-watermarks[classify-onnx]'`. Weights are not in git. First call downloads [`wiltodelta/raiw-photo-classify`](https://huggingface.co/wiltodelta/raiw-photo-classify) @@ -47,9 +51,18 @@ at the freeze revision pinned in `classify.py`, or reads `probe-weights-clip-l-ft.npz`, `detector.pt`, and `provider.pt`. The training catalog is not published with the package. Deployments that pre-cache the snapshot for offline resolution must request exactly -`classify.WEIGHTS_ALLOW_PATTERNS` (the runtime resolves that same list under -`HF_HUB_OFFLINE=1`; a partial pre-cache fails the whole resolution, even for -files the runtime would fall back on). +`classify.WEIGHTS_ALLOW_PATTERNS`. That exported list covers both backends; +the online runtime narrows its download to the selected model plus the shared +heads. A partial offline pre-cache fails resolution, even for files the runtime +would otherwise fall back on. + +Generate the optional vision-only graph from the exact frozen checkpoint: + +```bash +uv run --with onnx python scripts/export_photo_classify_onnx.py \ + --checkpoint /path/to/clip-l-ft.pt \ + --out /path/to/clip-l-ft-vision-fp32.onnx +``` ## What one call returns diff --git a/docs/python-api.md b/docs/python-api.md index b611ad0..af666b3 100644 --- a/docs/python-api.md +++ b/docs/python-api.md @@ -226,7 +226,9 @@ head abstains to `None`. Otherwise it is `None`, including when 124-d extraction refuses the file. `device` is a library parameter: `None` / `"auto"` detect, `"cpu"` or `"cuda"` -pin. It is not a CLI option. +pin. It is not a CLI option. `backend="onnx"` selects the optional CPU-only +FP32 vision runtime and requires `remove-ai-watermarks[classify-onnx]`; +`backend="torch"` remains the default. Missing extra raises `RuntimeError` with the quoted install command `'remove-ai-watermarks[classify]'`. Guide: diff --git a/docs/release-and-distribution.md b/docs/release-and-distribution.md index f3ac293..a78582c 100644 --- a/docs/release-and-distribution.md +++ b/docs/release-and-distribution.md @@ -79,8 +79,9 @@ release fan-out. It holds the photo-classify freeze weights. The library extra `RAIW_CLASSIFY_WEIGHTS`. Update the Hub repo only with `.github/workflows/publish-photo-classify-hf.yml` (`workflow_dispatch`, type `publish`). Full mode downloads the four weight files from the GitHub -Release tag `photo-classify-freeze-2026-08-31` (or the tag you pass) and -uploads them with the card in `docs/photo-classify-hf/`. Card mode updates +Release tag `photo-classify-freeze-2026-08-31` (or the tag you pass), exports +the static FP32 vision-only ONNX graph from `clip-l-ft.pt`, and uploads all five +files with the card in `docs/photo-classify-hf/`. Card mode updates the README, `operating-point.json`, and the receipt-gate head staged from the package assets; freeze weights are untouched. A library version bump alone does not upload new heads -- bump the pinned Hub revision when a new diff --git a/pyproject.toml b/pyproject.toml index 6f41954..5d5502c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -154,6 +154,12 @@ classify = [ "tokenizers>=0.22,<0.23", "huggingface-hub>=0.20.0", ] +# Optional vision-only CPU runtime for the photo classifier. Keep it separate +# until the ONNX artifact is published and the backend becomes the default. +classify-onnx = [ + "remove-ai-watermarks[classify]", + "onnxruntime>=1.24.0", +] # Universal region eraser backend -- big-LaMa via onnxruntime (Carve/LaMa-ONNX, # Apache-2.0). CPU, no torch. Model (~200 MB) is downloaded on first use and # cached by huggingface_hub; it is never bundled in this repo. The default cv2 diff --git a/scripts/export_photo_classify_onnx.py b/scripts/export_photo_classify_onnx.py new file mode 100644 index 0000000..2d0da1c --- /dev/null +++ b/scripts/export_photo_classify_onnx.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +"""Export the frozen CLIP-L vision tower as a static FP32 ONNX graph. + +Run from the repository root with the ONNX exporter installed temporarily: + + uv run --with onnx python scripts/export_photo_classify_onnx.py \ + --checkpoint /path/to/clip-l-ft.pt \ + --out /path/to/clip-l-ft-vision-fp32.onnx +""" + +from __future__ import annotations + +import argparse +import hashlib +import inspect +import logging +import sys +from pathlib import Path +from typing import Any + +import torch +from torch import nn + +REPO_ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(REPO_ROOT)) + +log = logging.getLogger(__name__) + + +class VisionEmbedding(nn.Module): + """Only the vision modules used by the classifier's embedding path.""" + + def __init__(self, headed: Any) -> None: + super().__init__() + self.vision_model = headed.clip.vision_model + self.visual_projection = headed.clip.visual_projection + + def forward(self, pixel_values: Any) -> Any: + pooled = self.vision_model(pixel_values=pixel_values).pooler_output + return nn.functional.normalize(self.visual_projection(pooled), dim=-1) + + +def export_onnx(checkpoint: Path, output: Path) -> None: + """Write the static batch-one graph consumed by the ONNX runtime backend.""" + from remove_ai_watermarks._internal.clip_l_ft import SIZE, load_headed_clip + + checkpoint = checkpoint.expanduser().resolve() + output = output.expanduser().resolve() + if not checkpoint.is_file(): + raise SystemExit(f"missing checkpoint: {checkpoint}") + if output.exists(): + raise SystemExit(f"refusing to overwrite existing output: {output}") + output.parent.mkdir(parents=True, exist_ok=True) + model = VisionEmbedding(load_headed_clip(checkpoint, torch.device("cpu"))).eval() + export_options: dict[str, Any] = {} + if "dynamo" in inspect.signature(torch.onnx.export).parameters: + export_options["dynamo"] = False + with torch.inference_mode(): + torch.onnx.export( + model, + (torch.zeros(1, 3, SIZE, SIZE, dtype=torch.float32),), + output, + input_names=["pixel_values"], + output_names=["embedding"], + opset_version=17, + do_constant_folding=True, + **export_options, + ) + with output.open("rb") as handle: + digest = hashlib.file_digest(handle, "sha256").hexdigest() + log.info("wrote %s bytes=%d sha256=%s", output, output.stat().st_size, digest) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + logging.basicConfig(level=logging.INFO, format="%(message)s") + export_onnx(args.checkpoint, args.out) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/publish_photo_classify_hf.py b/scripts/publish_photo_classify_hf.py index 1f6f323..cf17e08 100644 --- a/scripts/publish_photo_classify_hf.py +++ b/scripts/publish_photo_classify_hf.py @@ -2,8 +2,8 @@ """Publish the photo-classify freeze to Hugging Face wiltodelta/raiw-photo-classify. Manual. Not a library release step. Weights are not in git: pass --src, a -directory that holds clip-l-ft.pt and the probe, with detector.pt and -provider.pt either at the root or in run1/. +directory that holds clip-l-ft.pt, clip-l-ft-vision-fp32.onnx, and the probe, +with detector.pt and provider.pt either at the root or in run1/. uv run python scripts/publish_photo_classify_hf.py --src DIR uv run python scripts/publish_photo_classify_hf.py --card-only @@ -28,10 +28,11 @@ log = logging.getLogger(__name__) HUB_REPO = "wiltodelta/raiw-photo-classify" CLIP_FILE = "clip-l-ft.pt" +CLIP_ONNX_FILE = "clip-l-ft-vision-fp32.onnx" PROBE_FILE = "probe-weights-clip-l-ft.npz" DETECTOR_FILE = "detector.pt" PROVIDER_FILE = "provider.pt" -WEIGHT_FILES = (CLIP_FILE, PROBE_FILE, DETECTOR_FILE, PROVIDER_FILE) +WEIGHT_FILES = (CLIP_FILE, CLIP_ONNX_FILE, PROBE_FILE, DETECTOR_FILE, PROVIDER_FILE) GATE_FILE = "receipt-gate-2026-09-02.npz" GATE_STABLE_FILE = "receipt-gate.npz" REPO_ROOT = Path(__file__).resolve().parents[1] @@ -68,7 +69,7 @@ def stage_card(dest: Path) -> None: shutil.copy2(readme, dest / "README.md") shutil.copy2(operating, dest / "operating-point.json") # The gate head versions with the model: card mode carries it even when - # the four freeze weight files are not being re-uploaded. + # the five freeze weight files are not being re-uploaded. shutil.copy2(GATE_SOURCE, dest / GATE_FILE) diff --git a/src/remove_ai_watermarks/_internal/clip_l_ft.py b/src/remove_ai_watermarks/_internal/clip_l_ft.py index 4f0582a..7cdf750 100644 --- a/src/remove_ai_watermarks/_internal/clip_l_ft.py +++ b/src/remove_ai_watermarks/_internal/clip_l_ft.py @@ -79,6 +79,18 @@ def load_processor() -> Any: return processor +def load_onnx_vision(model: Path) -> Any: + """Load the portable CPU session for the vision-only FP32 export.""" + import onnxruntime as ort + + options = ort.SessionOptions() + # Runtime graph optimization rebuilds this 1.2 GB static graph on every cold + # process. The export is already constant-folded; disabling that repeated pass + # reduced local session creation without changing the graph or its outputs. + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL + return ort.InferenceSession(str(model), sess_options=options, providers=["CPUExecutionProvider"]) + + def embed_image(model: HeadedCLIP, processor: Any, image: Image.Image, device: torch.device) -> NDArray[Any]: """Return a 768-d L2-normalized CLIP-L-ft vector.""" import numpy as np @@ -91,3 +103,14 @@ def embed_image(model: HeadedCLIP, processor: Any, image: Image.Image, device: t with torch.inference_mode(): vector = model.embed(pixels).float().cpu().numpy()[0] return np.asarray(vector, dtype=np.float64) + + +def embed_image_onnx(session: Any, processor: Any, image: Image.Image) -> NDArray[Any]: + """Return the vision-only ONNX export's normalized 768-d vector.""" + import numpy as np + + boxed = letterbox(image) + pixels = processor(images=[boxed], return_tensors="np")["pixel_values"] + inputs = {"pixel_values": np.asarray(pixels, dtype=np.float32)} + vector = session.run(["embedding"], inputs)[0][0] + return np.asarray(vector, dtype=np.float64) diff --git a/src/remove_ai_watermarks/classify.py b/src/remove_ai_watermarks/classify.py index fc07bb6..4da08ea 100644 --- a/src/remove_ai_watermarks/classify.py +++ b/src/remove_ai_watermarks/classify.py @@ -19,15 +19,18 @@ from typing import Any, Literal log = logging.getLogger(__name__) CLASSIFY_EXTRA = "'remove-ai-watermarks[classify]'" +CLASSIFY_ONNX_EXTRA = "'remove-ai-watermarks[classify-onnx]'" WEIGHTS_ENV = "RAIW_CLASSIFY_WEIGHTS" WEIGHTS_REPO = "wiltodelta/raiw-photo-classify" WEIGHTS_REVISION = "4c2763766dd1c8c212d64e01e1cc3f166243c47d" CLIP_FILE = "clip-l-ft.pt" +CLIP_ONNX_FILE = "clip-l-ft-vision-fp32.onnx" PROBE_FILE = "probe-weights-clip-l-ft.npz" DETECTOR_FILE = "detector.pt" PROVIDER_FILE = "provider.pt" OPERATING_POINT_FILE = "operating-point.json" -_WEIGHT_FILES = (CLIP_FILE, PROBE_FILE, DETECTOR_FILE, PROVIDER_FILE) +_COMMON_WEIGHT_FILES = (PROBE_FILE, DETECTOR_FILE, PROVIDER_FILE) +_WEIGHT_FILES = (CLIP_FILE, *_COMMON_WEIGHT_FILES) # 2026-09-02 receipt-document gate: linear head on the same CLIP-L-ft # vector, trained on CORD-v2 train (800, CC BY 4.0) plus 200 synthetic @@ -55,6 +58,7 @@ RECEIPT_GATE_THRESHOLD = 2.0432573877459697 # drifting the next time a release adds one. WEIGHTS_ALLOW_PATTERNS: tuple[str, ...] = ( *_WEIGHT_FILES, + CLIP_ONNX_FILE, OPERATING_POINT_FILE, RECEIPT_GATE_STABLE_FILE, RECEIPT_GATE_FILE, @@ -87,6 +91,7 @@ DetectorLevel = Literal["definitely", "possibly", "likely_human"] PixelLabel = Literal["ai", "human", "unknown"] PixelDomain = Literal["photo"] PixelProvider = Literal["openai", "google", "muse-image", "bytedance"] +ClassifierBackend = Literal["torch", "onnx"] PUBLIC_PROVIDER: dict[str, PixelProvider | None] = { "openai": "openai", "google": "google", @@ -255,7 +260,12 @@ def classify_from_scores( return PixelClassification(label=label, domain="photo", detector=level, provider=provider) -def classify_pixels(path: Path, *, device: str | None = None) -> PixelClassification: +def classify_pixels( + path: Path, + *, + device: str | None = None, + backend: ClassifierBackend = "torch", +) -> PixelClassification: """Run Model 1 then, on DEFINITELY, the receipt gate and Model 2. Never called by ``identify``. A DEFINITELY verdict that the receipt gate @@ -263,7 +273,14 @@ def classify_pixels(path: Path, *, device: str | None = None) -> PixelClassifica """ if not is_available(): raise RuntimeError(f"pixel classification requires the classify extra. Install: pip install {CLASSIFY_EXTRA}") - runtime = _get_runtime(device) + if backend == "onnx": + from remove_ai_watermarks.optional_deps import module_available + + if not module_available("onnxruntime"): + raise RuntimeError( + f"the ONNX photo-classifier backend requires ONNX Runtime. Install: pip install {CLASSIFY_ONNX_EXTRA}" + ) + runtime = _get_runtime(device, backend) from PIL import Image with Image.open(path) as image: @@ -304,6 +321,7 @@ def classify_pixels(path: Path, *, device: str | None = None) -> PixelClassifica @dataclass class _Runtime: device: Any + backend: ClassifierBackend clip_model: Any processor: Any detector: Any @@ -331,7 +349,26 @@ def _resolve_device(device: str | None) -> Any: raise ValueError(f"unsupported classify device {device!r}") -def _weights_dir() -> Path: +def _resolve_backend(backend: ClassifierBackend, device: Any) -> ClassifierBackend: + if backend == "onnx" and device.type != "cpu": + raise ValueError("the ONNX photo-classifier backend supports CPU only") + if backend not in ("torch", "onnx"): + raise ValueError(f"unsupported classify backend {backend!r}") + return backend + + +def _backend_allow_patterns(backend: ClassifierBackend) -> tuple[str, ...]: + model_file = CLIP_ONNX_FILE if backend == "onnx" else CLIP_FILE + return ( + model_file, + *_COMMON_WEIGHT_FILES, + OPERATING_POINT_FILE, + RECEIPT_GATE_STABLE_FILE, + RECEIPT_GATE_FILE, + ) + + +def _weights_dir(backend: ClassifierBackend = "torch") -> Path: override = os.environ.get(WEIGHTS_ENV) if override: path = Path(override) @@ -345,7 +382,7 @@ def _weights_dir() -> Path: snapshot_download( repo_id=WEIGHTS_REPO, revision=WEIGHTS_REVISION, - allow_patterns=list(WEIGHTS_ALLOW_PATTERNS), + allow_patterns=list(_backend_allow_patterns(backend)), ) ) except Exception as exc: @@ -373,8 +410,10 @@ def _weight_path(folder: Path, name: str) -> Path: return found -def _require_files(folder: Path) -> None: - missing = [name for name in _WEIGHT_FILES if _find_weight(folder, name) is None] +def _require_files(folder: Path, *, backend: ClassifierBackend = "torch") -> None: + model_file = CLIP_ONNX_FILE if backend == "onnx" else CLIP_FILE + required = (model_file, *_COMMON_WEIGHT_FILES) + missing = [name for name in required if _find_weight(folder, name) is None] if missing: raise RuntimeError( "pixel classification weights are missing " @@ -417,14 +456,14 @@ def _build_provider_head() -> Any: ) -def _load_runtime(device: Any) -> _Runtime: +def _load_runtime(device: Any, backend: ClassifierBackend = "torch") -> _Runtime: import numpy as np import torch - from remove_ai_watermarks._internal.clip_l_ft import load_headed_clip, load_processor + from remove_ai_watermarks._internal.clip_l_ft import load_headed_clip, load_onnx_vision, load_processor - folder = _weights_dir() - _require_files(folder) + folder = _weights_dir(backend) + _require_files(folder, backend=backend) probe = np.load(_weight_path(folder, PROBE_FILE)) detector = _build_detector() detector.load_state_dict(torch.load(_weight_path(folder, DETECTOR_FILE), map_location="cpu", weights_only=True)) @@ -441,9 +480,15 @@ def _load_runtime(device: Any) -> _Runtime: ridge_threshold = float(probe["thr_oi_1pct"]) mlp_threshold, provider_margin = _operating_point(folder) log.info("loaded photo-classify freeze from %s", folder) + clip_model = ( + load_onnx_vision(_weight_path(folder, CLIP_ONNX_FILE)) + if backend == "onnx" + else load_headed_clip(_weight_path(folder, CLIP_FILE), device) + ) return _Runtime( device=device, - clip_model=load_headed_clip(_weight_path(folder, CLIP_FILE), device), + backend=backend, + clip_model=clip_model, processor=load_processor(), detector=detector, provider_heads=heads, @@ -457,18 +502,21 @@ def _load_runtime(device: Any) -> _Runtime: ) -def _get_runtime(device: str | None) -> _Runtime: +def _get_runtime(device: str | None, backend: ClassifierBackend = "torch") -> _Runtime: global _runtime resolved = _resolve_device(device) + resolved_backend = _resolve_backend(backend, resolved) with _runtime_lock: - if _runtime is None or _runtime.device != resolved: - _runtime = _load_runtime(resolved) + if _runtime is None or _runtime.device != resolved or _runtime.backend != resolved_backend: + _runtime = _load_runtime(resolved, resolved_backend) return _runtime def _embed(runtime: _Runtime, image: Any) -> Any: - from remove_ai_watermarks._internal.clip_l_ft import embed_image + from remove_ai_watermarks._internal.clip_l_ft import embed_image, embed_image_onnx + if runtime.backend == "onnx": + return embed_image_onnx(runtime.clip_model, runtime.processor, image) return embed_image(runtime.clip_model, runtime.processor, image, runtime.device) diff --git a/tests/test_classify.py b/tests/test_classify.py index 7778696..af484c2 100644 --- a/tests/test_classify.py +++ b/tests/test_classify.py @@ -13,6 +13,7 @@ from click.testing import CliRunner from PIL import Image from remove_ai_watermarks.classify import ( + CLIP_ONNX_FILE, MLP_THRESHOLD, PROVIDER_MARGIN, RECEIPT_GATE_FILE, @@ -57,13 +58,11 @@ def test_hub_snapshot_is_the_freeze_revision() -> None: assert WEIGHTS_REVISION == "4c2763766dd1c8c212d64e01e1cc3f166243c47d" -def test_runtime_requests_the_exported_allow_patterns(monkeypatch: pytest.MonkeyPatch) -> None: - """The exported pre-cache list and the runtime request must stay one list. +def test_runtime_requests_only_the_selected_backend_model(monkeypatch: pytest.MonkeyPatch) -> None: + """Online downloads stay narrow while the exported pre-cache covers both. An offline deploy pre-caches WEIGHTS_ALLOW_PATTERNS from the Hub; the - runtime then resolves the same snapshot offline. If the runtime's request - grows a file the export misses, every offline deployment breaks on a - four-fifths-present cache, so pin the seam in both directions. + runtime requests a subset of that snapshot for its selected backend. """ from remove_ai_watermarks.classify import _WEIGHT_FILES, _weights_dir @@ -81,8 +80,65 @@ def test_runtime_requests_the_exported_allow_patterns(monkeypatch: pytest.Monkey assert _weights_dir() == Path("/cached/snapshot") assert captured[0]["repo_id"] == WEIGHTS_REPO assert captured[0]["revision"] == WEIGHTS_REVISION - assert captured[0]["allow_patterns"] == list(WEIGHTS_ALLOW_PATTERNS) + torch_patterns = captured[0]["allow_patterns"] + assert "clip-l-ft.pt" in torch_patterns + assert CLIP_ONNX_FILE not in torch_patterns assert set(_WEIGHT_FILES) <= set(WEIGHTS_ALLOW_PATTERNS) + assert CLIP_ONNX_FILE in WEIGHTS_ALLOW_PATTERNS + + captured.clear() + assert _weights_dir("onnx") == Path("/cached/snapshot") + onnx_patterns = captured[0]["allow_patterns"] + assert CLIP_ONNX_FILE in onnx_patterns + assert "clip-l-ft.pt" not in onnx_patterns + assert set(torch_patterns) | set(onnx_patterns) == set(WEIGHTS_ALLOW_PATTERNS) + + +def test_backend_requires_its_own_model_file(tmp_path: Path) -> None: + from remove_ai_watermarks.classify import DETECTOR_FILE, PROBE_FILE, PROVIDER_FILE, _require_files + + for name in (CLIP_ONNX_FILE, PROBE_FILE, DETECTOR_FILE, PROVIDER_FILE): + (tmp_path / name).touch() + + _require_files(tmp_path, backend="onnx") + with pytest.raises(RuntimeError, match=r"clip-l-ft\.pt"): + _require_files(tmp_path, backend="torch") + + +def test_onnx_backend_is_cpu_only() -> None: + import torch + + from remove_ai_watermarks.classify import _resolve_backend + + assert _resolve_backend("onnx", torch.device("cpu")) == "onnx" + with pytest.raises(ValueError, match=r"ONNX.*CPU"): + _resolve_backend("onnx", torch.device("cuda")) + + +def test_onnx_backend_names_its_install_extra(monkeypatch: pytest.MonkeyPatch) -> None: + from remove_ai_watermarks import classify as classify_mod + from remove_ai_watermarks import optional_deps + + monkeypatch.setattr(classify_mod, "is_available", lambda: True) + monkeypatch.setattr(optional_deps, "module_available", lambda *_names: False) + + with pytest.raises(RuntimeError, match=r"remove-ai-watermarks\[classify-onnx\]"): + classify_mod.classify_pixels(Path("unused.png"), backend="onnx") + + +def test_embed_dispatches_to_the_selected_backend(monkeypatch: pytest.MonkeyPatch) -> None: + from remove_ai_watermarks import classify as classify_mod + from remove_ai_watermarks._internal import clip_l_ft + + calls: list[str] = [] + monkeypatch.setattr(clip_l_ft, "embed_image", lambda *_args: calls.append("torch") or "torch-vector") + monkeypatch.setattr(clip_l_ft, "embed_image_onnx", lambda *_args: calls.append("onnx") or "onnx-vector") + runtime = SimpleNamespace(backend="onnx", clip_model=object(), processor=object(), device=object()) + + assert classify_mod._embed(runtime, object()) == "onnx-vector" + runtime.backend = "torch" + assert classify_mod._embed(runtime, object()) == "torch-vector" + assert calls == ["onnx", "torch"] def test_shipped_operating_point_matches_the_runtime_defaults() -> None: @@ -332,7 +388,7 @@ def _patch_definitely_runtime(monkeypatch: pytest.MonkeyPatch, *, forensic_calls weights_dir=None, ) monkeypatch.setattr("remove_ai_watermarks.classify.is_available", lambda: True) - monkeypatch.setattr("remove_ai_watermarks.classify._get_runtime", lambda device: runtime) + monkeypatch.setattr("remove_ai_watermarks.classify._get_runtime", lambda device, backend="torch": runtime) monkeypatch.setattr("remove_ai_watermarks.classify._embed", lambda *args, **kwargs: np.full(768, 100.0)) monkeypatch.setattr("remove_ai_watermarks.classify._mlp_score", lambda *args, **kwargs: 10.0) monkeypatch.setattr( @@ -402,7 +458,7 @@ def test_124d_is_not_extracted_unless_definitely(monkeypatch: pytest.MonkeyPatch weights_dir=None, ) monkeypatch.setattr("remove_ai_watermarks.classify.is_available", lambda: True) - monkeypatch.setattr("remove_ai_watermarks.classify._get_runtime", lambda device: runtime) + monkeypatch.setattr("remove_ai_watermarks.classify._get_runtime", lambda device, backend="torch": runtime) monkeypatch.setattr("remove_ai_watermarks.classify._embed", lambda *args, **kwargs: np.zeros(768)) monkeypatch.setattr("remove_ai_watermarks.classify._mlp_score", lambda *args, **kwargs: 0.0) monkeypatch.setattr( diff --git a/tests/test_clip_l_ft.py b/tests/test_clip_l_ft.py index e6f0b02..e518236 100644 --- a/tests/test_clip_l_ft.py +++ b/tests/test_clip_l_ft.py @@ -2,9 +2,13 @@ from __future__ import annotations +import sys +from types import SimpleNamespace from typing import TYPE_CHECKING, Any +import numpy as np import pytest +from PIL import Image if TYPE_CHECKING: from pathlib import Path @@ -67,3 +71,66 @@ def test_load_headed_clip_skips_random_initialization_and_materializes_checkpoin ): assert name == expected_name assert torch.equal(buffer, expected_buffer) + + +def test_load_onnx_vision_uses_the_portable_cpu_provider( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from remove_ai_watermarks._internal import clip_l_ft + + captured: dict[str, Any] = {} + + class FakeOptions: + graph_optimization_level: Any = None + + class FakeSession: + def __init__(self, path: str, *, sess_options: FakeOptions, providers: list[str]) -> None: + captured.update(path=path, optimization=sess_options.graph_optimization_level, providers=providers) + + monkeypatch.setitem( + sys.modules, + "onnxruntime", + SimpleNamespace( + GraphOptimizationLevel=SimpleNamespace(ORT_DISABLE_ALL="disabled"), + InferenceSession=FakeSession, + SessionOptions=FakeOptions, + ), + ) + model = tmp_path / "vision.onnx" + + assert isinstance(clip_l_ft.load_onnx_vision(model), FakeSession) + assert captured == { + "path": str(model), + "optimization": "disabled", + "providers": ["CPUExecutionProvider"], + } + + +def test_embed_image_onnx_uses_the_freeze_preprocessor() -> None: + from remove_ai_watermarks._internal import clip_l_ft + + captured: dict[str, Any] = {} + + class FakeProcessor: + def __call__(self, *, images: list[Image.Image], return_tensors: str) -> dict[str, np.ndarray]: + captured["image_size"] = images[0].size + captured["return_tensors"] = return_tensors + return {"pixel_values": np.ones((1, 3, 224, 224), dtype=np.float64)} + + class FakeSession: + def run(self, outputs: list[str], inputs: dict[str, np.ndarray]) -> list[np.ndarray]: + captured["outputs"] = outputs + captured["input_dtype"] = inputs["pixel_values"].dtype + return [np.array([[3.0, 4.0]], dtype=np.float32)] + + vector = clip_l_ft.embed_image_onnx(FakeSession(), FakeProcessor(), Image.new("RGB", (40, 20))) + + assert captured == { + "image_size": (224, 224), + "return_tensors": "np", + "outputs": ["embedding"], + "input_dtype": np.dtype(np.float32), + } + assert vector.dtype == np.float64 + assert vector.tolist() == [3.0, 4.0] diff --git a/tests/test_export_photo_classify_onnx.py b/tests/test_export_photo_classify_onnx.py new file mode 100644 index 0000000..0202052 --- /dev/null +++ b/tests/test_export_photo_classify_onnx.py @@ -0,0 +1,40 @@ +"""Vision-only photo-classifier export helpers, without the real checkpoint.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest + + +def _exporter() -> Any: + path = Path(__file__).resolve().parents[1] / "scripts" / "export_photo_classify_onnx.py" + spec = importlib.util.spec_from_file_location("export_photo_classify_onnx", path) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_vision_embedding_drops_the_text_tower_and_normalizes() -> None: + torch = pytest.importorskip("torch") + + class FakeVision(torch.nn.Module): + def forward(self, *, pixel_values: Any) -> Any: + return SimpleNamespace(pooler_output=pixel_values) + + clip = SimpleNamespace( + vision_model=FakeVision(), + visual_projection=torch.nn.Identity(), + text_model=object(), + ) + model = _exporter().VisionEmbedding(SimpleNamespace(clip=clip)) + + output = model(torch.tensor([[3.0, 4.0]])) + + assert set(dict(model.named_children())) == {"vision_model", "visual_projection"} + assert torch.equal(output, torch.tensor([[0.6, 0.8]])) diff --git a/tests/test_packaging.py b/tests/test_packaging.py index 03655f6..43c485d 100644 --- a/tests/test_packaging.py +++ b/tests/test_packaging.py @@ -76,7 +76,7 @@ def test_trustmark_is_limited_to_its_numpy_compatible_python_range(): def test_extras_use_capability_names_without_legacy_aliases(): extras = set(metadata("remove-ai-watermarks").get_all("Provides-Extra") or []) - assert {"pixels", "heif", "visible", "video", "detect", "diffusion", "classify"} <= extras + assert {"pixels", "heif", "visible", "video", "detect", "diffusion", "classify", "classify-onnx"} <= extras assert {"gpu", "remove", "detect-pywavelets", "verify"}.isdisjoint(extras) @@ -85,6 +85,11 @@ def test_classify_extra_owns_the_photo_heads() -> None: assert {"torch", "transformers", "huggingface-hub", "tokenizers"} <= names +def test_classify_onnx_extra_extends_the_default_classifier() -> None: + names = _requirement_names("classify-onnx") + assert {"torch", "transformers", "huggingface-hub", "tokenizers", "onnxruntime"} <= names + + def test_qwen_zimage_extra_owns_every_profile_tokenizer_backend() -> None: assert "sentencepiece" in _requirement_names("qwen-zimage") diff --git a/tests/test_publish_photo_classify_hf.py b/tests/test_publish_photo_classify_hf.py index ad1eaad..a0b9d44 100644 --- a/tests/test_publish_photo_classify_hf.py +++ b/tests/test_publish_photo_classify_hf.py @@ -23,8 +23,19 @@ def test_hub_id_is_the_photo_classify_repo() -> None: assert _publish().HUB_REPO == "wiltodelta/raiw-photo-classify" +def test_publish_freeze_includes_both_clip_runtimes() -> None: + publish = _publish() + from remove_ai_watermarks.classify import _WEIGHT_FILES, CLIP_ONNX_FILE, WEIGHTS_ALLOW_PATTERNS + + assert publish.CLIP_FILE in publish.WEIGHT_FILES + assert publish.CLIP_ONNX_FILE in publish.WEIGHT_FILES + assert set(publish.WEIGHT_FILES) == {*_WEIGHT_FILES, CLIP_ONNX_FILE} + assert set(publish.WEIGHT_FILES) <= set(WEIGHTS_ALLOW_PATTERNS) + + def test_weight_paths_accepts_run1_layout(tmp_path: Path) -> None: (tmp_path / "clip-l-ft.pt").write_bytes(b"clip") + (tmp_path / "clip-l-ft-vision-fp32.onnx").write_bytes(b"onnx") (tmp_path / "probe-weights-clip-l-ft.npz").write_bytes(b"probe") run1 = tmp_path / "run1" run1.mkdir() @@ -37,6 +48,7 @@ def test_weight_paths_accepts_run1_layout(tmp_path: Path) -> None: def test_weight_paths_requires_every_file(tmp_path: Path) -> None: (tmp_path / "clip-l-ft.pt").write_bytes(b"clip") + (tmp_path / "clip-l-ft-vision-fp32.onnx").write_bytes(b"onnx") with pytest.raises(SystemExit, match=r"missing probe-weights-clip-l-ft\.npz"): _publish().weight_paths(tmp_path) diff --git a/uv.lock b/uv.lock index 312f214..718d71e 100644 --- a/uv.lock +++ b/uv.lock @@ -800,7 +800,7 @@ name = "cuda-bindings" version = "13.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-pathfinder", marker = "sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "cuda-pathfinder" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/51/6b/457ca12dad3ee9bfcc9a545cfd6b64b359ba49de40f776f6e028e678f262/cuda_bindings-13.3.1-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c5879712accf6e14bb01aa5e67440eb84998b8d104b509cc7a6dc0b8f656a474", size = 6053539, upload-time = "2026-05-29T23:11:43.19Z" }, @@ -833,43 +833,43 @@ wheels = [ [package.optional-dependencies] cublas = [ - { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cublas", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cudart = [ - { name = "nvidia-cuda-runtime", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cuda-runtime", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cufft = [ - { name = "nvidia-cufft", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cufft", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cufile = [ - { name = "nvidia-cufile", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cufile", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cupti = [ - { name = "nvidia-cuda-cupti", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cuda-cupti", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] curand = [ - { name = "nvidia-curand", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-curand", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cusolver = [ - { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-cusolver", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cublas", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cusolver", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cusparse = [ - { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] nvjitlink = [ - { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] nvrtc = [ - { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] nvtx = [ - { name = "nvidia-nvtx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-nvtx", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] [[package]] @@ -1066,7 +1066,7 @@ wheels = [ [package.optional-dependencies] http = [ - { name = "aiohttp", marker = "python_full_version < '3.13'" }, + { name = "aiohttp" }, ] [[package]] @@ -1198,8 +1198,8 @@ name = "httpcore2" version = "2.12.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "h11", marker = "sys_platform != 'emscripten'" }, - { name = "truststore", marker = "sys_platform != 'emscripten'" }, + { name = "h11" }, + { name = "truststore" }, ] sdist = { url = "https://files.pythonhosted.org/packages/be/ad/f4f0e57345f1870f3e8cb624e058d7eca6e5a27d33bcc3311d9b618734cd/httpcore2-2.12.0.tar.gz", hash = "sha256:9293522bba0aa7c4c8e9e3f040c16575bd8868e155a77fa30c7a9085a5eae648", size = 67548, upload-time = "2026-08-18T13:22:08.211Z" } wheels = [ @@ -1458,22 +1458,22 @@ wheels = [ [[package]] name = "lightning" -version = "2.6.5" +version = "2.6.6" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "fsspec", extra = ["http"], marker = "python_full_version < '3.13'" }, - { name = "lightning-utilities", marker = "python_full_version < '3.13'" }, - { name = "packaging", marker = "python_full_version < '3.13'" }, - { name = "pytorch-lightning", marker = "python_full_version < '3.13'" }, - { name = "pyyaml", marker = "python_full_version < '3.13'" }, - { name = "torch", marker = "python_full_version < '3.13'" }, - { name = "torchmetrics", marker = "python_full_version < '3.13'" }, - { name = "tqdm", marker = "python_full_version < '3.13'" }, - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "fsspec", extra = ["http"] }, + { name = "lightning-utilities" }, + { name = "packaging" }, + { name = "pytorch-lightning" }, + { name = "pyyaml" }, + { name = "torch" }, + { name = "torchmetrics" }, + { name = "tqdm" }, + { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/c9/1d/83be8536bec71a0173e762a9a1fd92a24a5ad0d0f74c59550c3c4e6c103b/lightning-2.6.5.tar.gz", hash = "sha256:16a30310ed69afde3748491feb5d13508908effd70390d2bfc203dc0812a4b4a", size = 659201, upload-time = "2026-05-27T14:33:41.806Z" } +sdist = { url = "https://files.pythonhosted.org/packages/39/b8/fbf1469d917c3ed3b22e00ab3a852fb9a5f85b9dc7f7060eedc3449098e7/lightning-2.6.6.tar.gz", hash = "sha256:1c6329e241f8a85f29fedead89c1ee138b36caa08de60c260c6a7eaee81c8295", size = 659902, upload-time = "2026-09-10T09:41:03.009Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e7/c5/fca7144236b6fa3279d0fb3172b32576c5ad8b84a63b9432ad6592d24847/lightning-2.6.5-py3-none-any.whl", hash = "sha256:3702fb7ef4ab51a8c3d4a140f4674514fe72973a9673dfa05e07a078e2767389", size = 848611, upload-time = "2026-05-27T14:33:39.714Z" }, + { url = "https://files.pythonhosted.org/packages/12/43/7da7fe42f254ce57777d1a86eece3f86d1b3329fc568d48e2bc16c8b7ebc/lightning-2.6.6-py3-none-any.whl", hash = "sha256:2bcbd6ee840071cd076c7dc9953d7ac1040df92f1c5fa76583aae644f037ab85", size = 849265, upload-time = "2026-09-10T09:41:00.882Z" }, ] [[package]] @@ -1481,8 +1481,8 @@ name = "lightning-utilities" version = "0.15.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "packaging", marker = "python_full_version < '3.13'" }, - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "packaging" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/f1/45/7fa8f56b17dc0f0a41ec70dd307ecd6787254483549843bef4c30ab5adce/lightning_utilities-0.15.3.tar.gz", hash = "sha256:792ae0204c79f6859721ac7f386c237a33b0ed06ba775009cb894e010a842033", size = 33553, upload-time = "2026-02-22T14:48:53.348Z" } wheels = [ @@ -1896,7 +1896,7 @@ name = "nvidia-cublas" version = "13.1.1.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cuda-nvrtc", marker = "sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-cuda-nvrtc" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/a7/a1/0bd24ee8c8d03adac032fd2909426a00c88f8c57961b1277ded97f91119f/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b7a210458267ac818974c53038fbec2e969d5c99f305ab15c72522fa9f001dd5", size = 542848918, upload-time = "2026-04-08T18:46:22.985Z" }, @@ -1935,7 +1935,7 @@ name = "nvidia-cudnn-cu13" version = "9.24.0.43" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas", marker = "sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-cublas" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/ca/30/7c257e3d5cb4fecb147b93895c66e29c93f8e76d74b45bb418ff0587c4ec/nvidia_cudnn_cu13-9.24.0.43-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:a6812a554a1ff0413e9c52b84c26c050380649ab9615f9c16bded368ce9f421f", size = 650976863, upload-time = "2026-07-02T16:23:39.248Z" }, @@ -1947,7 +1947,7 @@ name = "nvidia-cufft" version = "12.0.0.61" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink", marker = "sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-nvjitlink" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" }, @@ -1977,9 +1977,9 @@ name = "nvidia-cusolver" version = "12.0.4.66" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas", marker = "sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "nvidia-cusparse", marker = "sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "nvidia-nvjitlink", marker = "sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-cublas" }, + { name = "nvidia-cusparse" }, + { name = "nvidia-nvjitlink" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" }, @@ -1991,7 +1991,7 @@ name = "nvidia-cusparse" version = "12.6.3.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink", marker = "sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-nvjitlink" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" }, @@ -2146,7 +2146,7 @@ resolution-markers = [ "(python_full_version < '3.12' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32')", ] dependencies = [ - { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" }, + { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" } }, ] sdist = { url = "https://files.pythonhosted.org/packages/17/06/68c27a523103dad5837dc5b87e71285280c4f098c60e4fe8a8db6486ab09/opencv-python-4.11.0.86.tar.gz", hash = "sha256:03d60ccae62304860d232272e4a4fda93c39d595780cb40b161b310244b736a4", size = 95171956, upload-time = "2025-01-16T13:52:24.737Z" } wheels = [ @@ -2175,7 +2175,7 @@ resolution-markers = [ "(python_full_version == '3.13.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32')", ] dependencies = [ - { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" }, + { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" } }, ] sdist = { url = "https://files.pythonhosted.org/packages/ed/e8/c2959c0f969ea4ae96b55ad10c108c0410de4e8b3550d9941a4cfdd52703/opencv_python-4.14.0.94.tar.gz", hash = "sha256:afadc76a38a07713a4c2b4702da917cf0ed489be57e2a40a34e670189731de7d", size = 96400225, upload-time = "2026-07-28T19:23:09.546Z" } wheels = [ @@ -2206,7 +2206,7 @@ resolution-markers = [ "(python_full_version < '3.12' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32')", ] dependencies = [ - { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" }, + { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" } }, ] sdist = { url = "https://files.pythonhosted.org/packages/36/2f/5b2b3ba52c864848885ba988f24b7f105052f68da9ab0e693cc7c25b0b30/opencv-python-headless-4.11.0.86.tar.gz", hash = "sha256:996eb282ca4b43ec6a3972414de0e2331f5d9cda2b41091a49739c19fb843798", size = 95177929, upload-time = "2025-01-16T13:53:40.22Z" } wheels = [ @@ -2235,7 +2235,7 @@ resolution-markers = [ "(python_full_version == '3.13.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32')", ] dependencies = [ - { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" }, + { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" } }, ] sdist = { url = "https://files.pythonhosted.org/packages/67/8a/51ff1c6287be0b6cd3070cf095875da4ff7d6eda38507e6ada5608786a78/opencv_python_headless-4.14.0.94.tar.gz", hash = "sha256:4afa2ea1214453648be88259f035712454faa9039b686de7753569ba8eec1577", size = 96405664, upload-time = "2026-07-28T19:23:48.758Z" } wheels = [ @@ -3176,21 +3176,21 @@ wheels = [ [[package]] name = "pytorch-lightning" -version = "2.6.5" +version = "2.6.6" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "fsspec", extra = ["http"], marker = "python_full_version < '3.13'" }, - { name = "lightning-utilities", marker = "python_full_version < '3.13'" }, - { name = "packaging", marker = "python_full_version < '3.13'" }, - { name = "pyyaml", marker = "python_full_version < '3.13'" }, - { name = "torch", marker = "python_full_version < '3.13'" }, - { name = "torchmetrics", marker = "python_full_version < '3.13'" }, - { name = "tqdm", marker = "python_full_version < '3.13'" }, - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "fsspec", extra = ["http"] }, + { name = "lightning-utilities" }, + { name = "packaging" }, + { name = "pyyaml" }, + { name = "torch" }, + { name = "torchmetrics" }, + { name = "tqdm" }, + { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/52/2c/8e73a3929b4c4bd600cafd38a97aaf7242a8cf518fb9f33d27c274ec898f/pytorch_lightning-2.6.5.tar.gz", hash = "sha256:1c32cefa76a1a9c4c5250338272d961d1e48b180e68396849efe128538ddb28e", size = 661673, upload-time = "2026-05-27T14:33:41.961Z" } +sdist = { url = "https://files.pythonhosted.org/packages/97/e4/5025cb42311b4f8cc6d1b12cd81399ca433159f8e3831ab229c5a6c5177e/pytorch_lightning-2.6.6.tar.gz", hash = "sha256:52c52b20522dd3d3feb3bb3aae2dbe3f291e8e836d4111a74cfd72e9df018834", size = 662507, upload-time = "2026-09-10T09:41:16.684Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/8b/4d/5740c27110b83634d8491c3b5facf0111b3e554c3164f4fb953be9bddaf6/pytorch_lightning-2.6.5-py3-none-any.whl", hash = "sha256:62d9c8549b2278fedc3364f0a5607a56c6063d18635008f8cf3fae8d802b0d76", size = 852407, upload-time = "2026-05-27T14:33:39.856Z" }, + { url = "https://files.pythonhosted.org/packages/4f/cf/ba8e4cc1b982ff6dc6f9fdef0e4074605438b82640f6270888a10e7e42c3/pytorch_lightning-2.6.6-py3-none-any.whl", hash = "sha256:71f95c7b22f25c4f91cc659e1734bcdc6c69c8b66543f2ed2d78dbe29ce2f167", size = 853045, upload-time = "2026-09-10T09:41:14.923Z" }, ] [[package]] @@ -3205,7 +3205,7 @@ resolution-markers = [ "(python_full_version < '3.12' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32')", ] dependencies = [ - { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" } }, ] sdist = { url = "https://files.pythonhosted.org/packages/5a/75/50581633d199812205ea8cdd0f6d52f12a624886b74bf1486335b67f01ff/pywavelets-1.9.0.tar.gz", hash = "sha256:148d12203377772bea452a59211d98649c8ee4a05eff019a9021853a36babdc8", size = 3938340, upload-time = "2025-08-04T16:20:04.978Z" } wheels = [ @@ -3281,7 +3281,7 @@ resolution-markers = [ "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32')", ] dependencies = [ - { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.12.*'" }, + { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" }, { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/25/84/a9e80e33eb114e78bd2f1829a90d72d8d4d660109c721a8aa2f0d58d7592/pywavelets-1.10.0.tar.gz", hash = "sha256:f3cb8640225a8f3fdcbc3f6f3bd112c78927c9407cbfa65a5ff303990a6a3e32", size = 3948163, upload-time = "2026-09-07T17:15:15.637Z" } @@ -3492,6 +3492,17 @@ classify = [ { name = "torch" }, { name = "transformers" }, ] +classify-onnx = [ + { name = "huggingface-hub" }, + { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" }, + { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" }, + { name = "onnxruntime" }, + { name = "opencv-python-headless", version = "4.11.0.86", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" }, + { name = "opencv-python-headless", version = "4.14.0.94", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" }, + { name = "tokenizers" }, + { name = "torch" }, + { name = "transformers" }, +] detect = [ { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" }, { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" }, @@ -3630,6 +3641,7 @@ requires-dist = [ { name = "invisible-watermark", marker = "extra == 'dev'", specifier = ">=0.2.0" }, { name = "numpy", marker = "python_full_version >= '3.13' and extra == 'pixels'", specifier = ">=2" }, { name = "numpy", marker = "python_full_version < '3.13' and extra == 'pixels'", specifier = ">=1.24.0" }, + { name = "onnxruntime", marker = "extra == 'classify-onnx'", specifier = ">=1.24.0" }, { name = "onnxruntime", marker = "extra == 'lama'", specifier = ">=1.24.0" }, { name = "onnxruntime", marker = "extra == 'migan'", specifier = ">=1.24.0" }, { name = "openai", marker = "extra == 'dev'", specifier = ">=2.52.0" }, @@ -3647,6 +3659,7 @@ requires-dist = [ { name = "pytest-xdist", marker = "extra == 'dev'", specifier = ">=3.5.0" }, { name = "python-dotenv", specifier = ">=1.0.0" }, { name = "pywavelets", marker = "extra == 'detect'", specifier = ">=1.1.1" }, + { name = "remove-ai-watermarks", extras = ["classify"], marker = "extra == 'classify-onnx'" }, { name = "remove-ai-watermarks", extras = ["detect"], marker = "extra == 'dev'" }, { name = "remove-ai-watermarks", extras = ["diffusion"], marker = "extra == 'qwen-zimage'" }, { name = "remove-ai-watermarks", extras = ["pixels"], marker = "extra == 'classify'" }, @@ -3671,7 +3684,7 @@ requires-dist = [ { name = "transformers", marker = "extra == 'diffusion'", specifier = ">=5,<6" }, { name = "trustmark", marker = "python_full_version < '3.13' and extra == 'trustmark'", specifier = ">=0.8.0" }, ] -provides-extras = ["pixels", "heif", "visible", "video", "detect", "diffusion", "qwen-zimage", "text-restoration", "text-draft", "trustmark", "classify", "lama", "migan", "dev", "all"] +provides-extras = ["pixels", "heif", "visible", "video", "detect", "diffusion", "qwen-zimage", "text-restoration", "text-draft", "trustmark", "classify", "classify-onnx", "lama", "migan", "dev", "all"] [[package]] name = "requests" @@ -4054,10 +4067,10 @@ name = "torchmetrics" version = "1.9.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "lightning-utilities", marker = "python_full_version < '3.13'" }, - { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" }, - { name = "packaging", marker = "python_full_version < '3.13'" }, - { name = "torch", marker = "python_full_version < '3.13'" }, + { name = "lightning-utilities" }, + { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" } }, + { name = "packaging" }, + { name = "torch" }, ] sdist = { url = "https://files.pythonhosted.org/packages/81/34/39b8b749333db56c0585d7a11fa62a283c087bb1dfc897d69fb8cedbefb1/torchmetrics-1.9.0.tar.gz", hash = "sha256:a488609948600df52d3db4fcdab02e62aab2a85ef34da67037dc3e65b8512faa", size = 581765, upload-time = "2026-03-09T17:41:22.443Z" } wheels = [ @@ -4152,13 +4165,13 @@ name = "trustmark" version = "0.9.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "einops", marker = "python_full_version < '3.13'" }, - { name = "lightning", marker = "python_full_version < '3.13'" }, - { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" }, - { name = "omegaconf", marker = "python_full_version < '3.13'" }, - { name = "six", marker = "python_full_version < '3.13'" }, - { name = "torch", marker = "python_full_version < '3.13'" }, - { name = "torchvision", marker = "python_full_version < '3.13'" }, + { name = "einops" }, + { name = "lightning" }, + { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" } }, + { name = "omegaconf" }, + { name = "six" }, + { name = "torch" }, + { name = "torchvision" }, ] sdist = { url = "https://files.pythonhosted.org/packages/53/41/32e84596af94e8696c55060d55eb9eed9d99722a673d12212a1c0c4ef1d2/trustmark-0.9.2.tar.gz", hash = "sha256:bcaa3bc38ead1c0558df81932672f34e306efaed31389da57dbed100e236b28a", size = 63780, upload-time = "2026-09-08T10:58:11.601Z" }