pbc3: promote simplified codec

PBC3.py

@@ -1,10 +1,3 @@
-# ====================================================================================================
-#
-#           PBC v3.0 - Probabilistic Brush Compression
-#           Lossy Image Compression Algorithm by EgeEken (github.com/EgeEken)
-#           3.0 Update - 2026-06 - Whole algorithm overhaul
-#
-# ====================================================================================================
 import lzma
 import math
 import time
@@ -12,11 +5,43 @@ import time
 import numpy as np
 from PIL import Image, ImageOps
 
+import pbc3_stream as stream
 import pbc3_ops as ops
 from pbc3_heads import DownsampleInitHead, FillerHead, SearchHead
 from pbc3_types import BitReader, BitWriter, PBC3Config, PBC3Result
 
 
+def _config(config, kwargs):
+    if config is None:
+        return PBC3Config(**kwargs)
+    return PBC3Config(**{**config.__dict__, **kwargs}) if kwargs else config
+
+
+def _canvas_from_bases(shape, bases):
+    canvas = np.zeros(shape, dtype=np.int32)
+    for channel, base in enumerate(bases):
+        canvas[:, :, channel] = base
+    return canvas
+
+
+def _validate(prep, config):
+    h, w = prep["h"], prep["w"]
+    ow, oh = prep["original_w"], prep["original_h"]
+    if max(w, h, ow, oh) > 65535:
+        raise ValueError("this prototype stores dimensions as uint16")
+    if not 1 <= config.mask_size <= 1023:
+        raise ValueError("mask_size must be in 1..1023")
+    if config.auto_downsample_max_pixels < 1:
+        raise ValueError("auto_downsample_max_pixels must be >= 1")
+    if not (
+        1 <= config.downsample_palette_bitcount <= 9
+        and 1 <= config.patch_palette_bitcount <= 9
+    ):
+        raise ValueError("palette bitcounts must be in 1..9")
+    if str(config.channel_cycle).lower() not in {"sum", "mod"}:
+        raise ValueError('channel_cycle must be "Sum" or "Mod"')
+
+
 class PBC3:
     MAGIC = b"PBC3"
     VERSION = 0
@@ -30,8 +55,6 @@ class PBC3:
     RESAMPLE_FILTER = ops.RESAMPLE_FILTER
     RESAMPLE_REDUCING_GAP = ops.RESAMPLE_REDUCING_GAP
 
-    # ---- image conversion -------------------------------------------------
-
     @staticmethod
     def _to_image(image) -> Image.Image:
         """## Returns a PIL image from a path, PIL image, or image-like array"""
@@ -60,8 +83,6 @@ class PBC3:
             return color
         return Image.fromarray(arr, color_space).convert("RGB")
 
-    # ---- entropy + bitstream framing -------------------------------------
-
     @classmethod
     def _entropy_pack(cls, body: bytes, use_lzma: bool = True) -> tuple[int, bytes]:
         """## Returns the smaller of raw body or LZMA-compressed body"""
@@ -91,131 +112,18 @@ class PBC3:
             raise ValueError(f"unsupported PBC3 version {version}")
         return version, cls._entropy_unpack(data[5], data[6:])
 
-    @classmethod
-    def _write_grid(cls, bw: BitWriter, flat, bitcount: int) -> None:
-        """## Writes a flat grid of palette indices"""
-        for value in flat:
-            bw.write(int(value), bitcount)
-
-    @classmethod
-    def _read_grid(cls, br: BitReader, n: int, bitcount: int) -> np.ndarray:
-        """## Reads a flat grid of palette indices"""
-        flat = np.zeros(n, dtype=np.uint16)
-        for k in range(n):
-            flat[k] = br.read(bitcount)
-        return flat
-
-    @classmethod
-    def _write_header(
-        cls,
-        bw: BitWriter,
-        w: int,
-        h: int,
-        original_w: int,
-        original_h: int,
-        downsampled: bool,
-        color_id: int,
-        channels: int,
-        channel_bits: int,
-        positive_bias: bool,
-        has_alpha: bool,
-        patch_count: int,
-        base_values,
-        warmup=None,
-    ) -> None:
-        """## Writes the image-level stream header"""
-        bw.write(int(downsampled), 1)
-        if downsampled:
-            bw.write(original_w, 16)
-            bw.write(original_h, 16)
-        bw.write(w, 16)
-        bw.write(h, 16)
-        bw.write(color_id, 2)
-        bw.write(channels, 8)
-        bw.write(channel_bits, 4)
-        bw.write(int(positive_bias), 1)
-        bw.write(int(has_alpha), 1)
-        bw.write(patch_count, 32)
-        for base in base_values:
-            bw.write(base, 8)
-        bw.write(int(warmup is not None), 1)
-        if warmup is not None:
-            wm_w, wm_h, wm_split = warmup
-            bw.write(wm_w, 16)
-            bw.write(wm_h, 16)
-            bw.write(wm_split, 32)
-
-    @classmethod
-    def _read_header(cls, br: BitReader):
-        """## Reads the image-level stream header"""
-        downsampled = bool(br.read(1))
-        original_w = br.read(16) if downsampled else None
-        original_h = br.read(16) if downsampled else None
-        w = br.read(16)
-        h = br.read(16)
-        color_id = br.read(2)
-        channels = br.read(8)
-        channel_bits = br.read(4)
-        positive_bias = bool(br.read(1))
-        has_alpha = bool(br.read(1))
-        patch_count = br.read(32)
-        base_values = [br.read(8) for _ in range(channels)]
-        warmup_on = bool(br.read(1))
-        warm_w = warm_h = warmup_split = None
-        if warmup_on:
-            warm_w = br.read(16)
-            warm_h = br.read(16)
-            warmup_split = br.read(32)
-        return (
-            downsampled, original_w, original_h, w, h, cls.COLOR_SPACE_NAMES[color_id], channels,
-            channel_bits, positive_bias, has_alpha, patch_count, base_values, warmup_on, warm_w,
-            warm_h, warmup_split,
-        )
+    _write_grid = staticmethod(stream.write_grid)
+    _read_grid = staticmethod(stream.read_grid)
+    _write_header = staticmethod(stream.write_header)
+    _write_patch = staticmethod(stream.write_patch)
 
     @classmethod
