diff --git a/Dockerfile.markdiffusion b/Dockerfile.markdiffusion new file mode 100644 index 0000000..ddde6d3 --- /dev/null +++ b/Dockerfile.markdiffusion @@ -0,0 +1,48 @@ +# Optional local Docker image for the MarkDiffusion image-watermark harness. +# +# Build from the repository root: +# docker build -f Dockerfile.markdiffusion -t watermarks-remover-markdiffusion . +# +# The upstream package is installed from PyPI at build time (pinned in +# requirements-markdiffusion.txt) and is NOT redistributed by this repository. +# Upstream is Apache-2.0. +# +# Vendored fork hardening: +# - base image pinned by digest (no moving tag drift) +# - package version pinned exactly in requirements-markdiffusion.txt +# - pip itself pinned (no unpinned bootstrap step) +# - runs as an unprivileged user (a parser bug in a crafted image can no +# longer write files as root inside the container) + +# python:3.14-slim linux/amd64 digest. +FROM python:3.14-slim@sha256:ce40764625a4ff50df3548277632e7f96c4e77fe75fa848aae9885476e7df5a4 + +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + libgl1 \ + libglib2.0-0 \ + passwd \ + && rm -rf /var/lib/apt/lists/* + +COPY skills/remove-ai-marks/scripts/requirements-markdiffusion.txt /app/requirements-markdiffusion.txt +COPY skills/remove-ai-marks/scripts/markdiffusion_harness.py /app/markdiffusion_harness.py + +# torch is pinned (with everything else) in requirements-markdiffusion.txt; the +# CPU index keeps the image small. No GPU wheel index inside the image — CUDA +# users should run setup_markdiffusion.sh on the host instead. +RUN python3 -m pip install --no-cache-dir "pip==26.2.1" \ + && python3 -m pip install --no-cache-dir --index-url https://download.pytorch.org/whl/cpu "torch>=2.4,<2.11" \ + && python3 -m pip install --no-cache-dir -r /app/requirements-markdiffusion.txt + +# Unprivileged runtime user. The harness only reads input files and writes to +# stdout/tmp, so nothing under /opt, /app, or the mounted data dir needs root. +RUN useradd --create-home --uid 10001 --shell /usr/sbin/nologin markdiffusion +USER markdiffusion + +ENV HOME=/home/markdiffusion \ + PYTHONUNBUFFERED=1 \ + PYTHONDONTWRITEBYTECODE=1 \ + HF_HOME=/home/markdiffusion/.cache/huggingface + +WORKDIR /app +ENTRYPOINT ["python3", "/app/markdiffusion_harness.py"] diff --git a/Makefile b/Makefile index 274333a..656ebc5 100644 --- a/Makefile +++ b/Makefile @@ -1,6 +1,8 @@ .PHONY: test smoke smoke-synthid bootstrap-synthid docker-synthid-build docker-synthid-help \ smoke-ctrlregen bootstrap-ctrlregen docker-ctrlregen-build docker-ctrlregen-help \ - smoke-markllm bootstrap-markllm docker-markllm-build docker-markllm-help install-skill clean + smoke-markllm bootstrap-markllm docker-markllm-build docker-markllm-help \ + smoke-markdiffusion bootstrap-markdiffusion docker-markdiffusion-build docker-markdiffusion-help \ + install-skill clean SCRIPTS := skills/remove-ai-marks/scripts PYTHON ?= $(shell if [ -x .venv/bin/python ]; then echo .venv/bin/python; else echo python3; fi) @@ -66,6 +68,22 @@ docker-markllm-build: docker-markllm-help: docker run --rm watermarks-remover-markllm --help +smoke-markdiffusion: + @if [ -z "$(MARKDIFFUSION_DIR)" ]; then \ + echo "smoke-markdiffusion skipped (set MARKDIFFUSION_DIR)"; \ + else \ + $(PYTHON) $(SCRIPTS)/markdiffusion_harness.py --help >/dev/null && echo "markdiffusion_harness adapter present"; \ + fi + +bootstrap-markdiffusion: + ./skills/remove-ai-marks/scripts/setup_markdiffusion.sh + +docker-markdiffusion-build: + docker build -f Dockerfile.markdiffusion -t watermarks-remover-markdiffusion . + +docker-markdiffusion-help: + docker run --rm watermarks-remover-markdiffusion --help + install-skill: mkdir -p $(HOME)/.grok/skills ln -sfn $(CURDIR)/skills/remove-ai-marks $(HOME)/.grok/skills/remove-ai-marks diff --git a/README.md b/README.md index d493be2..c9e41d9 100644 --- a/README.md +++ b/README.md @@ -307,6 +307,83 @@ docker run --rm --user "$(id -u):$(id -g)" -v "$(pwd):/data" \ watermarks-remover-markllm detect /data/wm.txt --scheme kgw --json ``` +## Optional MarkDiffusion image-watermark harness + +For **controlled experiments on images**, an optional external harness wraps +[`THU-BPM/MarkDiffusion`](https://github.com/THU-BPM/MarkDiffusion) (Apache-2.0), +a *generative watermarking* toolkit for latent diffusion models (it embeds marks +— it does not remove them). We use it for three things: + +1. **Verification harness** (like MarkLLM, but for images): watermark a test + image with a scheme, run removal, and re-detect with the *same* scheme config + — e.g. prove a Tree-Ring-class mark clears under your pipeline. It is a + **verification harness, not an oracle**: detection requires the generating + model (and keys for key-based schemes), so it cannot certify a vendor + detector will fail on an arbitrary image. +2. **Optional pixel-removal engine**: its `DiffusionPurification` regeneration + attack is exposed as `clean_image.py --remove-pixel diffusion`, an + alternative to CtrlRegen. It is **blind** regeneration (no ControlNet + conditioning), so it drifts image content more than CtrlRegen — conservative + strength default (`0.3`), treated as a fallback/comparison, never a + guarantee. +3. **Local same-scheme detector** for Tree-Ring-class marks, partially filling + the "no local detector for StegaStamp/Tree-Ring/StableSignature" gap (it + covers Tree-Ring/Ring-ID/Gaussian-Shading etc., not StegaStamp / + StableSignature / SynthID-media). + +The backend is **not bundled**. `setup_markdiffusion.sh` creates a venv and +installs `markdiffusion==1.0.2` from PyPI (pinned), with torch installed from +the right platform index; `--checkout` installs an editable clone at a pinned +commit instead. The Stable Diffusion model (default +`huanzi05/stable-diffusion-2-1-base`) downloads from Hugging Face on first run. + +```bash +SCRIPTS=skills/remove-ai-marks/scripts + +# Bootstrap (PyPI pin default; creates ~/markdiffusion/.venv, installs deps). +"$SCRIPTS/setup_markdiffusion.sh" + +# 1. Generate a Tree-Ring watermarked image (+ unwatermarked control). +echo "a red fox in snow" > /tmp/prompt.txt +MARKDIFFUSION_DIR=~/markdiffusion \ + ~/markdiffusion/.venv/bin/python "$SCRIPTS/markdiffusion_harness.py" watermark \ + /tmp/prompt.txt -o wm.png -o2 plain.png --scheme tr --json + +# 2. Remove with the DiffusionPurification regeneration attack. +MARKDIFFUSION_DIR=~/markdiffusion \ + ~/markdiffusion/.venv/bin/python "$SCRIPTS/markdiffusion_harness.py" purify \ + wm.png -o wm.purified.png --purification-strength 0.3 --json + +# 3. Re-detect with the SAME scheme config. +MARKDIFFUSION_DIR=~/markdiffusion \ + ~/markdiffusion/.venv/bin/python "$SCRIPTS/markdiffusion_harness.py" detect \ + wm.purified.png --scheme tr --detector-type l1_distance --json +``` + +Or run purification as part of the normal image pipeline: + +```bash +MARKDIFFUSION_DIR=~/markdiffusion \ + ~/markdiffusion/.venv/bin/python "$SCRIPTS/clean_image.py" shot.png \ + -o shot.cleaned.png --remove-pixel diffusion +``` + +Hardening knobs mirror the MarkLLM harness: `--offline` loads the model from +the Hugging Face cache only (zero network egress, no remote code), `HF_TOKEN` +is env-only (never argv), algorithm configs are capped at 1 MiB, and the +subprocess gets the same higher resource caps as CtrlRegen. + +### Docker + +```bash +make docker-markdiffusion-build +docker run --rm --user "$(id -u):$(id -g)" -v "$(pwd):/data" \ + watermarks-remover-markdiffusion detect /data/wm.png --scheme tr --json +``` + +The image installs a CPU torch; CUDA users should run `setup_markdiffusion.sh` +on the host instead. Model downloads still hit the HF hub on first run. + ## Coverage matrix | Channel | Claude | Gemini/SynthID | OpenAI | Open-LLM | @@ -314,7 +391,7 @@ docker run --rm --user "$(id -u):$(id -g)" -v "$(pwd):/data" \ | Unicode / edit-based text | Layer A | Layer A | Layer A | Layer A | | Statistical sampling text | Layer B best-effort | Layer B best-effort | Layer B if present | Layer B best-effort | | C2PA / file metadata | Yes (listed formats) | Yes when present | Yes when present | Yes when present | -| Pixel image marks | Out of scope | Optional SynthID score + CtrlRegen removal (external) | Out of scope | Optional CtrlRegen removal (external) | +| Pixel image marks | Out of scope | Optional SynthID score + CtrlRegen removal (external); optional MarkDiffusion same-scheme detect + DiffusionPurification removal (external) | Out of scope | Optional CtrlRegen / MarkDiffusion removal (external) | | Training backdoors | Out of scope | Out of scope | Out of scope | Out of scope | Details: [`skills/remove-ai-marks/references/vendor-notes.md`](skills/remove-ai-marks/references/vendor-notes.md), [`mark-classes.md`](skills/remove-ai-marks/references/mark-classes.md). @@ -412,6 +489,7 @@ Industry two-layer context (C2PA + imperceptible watermark): [Institute of AI PM | Rewrite (Layer B) | Statistical token marks (best-effort) | Always offered by skill; costs style — see [Disclaimer](#disclaimer-what-removing-a-text-watermark-costs) | | Container/metadata strip | File provenance | See format table | | CtrlRegen pixel removal (optional) | Pixel-domain image marks (SynthID-class, StegaStamp, Tree-Ring, StableSignature) | External backend; heavy compute; conservative strength default | +| DiffusionPurification pixel removal (optional) | Pixel-domain image marks (Tree-Ring-class) | MarkDiffusion backend; blind regeneration (more drift than CtrlRegen); conservative strength default | | Open-weight local models | Avoid re-stamping with origin model | Operational alternative | Matrix: [`skills/remove-ai-marks/references/removal-matrix.md`](skills/remove-ai-marks/references/removal-matrix.md). @@ -434,6 +512,11 @@ make smoke # quick CLI smoke on fixtures ### Unreleased +- New optional MarkDiffusion image-watermark harness (external `THU-BPM/MarkDiffusion`, Apache-2.0): `markdiffusion_harness.py` with `watermark` / `detect` / `purify` subcommands for nine image schemes (Tree-Ring, Ring-ID, ROBIN, WIND, SFW, Gaussian-Shading, GaussMarker, PRC, SEAL) +- `clean_image.py --remove-pixel diffusion` runs the MarkDiffusion `DiffusionPurification` regeneration attack as an alternative pixel-removal engine (conservative strength 0.3 default) +- `setup_markdiffusion.sh` bootstrap (PyPI pin `1.0.2`; `--checkout` editable clone at pinned commit) + `requirements-markdiffusion.txt` + `Dockerfile.markdiffusion` and Makefile `bootstrap-markdiffusion` / `smoke-markdiffusion` / `docker-markdiffusion-build` / `docker-markdiffusion-help` +- Mock-based tests (`tests/test_markdiffusion_harness.py`) — no torch in CI; `references/markdiffusion.md` reference doc +- Docs: same-scheme-only verification caveat (not a vendor-detector oracle) and blind-regeneration drift caveat in README, SKILL.md, `removal-matrix.md`, `markdiffusion.md` - Add stdlib-only WebP inspection and metadata cleaning for RIFF `C2PA`, XMP, EXIF, and ICC profile chunks - New optional MarkLLM harness (external `THU-BPM/MarkLLM` checkout, Apache-2.0): `detect_text_watermark.py` with `detect` / `watermark` subcommands for KGW and SynthID schemes - `rewrite_text.py --markllm-scheme` runs before/after detection around a Layer B rewrite (env-gated; reports `cleared`) @@ -541,6 +624,7 @@ MIT — see [LICENSE](LICENSE). - [C2PA](https://c2pa.org/) / [c2patool](https://github.com/contentauth/c2pa-rs/tree/main/cli) - Kirchenbauer et al., [*A Watermark for Large Language Models*](https://arxiv.org/abs/2301.10226) - [THU-BPM/MarkLLM](https://github.com/THU-BPM/MarkLLM) (unified toolkit for evaluating LLM watermarking algorithms) +- Pan et al., [*MarkDiffusion: An Open-Source Toolkit for Generative Watermarking of Latent Diffusion Models*](https://arxiv.org/abs/2509.10569) (JMLR) — the embedding toolkit this repo's optional image-watermark harness wraps — [code](https://github.com/THU-BPM/MarkDiffusion), [docs](https://markdiffusion.readthedocs.io) - Zhang et al., [*Watermarks in the Sand: Impossibility of Strong Watermarking for Generative Models*](https://arxiv.org/abs/2311.04378) (ICML 2024) - [google-deepmind/synthid-text](https://github.com/google-deepmind/synthid-text) (research reference; not used for detection here) - [aloshdenny/reverse-SynthID](https://github.com/aloshdenny/reverse-SynthID) (research reference) diff --git a/skills/remove-ai-marks/SKILL.md b/skills/remove-ai-marks/SKILL.md index 6e2607c..09db387 100644 --- a/skills/remove-ai-marks/SKILL.md +++ b/skills/remove-ai-marks/SKILL.md @@ -21,6 +21,7 @@ Read if needed: - `references/removal-matrix.md` — which layer when - `references/ethics.md` — intended use - `references/how-claude-marks.md` — Anthropic-specific detail +- `references/markdiffusion.md` — optional MarkDiffusion image harness (schemes, honesty caveats) Scripts live in this skill’s `scripts/` directory. Resolve `SCRIPTS` to that folder (absolute path of this skill + `/scripts`). @@ -34,6 +35,8 @@ python3 "$SCRIPTS/inspect_image.py" ... python3 "$SCRIPTS/clean_image.py" ... python3 "$SCRIPTS/clean_ctrlregen.py" ... # optional external pixel removal (bootstrap first) "$SCRIPTS/setup_ctrlregen.sh" # one-command bootstrap (Windows: setup_ctrlregen.ps1) +python3 "$SCRIPTS/markdiffusion_harness.py" ... # optional MarkDiffusion image harness (bootstrap first) +"$SCRIPTS/setup_markdiffusion.sh" # one-command bootstrap (PyPI pin; --checkout for editable) python3 "$SCRIPTS/rewrite_text.py" ... python3 "$SCRIPTS/detect_text_watermark.py" ... # optional external MarkLLM verification harness (bootstrap first) "$SCRIPTS/setup_markllm.sh" # one-command bootstrap for the above @@ -80,7 +83,16 @@ local image with `make docker-synthid-build`. For pixel-domain **removal**, bootstrap the CtrlRegen backend with `scripts/setup_ctrlregen.sh` (or `make docker-ctrlregen-build`), then use `clean_image.py --remove-pixel ctrlregen`. See the README "Optional CtrlRegen -pixel removal" section for strength presets and the 512×512 size handling. +pixel removal" section for strength presets and the 512×512 size handling. An +alternative engine, MarkDiffusion's `DiffusionPurification`, is available as +`clean_image.py --remove-pixel diffusion` (bootstrap with +`scripts/setup_markdiffusion.sh`); it is **blind** regeneration and drifts more +than CtrlRegen, so it is a fallback/comparison path, not the default. + +For controlled before/after experiments on images, the MarkDiffusion harness +(`markdiffusion_harness.py watermark` → removal → `detect`) proves a specific +Tree-Ring-class scheme config clears. **Same-scheme/same-model only** — it is +not a vendor-detector oracle. ### Aggregate audits and confidence @@ -252,6 +264,7 @@ Always state: - PDF strip is best-effort without `exiftool`, and incomplete without `qpdf`: exiftool alone leaves the freed metadata objects in the byte stream. - Pixel-domain **image** watermarks can be removed optionally via the external CtrlRegen backend (`clean_image.py --remove-pixel ctrlregen`); audio/video watermarks remain out of scope for removal. - The CtrlRegen backend is external, all-rights-reserved (no LICENSE file), never bundled, heavy (~10 GB model downloads), and a regenerating remover — no local detector certifies StegaStamp/Tree-Ring/StableSignature removal. +- The MarkDiffusion backend (Apache-2.0, PyPI-pinned) adds a same-scheme detector for Tree-Ring-class marks and a blind-regeneration remover (`--remove-pixel diffusion`); detection is same-scheme/same-model-only and it is not a vendor-detector oracle. - The reverse-SynthID scorer is external, best-effort, and under a non-commercial Research License; it is not bundled and is not an official Google detector. - **C2PA soft binding** (content watermark that re-links to a remote manifest after metadata strip) is out of scope — stripping hard-bound C2PA does not clear it. - Data-driven / backdoor model marks (trigger phrases) are out of scope. @@ -281,6 +294,18 @@ NOAI_WATERMARK_DIR=~/noai-watermark \ ~/noai-watermark/.venv/bin/python scripts/clean_image.py shot.png \ -o shot.cleaned.png --remove-pixel ctrlregen +# Optional MarkDiffusion image harness (watermark -> purify -> detect) +scripts/setup_markdiffusion.sh +MARKDIFFUSION_DIR=~/markdiffusion \ + ~/markdiffusion/.venv/bin/python scripts/markdiffusion_harness.py \ + watermark prompt.txt -o wm.png -o2 plain.png --scheme tr --json +MARKDIFFUSION_DIR=~/markdiffusion \ + ~/markdiffusion/.venv/bin/python scripts/markdiffusion_harness.py \ + purify wm.png -o wm.purified.png --purification-strength 0.3 --json +MARKDIFFUSION_DIR=~/markdiffusion \ + ~/markdiffusion/.venv/bin/python scripts/markdiffusion_harness.py \ + detect wm.purified.png --scheme tr --detector-type l1_distance --json + # Optional MarkLLM verification (external backend; bootstrap first) scripts/setup_markllm.sh MARKLLM_DIR=~/MarkLLM \ diff --git a/skills/remove-ai-marks/references/markdiffusion.md b/skills/remove-ai-marks/references/markdiffusion.md new file mode 100644 index 0000000..06c0793 --- /dev/null +++ b/skills/remove-ai-marks/references/markdiffusion.md @@ -0,0 +1,80 @@ +# MarkDiffusion (THU-BPM) — reference + +External research backend: [`THU-BPM/MarkDiffusion`](https://github.com/THU-BPM/MarkDiffusion) +(JMLR; Apache-2.0). A **generative watermarking** toolkit for latent diffusion +models — it *embeds* marks, it does not remove them. This repo uses it as a +controlled-experiment harness and as an optional regeneration-removal engine. + +## What it covers (image-only scope) + +Nine image algorithms in two categories: + +| Category | Algorithms | Notes | +| --- | --- | --- | +| Pattern-based | Tree-Ring, Ring-ID, ROBIN, WIND, SFW | A fixed latent/FT pattern is injected at generation and inverted for detection | +| Key-based | Gaussian-Shading, GaussMarker, PRC, SEAL | A secret key modulates the noise/latent; detection needs the key | + +(Video algorithms VideoShield / VideoMark are out of scope for this project.) + +## What it gives the remover + +- **Same-scheme detector** for the above algorithms via `AutoWatermark.load(...).detect_watermark_in_media()`. Covers the Tree-Ring-class gap in + `removal-matrix.md`; it does **not** cover StegaStamp, StableSignature, or + SynthID-media. +- **`DiffusionPurification`** — a blind regeneration attack (DiffPure-style: + encode → partial noise → reverse-denoise) usable as a pixel-watermark remover. + Exposed as `--remove-pixel diffusion` in `clean_image.py`. +- **`NeuralCodecCompression`** — codec round-trip (compressai). Not wired here. + +## Hard honesty constraints + +1. **Detection is same-scheme and same-model only.** Inversion-based detection + requires the generating model (and for key-based schemes the key). It proves + "this image came from *this* model/config/params" — it cannot certify that a + vendor detector will fail on an arbitrary image. This is the same + same-config-only caveat as the MarkLLM text harness. +2. **`DiffusionPurification` is blind regeneration.** It reuses the *same* + pipeline, so it will not defeat a watermark that is robust to its own + regeneration path, and it drifts image content (more than CtrlRegen's + controllable ControlNet regeneration). Conservative strength default (0.3), + treated as a fallback/comparison engine, never a guarantee. +3. **Heavy stack.** torch ≥ 2.4,<2.11 + diffusers + a Stable Diffusion model + download (~4–10 GB). GPU strongly recommended. Some HF models are gated → + `HF_TOKEN` (env only, never argv). + +## License / hygiene + +Apache-2.0. Installed from PyPI at a pinned version +(`requirements-markdiffusion.txt`) or an editable checkout at a pinned commit +(`setup_markdiffusion.sh --checkout`). Never bundled into this repo. + +## Harness usage + +```bash +SCRIPTS=skills/remove-ai-marks/scripts +MD="$HOME/markdiffusion/.venv/bin/python" + +"$SCRIPTS/setup_markdiffusion.sh" # PyPI pin default +# or: "$SCRIPTS/setup_markdiffusion.sh" --checkout # editable pinned clone + +# 1. watermark a test image with a scheme +echo "a red fox in snow" > /tmp/prompt.txt +"$MD" "$SCRIPTS/markdiffusion_harness.py" watermark /tmp/prompt.txt \ + -o /tmp/wm.png -o2 /tmp/plain.png --scheme tr --json + +# 2. remove (blind regeneration) +"$MD" "$SCRIPTS/markdiffusion_harness.py" purify /tmp/wm.png \ + -o /tmp/wm.purified.png --purification-strength 0.3 --json + +# 3. re-detect with the SAME scheme config +"$MD" "$SCRIPTS/markdiffusion_harness.py" detect /tmp/wm.purified.png \ + --scheme tr --detector-type l1_distance --json +``` + +Exit codes: 0 ok · 1 runtime error · 2 bad input · 3 backend unavailable. + +## References + +- Paper: https://arxiv.org/abs/2509.10569 +- Docs: https://markdiffusion.readthedocs.io +- HF models: https://huggingface.co/Generative-Watermark-Toolkits diff --git a/skills/remove-ai-marks/references/removal-matrix.md b/skills/remove-ai-marks/references/removal-matrix.md index cce3c5b..f9c675b 100644 --- a/skills/remove-ai-marks/references/removal-matrix.md +++ b/skills/remove-ai-marks/references/removal-matrix.md @@ -11,7 +11,8 @@ | ODT meta:generator | Scrub `meta.xml` | `clean_file.py` | Loses generator tag | Yes | | HTML generator / JSON-LD provenance | Strip tags | `clean_file.py` | Loses meta | Yes | | Markdown AI frontmatter keys | Drop keys | `clean_file.py` | Loses YAML keys | Yes | -| Pixel image watermark (SynthID-media / StegaStamp / Tree-Ring / StableSignature) | CtrlRegen regeneration (external backend) | `clean_ctrlregen.py` / `clean_image.py --remove-pixel ctrlregen` | Regenerates pixels; heavy compute; detail drift at higher strength | No without official detector; reverse-SynthID score is a local surrogate | +| Pixel image watermark (SynthID-media / StegaStamp / Tree-Ring / StableSignature) | CtrlRegen regeneration (external backend) | `clean_ctrlregen.py` / `clean_image.py --remove-pixel ctrlregen` | Regenerates pixels; heavy compute; detail drift at higher strength | No without official detector; reverse-SynthID score is a local surrogate; **MarkDiffusion same-scheme harness** (`markdiffusion_harness.py detect`) verifies a Tree-Ring-class scheme config before/after | +| Pixel image watermark (Tree-Ring-class) | DiffusionPurification regeneration (external MarkDiffusion backend) | `clean_image.py --remove-pixel diffusion` | Blind regeneration; more drift than CtrlRegen; heavy compute | Same-scheme only via the MarkDiffusion harness (not a vendor-detector oracle) | | Audio / video watermarks (SynthID-media) | — | Out of scope | — | — | | C2PA soft binding (in-content link to manifest) | — | Out of scope (survives our metadata strip) | — | Vendor detector only | | Data-driven model backdoors | — | Out of scope | — | — | diff --git a/skills/remove-ai-marks/references/vendor-notes.md b/skills/remove-ai-marks/references/vendor-notes.md index 2c51820..479c755 100644 --- a/skills/remove-ai-marks/references/vendor-notes.md +++ b/skills/remove-ai-marks/references/vendor-notes.md @@ -35,7 +35,7 @@ Source: [How Claude marks AI-generated content](https://support.claude.com/en/ar - Optional external verification harness: [`THU-BPM/MarkLLM`](https://github.com/THU-BPM/MarkLLM) (Apache-2.0) reimplements SynthID-Text among other schemes with configurable keys; wired as `detect_text_watermark.py` / `rewrite_text.py --markllm-scheme`. Same-config-only — it verifies a mark you generated under a known config, not Google's production keying. - Current frontier production watermarks are **token-by-token** (streaming constraint); paragraph-level robust methods (SemStamp / PostMark) are not deployed yet, which keeps paraphrase-class attacks effective today. - Optional external reference: [`aloshdenny/reverse-SynthID`](https://github.com/aloshdenny/reverse-SynthID) provides a reverse-engineered pixel-domain scorer. It is **not bundled** here, is best-effort, and is under a non-commercial Research License; it is not the official Google detector. -- Optional pixel-domain removal: [`mertizci/noai-watermark`](https://github.com/mertizci/noai-watermark)'s CtrlRegen profile is wired through `clean_image.py --remove-pixel ctrlregen` / `clean_ctrlregen.py`. It is **not bundled** (no LICENSE file → all-rights-reserved), and no local detector certifies the result; the official Google check is the final authority. +- Optional pixel-domain removal: [`mertizci/noai-watermark`](https://github.com/mertizci/noai-watermark)'s CtrlRegen profile is wired through `clean_image.py --remove-pixel ctrlregen` / `clean_ctrlregen.py`. It is **not bundled** (no LICENSE file → all-rights-reserved), and no local detector certifies the result; the official Google check is the final authority. For Tree-Ring-class marks, the optional MarkDiffusion harness (`markdiffusion_harness.py`, Apache-2.0) adds a same-scheme detector and a blind-regeneration removal engine (`--remove-pixel diffusion`) — see `references/markdiffusion.md`. **Skill mapping:** same Layer B rewrite attacks (paraphrase / back-translate / structural) used in the literature against sampling watermarks. diff --git a/skills/remove-ai-marks/scripts/clean_image.py b/skills/remove-ai-marks/scripts/clean_image.py index 10ff91a..5f818e6 100755 --- a/skills/remove-ai-marks/scripts/clean_image.py +++ b/skills/remove-ai-marks/scripts/clean_image.py @@ -37,9 +37,11 @@ def main() -> int: ) p.add_argument( "--remove-pixel", - choices=["ctrlregen"], + choices=["ctrlregen", "diffusion"], default=None, - help="Run optional CtrlRegen pixel-watermark removal after metadata cleaning", + help="Run optional pixel-watermark removal after metadata cleaning " + "(ctrlregen = CtrlRegen regeneration; diffusion = MarkDiffusion " + "DiffusionPurification regeneration)", ) p.add_argument( "--ctrlregen-dir", @@ -77,6 +79,48 @@ def main() -> int: default=3600, help="CtrlRegen subprocess timeout in seconds (default: 3600)", ) + p.add_argument( + "--markdiffusion-dir", + type=str, + default=None, + help="MarkDiffusion bootstrap dir (default: $MARKDIFFUSION_DIR)", + ) + p.add_argument( + "--markdiffusion-strength", + type=float, + default=0.3, + help="DiffusionPurification strength in (0, 1] (default: 0.3)", + ) + p.add_argument( + "--markdiffusion-model", + type=str, + default=None, + help="Stable Diffusion model for purification (default: SD 2.1 base)", + ) + p.add_argument( + "--markdiffusion-size", + type=int, + default=512, + help="Purification working size in px (default: 512)", + ) + p.add_argument( + "--markdiffusion-steps", + type=int, + default=50, + help="Purification diffusion steps (default: 50)", + ) + p.add_argument( + "--markdiffusion-device", + type=str, + default=None, + help="Purification device: auto|cpu|cuda|mps (default: auto)", + ) + p.add_argument( + "--markdiffusion-timeout", + type=int, + default=3600, + help="Purification subprocess timeout in seconds (default: 3600)", + ) args = p.parse_args() if not args.path.is_file(): @@ -104,6 +148,13 @@ def main() -> int: ctrlregen_device=args.ctrlregen_device, ctrlregen_seed=args.ctrlregen_seed, ctrlregen_timeout=args.ctrlregen_timeout, + markdiffusion_dir=args.markdiffusion_dir, + markdiffusion_strength=args.markdiffusion_strength, + markdiffusion_model=args.markdiffusion_model, + markdiffusion_size=args.markdiffusion_size, + markdiffusion_steps=args.markdiffusion_steps, + markdiffusion_device=args.markdiffusion_device, + markdiffusion_timeout=args.markdiffusion_timeout, ) except Exception as e: eprint(f"error: {e}") @@ -135,9 +186,11 @@ def main() -> int: ) if pr is not None: if pr.get("available"): - eprint(f"CtrlRegen: removed on {pr.get('device', 'unknown device')}") + engine = "CtrlRegen" if args.remove_pixel == "ctrlregen" else "DiffusionPurification" + eprint(f"{engine}: removed on {pr.get('device', 'unknown device')}") else: - eprint(f"CtrlRegen: unavailable: {pr.get('error', 'unknown error')}") + engine = "CtrlRegen" if args.remove_pixel == "ctrlregen" else "DiffusionPurification" + eprint(f"{engine}: unavailable: {pr.get('error', 'unknown error')}") if residual: eprint("warning: residual C2PA/AI signals may remain") for f in result.get("post_findings") or []: diff --git a/skills/remove-ai-marks/scripts/image_meta.py b/skills/remove-ai-marks/scripts/image_meta.py index 6a7a586..cbefcb9 100755 --- a/skills/remove-ai-marks/scripts/image_meta.py +++ b/skills/remove-ai-marks/scripts/image_meta.py @@ -398,6 +398,100 @@ def _ctrlregen_python(upstream: Path) -> str: return sys.executable +def _markdiffusion_python(upstream: Path | None) -> str: + """Prefer the bootstrap venv so torch/diffusers/markdiffusion importable.""" + if upstream is not None: + if os.name == "nt": + venv = upstream / ".venv" / "Scripts" / "python.exe" + else: + venv = upstream / ".venv" / "bin" / "python" + if venv.is_file(): + return str(venv) + return sys.executable + + +def run_markdiffusion_purify( + path: Path, + output: Path, + *, + upstream_dir: str | None = None, + strength: float = 0.3, + model: str | None = None, + size: int = 512, + steps: int = 50, + device: str | None = None, + timeout: int = 3600, +) -> dict[str, Any]: + """Run the optional MarkDiffusion DiffusionPurification remover in a subprocess. + + Returns ``{"available": False, "error": ...}`` when the backend is not + configured, its dependencies are missing, or it fails at runtime; a + successful run is ``{"available": True, ...}``. + """ + if upstream_dir is None: + upstream_dir = os.environ.get("MARKDIFFUSION_DIR") + + upstream = ( + Path(upstream_dir).expanduser().resolve() if upstream_dir else None + ) + if upstream is not None and not upstream.is_dir(): + return { + "available": False, + "error": f"MarkDiffusion dir not found: {upstream}", + } + + script = SCRIPTS_DIR / "markdiffusion_harness.py" + cmd = [ + _markdiffusion_python(upstream), + str(script), + "purify", + str(path), + "-o", + str(output), + "--purification-strength", + str(strength), + "--size", + str(size), + "--steps", + str(steps), + "--json", + ] + if upstream is not None: + cmd += ["--upstream-dir", str(upstream)] + if model: + cmd += ["--model", str(model)] + if device: + cmd += ["--device", str(device)] + + try: + r = subprocess.run( + cmd, + capture_output=True, + text=True, + timeout=timeout, + preexec_fn=ctrlregen_subprocess_preexec_fn, + ) + except subprocess.TimeoutExpired: + return { + "available": False, + "error": f"DiffusionPurification timed out after {timeout}s", + } + except Exception as e: + return {"available": False, "error": str(e)} + + if r.returncode != 0: + return {"available": False, "error": (r.stderr or "").strip()[:2000]} + try: + payload = json.loads(r.stdout or "{}") + except json.JSONDecodeError as e: + return { + "available": False, + "error": f"bad MarkDiffusion adapter JSON: {e}", + } + payload["available"] = True + return payload + + def run_ctrlregen_clean( path: Path, output: Path, @@ -700,6 +794,13 @@ def clean_image( ctrlregen_device: str | None = None, ctrlregen_seed: int | None = None, ctrlregen_timeout: int = 3600, + markdiffusion_dir: str | None = None, + markdiffusion_strength: float = 0.3, + markdiffusion_model: str | None = None, + markdiffusion_size: int = 512, + markdiffusion_steps: int = 50, + markdiffusion_device: str | None = None, + markdiffusion_timeout: int = 3600, ) -> dict[str, Any]: synthid_before = run_synthid_score(path, synthid_dir) data = path.read_bytes() @@ -737,25 +838,48 @@ def clean_image( pixel_removal: dict[str, Any] | None = None if remove_pixel: - if remove_pixel != "ctrlregen": - raise ValueError(f"unknown pixel remover: {remove_pixel}") - pixel_removal = run_ctrlregen_clean( - dest, - dest, - upstream_dir=ctrlregen_dir, - strength=ctrlregen_strength, - steps=ctrlregen_steps, - device=ctrlregen_device, - seed=ctrlregen_seed, - timeout=ctrlregen_timeout, - ) - if pixel_removal.get("available"): - actions.append(f"CtrlRegen pixel removal (strength {ctrlregen_strength})") - else: - actions.append( - "CtrlRegen pixel removal skipped: " - f"{pixel_removal.get('error', 'unknown error')}" + if remove_pixel == "ctrlregen": + pixel_removal = run_ctrlregen_clean( + dest, + dest, + upstream_dir=ctrlregen_dir, + strength=ctrlregen_strength, + steps=ctrlregen_steps, + device=ctrlregen_device, + seed=ctrlregen_seed, + timeout=ctrlregen_timeout, ) + if pixel_removal.get("available"): + actions.append(f"CtrlRegen pixel removal (strength {ctrlregen_strength})") + else: + actions.append( + "CtrlRegen pixel removal skipped: " + f"{pixel_removal.get('error', 'unknown error')}" + ) + elif remove_pixel == "diffusion": + pixel_removal = run_markdiffusion_purify( + dest, + dest, + upstream_dir=markdiffusion_dir, + strength=markdiffusion_strength, + model=markdiffusion_model, + size=markdiffusion_size, + steps=markdiffusion_steps, + device=markdiffusion_device, + timeout=markdiffusion_timeout, + ) + if pixel_removal.get("available"): + actions.append( + f"DiffusionPurification pixel removal " + f"(strength {markdiffusion_strength})" + ) + else: + actions.append( + "DiffusionPurification pixel removal skipped: " + f"{pixel_removal.get('error', 'unknown error')}" + ) + else: + raise ValueError(f"unknown pixel remover: {remove_pixel}") after = inspect_image(dest, synthid_dir=synthid_dir) return { diff --git a/skills/remove-ai-marks/scripts/markdiffusion_harness.py b/skills/remove-ai-marks/scripts/markdiffusion_harness.py new file mode 100644 index 0000000..8d87896 --- /dev/null +++ b/skills/remove-ai-marks/scripts/markdiffusion_harness.py @@ -0,0 +1,524 @@ +#!/usr/bin/env python3 +"""Optional MarkDiffusion image-watermark harness backed by the external +THU-BPM/MarkDiffusion package (Apache-2.0). + +This script does NOT vendor upstream code. It imports ``markdiffusion`` either +from a user-provided checkout (--upstream-dir / $MARKDIFFUSION_DIR) or from the +environment's installed package (``pip install markdiffusion[optional]``). The +checkout/venv environment must supply torch + diffusers. + +MarkDiffusion is a *generative watermarking* toolkit (it embeds marks). We use +it as a research/verification harness for controlled experiments on images you +own: watermark a test image, run removal, and re-detect with the SAME scheme +config. Detection is only valid against the same scheme config, model, and keys +used at generation — it cannot certify that a vendor detector will fail on the +given image. + +Subcommands: + watermark generate a watermarked (and optionally unwatermarked) image from + a prompt, for controlled before/after experiments + detect run same-scheme detection on an image + purify run the DiffusionPurification regeneration attack on an image + (optional pixel-watermark removal engine) + +Exit codes: + 0 success + 1 runtime error (model load, detection/generation failure) + 2 bad input (missing/unreadable file, binary input, bad args) + 3 unavailable (not configured / missing package / missing deps) +""" + +from __future__ import annotations + +import argparse +import io +import json +import os +import sys +from pathlib import Path +from typing import Any + +SCRIPTS_DIR = Path(__file__).resolve().parent +sys.path.insert(0, str(SCRIPTS_DIR)) + +from common import eprint, emit_json, read_text_input, safe_write_bytes # noqa: E402 + +# User-facing scheme names -> MarkDiffusion algorithm names (image-only; the +# video algorithms VideoShield/VideoMark are out of scope here). Canonical +# upstream names are accepted as-is. +SCHEMES = { + "tr": "TR", + "ringid": "RI", + "robin": "ROBIN", + "wind": "WIND", + "sfw": "SFW", + "gaussianshading": "GS", + "gaussmarker": "GM", + "prc": "PRC", + "seal": "SEAL", +} + +IMAGE_SCHEMES = {"TR", "RI", "ROBIN", "WIND", "SFW", "GS", "GM", "PRC", "SEAL"} + +DEFAULT_MODEL = "huanzi05/stable-diffusion-2-1-base" + +# Algorithm configs are a few hundred bytes (TR.json/GS.json). Cap well above +# that so a crafted or accidental huge file is refused before either this script +# or upstream reads it into memory. +MAX_CONFIG_BYTES = 1 << 20 + + +class _Unavailable(RuntimeError): + """Backend present but unusable (missing package/checkout, missing deps).""" + + +def resolve_upstream(raw: str | None) -> Path | None: + if not raw: + return None + upstream = Path(raw).expanduser().resolve() + if not upstream.is_dir(): + return None + return upstream + + +def resolve_device(raw: str | None) -> str: + """Resolve the ``auto`` device hint to a concrete torch device.""" + if raw and raw != "auto": + return raw + try: + import torch + + if torch.cuda.is_available(): + return "cuda" + mps = getattr(torch.backends, "mps", None) + if mps is not None and mps.is_available(): + return "mps" + except Exception: + pass + return "cpu" + + +def normalize_scheme(raw: str) -> str: + up = raw.upper() + if up in IMAGE_SCHEMES: + return up + if raw.lower() in SCHEMES: + return SCHEMES[raw.lower()] + raise ValueError( + f"unknown scheme {raw!r}; image schemes: " + + ", ".join(sorted(SCHEMES)) + ) + + +def _import_markdiffusion(upstream: Path | None) -> Any: + """Import the ``markdiffusion`` module from a checkout or the environment.""" + if upstream is not None: + sys.path.insert(0, str(upstream)) + try: + import markdiffusion + except ImportError as e: + raise _Unavailable( + "markdiffusion not importable: " + str(e) + + " (set MARKDIFFUSION_DIR / --upstream-dir to a checkout, or " + "pip install markdiffusion[optional])" + ) from e + return markdiffusion + + +def _load_diffusion(model: str, device: str, offline: bool, size: int): + """Load the Stable Diffusion pipeline and scheduler used by the harness.""" + if offline: + os.environ.setdefault("HF_HUB_OFFLINE", "1") + load_kwargs = {"local_files_only": True} if offline else {} + + import torch + from diffusers import DPMSolverMultistepScheduler, StableDiffusionPipeline + + scheduler = DPMSolverMultistepScheduler.from_pretrained( + model, subfolder="scheduler", **load_kwargs + ) + dtype = torch.float16 if device == "cuda" else torch.float32 + pipe = StableDiffusionPipeline.from_pretrained( + model, + scheduler=scheduler, + torch_dtype=dtype, + safety_checker=None, + **load_kwargs, + ).to(device) + return pipe, scheduler + + +def _resolve_config(upstream: Path | None, config: str | None) -> str | None: + if not config: + return None + path = Path(config).expanduser().resolve() + if not path.is_file(): + raise _Unavailable(f"MarkDiffusion config not found: {path}") + try: + size = path.stat().st_size + except OSError as e: + raise _Unavailable(f"cannot stat MarkDiffusion config {path}: {e}") from e + if size > MAX_CONFIG_BYTES: + raise _Unavailable( + f"MarkDiffusion config too large ({size} bytes > {MAX_CONFIG_BYTES}): {path}" + ) + return str(path) + + +def _json_safe(obj: Any) -> Any: + """Convert numpy/torch scalars so the payload is JSON-serializable.""" + if isinstance(obj, dict): + return {str(k): _json_safe(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_json_safe(v) for v in obj] + if isinstance(obj, (float, int, str, bool)) or obj is None: + return obj + try: + import numpy as np + + if isinstance(obj, np.generic): + return obj.item() + except Exception: + pass + try: + import torch + + if isinstance(obj, torch.Tensor): + return obj.item() + except Exception: + pass + try: + return float(obj) + except Exception: + return str(obj) + + +def _detect_payload(result: dict, config_dict: dict | None) -> dict[str, Any]: + is_wm = bool(result.get("is_watermarked", False)) + metrics = _json_safe(result) + score: float | None = None + for key in ("p_value", "l1_distance", "distance", "score", "acc", "corr"): + v = metrics.get(key) + if isinstance(v, (int, float)): + score = float(v) + break + if score is None: + for v in metrics.values(): + if isinstance(v, (int, float)) and str(v) != str(is_wm): + score = float(v) + break + cfg = config_dict or {} + threshold = cfg.get("threshold") + if not isinstance(threshold, (int, float)): + threshold = None + threshold_p = cfg.get("threshold_p_value") + if not isinstance(threshold_p, (int, float)): + threshold_p = None + return { + "is_watermarked": is_wm, + "score": score, + "threshold": threshold, + "threshold_p_value": threshold_p, + "metrics": metrics, + } + + +def _save_png(image: Any, path: str) -> None: + buf = io.BytesIO() + image.save(buf, format="PNG") + safe_write_bytes(path, buf.getvalue()) + + +def _cmd_watermark(args: argparse.Namespace, upstream: Path | None, scheme: str) -> int: + prompt = read_text_input(args.prompt, allow_binary=args.force_text) + + device = resolve_device(args.device) + config_path = _resolve_config(upstream, args.config) + + try: + _import_markdiffusion(upstream) + from markdiffusion.watermark import AutoWatermark + from markdiffusion.utils import DiffusionConfig + + pipe, scheduler = _load_diffusion(args.model, device, args.offline, args.size) + diffusion_config = DiffusionConfig( + scheduler=scheduler, + pipe=pipe, + device=device, + image_size=(args.size, args.size), + num_inference_steps=args.steps, + guidance_scale=args.guidance, + gen_seed=args.seed, + inversion_type="ddim", + ) + wm = AutoWatermark.load( + scheme, + algorithm_config=config_path, + diffusion_config=diffusion_config, + ) + if args.seed is not None: + import torch + + torch.manual_seed(args.seed) + watermarked = wm.generate_watermarked_media( + prompt, + guidance_scale=args.guidance, + num_inference_steps=args.steps, + height=args.size, + width=args.size, + ) + unwatermarked = None + if args.unwatermarked_output: + unwatermarked = wm.generate_unwatermarked_media(prompt) + except _Unavailable as e: + eprint(str(e)) + return 3 + except Exception as e: + eprint(f"generation error: {e}") + return 1 + + wm_out = args.watermarked_output + _save_png(watermarked, wm_out) + if unwatermarked is not None: + _save_png(unwatermarked, args.unwatermarked_output) + + payload = { + "available": True, + "upstream_dir": str(upstream) if upstream else None, + "scheme": scheme, + "config": config_path, + "model": args.model, + "device": device, + "watermarked_output": wm_out, + "unwatermarked_output": args.unwatermarked_output, + } + + if args.json: + emit_json(payload) + else: + print(f"{scheme}: watermarked image -> {wm_out}") + if unwatermarked is not None: + print(f" unwatermarked image -> {args.unwatermarked_output}") + return 0 + + +def _cmd_detect(args: argparse.Namespace, upstream: Path | None, scheme: str) -> int: + if args.path != "-" and not Path(args.path).is_file(): + eprint(f"not a file: {args.path}") + return 2 + + device = resolve_device(args.device) + config_path = _resolve_config(upstream, args.config) + + try: + _import_markdiffusion(upstream) + from PIL import Image + + from markdiffusion.watermark import AutoWatermark + from markdiffusion.utils import DiffusionConfig + + pipe, scheduler = _load_diffusion(args.model, device, args.offline, args.size) + diffusion_config = DiffusionConfig( + scheduler=scheduler, + pipe=pipe, + device=device, + image_size=(args.size, args.size), + num_inference_steps=args.steps, + guidance_scale=args.guidance, + gen_seed=args.seed, + inversion_type="ddim", + ) + wm = AutoWatermark.load( + scheme, + algorithm_config=config_path, + diffusion_config=diffusion_config, + ) + image = Image.open(args.path).convert("RGB") + kwargs: dict[str, Any] = {} + if args.detector_type: + kwargs["detector_type"] = args.detector_type + result = wm.detect_watermark_in_media( + image, prompt=args.prompt or "", **kwargs + ) + config_dict = getattr(getattr(wm, "config", None), "config_dict", None) + except _Unavailable as e: + eprint(str(e)) + return 3 + except Exception as e: + eprint(f"detection error: {e}") + return 1 + + det = _detect_payload(result, config_dict) + payload = { + "available": True, + "upstream_dir": str(upstream) if upstream else None, + "scheme": scheme, + "config": config_path, + "model": args.model, + "device": device, + **det, + } + + if args.json: + emit_json(payload) + else: + label = "watermarked" if det["is_watermarked"] else "not watermarked" + score_txt = f"{det['score']:.4f}" if det["score"] is not None else "n/a" + print(f"{scheme}: {label} (score {score_txt})") + return 0 + + +def _cmd_purify(args: argparse.Namespace, upstream: Path | None) -> int: + if not Path(args.path).is_file(): + eprint(f"not a file: {args.path}") + return 2 + + device = resolve_device(args.device) + + try: + _import_markdiffusion(upstream) + from PIL import Image + + from markdiffusion.evaluation.tools.image_editor import DiffusionPurification + from markdiffusion.utils import DiffusionConfig + + pipe, scheduler = _load_diffusion(args.model, device, args.offline, args.size) + diffusion_config = DiffusionConfig( + scheduler=scheduler, + pipe=pipe, + device=device, + image_size=(args.size, args.size), + num_inference_steps=args.steps, + guidance_scale=args.guidance, + inversion_type="ddim", + ) + purifier = DiffusionPurification( + diffusion_config, + purification_strength=args.purification_strength, + prompt=args.prompt or "", + ) + image = Image.open(args.path).convert("RGB") + purified = purifier.edit(image) + except _Unavailable as e: + eprint(str(e)) + return 3 + except Exception as e: + eprint(f"purification error: {e}") + return 1 + + _save_png(purified, args.output) + payload = { + "available": True, + "upstream_dir": str(upstream) if upstream else None, + "model": args.model, + "device": device, + "output": args.output, + "purification_strength": args.purification_strength, + } + + if args.json: + emit_json(payload) + else: + print(f"purified image (strength {args.purification_strength}) -> {args.output}") + return 0 + + +def _add_common(p: argparse.ArgumentParser) -> None: + p.add_argument( + "--upstream-dir", + type=Path, + default=None, + help="MarkDiffusion checkout root (default: $MARKDIFFUSION_DIR); " + "when unset the installed markdiffusion package is used", + ) + p.add_argument( + "--model", + default=os.environ.get("MARKDIFFUSION_MODEL", DEFAULT_MODEL), + help=f"HF Stable Diffusion model (default: $MARKDIFFUSION_MODEL or {DEFAULT_MODEL})", + ) + p.add_argument( + "--device", + default="auto", + help="auto|cpu|cuda|mps (default: auto)", + ) + p.add_argument( + "--offline", + action="store_true", + help="Never contact the HF hub: load the model from the local cache only", + ) + p.add_argument( + "--force-text", + action="store_true", + help="Process input even when it looks like a binary container", + ) + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__) + sub = p.add_subparsers(dest="cmd", required=True) + + wm = sub.add_parser("watermark", help="Generate a watermarked sample image") + wm.add_argument("prompt", help="Prompt file, or - for stdin") + wm.add_argument("-o", "--watermarked-output", required=True, + help="Output PNG path (images cannot go to stdout)") + wm.add_argument("-o2", "--unwatermarked-output", default=None, + help="Also write an unwatermarked sample to this path") + wm.add_argument("--scheme", default="tr", + help="Scheme: " + ", ".join(sorted(SCHEMES)) + " (default: tr)") + wm.add_argument("--config", default=None, help="Algorithm config JSON (default: bundled)") + wm.add_argument("--size", type=int, default=512, help="Image size in px (default: 512)") + wm.add_argument("--steps", type=int, default=50, help="Diffusion steps (default: 50)") + wm.add_argument("--guidance", type=float, default=7.5, help="Guidance scale (default: 7.5)") + wm.add_argument("--seed", type=int, default=None, help="Optional RNG seed") + _add_common(wm) + wm.add_argument("--json", action="store_true", help="Emit JSON on stdout") + wm.set_defaults(handler=_cmd_watermark) + + det = sub.add_parser("detect", help="Same-scheme detection on an image") + det.add_argument("path", help="Image file to detect on") + det.add_argument("--scheme", default="tr", + help="Scheme: " + ", ".join(sorted(SCHEMES)) + " (default: tr)") + det.add_argument("--config", default=None, help="Algorithm config JSON (default: bundled)") + det.add_argument("--prompt", default=None, help="Optional prompt used at generation") + det.add_argument("--detector-type", default=None, + help="Detector variant (e.g. l1_distance, p_value; scheme-dependent)") + det.add_argument("--size", type=int, default=512, help="Image size in px (default: 512)") + det.add_argument("--steps", type=int, default=50, help="Diffusion steps (default: 50)") + det.add_argument("--guidance", type=float, default=7.5, help="Guidance scale (default: 7.5)") + det.add_argument("--seed", type=int, default=None, help="Optional RNG seed") + _add_common(det) + det.add_argument("--json", action="store_true", help="Emit JSON on stdout") + det.set_defaults(handler=_cmd_detect) + + pf = sub.add_parser("purify", help="Run the DiffusionPurification regeneration attack") + pf.add_argument("path", help="Image file to purify") + pf.add_argument("-o", "--output", required=True, help="Output PNG path") + pf.add_argument("--purification-strength", type=float, default=0.3, + help="Fraction of the diffusion schedule to regenerate in (0, 1] " + "(default: 0.3)") + pf.add_argument("--prompt", default=None, help="Optional prompt for denoising (default: '')") + pf.add_argument("--size", type=int, default=512, help="Image size in px (default: 512)") + pf.add_argument("--steps", type=int, default=50, help="Diffusion steps (default: 50)") + pf.add_argument("--guidance", type=float, default=7.5, help="Guidance scale (default: 7.5)") + _add_common(pf) + pf.add_argument("--json", action="store_true", help="Emit JSON on stdout") + pf.set_defaults(handler=_cmd_purify) + + args = p.parse_args() + + try: + scheme = normalize_scheme(args.scheme) if args.cmd in ("watermark", "detect") else None + except ValueError as e: + eprint(str(e)) + return 2 + + raw_upstream = args.upstream_dir or os.environ.get("MARKDIFFUSION_DIR") + upstream = resolve_upstream(str(raw_upstream) if raw_upstream else None) + + if args.cmd == "purify": + return args.handler(args, upstream) + return args.handler(args, upstream, scheme) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/skills/remove-ai-marks/scripts/requirements-markdiffusion.txt b/skills/remove-ai-marks/scripts/requirements-markdiffusion.txt new file mode 100644 index 0000000..64bf7dc --- /dev/null +++ b/skills/remove-ai-marks/scripts/requirements-markdiffusion.txt @@ -0,0 +1,8 @@ +# Dependencies for the optional MarkDiffusion image-watermark harness +# (THU-BPM/MarkDiffusion, Apache-2.0). The package is installed from PyPI at a +# pinned version; the harness imports it at runtime and it is never bundled. +# +# torch is installed separately in setup_markdiffusion.sh with the correct +# platform index (CUDA or CPU) and satisfies markdiffusion's own +# torch>=2.4,<2.11 range, so it is intentionally not listed here. +markdiffusion==1.0.2 diff --git a/skills/remove-ai-marks/scripts/setup_markdiffusion.sh b/skills/remove-ai-marks/scripts/setup_markdiffusion.sh new file mode 100644 index 0000000..ba91ddb --- /dev/null +++ b/skills/remove-ai-marks/scripts/setup_markdiffusion.sh @@ -0,0 +1,137 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Bootstrap the optional MarkDiffusion image-watermark harness backend. +# +# THU-BPM/MarkDiffusion (https://github.com/THU-BPM/MarkDiffusion) is +# Apache-2.0 and is NOT bundled in this repository. By default this script +# creates a venv and installs the package (plus the model deps it needs) from +# PyPI at a pinned version. Contributors can use --checkout to install an +# editable checkout of the upstream repo at a pinned commit instead. +# +# torch is installed separately so the right platform wheel index (CUDA or CPU) +# is used; markdiffusion's own torch range is then already satisfied. + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +DEFAULT_DIR="${MARKDIFFUSION_DIR:-$HOME/markdiffusion}" +DIR="" +# Pinned upstream commit for --checkout mode (v1.0.2, peeled tag commit). Do not +# point at a moving branch. +REF="cefdb320890acec728e3495b48824607d30329d3" +# Pinned PyPI release (also pinned in requirements-markdiffusion.txt). +PYPI_VERSION="1.0.2" +PYTHON="${PYTHON:-python3}" +CHECKOUT=0 + +usage() { + cat </.venv and installs MarkDiffusion for the optional +markdiffusion_harness.py backend. + +Options: + --dir PATH venv (and optional checkout) directory (default: \$MARKDIFFUSION_DIR or ~/markdiffusion) + --checkout install an editable checkout of THU-BPM/MarkDiffusion at a pinned commit + --ref REF git ref for --checkout (default: pinned commit SHA) + --python PY Python interpreter used to create the venv (default: python3) +EOF +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --dir) + DIR="${2:?--dir requires a value}" + shift 2 + ;; + --ref) + REF="${2:?--ref requires a value}" + shift 2 + ;; + --checkout) + CHECKOUT=1 + shift + ;; + --python) + PYTHON="${2:?--python requires a value}" + shift 2 + ;; + -h|--help) + usage + exit 0 + ;; + *) + echo "unknown option: $1" >&2 + usage >&2 + exit 2 + ;; + esac +done + +DIR="${DIR:-$DEFAULT_DIR}" +mkdir -p "$(dirname "$DIR")" +if command -v realpath >/dev/null 2>&1; then + DIR="$(realpath -m "$DIR")" +else + DIR="$(cd "$(dirname "$DIR")" && pwd)/$(basename "$DIR")" +fi + +if [[ "$CHECKOUT" -eq 1 ]]; then + if [[ ! -d "$DIR/.git" ]]; then + echo "Cloning THU-BPM/MarkDiffusion into $DIR (pinned ref: $REF)" + git clone --depth 1 --filter=blob:none --sparse \ + https://github.com/THU-BPM/MarkDiffusion.git "$DIR" + git -C "$DIR" fetch --depth 1 origin "$REF" + git -C "$DIR" checkout --detach "$REF" + HEAD_SHA="$(git -C "$DIR" rev-parse HEAD)" + if [[ "$HEAD_SHA" != "$REF" ]]; then + echo "error: expected pinned ref $REF, got $HEAD_SHA" >&2 + exit 1 + fi + else + echo "Using existing checkout: $DIR" + fi +fi + +if [[ ! -x "$DIR/.venv/bin/python" ]]; then + echo "Creating venv at $DIR/.venv" + "$PYTHON" -m venv "$DIR/.venv" +fi + +echo "Installing Python dependencies" +# Pin pip itself (unpinned --upgrade pip was a supply-chain drift point). +"$DIR/.venv/bin/python" -m pip install --upgrade "pip==26.2.1" + +# Install torch with the right platform index before the rest. markdiffusion +# pins torch>=2.4,<2.11, so this satisfies its range and pip won't re-resolve. +if command -v nvidia-smi >/dev/null 2>&1; then + cuda="$(nvidia-smi 2>/dev/null | sed -n 's/.*CUDA Version: \([0-9]*\.[0-9]*\).*/\1/p' | head -1)" + if [[ -n "$cuda" ]]; then + tag="cu${cuda/./}" + index="https://download.pytorch.org/whl/$tag" + echo "NVIDIA GPU detected (CUDA $cuda); installing torch from $index" + "$DIR/.venv/bin/python" -m pip install "torch>=2.4,<2.11" --index-url "$index" + else + echo "nvidia-smi present but no CUDA version found; installing default torch" + "$DIR/.venv/bin/python" -m pip install "torch>=2.4,<2.11" + fi +else + echo "No NVIDIA GPU detected; installing default torch (CPU/MPS)" + "$DIR/.venv/bin/python" -m pip install "torch>=2.4,<2.11" +fi + +if [[ "$CHECKOUT" -eq 1 ]]; then + echo "Installing MarkDiffusion editable checkout" + "$DIR/.venv/bin/python" -m pip install -e "$DIR" +else + echo "Installing MarkDiffusion $PYPI_VERSION from PyPI" + "$DIR/.venv/bin/python" -m pip install -r "$SCRIPT_DIR/requirements-markdiffusion.txt" +fi + +cat < Path: + upstream = tmp_path / "markdiffusion" + markdiffusion_pkg = upstream / "markdiffusion" + for sub in ( + "", + "utils", + "watermark", + "evaluation", + "evaluation/tools", + ): + (markdiffusion_pkg / sub).mkdir(parents=True, exist_ok=True) + (markdiffusion_pkg / sub / "__init__.py").write_text("") + (upstream / "PIL").mkdir(parents=True) + (upstream / "PIL" / "__init__.py").write_text(FAKE_PIL) + (upstream / "markdiffusion" / "_fake.py").write_text(FAKE_MARKDIFFUSION) + (markdiffusion_pkg / "watermark" / "__init__.py").write_text( + "from markdiffusion._fake import AutoWatermark\n" + ) + (markdiffusion_pkg / "utils" / "__init__.py").write_text( + "from markdiffusion._fake import DiffusionConfig\n" + ) + (markdiffusion_pkg / "evaluation" / "tools" / "image_editor.py").write_text( + "from markdiffusion._fake import DiffusionPurification\n" + ) + (markdiffusion_pkg / "evaluation" / "tools" / "__init__.py").write_text( + "from markdiffusion._fake import DiffusionPurification\n" + ) + return upstream + + +def _minimal_png() -> bytes: + def chunk(ctype: bytes, payload: bytes) -> bytes: + crc = zlib.crc32(ctype) + crc = zlib.crc32(payload, crc) & 0xFFFFFFFF + return struct.pack(">I", len(payload)) + ctype + payload + struct.pack(">I", crc) + + sig = b"\x89PNG\r\n\x1a\n" + ihdr = struct.pack(">IIBBBBB", 1, 1, 8, 2, 0, 0, 0) + idat = zlib.compress(b"\x00\x00\x00") + return sig + chunk(b"IHDR", ihdr) + chunk(b"IDAT", idat) + chunk(b"IEND", b"") + + +def _import_harness(): + import importlib.util + + spec = importlib.util.spec_from_file_location("mdh", str(HARNESS_SCRIPT)) + assert spec and spec.loader + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +# ---- CLI exit-code paths that need no torch/diffusers ----------------------- + + +def test_cli_unavailable_without_upstream(tmp_path: Path): + env = os.environ.copy() + env.pop("MARKDIFFUSION_DIR", None) + img = tmp_path / "img.png" + img.write_bytes(b"x") + r = subprocess.run( + [sys.executable, str(HARNESS_SCRIPT), "detect", str(img), "--json"], + capture_output=True, + text=True, + env=env, + ) + assert r.returncode == 3 + assert "markdiffusion not importable" in (r.stderr or "") + + +def test_cli_bad_input_missing_file(tmp_path: Path): + r = subprocess.run( + [sys.executable, str(HARNESS_SCRIPT), "detect", str(tmp_path / "missing.png")], + capture_output=True, + text=True, + ) + assert r.returncode == 2 + + +def test_cli_bad_scheme(tmp_path: Path): + img = tmp_path / "img.png" + img.write_bytes(b"x") + r = subprocess.run( + [sys.executable, str(HARNESS_SCRIPT), "detect", str(img), "--scheme", "bogus"], + capture_output=True, + text=True, + ) + assert r.returncode == 2 + assert "unknown scheme" in (r.stderr or "") + + +def test_cli_missing_purify_output(tmp_path: Path): + img = tmp_path / "img.png" + img.write_bytes(b"x") + r = subprocess.run( + [sys.executable, str(HARNESS_SCRIPT), "purify", str(img)], + capture_output=True, + text=True, + ) + assert r.returncode == 2 + + +# ---- Direct function tests (monkeypatched model loading) -------------------- + + +def _build_args(mod, cmd: str, **overrides) -> object: + import argparse + + args = argparse.Namespace() + args.cmd = cmd + args.upstream_dir = None + args.model = "fake/model" + args.device = "cpu" + args.offline = False + args.force_text = False + args.config = None + args.size = 512 + args.steps = 50 + args.guidance = 7.5 + args.seed = None + args.json = True + args.scheme = "tr" + args.prompt = None + args.detector_type = None + args.path = None + args.output = None + args.watermarked_output = None + args.unwatermarked_output = None + args.purification_strength = 0.3 + for k, v in overrides.items(): + setattr(args, k, v) + return args + + +def test_cmd_detect_with_fake_upstream( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +): + mod = _import_harness() + upstream = _make_fake_upstream(tmp_path) + img = tmp_path / "img.png" + img.write_bytes(_minimal_png()) + monkeypatch.setattr( + mod, + "_load_diffusion", + lambda model, device, offline, size: (object(), object()), + ) + args = _build_args(mod, "detect", path=str(img)) + rc = mod._cmd_detect(args, upstream, "TR") + assert rc == 0 + payload = json.loads(capsys.readouterr().out) + assert payload["available"] is True + assert payload["scheme"] == "TR" + assert payload["is_watermarked"] is False + assert payload["score"] == 83.5 + assert payload["threshold"] == 50 + + +def test_cmd_detect_p_value(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + mod = _import_harness() + upstream = _make_fake_upstream(tmp_path) + img = tmp_path / "img.png" + img.write_bytes(_minimal_png()) + monkeypatch.setattr( + mod, + "_load_diffusion", + lambda model, device, offline, size: (object(), object()), + ) + args = _build_args( + mod, "detect", path=str(img), detector_type="p_value" + ) + assert mod._cmd_detect(args, upstream, "tr") == 0 + + +def test_cmd_watermark_with_fake_upstream(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + mod = _import_harness() + upstream = _make_fake_upstream(tmp_path) + prompt = tmp_path / "prompt.txt" + prompt.write_text("a red fox", "utf-8") + out = tmp_path / "wm.png" + out2 = tmp_path / "plain.png" + monkeypatch.setattr( + mod, + "_load_diffusion", + lambda model, device, offline, size: (object(), object()), + ) + args = _build_args( + mod, + "watermark", + prompt=str(prompt), + watermarked_output=str(out), + unwatermarked_output=str(out2), + scheme="ringid", + ) + assert mod._cmd_watermark(args, upstream, "RI") == 0 + assert out.read_bytes() == b"FAKEPNG" + assert out2.read_bytes() == b"FAKEPNG" + + +def test_cmd_purify_with_fake_upstream(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + mod = _import_harness() + upstream = _make_fake_upstream(tmp_path) + img = tmp_path / "img.png" + img.write_bytes(_minimal_png()) + out = tmp_path / "purified.png" + monkeypatch.setattr( + mod, + "_load_diffusion", + lambda model, device, offline, size: (object(), object()), + ) + args = _build_args(mod, "purify", path=str(img), output=str(out)) + assert mod._cmd_purify(args, upstream) == 0 + assert out.read_bytes() == b"FAKEPNG" + + +def test_cmd_purify_runtime_error(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + mod = _import_harness() + upstream = _make_fake_upstream(tmp_path) + img = tmp_path / "img.png" + img.write_bytes(_minimal_png()) + + def _boom(*a, **k): + raise RuntimeError("model missing") + + monkeypatch.setattr(mod, "_load_diffusion", _boom) + args = _build_args( + mod, "purify", path=str(img), output=str(tmp_path / "out.png") + ) + assert mod._cmd_purify(args, upstream) == 1 + + +# ---- run_markdiffusion_purify wiring --------------------------------------- + + +def test_run_purify_unconfigured_returns_unavailable(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("MARKDIFFUSION_DIR", raising=False) + result = run_markdiffusion_purify(Path("x.png"), Path("y.png")) + assert result["available"] is False + + +def test_run_purify_success_parses_json( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + upstream = tmp_path / "upstream" + upstream.mkdir() + payload = {"available": True, "output": "/tmp/y.png", "device": "cpu"} + captured: dict = {} + + def fake_run(cmd, **kwargs): + captured["cmd"] = cmd + captured["kwargs"] = kwargs + return SimpleNamespace(returncode=0, stdout=json.dumps(payload), stderr="") + + monkeypatch.setattr(image_meta.subprocess, "run", fake_run) + result = run_markdiffusion_purify( + Path("x.png"), + Path("y.png"), + upstream_dir=str(upstream), + strength=0.4, + size=512, + steps=40, + device="cpu", + timeout=99, + ) + + assert result["available"] is True + assert result["device"] == "cpu" + cmd = captured["cmd"] + assert "purify" in cmd + assert "--purification-strength" in cmd and "0.4" in cmd + assert "--upstream-dir" in cmd and str(upstream) in cmd + assert captured["kwargs"]["timeout"] == 99 + if os.name == "posix": + assert ( + captured["kwargs"]["preexec_fn"] + is image_meta.ctrlregen_subprocess_preexec_fn + ) + + +def test_run_purify_runtime_error_is_reported( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + upstream = tmp_path / "upstream" + upstream.mkdir() + + def fake_run(*args, **kwargs): + return SimpleNamespace(returncode=1, stdout="", stderr="boom") + + monkeypatch.setattr(image_meta.subprocess, "run", fake_run) + result = run_markdiffusion_purify( + Path("x.png"), Path("y.png"), upstream_dir=str(upstream) + ) + assert result["available"] is False + assert "boom" in result["error"] + + +def test_run_purify_prefers_venv_python( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + upstream = tmp_path / "upstream" + if os.name == "nt": + venv_python = upstream / ".venv" / "Scripts" / "python.exe" + else: + venv_python = upstream / ".venv" / "bin" / "python" + venv_python.parent.mkdir(parents=True) + venv_python.write_text("") + captured: dict = {} + + def fake_run(cmd, **kwargs): + captured["cmd"] = cmd + return SimpleNamespace(returncode=0, stdout="{}", stderr="") + + monkeypatch.setattr(image_meta.subprocess, "run", fake_run) + run_markdiffusion_purify( + Path("x.png"), Path("y.png"), upstream_dir=str(upstream) + ) + assert captured["cmd"][0] == str(venv_python) + + +def test_clean_image_diffusion_flag( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.delenv("MARKDIFFUSION_DIR", raising=False) + monkeypatch.delenv("REVERSE_SYNTHID_DIR", raising=False) + src = tmp_path / "t.png" + src.write_bytes(_minimal_png()) + dest = tmp_path / "t.cleaned.png" + captured: dict = {} + + def fake_purify(path, output, **kwargs): + captured["path"] = path + captured["output"] = output + captured["kwargs"] = kwargs + return {"available": True, "device": "cpu"} + + monkeypatch.setattr(image_meta, "run_markdiffusion_purify", fake_purify) + result = image_meta.clean_image( + src, + dest, + remove_pixel="diffusion", + markdiffusion_dir="/tmp/upstream", + markdiffusion_strength=0.3, + ) + + assert result["pixel_removal"]["available"] is True + assert any("DiffusionPurification pixel removal" in a for a in result["actions"]) + assert captured["path"] == dest + assert captured["output"] == dest + assert captured["kwargs"]["strength"] == 0.3