mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-10-14 05:20:46 +02:00
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.
This commit is contained in:
31 files changed
+3204
-1066
No files matched your search
@@ -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
|
||||
|
||||
@@ -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]"
|
||||
@@ -416,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
|
||||
|
||||
@@ -0,0 +1,306 @@
|
||||
# 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; option 2's adaptive policy 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.
|
||||
+1
-1
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -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`.
|
||||
@@ -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,23 @@ 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` (above qwen's `0.27`; at this floor the regeneration destroys
|
||||
dense text and collapses face identity -- the tradeoff that motivates the
|
||||
documented content-adaptive follow-up);
|
||||
- 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 +285,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`.
|
||||
|
||||
+58
-15
@@ -940,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,
|
||||
@@ -963,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"`
|
||||
@@ -991,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
|
||||
@@ -1039,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
|
||||
@@ -1140,13 +1153,43 @@ 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 and Meta 0.17 (both 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. A content-adaptive strength policy is the documented follow-up.
|
||||
|
||||
### 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
|
||||
@@ -1184,7 +1227,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
|
||||
@@ -1271,7 +1314,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
@@ -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:
|
||||
|
||||
@@ -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
@@ -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.0"
|
||||
description = "AI watermark remover for visible, invisible, and provenance marks in images and video"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11,<3.15"
|
||||
|
||||
Executable
+114
@@ -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))
|
||||
Executable
+109
@@ -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))
|
||||
Executable
+142
@@ -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))
|
||||
Executable
+152
@@ -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))
|
||||
Executable
+188
@@ -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."
|
||||
)
|
||||
Executable
+272
@@ -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.")
|
||||
Executable
+249
@@ -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")
|
||||
Executable
+76
@@ -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()
|
||||
@@ -32,7 +32,7 @@ _os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
|
||||
_warnings.filterwarnings("ignore", message=r".*ImageProcessorFast.*")
|
||||
|
||||
|
||||
__version__ = "0.34.1"
|
||||
__version__ = "0.35.0"
|
||||
|
||||
__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,44 @@ _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
|
||||
|
||||
_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 +173,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 +192,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)"
|
||||
)
|
||||
|
||||
|
||||
@@ -173,14 +218,17 @@ 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:
|
||||
return _CHROMA_ZIMAGE_STRENGTH_BY_VENDOR.get((vendor or "").casefold(), 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,
|
||||
@@ -194,7 +203,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.
|
||||
"""
|
||||
@@ -207,7 +216,7 @@ class WatermarkRemover:
|
||||
resolved_strength = resolve_strength(strength, vendor, self.model_profile, size=source.size)
|
||||
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")
|
||||
|
||||
@@ -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.")
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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,232 @@ 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."""
|
||||
import 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)))
|
||||
Reference in new issue
Block a user