-    def _write_patch(cls, bw: BitWriter, patch, channel_bits: int) -> None:
-        """## Writes one generated-palette patch"""
-        bw.write(patch["channel"], channel_bits)
-        bw.write(patch["x"], 16)
-        bw.write(patch["y"], 16)
-        bw.write(patch["w"], 16)
-        bw.write(patch["h"], 16)
-        bw.write(cls.PALETTE_GENERATED, 1)
-        mask = patch["mask"]
-        bw.write(len(mask), 10)
-        for bit in mask:
-            bw.write(bit, 1)
-        bw.write(patch["neg"], 8)
-        bw.write(patch["pos"], 8)
-        bw.write(patch["max_bitcount"], 4)
-        bw.write(patch["cell_size"], 16)
-        cls._write_grid(bw, patch["indices"].ravel().astype(np.int64), patch["bitcount"])
+    def _read_header(cls, br):
+        return stream.read_header(br, cls.COLOR_SPACE_NAMES)
 
-    @classmethod
-    def _read_patch(cls, br: BitReader, channel_bits: int, positive_bias: bool = True):
-        """## Reads one generated-palette patch and returns its decoded values"""
-        channel = br.read(channel_bits)
-        x = br.read(16)
-        y = br.read(16)
-        w = br.read(16)
-        h = br.read(16)
-        pm = br.read(1)
-        if pm != cls.PALETTE_GENERATED:
-            raise ValueError("explicit palette patches were removed in PBC3 3.0 release cleanup")
-        mask_size = br.read(10)
-        mask = [br.read(1) for _ in range(mask_size)]
-        negative_max = br.read(8)
-        positive_max = br.read(8)
-        max_bitcount = br.read(4)
-        bitcount = ops.resolve_palette_bitcount(mask, max_bitcount, negative_max, positive_max, positive_bias)
-        pal = ops.palette_generator(mask, max_bitcount, negative_max, positive_max, positive_bias)
-        cell_size = br.read(16)
-        gw = ops.ceil_div(w, cell_size)
-        gh = ops.ceil_div(h, cell_size)
-        indices = cls._read_grid(br, gh * gw, bitcount).reshape(gh, gw)
-        return channel, x, y, w, h, cell_size, pal[indices], bitcount
-
-    # ---- image preparation -----------------------------------------------
+    @staticmethod
+    def _read_patch(br, channel_bits: int, positive_bias: bool = True):
+        return stream.read_patch(br, channel_bits, positive_bias)
 
     @classmethod
     def _auto_downsample_rate(cls, image_size, downsample_rate: float, max_pixels: int) -> float:
@@ -268,10 +176,7 @@ class PBC3:
     @classmethod
     def prepare(cls, image, config: PBC3Config = None, **kwargs) -> dict:
         """## Prepares the source image and reusable encoder arrays"""
-        if config is None:
-            config = PBC3Config(**kwargs)
-        elif kwargs:
-            config = PBC3Config(**{**config.__dict__, **kwargs})
+        config = _config(config, kwargs)
 
         src = cls._to_image(image)
         has_alpha = cls._has_alpha(src)
@@ -286,12 +191,17 @@ class PBC3:
             orig_compare = src.convert("RGB")
 
         original_w, original_h = color_img.size
-        rate = cls._auto_downsample_rate(color_img.size, config.downsample_rate, config.auto_downsample_max_pixels)
+        rate = cls._auto_downsample_rate(
+            color_img.size, config.downsample_rate, config.auto_downsample_max_pixels
+        )
         color_ds = cls._downsample_image(color_img, rate)
         downsampled = color_ds.size != color_img.size
         arr = np.asarray(color_ds, dtype=np.uint8)
         if has_alpha:
-            alpha_ds = alpha_img.resize(color_ds.size, cls.RESAMPLE_FILTER, reducing_gap=cls.RESAMPLE_REDUCING_GAP) if downsampled else alpha_img
+            alpha_ds = (
+                alpha_img.resize(color_ds.size, cls.RESAMPLE_FILTER, reducing_gap=cls.RESAMPLE_REDUCING_GAP)
+                if downsampled else alpha_img
+            )
             arr = np.dstack([arr, np.asarray(alpha_ds, dtype=np.uint8)])
 
         warm_plan = cls._warmup_plan(config, color_img.size, rate)
@@ -302,7 +212,9 @@ class PBC3:
             warm_w, warm_h = warm_color_ds.size
             warm_arr = np.asarray(warm_color_ds, dtype=np.uint8)
             if has_alpha:
