Nemotron-Labs-Audex-2B — MLX (bf16)
An MLX conversion of NVIDIA's Nemotron-Labs-Audex-2B that runs natively on Apple Silicon. Weights are bf16.
Audex is a unified audio-text model: a single Transformer decoder built on a text-only LLM (Nemotron-Cascade-2) that gains speech understanding and generation without regressing on text intelligence. This repo hosts the converted weights for four capabilities, all on one model:
- Text generation — it is still a capable instruct LLM
- Text-to-speech (TTS) — single-codebook FSQ speech tokens → 16 kHz waveform
- Speech-to-text (ASR / audio understanding) — NV-Whisper encoder projected into the LLM
- Speech-to-speech — talk to it, it replies in speech (cascaded: ASR → text reply → TTS)
A 4-bit build (~2× smaller LM, ~2× faster synthesis) is at
Nemotron-Labs-Audex-2B-MLX-4bit.
Usage
The loader/inference code lives in the open-source repo
akashicMarga/mlx-audio-train under models/audex/.
git clone https://github.com/akashicMarga/mlx-audio-train
pip install mlx mlx-lm mlx-whisper soundfile transformers
huggingface-cli download akashicmarga/Nemotron-Labs-Audex-2B-MLX-bf16 --local-dir audex_mlx
from models.audex import load_model, text_generate, tts_generate, s2s_generate
import soundfile as sf
m = load_model("audex_mlx")
print(text_generate(m, "Explain RoPE in one sentence."))
sf.write("out.wav", tts_generate(m, "Hello from Apple Silicon."), 16000)
reply, wav = s2s_generate(m, "question.wav") # speech in -> speech out
sf.write("reply.wav", wav, 16000)
# CLIs
python scripts/infer_audex.py --model audex_mlx --tts "..." --output out.wav
python scripts/audex_s2s.py --model audex_mlx --input q.wav --output reply.wav
Architecture (as ported)
- LM — NemotronDense: 28-layer decoder, hidden 2048, GQA 16/8, head_dim 128, squared-ReLU MLP, RMSNorm, full RoPE (θ=1e8), vocab 205312 (text + 65536 speech-codec + 8192 audio-codec).
- Speech codec — Audex causal decoder: a single FSQ codebook (one token/frame @ 50 Hz) → 12-layer Vocos transformer → 320 samples/frame @ 16 kHz.
- Audio encoder — NV-Whisper (Qwen2-Audio encoder): 128-mel → 32 layers → 750 embeddings
per 30 s clip, projected into the LLM embedding space and injected at
<so_embedding>slots.
License & attribution
Derived from nvidia/Nemotron-Labs-Audex-2B and distributed under the same
NVIDIA OneWay Noncommercial License — non-commercial use only. See the base model for the
full license and paper (arXiv:2607.05196). All credit for
the model and training to NVIDIA; this repo only provides an MLX weight conversion.