Compare commits

...
Author SHA1 Message Date
Claude 05883ff470 [CI] Fix moss_tts_nano test import paths to unblock collection
The test files imported from `tests.utils` (non-existent) and
`tests.conftest.OmniServerParams` (not lazy-exported), causing pytest
collection to fail for the whole suite. This broke unrelated CI jobs
(Diffusion Cache Backend Test, Simple Unit Test) because collection
runs across all test files even when they are deselected by marker.

Align with the pattern used by sibling TTS tests
(e.g. test_cosyvoice3_tts.py, test_voxcpm.py):
- hardware_test -> tests.helpers.mark
- OmniServerParams -> tests.helpers.runtime

https://claude.ai/code/session_01JYESRH7Ua1mc1wqKCQQs5q
2026-04-24 15:44:03 +00:00
Yueqian Lin 74522ca3ec Merge branch 'main' into feat/moss-tts-nano
Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-04-24 11:00:13 -04:00
Yueqian Lin 0a43fb94ac style(moss_tts_nano): ruff-format on_requests_finished log call
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-22 10:12:07 -04:00
Yueqian Lin 123dfc0772 fix(moss_tts_nano): release stream generators on abnormal termination
forward() only popped _stream_gens on normal completion, so cancelled,
timed-out, or preempted requests leaked their generator and skipped the
finally block that unlinks the temp WAV files. Implement
on_requests_finished to close each finished generator, which raises
GeneratorExit inside it and triggers the existing cleanup.

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-22 10:09:58 -04:00
Yueqian Lin b03a41a7f9 refactor(moss_tts_nano): migrate to Pipeline + Deploy schema
Follow-up to #2958. MOSS-TTS-Nano was authored before the schema refactor
and still shipped as a legacy stage_configs yaml; this aligns it with the
same layout the 5 migrated TTS models now use.

- vllm_omni/model_executor/models/moss_tts_nano/pipeline.py declares
  MOSS_TTS_NANO_PIPELINE (single LLM_AR stage, owns_tokenizer, audio
  output, stop_token_ids=[2] as a hard EOS backstop).
- vllm_omni/deploy/moss_tts_nano.yaml holds runtime knobs (max_num_seqs,
  gpu_memory_utilization, enforce_eager, default_sampling_params,
  skip_mm_profiling); trust_remote_code stays at deploy top-level.
- vllm_omni/config/pipeline_registry.py registers the entry so the lazy
  registry can resolve it.
- moss_tts_nano/__init__.py exports MossTTSNanoForGeneration (VoxCPM2 pattern).
- Removed vllm_omni/model_executor/stage_configs/moss_tts_nano.yaml.

Examples, shell scripts, READMEs, and buildkite-invoked tests are
updated to use `vllm serve <model> --omni` / `--deploy-config`
(auto-load kicks in; no --stage-configs-path or --trust-remote-code).

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-21 18:07:50 -04:00
Yueqian Lin 26bbe7b527 fix(moss_tts_nano): address PR review feedback
- Wire ref_audio through /v1/audio/speech: _build_moss_tts_params is now
  async, resolves ref_audio via MediaConnector and passes the waveform as
  prompt_audio_array so the model can materialise a temp WAV for upstream
  inference_stream (previously voice cloning via the REST endpoint was a
  no-op). Serving now also requires ref_text when ref_audio is provided.
- Fix per-request EOS in batched decode: replace the shared
  _ar_emit_stop_token bool with a _ar_last_chunk_flags list so
  compute_logits emits EOS per row; finished requests no longer wait for
  the slowest peer in a max_num_seqs=4 batch.
- Snapshot and restore CPU+CUDA RNG state around torch.manual_seed to
  limit global-state bleed; add comment noting deterministic output
  under concurrent batching is best-effort (upstream inference_stream
  uses the global RNG).
- Align _DEFAULT_MODE with the offline example and tests ("voice_clone").
- Clean up outdated "buffer first" comment in _create_stream_gen; document
  single-threaded AR-worker assumption for _stream_gens; add one-liner
  explaining the load_weights iterator drain.

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-21 17:58:50 -04:00
Yueqian Lin 3bd77d3f8f Merge branch 'main' into feat/moss-tts-nano
Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-04-21 17:50:21 -04:00
Yueqian Lin 033e4041be fix(moss-tts-nano): remove unused voice variable (ruff F841)
The upstream inference_stream() API has no voice/spk parameter;
voice preset selection is not yet wired into the call. Remove the
dead assignment to silence ruff F841.

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-14 13:11:43 -04:00
Yueqian Lin 38aaaebeda feat(moss-tts-nano): enable progressive streaming via AR runner
Switch from GPUGenerationWorker to AR worker with OmniARScheduler,
using the VoxCPM-style generator pattern for streaming:

