From da6b1b5741b1364efc2f7d684f96bbcdb457fa0b Mon Sep 17 00:00:00 2001 From: Yueqian Lin Date: Mon, 13 Apr 2026 14:30:37 -0400 Subject: [PATCH] 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 --- .../moss_tts_nano/end2end.py | 184 +++++++++ .../offline_inference/test_moss_tts_nano.py | 132 +++++++ .../models/moss_tts_nano/__init__.py | 2 + .../configuration_moss_tts_nano.py | 60 +++ .../moss_tts_nano/modeling_moss_tts_nano.py | 361 ++++++++++++++++++ vllm_omni/model_executor/models/registry.py | 6 + .../stage_configs/moss_tts_nano.yaml | 61 +++ 7 files changed, 806 insertions(+) create mode 100644 examples/offline_inference/moss_tts_nano/end2end.py create mode 100644 tests/e2e/offline_inference/test_moss_tts_nano.py 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/stage_configs/moss_tts_nano.yaml 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..9eb0be671 --- /dev/null +++ b/examples/offline_inference/moss_tts_nano/end2end.py @@ -0,0 +1,184 @@ +# 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" + +_REPO_ROOT = Path(__file__).resolve().parents[4] +DEFAULT_STAGE_CONFIG = str( + _REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "moss_tts_nano.yaml" +) + +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, + stage_configs_path=args.stage_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("--stage-config", default=DEFAULT_STAGE_CONFIG) + parser.add_argument("--stage-init-timeout", type=int, default=120) + return parser.parse_args() + + +if __name__ == "__main__": + main(parse_args()) 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..ea1236c62 --- /dev/null +++ b/tests/e2e/offline_inference/test_moss_tts_nano.py @@ -0,0 +1,132 @@ +# 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.utils import hardware_test +from vllm_omni import Omni + +MODEL_NAME = "OpenMOSS-Team/MOSS-TTS-Nano" +STAGE_CONFIG = "vllm_omni/model_executor/stage_configs/moss_tts_nano.yaml" +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_configs_path=STAGE_CONFIG, 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/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..208f01a7c --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts_nano/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project 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..e0f3a9790 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py @@ -0,0 +1,361 @@ +# 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 GPUGenerationWorker stage. The 0.1B AR LM and the +MOSS-Audio-Tokenizer-Nano codec are both loaded here; the full +text → audio-codes → waveform pipeline executes inside forward(). + +Streaming (realtime) generation is supported via model.inference_stream(): +audio chunks are collected as the LM emits tokens and returned as a +concatenated waveform. Future work can expose per-chunk streaming via +vLLM-Omni's SSE audio endpoint. + +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 logging +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__) + +# 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" +_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 generation model. + + Integrates the 0.1B MOSS-TTS-Nano LM and the MOSS-Audio-Tokenizer-Nano + codec into a single vLLM-Omni generation stage (GPUGenerationWorker). + + The model uses trust_remote_code=True to load the upstream HuggingFace + model classes, keeping our implementation thin and easy to update as + the upstream model evolves. + + Supported request fields in additional_information: + text (str) – text to synthesize [required] + voice (str) – built-in voice preset name, default "Junhao" + mode (str) – "voice_clone" (default) or "continuation" + prompt_audio_path (str) – path to reference WAV/MP3 for custom voice clone + prompt_text (str) – reference transcript (for continuation mode) + max_new_frames (int) – max AR frames, default 375 (~14s audio) + text_temperature (float) – LM text-layer temperature + text_top_p (float) – LM text-layer top-p + text_top_k (int) – LM text-layer top-k + audio_temperature (float) – LM audio-layer temperature + audio_top_p (float) – LM audio-layer top-p + audio_top_k (int) – LM audio-layer top-k + audio_repetition_penalty (float) – LM audio repetition penalty + seed (int) – optional random seed for reproducibility + """ + + requires_raw_input_tokens = True + have_multimodal_outputs = True + has_preprocess = False + has_postprocess = False + enable_update_additional_information = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.config = vllm_config.model_config.hf_config + self.model_path: str = vllm_config.model_config.model + + # Populated by load_weights(); kept None until then so __init__ is + # allocation-free (vLLM may call __init__ before CUDA is ready). + self._lm: nn.Module | None = None + self._audio_tokenizer: nn.Module | None = None + self._device: torch.device | None = None + self._lock = threading.Lock() + + # Dummy parameter so vLLM can resolve the device via next(parameters()). + self._sentinel = nn.Parameter(torch.zeros(1), requires_grad=False) + + # ------------------------------------------------------------------ + # Weight loading + # ------------------------------------------------------------------ + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load the MOSS-TTS-Nano LM and MOSS-Audio-Tokenizer-Nano codec. + + Both models are loaded via trust_remote_code so we stay in sync + with any upstream architecture changes without reimplementing + the model code ourselves. + """ + with self._lock: + if self._lm is not None: + return set() + + try: + device = next(self.parameters()).device + except StopIteration: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + self._device = device + + # Select compute dtype: bfloat16 on capable CUDA, float32 elsewhere. + 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 + + # --- Load the AR language model --- + 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": + # Use SDPA when flash_attention_2 is unavailable. + try: + lm._set_attention_implementation("flash_attention_2") + except Exception: + lm._set_attention_implementation("sdpa") + lm.to(device=device) + lm.eval() + self._lm = lm + logger.info("MOSS-TTS-Nano LM loaded on %s", device) + + # --- Load the audio tokenizer (codec) --- + 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, # codec runs in fp32 for stability + ) + audio_tokenizer.to(device=device) + audio_tokenizer.eval() + self._audio_tokenizer = audio_tokenizer + logger.info("MOSS-Audio-Tokenizer-Nano loaded on %s", device) + + # Consume the vLLM weight iterator (we loaded weights ourselves above). + for _ in weights: + pass + return set() + + # ------------------------------------------------------------------ + # Dummy run support (vLLM profiling / KV-cache estimation) + # ------------------------------------------------------------------ + + def get_dummy_runtime_additional_information(self, num_reqs: int) -> list[dict]: + return [{"text": "hello", "voice": _DEFAULT_VOICE, "_is_dummy": True}] * num_reqs + + # ------------------------------------------------------------------ + # Core forward pass + # ------------------------------------------------------------------ + + @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: + """Run full MOSS-TTS-Nano pipeline: text → audio codes → waveform. + + Each request in the batch is processed independently. Results are + returned as lists so the generation runner can demultiplex them. + """ + sr = getattr(self.config, "audio_tokenizer_sample_rate", 48000) + sr_tensor = torch.tensor(sr, dtype=torch.int32) + empty = torch.zeros((0,), dtype=torch.float32) + + if not runtime_additional_information: + # Profiling / dummy run — return empty audio. + n = 1 if input_ids is None else max(1, input_ids.shape[0]) + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={ + "model_outputs": [empty] * n, + "sr": [sr_tensor] * n, + }, + ) + + if self._lm is None or self._audio_tokenizer is None: + raise RuntimeError("MOSS-TTS-Nano model not loaded. Was load_weights() called?") + + waveforms: list[torch.Tensor] = [] + srs: list[torch.Tensor] = [] + + for info in runtime_additional_information: + if info.get("_is_dummy"): + waveforms.append(empty) + srs.append(sr_tensor) + continue + + waveform = self._run_single_inference(info, sr) + waveforms.append(waveform) + srs.append(sr_tensor) + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": waveforms, "sr": srs}, + ) + + def _run_single_inference(self, info: dict[str, Any], sr: int) -> torch.Tensor: + """Run inference for a single request and return the waveform tensor.""" + text: str = str(_pick(info, "text", "") or "") + if not text.strip(): + logger.warning("MOSS-TTS-Nano received empty text; returning silence.") + return torch.zeros((sr,), dtype=torch.float32) + + voice: str = str(_pick(info, "voice", _DEFAULT_VOICE)) + 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) + 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: + seed = int(seed) + + 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) + ), + } + + if seed is not None: + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + device = self._device or torch.device("cpu") + + # Use a temp file as the output path (the upstream API requires a path). + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: + output_path = tmp.name + + 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() + audio_chunks.append(chunk) + elif event_type == "result": + # Final event -- prefer the full waveform from the result + # if no intermediate chunks were collected. + if not audio_chunks: + waveform = event.get("waveform") + if waveform is not None: + audio_chunks.append( + torch.as_tensor(waveform, dtype=torch.float32).cpu() + ) + except Exception: + logger.exception("MOSS-TTS-Nano inference failed for text=%r", text[:80]) + return torch.zeros((sr,), dtype=torch.float32) + finally: + # Clean up temp file. + try: + Path(output_path).unlink(missing_ok=True) + except Exception: + pass + + if not audio_chunks: + logger.warning("MOSS-TTS-Nano produced no audio for text=%r", text[:80]) + return torch.zeros((sr,), dtype=torch.float32) + + # Concatenate all streaming chunks into a single waveform. + # inference_stream yields chunks that are 1D or 2D [samples] / [channels, samples]. + # Flatten to 1D for the vLLM-Omni audio output format. + processed: list[torch.Tensor] = [] + for chunk in audio_chunks: + if chunk.ndim == 2: + # [channels, samples] → interleave to [samples * channels] for stereo + processed.append(chunk.T.reshape(-1)) + else: + processed.append(chunk.reshape(-1)) + + return torch.cat(processed, dim=0) + + # ------------------------------------------------------------------ + # vLLM boilerplate: embed_input_ids required by generation runner + # ------------------------------------------------------------------ + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings=None, + is_multimodal=None, + ) -> torch.Tensor: + # Generation stage does not use token embeddings -- return a dummy + # zero tensor of the expected shape so the runner's prefill works. + 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/registry.py b/vllm_omni/model_executor/models/registry.py index 089408800..f8f11cb81 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -163,6 +163,12 @@ _OMNI_MODELS = { "VoxtralTTSAudioGenerationForConditionalGeneration", ), "VoxtralTTSAudioTokenizer": ("voxtral_tts", "voxtral_tts_audio_tokenizer", "VoxtralTTSAudioTokenizer"), + ## MOSS-TTS-Nano + "MossTTSNanoForCausalLM": ( + "moss_tts_nano", + "moss_tts_nano", + "MossTTSNanoForGeneration", + ), "DyninOmniForConditionalGeneration": ( "dynin_omni", "dynin_omni", diff --git a/vllm_omni/model_executor/stage_configs/moss_tts_nano.yaml b/vllm_omni/model_executor/stage_configs/moss_tts_nano.yaml new file mode 100644 index 000000000..1155247e9 --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/moss_tts_nano.yaml @@ -0,0 +1,61 @@ +# Stage config for MOSS-TTS-Nano (OpenMOSS-Team/MOSS-TTS-Nano). +# +# Single generation stage: the 0.1B AR LM and MOSS-Audio-Tokenizer-Nano codec +# both run inside MossTTSNanoForGeneration.forward(). No inter-stage connector +# is needed. Streaming (realtime) audio is supported via inference_stream() +# inside the model, which progressively collects audio chunks during AR decoding. +# +# Verified on: 1x GPU with ≥2 GB VRAM, or CPU (slower). +# Output: 48 kHz, 2-channel stereo audio. +# +# Request additional_information fields: +# text (str) – text to synthesize [required] +# voice (str) – built-in voice preset, default "Junhao" +# mode (str) – "voice_clone" (default) | "continuation" +# prompt_audio_path (str) – path to reference audio for custom voice clone +# prompt_text (str) – reference transcript (continuation mode) +# max_new_frames (int) – max AR frames, default 375 (~14 s) +# text_temperature (float) – default 1.0 +# audio_temperature (float) – default 0.8 +# audio_top_k (int) – default 25 +# audio_top_p (float) – default 0.95 +# seed (int) – optional RNG seed + +async_chunk: false + +stage_args: + - stage_id: 0 + stage_type: llm + is_comprehension: false + runtime: + devices: "0" + engine_args: + model_stage: moss_tts_nano + model_arch: MossTTSNanoForCausalLM + worker_type: generation + worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + trust_remote_code: true + enforce_eager: true + engine_output_type: audio + gpu_memory_utilization: 0.3 + enable_prefix_caching: false + skip_mm_profiling: true + max_num_seqs: 4 + max_model_len: 4096 + max_num_batched_tokens: 4096 + final_output: true + final_output_type: audio + default_sampling_params: + temperature: 1.0 + top_p: 1.0 + top_k: 50 + max_tokens: 4096 + seed: 42 + detokenize: false + +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 4