Compare commits

...
11 Commits
Author SHA1 Message Date
Victor Kuznetsov 67d948f57d Release 0.35.1 2026-08-30 18:39:39 -07:00
Victor Kuznetsov 1d0e6b54dd Add content-adaptive Google floor for chroma-zimage: face content resolves to 0.125 instead of 0.40
The four-fixture Google calibration showed a clean YuNet face-count split:
both zero-face text cards need 0.25 first-clean, both face fixtures clear
at 0.12 with zero cross-source spread. The adaptive arm routes face
content (face_count > 0) to the measured 0.125 rung, saving 6.4 dB PSNR
and doubling face identity cosine vs the flat 0.40 floor. The split is
Google-SynthID-specific: OpenAI's face fixture was harder than its text
fixture, so no other cohort gets an adaptive arm.
2026-08-30 17:50:56 -07:00
Victor Kuznetsov fa4748a6a8 Fix uvx invocation in clean-install step 2026-08-30 17:33:44 -07:00
Victor Kuznetsov 1e9572e2d4 Install uv before the clean-install verification step 2026-08-30 17:31:44 -07:00
Victor Kuznetsov 1a6f50db75 Add automated post-release verification workflow for all four published surfaces 2026-08-30 17:30:13 -07:00
Victor Kuznetsov b684423fc4 Skip chroma call-shape test when diffusers is absent on CI 2026-08-30 16:45:05 -07:00
Victor Kuznetsov 16bc9879e9 Ship chroma-zimage: Apache-2.0 Chroma1 global stage with measured four-cohort oracle floors
Extract TwoStageZImagePipeline shared base from QwenZImagePipeline so all three
profiles are siblings, not a chain. Add ChromaZImagePipeline on that base with
diffusers ChromaImg2ImgPipeline bf16, neutral calibrated prompt, guidance 5.0,
four effective steps with Diffusers truncation compensation, no Canny.
Flat vendor floors from the 2026-08-29/30 oracle calibration across OpenAI
(0.09), Microsoft InvisMark (0.125), Google (0.40), Meta Content Seal (0.17),
seed-stability verified on Meta, full-path face-stage interaction clean through
the Gemini oracle. Research record in docs/chroma1-engine-research.md with the
matched-strength addendum showing Chroma1 regenerates better per unit strength
but its boundaries scatter wider than qwen's. Fix text-manifest guard to reject
chroma-zimage at the boundary, update all profile references from 'both' to
'all three' across src, docs, README, and python-api.
2026-08-30 16:37:20 -07:00
Victor Kuznetsov 343da7f87d Add a synthetic Ideogram C2PA attribution fixture
A throwaway test CA with leaf subject CN=Ideogram, Inc signs a flat PNG
carrying a trainedAlgorithmicMedia claim. The fixture exercises the
issuer-to-platform attribution end to end through the real reader and pins
the second dimension the project separates: attribution names Ideogram
while signer trust stays untrusted. Recipe documented in the fixtures
README; no keys, vendor bytes, or user uploads are committed.
2026-08-28 20:00:18 -07:00
Victor Kuznetsov 6f2d50cde8 Document the parked video-mark candidates with their measurements
Jimeng CN video wordmark: no separation (1/5 positives vs 2/25 negatives hold
a stable 0.35 run). RunningHub video: separates but the cohort is two
positives, below the registration bar. The top-left pill on video is parked
by analysis: its false-fire mode is a static textured corner, which temporal
recurrence cannot screen.
2026-08-28 19:52:53 -07:00
Victor Kuznetsov 95f467b4ad Register the Doubao video mark with a corpus-calibrated single-floor policy
detect_doubao_frame reuses the image engine's synthetic doubao alpha as its
video template, so the glyph cannot drift between registries. The policy is
seedance-style single-floor, calibrated by sequential decode over 39
TC260-confirmed Doubao corpus videos and 25 negatives: a 12-frame stable run
at 0.35 accepts 33/39 positives and 0/25 negatives (0.30 already accepted a
negative), and through the shipped arbiter identify_video selects 32/39 with
no false selection. Also documents the measured impossibility of a bare-pill
removal arm with the current detector (68 OCR-confirmed pills max 0.317 vs
clean negatives max 0.353 -- no separating threshold exists).
2026-08-28 19:44:25 -07:00
Victor Kuznetsov 34505e5cc2 Synchronize visible-mark content and provider labels 2026-08-28 11:51:59 -07:00
50 changed files with 3636 additions and 1106 deletions

No files matched your search

