ARCHITECTURE.md

Architecture: Streaming Speech Translation Pipeline

Overview

End-to-end streaming speech translation pipeline:

Audio Input → ASR → NMT → TTS → Audio Output
  (PCM16)   (ONNX) (GGUF) (ONNX)  (PCM16)

Translates spoken English into spoken Russian in real-time with streaming output.

Pipeline Stages

1. ASR — Cache-Aware Streaming ASR (NeMo Conformer RNN-T)

Model: NVIDIA NeMo Conformer RNN-T, exported to ONNX.

Architecture: 3 internal threads connected by queue.Queue:

Audio → [Preprocess Thread] → [Encoder Thread] → [Decoder Thread] → text deltas

• Preprocess: Buffers raw audio, extracts mel-spectrogram features, chunks them per CacheAwareStreamingConfig (chunk_size=[49,56], shift_size=[49,56])
• Encoder: Runs ONNX encoder inference on feature chunks, maintains encoder cache state
• Decoder: Runs RNN-T decoder (joint network), produces incremental text tokens

Key classes: CacheAwareStreamingAudioBuffer, CacheAwareStreamingASR, ASRModelPackage Wrapper: StreamingASR — exposes push_audio_chunk() / get_transcript_chunk() API

2. NMT — Streaming Segmented Translation (TranslateGemma)

Model: TranslateGemma 4B (GGUF, Q8_0) via llama-cpp-python.

Architecture: Single-threaded, three internal components:

text deltas → [Segmenter] → text segments → [Translator] → raw translations → [Merger] → display text

• StreamingSegmenter: Batches ASR tokens into word-groups (max 5 words + 2 hold-back). Triggers on punctuation, pause (>700ms), or max-token boundaries (min 3 words)
• StreamingTranslator: Multi-turn translation using init/continuation prompt templates with KV cache warming
• StreamingTranslationMerger: Handles revision/append/continuation logic for incremental translations. Detects trailing ellipsis (incomplete), leading ellipsis (continuation), and word-level LCP revision

Wrapper: StreamingNMT — exposes push_text_chunk() / flush() / check_pause() API

3. TTS — Streaming XTTS v2 (ONNX)

Model: XTTSv2 with ONNX-exported GPT-2 AR model + HiFi-GAN vocoder.

Architecture: Sequential within a single call:

text → [BPE Tokenizer] → [GPT-2 AR Loop] → mel latents → [HiFi-GAN Vocoder] → audio chunks

• Speaker conditioning: One-time compute from reference audio → gpt_cond_latent [1,32,1024] + speaker_embedding [1,512,1]
• AR generation: GPT-2 autoregressive loop producing audio token latents. Accumulates stream_chunk_size (default 20) tokens before running vocoder
• Vocoder: HiFi-GAN converts accumulated latents to waveform
• Crossfade stitching: Linear fade-in/fade-out between consecutive vocoder output chunks for seamless playback

Output: 24kHz float32 audio chunks Wrapper: StreamingTTS — exposes synthesize_stream() generator API

Concurrency Model

┌─────────────────────────────────────────────────────────────────┐
│                     asyncio Event Loop                          │
│                                                                 │
│  ┌──────────────┐    ┌──────────────┐    ┌──────────────────┐  │
│  │ WebSocket I/O │    │ asr_to_nmt   │    │ tts_synthesis    │  │
│  │   handler     │    │   loop       │    │    loop          │  │
│  └──────┬───────┘    └──────┬───────┘    └────────┬─────────┘  │
│         │                   │                      │            │
│         │ run_in_executor() │ run_in_executor()    │ run_in_ex… │
│         ▼                   ▼                      ▼            │
│  ┌──────────────────────────────────────────────────────────┐  │
│  │              ThreadPoolExecutor (4 workers)               │  │
│  │  ┌──────────────┐ ┌──────────────┐ ┌──────────────────┐  │  │
│  │  │ ASR internal  │ │ NMT blocking │ │ TTS blocking     │  │  │
│  │  │ threads (3)   │ │ llama-cpp    │ │ ONNX inference   │  │  │
│  │  └──────────────┘ └──────────────┘ └──────────────────┘  │  │
│  └──────────────────────────────────────────────────────────┘  │
└─────────────────────────────────────────────────────────────────┘

