Blame for worker-3b.js

// d1-3B in the browser: Transformers.js (WebGPU, WASM fallback) over onnx-community/d1-3B-ONNX.
import { AutoModelForImageTextToText, AutoProcessor, RawImage, env } from "https://cdn.jsdelivr.net/npm/@huggingface/transformers@4.3.1/dist/transformers.min.js";
import { answer, readoutIds3b, render3b, softmax } from "./prompt.js";

env.allowLocalModels = false;
const REPO = "onnx-community/d1-3B-ONNX";
const post = (m) => self.postMessage(m);
let processor = null, model = null, tok = null, device = "wasm";

async function init() {
  if (navigator.gpu) {
    try { if (await navigator.gpu.requestAdapter()) device = "webgpu"; } catch {}
  }
  const files = {};
  const progress_callback = (p) => {
    if (p.status === "progress" && p.file && p.total) {
      files[p.file] = [p.loaded, p.total];
      let loaded = 0, total = 0;
      for (const [l, t] of Object.values(files)) { loaded += l; total += t; }
      post({ type: "progress", label: "text", file: p.file, loaded, total, aggregate: true });
    }
  };
  processor = await AutoProcessor.from_pretrained(REPO);
  tok = processor.tokenizer;
  model = await AutoModelForImageTextToText.from_pretrained(REPO, {
    device,
    dtype: { embed_tokens: "q4", vision_encoder: "q8", decoder_model_merged: "q4" },
    progress_callback,
  });
  post({ type: "ready", device });
}

async function loadImage(blob) {
  let img = await RawImage.fromBlob(blob);
  img = img.rgb();
  const max = 1024 * 1024; // runner.py `cap_pixels`
  if (img.width * img.height > max) {
    const s = Math.sqrt(max / (img.width * img.height));
    img = await img.resize(Math.max(1, Math.floor(img.width * s)), Math.max(1, Math.floor(img.height * s)), { resample: 3 });
  }
  return img;
}

async function decide({ state, questions, image }) {
  const t0 = performance.now();
  const img = image ? await loadImage(image) : null;
  const answers = {}, probs = {};
  for (const q of questions) {
    const prompt = render3b(tok, state, q, img ? "<image>" : "");
    const inputs = img ? await processor(img, prompt, { add_special_tokens: false }) : tok(prompt, { add_special_tokens: false });
    const { logits } = await model(inputs);
    const [, T, V] = logits.dims;
    const row = logits.data.subarray((T - 1) * V, T * V);
    const scores = readoutIds3b(tok, q).map((g) => Math.max(...g.map((i) => row[i])));
    const p = softmax(scores);
    answers[q.name] = answer(q, p);
    probs[q.name] = p;
  }
  return { answers, probs, ms: performance.now() - t0 };
}

self.onmessage = async ({ data }) => {
  try {
    if (data.type === "init") await init();
    else if (data.type === "decide") post({ type: "result", id: data.id, ...(await decide(data)) });
  } catch (e) {
    post({ type: "error", id: data.id, message: String(e && e.message ? e.message : e) });
  }
};