+8 -3
View File
@@ -31,9 +31,14 @@ Every registered mark carries a committed example: `data/fixtures/visible/<key>/
(image marks) and `example.mp4` (video marks), regenerated by
`scripts/render_visible_examples.py`. `tests/test_visible_examples.py` is the canary: it
fails when a mark is registered without an example and when an engine stops detecting its
own example. The examples are synthetic composites of the committed silhouettes -- user
uploads never enter the repository. When the generator fails after a geometry or gate
change, fix the generator (and the engine) together; do not hand-edit the binaries.
own example. The examples are deterministic synthetic composites of the committed
silhouettes. When the generator fails after a geometry or gate change, fix the generator
(and the engine) together; do not hand-edit the binaries.
Visible-mark content must distinguish the registry key, the provider's canonical
product name, and the literal text rendered in the mark whenever those differ. Verify
external provider names against the provider's official source rather than inferring
them from an asset filename or OCR text.
## Model-adjacent tests
+151
View File
@@ -0,0 +1,151 @@
name: Verify release surfaces
# Post-release verification: after distribute.yml fans a release out to
# Homebrew, the HF Space, and the ComfyUI Registry, confirm every published
# surface actually landed and works. This workflow is the automated twin of
# the manual checklist in docs/release-and-distribution.md ("Release
# verification"), not a replacement for the library's own test suite (which
# already ran in test.yml and publish.yml).
on:
workflow_dispatch:
inputs:
version:
description: "Version to verify (no leading v). Blank = latest on PyPI."
required: false
workflow_run:
workflows: ["Distribute on release"]
types: [completed]
jobs:
verify:
# Only run on a successful distribute (not on manual dispatch, where the
# version input or PyPI-latest provides the target).
if: ${{ github.event_name == 'workflow_dispatch' || github.event.workflow_run.conclusion == 'success' }}
runs-on: ubuntu-latest
timeout-minutes: 10
steps:
- name: Resolve target version
id: v
run: |
set -euo pipefail
if [ -n "${{ github.event.inputs.version }}" ]; then
v="${{ github.event.inputs.version }}"
else
v=$(curl -fsSL https://pypi.org/pypi/remove-ai-watermarks/json \
| python3 -c "import sys,json;print(json.load(sys.stdin)['info']['version'])")
fi
echo "version=$v" >> "$GITHUB_OUTPUT"
echo "Verifying release $v"
# -- PyPI ---------------------------------------------------------------
- name: Verify PyPI
env:
VERSION: ${{ steps.v.outputs.version }}
run: |
set -euo pipefail
echo "Checking PyPI for $VERSION..."
rel=$(curl -fsSL "https://pypi.org/pypi/remove-ai-watermarks/$VERSION/json")
echo "$rel" | python3 -c "
import json, sys
d = json.load(sys.stdin)
assert d['info']['version'] == '$VERSION', f\"PyPI version mismatch: {d['info']['version']}\"
assert d['urls'], 'No files on PyPI'
has_wheel = any(u['packagetype'] == 'bdist_wheel' for u in d['urls'])
has_sdist = any(u['packagetype'] == 'sdist' for u in d['urls'])
assert has_wheel, 'No wheel on PyPI'
assert has_sdist, 'No sdist on PyPI'
print(f' PyPI: {len(d[\"urls\"])} files (wheel + sdist) for $VERSION')
"
# -- Clean install from PyPI ---------------------------------------------
- name: Install uv
run: curl -LsSf https://astral.sh/uv/install.sh | sh
- name: Verify clean install
env:
VERSION: ${{ steps.v.outputs.version }}
run: |
set -euo pipefail
echo "Installing remove-ai-watermarks==$VERSION in a clean environment..."
export PATH="$HOME/.local/bin:$PATH"
result=$(uvx --from "remove-ai-watermarks==$VERSION" remove-ai-watermarks --version 2>&1) \
|| { echo "::error::Clean install failed: $result"; exit 1; }
echo " Clean install: $result"
# -- Homebrew --------------------------------------------------------------
- name: Verify Homebrew formula
env:
VERSION: ${{ steps.v.outputs.version }}
run: |
set -euo pipefail
echo "Checking Homebrew tap formula..."
formula=$(curl -fsSL \
"https://raw.githubusercontent.com/wiltodelta/homebrew-tap/main/Formula/remove-ai-watermarks.rb")
echo "$formula" | grep -q "remove_ai_watermarks-$VERSION.tar.gz" \
|| { echo "::error::Homebrew formula does not reference $VERSION"; \
echo "Formula URL line:"; echo "$formula" | grep '^url'; exit 1; }
echo " Homebrew: formula points to $VERSION"
# -- ComfyUI ------------------------------------------------------------------
- name: Verify ComfyUI sync succeeded
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: |
set -euo pipefail
echo "Checking ComfyUI node sync..."
# The registry itself has no public JSON API; the sync workflow in
# the node repo already tests compatibility and publishes only on
# success, so verifying its latest run succeeded is the proxy.
latest=$(gh run list \
--repo wiltodelta/ComfyUI-remove-ai-watermarks \
--workflow sync-library.yml \
--limit 1 \
--json conclusion \
| python3 -c "
import json, sys
runs = json.load(sys.stdin)
r = runs[0] if runs else {}
print(r.get('conclusion', 'unknown'))
")
echo " ComfyUI sync: $latest"
if [ "$latest" != "success" ]; then
echo "::warning::ComfyUI sync latest run is '$latest', not success."
echo " Check: https://github.com/wiltodelta/ComfyUI-remove-ai-watermarks/actions"
else
echo " ComfyUI: sync succeeded, node published"
fi
# -- HF Space -------------------------------------------------------------------
- name: Verify HF Space is alive
run: |
set -euo pipefail
echo "Checking HF Space..."
space_url="https://wiltodelta-remove-ai-watermarks.hf.space"
# The Space may be sleeping after a factory rebuild; wake it.
for attempt in $(seq 1 12); do
status=$(curl -s -o /dev/null -w "%{http_code}" --max-time 30 "$space_url/" || echo "000")
if [ "$status" = "200" ]; then
echo " HF Space: alive (HTTP $status)"
exit 0
fi
echo " HF Space: HTTP $status (attempt $attempt), waiting for rebuild..."
sleep 15
done
echo "::warning::HF Space did not respond with HTTP 200 after 3 minutes."
echo " The Space may still be rebuilding after a factory restart."
echo " Verify manually at $space_url"
- name: Summary
env:
VERSION: ${{ steps.v.outputs.version }}
run: |
cat <<EOF
Release $VERSION verified:
PyPI: wheel + sdist present
Clean install: --version returns $VERSION
Homebrew: formula points to $VERSION
ComfyUI: sync succeeded
HF Space: alive and responding
EOF
+4
View File
@@ -30,6 +30,10 @@ _refs/
yolov8n.pt
.coverage
# Local prototype and evaluation outputs (scripts/chroma_scrub_prototype.py and
# friends); the durable record of an oracle verdict lives under data/evaluations/
out/
# Claude Code local settings
.claude/settings.local.json
.claude/scheduled_tasks.lock
+10 -5
View File
@@ -192,7 +192,7 @@ remains useful for unusually important files or after provider changes, but it
is not a product result state.
For invisible watermark removal, install the `qwen-zimage` extra. **An NVIDIA GPU
is required**: both profiles are CUDA-only, and there is no CPU or MPS fallback.
is required**: all profiles are CUDA-only, and there is no CPU or MPS fallback.
```bash
uv tool install --force "remove-ai-watermarks[qwen-zimage]"
@@ -243,7 +243,9 @@ uv run remove-ai-watermarks visible demo_banana_before.png \
`qwen-zimage` is the default profile: a Qwen-Image-2512 Lightning pass under Canny
ControlNet, followed by SAM-masked Z-Image repair of any detected face. The
alternative, `sdxl-zimage`, swaps the global stage for SDXL and keeps the same face
stage. Both are CUDA only.
stage. A third profile, `chroma-zimage`, uses the Apache-2.0 Chroma1 global pass
with its own flat vendor floors; see `docs/chroma1-engine-research.md` for the
calibration. All are CUDA only.
```bash
uv tool install --force "remove-ai-watermarks[qwen-zimage]"
@@ -320,7 +322,8 @@ remove-ai-watermarks batch ./images --mode all
Visible mark support includes:
- Google Gemini and Nano Banana visible sparkle watermark;
- Doubao, Jimeng, Qwen, Kling AI, Yuanbao, Baidu, LiblibAI, and RunningHub labels;
- Doubao, the Jimeng wordmark and top-left `AI生成` pill, Qwen, Kling AI,
Yuanbao, Baidu, LiblibAI, and RunningHub labels;
- one calibrated Microsoft top-right white AI-badge variant;
- one calibrated Samsung Galaxy AI label variant.
@@ -330,7 +333,9 @@ signals. Optional decoders add support for open DWT-DCT watermarks and Adobe
TrustMark.
The exact support matrix, including important locale and detector limits, lives
in [supported signals](docs/supported-signals.md).
in [supported signals](docs/supported-signals.md). The
[synthetic example gallery](data/fixtures/visible/README.md) carries one canonical
sample for every registered image and video mark.
## How it works
@@ -413,7 +418,7 @@ invisible removal.
The shipped profile is oracle-certified, but no public local decoder can
certify an arbitrary output at runtime. Recheck unusually important outputs
after provider changes.
- Invisible-watermark removal requires CUDA. Both profiles refuse any other
- Invisible-watermark removal requires CUDA. All profiles refuse any other
device at construction rather than falling back to one that cannot run them.
Visible removal, metadata stripping and `identify` still run anywhere.
- Provider watermark systems can change. Validate important outputs with the
+20
View File
@@ -13,3 +13,23 @@ example from [`adobe/trustmark`](https://github.com/adobe/trustmark) at commit
- SHA-256: `e58c5825ed7e5d9fb04710ea541b61bd55879cad65554c0d46260aa24b3d0755`
- Expected signal: TrustMark Variant P, schema 1
- License: MIT, reproduced in `../licenses/adobe-trustmark-MIT.txt`
## Ideogram C2PA attribution
`ideogram-c2pa.png` is a SYNTHETIC signed manifest, not a vendor download: a
64x64 flat PNG carrying a C2PA claim (`trainedAlgorithmicMedia`,
`claim_generator "Ideogram/3.0"`) signed by a throwaway test CA whose leaf
certificate subject is `O=Ideogram, Inc, CN=Ideogram, Inc`. It exercises the
issuer-to-platform ATTRIBUTION path end to end through the real reader; the
signer is deliberately untrusted (no trust anchors ship), which is also the
fixture's second assertion: attribution and signer trust are separate
dimensions.
- Expected signal: C2PA, platform `Ideogram`, `ai_from_metadata` True
- Expected signer: `untrusted` (test CA; attribution must not imply trust)
- Recipe (no keys are committed): create a CA and an ES256 leaf with
`digitalSignature` key usage and subject `CN=Ideogram, Inc`, then
`c2pa.Builder.from_json(...).sign_file(png, out, Signer.from_info(
C2paSignerInfo(b"es256", leaf_pem + ca_pem, pkcs8_key, None)))`.
- No user upload and no Ideogram artifact contributes pixels or bytes.
Binary file not shown.

After

Width:  |  Height:  |  Size: 43 KiB

+4 -4
View File
@@ -10,13 +10,13 @@ does an engine that stops detecting its own example.
Every example is SYNTHETIC: `scripts/render_visible_examples.py` composites the
mark's committed silhouette (the same font-rendered asset the detector matches)
onto a deterministic generated base photo at the engine's measured geometry.
No user upload and no vendor asset enters the repository: corpus files under
`data/spaces/` are user content and stay out of git by policy, and the
silhouettes themselves are our own renders (`scripts/render_vendor_silhouettes.py`).
The silhouettes are repository-owned font renders generated by
`scripts/render_vendor_silhouettes.py`; the gallery does not copy vendor raster
assets.
The examples demonstrate DETECTION geometry and house style, not vendor raster
fidelity; real-world variants (fonts, opacities, sizes) are covered by the
engines' calibration cohorts, which are local-only.
engines' recorded calibration evidence.
## Regeneration
Binary file not shown.
+308
View File
@@ -0,0 +1,308 @@
# Chroma1 engine research (2026-08-29/30)
> Research archive. This page records experiments and decisions from the dates
> above. Current package behavior is defined by the
> [supported signals](supported-signals.md), [known limitations](known-limitations.md),
> and [module internals](module-internals.md). Nothing here is shipped.
Cited experiment behind issue #88 ("Please integrate Flux 2"): FLUX.2 has no
native strength img2img anywhere (BFL reference code or diffusers; its image
inputs are reference conditioning and outputs start from pure noise), so the
requestable scrub lane in the FLUX family is **Chroma1** -- the FLUX.1
architecture re-trained on open data, Apache-2.0 on weights, with a native
`ChromaImg2ImgPipeline` strength parameter. This page records what was measured
before any integration decision.
## Setup
All Chroma1 runs: `lodestones/Chroma1-HD` bf16 on Modal H100, seed 0, the
neutral scrub prompt (`high quality, sharp, detailed, faithful to the original`
/ `blurry, lowres, distorted text, garbled text, artifacts`), guidance 5.0, and
**step compensation**: Diffusers truncates the step COUNT (`int(steps *
strength)`), so the requested count is scaled to always spend 4 effective
denoising steps, the same semantics as `sdxl_zimage_pipeline.requested_steps`.
Without it, strength 0.05 runs a single step and the floor reads artificially
high (the first, uncompensated sweep put the OpenAI boundary at (0.05, 0.10];
compensated it is (0.05, 0.06]).
Two Flux-family traps found while wiring the prototype:
- without explicit `width`/`height`, the img2img pipeline silently generates
1024x1024 regardless of input aspect;
- dimensions floor to the /16 latent-patch grid, so the result must be resized
back to the input size for pixel-exact comparison.
Scripts: `scripts/chroma_scrub_prototype.py` (first sweep),
`scripts/engine_quality_price_probe.py` (head-to-head with the production
qwen-zimage path), `scripts/engine_cohort_calibration.py` (the cohort ladders
below), `scripts/chroma_meta_high_rungs.py` (the Meta rungs above 0.08).
Generated outputs stay outside the repository (`out/`, gitignored); the durable
verdicts are the tables on this page.
## Measured floors
Oracle per cohort: OpenAI via openai.com/verify (SynthID read separately from
C2PA); Meta via the anonymous meta.ai/identification oracle. Each fixture's
rungs were checked ascending and stopped at the first clean verdict.
| cohort | fixture | detected at | clean from |
|---|---|---|---|
| OpenAI | typography card 1122x1402 | 0.05 | 0.06 |
| OpenAI | 9-face grid 1448x1086 | 0.06 | 0.075 |
| OpenAI | full-pipeline quality fixture 1122x1402 | 0.05 | 0.06 |
| Meta | gen_lighthouse_watercolor 1600x1600 | 0.08 | 0.10 |
| Meta | gen_fox_forest 1600x1600 | 0.06 | 0.08 |
| Meta | gen_night_city 1920x1280 | 0.045 | 0.06 |
| Meta | gen_studio_mug 1600x1600 | 0.03 | 0.045 |
| Meta | gen_text_poster 1600x1600 | -- | <=0.03 |
Derived operating points by the shipped rule (worst first-clean boundary plus
one observed cross-source spread):
- **OpenAI: 0.09** (0.075 + (0.075 - 0.06)). Against qwen-zimage 0.07675 and
sdxl-zimage 0.15. The typography boundary (0.05, 0.06] is seed-stable across
seeds 0, 1, 2.
- **Meta: 0.17** (0.10 + (0.10 - 0.03)), oracle-verified clean on the worst
fixture at 0.17. Against qwen-zimage 0.1. Chroma needs ~1.7x qwen's strength
on Content Seal.
- **Google: 0.40, measured 2026-08-30 across all four fixtures.** A
metadata-stripped original verified as pixel-DETECTED through the @synthid
Verify AI agent (the C2PA-carrying original answers about C2PA instead,
which is not a pixel verdict).
| fixture | detected at | clean from |
|---|---|---|
| Gemini_Generated_Image_633uuy | 0.20 | 0.25 |
| Gemini_Generated_Image_akdbei | 0.20 | 0.25 |
| Gemini_Generated_Image_y48j3c | 0.08 | 0.12 |
| Gemini_Generated_Image_3mc4t9 | 0.08 | 0.12 |
Derived by the shipped rule: worst first-clean 0.25 + spread (0.25 - 0.12
= 0.13) = 0.38, rounded up to the measured rung **0.40**, then
oracle-verified clean on the worst fixture. Against qwen-zimage 0.27:
Chroma needs ~1.5x the strength on Google SynthID. At 0.15-0.20 the tool
still reads positive while its wording degrades from "most or all" to
"part of", so the boundary band is real signal decay, not noise.
Oracle mechanics learned while bracketing: the quota is per ACCOUNT and
additionally dedupes SIMILAR images ("quota for checking images similar to
this one"), roughly three checks per account window; switching accounts
inside one browser works. A fresh Playwright context with the same Google
session cookies (exported browser-wide via CDP `Storage.getCookies`, never
through the transcript) carries working sign-in and its own check budget.
The agent's own instructions leak into the answer on some renders and
include an abstention rule for text-dominated images - always read the TOOL
outcome, never the prose.
- **Microsoft: 0.125**, measured 2026-08-30 on the SAME three valid Paint
carriers behind the qwen floor (staged under
`out/cohort-calibration/microsoft/` with pixel-identical metadata-stripped
controls; they are real raiw-corpus user uploads and are never committed).
Controls re-verified `Microsoft AI detected` on the day of the run. Against
qwen-zimage 0.15: Chroma needs LESS strength on InvisMark.
| source (Paint) | detected at | clean from | qwen first-clean |
|---|---|---|---|
| paint-1 (47d87e99, 1536x1024) | 0.06 | 0.08 | 0.055 |
| paint-2 (a579fb70, 1536x1024) | -- | <=0.04 | 0.04125 |
| paint-3 (883191b5, worst qwen source) | 0.06 | 0.08 | 0.095 |
Derived by the shipped rule: 0.08 + (0.08 - 0.04) = 0.12, rounded to the
measured rung **0.125**, which was then oracle-verified clean on both worst
sources. Oracle: the public `ai.azure.com/nextgen/validate` page; its
`Inconclusive` verdict is the site's watermark-negative, weaker than an
API-level `Watermark: false`, the same caveat the qwen cohort carries.
## Fidelity at the floors (scripts/fidelity_metrics.py)
qwen-zimage rows ran through the package's own `QwenZImagePipeline` on the same
H100 (production path, both stacks resident, prompt cache warm).
OpenAI typography (chroma 0.06 vs qwen 0.07675):
| | CER | img LPIPS | SSIM | PSNR |
|---|---|---|---|---|
| qwen-zimage | 0.241 | 0.083 | 0.834 | 27.7 |
| chroma1 | **0.138** | **0.049** | **0.881** | **31.2** |
OpenAI full-pipeline fixture (chroma 0.09 vs qwen 0.07675): chroma wins every
metric (CER 0.155 vs 0.259, LPIPS 0.056 vs 0.094, PSNR 29.6 vs 26.1).
OpenAI 9-face grid (chroma 0.075 global-only vs qwen 0.07675 with its Z-Image
face stage):
| | ID cos | face LPIPS | img LPIPS | SSIM | PSNR |
|---|---|---|---|---|---|
| qwen-zimage | **0.857** | **0.051** | 0.065 | 0.865 | 31.1 |
| chroma1 (no face stage) | 0.788 | 0.067 | **0.060** | **0.890** | **32.4** |
The face stage regenerates from the ORIGINAL crops and is engine-independent,
so an integrated `chroma-zimage` would composite the same faces over the
chroma global; whether that closes the identity gap is an integration-time
measurement, not a given.
Meta at the respective floors (chroma 0.17 vs qwen 0.1): qwen wins perceptual
fidelity on textured content (LPIPS: fox 0.163 vs 0.211, lighthouse 0.173 vs
0.248, night city 0.187 vs 0.277), chroma wins the text poster decisively
(LPIPS 0.008 vs 0.012, CER 0.000 vs 0.021) and the studio mug is a split. The
extra 0.07 of strength costs real fidelity on photoreal content.
Google at the respective floors (chroma 0.40 vs qwen 0.27, qwen row includes
its Z-Image face stage):
| fixture | qwen LPIPS/SSIM/ID cos | chroma LPIPS/SSIM/ID cos |
|---|---|---|
| 633uuy | **0.405 / 0.577** / - | 0.499 / 0.576 / - (CER 1.000 vs 0.333) |
| akdbei | **0.440 / 0.552** / - | 0.540 / 0.557 / - (CER 0.796 vs 0.367) |
| y48j3c (18 faces) | **0.436 / 0.608 / 0.801** | 0.560 / 0.548 / 0.279 |
| 3mc4t9 (5 faces) | **0.422 / 0.529 / 0.921** | 0.562 / 0.461 / 0.143 |
qwen wins every axis that matters on every fixture; at 0.40 the Chroma global
regeneration itself collapses face identity (0.14-0.28 cosine) and destroys
dense text (CER up to 1.0). The extra 0.13 of strength is simply too much
regeneration. This is the Meta pattern amplified.
Microsoft at the respective floors (chroma 0.125 vs qwen 0.15): chroma wins
EVERY metric on EVERY source while scrubbing at lower strength --
| source | CER q/ch | img LPIPS q/ch | SSIM q/ch | PSNR q/ch |
|---|---|---|---|---|
| paint-1 | 0.741 / 0.241 | 0.164 / 0.106 | 0.742 / 0.847 | 24.8 / 29.3 |
| paint-2 | 0.826 / 0.391 | 0.253 / 0.108 | 0.536 / 0.633 | 22.5 / 25.5 |
| paint-3 | 0.000 / 0.000 | 0.038 / 0.020 | 0.762 / 0.953 | 36.3 / 42.2 |
The Microsoft cohort therefore behaves like the OpenAI cohort, not like Meta:
for InvisMark carriers Chroma1 is the better engine at a lower operating point.
## Price (warm H100 container)
| | qwen-zimage | chroma1 |
|---|---|---|
| seconds per image | 2.44 | 3.62 |
| VRAM peak | 31.9 GiB | 29.2 GiB |
| USD per 1000 images ($3.95/h) | $2.68 | $3.97 |
| weights in cache | ~70 GB | ~28 GB |
The chroma number includes a T5 encode per call; the shipped qwen path caches
prompt embeddings, and the same treatment would put chroma near 1.5-2 s/image.
Both stacks fit an 80 GB card; chroma's smaller footprint likely fits smaller
cards untested.
## Matched-strength comparison: the floors hide the engine (2026-08-30 addendum)
The floor-based tables above compare chroma at its DERIVED floor (worst case
+ spread) against qwen at its shipped floor - different perturbation levels.
At the strength where each fixture ACTUALLY turns clean (chroma first-clean
rungs, all oracle-verified above), the picture reverses:
Google (chroma@first-clean vs qwen@0.27):
| fixture | chroma rung | LPIPS q/ch | SSIM q/ch | PSNR q/ch | CER q/ch | ID q/ch |
|---|---|---|---|---|---|---|
| 633uuy | 0.25 | 0.405/**0.375** | 0.577/**0.673** | 18.8/**22.3** | **0.333**/0.667 | - |
| akdbei | 0.25 | 0.440/**0.409** | 0.552/**0.650** | 18.6/**22.5** | 0.367/**0.163** | - |
| y48j3c | 0.12 | 0.436/**0.329** | 0.608/**0.700** | 21.0/**25.8** | 1.0/1.0 | **0.801**/0.611 |
| 3mc4t9 | 0.12 | 0.422/**0.318** | 0.529/**0.644** | 17.7/**21.3** | 0.443/**0.229** | **0.921**/0.410 |
Meta (chroma@first-clean vs qwen@0.1):
| fixture | chroma rung | LPIPS q/ch | SSIM q/ch | PSNR q/ch | CER q/ch |
|---|---|---|---|---|---|
| lighthouse | 0.10 (equal) | 0.173/**0.097** | 0.783/**0.886** | 25.6/**30.8** | - |
| fox | 0.08 | 0.163/**0.083** | 0.778/**0.907** | 25.6/**31.4** | - |
| night_city | 0.06 | 0.187/**0.092** | 0.748/**0.893** | 23.9/**29.6** | 0.850/**0.250** |
| studio_mug | 0.045 | 0.024/**0.009** | 0.984/**0.992** | **39.0**/37.8 | - |
| text_poster | 0.03 | 0.012/**0.002** | 0.984/**0.994** | 34.0/**41.2** | 0.021/**0.007** |
At effective-clean strength Chroma1 wins whole-image fidelity on essentially
every fixture (often by 4-6 dB) and text on five of six; qwen's remaining
edge is face identity, which its Z-Image face stage buys (0.80-0.92 vs
0.41-0.61 cosine) and a chroma-zimage profile would inherit on top of its own
global pass.
**Why the floors still matter:** a shipped strength must clear the worst
unseen content, and Chroma1's per-fixture boundaries scatter much wider
(Google 0.12-0.25, Meta 0.03-0.10) than the spread the flat floors were built
to absorb - so its derived floors (0.40 / 0.17) pay a large margin penalty
and at those floors the extra strength genuinely destroys quality. The
engines differ less in "quality per strength" than in how predictable their
boundary is. The integration lever this exposes: a content-adaptive strength
policy for chroma-zimage (the project already ships resolution-adaptive
denoise for qwen) could capture chroma's low per-content boundaries; the two
hard Google fixtures are both dense-text cards and the easy ones are
photoreal, so a content-class signal may predict the rung.
## Pre-ship validation (2026-08-30)
With the `chroma-zimage` profile implemented
(`_internal/chroma_zimage_pipeline.py` on `TwoStageZImagePipeline`, floors in
`watermark_profiles`), two ship-gate checks were run:
1. **Seed stability at the worst Meta boundary.** lighthouse @ 0.10 with
seeds 1 and 2: both CLEAN on the meta.ai oracle. The Meta boundary is not a
seed-0 fluke. The Google seed probes (633uuy @ 0.25, seeds 1-2) were
generated but the oracle's per-image-class quota blocked them on every
available account; they remain queued, and the shipped 0.40 Google floor
carries 0.15 of margin over that boundary regardless.
2. **Full-path face-stage interaction.** The complete profile run (Chroma1
global at the shipped Google floor 0.40, then the inherited
YuNet/SAM/Z-Image face stage) on the 5-face fixture 3mc4t9: CLEAN on the
@synthid oracle. The composited face regions do not re-introduce a
detectable signal on top of the floor.
Generation script: `scripts/chroma_preship_validation.py`; outputs under
`out/preship-chroma/` (gitignored).
## Verdict
Engine quality is cohort-dependent, and after measuring all four cohorts the
split is even and sharp:
- On OpenAI-provenance content Chroma1 is better at every measured fidelity
axis at its floor (0.09 vs 0.07675), for comparable money.
- On Microsoft InvisMark content Chroma1 is better at every measured fidelity
axis at a LOWER floor than qwen (0.125 vs 0.15).
- On Meta Content Seal Chroma1 needs 1.7x the strength (0.17 vs 0.1) and loses
LPIPS on textured content at its floor.
- On Google SynthID Chroma1 needs ~1.5x the strength (0.40 vs 0.27) and loses
every fidelity axis on every fixture, with face identity collapsing at 0.40
(0.14-0.28 cosine vs qwen's 0.80-0.92).
A single-engine `chroma-zimage` default at a flat per-vendor floor is NOT
supported: Google and Meta floors land too high and the quality at those
floors is worse than qwen's. But the matched-strength addendum changes what
IS supported: at the strength each image actually needs, Chroma1 regenerates
better than qwen almost everywhere except face identity (which the shared
face stage would supply). The two honest integration options:
1. Per-cohort engine router at flat floors: chroma for OpenAI/Microsoft,
qwen for Google/Meta - conservative, ships on the measurements as they
stand.
2. `chroma-zimage` with a content-adaptive strength policy that predicts the
per-image rung - potentially better everywhere, but the predictor is new
research (the measured signal so far: dense-text content needs the high
rungs, photoreal the low ones) and needs its own oracle validation before
any ship.
Both need the seed-sensitivity and face-stage composition checks called out
above before shipping. Both checks have now been run (Pre-ship validation
above): the Meta boundary is seed-stable across seeds 0-2 and the full
profile path stays clean through the face stage on the shipped Google floor.
The profile ships with option 1's floors plus a content-adaptive Google arm
(face content resolves to 0.125 instead of 0.40, keyed on YuNet detection).
Option 2's full adaptive policy for the remaining cohorts remains the
documented follow-up.
## Oracle session notes
- openai.com/verify throttles anonymous sessions through Cloudflare Turnstile
after ~15 rapid checks; a fresh browser context clears it. The programmatic
`POST /v1/content_provenance_checks` API is the better oracle (separate
SynthID/C2PA outcomes) but the checked `OPENAI_API_KEY` returned 401
org-wide, including `/v1/models`.
- meta.ai/identification enforces a sliding per-IP window: a burst of ~15
checks exhausted it twice; it reopens minutes later, and the page sometimes
overstates this as a daily limit. Pacing of roughly 2-3 checks per window
worked.
- The Gemini app oracle requires the user's logged-in session; driving the
user's live Chrome while it is in active use was refused (wrong-tab risk),
so those 20 checks remain queued.
+3 -3
View File
@@ -327,8 +327,8 @@ provider provenance may relax the visual score only for registered
provenance-aware marks; metadata alone never creates a detection.
`--mark auto` is the default. It evaluates all providers in one decode pass and
selects the first stable match in specificity order: Sora, Veo, Seedance, Dola,
Hailuo AI, then Kling AI. Their confidence scores are independently calibrated and
selects the first stable match in specificity order: Sora, Veo, Seedance,
Doubao, Dola, Hailuo AI, then Kling AI. Their confidence scores are independently calibrated and
are not compared across providers. Pass an explicit `--mark` to scan only that
provider.
@@ -367,7 +367,7 @@ failed encode does not overwrite an existing result.
## Remove invisible watermarks
Install the removal dependencies first. Both profiles are CUDA-only and both
Install the removal dependencies first. All profiles are CUDA-only and all
run the DiffSynth Z-Image face stage, so this is the extra either one needs:
```bash
+1
View File
@@ -39,3 +39,4 @@ The current behavior is defined by the code, tests, README, and user guides.
- [SynthID identity follow-up](synthid-robust-identity-research-2026-06-08.md)
- [Video SynthID quality research](video-synthid-quality-research.md)
- [Text protection research](text-protection-research.md)
- [Chroma1 engine research](chroma1-engine-research.md)
+8 -5
View File
@@ -75,11 +75,14 @@ uv tool install --force "remove-ai-watermarks[qwen-zimage]"
```
Both remaining profiles run a Z-Image face stage on the DiffSynth runtime, so
both need this extra. It includes the `diffusion` dependencies; `diffusion` on its
own covers the torch and diffusers imports but not the face stage, so it is not
enough to run a removal.
both need this extra. The third profile, `chroma-zimage`, uses diffusers'
ChromaImg2ImgPipeline instead of DiffSynth for its global stage but inherits
the same face stage, so it needs this extra too. It includes the `diffusion`
dependencies; `diffusion` on its own covers the torch and diffusers imports
but not the face stage, so it is not enough to run a removal.
**An NVIDIA GPU is required.** `qwen-zimage` and `sdxl-zimage` are CUDA-only, and
**An NVIDIA GPU is required.** `qwen-zimage`, `sdxl-zimage` and `chroma-zimage`
are CUDA-only, and
construction refuses any other device rather than falling back to a slow or broken
one. There is no CPU, MPS or XPU path for invisible-watermark removal. Visible-mark
removal, metadata stripping and every `identify` command still run anywhere.
@@ -241,5 +244,5 @@ image came from a relevant generator, use `--force`.
If the CLI reports that the removal dependencies are unavailable, install the
`qwen-zimage` extra. `diffusion` alone covers Torch and Diffusers but not the
DiffSynth face stage that both profiles run. Video SynthID removal is a separate
DiffSynth face stage that all profiles run. Video SynthID removal is a separate
path and needs `video` and `diffusion`.
+21 -6
View File
@@ -185,7 +185,7 @@ path has not yet answered.
### Strength is content and seed dependent
The two profiles resolve an unset strength differently, because different things
The profiles resolve an unset strength differently, because different things
were measured for each.
`qwen-zimage` reads unknown content from the resolution-adaptive denoise curve.
@@ -208,10 +208,25 @@ values are flat rather than a curve because flat values are what was measured: e
verdict came from a fixed strength at one size, and no size dependence has been
established for that stage.
An explicit `--strength` overrides both. The defaults are operating points, not
universal guarantees. Near a removal threshold, different content or a different
random seed may change the verifier result, which is why both profiles are
certified at a fixed seed. The live resolver is
`chroma-zimage` also reads a flat ladder from the provenance cohort, measured
2026-08-29/30 across the four vendor oracles (full record:
[`chroma1-engine-research.md`](chroma1-engine-research.md)):
- OpenAI: `0.09` (below qwen's `0.07675` is false economy neither way -- the
matched-strength fidelity favors Chroma1 on this cohort);
- Microsoft InvisMark: `0.125` (below qwen's `0.15`);
- Google: `0.40` for zero-face content (above qwen's `0.27`; at this floor the
regeneration destroys dense text and collapses face identity -- the tradeoff
that motivates the adaptive arm below), or `0.125` when YuNet detects at
least one face (both measured face fixtures clear at 0.12 with zero
cross-source spread; oracle-verified 2026-08-30);
- Meta Content Seal: `0.17` (above qwen's `0.1`);
- unknown: `0.40`, following the strictest measured cohort.
An explicit `--strength` overrides all three. The defaults are operating points,
not universal guarantees. Near a removal threshold, different content or a
different random seed may change the verifier result, which is why the profiles
are certified at a fixed seed. The live resolver is
[`watermark_profiles.py`](../src/remove_ai_watermarks/_internal/watermark_profiles.py).
### Pipelines have different quality tradeoffs
@@ -272,7 +287,7 @@ device and re-reads every parameter from disk on each stage transition.
Invisible-watermark removal refuses any device but CUDA at construction. There is
no MPS out-of-memory fallback that continues on CPU, and no lighter profile to
drop to: both remaining profiles need an NVIDIA GPU.
drop to: all three profiles need an NVIDIA GPU.
When a run does not fit, the levers are `--tile` (native geometry, tiled
diffusion), `--max-resolution` (an explicit downscale), and `--cpu-offload`.
+107 -18
View File
@@ -296,7 +296,7 @@ from the gate. The full-clip oracle floor is
detected while `0.15` did not.
[`video_visible.py`](../src/remove_ai_watermarks/video_visible.py) implements
the first pixel stages for Sora, Veo, Seedance, Dola, Hailuo AI, and Kling AI. The
the first pixel stages for Sora, Veo, Seedance, Doubao, Dola, Hailuo AI, and Kling AI. The
Sora detector searches a normalized frame with a fully synthetic
mascot-and-text silhouette at several scales. The Veo detector uses separate
synthetic silhouettes for the current four-point diamond and legacy `Veo`
@@ -327,8 +327,8 @@ carries `accepts_provenance`, which forces `provenance=False` for Hailuo AI and
— they have no metadata that could confirm them, and the guarantee used to be
structural (their wrappers took no `provenance` parameter at all). Provenance can relax a low-contrast run only
after recurring visual evidence exists. Sora transition frames follow the
nearest confirmed moving position only with Sora provenance. Seedance, Dola,
Hailuo AI, and Kling AI additionally require candidates to remain anchored to the
nearest confirmed moving position only with Sora provenance. Seedance, Doubao,
Dola, Hailuo AI, and Kling AI additionally require candidates to remain anchored to the
start of a run. This rejects slowly drifting scene details that still have
high frame-to-frame overlap. Hailuo AI and Kling AI do not infer provenance from
technical encoder tags; their confirmed public samples carried no provider
@@ -855,6 +855,22 @@ two-line mark can be light on dark scenes or dark on light scenes. Its detector
and footprint both use the same best-match box. The separate one-line overlay
variant is not covered.
#### Bare-upload pill arm: measured impossible with the current detector (2026-08-28)
`_keep_pill` never removes a metadata-bare pill, and the question "can a bare arm
(score + footprint-flatness) be opened" was measured to a closed NO over the local
spaces corpus: 68 OCR-confirmed bare pills (an independent pixel-level label: the
band sweep read a flush-left `AI生成`) score p50 0.181 / max 0.317, while 799 clean
no-signal negatives reach 0.353 and 18 of them already pass the shipped 0.22 raw
gate. The true-pill and clean distributions overlap completely -- there is no
threshold -- and the engine's own high scores (0.34-0.48 on the flagged cohort)
correlate with textured corners, not with pills: the edge-NCC keys on "text-like
structure top-left". As with the Jimeng wordmark's silhouette, no threshold repairs
this; a bare arm needs a detector that keys on the pill's rounded-rectangle
geometry, and until one exists the bare-arm question is settled by this
measurement, not by gate tuning. Harness: `data/spaces/_pill_bare_measure.py`
(cohorts A/B/C over the no-signal pool).
The capture-less Jimeng pill lives in
[`pill_engine.py`](../src/remove_ai_watermarks/pill_engine.py). It uses a
synthetic silhouette for detection and a fixed top-left footprint.
@@ -866,6 +882,32 @@ Shared behavior is covered by:
- [`test_text_mark_faint_mask.py`](../tests/test_text_mark_faint_mask.py)
- [`test_text_mark_memory.py`](../tests/test_text_mark_memory.py)
#### Doubao video mark (2026-08-28 registration)
`detect_doubao_frame` reuses the image engine's synthetic `doubao_alpha.png` as
its video template (same `豆包AI生成` run, bottom-right) through
`_asset_template`, so image and video cannot drift apart on the glyph. Search
profile: heights 3.2-6.0 percent of the short side from origin (0.60, 0.76),
measured on the corpus. Policy is seedance-style single-floor: weak 0.35 /
strong 0.55 / anchor IoU 0.80, calibrated by sequential decode over 39
TC260-confirmed Doubao videos and 25 negatives -- a 12-frame stable run at 0.35
accepts 33/39 positives and 0/25 negatives, while 0.30 already accepted one
negative, so the bar stays at 0.35. Through the shipped arbiter
(`identify_video`): 32/39 selected as doubao, 0/25 false; the remainder split
between no stable run (5) and cross-template ties handed to earlier table
entries (sora 1, kling 1) -- the registry's known order-decides behavior, not a
doubao-specific defect. Harness: `data/spaces/_doubao_seq_measure.py`.
Two further video candidates were measured and PARKED the same day (2026-08-28,
`data/spaces/_rh_jm_video_calibrate.py`): the Jimeng CN wordmark does not
separate on video (1/5 positives vs 2/25 negatives hold a 12-frame 0.35 run --
the same silhouette-generality the image jimeng engine documents), and
RunningHub video separates (1/2 positives, 0/25 negatives) but a two-positive
cohort is below the registration bar; both await corpus growth. The top-left
"AI生成" pill on video is parked by analysis, not measurement: its image-engine
false-fire mode is a STATIC textured corner, which is temporally stable, so the
recurrence arbiter cannot screen it.
### Fill backends and region erasing
[`region_eraser.py`](../src/remove_ai_watermarks/region_eraser.py) implements the
@@ -898,12 +940,12 @@ is the source of truth for:
- profile names and their underscore spellings;
- the fixed seed;
- the SDXL global-stage checkpoint id (`SDXL_MODEL_ID`) and the Canny ControlNet id;
- strength resolution for both profiles.
- strength resolution for every profile.
The current profiles are `qwen-zimage` (the default) and `sdxl-zimage`, and both
are CUDA-only. `controlnet`, `sdxl`, `qwen` and `default` were removed rather than
kept as a CPU path, and are rejected rather than aliased onward. There is no
content-dependent automatic router.
The current profiles are `qwen-zimage` (the default), `sdxl-zimage`, and
`chroma-zimage`, and all three are CUDA-only. `controlnet`, `sdxl`, `qwen` and
`default` were removed rather than kept as a CPU path, and are rejected rather
than aliased onward. There is no content-dependent automatic router.
`qwen-zimage` normally resolves global denoise from image area for unknown content.
Measured provider cohorts bypass that curve with flat operating points. The values,
@@ -921,8 +963,10 @@ by the profile, so none of them appears in `WatermarkRemover.__init__`,
then rejected several frames down; a signature that refuses the argument outright
fails where the caller can act on it, and stops a wrapper from threading a value
that would silently do nothing. The step count and CFG live with the stage that
runs them (`GLOBAL_STEPS`, `FACE_STEPS`, `GLOBAL_CFG`, `FACE_CFG` in
`qwen_zimage_pipeline.py`). The dtype is likewise profile-owned: see "Face-stage
runs them: `GLOBAL_STEPS`/`GLOBAL_CFG` with the Qwen global stage
(`qwen_zimage_pipeline.py`), `SDXL_STEPS` with the SDXL one
(`sdxl_zimage_pipeline.py`), and `FACE_STEPS`/`FACE_CFG` with the shared face
stage (`two_stage_pipeline.py`). The dtype is likewise profile-owned: see "Face-stage
dtype" for what an override cost the last time one existed.
`device` is the exception and remains a library parameter: `None` or `"auto"`
@@ -949,7 +993,7 @@ same pixels. They diverged before, in opposite directions, for exactly this knob
The global and face prompts are calibrated model inputs, and the Canny edge map
uses fixed thresholds of `_CANNY_LOW = 13` / `_CANNY_HIGH = 64`
(`qwen_zimage_pipeline.py`). Treat those values as behavioral compatibility
(`two_stage_pipeline.py`). Treat those values as behavioral compatibility
contracts: a refactor must preserve them, and any deliberate change requires
image-quality evaluation rather than only a unit-test pass. The prompt and
edge-map regression guards are
@@ -997,8 +1041,19 @@ Regression coverage:
### Qwen plus Z-Image
The profiles are siblings of one shared recipe rather than a chain:
[`_internal/two_stage_pipeline.py`](../src/remove_ai_watermarks/_internal/two_stage_pipeline.py)
holds `TwoStageZImagePipeline`, the base that owns everything profile-independent
(the Z-Image face stage with YuNet detection and SAM masks, sizing, compositing,
the prompt cache, and the `run`/`preload` orchestration). A profile subclass
implements `_load_global` and `_run_global` for its regeneration model and may
implement `_vae_roundtrip` to expose that stack's VAE as a verified-text donor;
nothing else differs, and a test asserts the shared methods are the same objects
on every subclass. Adding a third global stage means implementing those hooks,
not inheriting from another profile's pipeline.
[`_internal/qwen_zimage_pipeline.py`](../src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py)
implements the fixed CUDA-only two-stage profile:
implements the fixed CUDA-only two-stage profile on that base:
1. Qwen Image with Canny conditioning regenerates the frame.
2. YuNet locates faces, SAM builds masks, and Z-Image regenerates the selected
@@ -1098,13 +1153,47 @@ oracle outcomes, and the DiffSynth/GGUF, scheduler, detector, and seed caveats a
recorded in
[`data/evaluations/fidelity/upstream-v2-reproduction-2026-08-13.csv`](../data/evaluations/fidelity/upstream-v2-reproduction-2026-08-13.csv).
### Chroma1 plus Z-Image
[`_internal/chroma_zimage_pipeline.py`](../src/remove_ai_watermarks/_internal/chroma_zimage_pipeline.py)
runs the same two-stage recipe on a Chroma1-HD (`lodestones/Chroma1-HD`,
Apache-2.0) global pass through diffusers' `ChromaImg2ImgPipeline`. It is the
answer to issue #88's FLUX.2 request through the model that actually exposes
strength-controlled img2img in that family; the full research record,
including why FLUX.2 itself is not integrable this way, is
[`chroma1-engine-research.md`](chroma1-engine-research.md).
The profile swaps only the global stage (same inheritance invariant as
sdxl-zimage, asserted by the same test shape). Three things are bound to the
calibration and must not drift: the NEUTRAL prompt (`high quality, sharp,
detailed, faithful to the original` -- deliberately NOT the shared canny-stage
prompt, which was never calibrated against Chroma1), guidance 5.0, and the
four-EFFECTIVE-step schedule with `requested_steps` compensation (the same
Diffusers truncation trap sdxl-zimage documents). There is no Canny
conditioning: the floors were measured on a plain strength pass.
The flat vendor floors (`CHROMA_ZIMAGE_*_STRENGTH` in watermark_profiles):
OpenAI 0.09 and Microsoft 0.125 (both BELOW qwen's, with better
matched-strength fidelity), Google 0.40 for zero-face content / 0.125 when
faces are detected (a content-adaptive arm from the clean face-count split
in the calibration: both text cards need 0.25, both face fixtures clear at
0.12, so the YuNet detector -- already loaded for the face stage -- routes
the operating point), and Meta 0.17 (ABOVE qwen's, because Chroma1's
per-fixture boundaries scatter wider). The matched-strength addendum in the
research doc is the honest read: at the strength each image actually needs,
Chroma1 regenerates better almost everywhere except face identity (which the
inherited Z-Image face stage supplies); at the flat worst-case floors it
destroys dense text and face identity. The Google face-content arm is the
first shipped piece of the content-adaptive policy.
### SDXL plus Z-Image
[`_internal/sdxl_zimage_pipeline.py`](../src/remove_ai_watermarks/_internal/sdxl_zimage_pipeline.py)
runs the same two-stage recipe on an SDXL global pass. `SdxlZImagePipeline` subclasses
`QwenZImagePipeline` and overrides only `_run_global` and `preload`, so the face stage
is inherited rather than copied and cannot drift between the profiles; a test asserts
the shared methods are the same objects.
runs the same two-stage recipe on an SDXL global pass. `SdxlZImagePipeline`
subclasses `TwoStageZImagePipeline` and implements only the global stage
(`_load_global`, `_run_global`), so the face stage is inherited rather than copied
and cannot drift between the profiles; a test asserts the shared methods are the
same objects.
Four things are architecture-bound and swap with the model: the ControlNet
(`xinsir/controlnet-canny-sdxl-1.0`), the four-step distillation LoRA
@@ -1142,7 +1231,7 @@ avoid model downloads, so nothing exercised the loader. The lesson is narrower t
was a dtype the subclass silently changed out from under it.
Note what the seam is, because it decides where the fix belongs.
`SdxlZImagePipeline._load_sdxl` hardcodes fp16 for its own ControlNet, VAE and
`SdxlZImagePipeline._load_global` hardcodes fp16 for its own ControlNet, VAE and
pipeline, so `self.torch_dtype` was never actually the global stage's dtype on this
profile -- its only remaining readers were face-stage code. SAM was the second one:
it never crashed, because it casts its own inputs and leaves through `.float()`, but it
@@ -1229,7 +1318,7 @@ within one pipeline.
`_cache_static_prompt_embeddings` therefore persists what the text encoder produced
under `_model_cache_dir()/prompt-embeddings`, keyed by cache version, model id,
pipeline output params, and the exact prompt string. Once that file exists,
`_load_qwen` and `_load_zimage` drop the text-encoder `ModelConfig` from the model
`_load_global` and `_load_zimage` drop the text-encoder `ModelConfig` from the model
stack entirely and serve the stored tensors instead. Measured on an H100 volume in
August 2026, that removes **15.45 GiB** (Qwen2.5-VL) and **7.49 GiB** (Z-Image) of a
**87.6 GiB** per-request read, worth a median **11.76 s** and **4.10 s** of load time
+4 -4
View File
@@ -587,7 +587,7 @@ to 8-bit SDR.
## Remove invisible watermarks
Install `remove-ai-watermarks[qwen-zimage]`. Both profiles need it, and both
Install `remove-ai-watermarks[qwen-zimage]`. All three profiles need it, and all
need an NVIDIA GPU.
```python
@@ -596,7 +596,7 @@ from pathlib import Path
from remove_ai_watermarks.invisible_engine import InvisibleEngine
engine = InvisibleEngine(
pipeline="qwen-zimage", # the default; the only other value is "sdxl-zimage"
pipeline="qwen-zimage", # the default; also "sdxl-zimage" or "chroma-zimage"
device=None,
cpu_offload=False,
)
@@ -620,7 +620,7 @@ engine = InvisibleEngine(
)
```
Both profiles are CUDA-only, so on a machine without an NVIDIA GPU `device=None`
All profiles are CUDA-only, so on a machine without an NVIDIA GPU `device=None`
resolves to `cpu` and construction raises. For the SDXL global stage instead of
Qwen:
@@ -628,7 +628,7 @@ Qwen:
engine = InvisibleEngine(pipeline="sdxl-zimage")
```
The `qwen-zimage` extra is required for both profiles: each runs the same
The `qwen-zimage` extra is required for all three profiles: each runs the same
DiffSynth Z-Image face stage.
The opt-in verified-text stage uses the same `text_manifest` argument as the CLI:
+5 -1
View File
@@ -121,7 +121,11 @@ without removing schema 1 so long-lived consumers can update separately. A packa
release must never silently substitute its latest schema when a caller explicitly
requests an older supported one.
After publication, verify:
After publication, the `verify-release.yml` workflow runs automatically once
`distribute.yml` completes and checks every surface below. It can also be
dispatched manually with a specific version. The ComfyUI check verifies the
sync workflow's latest run succeeded rather than querying the registry directly
(the registry has no public JSON API). The full checklist for completeness:
- both wheel and source distribution exist on PyPI;
- the package version matches the tag;
+6 -1
View File
@@ -20,7 +20,7 @@ The `visible` command registers these mark keys:
| `runninghub` | `RunningHub AI生成` | Top left | Strict visual and position gates. |
| `baidu` | `百度 AI生成` | Bottom right | Detector and extended removal footprint. |
| `liblib` | `LiblibAI` | Bottom center | Includes a minimum image size gate. |
| `microsoft` | One Microsoft white AI-badge variant | Top right | Strict gate; other documented icon, text, and position variants are not covered. |
| `microsoft` | One Microsoft white AI-badge variant | Top right | Strict uses the visual gate; auto can use Microsoft provenance for the measured [relaxed gate](module-internals.md#visible-mark-removal). Other documented icon, text, and position variants are not covered. |
| `jimeng_pill` | `AI生成` pill | Top left | Weak detector with additional product and background gates. |
`--mark auto` evaluates all registered marks and removes every selected match.
@@ -30,12 +30,17 @@ the masked area.
Marks from other vendors are not detected automatically. Use `erase --region`
when you can select the affected area yourself.
Synthetic examples for every registered image and video mark live in the
[visible-mark gallery](../data/fixtures/visible/README.md). They demonstrate the
detectors' canonical geometry and house style, not vendor raster fidelity.
### Visible video marks
| Key | Mark | Motion | Important limit |
| --- | --- | --- | --- |
| `sora` | Sora 2 mascot and wordmark | Moves among frame positions | Requires a temporally recurring visual match; the older Sora Turbo corner swirl is a different unsupported mark. |
| `veo` | Current four-point diamond and legacy `Veo` text | Fixed bottom-right corner | Uses separate silhouettes and requires a recurring match; learned fill is preferable on structured backgrounds. |
| `doubao` | `豆包AI生成` text run | Fixed bottom-right corner | Reuses the image engine's synthetic alpha as the template; a stable 12-frame run above 0.35 is required. |
| `seedance` | Boxed `AI` label | Fixed bottom-right corner | Requires an anchored recurring match; the full localized box is filled because a thinner synthetic shape mask leaves the real translucent rim behind. |
| `dola` | `Dola AI` text | Fixed bottom-right corner | Requires an anchored recurring match; ByteDance or BytePlus provenance can relax only an existing visual run. |
| `hailuo` | `MINIMAX \| hailuo AI` composite label | Fixed lower edge | Uses a synthetic waveform, text, separator, and ring silhouette; the complete recurring label box is filled. A TC260 label naming MiniMax as producer can relax only an existing stable run. |
+1 -1
View File
@@ -248,7 +248,7 @@ from the test set + this doc).
watermark. Stop and reconsider.
2. **PhotoMaker-V1 prototype** in the existing `controlnet` pipeline:
- Mirror the `_load_controlnet_pipeline` path (removed in 0.24.0 with the
controlnet profile; the equivalent seam is now `SdxlZImagePipeline._load_sdxl`):
controlnet profile; the equivalent seam is now `SdxlZImagePipeline._load_global`):
add a PhotoMaker variant that
loads SDXL + canny ControlNet + PhotoMaker adapter on the same engine.
- Extract the OpenCLIP face embedding from the watermarked face crops (use
+1 -1
View File
@@ -4,7 +4,7 @@
# on it, including the ComfyUI node package. The console script below carries
# the same weight, since users have it on PATH.
name = "remove-ai-watermarks"
version = "0.34.1"
version = "0.35.1"
description = "AI watermark remover for visible, invisible, and provenance marks in images and video"
readme = "README.md"
requires-python = ">=3.11,<3.15"
@@ -12,6 +12,11 @@ provenance cohorts, not automatically visible-mark positives; comparison control
are independently selected no-signal images and are not adjudicated negatives.
Each needs capture-solved alphas or vendor-accurate font work before it can
ship; `scripts/vendor_mark_calibrate.py` is the candidate-detector harness.
Rows through `mindvideo_alpha.png` and the parked candidates were recorded on
2026-08-27; rows from `higgsfield_alpha.png` through `gendo_alpha.png` were
recorded on 2026-08-28. These are repository record dates, not inferred capture
dates. The Mark column quotes the candidate raster text; where that text differs
from the provider's canonical name, the provider is named separately.
| Asset | Mark | Evidence | Result |
| --- | --- | --- | --- |
@@ -20,13 +25,13 @@ ship; `scripts/vendor_mark_calibrate.py` is the candidate-detector harness.
| `notebooklm_alpha.png` | NotebookLM wordmark (bottom-right) | 12 corpus files | locate geometry not yet fitted; POS max 0.12. |
| `dola_alpha.png` | DolaAI on images (the video mark is registered) | 12 corpus files | POS 0.11-0.21 vs NEG max 0.30. |
| `mindvideo_alpha.png` | MindVideo.AI (top-right) | 11 corpus files | POS 0.29-0.32 vs NEG max 0.30 -- borderline overlap, not shippable. |
| `higgsfield_alpha.png` | HIGGSFIELD AI wordmark (bottom-right; the boxed `AI` variant shares the cohort) | 5 wordmark files (16 in the boxed-AI OCR cluster) | POS max 0.26 vs NEG max 0.22 -- no separation; the mark may be two-part (wordmark + boxed AI) and needs a composed template. |
| `jianying_alpha.png` | 剪映AI (CapCut's CN sibling, bottom-right) | 2 corpus files | POS 0.29 vs NEG max 0.35. |
| `higgsfield_alpha.png` | `HIGGSFIELD AI` wordmark (Higgsfield AI, bottom-right; the boxed `AI` variant shares the cohort) | 5 wordmark files (16 in the boxed-AI OCR cluster) | POS max 0.26 vs NEG max 0.22 -- no separation; the mark may be two-part (wordmark + boxed AI) and needs a composed template. |
| `jianying_alpha.png` | `剪映AI` wordmark (Jianying / 剪映, the China product in the CapCut family; bottom-right) | 2 corpus files | POS 0.29 vs NEG max 0.35. |
| `capcut_alpha.png` | CapCut AI pill (top-left; likely pill class, not plain text) | 3 corpus files | POS max 0.15 vs NEG max 0.33 -- locate geometry not yet fitted for the pill form. |
| `zsky_alpha.png` | MADE WITH zsky.ai (bottom-right) | 2 corpus files | POS 0.10 vs NEG max 0.27. |
| `chromastudio_alpha.png` | ChromaStudio.ai (bottom-right) | 2 corpus files | POS 0.11 vs NEG max 0.29. |
| `digenai_alpha.png` | DIGENAI (bottom-right) | 3 corpus files (one 2026-07-24 batch) | POS 0.16 vs NEG max 0.31. |
| `gendo_alpha.png` | GendoAI (bottom-left) | 3 corpus files | POS 0.08 vs NEG max 0.32. |
| `zsky_alpha.png` | `MADE WITH zsky.ai` wordmark (ZSky AI, bottom-right) | 2 corpus files | POS 0.10 vs NEG max 0.27. |
| `chromastudio_alpha.png` | `ChromaStudio.ai` wordmark (ChromaStudio, bottom-right) | 2 corpus files | POS 0.11 vs NEG max 0.29. |
| `digenai_alpha.png` | `DIGENAI` wordmark (Digen AI, bottom-right) | 3 corpus files (one 2026-07-24 batch) | POS 0.16 vs NEG max 0.31. |
| `gendo_alpha.png` | `GendoAI` wordmark (Gendo, bottom-left) | 3 corpus files | POS 0.08 vs NEG max 0.32. |
| `xinghui_alpha.png` | 星绘AI生成 (parked before this set) | -- | prior parking, unchanged. |
| `qingyan_alpha.png` | 清言·AI生成 (parked before this set) | -- | prior parking, unchanged. |
| `hailuo_alpha.png` | Hailuo AI image wordmark (parked before this set; the VIDEO label is registered) | -- | prior parking, unchanged. |
+114
View File
@@ -0,0 +1,114 @@
# /// script
# requires-python = ">=3.11"
# dependencies = ["modal>=1.5"]
# ///
"""Extra high Chroma1 rungs (0.25-0.40) for the Google cohort fixtures.
The 0.08-0.20 ladder came back detected end-to-end on the first fixture, so
this appends the high rungs for every Google fixture without touching the
already-oracled lower outputs.
"""
from __future__ import annotations
import json
import math
from pathlib import Path
import modal
_ROOT = Path(__file__).resolve().parent.parent
app = modal.App("chroma-google-high-rungs")
image = (
modal.Image.debian_slim(python_version="3.12")
.pip_install(
"torch",
"diffusers>=0.38,<1",
"transformers>=4.53",
"accelerate",
"sentencepiece",
"protobuf",
"pillow",
"numpy",
)
.env({"HF_HOME": "/cache"})
)
cache = modal.Volume.from_name("chroma-hf-cache", create_if_missing=True)
CHROMA_MODEL_ID = "lodestones/Chroma1-HD"
RUNGS = (0.25, 0.30, 0.35, 0.40)
EFFECTIVE_STEPS = 4
_PROMPT = "high quality, sharp, detailed, faithful to the original"
_NEGATIVE = "blurry, lowres, distorted text, garbled text, artifacts"
FIXTURES = (
"Gemini_Generated_Image_633uuy633uuy633u.png",
"Gemini_Generated_Image_akdbeiakdbeiakdb.png",
"Gemini_Generated_Image_y48j3cy48j3cy48j.png",
"Gemini_Generated_Image_3mc4t93mc4t93mc4.png",
)
MAX_SIDE = 1536
@app.function(image=image, gpu="H100", volumes={"/cache": cache}, timeout=2400)
def run(payload: dict[str, bytes]) -> dict:
import io
import time
import torch
from diffusers import ChromaImg2ImgPipeline
from PIL import Image
chroma = ChromaImg2ImgPipeline.from_pretrained(CHROMA_MODEL_ID, torch_dtype=torch.bfloat16).to("cuda")
out: dict = {}
for name, data in payload.items():
source = Image.open(io.BytesIO(data)).convert("RGB")
if MAX_SIDE and max(source.size) > MAX_SIDE:
scale = MAX_SIDE / max(source.size)
source = source.resize(
(round(source.width * scale), round(source.height * scale)),
Image.Resampling.LANCZOS,
)
width = max(16, (source.width // 16) * 16)
height = max(16, (source.height // 16) * 16)
for strength in RUNGS:
started = time.perf_counter()
generator = torch.Generator(device="cpu").manual_seed(0)
image = chroma(
prompt=_PROMPT,
negative_prompt=_NEGATIVE,
image=source,
width=width,
height=height,
strength=strength,
num_inference_steps=max(1, math.ceil(EFFECTIVE_STEPS / strength)),
guidance_scale=5.0,
generator=generator,
).images[0]
if image.size != source.size:
image = image.resize(source.size, Image.Resampling.LANCZOS)
buffer = io.BytesIO()
image.save(buffer, format="PNG")
out[f"{Path(name).stem}/chroma_s{strength:.2f}"] = {
"png": buffer.getvalue(),
"seconds": round(time.perf_counter() - started, 2),
}
print(f"{name} s={strength} done", flush=True)
cache.commit()
return out
@app.local_entrypoint()
def main(out: str = "out/cohort-calibration/google") -> None:
payload = {name: (_ROOT / "data" / "synthid" / "originals" / name).read_bytes() for name in FIXTURES}
results = run.remote(payload)
timing = {}
for key, item in results.items():
destination = Path(out) / f"{key}.png"
destination.parent.mkdir(parents=True, exist_ok=True)
destination.write_bytes(item["png"])
timing[key] = item["seconds"]
print(json.dumps(timing, indent=2))
+109
View File
@@ -0,0 +1,109 @@
# /// script
# requires-python = ">=3.11"
# dependencies = ["modal>=1.5"]
# ///
"""Extra high Chroma1 rungs for the Meta Content Seal cohort.
The first calibration ladder topped out at 0.08 and the worst fixture
(gen_lighthouse_watercolor) was still detected there, so this appends 0.10,
0.12 and 0.15 for every Meta fixture without touching the already-oracled
lower rungs.
"""
from __future__ import annotations
import json
import math
from pathlib import Path
import modal
app = modal.App("chroma-meta-high-rungs")
_ROOT = Path(__file__).resolve().parent.parent
image = (
modal.Image.debian_slim(python_version="3.12")
.pip_install(
"torch",
"diffusers>=0.38,<1",
"transformers>=4.53",
"accelerate",
"sentencepiece",
"protobuf",
"pillow",
"numpy",
)
.env({"HF_HOME": "/cache"})
)
cache = modal.Volume.from_name("chroma-hf-cache", create_if_missing=True)
CHROMA_MODEL_ID = "lodestones/Chroma1-HD"
RUNGS = (0.17,)
EFFECTIVE_STEPS = 4
_PROMPT = "high quality, sharp, detailed, faithful to the original"
_NEGATIVE = "blurry, lowres, distorted text, garbled text, artifacts"
META_FIXTURES = (
"gen_fox_forest.webp",
"gen_text_poster.webp",
"gen_lighthouse_watercolor.webp",
"gen_night_city.webp",
"gen_studio_mug.webp",
)
@app.function(image=image, gpu="H100", volumes={"/cache": cache}, timeout=1800)
def run(payload: dict[str, bytes]) -> dict:
import io
import time
import torch
from diffusers import ChromaImg2ImgPipeline
from PIL import Image
chroma = ChromaImg2ImgPipeline.from_pretrained(CHROMA_MODEL_ID, torch_dtype=torch.bfloat16).to("cuda")
out: dict = {}
for name, data in payload.items():
source = Image.open(io.BytesIO(data)).convert("RGB")
width = max(16, (source.width // 16) * 16)
height = max(16, (source.height // 16) * 16)
for strength in RUNGS:
started = time.perf_counter()
generator = torch.Generator(device="cpu").manual_seed(0)
image = chroma(
prompt=_PROMPT,
negative_prompt=_NEGATIVE,
image=source,
width=width,
height=height,
strength=strength,
num_inference_steps=max(1, math.ceil(EFFECTIVE_STEPS / strength)),
guidance_scale=5.0,
generator=generator,
).images[0]
if image.size != source.size:
image = image.resize(source.size, Image.Resampling.LANCZOS)
buffer = io.BytesIO()
image.save(buffer, format="PNG")
out[f"{Path(name).stem}/chroma_s{strength:.3f}"] = {
"png": buffer.getvalue(),
"seconds": round(time.perf_counter() - started, 2),
}
print(f"{name} s={strength} done", flush=True)
cache.commit()
return out
@app.local_entrypoint()
def main(out: str = "out/cohort-calibration/meta") -> None:
payload = {name: (_ROOT / "data" / "contentseal" / "originals" / name).read_bytes() for name in META_FIXTURES}
results = run.remote(payload)
timing = {}
for key, item in results.items():
destination = Path(out) / f"{key}.png"
destination.parent.mkdir(parents=True, exist_ok=True)
destination.write_bytes(item["png"])
timing[key] = item["seconds"]
print(json.dumps(timing, indent=2))
+142
View File
@@ -0,0 +1,142 @@
# /// script
# requires-python = ">=3.11"
# dependencies = ["modal>=1.5"]
# ///
"""Chroma1 ladder for the Microsoft InvisMark cohort (three valid Paint sources)
plus the matching qwen-zimage production runs at the shipped 0.15 floor.
The sources are the same three Paint carriers behind the qwen cohort floor
(staged under out/cohort-calibration/microsoft/<label>/original.png with
pixel-identical metadata-stripped controls). They are real user uploads from
the raiw corpus and must never be committed.
"""
from __future__ import annotations
import json
import math
from pathlib import Path
import modal
_ROOT = Path(__file__).resolve().parent.parent
app = modal.App("chroma-microsoft-ladder")
image = (
modal.Image.debian_slim(python_version="3.12")
.pip_install(
"torch",
"torchvision>=0.20.0",
"diffusers>=0.38,<1",
"diffsynth>=2.0.17,<3",
"transformers>=4.53",
"accelerate",
"sentencepiece",
"protobuf",
"pillow",
"numpy",
"opencv-python-headless",
)
.env({"HF_HOME": "/cache"})
.add_local_dir(_ROOT / "src" / "remove_ai_watermarks", remote_path="/pkg/remove_ai_watermarks")
)
cache = modal.Volume.from_name("chroma-hf-cache", create_if_missing=True)
CHROMA_MODEL_ID = "lodestones/Chroma1-HD"
CHROMA_RUNGS = (0.04, 0.06, 0.08, 0.10, 0.125, 0.15, 0.20)
QWEN_FLOOR = 0.15
EFFECTIVE_STEPS = 4
_SOURCES = ("paint-1", "paint-2", "paint-3")
_PROMPT = "high quality, sharp, detailed, faithful to the original"
_NEGATIVE = "blurry, lowres, distorted text, garbled text, artifacts"
@app.function(image=image, gpu="H100", volumes={"/cache": cache}, timeout=2400)
def run(payload: dict[str, bytes]) -> dict:
import io
import sys
import time
import torch
from PIL import Image
sys.path.insert(0, "/pkg")
def chroma_call(chroma: object, source: Image.Image, strength: float, seed: int) -> Image.Image:
width = max(16, (source.width // 16) * 16)
height = max(16, (source.height // 16) * 16)
generator = torch.Generator(device="cpu").manual_seed(seed)
out = chroma(
prompt=_PROMPT,
negative_prompt=_NEGATIVE,
image=source,
width=width,
height=height,
strength=strength,
num_inference_steps=max(1, math.ceil(EFFECTIVE_STEPS / strength)),
guidance_scale=5.0,
generator=generator,
).images[0]
return out if out.size == source.size else out.resize(source.size, Image.Resampling.LANCZOS)
def png_bytes(image: Image.Image) -> bytes:
buffer = io.BytesIO()
image.save(buffer, format="PNG")
return buffer.getvalue()
from diffusers import ChromaImg2ImgPipeline
chroma = ChromaImg2ImgPipeline.from_pretrained(CHROMA_MODEL_ID, torch_dtype=torch.bfloat16).to("cuda")
from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline
qwen = QwenZImagePipeline(
device="cuda",
torch_dtype=torch.bfloat16,
keep_global_models_on_device=True,
keep_face_models_on_device=True,
)
out: dict = {}
for label in _SOURCES:
source = Image.open(io.BytesIO(payload[label])).convert("RGB")
rungs: dict = {}
for strength in CHROMA_RUNGS:
started = time.perf_counter()
image = chroma_call(chroma, source, strength, 0)
rungs[f"{strength:.4f}"] = {
"png": png_bytes(image),
"seconds": round(time.perf_counter() - started, 2),
}
print(f"{label} chroma s={strength} done", flush=True)
started = time.perf_counter()
qwen_out = qwen.run(source, strength=QWEN_FLOOR, seed=0)
out[label] = {
"chroma": rungs,
"qwen_seconds": round(time.perf_counter() - started, 2),
"qwen_png": png_bytes(qwen_out),
}
print(f"{label} qwen s={QWEN_FLOOR} done", flush=True)
cache.commit()
return out
@app.local_entrypoint()
def main(staged: str = "out/cohort-calibration/microsoft") -> None:
staged_root = Path(staged)
payload = {label: (staged_root / label / "original.png").read_bytes() for label in _SOURCES}
results = run.remote(payload)
summary: dict = {}
for label, entry in results.items():
target = staged_root / label
for strength, item in entry["chroma"].items():
(target / f"chroma_s{float(strength):.4f}.png").write_bytes(item["png"])
(target / "qwen_floor.png").write_bytes(entry["qwen_png"])
summary[label] = {
"chroma_seconds": {k: v["seconds"] for k, v in entry["chroma"].items()},
"qwen_seconds": entry["qwen_seconds"],
}
print(json.dumps(summary, indent=2))
+152
View File
@@ -0,0 +1,152 @@
# /// script
# requires-python = ">=3.11"
# dependencies = ["modal>=1.5"]
# ///
"""Pre-ship validation outputs for the chroma-zimage profile.
1. Seed sensitivity of the measured boundaries: chroma at the worst-fixture
first-clean rungs (google 633uuy @ 0.25, meta lighthouse @ 0.10) with
seeds 1 and 2, to check the boundary is not a seed-0 fluke.
2. Face-stage interaction: the FULL chroma-zimage pipeline (Chroma1 global at
the shipped Google floor 0.40, then the inherited YuNet/SAM/Z-Image face
stage) on the 5-face fixture 3mc4t9, to verify the composited face regions
do not re-introduce a detectable signal on top of the floor.
"""
from __future__ import annotations
import json
import math
import sys
from pathlib import Path
import modal
_ROOT = Path(__file__).resolve().parent.parent
app = modal.App("chroma-preship-validation")
image = (
modal.Image.debian_slim(python_version="3.12")
.pip_install(
"torch",
"torchvision>=0.20.0",
"diffusers>=0.38,<1",
"diffsynth>=2.0.17,<3",
"transformers>=4.53",
"accelerate",
"sentencepiece",
"protobuf",
"pillow",
"numpy",
"opencv-python-headless",
)
.env({"HF_HOME": "/cache"})
.add_local_dir(_ROOT / "src" / "remove_ai_watermarks", remote_path="/pkg/remove_ai_watermarks")
)
cache = modal.Volume.from_name("chroma-hf-cache", create_if_missing=True)
CHROMA_MODEL_ID = "lodestones/Chroma1-HD"
_PROMPT = "high quality, sharp, detailed, faithful to the original"
_NEGATIVE = "blurry, lowres, distorted text, garbled text, artifacts"
GOOGLE_FLOOR = 0.40
MAX_SIDE = 1536
@app.function(image=image, gpu="H100", volumes={"/cache": cache}, timeout=2400)
def run(payload: dict[str, bytes]) -> dict:
import io
import time
import torch
from diffusers import ChromaImg2ImgPipeline
from PIL import Image
sys.path.insert(0, "/pkg")
def chroma_call(pipe: object, source: Image.Image, strength: float, seed: int) -> Image.Image:
width = max(16, (source.width // 16) * 16)
height = max(16, (source.height // 16) * 16)
generator = torch.Generator(device="cpu").manual_seed(seed)
out = pipe(
prompt=_PROMPT,
negative_prompt=_NEGATIVE,
image=source,
width=width,
height=height,
strength=strength,
num_inference_steps=max(1, math.ceil(4 / strength)),
guidance_scale=5.0,
generator=generator,
).images[0]
return out if out.size == source.size else out.resize(source.size, Image.Resampling.LANCZOS)
def png_bytes(image: Image.Image) -> bytes:
buffer = io.BytesIO()
image.save(buffer, format="PNG")
return buffer.getvalue()
out: dict = {}
chroma = ChromaImg2ImgPipeline.from_pretrained(CHROMA_MODEL_ID, torch_dtype=torch.bfloat16).to("cuda")
# Seed sensitivity at the worst-fixture boundaries.
for key, strength in (("google_633uuy", 0.25), ("meta_lighthouse", 0.10)):
source = Image.open(io.BytesIO(payload[key])).convert("RGB")
if MAX_SIDE and max(source.size) > MAX_SIDE:
scale = MAX_SIDE / max(source.size)
source = source.resize(
(round(source.width * scale), round(source.height * scale)),
Image.Resampling.LANCZOS,
)
for seed in (1, 2):
started = time.perf_counter()
image = chroma_call(chroma, source, strength, seed)
out[f"seedprobe_{key}_s{strength}_seed{seed}"] = png_bytes(image)
print(f"{key} s={strength} seed={seed} ({time.perf_counter() - started:.1f}s)", flush=True)
del chroma
import gc
gc.collect()
torch.cuda.empty_cache()
# Full chroma-zimage pipeline on the 5-face fixture at the shipped floor.
from remove_ai_watermarks._internal.chroma_zimage_pipeline import ChromaZImagePipeline
full = ChromaZImagePipeline(
device="cuda",
torch_dtype=torch.bfloat16,
keep_global_models_on_device=True,
keep_face_models_on_device=True,
)
source = Image.open(io.BytesIO(payload["google_3mc4t9"])).convert("RGB")
if MAX_SIDE and max(source.size) > MAX_SIDE:
scale = MAX_SIDE / max(source.size)
source = source.resize(
(round(source.width * scale), round(source.height * scale)),
Image.Resampling.LANCZOS,
)
started = time.perf_counter()
result = full.run(source, strength=GOOGLE_FLOOR, seed=0)
out["fullpath_3mc4t9_s0.40"] = png_bytes(result)
print(f"full pipeline 3mc4t9 @ {GOOGLE_FLOOR} ({time.perf_counter() - started:.1f}s)", flush=True)
cache.commit()
return out
@app.local_entrypoint()
def main(out: str = "out/preship-chroma") -> None:
sources = {
"google_633uuy": _ROOT / "data/synthid/originals/Gemini_Generated_Image_633uuy633uuy633u.png",
"meta_lighthouse": _ROOT / "data/contentseal/originals/gen_lighthouse_watercolor.webp",
"google_3mc4t9": _ROOT / "data/synthid/originals/Gemini_Generated_Image_3mc4t93mc4t93mc4.png",
}
results = run.remote({k: p.read_bytes() for k, p in sources.items()})
root = Path(out)
root.mkdir(parents=True, exist_ok=True)
for key, png in results.items():
(root / f"{key}.png").write_bytes(png)
print(json.dumps(sorted(results), indent=2))
+188
View File
@@ -0,0 +1,188 @@
# /// script
# requires-python = ">=3.11"
# dependencies = ["modal>=1.5"]
# ///
"""Isolated Modal GPU prototype: does a low-strength Chroma1 img2img pass scrub
the invisible watermark while keeping text/faces faithful?
Chroma1 (lodestones, Apache-2.0) is the FLUX.1 architecture re-trained on open
data, and diffusers ships it a native strength img2img pipeline, which is exactly
what FLUX.2 lacks (see docs/synthid.md and issue #88). This script is the same
shape as qwen_scrub_prototype.py: standalone, not imported by the package and
not in uv.lock, so the locked environment never moves for an experiment.
Run from the repository root (weights ~34 GB download once into a Modal volume,
then every later run starts warm):
uvx modal run scripts/chroma_scrub_prototype.py \
--source data/synthid/originals/'ChatGPT Image May 31, 2026, 02_03_55 PM.png' \
--out out/chroma --strengths 0.05,0.1,0.15,0.2
A different GPU (default H100, ~$3.95/h; the 34 GB bf16 stack also fits
A100-80GB at ~$2.50/h):
CHROMA_GPU=A100-80GB uvx modal run scripts/chroma_scrub_prototype.py ...
What to look for (same protocol as the Qwen study, docs/synthid.md 2.2):
* SCRUB: the oracle no longer reports the watermark at some strength
(openai.com/verify for OpenAI carriers, the Gemini app for Google ones).
* FIDELITY: text and faces stay faithful at that same strength; score with
scripts/fidelity_metrics.py after the oracle pass.
The smallest strength that clears the oracle while keeping fidelity is the
number to compare against the qwen-zimage floors (OpenAI 0.07675 / Google 0.25).
"""
from __future__ import annotations
import os
from pathlib import Path
from typing import Any
import modal
_GPU = os.environ.get("CHROMA_GPU", "H100")
app = modal.App("chroma-scrub-proto")
image = (
modal.Image.debian_slim(python_version="3.12")
.pip_install(
"torch",
"diffusers>=0.38,<1",
"transformers>=4.53",
"accelerate",
"sentencepiece",
"protobuf",
"pillow",
"numpy",
)
.env({"HF_HOME": "/cache"})
)
# Model weights persist here across runs, so only the first container downloads.
cache = modal.Volume.from_name("chroma-hf-cache", create_if_missing=True)
CHROMA_MODEL_ID = "lodestones/Chroma1-HD"
# A neutral, faithful-regeneration prompt (scrub, not restyle); the same intent
# as the Qwen prototype and the shipped global prompts.
_PROMPT = "high quality, sharp, detailed, faithful to the original"
_NEGATIVE = "blurry, lowres, distorted text, garbled text, artifacts"
_pipe = None
def _load_pipe() -> Any:
global _pipe
if _pipe is not None:
return _pipe
import torch
from diffusers import ChromaImg2ImgPipeline
print(f"loading {CHROMA_MODEL_ID} (bf16)...", flush=True)
_pipe = ChromaImg2ImgPipeline.from_pretrained(CHROMA_MODEL_ID, torch_dtype=torch.bfloat16)
_pipe = _pipe.to("cuda")
return _pipe
def _psnr(a: Any, b: Any) -> float:
import numpy as np
x = np.asarray(a.convert("RGB"), dtype=np.float64)
y = np.asarray(b.convert("RGB"), dtype=np.float64)
mse = float(((x - y) ** 2).mean())
return 99.0 if mse <= 1e-12 else 10.0 * float(np.log10(255.0**2 / mse))
@app.function(
image=image,
gpu=_GPU,
volumes={"/cache": cache},
timeout=2400,
)
def scrub(
source_bytes: bytes,
strengths: list[float],
steps: int,
guidance: float,
seed: int,
max_side: int,
) -> list[tuple[str, bytes, float]]:
"""Sweep one fixture over the strengths; return name, PNG bytes, and PSNR."""
import io
import torch
from PIL import Image
pipe = _load_pipe()
source = Image.open(io.BytesIO(source_bytes)).convert("RGB")
if max_side and max(source.size) > max_side:
scale = max_side / max(source.size)
source = source.resize(
(round(source.width * scale), round(source.height * scale)),
Image.Resampling.LANCZOS,
)
# The Flux-family img2img pipelines default to 1024x1024 when height/width
# are not passed, silently discarding the input aspect. Floor to the /16
# latent-patch grid, the same rule the shipped profiles use.
width = max(16, (source.width // 16) * 16)
height = max(16, (source.height // 16) * 16)
results: list[tuple[str, bytes, float]] = []
for strength in strengths:
generator = torch.Generator(device="cpu").manual_seed(seed)
print(f"strength={strength:.2f} on {width}x{height}...", flush=True)
output = pipe(
prompt=_PROMPT,
negative_prompt=_NEGATIVE,
image=source,
width=width,
height=height,
strength=strength,
num_inference_steps=steps,
guidance_scale=guidance,
generator=generator,
).images[0]
if output.size != source.size:
output = output.resize(source.size, Image.Resampling.LANCZOS)
buffer = io.BytesIO()
output.save(buffer, format="PNG")
results.append((f"s{strength:.2f}", buffer.getvalue(), _psnr(source, output)))
print(f" psnr vs source: {results[-1][2]:.2f} dB", flush=True)
cache.commit()
return results
@app.local_entrypoint()
def main(
source: str,
strengths: str = "0.05,0.1,0.15,0.2",
out: str = "out/chroma",
steps: int = 35,
guidance: float = 5.0,
seed: int = 0,
max_side: int = 1536,
) -> None:
"""Sweep Chroma1 img2img strength over SOURCE and save one output per strength."""
source_path = Path(source)
if not source_path.exists():
raise SystemExit(f"source not found: {source_path}")
values = [float(s) for s in strengths.split(",") if s.strip()]
stem = source_path.stem
print(f"gpu={_GPU} model={CHROMA_MODEL_ID} steps={steps} guidance={guidance} seed={seed}")
outputs = scrub.remote(source_path.read_bytes(), values, steps, guidance, seed, max_side)
out_dir = Path(out)
out_dir.mkdir(parents=True, exist_ok=True)
for name, payload, psnr in outputs:
destination = out_dir / f"{stem}_chroma_{name}.png"
destination.write_bytes(payload)
print(f"saved {destination} (psnr {psnr:.2f} dB vs source)")
print(
"\nDone. Submit each output to the matching oracle (openai.com/verify for "
"OpenAI carriers, the Gemini app for Google ones). The smallest strength "
"that clears the oracle while keeping fidelity is the Chroma1 floor."
)
+272
View File
@@ -0,0 +1,272 @@
# /// script
# requires-python = ">=3.11"
# dependencies = ["modal>=1.5"]
# ///
"""Chroma1 floor calibration across the committed vendor cohorts, plus the
matching qwen-zimage production runs at each cohort's shipped floor, for the
engine quality/price comparison.
Every Chroma1 run spends a fixed 4 effective denoising steps (the requested
count is scaled up, because Diffusers truncates the step COUNT), matching the
semantics the shipped profiles use. Google fixtures are capped to a 1536 px
long side, the same practical path the qwen floors were certified on; Meta
sources run at native size like their cohort measurement.
The Microsoft InvisMark cohort is NOT here: its carrier fixtures are not
committed to this repository (they were private sources), so no Chroma1 floor
can be measured for it until carriers exist. Until then a chroma profile would
treat Microsoft like sdxl-zimage treats Meta today: fall back to the
conservative unknown floor.
Run from the repository root (both stacks are warm in the shared volume):
uvx modal run scripts/engine_cohort_calibration.py
Outputs land in out/cohort-calibration/<vendor>/<fixture>/ . Oracle each
fixture's rungs in ASCENDING strength order and stop at the first clean one:
that pair brackets the floor with the fewest checks.
"""
from __future__ import annotations
import json
import math
import sys
from pathlib import Path
from typing import Any
import modal
_GPU = "H100"
app = modal.App("engine-cohort-calibration")
_ROOT = Path(__file__).resolve().parent.parent
image = (
modal.Image.debian_slim(python_version="3.12")
.pip_install(
"torch",
"torchvision>=0.20.0",
"diffusers>=0.38,<1",
"diffsynth>=2.0.17,<3",
"transformers>=4.53",
"accelerate",
"sentencepiece",
"protobuf",
"pillow",
"numpy",
"opencv-python-headless",
)
.env({"HF_HOME": "/cache"})
.add_local_dir(_ROOT / "src" / "remove_ai_watermarks", remote_path="/pkg/remove_ai_watermarks")
)
cache = modal.Volume.from_name("chroma-hf-cache", create_if_missing=True)
CHROMA_MODEL_ID = "lodestones/Chroma1-HD"
QWEN_OPENAI_FLOOR = 0.07675
QWEN_GOOGLE_FLOOR = 0.27
QWEN_META_FLOOR = 0.1
CHROMA_EFFECTIVE_STEPS = 4
_PROMPT = "high quality, sharp, detailed, faithful to the original"
_NEGATIVE = "blurry, lowres, distorted text, garbled text, artifacts"
# vendor -> fixture -> {chroma rungs, qwen floor strength, max side}
FIXTURES: dict[str, dict[str, dict]] = {
"openai": {
"ChatGPT Image May 30, 2026, 10_31_08 AM.png": { # June study 9-face grid, photoreal
"chroma": (0.04, 0.05, 0.06, 0.075, 0.09, 0.12),
"qwen": QWEN_OPENAI_FLOOR,
"max_side": 1536,
},
"ChatGPT Image May 31, 2026, 02_02_23 PM.png": { # full-pipeline quality fixture
"chroma": (0.04, 0.05, 0.06, 0.075, 0.09),
"qwen": QWEN_OPENAI_FLOOR,
"max_side": 1536,
},
},
"google": {
"Gemini_Generated_Image_633uuy633uuy633u.png": {
"chroma": (0.08, 0.10, 0.12, 0.15, 0.20),
"qwen": QWEN_GOOGLE_FLOOR,
"max_side": 1536,
},
"Gemini_Generated_Image_akdbeiakdbeiakdb.png": {
"chroma": (0.08, 0.10, 0.12, 0.15, 0.20),
"qwen": QWEN_GOOGLE_FLOOR,
"max_side": 1536,
},
"Gemini_Generated_Image_y48j3cy48j3cy48j.png": {
"chroma": (0.08, 0.10, 0.12, 0.15, 0.20),
"qwen": QWEN_GOOGLE_FLOOR,
"max_side": 1536,
},
"Gemini_Generated_Image_3mc4t93mc4t93mc4.png": { # the 18-face portrait grid
"chroma": (0.08, 0.10, 0.12, 0.15, 0.20),
"qwen": QWEN_GOOGLE_FLOOR,
"max_side": 1536,
},
},
"meta": {
name: {"chroma": (0.03, 0.045, 0.06, 0.08), "qwen": QWEN_META_FLOOR, "max_side": 0}
for name in (
"gen_fox_forest.webp",
"gen_text_poster.webp",
"gen_lighthouse_watercolor.webp",
"gen_night_city.webp",
"gen_studio_mug.webp",
)
},
}
# Seed sensitivity of the already-bracketed OpenAI typography fixture: the
# boundary was (0.05, 0.06] at seed 0; probe the same rungs at two more seeds.
SEED_PROBE = {
"source": "data/synthid/originals/ChatGPT Image May 31, 2026, 02_03_55 PM.png",
"rungs": (0.05, 0.06),
"seeds": (1, 2),
"max_side": 1536,
}
def requested_steps(effective_steps: int, strength: float) -> int:
return max(1, math.ceil(effective_steps / max(float(strength), 1e-6)))
@app.function(image=image, gpu=_GPU, volumes={"/cache": cache}, timeout=3000)
def calibrate(payload: dict[str, bytes]) -> dict:
"""Run the chroma ladders, the qwen floors, and the seed probe."""
import io
import time
import torch
from PIL import Image
sys.path.insert(0, "/pkg")
report: dict = {}
def load(path_key: str, max_side: int) -> Image.Image:
source = Image.open(io.BytesIO(payload[path_key])).convert("RGB")
if max_side and max(source.size) > max_side:
scale = max_side / max(source.size)
source = source.resize(
(round(source.width * scale), round(source.height * scale)),
Image.Resampling.LANCZOS,
)
return source
def chroma_call(chroma: Any, source: Image.Image, strength: float, seed: int) -> Image.Image:
width = max(16, (source.width // 16) * 16)
height = max(16, (source.height // 16) * 16)
generator = torch.Generator(device="cpu").manual_seed(seed)
out = chroma(
prompt=_PROMPT,
negative_prompt=_NEGATIVE,
image=source,
width=width,
height=height,
strength=strength,
num_inference_steps=requested_steps(CHROMA_EFFECTIVE_STEPS, strength),
guidance_scale=5.0,
generator=generator,
).images[0]
return out if out.size == source.size else out.resize(source.size, Image.Resampling.LANCZOS)
def png_bytes(image: Image.Image) -> bytes:
buffer = io.BytesIO()
image.save(buffer, format="PNG")
return buffer.getvalue()
from diffusers import ChromaImg2ImgPipeline
print(f"[chroma] loading {CHROMA_MODEL_ID}...", flush=True)
chroma = ChromaImg2ImgPipeline.from_pretrained(CHROMA_MODEL_ID, torch_dtype=torch.bfloat16).to("cuda")
from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline
print("[qwen] loading the production two-stage pipeline...", flush=True)
qwen = QwenZImagePipeline(
device="cuda",
torch_dtype=torch.bfloat16,
keep_global_models_on_device=True,
keep_face_models_on_device=True,
)
for vendor, fixtures in FIXTURES.items():
report[vendor] = {}
for name, spec in fixtures.items():
print(f"[{vendor}] {name}", flush=True)
source = load(name, spec["max_side"])
entry: dict = {"chroma": {}, "qwen_seconds": None}
for strength in spec["chroma"]:
started = time.perf_counter()
out = chroma_call(chroma, source, strength, 0)
seconds = time.perf_counter() - started
entry["chroma"][f"{strength:.4f}"] = {
"png": png_bytes(out),
"seconds": round(seconds, 2),
}
print(f" chroma s={strength:.3f} {seconds:.1f}s", flush=True)
started = time.perf_counter()
qwen_out = qwen.run(source, strength=spec["qwen"], seed=0)
entry["qwen_seconds"] = round(time.perf_counter() - started, 2)
entry["qwen_png"] = png_bytes(qwen_out)
print(f" qwen s={spec['qwen']} {entry['qwen_seconds']}s", flush=True)
report[vendor][name] = entry
print("[seeds] openai typography boundary probe...", flush=True)
seed_source = load(SEED_PROBE["source"], SEED_PROBE["max_side"])
report["seed_probe"] = {}
for seed in SEED_PROBE["seeds"]:
for strength in SEED_PROBE["rungs"]:
out = chroma_call(chroma, seed_source, strength, seed)
report["seed_probe"][f"s{strength:.2f}_seed{seed}"] = png_bytes(out)
print(f" seed={seed} s={strength:.2f}", flush=True)
cache.commit()
return report
@app.local_entrypoint()
def main(out: str = "out/cohort-calibration") -> None:
payload: dict[str, bytes] = {}
for vendor, fixtures in FIXTURES.items():
for name in fixtures:
if vendor == "meta":
path = _ROOT / "data" / "contentseal" / "originals" / name
else:
path = _ROOT / "data" / "synthid" / "originals" / name
if not path.exists():
raise SystemExit(f"fixture not found: {path}")
payload[name] = path.read_bytes()
payload[SEED_PROBE["source"]] = (_ROOT / SEED_PROBE["source"]).read_bytes()
report = calibrate.remote(payload)
out_root = Path(out)
summary: dict = {}
for vendor, fixtures in report.items():
if vendor == "seed_probe":
for key, png in fixtures.items():
destination = out_root / "seed-probe" / f"chroma_{key}.png"
destination.parent.mkdir(parents=True, exist_ok=True)
destination.write_bytes(png)
continue
summary[vendor] = {}
for name, entry in fixtures.items():
fixture_dir = out_root / vendor / Path(name).stem
fixture_dir.mkdir(parents=True, exist_ok=True)
(fixture_dir / "qwen_floor.png").write_bytes(entry["qwen_png"])
rungs = {}
for strength, run in entry["chroma"].items():
(fixture_dir / f"chroma_s{float(strength):.3f}.png").write_bytes(run["png"])
rungs[strength] = run["seconds"]
summary[vendor][name] = {
"chroma_rungs_seconds": rungs,
"qwen_seconds": entry["qwen_seconds"],
}
(out_root / "timing.json").write_text(json.dumps(summary, indent=2))
print(json.dumps(summary, indent=2))
print(f"\noutputs in {out_root}. Oracle each fixture ascending, stop at first clean.")
+249
View File
@@ -0,0 +1,249 @@
# /// script
# requires-python = ">=3.11"
# dependencies = ["modal>=1.5"]
# ///
"""Head-to-head probe: qwen-zimage (the shipped production path) vs Chroma1
img2img on one fixture -- warm per-image wall time, VRAM peak, and outputs for
the local fidelity metrics, so quality and price can be compared on the same
pixels.
qwen-zimage runs through the package's own ``QwenZImagePipeline`` (DiffSynth,
Lightning LoRA, Canny ControlNet, face stage) at its certified OpenAI operating
point (0.07675) with both stacks resident, exactly as the deployed worker does.
Chroma1 runs the same neutral prompt with the SDXL-style step compensation
(``requested_steps`` semantics: always spend 4 effective denoising steps), so a
low strength is not silently 1 step of 35.
Standalone like the other prototypes: not imported by the package, not in
uv.lock. The package source is mounted at runtime, never installed.
Run from the repository root (first run downloads both stacks, ~100 GB total,
into the shared volume; later runs start warm):
uvx modal run scripts/engine_quality_price_probe.py \
--source data/synthid/originals/'ChatGPT Image May 31, 2026, 02_03_55 PM.png'
Oracle the Chroma ladder outputs afterwards (openai.com/verify) and score both
engines' floor outputs with scripts/fidelity_metrics.py locally.
"""
from __future__ import annotations
import json
import sys
from pathlib import Path
from typing import Any
import modal
_GPU = "H100"
# Verified 2026-08-29 from modal.com/pricing.
_H100_PER_HOUR = 3.95
_A100_80_PER_HOUR = 2.50
app = modal.App("engine-quality-price-probe")
_ROOT = Path(__file__).resolve().parent.parent
image = (
modal.Image.debian_slim(python_version="3.12")
.pip_install(
"torch",
"torchvision>=0.20.0",
"diffusers>=0.38,<1",
"diffsynth>=2.0.17,<3",
"transformers>=4.53",
"accelerate",
"sentencepiece",
"protobuf",
"pillow",
"numpy",
"opencv-python-headless",
)
.env({"HF_HOME": "/cache"})
.add_local_dir(_ROOT / "src" / "remove_ai_watermarks", remote_path="/pkg/remove_ai_watermarks")
)
# Carries both engines' weights; named for the first prototype that filled it.
cache = modal.Volume.from_name("chroma-hf-cache", create_if_missing=True)
QWEN_OPENAI_FLOOR = 0.07675 # watermark_profiles.QWEN_ZIMAGE_OPENAI_STRENGTH
CHROMA_MODEL_ID = "lodestones/Chroma1-HD"
CHROMA_STRENGTHS = (0.05, 0.06, 0.075, 0.09, 0.10)
CHROMA_EFFECTIVE_STEPS = 4
_PROMPT = "high quality, sharp, detailed, faithful to the original"
_NEGATIVE = "blurry, lowres, distorted text, garbled text, artifacts"
def requested_steps(effective_steps: int, strength: float) -> int:
"""Diffusers truncates the step COUNT; spend effective_steps regardless."""
import math
return max(1, math.ceil(effective_steps / max(float(strength), 1e-6)))
def _psnr(a: Any, b: Any) -> float:
import numpy as np
x = np.asarray(a.convert("RGB"), dtype=np.float64)
y = np.asarray(b.convert("RGB"), dtype=np.float64)
mse = float(((x - y) ** 2).mean())
return 99.0 if mse <= 1e-12 else 10.0 * float(np.log10(255.0**2 / mse))
def _png(image: Any) -> bytes:
import io
buffer = io.BytesIO()
image.save(buffer, format="PNG")
return buffer.getvalue()
@app.function(image=image, gpu=_GPU, volumes={"/cache": cache}, timeout=3000)
def probe(source_bytes: bytes, max_side: int) -> dict:
"""Run both engines; return outputs, timings, and VRAM peaks."""
import gc
import io
import time
import torch
from PIL import Image
sys.path.insert(0, "/pkg")
source = Image.open(io.BytesIO(source_bytes)).convert("RGB")
if max_side and max(source.size) > max_side:
scale = max_side / max(source.size)
source = source.resize(
(round(source.width * scale), round(source.height * scale)),
Image.Resampling.LANCZOS,
)
report: dict = {"source_size": source.size, "engines": {}}
# --- qwen-zimage, the exact production path ---------------------------------
print("[qwen] loading the production two-stage pipeline (resident stacks)...", flush=True)
from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline
qwen = QwenZImagePipeline(
device="cuda",
torch_dtype=torch.bfloat16,
keep_global_models_on_device=True,
keep_face_models_on_device=True,
)
print(f"[qwen] warmup pass at strength={QWEN_OPENAI_FLOOR}...", flush=True)
qwen.run(source, strength=QWEN_OPENAI_FLOOR, seed=0)
torch.cuda.reset_peak_memory_stats()
started = time.perf_counter()
qwen_out = qwen.run(source, strength=QWEN_OPENAI_FLOOR, seed=0)
qwen_seconds = time.perf_counter() - started
qwen_vram_gib = torch.cuda.max_memory_allocated() / (1024**3)
report["engines"]["qwen-zimage"] = {
"strength": QWEN_OPENAI_FLOOR,
"warm_seconds": round(qwen_seconds, 2),
"vram_peak_gib": round(qwen_vram_gib, 1),
"psnr": round(_psnr(source, qwen_out), 2),
"png": _png(qwen_out),
}
qwen_report = report["engines"]["qwen-zimage"]
print(
f"[qwen] warm {qwen_seconds:.2f}s, vram peak {qwen_vram_gib:.1f} GiB, psnr {qwen_report['psnr']:.2f} dB",
flush=True,
)
del qwen
gc.collect()
torch.cuda.empty_cache()
# --- Chroma1, compensated-step ladder ---------------------------------------
from diffusers import ChromaImg2ImgPipeline
print(f"[chroma] loading {CHROMA_MODEL_ID} (bf16)...", flush=True)
chroma = ChromaImg2ImgPipeline.from_pretrained(CHROMA_MODEL_ID, torch_dtype=torch.bfloat16).to("cuda")
width = max(16, (source.width // 16) * 16)
height = max(16, (source.height // 16) * 16)
ladder: list[dict] = []
for strength in CHROMA_STRENGTHS:
steps = requested_steps(CHROMA_EFFECTIVE_STEPS, strength)
torch.cuda.reset_peak_memory_stats()
started = time.perf_counter()
generator = torch.Generator(device="cpu").manual_seed(0)
out = chroma(
prompt=_PROMPT,
negative_prompt=_NEGATIVE,
image=source,
width=width,
height=height,
strength=strength,
num_inference_steps=steps,
guidance_scale=5.0,
generator=generator,
).images[0]
seconds = time.perf_counter() - started
vram_gib = torch.cuda.max_memory_allocated() / (1024**3)
if out.size != source.size:
out = out.resize(source.size, Image.Resampling.LANCZOS)
ladder.append(
{
"strength": strength,
"requested_steps": steps,
"warm_seconds": round(seconds, 2),
"vram_peak_gib": round(vram_gib, 1),
"psnr": round(_psnr(source, out), 2),
"png": _png(out),
}
)
print(
f"[chroma] s={strength:.3f} steps={steps}/{steps} warm {seconds:.2f}s "
f"vram {vram_gib:.1f} GiB psnr {ladder[-1]['psnr']:.2f} dB",
flush=True,
)
report["engines"]["chroma1"] = {"ladder": ladder}
cache.commit()
return report
@app.local_entrypoint()
def main(
source: str,
out: str = "out/engine-comparison",
max_side: int = 1536,
) -> None:
source_path = Path(source)
if not source_path.exists():
raise SystemExit(f"source not found: {source_path}")
report = probe.remote(source_path.read_bytes(), max_side)
out_dir = Path(out) / source_path.stem
out_dir.mkdir(parents=True, exist_ok=True)
for engine, payload in report["engines"].items():
if engine == "chroma1":
for run in payload["ladder"]:
(out_dir / f"chroma_s{run['strength']:.3f}.png").write_bytes(run["png"])
else:
(out_dir / f"{engine}_s{payload['strength']}.png").write_bytes(payload["png"])
summary = {
"source": str(source_path),
"source_size": report["source_size"],
"gpu": _GPU,
"usd_per_hour": {"H100": _H100_PER_HOUR, "A100-80GB": _A100_80_PER_HOUR},
"engines": {
engine: (
{k: v for k, v in payload.items() if k != "png"}
if engine != "chroma1"
else {"ladder": [{k: v for k, v in run.items() if k != "png"} for run in payload["ladder"]]}
)
for engine, payload in report["engines"].items()
},
}
(out_dir / "timing.json").write_text(json.dumps(summary, indent=2))
print(json.dumps(summary, indent=2))
print(f"\noutputs in {out_dir}; oracle the chroma ladder, then score locally with:")
print(f" uv run scripts/fidelity_metrics.py compare --original '{source}' \\")
print(f" --variant qwen={out_dir / 'qwen-zimage_s0.0768.png'} --variant chroma={out_dir}/chroma_sXXX.png")
+76
View File
@@ -0,0 +1,76 @@
#!/usr/bin/env python3
"""One-shot OpenAI Content Provenance oracle check per file.
Usage: python3 scripts/openai_provenance_check.py FILE [FILE...]
Prints ``basename<TAB>synthid_outcome`` per file (DETECTED / NOT_DETECTED).
Reads OPENAI_API_KEY from the environment or the main checkout's .env.
This is the documented programmatic twin of openai.com/verify and reports the
SynthID and C2PA outcomes separately, which the web page collapses.
"""
from __future__ import annotations
import json
import os
import pathlib
import sys
import urllib.request
import uuid
API_URL = "https://api.openai.com/v1/content_provenance_checks"
def _api_key() -> str:
key = os.environ.get("OPENAI_API_KEY")
if key:
return key
env_file = pathlib.Path.home() / "Documents/GitHub/remove-ai-watermarks/.env"
if env_file.exists():
for line in env_file.read_text().splitlines():
if line.startswith("OPENAI_API_KEY="):
return line.split("=", 1)[1].strip()
raise SystemExit("no OPENAI_API_KEY in env and no main-checkout .env")
def check(path: pathlib.Path) -> str:
boundary = uuid.uuid4().hex
payload = path.read_bytes()
body = b"".join(
(
f"--{boundary}\r\n"
f'Content-Disposition: form-data; name="file"; filename="{path.name}"\r\n'
"Content-Type: image/png\r\n\r\n".encode(),
payload,
f"\r\n--{boundary}--\r\n".encode(),
)
)
request = urllib.request.Request(
API_URL,
data=body,
headers={
"Authorization": f"Bearer {_api_key()}",
"Content-Type": f"multipart/form-data; boundary={boundary}",
},
)
with urllib.request.urlopen(request, timeout=120) as response:
data = json.load(response)
for entry in data.get("results", []):
if entry.get("type") == "synthid":
return str(entry.get("outcome", "?")).upper()
return "NO_SYNTHID_ENTRY"
def main() -> None:
for name in sys.argv[1:]:
path = pathlib.Path(name)
try:
outcome = check(path)
except Exception as exc:
outcome = f"ERROR:{exc}"
print(f"{path.name}\t{outcome}", flush=True)
if __name__ == "__main__":
main()
+1 -1
View File
@@ -34,7 +34,7 @@ DATA SAFETY
Read-only over a local dataset. Writes only the given output path, which belongs
outside the repository. Resumable: rerunning skips files already recorded.
uv run python scripts/record_parity_audit.py data/spaces/originals .local-eval/parity.jsonl
uv run python scripts/record_parity_audit.py /path/to/dataset .local-eval/parity.jsonl
uv run python scripts/record_parity_audit.py <dataset> <out> --baseline .local-eval/previous.jsonl
"""
+1 -1
View File
@@ -6,7 +6,7 @@ the committed asset is a font-rendered binary SILHOUETTE (mark shape only, zero
content), used ONLY to (a) detect the pill by edge-NCC in the top-left corner and
(b) build the inpaint mask. It is NOT an alpha map -- removal quality comes from the
inpaint backend (MI-GAN/cv2), so the silhouette need not be pixel-accurate, and the
synthetic render keeps corpus/user content out of the tracked repo (data-safety).
synthetic render keeps the asset deterministic and reproducible.
The calibrated edge-NCC threshold is 0.22 in the top-left ROI. Re-run to regenerate
the asset: uv run python scripts/render_pill_silhouette.py
+2 -2
View File
@@ -75,8 +75,8 @@ MARKS = {
"chromastudio_alpha.png": "ChromaStudio.ai",
"digenai_alpha.png": "DIGENAI",
"gendo_alpha.png": "GendoAI",
# CapCut's Chinese sibling, JianYing, stamps 剪映AI bottom-right (the
# international CapCut pill sits top-left).
# Jianying / 剪映, the China product in the CapCut family, stamps 剪映AI
# bottom-right (the international CapCut pill sits top-left).
"jianying_alpha.png": "剪映AI",
}
_REGISTERED = {f"{key}_alpha.png" for key in mark_keys()} & MARKS.keys()
+4 -4
View File
@@ -5,10 +5,8 @@ sample of everything it supports and a canary test can hold both sides to it
(mark registered without example; engine regressed on its canonical example).
Every example is SYNTHETIC: a deterministic generated base photo with the mark's
own committed silhouette composited at the engine's measured geometry. User
uploads never enter the repository (data/spaces stays out of git), and no vendor
asset is copied -- the silhouettes are the same font-rendered templates the
detectors match against.
own committed silhouette composited at the engine's measured geometry. The
silhouettes are repository-owned font renders, not copied vendor raster assets.
Regenerate with:
uv run python scripts/render_visible_examples.py
@@ -189,6 +187,8 @@ def _video_mark_frame(key: str, w: int, h: int) -> np.ndarray:
return _composite_light(base, cv2.resize(tmpl, (tw, th)).astype(np.float32) / 255.0, x, y, 250)
elif key == "seedance":
tmpl, scale, x, y = templates["seedance"], 0.095, int(w * 0.74), int(h * 0.80)
elif key == "doubao":
tmpl, scale, x, y = templates["doubao"], 0.048, int(w * 0.72), int(h * 0.82)
elif key == "dola":
tmpl, scale, x, y = templates["dola"], 0.036, int(w * 0.70), int(h * 0.88)
elif key == "hailuo":
+1 -1
View File
@@ -32,7 +32,7 @@ _os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
_warnings.filterwarnings("ignore", message=r".*ImageProcessorFast.*")
__version__ = "0.34.1"
__version__ = "0.35.1"
__all__ = [
"BatchSummary",
@@ -0,0 +1,128 @@
"""The chroma-zimage recipe with a Chroma1 global stage.
Only the global regeneration model changes. The face stage is inherited verbatim
from :class:`TwoStageZImagePipeline` -- same YuNet detection, same SAM masks, same
Z-Image Turbo repair of the original crops, same feathered compositing -- so a
change there cannot silently diverge between the profiles.
Pieces bound to this architecture: the diffusers ``ChromaImg2ImgPipeline``
(no DiffSynth here), the neutral faithful-regeneration prompt the floors were
calibrated with, guidance 5.0, and the step-count compensation -- diffusers
truncates the step COUNT (``int(steps * strength)``), so the requested count is
scaled to always spend four effective denoising steps, the same semantics as
``sdxl_zimage_pipeline.requested_steps``. Strength is bound to it too: the flat
vendor floors in ``watermark_profiles`` come from the 2026-08-29/30 oracle
calibration (docs/chroma1-engine-research.md) and do not transfer to another
prompt, guidance, or effective step count.
No Canny conditioning: the floors were measured on a plain strength-controlled
img2img pass, so that is the calibrated path.
"""
# Diffusers and torch expose mostly untyped tensor APIs. Keep the relaxation local
# to this optional ML boundary.
# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportAttributeAccessIssue=false, reportPrivateUsage=false, reportPrivateImportUsage=false
from __future__ import annotations
import logging
import math
from dataclasses import dataclass
from typing import Any, ClassVar
from PIL import Image
from remove_ai_watermarks._internal.two_stage_pipeline import TwoStageZImagePipeline
log = logging.getLogger(__name__)
CHROMA_MODEL_ID = "lodestones/Chroma1-HD"
# The Chroma1 global stage spends four effective denoising steps at every
# strength. The calibration ladder (docs/chroma1-engine-research.md) held this
# fixed; a different count is a different engine.
CHROMA_STEPS = 4
CHROMA_GUIDANCE = 5.0
# The neutral faithful-regeneration prompt the floors were measured with. It is
# deliberately NOT the canny-stage prompt the Qwen/SDXL profiles use: those were
# never calibrated against Chroma1, and swapping prompts silently moves the
# oracle boundaries.
CHROMA_PROMPT = "high quality, sharp, detailed, faithful to the original"
CHROMA_NEGATIVE = "blurry, lowres, distorted text, garbled text, artifacts"
# The latent grid the calibration floors were generated on. The source is floored
# to this grid, generated, then resized back to the exact input size -- the same
# shape the prototype scripts used.
_LATENT_GRID = 16
def requested_steps(effective_steps: int, strength: float) -> int:
"""Spend ``effective_steps`` regardless of Diffusers' step-count truncation.
Diffusers img2img truncates the step COUNT (``init_timestep = int(steps *
strength)``), so at the floors this profile uses a naive request would run
zero or one steps. Scale the request so the effective count is exact.
"""
return max(1, math.ceil(effective_steps / max(float(strength), 1e-6)))
def chroma_target_size(width: int, height: int) -> tuple[int, int]:
"""Floor dimensions to the latent grid without changing aspect."""
return max(_LATENT_GRID, (width // _LATENT_GRID) * _LATENT_GRID), max(
_LATENT_GRID, (height // _LATENT_GRID) * _LATENT_GRID
)
@dataclass
class ChromaZImagePipeline(TwoStageZImagePipeline):
"""Lazy runtime for the Chroma1 global stage plus the inherited face stage."""
profile_name: ClassVar[str] = "chroma-zimage"
def __post_init__(self) -> None:
super().__post_init__()
self._chroma_pipe: Any = None
def _load_global(self) -> Any:
if self._chroma_pipe is not None:
return self._chroma_pipe
self._require_cuda()
try:
import torch
from diffusers import ChromaImg2ImgPipeline
except ImportError as exc:
raise ImportError(
"The chroma-zimage pipeline needs the optional dependency group. "
"Install: pip install 'remove-ai-watermarks[qwen-zimage]'"
) from exc
self._progress("Loading Chroma1-HD (bf16)...")
token = {"token": self.hf_token} if self.hf_token else {}
pipe = ChromaImg2ImgPipeline.from_pretrained(CHROMA_MODEL_ID, torch_dtype=torch.bfloat16, **token)
pipe = pipe.to(self.device)
self._chroma_pipe = pipe
return pipe
def _run_global(self, image: Image.Image, strength: float, seed: int | None) -> Image.Image:
import torch
pipe = self._load_global()
target = chroma_target_size(image.width, image.height)
prepared = image if image.size == target else image.resize(target, Image.Resampling.LANCZOS)
steps = requested_steps(CHROMA_STEPS, strength)
self._progress(f"Running Chroma1 pass: strength={strength:.4f}, steps={CHROMA_STEPS} of {steps}...")
generator = torch.Generator(device=self.device).manual_seed(seed) if seed is not None else None
result = pipe(
prompt=CHROMA_PROMPT,
negative_prompt=CHROMA_NEGATIVE,
image=prepared,
width=target[0],
height=target[1],
strength=float(strength),
num_inference_steps=steps,
guidance_scale=CHROMA_GUIDANCE,
generator=generator,
).images[0]
if result.size != image.size:
result = result.resize(image.size, Image.Resampling.LANCZOS)
return result.convert("RGB")
File diff suppressed because it is too large. Load diff
@@ -1,9 +1,9 @@
"""The qwen-zimage recipe with an SDXL global stage.
"""The sdxl-zimage recipe with an SDXL global stage.
Only the global regeneration model changes. The face stage is inherited verbatim
from :class:`QwenZImagePipeline` -- same YuNet detection, same SAM masks, same
from :class:`TwoStageZImagePipeline` -- same YuNet detection, same SAM masks, same
Z-Image Turbo repair of the original crops, same feathered compositing -- so a
change there cannot silently diverge between the two profiles.
change there cannot silently diverge between the profiles.
Three pieces cannot be shared, because they are bound to the architecture: the
ControlNet, the four-step distillation LoRA, and the sampler. Strength is bound to
@@ -19,13 +19,14 @@ from __future__ import annotations
import logging
import math
from dataclasses import dataclass
from typing import Any
from typing import Any, ClassVar
from PIL import Image
from remove_ai_watermarks._internal.qwen_zimage_pipeline import (
GLOBAL_STEPS,
QwenZImagePipeline,
from remove_ai_watermarks._internal.two_stage_pipeline import (
_GLOBAL_NEGATIVE,
_GLOBAL_PROMPT,
TwoStageZImagePipeline,
build_canny_control_image,
)
from remove_ai_watermarks._internal.watermark_profiles import (
@@ -40,6 +41,10 @@ log = logging.getLogger(__name__)
SDXL_VAE_MODEL_ID = "madebyollin/sdxl-vae-fp16-fix"
# SDXL aligns to an 8-pixel latent grid, against Qwen's 16.
_LATENT_GRID = 8
# SDXL-Lightning's own four-step distillation, at the strength its authors document.
# The reference graph loads the Qwen LoRA at 0.8; carrying that number to a different
# LoRA on a different architecture would be imitation, not parity.
SDXL_STEPS = 4
def sdxl_target_size(width: int, height: int) -> tuple[int, int]:
@@ -63,14 +68,16 @@ def requested_steps(effective_steps: int, strength: float) -> int:
@dataclass
class SdxlZImagePipeline(QwenZImagePipeline):
class SdxlZImagePipeline(TwoStageZImagePipeline):
"""Lazy runtime for the SDXL global stage plus the inherited face stage."""
profile_name: ClassVar[str] = "sdxl-zimage"
def __post_init__(self) -> None:
super().__post_init__()
self._sdxl_pipe: Any = None
def _load_sdxl(self) -> Any:
def _load_global(self) -> Any:
if self._sdxl_pipe is not None:
return self._sdxl_pipe
self._require_cuda()
@@ -96,9 +103,6 @@ class SdxlZImagePipeline(QwenZImagePipeline):
add_watermarker=False,
**token,
).to(self.device)
# SDXL's own four-step distillation, at the strength its authors document.
# The reference graph loads the Qwen LoRA at 0.8; carrying that number to a
# different LoRA on a different architecture would be imitation, not parity.
pipe.load_lora_weights(hf_hub_download(SDXL_LIGHTNING_MODEL_ID, SDXL_LIGHTNING_PATTERN, **token))
pipe.fuse_lora()
# SDXL-Lightning is distilled against trailing timestep spacing.
@@ -106,29 +110,19 @@ class SdxlZImagePipeline(QwenZImagePipeline):
self._sdxl_pipe = pipe
return pipe
def preload(self, *, global_only: bool = False) -> None:
"""Eagerly load the mandatory stage and, by default, the face stack."""
from remove_ai_watermarks._internal.qwen_zimage_pipeline import _yunet_model_path
self._load_sdxl()
_yunet_model_path()
if not global_only:
self._load_zimage()
self._load_sam()
def _run_global(self, image: Image.Image, strength: float, seed: int | None) -> Image.Image:
import torch
pipe = self._load_sdxl()
pipe = self._load_global()
target = sdxl_target_size(image.width, image.height)
prepared = image if image.size == target else image.resize(target, Image.Resampling.LANCZOS)
control = build_canny_control_image(prepared)
steps = requested_steps(GLOBAL_STEPS, strength)
self._progress(f"Running SDXL Canny pass: strength={strength:.4f}, steps={GLOBAL_STEPS} of {steps}...")
steps = requested_steps(SDXL_STEPS, strength)
self._progress(f"Running SDXL Canny pass: strength={strength:.4f}, steps={SDXL_STEPS} of {steps}...")
generator = torch.Generator(device=self.device).manual_seed(seed) if seed is not None else None
result = pipe(
prompt=self._global_prompt(),
negative_prompt=self._global_negative(),
prompt=_GLOBAL_PROMPT,
negative_prompt=_GLOBAL_NEGATIVE,
image=prepared,
control_image=control,
controlnet_conditioning_scale=float(self.controlnet_conditioning_scale),
@@ -140,15 +134,3 @@ class SdxlZImagePipeline(QwenZImagePipeline):
if result.size != image.size:
result = result.resize(image.size, Image.Resampling.LANCZOS)
return result.convert("RGB")
@staticmethod
def _global_prompt() -> str:
from remove_ai_watermarks._internal.qwen_zimage_pipeline import _GLOBAL_PROMPT
return _GLOBAL_PROMPT
@staticmethod
def _global_negative() -> str:
from remove_ai_watermarks._internal.qwen_zimage_pipeline import _GLOBAL_NEGATIVE
return _GLOBAL_NEGATIVE
@@ -0,0 +1,955 @@
"""The shared two-stage regeneration recipe: a profile global stage plus the Z-Image face stage.
Every removal profile implements this base rather than another profile:
``qwen-zimage``, ``sdxl-zimage`` and ``chroma-zimage`` each own one global regeneration stage, and a
future profile (a Flux 2 pass, for instance) plugs in the same way. Everything
here is shared verbatim, so the face stage, sizing, compositing, the prompt
cache, and the run orchestration cannot silently diverge between profiles.
"""
# DiffSynth, torch, transformers, and cv2 expose mostly untyped tensor/array APIs.
# Keep the relaxation local to this optional ML boundary.
# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalCall=false, reportOptionalSubscript=false, reportOptionalOperand=false, reportAttributeAccessIssue=false, reportPrivateImportUsage=false, reportPrivateUsage=false, reportUnnecessaryComparison=false
from __future__ import annotations
import contextlib
import hashlib
import logging
import os
import tempfile
import urllib.request
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any, ClassVar
import numpy as np
from PIL import Image
from remove_ai_watermarks._internal.watermark_profiles import resolve_seed
if TYPE_CHECKING:
from collections.abc import Callable
from remove_ai_watermarks._internal.text_restoration import VerifiedTextManifest
log = logging.getLogger(__name__)
ZIMAGE_TURBO_MODEL_ID = "Tongyi-MAI/Z-Image-Turbo"
SAM_MODEL_ID = "facebook/sam-vit-base"
YUNET_MODEL_URL = (
"https://media.githubusercontent.com/media/opencv/opencv_zoo/main/"
"models/face_detection_yunet/face_detection_yunet_2023mar.onnx"
)
YUNET_MODEL_NAME = "face_detection_yunet_2023mar.onnx"
YUNET_MODEL_SHA256 = "8f2383e4dd3cfbb4553ea8718107fc0423210dc964f9f4280604804ed2552fa4"
# This threshold retained the faces in the public validation set without admitting
# decorative background regions as faces.
YUNET_SCORE_THRESHOLD = 0.5
FACE_STEPS = 8
FACE_CFG = 1.0
RESIDENT_FACE_MODEL_MIN_VRAM_GIB = 64.0
FACE_DENOISE_SCALE = 0.5
GLOBAL_CONTROLNET_SCALE = 1.0
# Both stages prompt with compile-time constants, and at CFG 1.0 DiffSynth reuses the
# positive embedding for the negative side instead of encoding it, so exactly one
# embedding per stage is ever needed. Persisting it lets the next container load the
# stack without its text encoder at all: 15.45 GiB for Qwen and 7.49 GiB for Z-Image
# of a measured 87.6 GiB read per request. The stored tensors are what the encoder
# itself produced, so the output stays byte-identical. Bump the version when the
# stored shape changes; the model id and prompt are already part of the key, so a
# model or prompt change invalidates itself.
_PROMPT_CACHE_VERSION = 1
_PROMPT_CACHE_DIRNAME = "prompt-embeddings"
# DiffSynth identifies a pipeline unit by what it produces, so these tuples are how
# the prompt stage is located inside each pipeline and how its cache file is keyed.
_ZIMAGE_PROMPT_OUTPUTS = ("prompt_embeds",)
_CANNY_LOW = 13
_CANNY_HIGH = 64
# These short model inputs are retained as calibrated compatibility parameters.
# Changing them requires the same provider-oracle and identity evaluation as a model change.
_GLOBAL_PROMPT = "ultra clear and smoothe skin, spotless skin"
_GLOBAL_NEGATIVE = "moles, freckes, high detail skin"
_FACE_PROMPT = ""
_FACE_NEGATIVE = "blurry, ugly, bad quality,"
def resolve_face_model_residency(
requested: bool | None,
*,
total_memory_gib: float,
) -> bool:
"""Keep the Z-Image stack resident when explicitly requested or safely sized."""
if requested is not None:
return requested
return total_memory_gib >= RESIDENT_FACE_MODEL_MIN_VRAM_GIB
def _pin_vram_managed_models(pipe: Any) -> None:
"""Move the managed Z-Image stack to CUDA once and make offload a no-op."""
model_names = ["text_encoder", "dit", "vae_encoder", "vae_decoder"]
for name in model_names:
model = getattr(pipe, name, None)
if model is None:
continue
for module in model.modules():
if not all(
hasattr(module, attribute)
for attribute in (
"offload_dtype",
"offload_device",
"onload_dtype",
"onload_device",
"preparing_dtype",
"preparing_device",
"computation_dtype",
"computation_device",
)
):
continue
module.offload_dtype = module.computation_dtype
module.offload_device = module.computation_device
module.onload_dtype = module.computation_dtype
module.onload_device = module.computation_device
module.preparing_dtype = module.computation_dtype
module.preparing_device = module.computation_device
pipe.load_models_to_device(model_names)
def _cast_prompt_payload(payload: Any, device: Any, dtype: Any) -> Any:
"""Move a stored embedding payload onto the runtime device.
Only floating tensors take the pipeline dtype. The attention masks travel in the
same payload and are integer, so casting them would corrupt the prompt.
"""
import torch
if isinstance(payload, torch.Tensor):
if dtype is not None and payload.is_floating_point():
return payload.to(device=device, dtype=dtype)
return payload.to(device=device)
if isinstance(payload, dict):
return {key: _cast_prompt_payload(value, device, dtype) for key, value in payload.items()}
if isinstance(payload, (list, tuple)):
return type(payload)(_cast_prompt_payload(value, device, dtype) for value in payload)
return payload
def _prompt_cache_path(model_id: str, output_params: tuple[str, ...], prompt: str) -> Path:
"""Locate the stored embedding for one model and one exact prompt string."""
key = "\x1f".join((str(_PROMPT_CACHE_VERSION), model_id, *output_params, prompt))
digest = hashlib.sha256(key.encode("utf-8")).hexdigest()[:32]
return _model_cache_dir() / _PROMPT_CACHE_DIRNAME / f"{digest}.pt"
def _store_prompt_payload(path: Path, payload: Any) -> None:
"""Write the embedding atomically so a torn write can never read back as a hit."""
import torch
path.parent.mkdir(parents=True, exist_ok=True)
with tempfile.NamedTemporaryFile(dir=path.parent, suffix=".pt", delete=False) as handle:
torch.save(_cast_prompt_payload(payload, "cpu", None), handle)
temporary = Path(handle.name)
temporary.replace(path)
def _load_prompt_payload(path: Path, device: Any, dtype: Any) -> Any:
import torch
return _cast_prompt_payload(torch.load(path, map_location="cpu", weights_only=True), device, dtype)
def _cached_prompt_process(
original_process: Any,
*,
model_id: str | None,
output_params: tuple[str, ...],
require_cache: bool,
) -> Any:
cache: dict[str, dict[str, Any]] = {}
def cached_process(
runtime_pipe: Any,
prompt: str,
edit_image: Any = None,
) -> dict[str, Any]:
if edit_image is not None:
return original_process(runtime_pipe, prompt, edit_image=edit_image)
if prompt in cache:
return cache[prompt]
path = None if model_id is None else _prompt_cache_path(model_id, output_params, prompt)
if path is not None and path.exists():
try:
cache[prompt] = _load_prompt_payload(path, runtime_pipe.device, runtime_pipe.torch_dtype)
except Exception:
log.warning("Discarding unreadable prompt-embedding cache %s", path, exc_info=True)
else:
return cache[prompt]
if require_cache:
# The text encoder was left out of the model stack on the strength of
# this file, so there is nothing left to fall back to.
raise RuntimeError(
f"The cached prompt embedding {path} disappeared after the stack was loaded without its text encoder."
)
produced = original_process(runtime_pipe, prompt, edit_image=None)
if path is not None:
try:
_store_prompt_payload(path, produced)
except OSError:
log.warning("Could not persist the prompt embedding to %s", path, exc_info=True)
cache[prompt] = produced
return cache[prompt]
return cached_process
def _cache_static_prompt_embeddings(
pipe: Any,
output_params: tuple[str, ...],
*,
model_id: str | None = None,
require_cache: bool = False,
) -> bool:
"""Memoize a prompt unit when its embedding depends only on static text.
With ``model_id`` the memo is also persisted, which is what lets the next
container skip the text encoder entirely; without it the memo lives only as long
as the pipeline. ``require_cache`` says the stack was already built without a
text encoder, so a miss must fail loudly rather than call a model that is absent.
"""
for unit in pipe.units:
if tuple(getattr(unit, "output_params", ())) == output_params:
unit.process = _cached_prompt_process(
unit.process,
model_id=model_id,
output_params=output_params,
require_cache=require_cache,
)
return True
return False
def _clamp(value: float, minimum: float, maximum: float) -> float:
if maximum < minimum:
minimum, maximum = maximum, minimum
return max(minimum, min(maximum, value))
def resolution_adaptive_denoise(
width: int,
height: int,
*,
adaptive_level: int = 6,
denoise_min: float = 0.08,
denoise_max: float = 0.15,
) -> float:
"""Choose the calibrated global strength from image area and operator level."""
low, high = sorted((float(denoise_min), float(denoise_max)))
megapixels = max(1.0, float(width) * float(height)) * 1e-6
area_fraction = _clamp((megapixels - 0.30) / 3.40, 0.0, 1.0)
strength_range = high - low
strength = low + strength_range * area_fraction
level_delta = float(int(adaptive_level)) - 5.0
if level_delta >= 0.0:
strength += level_delta * strength_range * (0.285714 / 5.0)
else:
strength += level_delta * strength_range * (0.257143 / 4.0)
return _clamp(strength, 0.0001, 1.0)
def largest_face_denoise(
boxes: list[tuple[int, int, int, int]],
image_size: tuple[int, int],
*,
base_denoise: float = 0.10,
adaptive_ratio: float = 0.03,
denoise_min: float = 0.05,
denoise_max: float = 0.28,
) -> float:
"""Choose the calibrated face strength from the largest detected face area."""
width, height = image_size
image_area = max(1.0, float(width) * float(height))
largest_ratio = 0.0
for x1, y1, x2, y2 in boxes:
box_area = max(0.0, float(x2 - x1)) * max(0.0, float(y2 - y1))
largest_ratio = max(largest_ratio, box_area / image_area)
if largest_ratio <= 0.0:
return _clamp(base_denoise, denoise_min, denoise_max)
scaled = float(base_denoise) * largest_ratio / max(1e-6, float(adaptive_ratio))
return _clamp(scaled, denoise_min, denoise_max)
def _target_size(width: int, height: int) -> tuple[int, int]:
"""Floor image dimensions to the /16 latent grid without changing aspect."""
return max(16, (width // 16) * 16), max(16, (height // 16) * 16)
def _resize_to_target(image: Image.Image) -> Image.Image:
"""Resize pixels to the exact latent-grid dimensions passed to DiffSynth."""
target = _target_size(image.width, image.height)
if image.size == target:
return image
return image.resize(target, Image.Resampling.LANCZOS)
def build_canny_control_image(image: Image.Image) -> Image.Image:
"""Build the calibrated three-channel Canny conditioning map."""
import cv2
rgb = np.asarray(image.convert("RGB"))
gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
edges = cv2.Canny(gray, _CANNY_LOW, _CANNY_HIGH)
return Image.fromarray(np.repeat(edges[:, :, None], 3, axis=2))
def build_face_kwargs(crop: Image.Image, *, strength: float, seed: int | None) -> dict[str, Any]:
"""Build the DiffSynth Z-Image face-detail call shape."""
input_image = _resize_to_target(crop)
width, height = input_image.size
return {
"prompt": _FACE_PROMPT,
"negative_prompt": _FACE_NEGATIVE,
"cfg_scale": FACE_CFG,
"input_image": input_image,
"denoising_strength": float(strength),
"height": height,
"width": width,
"seed": seed,
"rand_device": "cpu",
"num_inference_steps": FACE_STEPS,
}
def composite_face(base: np.ndarray, detail: np.ndarray, mask: np.ndarray, *, feather: int = 10) -> np.ndarray:
"""Feather ``detail`` into ``base`` while preserving every zero-mask pixel."""
import cv2
if base.shape != detail.shape:
raise ValueError("base and detail must have identical shapes")
if mask.shape != base.shape[:2]:
raise ValueError("mask must match the image height and width")
alpha = mask.astype(np.float32) / 255.0
if feather > 0:
sigma = max(0.1, float(feather) / 3.0)
alpha = cv2.GaussianBlur(alpha, (0, 0), sigmaX=sigma, sigmaY=sigma)
# Blurring may introduce tiny values far outside the intended mask. Keep an
# explicit support dilation so pixels beyond the feather radius stay exact.
support = cv2.dilate((mask > 0).astype(np.uint8), np.ones((2 * feather + 1, 2 * feather + 1), np.uint8))
alpha *= support
alpha = np.clip(alpha, 0.0, 1.0)[:, :, None]
merged = base.astype(np.float32) * (1.0 - alpha) + detail.astype(np.float32) * alpha
return np.clip(np.rint(merged), 0, 255).astype(np.uint8)
def _expanded_box(
box: tuple[int, int, int, int],
image_size: tuple[int, int],
*,
factor: float = 2.5,
) -> tuple[int, int, int, int]:
"""Expand a face box around its center to include local lighting context."""
x1, y1, x2, y2 = box
image_width, image_height = image_size
center_x = (x1 + x2) / 2.0
center_y = (y1 + y2) / 2.0
width = max(1.0, (x2 - x1) * factor)
height = max(1.0, (y2 - y1) * factor)
return (
max(0, round(center_x - width / 2.0)),
max(0, round(center_y - height / 2.0)),
min(image_width, round(center_x + width / 2.0)),
min(image_height, round(center_y + height / 2.0)),
)
def _model_cache_dir() -> Path:
"""Where this library persists downloaded and derived model assets.
``HF_HOME`` comes first because on a scale-to-zero runner it is the one path
mounted persistently; anything under a container-local cache is re-derived on
every request, which defeats both the YuNet download and the prompt cache.
"""
root = os.environ.get("HF_HOME") or os.environ.get("XDG_CACHE_HOME")
base = Path(root) if root else Path.home() / ".cache"
return base / "remove-ai-watermarks"
def _yunet_model_path() -> Path:
"""Download the small MIT-licensed YuNet ONNX model on first use."""
model_path = _model_cache_dir() / YUNET_MODEL_NAME
if model_path.exists() and hashlib.sha256(model_path.read_bytes()).hexdigest() == YUNET_MODEL_SHA256:
return model_path
model_path.parent.mkdir(parents=True, exist_ok=True)
log.info("Downloading YuNet face detector: %s", YUNET_MODEL_URL)
request = urllib.request.Request(
YUNET_MODEL_URL,
headers={"User-Agent": "remove-ai-watermarks"},
)
with urllib.request.urlopen(request, timeout=60) as response: # noqa: S310 - fixed HTTPS source
payload = response.read()
log.info(
"YuNet download response: status=%s content_length=%s",
getattr(response, "status", None),
len(payload),
)
digest = hashlib.sha256(payload).hexdigest()
if digest != YUNET_MODEL_SHA256:
raise OSError(
"YuNet download failed integrity verification: "
f"expected {YUNET_MODEL_SHA256}, got {digest} ({len(payload)} bytes)"
)
with tempfile.NamedTemporaryFile(dir=model_path.parent, suffix=".onnx", delete=False) as handle:
handle.write(payload)
temporary = Path(handle.name)
temporary.replace(model_path)
return model_path
def _nms_boxes(
boxes: list[tuple[int, int, int, int]],
scores: list[float],
*,
threshold: float = 0.3,
) -> list[tuple[int, int, int, int]]:
"""Apply OpenCV NMS to detections collected at multiple image scales."""
if not boxes:
return []
import cv2
xywh = [[x1, y1, x2 - x1, y2 - y1] for x1, y1, x2, y2 in boxes]
indices = cv2.dnn.NMSBoxes(xywh, scores, score_threshold=0.2, nms_threshold=threshold)
if len(indices) == 0:
return []
return [boxes[int(index)] for index in np.asarray(indices).reshape(-1)]
def detect_faces(image: Image.Image) -> list[tuple[int, int, int, int]]:
"""Detect face boxes with YuNet at two scales for large and small faces."""
import cv2
rgb = np.asarray(image.convert("RGB"))
original_height, original_width = rgb.shape[:2]
detections: list[tuple[int, int, int, int]] = []
scores: list[float] = []
model_path = _yunet_model_path()
for long_side in (640, 1280):
scale = min(1.0, long_side / max(original_width, original_height))
width = max(1, round(original_width * scale))
height = max(1, round(original_height * scale))
resized = cv2.resize(rgb, (width, height), interpolation=cv2.INTER_AREA) if scale < 1.0 else rgb
bgr = cv2.cvtColor(resized, cv2.COLOR_RGB2BGR)
detector = cv2.FaceDetectorYN.create(
str(model_path),
"",
(width, height),
YUNET_SCORE_THRESHOLD,
0.3,
5000,
)
_, rows = detector.detect(bgr)
if rows is None:
continue
inverse = 1.0 / scale
for row in rows:
x, y, box_width, box_height = (float(value) for value in row[:4])
x1 = max(0, round(x * inverse))
y1 = max(0, round(y * inverse))
x2 = min(original_width, round((x + box_width) * inverse))
y2 = min(original_height, round((y + box_height) * inverse))
if x2 <= x1 or y2 <= y1:
continue
detections.append((x1, y1, x2, y2))
scores.append(float(row[-1]))
if scale == 1.0:
break
return _nms_boxes(detections, scores)
def _ellipse_masks(
boxes: list[tuple[int, int, int, int]],
image_size: tuple[int, int],
) -> list[np.ndarray]:
"""Safe fallback masks when SAM is unavailable."""
import cv2
width, height = image_size
masks: list[np.ndarray] = []
for x1, y1, x2, y2 in boxes:
mask = np.zeros((height, width), dtype=np.uint8)
center = ((x1 + x2) // 2, (y1 + y2) // 2)
axes = (max(1, int((x2 - x1) * 0.55)), max(1, int((y2 - y1) * 0.62)))
cv2.ellipse(mask, center, axes, 0, 0, 360, 255, -1)
masks.append(mask)
return masks
def _prepare_sam_inputs(inputs: Any, device: str, dtype: Any) -> Any:
"""Move SAM inputs to the target device without casting geometric prompts."""
prepared = inputs.to(device)
if "pixel_values" in prepared:
prepared["pixel_values"] = prepared["pixel_values"].to(dtype=dtype)
return prepared
def _sam_point_prompts(
boxes: list[tuple[int, int, int, int]],
) -> tuple[list[list[list[list[float]]]], list[list[list[int]]]]:
"""Build Impact Pack's center-1 positive prompt for every face box."""
points = [[[[(x1 + x2) / 2.0, (y1 + y2) / 2.0]] for x1, y1, x2, y2 in boxes]]
labels = [[[1] for _box in boxes]]
return points, labels
def _clip_sam_masks_to_boxes(
masks: list[np.ndarray],
boxes: list[tuple[int, int, int, int]],
image_size: tuple[int, int],
) -> list[np.ndarray]:
"""Match Impact Pack by intersecting each SAM mask with its detector box."""
# Same rectangle primitive the shared fill uses, rather than a private zeros/fill
# copy. `dilate=0` because this box CLIPS a SAM mask -- growing it would admit the
# pixels the clip exists to exclude. The function-local import keeps region_eraser's
# module-scope cv2 off this module's import path.
from remove_ai_watermarks.region_eraser import boxes_to_mask
width, height = image_size
clipped: list[np.ndarray] = []
for mask, (x1, y1, x2, y2) in zip(masks, boxes, strict=True):
# (x, y, w, h) with w = x2 - x1 keeps x + w == x2, so a negative origin clamps
# to the same span the explicit max/min pair produced.
box_mask = boxes_to_mask((height, width), [(x1, y1, x2 - x1, y2 - y1)], dilate=0)
clipped.append(np.bitwise_and(mask.astype(np.uint8), box_mask))
return clipped
def _select_sam_masks(
masks: np.ndarray,
scores: np.ndarray,
*,
threshold: float = 0.93,
) -> list[np.ndarray]:
"""Select and combine SAM proposals like Impact Pack's ``sub_threshold``."""
mask_array = np.asarray(masks)
score_array = np.asarray(scores)
if mask_array.ndim == 5 and mask_array.shape[0] == 1:
mask_array = mask_array[0]
if score_array.ndim == 3 and score_array.shape[0] == 1:
score_array = score_array[0]
if mask_array.ndim == 3:
mask_array = mask_array[:, None, :, :]
if score_array.ndim == 1:
score_array = score_array[:, None]
if mask_array.ndim != 4 or score_array.ndim != 2:
raise ValueError("SAM masks and scores have unexpected dimensions")
if mask_array.shape[:2] != score_array.shape:
raise ValueError("SAM mask proposals and IoU scores do not align")
selected_masks: list[np.ndarray] = []
for candidates, candidate_scores in zip(mask_array, score_array, strict=True):
selected = np.flatnonzero(candidate_scores >= threshold)
if selected.size == 0:
selected = np.asarray([int(np.argmax(candidate_scores))])
combined = np.any(candidates[selected] > 0, axis=0)
selected_masks.append(combined.astype(np.uint8) * 255)
return selected_masks
def _sam_outputs_to_numpy(masks: Any, scores: Any) -> tuple[np.ndarray, np.ndarray]:
"""Convert SAM tensors through float32 because NumPy rejects bfloat16."""
mask_array = masks.detach().float().cpu().numpy()
score_array = scores.detach().float().cpu().numpy()
return mask_array, score_array
@dataclass
class TwoStageZImagePipeline:
"""Lazy runtime for the shared two-stage regeneration recipe.
A profile subclass owns the global stage: ``_load_global`` loads and caches its
model stack, ``_run_global`` regenerates a frame, and ``_vae_roundtrip``
optionally exposes that stack's VAE as a verified-text donor. The Z-Image face
stage with YuNet detection and SAM masks, the prompt cache, sizing, and the run
orchestration live here and are inherited verbatim by every profile.
"""
# Names the CUDA-only refusal after the concrete profile rather than after this
# base, so a directly constructed subclass reports itself.
profile_name: ClassVar[str] = ""
device: str
torch_dtype: Any
hf_token: str | None = None
progress_callback: Callable[[str], None] | None = None
controlnet_conditioning_scale: float = GLOBAL_CONTROLNET_SCALE
keep_face_models_on_device: bool | None = None
keep_global_models_on_device: bool | None = None
cache_prompt_embeddings: bool = True
def __post_init__(self) -> None:
self._zimage_pipe: Any = None
self._sam_model: Any = None
self._sam_processor: Any = None
def _progress(self, message: str) -> None:
if self.progress_callback is not None:
with contextlib.suppress(Exception):
self.progress_callback(message)
def _require_cuda(self) -> None:
if self.device != "cuda":
name = self.profile_name or type(self).__name__
raise RuntimeError(
f"The {name} pipeline is CUDA-only; there is no CPU or MPS path "
"for it. WatermarkRemover already refuses a non-CUDA device, so "
"reaching this guard means the pipeline was constructed directly."
)
def _vram_limit(self) -> float | None:
import torch
with contextlib.suppress(Exception):
return max(1.0, torch.cuda.mem_get_info("cuda")[1] / (1024**3) - 0.5)
return None
def _total_vram_gib(self) -> float:
"""Card capacity for the residency gates, 0.0 when it cannot be read.
Separate from ``_vram_limit``, which answers a different question (the budget
handed to DiffSynth) and must stay ``None`` rather than 0.0 when unknown, so
an unreadable device means "no limit" there and "assume small" here.
"""
import torch
with contextlib.suppress(Exception):
return torch.cuda.get_device_properties("cuda").total_memory / (1024**3)
return 0.0
def _keep_face_models_resident(self) -> bool:
return resolve_face_model_residency(
self.keep_face_models_on_device,
total_memory_gib=self._total_vram_gib(),
)
def _prompt_is_cached(self, model_id: str, output_params: tuple[str, ...], prompt: str) -> bool:
"""Whether this stack can be built without its text encoder at all."""
return self.cache_prompt_embeddings and _prompt_cache_path(model_id, output_params, prompt).exists()
@classmethod
def _face_stage_dtype(cls) -> Any:
"""The dtype every face-stage model loads and computes in.
Deliberately independent of ``self.torch_dtype``, which belongs to the global
stage. ``sdxl-zimage`` runs its global model in fp16, and inheriting that here
built the Z-Image modules bf16 (per the VRAM config below) while handing them
fp16 latents -- a Half/BFloat16 conv mismatch that crashed every face image
while zero-face inputs passed, so no unit test could see it.
Read by Z-Image and by SAM alike, so the whole stage moves together. SAM never
crashed, because it casts its own inputs and leaves through ``.float()``, but it
was reading the same wrong field and would re-land the bug for the next profile
that changes its global dtype.
"""
return cls._zimage_vram_config()["computation_dtype"]
@staticmethod
def _zimage_vram_config() -> dict[str, Any]:
import torch
return {
"offload_dtype": torch.bfloat16,
"offload_device": "cpu",
"onload_dtype": torch.bfloat16,
"onload_device": "cpu",
"preparing_dtype": torch.bfloat16,
"preparing_device": "cuda",
"computation_dtype": torch.bfloat16,
"computation_device": "cuda",
}
def _load_zimage(self) -> Any:
if self._zimage_pipe is not None:
return self._zimage_pipe
self._require_cuda()
os.environ.setdefault("DIFFSYNTH_DOWNLOAD_SOURCE", "huggingface")
if self.hf_token:
os.environ.setdefault("HF_TOKEN", self.hf_token)
try:
from diffsynth.pipelines.z_image import ModelConfig, ZImagePipeline
except ImportError as exc:
raise ImportError(
"The two-stage pipelines need the optional dependency group. "
"Install: pip install 'remove-ai-watermarks[qwen-zimage]'"
) from exc
self._progress("Loading Z-Image Turbo face-detail model...")
keep_on_device = self._keep_face_models_resident()
config = self._zimage_vram_config()
text_encoder_config = ModelConfig(
model_id=ZIMAGE_TURBO_MODEL_ID,
origin_file_pattern="text_encoder/*.safetensors",
**config,
)
model_configs = [
ModelConfig(
model_id=ZIMAGE_TURBO_MODEL_ID,
origin_file_pattern="transformer/*.safetensors",
**config,
),
text_encoder_config,
ModelConfig(
model_id=ZIMAGE_TURBO_MODEL_ID,
origin_file_pattern="vae/diffusion_pytorch_model.safetensors",
**config,
),
]
prompt_cached = self._prompt_is_cached(ZIMAGE_TURBO_MODEL_ID, _ZIMAGE_PROMPT_OUTPUTS, _FACE_PROMPT)
if prompt_cached:
model_configs.remove(text_encoder_config)
log.info("Z-Image prompt embedding is cached; loading the stack without its text encoder")
pipe = ZImagePipeline.from_pretrained(
torch_dtype=self._face_stage_dtype(),
device=self.device,
model_configs=model_configs,
tokenizer_config=ModelConfig(
model_id=ZIMAGE_TURBO_MODEL_ID,
origin_file_pattern="tokenizer/",
),
vram_limit=self._vram_limit(),
)
if keep_on_device:
_pin_vram_managed_models(pipe)
if self.cache_prompt_embeddings:
_cache_static_prompt_embeddings(
pipe,
_ZIMAGE_PROMPT_OUTPUTS,
model_id=ZIMAGE_TURBO_MODEL_ID,
require_cache=prompt_cached,
)
self._zimage_pipe = pipe
return pipe
def _load_sam(self) -> tuple[Any, Any]:
if self._sam_model is not None and self._sam_processor is not None:
return self._sam_model, self._sam_processor
self._progress("Loading SAM face-mask model...")
try:
from transformers import AutoModelForMaskGeneration, AutoProcessor
except ImportError as exc:
raise ImportError("SAM needs transformers and torchvision from the qwen-zimage extra.") from exc
kwargs: dict[str, Any] = {}
if self.hf_token:
kwargs["token"] = self.hf_token
processor = AutoProcessor.from_pretrained(SAM_MODEL_ID, **kwargs)
model = AutoModelForMaskGeneration.from_pretrained(
SAM_MODEL_ID,
torch_dtype=self._face_stage_dtype(),
**kwargs,
).to(self.device)
model.eval()
self._sam_model = model
self._sam_processor = processor
return model, processor
def _sam_masks(
self,
image: Image.Image,
boxes: list[tuple[int, int, int, int]],
) -> list[np.ndarray]:
if not boxes:
return []
import torch
try:
model, processor = self._load_sam()
input_points, input_labels = _sam_point_prompts(boxes)
inputs = processor(
images=image,
input_boxes=[[list(box) for box in boxes]],
input_points=input_points,
input_labels=input_labels,
return_tensors="pt",
)
original_sizes = inputs["original_sizes"].clone()
reshaped_sizes = inputs["reshaped_input_sizes"].clone()
inputs = _prepare_sam_inputs(inputs, self.device, self._face_stage_dtype())
with torch.inference_mode():
outputs = model(**inputs, multimask_output=True)
processed = processor.post_process_masks(
outputs.pred_masks.detach().cpu(),
original_sizes,
reshaped_sizes,
)[0]
mask_array, score_array = _sam_outputs_to_numpy(processed, outputs.iou_scores)
binary_masks = _select_sam_masks(
mask_array,
score_array,
)
return _clip_sam_masks_to_boxes(binary_masks, boxes, image.size)
except Exception as exc:
log.warning("SAM face-mask refinement failed (%s); using box-derived ellipse masks", exc)
return _ellipse_masks(boxes, image.size)
def preload(self, *, global_only: bool = False) -> None:
"""Eagerly load the mandatory stage and, by default, the face stack."""
self._load_global()
_yunet_model_path()
if not global_only:
self._load_zimage()
self._load_sam()
def _load_global(self) -> Any:
"""Load and cache the profile's global model stack."""
raise NotImplementedError("A profile pipeline must implement the global stage loader")
def _run_global(self, image: Image.Image, strength: float, seed: int | None) -> Image.Image:
"""Regenerate one frame with the profile's global stage."""
raise NotImplementedError("A profile pipeline must implement the global stage")
def _vae_roundtrip(self, image: Image.Image) -> Image.Image:
"""Reconstruct source pixels through the profile's already loaded VAE."""
raise NotImplementedError(f"{type(self).__name__} does not expose a VAE donor for verified text restoration")
@staticmethod
def _detail_size(
crop_size: tuple[int, int],
face_size: tuple[int, int],
) -> tuple[int, int]:
"""Scale a crop toward a 768px face guide while capping it at 1024px."""
crop_width, crop_height = crop_size
face_width, face_height = face_size
scale_for_face = 768.0 / max(1, max(face_width, face_height))
scale_for_crop = 1024.0 / max(1, max(crop_width, crop_height))
scale = min(scale_for_face, scale_for_crop)
# Never shrink below the crop's current size unless the cap requires it.
if max(crop_width, crop_height) <= 1024:
scale = max(1.0, scale)
width = max(16, round(crop_width * scale / 16.0) * 16)
height = max(16, round(crop_height * scale / 16.0) * 16)
return width, height
def _run_faces(
self,
original: Image.Image,
global_result: Image.Image,
boxes: list[tuple[int, int, int, int]],
masks: list[np.ndarray],
*,
strength: float,
seed: int | None,
) -> Image.Image:
if not boxes:
return global_result
pipe = self._load_zimage()
base = np.asarray(global_result.convert("RGB")).copy()
source = np.asarray(original.convert("RGB"))
detail_seed = None if seed is None else seed + 1
for index, (box, mask) in enumerate(zip(boxes, masks, strict=True), start=1):
crop_box = _expanded_box(box, original.size)
cx1, cy1, cx2, cy2 = crop_box
crop_source = source[cy1:cy2, cx1:cx2]
crop_mask = mask[cy1:cy2, cx1:cx2]
if crop_source.size == 0 or not np.any(crop_mask):
continue
face_width = box[2] - box[0]
face_height = box[3] - box[1]
process_size = self._detail_size((cx2 - cx1, cy2 - cy1), (face_width, face_height))
crop_image = Image.fromarray(crop_source).resize(process_size, Image.Resampling.LANCZOS)
self._progress(
f"Regenerating face {index}/{len(boxes)} with Z-Image: strength={strength:.4f}, steps={FACE_STEPS}..."
)
detailed = pipe(**build_face_kwargs(crop_image, strength=strength, seed=detail_seed))
detailed = detailed.convert("RGB").resize((cx2 - cx1, cy2 - cy1), Image.Resampling.LANCZOS)
base_crop = base[cy1:cy2, cx1:cx2]
base[cy1:cy2, cx1:cx2] = composite_face(
base_crop,
np.asarray(detailed),
crop_mask,
feather=10,
)
return Image.fromarray(base)
def run(
self,
image: Image.Image,
*,
strength: float | None,
seed: int | None,
tile: bool = False,
tile_size: int = 1024,
tile_overlap: int = 128,
text_manifest: VerifiedTextManifest | None = None,
fidelity_anchor: bool = False,
) -> Image.Image:
"""Execute global regeneration and masked face repair."""
if text_manifest is not None and tile:
raise ValueError("Verified text restoration is not calibrated with tiled diffusion")
self._require_cuda()
seed = resolve_seed(seed)
donor = None
if text_manifest is not None:
self._progress("Reconstructing the verified text donor...")
donor = self._vae_roundtrip(image)
global_strength = (
resolution_adaptive_denoise(image.width, image.height) if strength is None else float(strength)
)
if tile and max(image.size) > tile_size:
from remove_ai_watermarks._internal.tiling import run_tiled
global_result = run_tiled(
lambda tile_image: self._run_global(tile_image, global_strength, seed),
image,
tile_size,
tile_overlap,
self._progress,
)
else:
global_result = self._run_global(image, global_strength, seed)
self._progress("Detecting faces on the original image...")
boxes = detect_faces(image)
if not boxes:
self._progress("No faces detected; keeping the global result.")
result = global_result
else:
masks = self._sam_masks(image, boxes)
face_strength = largest_face_denoise(boxes, image.size) * FACE_DENOISE_SCALE
result = self._run_faces(
image,
global_result,
boxes,
masks,
strength=face_strength,
seed=seed,
)
if text_manifest is None:
return result
if donor is None:
raise RuntimeError("Verified text restoration requires a VAE donor")
from remove_ai_watermarks._internal.text_restoration import (
blend_fidelity_anchor,
restore_verified_text,
)
# The anchor is OFF by default since 0.27.1: blending 15% of the Qwen-VAE
# donor ACROSS THE WHOLE FRAME returned detector-visible OpenAI SynthID on
# poster-scale manifests (official Content Provenance API, 2026-08-19:
# detected x6 with the anchor, clean x6 without it, base clean throughout;
# see docs/text-protection-research.md). ``fidelity_anchor=True`` keeps the
# 0.27.0 research behavior for reproduction.
if fidelity_anchor:
self._progress("Blending the Qwen-VAE fidelity anchor...")
anchor = blend_fidelity_anchor(result, donor)
else:
anchor = result
self._progress(f"Restoring {len(text_manifest.lines)} verified text lines...")
return restore_verified_text(image, anchor, donor, text_manifest.lines)
@@ -1,6 +1,6 @@
"""Project-owned configuration for invisible-watermark regeneration profiles.
Two profiles remain, and both are CUDA-only: ``qwen-zimage`` (the default) and
Three profiles remain, and all are CUDA-only: ``qwen-zimage`` (the default),
``sdxl-zimage``. The older ``controlnet``, ``sdxl``, ``qwen`` and ``default`` profiles
were removed rather than kept as a CPU path, because none of them matched the two-stage
recipe's face preservation and keeping them implied a quality this library no longer
@@ -18,14 +18,15 @@ if TYPE_CHECKING:
# SDXL base is no longer a profile of its own, but it is still the global stage of
# sdxl-zimage, so the checkpoint id stays. Named for what it is rather than
# ``DEFAULT_MODEL_ID``: there is no user-selectable model any more, so "default"
# implied an override that both profiles reject.
# implied an override that every profile rejects.
SDXL_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
CONTROLNET_CANNY_MODEL = "xinsir/controlnet-canny-sdxl-1.0"
QWEN_ZIMAGE_PROFILE = "qwen-zimage"
SDXL_ZIMAGE_PROFILE = "sdxl-zimage"
CHROMA_ZIMAGE_PROFILE = "chroma-zimage"
DEFAULT_PROFILE = QWEN_ZIMAGE_PROFILE
PROFILE_CHOICES = (QWEN_ZIMAGE_PROFILE, SDXL_ZIMAGE_PROFILE)
PROFILE_CHOICES = (QWEN_ZIMAGE_PROFILE, SDXL_ZIMAGE_PROFILE, CHROMA_ZIMAGE_PROFILE)
# The modules a real removal run needs, and the extra that installs them. Both live
# here, in the only profile module that imports nothing heavy, because the CLI's
@@ -43,12 +44,19 @@ INVISIBLE_EXTRA = "'remove-ai-watermarks[qwen-zimage]'"
# no-op at best. sdxl-zimage's global pass leaves the softer output the polish exists
# for. This is per-profile data, not a CLI concern: the flag defaults to None so that
# "the user did not choose" stays a value rather than an inference from Click state.
PROFILE_ADAPTIVE_POLISH = {QWEN_ZIMAGE_PROFILE: False, SDXL_ZIMAGE_PROFILE: True}
PROFILE_ADAPTIVE_POLISH = {
QWEN_ZIMAGE_PROFILE: False,
SDXL_ZIMAGE_PROFILE: True,
# Chroma1's output matches the input's detail level, like qwen-zimage: the
# matched-strength calibration measured LPIPS/SSIM parity or better, so
# polishing would add grain to an already-faithful output.
CHROMA_ZIMAGE_PROFILE: False,
}
SDXL_LIGHTNING_MODEL_ID = "ByteDance/SDXL-Lightning"
SDXL_LIGHTNING_PATTERN = "sdxl_lightning_4step_lora.safetensors"
# Both profiles are certified at a fixed seed because SynthID removal near the
# All profiles are certified at a fixed seed because SynthID removal near the
# strength floor is seed-dependent. The step count and CFG are not settable at all --
# each stage owns them (``GLOBAL_STEPS`` / ``FACE_STEPS`` in qwen_zimage_pipeline).
PROFILE_SEED = 0
@@ -117,9 +125,57 @@ _SDXL_ZIMAGE_STRENGTH_BY_VENDOR: dict[str, float] = {
"openai": SDXL_ZIMAGE_OPENAI_STRENGTH,
"google": SDXL_ZIMAGE_GEMINI_STRENGTH,
}
# chroma-zimage floors, from the 2026-08-29/30 four-cohort oracle calibration
# (docs/chroma1-engine-research.md; ChromaImg2ImgPipeline, neutral prompt,
# guidance 5.0, four effective steps, seed 0). Flat per-vendor values derived
# by the same worst-clean-plus-one-spread rule as the other profiles:
#
# - OpenAI: first-clean boundaries 0.06 / 0.06 / 0.075 on three fixtures ->
# 0.075 + (0.075 - 0.06) = 0.09. Chroma1's OpenAI floor is BELOW qwen's
# 0.07675 while its matched-strength fidelity beats qwen on every metric.
# - Microsoft InvisMark: paint-1 (0.06, 0.08], paint-2 <= 0.04, paint-3
# (0.06, 0.08] -> 0.08 + (0.08 - 0.04) = 0.12, rounded up to the measured
# rung 0.125, oracle-verified clean on both worst sources. BELOW qwen's 0.15.
# - Google: 633uuy (0.20, 0.25], akdbei (0.20, 0.25], y48j3c (0.08, 0.12],
# 3mc4t9 (0.08, 0.12] -> 0.25 + (0.25 - 0.12) = 0.38, rounded up to 0.40,
# oracle-verified clean on the worst fixture. ABOVE qwen's 0.27; at 0.40 the
# regeneration destroys dense text and collapses face identity, which is why
# the matched-strength addendum recommends the adaptive-strength follow-up.
# - Meta Content Seal: lighthouse (0.08, 0.10], fox (0.06, 0.08], night_city
# (0.045, 0.06], studio_mug (0.03, 0.045], text_poster <= 0.03 ->
# 0.10 + (0.10 - 0.03) = 0.17. ABOVE qwen's 0.1; the floors exist because
# Chroma1's boundaries scatter wider than qwen's, not because the engine is
# worse per unit strength.
CHROMA_ZIMAGE_OPENAI_STRENGTH = 0.09
CHROMA_ZIMAGE_MICROSOFT_STRENGTH = 0.125
CHROMA_ZIMAGE_GOOGLE_STRENGTH = 0.40
CHROMA_ZIMAGE_META_STRENGTH = 0.17
CHROMA_ZIMAGE_UNKNOWN_STRENGTH = CHROMA_ZIMAGE_GOOGLE_STRENGTH
# Content-adaptive Google floor: the four-fixture calibration showed a clean
# face-count split. Both zero-face fixtures (dense-text cards 633uuy and akdbei)
# need 0.25 first-clean; both face fixtures (y48j3c with 17 detected faces and
# 3mc4t9 with 7) clear at 0.12. The flat 0.40 floor is the zero-face policy
# (worst 0.25 + spread 0.13 = 0.38, rounded to 0.40). Face content can use the
# measured 0.125 rung instead: both face fixtures first-cleaned at 0.12, and
# with identical boundaries the cross-source spread is zero, so 0.125 (the next
# measured rung above 0.12) is the operating point. This split is
# Google-SynthID-specific: OpenAI's face fixture was HARDER than its text
# fixture (0.075 vs 0.06), so no other cohort gets an adaptive arm.
# Oracle-verified 2026-08-30; see docs/chroma1-engine-research.md.
CHROMA_ZIMAGE_GOOGLE_FACE_STRENGTH = 0.125
_CHROMA_ZIMAGE_STRENGTH_BY_VENDOR: dict[str, float] = {
"openai": CHROMA_ZIMAGE_OPENAI_STRENGTH,
"google": CHROMA_ZIMAGE_GOOGLE_STRENGTH,
"microsoft": CHROMA_ZIMAGE_MICROSOFT_STRENGTH,
"meta": CHROMA_ZIMAGE_META_STRENGTH,
}
_ALIASES = {
"qwen_zimage": QWEN_ZIMAGE_PROFILE,
"sdxl_zimage": SDXL_ZIMAGE_PROFILE,
"chroma_zimage": CHROMA_ZIMAGE_PROFILE,
}
@@ -130,7 +186,7 @@ def normalize_profile(profile: str) -> str:
def resolve_seed(seed: int | None) -> int:
"""Keep both profiles reproducible by default."""
"""Keep every profile reproducible by default."""
return PROFILE_SEED if seed is None else seed
@@ -149,7 +205,9 @@ def strength_default_help() -> str:
f"Microsoft InvisMark {QWEN_ZIMAGE_MICROSOFT_STRENGTH} / Meta Content Seal "
f"{QWEN_ZIMAGE_META_STRENGTH} floors; sdxl-zimage "
f"uses OpenAI {SDXL_ZIMAGE_OPENAI_STRENGTH} / Google {SDXL_ZIMAGE_GEMINI_STRENGTH} / "
f"unknown {SDXL_ZIMAGE_UNKNOWN_STRENGTH}, from the C2PA issuer)"
f"unknown {SDXL_ZIMAGE_UNKNOWN_STRENGTH}; chroma-zimage uses OpenAI "
f"{CHROMA_ZIMAGE_OPENAI_STRENGTH} / Microsoft {CHROMA_ZIMAGE_MICROSOFT_STRENGTH} / "
f"Google {CHROMA_ZIMAGE_GOOGLE_STRENGTH} / Meta {CHROMA_ZIMAGE_META_STRENGTH}, from the C2PA issuer)"
)
@@ -159,6 +217,7 @@ def resolve_strength(
pipeline: str | None = None,
*,
size: tuple[int, int] | None = None,
face_count: int | None = None,
) -> float:
"""Resolve a user override or the calibrated policy for a profile and vendor.
@@ -173,14 +232,20 @@ def resolve_strength(
"""
if strength is not None:
return strength
if normalize_profile(pipeline or "") == SDXL_ZIMAGE_PROFILE:
normalized = normalize_profile(pipeline or "")
if normalized == SDXL_ZIMAGE_PROFILE:
return _SDXL_ZIMAGE_STRENGTH_BY_VENDOR.get((vendor or "").casefold(), SDXL_ZIMAGE_UNKNOWN_STRENGTH)
if normalized == CHROMA_ZIMAGE_PROFILE:
vendor_key = (vendor or "").casefold()
if vendor_key == "google" and face_count is not None and face_count > 0:
return CHROMA_ZIMAGE_GOOGLE_FACE_STRENGTH
return _CHROMA_ZIMAGE_STRENGTH_BY_VENDOR.get(vendor_key, CHROMA_ZIMAGE_UNKNOWN_STRENGTH)
if size is None:
raise ValueError("qwen-zimage resolves strength from image area, so size is required")
vendor_strength = _QWEN_ZIMAGE_FLAT_STRENGTH_BY_VENDOR.get((vendor or "").casefold())
if vendor_strength is not None:
return vendor_strength
from remove_ai_watermarks._internal.qwen_zimage_pipeline import resolution_adaptive_denoise
from remove_ai_watermarks._internal.two_stage_pipeline import resolution_adaptive_denoise
return resolution_adaptive_denoise(*size)
@@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any
from PIL import Image
from remove_ai_watermarks._internal.watermark_profiles import (
CHROMA_ZIMAGE_PROFILE,
DEFAULT_PROFILE,
INVISIBLE_EXTRA,
PROFILE_CHOICES,
@@ -80,7 +81,7 @@ def _cuda_works() -> bool:
def get_device() -> str:
"""Return ``"cuda"`` when a usable CUDA backend is present, else ``"cpu"``.
Deliberately binary. Both profiles are CUDA-only, so an XPU or MPS answer would
Deliberately binary. All profiles are CUDA-only, so an XPU or MPS answer would
only travel one frame further to the same refusal in :class:`WatermarkRemover`,
while costing a probe on each. ``"cpu"`` here means "no CUDA", which is exactly
what that refusal reports.
@@ -117,13 +118,13 @@ class WatermarkRemover:
_ensure_watermark_deps()
selected_device = (device or get_device()).casefold()
self.device = get_device() if selected_device == "auto" else selected_device
# CUDA is a precondition of the object, not of the run. Both profiles raise on
# CUDA is a precondition of the object, not of the run. All profiles raise on
# any other device, so accepting one here only defers a guaranteed failure to
# model-load time, several layers down and under the wrong profile's name.
if self.device != "cuda":
raise ValueError(
f"Invisible-watermark removal is CUDA-only, so '{self.device}' cannot run it. "
"Both remaining profiles need an NVIDIA GPU. Visible-mark removal and "
"All profiles need an NVIDIA GPU. Visible-mark removal and "
"every identify command still run on CPU."
)
@@ -131,6 +132,10 @@ class WatermarkRemover:
# SDXL ships fp16 weights and an fp16-safe VAE; bf16 would give up the
# variant without buying anything on this architecture.
self.torch_dtype = torch.float16 # type: ignore[union-attr]
elif self.model_profile == CHROMA_ZIMAGE_PROFILE:
# Chroma1-HD is calibrated bf16 end to end (the floors were measured
# with bf16 weights); the inherited face stage is bf16 regardless.
self.torch_dtype = torch.bfloat16 # type: ignore[union-attr]
else:
self.torch_dtype = torch.bfloat16 # type: ignore[union-attr]
@@ -150,6 +155,10 @@ class WatermarkRemover:
from remove_ai_watermarks._internal.sdxl_zimage_pipeline import (
SdxlZImagePipeline as _Pipeline,
)
elif self.model_profile == CHROMA_ZIMAGE_PROFILE:
from remove_ai_watermarks._internal.chroma_zimage_pipeline import (
ChromaZImagePipeline as _Pipeline,
)
else:
from remove_ai_watermarks._internal.qwen_zimage_pipeline import (
QwenZImagePipeline as _Pipeline,
@@ -179,6 +188,26 @@ class WatermarkRemover:
remove_ai_metadata(output_path, output_path, keep_standard=True)
def _resolve_chroma_strength(
self,
strength: float | None,
vendor: str | None,
source: Image.Image,
) -> float:
"""Resolve strength, running face detection for the chroma Google arm.
The content-adaptive Google floor needs a face count. YuNet is fast
(~50 ms) and the model download is already managed by the shared base.
"""
if self.model_profile != CHROMA_ZIMAGE_PROFILE or (vendor or "").casefold() != "google":
return resolve_strength(strength, vendor, self.model_profile, size=source.size)
if strength is not None:
return strength
from remove_ai_watermarks._internal.two_stage_pipeline import detect_faces
face_count = len(detect_faces(source))
return resolve_strength(strength, vendor, self.model_profile, size=source.size, face_count=face_count)
def remove_watermark(
self,
image_path: Path,
@@ -194,7 +223,7 @@ class WatermarkRemover:
) -> Path:
"""Regenerate image pixels and write the result without AI metadata.
Step count and CFG are not parameters. Each stage of both profiles is a
Step count and CFG are not parameters. Each stage of every profile is a
distilled schedule that owns its own, so the only thing a caller-supplied
value could do is break the run or be rejected.
"""
@@ -204,10 +233,10 @@ class WatermarkRemover:
with Image.open(image_path) as opened:
source = opened.convert("RGB")
resolved_strength = resolve_strength(strength, vendor, self.model_profile, size=source.size)
resolved_strength = self._resolve_chroma_strength(strength, vendor, source)
if not 0.0 <= resolved_strength <= 1.0:
raise ValueError(f"Strength must be between 0.0 and 1.0, got {resolved_strength}")
if text_manifest is not None and self.model_profile == SDXL_ZIMAGE_PROFILE:
if text_manifest is not None and self.model_profile in (SDXL_ZIMAGE_PROFILE, CHROMA_ZIMAGE_PROFILE):
raise ValueError("Verified text restoration is supported only by the qwen-zimage profile")
if text_manifest is not None and tile:
raise ValueError("Verified text restoration is not calibrated with tiled diffusion")
+6 -4
View File
@@ -162,7 +162,7 @@ def _resolved_strength_for_display(
) -> float:
"""Resolve the same profile-specific strength the engine will execute.
One call for both profiles, so the printed value cannot drift from the executed
One call for every profile, so the printed value cannot drift from the executed
one; the size is what qwen-zimage derives its strength from.
"""
from PIL import Image
@@ -254,8 +254,10 @@ _PIPELINE_CHOICES = list(PROFILE_CHOICES)
_PIPELINE_HELP = (
"Pipeline profile. qwen-zimage (DEFAULT) = Qwen-Image-2512 + Lightning + Canny, "
"followed by SAM-masked Z-Image face repair; sdxl-zimage = the same recipe and the "
"same face stage on an SDXL global pass, which needs more denoise. Both are "
"CUDA-ONLY -- install the qwen-zimage extra. There is no CPU or MPS profile for "
"same face stage on an SDXL global pass, which needs more denoise; chroma-zimage = "
"the same face stage on an Apache-2.0 Chroma1 global pass with lower OpenAI and "
"Microsoft floors and higher Google/Meta floors (docs/chroma1-engine-research.md). "
"All are CUDA-ONLY -- install the qwen-zimage extra. There is no CPU or MPS profile for "
"invisible-watermark removal."
)
@@ -304,7 +306,7 @@ _seed_option = click.option(
"--seed",
type=int,
default=None,
help="Random seed for reproducibility. Default 0: both profiles are certified "
help="Random seed for reproducibility. Default 0: all profiles are certified "
"at a fixed seed, because SynthID removal near the strength floor is seed-dependent.",
)
_hf_token_option = click.option("--hf-token", type=str, default=None, help="Hugging Face API token.")
+10 -8
View File
@@ -95,19 +95,21 @@ class InvisibleEngine:
"""Initialize the invisible watermark removal engine.
Args:
device: Device for inference. Both profiles are CUDA-only, so the
device: Device for inference. All profiles are CUDA-only, so the
usable values are "cuda" and None/"auto" (which detects it);
anything else raises rather than falling back.
pipeline: Pipeline profile, one of "qwen-zimage" (DEFAULT;
Qwen-Image-2512 Lightning + Canny, then SAM-masked Z-Image face repair)
or "sdxl-zimage" (the same recipe and the same face stage on an SDXL
global pass, vendor-adaptive strength because an SDXL global stage
needs more of it). BOTH ARE CUDA-ONLY -- there is no CPU or MPS path
for invisible-watermark removal.
needs more of it) or "chroma-zimage" (the same face stage on an
Apache-2.0 Chroma1 global pass with its own flat vendor floors;
see docs/chroma1-engine-research.md). ALL ARE CUDA-ONLY -- there is
no CPU or MPS path for invisible-watermark removal.
hf_token: Hugging Face API token.
progress_callback: Optional callback for progress messages.
controlnet_conditioning_scale: Canny ControlNet structure-preservation
strength on the global stage of both profiles.
strength on the global stage (not used by chroma-zimage).
cpu_offload: Offload model components to CPU between CUDA calls instead
of keeping the whole pipeline in VRAM, at the cost of speed. For
qwen-zimage, force the face stack to offload instead of using automatic
@@ -158,9 +160,9 @@ class InvisibleEngine:
image_path: Path to the watermarked image.
output_path: Output path (None = overwrite source).
strength: Denoising strength (0.0-1.0). None -> the profile's calibrated
default (resolution-adaptive for qwen-zimage, vendor-adaptive for
sdxl-zimage).
seed: Random seed for reproducibility. None resolves to 0, because both
default (resolution-adaptive for qwen-zimage, vendor-adaptive flat
floors for sdxl-zimage and chroma-zimage).
seed: Random seed for reproducibility. None resolves to 0, because all
profiles are certified at a fixed seed.
humanize: Intensity of Analog Humanizer film grain (0 = off).
unsharp: Final unsharp-mask sharpening strength (0 = off, default).
@@ -246,7 +248,7 @@ class InvisibleEngine:
verified_text = load_verified_text_manifest(text_manifest, reference_pil)
# Both profiles run at the input's native geometry, so only the explicit max
# All profiles run at the input's native geometry, so only the explicit max
# cap can move it, and it can only ever scale down.
target = _target_size(image.width, image.height, max_resolution)
if target is not None:
+2 -1
View File
@@ -25,7 +25,7 @@ if TYPE_CHECKING:
from remove_ai_watermarks.video_visible import VideoScan
VIDEO_EXTENSIONS: frozenset[str] = frozenset({".mp4", ".mov", ".m4v", ".webm", ".mkv", ".avi", ".flv"})
VIDEO_VISIBLE_MARKS = ("sora", "veo", "seedance", "dola", "hailuo", "kling")
VIDEO_VISIBLE_MARKS = ("sora", "veo", "seedance", "doubao", "dola", "hailuo", "kling")
_ISOBMFF_VIDEO_EXTENSIONS: frozenset[str] = frozenset({".mp4", ".mov", ".m4v"})
_EBML_VIDEO_EXTENSIONS: frozenset[str] = frozenset({".webm", ".mkv"})
_RIFF_VIDEO_EXTENSIONS: frozenset[str] = frozenset({".avi"})
@@ -178,6 +178,7 @@ _VISIBLE_PLATFORM = {
"sora": "OpenAI Sora",
"veo": "Google Veo",
"seedance": "ByteDance Seedance",
"doubao": "ByteDance Doubao",
"dola": "ByteDance Dola",
"hailuo": "MiniMax Hailuo AI",
"kling": "Kuaishou Kling AI",
+68
View File
@@ -313,6 +313,28 @@ def _top_hat(gray: NDArray[Any]) -> NDArray[Any]:
@lru_cache(maxsize=1)
def _asset_template(name: str) -> NDArray[Any]:
"""Load a package asset as a video detect template (white glyph on black).
The image-registry alpha maps are the same synthetic silhouettes the image
engines match; a video mark that shares the glyph reuses the asset instead of
re-rendering it. No source frame contributes pixels.
"""
from pathlib import Path as _P
from remove_ai_watermarks.image_io import imread as _imread
at = _imread(str(_P(__file__).parent / "assets" / name), 2)
if at is None:
raise RuntimeError(f"missing video template asset: {name}")
return at
@lru_cache(maxsize=1)
def _doubao_template() -> NDArray[Any]:
return _asset_template("doubao_alpha.png")
def _template_sources() -> dict[str, NDArray[Any]]:
"""Return every immutable synthetic template by detector-local key."""
sora_word, sora_icon = _sora_templates()
@@ -326,6 +348,9 @@ def _template_sources() -> dict[str, NDArray[Any]]:
"dola": _dola_template(),
"hailuo": _hailuo_template(),
"kling-logo": _kling_logo_template(),
# The Doubao video mark is the same "豆包AI生成" run the image engine
# matches (its synthetic alpha asset), bottom-right on video frames.
"doubao": _doubao_template(),
}
sources.update({f"kling-{index}": template for index, template in enumerate(_kling_templates())})
return sources
@@ -606,6 +631,32 @@ def detect_seedance_frame(
)
def detect_doubao_frame(
image_bgr: NDArray[Any],
*,
frame_index: int = 0,
prepared: _PreparedFrame | None = None,
) -> FrameLocalization:
"""Locate the strongest fixed Doubao "豆包AI生成" text candidate.
The ByteDance Doubao mark registered for IMAGES also appears burned into the
bottom-right of Doubao VIDEO exports; this detector reuses the image engine's
synthetic alpha as the template. Search profile measured on the local spaces
corpus (48 TC260-confirmed Doubao videos, 2026-08-28): the run sits at
cx 0.83-0.90 / bottom-flush, heights swept 3.2-6.0 percent of the short side.
"""
return _detect_fixed_mark(
image_bgr,
"doubao",
relative_heights=tuple(value / 1000 for value in range(32, 61, 4)),
search_origin=(0.60, 0.76),
kernel_fraction=0.14,
prefer_larger_within=0.04,
normalized=None if prepared is None else (prepared.normalized_gray, prepared.normalized_scale),
frame_index=frame_index,
)
def detect_dola_frame(
image_bgr: NDArray[Any],
*,
@@ -952,6 +1003,22 @@ VISIBLE_MARK_POLICIES: dict[str, VisibleMarkPolicy] = {
padding_fraction=0.0,
mask_style="box",
),
"doubao": VisibleMarkPolicy(
# One floor at either trust level (seedance-style: the TC260 producer
# confirms the vendor inside the arbiter; the confidence bar does not
# move). Floors measured 2026-08-28 on the local spaces corpus with
# sequential decode: a >=0.35 stable run of 12+ frames accepts 33/39
# TC260-confirmed Doubao videos and 0/25 no-Doubao negatives; 0.30
# already accepted 1 negative, so the bar stays at 0.35.
weak_floor=0.35,
strong_floor=0.55,
transition_floor=0.30,
min_stable_frames=_MIN_FIXED_MARK_STABLE_FRAMES,
cover_after_confirmation=True,
anchor_iou=0.80,
padding_fraction=0.0,
mask_style="box",
),
"dola": VisibleMarkPolicy(
weak_floor=_DOLA_STRICT_WEAK_CONFIDENCE,
provenance_weak_floor=_DOLA_PROVENANCE_WEAK_CONFIDENCE,
@@ -1181,6 +1248,7 @@ def scan_video_marks(
detect_sora_frame,
detect_veo_frame,
detect_seedance_frame,
detect_doubao_frame,
detect_dola_frame,
detect_hailuo_frame,
detect_kling_frame,
+12
View File
@@ -1489,6 +1489,18 @@ class TestVendorOf:
assert _vendor_of("Eleven Labs Inc.") == "ElevenLabs"
assert _vendor_of("Ideogram") == "Ideogram"
def test_ideogram_fixture_attributed_end_to_end(self):
# Committed synthetic manifest signed by a test CA with CN "Ideogram, Inc":
# attribution must name the platform while signer trust stays untrusted.
# Reachability against the fixture, not a synthesized status set.
fixture = Path(__file__).resolve().parents[1] / "data" / "fixtures" / "provenance" / "ideogram-c2pa.png"
if not fixture.exists(): # the sdist ships no data/
pytest.skip("provenance fixture not present")
report = identify(fixture, check_visible=False, check_invisible=False)
assert report.ai_from_metadata is True
assert report.platform == "Ideogram"
assert any("Ideogram" in w for w in report.watermarks)
def test_ideogram_issuer_attributed(self):
# Corpus evidence 2026-08-08: four uploads signed "Ideogram, Inc" read as
# unknown-signer C2PA with no platform. The issuer token is the org prefix.
+1 -2
View File
@@ -359,8 +359,7 @@ class TestC2paMarkerIn:
class TestSamsungGenai:
"""Samsung Galaxy AI editing marker (genAIType in PhotoEditor_Re_Edit_Data).
Synthetic byte blobs -- real Galaxy files are user content and not shipped
(public repo), same discipline as the Grok/Doubao fixtures.
Synthetic byte blobs keep this fixture self-contained and deterministic.
"""
@staticmethod
+1 -1
View File
@@ -194,7 +194,7 @@ class TestNoReembeddedWatermark:
pipeline = SdxlZImagePipeline(device="cuda", torch_dtype=None)
monkeypatch.setattr(type(pipeline), "_require_cuda", lambda self: None)
pipeline._load_sdxl()
pipeline._load_global()
assert calls["pipeline"].get("add_watermarker") is False
assert "add_watermarker" not in calls["controlnet"]
+310 -57
View File
@@ -24,14 +24,14 @@ def _mock_watermark_runtime_deps(monkeypatch):
def test_resolution_adaptive_denoise_preserves_calibrated_values():
from remove_ai_watermarks._internal.qwen_zimage_pipeline import resolution_adaptive_denoise
from remove_ai_watermarks._internal.two_stage_pipeline import resolution_adaptive_denoise
assert resolution_adaptive_denoise(600, 500, adaptive_level=6) == pytest.approx(0.084)
assert resolution_adaptive_denoise(2000, 1850, adaptive_level=6) == pytest.approx(0.154)
def test_largest_face_denoise_preserves_calibrated_values():
from remove_ai_watermarks._internal.qwen_zimage_pipeline import largest_face_denoise
from remove_ai_watermarks._internal.two_stage_pipeline import largest_face_denoise
image_size = (1000, 1000)
assert largest_face_denoise([(0, 0, 300, 100)], image_size) == pytest.approx(0.10)
@@ -59,7 +59,7 @@ def test_global_kwargs_use_lightning_and_diffsynth_controlnet_shape():
def test_face_kwargs_use_project_zimage_settings():
from remove_ai_watermarks._internal.qwen_zimage_pipeline import build_face_kwargs
from remove_ai_watermarks._internal.two_stage_pipeline import build_face_kwargs
crop = Image.new("RGB", (713, 941))
kwargs = build_face_kwargs(crop, strength=0.17, seed=9)
@@ -76,7 +76,7 @@ def test_face_kwargs_use_project_zimage_settings():
def test_canny_control_image_is_three_channel_and_detects_an_edge():
from remove_ai_watermarks._internal.qwen_zimage_pipeline import build_canny_control_image
from remove_ai_watermarks._internal.two_stage_pipeline import build_canny_control_image
source = np.zeros((64, 80, 3), dtype=np.uint8)
source[:, 40:] = 255
@@ -92,14 +92,15 @@ def test_canny_control_image_is_three_channel_and_detects_an_edge():
def test_face_crop_geometry_preserves_calibrated_values():
from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline, _expanded_box
from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline
from remove_ai_watermarks._internal.two_stage_pipeline import _expanded_box
assert _expanded_box((100, 100, 200, 200), (500, 500)) == (25, 25, 275, 275)
assert QwenZImagePipeline._detail_size((500, 400), (100, 80)) == (1024, 816)
def test_yunet_download_targets_verified_lfs_artifact():
from remove_ai_watermarks._internal.qwen_zimage_pipeline import (
from remove_ai_watermarks._internal.two_stage_pipeline import (
YUNET_MODEL_SHA256,
YUNET_MODEL_URL,
YUNET_SCORE_THRESHOLD,
@@ -239,8 +240,8 @@ def test_face_stage_loads_in_its_own_dtype_when_the_global_stage_differs(monkeyp
def test_resident_face_models_disable_vram_offload():
from remove_ai_watermarks._internal.qwen_zimage_pipeline import (
QwenZImagePipeline,
from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline
from remove_ai_watermarks._internal.two_stage_pipeline import (
_pin_vram_managed_models,
resolve_face_model_residency,
)
@@ -291,7 +292,7 @@ def test_resident_face_models_disable_vram_offload():
def test_static_prompt_cache_reuses_embeddings_without_caching_image_edits():
from remove_ai_watermarks._internal.qwen_zimage_pipeline import _cache_static_prompt_embeddings
from remove_ai_watermarks._internal.two_stage_pipeline import _cache_static_prompt_embeddings
class PromptUnit:
output_params = ("prompt_embeds",)
@@ -322,24 +323,24 @@ def test_static_prompt_cache_reuses_embeddings_without_caching_image_edits():
def test_prompt_cache_path_is_keyed_by_version_model_outputs_and_prompt(monkeypatch, tmp_path):
"""A model, prompt or format change must not read a stale embedding."""
from remove_ai_watermarks._internal import qwen_zimage_pipeline as qz
from remove_ai_watermarks._internal import two_stage_pipeline as ts
monkeypatch.setenv("HF_HOME", str(tmp_path))
baseline = qz._prompt_cache_path("model/a", ("prompt_emb",), "text")
baseline = ts._prompt_cache_path("model/a", ("prompt_emb",), "text")
assert baseline.parent == tmp_path / "remove-ai-watermarks" / "prompt-embeddings"
assert baseline == qz._prompt_cache_path("model/a", ("prompt_emb",), "text")
assert baseline != qz._prompt_cache_path("model/b", ("prompt_emb",), "text")
assert baseline != qz._prompt_cache_path("model/a", ("prompt_embeds",), "text")
assert baseline != qz._prompt_cache_path("model/a", ("prompt_emb",), "other")
assert baseline == ts._prompt_cache_path("model/a", ("prompt_emb",), "text")
assert baseline != ts._prompt_cache_path("model/b", ("prompt_emb",), "text")
assert baseline != ts._prompt_cache_path("model/a", ("prompt_embeds",), "text")
assert baseline != ts._prompt_cache_path("model/a", ("prompt_emb",), "other")
monkeypatch.setattr(qz, "_PROMPT_CACHE_VERSION", qz._PROMPT_CACHE_VERSION + 1)
assert baseline != qz._prompt_cache_path("model/a", ("prompt_emb",), "text")
monkeypatch.setattr(ts, "_PROMPT_CACHE_VERSION", ts._PROMPT_CACHE_VERSION + 1)
assert baseline != ts._prompt_cache_path("model/a", ("prompt_emb",), "text")
def test_model_cache_dir_prefers_the_persistent_hugging_face_root(monkeypatch, tmp_path):
"""A scale-to-zero runner only mounts HF_HOME, so it must win over XDG."""
from remove_ai_watermarks._internal.qwen_zimage_pipeline import _model_cache_dir
from remove_ai_watermarks._internal.two_stage_pipeline import _model_cache_dir
monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "xdg"))
monkeypatch.delenv("HF_HOME", raising=False)
@@ -353,16 +354,17 @@ def test_stored_prompt_embedding_round_trips_without_casting_the_mask(monkeypatc
"""The mask rides in the same payload and is integer; casting it corrupts the prompt."""
import torch
from remove_ai_watermarks._internal import qwen_zimage_pipeline as qz
from remove_ai_watermarks._internal import two_stage_pipeline as ts
from remove_ai_watermarks._internal.qwen_zimage_pipeline import _QWEN_PROMPT_OUTPUTS
monkeypatch.setenv("HF_HOME", str(tmp_path))
path = qz._prompt_cache_path("model/a", qz._QWEN_PROMPT_OUTPUTS, "text")
path = ts._prompt_cache_path("model/a", _QWEN_PROMPT_OUTPUTS, "text")
payload = {
"prompt_emb": torch.ones((1, 2, 3), dtype=torch.float32),
"prompt_emb_mask": torch.ones((1, 2), dtype=torch.int64),
}
qz._store_prompt_payload(path, payload)
restored = qz._load_prompt_payload(path, "cpu", torch.bfloat16)
ts._store_prompt_payload(path, payload)
restored = ts._load_prompt_payload(path, "cpu", torch.bfloat16)
assert restored["prompt_emb"].dtype == torch.bfloat16
assert restored["prompt_emb_mask"].dtype == torch.int64
@@ -373,12 +375,12 @@ def test_persisted_prompt_cache_lets_a_second_pipeline_skip_the_text_encoder(mon
"""The whole point: container two must not call the encoder container one ran."""
import torch
from remove_ai_watermarks._internal import qwen_zimage_pipeline as qz
from remove_ai_watermarks._internal import two_stage_pipeline as ts
monkeypatch.setenv("HF_HOME", str(tmp_path))
class PromptUnit:
output_params = qz._ZIMAGE_PROMPT_OUTPUTS
output_params = ts._ZIMAGE_PROMPT_OUTPUTS
def __init__(self):
self.calls = 0
@@ -393,13 +395,13 @@ def test_persisted_prompt_cache_lets_a_second_pipeline_skip_the_text_encoder(mon
return unit, pipe
first_unit, first_pipe = build()
qz._cache_static_prompt_embeddings(first_pipe, qz._ZIMAGE_PROMPT_OUTPUTS, model_id="model/a", require_cache=False)
first_unit.process(first_pipe, qz._FACE_PROMPT)
ts._cache_static_prompt_embeddings(first_pipe, ts._ZIMAGE_PROMPT_OUTPUTS, model_id="model/a", require_cache=False)
first_unit.process(first_pipe, ts._FACE_PROMPT)
assert first_unit.calls == 1
second_unit, second_pipe = build()
qz._cache_static_prompt_embeddings(second_pipe, qz._ZIMAGE_PROMPT_OUTPUTS, model_id="model/a", require_cache=True)
restored = second_unit.process(second_pipe, qz._FACE_PROMPT)
ts._cache_static_prompt_embeddings(second_pipe, ts._ZIMAGE_PROMPT_OUTPUTS, model_id="model/a", require_cache=True)
restored = second_unit.process(second_pipe, ts._FACE_PROMPT)
assert second_unit.calls == 0
assert torch.equal(restored["prompt_embeds"][0], torch.ones((2, 2)))
@@ -409,19 +411,20 @@ def test_a_missing_cache_fails_loudly_once_the_text_encoder_was_left_out(monkeyp
"""Silently calling an absent text encoder would surface as an opaque crash."""
import torch
from remove_ai_watermarks._internal import qwen_zimage_pipeline as qz
from remove_ai_watermarks._internal import two_stage_pipeline as ts
from remove_ai_watermarks._internal.qwen_zimage_pipeline import _QWEN_PROMPT_OUTPUTS
monkeypatch.setenv("HF_HOME", str(tmp_path))
class PromptUnit:
output_params = qz._QWEN_PROMPT_OUTPUTS
output_params = _QWEN_PROMPT_OUTPUTS
def process(self, _pipe, prompt, edit_image=None):
raise AssertionError("the text encoder is not loaded")
unit = PromptUnit()
pipe = MagicMock(units=[unit], device="cpu", torch_dtype=torch.float32)
qz._cache_static_prompt_embeddings(pipe, qz._QWEN_PROMPT_OUTPUTS, model_id="model/a", require_cache=True)
ts._cache_static_prompt_embeddings(pipe, _QWEN_PROMPT_OUTPUTS, model_id="model/a", require_cache=True)
with pytest.raises(RuntimeError, match="disappeared"):
unit.process(pipe, "never cached")
@@ -430,7 +433,7 @@ def test_a_missing_cache_fails_loudly_once_the_text_encoder_was_left_out(monkeyp
def test_sam_pixels_match_model_dtype_without_casting_boxes():
import torch
from remove_ai_watermarks._internal.qwen_zimage_pipeline import _prepare_sam_inputs
from remove_ai_watermarks._internal.two_stage_pipeline import _prepare_sam_inputs
class Inputs(dict[str, torch.Tensor]):
def to(self, device: str):
@@ -450,7 +453,7 @@ def test_sam_pixels_match_model_dtype_without_casting_boxes():
def test_sam_prompts_match_impact_center_and_clip_masks_to_boxes():
from remove_ai_watermarks._internal.qwen_zimage_pipeline import (
from remove_ai_watermarks._internal.two_stage_pipeline import (
_clip_sam_masks_to_boxes,
_sam_point_prompts,
)
@@ -470,7 +473,7 @@ def test_sam_prompts_match_impact_center_and_clip_masks_to_boxes():
def test_sam_proposal_selection_matches_impact_sub_threshold():
from remove_ai_watermarks._internal.qwen_zimage_pipeline import _select_sam_masks
from remove_ai_watermarks._internal.two_stage_pipeline import _select_sam_masks
masks = np.zeros((2, 3, 8, 8), dtype=np.float32)
masks[0, 0, 1:3, 1:3] = 1.0
@@ -500,7 +503,7 @@ def test_sam_proposal_selection_matches_impact_sub_threshold():
def test_sam_bfloat16_outputs_convert_to_numpy_float32():
import torch
from remove_ai_watermarks._internal.qwen_zimage_pipeline import _sam_outputs_to_numpy
from remove_ai_watermarks._internal.two_stage_pipeline import _sam_outputs_to_numpy
masks = torch.ones((1, 2, 3, 4, 4), dtype=torch.bfloat16)
scores = torch.tensor([[[0.95, 0.75, 0.50], [0.99, 0.80, 0.60]]], dtype=torch.bfloat16)
@@ -513,7 +516,7 @@ def test_sam_bfloat16_outputs_convert_to_numpy_float32():
def test_face_composite_preserves_every_pixel_outside_mask():
from remove_ai_watermarks._internal.qwen_zimage_pipeline import composite_face
from remove_ai_watermarks._internal.two_stage_pipeline import composite_face
base = np.full((32, 32, 3), 10, dtype=np.uint8)
detail = np.full((32, 32, 3), 240, dtype=np.uint8)
@@ -618,7 +621,7 @@ def test_cli_reports_verified_text_manifest_errors(tmp_image_path, tmp_path, mon
def test_no_face_path_still_runs_verified_text_restoration(monkeypatch):
from remove_ai_watermarks._internal import qwen_zimage_pipeline, text_restoration
from remove_ai_watermarks._internal import text_restoration, two_stage_pipeline
from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline
from remove_ai_watermarks._internal.text_restoration import VerifiedTextLine, VerifiedTextManifest
@@ -630,9 +633,9 @@ def test_no_face_path_still_runs_verified_text_restoration(monkeypatch):
global_result = Image.new("RGB", (32, 32), (70, 80, 90))
anchor = Image.new("RGB", (32, 32), (100, 110, 120))
restored = Image.new("RGB", (32, 32), (130, 140, 150))
pipeline._qwen_vae_roundtrip = MagicMock(return_value=donor)
pipeline._vae_roundtrip = MagicMock(return_value=donor)
pipeline._run_global = MagicMock(return_value=global_result)
monkeypatch.setattr(qwen_zimage_pipeline, "detect_faces", lambda _image: [])
monkeypatch.setattr(two_stage_pipeline, "detect_faces", lambda _image: [])
blend = MagicMock(return_value=anchor)
restore = MagicMock(return_value=restored)
monkeypatch.setattr(text_restoration, "blend_fidelity_anchor", blend)
@@ -767,11 +770,9 @@ def test_watermark_remover_rejects_tiled_verified_text_before_pipeline(tmp_path,
def test_qwen_tiling_runs_global_tiles_then_one_full_frame_face_stage(monkeypatch):
from remove_ai_watermarks._internal.qwen_zimage_pipeline import (
QwenZImagePipeline,
resolution_adaptive_denoise,
)
from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline
from remove_ai_watermarks._internal.tiling import plan_tiles
from remove_ai_watermarks._internal.two_stage_pipeline import resolution_adaptive_denoise
image = Image.new("RGB", (1500, 1500), (20, 30, 40))
runtime = QwenZImagePipeline(device="cuda", torch_dtype="bf16")
@@ -787,7 +788,7 @@ def test_qwen_tiling_runs_global_tiles_then_one_full_frame_face_stage(monkeypatc
monkeypatch.setattr(runtime, "_run_global", fake_global)
monkeypatch.setattr(runtime, "_run_faces", face_stage)
monkeypatch.setattr(
"remove_ai_watermarks._internal.qwen_zimage_pipeline.detect_faces",
"remove_ai_watermarks._internal.two_stage_pipeline.detect_faces",
lambda _image: [(100, 100, 300, 300)],
)
monkeypatch.setattr(runtime, "_sam_masks", lambda _image, _boxes: [np.ones((1500, 1500), dtype=np.uint8)])
@@ -817,21 +818,21 @@ def test_global_only_preload_skips_face_models(monkeypatch):
from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline
runtime = QwenZImagePipeline(device="cuda", torch_dtype="bf16")
qwen = MagicMock()
global_stage = MagicMock()
zimage = MagicMock()
sam = MagicMock()
yunet = MagicMock()
monkeypatch.setattr(runtime, "_load_qwen", qwen)
monkeypatch.setattr(runtime, "_load_global", global_stage)
monkeypatch.setattr(runtime, "_load_zimage", zimage)
monkeypatch.setattr(runtime, "_load_sam", sam)
monkeypatch.setattr(
"remove_ai_watermarks._internal.qwen_zimage_pipeline._yunet_model_path",
"remove_ai_watermarks._internal.two_stage_pipeline._yunet_model_path",
yunet,
)
runtime.preload(global_only=True)
qwen.assert_called_once_with()
global_stage.assert_called_once_with()
zimage.assert_not_called()
sam.assert_not_called()
yunet.assert_called_once_with()
@@ -841,21 +842,21 @@ def test_full_preload_still_loads_face_models(monkeypatch):
from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline
runtime = QwenZImagePipeline(device="cuda", torch_dtype="bf16")
qwen = MagicMock()
global_stage = MagicMock()
zimage = MagicMock()
sam = MagicMock()
yunet = MagicMock()
monkeypatch.setattr(runtime, "_load_qwen", qwen)
monkeypatch.setattr(runtime, "_load_global", global_stage)
monkeypatch.setattr(runtime, "_load_zimage", zimage)
monkeypatch.setattr(runtime, "_load_sam", sam)
monkeypatch.setattr(
"remove_ai_watermarks._internal.qwen_zimage_pipeline._yunet_model_path",
"remove_ai_watermarks._internal.two_stage_pipeline._yunet_model_path",
yunet,
)
runtime.preload()
qwen.assert_called_once_with()
global_stage.assert_called_once_with()
zimage.assert_called_once_with()
sam.assert_called_once_with()
yunet.assert_called_once_with()
@@ -964,8 +965,8 @@ def test_sdxl_requested_steps_compensate_for_the_diffusers_truncation():
def test_sdxl_zimage_floors_to_its_own_latent_grid():
"""SDXL aligns to 8 pixels where Qwen aligns to 16."""
from remove_ai_watermarks._internal.qwen_zimage_pipeline import _target_size
from remove_ai_watermarks._internal.sdxl_zimage_pipeline import sdxl_target_size
from remove_ai_watermarks._internal.two_stage_pipeline import _target_size
assert sdxl_target_size(1122, 1402) == (1120, 1400)
assert _target_size(1122, 1402) == (1120, 1392)
@@ -973,12 +974,264 @@ def test_sdxl_zimage_floors_to_its_own_latent_grid():
def test_sdxl_zimage_inherits_the_face_stage_rather_than_copying_it():
"""The face stage must not be able to diverge between the two profiles."""
"""The face stage must not be able to diverge between profiles.
Both profiles are siblings of the shared two-stage base, not a chain: an SDXL or
Flux profile must never inherit from the Qwen profile, and the shared stages live
on the base so a change there reaches every profile at once.
"""
from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline
from remove_ai_watermarks._internal.sdxl_zimage_pipeline import SdxlZImagePipeline
from remove_ai_watermarks._internal.two_stage_pipeline import TwoStageZImagePipeline
assert issubclass(SdxlZImagePipeline, QwenZImagePipeline)
for shared in ("_run_faces", "_sam_masks", "_load_zimage", "_load_sam", "run"):
assert getattr(SdxlZImagePipeline, shared) is getattr(QwenZImagePipeline, shared)
assert issubclass(SdxlZImagePipeline, TwoStageZImagePipeline)
assert issubclass(QwenZImagePipeline, TwoStageZImagePipeline)
assert QwenZImagePipeline not in SdxlZImagePipeline.__mro__
for shared in ("_run_faces", "_sam_masks", "_load_zimage", "_load_sam", "run", "preload"):
base_method = getattr(TwoStageZImagePipeline, shared)
assert getattr(SdxlZImagePipeline, shared) is base_method
assert getattr(QwenZImagePipeline, shared) is base_method
# Only the global stage and what it needs may differ.
assert SdxlZImagePipeline._run_global is not QwenZImagePipeline._run_global
assert SdxlZImagePipeline._load_global is not QwenZImagePipeline._load_global
def test_chroma_zimage_strength_uses_measured_flat_floors():
"""The chroma floors come from the four-cohort oracle calibration; they must not
drift from their derivation in watermark_profiles."""
from remove_ai_watermarks._internal.watermark_profiles import (
CHROMA_ZIMAGE_GOOGLE_STRENGTH,
CHROMA_ZIMAGE_META_STRENGTH,
CHROMA_ZIMAGE_MICROSOFT_STRENGTH,
CHROMA_ZIMAGE_OPENAI_STRENGTH,
CHROMA_ZIMAGE_UNKNOWN_STRENGTH,
normalize_profile,
resolve_strength,
)
assert normalize_profile("chroma_zimage") == "chroma-zimage"
assert resolve_strength(None, "openai", "chroma-zimage") == pytest.approx(CHROMA_ZIMAGE_OPENAI_STRENGTH)
assert pytest.approx(0.09) == CHROMA_ZIMAGE_OPENAI_STRENGTH
assert resolve_strength(None, "microsoft", "chroma-zimage") == pytest.approx(CHROMA_ZIMAGE_MICROSOFT_STRENGTH)
assert pytest.approx(0.125) == CHROMA_ZIMAGE_MICROSOFT_STRENGTH
assert resolve_strength(None, "google", "chroma-zimage") == pytest.approx(CHROMA_ZIMAGE_GOOGLE_STRENGTH)
assert pytest.approx(0.40) == CHROMA_ZIMAGE_GOOGLE_STRENGTH
assert resolve_strength(None, "meta", "chroma-zimage") == pytest.approx(CHROMA_ZIMAGE_META_STRENGTH)
assert pytest.approx(0.17) == CHROMA_ZIMAGE_META_STRENGTH
# Unknown follows the strictest measured cohort, as with the other profiles.
assert resolve_strength(None, None, "chroma-zimage") == pytest.approx(CHROMA_ZIMAGE_UNKNOWN_STRENGTH)
assert CHROMA_ZIMAGE_UNKNOWN_STRENGTH == CHROMA_ZIMAGE_GOOGLE_STRENGTH
# An explicit strength still wins, and other profiles are untouched.
assert resolve_strength(0.2, "google", "chroma-zimage") == pytest.approx(0.2)
assert resolve_strength(None, "openai", "qwen-zimage", size=(2000, 1850)) == pytest.approx(0.07675)
def test_chroma_zimage_inherits_the_shared_stages_and_only_swaps_the_global():
"""Like sdxl-zimage: a sibling of the base, never a chain, with the shared
stages inherited verbatim."""
from remove_ai_watermarks._internal.chroma_zimage_pipeline import ChromaZImagePipeline
from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline
from remove_ai_watermarks._internal.two_stage_pipeline import TwoStageZImagePipeline
assert issubclass(ChromaZImagePipeline, TwoStageZImagePipeline)
assert QwenZImagePipeline not in ChromaZImagePipeline.__mro__
for shared in ("_run_faces", "_sam_masks", "_load_zimage", "_load_sam", "run", "preload"):
base_method = getattr(TwoStageZImagePipeline, shared)
assert getattr(ChromaZImagePipeline, shared) is base_method
assert ChromaZImagePipeline._run_global is not QwenZImagePipeline._run_global
assert ChromaZImagePipeline._load_global is not QwenZImagePipeline._load_global
assert ChromaZImagePipeline.profile_name == "chroma-zimage"
def test_chroma_requested_steps_compensate_for_the_diffusers_truncation():
from remove_ai_watermarks._internal.chroma_zimage_pipeline import CHROMA_STEPS, requested_steps
for strength in (0.09, 0.125, 0.17, 0.40):
steps = requested_steps(CHROMA_STEPS, strength)
assert int(steps * strength) >= CHROMA_STEPS
assert int(CHROMA_STEPS * strength) < CHROMA_STEPS
def test_chroma_zimage_uses_its_own_calibrated_prompts():
"""The floors were measured with the neutral faithful prompt, not the shared
canny-stage prompt; swapping them silently moves the oracle boundaries."""
from remove_ai_watermarks._internal import chroma_zimage_pipeline as chroma_mod
from remove_ai_watermarks._internal.two_stage_pipeline import _GLOBAL_PROMPT
assert chroma_mod.CHROMA_PROMPT != _GLOBAL_PROMPT
assert chroma_mod.CHROMA_PROMPT == "high quality, sharp, detailed, faithful to the original"
assert chroma_mod.CHROMA_GUIDANCE == 5.0
assert chroma_mod.CHROMA_STEPS == 4
def test_chroma_global_stage_calls_diffusers_with_the_calibrated_shape(monkeypatch):
"""The calibrated call shape must not drift: prompt, guidance, width/height
on the /16 grid, effective-step count, and no Canny conditioning. A silent
parameter change moves the oracle floors without any test seeing it."""
# CI's base job installs the library without the qwen-zimage extra.
diffusers = pytest.importorskip("diffusers")
import torch
from remove_ai_watermarks._internal.chroma_zimage_pipeline import (
CHROMA_GUIDANCE,
ChromaZImagePipeline,
)
calls: dict[str, dict] = {}
def fake_chroma(*_args, **kwargs):
calls["chroma"] = kwargs
from PIL import Image
return type("R", (), {"images": [Image.new("RGB", (kwargs["width"], kwargs["height"]))]})()
class FakeChromaPipe:
def to(self, _device):
return self
@staticmethod
def __call__(*args, **kwargs):
return fake_chroma(*args, **kwargs)
monkeypatch.setattr(
diffusers.ChromaImg2ImgPipeline,
"from_pretrained",
staticmethod(lambda *a, **k: FakeChromaPipe()),
)
# _run_global creates a torch.Generator(device="cuda"), which needs a real
# CUDA library; patch the module's torch reference to yield a CPU-safe stub.
class FakeGenerator:
def manual_seed(self, _seed):
return self
monkeypatch.setattr("torch.Generator", lambda device=None: FakeGenerator())
pipeline = ChromaZImagePipeline(device="cpu", torch_dtype=torch.bfloat16)
monkeypatch.setattr(type(pipeline), "_require_cuda", lambda self: None)
result = pipeline._run_global(Image.new("RGB", (1122, 1402)), 0.09, 0)
assert result.size == (1122, 1402)
call = calls["chroma"]
# The calibrated prompt, NOT the shared canny-stage prompt.
assert call["prompt"] == "high quality, sharp, detailed, faithful to the original"
assert call["negative_prompt"] == "blurry, lowres, distorted text, garbled text, artifacts"
assert call["guidance_scale"] == CHROMA_GUIDANCE
# The /16 latent grid, not the raw input size.
assert call["width"] == 1120
assert call["height"] == 1392
# Effective-step compensation: ceil(4 / 0.09) = 45 requested.
assert call["num_inference_steps"] == 45
assert call["strength"] == 0.09
# No Canny conditioning: the calibrated path is a plain strength pass.
assert "control_image" not in call
assert "controlnet_conditioning_scale" not in call
def test_chroma_zimage_floors_to_its_own_latent_grid():
"""Chroma1 floors to /16, the same grid as Qwen but unlike SDXL's /8."""
from remove_ai_watermarks._internal.chroma_zimage_pipeline import chroma_target_size
from remove_ai_watermarks._internal.sdxl_zimage_pipeline import sdxl_target_size
assert chroma_target_size(1122, 1402) == (1120, 1392)
assert sdxl_target_size(1122, 1402) == (1120, 1400)
assert chroma_target_size(3, 3) == (16, 16)
def test_watermark_remover_rejects_text_manifest_for_chroma(tmp_path, monkeypatch):
"""Verified text restoration needs the Qwen VAE donor; chroma-zimage and
sdxl-zimage must both be rejected at the boundary, not deep inside."""
from remove_ai_watermarks._internal.text_restoration import VerifiedTextLine, VerifiedTextManifest
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
_mock_watermark_runtime_deps(monkeypatch)
source = tmp_path / "source.png"
Image.new("RGB", (96, 80)).save(source)
manifest = VerifiedTextManifest("0" * 64, 96, 80, (VerifiedTextLine((4, 4, 20, 16), "x", "alphabetic"),))
for profile in ("chroma-zimage", "sdxl-zimage"):
remover = WatermarkRemover(device="cuda", pipeline=profile)
with pytest.raises(ValueError, match="only by the qwen-zimage profile"):
remover.remove_watermark(source, text_manifest=manifest)
def test_watermark_remover_dispatches_to_chroma_pipeline(monkeypatch):
"""The remover must construct ChromaZImagePipeline, not fall through to qwen."""
from remove_ai_watermarks._internal import chroma_zimage_pipeline as chroma_module
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
captured: dict[str, object] = {}
class Recorder:
def __init__(self, **kwargs):
captured.update(kwargs)
monkeypatch.setattr(chroma_module, "ChromaZImagePipeline", Recorder)
remover = WatermarkRemover.__new__(WatermarkRemover)
remover.model_profile = "chroma-zimage"
remover.device = "cuda"
remover.torch_dtype = None
remover.hf_token = None
remover._progress_callback = None
remover.controlnet_conditioning_scale = 1.0
remover.cpu_offload = False
remover._qwen_zimage_pipeline = None
remover._load_qwen_zimage_pipeline()
assert captured["device"] == "cuda"
assert captured["hf_token"] is None
def test_chroma_adaptive_polish_defaults_off():
"""Like qwen-zimage: Chroma1's output already matches the input's detail
level, so polishing would add grain to an already-faithful output."""
from remove_ai_watermarks._internal.watermark_profiles import resolve_adaptive_polish
assert resolve_adaptive_polish(None, "chroma-zimage") is False
assert resolve_adaptive_polish(True, "chroma-zimage") is True
def test_chroma_vae_roundtrip_raises_not_implemented():
"""The base class provides no VAE donor for chroma-zimage; the failure must
be the explicit NotImplementedError, not an opaque crash from a missing
Qwen stack."""
import torch
from remove_ai_watermarks._internal.chroma_zimage_pipeline import ChromaZImagePipeline
pipeline = ChromaZImagePipeline(device="cuda", torch_dtype=torch.bfloat16)
with pytest.raises(NotImplementedError, match="does not expose a VAE donor"):
pipeline._vae_roundtrip(Image.new("RGB", (32, 32)))
def test_chroma_google_face_content_gets_the_lower_adaptive_floor():
"""Face content on Google SynthID clears at 0.125, not 0.40.
The four-fixture calibration showed a clean face-count split: both zero-face
text cards need 0.25, both face fixtures clear at 0.12. The adaptive arm
saves face identity (0.611 vs 0.279 cosine) and 6.4 dB PSNR on face content.
"""
from remove_ai_watermarks._internal.watermark_profiles import (
CHROMA_ZIMAGE_GOOGLE_FACE_STRENGTH,
CHROMA_ZIMAGE_GOOGLE_STRENGTH,
resolve_strength,
)
assert pytest.approx(0.125) == CHROMA_ZIMAGE_GOOGLE_FACE_STRENGTH
# Face content gets the lower adaptive floor.
assert resolve_strength(None, "google", "chroma-zimage", face_count=17) == pytest.approx(0.125)
assert resolve_strength(None, "google", "chroma-zimage", face_count=1) == pytest.approx(0.125)
# Zero-face content gets the flat floor (text/graphic cohort).
assert resolve_strength(None, "google", "chroma-zimage", face_count=0) == pytest.approx(0.40)
assert pytest.approx(0.40) == CHROMA_ZIMAGE_GOOGLE_STRENGTH
# No face_count hint: conservative (flat floor).
assert resolve_strength(None, "google", "chroma-zimage") == pytest.approx(0.40)
# Other vendors ignore face_count: the split is Google-SynthID-specific.
assert resolve_strength(None, "openai", "chroma-zimage", face_count=5) == pytest.approx(0.09)
assert resolve_strength(None, "meta", "chroma-zimage", face_count=5) == pytest.approx(0.17)
# An explicit strength still wins over the adaptive arm.
assert resolve_strength(0.20, "google", "chroma-zimage", face_count=17) == pytest.approx(0.20)
# Other profiles are untouched by the face_count parameter.
assert resolve_strength(None, "google", "qwen-zimage", size=(2000, 1850), face_count=17) == pytest.approx(0.27)
+5 -5
View File
@@ -924,7 +924,7 @@ class TestVideoProvenanceApi:
assert report.visible_mark == "sora"
assert report.visible_detected_frames == 5
assert report.total_frames == 5
assert tuple(VIDEO_VISIBLE_MARKS) == ("sora", "veo", "seedance", "dola", "hailuo", "kling")
assert tuple(VIDEO_VISIBLE_MARKS) == ("sora", "veo", "seedance", "doubao", "dola", "hailuo", "kling")
def test_reports_unknown_instead_of_clean(self, tmp_path: Path):
from remove_ai_watermarks.video import identify_video
@@ -2085,10 +2085,10 @@ class TestVideoVisibleScan:
scans = scan_video_marks(
tmp_path / "synthetic.mp4",
("sora", "veo", "seedance", "dola", "hailuo", "kling"),
("sora", "veo", "seedance", "doubao", "dola", "hailuo", "kling"),
)
assert set(scans) == {"sora", "veo", "seedance", "dola", "hailuo", "kling"}
assert set(scans) == {"sora", "veo", "seedance", "doubao", "dola", "hailuo", "kling"}
assert all(scan.timestamps == (0.25,) for scan in scans.values())
assert len(prepared_ids) == 6
assert len(set(prepared_ids)) == 1
@@ -2576,7 +2576,7 @@ class TestVideoVisibleApi:
)
def fake_scan(_source: Path, marks: tuple[str, ...]):
assert marks == ("sora", "veo", "seedance", "dola", "hailuo", "kling")
assert marks == ("sora", "veo", "seedance", "doubao", "dola", "hailuo", "kling")
return {
"sora": sora_scan,
"veo": sora_scan,
@@ -2792,7 +2792,7 @@ class TestVideoVisibleCli:
assert result.exit_code == 0, result.output
assert "temporally stable" in result.output
assert "auto|sora|veo|seedance|dola|hailuo|kling" in result.output
assert "auto|sora|veo|seedance|doubao|dola|hailuo|kling" in result.output
assert "--temporal-consistency" in result.output
def test_reports_removed_frames(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+1 -1
View File
@@ -6,7 +6,7 @@ Two failures this suite exists to catch:
generated from the engines' measured geometry, so this is a regression tripwire).
The examples are SYNTHETIC (``scripts/render_visible_examples.py`` composites the
committed silhouettes onto a generated base). User uploads never enter the repo.
committed silhouettes onto a generated base).
"""
from __future__ import annotations
Generated
+1 -1
View File
@@ -3169,7 +3169,7 @@ wheels = [
[[package]]
name = "remove-ai-watermarks"
version = "0.34.1"
version = "0.35.1"
source = { editable = "." }
dependencies = [
{ name = "c2pa-python" },