Halve decoder memory with strip processing and a raveled Haar pass

Each Haar pass is one flat pywt.downcoef call over a raveled strip instead of
pywt.dwt(..., axis=1)[0], and the plane is walked in strips so no full-plane
float64 intermediate exists. Exact only while the last axis is even, so
_approximation raises on an odd width rather than returning wrong bits, and
TestRaveledHaarPass pins both that raise and the downcoef/dwt equivalence a
pywt upgrade could take away.

Drops the block constructor knob: the fold chains are written for 4, nothing
ever passed another value, and a knob that silently decodes wrong is worse
than no knob.

Peak RSS 111 MB to 21 MB on a 4.3 MP image; the decoder itself 0.011s to
0.007s, which is only 0.4% of identify() now that it is under 2% of the run.
Output bits and detector verdicts over 200 sampled data/ images, two
synthesized carriers and eight degenerate shapes are byte-identical to the
pre-vectorization decoder.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

docs/module-internals.md

@@ -595,18 +595,32 @@ 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.
+Each pass is one flat `pywt.downcoef` call over a raveled strip rather than
+`pywt.dwt(..., axis=1)[0]`, and the plane is processed in strips of `_STRIP`
+block-rows so no full-plane float64 intermediate is ever materialized. The strip
+height is not a tuned value: 8 through 64 all scored inside each other's noise
+with unstable ordering, and only "strips at all" versus whole-plane matters.
 
-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.
+`_approximation`'s even-last-axis check is load-bearing, not defensive. Haar's
+filter is length 2, so an even row length keeps every pair inside its own row;
+on an odd width the pairs walk across row boundaries and the reshape still
+succeeds whenever the total is even, which would be wrong bits with no
+exception. Both call sites are even by construction today, and
+`TestRaveledHaarPass` pins both halves -- the `downcoef`/`dwt` equivalence,
+which a pywt upgrade could take away, and the raise on an odd width.
+
+Measured on a 1536x2816 image, all arms timed in one process: the decoder went
+0.112 s (the original per-block Python loop) to 0.011 s vectorized to 0.007 s
+with strips, and a warm `identify()` 1.757 s to 1.365 s on the first step and a
+further 0.4% on the second. That last figure is the point at which this target
+is finished: the decoder is now under 2% of `identify()`, so speed here has
+stopped buying anything. What the strips buy is peak RSS in the stage, 111 MB to
+21 MB on a 4.3 MP image, which is what matters on the memory-limited Space.
+
+Both steps were verified by recording decoder output and detector verdict over
+200 sampled `data/` images plus two synthesized carriers before and after: the
+record is byte-identical, as are seven degenerate shapes (`1x65536` through
+`8x8192`, plus an odd width) that clear the caller's area check.
 
 Regression coverage:
 

src/remove_ai_watermarks/dwt_dct.py

@@ -27,7 +27,39 @@ if TYPE_CHECKING:
     from numpy.typing import NDArray
 
 _DEFAULT_SCALES = (0, 36, 36)
-_DEFAULT_BLOCK = 4
+# Fixed by the format being decoded, and the 4-way fold chains in `_frame_bits`
+# are written for it. It was a constructor parameter while the block scan was a
+# generic Python loop; nothing ever passed another value, and a knob that now
+# silently returns wrong bits is worse than no knob.
+_BLOCK = 4
+
+# Block-rows of the approximation band handled per strip. Not a tuned value:
+# every height measured landed inside the others' noise. What matters is strips
+# at all rather than a full-plane intermediate, not this number.
+_STRIP = 16
+
+
+def _approximation(rows: NDArray[Any]) -> NDArray[Any]:
+    """One Haar pass along the last axis, approximation band only.
+
+    ``pywt.dwt(x, "haar", axis=1)[0]`` computes and allocates the detail band as
+    well, and dispatches per row. Flattening lets one ``downcoef`` call do the
+    whole plane, and it is the same numbers in the same order **only while the
+    last axis is even**: Haar's filter is length 2, so an even row length makes
+    every pair fall inside its own row with no boundary extension. On an odd
+    width the pairs walk across row boundaries and the reshape below still
+    succeeds whenever the total is even -- wrong bits, no exception, and the
+    upstream-parity test is `skipif`-gated. Hence the explicit check.
+
+    The transform itself stays inside pywt however slow that looks. Its C
+    convolution contracts into an FMA that no numpy expression reproduces, and
+    the caller's threshold is ``peak % 36 > 18.0`` against values that are exact
+    multiples of 0.5, so a 1-ulp difference flips real bits.
+    """
+    height, width = rows.shape
+    if width % 2:
+        raise RuntimeError(f"row length {width} is odd; the raveled Haar pass requires an even last axis")
+    return pywt.downcoef("a", rows.ravel(), "haar").reshape(height, width // 2)
 
 
 class _DecodeMaxDct:
@@ -37,11 +69,9 @@ class _DecodeMaxDct:
         self,
         wm_lengths: tuple[int, ...],
         scales: tuple[int, int, int] = _DEFAULT_SCALES,
-        block: int = _DEFAULT_BLOCK,
     ) -> None:
         self._wm_lengths = wm_lengths
         self._scales = scales
-        self._block = block
 
     def decode(self, bgr: NDArray[Any]) -> dict[int, NDArray[Any]]:
         row, col, _channels = bgr.shape
