jax-realtime

A real-time, full-duplex voice assistant that runs entirely in your browser on WebGPU, built with jax-js.

Every stage (speech → ASR → LLM → TTS → speech, plus optional vision) runs locally in the tab; nothing is sent to a server.

It's inspired by the Thinking Machines interaction model and GPT-Live, rebuilt as a small-model cascade that fits in a browser tab. The goal is a conversation that feels live: you can interrupt it mid-sentence, pause mid-thought without losing your turn, and it keeps searching in the background while you talk.

Everything shares the single WebGPU device. The streaming ASR lane is paused while the assistant speaks so it doesn't contend with TTS for the GPU; barge-in is therefore energy-based (see below), and captions resume the moment the assistant stops.

Interaction

• Full-duplex micro-turns: a ~150 ms tick loop drives a deterministic, priority-ordered policy: adaptive barge-in (talk over the assistant, including its tool narrations, and the audio cuts in ~300 ms; the threshold auto-calibrates to the echo floor of each reply), adaptive endpointing, and time-awareness timers. A watchdog force-recovers the session if a reply ever stalls, so it can't wedge.
• Continuation-merge: if the endpoint fires on a mid-thought pause and you resume speaking before the reply's first audio, the unheard reply is aborted and both halves are answered as one turn ("append, don't restart").
• Learned turn signal: a pure-TypeScript port of Silero VAD v5 scores P(speech) every 32 ms on the CPU and drives speech onset, silence tracking, and the phantom-turn guard: keyboard noise and ambient swells never even latch an utterance, near-silence never reaches Whisper (so no hallucinated "thank you"s), and quiet speech still passes. A repetition-degeneracy gate drops decoder loops on top.
• Eye (vision): enabled by default for webcam context, with a pre-load toggle to skip its 42 MB model, camera access, and GPU residency. D-FINE runs low-priority object detection (it yields the GPU to audio), smooths the person count, and answers "what do you see?" / "how many people?" / "tell me about the person" directly from the measurements. Proactive interjections (stepped away, phone spotted, slouching) are best-effort rule heuristics. The webcam shows as a corner PiP with detection boxes.
• Typed conversation memory: bounded facts the user explicitly states (name, trip, pet, favorite, plans, relationships) are retained and injected only when relevant; exact recall bypasses small-model guessing.
• Two-tier tools: factual asks are delegated so the small on-device model isn't left guessing: weather ("what's the weather in Tokyo" → open-meteo, in °F/mph), facts ("who is Ada Lovelace" → Wikipedia), plus instant offline calculator and clock/date. Web lookups speak a holding line and fetch in the background, then answer on the next silence and render a card; the card clears when the conversation moves on, and the spoken answer can be interrupted like any reply.

Performance (all on jax-js / WebGPU)

The turn-latency floor is set by the single GPU, so the work went into cutting GPU cost per token/frame rather than overlapping stages (which a single device can't do; see docs/BENCHMARKS.md for the full map-reduce campaign log, including the negative results):

• Fused decode: the LLM decode step is fused from dozens of per-layer jit dispatches into one, and Pocket TTS from ~11 into two, cutting the command-buffer submit overhead that dominated per-step cost (~22% each).
• GPU top-k sampling: the LLM samples from a device-side top-64 (one small readback) instead of transferring the full vocab logits every token, folded into the fused step's single dispatch.
• Stable prefill shapes: every turn has a different prompt length, which otherwise forces jax-js to compile new traces mid-conversation. Prompts use 256-token buckets, the common buckets are warmed during loading, and the KV cache has one fixed capacity, so long conversations never hit a multi-second recompile stall mid-turn.
• Faster confidence-aware ASR: timestamp-gate candidate reductions are reused for confidence scoring instead of scanning the vocabulary again. ASR runs 5–7% faster while preserving all 21 paired clean/quiet/distorted transcripts; low-confidence failures request a repeat before invoking the LLM.
• Bounded history window: the prompt is capped at 8 messages so it never leaves the warm prefill buckets: first-token latency stays flat (~500–610 ms) across long sessions instead of doubling around turn 6, and the worst turn of a 12-turn session sits within ~150 ms of the median.
• Deterministic memory fast paths: exact recall and bounded trip, pet, and activity follow-ups can answer in a few milliseconds without model generation.
• Smaller download: Whisper ships a per-row int8 build (73 MB instead of 144 MB), while the TTS checkpoint omits 35 MB of audio-encoder weights never used for synthesis. The SmolLM2 brain stays full fp16 for conversation quality.

Runtime behaviour is tunable at src/tunables.ts (read live, so A/B experiments don't need a rebuild).

Hard-won details

• The assistant used to interrupt itself. The barge-in threshold calibrated its echo floor during the reply's first ticks, which fall in the silent LLM/TTS latency gap, so the floor was ambient-level and the assistant's own playback echo tripped it. Calibration now only counts ticks where the TTS output analyser confirms audio is actually playing.
• Whisper invents "Thank you." from silence. Every Whisper size does this (silent-outro captions in its training data), and it decodes with high confidence, so no model swap or confidence gate catches it. The fix is signal-side: unvoiced audio never reaches the decoder, on every endpoint path, including barge-in continuations, which originally skipped the guard.
• A mid-sentence TTS cut sounds like a full stop. Each chunk is synthesized as a complete utterance, so flushing at a bare word boundary produces sentence-final falling intonation plus a pause in the middle of your sentence. Chunks now only split at real punctuation: a measured latency cost, paid for prosody.

Run it

npm install
npm run dev

Open http://localhost:5173 in a WebGPU-capable browser (Chrome/Edge on desktop, Safari 26+). Click Load models: ~1.0 GB on first run, cached in OPFS afterwards. The Eye is on by default and requests camera access; uncheck it before loading to skip it. Then press the orb once and just talk, hands-free. Press it again to end.

Weights come from Hugging Face: SmolLM2 at full fp16 for conversation quality, Whisper as a per-row int8 build dequantized at load (it transcribes the bench suite identically to fp16).

How it works

The pipeline stages, from microphone to speaker:

License

MIT. Model inference code is adapted from the jax-js repository by Eric Zhang (MIT licensed); model weights remain under their respective licenses.