-                warm_alpha = alpha_img.resize(warm_color_ds.size, cls.RESAMPLE_FILTER, reducing_gap=cls.RESAMPLE_REDUCING_GAP)
+                warm_alpha = alpha_img.resize(
+                    warm_color_ds.size, cls.RESAMPLE_FILTER, reducing_gap=cls.RESAMPLE_REDUCING_GAP
+                )
                 warm_arr = np.dstack([warm_arr, np.asarray(warm_alpha, dtype=np.uint8)])
             warm_target = warm_arr.astype(np.int32)
 
@@ -327,12 +239,14 @@ class PBC3:
             "warm_target": warm_target,
         }
 
-    # ---- encode -----------------------------------------------------------
-
     @staticmethod
     def _choose_channel(scores, step: int, channels: int, mode: str) -> int:
         """## Chooses the next channel by round-robin or current total error"""
-        return (step - 1) % channels if str(mode).lower() == "mod" else int(max(range(channels), key=lambda c: scores[c]))
+        return (
+            (step - 1) % channels
+            if str(mode).lower() == "mod"
+            else int(max(range(channels), key=lambda c: scores[c]))
+        )
 
     @staticmethod
     def _channel_sum_error(target, canvas, c: int) -> float:
@@ -351,10 +265,7 @@ class PBC3:
     @classmethod
     def compress_stream(cls, image, config: PBC3Config = None, *, reuse=None, frame_every: int = 25, **kwargs):
         """## Compresses an image and yields optional preview frames plus the final result"""
-        if config is None:
-            config = PBC3Config(**kwargs)
-        elif kwargs:
-            config = PBC3Config(**{**config.__dict__, **kwargs})
+        config = _config(config, kwargs)
 
         t0 = time.perf_counter()
         debug_lines = []
@@ -370,37 +281,37 @@ class PBC3:
         did_warmup = False
         warmup_split = None
 
-        if w > 65535 or h > 65535 or original_w > 65535 or original_h > 65535:
-            raise ValueError("this prototype stores dimensions as uint16")
-        if config.mask_size < 1 or config.mask_size > 1023:
-            raise ValueError("mask_size must be in 1..1023")
-        if config.auto_downsample_max_pixels < 1:
-            raise ValueError("auto_downsample_max_pixels must be >= 1")
-        if not (1 <= config.downsample_palette_bitcount <= 9 and 1 <= config.patch_palette_bitcount <= 9):
-            raise ValueError("palette bitcounts must be in 1..9")
-        if str(config.channel_cycle).lower() not in {"sum", "mod"}:
-            raise ValueError('channel_cycle must be "Sum" or "Mod"')
+        _validate(prep, config)
 
         channel_bits = max(1, math.ceil(math.log2(channels)))
         base_values = [int(round(float(np.mean(arr[:, :, c])))) for c in range(channels)]
-        canvas = np.zeros((h, w, channels), dtype=np.int32)
-        for c, base in enumerate(base_values):
-            canvas[:, :, c] = base
+        canvas = _canvas_from_bases((h, w, channels), base_values)
         if frame_every:
-            yield {"event": "frame", "step": 0, "total": int(config.patch_count), "image": cls._canvas_to_image(canvas, config.color_space, has_alpha)}
+            yield {
+                "event": "frame", "step": 0, "total": int(config.patch_count),
+                "image": cls._canvas_to_image(canvas, config.color_space, has_alpha),
+            }
 
         patches = []
         init_head = DownsampleInitHead()
         for c in range(channels):
-            patch, values, init_cell, init_bits = init_head.select(c, target, canvas, w, h, config, channel_bits)
+            patch, values, init_cell, init_bits = init_head.select(
+                c, target, canvas, w, h, config, channel_bits
+            )
             if config.debug_print:
                 print(f"[auto-init] channel {c}: cell={init_cell}, bitcount={init_bits}")
             ops.apply_grid(canvas[:, :, c], 0, 0, w, h, init_cell, values)
             patches.append(patch)
             if config.debug_mode:
-                debug_lines.append(ops.debug_line("INIT", stream_patch=len(patches), channel=c, x=0, y=0, w=w, h=h, cell_size=init_cell, bitcount=init_bits))
+                debug_lines.append(ops.debug_line(
+                    "INIT", stream_patch=len(patches), channel=c, x=0, y=0, w=w, h=h,
+                    cell_size=init_cell, bitcount=init_bits,
+                ))
         if frame_every:
-            yield {"event": "frame", "step": 0, "total": int(config.patch_count), "image": cls._canvas_to_image(canvas, config.color_space, has_alpha)}
+            yield {
+                "event": "frame", "step": 0, "total": int(config.patch_count),
+                "image": cls._canvas_to_image(canvas, config.color_space, has_alpha),
+            }
 
         channel_scores = [cls._channel_sum_error(target, canvas, c) for c in range(channels)]
         quality_target = float(config.quality_target_mae)
@@ -410,13 +321,22 @@ class PBC3:
         applied = 0
         for step in range(1, max(0, int(config.patch_count)) + 1):
             current_channel = cls._choose_channel(channel_scores, step, channels, config.channel_cycle)
-            boxes = None if filler.learned is not None else search.propose(target, canvas, config, rng, step, current_channel)
-            patch, values = filler.select(target, canvas, config, rng, channel_bits, step, current_channel, boxes, len(patches), debug_lines)
+            boxes = (
+                None if filler.learned is not None
+                else search.propose(target, canvas, config, rng, step, current_channel)
+            )
+            patch, values = filler.select(
+                target, canvas, config, rng, channel_bits, step, current_channel,
+                boxes, len(patches), debug_lines,
+            )
             if patch is None:
                 break
 
             c = patch["channel"]
