mlx-bun 0.0.10 → 0.0.12

package/src/train/trainer.ts

@@ -20,7 +20,6 @@ import type { RuntimeModel } from "../model/factory";
 import { Gemma4Model, type GradCheckpointCtx } from "../model/gemma4";
 import { MiniCPM5Model } from "../model/minicpm5";
 import { setTrainingAttn } from "../model/flash-attention";
-import { setFusedGeluTraining } from "../model/fused-geglu-kernel";
 import { SegmentedBackward, SegmentedBackwardGemma4, SegmentedBackwardOrpo, SegmentedBackwardOrpoGemma4, SegmentedBackwardOrpoPrefix, SegmentedBackwardOrpoPrefixGemma4, planSegmentsBySize } from "./segmented";
 import type { LoadedTokenizer } from "../tokenizer";
 import type { ChatTemplate } from "../chat-template";
@@ -203,6 +202,51 @@ export const DEFAULT_TRAIN_CONFIG: TrainConfig = {
   baseModel: "",
 };
 
+/** Build the per-leaf LoRA+ LR multipliers for flatParams' stable
+ *  `[...A, ...B]` order. `undefined` keeps AdamW's uniform-LR fast path. */
+export function loraPlusLrScale(
+  paramCount: number,
+  targetCount: number,
+  ratio: number,
+): number[] | undefined {
+  if (!Number.isFinite(ratio) || ratio < 1)
+    throw new RangeError(`loraPlusRatio must be a finite number >= 1 (got ${ratio})`);
+  if (paramCount !== targetCount * 2)
+    throw new Error(
+      `LoRA+ expected ${targetCount * 2} A/B leaves for ${targetCount} targets, got ${paramCount}`,
+    );
+  if (ratio === 1) return undefined;
+  return Array.from({ length: paramCount }, (_, i) => (i < targetCount ? 1 : ratio));
+}
+
+/** One optimizer policy for SFT, DPO, and ORPO. Keeping LoRA+ here prevents a
+ *  method-specific loop from silently accepting the knob without applying it. */
+function createLoraOptimizer(
+  params: MlxArray[],
+  lora: TrainableLora,
+  cfg: TrainConfig,
+  emit: Emit,
+): AdamW {
+  const lrScale = loraPlusLrScale(params.length, lora.targets.length, cfg.loraPlusRatio);
+  if (lrScale) emit({
+    type: "stage",
+    stage: "setup",
+    progress: 0.04,
+    message: `LoRA+ : B leaves at ${cfg.loraPlusRatio}× LR`,
+  });
+  return new AdamW(
+    params,
+    {
+      lr: cfg.learningRate,
+      betas: cfg.betas,
+      eps: 1e-8,
+      weightDecay: cfg.weightDecay,
+      lrScale,
+    },
+    (i, p) => writeParam(lora, i, p),
+  );
+}
+
 /** Run a LoRA fine-tune. Emits per-step train metrics, periodic val metrics,
  *  and a final stage:done with the adapter path + applied ranks. Saves the
  *  adapter (last + best-on-val) and returns where it landed. */
