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>
This commit is contained in:
Yueqian Lin authored and Yueqian Lin committed
1 parent
bc4a659f03
commit
da6b1b5741
7 files changed
+806
No files matched your search
@@ -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())
|
||||
@@ -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
|
||||
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -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,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,
|
||||
)
|
||||
@@ -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",
|
||||
|
||||
@@ -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
|
||||
Reference in new issue
Block a user