repository·main·Indexed 21 days ago
https://github.com/trevors/voxtral-mini-realtime-rsA pure Rust implementation of Mistral's Voxtral Mini 4B Realtime (ASR) and Voxtral 4B TTS models using the Burn ML framework. It supports high-performance inference via WebGPU/WGPU natively and in the browser via WASM. The project includes a CLI for streaming speech recognition and text-to-speech synthesis, supporting both BF16 and Q4 GGUF quantized models.
The repository is organized as follows:
src/lib.rs: The public API entry point providing VoxtralRealtime<B>.src/models/: Contains the core model components (encoder.rs, decoder.rs, adapter.rs, voxtral.rs) and the VoxtralModelLoader.src/audio/: Handles audio processing, including MelSpectrogram extraction, resampling, and chunking.src/tokenizer/: A wrapper around the Tekken tokenizer.scripts/: Python utilities for weight inspection, reference generation, and validation.web/: Contains the browser demo, including worker.js for off-main-thread inference and voxtral-client.js for the high-level API.To prevent merge conflicts during concurrent development, the src/tts/mod.rs file acts as the central module root. It declares all submodules upfront. When adding new functionality, ensure the corresponding module is declared in this file:
pub mod config;
pub mod backbone;
pub mod voice;
pub mod embeddings; // flow-matching
pub mod sequence; // flow-matching
pub mod flow_matching; // flow-matching
pub mod codec; // integration
pub mod pipeline; // integrationWhen performing autoregressive streaming inference, you must use a prefix length of 38 instead of the standard 39. Using 39 tokens causes an anomaly at position 38 that results in the model outputting only padding tokens.
Correct Prefix Pattern:
[1] (BOS) followed by 37 [STREAMING_PAD] (token 32) tokens.Token Reference:
[STREAMING_PAD] (token 32): Marks pauses.[STREAMING_WORD] (token 33): Starts words.token_id - 1000).# Correct prefix for autoregressive generation
prefix_tokens = [1] + [32] * 37 # BOS + 37 STREAMING_PAD = 38 tokensThe Flow-Matching Transformer predicts acoustic tokens using an Euler ODE solver. During inference, the model performs 8 steps (dt = 1/8) to transition from noise to the target state.
Inference Steps:
x_1 ~ N(0, 1) in $\mathbb{R}^{36}$.t in [1.0, 0.875, 0.75, ..., 0.125]:v_cond = FM(x_t, t, h)v_uncond = FM(x_t, t, zeros)v = 1.2 * v_cond + (1 - 1.2) * v_uncondx_{t-dt} = x_t - v * dtx_0 to 21 FSQ levels per dimension.Transformer Inputs per frame:
h (projected via llm_projection [3072, 3072])t (projected via time_projection [3072, 3072])x_t (36-dim, projected via input_projection [3072, 36])To enable streaming in the audio encoder, you must use causal attention instead of bidirectional attention. For optimized streaming, combine causal attention with a sliding window constraint to limit the attention span.
Causal Mask: Ensures position i can only attend to positions 0..=i.
Sliding Causal Mask: Ensures position i can only attend to positions within a specific window size behind it.
fn create_causal_mask(seq_len: usize, device: &Device) -> Tensor {
let mut mask = vec![f32::NEG_INFINITY; seq_len * seq_len];
for i in 0..seq_len {
for j in 0..=i {
mask[i * seq_len + j] = 0.0;
}
}
Tensor::from_vec(mask, (seq_len, seq_len), device)
}
fn create_sliding_causal_mask(seq_len: usize, window: usize, device: &Device) -> Tensor {
let mut mask = vec![f32::NEG_INFINITY; seq_len * seq_len];
for i in 0..seq_len {
let start = i.saturating_sub(window);
for j in start..=i {
mask[i * seq_len + j] = 0.0;
}
}
Tensor::from_vec(mask, (seq_len, seq_len), device)
}The language model is a decoder-only transformer based on Ministral-3B that processes audio embeddings.
embed_tokens = lm_head.T)To reduce KV cache memory, the model uses a 4:1 GQA ratio:
Voxtral Mini 4B Realtime is a pure Rust implementation of a streaming ASR (Automatic Speech Recognition) model built using the Burn ML framework. The model consists of three primary components:
Key Timing Metric: 1 text token corresponds to approximately 80ms of audio (1280 samples @ 16kHz).
Finite Scalar Quantization (FSQ) is used to map continuous 36-dimensional vectors to discrete levels. This is applied after the Euler ODE solver to prepare data for the codec decoder.
Quantization Logic:
[-1, 1].idx = argmin(|x - levels|) per dimension, producing 36 integer indices in the range [0..20].x = levels[idx] (used for codec decoder input).This process is purely arithmetic and uses no learned parameters.
The audio encoder is a causal transformer designed for streaming mel spectrograms.
clamp(log(max(mel, 1e-10)) / 1.5, -1.0, 1.0).The audio codebook embedding layer maps semantic and acoustic token indices to a single 3072-dim vector by summing all 37 embeddings (1 semantic + 36 acoustic).
Index Layout:
EMPTY_AUDIO=0, END_AUDIO=1).Index Arithmetic:
semantic_global = raw_semantic_idx + 2acoustic_global = 8194 + cb * 23 + (level + 2)