mlx-bun 0.5.0 → 1.0.0

package/src/vision/embedder.ts

@@ -1,176 +0,0 @@
-// Vision feature pipeline — port of optiq's gemma4_unified VisionEmbedder
-// + MultimodalEmbedder (vendored from mlx-vlm, BSD-3).
-//
-// Encoder-free: patches [n, 6912] → LayerNorm → Linear(+bias) → LayerNorm
-// → + 2D pos embedding → LayerNorm → RMSNormNoScale → Linear(no bias)
-// → divide by embed_scale (the language model re-multiplies it).
-//
-// Sidecar weights stay bf16; features come out bf16 like text embeddings.
-
-import { tensorFingerprint } from "../model/fingerprint";
-import { ptr, read } from "bun:ffi";
-import { MlxArray, cpuStream } from "../mlx/array";
-import { C } from "../mlx/ffi";
-import * as ops from "../mlx/ops";
-import type { PreprocessedImage } from "./preprocess";
-import { NUM_SOFT_TOKENS, preprocessImage } from "./preprocess";
-
-const cstr = (s: string) => Buffer.from(s + "\0", "utf8");
-
-export class VisionTower {
-  #weights = new Map<string, MlxArray>();
-  #cacheIdentity?: string;
-  /** Immutable encoder weights and computation settings, hashed once on use. */
-  get cacheIdentity(): string {
-    return this.#cacheIdentity ??= new Bun.CryptoHasher("sha256")
-      .update("gemma-vision-v1").update(JSON.stringify([this.embedScale, this.rmsNormEps]))
-      .update(tensorFingerprint(this.#weights)).digest("hex");
-  }
-  readonly embedScale: number;
-  readonly rmsNormEps: number;
-
-  private constructor(embedScale: number, rmsNormEps: number) {
-    this.embedScale = embedScale;
-    this.rmsNormEps = rmsNormEps;
-  }
-
-  static load(modelDir: string, embedScale: number, rmsNormEps = 1e-6): VisionTower {
-    const self = new VisionTower(embedScale, rmsNormEps);
-    // out-param slots read back via read.u64, not [0] (DFG stale-read bug
-    // — see outArray in mlx/ffi.ts).
-    const arrMap = new BigUint64Array([C.mlx_map_string_to_array_new()]);
-    const metaMap = new BigUint64Array([C.mlx_map_string_to_string_new()]);
-    const arrMapPtr = ptr(arrMap);
-    const metaMapPtr = ptr(metaMap);
-    const status = C.mlx_load_safetensors(
-      arrMapPtr, metaMapPtr, ptr(cstr(`${modelDir}/optiq_vision.safetensors`)), cpuStream,
-    );
-    C.mlx_map_string_to_string_free(read.u64(metaMapPtr, 0));
-    if (status !== 0) throw new Error(`failed to load vision sidecar from ${modelDir}`);
-    const arrMapHandle = read.u64(arrMapPtr, 0);
-    const names = [
-      "vision_embedder.patch_ln1.weight", "vision_embedder.patch_ln1.bias",
-      "vision_embedder.patch_dense.weight", "vision_embedder.patch_dense.bias",
-      "vision_embedder.patch_ln2.weight", "vision_embedder.patch_ln2.bias",
-      "vision_embedder.pos_embedding",
-      "vision_embedder.pos_norm.weight", "vision_embedder.pos_norm.bias",
-      "embed_vision.embedding_projection.weight",
-    ];
-    for (const name of names) {
-      const slot = new BigUint64Array([C.mlx_array_new()]);
-      const slotPtr = ptr(slot);
-      if (C.mlx_map_string_to_array_get(slotPtr, arrMapHandle, ptr(cstr(name))) !== 0)
-        throw new Error(`vision sidecar missing tensor ${name}`);
-      self.#weights.set(name, new MlxArray(read.u64(slotPtr, 0)));
-    }
-    C.mlx_map_string_to_array_free(arrMapHandle);
-    return self;
-  }
-
-  async preprocess(bytes: Uint8Array): Promise<PreprocessedImage> {
-    return preprocessImage(bytes);
-  }
-
-  #w(name: string): MlxArray {
-    return this.#weights.get(name)!;
-  }
-
-  /** nn.Linear: x @ W.T (+ bias). */
-  #linear(x: MlxArray, weight: MlxArray, bias: MlxArray | null): MlxArray {
-    const wT = ops.transposeAxes(weight, [1, 0]);
-    let out = ops.matmul(x, wT);
-    wT.dispose();
-    if (bias) {
-      const withBias = ops.add(out, bias);
-      out.dispose();
-      out = withBias;
-    }
-    return out;
-  }
-
-  /** One image's preprocessed patches → language-space features
-   *  [softTokens, hidden] (pre-divided by embed_scale). */
-  features(img: PreprocessedImage): MlxArray {
-    const eps = this.rmsNormEps;
-    const patchDim = img.patches.length / NUM_SOFT_TOKENS;
-    let h: MlxArray = MlxArray.fromFloat32(img.patches, [1, NUM_SOFT_TOKENS, patchDim]);
-
-    h = disposing(h, ops.layerNorm(h,
-      this.#w("vision_embedder.patch_ln1.weight"),
-      this.#w("vision_embedder.patch_ln1.bias"), eps));
-    h = disposing(h, this.#linear(h,
-      this.#w("vision_embedder.patch_dense.weight"),
-      this.#w("vision_embedder.patch_dense.bias")));
-    h = disposing(h, ops.layerNorm(h,
-      this.#w("vision_embedder.patch_ln2.weight"),
-      this.#w("vision_embedder.patch_ln2.bias"), eps));
-
-    // + 2D position embedding: posemb[clamped[:,0], 0] + posemb[clamped[:,1], 1]
-    // (clamp + validity mask precomputed host-side; padding rows get -1)
-    const n = NUM_SOFT_TOKENS;
-    const xIdx = new Int32Array(n);
-    const yIdx = new Int32Array(n);
-    const valid = new Float32Array(2 * n);
-    for (let i = 0; i < n; i++) {
-      const px = img.positions[i * 2]!;
-      const py = img.positions[i * 2 + 1]!;
-      xIdx[i] = Math.max(px, 0);
-      yIdx[i] = Math.max(py, 0);
-      valid[i] = px === -1 ? 0 : 1;
-      valid[n + i] = py === -1 ? 0 : 1;
-    }
-    const posemb = this.#w("vision_embedder.pos_embedding"); // [1120, 2, 3840]
-    const addPos = (idx: Int32Array, dim: number, validOff: number): MlxArray => {
-      const idxArr = MlxArray.fromInt32(idx, [n]);
-      const rows = ops.takeAxis(posemb, idxArr, 0); // [n, 2, 3840]
-      idxArr.dispose();
-      const sel = rows.slice([0, dim, 0], [n, dim + 1, 3840]);
-      rows.dispose();
-      const flat = ops.reshape(sel, [1, n, 3840]);
-      sel.dispose();
-      const v = MlxArray.fromFloat32(valid.subarray(validOff, validOff + n).slice(), [1, n, 1]);
-      const vCast = v.astype(flat.dtype);
-      v.dispose();
-      const masked = ops.mul(flat, vCast);
-      flat.dispose();
-      vCast.dispose();
-      return masked;
-    };
-    const xPos = addPos(xIdx, 0, 0);
-    const yPos = addPos(yIdx, 1, n);
-    const posSum = ops.add(xPos, yPos);
-    xPos.dispose();
-    yPos.dispose();
-    h = disposing(h, ops.add(h, posSum));
-    posSum.dispose();
-
-    h = disposing(h, ops.layerNorm(h,
-      this.#w("vision_embedder.pos_norm.weight"),
-      this.#w("vision_embedder.pos_norm.bias"), eps));
-
-    // MultimodalEmbedder: RMSNormNoScale → Linear(no bias)
-    h = disposing(h, ops.rmsNorm(h, null, eps));
-    h = disposing(h, this.#linear(h,
-      this.#w("embed_vision.embedding_projection.weight"), null));
-
-    // keep only real (unpadded) soft tokens; pre-divide by embed_scale
-    const hidden = h.shape[2]!;
-    const real = h.slice([0, 0, 0], [1, img.softTokens, hidden]);
-    h.dispose();
-    const scale = ops.scalarLike(this.embedScale, real);
-    const out = ops.div(real, scale);
-    real.dispose();
-    scale.dispose();
-    return out; // [1, softTokens, hidden]
-  }
-
-  dispose(): void {
-    for (const w of this.#weights.values()) w.dispose();
-    this.#weights.clear();
-  }
-}
-
-function disposing(old: MlxArray, next: MlxArray): MlxArray {
-  old.dispose();
-  return next;
-}

package/src/vision/preprocess.ts

@@ -1,247 +0,0 @@
-// 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 };
-}