Add PersonaPlex / Moshi full-duplex S2S support (Mimi codec + Moshi LM + ORT example) - #368
Merged
Conversation
Add support for the Kyutai Mimi neural audio codec used by
nvidia/personaplex-7b-v1, the first phase of full-duplex
speech-to-speech (Moshi) support.
What:
- src/mobius/models/mimi.py: MimiModel building two ONNX graphs via
CodecTask — encoder (waveform -> 8 RVQ codes @12.5Hz) and decoder
(codes -> 24kHz waveform). Reuses CodecEncoderTransformerModel and
SplitResidualVectorQuantizer; adds the SEANet encoder/decoder,
learnt downsample/upsample, and the encoder-side argmin quantizer.
Module nesting reproduces Kyutai weight names so conv/transformer
weights need no renames; preprocess_weights handles the fused QKV
split, interleaved->half-split RoPE permutation, and codebook table
reconstruction from embedding_sum/cluster_usage.
- src/mobius/integrations/moshi/: native Kyutai checkpoint loader
(build_mimi) mirroring the NeMo integration; resolves the
tokenizer-*.safetensors file from a local path or HF Hub.
- scripts/generate_mimi_golden.py + testdata/golden/audio/
mimi-personaplex.json: tiny JSON golden (exact codes + decoded
slices) from a deterministic in-code waveform — no committed audio.
- Tests: L1 graph-build tests and an integration parity test.
Why the LayerNorm eps change:
- CodecEncoderTransformerLayer now takes layer_norm_eps (default 1e-6,
unchanged for existing callers). Mimi passes 1e-5 to match Kyutai's
create_norm_fn("layer_norm"); the 1e-6 vs 1e-5 mismatch produced a
~0.004 per-element error in the encoder transformer that cascaded to
wrong RVQ codes (masked on the decoder side by LayerScale 0.01).
Parity vs the Kyutai moshi reference: encoder codes match exactly;
decoded waveform matches to ~2e-7 (atol=1e-4).
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
…PersonaPlex P2+P3)
Adds end-to-end support for nvidia/personaplex-7b-v1, the Kyutai Moshi
full-duplex speech-to-speech architecture, building on the Mimi codec
export landed in P1.
P2 — Moshi LM:
* models/moshi.py: MoshiTemporalModel (dim=4096, 32L, RoPE, sliding-window
causal, SwiGLU via FusedGateUpMLP) and MoshiDepformerModel (dim=1024, 6L,
no RoPE, full causal, weights_per_step=16 per-substep linears, one-substep
graph selected by substep_index). Both verified to exact argmax parity
against the reference (CPU exact; CUDA exact with use_tf32=0).
* tasks/_moshi.py: MoshiTemporalTask (input_frame[B,17,S] + KV -> hidden +
text_logits + KV) and MoshiDepformerTask (hidden + prev_token +
substep_index + KV -> logits + KV).
* integrations/moshi: build_moshi_lm() native Kyutai-format loader returning
{"temporal", "depformer"} ModelPackages.
* L1 build tests, committed golden (testdata/golden/audio/
moshi-lm-personaplex.json) + generator script, and an integration parity
test (CPU-only; H200 TF32 flips greedy sampling).
P3 — runtime example:
* examples/personaplex_moshi.py: faithful NumPy port of Kyutai LMGen.step
(ring cache, per-codebook delays, temporal step, greedy text sampling,
16-substep autoregressive depformer, delayed output collection) driving
four ONNX Runtime sessions (Mimi encoder/decoder + temporal + depformer)
for full-duplex S2S generation. Verified end to end (5 frames -> 3
assistant frames -> decoded waveform).
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Performance Comparison
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
| _TR_HEAD_DIM = 64 | ||
| _TR_FFN = 2048 | ||
| _TR_THETA = 10000.0 | ||
| _TR_CONTEXT = 250 |
| _ROPE_THETA = 10000.0 | ||
|
|
||
| # --- Token vocabularies --- | ||
| _NUM_CH = 17 # 1 text + 16 audio codebooks |
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
Pull request overview
Adds end-to-end ONNX graph construction, native-checkpoint loading, parity tests, and an ONNX Runtime driver example for the PersonaPlex / Kyutai Moshi full-duplex speech-to-speech stack (Mimi codec + Moshi temporal LM + depformer).
Changes:
- Introduces Mimi codec (encoder/decoder) and Moshi LM (temporal + depformer) ONNXScript model implementations and Moshi-specific task builders.
- Adds native Kyutai/HF-checkpoint builder utilities (
mobius.integrations.moshi) plus an ONNX Runtime example that runs the full generation loop. - Adds build-graph unit tests and CPU integration parity tests backed by committed JSON goldens + golden regeneration scripts.
Reviewed changes
Copilot reviewed 16 out of 16 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
tests/moshi_lm_integration_test.py |
CPU integration parity test for Moshi temporal + depformer against committed golden. |
tests/moshi_integration_test.py |
CPU integration parity test for Mimi encoder/decoder against committed golden. |
tests/build_graph_test.py |
Adds graph-build tests for Mimi and Moshi task I/O contracts. |
testdata/golden/audio/moshi-lm-personaplex.json |
Committed Moshi LM parity golden (hidden/logit slices + argmax). |
testdata/golden/audio/mimi-personaplex.json |
Committed Mimi codec parity golden (codes + decoded waveform slices). |
src/mobius/tasks/_moshi.py |
New MoshiTemporalTask and MoshiDepformerTask graph builders with KV cache wiring. |
src/mobius/tasks/__init__.py |
Exposes Moshi tasks and registers task-name mappings. |
src/mobius/models/moshi.py |
Implements Moshi temporal + depformer ONNXScript models and weight preprocessors. |
src/mobius/models/mimi.py |
Implements Mimi codec ONNXScript models, SEANet, RVQ encode, and weight preprocessing. |
src/mobius/models/__init__.py |
Exports Mimi/Moshi models and config helpers from the public models package. |
src/mobius/integrations/moshi/_builder.py |
Adds native-checkpoint resolution + build/apply-weights helpers for Mimi and Moshi LM. |
src/mobius/integrations/moshi/__init__.py |
Public integration entrypoints (build_mimi, build_moshi_lm). |
src/mobius/components/_codec_transformer.py |
Adds layer_norm_eps plumbed into Mimi-style encoder transformer layers. |
scripts/generate_moshi_lm_golden.py |
Script to regenerate Moshi LM golden from the Kyutai reference. |
scripts/generate_mimi_golden.py |
Script to regenerate Mimi golden from the Kyutai reference. |
examples/personaplex_moshi.py |
ORT example driving Mimi + Moshi temporal + depformer for full-duplex S2S. |
The graph optimizer prunes the unused temporal `position_ids` input on the fp16/CUDA build (RoPE derives its offset from the KV-cache length), and the KV cache dtype follows the exported model dtype. Read the declared inputs and KV dtype from the session so the loop works across fp32/fp16 builds and both CPU/CUDA EPs, instead of always feeding position_ids and float32 KV (which crashes the CUDA path). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
…xample
Adds full-duplex real-time generation on top of the offline loop:
* StreamingMimiEncoder / StreamingMimiDecoder: rolling-window wrappers around
the whole-utterance Mimi graphs that encode/decode one 12.5Hz frame at a
time using the codec's causal convolutions (keep only the newest frame).
* --stream: simulated real-time from a wav (or silence) that reports per-frame
compute, RTF, and how many frames exceed the 80ms budget; optional --pace
sleeps to emulate a live 12.5Hz source. Writes assistant_stream.wav.
* --mic: live microphone -> speaker full-duplex via sounddevice (best-effort;
needs audio hardware), with a clean error when unavailable.
* --lm-dtype {f32,f16}: build the LM in fp16 (Mimi stays fp32 due to its fp16
Conv export bug) for real-time speed on CUDA.
* MoshiORT.warmup(): runs a few frames to trigger CUDA/codec autotune then
resets LM state, removing the first-frame stall (audio glitch).
Verified on H200 (fp16 LM + fp32 Mimi, use_tf32=0): RTF=0.63, per-frame
mean=51ms / p90=57ms / max=65ms, 0/60 frames over the 80ms budget -> real-time
full-duplex confirmed. CPU fp32 stays ~3.3s/frame (offline only).
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
build_mimi(dtype="f16") produced a graph that failed to load in ORT with "Type parameter (T) of Optype (Conv) bound to different types (float and float16)": the encoder's first Conv received the float32 `waveform` input while its weights were float16. Keep the ONNX I/O contract as float32 PCM (what external audio consumers feed/expect) and cast to the compute dtype inside the model: encoder casts the input waveform float32 -> dtype; decoder casts its output waveform dtype -> float32. No-op for fp32 builds. Verified on H200 (CUDA, use_tf32=0): both encoder/decoder now load with float32 I/O; fp16 decoder on fp32 codes vs fp32 reference = 71 dB SNR (near-lossless); encoder VQ argmax agrees 92.7% with fp32 (boundary flips). fp32 codec golden parity and L1 tests unchanged. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Move the ORT example into examples/personaplex/ and add a real-time browser client so the model can be driven over a port-forwarded WebSocket. - examples/personaplex/moshi_ort.py (moved from examples/personaplex_moshi.py): add MoshiORT.reset_stream() + process_frame() so the streaming sessions can be reused with resettable per-conversation state; refactor --stream/--mic to use the shared process_frame() front door. - examples/personaplex/server.py: aiohttp WebSocket server that serves the static page and streams full-duplex 24 kHz float32 frames to/from a single pre-built MoshiORT (no mobius import, runs in onnxruntime-gpu envs). Single- user lock, warmup on connect, per-frame budget reporting. - examples/personaplex/static/index.html: browser client capturing mic at 24 kHz mono into 1920-sample frames and playing back assistant frames via a small jitter-buffered AudioBufferSourceNode queue. - examples/personaplex/README.md: build / serve / SSH port-forward guide. Verified headless on H200: static page serves (HTTP 200), WebSocket round-trip returns assistant audio frames in real time, single-user busy lock works. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Port 8080 is commonly occupied; switch the server default (and README examples) to 7681 to avoid bind conflicts on shared hosts. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
The streaming step() forced the text codebook to ZERO_TEXT_CODE (pad) every frame and marked it provided=True, so the model's generated text was always discarded. In Moshi the text stream drives the assistant's speech, so forcing it to pad made the model emit silence forever (observed out_rms stuck at the ~2e-4 noise floor regardless of input). Match the Kyutai LMGen reference instead: - Do not force/provide the text token; let the model generate it. Only honour an explicitly supplied text_token (e.g. to force-feed a transcript). - Replace greedy argmax with temperature + top-k sampling (Kyutai defaults: text temp=0.7/top_k=25, audio temp=0.8/top_k=250). Greedy on the text stream keeps picking the pad token, which also suppresses speech. Verified offline (CUDA, fp16 LM + fp32 Mimi): assistant out_rms now reaches ~0.11 peak with ~6% non-silent frames from silence input (the model self-initiates speech), versus a constant ~2e-4 before. Also: - server.py: log per-second input/output RMS so silence vs speech is visible. - static/index.html: resample mic <-> 24 kHz when the browser AudioContext does not honour the 24 kHz hint (otherwise the model gets wrong-rate audio). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Each assistant frame is 80 ms but the server's per-frame latency jitters (occasionally >200 ms under GPU/network load), so the previous 50 ms playback buffer underran and clicked. Maintain ~250 ms of buffered audio with a contiguous schedule head, only rebuilding the buffer on a true underrun, and surface the buffered depth / rebuffer count in the status line. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
PersonaPlex supports conditioning the assistant on a reference voice and a persona/system-prompt before the conversation (LMGen.step_system_prompts). This wires that into the browser demo: - moshi_ort.py: add SINE_TOKENS; extend step() with a moshi_tokens arg to teacher-force the assistant audio stream (k=1..8); add MoshiORT.prime() that runs the 4-phase priming (voice prompt -> silence -> persona text -> silence); add SentencePiece persona tokenizer helpers (load_persona_tokenizer, encode_persona with <system> ... <system> wrapping). - server.py: new handshake -- reply 'config', receive persona JSON + optional binary voice PCM, tokenize + Mimi-encode + prime(), then 'ready'; load the tokenizer once (optional, --tokenizer to point at a local model). - static/index.html: persona textarea, ~6s voice recorder (resampled to 24kHz), Start session / Restart session handshake before live streaming. - README: document persona + voice customization and the new deps. Verified headless on H200 (GPU2): config -> persona+voice -> primed in ~2s -> ready -> non-silent assistant audio; steady-state RTF 0.72 (58 ms/frame). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
| d = json.loads(cfg.data) | ||
| persona = (d.get("persona") or "").strip() or DEFAULT_PERSONA | ||
| expect_voice = bool(d.get("hasVoice")) | ||
| except (ValueError, TypeError): |
The voice prompt can now come from three sources, all normalized to a 24 kHz mono reference clip before priming: - Record ~6 s from the mic (existing) - Upload an audio file (wav/mp3/flac/…); decoded via Web Audio decodeAudioData, resampled to 24 kHz, capped to ~10 s - Pick a server-side preset from static/voices/ (new /voices endpoint lists the clips; the browser fetches and decodes the chosen file) index.html unifies all three through setVoice()/decodeToVoice(); server.py adds the /voices listing endpoint and serves the clips via the existing /static route. No presets are bundled (PersonaPlex voices.tgz are .pt embeddings, incompatible with token forcing); drop clips into static/voices/ to populate the dropdown. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Two issues made the 'Connecting…' button fail to become 'Start session': - The server only sent 'config' after warmup (a few seconds), so the button looked stuck. Now it sends 'config' immediately on connect and warms up in the background, joining before priming — the UI unlocks right away and warmup overlaps with the user choosing a persona/voice. - On a page refresh the new WebSocket races the old connection's lock release and gets 'busy', leaving the button stuck. The client now auto-retries on 'busy' (up to 10×, 1 s apart) so it recovers once the stale connection frees the single-user lock. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Add a stats panel to the browser client and stream metrics from the server:
- server.py: every ~1 s pushes a 'stats {json}' message with mean compute
ms/frame, compute RTF (vs the 80 ms budget), in/out RMS, frame count, and
over-budget count.
- index.html: a stats grid shows jitter-buffer depth (ms) and rebuffer count
(updated locally per playback frame) plus the server's compute ms/frame, RTF,
mic/out RMS, frames, and over-budget frames. Values turn amber/red as they
approach or exceed the real-time budget so glitches are easy to diagnose.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
| n_voice = 0 if voice_pcm is None else voice_pcm.size // FRAME_SIZE | ||
| print(f"[ws] priming: persona={len(text_tokens or [])} toks, voice={n_voice} frames") | ||
| # Make sure warmup + reset finished before we prime / generate. | ||
| await warm_task |
The per-frame AudioBufferSourceNode scheduling stuttered under network/compute jitter (each 80 ms frame scheduled against a moving playHead clicks/gaps when a frame arrives late). Replace it with an AudioWorklet that pulls from a ring buffer every render quantum and fills silence on underrun -- gapless and immune to main-thread/scheduling jitter. - Inline 'pcm-player' AudioWorklet (Blob module): ring buffer, pre-buffers 'prime' seconds, re-primes on full drain, reports depth/underruns at ~10 Hz. - Add a 'Playback buffer' slider (60-400 ms, default 180) so the user can trade latency vs smoothness live; setPrime updates without clearing the buffer. - Stats panel jitter-buffer/rebuffer rows now driven by the worklet. Compute stays ~55 ms/frame (RTF 0.69) on H200; this targets the playback path, which was the stutter source. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Separate the perf/jitter stats from the live status line so they no longer visually compete. Stats now live in a collapsible 'Live metrics' panel; the status line keeps only connection/session state and the live indicator dot. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
The browser voice-chat demo stuttered progressively as the conversation grew: per-frame compute climbed from ~60ms to 200ms+ over ~1-2 minutes, crossing the 80ms real-time budget after ~20s. Root cause: the temporal transformer's KV cache was fed as numpy each frame, so ORT copied the entire (unbounded, growing) cache host->device every step -- an O(N) cost that dominates as the cache lengthens (~0.5MB/frame; ~0.5GB host copy at ctx=1000). GQA compute itself is only ~6ms; the host copy was the whole problem. Fix: keep the temporal KV cache resident on the inference device via ORT IO binding. present.* outputs are bound to device and reused as the next frame's past.* inputs (double-buffering); the cache never touches host memory. hidden/text_logits still come back to host for the depformer and sampling. Per-frame temporal cost is now flat (~6ms) regardless of conversation length. Measured end-to-end (H200, fp16 LM + fp32 Mimi, 600 frames): per-frame total stays flat at ~49ms (was 60->200ms), p90 53ms, RTF ~0.61 -- well under the 80ms budget for the full session. Note: a sliding-window alternative (re-export with position_ids, non-GQA op.Attention) was prototyped but rejected -- standard opset-23 Attention on CUDA is ~13x slower than the fused GroupQueryAttention contrib op, so it was a net perf regression. IO binding alone removes the real bottleneck. The device cache grows ~0.5MB/frame (fine for multi-minute demos; very long sessions would need a periodic reset+re-prime). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
3 tasks
justinchuby
commented
Jun 22, 2026
- server.py/README: default --host to 127.0.0.1 (localhost) instead of 0.0.0.0, matching the documented SSH port-forward workflow and avoiding binding to all interfaces by default. - Fix ruff warnings across the Moshi branch: D205 docstring summary/blank line (mimi.py), RUF005 list concatenation -> unpacking (golden script + integration test). Import sort/format auto-applied by lintrunner. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds end-to-end ONNX support for
nvidia/personaplex-7b-v1, the Kyutai Moshi full-duplex speech-to-speech architecture, built declaratively withonnxscript.nnfrom the native Kyutaisafetensorscheckpoints (no HFconfig.json).Three phases, one PR:
P1 — Mimi neural codec (commit
0fc84c1)SEANet encoder/decoder + split-RVQ + codec transformer. Encoder codes exact, decode ~2e-7 vs reference. Root cause fixed: Kyutai LayerNorm eps is
1e-5(mobius default was1e-6).P2 — Moshi LM (temporal + depformer)
models/moshi.py:FusedGateUpMLP, RMSNorm eps=1e-8.weights_per_step=16per-substep linears; emitted as a one-substep graph selected bysubstep_index(Gather), embedding select viaWhere(noIf), looped 16× externally.tasks/_moshi.py:MoshiTemporalTask(input_frame[B,17,S]+ KV → hidden + text_logits + KV) andMoshiDepformerTask(hidden+prev_token+substep_index+ KV → logits + KV).integrations/moshi:build_moshi_lm()native loader →{"temporal", "depformer"}ModelPackages.use_tf32=0— H200 TF32 otherwise flips greedy sampling). Committed golden + generator script; integration parity test (CPU-only) passes.P3 — ONNX Runtime example
examples/personaplex_moshi.py: faithful NumPy port of KyutaiLMGen.step— ring cache with per-codebook delays[0,0,1,1,1,1,1,1,1,0,1,1,1,1,1,1,1], temporal step, greedy text sampling, 16-substep autoregressive depformer, delayed output collection — driving four ORT sessions (Mimi encoder/decoder + temporal + depformer) for full-duplex S2S. Verified end to end (5 input frames → 3 assistant frames → decoded waveform).Testing
tests/moshi_lm_integration_test.py, ~280s, CPU).assistant.wav).Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com