From afa2b096846dcca282db3b048a3a499c7e393034 Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Sun, 26 Apr 2026 22:26:26 -0400 Subject: [PATCH] feat: add MOSS-TTS-Nano single-stage TTS support (#2753) Signed-off-by: Yueqian Lin Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Co-authored-by: Hongsheng Liu --- .buildkite/test-merge.yml | 22 + .../offline_inference/moss_tts_nano/README.md | 101 +++ .../moss_tts_nano/end2end.py | 180 +++++ .../online_serving/moss_tts_nano/README.md | 158 ++++ .../moss_tts_nano/gradio_demo.py | 690 ++++++++++++++++++ .../moss_tts_nano/run_gradio_demo.sh | 43 ++ .../moss_tts_nano/run_server.sh | 19 + .../test_seed_tts_dataset_variants.py | 1 + .../offline_inference/test_moss_tts_nano.py | 131 ++++ .../e2e/online_serving/test_moss_tts_nano.py | 117 +++ vllm_omni/config/pipeline_registry.py | 4 + vllm_omni/deploy/moss_tts_nano.yaml | 31 + .../entrypoints/openai/serving_speech.py | 47 ++ .../models/moss_tts_nano/__init__.py | 5 + .../configuration_moss_tts_nano.py | 60 ++ .../moss_tts_nano/modeling_moss_tts_nano.py | 515 +++++++++++++ .../models/moss_tts_nano/pipeline.py | 40 + vllm_omni/model_executor/models/registry.py | 6 + 18 files changed, 2170 insertions(+) create mode 100644 examples/offline_inference/moss_tts_nano/README.md create mode 100644 examples/offline_inference/moss_tts_nano/end2end.py create mode 100644 examples/online_serving/moss_tts_nano/README.md create mode 100644 examples/online_serving/moss_tts_nano/gradio_demo.py create mode 100755 examples/online_serving/moss_tts_nano/run_gradio_demo.sh create mode 100755 examples/online_serving/moss_tts_nano/run_server.sh create mode 100644 tests/e2e/offline_inference/test_moss_tts_nano.py create mode 100644 tests/e2e/online_serving/test_moss_tts_nano.py create mode 100644 vllm_omni/deploy/moss_tts_nano.yaml create mode 100644 vllm_omni/model_executor/models/moss_tts_nano/__init__.py create mode 100644 vllm_omni/model_executor/models/moss_tts_nano/configuration_moss_tts_nano.py create mode 100644 vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py create mode 100644 vllm_omni/model_executor/models/moss_tts_nano/pipeline.py diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index 691f3f876..848297cec 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/test-merge.yml @@ -235,6 +235,28 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" + - label: "MOSS-TTS-Nano E2E Test" + depends_on: upload-merge-pipeline + commands: + - | + timeout 20m bash -c ' + export VLLM_LOGGING_LEVEL=DEBUG + export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" + pytest -s -v tests/e2e/online_serving/test_moss_tts_nano.py tests/e2e/offline_inference/test_moss_tts_nano.py -m "advanced_model" --run-level "advanced_model" + ' + agents: + queue: "gpu_1_queue" + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" + - label: "Omni Model Test with H100" timeout_in_minutes: 30 depends_on: upload-merge-pipeline diff --git a/examples/offline_inference/moss_tts_nano/README.md b/examples/offline_inference/moss_tts_nano/README.md new file mode 100644 index 000000000..8534d0549 --- /dev/null +++ b/examples/offline_inference/moss_tts_nano/README.md @@ -0,0 +1,101 @@ +# MOSS-TTS-Nano Offline Inference + +## Overview + +Single-stage offline TTS pipeline using the 0.1B MOSS-TTS-Nano AR LM and MOSS-Audio-Tokenizer-Nano codec. Outputs 48 kHz stereo WAV. + +## Quick Start + +```bash +python end2end.py --text "Hello, this is MOSS-TTS-Nano." +``` + +The first run downloads `OpenMOSS-Team/MOSS-TTS-Nano` and `OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano` from Hugging Face. + +## Usage + +``` +python end2end.py [OPTIONS] + +Options: + --text TEXT Text to synthesize (default: "Hello, this is MOSS-TTS-Nano speaking.") + --voice VOICE Built-in voice preset (default: Junhao). See voice table below. + --mode MODE voice_clone (default) or continuation + --prompt-audio PATH Reference WAV/MP3 for custom voice cloning + --prompt-text TEXT Reference transcript (continuation mode) + --max-new-frames N Max AR frames, default 375 (~14 s audio) + --seed INT Random seed for reproducibility + --audio-temperature F Audio sampling temperature (default: 0.8) + --audio-top-k N Audio top-k sampling (default: 25) + --audio-top-p F Audio top-p sampling (default: 0.95) + --text-temperature F Text layer temperature (default: 1.0) + --batch Run a built-in batch of diverse samples (ZH/EN/FR) + --output-dir DIR Directory for WAV outputs (default: /tmp/moss_tts_nano_output) + --deploy-config PATH Override deploy YAML (defaults to vllm_omni/deploy/moss_tts_nano.yaml) + --stage-init-timeout INT Timeout in seconds for stage init (default: 120) +``` + +## Examples + +```bash +# Built-in Chinese voice +python end2end.py --text "你好，这是MOSS-TTS-Nano的语音合成演示。" --voice Junhao + +# Built-in English voice +python end2end.py --text "Hello from MOSS-TTS-Nano." --voice Ava + +# Custom voice clone +python end2end.py \ + --text "Hello, this is a cloned voice." \ + --prompt-audio /path/to/reference.wav \ + --prompt-text "Exact transcript of the reference audio." + +# Batch synthesis (ZH/EN/FR) +python end2end.py --batch --output-dir /tmp/batch_output + +# Reproducible output +python end2end.py --text "Deterministic test." --seed 42 +``` + +## Built-in Voice Presets + +| Voice | Language | +|-------|----------| +| `Junhao` | ZH | +| `Zhiming` | ZH | +| `Weiguo` | ZH | +| `Xiaoyu` | ZH | +| `Yuewen` | ZH | +| `Lingyu` | ZH | +| `Ava` | EN | +| `Bella` | EN | +| `Adam` | EN | +| `Nathan` | EN | +| `Sakura` | JA | +| `Yui` | JA | +| `Aoi` | JA | +| `Hina` | JA | +| `Mei` | JA | + +## Deploy Config + +Runtime knobs live in `vllm_omni/deploy/moss_tts_nano.yaml` (auto-loaded; +override with `--deploy-config PATH`). Key stage-level settings: + +```yaml +stages: + - stage_id: 0 + gpu_memory_utilization: 0.3 # ~2 GB VRAM; increase for faster init + max_num_seqs: 4 # concurrent requests + max_model_len: 4096 +``` + +## Output Format + +WAV files, 48 kHz, stereo (2-channel). The codec interleaves stereo as `[L, R, L, R, ...]` in the flat tensor returned by the model. + +## Troubleshooting + +- **`libnvrtc.so.13: cannot open shared object file`**: torchaudio 2.10+ torchcodec backend requires NVRTC. The model patches `torchaudio.load/save` automatically at load time to fall back to soundfile. +- **`flash_attn not installed`**: The model falls back to `sdpa` attention automatically. +- **Empty audio**: Check that `--text` is non-empty and the model loaded successfully (look for "MOSS-TTS-Nano LM loaded" in logs). diff --git a/examples/offline_inference/moss_tts_nano/end2end.py b/examples/offline_inference/moss_tts_nano/end2end.py new file mode 100644 index 000000000..177e03791 --- /dev/null +++ b/examples/offline_inference/moss_tts_nano/end2end.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Offline inference example for MOSS-TTS-Nano via vLLM-Omni. + +Single-stage pipeline: the 0.1B AR LM and MOSS-Audio-Tokenizer-Nano codec +both run inside one generation stage. Output is 48 kHz stereo WAV. + +Supports: + - Built-in voice presets (--voice) + - Custom voice cloning via reference audio (--prompt-audio) + - Batch synthesis of multiple texts + - Reproducible output via --seed + +Usage examples: + # Built-in voice + python end2end.py --text "Hello from MOSS-TTS-Nano." + + # Voice clone with reference audio + python end2end.py --text "Hello!" --prompt-audio /path/to/ref.wav + + # Batch with different voices + python end2end.py --batch --output-dir /tmp/moss_output +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import soundfile as sf +import torch +from vllm import SamplingParams +from vllm.utils.argparse_utils import FlexibleArgumentParser + +# Prevent multiprocessing from re-importing CUDA in the wrong context. +os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + +from vllm_omni import Omni # noqa: E402 + +MODEL = "OpenMOSS-Team/MOSS-TTS-Nano" + +BATCH_SAMPLES = [ + {"text": "Hello, this is a test of MOSS-TTS-Nano.", "voice": "Ava", "label": "en_hello"}, + {"text": "你好，这是 MOSS-TTS-Nano 的语音合成测试。", "voice": "Junhao", "label": "zh_hello"}, + {"text": "Bonjour, ceci est un test de synthèse vocale.", "voice": "Bella", "label": "fr_hello"}, +] + + +def build_request( + text: str, + voice: str = "Junhao", + mode: str = "voice_clone", + prompt_audio_path: str | None = None, + prompt_text: str | None = None, + max_new_frames: int = 375, + seed: int | None = None, + audio_temperature: float = 0.8, + audio_top_k: int = 25, + audio_top_p: float = 0.95, + text_temperature: float = 1.0, +) -> dict: + """Build an Omni request payload for MOSS-TTS-Nano.""" + additional: dict = { + "text": [text], + "voice": [voice], + "mode": [mode], + "max_new_frames": [max_new_frames], + "audio_temperature": [audio_temperature], + "audio_top_k": [audio_top_k], + "audio_top_p": [audio_top_p], + "text_temperature": [text_temperature], + } + if prompt_audio_path: + additional["prompt_audio_path"] = [str(prompt_audio_path)] + if prompt_text: + additional["prompt_text"] = [prompt_text] + if seed is not None: + additional["seed"] = [seed] + + return { + "prompt": "<|im_start|>assistant\n", # minimal placeholder prompt + "additional_information": additional, + } + + +def save_audio(waveform: torch.Tensor, path: str, sample_rate: int = 48000) -> None: + audio_np = waveform.float().numpy() + # Reshape stereo: inference_stream yields interleaved [samples*2]; reshape to [samples, 2] + if audio_np.ndim == 1 and audio_np.shape[0] % 2 == 0: + audio_np = audio_np.reshape(-1, 2) + sf.write(path, audio_np, sample_rate) + print(f" Saved {path} ({audio_np.shape}, {sample_rate} Hz)") + + +def main(args) -> None: + omni = Omni( + model=MODEL, + deploy_config=args.deploy_config, + stage_init_timeout=args.stage_init_timeout, + ) + + sampling_params = SamplingParams( + temperature=1.0, + top_p=1.0, + top_k=50, + max_tokens=4096, + seed=args.seed if args.seed is not None else 42, + detokenize=False, + ) + + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + if args.batch: + print(f"Running batch synthesis ({len(BATCH_SAMPLES)} samples)...") + inputs = [build_request(s["text"], voice=s["voice"], seed=args.seed) for s in BATCH_SAMPLES] + params_list = [sampling_params] * len(inputs) + else: + print(f"Synthesizing: {args.text!r}") + inputs = build_request( + text=args.text, + voice=args.voice, + mode=args.mode, + prompt_audio_path=args.prompt_audio, + prompt_text=args.prompt_text, + max_new_frames=args.max_new_frames, + seed=args.seed, + audio_temperature=args.audio_temperature, + audio_top_k=args.audio_top_k, + audio_top_p=args.audio_top_p, + text_temperature=args.text_temperature, + ) + params_list = sampling_params + + for stage_outputs in omni.generate(inputs, params_list): + for i, req_output in enumerate(stage_outputs.request_output): + for j, out in enumerate(req_output.outputs): + mm = out.multimodal_output + if mm is None: + print(f" [req {i}] No audio output.") + continue + audio = mm.get("audio") + sr_tensor = mm.get("sr") + if audio is None: + print(f" [req {i}] No waveform in multimodal_output.") + continue + sr = int(sr_tensor.item()) if sr_tensor is not None else 48000 + label = BATCH_SAMPLES[i]["label"] if args.batch else f"output_{i}_{j}" + out_path = str(output_dir / f"{label}.wav") + save_audio(audio.cpu(), out_path, sr) + + print("Done.") + + +def parse_args(): + parser = FlexibleArgumentParser(description="MOSS-TTS-Nano offline inference") + parser.add_argument("--text", default="Hello, this is MOSS-TTS-Nano speaking.", help="Text to synthesize.") + parser.add_argument("--voice", default="Junhao", help="Built-in voice preset name.") + parser.add_argument("--mode", default="voice_clone", choices=["voice_clone", "continuation"]) + parser.add_argument("--prompt-audio", default=None, help="Path to reference audio for voice cloning.") + parser.add_argument("--prompt-text", default=None, help="Reference transcript (continuation mode).") + parser.add_argument("--max-new-frames", type=int, default=375, help="Max AR frames (~14s at default).") + parser.add_argument("--seed", type=int, default=None, help="Random seed.") + parser.add_argument("--audio-temperature", type=float, default=0.8) + parser.add_argument("--audio-top-k", type=int, default=25) + parser.add_argument("--audio-top-p", type=float, default=0.95) + parser.add_argument("--text-temperature", type=float, default=1.0) + parser.add_argument("--batch", action="store_true", help="Run built-in batch of diverse samples.") + parser.add_argument("--output-dir", default="/tmp/moss_tts_nano_output", help="Directory for WAV outputs.") + parser.add_argument( + "--deploy-config", + default=None, + help="Path to a deploy YAML; leave unset to auto-load vllm_omni/deploy/moss_tts_nano.yaml.", + ) + parser.add_argument("--stage-init-timeout", type=int, default=120) + return parser.parse_args() + + +if __name__ == "__main__": + main(parse_args()) diff --git a/examples/online_serving/moss_tts_nano/README.md b/examples/online_serving/moss_tts_nano/README.md new file mode 100644 index 000000000..c366ce02e --- /dev/null +++ b/examples/online_serving/moss_tts_nano/README.md @@ -0,0 +1,158 @@ +# MOSS-TTS-Nano + +## Model checkpoint + +| Model | Description | +|-------|-------------| +| `OpenMOSS-Team/MOSS-TTS-Nano` | 0.1B AR LM + MOSS-Audio-Tokenizer-Nano codec, 48 kHz stereo, ZH/EN/JA | + +## Gradio Demo + +An interactive Gradio demo is available with multilingual voice presets, custom voice cloning, and streaming support. + +```bash +# Option 1: Launch server + Gradio together +./run_gradio_demo.sh + +# Option 2: If server is already running +python gradio_demo.py --api-base http://localhost:8091 +``` + +Then open http://localhost:7860 in your browser. + +Features: + +- 15 built-in voice presets (6 ZH · 4 EN · 5 JA) +- Custom voice cloning from uploaded reference audio +- Streaming mode (progressive PCM playback) + +## Launch the Server + +```bash +vllm serve OpenMOSS-Team/MOSS-TTS-Nano --omni --port 8091 +``` + +The deploy config at `vllm_omni/deploy/moss_tts_nano.yaml` auto-loads; no +`--stage-configs-path`, `--trust-remote-code`, or `--enforce-eager` flags +are needed. + +Or use the convenience script: + +```bash +./run_server.sh +``` + +## Send TTS Request + +### Using curl + +```bash +# Built-in voice preset +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "你好，这是语音合成测试。", + "voice": "Junhao", + "response_format": "wav" + }' --output output.wav + +# English preset +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, this is MOSS-TTS-Nano.", + "voice": "Ava", + "response_format": "wav" + }' --output output_en.wav +``` + +### Custom Voice Cloning + +Provide a reference audio (base64 data URL) and its transcript: + +```bash +REF_AUDIO=$(base64 -w 0 /path/to/reference.wav) +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d "{ + \"input\": \"Hello, this is a cloned voice.\", + \"voice\": \"Junhao\", + \"ref_audio\": \"data:audio/wav;base64,${REF_AUDIO}\", + \"ref_text\": \"Exact transcript of the reference audio.\" + }" --output cloned.wav +``` + +### Using Python + +```python +import httpx + +response = httpx.post( + "http://localhost:8091/v1/audio/speech", + json={ + "input": "你好，这是语音合成测试。", + "voice": "Junhao", + "response_format": "wav", + }, + timeout=300.0, +) + +with open("output.wav", "wb") as f: + f.write(response.content) +``` + +### Streaming + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, streaming output from MOSS-TTS-Nano.", + "voice": "Ava", + "stream": true, + "response_format": "pcm" + }' --no-buffer | play -t raw -r 48000 -e signed -b 16 -c 2 - +``` + +**Note:** MOSS-TTS-Nano outputs at 48 kHz stereo (2-channel). + +## API Parameters + +MOSS-TTS-Nano uses the standard `/v1/audio/speech` endpoint. + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `input` | string | **required** | Text to synthesize (ZH / EN / JA) | +| `voice` | string | `"Junhao"` | Built-in voice name (see table below) | +| `response_format` | string | `"wav"` | Audio format: wav, mp3, flac, pcm | +| `ref_audio` | string | null | Base64 data URL for custom voice cloning | +| `ref_text` | string | null | Transcript of reference audio (required when `ref_audio` is set) | +| `stream` | bool | false | Stream raw PCM chunks | +| `max_new_tokens` | int | 4096 | Maximum tokens to generate | + +## Built-in Voice Presets + +| Voice | Language | Description | +|-------|----------|-------------| +| `Junhao` | ZH | Male, standard Mandarin | +| `Zhiming` | ZH | Male | +| `Weiguo` | ZH | Male | +| `Xiaoyu` | ZH | Female | +| `Yuewen` | ZH | Female | +| `Lingyu` | ZH | Female | +| `Ava` | EN | Female, American English | +| `Bella` | EN | Female | +| `Adam` | EN | Male | +| `Nathan` | EN | Male | +| `Sakura` | JA | Female | +| `Yui` | JA | Female | +| `Aoi` | JA | Female | +| `Hina` | JA | Female | +| `Mei` | JA | Female | + +## Troubleshooting + +1. **`libnvrtc.so.13: cannot open shared object file`**: torchaudio 2.10+ defaults to torchcodec which requires NVRTC. vLLM-Omni patches this automatically at model load time to use soundfile instead. +2. **Connection refused**: Ensure the server is running on the correct port. +3. **Flashinfer version mismatch**: Set `FLASHINFER_DISABLE_VERSION_CHECK=1` if you see version warnings. +4. **Out of memory**: The default `gpu_memory_utilization=0.3` is conservative. Increase it in the stage config if you have more VRAM available. diff --git a/examples/online_serving/moss_tts_nano/gradio_demo.py b/examples/online_serving/moss_tts_nano/gradio_demo.py new file mode 100644 index 000000000..f120f838c --- /dev/null +++ b/examples/online_serving/moss_tts_nano/gradio_demo.py @@ -0,0 +1,690 @@ +"""Gradio demo for MOSS-TTS-Nano with gapless streaming audio playback. + +Uses a custom AudioWorklet-based player (adapted from the Qwen3-TTS demo) +for gap-free streaming. Audio is streamed from the vLLM-Omni server +through a same-origin proxy and played via the Web Audio API's AudioWorklet. + +MOSS-TTS-Nano outputs 48 kHz stereo (2-channel interleaved L,R,L,R,...). +The AudioWorklet downmixes stereo to mono during playback. + +Also supports non-streaming mode (full audio download) via gr.Audio. + +Usage: + # Start the server first (see run_server.sh), then: + python gradio_demo.py --api-base http://localhost:8091 + + # Or use run_gradio_demo.sh to start both server and demo together. +""" + +from __future__ import annotations + +import argparse +import base64 +import io +import json +import logging + +try: + import gradio as gr +except ImportError: + raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None + +import httpx +import numpy as np +import soundfile as sf +from fastapi import FastAPI, Request +from fastapi.responses import Response, StreamingResponse + +logger = logging.getLogger(__name__) + +# MOSS-TTS-Nano outputs 48 kHz stereo (2-channel interleaved). +PCM_SAMPLE_RATE = 48000 +PCM_CHANNELS = 2 +DEFAULT_API_BASE = "http://localhost:8091" + +# Built-in voice presets (matches the MOSS-TTS-Nano GitHub repo assets/) +VOICE_PRESETS = { + # Chinese + "Junhao (ZH)": "Junhao", + "Zhiming (ZH)": "Zhiming", + "Weiguo (ZH)": "Weiguo", + "Xiaoyu (ZH)": "Xiaoyu", + "Yuewen (ZH)": "Yuewen", + "Lingyu (ZH)": "Lingyu", + # English + "Ava (EN)": "Ava", + "Bella (EN)": "Bella", + "Adam (EN)": "Adam", + "Nathan (EN)": "Nathan", + # Japanese + "Sakura (JA)": "Sakura", + "Yui (JA)": "Yui", + "Aoi (JA)": "Aoi", + "Hina (JA)": "Hina", + "Mei (JA)": "Mei", +} +PRESET_LABELS = list(VOICE_PRESETS.keys()) +DEFAULT_PRESET = "Junhao (ZH)" + +# ── AudioWorklet processor (loaded in browser via Blob URL) ────────── +# Downmixes stereo-interleaved int16 PCM to mono float32 for playback. +WORKLET_JS = r""" +class TTSPlaybackProcessor extends AudioWorkletProcessor { + constructor() { + super(); + this.queue = []; + this.buf = null; + this.pos = 0; + this.playing = false; + this.played = 0; + this.port.onmessage = (e) => { + if (e.data && e.data.type === 'clear') { + this.queue = []; this.buf = null; this.pos = 0; this.played = 0; + if (this.playing) { this.playing = false; this.port.postMessage({type:'stopped'}); } + return; + } + // Input: Int16Array of stereo-interleaved samples [L0,R0,L1,R1,...] + // Convert to mono float32 by averaging L+R channels. + const pcm = e.data; + const monoLen = Math.floor(pcm.length / 2); + const mono = new Float32Array(monoLen); + for (let i = 0; i < monoLen; i++) { + mono[i] = (pcm[i * 2] + pcm[i * 2 + 1]) / 65536.0; + } + this.queue.push(mono); + }; + } + process(inputs, outputs) { + const out = outputs[0][0]; + for (let i = 0; i < out.length; i++) { + if (!this.buf || this.pos >= this.buf.length) { + if (this.queue.length > 0) { + this.buf = this.queue.shift(); this.pos = 0; + } else { + for (let j = i; j < out.length; j++) out[j] = 0; + if (this.playing) { this.playing = false; this.port.postMessage({type:'stopped', played:this.played}); } + return true; + } + } + out[i] = this.buf[this.pos++]; + this.played++; + } + if (!this.playing) { this.playing = true; this.port.postMessage({type:'started'}); } + return true; + } +} +registerProcessor('tts-playback-processor', TTSPlaybackProcessor); +""" + +# ── Player HTML (container with metric cards) ──────────────────────── +PLAYER_HTML = """ +
+
+
+ Ready + +
+ + + +
+""" + + +def _build_player_js(sample_rate: int) -> str: + """Build the JavaScript that powers the AudioWorklet player.""" + return f""" + +""" + + +def encode_audio_to_base64(audio_data: tuple) -> str: + """Encode Gradio audio input (sample_rate, numpy_array) to base64 data URL.""" + sample_rate, audio_np = audio_data + if audio_np.dtype != np.int16: + if audio_np.dtype in (np.float32, np.float64): + audio_np = np.clip(audio_np, -1.0, 1.0) + audio_np = (audio_np * 32767).astype(np.int16) + else: + audio_np = audio_np.astype(np.int16) + buf = io.BytesIO() + sf.write(buf, audio_np, sample_rate, format="WAV") + wav_b64 = base64.b64encode(buf.getvalue()).decode("utf-8") + return f"data:audio/wav;base64,{wav_b64}" + + +def build_payload( + text: str, + voice_preset: str, + ref_audio: tuple | None, + ref_text: str, + response_format: str = "pcm", + stream: bool = True, +) -> dict: + """Build the /v1/audio/speech request payload.""" + if not text or not text.strip(): + raise gr.Error("Please enter text to synthesize.") + + voice_name = VOICE_PRESETS.get(voice_preset, "Junhao") + + payload: dict = { + "input": text.strip(), + "voice": voice_name, + "response_format": "pcm" if stream else response_format, + "stream": stream, + } + + has_ref = ref_audio is not None + if has_ref: + payload["ref_audio"] = encode_audio_to_base64(ref_audio) + ref_text_stripped = ref_text.strip() if ref_text else "" + if ref_text_stripped: + payload["ref_text"] = ref_text_stripped + else: + raise gr.Error("Voice cloning requires a transcript of the reference audio.") + + return payload + + +def generate_speech( + api_base: str, + text: str, + voice_preset: str, + ref_audio: tuple | None, + ref_text: str, + response_format: str, +) -> tuple: + """Non-streaming: call /v1/audio/speech and return full audio as (sr, np_array).""" + payload = build_payload(text, voice_preset, ref_audio, ref_text, response_format, stream=False) + + try: + with httpx.Client(timeout=300.0) as client: + resp = client.post( + f"{api_base}/v1/audio/speech", + json=payload, + headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"}, + ) + except httpx.TimeoutException: + raise gr.Error("Request timed out. The server may be busy.") + except httpx.ConnectError: + raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?") + + if resp.status_code != 200: + raise gr.Error(f"Server error ({resp.status_code}): {resp.text}") + + try: + if response_format == "pcm": + # PCM is stereo interleaved; downmix to mono for Gradio. + samples = np.frombuffer(resp.content, dtype=np.int16) + n = len(samples) - (len(samples) % 2) + left = samples[:n:2].astype(np.float32) + right = samples[1:n:2].astype(np.float32) + mono = (left + right) / 65536.0 + return (PCM_SAMPLE_RATE, mono) + # WAV/MP3/FLAC: soundfile handles header parsing. + audio_np, sample_rate = sf.read(io.BytesIO(resp.content)) + if audio_np.ndim > 1: + audio_np = audio_np.mean(axis=1) + return (int(sample_rate), audio_np.astype(np.float32)) + except Exception as e: + raise gr.Error(f"Failed to decode audio: {e}") + + +def create_app(api_base: str): + """Create the FastAPI app with streaming proxy + Gradio UI.""" + fastapi_app = FastAPI() + + # ── Streaming proxy (same-origin, no CORS issues) ──────────── + @fastapi_app.post("/proxy/v1/audio/speech") + async def proxy_speech(request: Request): + body = await request.json() + logger.info( + "Proxy request: input=%r voice=%s stream=%s", + body.get("input", "")[:50], + body.get("voice"), + body.get("stream"), + ) + try: + client = httpx.AsyncClient(timeout=300) + resp = await client.send( + client.build_request( + "POST", + f"{api_base}/v1/audio/speech", + json=body, + headers={"Authorization": "Bearer EMPTY", "Content-Type": "application/json"}, + ), + stream=True, + ) + except Exception as exc: + logger.exception("Proxy connection error") + await client.aclose() + return Response(content=str(exc), status_code=502) + + if resp.status_code != 200: + content = await resp.aread() + logger.error("Proxy upstream error %d: %s", resp.status_code, content[:200]) + await resp.aclose() + await client.aclose() + return Response(content=content, status_code=resp.status_code) + + async def relay(): + total = 0 + try: + async for chunk in resp.aiter_bytes(): + total += len(chunk) + yield chunk + except Exception: + logger.exception("Proxy relay error after %d bytes", total) + finally: + logger.info("Proxy relay done: %d bytes", total) + await resp.aclose() + await client.aclose() + + return StreamingResponse(relay(), media_type="application/octet-stream") + + # ── Gradio UI ──────────────────────────────────────────────── + css = """ + #generate-btn button { width: 100%; } + #streaming-player { border: 1px solid var(--border-color-primary) !important; border-radius: var(--block-radius) !important; padding: var(--block-padding) !important; } + """ + + theme = gr.themes.Default( + primary_hue=gr.themes.Color( + c50="#E8F5E9", + c100="#C8E6C9", + c200="#A5D6A7", + c300="#81C784", + c400="#66BB6A", + c500="#4CAF50", + c600="#43A047", + c700="#388E3C", + c800="#2E7D32", + c900="#1B5E20", + c950="#0D3B10", + ), + ) + + with gr.Blocks(title="MOSS-TTS-Nano Demo") as demo: + gr.HTML(f""" +
+ vLLM-Omni +
+

MOSS-TTS-Nano Streaming Demo

+ + 0.1B AR LM + MOSS-Audio-Tokenizer-Nano · 48 kHz · ZH / EN / JA +  ·  Served by + vLLM-Omni +  ·  {api_base} + +
+
+ """) + + with gr.Row(): + with gr.Column(scale=3): + text_input = gr.Textbox( + label="Text to Synthesize", + placeholder="Enter text in Chinese, English, or Japanese...", + lines=4, + ) + + voice_preset = gr.Dropdown( + choices=PRESET_LABELS, + value=DEFAULT_PRESET, + label="Voice Preset", + info="Built-in multilingual voices. ZH = Chinese, EN = English, JA = Japanese.", + ) + + with gr.Accordion("Custom Voice Clone (optional)", open=False): + gr.Markdown( + "Upload a reference audio clip (10-30 s) and provide its transcript " + "to clone that voice instead of using the built-in preset." + ) + ref_audio = gr.Audio( + label="Reference Audio", + type="numpy", + sources=["upload", "microphone"], + ) + ref_text = gr.Textbox( + label="Reference Audio Transcript (required for cloning)", + placeholder="Exact transcript of the reference audio...", + lines=2, + ) + + with gr.Row(): + response_format = gr.Dropdown( + choices=["wav", "mp3", "flac", "pcm"], + value="pcm", + label="Audio Format", + interactive=False, + scale=1, + ) + stream_checkbox = gr.Checkbox( + label="Stream (gapless)", + value=True, + info="AudioWorklet streaming (recommended)", + scale=1, + ) + + with gr.Row(): + generate_btn = gr.Button( + "Generate Speech", + variant="primary", + size="lg", + elem_id="generate-btn", + scale=3, + ) + reset_btn = gr.Button( + "Reset", + variant="secondary", + size="lg", + scale=1, + ) + + with gr.Column(scale=2): + player_html = gr.HTML( + value=PLAYER_HTML, + visible=True, + label="streaming player", + elem_id="streaming-player", + ) + audio_output = gr.Audio( + label="Generated Audio", + interactive=False, + autoplay=True, + visible=False, + ) + gr.Examples( + examples=[ + ["Hello, this is MOSS-TTS-Nano speaking.", "Ava (EN)"], + ["The quick brown fox jumps over the lazy dog.", "Adam (EN)"], + ["MOSS-TTS-Nanoは、軽量なテキスト読み上げモデルです。", "Sakura (JA)"], + ], + inputs=[text_input, voice_preset], + label="Examples", + ) + gr.HTML(""" +
+ + vLLM-Omni + +
+ """) + + # Hidden textbox to pass payload from Python -> JavaScript + hidden_payload = gr.Textbox(visible=False, elem_id="tts-payload") + + # ── Event wiring ───────────────────────────────────────── + def on_stream_change(stream: bool): + if stream: + return ( + gr.update(value="pcm", interactive=False), + gr.update(visible=True), # player + gr.update(visible=False), # audio + ) + return ( + gr.update(value="wav", interactive=True), + gr.update(visible=False), + gr.update(visible=True), + ) + + stream_checkbox.change( + fn=on_stream_change, + inputs=[stream_checkbox], + outputs=[response_format, player_html, audio_output], + ) + + def on_reset(): + return ( + "", # text + None, # audio_output + "", # hidden_payload + PLAYER_HTML, # reset player + ) + + reset_btn.click( + fn=on_reset, + outputs=[text_input, audio_output, hidden_payload, player_html], + js="() => { if (window.ttsStop) window.ttsStop(); }", + ) + + all_inputs = [text_input, voice_preset, ref_audio, ref_text, response_format] + + def on_generate(stream_enabled, *args): + import time as _time + + text, voice, ref_a, ref_t, _fmt = args + if stream_enabled: + payload = build_payload(text, voice, ref_a, ref_t, stream=True) + payload["_nonce"] = int(_time.time() * 1000) + return json.dumps(payload), gr.update() + else: + audio = generate_speech(api_base, text, voice, ref_a, ref_t, _fmt) + return "", audio + + generate_btn.click( + fn=on_generate, + inputs=[stream_checkbox] + all_inputs, + outputs=[hidden_payload, audio_output], + ).then( + fn=lambda p: p, + inputs=[hidden_payload], + outputs=[hidden_payload], + js="(p) => { if (p && p.trim()) { const d = JSON.parse(p); delete d._nonce; window.ttsGenerate(d); } return p; }", + ) + + demo.queue() + + return gr.mount_gradio_app( + fastapi_app, + demo, + path="/", + css=css, + theme=theme, + head=_build_player_js(PCM_SAMPLE_RATE), + ) + + +def main(): + parser = argparse.ArgumentParser( + description="Gradio demo for MOSS-TTS-Nano with gapless AudioWorklet streaming.", + ) + parser.add_argument("--api-base", default=DEFAULT_API_BASE, help=f"API base URL (default: {DEFAULT_API_BASE})") + parser.add_argument("--host", default="0.0.0.0", help="Gradio host (default: 0.0.0.0)") + parser.add_argument("--port", type=int, default=7860, help="Gradio port (default: 7860)") + parser.add_argument("--share", action="store_true", help="Share publicly via Gradio tunnel") + args = parser.parse_args() + + logging.basicConfig(level=logging.INFO) + print(f"Connecting to vLLM-Omni server at: {args.api_base}") + + app = create_app(args.api_base) + + import uvicorn + + uvicorn.run(app, host=args.host, port=args.port) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/moss_tts_nano/run_gradio_demo.sh b/examples/online_serving/moss_tts_nano/run_gradio_demo.sh new file mode 100755 index 000000000..cd96c6b4f --- /dev/null +++ b/examples/online_serving/moss_tts_nano/run_gradio_demo.sh @@ -0,0 +1,43 @@ +#!/bin/bash +# Launch MOSS-TTS-Nano server + Gradio demo together. +# +# Usage: +# ./run_gradio_demo.sh +# CUDA_VISIBLE_DEVICES=0 PORT=8091 GRADIO_PORT=7860 ./run_gradio_demo.sh + +set -e + +MODEL="${MODEL:-OpenMOSS-Team/MOSS-TTS-Nano}" +PORT="${PORT:-8091}" +GRADIO_PORT="${GRADIO_PORT:-7860}" +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +echo "Starting MOSS-TTS-Nano server (port $PORT)..." +FLASHINFER_DISABLE_VERSION_CHECK=1 \ +vllm serve "$MODEL" \ + --host 0.0.0.0 \ + --port "$PORT" \ + --omni & +SERVER_PID=$! + +cleanup() { + echo "Stopping server (PID $SERVER_PID)..." + kill $SERVER_PID 2>/dev/null + wait $SERVER_PID 2>/dev/null +} +trap cleanup EXIT + +# Wait for server to be ready. +echo "Waiting for server to start..." +for i in $(seq 1 120); do + if curl -s "http://localhost:$PORT/health" > /dev/null 2>&1; then + echo "Server ready." + break + fi + sleep 2 +done + +echo "Starting Gradio demo (port $GRADIO_PORT)..." +python "$SCRIPT_DIR/gradio_demo.py" \ + --api-base "http://localhost:$PORT" \ + --port "$GRADIO_PORT" diff --git a/examples/online_serving/moss_tts_nano/run_server.sh b/examples/online_serving/moss_tts_nano/run_server.sh new file mode 100755 index 000000000..5ddbac09c --- /dev/null +++ b/examples/online_serving/moss_tts_nano/run_server.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Launch vLLM-Omni server for MOSS-TTS-Nano +# +# Usage: +# ./run_server.sh +# CUDA_VISIBLE_DEVICES=0 PORT=8091 ./run_server.sh + +set -e + +MODEL="${MODEL:-OpenMOSS-Team/MOSS-TTS-Nano}" +PORT="${PORT:-8091}" + +echo "Starting MOSS-TTS-Nano server with model: $MODEL" + +FLASHINFER_DISABLE_VERSION_CHECK=1 \ +vllm serve "$MODEL" \ + --host 0.0.0.0 \ + --port "$PORT" \ + --omni diff --git a/tests/benchmarks/test_seed_tts_dataset_variants.py b/tests/benchmarks/test_seed_tts_dataset_variants.py index 7fa5747bd..ff7b40c48 100644 --- a/tests/benchmarks/test_seed_tts_dataset_variants.py +++ b/tests/benchmarks/test_seed_tts_dataset_variants.py @@ -59,6 +59,7 @@ def mock_tokenizer(mocker): tokenizer.all_special_ids = [] tokenizer.all_special_tokens = [] tokenizer.vocab_size = 1 + tokenizer.__len__.return_value = 1 return tokenizer diff --git a/tests/e2e/offline_inference/test_moss_tts_nano.py b/tests/e2e/offline_inference/test_moss_tts_nano.py new file mode 100644 index 000000000..40561d53f --- /dev/null +++ b/tests/e2e/offline_inference/test_moss_tts_nano.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E offline inference tests for MOSS-TTS-Nano single-stage pipeline.""" + +from __future__ import annotations + +import pytest +import torch +from vllm import SamplingParams + +from tests.helpers.mark import hardware_test +from vllm_omni import Omni + +MODEL_NAME = "OpenMOSS-Team/MOSS-TTS-Nano" +SAMPLE_RATE = 48000 + +DEFAULT_SAMPLING = SamplingParams( + temperature=1.0, + top_p=1.0, + top_k=50, + max_tokens=4096, + seed=42, + detokenize=False, +) + + +def _build_request( + text: str, + voice: str = "Junhao", + mode: str = "voice_clone", + max_new_frames: int = 100, # short for tests + seed: int = 42, +) -> dict: + return { + "prompt": "<|im_start|>assistant\n", + "additional_information": { + "text": [text], + "voice": [voice], + "mode": [mode], + "max_new_frames": [max_new_frames], + "seed": [seed], + }, + } + + +def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]: + """Run a single request and return (waveform, sample_rate).""" + for stage_outputs in omni.generate(request, DEFAULT_SAMPLING): + for req_output in stage_outputs.request_output: + mm = req_output.outputs[0].multimodal_output + assert mm is not None, "Expected multimodal_output to be non-None" + audio = mm.get("audio") + sr = mm.get("sr") + assert audio is not None, "Expected 'audio' key in multimodal_output" + assert isinstance(audio, torch.Tensor), f"audio should be Tensor, got {type(audio)}" + return audio.cpu(), int(sr.item()) if sr is not None else SAMPLE_RATE + raise AssertionError("No stage outputs received") + + +@pytest.fixture(scope="module") +def omni_engine(): + """Module-scoped Omni engine to avoid re-loading the model for each test.""" + return Omni(model=MODEL_NAME, stage_init_timeout=180) + + +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}) +def test_moss_tts_nano_english(omni_engine): + """English TTS produces non-empty 48 kHz stereo audio.""" + req = _build_request("Hello, this is a test of MOSS-TTS-Nano.", voice="Ava") + audio, sr = _collect_audio(omni_engine, req) + + assert sr == SAMPLE_RATE, f"Expected sample_rate={SAMPLE_RATE}, got {sr}" + assert audio.numel() > 0, "Audio tensor should not be empty" + assert not torch.all(audio == 0), "Audio should not be all-zeros (silence)" + + +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}) +def test_moss_tts_nano_chinese(omni_engine): + """Chinese TTS produces non-empty audio.""" + req = _build_request("你好，这是语音合成测试。", voice="Junhao") + audio, sr = _collect_audio(omni_engine, req) + + assert sr == SAMPLE_RATE + assert audio.numel() > 0 + assert not torch.all(audio == 0) + + +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}) +def test_moss_tts_nano_deterministic(omni_engine): + """Same seed produces identical waveforms.""" + req = _build_request("Reproducible output test.", voice="Adam", seed=123) + audio1, _ = _collect_audio(omni_engine, req) + audio2, _ = _collect_audio(omni_engine, req) + + assert audio1.shape == audio2.shape, "Waveform shapes should match with same seed" + assert torch.allclose(audio1, audio2, atol=1e-4), "Waveforms should match with same seed" + + +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}) +def test_moss_tts_nano_batch(omni_engine): + """Batch of two requests returns audio for each.""" + requests = [ + _build_request("First request.", voice="Ava"), + _build_request("Second request.", voice="Bella"), + ] + results = [] + for stage_outputs in omni_engine.generate(requests, [DEFAULT_SAMPLING] * 2): + for req_output in stage_outputs.request_output: + mm = req_output.outputs[0].multimodal_output + assert mm is not None + results.append(mm["audio"].cpu()) + + assert len(results) == 2, f"Expected 2 outputs, got {len(results)}" + for i, audio in enumerate(results): + assert audio.numel() > 0, f"Audio {i} is empty" + + +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}) +def test_moss_tts_nano_voice_presets(omni_engine): + """Multiple built-in voice presets all produce valid audio.""" + voices = ["Junhao", "Ava", "Bella"] + for voice in voices: + req = _build_request(f"Testing voice {voice}.", voice=voice) + audio, sr = _collect_audio(omni_engine, req) + assert audio.numel() > 0, f"Voice {voice} produced empty audio" + assert sr == SAMPLE_RATE diff --git a/tests/e2e/online_serving/test_moss_tts_nano.py b/tests/e2e/online_serving/test_moss_tts_nano.py new file mode 100644 index 000000000..22e305345 --- /dev/null +++ b/tests/e2e/online_serving/test_moss_tts_nano.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +E2E Online tests for MOSS-TTS-Nano model with text input and audio output. + +These tests verify the /v1/audio/speech endpoint works correctly with +actual model inference, not mocks. +""" + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path + +MODEL = "OpenMOSS-Team/MOSS-TTS-Nano" + + +def get_prompt(prompt_type="text"): + """Text prompt for text-to-audio tests.""" + prompts = { + "text": "Hello, this is a test of MOSS-TTS-Nano text to speech synthesis.", + "chinese": "你好，这是MOSS-TTS-Nano的语音合成测试。", + } + return prompts.get(prompt_type, prompts["text"]) + + +tts_server_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=get_deploy_config_path("moss_tts_nano.yaml"), + server_args=["--disable-log-stats"], + ), + id="moss_tts_nano", + ) +] + + +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_text_to_audio_001(omni_server, openai_client) -> None: + """ + Test text input processing and audio output via /v1/audio/speech. + Deploy Setting: default yaml + Input Modal: text + Output Modal: audio (48 kHz) + Input Setting: stream=False, voice=Junhao + Datasets: single request + """ + request_config = { + "model": omni_server.model, + "input": get_prompt(), + "stream": False, + "response_format": "wav", + "voice": "Junhao", + } + + openai_client.send_audio_speech_request(request_config) + + +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_text_to_audio_002(omni_server, openai_client) -> None: + """ + Test streaming text-to-audio via /v1/audio/speech. + Deploy Setting: default yaml + Input Modal: text + Output Modal: audio (48 kHz, PCM stream) + Input Setting: stream=True, voice=Ava + Datasets: single request + """ + request_config = { + "model": omni_server.model, + "input": get_prompt(), + "stream": True, + "response_format": "pcm", + "voice": "Ava", + } + + openai_client.send_audio_speech_request(request_config) + + +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_text_to_audio_003(omni_server, openai_client) -> None: + """ + Test Chinese text-to-audio via /v1/audio/speech. + Deploy Setting: default yaml + Input Modal: text (Chinese) + Output Modal: audio (48 kHz) + Input Setting: stream=False, voice=Junhao + Datasets: single request + """ + request_config = { + "model": omni_server.model, + "input": get_prompt("chinese"), + "stream": False, + "response_format": "wav", + "voice": "Junhao", + } + + openai_client.send_audio_speech_request(request_config) diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 937220914..6f5c072a3 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -85,4 +85,8 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = { "vllm_omni.model_executor.models.fish_speech.pipeline", "FISH_SPEECH_PIPELINE", ), + "moss_tts_nano": ( + "vllm_omni.model_executor.models.moss_tts_nano.pipeline", + "MOSS_TTS_NANO_PIPELINE", + ), } diff --git a/vllm_omni/deploy/moss_tts_nano.yaml b/vllm_omni/deploy/moss_tts_nano.yaml new file mode 100644 index 000000000..585e244ca --- /dev/null +++ b/vllm_omni/deploy/moss_tts_nano.yaml @@ -0,0 +1,31 @@ +# MOSS-TTS-Nano deploy: single-stage AR TTS (0.1B). +# Verified on 1x L4 24GB (and sufficient for H20 / A100 / 3090). +# +# Fields omitted from a stage fall back to StageDeployConfig defaults (see +# vllm_omni/config/stage_config.py). Notable choices: +# * enforce_eager=true — the model uses upstream inference_stream() with +# trust_remote_code; CUDA graph capture is not wired up. +# * gpu_memory_utilization=0.3 — the AR LM + codec fit in ~2 GiB. +# * max_num_seqs=4 — mirrors the legacy stage_configs contract; per-request +# streaming generators are stored in self._stream_gens. +# +# MOSS-TTS-Nano has no multimodal input preprocessing, so skip_mm_profiling +# bypasses the dummy-MM pass at profile time. +async_chunk: false +trust_remote_code: true + +stages: + - stage_id: 0 + max_num_seqs: 4 + gpu_memory_utilization: 0.3 + enforce_eager: true + max_num_batched_tokens: 4096 + max_model_len: 4096 + devices: "0" + skip_mm_profiling: true + default_sampling_params: + temperature: 1.0 + top_p: 1.0 + top_k: 50 + max_tokens: 4096 + seed: 42 diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 08cf44d81..e1e165637 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -65,6 +65,7 @@ _OMNIVOICE_TTS_MODEL_STAGES = {"omnivoice_generator"} _VOXCPM_TTS_MODEL_STAGES = {"latent_generator", "vae"} _VOXCPM2_TTS_MODEL_STAGES = {"latent_generator"} _MING_TTS_MODEL_STAGES = {"ming_tts"} +_MOSS_TTS_MODEL_STAGES = {"moss_tts_nano"} _TTS_MODEL_STAGES: set[str] = ( _VOXTRAL_TTS_MODEL_STAGES | _QWEN3_TTS_MODEL_STAGES @@ -74,6 +75,7 @@ _TTS_MODEL_STAGES: set[str] = ( | _VOXCPM_TTS_MODEL_STAGES | _VOXCPM2_TTS_MODEL_STAGES | _MING_TTS_MODEL_STAGES + | _MOSS_TTS_MODEL_STAGES ) _TTS_LANGUAGES: set[str] = { "Auto", @@ -324,6 +326,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return "voxcpm" if has_vae_stage or model_stage == "vae" else "voxcpm2" if model_stage in _MING_TTS_MODEL_STAGES: return "ming_flash_omni_tts" + if model_stage in _MOSS_TTS_MODEL_STAGES: + return "moss_tts_nano" return None def _compute_max_instructions_length(self) -> int: @@ -855,6 +859,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return None # VoxCPM2 accepts any text input if self._tts_model_type == "ming_flash_omni_tts": return self._validate_ming_tts_request(request) + if self._tts_model_type == "moss_tts_nano": + return self._validate_moss_tts_request(request) return self._validate_qwen_tts_request(request) def _voxcpm2_encode(self, text: str) -> list[int]: @@ -1110,6 +1116,42 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return None + def _validate_moss_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: + """Validate MOSS-TTS-Nano request. Accepts optional ref_audio for voice cloning.""" + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + if request.ref_audio is not None: + fmt_err = self._validate_ref_audio_format(request.ref_audio) + if fmt_err: + return fmt_err + if not request.ref_text or not request.ref_text.strip(): + return "Voice cloning requires 'ref_text' (transcript of the reference audio)" + return None + + async def _build_moss_tts_params(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]: + """Build additional_information for MOSS-TTS-Nano. + + Maps the standard /v1/audio/speech fields to MOSS-TTS-Nano's + additional_information keys (text, voice, mode, prompt_audio_array, + prompt_text, ...). Values are wrapped in lists to match the + vllm-omni convention. When ``request.ref_audio`` is provided, it is + resolved via MediaConnector and passed as a (wav_list, sample_rate) + tuple so the model owns temp-file lifecycle. + """ + params: dict[str, Any] = { + "text": [request.input], + } + if request.voice is not None: + params["voice"] = [request.voice] + if request.ref_text is not None: + params["prompt_text"] = [request.ref_text] + if request.max_new_tokens is not None: + params["max_new_frames"] = [request.max_new_tokens] + if request.ref_audio is not None: + wav_list, sr = await self._resolve_ref_audio(request.ref_audio) + params["prompt_audio_array"] = [[wav_list, sr]] + return params + def _validate_fish_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: """Validate Fish Speech request parameters. Returns error message or None. @@ -1691,6 +1733,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): elif self._tts_model_type == "ming_flash_omni_tts": prompt = self._build_ming_prompt(request) tts_params = {} + elif self._tts_model_type == "moss_tts_nano": + tts_params = await self._build_moss_tts_params(request) + prompt = {"prompt_token_ids": [1], "additional_information": tts_params} else: tts_params = self._build_tts_params(request) # Resolve ref_audio (explicit or auto-set for uploaded voices) @@ -1740,6 +1785,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): model_type = "voxcpm2" elif self._tts_model_type == "ming_flash_omni_tts": model_type = "ming_flash_omni_tts" + elif self._tts_model_type == "moss_tts_nano": + model_type = "moss_tts_nano" elif self._is_tts: model_type = tts_params.get("task_type", ["unknown"])[0] else: diff --git a/vllm_omni/model_executor/models/moss_tts_nano/__init__.py b/vllm_omni/model_executor/models/moss_tts_nano/__init__.py new file mode 100644 index 000000000..896ead57d --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts_nano/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from .modeling_moss_tts_nano import MossTTSNanoForGeneration + +__all__ = ["MossTTSNanoForGeneration"] diff --git a/vllm_omni/model_executor/models/moss_tts_nano/configuration_moss_tts_nano.py b/vllm_omni/model_executor/models/moss_tts_nano/configuration_moss_tts_nano.py new file mode 100644 index 000000000..a5a9f25ac --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts_nano/configuration_moss_tts_nano.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Configuration for MOSS-TTS-Nano in vLLM-Omni single-stage pipeline.""" + +from transformers.configuration_utils import PretrainedConfig + + +class MossTTSNanoConfig(PretrainedConfig): + """Config for MOSS-TTS-Nano (OpenMOSS-Team/MOSS-TTS-Nano). + + The model is a Global-Local GPT2-based AR LM (0.1B) paired with the + MOSS-Audio-Tokenizer-Nano codec. Both components run in a single + vLLM-Omni generation stage. + + Relevant fields from config.json: + gpt2_config – backbone config (n_layer=12, n_embd=768, n_head=12, ...) + n_vq – number of RVQ codebooks (16) + audio_vocab_size – per-codebook vocabulary size (1024) + audio_tokenizer_pretrained_name_or_path – HF hub path for codec model + """ + + model_type = "moss_tts_nano" + + def __init__(self, **kwargs): + gpt2_cfg = kwargs.pop("gpt2_config", None) or {} + if hasattr(gpt2_cfg, "to_dict"): + gpt2_cfg = gpt2_cfg.to_dict() + + super().__init__(**kwargs) + + # --- GPT2 backbone parameters (exposed at top level for vLLM) --- + self.hidden_size = gpt2_cfg.get("n_embd", 768) + self.num_hidden_layers = gpt2_cfg.get("n_layer", 12) + self.num_attention_heads = gpt2_cfg.get("n_head", 12) + self.num_key_value_heads = self.num_attention_heads # no GQA in GPT2 + self.head_dim = self.hidden_size // self.num_attention_heads + self.vocab_size = gpt2_cfg.get("vocab_size", 16384) + self.max_position_embeddings = gpt2_cfg.get("n_positions", 32768) + self.intermediate_size = gpt2_cfg.get("n_inner", self.hidden_size * 4) + + # --- Audio codec parameters --- + self.n_vq: int = getattr(self, "n_vq", 16) + self.audio_vocab_size: int = getattr(self, "audio_vocab_size", 1024) + self.audio_start_token_id: int = getattr(self, "audio_start_token_id", 6) + self.audio_end_token_id: int = getattr(self, "audio_end_token_id", 7) + self.audio_user_slot_token_id: int = getattr(self, "audio_user_slot_token_id", 8) + self.audio_assistant_slot_token_id: int = getattr(self, "audio_assistant_slot_token_id", 9) + self.audio_tokenizer_pretrained_name_or_path: str = getattr( + self, + "audio_tokenizer_pretrained_name_or_path", + "OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano", + ) + self.audio_tokenizer_sample_rate: int = getattr(self, "audio_tokenizer_sample_rate", 48000) + + # vLLM requires speculative_config to be absent or None + self.speculative_config = None + + def get_text_config(self, **kwargs): + """Return self so vLLM uses our top-level config.""" + return self diff --git a/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py b/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py new file mode 100644 index 000000000..795a57733 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py @@ -0,0 +1,515 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MOSS-TTS-Nano single-stage model for vLLM-Omni. + +Runs in a single AR worker stage. The 0.1B AR LM and the +MOSS-Audio-Tokenizer-Nano codec are both loaded here. + +Streaming is supported via the VoxCPM-style generator pattern: + - On first forward() for a request, inference_stream() is started as + a Python generator and stored in self._stream_gens[request_key]. + - Each subsequent forward() call pops one audio chunk from the generator + and returns it as multimodal_outputs. + - compute_logits() emits EOS only when the last chunk has been yielded, + telling the AR scheduler to finish the request. + +Weight loading deliberately happens inside load_weights() -- NOT __init__ -- +so that vLLM initialises distributed state before any CUDA allocations occur. +""" + +from __future__ import annotations + +import tempfile +import threading +from collections.abc import Iterable +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.output_templates import OmniOutput + +logger = init_logger(__name__) + + +def _patch_torchaudio_load() -> None: + """Patch torchaudio.load to use soundfile if torchcodec is unavailable.""" + try: + import torchaudio + + torchaudio # noqa + import torchcodec # noqa: F401 + + return + except Exception: + pass + + import soundfile as sf + + def _soundfile_load(path, frame_offset=0, num_frames=-1, normalize=True, channels_first=True, format=None): + data, sr = sf.read(str(path), dtype="float32", always_2d=True) + if frame_offset > 0: + data = data[frame_offset:] + if num_frames > 0: + data = data[:num_frames] + waveform = torch.from_numpy(data) + if channels_first: + waveform = waveform.T + return waveform, sr + + def _soundfile_save(path, src, sample_rate, channels_first=True, **kwargs): + wav = src.detach().cpu().float().numpy() + if channels_first and wav.ndim == 2: + wav = wav.T + sf.write(str(path), wav, sample_rate) + + try: + import torchaudio + + torchaudio.load = _soundfile_load + torchaudio.save = _soundfile_save + logger.info("Patched torchaudio.load/save to use soundfile (torchcodec unavailable)") + except Exception as e: + logger.warning("Could not patch torchaudio: %s", e) + + +# Default sampling parameters matching the upstream demo defaults. +_DEFAULT_TEXT_TEMPERATURE = 1.0 +_DEFAULT_TEXT_TOP_P = 1.0 +_DEFAULT_TEXT_TOP_K = 50 +_DEFAULT_AUDIO_TEMPERATURE = 0.8 +_DEFAULT_AUDIO_TOP_P = 0.95 +_DEFAULT_AUDIO_TOP_K = 25 +_DEFAULT_AUDIO_REPETITION_PENALTY = 1.2 +_DEFAULT_MAX_NEW_FRAMES = 375 +_DEFAULT_VOICE = "Junhao" +# "voice_clone" mirrors the mode used by offline examples and tests: +# built-in voices (e.g. "Junhao") are backed by preset reference audio +# upstream, so the default produces the expected voice timbre. +_DEFAULT_MODE = "voice_clone" + + +def _pick(info: dict, key: str, default): + """Extract scalar from additional_information dict (list or plain value).""" + val = info.get(key, default) + if isinstance(val, (list, tuple)) and len(val) > 0: + return val[0] + return val if val is not None else default + + +class MossTTSNanoForGeneration(nn.Module): + """Single-stage MOSS-TTS-Nano model with streaming audio output. + + Uses the VoxCPM-style generator pattern: inference_stream() is stored + per-request and yields one audio chunk per forward() call. The AR + scheduler keeps the request alive until compute_logits() emits EOS. + """ + + requires_raw_input_tokens = True + have_multimodal_outputs = True + has_preprocess = False + has_postprocess = False + enable_update_additional_information = True + inject_omni_request_id_into_runtime_info = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.vllm_config = vllm_config + self.config = vllm_config.model_config.hf_config + self.model_path: str = vllm_config.model_config.model + + self._lm: nn.Module | None = None + self._audio_tokenizer: nn.Module | None = None + self._device: torch.device | None = None + self._lock = threading.Lock() + + # Per-request streaming generators (VoxCPM pattern). + # Single-threaded AR worker: mutations happen only inside forward(), + # so we rely on the worker's serial execution instead of self._lock. + self._stream_gens: dict[str, Any] = {} + # Per-row EOS mask aligned with the most recent forward() batch. + # compute_logits() emits EOS for rows flagged True, keeps others alive. + self._ar_last_chunk_flags: list[bool] = [] + + # ------------------------------------------------------------------ + # Weight loading + # ------------------------------------------------------------------ + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + with self._lock: + if self._lm is not None: + return set() + _patch_torchaudio_load() + + try: + device = next(self.parameters()).device + except StopIteration: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self._device = device + + if device.type == "cuda" and torch.cuda.is_bf16_supported(): + tts_dtype = torch.bfloat16 + elif device.type == "cuda": + tts_dtype = torch.float16 + else: + tts_dtype = torch.float32 + + logger.info("Loading MOSS-TTS-Nano LM from %s (dtype=%s)", self.model_path, tts_dtype) + from transformers import AutoModelForCausalLM + + lm = AutoModelForCausalLM.from_pretrained( + self.model_path, + trust_remote_code=True, + torch_dtype=tts_dtype, + ) + if device.type == "cuda": + try: + import flash_attn # noqa: F401 + + lm._set_attention_implementation("flash_attention_2") + logger.info("MOSS-TTS-Nano using flash_attention_2") + except ImportError: + lm._set_attention_implementation("sdpa") + logger.info("MOSS-TTS-Nano using sdpa (flash_attn not installed)") + lm.to(device=device) + lm.eval() + self._lm = lm + logger.info("MOSS-TTS-Nano LM loaded on %s", device) + + codec_path: str = getattr( + self.config, + "audio_tokenizer_pretrained_name_or_path", + "OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano", + ) + logger.info("Loading MOSS-Audio-Tokenizer-Nano from %s", codec_path) + from transformers import AutoModel + + audio_tokenizer = AutoModel.from_pretrained( + codec_path, + trust_remote_code=True, + torch_dtype=torch.float32, + ) + audio_tokenizer.to(device=device) + audio_tokenizer.eval() + self._audio_tokenizer = audio_tokenizer + logger.info("MOSS-Audio-Tokenizer-Nano loaded on %s", device) + + # MOSS-TTS-Nano loads weights inline via from_pretrained() above; + # vLLM's weight-loading protocol still requires us to exhaust the + # iterator so the underlying stream is closed cleanly. + for _ in weights: + pass + return set() + + # ------------------------------------------------------------------ + # Dummy run support + # ------------------------------------------------------------------ + + def get_dummy_runtime_additional_information(self, num_reqs: int) -> list[dict]: + return [{"text": "hello", "voice": _DEFAULT_VOICE, "_is_dummy": True}] * num_reqs + + # ------------------------------------------------------------------ + # Streaming generator management + # ------------------------------------------------------------------ + + def _create_stream_gen(self, info: dict[str, Any]): + """Create an inference_stream() generator for a request. + + Yields (waveform_tensor, is_last) tuples. + """ + text: str = str(_pick(info, "text", "") or "") + if not text.strip(): + logger.warning("MOSS-TTS-Nano received empty text; yielding silence.") + sr = getattr(self.config, "audio_tokenizer_sample_rate", 48000) + yield torch.zeros((sr,), dtype=torch.float32), True + return + + mode: str = str(_pick(info, "mode", _DEFAULT_MODE)) + prompt_audio_path: str | None = _pick(info, "prompt_audio_path", None) + if prompt_audio_path is not None: + prompt_audio_path = str(prompt_audio_path) + # Voice-cloning path: serving layer passes the resolved waveform as + # (wav_list, sample_rate) so we can avoid re-decoding base64 and the + # model owns temp-file lifecycle. + prompt_audio_array = _pick(info, "prompt_audio_array", None) + prompt_text: str | None = _pick(info, "prompt_text", None) + if prompt_text is not None: + prompt_text = str(prompt_text) + max_new_frames: int = int(_pick(info, "max_new_frames", _DEFAULT_MAX_NEW_FRAMES)) + seed: int | None = _pick(info, "seed", None) + if seed is not None: + # Upstream inference_stream relies on global RNG state. With + # max_num_seqs > 1, concurrent seeded requests will race to set + # the global seed; snapshot and restore on exit so we at least + # don't leak state back to other components of the engine. + seed = int(seed) + _cpu_rng_state = torch.get_rng_state() + _cuda_rng_state = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + else: + _cpu_rng_state = None + _cuda_rng_state = None + + sampling = { + "text_temperature": float(_pick(info, "text_temperature", _DEFAULT_TEXT_TEMPERATURE)), + "text_top_p": float(_pick(info, "text_top_p", _DEFAULT_TEXT_TOP_P)), + "text_top_k": int(_pick(info, "text_top_k", _DEFAULT_TEXT_TOP_K)), + "audio_temperature": float(_pick(info, "audio_temperature", _DEFAULT_AUDIO_TEMPERATURE)), + "audio_top_p": float(_pick(info, "audio_top_p", _DEFAULT_AUDIO_TOP_P)), + "audio_top_k": int(_pick(info, "audio_top_k", _DEFAULT_AUDIO_TOP_K)), + "audio_repetition_penalty": float( + _pick(info, "audio_repetition_penalty", _DEFAULT_AUDIO_REPETITION_PENALTY) + ), + } + + device = self._device or torch.device("cpu") + + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: + output_path = tmp.name + + # Materialise resolved reference audio to a temp WAV so upstream's + # prompt_audio_path (which expects a filesystem path) can consume it. + prompt_audio_tmp: str | None = None + if prompt_audio_array is not None and prompt_audio_path is None: + try: + import numpy as np + import soundfile as sf + + wav_list, sr_in = prompt_audio_array + wav_np = np.asarray(wav_list, dtype=np.float32) + if wav_np.ndim > 1: + wav_np = np.mean(wav_np, axis=-1) + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_in: + prompt_audio_tmp = tmp_in.name + sf.write(prompt_audio_tmp, wav_np, int(sr_in)) + prompt_audio_path = prompt_audio_tmp + except Exception: + logger.exception("MOSS-TTS-Nano failed to stage prompt_audio_array to temp file") + prompt_audio_tmp = None + + # Stream audio chunks from inference_stream as they arrive. A final + # "result" event is only used as a fallback when no audio events + # were emitted (e.g. very short utterances on some backends). + audio_chunks: list[torch.Tensor] = [] + try: + for event in self._lm.inference_stream( + text=text, + output_audio_path=output_path, + mode=mode, + prompt_text=prompt_text, + prompt_audio_path=prompt_audio_path, + text_tokenizer_path=self.model_path, + audio_tokenizer=self._audio_tokenizer, + device=device, + nq=None, + max_new_frames=max_new_frames, + do_sample=True, + use_kv_cache=True, + **sampling, + ): + event_type = str(event.get("type", "")) + if event_type == "audio": + waveform = event.get("waveform") + if waveform is not None: + chunk = torch.as_tensor(waveform, dtype=torch.float32).cpu() + if chunk.ndim == 2: + chunk = chunk.T.reshape(-1) + else: + chunk = chunk.reshape(-1) + audio_chunks.append(chunk) + # Yield each chunk as it arrives (is_last=False). + yield chunk, False + elif event_type == "result": + if not audio_chunks: + waveform = event.get("waveform") + if waveform is not None: + chunk = torch.as_tensor(waveform, dtype=torch.float32).cpu().reshape(-1) + yield chunk, True + return + except Exception: + logger.exception("MOSS-TTS-Nano inference failed for text=%r", text[:80]) + finally: + try: + Path(output_path).unlink(missing_ok=True) + except Exception: + pass + if prompt_audio_tmp is not None: + try: + Path(prompt_audio_tmp).unlink(missing_ok=True) + except Exception: + pass + if _cpu_rng_state is not None: + torch.set_rng_state(_cpu_rng_state) + if _cuda_rng_state is not None: + torch.cuda.set_rng_state_all(_cuda_rng_state) + + # Signal completion. If we yielded audio chunks above, the last + # yield was is_last=False, so emit a final empty sentinel. + yield torch.zeros((0,), dtype=torch.float32), True + + # ------------------------------------------------------------------ + # Core forward pass (streaming, VoxCPM pattern) + # ------------------------------------------------------------------ + + def _make_dummy_hidden(self, input_ids: torch.Tensor | None) -> torch.Tensor: + """Return a dummy hidden_states tensor for the AR runner.""" + device = self._device or torch.device("cpu") + hidden = int(getattr(self.config, "hidden_size", 768)) + n = 1 if input_ids is None else max(1, input_ids.shape[0]) + return torch.zeros((n, hidden), device=device, dtype=torch.float32) + + @torch.inference_mode() + def forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + intermediate_tensors: Any = None, + inputs_embeds: torch.Tensor | None = None, + runtime_additional_information: list[dict[str, Any]] | None = None, + **kwargs: Any, + ) -> OmniOutput: + sr = getattr(self.config, "audio_tokenizer_sample_rate", 48000) + sr_tensor = torch.tensor(sr, dtype=torch.int32) + empty = torch.zeros((0,), dtype=torch.float32) + hidden = self._make_dummy_hidden(input_ids) + + infos = runtime_additional_information or [{}] + + if not runtime_additional_information or all(info.get("_is_dummy") for info in infos): + # Dummy/warmup path: finish immediately for every row. + self._ar_last_chunk_flags = [True] * len(infos) + return OmniOutput( + text_hidden_states=hidden, + multimodal_outputs={ + "model_outputs": [empty] * len(infos), + "sr": [sr_tensor] * len(infos), + }, + ) + + if self._lm is None or self._audio_tokenizer is None: + raise RuntimeError("MOSS-TTS-Nano model not loaded. Was load_weights() called?") + + outputs: list[torch.Tensor] = [] + srs: list[torch.Tensor] = [] + last_chunk_flags: list[bool] = [] + + for info in infos: + if info.get("_is_dummy"): + outputs.append(empty) + srs.append(sr_tensor) + last_chunk_flags.append(True) + continue + + request_key = str(info.get("_omni_req_id", "0")) + + # Create generator on first call for this request. + if request_key not in self._stream_gens: + self._stream_gens[request_key] = self._create_stream_gen(info) + + generator = self._stream_gens[request_key] + try: + chunk, is_last = next(generator) + except StopIteration: + self._stream_gens.pop(request_key, None) + outputs.append(empty) + last_chunk_flags.append(True) + else: + if is_last: + self._stream_gens.pop(request_key, None) + outputs.append(chunk) + last_chunk_flags.append(bool(is_last)) + + srs.append(sr_tensor) + + # Per-row EOS mask: compute_logits() finishes each request the step + # its last chunk was yielded, without waiting for the slowest peer. + self._ar_last_chunk_flags = last_chunk_flags + + return OmniOutput( + text_hidden_states=hidden, + multimodal_outputs={"model_outputs": outputs, "sr": srs}, + ) + + def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None: + """Release streaming generators for requests the scheduler finished. + + forward() only pops generators on normal completion (is_last or + StopIteration). Abnormal termination (cancel, timeout, preempt) would + otherwise leak the generator and skip its ``finally`` block, stranding + temp WAV files. Closing here raises GeneratorExit inside the generator + so the cleanup block runs. + """ + for req_id in finished_req_ids: + gen = self._stream_gens.pop(str(req_id), None) + if gen is not None: + try: + gen.close() + except Exception: + logger.exception("MOSS-TTS-Nano failed to close stream gen for request %s", req_id) + + # ------------------------------------------------------------------ + # AR runner interface + # ------------------------------------------------------------------ + + def compute_logits( + self, + hidden_states: torch.Tensor | OmniOutput, + sampling_metadata: Any = None, + ) -> torch.Tensor: + """Emit per-row EOS / non-EOS logits to control AR scheduler lifetime. + + Rows whose ``_ar_last_chunk_flags`` entry is True get EOS-dominant + logits so the scheduler finishes that request; other rows get a + non-EOS token so they stay alive for the next streaming chunk. + """ + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + + if hidden_states is None: + device = self._device or torch.device("cpu") + hidden_states = torch.zeros((0, 1), device=device, dtype=torch.float32) + if hidden_states.ndim == 1: + hidden_states = hidden_states.unsqueeze(-1) + elif hidden_states.ndim > 2: + hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) + + vocab_size = int(getattr(self.config, "vocab_size", 32000)) + num_rows = int(hidden_states.shape[0]) + logits = torch.zeros( + (num_rows, vocab_size), + dtype=torch.float32, + device=hidden_states.device, + ) + eos_id = 2 if vocab_size > 2 else 0 + safe_id = 1 if vocab_size > 1 and 1 != eos_id else 0 + + flags = self._ar_last_chunk_flags + # If we lost alignment with the forward batch (e.g. scheduler dropped + # a row), conservatively treat missing rows as "finished" to avoid + # stranding requests. + for row in range(num_rows): + is_last = flags[row] if row < len(flags) else True + if is_last: + logits[row, eos_id] = 1.0e6 + else: + logits[row, eos_id] = -1.0e9 + logits[row, safe_id] = 1.0e6 + return logits + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings=None, + is_multimodal=None, + ) -> torch.Tensor: + hidden = int(getattr(self.config, "hidden_size", 768)) + return torch.zeros( + (input_ids.shape[0], hidden), + device=input_ids.device, + dtype=torch.float32, + ) diff --git a/vllm_omni/model_executor/models/moss_tts_nano/pipeline.py b/vllm_omni/model_executor/models/moss_tts_nano/pipeline.py new file mode 100644 index 000000000..749c4f980 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts_nano/pipeline.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MOSS-TTS-Nano pipeline topology (frozen). + +Single-stage AR TTS: text -> speech waveform in one pass. The 0.1B AR LM +and the MOSS-Audio-Tokenizer-Nano codec both run inside +``MossTTSNanoForGeneration.forward()``, which uses the VoxCPM-style +generator pattern (``inference_stream()`` stored per-request; one audio +chunk yielded per forward) to drive progressive streaming through the +AR scheduler. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +MOSS_TTS_NANO_PIPELINE = PipelineConfig( + model_type="moss_tts_nano", + model_arch="MossTTSNanoForCausalLM", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="moss_tts_nano", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=True, + final_output_type="audio", + owns_tokenizer=True, + engine_output_type="audio", + sampling_constraints={ + "detokenize": False, + # compute_logits() forces EOS (token id 2) when the last + # streaming chunk is yielded; keep a hard backstop here. + "stop_token_ids": [2], + }, + ), + ), +) diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 0cbe60610..5c2ef65d0 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -174,6 +174,12 @@ _OMNI_MODELS = { "VoxtralTTSAudioGenerationForConditionalGeneration", ), "VoxtralTTSAudioTokenizer": ("voxtral_tts", "voxtral_tts_audio_tokenizer", "VoxtralTTSAudioTokenizer"), + ## MOSS-TTS-Nano + "MossTTSNanoForCausalLM": ( + "moss_tts_nano", + "modeling_moss_tts_nano", + "MossTTSNanoForGeneration", + ), "DyninOmniForConditionalGeneration": ( "dynin_omni", "dynin_omni",