CalvinBackup/remove-ai-watermarks
1
0
Fork 0
mirror of https://github.com/wiltodelta/remove-ai-watermarks.git synced 2026-08-09 23:50:40 +02:00
Code Issues Packages Projects Releases Wiki Activity
Files
2f7225799654309959e232f52a10854c00e2dfa9
remove-ai-watermarks/docs/module-internals.md
T
Victor KuznetsovandClaude Opus 5 2f72257996 Fix what the verification pass found in the knob removal 
An adversarial review of 52b2c11 (five independent audits, each finding put to
two skeptics, plus a completeness critic) found four defects that commit
introduced and several stale claims it should have caught.

The install hint no longer installs -- again. Folding five hints into one
INVISIBLE_EXTRA constant dropped the shell quoting the originals had, so the
printed remediation was `pip install remove-ai-watermarks[qwen-zimage]`. Bare
brackets are a glob in zsh, the macOS default shell: it dies with "no matches
found" before pip runs. That is the exact failure 52b2c11 existed to stop
producing, reintroduced in a different form by a bulk replace. The constant is
quoted now, and a test asserts the quotes rather than the bare substring -- the
old assertions passed either way, which is why nothing caught it.

Three tests were not guarding what they claimed:

- The commit's headline behaviour change, per-profile polish resolution inside
  the engine, had no test at all. Rebinding resolve_adaptive_polish to the
  pre-commit `bool(value)` left the full suite green. Now covered by a test that
  drives the real engine and observes whether humanizer.adaptive_polish ran;
  that mutation now fails it.
- TestAvailability still asserted the pre-commit (torch, diffusers) contract, so
  in a diffusion-only environment it was simply wrong, and comparing each gate to
  a tuple copied from itself could never catch the two gates disagreeing -- the
  drift the shared REMOVAL_MODULES was introduced to prevent. Replaced with a
  test that simulates each module's absence and requires BOTH gates to close.
- Both CUDA-refusal guards skipped in every environment, including CI: they were
  gated on the diffusion stack, which no CI job installs. The refusal fires
  before any torch attribute is read, so they now run everywhere; only the dtype
  assertion keeps its skip.

Also: the retired-knob test covered `invisible` but not `all` or `batch`, though
all three declared those options separately; and smoke_matrix.py still called
remove_watermark(region=...), a parameter 52b2c11 deleted, with the resulting
TypeError swallowed into a skip by a broad except.

Stale documentation the previous sweep missed: known-limitations still described
an MPS out-of-memory fallback and a lighter-pipeline escape that no code can
produce; module-internals declared Canny thresholds of 100/200 as a compatibility
contract while the code uses 13/64, attributed enable_model_cpu_offload to
deleted profiles, and still warned that the engine and CLI defaults differ (this
commit's predecessor made them identical); cli.md gated `all` on the `diffusion`
extra; python-api claimed "cuda" was the only accepted explicit device when
"auto" is too. The claim that `device` is not a parameter was wrong in both
module-internals and .claude/rules/development.md -- it is one, deliberately, and
now says so. `--cpu-offload` help and the pipeline's CUDA guard both still
pointed at MPS.

Not fixed here, reported instead -- both are outside this repo:
- ComfyUI-remove-ai-watermarks nodes.py:332 passes num_inference_steps and
  guidance_scale (plus min_resolution/upscaler from bf4bfc1). distribute.yml's
  comfyui job runs on every release and fails the release if the node sync fails,
  so 0.25.0 needs that node updated first.
- raiw-app modal_app.py:422-425 forwards the same two kwargs into
  remove_watermark. Latent: it is pinned to 1a77e24 and nothing supplies a value
  today, so it fires on the next pin bump.

pre-commit: 1) maintain.sh - exit 0 (1093 tests, Pyright 0 errors, no
vulnerabilities); 2) /simplify - not re-run, this commit is the applied output of
a five-dimension adversarial review; 3) docs sync - grepped MPS/mps, the extras
names and every symbol touched across README, docs/, scripts/, .claude/; updated
6 docs; 4) CLAUDE.md - corrected the device claim in .claude/rules/development.md
and added the shell-quoting rule

Verified by execution, not assertion: smoke_matrix --quick 51 pass / 0 fail,
_knob_rows driven directly 10 pass / 0 fail / 7 skip (no CUDA), the install hint
rendered and round-tripped through zsh, and each new test confirmed to fail
under the mutation it is meant to catch.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-08-03 16:31:39 -07:00

830 lines
42 KiB
Markdown
Raw Blame History

