How krafton-ai/raon-speech works
1. What this repo is
- Raon-Speech is a bilingual speech-language model codebase with two product lines:
- Raon-Speech: offline SpeechLM for STT, TTS, TextQA, and SpeechChat
- Raon-SpeechChat: full-duplex conversational speech model for offline and realtime use
- Both share the same core package under src/raon/.
2. Core model boundary
- The main model implementation lives in src/raon/models/raon.py and src/raon/models/wrapper.py.
- The architecture centers on a shared LLM backbone plus speech-specific components: audio encoding/tokenization, speaker conditioning, code prediction, and audio decoding.
- The readme describes the trainable boundary as the alignment/adaptation stack around the frozen or partially frozen speech/audio modules.
3. Shared module stack
- src/raon/modules/ contains the reusable speech building blocks:
- audio_encoder, audio_tokenizer
- speaker_encoder
- code_predictor, concurrent_audio_decoder
- adaptor, embedding
- voxtral_encoder, voxtral_wrapper
- This is the main separation between generic language modeling and speech-specific I/O.
4. SpeechLM offline flow
- Offline usage flows from JSONL samples plus audio files into src/raon/pipeline.py or the CLI entrypoint src/raon/generate.py.
- The pipeline exposes task methods such as stt, tts, textqa, and speech_chat, and can also save generated audio.
- config/infer.yaml and scripts/infer.sh define the default inference wiring.
5. Full-duplex flow
- The duplex track uses src/raon/duplex_generate.py, src/raon/duplex_train.py, and config/duplex_infer.yaml.
- It consumes stereo dialogue audio plus timing annotations and produces interaction-aware output for turn-taking and interruption handling.
- The duplex path is architecturally distinct because it needs conversation-state handling, prompt selection, and timeline-based supervision.
6. Training subsystem
- src/raon/train.py handles SpeechLM training, while src/raon/duplex_train.py handles duplex training.
- Training is launched through shell wrappers in scripts/train.sh and scripts/duplex_train.sh, typically with torchrun for multi-GPU execution.
- src/raon/utils/loss.py, training_callbacks.py, delay.py, and state_machine.py support optimization and sequence/control logic.
7. Data and supervision
- Sample datasets are shipped in data/speechllm/ and data/duplex/.
- SpeechLM data is JSONL-based with fields like conversations, audios, speaker_ref_audios, and channel.
- Duplex data uses stereo wavs plus timing annotations via turns or scripts, along with metadata such as speak_first, include_in_training, language, and channel.
- This makes the data boundary very explicit: text-audio conversation records for offline tasks, and timestamped dual-channel dialogue records for duplex.
8. Realtime demo architecture
- demo/realtime/ is a separate runtime stack for the duplex demo.
- It is split into:
- api/ for the FastAPI app
- backends/sglang_backend.py for model serving
- runtime/ for session, config, prompt mapping, and artifacts
- protocol/messages.py for structured runtime messages
- web/ for browser-side streaming controls
- This demo is not just a UI; it is a small serving system with session/state orchestration and backend abstraction.
9. Serving/export boundary
- scripts/export.sh converts a Hugging Face checkpoint into an SGLang-friendly bundle for realtime serving.
- The realtime duplex demo depends on that exported artifact, plus sglang, gradio, fastapi, and uvicorn.
- This creates a clear deployment split between offline HF-style inference and low-latency interactive serving.
10. External integration and runtime assumptions
- The repo is built around the Hugging Face ecosystem: AutoModel.from_pretrained(..., trust_remote_code=True) is a first-class loading path.
- PyTorch and Torchaudio are the core runtime dependencies; CUDA and bfloat16/float16 are expected for practical performance.
- FlashAttention is optional and selected via attention backend flags, which affects training/inference performance but not the logical architecture.
11. Utilities and preprocessing
- src/raon/utils/ contains the support layer for audio I/O, mel features, data handling, prompt catalogs, special tokens, processor logic, and duplex-specific dataset helpers.
- These utilities sit between raw files and the model, handling normalization, feature extraction, prompt construction, and output formatting.
12. How to draw the architecture
- A good diagram should show four major zones:
- Data layer: packaged JSONL + wav samples
- Core model layer: shared Raon backbone and speech modules
- Workflows: train, offline inference, duplex inference
- Serving/demo layer: Gradio/FastAPI/SGLang realtime stack
- The key architectural idea is that SpeechLM and SpeechChat are two task families built on one shared model stack, with separate training/inference/runtime paths on top.
Main entry points: raon.py Raon model · app.py Realtime API · duplex_generate.py Duplex infer · pipeline.py Offline infer
Components 21
Data layer
Packaged speech and duplex samples used for training and evaluation.
- SpeechLM data · Offline SpeechLM training and evaluation samples for STT, TTS, TextQA, and SpeechChat.
data/speechllm - Duplex data · Full-duplex dialogue samples with stereo audio, turn timing, and supervision metadata.
data/duplex
Core model
Shared Raon backbone plus speech-specific modules and utilities.
- Raon model · Main model implementation that combines the LLM backbone with speech I/O and adaptation logic.
src/raon/models/raon.py - Model wrapper · Wrapper around model construction, loading, and runtime-specific adaptation.
src/raon/models/wrapper.py - Audio stack · Audio encoding and tokenization boundary for speech input and representation.
src/raon/modules/audio_encoder.py - Speaker encoder · Speaker reference encoding used to condition speech generation and dialogue control.
src/raon/modules/speaker_encoder.py - Code predictor · Predicts intermediate speech codes and coordinates concurrent audio decoding.
src/raon/modules/code_predictor.py - Audio decoder · Decodes predicted codes into audio output, including concurrent decoding paths.
src/raon/modules/concurrent_audio_decoder.py - Adaptation stack · Shared adaptation and embedding modules that bridge language and speech components.
src/raon/modules/adaptor.py - Voxtral bridge · Wrapper and encoder integration for Voxtral-style speech/language interoperability.
src/raon/modules/voxtral_wrapper.py
Workflows
Training and inference entrypoints for offline and duplex paths.
- Offline train · SpeechLM training entrypoint for the offline product line.
src/raon/train.py - Duplex train · Training entrypoint for turn-aware duplex supervision and conversation-state learning.
src/raon/duplex_train.py - Offline infer · Offline generation pipeline exposing STT, TTS, TextQA, and SpeechChat flows.
src/raon/pipeline.py - Generate CLI · Command-line entrypoint for offline inference and audio generation.
src/raon/generate.py - Duplex infer · Full-duplex inference path with prompt selection, turn handling, and timeline-based output.
src/raon/duplex_generate.py - Export bundle · Converts a checkpoint into a serving-friendly artifact for realtime deployment.
src/raon/export.py
Serving demo
Realtime duplex serving stack with API, backend, runtime, and web controls.
- Realtime API · HTTP service layer for the duplex realtime demo.
demo/realtime/api/app.py - SGLang backend · Serving backend used by the realtime duplex demo.
demo/realtime/backends/sglang_backend.py - Runtime state · Session, config, prompt mapping, and artifact coordination for realtime interactions.
demo/realtime/runtime/session.py - Message proto · Structured runtime message definitions for streaming and orchestration.
demo/realtime/protocol/messages.py - Web controls · Browser-side streaming and stop controls for the realtime demo.
demo/realtime/web/gradio_stream.js
Connections 21
- SpeechLM data → Offline train: offline samples · JSONL and audio samples feed SpeechLM training.
- SpeechLM data → Offline infer: eval inputs · Evaluation examples drive offline inference flows.
- Duplex data → Duplex train: duplex supervision · Stereo audio and timing annotations train the duplex model.
- Duplex data → Duplex infer: dialogue inputs · Duplex inference consumes turn-aware dialogue records and prompts.
- Offline train → Raon model: optimizes · Offline training updates the shared model boundary.
- Duplex train → Raon model: optimizes · Duplex training updates the shared model with interaction-aware supervision.
- Offline infer → Model wrapper: loads via · Offline inference uses the wrapper for model construction and runtime setup.
- Generate CLI → Offline infer: invokes · CLI entrypoint drives the offline pipeline.
- Duplex infer → Model wrapper: loads via · Duplex inference uses the same model loading boundary.
- Raon model → Audio stack: uses · The core model consumes speech audio through shared encoders and tokenizers.
- Raon model → Speaker encoder: uses · Speaker conditioning informs generation and dialogue control.
- Raon model → Code predictor: uses · Intermediate code prediction bridges language state to audio output.
- Code predictor → Audio decoder: feeds · Predicted codes are converted into audio by the decoder stack.
- Raon model → Adaptation stack: uses · Adaptation and embedding modules mediate speech-language alignment.
- Raon model → Voxtral bridge: integrates · Voxtral integration provides an alternate speech-language bridge.
- Export bundle → SGLang backend: prepares · Exported checkpoints are packaged for low-latency serving.
- SGLang backend → Realtime API: serves · The backend powers the realtime API layer.
- Realtime API → Runtime state: orchestrates · API requests are coordinated through session and runtime state.
- Runtime state → Message proto: formats · Runtime messaging uses structured protocol objects.
- Realtime API → Web controls: streams to · API responses drive browser-side streaming controls.
- Duplex infer → Realtime API: same domain · Duplex inference and realtime demo share the interaction-oriented serving boundary.