-            ops.apply_grid(canvas[:, :, c], patch["x"], patch["y"], patch["w"], patch["h"], patch["cell_size"], values)
+            ops.apply_grid(
+                canvas[:, :, c], patch["x"], patch["y"], patch["w"], patch["h"],
+                patch["cell_size"], values,
+            )
             patches.append(patch)
             channel_scores[c] = cls._channel_sum_error(target, canvas, c)
             applied += 1
@@ -428,12 +348,21 @@ class PBC3:
                 warmup_split = len(patches)
                 did_warmup = True
             if config.debug_mode:
-                debug_lines.append(ops.debug_line("APPLIED", patch_step=step, stream_patch=len(patches), channel=c, channel_score=f"{channel_scores[c]:.4f}", x=patch["x"], y=patch["y"], w=patch["w"], h=patch["h"], cell_size=patch["cell_size"]))
+                debug_lines.append(ops.debug_line(
+                    "APPLIED", patch_step=step, stream_patch=len(patches), channel=c,
+                    channel_score=f"{channel_scores[c]:.4f}", x=patch["x"], y=patch["y"],
+                    w=patch["w"], h=patch["h"], cell_size=patch["cell_size"],
+                ))
             if config.debug_print:
                 print("|", end="", flush=True)
             if frame_every and applied % frame_every == 0:
-                yield {"event": "frame", "step": step, "total": int(config.patch_count), "image": cls._canvas_to_image(canvas, config.color_space, has_alpha)}
-            if quality_target > 0 and float(np.mean(np.abs(target - np.clip(canvas, 0, 255)))) <= quality_target:
+                yield {
+                    "event": "frame", "step": step, "total": int(config.patch_count),
+                    "image": cls._canvas_to_image(canvas, config.color_space, has_alpha),
+                }
+            if quality_target > 0 and float(
+                np.mean(np.abs(target - np.clip(canvas, 0, 255)))
+            ) <= quality_target:
                 break
         if config.debug_print:
             print()
@@ -473,8 +402,6 @@ class PBC3:
             ),
         }
 
-    # ---- decode (deterministic, ML-free) ---------------------------------
-
     @classmethod
     def _decode_to_canvas(cls, data, max_patches: int = None):
         """## Decodes a PBC3 stream to the internal canvas without making a PIL image"""
@@ -484,17 +411,24 @@ class PBC3:
         version, body = cls._open_body(data)
         br = BitReader(body)
         header = cls._read_header(br)
-        downsampled, original_w, original_h, w, h, color_space, channels, channel_bits, positive_bias, has_alpha, patch_count, base_values, warmup_on, warm_w, warm_h, warmup_split = header
-        canvas = np.zeros((h, w, channels), dtype=np.int32)
-        for c, base in enumerate(base_values):
-            canvas[:, :, c] = base
+        (
+            downsampled, original_w, original_h, w, h, color_space, channels,
+            channel_bits, positive_bias, has_alpha, patch_count, base_values,
+            warmup_on, warm_w, warm_h, warmup_split,
+        ) = header
+        canvas = _canvas_from_bases((h, w, channels), base_values)
         patches_to_read = patch_count if max_patches is None else min(int(max_patches), patch_count)
         for idx in range(patches_to_read):
             if warmup_on and idx == warmup_split:
                 canvas = cls._resize_canvas(canvas, warm_w, warm_h)
-            channel, x, y, pw, ph, cell_size, values, _ = cls._read_patch(br, channel_bits, positive_bias)
+            channel, x, y, pw, ph, cell_size, values, _ = cls._read_patch(
+                br, channel_bits, positive_bias
+            )
             ops.apply_grid(canvas[:, :, channel], x, y, pw, ph, cell_size, values)
-        return canvas, color_space, downsampled, original_w, original_h, canvas.shape[1], canvas.shape[0], has_alpha, channels, patch_count
+        return (
+            canvas, color_space, downsampled, original_w, original_h,
+            canvas.shape[1], canvas.shape[0], has_alpha, channels, patch_count,
+        )
 
     @classmethod
     def decompress(cls, data, max_patches: int = None) -> PBC3Result:
@@ -503,17 +437,23 @@ class PBC3:
         if isinstance(data, str):
             with open(data, "rb") as f:
                 data = f.read()
-        canvas, color_space, downsampled, original_w, original_h, w, h, has_alpha, channels, patch_count = cls._decode_to_canvas(data, max_patches=max_patches)
+        (
+            canvas, color_space, downsampled, original_w, original_h, w, h,
+            has_alpha, channels, patch_count,
+        ) = cls._decode_to_canvas(data, max_patches=max_patches)
         img = cls._canvas_to_image(canvas, color_space, has_alpha)
         if downsampled and img.size != (original_w, original_h):
             img = img.resize((original_w, original_h), cls.RESAMPLE_FILTER, reducing_gap=cls.RESAMPLE_REDUCING_GAP)
         cfg = PBC3Config(color_space=color_space)
-        return PBC3Result(img, data, cfg, None, time.perf_counter() - t0, len(data) * 8, original_w or w, original_h or h, w, h, channels=channels)
-
-    # ---- file helpers -----------------------------------------------------
+        return PBC3Result(
+            img, data, cfg, None, time.perf_counter() - t0, len(data) * 8,
+            original_w or w, original_h or h, w, h, channels=channels,
+        )
 
     @classmethod
