rlx-inflect-nano

Inflect-Nano-v1 English text-to-speech for RLX — a ~4.6M-parameter two-stage TTS, ported to Rust with full numeric parity to the reference PyTorch implementation and a fully standalone text frontend (no Python at runtime).

Architecture

text ─▶ frontend ─▶ phone/tone/lang ids ─▶ MicroFastSpeech ─▶ 80-bin log-mel ─▶ Snake HiFi-GAN ─▶ 24 kHz wav

• Acoustic (MicroFastSpeech, 3.47M params): phone/tone/lang/speaker embeddings → 5 ConvFFN encoder blocks → duration/energy/bright/pitch prosody heads → length regulation → 6 ConvFFN decoder blocks → bidirectional GRU → mel head → postnet. Fully deterministic (no sampling).
• Vocoder (Snake HiFi-GAN snake_v2mid, 1.16M params): conv_pre → 4× (Snake → ConvTranspose1d → 3 ResBlocks) → Snake → conv_post → tanh. Weight-norm folded at export.
• Frontend (standalone Rust): clean → normalize (a faithful inflect number/time/currency
  • abbreviation port) → BERT WordPiece word grouping → CMUdict → g2p_en neural OOV (single-layer GRU seq2seq, greedy) + nltk averaged-perceptron POS tagger for homographs → phoneme ids → blanks.

Backends

The vocoder (the compute core) builds as an rlx-ir graph compiled per device, so it runs on every RLX backend. 1-D convs are NCHW convs with time in H; the transposed convs are rewritten as zero-insertion + regular conv (so no backend needs a native ConvTranspose kernel); Snake α is folded to constants. Runtime-validated on CPU, Metal, MLX, and wgpu (this Apple-Silicon host); CUDA/ROCm flags are wired but need their toolkits/hardware to build.

CoreML is reached through ONNX Runtime's CoreML execution provider (the workspace's standard CoreML route), not the native rlx-graph compiler. CoreML's MIL requires bounded dims, so the dynamic-axis vocoder.onnx would silently fall back to CPU — the CoreML path instead loads a static-shape vocoder_static.onnx and chunks the mel into fixed-length segments (with overlap; bit-exact stitch), so it genuinely executes on the ANE/GPU (fp16; ≈4% vs the f32 host, inaudible on a tanh-bounded waveform). The ORT session is cached (InflectNano::synthesize_coreml), so the one-time CoreML model compile (~4.5 s) is paid on first call and warm runs are ~1.8× faster than the host CPU pipeline. The acoustic stage stays host-eager like the other backends.

Feature flags mirror the other RLX model crates:

cargo run -p rlx-inflect-nano --features metal -- --device metal --text "..." # Apple GPU
cargo run -p rlx-inflect-nano --features mlx   -- --device mlx   --text "..." # Apple MLX
cargo build -p rlx-inflect-nano --features all-backends                       # wire all

A host-eager (ndarray, matmul-based) CPU path is always available and is the numeric reference. The acoustic stage runs host-eager (it is tiny); synthesize_on(device) runs the vocoder graph on the chosen backend.

Speed (2.77 s utterance, this Apple-Silicon host)

The Rust CPU path (frontend + acoustic + vocoder) is ~1.6× faster than the 10-thread PyTorch reference; the vocoder graph on Metal is far faster again.

Asset bundle

The PyTorch checkpoints + frontend assets are converted to an RLX bundle by the export script (weights/dicts are gitignored; regenerate locally):

python3 -m venv --system-site-packages .venv-inflect
.venv-inflect/bin/pip install g2p_en inflect nltk safetensors
git clone https://huggingface.co/owensong/Inflect-Nano-v1 /tmp/inflect-nano
.venv-inflect/bin/python scripts/export_inflect_nano.py --repo /tmp/inflect-nano --out weights/inflect-nano-rlx

Bundle layout: acoustic.safetensors, vocoder.safetensors, config.json, and frontend/ (CMUdict, g2p_en checkpoint, homographs, perceptron tagger, BERT tokenizer, symbol table).

