Build a FastVLM WebGPU Image App

import { checkBrowserCapability } from './capabilities';
import type { WorkerResponse } from './protocol';

const input = document.querySelector<HTMLInputElement>('#image')!;
const prompt = document.querySelector<HTMLTextAreaElement>('#prompt')!;
const preview = document.querySelector<HTMLImageElement>('#preview')!;
const status = document.querySelector<HTMLElement>('#status')!;
const answer = document.querySelector<HTMLElement>('#answer')!;
const run = document.querySelector<HTMLButtonElement>('#run')!;
const cancel = document.querySelector<HTMLButtonElement>('#cancel')!;
const sample = document.querySelector<HTMLButtonElement>('#sample')!;
const exportButton = document.querySelector<HTMLButtonElement>('#export')!;
let worker: Worker | null = null;
let image = '';
let busy = false;
let inputVersion = 0;
let result: Extract<WorkerResponse, { type: 'result' }> | null = null;
let resultPrompt = '';

function setBusy(value: boolean) {
  busy = value;
  input.disabled = prompt.disabled = sample.disabled = value;
  run.disabled = value || !image || !prompt.value.trim();
  cancel.hidden = !value;
  exportButton.disabled = value || !result;
}
function stop() { worker?.terminate(); worker = null; setBusy(false); }

async function selectImage(blob: Blob) {
  const version = ++inputVersion;
  result = null; image = ''; preview.hidden = true;
  answer.textContent = ''; status.textContent = '';
  setBusy(true);
  const url = URL.createObjectURL(blob);
  try {
    const decoded = new Image(); decoded.src = url; await decoded.decode();
    if (version !== inputVersion) return;
    const scale = Math.min(1, 1024 / Math.max(decoded.width, decoded.height));
    const canvas = document.createElement('canvas');
    canvas.width = Math.max(1, Math.round(decoded.width * scale));
    canvas.height = Math.max(1, Math.round(decoded.height * scale));
    const context = canvas.getContext('2d')!;
    context.fillStyle = 'white'; context.fillRect(0, 0, canvas.width, canvas.height);
    context.drawImage(decoded, 0, 0, canvas.width, canvas.height);
    image = canvas.toDataURL('image/png'); preview.src = image; preview.hidden = false;
  } catch { status.textContent = 'Cannot read this image. Choose a JPEG, PNG or WebP.'; }
  finally { URL.revokeObjectURL(url); if (version === inputVersion) setBusy(false); }
}

input.onchange = () => {
  const file = input.files?.[0]; input.value = '';
  if (!file) return;
  if (!['image/jpeg','image/png','image/webp'].includes(file.type) || file.size > 10 * 1024 * 1024) {
    image = ''; preview.hidden = true; result = null; answer.textContent = ''; setBusy(false);
    status.textContent = 'Use a JPEG, PNG or WebP of at most 10 MB.'; return;
  }
  void selectImage(file);
};
sample.onclick = async () => {
  const version = ++inputVersion;
  setBusy(true);
  try {
    const response = await fetch('/sample.png'); if (!response.ok) throw new Error('Sample unavailable');
    const blob = await response.blob(); if (version !== inputVersion) return;
    await selectImage(blob);
  } catch { if (version === inputVersion) { setBusy(false); status.textContent = 'Sample unavailable. Choose a local image.'; } }
};
prompt.oninput = () => setBusy(busy);
cancel.onclick = () => { inputVersion++; stop(); status.textContent = 'Cancelled. Run again when ready.'; };

run.onclick = async () => {
  if (busy || !image || !prompt.value.trim()) return;
  result = null; answer.textContent = ''; setBusy(true);
  status.textContent = 'Checking this browser…';
  const version = ++inputVersion;
  const capability = await checkBrowserCapability();
  if (version !== inputVersion) return;
  if (!capability.supported) { setBusy(false); status.textContent = 'WebGPU with shader-f16 is required. See the Python guide below.'; return; }
  resultPrompt = prompt.value.trim();
  try {
    worker ??= new Worker(new URL('./worker.ts', import.meta.url), { type: 'module' });
    worker.onmessage = ({ data }: MessageEvent<WorkerResponse>) => {
      if (version !== inputVersion) return;
      if (data.type === 'loading') status.textContent = 'Downloading/preparing model…';
      if (data.type === 'progress') status.textContent = `Model files read: ${data.loadedMB} MB; preparing model…`;
      if (data.type === 'running') status.textContent = 'Reading image…';
      if (data.type === 'text') answer.textContent = data.text;
      if (data.type === 'result') {
        result = data; answer.textContent = data.text; setBusy(false);
        status.textContent = `Preparation: ${data.timing.load_ms} ms · First text: ${data.timing.first_text_ms ?? '—'} ms · Image to answer: ${data.timing.inference_ms} ms`;
      }
      if (data.type === 'error') { status.textContent = `Failed during ${data.phase}: ${data.message}`; stop(); }
    };
    worker.onerror = () => { if (version === inputVersion) { status.textContent = 'Worker failed. Retry or use local Python.'; stop(); } };
    worker.postMessage({ image, prompt: resultPrompt });
  } catch { status.textContent = 'Could not start the worker.'; stop(); }
};
exportButton.onclick = () => {
  if (!result) return;
  const url = URL.createObjectURL(new Blob([JSON.stringify({ ...result, prompt: resultPrompt }, null, 2)], { type: 'application/json' }));
  const link = document.createElement('a'); link.href = url; link.download = 'fastvlm-result.json'; link.click();
  setTimeout(() => URL.revokeObjectURL(url), 1000);
};
void checkBrowserCapability().then(value => {
  document.querySelector('#capability')!.textContent = value.supported ? 'WebGPU and shader-f16 available. Free GPU memory is also required.' : 'This browser does not meet the WebGPU / shader-f16 requirements. See the Python guide below.';
});
window.addEventListener('pagehide', () => { inputVersion++; stop(); });

