mirror of
https://github.com/guillaumemeyer/watermarks-remover.git
synced 2026-08-22 13:11:57 +02:00
feat: optional MarkDiffusion image-watermark harness and purification removal
Add an optional harness around THU-BPM/MarkDiffusion (Apache-2.0) for controlled image-watermark experiments and an alternative pixel-removal engine: - markdiffusion_harness.py with watermark / detect / purify subcommands for nine image schemes (Tree-Ring, Ring-ID, ROBIN, WIND, SFW, Gaussian-Shading, GaussMarker, PRC, SEAL); same-scheme/same-model detection only - clean_image.py --remove-pixel diffusion runs the DiffusionPurification regeneration attack (blind regeneration; 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, Makefile targets, and mock-based tests (no torch in CI) - Docs: README section, SKILL.md, removal-matrix.md, vendor-notes.md, references/markdiffusion.md
This commit is contained in:
Guillaume Meyer (The Opinionated Man)
parent
be1dc9a83d
commit
67ab20afe7
@@ -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"]
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 \
|
||||
|
||||
@@ -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
|
||||
@@ -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 | — | — |
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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 []:
|
||||
|
||||
@@ -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,8 +838,7 @@ 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}")
|
||||
if remove_pixel == "ctrlregen":
|
||||
pixel_removal = run_ctrlregen_clean(
|
||||
dest,
|
||||
dest,
|
||||
@@ -756,6 +856,30 @@ def clean_image(
|
||||
"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 {
|
||||
|
||||
@@ -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())
|
||||
@@ -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
|
||||
@@ -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 <<EOF
|
||||
Usage: setup_markdiffusion.sh [--dir PATH] [--ref REF] [--checkout] [--python PYTHON]
|
||||
|
||||
Creates a venv at <dir>/.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 <<EOF
|
||||
|
||||
Done. Use the harness with:
|
||||
|
||||
export MARKDIFFUSION_DIR="$DIR"
|
||||
"$DIR/.venv/bin/python" "$SCRIPT_DIR/markdiffusion_harness.py" --help
|
||||
EOF
|
||||
@@ -0,0 +1,426 @@
|
||||
"""Tests for the optional MarkDiffusion image-watermark harness adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import struct
|
||||
import subprocess
|
||||
import sys
|
||||
import zlib
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
SCRIPTS = ROOT / "skills" / "remove-ai-marks" / "scripts"
|
||||
sys.path.insert(0, str(SCRIPTS))
|
||||
|
||||
import image_meta # noqa: E402
|
||||
from image_meta import run_markdiffusion_purify # noqa: E402
|
||||
|
||||
HARNESS_SCRIPT = SCRIPTS / "markdiffusion_harness.py"
|
||||
|
||||
FAKE_PIL = """\
|
||||
class Image:
|
||||
def __init__(self, mode="RGB", size=(10, 20)):
|
||||
self.mode = mode
|
||||
self.size = size
|
||||
@staticmethod
|
||||
def open(path):
|
||||
return Image()
|
||||
@staticmethod
|
||||
def new(mode, size):
|
||||
return Image(mode, size)
|
||||
def convert(self, mode):
|
||||
return self
|
||||
def save(self, fp, format=None, **kwargs):
|
||||
fp.write(b"FAKEPNG")
|
||||
"""
|
||||
|
||||
FAKE_MARKDIFFUSION = '''\
|
||||
class DiffusionConfig:
|
||||
def __init__(self, scheduler=None, pipe=None, device="cpu", **kwargs):
|
||||
self.device = device
|
||||
self.image_size = kwargs.get("image_size", (512, 512))
|
||||
self.num_inference_steps = kwargs.get("num_inference_steps", 50)
|
||||
|
||||
class AutoWatermark:
|
||||
def __init__(self):
|
||||
self.config = SimpleNamespace(config_dict={"threshold": 50})
|
||||
@staticmethod
|
||||
def load(scheme, algorithm_config=None, diffusion_config=None, **kwargs):
|
||||
return AutoWatermark()
|
||||
def generate_watermarked_media(self, prompt, **kwargs):
|
||||
from PIL import Image
|
||||
return Image.new("RGB", (16, 16))
|
||||
def generate_unwatermarked_media(self, prompt):
|
||||
from PIL import Image
|
||||
return Image.new("RGB", (16, 16))
|
||||
def detect_watermark_in_media(self, image, prompt="", **kwargs):
|
||||
if kwargs.get("detector_type") == "p_value":
|
||||
return {"is_watermarked": True, "p_value": 0.0001}
|
||||
return {"is_watermarked": False, "l1_distance": 83.5}
|
||||
|
||||
class DiffusionPurification:
|
||||
def __init__(self, diffusion_config, purification_strength=0.3, prompt="", purifier_pipe=None):
|
||||
self.strength = purification_strength
|
||||
def edit(self, image, prompt=None):
|
||||
from PIL import Image
|
||||
return Image.new("RGB", (16, 16))
|
||||
|
||||
from types import SimpleNamespace
|
||||
'''
|
||||
|
||||
|
||||
def _make_fake_upstream(tmp_path: Path) -> 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
|
||||
Reference in New Issue
Block a user