- inference_stream() stored per-request in _stream_gens dict
- Each forward() call yields one audio chunk via next(generator)
- compute_logits() emits EOS only when last chunk is yielded
- AR scheduler loops until EOS, enabling progressive audio output

TTFP reduced from ~3.1s to ~0.11s (30x improvement).

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-14 13:06:44 -04:00
Yueqian Lin 71db0c5ef5 fix(moss-tts-nano): add buildkite pre-merge test, remove misplaced docs
- Add MOSS-TTS-Nano E2E test to .buildkite/test-merge.yml
- Remove docs from docs/ (will be synced from examples/ READMEs)

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-14 13:06:44 -04:00
Yueqian Lin 45e8829f07 docs(moss-tts-nano): add e2e online serving test and user guide docs
- Add online serving e2e test (non-streaming, streaming, Chinese)
- Add online serving user guide doc (API, voices, Gradio, curl/Python)
- Add offline inference user guide doc

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-14 13:06:44 -04:00
Yueqian Lin fc81b8504a feat(moss-tts-nano): add online serving, Gradio demo, and fix /v1/audio/speech integration
- Register moss_tts_nano as TTS model stage in serving_speech.py
  (model detection, validation, prompt building)
- Fix registry mod_relname: moss_tts_nano -> modeling_moss_tts_nano
- Fix stage config: is_comprehension=true (required for generate task)
- Fix default mode: voice_clone -> continuation (built-in presets)
- Add compute_logits stub for VllmModelForTextGeneration protocol
- Remove unused _sentinel nn.Parameter
- Add Gradio demo with AudioWorklet streaming player (48kHz stereo)
- Add online/offline serving docs and launch scripts

TODO: Single-stage generation models don't support true streaming
(progressive audio chunks). Current TTFP = full generation time.
Needs multi-step scheduling support in GPUGenerationModelRunner.

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-14 13:06:44 -04:00
Yueqian Lin 571c8d4235 fix(moss-tts-nano): also patch torchaudio.save for soundfile fallback
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-14 13:06:12 -04:00
Yueqian Lin 0b8190d07f fix(moss-tts-nano): patch torchaudio.load to use soundfile when torchcodec unavailable
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-14 13:06:12 -04:00
Yueqian Lin da6b1b5741 feat(moss-tts-nano): add single-stage MOSS-TTS-Nano TTS support
Integrates OpenMOSS-Team/MOSS-TTS-Nano (0.1B) as a single
GPUGenerationWorker stage.  Both the AR LM and MOSS-Audio-Tokenizer-Nano
codec run inside MossTTSNanoForGeneration.forward(), removing the need for
an inter-stage connector.

Key design choices:
- Weights loaded in load_weights() not __init__ (avoids pre-CUDA alloc)
- trust_remote_code delegates to upstream HF model classes
- codec path read from config.audio_tokenizer_pretrained_name_or_path
- inference_stream() collects progressive audio chunks for low latency
- 48 kHz stereo output; voice clone + continuation modes supported


Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-04-14 13:06:12 -04:00
17 changed files with 2169 additions and 0 deletions

No files matched your search

