From 37789e02f695dc3e2bbc3f3c6a70c8f397af41e8 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Thu, 13 Aug 2026 11:59:56 -0700 Subject: [PATCH] Vectorize the DWT-DCT decode path, 15x on the decoder Output stays bit-identical: decoder bits and detector verdicts recorded over 200 sampled data/ images plus two synthesized carriers before and after, and the record is byte-identical. Measured on a 1536x2816 image -- decoder 0.280s to 0.016s, warm identify() 1.757s to 1.365s with both arms timed in one process. Co-Authored-By: Claude Opus 5 --- .claude/rules/development.md | 9 +++ docs/code-provenance.md | 5 ++ docs/module-internals.md | 48 +++++++++++++-- src/remove_ai_watermarks/dwt_dct.py | 95 ++++++++++++++++++----------- 4 files changed, 117 insertions(+), 40 deletions(-) diff --git a/.claude/rules/development.md b/.claude/rules/development.md index c2a95aa..9db0237 100644 --- a/.claude/rules/development.md +++ b/.claude/rules/development.md @@ -98,6 +98,15 @@ record is byte-identical, and a green test suite does not establish that on its change that is meant to FIX detection is the exception that proves the rule: the diff must then be exactly the files you intended to change, named in advance. +Two corollaries in `dwt_dct.py`, where "close enough" has an exact meaning. The bit test +is `peak % 36 > 18.0` and for uint8 input the exact Haar value is a multiple of 0.5, so it +lands ON the threshold once per 72 blocks and a 1-ulp difference flips real bits: leave the +transform to `pywt` however slow it looks, because its C convolution contracts into an FMA +that numpy has no ufunc for (`docs/module-internals.md` carries what that costs). And +`tests/test_invisible_watermark.py` is `skipif(not is_available())` -- without the `detect` +extra the upstream-parity test never runs, so a green suite is not evidence here and the +verdict record is not optional. + ## A certified operating point is data, not a constant The video SynthID default is only meaningful as a row in diff --git a/docs/code-provenance.md b/docs/code-provenance.md index d26ba2c..eb5de34 100644 --- a/docs/code-provenance.md +++ b/docs/code-provenance.md @@ -7,6 +7,11 @@ This page records notices required by source dependencies and licensed derivativ - The DWT-DCT implementation derives from ShieldMnt's [`invisible-watermark`](https://github.com/ShieldMnt/invisible-watermark), licensed under MIT. Its notice ships in `src/remove_ai_watermarks/licenses/invisible-watermark-MIT.txt`. + The decode path is a vectorized reformulation, not a transcription: it produces + the same bits and is checked against upstream's own decoder, but it no longer + corresponds line by line to `imwatermark/maxDct.py`. Diffing the two files will + show structurally different code, which is expected and does not mean the + derivation notice is stale. ## Licensed test fixtures diff --git a/docs/module-internals.md b/docs/module-internals.md index 69c7784..cfb2fa3 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -555,9 +555,11 @@ the function it replaces does — no provenance means no relaxation, and an unkn invisible target means scrub rather than skip. 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/`. +in-tree [`dwt_dct.py`](../src/remove_ai_watermarks/dwt_dct.py) decoder reproduces +the upstream algorithm's output bit for bit without installing Torch or +non-headless OpenCV; the block scan is vectorized rather than transcribed, so +the file no longer reads line by line against `maxDct.py`. 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 @@ -572,11 +574,49 @@ schemas 0-2 count as positives. Schema 3 is below the precision threshold: all same false payload after re-encoding. The official Adobe Variant P schema-1 fixture in `data/fixtures/provenance/` is the positive regression control. +#### Why the DWT-DCT decoder looks the way it does + +Two things in `dwt_dct.py` are load-bearing and neither is obvious from the code. + +`_approximation` calls `pywt.dwt` twice instead of `pywt.dwt2`, and transposes +before each pass. A factorial ablation over both axes separates the two: +skipping the three detail bands `dwt2` computes and this code discards is worth +**4%**, while the transposes are worth **2.8x**, because pywt walks the axis it +transforms and on axis 0 of a C-contiguous plane that is a column walk. The +arithmetic saving is the intuitive explanation and it is the small term; four +independent profiles named it as the mechanism before the ablation contradicted +them. + +Bit-identity is a hard requirement, not a preference. For uint8 input the exact +Haar LL value is a multiple of 0.5 and the bit test is `peak % 36 > 18.0`, a +threshold sitting exactly on a representable value that ~1 block in 72 lands on, +so a 1-ulp difference deterministically flips real bits. That is why the ~16x +available from a hand-rolled numpy Haar is unreachable rather than merely +untaken: pywt's C convolution contracts into an FMA that numpy has no ufunc for, +and `np.longdouble` is 64-bit on arm64 macOS. + +Measured on a 1536x2816 image: the decoder went 0.280 s to 0.019 s, and a warm +`identify()` on the same file 1.757 s to 1.365 s (-22.3%), both arms timed in one +process. Verified by recording decoder output and detector verdict over 200 +sampled `data/` images plus two synthesized carriers before and after the change: +the record is byte-identical, as are five degenerate shapes (`1x65536` through +`8x8192`) that clear the caller's area check. + +Two further optimizations were identified and deliberately not taken, because +each buys speed with a new precondition rather than with less work: +`pywt.downcoef("a", ...)` on a flattened plane (valid only while the last axis +stays even and C-contiguous), and a single-pass `np.abs(..., out=)` over the +block gather. Both would need their own ablation. + 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) +- [`test_invisible_watermark.py`](../tests/test_invisible_watermark.py) -- + note `test_in_tree_decoder_matches_upstream` is the parity guard against + upstream's own decoder, and the whole module is `skipif(not is_available())`. + A green run without the `detect` extra installed has not checked parity at + all, so a decoder change still owes the before/after verdict record. ## Visible mark removal diff --git a/src/remove_ai_watermarks/dwt_dct.py b/src/remove_ai_watermarks/dwt_dct.py index 7631a7d..5f45898 100644 --- a/src/remove_ai_watermarks/dwt_dct.py +++ b/src/remove_ai_watermarks/dwt_dct.py @@ -1,7 +1,11 @@ """DWT-DCT decoder compatible with invisible-watermark's ``dwtDct`` path. Derived from ShieldMnt/invisible-watermark ``imwatermark/maxDct.py`` (MIT), -trimmed to the matrix path used by Stable Diffusion, SDXL, and FLUX. +trimmed to the matrix path used by Stable Diffusion, SDXL, and FLUX. The block +scan is vectorized rather than transcribed, so the file no longer reads line by +line against upstream; what it preserves is the output, bit for bit. See +[`docs/module-internals.md`](../../docs/module-internals.md) for the +measurements and for why a faster hand-rolled transform is not available. Copyright (c) 2021 ShieldMnt @@ -42,45 +46,64 @@ class _DecodeMaxDct: def decode(self, bgr: NDArray[Any]) -> dict[int, NDArray[Any]]: row, col, _channels = bgr.shape yuv = cv2.cvtColor(bgr, cv2.COLOR_BGR2YUV) + trimmed = yuv[: row // 4 * 4, : col // 4 * 4] - scores_by_length = {wm_len: ([0] * wm_len, [0] * wm_len) for wm_len in self._wm_lengths} - for channel in range(2): - if self._scales[channel] <= 0: - continue - ca1, _detail = pywt.dwt2(yuv[: row // 4 * 4, : col // 4 * 4, channel], "haar") - self._decode_frame(ca1, self._scales[channel], scores_by_length) + per_channel = [ + self._frame_bits(self._approximation(trimmed, channel), self._scales[channel]) + for channel in range(2) + if self._scales[channel] > 0 + ] + # Each channel restarts the bit index at 0, so the buckets come from a + # per-channel arange rather than one running counter. + index = np.concatenate([np.arange(bits.size) for bits in per_channel] or [np.zeros(0, dtype=np.int64)]) + weights = np.concatenate(per_channel or [np.zeros(0)]) - return { - wm_len: np.asarray(sums) * 255 > np.asarray(counts) * 127 - for wm_len, (sums, counts) in scores_by_length.items() - } + decoded: dict[int, NDArray[Any]] = {} + for wm_len in self._wm_lengths: + bucket = index % wm_len + sums = np.bincount(bucket, weights=weights, minlength=wm_len) + counts = np.bincount(bucket, minlength=wm_len) + decoded[wm_len] = sums * 255 > counts * 127 + return decoded - def _decode_frame( - self, - frame: NDArray[Any], - scale: int, - scores_by_length: dict[int, tuple[list[int], list[int]]], - ) -> None: - row, col = frame.shape - bit_index = 0 - for i in range(row // self._block): - for j in range(col // self._block): - block = frame[ - i * self._block : i * self._block + self._block, - j * self._block : j * self._block + self._block, - ] - inferred = self._infer_bit(block, scale) - for wm_len, (sums, counts) in scores_by_length.items(): - bucket = bit_index % wm_len - sums[bucket] += inferred - counts[bucket] += 1 - bit_index += 1 + @staticmethod + def _approximation(trimmed: NDArray[Any], channel: int) -> NDArray[Any]: + """The Haar approximation band, and only it. - def _infer_bit(self, block: NDArray[Any], scale: int) -> int: - position = int(np.argmax(np.abs(block.flatten()[1:]))) + 1 - i, j = position // self._block, position % self._block - value = abs(float(block[i][j])) - return int((value % scale) > 0.5 * scale) + ``dwt2`` is ``dwtn``: it transforms along axis 0, then along axis 1 over + both halves, and three of the four bands it returns are discarded here. + Two ``dwt`` calls keeping ``[0]`` skip that, and transposing between them + lets pywt walk a contiguous axis instead of a column. + + The result must stay bit-identical to ``dwt2``'s, which is why the + transform is left to pywt however slow that is: the caller's threshold is + ``peak % 36 > 18.0``, and for uint8 input the exact value is a multiple + of 0.5, so it lands exactly on the threshold often enough that a 1-ulp + difference flips real bits. + """ + if trimmed.shape[0] == 0 or trimmed.shape[1] == 0: + # Reachable: a 1x65536 image clears the caller's area check. Left to + # dwt2 so the exception stays the one this module has always raised. + return pywt.dwt2(trimmed[:, :, channel], "haar")[0] + columns = cv2.transpose(cv2.extractChannel(trimmed, channel)) + along_rows = pywt.dwt(columns, "haar", axis=1)[0] + return pywt.dwt(cv2.transpose(along_rows), "haar", axis=1)[0] + + def _frame_bits(self, frame: NDArray[Any], scale: int) -> NDArray[Any]: + """One bit per 4x4 block, in row-major block order. + + Upstream's per-block loop, said to numpy once instead of to the + interpreter ~135k times per image. + """ + block = self._block + rows = frame.shape[0] // block + cols = frame.shape[1] // block + if rows == 0 or cols == 0: + return np.zeros(0, dtype=np.float64) + aligned = frame[: rows * block, : cols * block] + blocks = aligned.reshape(rows, block, cols, block).swapaxes(1, 2) + peak = np.abs(blocks.reshape(rows * cols, block * block)[:, 1:]).max(axis=1) + return ((peak % scale) > 0.5 * scale).astype(np.float64) def decode_dwt_dct(bgr: NDArray[Any], wm_len: int) -> NDArray[Any]: