[TTS][Model] support MOSS-TTS series (#3420)
Signed-off-by: Zhang <jianmusings@gmail.com> Signed-off-by: Zhang Jian <jianmusings@gmail.com> Signed-off-by: jian <jianmusings@gmail.com> Signed-off-by: Canlin Guo <961750412@qq.com> Co-authored-by: Canlin Guo <961750412@qq.com>
This commit is contained in:
Zhang JianandCanlin Guo authored and GitHub committed
1 parent
da9c45e1de
commit
b550709bba
32 files changed
+4766
-28
No files matched your search
@@ -0,0 +1,10 @@
|
||||
se001|ocean waves crashing on a rocky beach|10.0
|
||||
se002|heavy rain on a metal roof at night|8.0
|
||||
se003|a dog barking in the distance|6.0
|
||||
se004|a crackling campfire in a quiet forest|10.0
|
||||
se005|footsteps on dry leaves|6.0
|
||||
se006|thunder rolling across a wide valley|8.0
|
||||
se007|a busy coffee shop with chatter and espresso machines|10.0
|
||||
se008|wind howling through tall pine trees|8.0
|
||||
se009|a small stream flowing over pebbles|10.0
|
||||
se010|a creaking wooden door slowly opening|6.0
|
||||
@@ -0,0 +1,10 @@
|
||||
ttsd001|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Hello and welcome to the show today. [S2] Thanks for having me, it's a pleasure to be here. [S1] Let's dive right in.
|
||||
ttsd002|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Did you finish the report? [S2] Almost — I'm just polishing the conclusion. [S1] Great, send it over when you're done.
|
||||
ttsd003|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Have you tried that new restaurant downtown? [S2] Yes, the pasta is amazing. [S1] I'll book a table for Friday.
|
||||
ttsd004|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] What did you think of yesterday's meeting? [S2] Productive, but a bit long. [S1] Agreed, we should set an hour cap next time.
|
||||
ttsd005|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] The weather is beautiful today. [S2] Perfect for a hike. [S1] Want to drive out to the trails this afternoon?
|
||||
ttsd006|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] How did the demo go? [S2] Better than expected — the client signed off on the prototype. [S1] Fantastic news!
|
||||
ttsd007|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Coffee or tea this morning? [S2] Coffee please, with a splash of milk. [S1] Coming right up.
|
||||
ttsd008|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Did you read the email about the deadline change? [S2] I did, we now have until next Friday. [S1] That should give us enough buffer.
|
||||
ttsd009|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Are you joining the team dinner tonight? [S2] I'll be there around seven. [S1] See you then.
|
||||
ttsd010|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Have we backed up the database? [S2] Last backup ran two hours ago. [S1] Good, I'll trigger a fresh snapshot before the migration.
|
||||
@@ -53,6 +53,8 @@ _TASK_TO_DATASET: dict[str, str] = {
|
||||
"voice_clone": "seed-tts",
|
||||
"default_voice": "seed-tts-text",
|
||||
"voice_design": "seed-tts-design",
|
||||
"dialogue": "ttsd",
|
||||
"sound_effect": "sound-effect",
|
||||
}
|
||||
|
||||
# Default design dataset path (bundled with the repo)
|
||||
@@ -81,6 +83,7 @@ def build_bench_args(
|
||||
output_dir: str | None,
|
||||
result_filename: str | None,
|
||||
extra_cli_args: list[str],
|
||||
output_len: int | None = None,
|
||||
) -> list[str]:
|
||||
"""Build the ``vllm bench serve --omni`` command for one (task, concurrency) run."""
|
||||
dataset_name = _TASK_TO_DATASET[task]
|
||||
@@ -139,6 +142,9 @@ def build_bench_args(
|
||||
if wer_eval:
|
||||
cmd.append("--seed-tts-wer-eval")
|
||||
|
||||
if output_len is not None:
|
||||
cmd += ["--hf-output-len", str(output_len)]
|
||||
|
||||
if output_dir or result_filename:
|
||||
out_dir = output_dir or "."
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
@@ -238,6 +244,15 @@ def main() -> None:
|
||||
parser.add_argument("--host", default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--model-configs", default=str(_DEFAULT_MODEL_CONFIGS), help="Path to model_configs.yaml")
|
||||
parser.add_argument(
|
||||
"--output-len",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Override per-request max_new_tokens forwarded to the server "
|
||||
"(seed-tts datasets default to 2048; the MOSS-TTS full talker takes "
|
||||
"this as 'max_new_frames' so smaller values cap audio length per "
|
||||
"request and keep bench runtime tractable).",
|
||||
)
|
||||
parser.add_argument("extra", nargs=argparse.REMAINDER, help="Extra args passed directly to vllm bench serve")
|
||||
args = parser.parse_args()
|
||||
|
||||
@@ -292,6 +307,7 @@ def main() -> None:
|
||||
output_dir=args.output_dir,
|
||||
result_filename=result_filename,
|
||||
extra_cli_args=args.extra or [],
|
||||
output_len=args.output_len,
|
||||
)
|
||||
result = run_one_benchmark(cmd)
|
||||
if result is not None:
|
||||
|
||||
@@ -43,3 +43,61 @@ models:
|
||||
voice_clone: {}
|
||||
default_voice:
|
||||
voice: default
|
||||
|
||||
# MOSS-TTS family. After the serving_speech.py variant-detection refactor
|
||||
# all five variants now go through OpenAI /v1/audio/speech with the
|
||||
# variant-specific request fields (ambient_sound, ref_audio_2, instructions).
|
||||
OpenMOSS-Team/MOSS-TTS:
|
||||
supported_tasks: [voice_clone]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
task_extra_body:
|
||||
voice_clone: {}
|
||||
|
||||
# MOSS-TTS-v1.5: continued-training upgrade of MOSS-TTS 1.0 — same
|
||||
# MossTTSDelay-8B architecture and API, so identical wiring. Adds 31-language
|
||||
# synthesis via the `language` field and inline `[pause Xs]` markers.
|
||||
OpenMOSS-Team/MOSS-TTS-v1.5:
|
||||
supported_tasks: [voice_clone]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
task_extra_body:
|
||||
voice_clone: {}
|
||||
|
||||
OpenMOSS-Team/MOSS-TTS-Realtime:
|
||||
supported_tasks: [voice_clone]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
task_extra_body:
|
||||
voice_clone: {}
|
||||
|
||||
# MOSS-TTSD synthesizes [S1]/[S2] dialogue. The seed-tts voice_clone dataset
|
||||
# has single-speaker rows; bench piggybacks on it (both speakers share the
|
||||
# same ref_audio). For real dialogue evaluation use the ttsd dataset.
|
||||
OpenMOSS-Team/MOSS-TTSD-v1.0:
|
||||
supported_tasks: [voice_clone, dialogue]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
task_extra_body:
|
||||
voice_clone: {}
|
||||
dialogue: {}
|
||||
|
||||
# MOSS-SoundEffect ignores text/ref_audio; expects ambient_sound +
|
||||
# duration_seconds. The sound_effect dataset module emits those.
|
||||
OpenMOSS-Team/MOSS-SoundEffect:
|
||||
supported_tasks: [sound_effect]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
task_extra_body:
|
||||
sound_effect:
|
||||
duration_seconds: 10
|
||||
|
||||
# MOSS-VoiceGenerator: instructions-driven (no ref_audio). seed-tts-design
|
||||
# dataset already emits {instructions, target_text}, so voice_design works
|
||||
# out of the box once the serving validator accepts instructions.
|
||||
OpenMOSS-Team/MOSS-VoiceGenerator:
|
||||
supported_tasks: [voice_design]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
task_extra_body:
|
||||
voice_design: {}
|
||||
@@ -0,0 +1,377 @@
|
||||
"""Offline inference example for MOSS-TTS variants.
|
||||
|
||||
Covers all five models:
|
||||
- OpenMOSS-Team/MOSS-TTS (8B, general TTS)
|
||||
- OpenMOSS-Team/MOSS-TTSD-v1.0 (8B, dialogue TTS)
|
||||
- OpenMOSS-Team/MOSS-SoundEffect (8B, sound effect synthesis)
|
||||
- OpenMOSS-Team/MOSS-TTS-Realtime (1.7B, low-latency streaming)
|
||||
- OpenMOSS-Team/MOSS-VoiceGenerator(1.7B, voice design)
|
||||
|
||||
Usage
|
||||
-----
|
||||
# Standard TTS with voice cloning:
|
||||
python end2end.py \\
|
||||
--model OpenMOSS-Team/MOSS-TTS \\
|
||||
--text "Hello, this is a MOSS-TTS test." \\
|
||||
--ref-audio path/to/reference.wav \\
|
||||
--output output.wav
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import sys
|
||||
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
|
||||
def _load_ref_audio(path: str, target_sr: int = 24000) -> torch.Tensor:
|
||||
"""Load and resample reference audio to (1, T) at target_sr."""
|
||||
data, sr = sf.read(path, always_2d=True) # (T, C)
|
||||
wav = torch.from_numpy(data.T.astype("float32")) # (C, T)
|
||||
if wav.shape[0] > 1:
|
||||
wav = wav.mean(dim=0, keepdim=True)
|
||||
if sr != target_sr:
|
||||
import torchaudio
|
||||
|
||||
wav = torchaudio.functional.resample(wav, sr, target_sr)
|
||||
return wav # (1, T)
|
||||
|
||||
|
||||
def _build_unified_codes(
|
||||
model_id: str,
|
||||
text: str,
|
||||
ref_audio: torch.Tensor | None,
|
||||
*,
|
||||
instruction: str | None = None,
|
||||
tokens: int | None = None,
|
||||
quality: str | None = None,
|
||||
sound_event: str | None = None,
|
||||
ambient_sound: str | None = None,
|
||||
language: str | None = None,
|
||||
) -> tuple[list[int], torch.Tensor, int]:
|
||||
"""Build the upstream MossTTSDelay unified-codes prompt.
|
||||
|
||||
Returns (text_token_ids, audio_codes_LxNQ, n_vq) where
|
||||
``text_token_ids`` is the prefill prompt for the talker and
|
||||
``audio_codes_LxNQ`` is the matching delay-pattern audio code grid
|
||||
(``audio_pad_code`` everywhere except inside the reference-audio block).
|
||||
"""
|
||||
from transformers import AutoProcessor
|
||||
|
||||
proc = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
|
||||
|
||||
reference: list[torch.Tensor] | None = None
|
||||
if ref_audio is not None:
|
||||
# Pre-encode so the processor doesn't need torchaudio.load (which
|
||||
# now requires torchcodec on recent torchaudio releases).
|
||||
# encode_audios_from_wav expects (C, T) per item.
|
||||
wav_2d = ref_audio if ref_audio.dim() == 2 else ref_audio.unsqueeze(0)
|
||||
codes_list = proc.encode_audios_from_wav([wav_2d], sampling_rate=24000, n_vq=proc.model_config.n_vq)
|
||||
reference = [codes_list[0]]
|
||||
|
||||
user_msg_kwargs: dict = {"text": text}
|
||||
if reference is not None:
|
||||
user_msg_kwargs["reference"] = reference
|
||||
if instruction is not None:
|
||||
user_msg_kwargs["instruction"] = instruction
|
||||
if tokens is not None:
|
||||
user_msg_kwargs["tokens"] = tokens
|
||||
if quality is not None:
|
||||
user_msg_kwargs["quality"] = quality
|
||||
if sound_event is not None:
|
||||
user_msg_kwargs["sound_event"] = sound_event
|
||||
if ambient_sound is not None:
|
||||
user_msg_kwargs["ambient_sound"] = ambient_sound
|
||||
if language is not None:
|
||||
user_msg_kwargs["language"] = language
|
||||
|
||||
user_msg = proc.build_user_message(**user_msg_kwargs)
|
||||
batch = proc(conversations=[[user_msg]], mode="generation")
|
||||
unified = batch["input_ids"][0] # (L, 1+n_vq)
|
||||
n_vq = int(unified.shape[-1] - 1)
|
||||
text_ids = unified[:, 0].tolist()
|
||||
audio_codes = unified[:, 1:].contiguous().to(torch.int64) # (L, n_vq)
|
||||
|
||||
# Free the processor (incl. ~1.6B-param audio_tokenizer on CPU) before
|
||||
# we instantiate Omni — we only needed it for the one-shot encode.
|
||||
del proc
|
||||
gc.collect()
|
||||
|
||||
return text_ids, audio_codes, n_vq
|
||||
|
||||
|
||||
def _build_realtime_prompt(
|
||||
model_id: str,
|
||||
text: str,
|
||||
ref_audio: torch.Tensor | None,
|
||||
) -> tuple[list[int], torch.Tensor, int]:
|
||||
"""Build the prefill prompt for MOSS-TTS-Realtime.
|
||||
|
||||
Realtime uses a different prompt format than the delay variants: a system
|
||||
prompt + (optional) voice-clone context with the encoded reference audio
|
||||
inlined, followed by the user text and an assistant header. The format is
|
||||
a ``(L, channels+1=17)`` int grid where col 0 is text/special tokens and
|
||||
cols 1..16 are RVQ codebook entries (``audio_channel_pad`` outside the
|
||||
ref-audio block).
|
||||
|
||||
The upstream ``MossTTSRealtimeProcessor`` is shipped via
|
||||
``trust_remote_code`` but is not auto-discovered by ``AutoProcessor``
|
||||
(no ``processor_config.json``) — we import its module directly from the
|
||||
snapshot directory and use its ``make_ensemble``/``make_user_prompt``.
|
||||
"""
|
||||
import importlib.util
|
||||
import sys as _sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from huggingface_hub import snapshot_download
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
|
||||
snap_dir = Path(snapshot_download(repo_id=model_id))
|
||||
proc_module_path = snap_dir / "processing_mossttsrealtime.py"
|
||||
if not proc_module_path.exists():
|
||||
raise FileNotFoundError(f"realtime processor module missing at {proc_module_path}")
|
||||
|
||||
# Load the processor module without polluting sys.path globally.
|
||||
spec = importlib.util.spec_from_file_location("_moss_tts_realtime_proc", proc_module_path)
|
||||
mod = importlib.util.module_from_spec(spec) # type: ignore[arg-type]
|
||||
_sys.modules[spec.name] = mod # type: ignore[union-attr]
|
||||
spec.loader.exec_module(mod) # type: ignore[union-attr]
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
|
||||
realtime_processor = mod.MossTTSRealtimeProcessor(tokenizer=tokenizer)
|
||||
|
||||
# Encode reference audio via the standalone MOSS Audio Tokenizer (CPU is
|
||||
# fine — it runs once per request).
|
||||
audio_tokens = None
|
||||
if ref_audio is not None:
|
||||
codec = AutoModel.from_pretrained("OpenMOSS-Team/MOSS-Audio-Tokenizer", trust_remote_code=True).eval()
|
||||
wav_2d = ref_audio if ref_audio.dim() == 2 else ref_audio.unsqueeze(0)
|
||||
with torch.no_grad():
|
||||
enc = codec.batch_encode([wav_2d.squeeze(0)], num_quantizers=16)
|
||||
# enc.audio_codes: (NQ, B, T); we want (T, NQ) numpy.
|
||||
codes = enc.audio_codes[:, 0, : int(enc.audio_codes_lengths[0].item())]
|
||||
audio_tokens = codes.transpose(0, 1).contiguous().cpu().numpy() # (T, 16)
|
||||
del codec
|
||||
gc.collect()
|
||||
|
||||
# Build the inference-style prefill (mirrors upstream
|
||||
# ``MossTTSRealtimeInference._build_prefill_batch``):
|
||||
# system_prompt [+ voice_clone_block]
|
||||
# <|im_start|>assistant\n
|
||||
# <text_tokens...> (col 0 = text, cols 1..16 = audio_channel_pad)
|
||||
# <last text token> (col 1 = audio_bos at the last text position)
|
||||
# Generation continues from here; the local transformer emits real audio
|
||||
# codes for codebook 0 directly (no BOS frame at gen_step=0).
|
||||
system_grid = realtime_processor.make_ensemble(prompt_audio_tokens=audio_tokens) # (L_sys, 17)
|
||||
|
||||
# Append "<|im_start|>assistant\n" header (the snapshot's make_ensemble
|
||||
# omits it; inferencer.py's version includes it. We add it here.)
|
||||
assistant_header_ids = tokenizer.encode("<|im_start|>assistant\n", add_special_tokens=False)
|
||||
assistant_grid = np.full((len(assistant_header_ids), 17), 1024, dtype=np.int64)
|
||||
assistant_grid[:, 0] = assistant_header_ids
|
||||
|
||||
# Append the FIRST 12 text tokens with audio_bos (1025) at the last
|
||||
# included text position (audio column 0 == grid col 1). The remaining
|
||||
# text tokens are fed one-per-step during generation via the talker's
|
||||
# compute_logits → vLLM sampler pipeline (see ``ids.all`` in the
|
||||
# additional_information). This matches upstream
|
||||
# ``MossTTSRealtimeInference._build_prefill_batch`` (prefill_max_text=12).
|
||||
text_ids_only = tokenizer.encode(text, add_special_tokens=False)
|
||||
if not text_ids_only:
|
||||
raise ValueError("Empty text after tokenisation")
|
||||
PREFILL_MAX_TEXT = 12
|
||||
cur_len = min(len(text_ids_only), PREFILL_MAX_TEXT)
|
||||
text_grid = np.full((cur_len, 17), 1024, dtype=np.int64)
|
||||
text_grid[:, 0] = text_ids_only[:cur_len]
|
||||
text_grid[-1, 1] = 1025 # audio_bos_token at last *prefilled* text token
|
||||
|
||||
grid = np.concatenate([system_grid, assistant_grid, text_grid], axis=0)
|
||||
text_ids = grid[:, 0].tolist()
|
||||
audio_codes = torch.from_numpy(grid[:, 1:].astype(np.int64).copy()) # (L, 16)
|
||||
|
||||
# ``remaining_text_ids`` will be fed one-per-step by the talker.
|
||||
remaining_text_ids = list(text_ids_only[cur_len:])
|
||||
return text_ids, audio_codes, 16, remaining_text_ids
|
||||
|
||||
|
||||
def run_tts(args: argparse.Namespace) -> None:
|
||||
from vllm import SamplingParams
|
||||
|
||||
from vllm_omni import Omni
|
||||
|
||||
deploy_map = {
|
||||
"OpenMOSS-Team/MOSS-TTS": "moss_tts",
|
||||
"OpenMOSS-Team/MOSS-TTSD-v1.0": "moss_ttsd",
|
||||
"OpenMOSS-Team/MOSS-SoundEffect": "moss_sound_effect",
|
||||
"OpenMOSS-Team/MOSS-TTS-Realtime": "moss_tts_realtime",
|
||||
"OpenMOSS-Team/MOSS-VoiceGenerator": "moss_voice_generator",
|
||||
}
|
||||
deploy_config = deploy_map.get(args.model, "moss_tts")
|
||||
|
||||
is_sound_effect = "SoundEffect" in args.model
|
||||
is_voice_gen = "VoiceGenerator" in args.model
|
||||
is_realtime = "Realtime" in args.model
|
||||
|
||||
ref_audio = None
|
||||
builder_kwargs: dict = {}
|
||||
|
||||
remaining_text_ids: list[int] | None = None
|
||||
if is_realtime:
|
||||
# MOSS-TTS-Realtime uses a different prompt format and has its own
|
||||
# per-step depth transformer; the delay-style ``MossTTSDelayProcessor``
|
||||
# path doesn't apply.
|
||||
if not args.ref_audio:
|
||||
sys.exit("MOSS-TTS-Realtime currently requires --ref-audio (voice clone).")
|
||||
ref_audio = _load_ref_audio(args.ref_audio)
|
||||
text = args.text or ""
|
||||
text_ids, audio_codes, n_vq, remaining_text_ids = _build_realtime_prompt(args.model, text, ref_audio)
|
||||
print(
|
||||
f"Prefill prompt: {len(text_ids)} tokens, {n_vq}-quantizer audio block (realtime); "
|
||||
f"{len(remaining_text_ids)} text tokens to stream"
|
||||
)
|
||||
elif is_sound_effect:
|
||||
if not args.ambient_sound:
|
||||
sys.exit("MOSS-SoundEffect requires --ambient-sound")
|
||||
builder_kwargs["ambient_sound"] = args.ambient_sound
|
||||
# MOSS-SoundEffect emits ~12.5 codec frames per second; ``tokens``
|
||||
# tells the AR loop the target duration.
|
||||
if args.duration_tokens:
|
||||
builder_kwargs["tokens"] = args.duration_tokens
|
||||
elif args.duration_seconds:
|
||||
builder_kwargs["tokens"] = max(1, int(float(args.duration_seconds) * 12.5))
|
||||
text = args.text or ""
|
||||
elif is_voice_gen:
|
||||
# MOSS-VoiceGenerator needs an instruction (voice description) plus
|
||||
# the text to synthesise; no reference audio.
|
||||
if not args.instruction:
|
||||
sys.exit("MOSS-VoiceGenerator requires --instruction (voice description)")
|
||||
builder_kwargs["instruction"] = args.instruction
|
||||
text = args.text or ""
|
||||
else:
|
||||
if not args.ref_audio:
|
||||
sys.exit("Voice cloning requires --ref-audio")
|
||||
ref_audio = _load_ref_audio(args.ref_audio)
|
||||
text = args.text or ""
|
||||
|
||||
if not is_realtime:
|
||||
text_ids, audio_codes, n_vq = _build_unified_codes(args.model, text, ref_audio, **builder_kwargs)
|
||||
print(f"Prefill prompt: {len(text_ids)} tokens, {n_vq}-quantizer audio block")
|
||||
|
||||
omni = Omni(model=args.model, deploy_config=deploy_config, stage_init_timeout=600)
|
||||
|
||||
talker_sp = SamplingParams(
|
||||
temperature=1.7,
|
||||
top_p=0.8,
|
||||
top_k=25,
|
||||
max_tokens=2048,
|
||||
seed=42,
|
||||
detokenize=False,
|
||||
)
|
||||
codec_sp = SamplingParams(
|
||||
temperature=0.0,
|
||||
top_p=1.0,
|
||||
top_k=-1,
|
||||
max_tokens=65536,
|
||||
seed=42,
|
||||
detokenize=True,
|
||||
)
|
||||
|
||||
additional_info: dict = {
|
||||
# Reference-audio codes consumed by the talker prefill (additive
|
||||
# embedding); shape is (L, n_vq) where pad_code marks non-audio
|
||||
# positions. The OmniPayload schema reserves ``codes.ref`` for
|
||||
# exactly this kind of conditioning input.
|
||||
"codes": {"ref": audio_codes},
|
||||
}
|
||||
if remaining_text_ids:
|
||||
# Realtime variant: feed the remaining text tokens (everything past
|
||||
# the first ``PREFILL_MAX_TEXT``) one-per-step during decode via the
|
||||
# talker's compute_logits → vLLM sampler one-hot trick.
|
||||
additional_info["ids"] = {"all": remaining_text_ids}
|
||||
|
||||
request = {
|
||||
"prompt_token_ids": text_ids,
|
||||
"additional_information": additional_info,
|
||||
}
|
||||
|
||||
import time as _time
|
||||
|
||||
chunks: list[torch.Tensor] = []
|
||||
t_gen_start = _time.monotonic()
|
||||
t_first_audio: float | None = None
|
||||
for stage_outputs in omni.generate(request, [talker_sp, codec_sp]):
|
||||
# Each yielded ``stage_outputs`` is an OmniRequestOutput; the codec
|
||||
# stage attaches its waveform under ``multimodal_output["model_outputs"]``.
|
||||
# Avoid ``a or b`` when ``a`` may be a tensor (truthiness raises).
|
||||
mm = stage_outputs.multimodal_output or {}
|
||||
wav = mm.get("audio")
|
||||
if wav is None:
|
||||
wav = mm.get("model_outputs")
|
||||
if wav is None:
|
||||
continue
|
||||
if isinstance(wav, list):
|
||||
non_empty = [w.reshape(-1) for w in wav if isinstance(w, torch.Tensor) and w.numel() > 0]
|
||||
if non_empty and t_first_audio is None:
|
||||
t_first_audio = _time.monotonic()
|
||||
chunks.extend(non_empty)
|
||||
elif isinstance(wav, torch.Tensor) and wav.numel() > 0:
|
||||
if t_first_audio is None:
|
||||
t_first_audio = _time.monotonic()
|
||||
chunks.append(wav.reshape(-1))
|
||||
|
||||
t_gen_end = _time.monotonic()
|
||||
if not chunks:
|
||||
print("No audio in output — check model loading logs.")
|
||||
return
|
||||
|
||||
wav_t = torch.cat(chunks).float().cpu()
|
||||
sf.write(args.output, wav_t.numpy(), 24000)
|
||||
print(f"Saved {len(wav_t) / 24000:.2f}s of audio to {args.output}")
|
||||
ttfa_ms = (t_first_audio - t_gen_start) * 1000.0 if t_first_audio is not None else float("nan")
|
||||
gen_total_ms = (t_gen_end - t_gen_start) * 1000.0
|
||||
print(f"TTFA_MS={ttfa_ms:.1f} GEN_TOTAL_MS={gen_total_ms:.1f} AUDIO_S={len(wav_t) / 24000:.2f}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="MOSS-TTS offline inference example")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="OpenMOSS-Team/MOSS-TTS",
|
||||
help="HuggingFace model ID",
|
||||
)
|
||||
parser.add_argument("--text", default="", help="Text to synthesise (TTS variants)")
|
||||
parser.add_argument("--ref-audio", default=None, help="Reference audio for voice cloning")
|
||||
parser.add_argument(
|
||||
"--ambient-sound",
|
||||
default=None,
|
||||
help="Sound description for MOSS-SoundEffect",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--duration-tokens",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Target duration in tokens (1 s ≈ 12.5 tokens) for MOSS-SoundEffect",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--duration-seconds",
|
||||
type=float,
|
||||
default=None,
|
||||
help="Target duration in seconds (converted via ~12.5 tokens/s) for MOSS-SoundEffect",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--instruction",
|
||||
default=None,
|
||||
help="Voice description for MOSS-VoiceGenerator (e.g. 'a young woman with a clear voice')",
|
||||
)
|
||||
parser.add_argument("--output", default="output.wav", help="Output WAV path")
|
||||
args = parser.parse_args()
|
||||
run_tts(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,145 @@
|
||||
# MOSS-TTS
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: OpenMOSS
|
||||
- Models: `OpenMOSS-Team/MOSS-TTS` (8B), `OpenMOSS-Team/MOSS-TTS-v1.5` (8B),
|
||||
`OpenMOSS-Team/MOSS-TTS-Realtime` (1.7B), `OpenMOSS-Team/MOSS-TTSD-v1.0` (8B),
|
||||
`OpenMOSS-Team/MOSS-SoundEffect` (8B), `OpenMOSS-Team/MOSS-VoiceGenerator` (1.7B)
|
||||
- Task: Text-to-speech synthesis, sound effect generation, zero-shot voice design
|
||||
- Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API; offline inference
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe for 24 kHz multilingual TTS with voice cloning (20 languages including
|
||||
Chinese and English). Choose a variant based on your latency and quality requirements:
|
||||
|
||||
| Model | Params | Use case |
|
||||
|---|---|---|
|
||||
| MOSS-TTS | 8B | General TTS, highest quality |
|
||||
| MOSS-TTS-v1.5 | 8B | General TTS upgrade of 1.0: 31 languages, steadier cloning, `[pause Xs]` markers (set `language` for best results); same `MossTTSDelay` API |
|
||||
| MOSS-TTS-Realtime | 1.7B | Lowest latency (TTFB ~180 ms), streaming-first |
|
||||
| MOSS-TTSD-v1.0 | 8B | Multi-turn dialogue TTS |
|
||||
| MOSS-SoundEffect | 8B | Sound effect synthesis from text description |
|
||||
| MOSS-VoiceGenerator | 1.7B | Zero-shot voice design |
|
||||
|
||||
All variants share the same codec (`OpenMOSS-Team/MOSS-Audio-Tokenizer`, ~7 GB) and
|
||||
output 24 kHz mono audio.
|
||||
|
||||
## References
|
||||
|
||||
- Offline inference example: [`examples/offline_inference/text_to_speech/moss_tts/`](../../examples/offline_inference/text_to_speech/moss_tts/)
|
||||
- Deploy configs: [`vllm_omni/deploy/moss_tts.yaml`](../../vllm_omni/deploy/moss_tts.yaml) and variants
|
||||
- HuggingFace org: <https://huggingface.co/OpenMOSS-Team>
|
||||
|
||||
## Hardware Support
|
||||
|
||||
### GPU
|
||||
|
||||
#### 1x H100 80GB — MOSS-TTS (8B)
|
||||
|
||||
##### Environment
|
||||
|
||||
- OS: Linux
|
||||
- Python: 3.11+
|
||||
- CUDA 12.8
|
||||
- vLLM-Omni version: see `vllm_omni/__version__.py`
|
||||
|
||||
##### Command
|
||||
|
||||
```bash
|
||||
# The codec is loaded automatically from OpenMOSS-Team/MOSS-Audio-Tokenizer.
|
||||
# Override the path with MOSS_TTS_CODEC_PATH if you have a local copy.
|
||||
vllm serve OpenMOSS-Team/MOSS-TTS --omni --port 8091
|
||||
```
|
||||
|
||||
##### Verification
|
||||
|
||||
Voice cloning (provide a reference audio clip):
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "OpenMOSS-Team/MOSS-TTS",
|
||||
"input": "Hello, this is a voice cloning test.",
|
||||
"voice": "default",
|
||||
"ref_audio": "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/zh_1.wav",
|
||||
"response_format": "wav"
|
||||
}' --output output.wav
|
||||
```
|
||||
|
||||
##### Notes
|
||||
|
||||
- Peak GPU memory: ~18 GB for the talker (8B) + ~8 GB for the codec decoder on the same device.
|
||||
Use `gpu_memory_utilization: 0.85` in `moss_tts.yaml` (default).
|
||||
- Output: 24 kHz mono WAV.
|
||||
- The `MOSS_TTS_CODEC_PATH` environment variable overrides the codec checkpoint location.
|
||||
|
||||
---
|
||||
|
||||
#### 1x A10G 24GB — MOSS-TTS-Realtime (1.7B)
|
||||
|
||||
##### Environment
|
||||
|
||||
- OS: Linux
|
||||
- Python: 3.11+
|
||||
- CUDA 12.8
|
||||
|
||||
##### Command
|
||||
|
||||
```bash
|
||||
vllm serve OpenMOSS-Team/MOSS-TTS-Realtime --omni --port 8091
|
||||
```
|
||||
|
||||
##### Verification
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "OpenMOSS-Team/MOSS-TTS-Realtime",
|
||||
"input": "This is a low-latency streaming TTS test.",
|
||||
"voice": "default",
|
||||
"ref_audio": "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/zh_1.wav",
|
||||
"response_format": "wav",
|
||||
"stream": true
|
||||
}' --output output.wav
|
||||
```
|
||||
|
||||
##### Notes
|
||||
|
||||
- Peak GPU memory: ~6 GB for the talker (1.7B) + ~8 GB for the codec decoder.
|
||||
- First-audio latency (TTFB): ~180 ms on A10G.
|
||||
- `codec_chunk_frames: 15` in `moss_tts_realtime.yaml` for lower TTFA than the 8B variant.
|
||||
|
||||
---
|
||||
|
||||
#### 1x A10G 24GB — MOSS-SoundEffect (8B, sound effect synthesis)
|
||||
|
||||
##### Command
|
||||
|
||||
```bash
|
||||
vllm serve OpenMOSS-Team/MOSS-SoundEffect --omni --port 8091
|
||||
```
|
||||
|
||||
##### Verification
|
||||
|
||||
Sound effect synthesis takes a text description instead of reference audio:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "OpenMOSS-Team/MOSS-SoundEffect",
|
||||
"input": "Thunder rumbling, rain pattering on a tin roof.",
|
||||
"response_format": "wav"
|
||||
}' --output thunder.wav
|
||||
```
|
||||
|
||||
##### Notes
|
||||
|
||||
- No `ref_audio` required or accepted for MOSS-SoundEffect.
|
||||
- Input field maps to the `ambient_sound` parameter in the upstream processor.
|
||||
- Rate: ~12.5 tokens per second; longer descriptions produce longer audio.
|
||||
@@ -0,0 +1,230 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""E2E offline inference tests for MOSS-TTS two-stage pipeline.
|
||||
|
||||
Tests the MossTTSDelayModel path using MOSS-VoiceGenerator (1.7B) — the
|
||||
smallest MossTTSDelayModel variant — so the test runs on a single A10G or L4
|
||||
without requiring an 80 GB GPU.
|
||||
|
||||
A second test class covers MOSS-TTS-Realtime (1.7B, MossTTSRealtime) to
|
||||
exercise the local-transformer generation path.
|
||||
|
||||
Both classes use a module-scoped engine so the model is loaded once per test
|
||||
module execution. Do NOT add a second engine fixture in this file — that
|
||||
triggers mid-module teardown races (see skill invariant I4).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import urllib.request
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm import SamplingParams
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from vllm_omni import Omni
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Shared constants
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Selected by the nightly "TTS · Function Test" (-m "full_model and L4 and tts").
|
||||
# Without the tts + run-level marks these tests are collected but never run.
|
||||
pytestmark = [pytest.mark.full_model, pytest.mark.tts]
|
||||
|
||||
SAMPLE_RATE = 24_000 # All MOSS-TTS full variants output 24 kHz
|
||||
|
||||
REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/reference_zh_1.wav"
|
||||
|
||||
_DEFAULT_SAMPLING = SamplingParams(
|
||||
temperature=1.7,
|
||||
top_p=0.8,
|
||||
top_k=25,
|
||||
max_tokens=512, # short for tests (~20 s at 24 kHz, ~12.5 tokens/s)
|
||||
seed=42,
|
||||
detokenize=False,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Session fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def ref_audio_path(tmp_path_factory) -> str:
|
||||
"""Download the upstream reference clip once per session.
|
||||
|
||||
Set ``MOSS_TTS_SKIP_ON_NET_FAIL=1`` to skip in air-gapped environments.
|
||||
"""
|
||||
cache_dir = tmp_path_factory.mktemp("moss_tts_ref")
|
||||
target = cache_dir / "zh_1.wav"
|
||||
try:
|
||||
with urllib.request.urlopen(REF_AUDIO_URL, timeout=30) as resp:
|
||||
target.write_bytes(resp.read())
|
||||
except Exception as exc:
|
||||
msg = f"Cannot fetch reference clip {REF_AUDIO_URL}: {exc}"
|
||||
if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"):
|
||||
pytest.skip(msg)
|
||||
pytest.fail(msg)
|
||||
if not target.exists() or target.stat().st_size == 0:
|
||||
pytest.fail(f"Reference clip empty after download: {target}")
|
||||
return str(target)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _build_request(text: str, ref_audio_path: str, seed: int = 42) -> dict:
|
||||
return {
|
||||
"prompt": "<|im_start|>assistant\n",
|
||||
"additional_information": {
|
||||
"task_type": ["voice_clone"],
|
||||
"text": [text],
|
||||
"mode": ["voice_clone"],
|
||||
"prompt_audio_path": [ref_audio_path],
|
||||
"seed": [seed],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]:
|
||||
"""Run one request; return (waveform_cpu, 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 non-None multimodal_output"
|
||||
# Never use ``a or b`` on tensor values — a multi-element tensor
|
||||
# raises on truthiness. Check ``is None`` explicitly and handle the
|
||||
# pre-consolidation list form.
|
||||
audio = mm.get("audio")
|
||||
if audio is None:
|
||||
audio = mm.get("model_outputs")
|
||||
sr = mm.get("sr")
|
||||
assert audio is not None, "Expected 'audio' or 'model_outputs' in multimodal_output"
|
||||
if isinstance(audio, list):
|
||||
audio = torch.cat(
|
||||
[t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0],
|
||||
dim=0,
|
||||
)
|
||||
assert isinstance(audio, torch.Tensor), f"Expected Tensor, got {type(audio)}"
|
||||
return audio.reshape(-1).cpu(), int(sr.item()) if sr is not None else SAMPLE_RATE
|
||||
raise AssertionError("No stage outputs received")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# MossTTSDelayModel — MOSS-VoiceGenerator 1.7B
|
||||
# (tests the delay-pattern generation path on a small GPU)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_DELAY_MODEL = "OpenMOSS-Team/MOSS-VoiceGenerator"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def delay_engine():
|
||||
return Omni(model=_DELAY_MODEL, stage_init_timeout=300)
|
||||
|
||||
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "L4"})
|
||||
def test_moss_tts_delay_english(delay_engine, ref_audio_path):
|
||||
"""MossTTSDelayModel: English voice_clone produces non-empty 24 kHz audio."""
|
||||
req = _build_request("Hello, this is a MOSS-TTS voice cloning test.", ref_audio_path)
|
||||
audio, sr = _collect_audio(delay_engine, req)
|
||||
|
||||
assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}"
|
||||
assert audio.numel() > 0, "Audio tensor is empty"
|
||||
assert not torch.all(audio == 0), "Audio is silence"
|
||||
|
||||
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "L4"})
|
||||
def test_moss_tts_delay_chinese(delay_engine, ref_audio_path):
|
||||
"""MossTTSDelayModel: Chinese input produces non-empty audio."""
|
||||
req = _build_request("你好，这是语音合成测试。", ref_audio_path)
|
||||
audio, sr = _collect_audio(delay_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_delay_deterministic(delay_engine, ref_audio_path):
|
||||
"""MossTTSDelayModel: same seed yields identical waveforms."""
|
||||
req = _build_request("Reproducibility check.", ref_audio_path, seed=99)
|
||||
audio1, _ = _collect_audio(delay_engine, req)
|
||||
audio2, _ = _collect_audio(delay_engine, req)
|
||||
|
||||
assert audio1.shape == audio2.shape, "Shapes differ across identical seeds"
|
||||
assert torch.allclose(audio1, audio2, atol=1e-4), "Waveforms differ across identical seeds"
|
||||
|
||||
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "L4"})
|
||||
def test_moss_tts_delay_batch(delay_engine, ref_audio_path):
|
||||
"""MossTTSDelayModel: batch of two requests each returns non-empty audio."""
|
||||
requests = [
|
||||
_build_request("First sentence.", ref_audio_path),
|
||||
_build_request("Second sentence.", ref_audio_path),
|
||||
]
|
||||
results: list[torch.Tensor] = []
|
||||
for stage_outputs in delay_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
|
||||
audio = mm.get("audio")
|
||||
if audio is None:
|
||||
audio = mm.get("model_outputs")
|
||||
assert audio is not None
|
||||
if isinstance(audio, list):
|
||||
audio = torch.cat(
|
||||
[t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0],
|
||||
dim=0,
|
||||
)
|
||||
results.append(audio.reshape(-1).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"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# MossTTSRealtime — MOSS-TTS-Realtime 1.7B
|
||||
# (tests the local-transformer generation path)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_REALTIME_MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def realtime_engine():
|
||||
return Omni(model=_REALTIME_MODEL, stage_init_timeout=300)
|
||||
|
||||
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "L4"})
|
||||
def test_moss_tts_realtime_english(realtime_engine, ref_audio_path):
|
||||
"""MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio."""
|
||||
req = _build_request("This is a real-time TTS streaming test.", ref_audio_path)
|
||||
audio, sr = _collect_audio(realtime_engine, req)
|
||||
|
||||
assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}"
|
||||
assert audio.numel() > 0, "Audio tensor is empty"
|
||||
assert not torch.all(audio == 0), "Audio is silence"
|
||||
|
||||
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "L4"})
|
||||
def test_moss_tts_realtime_deterministic(realtime_engine, ref_audio_path):
|
||||
"""MossTTSRealtime: same seed yields identical waveforms."""
|
||||
req = _build_request("Determinism check for realtime model.", ref_audio_path, seed=7)
|
||||
audio1, _ = _collect_audio(realtime_engine, req)
|
||||
audio2, _ = _collect_audio(realtime_engine, req)
|
||||
|
||||
assert audio1.shape == audio2.shape
|
||||
assert torch.allclose(audio1, audio2, atol=1e-4)
|
||||
@@ -0,0 +1,113 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""E2E offline inference test for MOSS-TTS-v1.5 (MossTTSDelay-8B).
|
||||
|
||||
MOSS-TTS-v1.5 is a continued-training upgrade of MOSS-TTS 1.0 with the SAME
|
||||
``MossTTSDelay`` architecture and API, so it reuses the existing talker/codec
|
||||
path unchanged — this test just pins that the 8B v1.5 checkpoint loads and
|
||||
produces audio.
|
||||
|
||||
The model is 8B, so this lives in its own file (one module-scoped engine per
|
||||
file, per skill invariant I4) and is gated behind an 80 GB GPU; the small-GPU
|
||||
MossTTSDelay coverage stays in ``test_moss_tts.py`` (1.7B VoiceGenerator).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import urllib.request
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm import SamplingParams
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from vllm_omni import Omni
|
||||
|
||||
# H100-gated (8B). Picked up by the nightly H100 omni job
|
||||
# (-m "full_model and H100 and omni"); the tts mark keeps it consistent with
|
||||
# the other MOSS-TTS tests and any tts-scoped sweeps.
|
||||
pytestmark = [pytest.mark.full_model, pytest.mark.tts]
|
||||
|
||||
SAMPLE_RATE = 24_000
|
||||
REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/reference_zh_1.wav"
|
||||
_MODEL = "OpenMOSS-Team/MOSS-TTS-v1.5"
|
||||
|
||||
_DEFAULT_SAMPLING = SamplingParams(
|
||||
temperature=1.7,
|
||||
top_p=0.8,
|
||||
top_k=25,
|
||||
max_tokens=512, # short for tests (~20 s at 24 kHz, ~12.5 tokens/s)
|
||||
seed=42,
|
||||
detokenize=False,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def ref_audio_path(tmp_path_factory) -> str:
|
||||
"""Download the upstream reference clip once per session.
|
||||
|
||||
Set ``MOSS_TTS_SKIP_ON_NET_FAIL=1`` to skip in air-gapped environments.
|
||||
"""
|
||||
target = tmp_path_factory.mktemp("moss_tts_v15_ref") / "zh_1.wav"
|
||||
try:
|
||||
with urllib.request.urlopen(REF_AUDIO_URL, timeout=30) as resp:
|
||||
target.write_bytes(resp.read())
|
||||
except Exception as exc: # noqa: BLE001
|
||||
msg = f"Cannot fetch reference clip {REF_AUDIO_URL}: {exc}"
|
||||
if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"):
|
||||
pytest.skip(msg)
|
||||
pytest.fail(msg)
|
||||
if not target.exists() or target.stat().st_size == 0:
|
||||
pytest.fail(f"Reference clip empty after download: {target}")
|
||||
return str(target)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def v15_engine():
|
||||
return Omni(model=_MODEL, stage_init_timeout=600)
|
||||
|
||||
|
||||
def _build_request(text: str, ref_audio_path: str) -> dict:
|
||||
return {
|
||||
"prompt": "<|im_start|>assistant\n",
|
||||
"additional_information": {
|
||||
"task_type": ["voice_clone"],
|
||||
"text": [text],
|
||||
"mode": ["voice_clone"],
|
||||
"prompt_audio_path": [ref_audio_path],
|
||||
"seed": [42],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]:
|
||||
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 non-None multimodal_output"
|
||||
audio = mm.get("audio")
|
||||
if audio is None:
|
||||
audio = mm.get("model_outputs")
|
||||
sr = mm.get("sr")
|
||||
assert audio is not None, "Expected 'audio' or 'model_outputs' in multimodal_output"
|
||||
if isinstance(audio, list):
|
||||
audio = torch.cat(
|
||||
[t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0],
|
||||
dim=0,
|
||||
)
|
||||
assert isinstance(audio, torch.Tensor), f"Expected Tensor, got {type(audio)}"
|
||||
return audio.reshape(-1).cpu(), int(sr.item()) if sr is not None else SAMPLE_RATE
|
||||
raise AssertionError("No stage outputs received")
|
||||
|
||||
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "H100"})
|
||||
def test_moss_tts_v15_voice_clone(v15_engine, ref_audio_path):
|
||||
"""MOSS-TTS-v1.5 (8B): voice_clone produces non-empty 24 kHz audio."""
|
||||
req = _build_request("Hello, this is a MOSS-TTS v1.5 voice cloning test.", ref_audio_path)
|
||||
audio, sr = _collect_audio(v15_engine, req)
|
||||
|
||||
assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}"
|
||||
assert audio.numel() > 0, "Audio tensor is empty"
|
||||
assert not torch.all(audio == 0), "Audio is silence"
|
||||
@@ -0,0 +1,155 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
E2E online tests for the full MOSS-TTS family via /v1/audio/speech.
|
||||
|
||||
The offline suite (``tests/e2e/offline_inference/test_moss_tts.py`` and
|
||||
``test_moss_tts_v1_5.py``) covers the engine path; this file covers the
|
||||
*serving* path (``serving_speech.py`` → ``_detect_moss_variant`` → the
|
||||
delay/realtime prompt builders), which the Nano online test does not exercise.
|
||||
|
||||
We use MOSS-TTS-Realtime (1.7B, ``MossTTSRealtime``) because it is the
|
||||
smallest full-family variant and fits a single L4 — the 8B delay variants
|
||||
(MOSS-TTS / v1.5) are H100-gated and stay in the offline suite. Realtime uses
|
||||
the local-transformer generation path and, like the ``tts`` delay variant,
|
||||
requires ``ref_audio`` for voice cloning.
|
||||
|
||||
The reference clip is fetched from the upstream repo
|
||||
(``assets/audio/reference_zh_1.wav``, ~1.5 MB) and passed as a base64 data URL.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import os
|
||||
import urllib.request
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
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
|
||||
|
||||
pytestmark = [pytest.mark.full_model, pytest.mark.tts]
|
||||
|
||||
MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime"
|
||||
REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/HEAD/assets/audio/reference_zh_1.wav"
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def ref_audio_data_url() -> str:
|
||||
"""Fetch the upstream sample clip and return it as a base64 data URL.
|
||||
|
||||
The fetch failure is escalated to a hard failure (not pytest.skip) so a
|
||||
broken network path does not silently mask regressions in
|
||||
/v1/audio/speech. Set ``MOSS_TTS_SKIP_ON_NET_FAIL=1`` to opt into skipping
|
||||
in air-gapped environments.
|
||||
"""
|
||||
try:
|
||||
with urllib.request.urlopen(REF_AUDIO_URL, timeout=30) as resp:
|
||||
data = resp.read()
|
||||
except Exception as e: # noqa: BLE001
|
||||
msg = f"Cannot fetch upstream reference clip {REF_AUDIO_URL}: {e}"
|
||||
if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"):
|
||||
pytest.skip(msg)
|
||||
pytest.fail(msg)
|
||||
if not data:
|
||||
pytest.fail(f"Reference clip empty: {REF_AUDIO_URL}")
|
||||
return f"data:audio/wav;base64,{base64.b64encode(data).decode('ascii')}"
|
||||
|
||||
|
||||
def get_prompt(prompt_type: str = "text") -> str:
|
||||
"""Plain natural-sounding sentences for text-to-audio tests.
|
||||
|
||||
Avoid the model's own name in the input — the codec mispronounces it,
|
||||
which can trip transcript-similarity checks without a real regression.
|
||||
"""
|
||||
prompts = {
|
||||
"text": "Hello, this is a real-time voice cloning demo for testing.",
|
||||
"chinese": "你好，这是一段实时语音合成测试。",
|
||||
}
|
||||
return prompts.get(prompt_type, prompts["text"])
|
||||
|
||||
|
||||
tts_server_params = [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL,
|
||||
stage_config_path=get_deploy_config_path("moss_tts_realtime.yaml"),
|
||||
server_args=["--disable-log-stats"],
|
||||
),
|
||||
id="moss_tts_realtime",
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@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, ref_audio_data_url) -> None:
|
||||
"""
|
||||
Realtime voice_clone via /v1/audio/speech, non-streaming.
|
||||
Deploy Setting: moss_tts_realtime.yaml
|
||||
Input Modal: text + reference audio (voice clone)
|
||||
Output Modal: audio (24 kHz, WAV)
|
||||
Input Setting: stream=False
|
||||
Datasets: single request
|
||||
"""
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"input": get_prompt(),
|
||||
"stream": False,
|
||||
"response_format": "wav",
|
||||
"ref_audio": ref_audio_data_url,
|
||||
}
|
||||
|
||||
openai_client.send_audio_speech_request(request_config)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
|
||||
def test_text_to_audio_002_streaming(omni_server, openai_client, ref_audio_data_url) -> None:
|
||||
"""
|
||||
Realtime voice_clone via /v1/audio/speech, streaming PCM.
|
||||
Deploy Setting: moss_tts_realtime.yaml
|
||||
Input Modal: text + reference audio (voice clone)
|
||||
Output Modal: audio (24 kHz, PCM stream)
|
||||
Input Setting: stream=True
|
||||
Datasets: single request
|
||||
|
||||
NOTE: ``min_hnr_db=-5.0`` mirrors the Nano streaming case — MOSS
|
||||
voice_clone output is intrinsically noisy by the HNR metric even when it
|
||||
sounds correct; the relaxed threshold still catches catastrophic decoder
|
||||
failures (HNR far below 0) without flaking on normal output.
|
||||
"""
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"input": get_prompt(),
|
||||
"stream": True,
|
||||
"response_format": "pcm",
|
||||
"ref_audio": ref_audio_data_url,
|
||||
"min_hnr_db": -5.0,
|
||||
}
|
||||
|
||||
openai_client.send_audio_speech_request(request_config)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
|
||||
def test_text_to_audio_003_chinese(omni_server, openai_client, ref_audio_data_url) -> None:
|
||||
"""
|
||||
Realtime voice_clone via /v1/audio/speech, Chinese input.
|
||||
Deploy Setting: moss_tts_realtime.yaml
|
||||
Input Modal: text (Chinese) + reference audio (voice clone)
|
||||
Output Modal: audio (24 kHz, WAV)
|
||||
Input Setting: stream=False
|
||||
Datasets: single request
|
||||
"""
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"input": get_prompt("chinese"),
|
||||
"stream": False,
|
||||
"response_format": "wav",
|
||||
"ref_audio": ref_audio_data_url,
|
||||
}
|
||||
|
||||
openai_client.send_audio_speech_request(request_config)
|
||||
@@ -0,0 +1,128 @@
|
||||
"""MOSS-SoundEffect dataset for ``vllm bench serve``.
|
||||
|
||||
Loads ``meta.lst`` rows of the form::
|
||||
|
||||
utt_id|ambient_sound|duration_seconds
|
||||
|
||||
No reference audio. Builds one SampleRequest per row sending
|
||||
``ambient_sound`` + ``duration_seconds`` to ``/v1/audio/speech``. Speaker
|
||||
similarity (SIM) is not applicable; the ``seed_tts_ref_wav_path`` field is
|
||||
left empty so the seed-tts-eval pipeline skips it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from vllm.benchmarks.datasets import SampleRequest
|
||||
from vllm.tokenizers import TokenizerLike
|
||||
from vllm.tokenizers.hf import get_cached_tokenizer
|
||||
|
||||
from vllm_omni.benchmarks.data_modules.seed_tts_dataset import (
|
||||
SeedTTSDataset,
|
||||
SeedTTSSampleRequest,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _SoundEffectRow:
|
||||
utterance_id: str
|
||||
ambient_sound: str
|
||||
duration_seconds: float
|
||||
|
||||
|
||||
def _parse_se_line(line: str) -> _SoundEffectRow | None:
|
||||
line = line.strip()
|
||||
if not line or line.startswith("#"):
|
||||
return None
|
||||
parts = line.split("|")
|
||||
if len(parts) < 3:
|
||||
logger.warning("Skipping malformed sound-effect meta line (need 3 '|'-fields): %r", line[:120])
|
||||
return None
|
||||
try:
|
||||
duration = float(parts[2].strip())
|
||||
except ValueError:
|
||||
logger.warning("Bad duration_seconds %r in: %s", parts[2], line[:120])
|
||||
return None
|
||||
text = parts[1].strip()
|
||||
if not text:
|
||||
return None
|
||||
return _SoundEffectRow(utterance_id=parts[0].strip(), ambient_sound=text, duration_seconds=duration)
|
||||
|
||||
|
||||
class SoundEffectDataset(SeedTTSDataset):
|
||||
"""Ambient-sound prompts for MOSS-SoundEffect.
|
||||
|
||||
Reuses the SeedTTS root resolver. The ``input`` text in each request is
|
||||
set to an empty string — SoundEffect ignores it; ``ambient_sound`` is the
|
||||
actual conditioning signal.
|
||||
"""
|
||||
|
||||
DEFAULT_OUTPUT_LEN = 512 # ~40 s of audio at 12.5 fps
|
||||
|
||||
def load_data(self) -> None:
|
||||
meta = self._root / self.locale / "meta.lst"
|
||||
if not meta.is_file():
|
||||
raise FileNotFoundError(
|
||||
f"Sound-effect meta not found: {meta} (expected {self._root}/{self.locale}/meta.lst)"
|
||||
)
|
||||
text = meta.read_text(encoding="utf-8")
|
||||
rows: list[_SoundEffectRow] = []
|
||||
for line in text.splitlines():
|
||||
r = _parse_se_line(line)
|
||||
if r is not None:
|
||||
rows.append(r)
|
||||
if not rows:
|
||||
raise ValueError(f"No valid rows in {meta}")
|
||||
if not self.disable_shuffle:
|
||||
rng = random.Random(self.random_seed)
|
||||
rng.shuffle(rows)
|
||||
self._se_rows = rows
|
||||
self._rows = []
|
||||
self.data = self._se_rows
|
||||
logger.info("Loaded SoundEffect: root=%s locale=%s rows=%d", self._root, self.locale, len(rows))
|
||||
|
||||
def sample(
|
||||
self,
|
||||
tokenizer: TokenizerLike,
|
||||
num_requests: int,
|
||||
output_len: int | None = None,
|
||||
request_id_prefix: str = "",
|
||||
no_oversample: bool = False,
|
||||
**kwargs: Any,
|
||||
) -> list[SampleRequest]:
|
||||
if output_len is None:
|
||||
output_len = self.DEFAULT_OUTPUT_LEN
|
||||
tok = get_cached_tokenizer(tokenizer)
|
||||
out: list[SampleRequest] = []
|
||||
for i, row in enumerate(self._se_rows):
|
||||
if len(out) >= num_requests:
|
||||
break
|
||||
speech_extra: dict[str, Any] = {
|
||||
"ambient_sound": row.ambient_sound,
|
||||
"duration_seconds": row.duration_seconds,
|
||||
"max_new_tokens": output_len,
|
||||
}
|
||||
prompt_len = max(1, len(tok.encode(row.ambient_sound)))
|
||||
out.append(
|
||||
SeedTTSSampleRequest(
|
||||
prompt="", # SoundEffect ignores text
|
||||
prompt_len=prompt_len,
|
||||
expected_output_len=output_len,
|
||||
multi_modal_data=None,
|
||||
request_id=f"{request_id_prefix}{i}",
|
||||
seed_tts_speech_extra=speech_extra,
|
||||
seed_tts_utterance_id=row.utterance_id,
|
||||
seed_tts_locale=self.locale,
|
||||
seed_tts_system_prompt=self._system_prompt,
|
||||
seed_tts_ref_wav_path="", # SIM not applicable
|
||||
)
|
||||
)
|
||||
logger.info("SoundEffect: built %d requests (asked %d)", len(out), num_requests)
|
||||
self.maybe_oversample_requests(out, num_requests, request_id_prefix, no_oversample)
|
||||
return out
|
||||
@@ -0,0 +1,140 @@
|
||||
"""MOSS-TTSD dialogue dataset for ``vllm bench serve``.
|
||||
|
||||
Loads ``meta.lst`` rows of the form::
|
||||
|
||||
utt_id|ref1_wav_rel|ref2_wav_rel|dialogue_text
|
||||
|
||||
where ``dialogue_text`` contains ``[S1]`` / ``[S2]`` speaker tags. Both ref
|
||||
wavs live under ``{root}/<locale>/prompt-wavs/``. Builds one
|
||||
``SampleRequest`` per row with ``ref_audio`` (speaker 1) and ``ref_audio_2``
|
||||
(speaker 2) inlined as ``data:`` URLs.
|
||||
|
||||
Companion to ``seed_tts_dataset.py``; same on-disk layout convention so the
|
||||
existing ``--dataset-path`` flag works.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from vllm.benchmarks.datasets import SampleRequest
|
||||
from vllm.tokenizers import TokenizerLike
|
||||
from vllm.tokenizers.hf import get_cached_tokenizer
|
||||
|
||||
from vllm_omni.benchmarks.data_modules.seed_tts_dataset import (
|
||||
SeedTTSDataset,
|
||||
SeedTTSSampleRequest,
|
||||
_ref_audio_payload,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _TTSDRow:
|
||||
utterance_id: str
|
||||
ref1_rel: str
|
||||
ref2_rel: str
|
||||
target_text: str
|
||||
|
||||
|
||||
def _parse_ttsd_line(line: str) -> _TTSDRow | None:
|
||||
line = line.strip()
|
||||
if not line or line.startswith("#"):
|
||||
return None
|
||||
parts = line.split("|")
|
||||
if len(parts) < 4:
|
||||
logger.warning("Skipping malformed TTSD meta line (need 4 '|'-fields): %r", line[:120])
|
||||
return None
|
||||
return _TTSDRow(
|
||||
utterance_id=parts[0].strip(),
|
||||
ref1_rel=parts[1].strip(),
|
||||
ref2_rel=parts[2].strip(),
|
||||
target_text=parts[3].strip(),
|
||||
)
|
||||
|
||||
|
||||
class TTSDDataset(SeedTTSDataset):
|
||||
"""Two-speaker dialogue prompts for MOSS-TTSD-v1.0.
|
||||
|
||||
Reuses the SeedTTS root resolver so users can point ``--dataset-path`` at
|
||||
a local directory or a HF repo with the same layout convention.
|
||||
"""
|
||||
|
||||
DEFAULT_OUTPUT_LEN = 4096
|
||||
|
||||
def load_data(self) -> None:
|
||||
meta = self._root / self.locale / "meta.lst"
|
||||
if not meta.is_file():
|
||||
raise FileNotFoundError(
|
||||
f"TTSD meta not found: {meta} (expected layout: {self._root}/{self.locale}/meta.lst)"
|
||||
)
|
||||
text = meta.read_text(encoding="utf-8")
|
||||
rows: list[_TTSDRow] = []
|
||||
for line in text.splitlines():
|
||||
r = _parse_ttsd_line(line)
|
||||
if r is not None:
|
||||
rows.append(r)
|
||||
if not rows:
|
||||
raise ValueError(f"No valid rows in {meta}")
|
||||
if not self.disable_shuffle:
|
||||
rng = random.Random(self.random_seed)
|
||||
rng.shuffle(rows)
|
||||
self._ttsd_rows = rows
|
||||
# Parent's self._rows stays empty; sample() is overridden.
|
||||
self._rows = []
|
||||
self.data = self._ttsd_rows
|
||||
logger.info("Loaded TTSD: root=%s locale=%s rows=%d", self._root, self.locale, len(rows))
|
||||
|
||||
def sample(
|
||||
self,
|
||||
tokenizer: TokenizerLike,
|
||||
num_requests: int,
|
||||
output_len: int | None = None,
|
||||
request_id_prefix: str = "",
|
||||
no_oversample: bool = False,
|
||||
**kwargs: Any,
|
||||
) -> list[SampleRequest]:
|
||||
if output_len is None:
|
||||
output_len = self.DEFAULT_OUTPUT_LEN
|
||||
tok = get_cached_tokenizer(tokenizer)
|
||||
out: list[SampleRequest] = []
|
||||
for i, row in enumerate(self._ttsd_rows):
|
||||
if len(out) >= num_requests:
|
||||
break
|
||||
wav1 = (self._root / self.locale / row.ref1_rel).resolve()
|
||||
wav2 = (self._root / self.locale / row.ref2_rel).resolve()
|
||||
if not wav1.is_file() or not wav2.is_file():
|
||||
logger.warning("Missing TTSD refs for %s: %s / %s", row.utterance_id, wav1, wav2)
|
||||
continue
|
||||
ref1_uri = _ref_audio_payload(wav1, inline=self.inline_ref_audio)
|
||||
ref2_uri = _ref_audio_payload(wav2, inline=self.inline_ref_audio)
|
||||
speech_extra: dict[str, Any] = {
|
||||
"ref_audio": ref1_uri,
|
||||
"ref_audio_2": ref2_uri,
|
||||
"max_new_tokens": output_len,
|
||||
}
|
||||
target = row.target_text
|
||||
prompt_len = len(tok.encode(target))
|
||||
out.append(
|
||||
SeedTTSSampleRequest(
|
||||
prompt=target,
|
||||
prompt_len=prompt_len,
|
||||
expected_output_len=output_len,
|
||||
multi_modal_data=None,
|
||||
request_id=f"{request_id_prefix}{i}",
|
||||
seed_tts_speech_extra=speech_extra,
|
||||
seed_tts_utterance_id=row.utterance_id,
|
||||
seed_tts_locale=self.locale,
|
||||
seed_tts_system_prompt=self._system_prompt,
|
||||
# SIM computed against ref1 (speaker 1) for a stable signal;
|
||||
# per-speaker SIM is left for offline eval.
|
||||
seed_tts_ref_wav_path=str(wav1),
|
||||
)
|
||||
)
|
||||
logger.info("TTSD: built %d requests (asked %d)", len(out), num_requests)
|
||||
self.maybe_oversample_requests(out, num_requests, request_id_prefix, no_oversample)
|
||||
return out
|
||||
@@ -46,6 +46,8 @@ from vllm_omni.benchmarks.data_modules.seed_tts_dataset import (
|
||||
SeedTTSSampleRequest,
|
||||
SeedTTSTextDataset,
|
||||
)
|
||||
from vllm_omni.benchmarks.data_modules.sound_effect_dataset import SoundEffectDataset
|
||||
from vllm_omni.benchmarks.data_modules.ttsd_dataset import TTSDDataset
|
||||
|
||||
_AUDIO_CONTINUITY_THRESHOLD_ENV = "VLLM_OMNI_BENCH_AUDIO_CONTINUITY_THRESHOLD_S"
|
||||
_DEFAULT_AUDIO_CONTINUITY_THRESHOLD_S = 0.1
|
||||
@@ -166,7 +168,13 @@ def get_samples(args, tokenizer):
|
||||
is_daily_omni = args.dataset_name == "daily-omni" or (
|
||||
args.dataset_name == "hf" and _daily_omni_repo_from_args(args) is not None
|
||||
)
|
||||
is_seed_tts = args.dataset_name in ("seed-tts", "seed-tts-text", "seed-tts-design")
|
||||
is_seed_tts = args.dataset_name in (
|
||||
"seed-tts",
|
||||
"seed-tts-text",
|
||||
"seed-tts-design",
|
||||
"ttsd",
|
||||
"sound-effect",
|
||||
)
|
||||
|
||||
# Check if we need to handle omni-related backends/datasets
|
||||
is_omni_backend = args.backend in ["openai-chat-omni", "openai-audio-speech", "daily-omni"]
|
||||
@@ -283,6 +291,8 @@ def get_samples(args, tokenizer):
|
||||
"seed-tts": SeedTTSDataset,
|
||||
"seed-tts-text": SeedTTSTextDataset,
|
||||
"seed-tts-design": SeedTTSDesignDataset,
|
||||
"ttsd": TTSDDataset,
|
||||
"sound-effect": SoundEffectDataset,
|
||||
}
|
||||
DatasetCls = _cls_map[args.dataset_name]
|
||||
dataset = DatasetCls(
|
||||
|
||||
@@ -130,6 +130,14 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
|
||||
"vllm_omni.model_executor.models.moss_tts_nano.pipeline",
|
||||
"MOSS_TTS_NANO_PIPELINE",
|
||||
),
|
||||
"moss_tts_delay": (
|
||||
"vllm_omni.model_executor.models.moss_tts.pipeline",
|
||||
"MOSS_TTS_PIPELINE",
|
||||
),
|
||||
"moss_tts_realtime": (
|
||||
"vllm_omni.model_executor.models.moss_tts.pipeline",
|
||||
"MOSS_TTS_REALTIME_PIPELINE",
|
||||
),
|
||||
"minicpmo_4_5": (
|
||||
"vllm_omni.model_executor.models.minicpmo_4_5.pipeline",
|
||||
"MINICPMO_4_5_PIPELINE",
|
||||
|
||||
@@ -1270,6 +1270,17 @@ class StageConfigFactory:
|
||||
deploy_path = _DEPLOY_DIR / f"{model_type}.yaml"
|
||||
else:
|
||||
deploy_path = Path(deploy_config_path)
|
||||
# Bare-name lookup: if the value isn't a real path, try the
|
||||
# built-in deploy directory. This lets ``deploy_config="moss_tts"``
|
||||
# resolve to ``vllm_omni/deploy/moss_tts.yaml`` when several model
|
||||
# variants share an HF ``model_type`` and need different YAMLs.
|
||||
if not deploy_path.exists() and deploy_path.parent == Path("."):
|
||||
bare_name = deploy_path.name
|
||||
if not bare_name.endswith(".yaml"):
|
||||
bare_name = f"{bare_name}.yaml"
|
||||
candidate = _DEPLOY_DIR / bare_name
|
||||
if candidate.exists():
|
||||
deploy_path = candidate
|
||||
|
||||
if not deploy_path.exists():
|
||||
logger.warning(
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
# MOSS-SoundEffect (MossTTSDelayModel, 8B, n_vq=16, 24 kHz)
|
||||
# HF: OpenMOSS-Team/MOSS-SoundEffect
|
||||
# Sound effect synthesis — input is a text description of the sound,
|
||||
# no reference audio required.
|
||||
# Rate: ~12.5 tokens/second; pass `tokens` to control duration.
|
||||
async_chunk: true
|
||||
|
||||
connectors:
|
||||
shm:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536
|
||||
codec_streaming: true
|
||||
connector_get_sleep_s: 0.01
|
||||
connector_get_max_wait_first_chunk: 3000
|
||||
connector_get_max_wait: 300
|
||||
codec_chunk_frames: 25
|
||||
codec_left_context_frames: 0
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
max_num_seqs: 4
|
||||
gpu_memory_utilization: 0.85
|
||||
enable_prefix_caching: false
|
||||
async_scheduling: true
|
||||
max_num_batched_tokens: 4096
|
||||
max_model_len: 4096
|
||||
devices: "0"
|
||||
output_connectors:
|
||||
to_stage_1: shm
|
||||
default_sampling_params:
|
||||
temperature: 1.0
|
||||
top_p: 0.95
|
||||
top_k: 50
|
||||
max_tokens: 4096
|
||||
seed: 42
|
||||
repetition_penalty: 1.0
|
||||
|
||||
- stage_id: 1
|
||||
max_num_seqs: 1
|
||||
gpu_memory_utilization: 0.12
|
||||
enforce_eager: true
|
||||
enable_prefix_caching: false
|
||||
async_scheduling: true
|
||||
max_num_batched_tokens: 196608
|
||||
max_model_len: 196608
|
||||
devices: "0"
|
||||
input_connectors:
|
||||
from_stage_0: shm
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 65536
|
||||
seed: 42
|
||||
repetition_penalty: 1.0
|
||||
@@ -0,0 +1,58 @@
|
||||
# MOSS-TTS (MossTTSDelayModel, 8B, n_vq=32, 24 kHz)
|
||||
# HF: OpenMOSS-Team/MOSS-TTS
|
||||
# General-purpose TTS with voice cloning. Production-recommended variant.
|
||||
async_chunk: true
|
||||
|
||||
connectors:
|
||||
shm:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536
|
||||
codec_streaming: true
|
||||
connector_get_sleep_s: 0.01
|
||||
connector_get_max_wait_first_chunk: 3000
|
||||
connector_get_max_wait: 300
|
||||
codec_chunk_frames: 25
|
||||
codec_left_context_frames: 0
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
max_num_seqs: 4
|
||||
gpu_memory_utilization: 0.85
|
||||
enable_prefix_caching: false
|
||||
async_scheduling: true
|
||||
max_num_batched_tokens: 4096
|
||||
max_model_len: 4096
|
||||
devices: "0"
|
||||
output_connectors:
|
||||
to_stage_1: shm
|
||||
default_sampling_params:
|
||||
temperature: 1.7
|
||||
top_p: 0.8
|
||||
top_k: 25
|
||||
max_tokens: 4096
|
||||
seed: 42
|
||||
repetition_penalty: 1.0
|
||||
stop_token_ids: [151645] # <|im_end|> — terminates the talker AR loop
|
||||
|
||||
- stage_id: 1
|
||||
max_num_seqs: 1
|
||||
gpu_memory_utilization: 0.12
|
||||
enforce_eager: true
|
||||
enable_prefix_caching: false
|
||||
async_scheduling: true
|
||||
# Talker can emit up to max_tokens=4096 frames × n_vq=32 = 131072 codec
|
||||
# tokens per request; cap codec stage above that so a single long output
|
||||
# doesn't trigger "could not broadcast (>65536,) into (65536,)".
|
||||
max_num_batched_tokens: 196608
|
||||
max_model_len: 196608
|
||||
devices: "0"
|
||||
input_connectors:
|
||||
from_stage_0: shm
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 65536
|
||||
seed: 42
|
||||
repetition_penalty: 1.0
|
||||
@@ -0,0 +1,55 @@
|
||||
# MOSS-TTS-Realtime (MossTTSRealtime, 1.7B, n_vq=16, 24 kHz)
|
||||
# HF: OpenMOSS-Team/MOSS-TTS-Realtime
|
||||
# Real-time streaming TTS with TTFB ~180 ms and LLM+TTS first-sentence ~377 ms.
|
||||
# Uses local depth transformer (no delay warm-up) — lowest latency variant.
|
||||
async_chunk: true
|
||||
|
||||
connectors:
|
||||
shm:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536
|
||||
codec_streaming: true
|
||||
connector_get_sleep_s: 0.005
|
||||
connector_get_max_wait_first_chunk: 1000
|
||||
connector_get_max_wait: 300
|
||||
codec_chunk_frames: 15 # smaller chunks for lower TTFA
|
||||
codec_left_context_frames: 0
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
max_num_seqs: 8
|
||||
gpu_memory_utilization: 0.60
|
||||
enable_prefix_caching: false
|
||||
async_scheduling: true
|
||||
max_num_batched_tokens: 4096
|
||||
max_model_len: 4096
|
||||
devices: "0"
|
||||
output_connectors:
|
||||
to_stage_1: shm
|
||||
default_sampling_params:
|
||||
temperature: 1.0
|
||||
top_p: 0.95
|
||||
top_k: 50
|
||||
max_tokens: 4096
|
||||
seed: 42
|
||||
repetition_penalty: 1.1
|
||||
|
||||
- stage_id: 1
|
||||
max_num_seqs: 1
|
||||
gpu_memory_utilization: 0.12
|
||||
enforce_eager: true
|
||||
enable_prefix_caching: false
|
||||
async_scheduling: true
|
||||
max_num_batched_tokens: 65536
|
||||
max_model_len: 65536
|
||||
devices: "0"
|
||||
input_connectors:
|
||||
from_stage_0: shm
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 65536
|
||||
seed: 42
|
||||
repetition_penalty: 1.0
|
||||
@@ -0,0 +1,54 @@
|
||||
# MOSS-TTSD-v1.0 (MossTTSDelayModel, 8B, n_vq=16, 24 kHz)
|
||||
# HF: OpenMOSS-Team/MOSS-TTSD-v1.0
|
||||
# Dialogue TTS — optimised for multi-turn conversational speech generation.
|
||||
async_chunk: true
|
||||
|
||||
connectors:
|
||||
shm:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536
|
||||
codec_streaming: true
|
||||
connector_get_sleep_s: 0.01
|
||||
connector_get_max_wait_first_chunk: 3000
|
||||
connector_get_max_wait: 300
|
||||
codec_chunk_frames: 25
|
||||
codec_left_context_frames: 0
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
max_num_seqs: 4
|
||||
gpu_memory_utilization: 0.85
|
||||
enable_prefix_caching: false
|
||||
async_scheduling: true
|
||||
max_num_batched_tokens: 4096
|
||||
max_model_len: 4096
|
||||
devices: "0"
|
||||
output_connectors:
|
||||
to_stage_1: shm
|
||||
default_sampling_params:
|
||||
temperature: 1.7
|
||||
top_p: 0.8
|
||||
top_k: 25
|
||||
max_tokens: 4096
|
||||
seed: 42
|
||||
repetition_penalty: 1.0
|
||||
|
||||
- stage_id: 1
|
||||
max_num_seqs: 1
|
||||
gpu_memory_utilization: 0.12
|
||||
enforce_eager: true
|
||||
enable_prefix_caching: false
|
||||
async_scheduling: true
|
||||
max_num_batched_tokens: 196608
|
||||
max_model_len: 196608
|
||||
devices: "0"
|
||||
input_connectors:
|
||||
from_stage_0: shm
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 65536
|
||||
seed: 42
|
||||
repetition_penalty: 1.0
|
||||
@@ -0,0 +1,54 @@
|
||||
# MOSS-VoiceGenerator (MossTTSDelayModel, 1.7B, n_vq=16, 24 kHz)
|
||||
# HF: OpenMOSS-Team/MOSS-VoiceGenerator
|
||||
# Zero-shot voice design — create a new voice from a text description.
|
||||
async_chunk: true
|
||||
|
||||
connectors:
|
||||
shm:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536
|
||||
codec_streaming: true
|
||||
connector_get_sleep_s: 0.01
|
||||
connector_get_max_wait_first_chunk: 3000
|
||||
connector_get_max_wait: 300
|
||||
codec_chunk_frames: 25
|
||||
codec_left_context_frames: 0
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
max_num_seqs: 8
|
||||
gpu_memory_utilization: 0.60
|
||||
enable_prefix_caching: false
|
||||
async_scheduling: true
|
||||
max_num_batched_tokens: 4096
|
||||
max_model_len: 4096
|
||||
devices: "0"
|
||||
output_connectors:
|
||||
to_stage_1: shm
|
||||
default_sampling_params:
|
||||
temperature: 1.0
|
||||
top_p: 0.95
|
||||
top_k: 50
|
||||
max_tokens: 4096
|
||||
seed: 42
|
||||
repetition_penalty: 1.0
|
||||
|
||||
- stage_id: 1
|
||||
max_num_seqs: 1
|
||||
gpu_memory_utilization: 0.12
|
||||
enforce_eager: true
|
||||
enable_prefix_caching: false
|
||||
async_scheduling: true
|
||||
max_num_batched_tokens: 196608
|
||||
max_model_len: 196608
|
||||
devices: "0"
|
||||
input_connectors:
|
||||
from_stage_0: shm
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 65536
|
||||
seed: 42
|
||||
repetition_penalty: 1.0
|
||||
@@ -148,7 +148,16 @@ class OmniBenchmarkServingSubcommand(OmniBenchmarkSubcommandBase):
|
||||
for action in parser._actions:
|
||||
if action.dest == "dataset_name" and action.choices is not None:
|
||||
extra = [
|
||||
c for c in ("daily-omni", "seed-tts", "seed-tts-text", "seed-tts-design") if c not in action.choices
|
||||
c
|
||||
for c in (
|
||||
"daily-omni",
|
||||
"seed-tts",
|
||||
"seed-tts-text",
|
||||
"seed-tts-design",
|
||||
"ttsd",
|
||||
"sound-effect",
|
||||
)
|
||||
if c not in action.choices
|
||||
]
|
||||
if extra:
|
||||
action.choices = list(action.choices) + extra
|
||||
|
||||
@@ -54,6 +54,21 @@ class OpenAICreateSpeechRequest(BaseModel):
|
||||
default=None,
|
||||
description="Transcript of reference audio for voice cloning (Base task)",
|
||||
)
|
||||
ref_audio_2: str | None = Field(
|
||||
default=None,
|
||||
description="Second reference audio for two-speaker dialogue (MOSS-TTSD). "
|
||||
"URL, base64, or file URI. Ignored by single-speaker models.",
|
||||
)
|
||||
ambient_sound: str | None = Field(
|
||||
default=None,
|
||||
description="Sound description for ambient/effect synthesis (MOSS-SoundEffect). "
|
||||
"Natural language, e.g. 'ocean waves crashing on a rocky beach'.",
|
||||
)
|
||||
duration_seconds: float | None = Field(
|
||||
default=None,
|
||||
ge=0.0,
|
||||
description="Target audio duration in seconds (MOSS-SoundEffect). Converted to ~12.5 frames/s internally.",
|
||||
)
|
||||
x_vector_only_mode: bool | None = Field(
|
||||
default=None,
|
||||
description="Use speaker embedding only without in-context learning (Base task)",
|
||||
|
||||
@@ -74,6 +74,7 @@ _COVO_AUDIO_MODEL_STAGES = {"fused_thinker_talker"}
|
||||
_VOXCPM2_TTS_MODEL_STAGES = {"latent_generator"}
|
||||
_MING_TTS_MODEL_STAGES = {"ming_tts"}
|
||||
_MOSS_TTS_MODEL_STAGES = {"moss_tts_nano"}
|
||||
_MOSS_TTS_FULL_MODEL_STAGES = {"moss_tts", "moss_tts_codec"}
|
||||
_HIGGS_AUDIO_V2_TTS_MODEL_STAGES = {"higgs_audio_v2"}
|
||||
_GLM_TTS_MODEL_STAGES = {"glm_tts"}
|
||||
_TTS_MODEL_STAGES: set[str] = (
|
||||
@@ -87,6 +88,7 @@ _TTS_MODEL_STAGES: set[str] = (
|
||||
| _VOXCPM2_TTS_MODEL_STAGES
|
||||
| _MING_TTS_MODEL_STAGES
|
||||
| _MOSS_TTS_MODEL_STAGES
|
||||
| _MOSS_TTS_FULL_MODEL_STAGES
|
||||
| _GLM_TTS_MODEL_STAGES
|
||||
)
|
||||
_SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = {
|
||||
@@ -412,6 +414,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
self._tts_model_type = self._detect_tts_model_type()
|
||||
self.precomputed_speakers = self._load_precomputed_speakers()
|
||||
|
||||
# Sub-variant inside the full MOSS-TTS family. We collapse all five
|
||||
# variants onto the same _tts_model_type="moss_tts" because they share
|
||||
# the same stage layout, but request validation + param building
|
||||
# differ per HF repo (voice-clone vs dialogue vs ambient-sound vs
|
||||
# instruction vs streaming voice-clone).
|
||||
self._moss_variant = self._detect_moss_variant() if self._tts_model_type == "moss_tts" else None
|
||||
|
||||
# GLM-TTS lazy-cached resources (populated on first GLM-TTS request)
|
||||
self._glm_tts_text_tokenizer: object | None = None
|
||||
self._glm_tts_text_frontend: object | None = None
|
||||
@@ -581,6 +590,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
return "ming_flash_omni_tts"
|
||||
if model_stage in _MOSS_TTS_MODEL_STAGES:
|
||||
return "moss_tts_nano"
|
||||
if model_stage in _MOSS_TTS_FULL_MODEL_STAGES:
|
||||
return "moss_tts"
|
||||
if model_stage in _HIGGS_AUDIO_V2_TTS_MODEL_STAGES:
|
||||
return "higgs_audio_v2"
|
||||
if model_stage in _GLM_TTS_MODEL_STAGES:
|
||||
@@ -1268,7 +1279,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
return self._validate_voxcpm2_request(request)
|
||||
if self._tts_model_type == "ming_flash_omni_tts":
|
||||
return self._validate_ming_tts_request(request)
|
||||
if self._tts_model_type == "moss_tts_nano":
|
||||
if self._tts_model_type in ("moss_tts_nano", "moss_tts"):
|
||||
return self._validate_moss_tts_request(request)
|
||||
if self._tts_model_type == "glm_tts":
|
||||
return self._validate_glm_tts_request(request)
|
||||
@@ -1506,43 +1517,125 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
|
||||
return None
|
||||
|
||||
def _validate_moss_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
|
||||
"""Validate MOSS-TTS-Nano request.
|
||||
def _detect_moss_variant(self) -> str:
|
||||
"""Sub-classify a ``moss_tts``-stage server into the actual MOSS-TTS
|
||||
variant family (tts, ttsd, sound_effect, voice_generator, realtime).
|
||||
|
||||
Every request must include ``ref_audio``; the model has no built-in
|
||||
speaker presets, so the OpenAI ``voice`` field is accepted but
|
||||
ignored. ``ref_text`` is also accepted but ignored — upstream's
|
||||
``voice_clone`` (the only mode we expose, and the recommended
|
||||
workflow per its README/``infer.py``) does not consume a transcript,
|
||||
and its ``continuation`` mode produces near-silent output when given
|
||||
a reference clip + transcript pair, so routing there is not useful.
|
||||
Detection key is the HF repo path / model_name; matches
|
||||
``_try_resolve_omni_model_type`` in entrypoints/utils.py so users get
|
||||
consistent behaviour no matter how they launched the server (--model
|
||||
OpenMOSS-Team/MOSS-TTSD-v1.0 vs --deploy-config moss_ttsd.yaml).
|
||||
"""
|
||||
try:
|
||||
name = (self.engine_client.model_config.model or "").lower().replace("-", "").replace("_", "")
|
||||
except Exception:
|
||||
name = ""
|
||||
if "realtime" in name:
|
||||
return "realtime"
|
||||
if "ttsd" in name:
|
||||
return "ttsd"
|
||||
if "soundeffect" in name:
|
||||
return "sound_effect"
|
||||
if "voicegenerator" in name:
|
||||
return "voice_generator"
|
||||
return "tts"
|
||||
|
||||
def _validate_moss_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
|
||||
"""Validate any MOSS-TTS-family request (nano + 5 full variants).
|
||||
|
||||
Dispatches by ``self._moss_variant``:
|
||||
- ``tts``/``realtime``: require ``ref_audio`` (voice cloning).
|
||||
- ``ttsd``: require ``ref_audio`` (speaker 1); ``ref_audio_2``
|
||||
optional (defaults to the same ref for both speakers).
|
||||
- ``sound_effect``: require ``ambient_sound`` (no ref_audio).
|
||||
- ``voice_generator``: require ``instructions`` (no ref_audio).
|
||||
- For the legacy moss_tts_nano model_type the variant is None and
|
||||
we fall through to the original nano contract (ref_audio only).
|
||||
"""
|
||||
if not request.input or not request.input.strip():
|
||||
# SoundEffect can legitimately have empty input (just ambient_sound).
|
||||
if self._moss_variant != "sound_effect":
|
||||
return "Input text cannot be empty"
|
||||
|
||||
v = self._moss_variant
|
||||
if v in (None, "tts", "realtime"):
|
||||
if request.ref_audio is None:
|
||||
return (
|
||||
"MOSS-TTS-Nano requires 'ref_audio' (reference audio for voice cloning); "
|
||||
"the upstream model has no built-in voice presets."
|
||||
)
|
||||
label = "MOSS-TTS-Nano" if v is None else ("MOSS-TTS-Realtime" if v == "realtime" else "MOSS-TTS")
|
||||
return f"{label} requires 'ref_audio' (reference audio for voice cloning)."
|
||||
return self._validate_ref_audio_format(request.ref_audio)
|
||||
|
||||
if v == "ttsd":
|
||||
if request.ref_audio is None:
|
||||
return "MOSS-TTSD requires 'ref_audio' (speaker 1 reference)."
|
||||
fmt_err = self._validate_ref_audio_format(request.ref_audio)
|
||||
if fmt_err:
|
||||
return fmt_err
|
||||
if request.ref_audio_2 is not None:
|
||||
return self._validate_ref_audio_format(request.ref_audio_2)
|
||||
return None
|
||||
|
||||
async def _build_moss_tts_params(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]:
|
||||
"""Build additional_information for MOSS-TTS-Nano.
|
||||
if v == "sound_effect":
|
||||
if not request.ambient_sound or not request.ambient_sound.strip():
|
||||
return (
|
||||
"MOSS-SoundEffect requires 'ambient_sound' (natural language "
|
||||
"description of the sound effect to synthesise)."
|
||||
)
|
||||
return None
|
||||
|
||||
Always uses upstream's ``voice_clone`` mode (the recommended workflow
|
||||
per the README / ``infer.py`` default). Upstream's
|
||||
``_resolve_inference_mode`` rejects ``prompt_text`` in this mode, so
|
||||
we never forward it even if ``request.ref_text`` was supplied.
|
||||
``ref_audio`` is resolved via MediaConnector and passed as a
|
||||
``(wav_list, sample_rate)`` tuple so the model owns temp-file
|
||||
lifecycle. ``request.voice`` and ``request.ref_text`` are
|
||||
intentionally ignored — see ``_validate_moss_tts_request``.
|
||||
if v == "voice_generator":
|
||||
if not request.instructions or not request.instructions.strip():
|
||||
return (
|
||||
"MOSS-VoiceGenerator requires 'instructions' (natural language "
|
||||
"voice description, e.g. 'a warm female voice with an American accent')."
|
||||
)
|
||||
return None
|
||||
|
||||
return None # unreachable
|
||||
|
||||
def _get_moss_processor(self):
|
||||
"""Lazily load the upstream MOSS-TTS processor once per server.
|
||||
|
||||
Cached on ``self._moss_processor_cache``. The processor owns its own
|
||||
audio_tokenizer (~1.6 B params); we keep it on CPU so it doesn't
|
||||
compete with the talker (~8 GiB) and codec (~7 GiB) for our 96 GiB
|
||||
GPU — per-request ref-audio encoding is fast enough on CPU.
|
||||
"""
|
||||
cached = getattr(self, "_moss_processor_cache", None)
|
||||
if cached is not None:
|
||||
return cached
|
||||
from transformers import AutoProcessor
|
||||
|
||||
model_id = self.engine_client.model_config.model
|
||||
proc = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
|
||||
if hasattr(proc, "audio_tokenizer"):
|
||||
proc.audio_tokenizer = proc.audio_tokenizer.to("cpu").eval()
|
||||
self._moss_processor_cache = proc
|
||||
return proc
|
||||
|
||||
async def _build_moss_tts_params(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]:
|
||||
"""Build the talker prompt + ``additional_information`` payload for any
|
||||
MOSS-TTS-family request (nano + 5 full variants).
|
||||
|
||||
For the legacy ``moss_tts_nano`` model_type, keeps the original nano
|
||||
contract (``{text, mode=voice_clone, prompt_audio_array}``); the
|
||||
caller still uses a ``[1]`` placeholder prompt.
|
||||
|
||||
For the full MOSS-TTS family (``MossTTSDelay*`` / ``MossTTSRealtime*``)
|
||||
we **call the upstream processor server-side** to produce the unified
|
||||
``(text_ids, audio_codes)`` shape the talker actually consumes — same
|
||||
flow as ``examples/.../moss_tts/end2end.py:_build_unified_codes``.
|
||||
Returns ``{prompt_token_ids: list[int], codes.ref: torch.LongTensor,
|
||||
max_new_frames, ...}``. The caller treats ``prompt_token_ids`` as the
|
||||
prompt and forwards the rest as ``additional_information``.
|
||||
"""
|
||||
import torch # local to avoid pulling torch at module import time
|
||||
|
||||
v = self._moss_variant
|
||||
|
||||
# ---- Legacy nano path (unchanged) ----
|
||||
if v is None: # moss_tts_nano
|
||||
params: dict[str, Any] = {
|
||||
"text": [request.input],
|
||||
"text": [request.input or ""],
|
||||
"mode": ["voice_clone"],
|
||||
}
|
||||
if request.max_new_tokens is not None:
|
||||
@@ -1551,6 +1644,97 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
params["prompt_audio_array"] = [[wav_list, sr]]
|
||||
return params
|
||||
|
||||
# ---- MOSS-TTS-Realtime: keep the old prompt_audio_array path ----
|
||||
# ``AutoProcessor.from_pretrained`` doesn't auto-discover
|
||||
# ``MossTTSRealtimeProcessor`` (no ``processor_config.json`` in the
|
||||
# snapshot), and Realtime's prompt format diverges from MossTTSDelay
|
||||
# (16-channel grid, separate per-step text feed). The
|
||||
# ``prompt_audio_array`` shape lines up well enough with what the
|
||||
# talker reads for short prompts; full Realtime support needs a
|
||||
# separate processor.from_module path which we don't wire here.
|
||||
if v == "realtime":
|
||||
params: dict[str, Any] = {
|
||||
"text": [request.input or ""],
|
||||
"mode": ["voice_clone"],
|
||||
}
|
||||
if request.max_new_tokens is not None:
|
||||
params["max_new_frames"] = [request.max_new_tokens]
|
||||
wav_list, sr = await self._resolve_ref_audio(request.ref_audio)
|
||||
params["prompt_audio_array"] = [[wav_list, sr]]
|
||||
return params
|
||||
|
||||
# ---- MossTTSDelay family (tts/ttsd/sound_effect/voice_generator):
|
||||
# call the upstream processor server-side to produce unified codes ----
|
||||
proc = self._get_moss_processor()
|
||||
n_vq = int(getattr(proc.model_config, "n_vq", 32))
|
||||
sr_target = int(getattr(proc.model_config, "sampling_rate", 24000))
|
||||
|
||||
# Reference-audio encoding + speaker caching lives in the model package
|
||||
# (moss_tts.reference_encoder), mirroring Fish Speech / CosyVoice3 /
|
||||
# Qwen3-TTS which keep reference handling with the model rather than in
|
||||
# this shared serving file. Imported lazily so the API-server process
|
||||
# only pulls it on the delay-family path (alongside the upstream proc).
|
||||
from vllm_omni.model_executor.models.moss_tts.reference_encoder import encode_reference_codes
|
||||
|
||||
_voice = getattr(request, "voice", None)
|
||||
_voice = _voice.strip() if isinstance(_voice, str) else ""
|
||||
_voice_created = self._voice_created_at(_voice.lower()) if _voice else 0
|
||||
|
||||
async def _encode_ref(ref_str: str) -> torch.Tensor:
|
||||
return await encode_reference_codes(
|
||||
ref_str,
|
||||
processor=proc,
|
||||
resolve_ref_audio=self._resolve_ref_audio,
|
||||
speaker_cache=self._speaker_cache,
|
||||
variant=v,
|
||||
n_vq=n_vq,
|
||||
sr_target=sr_target,
|
||||
voice_name=_voice or None,
|
||||
voice_created_at=_voice_created,
|
||||
)
|
||||
|
||||
user_kwargs: dict[str, Any] = {"text": request.input or ""}
|
||||
if v in ("tts", "realtime"):
|
||||
user_kwargs["reference"] = [await _encode_ref(request.ref_audio)]
|
||||
elif v == "ttsd":
|
||||
refs = [await _encode_ref(request.ref_audio)]
|
||||
if request.ref_audio_2:
|
||||
refs.append(await _encode_ref(request.ref_audio_2))
|
||||
user_kwargs["reference"] = refs
|
||||
elif v == "sound_effect":
|
||||
user_kwargs["text"] = request.input or "" # may be empty
|
||||
user_kwargs["ambient_sound"] = request.ambient_sound or ""
|
||||
if request.duration_seconds is not None:
|
||||
user_kwargs["tokens"] = max(1, int(float(request.duration_seconds) * 12.5))
|
||||
elif request.max_new_tokens is not None:
|
||||
user_kwargs["tokens"] = int(request.max_new_tokens)
|
||||
elif v == "voice_generator":
|
||||
user_kwargs["instruction"] = request.instructions or ""
|
||||
|
||||
# Optional language tag for the spoken-text variants. MOSS-TTS-v1.5's
|
||||
# headline improvement is multilingual synthesis when the language is
|
||||
# given (build_user_message(..., language=...)); 1.0 ignores it
|
||||
# gracefully. Sound-effect output is non-verbal, so skip it there.
|
||||
if v in ("tts", "ttsd", "voice_generator") and getattr(request, "language", None):
|
||||
user_kwargs["language"] = request.language
|
||||
|
||||
# Build the unified-codes prompt: (L, 1+n_vq) where col 0 is text/special
|
||||
# tokens and cols 1..n_vq are the delay-pattern audio code grid (mostly
|
||||
# audio_pad_code outside the reference block).
|
||||
user_msg = proc.build_user_message(**user_kwargs)
|
||||
batch = proc(conversations=[[user_msg]], mode="generation")
|
||||
unified = batch["input_ids"][0] # torch.LongTensor (L, 1+n_vq)
|
||||
text_ids: list[int] = unified[:, 0].tolist()
|
||||
audio_codes: torch.Tensor = unified[:, 1:].contiguous().to(torch.int64)
|
||||
|
||||
params: dict[str, Any] = {
|
||||
"prompt_token_ids": text_ids,
|
||||
"codes": {"ref": audio_codes},
|
||||
}
|
||||
if request.max_new_tokens is not None:
|
||||
params["max_new_frames"] = [request.max_new_tokens]
|
||||
return params
|
||||
|
||||
def _validate_higgs_audio_v2_request(self, request: OpenAICreateSpeechRequest) -> str | None:
|
||||
"""Validate higgs_audio_v2 request parameters. Returns error message or None.
|
||||
|
||||
@@ -2655,7 +2839,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
elif self._tts_model_type == "ming_flash_omni_tts":
|
||||
prompt = self._build_ming_prompt(request)
|
||||
tts_params = {}
|
||||
elif self._tts_model_type == "moss_tts_nano":
|
||||
elif self._tts_model_type in ("moss_tts_nano", "moss_tts"):
|
||||
tts_params = await self._build_moss_tts_params(request)
|
||||
if request.voice:
|
||||
voice_lower = request.voice.lower()
|
||||
@@ -2669,6 +2853,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
# state and produces non-deterministic output.
|
||||
if sampling_params_list and getattr(sampling_params_list[0], "seed", None) is not None:
|
||||
tts_params["seed"] = [sampling_params_list[0].seed]
|
||||
# Full MOSS-TTS family returns real text_ids from the upstream
|
||||
# processor; use them as prompt_token_ids. Nano keeps the
|
||||
# ``[1]`` placeholder (it owns its own tokenisation internally).
|
||||
if isinstance(tts_params.get("prompt_token_ids"), list):
|
||||
prompt_token_ids = tts_params.pop("prompt_token_ids")
|
||||
prompt = tokens_input(prompt_token_ids=prompt_token_ids)
|
||||
else:
|
||||
prompt = tokens_input(prompt_token_ids=[1])
|
||||
prompt["additional_information"] = tts_params
|
||||
prompt["cache_salt"] = _conditioning_cache_salt(request, tts_params)
|
||||
@@ -2735,6 +2926,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
model_type = "ming_flash_omni_tts"
|
||||
elif self._tts_model_type == "moss_tts_nano":
|
||||
model_type = "moss_tts_nano"
|
||||
elif self._tts_model_type == "moss_tts":
|
||||
model_type = "moss_tts"
|
||||
elif self._tts_model_type == "higgs_audio_v2":
|
||||
model_type = "higgs_audio_v2"
|
||||
elif self._tts_model_type == "glm_tts":
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import (
|
||||
MossTTSCodecDecoder,
|
||||
)
|
||||
from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
|
||||
MossTTSDelayTalkerForGeneration,
|
||||
MossTTSRealtimeTalkerForGeneration,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"MossTTSDelayTalkerForGeneration",
|
||||
"MossTTSRealtimeTalkerForGeneration",
|
||||
"MossTTSCodecDecoder",
|
||||
]
|
||||
@@ -0,0 +1,578 @@
|
||||
# Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Vendored from OpenMOSS-Team/MOSS-Audio-Tokenizer (configuration_moss_audio_tokenizer.py
|
||||
# and modeling_moss_audio_tokenizer.py). Simplified for inference-only use:
|
||||
# - Streaming KV-cache infrastructure removed (single-pass batch decode only).
|
||||
# - Training-only methods (forward, encode, decode) removed.
|
||||
# - Dead branches removed: gating="none" always, weights_per_step=0 always,
|
||||
# positional_embedding="rope" always, norm="layer_norm" always in default config.
|
||||
"""MOSS Audio Tokenizer — inference-only codec (encode waveform ↔ RVQ codes)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, cast
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from transformers import PretrainedConfig
|
||||
from transformers.modeling_utils import PreTrainedModel
|
||||
from transformers.utils import ModelOutput
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Configuration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class MossAudioTokenizerConfig(PretrainedConfig):
|
||||
model_type = "moss-audio-tokenizer"
|
||||
attribute_map = {"sample_rate": "sampling_rate"}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
version: str | None = None,
|
||||
sampling_rate: int = 24000,
|
||||
downsample_rate: int = 1920,
|
||||
causal_transformer_context_duration: float = 10.0,
|
||||
encoder_kwargs: list[dict[str, Any]] | None = None,
|
||||
decoder_kwargs: list[dict[str, Any]] | None = None,
|
||||
quantizer_type: str = "rlfq",
|
||||
quantizer_kwargs: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
kwargs.pop("model_type", None)
|
||||
self.version = version
|
||||
self.sampling_rate = sampling_rate
|
||||
self.downsample_rate = downsample_rate
|
||||
self.causal_transformer_context_duration = causal_transformer_context_duration
|
||||
self.encoder_kwargs = encoder_kwargs or _default_encoder_kwargs()
|
||||
self.decoder_kwargs = decoder_kwargs or _default_decoder_kwargs()
|
||||
if quantizer_kwargs is None:
|
||||
quantizer_kwargs = {
|
||||
"input_dim": 768,
|
||||
"rvq_dim": 512,
|
||||
"output_dim": 768,
|
||||
"num_quantizers": 32,
|
||||
"codebook_size": 1024,
|
||||
"codebook_dim": 8,
|
||||
"quantizer_type": "rlfq",
|
||||
}
|
||||
kw_qtype = quantizer_kwargs.get("quantizer_type")
|
||||
self.quantizer_type = kw_qtype if kw_qtype is not None else quantizer_type
|
||||
quantizer_kwargs["quantizer_type"] = self.quantizer_type
|
||||
self.quantizer_kwargs = quantizer_kwargs
|
||||
super().__init__(**kwargs)
|
||||
|
||||
@property
|
||||
def num_quantizers(self) -> int:
|
||||
return int(self.quantizer_kwargs.get("num_quantizers", 32))
|
||||
|
||||
@property
|
||||
def codebook_size(self) -> int:
|
||||
return int(self.quantizer_kwargs.get("codebook_size", 4096))
|
||||
|
||||
@property
|
||||
def frame_rate(self) -> float:
|
||||
return self.sampling_rate / self.downsample_rate
|
||||
|
||||
|
||||
def _transformer_block(
|
||||
input_dim: int, output_dim: int, d_model: int, num_heads: int, num_layers: int
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"module_type": "Transformer",
|
||||
"input_dimension": input_dim,
|
||||
"output_dimension": output_dim,
|
||||
"d_model": d_model,
|
||||
"num_heads": num_heads,
|
||||
"num_layers": num_layers,
|
||||
"dim_feedforward": d_model * 4,
|
||||
"causal": True,
|
||||
"norm": "layer_norm",
|
||||
"positional_embedding": "rope",
|
||||
"max_period": 10000,
|
||||
"gating": "none",
|
||||
"layer_scale": 0.01,
|
||||
"conv_layout": True,
|
||||
}
|
||||
|
||||
|
||||
def _default_encoder_kwargs() -> list[dict[str, Any]]:
|
||||
return [
|
||||
{"module_type": "PatchedPretransform", "patch_size": 240},
|
||||
_transformer_block(240, 384, 768, 12, 12),
|
||||
{"module_type": "PatchedPretransform", "patch_size": 2},
|
||||
_transformer_block(768, 384, 768, 12, 12),
|
||||
{"module_type": "PatchedPretransform", "patch_size": 2},
|
||||
_transformer_block(768, 640, 768, 12, 12),
|
||||
{"module_type": "PatchedPretransform", "patch_size": 2},
|
||||
_transformer_block(1280, 768, 1280, 20, 32),
|
||||
]
|
||||
|
||||
|
||||
def _default_decoder_kwargs() -> list[dict[str, Any]]:
|
||||
return [
|
||||
_transformer_block(768, 1280, 1280, 20, 32),
|
||||
{"module_type": "PatchedPretransform", "patch_size": 2},
|
||||
_transformer_block(640, 768, 768, 12, 12),
|
||||
{"module_type": "PatchedPretransform", "patch_size": 2},
|
||||
_transformer_block(384, 768, 768, 12, 12),
|
||||
{"module_type": "PatchedPretransform", "patch_size": 2},
|
||||
_transformer_block(384, 768, 768, 12, 12),
|
||||
{"module_type": "PatchedPretransform", "patch_size": 2},
|
||||
_transformer_block(384, 240, 768, 12, 12),
|
||||
{"module_type": "PatchedPretransform", "patch_size": 240},
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Outputs
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class MossAudioTokenizerEncoderOutput(ModelOutput):
|
||||
audio_codes: torch.Tensor | None = None
|
||||
audio_codes_lengths: torch.Tensor | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class MossAudioTokenizerDecoderOutput(ModelOutput):
|
||||
audio: torch.Tensor | None = None
|
||||
audio_lengths: torch.Tensor | None = None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Building blocks
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _LayerScale(nn.Module):
|
||||
def __init__(self, channels: int, init: float = 1e-4, device=None, dtype=None) -> None:
|
||||
super().__init__()
|
||||
self.scale = nn.Parameter(torch.full((channels,), init, device=device, dtype=dtype))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.scale * x
|
||||
|
||||
|
||||
def _apply_rope(q: torch.Tensor, k: torch.Tensor, max_period: float = 10_000) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Rotary position embedding over sequence dimension (B, H, T, D).
|
||||
|
||||
Matches upstream MossAudioTokenizer's ``apply_rope``: pair the last dim
|
||||
as ``(re, im)`` interleaved (GPT-J style) — *not* GPT-NeoX split-halves.
|
||||
The two conventions are not interchangeable with the same checkpoint;
|
||||
using the wrong one silently rotates random subspaces of the K/V vectors.
|
||||
"""
|
||||
B, H, T, D = q.shape
|
||||
half = D // 2
|
||||
ds = torch.arange(half, device=q.device, dtype=torch.float32)
|
||||
freqs = torch.exp(ds * (-math.log(max_period) * 2 / D))
|
||||
ts = torch.arange(T, device=q.device, dtype=torch.float32).view(1, 1, -1, 1) # (1, 1, T, 1)
|
||||
rotr = torch.cos(freqs * ts) # (1, 1, T, D/2)
|
||||
roti = torch.sin(freqs * ts)
|
||||
|
||||
dims = q.shape[:-1]
|
||||
q_pair = q.view(*dims, half, 2)
|
||||
k_pair = k.view(*dims, half, 2)
|
||||
qr, qi = q_pair[..., 0].float(), q_pair[..., 1].float()
|
||||
kr, ki = k_pair[..., 0].float(), k_pair[..., 1].float()
|
||||
|
||||
qor = qr * rotr - qi * roti
|
||||
qoi = qr * roti + qi * rotr
|
||||
kor = kr * rotr - ki * roti
|
||||
koi = kr * roti + ki * rotr
|
||||
|
||||
qo = torch.stack([qor.to(q.dtype), qoi.to(q.dtype)], dim=-1).view(*dims, D)
|
||||
ko = torch.stack([kor.to(k.dtype), koi.to(k.dtype)], dim=-1).view(*dims, D)
|
||||
return qo, ko
|
||||
|
||||
|
||||
class _Attention(nn.Module):
|
||||
"""Causal multi-head self-attention with RoPE, no streaming KV cache."""
|
||||
|
||||
def __init__(
|
||||
self, embed_dim: int, num_heads: int, causal: bool, max_period: float, device=None, dtype=None
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.embed_dim = embed_dim
|
||||
self.num_heads = num_heads
|
||||
self.causal = causal
|
||||
self.max_period = max_period
|
||||
kw = {"device": device, "dtype": dtype}
|
||||
self.in_proj = nn.Linear(embed_dim, 3 * embed_dim, bias=False, **kw)
|
||||
self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False, **kw)
|
||||
# Remap legacy weight names produced by the upstream checkpoint saver.
|
||||
self._register_load_state_dict_pre_hook(self._remap_weights, with_module=True)
|
||||
|
||||
@staticmethod
|
||||
def _remap_weights(module, state_dict, prefix, *_):
|
||||
for old in ("in_proj_weight", "in_proj.weight"):
|
||||
key = prefix + old
|
||||
if key in state_dict and (prefix + "in_proj.weight") not in state_dict:
|
||||
state_dict[prefix + "in_proj.weight"] = state_dict.pop(key)
|
||||
for old in ("out_proj.weight",):
|
||||
key = prefix + old
|
||||
# If stored under in_projs.0 / out_projs.0 (multi-module layout), remap.
|
||||
for src in (prefix + "in_projs.0.weight", prefix + "out_projs.0.weight"):
|
||||
if src in state_dict:
|
||||
dst = src.replace("in_projs.0.", "in_proj.").replace("out_projs.0.", "out_proj.")
|
||||
state_dict[dst] = state_dict.pop(src)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
B, T, _ = x.shape
|
||||
H, D = self.num_heads, self.embed_dim // self.num_heads
|
||||
qkv = self.in_proj(x).reshape(B, T, 3, H, D).permute(2, 0, 3, 1, 4)
|
||||
q, k, v = qkv[0], qkv[1], qkv[2] # each (B, H, T, D)
|
||||
q, k = _apply_rope(q, k, self.max_period)
|
||||
out = F.scaled_dot_product_attention(q, k, v, is_causal=self.causal)
|
||||
out = out.transpose(1, 2).reshape(B, T, self.embed_dim)
|
||||
return self.out_proj(out)
|
||||
|
||||
|
||||
class _TransformerLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
d_model: int,
|
||||
num_heads: int,
|
||||
dim_feedforward: int,
|
||||
causal: bool,
|
||||
max_period: float,
|
||||
layer_scale: float | None,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
kw = {"device": device, "dtype": dtype}
|
||||
self.norm1 = nn.LayerNorm(d_model, eps=1e-5, **kw)
|
||||
self.norm2 = nn.LayerNorm(d_model, eps=1e-5, **kw)
|
||||
self.attn = _Attention(d_model, num_heads, causal, max_period, device, dtype)
|
||||
self.ff1 = nn.Linear(d_model, dim_feedforward, bias=False, **kw)
|
||||
self.ff2 = nn.Linear(dim_feedforward, d_model, bias=False, **kw)
|
||||
ls_kw = cast(dict[str, object], kw)
|
||||
if layer_scale is not None:
|
||||
self.ls1 = _LayerScale(d_model, layer_scale, **ls_kw)
|
||||
self.ls2 = _LayerScale(d_model, layer_scale, **ls_kw)
|
||||
else:
|
||||
self.ls1 = nn.Identity()
|
||||
self.ls2 = nn.Identity()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = x + self.ls1(self.attn(self.norm1(x)))
|
||||
x = x + self.ls2(self.ff2(F.gelu(self.ff1(self.norm2(x)))))
|
||||
return x
|
||||
|
||||
|
||||
class _Transformer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
d_model: int,
|
||||
num_heads: int,
|
||||
num_layers: int,
|
||||
dim_feedforward: int,
|
||||
causal: bool,
|
||||
max_period: float,
|
||||
layer_scale: float | None,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.layers = nn.ModuleList(
|
||||
[
|
||||
_TransformerLayer(d_model, num_heads, dim_feedforward, causal, max_period, layer_scale, device, dtype)
|
||||
for _ in range(num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
for layer in self.layers:
|
||||
x = layer(x)
|
||||
return x
|
||||
|
||||
|
||||
class _ProjectedTransformer(nn.Module):
|
||||
"""Input/output projection + causal transformer (conv-layout: B, D, T)."""
|
||||
|
||||
def __init__(
|
||||
self, input_dimension: int, output_dimension: int, d_model: int, *, module_type: str, **kwargs: Any
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.downsample_ratio: int = 1
|
||||
self.in_proj = nn.Linear(input_dimension, d_model, bias=False) if d_model != input_dimension else nn.Identity()
|
||||
self.transformer = _Transformer(d_model=d_model, **kwargs)
|
||||
self.out_proj = (
|
||||
nn.Linear(d_model, output_dimension, bias=False) if d_model != output_dimension else nn.Identity()
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor, lengths: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
x = self.in_proj(x.transpose(1, 2)) # (B, D, T) → (B, T, d_model)
|
||||
x = self.transformer(x)
|
||||
x = self.out_proj(x).transpose(1, 2) # (B, T, D) → (B, D, T)
|
||||
return x, lengths
|
||||
|
||||
|
||||
class _PatchedPretransform(nn.Module):
|
||||
"""Patch-based up/down-sampling (no learned weights)."""
|
||||
|
||||
def __init__(self, patch_size: int, is_downsample: bool, module_type: str, **_: Any) -> None:
|
||||
super().__init__()
|
||||
self.patch_size = patch_size
|
||||
self.downsample_ratio: int = patch_size
|
||||
self.is_downsample = is_downsample
|
||||
|
||||
def forward(self, x: torch.Tensor, lengths: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
b, d, _ = x.shape
|
||||
h = self.patch_size
|
||||
if self.is_downsample:
|
||||
x = x.reshape(b, d, -1, h).permute(0, 1, 3, 2).reshape(b, d * h, -1)
|
||||
return x, lengths // h
|
||||
else:
|
||||
# Upsample by ``h``: split channel dim, then flatten last two dims.
|
||||
# Was ``reshape(b, d // h, -1 * h)`` which evaluates to ``-h`` and
|
||||
# is rejected by torch (the failing path produced ``[1, 640, -2]``).
|
||||
x = x.reshape(b, d // h, h, -1).permute(0, 1, 3, 2).reshape(b, d // h, -1)
|
||||
return x, lengths * h
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Vector Quantization
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _wn_conv1d(*args: Any, **kwargs: Any) -> nn.Module:
|
||||
return nn.utils.parametrizations.weight_norm(nn.Conv1d(*args, **kwargs))
|
||||
|
||||
|
||||
class _VQ(nn.Module):
|
||||
"""Single RVQ codebook (inference: encode + decode)."""
|
||||
|
||||
def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int, **_: Any) -> None:
|
||||
super().__init__()
|
||||
self.in_proj = _wn_conv1d(input_dim, codebook_dim, 1) if input_dim != codebook_dim else nn.Identity()
|
||||
self.out_proj = _wn_conv1d(codebook_dim, input_dim, 1) if input_dim != codebook_dim else nn.Identity()
|
||||
self.codebook = nn.Embedding(codebook_size, codebook_dim)
|
||||
|
||||
@torch.no_grad()
|
||||
def encode(self, z: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
z_e = self.in_proj(z.float())
|
||||
enc = z_e.transpose(1, 2).reshape(-1, z_e.shape[1])
|
||||
w = self.codebook.weight.float()
|
||||
dist = enc.pow(2).sum(1, keepdim=True) - 2 * enc @ w.t() + w.pow(2).sum(1, keepdim=True).t()
|
||||
ids = (-dist).max(1)[1].reshape(z.size(0), -1)
|
||||
return self.out_proj(self.codebook(ids).transpose(1, 2).float()).float(), ids
|
||||
|
||||
def decode(self, ids: torch.Tensor) -> torch.Tensor:
|
||||
return self.out_proj(self.codebook(ids).transpose(1, 2).float()).float()
|
||||
|
||||
|
||||
class _LFQ(nn.Module):
|
||||
"""Single RLFQ codebook (inference: encode + decode)."""
|
||||
|
||||
def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int, **_: Any) -> None:
|
||||
super().__init__()
|
||||
self.in_proj = _wn_conv1d(input_dim, codebook_dim, 1) if input_dim != codebook_dim else nn.Identity()
|
||||
self.out_proj = _wn_conv1d(codebook_dim, input_dim, 1) if input_dim != codebook_dim else nn.Identity()
|
||||
self.codebook = nn.Embedding(codebook_size, codebook_dim)
|
||||
|
||||
@torch.no_grad()
|
||||
def encode(self, z: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
z_e = self.in_proj(z.float())
|
||||
enc = F.normalize(z_e.transpose(1, 2).reshape(-1, z_e.shape[1]).float())
|
||||
cb = F.normalize(self.codebook.weight.float())
|
||||
dist = enc.pow(2).sum(1, keepdim=True) - 2 * enc @ cb.t() + cb.pow(2).sum(1, keepdim=True).t()
|
||||
ids = (-dist).max(1)[1].reshape(z.size(0), -1)
|
||||
z_q = self.codebook(ids).transpose(1, 2).float()
|
||||
z_q = (z_e + (z_q - z_e).detach()).float()
|
||||
return self.out_proj(z_q).float(), ids
|
||||
|
||||
def decode(self, ids: torch.Tensor) -> torch.Tensor:
|
||||
return self.out_proj(self.codebook(ids).transpose(1, 2).float()).float()
|
||||
|
||||
|
||||
class _ResidualQ(nn.Module):
|
||||
"""Residual VQ or LFQ stack."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_dim: int,
|
||||
rvq_dim: int | None,
|
||||
output_dim: int | None,
|
||||
num_quantizers: int,
|
||||
codebook_size: int,
|
||||
codebook_dim: int,
|
||||
quantizer_type: str = "rlfq",
|
||||
**_: Any,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.rvq_dim = rvq_dim or input_dim
|
||||
self.output_dim = output_dim or input_dim
|
||||
QCls = _LFQ if quantizer_type in {"rlfq", "random_prefix_rlfq"} else _VQ
|
||||
self.input_proj = _wn_conv1d(input_dim, self.rvq_dim, 1) if input_dim != self.rvq_dim else nn.Identity()
|
||||
self.output_proj = (
|
||||
_wn_conv1d(self.rvq_dim, self.output_dim, 1) if self.rvq_dim != self.output_dim else nn.Identity()
|
||||
)
|
||||
self.quantizers = nn.ModuleList(
|
||||
[QCls(self.rvq_dim, codebook_size, codebook_dim) for _ in range(num_quantizers)]
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
def encode(
|
||||
self, z: torch.Tensor, lengths: torch.Tensor, n: int | None = None
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
z = self.input_proj(z)
|
||||
B, _, T = z.shape
|
||||
mask = torch.arange(T, device=z.device).expand(B, T) < lengths.unsqueeze(1)
|
||||
out = torch.zeros_like(z, dtype=torch.float32)
|
||||
residual = z.clone().float()
|
||||
codes: list[torch.Tensor] = []
|
||||
for i, q in enumerate(self.quantizers[: n or len(self.quantizers)]):
|
||||
zq, ids = q.encode(residual * mask.unsqueeze(1))
|
||||
out += zq * mask.unsqueeze(1)
|
||||
residual -= zq * mask.unsqueeze(1)
|
||||
codes.append(ids)
|
||||
return self.output_proj(out), torch.stack(codes), lengths
|
||||
|
||||
def decode_codes(self, codes: torch.Tensor) -> torch.Tensor:
|
||||
"""codes: (NQ, B, T) → (B, rvq_dim, T)."""
|
||||
nq, B, T = codes.shape
|
||||
emb = torch.zeros(B, self.rvq_dim, T, device=codes.device, dtype=torch.float32)
|
||||
for i, q in enumerate(self.quantizers[:nq]):
|
||||
emb += q.decode(codes[i])
|
||||
return self.output_proj(emb)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _build_modules(
|
||||
specs: list[dict[str, Any]], is_downsample: bool, context_duration: float, sampling_rate: float
|
||||
) -> nn.ModuleList:
|
||||
modules: list[nn.Module] = []
|
||||
rate = float(sampling_rate)
|
||||
for spec in specs:
|
||||
spec = dict(spec)
|
||||
if spec["module_type"] == "PatchedPretransform":
|
||||
m = _PatchedPretransform(**spec, is_downsample=is_downsample)
|
||||
else:
|
||||
spec.pop("conv_layout", None)
|
||||
spec.pop("module_type", None)
|
||||
spec.pop("gating", None)
|
||||
spec.pop("positional_embedding", None)
|
||||
spec.pop("norm", None)
|
||||
spec.pop("causal", None)
|
||||
m = _ProjectedTransformer(
|
||||
module_type="Transformer",
|
||||
causal=True,
|
||||
max_period=spec.pop("max_period", 10000),
|
||||
layer_scale=spec.pop("layer_scale", None),
|
||||
**spec,
|
||||
)
|
||||
modules.append(m)
|
||||
rate = rate / m.downsample_ratio if is_downsample else rate * m.downsample_ratio
|
||||
return nn.ModuleList(modules)
|
||||
|
||||
|
||||
class MossAudioTokenizerModel(PreTrainedModel):
|
||||
"""MOSS Audio Tokenizer — inference-only (batch_encode / batch_decode)."""
|
||||
|
||||
config_class = MossAudioTokenizerConfig
|
||||
base_model_prefix = ""
|
||||
supports_gradient_checkpointing = False
|
||||
|
||||
def __init__(self, config: MossAudioTokenizerConfig) -> None:
|
||||
super().__init__(config)
|
||||
self.sampling_rate = config.sampling_rate
|
||||
self.downsample_rate = config.downsample_rate
|
||||
|
||||
ctx = config.causal_transformer_context_duration
|
||||
sr = float(config.sampling_rate)
|
||||
self.encoder = _build_modules(config.encoder_kwargs, is_downsample=True, context_duration=ctx, sampling_rate=sr)
|
||||
self.decoder = _build_modules(
|
||||
copy.deepcopy(config.decoder_kwargs), is_downsample=False, context_duration=ctx, sampling_rate=sr
|
||||
)
|
||||
|
||||
kw = dict(config.quantizer_kwargs)
|
||||
self.quantizer = _ResidualQ(**kw)
|
||||
self.post_init()
|
||||
|
||||
@torch.no_grad()
|
||||
def batch_encode(
|
||||
self,
|
||||
wav_list: list[torch.Tensor],
|
||||
num_quantizers: int | None = None,
|
||||
) -> MossAudioTokenizerEncoderOutput:
|
||||
"""Encode a list of 1-D waveform tensors → RVQ codes (NQ, B, T)."""
|
||||
device = wav_list[0].device
|
||||
B = len(wav_list)
|
||||
max_len = max(w.shape[-1] for w in wav_list)
|
||||
x = torch.zeros(B, 1, max_len, device=device)
|
||||
lengths = torch.zeros(B, device=device, dtype=torch.long)
|
||||
for i, w in enumerate(wav_list):
|
||||
x[i, 0, : w.shape[-1]] = w
|
||||
lengths[i] = w.shape[-1]
|
||||
return self._encode(x, lengths, num_quantizers)
|
||||
|
||||
@torch.no_grad()
|
||||
def batch_decode(
|
||||
self,
|
||||
codes_list: list[torch.Tensor],
|
||||
num_quantizers: int | None = None,
|
||||
) -> MossAudioTokenizerDecoderOutput:
|
||||
"""Decode a list of (NQ, T) code tensors → waveforms."""
|
||||
device = codes_list[0].device
|
||||
B = len(codes_list)
|
||||
nq = num_quantizers or codes_list[0].shape[0]
|
||||
max_t = max(c.shape[-1] for c in codes_list)
|
||||
codes = torch.zeros(nq, B, max_t, device=device, dtype=torch.long)
|
||||
lengths = torch.zeros(B, device=device, dtype=torch.long)
|
||||
for i, c in enumerate(codes_list):
|
||||
codes[:nq, i, : c.shape[-1]] = c[:nq]
|
||||
lengths[i] = c.shape[-1]
|
||||
return self._decode(codes, lengths)
|
||||
|
||||
def _encode(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
lengths: torch.Tensor,
|
||||
n_quantizers: int | None = None,
|
||||
) -> MossAudioTokenizerEncoderOutput:
|
||||
if x.dim() == 2:
|
||||
x = x.unsqueeze(1)
|
||||
T = x.shape[-1]
|
||||
if T % self.downsample_rate:
|
||||
x = F.pad(x, (0, self.downsample_rate - T % self.downsample_rate))
|
||||
e, e_len = x, lengths
|
||||
for m in self.encoder:
|
||||
e, e_len = m(e, e_len)
|
||||
_, codes, code_len = self.quantizer.encode(e, e_len, n_quantizers)
|
||||
return MossAudioTokenizerEncoderOutput(audio_codes=codes, audio_codes_lengths=code_len)
|
||||
|
||||
def _decode(
|
||||
self,
|
||||
codes: torch.Tensor,
|
||||
lengths: torch.Tensor,
|
||||
) -> MossAudioTokenizerDecoderOutput:
|
||||
z = self.quantizer.decode_codes(codes)
|
||||
d, d_len = z, lengths
|
||||
for m in self.decoder:
|
||||
d, d_len = m(d, d_len)
|
||||
return MossAudioTokenizerDecoderOutput(audio=d, audio_lengths=d_len)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MossAudioTokenizerConfig",
|
||||
"MossAudioTokenizerModel",
|
||||
"MossAudioTokenizerEncoderOutput",
|
||||
"MossAudioTokenizerDecoderOutput",
|
||||
]
|
||||
@@ -0,0 +1,241 @@
|
||||
# Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""MOSS-TTS model configuration."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from transformers import AutoConfig
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.models.qwen3 import Qwen3Config
|
||||
|
||||
|
||||
class MossTTSDelayConfig(PretrainedConfig):
|
||||
"""Config for MossTTSDelayModel (8B: MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect;
|
||||
1.7B: MOSS-VoiceGenerator).
|
||||
|
||||
The HuggingFace checkpoint stores a nested ``language_config`` which is a
|
||||
Qwen3 config. We unwrap it here so that vLLM can use ``get_text_config()``
|
||||
to size KV caches and allocate memory correctly.
|
||||
"""
|
||||
|
||||
model_type = "moss_tts_delay"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
language_config: dict | None = None,
|
||||
n_vq: int = 32,
|
||||
audio_vocab_size: int = 1024,
|
||||
sampling_rate: int = 24000,
|
||||
audio_pad_code: int = 1024,
|
||||
audio_start_token_id: int = 151652,
|
||||
audio_end_token_id: int = 151653,
|
||||
audio_user_slot_token_id: int = 151654,
|
||||
audio_assistant_gen_slot_token_id: int = 151656,
|
||||
audio_assistant_delay_slot_token_id: int = 151662,
|
||||
pad_token_id: int = 151643,
|
||||
im_start_token_id: int = 151644,
|
||||
im_end_token_id: int = 151645,
|
||||
codec_model_name_or_path: str = "OpenMOSS-Team/MOSS-Audio-Tokenizer",
|
||||
**kwargs: object,
|
||||
) -> None:
|
||||
if language_config is None:
|
||||
language_config = {}
|
||||
if isinstance(language_config, dict):
|
||||
language_config.pop("model_type", None)
|
||||
self.language_config = Qwen3Config(**language_config)
|
||||
else:
|
||||
self.language_config = language_config
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
self.n_vq = n_vq
|
||||
self.audio_vocab_size = audio_vocab_size
|
||||
self.sampling_rate = sampling_rate
|
||||
self.audio_pad_code = audio_pad_code
|
||||
self.audio_start_token_id = audio_start_token_id
|
||||
self.audio_end_token_id = audio_end_token_id
|
||||
self.audio_user_slot_token_id = audio_user_slot_token_id
|
||||
self.audio_assistant_gen_slot_token_id = audio_assistant_gen_slot_token_id
|
||||
self.audio_assistant_delay_slot_token_id = audio_assistant_delay_slot_token_id
|
||||
self.pad_token_id = pad_token_id
|
||||
self.im_start_token_id = im_start_token_id
|
||||
self.im_end_token_id = im_end_token_id
|
||||
self.codec_model_name_or_path = codec_model_name_or_path
|
||||
|
||||
def get_text_config(self, **_: object) -> Qwen3Config:
|
||||
return self.language_config
|
||||
|
||||
|
||||
class MossTTSLocalTransformerConfig(PretrainedConfig):
|
||||
"""Config for the lightweight local depth transformer in MossTTSRealtime.
|
||||
|
||||
Mirrors upstream ``MossTTSRealtimeLocalTransformerConfig`` so the realtime
|
||||
checkpoint loads via ``AutoConfig`` without dropping fields silently.
|
||||
"""
|
||||
|
||||
model_type = "moss_tts_realtime_local_transformer"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
head_dim: int = 128,
|
||||
hidden_size: int = 2048,
|
||||
num_hidden_layers: int = 4,
|
||||
num_attention_heads: int = 16,
|
||||
num_key_value_heads: int = 8,
|
||||
intermediate_size: int = 6144,
|
||||
max_position_embeddings: int = 33,
|
||||
rms_norm_eps: float = 1e-6,
|
||||
rope_theta: float = 1_000_000.0,
|
||||
rope_type: str = "linear",
|
||||
attention_bias: bool = False,
|
||||
attention_dropout: float = 0.0,
|
||||
hidden_act: str = "silu",
|
||||
initializer_range: float = 0.02,
|
||||
use_cache: bool = True,
|
||||
tie_word_embeddings: bool = False,
|
||||
pad_token_id: int = 1024,
|
||||
audio_pad_token: int = 1024,
|
||||
audio_vocab_size: int = 1027,
|
||||
rvq: int = 16,
|
||||
**kwargs: object,
|
||||
) -> None:
|
||||
super().__init__(pad_token_id=pad_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs)
|
||||
self.head_dim = head_dim
|
||||
self.hidden_size = hidden_size
|
||||
self.num_hidden_layers = num_hidden_layers
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.num_key_value_heads = num_key_value_heads
|
||||
self.intermediate_size = intermediate_size
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.rms_norm_eps = rms_norm_eps
|
||||
self.rope_theta = rope_theta
|
||||
self.rope_type = rope_type
|
||||
self.attention_bias = attention_bias
|
||||
self.attention_dropout = attention_dropout
|
||||
self.hidden_act = hidden_act
|
||||
self.initializer_range = initializer_range
|
||||
self.use_cache = use_cache
|
||||
self.audio_pad_token = audio_pad_token
|
||||
self.audio_vocab_size = audio_vocab_size
|
||||
self.rvq = rvq
|
||||
# The shared ``CodePredictorBaseModel`` (reused from Qwen3-TTS /
|
||||
# Qwen3-Omni) reads ``num_code_groups`` and ``vocab_size`` to size its
|
||||
# codebook embeddings and causal mask. MossTTS names the same
|
||||
# quantities ``rvq`` and ``audio_vocab_size``; mirror them so the
|
||||
# common transformer body can be driven from this config unchanged.
|
||||
self.num_code_groups = rvq
|
||||
self.vocab_size = audio_vocab_size
|
||||
# rope_parameters used by transformers.modeling_rope_utils
|
||||
self.rope_parameters = {
|
||||
"rope_type": rope_type,
|
||||
"rope_theta": rope_theta,
|
||||
"factor": 1.0,
|
||||
}
|
||||
|
||||
|
||||
class MossTTSRealtimeConfig(PretrainedConfig):
|
||||
"""Config for MossTTSRealtime (1.7B, TTFB ~180 ms streaming model).
|
||||
|
||||
Mirrors the upstream layout: a nested ``language_config`` (Qwen3) plus a
|
||||
``local_config`` for the per-step depth transformer that decodes the RVQ
|
||||
codebooks. The flat fields here are the *outer* model's attributes.
|
||||
"""
|
||||
|
||||
model_type = "moss_tts_realtime"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
language_config: dict | None = None,
|
||||
local_config: dict | None = None,
|
||||
rvq: int = 16,
|
||||
audio_vocab_size: int = 1027,
|
||||
audio_pad_token: int = 1024,
|
||||
sampling_rate: int = 24000,
|
||||
reference_audio_pad: int = 151654,
|
||||
text_pad: int = 151655,
|
||||
vocab_size: int = 151936,
|
||||
bos_token_id: int = 151643,
|
||||
eos_token_id: int = 151645,
|
||||
initializer_range: float = 0.02,
|
||||
codec_model_name_or_path: str = "OpenMOSS-Team/MOSS-Audio-Tokenizer",
|
||||
**kwargs: object,
|
||||
) -> None:
|
||||
if language_config is None:
|
||||
language_config = {}
|
||||
if isinstance(language_config, dict):
|
||||
language_config.pop("model_type", None)
|
||||
self.language_config = Qwen3Config(**language_config)
|
||||
else:
|
||||
self.language_config = language_config
|
||||
|
||||
if local_config is None:
|
||||
local_config = {}
|
||||
if isinstance(local_config, dict):
|
||||
local_config.pop("model_type", None)
|
||||
self.local_config = MossTTSLocalTransformerConfig(**local_config)
|
||||
else:
|
||||
self.local_config = local_config
|
||||
|
||||
super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
|
||||
|
||||
self.rvq = rvq
|
||||
self.n_vq = rvq # alias for uniform access from talker code
|
||||
self.audio_vocab_size = audio_vocab_size
|
||||
self.audio_pad_token = audio_pad_token
|
||||
self.sampling_rate = sampling_rate
|
||||
self.reference_audio_pad = reference_audio_pad
|
||||
self.text_pad = text_pad
|
||||
self.vocab_size = vocab_size
|
||||
self.initializer_range = initializer_range
|
||||
self.codec_model_name_or_path = codec_model_name_or_path
|
||||
|
||||
# The talker still reads ``hidden_size`` etc. directly from the config in
|
||||
# places — proxy them through to ``language_config`` for backward compat.
|
||||
@property
|
||||
def hidden_size(self) -> int:
|
||||
return int(self.language_config.hidden_size)
|
||||
|
||||
@property
|
||||
def num_hidden_layers(self) -> int:
|
||||
return int(self.language_config.num_hidden_layers)
|
||||
|
||||
@property
|
||||
def num_attention_heads(self) -> int:
|
||||
return int(self.language_config.num_attention_heads)
|
||||
|
||||
@property
|
||||
def intermediate_size(self) -> int:
|
||||
return int(self.language_config.intermediate_size)
|
||||
|
||||
@property
|
||||
def max_position_embeddings(self) -> int:
|
||||
return int(self.language_config.max_position_embeddings)
|
||||
|
||||
@property
|
||||
def rms_norm_eps(self) -> float:
|
||||
return float(self.language_config.rms_norm_eps)
|
||||
|
||||
@property
|
||||
def rope_theta(self) -> float:
|
||||
return float(self.language_config.rope_theta)
|
||||
|
||||
# vLLM sizes KV cache from get_text_config(); return the inner Qwen3 cfg.
|
||||
def get_text_config(self, **_: object) -> Qwen3Config:
|
||||
return self.language_config
|
||||
|
||||
|
||||
AutoConfig.register("moss_tts_delay", MossTTSDelayConfig)
|
||||
AutoConfig.register("moss_tts_realtime", MossTTSRealtimeConfig)
|
||||
|
||||
__all__ = ["MossTTSDelayConfig", "MossTTSRealtimeConfig", "MossTTSLocalTransformerConfig"]
|
||||
@@ -0,0 +1,301 @@
|
||||
# Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License").
|
||||
"""MOSS-TTS Stage-1 codec decoder: RVQ codes → 24 kHz waveform."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from collections.abc import Iterable
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.model_loader import DefaultModelLoader
|
||||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
|
||||
from vllm_omni.model_executor.models.moss_tts.audio_tokenizer import (
|
||||
MossAudioTokenizerConfig,
|
||||
MossAudioTokenizerModel,
|
||||
)
|
||||
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class MossTTSCodecDecoder(nn.Module):
|
||||
"""Stage-1 decoder for all MOSS-TTS variants.
|
||||
|
||||
Consumes ``(T, NQ)`` audio VQ codes emitted by Stage 0 and decodes them
|
||||
to a 24 kHz mono waveform using the vendored
|
||||
``MossAudioTokenizerModel``.
|
||||
|
||||
All five variants share the same codec checkpoint
|
||||
``OpenMOSS-Team/MOSS-Audio-Tokenizer``. The number of quantizers
|
||||
(``n_vq``) is passed as a ``num_quantizers`` argument to ``batch_decode``
|
||||
so one codec instance handles both ``n_vq=32`` (MOSS-TTS) and ``n_vq=16``
|
||||
(all other variants) without swapping weights.
|
||||
|
||||
The codec checkpoint path defaults to the value stored in
|
||||
``vllm_config.model_config.hf_config.codec_model_name_or_path`` but can
|
||||
be overridden by setting the environment variable
|
||||
``MOSS_TTS_CODEC_PATH``.
|
||||
"""
|
||||
|
||||
input_modalities = "audio"
|
||||
|
||||
have_multimodal_outputs: bool = True
|
||||
has_preprocess: bool = False
|
||||
has_postprocess: bool = False
|
||||
enable_update_additional_information: bool = True
|
||||
requires_raw_input_tokens: bool = True
|
||||
|
||||
_OUTPUT_SAMPLE_RATE: int = 24_000
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
|
||||
super().__init__()
|
||||
self.vllm_config = vllm_config
|
||||
|
||||
cfg = vllm_config.model_config.hf_config
|
||||
self._n_vq: int = int(getattr(cfg, "n_vq", getattr(cfg, "rvq", 16)))
|
||||
self._codec_path: str = str(getattr(cfg, "codec_model_name_or_path", "OpenMOSS-Team/MOSS-Audio-Tokenizer"))
|
||||
|
||||
# Resolved at load_weights() time
|
||||
self._codec: MossAudioTokenizerModel | None = None
|
||||
self._sr_tensor = torch.tensor(self._OUTPUT_SAMPLE_RATE, dtype=torch.int32)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# vLLM-Omni stubs (codec has no AR loop)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
|
||||
if input_ids.numel() == 0:
|
||||
return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32)
|
||||
return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32)
|
||||
|
||||
def compute_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor | OmniOutput,
|
||||
sampling_metadata: Any = None,
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Decode
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@torch.no_grad()
|
||||
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:
|
||||
"""Decode audio VQ codes to waveform.
|
||||
|
||||
Stage 0 emits flat codebook-major ``[NQ * T_chunk]`` audio codes. The
|
||||
chunk transfer adapter assigns those to ``request.prompt_token_ids``,
|
||||
which arrives here as ``input_ids``. Per-request offsets are derived
|
||||
from ``runtime_additional_information``-attached request metadata; meta
|
||||
like ``left_context_size`` (overlap context for the causal decoder)
|
||||
also lives there.
|
||||
|
||||
Returns
|
||||
-------
|
||||
OmniOutput with:
|
||||
multimodal_outputs["model_outputs"] — list of (T_wav,) float32 tensors
|
||||
multimodal_outputs["sr"] — list of scalar int32 tensors
|
||||
"""
|
||||
sr_tensor = self._sr_tensor
|
||||
empty = torch.zeros((0,), dtype=torch.float32)
|
||||
info_list: list[dict[str, Any]] = runtime_additional_information or []
|
||||
num_req = max(len(info_list), 1)
|
||||
|
||||
if self._codec is None:
|
||||
logger.warning("MossTTSCodecDecoder called before load_weights(); returning silence.")
|
||||
return OmniOutput(
|
||||
text_hidden_states=None,
|
||||
multimodal_outputs={
|
||||
"model_outputs": [empty] * num_req,
|
||||
"sr": [sr_tensor] * num_req,
|
||||
},
|
||||
)
|
||||
|
||||
audios: list[torch.Tensor] = [empty] * num_req
|
||||
srs: list[torch.Tensor] = [sr_tensor] * num_req
|
||||
device = next(self._codec.parameters()).device
|
||||
|
||||
if input_ids is None or input_ids.numel() == 0:
|
||||
return OmniOutput(
|
||||
text_hidden_states=None,
|
||||
multimodal_outputs={"model_outputs": audios, "sr": srs},
|
||||
)
|
||||
|
||||
# ``input_ids`` is concatenated across all requests; split by
|
||||
# query_start_loc-style offsets from ``num_scheduled_tokens`` in
|
||||
# kwargs if available, else assume one request.
|
||||
ids_flat = input_ids.reshape(-1).to(dtype=torch.long)
|
||||
num_scheduled_tokens = kwargs.get("num_scheduled_tokens")
|
||||
if isinstance(num_scheduled_tokens, list) and len(num_scheduled_tokens) == num_req:
|
||||
offsets = [0]
|
||||
for n in num_scheduled_tokens:
|
||||
offsets.append(offsets[-1] + int(n))
|
||||
else:
|
||||
offsets = [0, int(ids_flat.shape[0])]
|
||||
|
||||
for i, info in enumerate(info_list):
|
||||
if i + 1 >= len(offsets):
|
||||
break
|
||||
seg = ids_flat[offsets[i] : offsets[i + 1]]
|
||||
if seg.numel() == 0:
|
||||
continue
|
||||
if seg.numel() % self._n_vq != 0:
|
||||
logger.warning(
|
||||
"MossTTS codec input length %d not divisible by n_vq %d; skipping.",
|
||||
int(seg.numel()),
|
||||
self._n_vq,
|
||||
)
|
||||
continue
|
||||
t_chunk = int(seg.numel() // self._n_vq)
|
||||
codes_nq_t = seg.reshape(self._n_vq, t_chunk).to(device=device)
|
||||
# Clamp out-of-range codes (the talker uses ``audio_pad_code``
|
||||
# =``codebook_size`` for delay-pattern padding; passing it to the
|
||||
# codec embedding lookup would be an OOB index. The downstream
|
||||
# de-delay step is intended to drop those rows, but until that's
|
||||
# wired we clamp so the kernel does not assert.)
|
||||
codebook_size = self._codec.config.codebook_size
|
||||
codes_nq_t = codes_nq_t.clamp_(0, int(codebook_size) - 1)
|
||||
|
||||
# left_context_size: number of left-context frames prepended to
|
||||
# this chunk by the stage input processor (overlap-add pattern,
|
||||
# same as Qwen3-TTS code2wav). We decode them together with the
|
||||
# real chunk for smooth boundaries, then trim them from the output.
|
||||
meta = (info.get("meta", {}) if isinstance(info, dict) else {}) or {}
|
||||
left_ctx = meta.get("left_context_size", 0)
|
||||
if isinstance(left_ctx, (list, tuple)):
|
||||
left_ctx = int(left_ctx[0]) if left_ctx else 0
|
||||
elif isinstance(left_ctx, torch.Tensor):
|
||||
left_ctx = int(left_ctx.reshape(-1)[0].item()) if left_ctx.numel() else 0
|
||||
left_ctx = int(left_ctx)
|
||||
|
||||
out = self._codec.batch_decode(codes_list=[codes_nq_t], num_quantizers=self._n_vq)
|
||||
|
||||
if out.audio is None:
|
||||
continue
|
||||
|
||||
wav = out.audio.reshape(-1).to(dtype=torch.float32).cpu()
|
||||
if out.audio_lengths is not None:
|
||||
wav = wav[: int(out.audio_lengths[0].item())]
|
||||
|
||||
# Trim left-context samples.
|
||||
if left_ctx > 0:
|
||||
trim = left_ctx * self._codec.downsample_rate
|
||||
wav = wav[trim:]
|
||||
|
||||
audios[i] = wav
|
||||
|
||||
return OmniOutput(
|
||||
text_hidden_states=None,
|
||||
multimodal_outputs={"model_outputs": audios, "sr": srs},
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Weight loading
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
"""Drain the Stage-0 weights iterator, then load the codec from its own checkpoint.
|
||||
|
||||
The codec lives in a separate HuggingFace repo
|
||||
(``OpenMOSS-Team/MOSS-Audio-Tokenizer``) and is loaded independently
|
||||
of the talker weights.
|
||||
"""
|
||||
# Drain the incoming weights iterator — all Stage-0 weights are
|
||||
# irrelevant to this stage.
|
||||
for _ in weights:
|
||||
pass
|
||||
|
||||
codec_path = os.environ.get("MOSS_TTS_CODEC_PATH", self._codec_path)
|
||||
logger.info("Loading MOSS Audio Tokenizer from %s", codec_path)
|
||||
|
||||
codec_cfg = MossAudioTokenizerConfig.from_pretrained(codec_path)
|
||||
codec = MossAudioTokenizerModel(codec_cfg)
|
||||
|
||||
model_loader = DefaultModelLoader(self.vllm_config.load_config)
|
||||
source = DefaultModelLoader.Source(
|
||||
model_or_path=codec_path,
|
||||
revision=None,
|
||||
subfolder=None,
|
||||
)
|
||||
codec_weights = model_loader._get_weights_iterator(source)
|
||||
params_dict = dict(codec.named_parameters())
|
||||
|
||||
# Upstream MossAudioTokenizer uses different submodule names than the
|
||||
# vendored re-implementation in ``audio_tokenizer.py``. Without this
|
||||
# remap only ~half the codec parameters load (codebooks + WN convs)
|
||||
# and the rest stay at their random init, which produces noise that
|
||||
# sounds correct in duration but is structurally garbage.
|
||||
_SUFFIX_REMAP: list[tuple[str, str]] = [
|
||||
(".self_attn.in_projs.0.", ".attn.in_proj."),
|
||||
(".self_attn.out_projs.0.", ".attn.out_proj."),
|
||||
(".linear1.", ".ff1."),
|
||||
(".linear2.", ".ff2."),
|
||||
(".layer_scale_1.", ".ls1."),
|
||||
(".layer_scale_2.", ".ls2."),
|
||||
(".input_proj.", ".in_proj."),
|
||||
(".output_proj.", ".out_proj."),
|
||||
]
|
||||
|
||||
def _remap(name: str) -> str:
|
||||
for src, dst in _SUFFIX_REMAP:
|
||||
if src in name:
|
||||
return name.replace(src, dst)
|
||||
return name
|
||||
|
||||
loaded_total = 0
|
||||
skipped: list[str] = []
|
||||
for name, tensor in codec_weights:
|
||||
# Try direct name first (e.g. ``quantizer.input_proj.*`` exists
|
||||
# under the same name in both layouts), then the remap (transformer
|
||||
# submodules need ``.linear1.``→``.ff1.`` etc.).
|
||||
tgt = name if name in params_dict else _remap(name)
|
||||
if tgt in params_dict:
|
||||
default_weight_loader(params_dict[tgt], tensor)
|
||||
loaded_total += 1
|
||||
else:
|
||||
skipped.append(name)
|
||||
logger.info(
|
||||
"MOSS Audio Tokenizer weights: loaded=%d/%d skipped=%d (first skipped: %s)",
|
||||
loaded_total,
|
||||
len(params_dict),
|
||||
len(skipped),
|
||||
skipped[:3] if skipped else "none",
|
||||
)
|
||||
|
||||
device = self.vllm_config.device_config.device
|
||||
codec.to(device=device, dtype=torch.float32)
|
||||
codec.eval()
|
||||
self._codec = codec
|
||||
|
||||
logger.info(
|
||||
"MOSS Audio Tokenizer loaded: sampling_rate=%d, n_vq=%d",
|
||||
codec_cfg.sampling_rate,
|
||||
codec_cfg.num_quantizers,
|
||||
)
|
||||
|
||||
# Move sr_tensor to the same device
|
||||
self._sr_tensor = self._sr_tensor.to(device=device)
|
||||
|
||||
# vLLM's track_weights_loading() compares the returned set against
|
||||
# ``self.named_parameters()``. After ``self._codec = codec`` above,
|
||||
# those parameters are registered with the ``_codec.`` prefix, so
|
||||
# mirror that here.
|
||||
return {f"_codec.{name}" for name, _ in codec.named_parameters()}
|
||||
|
||||
|
||||
__all__ = ["MossTTSCodecDecoder"]
|
||||
@@ -0,0 +1,159 @@
|
||||
# Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License").
|
||||
"""Local depth transformer for MossTTSRealtime.
|
||||
|
||||
A small (4-layer) Qwen3-style decoder that generates the rvq=16 RVQ codebook
|
||||
codes for one audio frame, autoregressively over codebooks. It runs inside the
|
||||
talker's per-step ``make_omni_output``, independent from vLLM's main scheduler.
|
||||
|
||||
The transformer body is shared with Qwen3-TTS and Qwen3-Omni via
|
||||
``common.qwen3_code_predictor.CodePredictorBaseModel`` (re-prefill, no KV cache,
|
||||
HF-compatible numerics). MossTTSRealtime differs in two ways:
|
||||
|
||||
* codebook 0 is generated here from ``backbone_last_hidden`` (the other models
|
||||
receive it from the talker's main LM head), so we run one extra step and own
|
||||
all ``rvq`` LM heads;
|
||||
* sampling adds top-p and a windowed repetition penalty on top of
|
||||
temperature + top-k, matching upstream ``MossTTSRealtimeInference.generate``.
|
||||
|
||||
Re-prefilling the (<=rvq) frame each step is numerically identical to the
|
||||
previous KV-cache loop -- causal attention, only the last position is read.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from vllm_omni.model_executor.models.common.qwen3_code_predictor import (
|
||||
CodePredictorBaseModel,
|
||||
)
|
||||
from vllm_omni.model_executor.models.moss_tts.configuration_moss_tts import (
|
||||
MossTTSLocalTransformerConfig,
|
||||
)
|
||||
|
||||
|
||||
class MossTTSRealtimeLocalTransformer(nn.Module):
|
||||
"""Per-frame depth transformer. Mirrors upstream ``...LocalTransformer``.
|
||||
|
||||
State per audio frame:
|
||||
- The first input token uses ``backbone_last_hidden_state`` as the
|
||||
embedding (codebook 0's "input" is the backbone hidden, not a token).
|
||||
- Subsequent tokens (1..rvq-1) embed via ``model.codec_embedding[idx-1]``.
|
||||
|
||||
Outputs:
|
||||
- One logit row per codebook position, projected through
|
||||
``local_lm_heads[codebook_idx]`` (passed in by the talker).
|
||||
"""
|
||||
|
||||
def __init__(self, cfg: MossTTSLocalTransformerConfig) -> None:
|
||||
super().__init__()
|
||||
self.config = cfg
|
||||
# Shared body. Its codec_embedding holds rvq-1 embeddings -- upstream's
|
||||
# embed_tokens for codebooks 1..rvq-1 (codebook 0 uses the backbone
|
||||
# hidden). embedding_dim == hidden_size, so no projection is needed.
|
||||
self.model = CodePredictorBaseModel(
|
||||
cfg,
|
||||
embedding_dim=cfg.hidden_size,
|
||||
use_parallel_embedding=False,
|
||||
prefix="model",
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
def generate_frame(
|
||||
self,
|
||||
backbone_last_hidden: torch.Tensor, # (B, H)
|
||||
lm_heads: nn.ModuleList, # ModuleList of rvq Linear(H -> audio_vocab_size)
|
||||
*,
|
||||
temperature: float = 1.0,
|
||||
top_k: int = 50,
|
||||
top_p: float = 0.95,
|
||||
do_sample: bool = True,
|
||||
repetition_penalty: float = 1.0,
|
||||
history_per_codebook: list[list[int]] | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Generate one audio frame (rvq codebook tokens) for batch B.
|
||||
|
||||
Returns a ``(B, rvq)`` LongTensor.
|
||||
|
||||
``history_per_codebook[i]`` is a list of recently-emitted token ids for
|
||||
codebook ``i``; when ``repetition_penalty != 1.0`` those tokens get
|
||||
their logits scaled down (mirrors upstream's rep-penalty behaviour).
|
||||
"""
|
||||
device = backbone_last_hidden.device
|
||||
B = backbone_last_hidden.shape[0]
|
||||
rvq = self.config.rvq
|
||||
hidden_size = self.config.hidden_size
|
||||
|
||||
codec_embeds = self.model.codec_embedding
|
||||
|
||||
# Re-prefill buffer: position 0 = backbone hidden, position s = embed of
|
||||
# code_{s-1}. At step s we forward positions [0..s] and read the last.
|
||||
embeds = backbone_last_hidden.new_zeros((B, rvq, hidden_size))
|
||||
embeds[:, 0, :] = backbone_last_hidden.to(embeds.dtype)
|
||||
|
||||
pos_ids_full = torch.arange(rvq, device=device, dtype=torch.long).unsqueeze(0).expand(B, -1)
|
||||
|
||||
codes = backbone_last_hidden.new_zeros((B, rvq), dtype=torch.long)
|
||||
|
||||
for step in range(rvq):
|
||||
seq_len = step + 1
|
||||
hidden = self.model(embeds[:, :seq_len, :], pos_ids_full[:, :seq_len])
|
||||
logits = lm_heads[step](hidden[:, step, :]).float()
|
||||
|
||||
if repetition_penalty != 1.0 and history_per_codebook is not None and step < len(history_per_codebook):
|
||||
hist = history_per_codebook[step]
|
||||
if hist:
|
||||
hist_t = torch.tensor(hist, dtype=torch.long, device=logits.device)
|
||||
sel = logits.index_select(-1, hist_t)
|
||||
pos = sel > 0
|
||||
sel = torch.where(pos, sel / repetition_penalty, sel * repetition_penalty)
|
||||
logits.index_copy_(-1, hist_t, sel)
|
||||
|
||||
codes[:, step] = _sample_token(logits, temperature, top_k, top_p, do_sample)
|
||||
|
||||
if step + 1 < rvq:
|
||||
embeds[:, step + 1, :] = codec_embeds[step](codes[:, step].view(B, 1)).view(B, hidden_size)
|
||||
|
||||
return codes
|
||||
|
||||
|
||||
def _sample_token(
|
||||
logits: torch.Tensor,
|
||||
temperature: float,
|
||||
top_k: int,
|
||||
top_p: float,
|
||||
do_sample: bool,
|
||||
) -> torch.Tensor:
|
||||
"""Top-k + top-p sampling (matches upstream's ``sample_token`` for the
|
||||
inference branch).
|
||||
"""
|
||||
if not do_sample or temperature <= 0:
|
||||
return logits.argmax(dim=-1)
|
||||
|
||||
logits = logits / max(temperature, 1e-6)
|
||||
if top_k and top_k > 0 and top_k < logits.shape[-1]:
|
||||
top_vals, _ = torch.topk(logits, top_k, dim=-1)
|
||||
thresh = top_vals[..., -1:].expand_as(logits)
|
||||
logits = torch.where(logits < thresh, torch.full_like(logits, float("-inf")), logits)
|
||||
|
||||
if 0.0 < top_p < 1.0:
|
||||
sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)
|
||||
probs = F.softmax(sorted_logits, dim=-1)
|
||||
cum = probs.cumsum(dim=-1)
|
||||
# Drop tail beyond top_p (keep at least one token).
|
||||
drop = cum > top_p
|
||||
drop[..., 1:] = drop[..., :-1].clone()
|
||||
drop[..., 0] = False
|
||||
sorted_logits = sorted_logits.masked_fill(drop, float("-inf"))
|
||||
logits = torch.full_like(logits, float("-inf")).scatter_(-1, sorted_idx, sorted_logits)
|
||||
|
||||
probs = F.softmax(logits, dim=-1)
|
||||
flat = probs.reshape(-1, probs.shape[-1])
|
||||
sampled = torch.multinomial(flat, num_samples=1).reshape(probs.shape[:-1])
|
||||
return sampled
|
||||
|
||||
|
||||
__all__ = ["MossTTSRealtimeLocalTransformer"]
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,88 @@
|
||||
# Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License").
|
||||
"""Pipeline topology for all MOSS-TTS variants (2-stage: talker → codec)."""
|
||||
|
||||
from vllm_omni.config.stage_config import (
|
||||
PipelineConfig,
|
||||
StageExecutionType,
|
||||
StagePipelineConfig,
|
||||
)
|
||||
|
||||
_PROC = "vllm_omni.model_executor.stage_input_processors.moss_tts"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Shared 2-stage pipeline (used by all 5 MOSS-TTS variants)
|
||||
#
|
||||
# Stage 0 (LLM_AR) — Qwen3 backbone + (n_vq+1) parallel heads
|
||||
# emits interleaved text + audio VQ codes
|
||||
# Stage 1 (LLM_GENERATION) — MOSS Audio Tokenizer decode
|
||||
# emits 24 kHz mono waveform chunks
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
MOSS_TTS_PIPELINE = PipelineConfig(
|
||||
model_type="moss_tts",
|
||||
model_arch="MossTTSDelayModel", # HF architectures string
|
||||
stages=(
|
||||
StagePipelineConfig(
|
||||
stage_id=0,
|
||||
model_stage="moss_tts",
|
||||
execution_type=StageExecutionType.LLM_AR,
|
||||
input_sources=(),
|
||||
owns_tokenizer=True,
|
||||
engine_output_type="latent",
|
||||
async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_async_chunk"),
|
||||
sampling_constraints={
|
||||
"detokenize": False,
|
||||
},
|
||||
),
|
||||
StagePipelineConfig(
|
||||
stage_id=1,
|
||||
model_stage="moss_tts_codec",
|
||||
execution_type=StageExecutionType.LLM_GENERATION,
|
||||
input_sources=(0,),
|
||||
final_output=True,
|
||||
final_output_type="audio",
|
||||
engine_output_type="audio",
|
||||
model_arch="MossTTSCodecDecoder",
|
||||
sync_process_input_func=f"{_PROC}.talker2codec",
|
||||
sampling_constraints={"detokenize": True},
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
MOSS_TTS_REALTIME_PIPELINE = PipelineConfig(
|
||||
model_type="moss_tts_realtime",
|
||||
model_arch="MossTTSRealtime", # different talker class from the delay variant
|
||||
stages=(
|
||||
StagePipelineConfig(
|
||||
stage_id=0,
|
||||
model_stage="moss_tts",
|
||||
execution_type=StageExecutionType.LLM_AR,
|
||||
input_sources=(),
|
||||
owns_tokenizer=True,
|
||||
engine_output_type="latent",
|
||||
async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_async_chunk"),
|
||||
sampling_constraints={"detokenize": False},
|
||||
),
|
||||
StagePipelineConfig(
|
||||
stage_id=1,
|
||||
model_stage="moss_tts_codec",
|
||||
execution_type=StageExecutionType.LLM_GENERATION,
|
||||
input_sources=(0,),
|
||||
final_output=True,
|
||||
final_output_type="audio",
|
||||
engine_output_type="audio",
|
||||
model_arch="MossTTSCodecDecoder",
|
||||
sync_process_input_func=f"{_PROC}.talker2codec",
|
||||
sampling_constraints={"detokenize": True},
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
# The pipeline config is otherwise the same for all 5 variants; the per-variant
|
||||
# differences (n_vq, backbone size, generation strategy) are encoded in the
|
||||
# HF config.json and the deploy YAML. Realtime is split out because it has a
|
||||
# different talker architecture (MossTTSRealtime vs MossTTSDelayModel).
|
||||
|
||||
__all__ = ["MOSS_TTS_PIPELINE", "MOSS_TTS_REALTIME_PIPELINE"]
|
||||
@@ -0,0 +1,100 @@
|
||||
"""Reference-audio encoding + speaker cache for the MOSS-TTS-family talker.
|
||||
|
||||
This lives in the model package (not the shared serving layer) so all
|
||||
MOSS-specific reference handling stays with the model — mirroring how Fish
|
||||
Speech (``dac_encoder.encode_reference_audio_codes``), CosyVoice3, and
|
||||
Qwen3-TTS keep their reference/speaker extraction next to the model rather than
|
||||
in ``serving_speech.py``. The serving layer just calls
|
||||
:func:`encode_reference_codes` with its generic helpers (the audio resolver and
|
||||
the process-wide speaker cache).
|
||||
|
||||
Kept import-light (only ``asyncio`` / ``hashlib`` / ``torch`` plus the logger)
|
||||
so importing it from the API-server process does not pull the talker/codec.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from vllm.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _encode_wav_sync(processor: Any, wav_list: list, sr: int, sr_target: int, n_vq: int) -> torch.Tensor:
|
||||
"""Blocking resample + CPU codec encode (the expensive bit)."""
|
||||
wav = torch.tensor(wav_list, dtype=torch.float32)
|
||||
if wav.dim() == 1:
|
||||
wav = wav.unsqueeze(0)
|
||||
if sr != sr_target:
|
||||
import torchaudio
|
||||
|
||||
wav = torchaudio.functional.resample(wav, sr, sr_target)
|
||||
with torch.no_grad():
|
||||
codes_list = processor.encode_audios_from_wav([wav], sampling_rate=sr_target, n_vq=n_vq)
|
||||
return codes_list[0]
|
||||
|
||||
|
||||
async def encode_reference_codes(
|
||||
ref_str: str,
|
||||
*,
|
||||
processor: Any,
|
||||
resolve_ref_audio: Callable[[str], Awaitable[tuple[list, int]]],
|
||||
speaker_cache: Any,
|
||||
variant: str,
|
||||
n_vq: int,
|
||||
sr_target: int,
|
||||
voice_name: str | None = None,
|
||||
voice_created_at: int = 0,
|
||||
) -> torch.Tensor:
|
||||
"""Encode one reference clip into MOSS RVQ codes, reusing the speaker cache.
|
||||
|
||||
The MOSS audio tokenizer sits on CPU (to spare ~6.7 GiB next to the 8B
|
||||
talker), so re-encoding the same reference is a fixed per-request cost that
|
||||
otherwise dominates the 8B voice-clone variants and serializes under
|
||||
concurrency. Mirror CosyVoice3 / Qwen3-TTS: cache by named voice when one is
|
||||
supplied (``voice_created_at`` invalidates on re-upload), else by a content
|
||||
hash of the reference. The blocking encode runs in a worker thread via
|
||||
``asyncio.to_thread`` so cold/anonymous encodes from concurrent requests
|
||||
overlap instead of serializing on the event loop.
|
||||
|
||||
Args:
|
||||
ref_str: The raw reference audio (URL / path / data URL) as received.
|
||||
processor: Upstream MOSS processor exposing ``encode_audios_from_wav``.
|
||||
resolve_ref_audio: Async callable mapping ``ref_str`` to ``(wav_list, sr)``.
|
||||
speaker_cache: Process-wide ``SpeakerEmbeddingCache``.
|
||||
variant: MOSS sub-variant (``tts`` / ``ttsd`` / ...), namespaces the cache key.
|
||||
n_vq: Number of RVQ codebooks (also namespaces the cache key).
|
||||
sr_target: Target sample rate for the codec.
|
||||
voice_name: Named/uploaded voice, if any (enables stable cross-request caching).
|
||||
voice_created_at: Upload timestamp; bumps the cache slot on re-upload.
|
||||
|
||||
Returns:
|
||||
The reference RVQ codes tensor (CPU), ready to pass to the processor.
|
||||
"""
|
||||
if voice_name:
|
||||
speaker_name = voice_name
|
||||
created_at = int(voice_created_at)
|
||||
else:
|
||||
speaker_name = "ref:" + hashlib.sha1((ref_str or "").encode("utf-8")).hexdigest()
|
||||
created_at = 0
|
||||
|
||||
cache_key = speaker_cache.make_cache_key(
|
||||
speaker_name,
|
||||
model_type=f"moss_tts_{variant}_nq{n_vq}",
|
||||
created_at=created_at,
|
||||
)
|
||||
cached = speaker_cache.get(cache_key)
|
||||
if cached is not None:
|
||||
logger.debug("Speaker cache HIT for MOSS-TTS reference '%s'", speaker_name)
|
||||
return cached["codes"].clone()
|
||||
|
||||
wav_list, sr = await resolve_ref_audio(ref_str)
|
||||
codes = await asyncio.to_thread(_encode_wav_sync, processor, wav_list, sr, sr_target, n_vq)
|
||||
speaker_cache.put(cache_key, {"codes": codes.detach().cpu()})
|
||||
logger.debug("Speaker cache STORE for MOSS-TTS reference '%s'", speaker_name)
|
||||
return codes
|
||||
@@ -223,6 +223,25 @@ _OMNI_MODELS = {
|
||||
"modeling_moss_tts_nano",
|
||||
"MossTTSNanoForGeneration",
|
||||
),
|
||||
## MOSS-TTS (full variants: Delay + Realtime)
|
||||
# MossTTSDelayModel: MOSS-TTS (8B), MOSS-TTSD (8B), MOSS-SoundEffect (8B), MOSS-VoiceGenerator (1.7B)
|
||||
"MossTTSDelayModel": (
|
||||
"moss_tts",
|
||||
"modeling_moss_tts_talker",
|
||||
"MossTTSDelayTalkerForGeneration",
|
||||
),
|
||||
# MossTTSRealtime: MOSS-TTS-Realtime (1.7B)
|
||||
"MossTTSRealtime": (
|
||||
"moss_tts",
|
||||
"modeling_moss_tts_talker",
|
||||
"MossTTSRealtimeTalkerForGeneration",
|
||||
),
|
||||
# Stage-1 codec decoder (shared by all 5 variants)
|
||||
"MossTTSCodecDecoder": (
|
||||
"moss_tts",
|
||||
"modeling_moss_tts_codec",
|
||||
"MossTTSCodecDecoder",
|
||||
),
|
||||
"DyninOmniForConditionalGeneration": (
|
||||
"dynin_omni",
|
||||
"dynin_omni",
|
||||
|
||||
@@ -0,0 +1,231 @@
|
||||
# Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License").
|
||||
"""Stage input processors: MOSS-TTS talker (Stage 0) → codec (Stage 1)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from vllm.inputs import TokensPrompt as OmniTokensPrompt
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _extract_audio_codes(stage_output: Any) -> torch.Tensor | None:
|
||||
"""Pull audio codes from a Stage-0 OmniOutput or raw tensor."""
|
||||
if stage_output is None:
|
||||
return None
|
||||
|
||||
# OmniOutput
|
||||
mm = getattr(stage_output, "multimodal_outputs", None)
|
||||
if mm is not None:
|
||||
codes_dict = mm.get("codes", {})
|
||||
if isinstance(codes_dict, dict):
|
||||
ac = codes_dict.get("audio")
|
||||
if isinstance(ac, torch.Tensor):
|
||||
return ac
|
||||
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Non-streaming (sync): called once after Stage 0 finishes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def talker2codec(
|
||||
stage_list: list[Any],
|
||||
engine_input_source: list[int],
|
||||
prompt: Any = None,
|
||||
requires_multimodal_data: bool = False,
|
||||
) -> list[Any]:
|
||||
"""Convert all talker codes to a single Stage-1 token sequence.
|
||||
|
||||
Stage 0 output contains ``codes["audio"]`` shaped ``(T, NQ)`` where T is
|
||||
the number of generated audio frames and NQ is n_vq. We flatten to
|
||||
``[NQ * T]`` as the Stage-1 ``input_ids`` so the codec can reshape back
|
||||
to ``(NQ, T)`` for decoding.
|
||||
"""
|
||||
results: list[Any] = []
|
||||
|
||||
for src_idx in engine_input_source:
|
||||
if src_idx >= len(stage_list):
|
||||
results.append(OmniTokensPrompt(prompt_token_ids=[]))
|
||||
continue
|
||||
|
||||
stage_out = stage_list[src_idx]
|
||||
audio_codes = _extract_audio_codes(stage_out)
|
||||
|
||||
if audio_codes is None or audio_codes.numel() == 0:
|
||||
logger.warning("talker2codec: no audio codes in stage output %d; emitting silence.", src_idx)
|
||||
results.append(OmniTokensPrompt(prompt_token_ids=[]))
|
||||
continue
|
||||
|
||||
# audio_codes: (T, NQ) → flatten to [NQ, T] → list[int]
|
||||
codes_nq_t = audio_codes.transpose(0, 1).contiguous() # (NQ, T)
|
||||
flat = codes_nq_t.reshape(-1).tolist()
|
||||
|
||||
results.append(
|
||||
OmniTokensPrompt(
|
||||
prompt_token_ids=flat,
|
||||
multi_modal_data={"codes": {"audio": codes_nq_t}},
|
||||
)
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Streaming (async chunk): called each time Stage 0 emits a new chunk
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def talker2codec_async_chunk(
|
||||
transfer_manager: Any,
|
||||
pooling_output: dict[str, Any] | None,
|
||||
request: Any,
|
||||
is_finished: bool = False,
|
||||
) -> OmniPayloadStruct | None:
|
||||
"""Emit accumulated audio codes to Stage 1 as they arrive from Stage 0.
|
||||
|
||||
State is maintained in ``transfer_manager`` keyed by request ID.
|
||||
A chunk is forwarded to Stage 1 when either:
|
||||
(a) ``is_finished`` is True (flush all remaining codes), or
|
||||
(b) the accumulated frame count reaches ``chunk_frames`` (default 25).
|
||||
|
||||
Returns a dict compatible with the Stage-1 input format, or None to
|
||||
signal "not enough data yet — wait for more frames".
|
||||
"""
|
||||
req_id: str = str(getattr(request, "request_id", id(request)))
|
||||
|
||||
# Initialise per-request accumulation state
|
||||
if not hasattr(transfer_manager, "_moss_tts_state"):
|
||||
transfer_manager._moss_tts_state = {}
|
||||
state = transfer_manager._moss_tts_state
|
||||
|
||||
if req_id not in state:
|
||||
state[req_id] = {
|
||||
"accumulated": None, # (T_acc, NQ) tensor or None
|
||||
"total_emitted": 0,
|
||||
}
|
||||
req_state = state[req_id]
|
||||
|
||||
# Extract new codes from this chunk. The talker emits the full per-request
|
||||
# ``audio_codes["accumulated"]`` snapshot every step, so we only append the
|
||||
# *new* tail rows (otherwise we'd duplicate history quadratically).
|
||||
# ``pooling_output`` carries the unflattened OmniPayload at the top level
|
||||
# (``codes.audio``), matching the qwen3_tts pattern.
|
||||
if pooling_output is not None:
|
||||
codes_dict = pooling_output.get("codes", {}) or {}
|
||||
snapshot = codes_dict.get("audio")
|
||||
if isinstance(snapshot, torch.Tensor) and snapshot.numel() > 0:
|
||||
snapshot_cpu = snapshot.cpu()
|
||||
prev_t = 0 if req_state["accumulated"] is None else int(req_state["accumulated"].shape[0])
|
||||
new_rows = snapshot_cpu[prev_t:]
|
||||
if new_rows.numel() > 0:
|
||||
if req_state["accumulated"] is None:
|
||||
req_state["accumulated"] = new_rows
|
||||
else:
|
||||
req_state["accumulated"] = torch.cat([req_state["accumulated"], new_rows], dim=0)
|
||||
# Realtime variant emits raw (un-delay-patterned) codes; record that
|
||||
# so the emit step below skips the apply_de_delay_pattern transform.
|
||||
# Realtime sets ``meta.finished`` to a 1-D bool tensor; the delay
|
||||
# variant leaves ``meta`` empty.
|
||||
meta_in = pooling_output.get("meta", {}) or {}
|
||||
flag = meta_in.get("finished")
|
||||
if isinstance(flag, torch.Tensor) and flag.numel() >= 1 and bool(flag.reshape(-1)[0].item()):
|
||||
req_state["skip_dedelay"] = True
|
||||
|
||||
acc = req_state["accumulated"]
|
||||
if acc is None or acc.numel() == 0:
|
||||
if is_finished:
|
||||
del state[req_id]
|
||||
return None
|
||||
|
||||
# The MOSS audio tokenizer's causal decoder doesn't yet have left-context
|
||||
# plumbing in this port, and a streaming chunk of 25 frames trips an
|
||||
# internal patched-pretransform reshape on the first chunk. Until we wire
|
||||
# left-context properly, accumulate all codes and emit only on finish.
|
||||
chunk_frames: int = 1 << 30
|
||||
left_context: int = 0
|
||||
|
||||
t_acc = int(acc.shape[0])
|
||||
should_emit = is_finished or (t_acc - req_state["total_emitted"] >= chunk_frames)
|
||||
|
||||
if not should_emit:
|
||||
return None
|
||||
|
||||
# Determine the slice to emit
|
||||
emit_start = max(0, req_state["total_emitted"] - left_context)
|
||||
chunk_codes = acc[emit_start:] # (T_chunk, NQ)
|
||||
req_state["total_emitted"] = t_acc
|
||||
|
||||
if is_finished:
|
||||
del state[req_id]
|
||||
|
||||
# Mirror upstream ``MossTTSDelayProcessor._parse_audio_codes``: the delay
|
||||
# talker samples codes in a delay pattern (a row is emitted every step, but
|
||||
# only the slot ``i == arange < audio_lengths`` carries a real code; the
|
||||
# rest is ``audio_pad_code``). Before sending to the codec we must
|
||||
# 1. de-delay ``(T+nq-1, nq)`` → ``(T, nq)``
|
||||
# 2. drop rows that are entirely pad (separators between audio segments,
|
||||
# and the leading text-mode rows that precede the first audio_start).
|
||||
# The realtime talker emits raw codes (no delay), so steps 1 is skipped.
|
||||
chunk_codes_long = chunk_codes.to(torch.long).cpu().contiguous() # (T_chunk, NQ)
|
||||
nq = int(chunk_codes_long.shape[1])
|
||||
t_chunk = int(chunk_codes_long.shape[0])
|
||||
audio_pad_code = 1024 # MOSS-TTS audio_pad_code; same value across variants.
|
||||
|
||||
if req_state.get("skip_dedelay"):
|
||||
de_delayed = chunk_codes_long
|
||||
elif t_chunk > nq:
|
||||
de_delayed = chunk_codes_long.new_zeros((t_chunk - nq + 1, nq))
|
||||
for i in range(nq):
|
||||
de_delayed[:, i] = chunk_codes_long[i : i + de_delayed.shape[0], i]
|
||||
else:
|
||||
de_delayed = chunk_codes_long.new_zeros((0, nq))
|
||||
|
||||
if de_delayed.shape[0] > 0:
|
||||
is_pad = (de_delayed == audio_pad_code).all(dim=1)
|
||||
non_pad = ~is_pad
|
||||
if bool(non_pad.any()):
|
||||
de_delayed = de_delayed[non_pad]
|
||||
else:
|
||||
de_delayed = de_delayed.new_zeros((0, nq))
|
||||
|
||||
if de_delayed.shape[0] == 0:
|
||||
# Nothing left after filtering — emit silence sentinel so the codec
|
||||
# request still completes cleanly.
|
||||
codec_flat: list[int] = []
|
||||
else:
|
||||
# Stage 1 (LLM_GENERATION codec) consumes ``codes.audio`` as a flat
|
||||
# codebook-major int list — chunk_transfer_adapter assigns it to
|
||||
# ``request.prompt_token_ids`` and the codec rebuilds the (NQ, T) grid.
|
||||
# Keep it a list (not a tensor): the receive path treats a list as the
|
||||
# token-id sequence directly, while a 1-D tensor would be ``.tolist()``-ed
|
||||
# anyway — the list keeps the downstream ``if not new_ids`` checks intact.
|
||||
codec_flat = de_delayed.transpose(0, 1).contiguous().reshape(-1).tolist()
|
||||
|
||||
# main migrated the inter-stage connector from plain dicts to typed
|
||||
# ``OmniPayloadStruct`` (the chunk_transfer_adapter sender reads
|
||||
# ``payload_data.meta``); return the struct so MOSS matches that schema.
|
||||
return OmniPayloadStruct(
|
||||
codes=CodesStruct(audio=codec_flat),
|
||||
meta=MetaStruct(
|
||||
left_context_size=left_context,
|
||||
finished=torch.tensor(is_finished, dtype=torch.bool),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["talker2codec", "talker2codec_async_chunk"]
|
||||
Reference in new issue
Block a user