# Module internals
This page documents the current implementation contract. It intentionally
avoids experiment logs, corpus counts, and calibration history. Those records
live in [the verification plan](verification-plan.md) and the research archive
listed in [the documentation index](index.md).
Read the relevant section before changing a subsystem. When this page and the
code disagree, the code and its tests are authoritative and this page must be
updated in the same change.
## Architecture
The package has four main paths:
```mermaid
flowchart LR
Input[Input file] --> Identify[Identify provenance]
Input --> Visible[Visible mark removal]
Input --> Invisible[Diffusion regeneration]
Input --> Metadata[Metadata stripping]
Identify --> Report[ProvenanceReport]
Visible --> VisibleOutput[Localized and filled image]
Invisible --> InvisibleOutput[Regenerated image]
Metadata --> MetadataOutput[Container with AI metadata removed]
```
The `all` command runs visible removal, optional invisible regeneration, and
metadata stripping in that order.
## Command line interface
[`cli.py`](../src/remove_ai_watermarks/cli.py) owns command parsing and
user-facing exit behavior.
Important contracts:
- Single-image arguments reject directories.
- `visible` writes no output when no registered mark is selected and exits with
`EXIT_NO_VISIBLE_MARK`.
- `invisible` writes no output when no supported local signal is found, unless
`--force` is supplied.
- The two no-signal conditions currently share exit code `2`.
- Hard processing and write failures exit with code `1`.
- `all` can still write the completed visible and metadata stages when the
diffusion dependencies are unavailable, but exits with code `1` so the
partial result is not reported as complete.
- `batch` counts per-file failures and exits nonzero if any file failed or an
applicable invisible stage was skipped because its dependencies were absent.
The decorators for diffusion options are shared by `invisible`, `all`, and
`batch`. The runtime help generated by Click is the source of truth for option
names and defaults.
`--adaptive-polish` is tri-state: it declares `default=None`, so "the user did not
choose" is a value the CLI passes through rather than a default it has to invent.
`resolve_adaptive_polish` in `watermark_profiles.py` turns that `None` into the
profile's answer (off for `qwen-zimage`, whose output already matches the input's
detail level; on for `sdxl-zimage`). The same call runs inside
`InvisibleEngine.remove_watermark`, so a library caller and a CLI caller on one
profile get the same output.
It used to read Click's parameter source in the CLI instead. That put per-profile
data in the argument-parsing layer, left the engine declaring the opposite default,
and silently lost the polish for anything supplying the flag non-interactively (an
envvar default or a wrapper calling `main()` with a defaulted list is classified
`DEFAULT`). The seed follows the same rule: the CLI does not pre-resolve it either.
Regression coverage:
- [`test_cli.py`](../tests/test_cli.py)
- [`test_cli_robustness.py`](../tests/test_cli_robustness.py)
- [`test_optional_deps.py`](../tests/test_optional_deps.py)
## High-level Python API
[`api.py`](../src/remove_ai_watermarks/api.py) provides:
- `remove_visible`
- `visible_provenance`
The package root exposes both lazily through
[`__getattr__`](../src/remove_ai_watermarks/__init__.py), keeping a plain package
import free of the heavier image and model imports.
For path inputs, `remove_visible` reads provenance metadata, preserves alpha,
and optionally writes and strips metadata. Array inputs are treated as BGR
arrays and have no file provenance or separate alpha plane.
When no visible mark is removed, a same-format path copy preserves the original
bytes. `write_noop=False` leaves the requested output path untouched instead.
Regression coverage:
- [`test_api.py`](../tests/test_api.py)
- [`test_image_io.py`](../tests/test_image_io.py)
[`video.py`](../src/remove_ai_watermarks/video.py) provides the high-level
video entry point:
- `identify_video`
- `inspect_video_metadata`
- `remove_video_all`
- `remove_video_batch`
- `remove_video_invisible`
- `remove_video_metadata`
- `remove_video_visible`
The video API validates both the supported extension and container signature,
then delegates all metadata detection and stripping to `metadata.py`. It
requires a separate same-container output, defaulting to `<source>_clean`, so
the product path does not overwrite an original. The package root exposes all
functions lazily.
`identify_video` runs the same stable-mark selection helper as
`remove_video_visible`, so a provenance report cannot authorize a mark that the
removal path would reject. It reports an empty local result as unknown rather
than clean. Identification skips the separate per-frame timestamp probe because
it never encodes frames. `remove_video_all` is the predictable-output
composition: visible removal plus verified metadata stripping by default, with
a same-container passthrough when neither signal exists. The lossy invisible
removal stage is an explicit opt-in through the oracle-certified profile.
`remove_video_batch` applies those contracts sequentially across a top-level
directory, returns every per-file failure, and byte-copies visible no-ops so a
successful output set has no silent holes. An invisible batch loads one VAE
runtime and reuses it across every compatible file; a failed model load is
reported per file without retrying the same multi-GB initialization.
Native MP4/MOV TC260 labels follow TC260-PG-20257A:
`moov.udta.meta.keys` maps an `AIGC` key to a raw JSON value in `ilst`.
[`_internal/isobmff.py`](../src/remove_ai_watermarks/_internal/isobmff.py) walks those
nested boxes by seeking, so detection reaches a tail `moov` without reading the
preceding `mdat`. The MP4/MOV/M4V/M4A removal path first validates the top-level
box walk, then copies the source to a sibling temporary file in bounded chunks.
Supported C2PA/JUMBF/AI-label boxes become same-size `free` boxes with blank
payloads; TC260 removal changes the four-byte key to `free` and blanks only the
validated JSON value with same-length spaces. This preserves every box size,
`stco`/`co64` offset, encoded stream byte, and source-sized memory bound.
Publication is atomic, and a malformed top-level walk is copied unchanged. A
generic `AIGC` key whose value has no TC260 field is ignored.
[`_internal/ebml.py`](../src/remove_ai_watermarks/_internal/ebml.py) provides the
corresponding bounded Matroska/WebM reader. It seeks over clusters and accepts
only a `Segment.Tags.Tag.SimpleTag` pairing `TagName=AIGC` with a JSON
`TagString` carrying a TC260 field. The existing ffmpeg stream-copy path removes
those container tags without transcoding the encoded streams.
[`_internal/riff.py`](../src/remove_ai_watermarks/_internal/riff.py) and
[`_internal/flv.py`](../src/remove_ai_watermarks/_internal/flv.py) implement the remaining
normative TC260 video placements. The RIFF walker reads only AVI
`LIST/INFO/AIGC` children. The FLV walker skips media tags and parses the AMF0
`script.onMetaData.AIGC` string. Both require a recognized TC260 JSON field and
use the verified ffmpeg stream-copy path for removal.
[`video_encoding.py`](../src/remove_ai_watermarks/video_encoding.py) owns the
ffmpeg command and pipe lifecycle shared by visible removal and invisible
regeneration. It centralizes container codecs, optional audio stream copying,
metadata/chapter policy, encode-failure reporting, and atomic same-directory
publication. Each mapped stream is allowed to reach its own end, so a copied
audio tail is not shortened to the frame-input duration.
Both the raw-BGR and timestamped-NUT stdin modes redirect ffmpeg stderr to a
temporary file while frames are written. Waiting to read diagnostics until
after stdin closed allowed stderr backpressure to stop ffmpeg's frame reads,
which in turn blocked the producer before it could close stdin. The file consumes
no pipe capacity or RAM while ffmpeg runs; completion reports a bounded head and
tail when diagnostics are unusually large. Aborts release it even when ffmpeg
has already exited. A real subprocess regression writes diagnostics beyond pipe
capacity while streaming frames, checks bounded failure reporting, and the Linux
full-clip CI job guards the complete path.
Frame encoding and source-audio copying run as two ffmpeg processes in sequence.
The streaming encoder has only the frame pipe as input, so input probing or
demux queues cannot deadlock the producer against a second input. After that
pipe reaches EOF, a finite stream-copy mux combines the encoded video with the
source audio and applies the requested metadata/chapter policy. Both stages use
sibling temporary files, and only the completed mux is published atomically.
The mux also redirects diagnostics to disk and reports only a bounded head and
tail. Command regressions assert the single-input encoder and final map targets;
failure regressions cover bounded mux diagnostics and atomic cleanup.
`probe_video_encode_profile` reads the first source video stream with ffprobe
and preserves the supported properties that survive the 8-bit BGR boundary:
`yuv420p`/`yuv422p`/`yuv444p` chroma sampling, recognized color tags, encoder
time base, MP4/MOV track timescale, source pixel format, and component depth.
Both raw-CFR and timestamped-NUT inputs use ffmpeg's passthrough FPS mode. This
keeps one encoded frame per supplied frame when an older ffmpeg receives a
fine-grained source encoder time base such as `1/90000`; implicit synchronization
can otherwise synthesize thousands of duplicate frames between CFR timestamps.
HDR transfer functions and component depths above 8 bits are rejected before
encoding so the OpenCV boundary cannot silently reduce them to SDR 8-bit.
`probe_video_timestamps` reads authoritative per-frame display PTS through
ffprobe. OpenCV timestamps are only a count-matched fallback when ffprobe is
unavailable or fails; this avoids decoder anomalies such as one spurious
negative first-frame timestamp turning a CFR clip into false VFR. A uniform sequence keeps the
cheap raw-BGR pipe unless the source starts at a non-zero PTS. A variable or
offset sequence is packetized by the lazy PyAV bridge as rawvideo in an
in-memory NUT stream with explicit PTS. System ffmpeg reads that stream with
`-fps_mode passthrough`; `-copyts` additionally retains a non-zero video start
and the corresponding copied-audio offset. No temporary frame sequence or
second video encoder is introduced.
[`video_temporal.py`](../src/remove_ai_watermarks/video_temporal.py) owns the
shared optical-flow maps and temporal residual metric. Visible removal uses
`stabilize_filled_frame` after the selected image backend: it works on a
bounded crop around adjacent masks, backward-warps the prior cleaned frame,
requires high warped-mask coverage, and gates blending on an unmasked
source-context ring. Only covered current-mask pixels change. Scene cuts,
disjoint marks, and poor motion matches therefore retain the independent
current-frame fill. The same module supplies the motion-compensated metric used
by the invisible-video sweep.
[`video_invisible.py`](../src/remove_ai_watermarks/video_invisible.py)
implements the oracle-certified video SynthID removal engine. It samples frames
uniformly, resizes to a VAE-aligned geometry, encodes each frame to latent
space, applies one seeded spatial-noise field across the entire sequence, and
decodes fresh pixels. Reusing a single noise field avoids independent
frame-to-frame noise. The shipped path retains only one configured frame batch,
updates PSNR and temporal residuals incrementally, and streams BGR frames
directly to the video-only ffmpeg encoder. A separate stream-copy mux then adds
optional source audio and drops all source metadata. The result is written
through same-directory temporary files and atomically replaced only after both
stages succeed.
The engine returns PSNR and a motion-compensated temporal-residual ratio as
quality measurements. Neither is a watermark detector. The high-level result
reports completed removal without a separate verification-status flag. The companion
`scripts/video_synthid_sweep.py` imports the same engine helpers to build a
matched control and candidate grid, preventing research and shipped
regeneration paths from drifting. The full-clip oracle floor is
`noise_std=0.15`: on the public eight-second Veo carrier, `0.10` remained
detected while `0.15` did not.
[`video_visible.py`](../src/remove_ai_watermarks/video_visible.py) implements
the first pixel stages for Sora, Veo, Seedance, Dola, Hailuo, and Kling. The
Sora detector searches a normalized frame with a fully synthetic
mascot-and-text silhouette at several scales. The Veo detector uses separate
synthetic silhouettes for the current four-point diamond and legacy `Veo`
text. Seedance uses a synthetic rounded boxed-`AI` silhouette, while Dola uses
an OpenCV-font `Dola AI` silhouette. Hailuo uses a synthetic waveform,
MINIMAX/Hailuo text, separator, and ring. Kling combines synthetic font
variants with a ring approximation of its swirl; the logo path rescues
wordmarks whose version or font differs, while the edge and white-label gates
reject recurring scene texture. All fixed-mark searches are bounded to the
expected lower-frame area and calibrated independently. A strong relocated Veo
diamond may bypass the known layout anchors, but weak free-corner matches never
enter the temporal arbiter.
The default `auto` route decodes each frame once, shares its grayscale and
normalized representations across all detectors, and caches resized synthetic
template features for the fixed stream geometry. Provider confidence scales
are not comparable: selection applies each provider's temporal arbiter and
takes the first stable result in specificity order (`sora`, `veo`, `seedance`,
`dola`, `hailuo`, `kling`). An explicit mark uses the same scan path with one
candidate. Removal also collects authoritative per-frame timestamps for the
encoder, while identification omits that unused ffprobe pass.
Every per-frame result is untrusted. The provider-specific stabilization
wrappers share one recurrence implementation, while retaining separate visual
floors and minimum-run policy. Provenance can relax a low-contrast run only
after recurring visual evidence exists. Sora transition frames follow the
nearest confirmed moving position only with Sora provenance. Veo, Seedance,
Dola, Hailuo, and Kling additionally require candidates to remain anchored to
the start of a run. This rejects slowly drifting scene details that still have
high frame-to-frame overlap. Hailuo and Kling do not infer provenance from
technical encoder tags; their confirmed public samples carried no provider
metadata.
Removal runs in a second decode pass. Sora, legacy Veo text, Dola text,
Seedance, Hailuo, and Kling use box masks. Seedance deliberately fills the
complete localized box: a synthetic outline mask passed repeat detection but
left part of the real translucent border visible during visual end-to-end
review. Hailuo expands beyond the matched core to cover both provider icons.
Kling expands around the wordmark or swirl to include the version and optional
`PRO` suffix. The square Veo diamond uses a synthetic shape mask so transparent
corners do not erase unrelated pixels. Every mask goes through the shared
`watermark_registry.fill` backends. ffmpeg encodes the changed video stream and
copies optional audio. The default OpenCV fill is the speed floor; structured
backgrounds need MI-GAN or LaMa for better reconstruction. Invisible video
stages must continue to reuse the image and metadata implementations rather
than copying their logic.
Regression coverage:
- [`test_video.py`](../tests/test_video.py), including a real ffmpeg full-clip
Sora/OpenCV path that generates a synthetic marked MP4 with AAC audio and
C2PA provenance, runs both `remove_video_visible` and the composed
`remove_video_all` API without mocks, and verifies complete removal, frame
count, frame rate, duration, untouched-region PSNR, paired temporal deltas
inside the filled region, byte-identical copied audio packets, source stream
properties, metadata stripping, and a large-`mdat` metadata case that rejects
any full-source `read_bytes()` call. CI installs ffmpeg explicitly for this
test so the integration gate cannot silently skip.
## Metadata and provenance
### C2PA
[`_internal/c2pa.py`](../src/remove_ai_watermarks/_internal/c2pa.py) reads C2PA with the
official `c2pa-python` reader first. Its byte-level PNG parser remains a fallback
for partial and synthetic fixtures that the official reader rejects.
Vendor attribution comes from the registry in
[`_internal/constants.py`](../src/remove_ai_watermarks/_internal/constants.py). Derived
issuer and platform maps should not be maintained separately.
### Metadata scanning and stripping
[`metadata.py`](../src/remove_ai_watermarks/metadata.py) contains the shared
metadata scanners and `remove_ai_metadata`.
Key contracts:
- `scan_head` is the shared cached input for bounded byte scans.
- JPEG stripping walks metadata segments and preserves the entropy-coded image
scan.
- ISOBMFF containers use
[`_internal/isobmff.py`](../src/remove_ai_watermarks/_internal/isobmff.py).
- Native MP4/MOV TC260 `AIGC` entries are read from
`moov.udta.meta.keys/ilst` and blanked without changing box sizes.
- Native MKV/WebM TC260 `AIGC` entries are read from
`Segment.Tags.Tag.SimpleTag` and removed through the ffmpeg stream-copy path.
- Native AVI and FLV TC260 entries are read from `LIST/INFO/AIGC` and
`script.onMetaData.AIGC`, respectively, then removed through ffmpeg stream
copying.
- Supported non-ISOBMFF audio and video containers use ffmpeg stream copying.
- The low-level remover is fail-safe and can copy an undecodable file through
unchanged.
- A caller that reports success must use `strip_and_verify`, which scans the
written output for surviving markers. If the metadata-preserving decoder
rejected the container but `image_io` can still decode its raster,
`strip_and_verify` normalizes that raster and scans again. A truly undecodable
file keeps the surviving-marker result.
Detection and removal must stay in parity. A new marker is incomplete until the
scanner can find it, the remover can reach every supported placement, and a
test proves that it no longer appears in the output.
Regression coverage:
- [`test_metadata.py`](../tests/test_metadata.py)
- [`test_metadata_internals.py`](../tests/test_metadata_internals.py)
- [`test_security_clamp.py`](../tests/test_security_clamp.py)
### Provenance report
[`identify.py`](../src/remove_ai_watermarks/identify.py) separates file-backed
metadata extraction from verdict logic:
- `extract_provenance_evidence` reads the supported metadata signals into
`ProvenanceEvidence`.
- `evidence_from_metadata_record` normalizes an externally collected nested
metadata record into the same evidence type without file access. Diagnostic
values under `error` and `kind` are excluded from evidence while nested raw
bytes remain available through encoded binary fields.
- `identify_from_evidence` evaluates that evidence without reopening the source.
- `identify` preserves the path-based API and adds the optional registered
visible-mark and open invisible-watermark decoders after extraction.
The `detect` extra composes the shared `pixels` runtime with PyWavelets. Its
in-tree [`dwt_dct.py`](../src/remove_ai_watermarks/dwt_dct.py) decoder preserves
the upstream matrix algorithm without installing Torch or non-headless OpenCV.
The upstream MIT notice ships inside the wheel under `licenses/`.
`is_ai_generated` is `True` or `None`; absence of evidence is not reported as a
human-made verdict. `ai_source_kind` distinguishes fully generated content from
AI-enhanced composites when the source metadata provides that distinction.
TrustMark is reported as a watermark signal but does not by itself assert AI
origin because it can also protect human-authored content.
Regression coverage:
- [`test_identify.py`](../tests/test_identify.py)
- [`test_trustmark_detector.py`](../tests/test_trustmark_detector.py)
- [`test_invisible_watermark.py`](../tests/test_invisible_watermark.py)
## Visible mark removal
### Registry and decision flow
[`watermark_registry.py`](../src/remove_ai_watermarks/watermark_registry.py) is
the only visible-mark registry. `mark_keys()` supplies the CLI choices, so the
CLI must not maintain a separate mark list.
Automatic removal has three distinct stages:
1. Perception: each registered detector produces strict and relaxed candidates.
2. Decision: the pure `decide` arbiter applies sensitivity and corroborating
provenance.
3. Action: each selected mark is localized to a mask and passed to the shared
fill function.
`sensitivity="strict"` never relaxes a detector. `sensitivity="auto"` can relax
one only when metadata or a sufficiently strong same-product sibling confirms
that product. The removed blanket `assume_ai` mode is rejected explicitly.
The Jimeng pill has an additional decision gate because its visual detector is
weaker than the other registered marks. Keep that policy in the registry, not
inside unrelated detector engines.
`remove_auto_marks` removes every selected mark, not only the strongest one.
This matters for images that carry marks in more than one corner.
Regression coverage:
- [`test_watermark_registry.py`](../tests/test_watermark_registry.py)
- [`test_api.py`](../tests/test_api.py)
### Gemini sparkle
[`gemini_engine.py`](../src/remove_ai_watermarks/gemini_engine.py) uses a
multi-scale shape search and a false-positive gate. Its captured sparkle assets
serve detection and mask geometry only. Pixel recovery is performed by the
shared fill backend.
`detect_sparkle_confidence` uses a process-wide shared engine because its loaded
assets and template ladder are immutable.
Regression coverage:
- [`test_gemini_engine.py`](../tests/test_gemini_engine.py)
### Text mark engines
[`_text_mark_engine.py`](../src/remove_ai_watermarks/_text_mark_engine.py)
provides common localization, detection front ends, template caching, rival
comparison, and footprint construction.
Each vendor module supplies a `TextMarkConfig` and only the behavior that cannot
be represented by the shared base:
- [`doubao_engine.py`](../src/remove_ai_watermarks/doubao_engine.py)
- [`jimeng_engine.py`](../src/remove_ai_watermarks/jimeng_engine.py)
- [`qwen_engine.py`](../src/remove_ai_watermarks/qwen_engine.py)
- [`kling_engine.py`](../src/remove_ai_watermarks/kling_engine.py)
- [`yuanbao_engine.py`](../src/remove_ai_watermarks/yuanbao_engine.py)
- [`samsung_engine.py`](../src/remove_ai_watermarks/samsung_engine.py)
- [`runninghub_engine.py`](../src/remove_ai_watermarks/runninghub_engine.py)
- [`baidu_engine.py`](../src/remove_ai_watermarks/baidu_engine.py)
- [`liblib_engine.py`](../src/remove_ai_watermarks/liblib_engine.py)
The detector and removal mask must use compatible geometry. A detector that
fires while producing an empty or misplaced mask is a removal failure even if
the detection test passes.
Yuanbao uses the polarity-independent `contrast` front end because its standard
two-line mark can be light on dark scenes or dark on light scenes. Its detector
and footprint both use the same best-match box. The separate one-line overlay
variant is not covered.
The capture-less Jimeng pill lives in
[`pill_engine.py`](../src/remove_ai_watermarks/pill_engine.py). It uses a
synthetic silhouette for detection and a fixed top-left footprint.
Each engine has a corresponding test module under [`tests/`](../tests/).
Shared behavior is covered by:
- [`test_text_mark_engine.py`](../tests/test_text_mark_engine.py)
- [`test_text_mark_faint_mask.py`](../tests/test_text_mark_faint_mask.py)
- [`test_text_mark_memory.py`](../tests/test_text_mark_memory.py)
### Fill backends and region erasing
[`region_eraser.py`](../src/remove_ai_watermarks/region_eraser.py) implements the
same backends used by visible removal and the user-directed `erase` command:
- `cv2`
- `migan`
- `lama`
`watermark_registry.resolve_backend` selects LaMa first, then MI-GAN, then
OpenCV for `auto`. A memory-constrained caller should explicitly select MI-GAN
or OpenCV instead of relying on `auto`.
MI-GAN and LaMa crop around the mask before model inference and paste back only
masked pixels. Their model sessions are loaded lazily. MI-GAN uses the inverse
mask polarity expected by its ONNX model.
Regression coverage:
- [`test_region_eraser.py`](../tests/test_region_eraser.py)
- [`test_inpaint_fallback.py`](../tests/test_inpaint_fallback.py)
## Invisible watermark regeneration
### Profiles and strength
[`_internal/watermark_profiles.py`](../src/remove_ai_watermarks/_internal/watermark_profiles.py)
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.
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.
For serverless cold starts, `InvisibleEngine.preload(global_only=True)` loads the
mandatory global stage and YuNet while leaving the optional Z-Image and SAM face
stack lazy until a face is detected. The default `preload()` still loads every
stage.
**What is deliberately not a parameter.** Model id, step count and CFG are fixed
by the profile, so none of them appears in `WatermarkRemover.__init__`,
`remove_watermark`, `InvisibleEngine`, or the CLI. They used to be accepted and
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
dtype" for what an override cost the last time one existed.
`device` is the exception and remains a library parameter: `None` or `"auto"`
detect, `"cuda"` pins without detecting (which is what a container that knows its
hardware wants), and any other value raises at construction. It is not a CLI
option, because the only useful value a user could type is the one detection
already returns.
[`invisible_engine.py`](../src/remove_ai_watermarks/invisible_engine.py) handles
image sizing, postprocessing, and the public engine
interface. It delegates model execution to
[`_internal/watermark_remover.py`](../src/remove_ai_watermarks/_internal/watermark_remover.py).
`get_device` in that module answers only `cuda` or `cpu`. An `mps` or `xpu` answer
would travel one frame to the same CUDA-only refusal while costing a device probe,
and reporting it implied an Apple-silicon or Intel-GPU path that does not exist.
The refusal names the *resolved* device, so `device=None` on a CUDA-less host says
`'cpu'` rather than `'None'`.
The Python engine and the CLI now resolve the same defaults: the CLI forwards an
unset `--adaptive-polish` and `--seed` as `None` and the engine applies the
profile's answer, so a library caller and a CLI caller on one profile produce the
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
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
`test_qwen_zimage_pipeline.py::test_global_kwargs_use_lightning_and_diffsynth_controlnet_shape`,
`::test_face_kwargs_use_project_zimage_settings` and
`::test_canny_control_image_is_three_channel_and_detects_an_edge`.
Regression coverage:
- [`test_invisible_engine.py`](../tests/test_invisible_engine.py)
- [`test_qwen_zimage_pipeline.py`](../tests/test_qwen_zimage_pipeline.py)
- [`test_platform.py`](../tests/test_platform.py)
### CPU offload
CPU offload is enabled only when requested. Nothing calls Diffusers'
`enable_model_cpu_offload` any more -- that belonged to the deleted single-stage
profiles. `--cpu-offload` now forces **both** stacks of the two-stage profiles out
of automatic device residency.
Residency is otherwise chosen from the card's total VRAM, once per stack:
`resolve_global_model_residency` gates the mandatory Qwen stack at
`RESIDENT_GLOBAL_MODEL_MIN_VRAM_GIB` and `resolve_face_model_residency` gates the
optional Z-Image stack at `RESIDENT_FACE_MODEL_MIN_VRAM_GIB`.
Below the global floor, `_qwen_vram_config` streams the stack from disk, which is
what makes a 20B model runnable on a consumer card. At or above it, streaming is
pure waste and the weights stay on the GPU. The difference is not marginal:
DiffSynth offloads by dropping the weights to the meta device and re-reading every
parameter through its `DiskMap` on the next onload, and the pipeline moves between
text encoder, transformer and VAE on each pass. Measured on an H100 (80 GiB) in
August 2026, a warm global pass took 37.3 s at 0.8 GiB resident with the streaming
config, against 2.2 s at 28.7 GiB with the stack resident; both stacks resident
peaked at 48.0 GiB. Faster storage cannot close that gap, because the cost is the
reload itself rather than the read.
The resident config deliberately passes no `"disk"` value anywhere. DiffSynth latches
`disk_offload` once, from `offload_dtype`, so leaving the sentinel in place while
pointing every device at CUDA would keep the meta-drop and re-read.
Regression coverage:
- [`test_cpu_offload.py`](../tests/test_cpu_offload.py)
- [`test_qwen_zimage_pipeline.py`](../tests/test_qwen_zimage_pipeline.py)
### Qwen plus Z-Image
[`_internal/qwen_zimage_pipeline.py`](../src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py)
implements the fixed CUDA-only two-stage profile:
1. Qwen Image with Canny conditioning regenerates the frame.
2. YuNet locates faces, SAM builds masks, and Z-Image regenerates the selected
face regions.
The profile rejects a custom model identifier. Its global and face model stack
is fixed by the implementation. When tiling is enabled, only the global stage
is tiled; the face stage runs once after the tiles are blended.
#### How the face stage actually composites
The mechanism is easy to misread from the parameter names, so state it plainly.
`_run_faces` crops the expanded box from the **original** image, resizes it toward the
768 px face guide, runs Z-Image over **the entire crop**, resizes back, and only then
merges through `composite_face`, which cross-fades on a Gaussian-blurred SAM mask with
`feather=10`. The base it merges into is the global Qwen result.
Two consequences follow, and both matter when tuning:
- **Everything inside the crop is regenerated, including the pixels the mask later
discards.** The generation is therefore conditioned on a fully noised neighbourhood,
not on an intact one. An alternative design passes the mask into the sampler as a
latent noise mask, so only masked pixels are ever denoised and the edge transition
happens inside the generation rather than as a post-hoc blend. That approach has
never been tried in this runtime and is an open lever, particularly since the face
stage is the largest measured quality contributor: removing it costs 3.5 dB inside
the face boxes on one fixture and 6.1 dB on another.
- **`FACE_DENOISE_SCALE = 0.5` is best understood as compensation for the above.**
Regenerating a whole crop and blending is a stronger operation than denoising only
inside a mask, so the halved strength brings the visible result back into range. Read
it as coupled to the compositing design rather than as an independently calibrated
constant: changing the compositing without revisiting the scale would change output
strength by roughly a factor of two.
The maintained implementation preserves the previously oracle-tested strength,
conditioning, crop, and sampler parameters as compatibility contracts. Its Python
orchestration, YuNet integration, SAM selection, masks, sizing helpers, and pixel
compositing are implemented for this runtime. Changing a calibrated model input
requires the same provider-oracle and identity evaluation as a model change.
### 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.
Four things are architecture-bound and swap with the model: the ControlNet
(`xinsir/controlnet-canny-sdxl-1.0`), the four-step distillation LoRA
(`ByteDance/SDXL-Lightning` at its documented strength 1.0, not the reference graph's
0.8, which belongs to a different LoRA), the sampler (Euler with trailing spacing, no
AuraFlow shift), and the latent grid (8 px against Qwen's 16).
**Strength is architecture-bound too, and that is the easy mistake.** An SDXL global
pass leaves SynthID at the strength Qwen needs: verified through the Gemini app on a
native 2816x1536 original, 0.154 is FOUND while 0.20, 0.25 and 0.30 are clean. So this
profile takes a vendor policy (`SDXL_ZIMAGE_OPENAI_STRENGTH` 0.15,
`SDXL_ZIMAGE_GEMINI_STRENGTH` 0.25, unknown following Gemini) rather than
`resolution_adaptive_denoise`. Flat values are what was measured; no size dependence
has been established for this stage, so none is asserted.
`requested_steps` exists because the two runtimes truncate differently. DiffSynth sets
`sigma_start = denoising_strength` and runs every requested step across the shortened
sigma range; Diffusers img2img truncates the step *count*
(`init_timestep = int(steps * strength)`), so asking it for four steps at 0.15 executes
**zero** and returns a bare VAE round-trip.
**The face stage keeps its own dtype, and 0.23.0 shipped without that.** The remover
gives this profile `torch.float16`, because SDXL ships fp16 weights and an fp16-safe
VAE. That dtype reached the inherited `_load_zimage`, while `_zimage_vram_config()`
hardcodes bfloat16 for its offload, onload and computation dtypes -- so the Z-Image
modules were built bf16 and handed fp16 latents, and every image containing a face
died in the VAE with `Input type (c10::Half) and bias type (c10::BFloat16) should be
the same`. Every face-stage loader now reads `_face_stage_dtype()`, which returns the
computation dtype of the VRAM config it is paired with, so the two cannot drift again.
Two things hid this. Zero-face inputs never enter `_run_faces`, so the profile looked
healthy on exactly the images used to time it; and the profile's tests deliberately
avoid model downloads, so nothing exercised the loader. The lesson is narrower than
"add a GPU test": inheriting a stage means inheriting its *invariants*, and this one
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
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
was reading the same wrong field and would have re-landed the bug for the next profile
with a different global dtype. It is routed through the same accessor, which for
`qwen-zimage` is the bfloat16 it already used.
The guard is `test_face_stage_loads_in_its_own_dtype_when_the_global_stage_differs`.
It asserts the dtype the Z-Image and SAM loaders actually receive, not the accessor
against the config it is derived from -- that comparison would restate the
implementation and pass for any consistently wrong value. Both assertions were
mutation-tested against the pre-fix line. For `qwen-zimage` the whole change is a
strict no-op: the remover already handed it bfloat16, the same value
`_face_stage_dtype()` returns.
This profile is not deployed. Before it could be, it needs the other three Gemini
originals, OpenAI re-verified at 0.15, a flat-graphic content class, and a low
resolution case -- every verdict so far comes from one fixture and one seed.
### Measured provider boundaries for qwen-zimage
Both ends of the shipped curve now have oracle verdicts, and the shipped curve clears
everything it has been tested at:
| oracle | fixture size | detected at | clean from |
|---|---|---|---|
| openai.com/verify | 1.57 MP | 0.06 | 0.08 |
| Gemini app | 4.33 MP | 0.08 | 0.10 |
| Gemini app | 0.57 MP | -- | 0.0896 (the curve's own value) |
| Gemini app | 1.40 MP | -- | 0.1066 (the curve's own value) |
Read the last two rows before concluding the curve's low end is under-driven. Against
the 4.33 MP Gemini boundary the sub-1 MP rungs of 0.084-0.094 look short, but at those
sizes the curve's own values verify clean, which is what a resolution-scaled
requirement would predict. There is no measured size at which the shipped curve fails,
so it is left alone.
### Static prompt embeddings
Both stages prompt with module constants, and at CFG 1.0 DiffSynth's
`PipelineUnitRunner` reuses the positive embedding for the negative side instead of
encoding it. So exactly one embedding per stage is ever computed, from text that
cannot vary at runtime, which makes it cacheable across containers rather than only
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
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
(paired within five containers). The output is **byte-identical** -- the stored
tensors are the encoder's own -- so this needs no provider-oracle re-verification.
Three properties are load-bearing:
- **The key self-heals.** A model bump or a prompt edit changes the key, so the next
container recomputes rather than reading a stale embedding. `_PROMPT_CACHE_VERSION`
covers a change to the stored shape itself.
- **The write is atomic.** A torn write must never be readable as a cache hit, so the
payload lands in a temp file and is renamed into place.
- **A miss after the encoder was dropped raises.** `require_cache` records that the
stack was built without a text encoder on the strength of the file; falling back
would call a model that is not loaded, which surfaces as an opaque crash.
`_model_cache_dir()` prefers `HF_HOME` for the same reason: on a scale-to-zero runner
that is the only persistently mounted path, and anything below it is re-derived per
request. The YuNet download follows the same root.
Regression coverage:
- [`test_qwen_zimage_pipeline.py`](../tests/test_qwen_zimage_pipeline.py)
- [`test_cpu_offload.py`](../tests/test_cpu_offload.py)
### Tiling
[`_internal/tiling.py`](../src/remove_ai_watermarks/_internal/tiling.py) contains pure
tile planning, feather weights, and tile orchestration.
Tiling engages only when requested and the long side exceeds the tile size.
It avoids an explicit full-image downscale but does not make diffusion
pixel-preserving. Each tile is still regenerated.
It also held a `feather_region_composite` for AI-*enhanced* composites, where only
the edited region should change. Nothing ever reached it: the `erase` command
inpaints through `region_eraser`, and the remover's `region` argument was only
reachable from a module-level convenience wrapper with no callers. Both went.
Regression coverage:
- [`test_tiling.py`](../tests/test_tiling.py)
### Postprocessing
[`humanizer.py`](../src/remove_ai_watermarks/humanizer.py) contains explicit
grain, unsharp masking, and adaptive polish helpers.
`upscaler.py` held an optional Real-ESRGAN path, reachable only when enlarging a
small image to the minimum-resolution floor. That floor existed to lift small
inputs toward SDXL's ~1024 training size; when the SDXL profiles were removed it
was forced to 0 on every path, so the module, the `--min-resolution` and
`--upscaler` options and the `esrgan` extra were all unreachable and went with
it. Only the `max_resolution` cap can move geometry now, and it only scales down.
Regression coverage:
- [`test_humanizer.py`](../tests/test_humanizer.py)
## Image input and output
[`image_io.py`](../src/remove_ai_watermarks/image_io.py) is the shared image
codec boundary.
Contracts:
- All package OpenCV file reads and writes use `image_io.imread` and
`image_io.imwrite`.
- `to_bgr` normalizes grayscale and alpha-bearing arrays.
- `read_bgr_and_alpha` and `write_bgr_with_alpha` preserve the alpha plane.
- `imwrite` returns a success flag; every caller must check it.
- HEIC, HEIF, and AVIF pixel reads fall back to Pillow plus `pillow-heif` from
the independent `heif` extra. Metadata scanning does not require that plugin.
- A visible no-op can preserve the original file bytes.
Regression coverage:
- [`test_image_io.py`](../tests/test_image_io.py)
- [`test_cli_robustness.py`](../tests/test_cli_robustness.py)
## Adding or changing behavior
For a new visible mark:
1. create a synthetic detection silhouette;
2. add or extend a vendor engine;
3. add one registry entry;
4. test detection, false positives, localization, and actual pixel change;
5. update [supported signals](supported-signals.md).
For a new metadata signal:
1. add the scanner;
2. add every supported removal placement;
3. verify the output through `strip_and_verify`;
4. add identification and removal tests;
5. update [supported signals](supported-signals.md) and, when relevant,
[the watermarking landscape](watermarking-landscape.md).
For a diffusion change:
1. keep model-free logic in pure helpers where possible;
2. test option propagation and dispatch without downloading models;
3. run a real model smoke for the changed model path;
4. treat provider-verifier results as specific to the exact checked output;
5. update [known limitations](known-limitations.md).
Reference in New Issue View Git Blame Copy Permalink