+22
View File
@@ -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
@@ -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).
@@ -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())
@@ -0,0 +1,158 @@
# MOSS-TTS-Nano
## Model
| 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.
@@ -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 = """
<div id="tts-player">
<div style="display:flex; align-items:center; gap:10px;">
<div id="tts-status-dot" style="width:10px;height:10px;border-radius:50%;background:#ccc;flex-shrink:0;"></div>
<span id="tts-status" style="font-weight:600;font-size:1.05em;">Ready</span>
<button id="tts-stop-btn" onclick="window.ttsStop()"
style="display:none; margin-left:auto; padding:5px 16px; border-radius:6px; border:1px solid #EF5552;
background:#fff; color:#EF5552; cursor:pointer; font-size:0.85em; transition:all 0.15s;">Stop</button>
</div>
<div id="tts-metrics" style="display:none; grid-template-columns:repeat(4,1fr); gap:10px; margin-top:12px;">
<div style="background:#f8f9fa;border-radius:6px;padding:8px 10px;text-align:center;">
<div style="font-size:0.7em;text-transform:uppercase;color:#888;letter-spacing:0.5px;margin-bottom:2px;">TTFP</div>
<div id="tts-m-ttfp" style="font-size:1.2em;font-weight:700;color:#333;">—</div>
</div>
<div style="background:#f8f9fa;border-radius:6px;padding:8px 10px;text-align:center;">
<div style="font-size:0.7em;text-transform:uppercase;color:#888;letter-spacing:0.5px;margin-bottom:2px;">RTF</div>
<div id="tts-m-rtf" style="font-size:1.2em;font-weight:700;color:#333;">—</div>
</div>
<div style="background:#f8f9fa;border-radius:6px;padding:8px 10px;text-align:center;">
<div style="font-size:0.7em;text-transform:uppercase;color:#888;letter-spacing:0.5px;margin-bottom:2px;">Audio</div>
<div id="tts-m-dur" style="font-size:1.2em;font-weight:700;color:#333;">—</div>
</div>
<div style="background:#f8f9fa;border-radius:6px;padding:8px 10px;text-align:center;">
<div style="font-size:0.7em;text-transform:uppercase;color:#888;letter-spacing:0.5px;margin-bottom:2px;">Speed</div>
<div id="tts-m-speed" style="font-size:1.2em;font-weight:700;color:#333;">—</div>
</div>
</div>
<div id="tts-rtf-bar-wrap" style="display:none; background:#e8ecf1; border-radius:4px; height:20px; overflow:hidden; position:relative; margin-top:10px;">
<div id="tts-rtf-bar" style="height:100%; border-radius:4px; transition:width 0.3s ease, background 0.3s ease; width:0%;"></div>
<span id="tts-rtf-label" style="position:absolute;top:50%;left:50%;transform:translate(-50%,-50%);font-size:0.75em;font-weight:600;color:#444;"></span>
</div>
<div id="tts-elapsed" style="display:none; margin-top:6px; font-size:0.8em; color:#999; text-align:right;"></div>
</div>
"""
def _build_player_js(sample_rate: int) -> str:
"""Build the JavaScript that powers the AudioWorklet player."""
return f"""
<script>
const SR = {sample_rate};
const WC = {json.dumps(WORKLET_JS)};
let ctx = null, node = null, abort = null, gen = false, st = {{}};
async function init() {{
if (ctx) return;
ctx = new AudioContext({{ sampleRate: SR }});
const b = new Blob([WC], {{ type: 'application/javascript' }});
const u = URL.createObjectURL(b);
await ctx.audioWorklet.addModule(u);
URL.revokeObjectURL(u);
node = new AudioWorkletNode(ctx, 'tts-playback-processor');
node.connect(ctx.destination);
node.port.onmessage = (e) => {{
if (e.data.type === 'started') setStatus('Playing...', '#2E7D32');
else if (e.data.type === 'stopped' && !gen) {{
setStatus('Done', '#2E7D32'); showStats(true);
const btn = document.getElementById('tts-stop-btn');
if (btn) btn.style.display = 'none';
}}
}};
}}
function setStatus(text, color) {{
const s = document.getElementById('tts-status');
const d = document.getElementById('tts-status-dot');
if (s) s.textContent = text;
if (d) d.style.background = color || '#ccc';
}}
function showStats(fin) {{
if (!st.t0) return;
const elapsed = (fin && st.streamEnd ? (st.streamEnd - st.t0) : (performance.now() - st.t0)) / 1000;
// st.samples counts stereo int16 samples; mono frames = samples / 2
const dur = (st.samples / 2) / SR;
const mTtfp = document.getElementById('tts-m-ttfp');
const mRtf = document.getElementById('tts-m-rtf');
const mDur = document.getElementById('tts-m-dur');
const mSpeed = document.getElementById('tts-m-speed');
const bar = document.getElementById('tts-rtf-bar');
const barLabel = document.getElementById('tts-rtf-label');
const elapsedEl = document.getElementById('tts-elapsed');
if (mTtfp && st.ttfp != null) mTtfp.textContent = st.ttfp.toFixed(0) + 'ms';
if (mDur) mDur.textContent = dur.toFixed(1) + 's';
if (dur > 0 && elapsed > 0) {{
const rtf = elapsed / dur;
const speed = 1 / rtf;
if (mRtf) {{
mRtf.textContent = rtf.toFixed(2) + 'x';
mRtf.style.color = rtf < 1 ? '#2E7D32' : rtf < 1.5 ? '#e8a317' : '#EF5552';
}}
if (mSpeed) {{
mSpeed.textContent = speed.toFixed(1) + 'x';
mSpeed.style.color = speed > 1 ? '#2E7D32' : speed > 0.7 ? '#e8a317' : '#EF5552';
}}
if (bar) {{
const pct = Math.min(speed / 3 * 100, 100);
bar.style.width = pct + '%';
bar.style.background = speed > 1 ? 'linear-gradient(90deg,#66BB6A,#2E7D32)' : speed > 0.7 ? 'linear-gradient(90deg,#e8a317,#f0b866)' : 'linear-gradient(90deg,#EF5552,#f87171)';
}}
if (barLabel) barLabel.textContent = speed.toFixed(1) + 'x realtime';
}}
if (elapsedEl) {{
elapsedEl.style.display = 'block';
elapsedEl.textContent = fin ? 'Completed in ' + elapsed.toFixed(1) + 's (' + st.chunks + ' chunks)' : elapsed.toFixed(1) + 's elapsed (' + st.chunks + ' chunks)';
}}
}}
window.ttsStop = function() {{
if (abort) abort.abort();
if (node) node.port.postMessage({{ type: 'clear' }});
gen = false;
setStatus('Stopped', '#999');
const btn = document.getElementById('tts-stop-btn');
if (btn) btn.style.display = 'none';
}};
window.ttsGenerate = async function(payload) {{
try {{ await init(); if (ctx.state === 'suspended') await ctx.resume(); }}
catch (e) {{ const s = document.getElementById('tts-status'); if (s) s.textContent = 'Audio init error: ' + e.message; return; }}
if (abort) abort.abort();
node.port.postMessage({{ type: 'clear' }});
await new Promise(r => setTimeout(r, 50));
node.port.postMessage({{ type: 'clear' }});
gen = true;
st = {{ t0: null, chunks: 0, samples: 0, ttfp: null }};
setStatus('Connecting...', '#2E7D32');
const bEl = document.getElementById('tts-stop-btn');
if (bEl) bEl.style.display = 'inline-block';
const mp = document.getElementById('tts-metrics');
if (mp) {{ mp.style.display = 'grid'; ['tts-m-ttfp','tts-m-rtf','tts-m-dur','tts-m-speed'].forEach(id => {{ const e = document.getElementById(id); if(e) {{ e.textContent = '\\u2014'; e.style.color = '#333'; }} }}); }}
const bw = document.getElementById('tts-rtf-bar-wrap');
if (bw) bw.style.display = 'block';
const bar = document.getElementById('tts-rtf-bar');
if (bar) bar.style.width = '0%';
const bl = document.getElementById('tts-rtf-label');
if (bl) bl.textContent = '';
const ee = document.getElementById('tts-elapsed');
if (ee) {{ ee.style.display = 'none'; ee.textContent = ''; }}
abort = new AbortController();
try {{
st.t0 = performance.now();
const r = await fetch('/proxy/v1/audio/speech', {{
method: 'POST',
headers: {{ 'Content-Type': 'application/json' }},
body: JSON.stringify(payload),
signal: abort.signal,
}});
if (!r.ok) {{ const t = await r.text(); throw new Error('Server ' + r.status + ': ' + t.slice(0, 200)); }}
setStatus('Streaming...', '#43A047');
const reader = r.body.getReader();
let left = new Uint8Array(0);
while (true) {{
const {{ done, value }} = await reader.read();
if (done) break;
let raw;
if (left.length > 0) {{
raw = new Uint8Array(left.length + value.length);
raw.set(left); raw.set(value, left.length);
}} else {{ raw = value; }}
const usable = raw.length - (raw.length % 2);
left = usable < raw.length ? raw.slice(usable) : new Uint8Array(0);
if (usable > 0) {{
const ab = new ArrayBuffer(usable);
new Uint8Array(ab).set(raw.subarray(0, usable));
const pcm = new Int16Array(ab);
node.port.postMessage(pcm);
st.chunks++;
st.samples += pcm.length;
if (st.ttfp == null) st.ttfp = performance.now() - st.t0;
showStats(false);
}}
}}
}} catch (e) {{
if (e.name !== 'AbortError') {{
setStatus('Error: ' + e.message, '#EF5552');
console.error('TTS error:', e);
}}
}} finally {{
st.streamEnd = performance.now();
showStats(true);
gen = false;
if (st.samples > 0) setStatus('Finishing playback...', '#2E7D32');
else {{
setStatus('No audio received', '#999');
if (bEl) bEl.style.display = 'none';
}}
}}
}};
</script>
"""
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"""
<div style="display:flex; align-items:center; gap:16px; margin-bottom:8px;">
<img src="https://raw.githubusercontent.com/vllm-project/vllm-omni/main/docs/source/logos/vllm-omni-logo.png"
alt="vLLM-Omni" style="height:42px;">
<div>
<h1 style="margin:0; font-size:1.5em;">MOSS-TTS-Nano Streaming Demo</h1>
<span style="font-size:0.85em; color:#666;">
0.1B AR LM + MOSS-Audio-Tokenizer-Nano &middot; 48 kHz &middot; ZH / EN / JA
&nbsp;&middot;&nbsp; Served by
<a href="https://github.com/vllm-project/vllm-omni" target="_blank"
style="color:#43A047; text-decoration:none; font-weight:600;">vLLM-Omni</a>
&nbsp;&middot;&nbsp; <code style="background:#eef2f7; padding:2px 6px; border-radius:4px; font-size:0.9em;">{api_base}</code>
</span>
</div>
</div>
""")
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("""
<div style="text-align:center; padding:8px 0; margin-top:4px;">
<a href="https://github.com/vllm-project/vllm-omni" target="_blank">
<img src="https://raw.githubusercontent.com/vllm-project/vllm-omni/main/docs/source/logos/vllm-omni-logo.png"
alt="vLLM-Omni" style="height:28px; opacity:0.7;">
</a>
</div>
""")
# 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()
+43
View File
@@ -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"
+19
View File
@@ -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
@@ -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
@@ -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)
+4
View File
@@ -79,6 +79,10 @@ _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",
),
}
# --- Single-stage diffusion pipelines (populated in PR 3/N) ---
+31
View File
@@ -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
@@ -63,6 +63,7 @@ _COSYVOICE3_TTS_MODEL_STAGES = {"cosyvoice3_talker"}
_OMNIVOICE_TTS_MODEL_STAGES = {"omnivoice_generator"}
_VOXCPM_TTS_MODEL_STAGES = {"latent_generator", "vae"}
_VOXCPM2_TTS_MODEL_STAGES = {"latent_generator"}
_MOSS_TTS_MODEL_STAGES = {"moss_tts_nano"}
_MING_TTS_MODEL_STAGES = {"ming_tts"}
_TTS_MODEL_STAGES: set[str] = (
_VOXTRAL_TTS_MODEL_STAGES
@@ -72,6 +73,7 @@ _TTS_MODEL_STAGES: set[str] = (
| _OMNIVOICE_TTS_MODEL_STAGES
| _VOXCPM_TTS_MODEL_STAGES
| _VOXCPM2_TTS_MODEL_STAGES
| _MOSS_TTS_MODEL_STAGES
| _MING_TTS_MODEL_STAGES
)
_TTS_LANGUAGES: set[str] = {
@@ -315,6 +317,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
return "cosyvoice3"
if model_stage in _OMNIVOICE_TTS_MODEL_STAGES:
return "omnivoice"
if model_stage in _VOXCPM2_TTS_MODEL_STAGES:
return "voxcpm2"
if model_stage in _MOSS_TTS_MODEL_STAGES:
return "moss_tts_nano"
if model_stage in (_VOXCPM_TTS_MODEL_STAGES | _VOXCPM2_TTS_MODEL_STAGES):
has_vae_stage = any(
getattr(getattr(stage, "engine_args", None), "model_stage", None) == "vae"
@@ -852,6 +858,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
return self._validate_voxcpm_request(request)
if self._tts_model_type == "voxcpm2":
return None # VoxCPM2 accepts any text input
if self._tts_model_type == "moss_tts_nano":
return self._validate_moss_tts_request(request)
if self._tts_model_type == "ming_flash_omni_tts":
return self._validate_ming_tts_request(request)
return self._validate_qwen_tts_request(request)
@@ -1109,6 +1117,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.
@@ -1680,6 +1724,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
elif self._tts_model_type == "cosyvoice3":
prompt = await self._build_cosyvoice3_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}
elif self._tts_model_type == "ming_flash_omni_tts":
prompt = self._build_ming_prompt(request)
tts_params = {}
@@ -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"]
@@ -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
@@ -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,
)
@@ -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],
},
),
),
)
@@ -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",