-    def encode_file(cls, input_path: str, output_path: str, config: PBC3Config = None, **kwargs) -> PBC3Result:
+    def encode_file(
+        cls, input_path: str, output_path: str, config: PBC3Config = None, **kwargs
+    ) -> PBC3Result:
         """## Compresses a file and writes the .pbc3 output"""
         result = cls.compress(Image.open(input_path), config=config, **kwargs)
         with open(output_path, "wb") as f:
@@ -545,4 +485,4 @@ if __name__ == "__main__":
     else:
         preload_numba()
         res = PBC3.encode_file(sys.argv[1], sys.argv[2])
-        print(f"MSE: {res.mse:.2f} | Size: {len(res.data) / 1024:.2f} KB | Rate: {res.compression_rate:.2f}x | Time: {res.encode_seconds:.3f}s")
+        print(f"MSE: {res.mse:.2f} | Size: {len(res.data) / 1024:.2f} KB | Rate: {res.compression_rate:.2f}x | Time: {res.encode_seconds:.3f}s")

learned_filler.py

@@ -47,27 +47,25 @@ def _score_patch(target, canvas, box, channel_bits: int, patch, values) -> tuple
     return reduction, bits
 
 
-def build_action(target, canvas, box, config, channel_bits: int, idx: int):
-    """## Builds and scores one legacy single-index action"""
+def _build_action(target, canvas, box, config, channel_bits, cell_size, bitcount):
     c, x, y, bw, bh = box
-    cs, bc = ACTIONS[idx]
-    cell = max(1, min(cs, bw, bh))
+    cell = max(1, min(cell_size, bw, bh))
     residual = target[y:y + bh, x:x + bw, c] - canvas[y:y + bh, x:x + bw, c]
-    patch, values = ops.make_patch(c, x, y, bw, bh, cell, residual, config, bc)
+    patch, values = ops.make_patch(c, x, y, bw, bh, cell, residual, config, bitcount)
     reduction, bits = _score_patch(target, canvas, box, channel_bits, patch, values)
     return patch, values, reduction, bits
 
 
+def build_action(target, canvas, box, config, channel_bits: int, idx: int):
+    """## Builds and scores one legacy single-index action"""
+    cell_size, bitcount = ACTIONS[idx]
+    return _build_action(target, canvas, box, config, channel_bits, cell_size, bitcount)
+
+
 def build_action_factored(target, canvas, box, config, channel_bits: int, cs_idx: int, bc_idx: int, ms_idx: int):
     """## Builds and scores one factored action: cell size, bitcount, and mask size"""
-    c, x, y, bw, bh = box
-    cs, bc, ms = CELL_SIZES[cs_idx], BITCOUNTS[bc_idx], MASK_SIZES[ms_idx]
-    cell = max(1, min(cs, bw, bh))
-    cfg = dataclasses.replace(config, mask_size=ms)
-    residual = target[y:y + bh, x:x + bw, c] - canvas[y:y + bh, x:x + bw, c]
-    patch, values = ops.make_patch(c, x, y, bw, bh, cell, residual, cfg, bc)
-    reduction, bits = _score_patch(target, canvas, box, channel_bits, patch, values)
-    return patch, values, reduction, bits
+    cfg = dataclasses.replace(config, mask_size=MASK_SIZES[ms_idx])
+    return _build_action(target, canvas, box, cfg, channel_bits, CELL_SIZES[cs_idx], BITCOUNTS[bc_idx])
 
 
 def extract_state(target, canvas, box, step: int, image_w: int, image_h: int, patch_count: int, channels: int, lam: float, bits_spent: float, pixels: int) -> np.ndarray:
