shreyask/Gemma-4-WebGPU

1// web/src/worker.js2 3import {4  AutoProcessor,5  Gemma4ForConditionalGeneration,6  TextStreamer,7  InterruptableStoppingCriteria,8  load_image,9  RawImage,10} from "@huggingface/transformers";11 12const MODEL_ID = "onnx-community/gemma-4-E2B-it-ONNX";13 14const THINK_START = "‹‹THINK››";15const THINK_END = "‹‹/THINK››";16 17function cleanGemmaOutput(raw) {18  return raw19    .replace(/<\|?channel\|?>?\s*thought\s*/gi, THINK_START)20    .replace(/<\|?channell?\|?>/gi, THINK_END)21    .replace(/<\|?[a-z_]+\|?>/gi, "")22    .trim();23}24 25 26let processor = null;27let model = null;28 29const stoppingCriteria = new InterruptableStoppingCriteria();30 31async function checkWebGPU() {32  try {33    const adapter = await navigator.gpu?.requestAdapter();34    self.postMessage({35      type: "status",36      status: adapter ? "webgpu-available" : "webgpu-unavailable",37    });38  } catch {39    self.postMessage({ type: "status", status: "webgpu-unavailable" });40  }41}42 43async function loadModel() {44  try {45    self.postMessage({ type: "status", status: "loading" });46 47    const progress_callback = (p) => self.postMessage({ type: "progress", ...p });48 49    processor = await AutoProcessor.from_pretrained(MODEL_ID, { progress_callback });50 51    model = await Gemma4ForConditionalGeneration.from_pretrained(MODEL_ID, {52      dtype: "q4f16",53      device: "webgpu",54      progress_callback,55    });56 57    self.postMessage({ type: "status", status: "ready" });58  } catch (err) {59    self.postMessage({ type: "error", message: err.message });60  }61}62 63async function generate({ messages, imageUrl, videoData, audioData, enableThinking }) {64  if (!model || !processor) {65    self.postMessage({ type: "error", message: "Model not loaded" });66    return;67  }68 69  try {70    self.postMessage({ type: "status", status: "generating" });71    stoppingCriteria.reset();72 73    const prompt = processor.apply_chat_template(messages, {74      enable_thinking: enableThinking,75      add_generation_prompt: true,76    });77 78    // Gemma4ImageProcessor expects RawImage | RawImage[], not RawVideo79    let image = null;80    if (videoData) {81      image = videoData.frames.map((f) =>82        new RawImage(new Uint8ClampedArray(f.data), f.width, f.height, f.channels)83      );84    } else if (imageUrl) {85      image = await load_image(imageUrl);86    }87    const audio = audioData ?? null;88 89    const inputs = await processor(prompt, image, audio, {90      add_special_tokens: false,91    });92 93    let fullText = "";94    const streamer = new TextStreamer(processor.tokenizer, {95      skip_prompt: true,96      skip_special_tokens: false,97      callback_function: (text) => {98        fullText += text;99        const cleaned = cleanGemmaOutput(fullText);100        self.postMessage({ type: "update", text: cleaned });101      },102    });103 104    await model.generate({105      ...inputs,106      max_new_tokens: 512,107      do_sample: false,108      streamer,109      stopping_criteria: [stoppingCriteria],110    });111 112    self.postMessage({ type: "complete", text: cleanGemmaOutput(fullText) });113  } catch (err) {114    self.postMessage({ type: "error", message: err.message });115  }116}117 118self.onmessage = (e) => {119  switch (e.data.type) {120    case "check":121      checkWebGPU();122      break;123    case "load":124      loadModel();125      break;126    case "generate":127      generate(e.data);128      break;129    case "interrupt":130      stoppingCriteria.interrupt();131      break;132  }133};134