MedGemma 4B-it with images: ONNX for Transformers.js

ONNX export of google/medgemma-4b-it with image input, for Transformers.js in the browser. The 4-bit (q4) files use fp32 math, so they run on GPUs without shader-f16 support.

Files

part q4 (recommended) fp32 input -> output
text embeddings onnx/embed_tokens_q4.onnx - input_ids -> inputs_embeds
vision encoder (SigLIP + projector) onnx/vision_encoder_q4.onnx onnx/vision_encoder.onnx (~1.7 GB) pixel_values -> image_features
language model onnx/decoder_model_merged_q4.onnx - inputs_embeds + KV cache -> logits

Each image becomes 256 image tokens (896x896 input). The q4 vision encoder's image features have a mean cosine similarity of 0.967 to the PyTorch model on a test chest X-ray. For the most faithful image reading, load the fp32 vision encoder instead (dtype: { vision_encoder: "fp32", ... }); it is a larger download (~1.7 GB).

Usage (browser)

import { AutoProcessor, Gemma3ForConditionalGeneration, RawImage, TextStreamer } from "@huggingface/transformers";

const id = "geeek/medgemma-4b-it-ONNX";
const processor = await AutoProcessor.from_pretrained(id);
const model = await Gemma3ForConditionalGeneration.from_pretrained(id, {
  device: "webgpu",
  dtype: { embed_tokens: "q4", vision_encoder: "q4", decoder_model_merged: "q4" },
});

const image = await RawImage.read("https://upload.wikimedia.org/wikipedia/commons/c/c8/Chest_Xray_PA_3-8-2010.png");
const messages = [{ role: "user", content: [{ type: "image" }, { type: "text", text: "Describe this X-ray." }] }];
const prompt = processor.apply_chat_template(messages, { add_generation_prompt: true });
const inputs = await processor(prompt, image);

const output = await model.generate({
  ...inputs,
  max_new_tokens: 256,
  do_sample: false,
  streamer: new TextStreamer(processor.tokenizer, { skip_prompt: true }),
});
console.log(processor.batch_decode(output.slice(null, [inputs.input_ids.dims.at(-1), null]), { skip_special_tokens: true })[0]);

Text-only chat also works with the same repo: pipeline("text-generation", id, { device: "webgpu", dtype: "q4" }) loads only the embedding and language model files.

Download on first load: about 3.2 GB with all-q4 files (cached afterwards). Needs roughly 4–5 GB of free GPU memory.

How it was made

  1. The language model is the browser-verified q4 text model previously published in this repo (4-bit weights, fp32 math, no float64, constant Sqrt folded), split in the ONNX graph into the embedding lookup (embed_tokens) and the rest (decoder_model_merged), so the decoder takes embeddings and image features can be inserted.
  2. The vision tower and multimodal projector were exported with torch.onnx in fp32, then quantized to 4-bit (MatMulNBits, block size 32) with ONNX Runtime and made browser-compatible (float64 -> float32).
  3. Verified: generation with ONNX Runtime (CPU) for text and for an image, image features compared to PyTorch, and session creation of every file with onnxruntime-web (WASM).

Limitations

  • In the original model, image tokens attend to each other in both directions. Transformers.js runs the decoder with causal attention only, so image understanding can be somewhat weaker than the original.
  • 4-bit quantization can reduce answer quality.
  • MedGemma is a developer model, not a medical device; outputs are not clinical advice and must be verified independently. See the original model card.

License

Governed by the Health AI Developer Foundations terms of use, the same terms as the original MedGemma model.

Downloads last month
939
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for geeek/medgemma-4b-it-ONNX

Quantized
(56)
this model