• asyncio loop: Handles WebSocket I/O and coordinates pipeline stages
• ASR threads: 3 dedicated daemon threads (preprocess, encoder, decoder)
• NMT: Blocking llama-cpp inference bridged via run_in_executor()
• TTS: Blocking ONNX inference bridged via run_in_executor()
• Concurrency: All per-session state (including NMT turn context: prev_translation, prev_query) is held in per-session wrapper objects. The shared model weights (ASRModelPackage, StreamingTranslator, StreamingTTSPipeline) hold no mutable per-session state after initialization, making concurrent sessions safe.

Data Flow & Queues

WebSocket binary → push_audio()
                      │
                      ▼
              ┌──────────────┐
              │  audio_queue  │ (maxsize=256, queue.Queue)
              └──────┬───────┘
                     ▼
            [ASR Internal Threads]
                     │
                     ▼
              ┌──────────────┐
              │ output_queue  │ (maxsize=64, queue.Queue)
              └──────┬───────┘
                     ▼
         [asr_to_nmt_loop via executor]
                     │
              ┌──────┴───────┐
              ▼              ▼
     transcript_queue   tts_text_queue (maxsize=16, asyncio.Queue)
     (→ WebSocket)           │
                             ▼
              [tts_synthesis_loop via executor]
                             │
                             ▼
                    ┌──────────────┐
                    │audio_out_queue│ (maxsize=32, asyncio.Queue)
                    └──────┬───────┘
                           ▼
                   WebSocket binary

Backpressure Strategy

• Audio input: put_nowait with drop-on-full (acceptable to lose frames vs. building latency)
• ASR→NMT: put_nowait with drop warning on encoder/decoder output queues
• NMT→TTS: put_nowait with drop warning (translations can be reconstructed from next segment)
• TTS→Output: put_nowait with drop warning per audio chunk
• All queue sizes are configurable via PipelineConfig

Model Loading & Session Lifecycle

Startup: Models loaded ONCE in TranslationServer._load_models():

• ASR ONNX sessions (ASRModelPackage)
• NMT GGUF model (StreamingTranslator) + KV cache warmup
• TTS ONNX sessions (StreamingTTSPipeline)

Per-session: Each WebSocket connection creates:

• StreamingASR — own audio buffers, streaming state, thread pool
• StreamingNMT — own segmenter, merger, and translation context (prev_translation, prev_query); shares model weights only
• StreamingTTS — own speaker conditioning; shares ONNX sessions

Cleanup: On disconnect, orchestrator flushes remaining NMT text through TTS, then stops all threads and resets state.

WebSocket Protocol

Direction	Type	Format	Description
Client→	Binary	PCM16	Raw audio at declared sample rate
Client→	Text	JSON	{"action": "start", "sample_rate": 16000}
Client→	Text	JSON	{"action": "stop"}
→Client	Binary	PCM16	Synthesized audio at 24kHz
→Client	Text	JSON	{"type": "transcript", "text": "..."}
→Client	Text	JSON	{"type": "translation", "text": "..."}
→Client	Text	JSON	{"type": "status", "status": "started"}

Configuration

All tunables are in PipelineConfig (dataclass) and exposed as CLI args:

Parameter	Default	Description
asr_chunk_duration_ms	10	Audio chunk duration for ASR
nmt_n_threads	4	CPU threads for llama-cpp
tts_stream_chunk_size	20	AR tokens per vocoder chunk
audio_queue_maxsize	256	Audio input queue bound
tts_queue_maxsize	16	NMT→TTS text queue bound
audio_out_queue_maxsize	32	TTS→output audio queue bound

File Structure

src/
├── asr/
│   ├── streaming_asr.py          # StreamingASR wrapper
│   ├── pipeline.py               # ThreadedSpeechTranslator (reference)
│   ├── cache_aware_modules.py    # Audio buffer + streaming ASR
│   ├── modules.py                # ONNX model loading
│   └── utils.py                  # Audio utilities
├── nmt/
│   ├── streaming_nmt.py          # StreamingNMT wrapper
│   ├── streaming_segmenter.py    # Word-group segmentation
│   ├── streaming_translation_merger.py  # Translation merging
│   └── translator_module.py      # TranslateGemma via llama-cpp
├── tts/
│   ├── streaming_tts.py          # StreamingTTS wrapper
│   ├── xtts_streaming_pipeline.py  # Full TTS pipeline
│   ├── xtts_onnx_orchestrator.py   # GPT-2 AR + vocoder
│   ├── xtts_tokenizer.py          # BPE tokenizer
│   └── zh_num2words.py            # Chinese text normalization
├── pipeline/
│   ├── orchestrator.py           # PipelineOrchestrator
│   └── config.py                 # PipelineConfig
└── server/
    └── websocket_server.py       # WebSocket server