@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/input/vision/index.ts

@@ -0,0 +1,3 @@
+export * from "./preprocess";
+export * from "./prompt";
+export * from "./video-frames";

package/src/input/vision/preprocess.ts

@@ -0,0 +1,247 @@
+// Gemma-4 unified image preprocessing — port of optiq's
+// vlm/gemma4_unified/image_processing.py (vendored from mlx-vlm, BSD-3).
+//
+// decode → aspect-ratio-preserving resize (PIL-style convolution
+// resampling, bicubic kernel with antialias-scaled support) → rescale
+// 1/255 → patchify into 48×48 patch vectors with 2D grid positions →
+// pad to the 280 soft-token budget.
+//
+// Decode strategy (Bun ≥1.4.0): PNG goes straight through fast-png
+// (exact, lossless — the parity-golden path). Everything else (JPEG,
+// HEIC, AVIF, WebP, TIFF, GIF, BMP) is transcoded to lossless PNG by
+// Bun.Image — native OS codecs (ImageIO on macOS), EXIF auto-orient,
+// off-thread — then decoded by fast-png. Bun.Image has no raw-pixel
+// terminal, so the PNG bridge is the supported path to pixels.
+//
+// Resize note: Bun.Image.resize is NOT used — its kernels (lanczos3,
+// cubic≠PIL-bicubic) don't match PIL's antialiased bicubic, which the
+// vision tower was trained behind. PIL's exact float pipeline isn't
+// bit-reproducible here either, so vision logits can differ slightly
+// from the python stack on images that need resizing. Images already
+// sized to multiples of 48 (≤2520 patches) skip the resize entirely and
+// are bit-identical through preprocessing.
+
+import { decode as decodePng } from "fast-png";
+
+export const PATCH_SIZE = 16;
+export const POOLING_KERNEL_SIZE = 3;
+export const MODEL_PATCH_SIZE = 48; // patch_size * pooling_kernel_size
+export const NUM_SOFT_TOKENS = 280;
+
+export interface RGBImage {
+  width: number;
+  height: number;
+  /** Row-major RGB, 3 bytes/px. */
+  data: Uint8Array;
+}
+
+export interface PreprocessedImage {
+  /** [numSoftTokens, 6912] f32 patch vectors (padded). */
+  patches: Float32Array;
+  /** [numSoftTokens, 2] (x, y) grid positions; -1 padding. */
+  positions: Int32Array;
+  /** Real (unpadded) soft-token count. */
+  softTokens: number;
+}
+
+function pngToRgb(bytes: Uint8Array): RGBImage {
+  const png = decodePng(bytes);
+  const ch = png.channels;
+  const out = new Uint8Array(png.width * png.height * 3);
+  const src = png.data as Uint8Array;
+  const n = png.width * png.height;
+  if (png.depth === 16) {
+    const src16 = png.data as Uint16Array;
+    for (let i = 0; i < n; i++)
+      for (let c = 0; c < 3; c++)
+        out[i * 3 + c] = src16[i * ch + Math.min(c, ch - 1)]! >> 8;
+  } else {
+    for (let i = 0; i < n; i++)
+      for (let c = 0; c < 3; c++)
+        out[i * 3 + c] = src[i * ch + Math.min(c, ch - 1)]!;
+  }
+  return { width: png.width, height: png.height, data: out };
+}
+
+export async function decodeImage(bytes: Uint8Array): Promise<RGBImage> {
+  // PNG: direct, exact (the parity-golden path)
+  if (bytes.length > 8 && bytes[0] === 0x89 && bytes[1] === 0x50)
+    return pngToRgb(bytes);
+  // everything else: native transcode to lossless PNG (HEIC/AVIF/WebP/
+  // JPEG/TIFF/GIF/BMP — whatever the OS codecs decode), EXIF auto-orient
+  try {
+    const png = await new Bun.Image(bytes).png({ compressionLevel: 1 }).bytes();
+    return pngToRgb(png);
+  } catch (e) {
+    const code = (e as { code?: string }).code;
+    throw new Error(
+      code === "ERR_IMAGE_UNKNOWN_FORMAT"
+        ? "unsupported image format"
+        : `image decode failed: ${(e as Error).message}`,
+    );
+  }
+}
+
+// --- PIL-style convolution resize (bicubic, a = -0.5) ---------------------
+
+function bicubic(x: number): number {
+  const a = -0.5;
+  x = Math.abs(x);
+  if (x < 1) return ((a + 2) * x - (a + 3)) * x * x + 1;
+  if (x < 2) return (((x - 5) * x + 8) * x - 4) * a;
+  return 0;
+}
+
+interface ResampleCoeffs {
+  bounds: Int32Array; // [outSize × 2]: (start, count)
+  weights: Float64Array; // [outSize × kmax]
+  kmax: number;
+}
+
+/** Port of PIL's precompute_coeffs (antialias: support scales with ratio). */
+function precomputeCoeffs(inSize: number, outSize: number): ResampleCoeffs {
+  const SUPPORT = 2.0;
+  const scale = inSize / outSize;
+  const filterscale = Math.max(scale, 1.0);
+  const support = SUPPORT * filterscale;
+  const kmax = Math.ceil(support) * 2 + 1;
+  const bounds = new Int32Array(outSize * 2);
+  const weights = new Float64Array(outSize * kmax);
+
+  for (let xx = 0; xx < outSize; xx++) {
+    const center = (xx + 0.5) * scale;
+    let xmin = Math.floor(center - support);
+    if (xmin < 0) xmin = 0;
+    let xmax = Math.ceil(center + support);
+    if (xmax > inSize) xmax = inSize;
+    const count = xmax - xmin;
+    let sum = 0;
+    for (let k = 0; k < count; k++) {
+      const w = bicubic((xmin + k - center + 0.5) / filterscale);
+      weights[xx * kmax + k] = w;
+      sum += w;
+    }
+    if (sum !== 0)
+      for (let k = 0; k < count; k++) weights[xx * kmax + k]! /= sum;
+    bounds[xx * 2] = xmin;
+    bounds[xx * 2 + 1] = count;
+  }
+  return { bounds, weights, kmax };
+}
+
+/** Separable resize on float RGB planes (clamped back to [0,255] like PIL). */
+export function resizeBicubic(img: RGBImage, outW: number, outH: number): RGBImage {
+  const { width: inW, height: inH, data } = img;
+  // horizontal pass → [inH, outW]
+  const hCoef = precomputeCoeffs(inW, outW);
+  const tmp = new Float64Array(inH * outW * 3);
+  for (let y = 0; y < inH; y++) {
+    for (let x = 0; x < outW; x++) {
+      const xmin = hCoef.bounds[x * 2]!;
+      const count = hCoef.bounds[x * 2 + 1]!;
+      let r = 0, g = 0, b = 0;
+      for (let k = 0; k < count; k++) {
+        const w = hCoef.weights[x * hCoef.kmax + k]!;
+        const src = (y * inW + xmin + k) * 3;
+        r += data[src]! * w;
+        g += data[src + 1]! * w;
+        b += data[src + 2]! * w;
+      }
+      const dst = (y * outW + x) * 3;
+      tmp[dst] = r;
+      tmp[dst + 1] = g;
+      tmp[dst + 2] = b;
+    }
+  }
+  // vertical pass → [outH, outW]
+  const vCoef = precomputeCoeffs(inH, outH);
+  const out = new Uint8Array(outH * outW * 3);
+  for (let y = 0; y < outH; y++) {
+    const ymin = vCoef.bounds[y * 2]!;
+    const count = vCoef.bounds[y * 2 + 1]!;
+    for (let x = 0; x < outW; x++) {
+      let r = 0, g = 0, b = 0;
+      for (let k = 0; k < count; k++) {
+        const w = vCoef.weights[y * vCoef.kmax + k]!;
+        const src = ((ymin + k) * outW + x) * 3;
+        r += tmp[src]! * w;
+        g += tmp[src + 1]! * w;
+        b += tmp[src + 2]! * w;
+      }
+      const dst = (y * outW + x) * 3;
+      out[dst] = Math.min(255, Math.max(0, Math.round(r)));
+      out[dst + 1] = Math.min(255, Math.max(0, Math.round(g)));
+      out[dst + 2] = Math.min(255, Math.max(0, Math.round(b)));
+    }
+  }
+  return { width: outW, height: outH, data: out };
+}
+
+/** Port of _aspect_ratio_preserving_resize. */
+export function targetSize(
+  width: number, height: number,
+  patchSize = PATCH_SIZE, poolingKernel = POOLING_KERNEL_SIZE,
+): { width: number; height: number } {
+  const maxPatches = NUM_SOFT_TOKENS * poolingKernel ** 2;
+  const targetPx = maxPatches * patchSize ** 2;
+  const factor = Math.sqrt(targetPx / (height * width));
+  const sideMult = poolingKernel * patchSize; // 48
+
+  let tH = Math.floor((factor * height) / sideMult) * sideMult;
+  let tW = Math.floor((factor * width) / sideMult) * sideMult;
+  if (tH === 0 && tW === 0) throw new Error("image resizes to 0×0");
+
+  const maxSide = Math.floor(maxPatches / poolingKernel ** 2) * sideMult;
+  if (tH === 0) {
+    tH = sideMult;
+    tW = Math.min(Math.floor(width / height) * sideMult, maxSide);
+  } else if (tW === 0) {
+    tW = sideMult;
+    tH = Math.min(Math.floor(height / width) * sideMult, maxSide);
+  }
+  return { width: tW, height: tH };
+}
+
+/** Full pipeline: bytes → padded patch vectors + grid positions. */
+export async function preprocessImage(bytes: Uint8Array): Promise<PreprocessedImage> {
+  let img = await decodeImage(bytes);
+  const t = targetSize(img.width, img.height);
+  if (t.width !== img.width || t.height !== img.height)
+    img = resizeBicubic(img, t.width, t.height);
+
+  const { width, height, data } = img;
+  const p = MODEL_PATCH_SIZE;
+  const pH = Math.floor(height / p);
+  const pW = Math.floor(width / p);
+  const n = pH * pW;
+  if (n === 0) throw new Error("image smaller than one 48×48 patch");
+  if (n > NUM_SOFT_TOKENS)
+    throw new Error(`image yields ${n} patches > budget ${NUM_SOFT_TOKENS}`);
+
+  const patchDim = p * p * 3;
+  const patches = new Float32Array(NUM_SOFT_TOKENS * patchDim);
+  const positions = new Int32Array(NUM_SOFT_TOKENS * 2).fill(-1);
+
+  // patch vector layout matches the python reshape/transpose:
+  // [pH, pW, p, p, C] flattened — i.e. for patch (py, px):
+  // index = ((dy * p) + dx) * 3 + c over [C=channel-last]... python does
+  // transpose(1,3,2,4,0) from [C, pH, p, pW, p] → [pH, pW, p, p, C]
+  for (let py = 0; py < pH; py++) {
+    for (let px = 0; px < pW; px++) {
+      const base = (py * pW + px) * patchDim;
+      for (let dy = 0; dy < p; dy++) {
+        for (let dx = 0; dx < p; dx++) {
+          const src = ((py * p + dy) * width + px * p + dx) * 3;
+          const dst = base + (dy * p + dx) * 3;
+          // python rescales then patchifies in CHW; layout [p, p, C]
+          patches[dst] = data[src]! / 255;
+          patches[dst + 1] = data[src + 1]! / 255;
+          patches[dst + 2] = data[src + 2]! / 255;
+        }
+      }
+      positions[(py * pW + px) * 2] = px; // x
+      positions[(py * pW + px) * 2 + 1] = py; // y
+    }
+  }
+  return { patches, positions, softTokens: n };
+}

