lesiondiffusion-mlx 0.2.1__py3-none-any.whl

lesiondiffusion_mlx/mlx_model.py

@@ -0,0 +1,456 @@
+"""Native Metal networks for LesionDiffusion source 880736e2.
+
+Weights retain their source names/layouts. Shared MedMLX array primitives own
+the NCDHW to MLX convolution boundary; the architecture here is the upstream
+3-D LDM architecture, including its distinct normalization and attention rules.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, cast
+
+if TYPE_CHECKING:
+    from lesiondiffusion_mlx._typing import Array, Mlx
+
+from medmlx_core.ops import (
+    conv3d_ncdhw,
+    group_norm_ncdhw,
+    linear,
+    silu,
+)
+
+from lesiondiffusion_mlx._source_half import conv3d_half, group_norm_half
+from lesiondiffusion_mlx._spatial import upsample_nearest2
+from lesiondiffusion_mlx._typing import ArrayConstructor
+
+
+@dataclass(frozen=True)
+class UnetConfig:
+    """Architecture of one upstream 3-D UNet (LMNet or LINet).
+
+    Attributes:
+        model_channels: Base channel width; also sets the sinusoidal time-embedding size.
+        num_res_blocks: Residual blocks per resolution level.
+        attention_resolutions: Downsampling factors (1, 2, 4, ...) at which residual
+            blocks are followed by a spatial transformer. The middle block always has one.
+        channel_mult: Per-level multipliers of ``model_channels``; each level after
+            the first halves the spatial size.
+        num_head_channels: Channels per attention head.
+        skip_context: Drop the text context, so cross-attention layers attend to
+            their own input (upstream's ``skip_context``).
+    """
+
+    model_channels: int
+    num_res_blocks: int
+    attention_resolutions: tuple[int, ...]
+    channel_mult: tuple[int, ...] = (1, 2, 4)
+    num_head_channels: int = 32
+    skip_context: bool = False
+
+
+LMNET_CONFIG = UnetConfig(16, 1, (4,))
+# skip_context=True follows the pinned inference YAML; the runner overrides it per
+# inference mode (context is enabled only in "conditioned" mode).
+LINET_CONFIG = UnetConfig(128, 2, (4, 8), skip_context=True)
+
+
+class SourceLayers:
+    """Source-keyed native MLX layers without a host execution path.
+
+    ``weights`` maps upstream state-dict names (for example ``input_blocks.1.0.in_layers.2``)
+    to MLX arrays in upstream's NCDHW layout. With ``source_arithmetic=True``, float16
+    convolutions use the Torch-CPU-ordered Metal kernel (``conv3d_half``); otherwise
+    they use the shared MedMLX convolution. Float16 group normalization always uses
+    the source-ordered kernel.
+    """
+
+    def __init__(self, weights: Mapping[str, Array], mx: Mlx, *, source_arithmetic: bool = True):
+        self.weights = weights
+        self.mx = mx
+        self.source_arithmetic = source_arithmetic
+
+    def conv(self, x: Array, key: str, *, padding: int = 1, stride: int = 1) -> Array:
+        weight, bias = self.weights[key + ".weight"], self.weights.get(key + ".bias")
+        if x.dtype == self.mx.float16 and self.source_arithmetic:
+            return conv3d_half(x, weight, bias, padding=padding, stride=stride, mx=self.mx)
+        return conv3d_ncdhw(
+            x,
+            weight,
+            bias,
+            padding=padding,
+            stride=stride,
+            mx=self.mx,
+        )
+
+    def dense(self, x: Array, key: str) -> Array:
+        dtype = x.dtype
+        x, weight, bias = self._affine_inputs(x, key)
+        return linear(x, weight, bias, self.mx).astype(dtype)
+
+    def _affine_inputs(self, x: Array, key: str) -> tuple[Array, Array, Array | None]:
+        weight = self.weights[key + ".weight"]
+        bias = self.weights.get(key + ".bias")
+        if x.dtype == self.mx.float16:
+            # Torch's CPU half kernels add bias to the FP32 accumulator before
+            # rounding once. A half convolution/matmul followed by bias rounds twice.
+            x, weight = x.astype(self.mx.float32), weight.astype(self.mx.float32)
+            if bias is not None:
+                bias = bias.astype(self.mx.float32)
+        return x, weight, bias
+
+    def norm(self, x: Array, key: str, *, groups: int, eps: float) -> Array:
+        if x.dtype == self.mx.float16:
+            return group_norm_half(
+                x,
+                self.weights[key + ".weight"],
+                self.weights[key + ".bias"],
+                groups=groups,
+                eps=eps,
+                mx=self.mx,
+            )
+        # Torch GroupNorm accumulates statistics and applies its affine in
+        # float32 before rounding the result back to the input precision.
+        return group_norm_ncdhw(
+            x.astype(self.mx.float32),
+            self.weights[key + ".weight"].astype(self.mx.float32),
+            self.weights[key + ".bias"].astype(self.mx.float32),
+            num_groups=groups,
+            eps=eps,
+            mx=self.mx,
+        ).astype(x.dtype)
+
+    def layer_norm(self, x: Array, key: str, *, eps: float = 1e-5) -> Array:
+        return self.mx.fast.layer_norm(
+            x,
+            self.weights[key + ".weight"],
+            self.weights[key + ".bias"],
+            eps,
+        )
+
+    def gelu(self, x: Array) -> Array:
+        # Source nn.GELU uses the exact erf formula, not tanh approximation.
+        dtype = x.dtype
+        x = x.astype(self.mx.float32)
+        return (x * (1 + self.mx.erf(x / math.sqrt(2))) * 0.5).astype(dtype)
+
+    def resize_nearest(self, x: Array, shape: tuple[int, ...]) -> Array:
+        """Nearest-neighbour resize of ``(B, C, D, H, W)`` to spatial ``shape`` (D, H, W).
+
+        Output index ``i`` reads input index ``floor(i * in_size / out_size)``.
+        """
+        if len(shape) != 3:
+            raise ValueError("LesionDiffusion nearest resizing needs three spatial axes")
+        if x.ndim == 5 and min(x.shape) > 0 and shape == tuple(2 * s for s in x.shape[2:]):
+            # Both source decoders double every spatial axis. A single strided
+            # copy preserves their floor-index rule without three intermediates.
+            return upsample_nearest2(x, mx=self.mx)
+        for axis, target in enumerate(shape, start=2):
+            indices = self.mx.floor(self.mx.arange(target) * (x.shape[axis] / target))
+            x = self.mx.take(x, indices.astype(self.mx.int32), axis=axis)
+        return x
+
+
+class LdmUnet(SourceLayers):
+    """LMNet/LINet: upstream's 3-D OpenAI-style UNet with spatial transformers.
+
+    ``fused_attention=True`` replaces explicit softmax attention with Metal's fused
+    kernel (used for LMNet in "conditioned" mode); source mode keeps the explicit
+    form to stay on upstream's arithmetic.
+    """
+
+    def __init__(
+        self,
+        weights: Mapping[str, Array],
+        mx: Mlx,
+        config: UnetConfig,
+        *,
+        source_arithmetic: bool = True,
+        fused_attention: bool = False,
+    ):
+        super().__init__(weights, mx, source_arithmetic=source_arithmetic)
+        self.config = config
+        self.fused_attention = fused_attention
+
+    def silu(self, x: Array) -> Array:
+        # Source nn.SiLU rounds once after the activation. VQ's separate
+        # sigmoid/product deliberately retains its two half rounding steps.
+        return silu(x.astype(self.mx.float32), self.mx).astype(x.dtype)
+
+    def softmax(self, x: Array) -> Array:
+        # Torch softmax accumulates in float32, then returns the input dtype.
+        return self.mx.softmax(x, axis=-1, precise=True)
+
+    def residual(self, x: Array, emb: Array, key: str) -> Array:
+        channels = x.shape[1]
+        groups = 32 if channels % 32 == 0 else channels
+        h = self.conv(
+            self.silu(self.norm(x, key + ".in_layers.0", groups=groups, eps=1e-5)),
+            key + ".in_layers.2",
+        )
+        h = h + self.dense(self.silu(emb), key + ".emb_layers.1")[:, :, None, None, None]
+        groups = 32 if h.shape[1] % 32 == 0 else h.shape[1]
+        h = self.conv(
+            self.silu(self.norm(h, key + ".out_layers.0", groups=groups, eps=1e-5)),
+            key + ".out_layers.3",
+        )
+        if key + ".skip_connection.weight" in self.weights:
+            x = self.conv(x, key + ".skip_connection", padding=0)
+        return x + h
+
+    def attention(self, x: Array, context: Array | None, key: str, heads: int) -> Array:
+        mx = self.mx
+        context = x if context is None else context
+        q, k, v = (
+            self.dense(x, key + ".to_q"),
+            self.dense(context, key + ".to_k"),
+            self.dense(context, key + ".to_v"),
+        )
+        depth = q.shape[-1] // heads
+
+        def split(tensor: Array) -> Array:
+            return mx.transpose(
+                mx.reshape(tensor, (tensor.shape[0], tensor.shape[1], heads, depth)), (0, 2, 1, 3)
+            )
+
+        if self.fused_attention:
+            q, k, v = split(q), split(k), split(v)
+            # Metal's fused full-attention kernel starts at 64-wide heads.
+            # Zero padding retains the 32-wide model's original dot products
+            # and scale; discard the padded value channels after attention.
+            if depth < 64:
+                padding = [(0, 0), (0, 0), (0, 0), (0, 64 - depth)]
+                q, k, v = (mx.pad(tensor, padding) for tensor in (q, k, v))
+            result = mx.fast.scaled_dot_product_attention(
+                q, k, v, scale=depth**-0.5, force_fused=True
+            )[..., :depth]
+        else:
+            # Source half attention rounds similarity and probability tensors;
+            # LINet retains those boundaries rather than using fused SDPA.
+            similarity = (split(q) @ mx.swapaxes(split(k), -1, -2)) * depth**-0.5
+            result = self.softmax(similarity) @ split(v)
+        result = mx.reshape(mx.transpose(result, (0, 2, 1, 3)), x.shape)
+        return self.dense(result, key + ".to_out.0")
+
+    def transformer(self, x: Array, context: Array | None, key: str) -> Array:
+        residual = x
+        x = self.conv(self.norm(x, key + ".norm", groups=32, eps=1e-6), key + ".proj_in", padding=0)
+        shape = x.shape
+        x = self.mx.reshape(self.mx.transpose(x, (0, 2, 3, 4, 1)), (shape[0], -1, shape[1]))
+        heads = shape[1] // self.config.num_head_channels
+        block = key + ".transformer_blocks.0"
+        x = x + self.attention(self.layer_norm(x, block + ".norm1"), None, block + ".attn1", heads)
+        x = x + self.attention(
+            self.layer_norm(x, block + ".norm2"), context, block + ".attn2", heads
+        )
+        ff = self.gelu(self.dense(self.layer_norm(x, block + ".norm3"), block + ".ff.net.0.0"))
+        x = x + self.dense(ff, block + ".ff.net.2")
+        x = self.mx.transpose(self.mx.reshape(x, (shape[0], *shape[2:], shape[1])), (0, 4, 1, 2, 3))
+        return residual + self.conv(x, key + ".proj_out", padding=0)
+
+    def __call__(self, x: Array, timesteps: Array, context: Array | None) -> Array:
+        """Predict noise for one denoising step.
+
+        Args:
+            x: Network input ``(B, C_in, D, H, W)``, channels first. Spatial sizes
+                must be divisible by ``2 ** (levels - 1)`` (4 for the shipped configs).
+            timesteps: Integer diffusion timesteps, shape ``(B,)``.
+            context: Text embeddings ``(B, N, D_ctx)``, or ``None``. Ignored when the
+                config sets ``skip_context``.
+
+        Returns:
+            Prediction ``(B, C_out, D, H, W)``, computed in the weights' dtype
+            (float32 for LMNet, float16 for LINet in the runner).
+        """
+        mx, cfg = self.mx, self.config
+        half = cfg.model_channels // 2
+        frequency = mx.exp(-math.log(10000) * mx.arange(half, dtype=mx.float32) / half)
+        phase = timesteps[:, None].astype(mx.float32) * frequency[None]
+        emb = mx.concatenate([mx.cos(phase), mx.sin(phase)], axis=1)
+        dtype = self.weights["time_embed.0.weight"].dtype
+        emb = self.dense(self.silu(self.dense(emb.astype(dtype), "time_embed.0")), "time_embed.2")
+        if context is not None:
+            context = (
+                None
+                if cfg.skip_context
+                else self.dense(context.astype(dtype), "context_linear_projection")
+            )
+        x = self.conv(x.astype(dtype), "input_blocks.0.0")
+        skips = [x]
+        block_index, resolution = 1, 1
+        for level in range(len(cfg.channel_mult)):
+            for _ in range(cfg.num_res_blocks):
+                key = f"input_blocks.{block_index}"
+                x = self.residual(x, emb, key + ".0")
+                if resolution in cfg.attention_resolutions:
+                    x = self.transformer(x, context, key + ".1")
+                skips.append(x)
+                block_index += 1
+            if level < len(cfg.channel_mult) - 1:
+                x = self.conv(x, f"input_blocks.{block_index}.0.op", stride=2)
+                skips.append(x)
+                block_index += 1
+                resolution *= 2
+        x = self.residual(x, emb, "middle_block.0")
+        x = self.transformer(x, context, "middle_block.1")
+        x = self.residual(x, emb, "middle_block.2")
+        block_index = 0
+        for level in reversed(range(len(cfg.channel_mult))):
+            for block in range(cfg.num_res_blocks + 1):
+                key = f"output_blocks.{block_index}"
+                x = self.residual(mx.concatenate([x, skips.pop()], axis=1), emb, key + ".0")
+                has_attention = resolution in cfg.attention_resolutions
+                if has_attention:
+                    x = self.transformer(x, context, key + ".1")
+                if level and block == cfg.num_res_blocks:
+                    x = self.resize_nearest(x, tuple(int(s * 2) for s in x.shape[2:]))
+                    x = self.conv(x, key + f".{2 if has_attention else 1}.conv")
+                    resolution //= 2
+                block_index += 1
+        return self.conv(
+            self.silu(
+                self.norm(x, "out.0", groups=32 if x.shape[1] % 32 == 0 else x.shape[1], eps=1e-5)
+            ),
+            "out.2",
+        )
+
+
+class VqAutoencoder(SourceLayers):
+    """Upstream's 3-D VQ autoencoder, as used by LINet: factor-two spatial latent.
+
+    Tensors are channels first. ``encode`` maps ``(B, C, D, H, W)`` to a latent of
+    half the spatial size (rounded down); ``decode`` snaps latent vectors to their
+    nearest codebook entry and returns twice the latent's spatial size.
+    """
+
+    def residual(self, x: Array, key: str) -> Array:
+        h = self.conv(
+            silu(self.norm(x, key + ".norm1", groups=16, eps=1e-6), self.mx), key + ".conv1"
+        )
+        h = self.conv(
+            silu(self.norm(h, key + ".norm2", groups=16, eps=1e-6), self.mx), key + ".conv2"
+        )
+        if key + ".nin_shortcut.weight" in self.weights:
+            x = self.conv(x, key + ".nin_shortcut", padding=0)
+        return x + h
+
+    def encode(self, x: Array) -> Array:
+        x = self.conv(x, "encoder.conv_in")
+        for level in range(3):
+            x = self.residual(x, f"encoder.down.{level}.block.0")
+            if level == 0:
+                x = self.mx.pad(x, [(0, 0), (0, 0), (0, 1), (0, 1), (0, 1)])
+                x = self.conv(x, "encoder.down.0.downsample.conv", padding=0, stride=2)
+        x = self.residual(x, "encoder.mid.block_1")
+        x = self.residual(x, "encoder.mid.block_2")
+        x = self.conv(
+            silu(self.norm(x, "encoder.norm_out", groups=16, eps=1e-6), self.mx), "encoder.conv_out"
+        )
+        return self.conv(x, "quant_conv", padding=0)
+
+    def quantize(self, z: Array, *, chunk_size: int = 4096) -> Array:
+        """Replace each latent vector of ``z`` ``(B, C, D, H, W)`` by its nearest codebook entry.
+
+        Nearest means smallest squared Euclidean distance; positions are processed in
+        chunks of ``chunk_size`` to bound the size of the distance matrix. Selection
+        is discontinuous, so tiny rounding differences can change a chosen code.
+        """
+        mx = self.mx
+        shape = z.shape
+        flat = mx.reshape(mx.transpose(z, (0, 2, 3, 4, 1)), (-1, shape[1]))
+        codebook = self.weights["quantize.embedding.weight"]
+        # Source half multiplication rounds each square, then torch.sum
+        # accumulates in float32 and rounds once back to half.
+        codebook_norm = mx.sum((codebook * codebook).astype(mx.float32), axis=1).astype(
+            codebook.dtype
+        )
+        pieces: list[Array] = []
+        for start in range(0, flat.shape[0], chunk_size):
+            part = flat[start : start + chunk_size]
+            distance = (
+                mx.sum((part * part).astype(mx.float32), axis=1, keepdims=True).astype(part.dtype)
+                + codebook_norm[None]
+                - 2 * (part @ codebook.T)
+            )
+            selected = codebook[mx.argmin(distance, axis=1)]
+            # Upstream VectorQuantizer retains its straight-through arithmetic
+            # during inference too; fp16 cancellation changes the returned values.
+            pieces.append(part + (selected - part))
+        return mx.transpose(
+            mx.reshape(mx.concatenate(list(pieces)), (shape[0], *shape[2:], shape[1])),
+            (0, 4, 1, 2, 3),
+        )
+
+    def decode(self, z: Array, *, quantize: bool = True) -> Array:
+        """Decode a latent ``(B, C, d, h, w)`` to a volume ``(B, C_out, 2d, 2h, 2w)``.
+
+        Set ``quantize=False`` to skip the codebook lookup and decode ``z`` as given.
+        """
+        x = self.conv(self.quantize(z) if quantize else z, "post_quant_conv", padding=0)
+        x = self.conv(x, "decoder.conv_in")
+        x = self.residual(x, "decoder.mid.block_1")
+        x = self.residual(x, "decoder.mid.block_2")
+        for level in reversed(range(3)):
+            for block in range(2):
+                x = self.residual(x, f"decoder.up.{level}.block.{block}")
+            if level == 1:
+                x = self.resize_nearest(x, tuple(int(s * 2) for s in x.shape[2:]))
+                x = self.conv(x, "decoder.up.1.upsample.conv")
+        return self.conv(
+            silu(self.norm(x, "decoder.norm_out", groups=16, eps=1e-6), self.mx), "decoder.conv_out"
+        )
+
+
+class BiomedTextEncoder(SourceLayers):
+    """PubMedBERT CLS pooling and OpenCLIP MLP projection on Metal."""
+
+    def __call__(self, token_ids: Array) -> Array:
+        """Embed token id sequences ``(N, L)`` as vectors ``(N, D_ctx)``.
+
+        Padding positions (id 0) are masked out of attention; the vector is the
+        projected first-token (CLS) state. ``L`` must not exceed the position-embedding
+        table (256 for BiomedCLIP).
+        """
+        mx = self.mx
+        prefix = "transformer.embeddings"
+        x = self.weights[prefix + ".word_embeddings.weight"][token_ids]
+        x = (
+            x
+            + self.weights[prefix + ".position_embeddings.weight"][mx.arange(token_ids.shape[1])][
+                None
+            ]
+        )
+        x = x + self.weights[prefix + ".token_type_embeddings.weight"][0]
+        x = self.layer_norm(x, prefix + ".LayerNorm", eps=1e-12)
+        mask = mx.where(
+            token_ids[:, None, None, :] != 0,
+            cast(ArrayConstructor, mx.array)(0, dtype=x.dtype),
+            cast(ArrayConstructor, mx.array)(-float("inf"), dtype=x.dtype),
+        )
+        for index in range(12):
+            key = f"transformer.encoder.layer.{index}"
+            q, k, v = (
+                self.dense(x, key + ".attention.self." + name) for name in ("query", "key", "value")
+            )
+
+            def split(tensor: Array) -> Array:
+                return mx.transpose(
+                    mx.reshape(tensor, (tensor.shape[0], tensor.shape[1], 12, 64)), (0, 2, 1, 3)
+                )
+
+            h = mx.fast.scaled_dot_product_attention(
+                split(q), split(k), split(v), scale=0.125, mask=mask
+            )
+            h = mx.reshape(mx.transpose(h, (0, 2, 1, 3)), x.shape)
+            x = self.layer_norm(
+                x + self.dense(h, key + ".attention.output.dense"),
+                key + ".attention.output.LayerNorm",
+                eps=1e-12,
+            )
+            h = self.gelu(self.dense(x, key + ".intermediate.dense"))
+            x = self.layer_norm(
+                x + self.dense(h, key + ".output.dense"), key + ".output.LayerNorm", eps=1e-12
+            )
+        return self.dense(self.gelu(self.dense(x[:, 0], "proj.0")), "proj.2")

lesiondiffusion_mlx/mlx_pipeline.py

@@ -0,0 +1,136 @@
+"""Array preprocessing for the pinned upstream LesionDiffusion pipeline."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from pathlib import Path
+from typing import TYPE_CHECKING, cast
+
+if TYPE_CHECKING:
+    from lesiondiffusion_mlx._typing import Array, ArrayValue, HostArray, Mlx, TokenizerContract
+
+import numpy as np
+from medmlx_core.errors import InvalidInputError, ModelExecutionError
+
+from lesiondiffusion_mlx._typing import ArrayConstructor
+from lesiondiffusion_mlx.mlx_model import BiomedTextEncoder, VqAutoencoder
+
+# Crop extents in source axis order (Z, Y, X), in voxels of the 1 mm grid: LMNet
+# generates the mask on LMNET_CROP; LINet inpaints on LINET_CROP.
+LMNET_CROP = (64, 152, 152)
+LINET_CROP = (128, 128, 128)
+
+
+def ras_to_source(array: HostArray) -> HostArray:
+    """Convert a RAS-ordered ``(X, Y, Z)`` array to the source layout ``(Z, Y, X)``.
+
+    The pinned upstream loader reorients to LAS before transposing, so the first
+    (X) axis is also reversed. Expects a 3-D array and returns a C-contiguous copy.
+    """
+    return np.ascontiguousarray(array[::-1].transpose(2, 1, 0))
+
+
+def source_to_ras(array: HostArray) -> HostArray:
+    """Invert ``ras_to_source``: ``(Z, Y, X)`` LAS-ordered data back to RAS ``(X, Y, Z)``."""
+    return np.ascontiguousarray(array.transpose(2, 1, 0)[::-1])
+
+
+def source_crop(mask: HostArray, shape: tuple[int, int, int]) -> tuple[slice, slice, slice]:
+    """Return the fixed-size crop centred on a mask's bounding box, in source axes.
+
+    Args:
+        mask: Boolean or label volume in source ``(Z, Y, X)`` order; non-zero voxels
+            define the extent (the box for LMNet, the lesion for LINet).
+        shape: Crop size ``(Z, Y, X)`` in voxels.
+
+    Returns:
+        Slices of length ``shape`` per axis, centred on the midpoint of the mask's
+        extent (integer floor) and shifted to stay inside the volume.
+
+    Raises:
+        ModelExecutionError: If ``mask`` is empty.
+        InvalidInputError: If any volume axis is shorter than the crop.
+    """
+    points = np.argwhere(mask)
+    if points.size == 0:
+        raise ModelExecutionError("LesionDiffusion cannot crop an empty lesion mask")
+    slices: list[slice] = []
+    for axis, requested in enumerate(shape):
+        center = int((points[:, axis].min() + points[:, axis].max()) // 2)
+        length = mask.shape[axis]
+        if length < requested:
+            raise InvalidInputError(
+                f"Pinned LesionDiffusion crop requires source axis {axis} >= {requested}; "
+                f"got {length}"
+            )
+        start = min(max(center - requested // 2, 0), length - requested)
+        slices.append(slice(start, start + requested))
+    return slices[0], slices[1], slices[2]
+
+
+def piecewise_ct(array: Array, mx: Mlx, *, inverse: bool = False) -> Array:
+    """Apply the source's four-knot CT intensity transform on the Metal GPU.
+
+    Map HU landmarks (-1000, -200, 200, 1000) to (0, 0.2, 0.8, 1) with
+    piecewise-linear interpolation and endpoint clipping. ``inverse=True``
+    swaps those landmarks to return HU, also clipping outside its input range.
+    The input array may have any shape; the result is float32 with the same shape.
+    Clipping means HU outside [-1000, 1000] cannot survive a round trip.
+    """
+    inputs, outputs = (
+        ((0.0, 0.2, 0.8, 1.0), (-1000.0, -200.0, 200.0, 1000.0))
+        if inverse
+        else ((-1000.0, -200.0, 200.0, 1000.0), (0.0, 0.2, 0.8, 1.0))
+    )
+    result = mx.full(array.shape, outputs[0], dtype=mx.float32)
+    for index in range(3):
+        slope = (outputs[index + 1] - outputs[index]) / (inputs[index + 1] - inputs[index])
+        result = mx.where(
+            (array >= inputs[index]) & (array < inputs[index + 1]),
+            slope * (array - inputs[index]) + outputs[index],
+            result,
+        )
+    return mx.where(array >= inputs[-1], outputs[-1], result)
+
+
+def encode_source(model: VqAutoencoder, array: Array, mx: Mlx) -> Array:
+    """Encode a normalized volume ``(B, C, D, H, W)`` to its VQ latent, evaluated eagerly.
+
+    ``array`` holds ``piecewise_ct`` output in [0, 1]. The latent has half the
+    spatial size, rounded down, so decoding an odd extent returns one voxel fewer;
+    the runner pads to even extents first. The latent is not quantized here;
+    ``VqAutoencoder.decode`` does that.
+    """
+    latent = model.encode(array)
+    mx.eval(latent)
+    return latent
+
+
+def encode_attributes(
+    attributes: Mapping[str, object],
+    model: BiomedTextEncoder,
+    tokenizer_path: Path,
+    mx: Mlx,
+) -> Array:
+    """Embed each attribute value as one BiomedCLIP text vector for cross-attention.
+
+    Only the values of ``attributes`` are encoded, one sequence per entry in
+    insertion order; keys are ignored. Tokenization runs on the host (truncated and
+    padded to 256 tokens, pad id 0) and every learned text layer runs in MLX.
+
+    Returns:
+        Context of shape ``(1, len(attributes), D_ctx)``, in the text encoder's dtype.
+    """
+    from tokenizers import Tokenizer
+
+    tokenizer = cast("TokenizerContract", Tokenizer.from_file(str(tokenizer_path)))
+    tokenizer.enable_truncation(max_length=256)
+    tokenizer.enable_padding(length=256, pad_id=0, pad_token="[PAD]")
+    tokens = tokenizer.encode_batch([str(value) for value in attributes.values()])
+    # Tokenizer IDs are nested numeric lists; list invariance loses the recursive type.
+    token_ids = cast(ArrayConstructor, mx.array)(
+        cast("ArrayValue", [token.ids for token in tokens]), dtype=mx.int32
+    )
+    encoded = model(token_ids)
+    mx.eval(encoded)
+    return encoded[None]