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>

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)