Samples

Generated by the example below (CPU path, this crate, 24 kHz):

Regenerate (the example writes WAV; ffmpeg -i out.wav -c:a aac out.mp4 to repackage):

cargo run -p rlx-inflect-nano --release --example synthesize -- \
    --data weights/inflect-nano-rlx --text "The weather is nice today, and I feel very relaxed." \
    --out out.wav

Usage

cargo run -p rlx-inflect-nano --release -- --data weights/inflect-nano-rlx \
    --text "The weather is nice today, and I feel very relaxed." --out out.wav

let model = InflectNano::load_from_dir(Path::new("weights/inflect-nano-rlx"))?;
let wav = model.synthesize("Hello, world!", &InferOpts::default())?;
audio::write_wav(Path::new("out.wav"), &wav.samples, wav.sample_rate)?;

Execution modes

synthesize_mode(text, &opts, mode) (CLI --mode …) picks the compute path. The acoustic stage is tiny and always runs host-eager; the mode selects how the vocoder (the compute core) runs:

cargo run -p rlx-inflect-nano --features apple-silicon --release -- \
    --data weights/inflect-nano-rlx --mode hybrid --text "Hello!" --out out.wav

let wav = model.synthesize_mode("Hello!", &InferOpts::default(), ExecutionMode::Latency)?;

Latency/Hybrid fall back to the (fast) CPU path when no GPU backend is compiled in or available, so they are always safe to request.

Real-time streaming (1 s of audio in < 1 s of compute)

synthesize_stream (CLI --stream [--chunk-secs 1.0]) emits the waveform in fixed-length chunks so playback can start after the first chunk and a long utterance sustains real-time with bounded latency. The acoustic model runs once; the vocoder runs per chunk over a small overlap context and the trimmed chunks concatenate bit-identically to a full-utterance vocode (validated in tests/streaming_parity.rs, max-diff 0.0).

cargo run -p rlx-inflect-nano --release -- --data weights/inflect-nano-rlx --stream --chunk-secs 1.0 \
    --text "Streaming one-second chunks in real time." --out out.wav

let report = model.synthesize_stream(text, &opts, 1.0, |chunk| audio_sink.write(chunk))?;
assert!(report.sustains_realtime()); // every 1 s chunk produced in < 1 s

Each 1-second chunk is produced in ~0.26 s on CPU alone (worst-chunk RTF ~3.8×), so the stream stays comfortably real-time with margin to spare (and far more on MLX/Metal). Streamed chunks are the raw tanh-bounded vocoder output — the whole-clip RMS normalization in synthesize needs the full signal and is intentionally skipped when streaming.

Per-mode benchmark

cargo run --features <backend> --release --example bench_modes — 9.27 s utterance, warmed iters, on this Apple-Silicon host. RTF (audio-seconds / wall-seconds), higher is better:

• The acoustic stage always runs host (it's tiny), so latency/hybrid only accelerate the vocoder. MLX gives a ~9× vocoder speedup (48× RTF); Metal ~2–3× (kernel JIT warmup adds variance — best-case ~17×); wgpu's dispatch/readback overhead makes it marginal for a model this small.
• Peak RSS is ~300 MB in every mode (297–302 MB), dominated by the text-frontend dictionaries + model weights rather than the backend — so mode choice barely affects memory here.
• CoreML / ANE: supported via ONNX Runtime's CoreML EP (--features coreml, --device coreml) — not the native rlx-graph compiler (rlx-runtime only has an unimplemented ane stub), so it runs the exported vocoder.onnx rather than the rlx graph.

Parity

Validated against Python reference fixtures (scripts/inflect_nano_reference.py, scripts/inflect_nano_frontend_fixtures.py):

Run the tests (the bundle dir defaults to weights/inflect-nano-rlx, override with RLX_INFLECT_NANO_DATA):

cargo test -p rlx-inflect-nano --release