@@ -84,30 +82,8 @@ def extract_state(target, canvas, box, step: int, image_w: int, image_h: int, pa
 
 def propose_boxes(target, canvas, config, rng, channel: int, step: int) -> list[tuple[int, int, int, int, int]]:
     """## Returns candidate boxes sorted by the learned filler's prescore front-end"""
-    h_img, w_img, _ = target.shape
     search_q = ops.interp(config.search_q_start, config.search_q_end, step, config.patch_count)
-    visible = np.clip(canvas[:, :, channel], 0, 255).astype(np.int32)
-    abs_error = np.abs(target[:, :, channel] - visible).astype(np.int64)
-    integral_abs = ops.integral(abs_error)
-    anchors = ops.top_anchors(abs_error.astype(np.float32), config.top_k, config.anchor_block_size, channel)
-    if not anchors:
-        return []
-
-    specs, sums, areas = [], [], []
-    for i in range(max(1, int(config.search_depth))):
-        c, x, y, bw, bh, _, _ = ops.sample_box(rng, anchors[i % len(anchors)], w_img, h_img, config)
-        s = integral_abs[y + bh, x + bw] - integral_abs[y, x + bw] - integral_abs[y + bh, x] + integral_abs[y, x]
-        if s <= 0:
-            continue
-        sums.append(float(s))
-        areas.append(float(bw * bh))
-        specs.append((c, x, y, bw, bh))
-    if not specs:
-        return []
-
-    pre = search_q * ops.norm(sums) - (1.0 - search_q) * ops.norm(areas)
-    keep = sorted(ops.select_top_indices(pre, config.proposal_depth), key=lambda i: pre[i], reverse=True)
-    return [specs[i] for i in keep]
+    return ops.search_boxes(target, canvas, config, rng, channel, config.search_depth, search_q, ordered=True)
 
 
 def _load_npz_cached(path: str) -> dict:
@@ -126,6 +102,12 @@ def _silu(x):
     return x / (1.0 + np.exp(-x))
 
 
+def _mlp(x, weights):
+    """## Runs the shared two-layer SiLU trunk"""
+    W0, b0, W2, b2 = weights
+    return _silu(_silu(x @ W0.T + b0) @ W2.T + b2)
+
+
 class LearnedFiller:
     """## Torch-free numpy inference for the legacy single-action learned filler"""
 
@@ -154,10 +136,7 @@ class LearnedFiller:
 
     def _forward(self, x) -> tuple[np.ndarray, np.ndarray]:
         """## Runs the legacy action and value heads"""
-        h = x
-        if self.hidden > 0:
-            h = _silu(h @ self.W0.T + self.b0)
-            h = _silu(h @ self.W2.T + self.b2)
+        h = _mlp(x, (self.W0, self.b0, self.W2, self.b2)) if self.hidden > 0 else x
         return h @ self.Wa.T + self.ba, h @ self.Wv.T + self.bv[None, :]
 
     def _featurize(self, target, canvas, boxes, step: int, q: float, w_img: int, h_img: int, patch_count: int, channels: int) -> np.ndarray:
@@ -234,8 +213,7 @@ class FactoredFiller:
 
     def _forward(self, x) -> list[np.ndarray]:
         """## Runs the factored policy heads"""
-        h = _silu(x @ self.W0.T + self.b0)
-        h = _silu(h @ self.W2.T + self.b2)
+        h = _mlp(x, (self.W0, self.b0, self.W2, self.b2))
         return [h @ W.T + b for W, b in self.heads]
 
     def select_patch(self, target, canvas, config, rng, channel_bits: int, step: int, current_channel: int):
@@ -263,4 +241,4 @@ class FactoredFiller:
         if reduction <= 0:
             return None, None
         self._bits += bits
-        return patch, values
+        return patch, values

pbc3_features.py

@@ -132,4 +132,4 @@ def extract_cheap(names, target, canvas, box, step: int, q: float, image_w: int,
         d["before_mse"] = float((before ** 2).mean())
     if "before_mae" in names:
         d["before_mae"] = float(np.abs(before).mean())
-    return np.array([d[nm] for nm in names], dtype=np.float32)
+    return np.array([d[nm] for nm in names], dtype=np.float32)

pbc3_heads.py

@@ -76,31 +76,7 @@ class SearchHead:
 
     def search(self, target, canvas, config, rng, channel: int, depth: int, search_q: float) -> list[tuple[int, int, int, int, int]]:
         """## Samples boxes around strong error anchors and keeps the best prescores"""
-        visible_canvas_channel = np.clip(canvas[:, :, channel], 0, 255).astype(np.int32)
-        visible_error = (target[:, :, channel] - visible_canvas_channel).astype(np.int64)
-        abs_error = np.abs(visible_error)
-        integral_abs = ops.integral(abs_error)
-
-        anchors = ops.top_anchors(abs_error.astype(np.float32), config.top_k, config.anchor_block_size, channel)
-        if not anchors:
-            return []
-
-        h, w, _ = target.shape
-        box_sums, box_areas, box_specs = [], [], []
-        for i in range(max(1, int(depth))):
-            c, x, y, bw, bh, ax, ay = ops.sample_box(rng, anchors[i % len(anchors)], w, h, config)
-            box_sum = integral_abs[y + bh, x + bw] - integral_abs[y, x + bw] - integral_abs[y + bh, x] + integral_abs[y, x]
-            if box_sum <= 0:
-                continue
-            box_sums.append(float(box_sum))
-            box_areas.append(float(bw * bh))
-            box_specs.append((c, x, y, bw, bh))
-        if not box_specs:
-            return []
-
-        pre_scores = search_q * ops.norm(box_sums) - (1.0 - search_q) * ops.norm(box_areas)
-        keep = ops.select_top_indices(pre_scores, config.proposal_depth)
-        return [box_specs[i] for i in keep]
+        return ops.search_boxes(target, canvas, config, rng, channel, depth, search_q)
 
 
 class FillerHead:

pbc3_ops.py

@@ -10,8 +10,6 @@ RESAMPLE_REDUCING_GAP = None
 PALETTE_GENERATED = 0
 
 
-# ===== DETERMINISTIC RNG + LOCAL BICUBIC RESAMPLER =====
-
 _UINT64_MASK = (1 << 64) - 1
 _UINT32_SCALE = 1.0 / 4294967296.0
 _RESAMPLE_COEFF_CACHE = {}
@@ -105,31 +103,6 @@ def bicubic_resample_2d(values, out_h: int, out_w: int) -> np.ndarray:
     return out
 
 
-def resize_array_bicubic(values, out_h: int, out_w: int):
-    arr = np.asarray(values)
-    if arr.ndim == 2:
-        out = bicubic_resample_2d(arr, out_h, out_w)
-    elif arr.ndim == 3:
-        chans = [bicubic_resample_2d(arr[:, :, c], out_h, out_w) for c in range(arr.shape[2])]
-        out = np.stack(chans, axis=2)
-    else:
-        raise ValueError(f"expected 2D or 3D array, got shape {arr.shape}")
-
-    if np.issubdtype(arr.dtype, np.integer):
-        info = np.iinfo(arr.dtype)
-        return np.clip(np.rint(out), info.min, info.max).astype(arr.dtype)
-    return out.astype(arr.dtype, copy=False)
-
-
-def resize_image_bicubic(img: Image.Image, out_h: int, out_w: int) -> Image.Image:
-    if img.size == (int(out_w), int(out_h)):
-        return img.copy()
-    arr = resize_array_bicubic(np.asarray(img), out_h, out_w)
-    return Image.fromarray(arr, img.mode)
-
-
-# ===== CORE MATH =====
-
 def ceil_div(a: int, b: int) -> int:
     """## Returns ceil(a / b) for integer cell/grid math"""
     return (int(a) + int(b) - 1) // int(b)
@@ -157,16 +130,6 @@ def integral(a) -> np.ndarray:
     return np.pad(a.astype(np.int64).cumsum(0).cumsum(1), ((1, 0), (1, 0)))
 
 
-def cell_edges(start: int, length: int, cell_size: int) -> np.ndarray:
-    """## Returns the edge coordinates of patch cells, clamped to the patch end"""
-    n = ceil_div(length, cell_size)
-    edges = start + np.arange(n + 1) * cell_size
-    edges[n] = start + length
-    return edges
-
-
-# ===== NUMBA KERNELS =====
-
 @njit(cache=True)
 def _box_cell_bound_kernel(integral_arr, x, y, bw, bh, cell_size):
     nx = (bw + cell_size - 1) // cell_size
@@ -345,8 +308,6 @@ def _quantize_signed_kernel(vals, pal):
         out[i] = best_i
     return out
 
-# ===== PALETTE OPS =====
-
 
 def palette_bounds(values) -> tuple[int, int]:
     """## Returns negative max (min or 0) and positive max (true max)
@@ -364,14 +325,10 @@ def range_counts(mask_size, negative_max=255, positive_max=255, positive_bias=Tr
     negative_max = max(0, int(negative_max))
     positive_max = max(0, int(positive_max))
     if side_bits == 0 or (negative_max == 0 and positive_max == 0):
-        # The reason it's 0,0 if mask size is 1 is that the only range that can be represented by it is the 0 bin
-        # This actually just makes mask size 1 useless for palette generation.
         return 0, 0
     if negative_max == 0:
-        # Obviously if there are no negative values, all bins should be positive
         return min(side_bits, positive_max), 0
     if positive_max == 0:
-        # Obviously same
         return 0, min(side_bits, negative_max)
     raw_pos = side_bits * positive_max / (positive_max + negative_max)
     pos_count = math.ceil(raw_pos) if positive_bias else math.floor(raw_pos)
@@ -488,8 +445,6 @@ def quantize_signed(values, palette):
 
 
 
-# ===== IMAGE / GRID OPS =====
-
 def signed_resample(values, out_h: int, out_w: int) -> np.ndarray:
     """## Resizes a signed grid to a patch-sized int16 delta image"""
     return np.rint(bicubic_resample_2d(values, int(out_h), int(out_w))).astype(np.int16)
@@ -520,8 +475,6 @@ def final_mse(reference_image, reconstructed_image) -> float:
     return image_mse(np.asarray(reference_image, dtype=np.float32), np.asarray(reconstructed_image, dtype=np.float32))
 
 
-# ===== SEARCH OPS =====
-
 def box_cell_bound(integral_arr, x: int, y: int, bw: int, bh: int, cell_size: int) -> float:
     """## Returns the best-case signed-error energy for a box/cell-size pair"""
     return float(_box_cell_bound_kernel(
@@ -589,6 +542,35 @@ def sample_box(rng, anchor, image_w: int, image_h: int, config) -> tuple[int, in
     return c, x, y, w, h, ax, ay
 
 
+def search_boxes(target, canvas, config, rng, channel: int, depth: int, search_q: float, ordered=False):
+    """Samples error-centered boxes and keeps the best prescores."""
+    h_img, w_img, _ = target.shape
+    visible = np.clip(canvas[:, :, channel], 0, 255).astype(np.int32)
+    abs_error = np.abs(target[:, :, channel] - visible).astype(np.int64)
+    integral_abs = integral(abs_error)
+    anchors = top_anchors(abs_error.astype(np.float32), config.top_k, config.anchor_block_size, channel)
+    if not anchors:
+        return []
+
+    specs, sums, areas = [], [], []
+    for i in range(max(1, int(depth))):
+        c, x, y, bw, bh, _, _ = sample_box(rng, anchors[i % len(anchors)], w_img, h_img, config)
+        total = integral_abs[y + bh, x + bw] - integral_abs[y, x + bw] - integral_abs[y + bh, x] + integral_abs[y, x]
+        if total <= 0:
+            continue
+        sums.append(float(total))
+        areas.append(float(bw * bh))
+        specs.append((c, x, y, bw, bh))
+    if not specs:
+        return []
+
+    scores = search_q * norm(sums) - (1.0 - search_q) * norm(areas)
+    keep = select_top_indices(scores, config.proposal_depth)
+    if ordered:
+        keep = sorted(keep, key=lambda i: scores[i], reverse=True)
+    return [specs[i] for i in keep]
+
+
 def base_cell_size(residual_patch, config) -> int:
     """## Returns the rough cell size suggested by local residual frequency"""
     if residual_patch.size == 0:
@@ -596,8 +578,6 @@ def base_cell_size(residual_patch, config) -> int:
     return int(_base_cell_size_kernel(np.ascontiguousarray(residual_patch, dtype=np.float64), int(config.max_cell_size)))
 
 
-# ===== PATCH OPS =====
-
 def patch_bits_for(patch, channel_bits: int) -> int:
     """## Returns the serialized bit cost of a complete patch"""
     grid_bits = ceil_div(patch["w"], patch["cell_size"]) * ceil_div(patch["h"], patch["cell_size"]) * patch["bitcount"]
@@ -639,4 +619,4 @@ def make_patch(channel: int, x: int, y: int, w: int, h: int, cell_size: int, res
 
 def debug_line(kind: str, **items) -> str:
     """## Returns one plain-text debug line with key=value fields"""
-    return kind + " " + " ".join(f"{k}={v}" for k, v in items.items())
+    return kind + " " + " ".join(f"{k}={v}" for k, v in items.items())

pbc3_stream.py

@@ -0,0 +1,109 @@
+"""PBC3 bitstream framing, separate from the image algorithm."""
+
+import numpy as np
+
+import pbc3_ops as ops
+from pbc3_types import BitReader, BitWriter
+
+
+def write_grid(bw: BitWriter, values, bitcount: int) -> None:
+    """Writes a flat grid of palette indices."""
+    for value in values:
+        bw.write(int(value), bitcount)
+
+
+def read_grid(br: BitReader, count: int, bitcount: int) -> np.ndarray:
+    """Reads a flat grid of palette indices."""
+    values = np.zeros(count, dtype=np.uint16)
+    for i in range(count):
+        values[i] = br.read(bitcount)
+    return values
+
+
+def write_header(bw, w, h, original_w, original_h, downsampled, color_id, channels,
+                 channel_bits, positive_bias, has_alpha, patch_count, base_values, warmup=None):
+    """Writes the image-level stream header."""
+    bw.write(int(downsampled), 1)
+    if downsampled:
+        bw.write(original_w, 16)
+        bw.write(original_h, 16)
+    bw.write(w, 16)
+    bw.write(h, 16)
+    bw.write(color_id, 2)
+    bw.write(channels, 8)
+    bw.write(channel_bits, 4)
+    bw.write(int(positive_bias), 1)
+    bw.write(int(has_alpha), 1)
+    bw.write(patch_count, 32)
+    for base in base_values:
+        bw.write(base, 8)
+    bw.write(int(warmup is not None), 1)
+    if warmup is not None:
+        warm_w, warm_h, warm_split = warmup
+        bw.write(warm_w, 16)
+        bw.write(warm_h, 16)
+        bw.write(warm_split, 32)
+
+
+def read_header(br, color_space_names):
+    """Reads the image-level stream header."""
+    downsampled = bool(br.read(1))
+    original_w = br.read(16) if downsampled else None
+    original_h = br.read(16) if downsampled else None
+    w = br.read(16)
+    h = br.read(16)
+    color_id = br.read(2)
+    channels = br.read(8)
+    channel_bits = br.read(4)
+    positive_bias = bool(br.read(1))
+    has_alpha = bool(br.read(1))
+    patch_count = br.read(32)
+    base_values = [br.read(8) for _ in range(channels)]
+    warmup_on = bool(br.read(1))
+    warm_w = warm_h = warmup_split = None
+    if warmup_on:
+        warm_w = br.read(16)
+        warm_h = br.read(16)
+        warmup_split = br.read(32)
+    return (
+        downsampled, original_w, original_h, w, h, color_space_names[color_id], channels,
+        channel_bits, positive_bias, has_alpha, patch_count, base_values, warmup_on, warm_w,
+        warm_h, warmup_split,
+    )
+
+
+def write_patch(bw, patch, channel_bits: int) -> None:
+    """Writes one generated-palette patch."""
+    bw.write(patch["channel"], channel_bits)
+    bw.write(patch["x"], 16)
+    bw.write(patch["y"], 16)
+    bw.write(patch["w"], 16)
+    bw.write(patch["h"], 16)
+    bw.write(0, 1)
+    mask = patch["mask"]
+    bw.write(len(mask), 10)
+    for bit in mask:
+        bw.write(bit, 1)
+    bw.write(patch["neg"], 8)
+    bw.write(patch["pos"], 8)
+    bw.write(patch["max_bitcount"], 4)
+    bw.write(patch["cell_size"], 16)
+    write_grid(bw, patch["indices"].ravel().astype(np.int64), patch["bitcount"])
+
+
+def read_patch(br, channel_bits: int, positive_bias: bool = True):
+    """Reads one generated-palette patch and returns its decoded values."""
+    channel = br.read(channel_bits)
+    x, y = br.read(16), br.read(16)
+    w, h = br.read(16), br.read(16)
+    if br.read(1) != 0:
+        raise ValueError("explicit palette patches were removed in PBC3 3.0 release cleanup")
+    mask = [br.read(1) for _ in range(br.read(10))]
+    negative_max, positive_max = br.read(8), br.read(8)
+    max_bitcount = br.read(4)
+    bitcount = ops.resolve_palette_bitcount(mask, max_bitcount, negative_max, positive_max, positive_bias)
+    palette = ops.palette_generator(mask, max_bitcount, negative_max, positive_max, positive_bias)
+    cell_size = br.read(16)
+    gw, gh = ops.ceil_div(w, cell_size), ops.ceil_div(h, cell_size)
+    indices = read_grid(br, gh * gw, bitcount).reshape(gh, gw)
+    return channel, x, y, w, h, cell_size, palette[indices], bitcount

pbc3_types.py

@@ -283,4 +283,4 @@ class PBC3Result:
             bbox=dict(boxstyle="round,pad=0.5", facecolor="black", alpha=0.72, edgecolor="none"),
         )
         image_ax.imshow(self.image)
-        plt.show()
+        plt.show()