@@ -49,7 +79,7 @@ class _DecodeMaxDct:
         trimmed = yuv[: row // 4 * 4, : col // 4 * 4]
 
         per_channel = [
-            self._frame_bits(self._approximation(trimmed, channel), self._scales[channel])
+            self._plane_bits(trimmed, channel, self._scales[channel])
             for channel in range(2)
             if self._scales[channel] > 0
         ]
@@ -66,43 +96,57 @@ class _DecodeMaxDct:
             decoded[wm_len] = sums * 255 > counts * 127
         return decoded
 
-    @staticmethod
-    def _approximation(trimmed: NDArray[Any], channel: int) -> NDArray[Any]:
-        """The Haar approximation band, and only it.
+    def _plane_bits(self, trimmed: NDArray[Any], channel: int, scale: int) -> NDArray[Any]:
+        """Block bits for one colour plane, a strip of block-rows at a time.
 
         ``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.
+        Only the approximation is ever asked for, and transposing between the
+        two passes lets pywt walk a contiguous axis instead of a column -- that
+        access pattern, not the arithmetic saved, is where the time goes.
 
-        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.
+        Strips mean no full-plane float64 intermediate is ever materialized, and
+        they are seam-free for the reason ``_approximation`` documents: output
+        row ``k`` reads input rows ``2k`` and ``2k + 1`` only. Strips start on
+        multiples of ``2 * _BLOCK``, so neither a pair nor a 4x4 block straddles
+        one.
         """
         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.
+            # Reachable: a 1x65536 image clears the caller's area check and
+            # trims to an empty plane. Left to dwt2 so the exception stays the
+            # one this module has always raised -- returning empty bits here
+            # instead would silently turn a raise into an all-false verdict.
             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]
+        rows = trimmed.shape[0] // (2 * _BLOCK)
+        cols = trimmed.shape[1] // (2 * _BLOCK)
+        if rows == 0 or cols == 0:
+            return np.zeros(0, dtype=np.float64)
+
+        width = cols * 2 * _BLOCK
+        pieces: list[NDArray[Any]] = []
+        for start in range(0, rows, _STRIP):
+            stop = min(start + _STRIP, rows)
+            strip = trimmed[start * 2 * _BLOCK : stop * 2 * _BLOCK, :width]
+            columns = cv2.transpose(cv2.extractChannel(strip, channel))
+            band = _approximation(cv2.transpose(_approximation(columns)))
+            pieces.append(self._frame_bits(band, scale))
+        return np.concatenate(pieces)
 
     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.
+        interpreter ~135k times per image. Zeroing the DC term in the absolute
+        band says "ignore index 0" without materializing a ``(nblocks, 16)``
+        copy, and the 4-way ``np.maximum`` chains reduce over contiguous rows.
         """
-        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)
+        rows = frame.shape[0] // _BLOCK
+        cols = frame.shape[1] // _BLOCK
+        band = np.abs(frame[: rows * _BLOCK, : cols * _BLOCK])
+        band[::_BLOCK, ::_BLOCK] = 0.0
+        folded = np.maximum(np.maximum(band[0::4], band[1::4]), np.maximum(band[2::4], band[3::4]))
+        folded = folded.reshape(rows * cols, _BLOCK)
+        peak = np.maximum(np.maximum(folded[:, 0], folded[:, 1]), np.maximum(folded[:, 2], folded[:, 3]))
         return ((peak % scale) > 0.5 * scale).astype(np.float64)
 
 

tests/test_invisible_watermark.py

@@ -59,6 +59,47 @@ class TestHelpers:
         assert _bytes_match_frac(b"abc", b"abcd") == 0.0
 
 
+class TestRaveledHaarPass:
+    """The precondition that makes the decoder's flat Haar pass legitimate.
+
+    `_approximation` replaces `pywt.dwt(x, "haar", axis=1)[0]` with one
+    `downcoef` call over `x.ravel()`. That is exact only while the last axis is
+    even. Neither half of this is checked anywhere else: the equivalence is a
+    property of pywt's implementation that an upgrade could take away, and an
+    odd width produces wrong bits with no exception, since the reshape still
+    succeeds whenever the total length is even.
+    """
+
+    @pytest.mark.parametrize("shape", [(64, 64), (7, 128), (129, 2), (2, 2), (33, 400)])
+    def test_matches_pywt_dwt_on_even_widths(self, shape: tuple[int, int]):
+        import pywt
+
+        from remove_ai_watermarks.dwt_dct import _approximation
+
+        rng = np.random.default_rng(0)
+        # uint8 is what the FIRST production pass receives -- `decode` builds its
+        # plane with cvtColor, and extractChannel and transpose preserve the
+        # dtype -- so a divergence in how downcoef coerces integers would be
+        # invisible to a float-only parametrization.
+        for array in (
+            rng.random(shape),
+            rng.integers(0, 256, shape).astype(np.float64),
+            rng.integers(0, 256, shape).astype(np.uint8),
+        ):
+            expected = pywt.dwt(array, "haar", axis=1)[0]
+            got = _approximation(array)
+            assert got.shape == expected.shape
+            assert np.array_equal(got, expected), "downcoef diverged from dwt -- a pywt upgrade may have changed it"
+
+    def test_odd_width_raises_instead_of_returning_wrong_bits(self):
+        from remove_ai_watermarks.dwt_dct import _approximation
+
+        # 4x6 ravels to 24, an even total, so the reshape would happily produce
+        # a 4x3 array of numbers that pair across row boundaries.
+        with pytest.raises(RuntimeError, match="odd"):
+            _approximation(np.zeros((4, 6))[:, :5])
+
+
 class TestDetect:
     def test_in_tree_decoder_matches_upstream(self, tmp_path: Path):
         from imwatermark import WatermarkDecoder