@@ -260,7 +304,6 @@ export async function trainLora(
   // autograd flows and the backward recomputes the gelu from the primal instead
   // of retaining the spelled-out intermediates. No-op for non-Gemma models.
   // Cleared in the finally below.
-  setFusedGeluTraining(true);
 
   const saveCfg: SaveAdapterConfig = {
     rank: cfg.rank,
@@ -288,6 +331,7 @@ export async function trainLora(
   // weak proxy. Capture them when the loop emits them so metrics.json lets you
   // pick the best-MARGIN checkpoint (overfit shows as train↑ / val-margin flat).
   const valHistory: { step: number; loss: number; margin?: number; accuracy?: number; checkpoint: string | null }[] = [];
+  const checkpointSaves: Promise<void>[] = [];
   const startedAt = Date.now();
 
   // Live, append-only metrics stream (the source `mlx-bun train-watch` tails).
@@ -321,7 +365,11 @@ export async function trainLora(
       if (cfg.saveCheckpoints) {
         const tag = `step-${String(e.step).padStart(5, "0")}-val${e.loss.toFixed(4)}`;
         checkpoint = `${cfg.adapterPath}/checkpoints/${tag}`;
-        saveAdapter(lora, checkpoint, saveCfg, appliedRanks);
+        const save = saveAdapter(lora, checkpoint, saveCfg, appliedRanks);
+        checkpointSaves.push(save);
+        // Keep a handler attached while training continues. Promise.all below
+        // still surfaces the original rejection before the run can succeed.
+        void save.catch(() => {});
         emit({ type: "stage", stage: "checkpoint", progress: e.step / cfg.iters,
                message: `checkpoint ${tag}` });
       }
@@ -339,9 +387,11 @@ export async function trainLora(
           ? await orpoLoop(model, tok, tmpl, dataDir, cfg, lora, collect)
           : await sftLoop(model, tok, tmpl, dataDir, cfg, lora, collect);
 
+    await Promise.all(checkpointSaves);
+
     // Final save (last adapter).
     detachTraining(model, lora);
-    saveAdapter(lora, cfg.adapterPath, saveCfg, appliedRanks);
+    await saveAdapter(lora, cfg.adapterPath, saveCfg, appliedRanks);
 
     // Durable, structured run record alongside the adapter (only when we kept
     // checkpoints). Pick the best-val one for convenience, but ship ALL of them
@@ -383,7 +433,6 @@ export async function trainLora(
     });
     return { adapterPath: cfg.adapterPath, appliedRanks, numIters: result.numIters };
   } finally {
-    setFusedGeluTraining(false);
     disposeLora(lora);
   }
 }
@@ -418,11 +467,7 @@ async function sftLoop(
     : [];
 
   const params = flatParams(lora);
-  const opt = new AdamW(
-    params,
-    { lr: cfg.learningRate, betas: cfg.betas, eps: 1e-8, weightDecay: cfg.weightDecay },
-    (i, p) => writeParam(lora, i, p),
-  );
+  const opt = createLoraOptimizer(params, lora, cfg, emit);
 
   // Segmented backward: stream the SFT backward segment-by-segment (only one
   // segment's activations live at a time). Phase A is MiniCPM5 SFT B=1 only;
@@ -489,7 +534,7 @@ async function sftLoop(
   // scripts/experiments/flash-fd-check.ts), though ~30× slower than ops.sdpa.
   // Cleared in finally.
   //
-  // GEMMA GUARD (same convention as the PERF_KERNEL/FUSED_GELU sanitization in
+  // GEMMA GUARD (same convention as the launcher's env sanitization in
   // cli.ts / job.ts, but enforced at the trainer): e4b on this path SIGTRAPed
   // (uncatchable native crash) at multi-K sequence lengths (>=2K, reproduced in
   // scripts/experiments/segmented-grad-test-e4b.ts; docs/reference/training.md)
@@ -661,11 +706,7 @@ async function dpoLoop(
     : [];
 
   const params = flatParams(lora);
-  const opt = new AdamW(
-    params,
-    { lr: cfg.learningRate, betas: cfg.betas, eps: 1e-8, weightDecay: cfg.weightDecay },
-    (i, p) => writeParam(lora, i, p),
-  );
+  const opt = createLoraOptimizer(params, lora, cfg, emit);
 
   const schedule =
     cfg.dpoLrSchedule === "cosine" || cfg.dpoWarmupIters > 0
@@ -792,18 +833,7 @@ async function orpoLoop(
     : [];
 
   const params = flatParams(lora);
-  // LoRA+: B leaves (the second half of flatParams) get a higher LR than A.
-  const nT = lora.targets.length;
-  const lrScale = cfg.loraPlusRatio !== 1
-    ? params.map((_, i) => (i < nT ? 1 : cfg.loraPlusRatio))
-    : undefined;
-  if (lrScale) emit({ type: "stage", stage: "setup", progress: 0.04,
-    message: `LoRA+ : B leaves at ${cfg.loraPlusRatio}× LR` });
-  const opt = new AdamW(
-    params,
-    { lr: cfg.learningRate, betas: cfg.betas, eps: 1e-8, weightDecay: cfg.weightDecay, lrScale },
-    (i, p) => writeParam(lora, i, p),
-  );
+  const opt = createLoraOptimizer(params, lora, cfg, emit);
 
   const schedule =
     cfg.orpoLrSchedule === "cosine" || cfg.orpoWarmupIters > 0
@@ -1175,11 +1205,11 @@ function clipGradsByNorm(grads: MlxArray[], normVal: number, maxNorm: number): v
 
 /** Supervised (response) tokens across ALL rows of the batch — the count
  *  tokens/sec is reported over (mlx-lm's n_tokens counts every row too). */
-function countResponseTokens(batch: SftBatch): number {
+export function countResponseTokens(batch: SftBatch): number {
   let total = 0;
   for (let r = 0; r < batch.ids.length; r++) {
     const promptLen = batch.promptLens[r]!;
-    const len = batch.ids[r]!.length;
+    const len = rowLength(batch, r);
     total += Math.max(0, len - Math.max(promptLen, 1));
   }
   return total;

package/src/vision/prompt.ts

@@ -1,10 +1,21 @@
-// Multimodal prompt assembly: OpenAI image_url content parts → spliced
-// token ids + merged input embeddings + image-token mask.
+// 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;
-// we expand each into <boi> + image_token × softTokens + <eoi> (matching
-// optiq's frontend), embed the full sequence, then overwrite the
-// image-token rows with the vision tower's features.
+// 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 (audio-input-plan.md §3.3, 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/array";
 import { Dtype } from "../mlx/ffi";
@@ -12,6 +23,10 @@ import * as ops from "../mlx/ops";
 import type { Gemma4Model } from "../model/gemma4";
 import type { ChatTemplate, ChatMessage, ToolDefinition } from "../chat-template";
 import type { LoadedTokenizer } from "../tokenizer";
+import type { AudioTower } from "../audio/conformer";
+import { audioSoftTokenCount, decodeAudio } from "../audio/decode";
+import { extractMelFeatures } from "../audio/features";
+import { fetchMediaBytes } from "../media-fetch";
 
 /** Common contract for both vision towers (encoder-free gemma4_unified in
  *  ./embedder.ts and the SigLIP encoder in ./siglip.ts): preprocess image
@@ -33,12 +48,38 @@ export interface VisionPrompt {
   imageMask: MlxArray;
 }
 
+export interface MultimodalPrompt {
+  ids: number[];
+  /** 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 };
+  audio?: { tower: AudioTower; tokenIds: AudioTokenIds };
+}
+
 /** Extract image bytes from OpenAI-style content parts, rewriting the
  *  parts to the template's {type:"image"} form. */
 export async function extractImages(
@@ -57,7 +98,7 @@ export async function extractImages(
         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 fetchImageBytes(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
@@ -76,83 +117,211 @@ export async function extractImages(
   return { messages: out, images };
 }
 
-async function fetchImageBytes(url: string): Promise<Uint8Array> {
-  if (url.startsWith("data:")) {
-    const comma = url.indexOf(",");
-    if (comma === -1) throw new Error("malformed data: URL");
-    const meta = url.slice(0, comma);
-    const body = url.slice(comma + 1);
-    if (!meta.includes("base64")) throw new Error("data: URL must be base64");
-    return Uint8Array.from(Buffer.from(body, "base64"));
-  }
-  if (url.startsWith("http://") || url.startsWith("https://")) {
-    const res = await fetch(url);
-    if (!res.ok) throw new Error(`image fetch failed: ${res.status} ${url}`);
-    return new Uint8Array(await res.arrayBuffer());
+/** 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 });
   }
-  throw new Error(`unsupported image url scheme: ${url.slice(0, 16)}`);
+  return { messages: out, audio };
 }
 
-export async function buildVisionPrompt<P extends { softTokens: number }>(
+/** Decoded + featurized audio clip, ready to splice. */
+interface PreparedAudio {
+  mel: Float32Array;
+  frames: number;
+  softTokens: number;
+}
+
+export async function buildMultimodalPrompt<P extends { softTokens: number }>(
   model: Gemma4Model,
-  tower: VisionEncoder<P>,
+  towers: MultimodalTowers<P>,
   tokenizer: LoadedTokenizer,
   template: ChatTemplate,
   messages: ChatMessage[],
   images: Uint8Array[],
-  tokenIds: VisionTokenIds,
+  audio: Uint8Array[],
   tools: ToolDefinition[] | null = null,
-): Promise<VisionPrompt> {
+): 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 + embed every image (in order of appearance), via the
-  // tower's own preprocessing (encoder-free vs SigLIP differ here).
-  const pre = await Promise.all(images.map((bytes) => tower.preprocess(bytes)));
+  // 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>
+  // 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 }[] = [];
+  const runs: { start: number; length: number; kind: "image" | "audio"; index: number }[] = [];
   let imgIdx = 0;
+  let audIdx = 0;
   for (const id of ids) {
-    if (id === tokenIds.imageTokenId) {
+    if (vIds && id === vIds.imageTokenId) {
       const p = pre[imgIdx];
       if (!p) throw new Error("more <|image|> markers than images");
-      spliced.push(tokenIds.boiTokenId);
-      runs.push({ start: spliced.length, length: p.softTokens });
-      for (let k = 0; k < p.softTokens; k++) spliced.push(tokenIds.imageTokenId);
-      spliced.push(tokenIds.eoiTokenId);
+      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`);
 
-  // embed text tokens, then overwrite image-token rows with vision features
+  // embed text tokens, then overwrite soft-token rows with tower features
   const idsArr = ops.fromInt32(spliced, [1, spliced.length]);
   let embeds = model.embed.encode(idsArr);
   idsArr.dispose();
   const hidden = embeds.shape[2]!;
-  for (let i = 0; i < runs.length; i++) {
-    const { start, length } = runs[i]!;
-    const feats = tower.features(pre[i]!); // [1, soft, hidden]
-    const cast = feats.astype(embeds.dtype);
-    feats.dispose();
+  for (const run of runs) {
+    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 feats = towers.vision!.tower.features(pre[run.index]!); // [1, soft, hidden]
+      cast = feats.astype(embeds.dtype);
+      feats.dispose();
+    } else {
+      // audio mirrors 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 raw = towers.audio!.tower.features(clip.mel, clip.frames, false); // [1, n, hidden] f32
+      if (raw.shape[1] !== clip.softTokens) {
+        const n = raw.shape[1];
+        raw.dispose();
+        embeds.dispose();
+        throw new Error(`audio tower produced ${n} frames but splice expects ${clip.softTokens}`);
+      }
+      const bf = raw.astype(embeds.dtype);
+      raw.dispose();
+      const scale = ops.scalarLike(towers.audio!.tower.embedScale, bf);
+      cast = ops.div(bf, scale);
+      bf.dispose();
+      scale.dispose();
+    }
+    const { start, length } = run;
     const updated = ops.sliceUpdate(embeds, cast, [0, start, 0], [1, start + length, hidden]);
     cast.dispose();
     embeds.dispose();
     embeds = updated;
   }
 
-  const maskInts = new Int32Array(spliced.length);
-  for (let i = 0; i < spliced.length; i++)
-    maskInts[i] = spliced[i] === tokenIds.imageTokenId ? 1 : 0;
-  const maskI32 = MlxArray.fromInt32(maskInts, [spliced.length]);
-  const imageMask = maskI32.astype(Dtype.bool);
-  maskI32.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();
 
-  return { ids: spliced, embeddings: embeds, imageMask };
+  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 };
+}
+
+/** 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: Gemma4Model,
+  tower: VisionEncoder<P>,
+  tokenizer: LoadedTokenizer,
+  template: ChatTemplate,
+  messages: ChatMessage[],
+  images: Uint8Array[],
+  tokenIds: VisionTokenIds,
+  tools: ToolDefinition[] | null = null,
+): Promise<VisionPrompt> {
+  const mp = await buildMultimodalPrompt(
+    model, { vision: { tower, tokenIds } }, tokenizer, template,
+    messages, images, [], tools,
+  );
+  // 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 };
 }