import { AutoModelForImageTextToText, AutoProcessor, RawImage, TextStreamer, env } from '@huggingface/transformers';
import type { LlavaProcessor, PreTrainedModel, Tensor } from '@huggingface/transformers';
import type { RunRequest, WorkerResponse } from './protocol';

env.allowLocalModels = false;
const modelId = 'onnx-community/FastVLM-0.5B-ONNX';
let model: PreTrainedModel | null = null;
let processor: LlavaProcessor | null = null;
let busy = false;
const send = (message: WorkerResponse) => self.postMessage(message);

self.onmessage = async ({ data }: MessageEvent<RunRequest>) => {
  if (busy) return;
  busy = true;
  let phase: 'load' | 'inference' = 'load';
  const loadStart = performance.now();
  try {
    if (!model || !processor) {
      send({ type: 'loading' });
      const files = new Map<string, number>();
      let lastMB = -1;
      const progress_callback = (info: { status: string; file?: string; loaded?: number }) => {
        if (info.status !== 'progress' || !info.file || typeof info.loaded !== 'number') return;
        files.set(info.file, info.loaded);
        const loadedMB = Math.floor([...files.values()].reduce((sum, bytes) => sum + bytes, 0) / 1_000_000);
        if (loadedMB !== lastMB) { lastMB = loadedMB; send({ type: 'progress', loadedMB }); }
      };
      processor = await AutoProcessor.from_pretrained(modelId, { progress_callback }) as LlavaProcessor;
      model = await AutoModelForImageTextToText.from_pretrained(modelId, {
        progress_callback,
        device: 'webgpu',
        dtype: { embed_tokens: 'fp16', vision_encoder: 'q4', decoder_model_merged: 'q4' },
      });
    }
    const loadMs = Math.round(performance.now() - loadStart);
    send({ type: 'loaded', duration: loadMs });
    phase = 'inference';
    send({ type: 'running' });
    const start = performance.now();
    const image = await RawImage.fromURL(data.image);
    const prompt = processor.apply_chat_template([
      { role: 'system', content: 'You are a helpful visual assistant. Answer the question accurately and concisely.' },
      { role: 'user', content: `<image>${data.prompt}` },
    ], { add_generation_prompt: true });
    const inputs = await processor(image, prompt, { add_special_tokens: false });
    let firstText: number | null = null;
    let streamed = '';
    const streamer = new TextStreamer(processor.tokenizer!, {
      skip_prompt: true,
      skip_special_tokens: true,
      callback_function: (chunk: string) => {
        if (chunk && firstText === null) firstText = performance.now() - start;
        streamed += chunk;
        send({ type: 'text', text: streamed });
      },
    });
    const output = await model.generate({ ...inputs, max_new_tokens: 192, do_sample: false, repetition_penalty: 1.2, streamer }) as Tensor;
    const inputLength = inputs.input_ids.dims.at(-1)!;
    const generated = output.slice(null, [inputLength, null]);
    const text = processor.batch_decode(generated, { skip_special_tokens: true })[0].trim();
    send({ type: 'result', model: modelId, text, timing: {
      load_ms: loadMs,
      inference_ms: Math.round(performance.now() - start),
      first_text_ms: firstText === null ? null : Math.round(firstText),
      generated_tokens: generated.dims.at(-1) ?? null,
    } });
  } catch (error) {
    // No input or raw exception text leaves the worker in telemetry.
    send({ type: 'error', phase, message: error instanceof Error ? error.message : 'Inference failed' });
    if (phase === 'load') { model = null; processor = null; }
  } finally { busy = false; }
};