package/src/input/vision/prompt.ts

@@ -0,0 +1,398 @@
+// Multimodal prompt assembly: OpenAI image/audio content parts → spliced
+// token ids + merged input embeddings + attention/zeroing masks.
+//
+// The chat template renders one <|image|> token (258880) per image part and
+// one <|audio|> token (258881) per audio part; we expand each in document
+// order — <boi> + image_token × softTokens + <eoi> (matching optiq's vision
+// frontend) and <boa> + audio_token × n + <eoa> with
+// n = min(ceil(duration_ms/40), 750) (processing_gemma4.py) — embed the full
+// sequence, then overwrite the soft-token rows with the towers' features.
+//
+// Attention semantics (`02d723a:docs/design/generic-model-support.md` §6.6, resolved vs the oracle):
+//   - image runs attend bidirectionally among themselves (bidirMask), BUT
+//     any audio token in the prompt disables the vision overlay entirely —
+//     mixed image+audio prompts run fully causal (bidirMask = null);
+//   - audio runs are strictly causal, always;
+//   - per-layer-input id zeroing (e2b/e4b) applies to the UNION of
+//     multimodal soft tokens (image | audio) — multimodalMask, decoupled
+//     from the bidirectional mask.
+
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { AudioEncoder,TextEmbeddingModel,VisionEncoder } from "../../contracts/mlx/media";
+import type { PromptNativeWork } from "../../contracts/portable/prompt";
+import type { EncoderCache } from "../../state/encoder-cache";
+import type { ChatMessage,ChatTemplate,ToolDefinition } from "../chat-template";
+import type { LoadedTokenizer } from "../tokenizer";
+
+import { audioSoftTokenCount,decodeAudio } from "../audio/decode";
+import { extractMelFeatures } from "../audio/features";
+import { fetchMediaBytes,videoMediaFetchPolicy } from "../media-fetch";
+
+export type { VisionEncoder } from "../../contracts/mlx/media";
+
+export interface VisionPrompt {
+  ids: number[];
+  prefixIdentity?: string;
+  /** Unscaled merged embeddings [1, L, hidden] — caller disposes. */
+  embeddings: MlxArray;
+  /** bool [L] image-token mask — caller disposes. */
+  imageMask: MlxArray;
+}
+
+export interface MultimodalPrompt {
+  ids: number[];
+  /** Includes every media item and the exact token prefix through the last one. */
+  prefixIdentity?: string;
+  /** Unscaled merged embeddings [1, L, hidden] — caller disposes. */
+  embeddings: MlxArray;
+  /** bool [L] image-token mask for the bidirectional attention overlay.
+   *  null when the prompt contains ANY audio (§3.3 Q1: mixed image+audio
+   *  runs fully causal) or no images. Caller disposes when non-null. */
+  bidirMask: MlxArray | null;
+  /** bool [L] union multimodal soft-token mask (image | audio) for
+   *  per-layer-input id zeroing (§3.3 Q2). Caller disposes. */
+  multimodalMask: MlxArray;
+}
+
+export interface VisionTokenIds {
+  imageTokenId: number;
+  boiTokenId: number;
+  eoiTokenId: number;
+}
+
+export interface AudioTokenIds {
+  audioTokenId: number;
+  boaTokenId: number;
+  eoaTokenId: number;
+}
+
+/** The towers a multimodal prompt may need; pass only what the request uses
+ *  (vision-only and audio-only fall out as special cases). */
+export interface MultimodalTowers<P extends { softTokens: number }> {
+  vision?: { tower: VisionEncoder<P>; tokenIds: VisionTokenIds; cache?: EncoderCache };
+  audio?: { tower: AudioEncoder; tokenIds: AudioTokenIds; cache?: EncoderCache };
+}
+
+/** Extract image bytes from OpenAI-style content parts, rewriting the
+ *  parts to the template's {type:"image"} form. */
+export async function extractImages(
+  messages: ChatMessage[],
+): Promise<{ messages: ChatMessage[]; images: Uint8Array[] }> {
+  const images: Uint8Array[] = [];
+  const out: ChatMessage[] = [];
+  for (const m of messages) {
+    if (!Array.isArray(m.content)) {
+      out.push(m);
+      continue;
+    }
+    const parts: Array<Record<string, unknown>> = [];
+    for (const part of m.content) {
+      if (part.type === "image_url") {
+        const url = (part.image_url as { url: string } | undefined)?.url
+          ?? (part.image_url as unknown as string);
+        if (typeof url !== "string") throw new Error("image_url part missing url");
+        images.push(await fetchMediaBytes(url, "image"));
+        parts.push({ type: "image" });
+      } else if (part.type === "image") {
+        // already template-form: expect base64 `data` or pass-through marker
+        if (typeof part.data === "string") {
+          images.push(Uint8Array.from(Buffer.from(part.data, "base64")));
+          parts.push({ type: "image" });
+        } else {
+          throw new Error("image part requires base64 `data`");
+        }
+      } else {
+        parts.push(part);
+      }
+    }
+    out.push({ ...m, content: parts });
+  }
+  return { messages: out, images };
+}
+
+/** Extract video bytes from content parts, rewriting to the template's
+ *  {type:"video"} form. Accepts {type:"video_url", video_url:{url}} (data:
+ *  or http(s), same SSRF guard as images with the larger video body cap)
+ *  and {type:"video"} with base64 `data`. Decode to frames happens later
+ *  (packages/inference/src/input/vision/video-frames.ts — the AVFoundation sidecar). */
+export async function extractVideos(
+  messages: ChatMessage[],
+): Promise<{ messages: ChatMessage[]; videos: Uint8Array[] }> {
+  const videos: Uint8Array[] = [];
+  const out: ChatMessage[] = [];
+  for (const m of messages) {
+    if (!Array.isArray(m.content)) {
+      out.push(m);
+      continue;
+    }
+    const parts: Array<Record<string, unknown>> = [];
+    for (const part of m.content) {
+      if (part.type === "video_url") {
+        const url = (part.video_url as { url: string } | undefined)?.url
+          ?? (part.video_url as unknown as string);
+        if (typeof url !== "string") throw new Error("video_url part missing url");
+        videos.push(await fetchMediaBytes(url, "video"));
+        parts.push({ type: "video" });
+      } else if (part.type === "video") {
+        if (typeof part.data === "string") {
+          // Same cap as fetched/data:-URL video bodies — the inline base64
+          // form must not be the uncapped route.
+          const cap = videoMediaFetchPolicy().maxBytes;
+          if (part.data.length * 0.75 > cap)
+            throw new Error(
+              `video part exceeds the ${Math.round(cap / 1024 / 1024)} MB cap`,
+            );
+          videos.push(Uint8Array.from(Buffer.from(part.data, "base64")));
+          parts.push({ type: "video" });
+        } else {
+          throw new Error("video part requires base64 `data`");
+        }
+      } else {
+        parts.push(part);
+      }
+    }
+    out.push({ ...m, content: parts });
+  }
+  return { messages: out, videos };
+}
+
+/** Extract audio bytes from OpenAI-style content parts, rewriting the parts
+ *  to the template's {type:"audio"} form. Accepts the OpenAI-canonical
+ *  {type:"input_audio", input_audio:{data:<b64>, format:...}} plus optiq's
+ *  {type:"audio"} (base64 `data`) and {type:"audio_url"} aliases. */
+export async function extractAudio(
+  messages: ChatMessage[],
+): Promise<{ messages: ChatMessage[]; audio: Uint8Array[] }> {
+  const audio: Uint8Array[] = [];
+  const out: ChatMessage[] = [];
+  for (const m of messages) {
+    if (!Array.isArray(m.content)) {
+      out.push(m);
+      continue;
+    }
+    const parts: Array<Record<string, unknown>> = [];
+    for (const part of m.content) {
+      if (part.type === "input_audio") {
+        const data = (part.input_audio as { data?: string } | undefined)?.data;
+        if (typeof data !== "string")
+          throw new Error("input_audio part requires base64 `input_audio.data`");
+        audio.push(Uint8Array.from(Buffer.from(data, "base64")));
+        parts.push({ type: "audio" });
+      } else if (part.type === "audio") {
+        if (typeof part.data === "string") {
+          audio.push(Uint8Array.from(Buffer.from(part.data, "base64")));
+          parts.push({ type: "audio" });
+        } else {
+          throw new Error("audio part requires base64 `data`");
+        }
+      } else if (part.type === "audio_url") {
+        const url = (part.audio_url as { url: string } | undefined)?.url
+          ?? (part.audio_url as unknown as string);
+        if (typeof url !== "string") throw new Error("audio_url part missing url");
+        audio.push(await fetchMediaBytes(url, "audio"));
+        parts.push({ type: "audio" });
+      } else {
+        parts.push(part);
+      }
+    }
+    out.push({ ...m, content: parts });
+  }
+  return { messages: out, audio };
+}
+
+/** Decoded + featurized audio clip, ready to splice. */
+interface PreparedAudio {
+  mel: Float32Array;
+  frames: number;
+  softTokens: number;
+}
+
+export async function buildMultimodalPrompt<P extends { softTokens: number }>(
+  model: TextEmbeddingModel,
+  towers: MultimodalTowers<P>,
+  tokenizer: LoadedTokenizer,
+  template: ChatTemplate,
+  messages: ChatMessage[],
+  images: Uint8Array[],
+  audio: Uint8Array[],
+  tools: ToolDefinition[] | null = null,
+  nativeWork: PromptNativeWork = (work) => work(),
+): Promise<MultimodalPrompt> {
+  if (images.length > 0 && !towers.vision)
+    throw new Error("images present but no vision tower configured");
+  if (audio.length > 0 && !towers.audio)
+    throw new Error("audio present but no audio tower configured");
+
+  const rendered = template.render(messages, { tools });
+  let ids = tokenizer.encode(rendered);
+  if (ids[0] === ids[1] && ids[0] === tokenizer.bosTokenId) ids = ids.slice(1);
+
+  // preprocess every clip/image up front (splice lengths depend on them);
+  // audio soft-token counts come from the DECODED samples, never metadata.
+  const pre = await Promise.all(
+    images.map((bytes) => towers.vision!.tower.preprocess(bytes)),
+  );
+  const clips: PreparedAudio[] = audio.map((bytes) => {
+    const samples = decodeAudio(bytes);
+    const softTokens = audioSoftTokenCount(samples.length);
+    const m = extractMelFeatures(samples);
+    return { mel: m.features, frames: m.frames, softTokens };
+  });
+
+  const vIds = towers.vision?.tokenIds;
+  const aIds = towers.audio?.tokenIds;
+
+  // expand each <|image|> into <boi> + image_token×soft + <eoi> and each
+  // <|audio|> into <boa> + audio_token×n + <eoa>, in document order
+  const spliced: number[] = [];
+  const runs: { start: number; length: number; kind: "image" | "audio"; index: number }[] = [];
+  let imgIdx = 0;
+  let audIdx = 0;
+  for (const id of ids) {
+    if (vIds && id === vIds.imageTokenId) {
+      const p = pre[imgIdx];
+      if (!p) throw new Error("more <|image|> markers than images");
+      spliced.push(vIds.boiTokenId);
+      runs.push({ start: spliced.length, length: p.softTokens, kind: "image", index: imgIdx });
+      for (let k = 0; k < p.softTokens; k++) spliced.push(vIds.imageTokenId);
+      spliced.push(vIds.eoiTokenId);
+      imgIdx++;
+    } else if (aIds && id === aIds.audioTokenId) {
+      const clip = clips[audIdx];
+      if (!clip) throw new Error("more <|audio|> markers than audio clips");
+      spliced.push(aIds.boaTokenId);
+      runs.push({ start: spliced.length, length: clip.softTokens, kind: "audio", index: audIdx });
+      for (let k = 0; k < clip.softTokens; k++) spliced.push(aIds.audioTokenId);
+      spliced.push(aIds.eoaTokenId);
+      audIdx++;
+    } else {
+      spliced.push(id);
+    }
+  }
+  if (imgIdx !== images.length)
+    throw new Error(`${images.length} images but ${imgIdx} <|image|> markers`);
+  if (audIdx !== audio.length)
+    throw new Error(`${audio.length} audio clips but ${audIdx} <|audio|> markers`);
+
+  // Exact encoded objects contain no request positions or template state.
+  // CPU preprocessing and SSD reads run before acquiring native execution.
+  const keys = runs.map(run => towers[run.kind === "image" ? "vision" : "audio"]?.cache
+    ? new Bun.CryptoHasher("sha256").update(`gemma-${run.kind}-input-v1`)
+      .update(run.kind === "image" ? images[run.index]! : audio[run.index]!).digest("hex")
+    : "");
+  const lastMedia = runs.at(-1);
+  // A matching namespace guarantees any reused prefix includes all media.
+  // The audio-presence bit also identifies Gemma's causal/bidirectional policy.
+  const prefixIdentity = lastMedia && keys.every(Boolean)
+    ? new Bun.CryptoHasher("sha256").update("gemma-prepared-prefix-v1")
+      .update(JSON.stringify({ audio: audio.length > 0,
+        tokens: spliced.slice(0, lastMedia.start + lastMedia.length),
+        media: runs.map((run, i) => [run.kind, keys[i],
+          (run.kind === "image" ? towers.vision : towers.audio)!.cache!.identity]),
+      })).digest("hex") : undefined;
+  const reused: (MlxArray | null)[] = [];
+  try {
+    for (const [index, run] of runs.entries()) {
+      const cache = run.kind === "image" ? towers.vision?.cache : towers.audio?.cache;
+      reused.push(cache ? await cache.take(keys[index]!) : null);
+    }
+    return await nativeWork(async () => {
+      // embed text tokens, then overwrite soft-token rows with tower features
+      using idsArr = ops.fromInt32(spliced, [1, spliced.length]);
+      let embeds = model.embed.encode(idsArr);
+      try {
+        const hidden = embeds.shape[2]!;
+        for (const [runIndex, run] of runs.entries()) {
+          let cast: MlxArray;
+          if (run.kind === "image") {
+            // vision towers return pre-divided by embed_scale (their oracle's
+            // frontend convention) — cast to the embedding dtype and splice.
+            const cached = reused[runIndex];
+            using feats = cached ?? towers.vision!.tower.features(pre[run.index]!); // [1, soft, hidden]
+            reused[runIndex] = null;
+            if (!cached) towers.vision!.cache?.put(keys[runIndex]!, feats);
+            cast = feats.astype(embeds.dtype);
+          } else {
+            // audio mirrors `02d723a:scripts/oracle/gen-e4b-audio-golden.py` EXACTLY:
+            //   features.astype(embeds.dtype) / embed_scale
+            // — raw f32 embed_audio output, cast to bf16 FIRST, then divided by a
+            // weak (dtype-following, i.e. bf16) embed_scale scalar. The order is
+            // load-bearing for the bit-exact greedy gate (see AudioTower.features).
+            const clip = clips[run.index]!;
+            const cached = reused[runIndex];
+            using raw = cached ?? towers.audio!.tower.features(clip.mel, clip.frames, false); // [1, n, hidden] f32
+            reused[runIndex] = null;
+            if (raw.shape[1] !== clip.softTokens) {
+              const n = raw.shape[1];
+              throw new Error(`audio tower produced ${n} frames but splice expects ${clip.softTokens}`);
+            }
+            if (!cached) towers.audio!.cache?.put(keys[runIndex]!, raw);
+            using bf = raw.astype(embeds.dtype);
+            using scale = ops.scalarLike(towers.audio!.tower.embedScale, bf);
+            cast = ops.div(bf, scale);
+          }
+          const { start, length } = run;
+          try {
+            const updated = ops.sliceUpdate(embeds, cast, [0, start, 0], [1, start + length, hidden]);
+            embeds.dispose();
+            embeds = updated;
+          } finally { cast.dispose(); }
+        }
+
+        // masks over the spliced sequence: union multimodal soft tokens (per-layer
+        // id zeroing) always; image-only bidirectional overlay only when NO audio
+        // is present (§3.3 Q1).
+        const mmInts = new Int32Array(spliced.length);
+        for (let i = 0; i < spliced.length; i++) {
+          const id = spliced[i]!;
+          mmInts[i] = (vIds && id === vIds.imageTokenId) || (aIds && id === aIds.audioTokenId) ? 1 : 0;
+        }
+        const mmI32 = MlxArray.fromInt32(mmInts, [spliced.length]);
+        const multimodalMask = mmI32.astype(Dtype.bool);
+        mmI32.dispose();
+
+        let bidirMask: MlxArray | null = null;
+        if (images.length > 0 && audio.length === 0) {
+          const imgInts = new Int32Array(spliced.length);
+          for (let i = 0; i < spliced.length; i++)
+            imgInts[i] = spliced[i] === vIds!.imageTokenId ? 1 : 0;
+          const imgI32 = MlxArray.fromInt32(imgInts, [spliced.length]);
+          bidirMask = imgI32.astype(Dtype.bool);
+          imgI32.dispose();
+        }
+
+        return { ids: spliced, embeddings: embeds, bidirMask, multimodalMask, prefixIdentity };
+      } catch (error) { embeds.dispose(); throw error; }
+    });
+  } finally { for (const tensor of reused) tensor?.dispose(); }
+}
+
+/** Vision-only prompt (back-compat wrapper over buildMultimodalPrompt —
+ *  behavior byte-identical to the pre-audio builder: the returned imageMask
+ *  serves as BOTH the bidirectional-attention mask and the per-layer-input
+ *  zeroing mask, which forwardEmbeddings derives from `bidir` when no
+ *  separate multimodal mask is passed). */
+export async function buildVisionPrompt<P extends { softTokens: number }>(
+  model: TextEmbeddingModel,
+  tower: VisionEncoder<P>,
+  tokenizer: LoadedTokenizer,
+  template: ChatTemplate,
+  messages: ChatMessage[],
+  images: Uint8Array[],
+  tokenIds: VisionTokenIds,
+  tools: ToolDefinition[] | null = null,
+  nativeWork: PromptNativeWork = (work) => work(),
+  cache?: EncoderCache,
+): Promise<VisionPrompt> {
+  const mp = await buildMultimodalPrompt(
+    model, { vision: { tower, tokenIds, cache } }, tokenizer, template,
+    messages, images, [], tools, nativeWork,
+  );
+  // no audio was passed, so the union mask IS the image mask — reuse it for
+  // the zero-image edge (the original builder returned an all-false mask).
+  const imageMask = mp.bidirMask ?? mp.multimodalMask;
+  if (mp.bidirMask) mp.multimodalMask.dispose();
+  return { ids: mp.ids, embeddings: mp.embeddings, imageMask, prefixIdentity: mp.prefixIdentity };
+}