diff --git a/benchmarks/build_dataset/sound_effect_smoke/en/meta.lst b/benchmarks/build_dataset/sound_effect_smoke/en/meta.lst new file mode 100644 index 000000000..d75bbeb36 --- /dev/null +++ b/benchmarks/build_dataset/sound_effect_smoke/en/meta.lst @@ -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 diff --git a/benchmarks/build_dataset/ttsd_smoke/en/meta.lst b/benchmarks/build_dataset/ttsd_smoke/en/meta.lst new file mode 100644 index 000000000..23a61ac41 --- /dev/null +++ b/benchmarks/build_dataset/ttsd_smoke/en/meta.lst @@ -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. diff --git a/benchmarks/tts/bench_tts.py b/benchmarks/tts/bench_tts.py index 78608e844..c79d74a31 100644 --- a/benchmarks/tts/bench_tts.py +++ b/benchmarks/tts/bench_tts.py @@ -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: diff --git a/benchmarks/tts/model_configs.yaml b/benchmarks/tts/model_configs.yaml index 804e60be8..0a26bc0cb 100644 --- a/benchmarks/tts/model_configs.yaml +++ b/benchmarks/tts/model_configs.yaml @@ -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: {} diff --git a/examples/offline_inference/text_to_speech/moss_tts/end2end.py b/examples/offline_inference/text_to_speech/moss_tts/end2end.py new file mode 100644 index 000000000..fb47ad8c9 --- /dev/null +++ b/examples/offline_inference/text_to_speech/moss_tts/end2end.py @@ -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 + # (col 0 = text, cols 1..16 = audio_channel_pad) + # (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() diff --git a/recipes/OpenMOSS/MOSS-TTS.md b/recipes/OpenMOSS/MOSS-TTS.md new file mode 100644 index 000000000..c88bf7281 --- /dev/null +++ b/recipes/OpenMOSS/MOSS-TTS.md @@ -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: + +## 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. diff --git a/tests/e2e/offline_inference/test_moss_tts.py b/tests/e2e/offline_inference/test_moss_tts.py new file mode 100644 index 000000000..e24872b8d --- /dev/null +++ b/tests/e2e/offline_inference/test_moss_tts.py @@ -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) diff --git a/tests/e2e/offline_inference/test_moss_tts_v1_5.py b/tests/e2e/offline_inference/test_moss_tts_v1_5.py new file mode 100644 index 000000000..e3f4d669d --- /dev/null +++ b/tests/e2e/offline_inference/test_moss_tts_v1_5.py @@ -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" diff --git a/tests/e2e/online_serving/test_moss_tts_expansion.py b/tests/e2e/online_serving/test_moss_tts_expansion.py new file mode 100644 index 000000000..3d7559717 --- /dev/null +++ b/tests/e2e/online_serving/test_moss_tts_expansion.py @@ -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) diff --git a/vllm_omni/benchmarks/data_modules/sound_effect_dataset.py b/vllm_omni/benchmarks/data_modules/sound_effect_dataset.py new file mode 100644 index 000000000..645756d68 --- /dev/null +++ b/vllm_omni/benchmarks/data_modules/sound_effect_dataset.py @@ -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 diff --git a/vllm_omni/benchmarks/data_modules/ttsd_dataset.py b/vllm_omni/benchmarks/data_modules/ttsd_dataset.py new file mode 100644 index 000000000..799a08605 --- /dev/null +++ b/vllm_omni/benchmarks/data_modules/ttsd_dataset.py @@ -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}//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 diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index 12aeef7e4..7238cdf88 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -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( diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index d5ea7b38f..761373478 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -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", diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 93453006f..5254a5396 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -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( diff --git a/vllm_omni/deploy/moss_sound_effect.yaml b/vllm_omni/deploy/moss_sound_effect.yaml new file mode 100644 index 000000000..4ef5fa7af --- /dev/null +++ b/vllm_omni/deploy/moss_sound_effect.yaml @@ -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 diff --git a/vllm_omni/deploy/moss_tts.yaml b/vllm_omni/deploy/moss_tts.yaml new file mode 100644 index 000000000..acdb87721 --- /dev/null +++ b/vllm_omni/deploy/moss_tts.yaml @@ -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 diff --git a/vllm_omni/deploy/moss_tts_realtime.yaml b/vllm_omni/deploy/moss_tts_realtime.yaml new file mode 100644 index 000000000..124b3b8be --- /dev/null +++ b/vllm_omni/deploy/moss_tts_realtime.yaml @@ -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 diff --git a/vllm_omni/deploy/moss_ttsd.yaml b/vllm_omni/deploy/moss_ttsd.yaml new file mode 100644 index 000000000..3907941db --- /dev/null +++ b/vllm_omni/deploy/moss_ttsd.yaml @@ -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 diff --git a/vllm_omni/deploy/moss_voice_generator.yaml b/vllm_omni/deploy/moss_voice_generator.yaml new file mode 100644 index 000000000..8045dcbcf --- /dev/null +++ b/vllm_omni/deploy/moss_voice_generator.yaml @@ -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 diff --git a/vllm_omni/entrypoints/cli/benchmark/serve.py b/vllm_omni/entrypoints/cli/benchmark/serve.py index cee900f61..b969e5146 100644 --- a/vllm_omni/entrypoints/cli/benchmark/serve.py +++ b/vllm_omni/entrypoints/cli/benchmark/serve.py @@ -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 diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index 0325017b4..9a858868e 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -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)", diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 0ca4cd319..57bed35eb 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -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,49 +1517,222 @@ 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(): - return "Input text cannot be empty" - 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." - ) - fmt_err = self._validate_ref_audio_format(request.ref_audio) - if fmt_err: - return fmt_err - return None + # 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: + 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 + + 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 + + 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 additional_information for MOSS-TTS-Nano. + """Build the talker prompt + ``additional_information`` payload for any + MOSS-TTS-family request (nano + 5 full variants). - 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``. + 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 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 + + # ---- 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] = { - "text": [request.input], - "mode": ["voice_clone"], + "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] - wav_list, sr = await self._resolve_ref_audio(request.ref_audio) - params["prompt_audio_array"] = [[wav_list, sr]] return params def _validate_higgs_audio_v2_request(self, request: OpenAICreateSpeechRequest) -> str | 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,7 +2853,14 @@ 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] - prompt = tokens_input(prompt_token_ids=[1]) + # 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) else: @@ -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": diff --git a/vllm_omni/model_executor/models/moss_tts/__init__.py b/vllm_omni/model_executor/models/moss_tts/__init__.py new file mode 100644 index 000000000..23a73b3fe --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/__init__.py @@ -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", +] diff --git a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py new file mode 100644 index 000000000..a2cf75b07 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py @@ -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", +] diff --git a/vllm_omni/model_executor/models/moss_tts/configuration_moss_tts.py b/vllm_omni/model_executor/models/moss_tts/configuration_moss_tts.py new file mode 100644 index 000000000..46147f27c --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/configuration_moss_tts.py @@ -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"] diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py new file mode 100644 index 000000000..fa731fa6f --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py @@ -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"] diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py new file mode 100644 index 000000000..bf01a83bf --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py @@ -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"] diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py new file mode 100644 index 000000000..e5fcdb234 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py @@ -0,0 +1,1103 @@ +# Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). +"""MOSS-TTS Stage-0 talker: Qwen3 backbone + (n_vq+1) parallel AR heads.""" + +from __future__ import annotations + +import copy +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.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.qwen3 import Qwen3Model +from vllm.sequence import IntermediateTensors +from vllm.v1.sample.metadata import SamplingMetadata + +from vllm_omni.model_executor.models.moss_tts.configuration_moss_tts import ( + MossTTSDelayConfig, + MossTTSRealtimeConfig, +) +from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local import ( + MossTTSRealtimeLocalTransformer, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + +logger = init_logger(__name__) + +_AUDIO_PAD_ID = 0 # padding id used before delay slot fires + + +def _maybe_prefix(prefix: str, name: str) -> str: + return f"{prefix}.{name}" if prefix else name + + +# --------------------------------------------------------------------------- +# MossTTSDelayTalkerForGeneration +# --------------------------------------------------------------------------- + + +class MossTTSDelayTalkerForGeneration(nn.Module): + """Stage-0 talker for MossTTSDelayModel variants. + + Covers all four repos that ship ``architectures: ["MossTTSDelayModel"]``: + - MOSS-TTS (8B, n_vq=32) + - MOSS-TTSD-v1.0 (8B, n_vq=16) + - MOSS-SoundEffect (8B, n_vq=16) + - MOSS-VoiceGenerator (1.7B, n_vq=16) + + Architecture + ~~~~~~~~~~~~ + * Backbone: Qwen3 transformer (hidden_size, num_hidden_layers, etc. from + ``config.language_config``). + * Embedding: text_embed(t) + Σᵢ audio_embed_i(aᵢ) — additive fusion, + no cross-attention or concatenation. + * Heads: (n_vq + 1) parallel linear heads over the final hidden state. + – Head 0 → text logits (drives AR scheduler) + – Heads 1…n_vq → audio VQ logits (one per RVQ codebook) + + Delay pattern + ~~~~~~~~~~~~~ + Audio heads are only active after the model emits the delay-slot token + (``audio_assistant_delay_slot_token_id``). Before the slot fires all + audio heads output a pad token (``audio_pad_code``). After the slot: + + step t: collect audio_codebook_0 for frame t + step t+1: collect audio_codebook_0 for frame t+1 + collect audio_codebook_1 for frame t (1-step lag) + … + step t+k: all k codebooks active; emit frame t + + The per-request ``delay_step`` counter (stored in the per-request info + dict) tracks this. Stage-1 receives the codes in (T, NQ) shape and the + codec's ``batch_decode`` handles the de-interleaving internally. + """ + + # vLLM-Omni integration flags + have_multimodal_outputs: bool = True + has_preprocess: bool = True + has_postprocess: bool = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.vllm_config = vllm_config + self.config: MossTTSDelayConfig = vllm_config.model_config.hf_config + + lm_cfg = self.config.language_config + hidden_size: int = lm_cfg.hidden_size + self.hidden_size = hidden_size + self.n_vq: int = self.config.n_vq + self.audio_vocab_size: int = self.config.audio_vocab_size + self.audio_pad_code: int = self.config.audio_pad_code + + # Token IDs (mirrors upstream MossTTSDelayConfig defaults). + self.audio_start_token_id: int = self.config.audio_start_token_id + self.audio_end_token_id: int = self.config.audio_end_token_id + self.audio_assistant_gen_slot_token_id: int = self.config.audio_assistant_gen_slot_token_id + self.audio_assistant_delay_slot_token_id: int = self.config.audio_assistant_delay_slot_token_id + self.pad_token_id: int = getattr(self.config, "pad_token_id", 151643) + self.im_end_token_id: int = getattr(self.config, "im_end_token_id", 151645) + + # Qwen3 backbone — weights live under ``language_model.*`` + self.model = Qwen3Model( + vllm_config=vllm_config, + prefix=_maybe_prefix(prefix, "model"), + ) + + # Text LM head (head 0) + self.text_lm_head = ParallelLMHead( + lm_cfg.vocab_size, + hidden_size, + bias=False, + prefix=_maybe_prefix(prefix, "text_lm_head"), + ) + self.logits_processor = LogitsProcessor(lm_cfg.vocab_size) + + # Audio VQ heads (heads 1…n_vq in upstream naming) + # Each head predicts one RVQ codebook: vocab = audio_vocab_size + 1 + self.audio_heads = nn.ModuleList( + [nn.Linear(hidden_size, self.audio_vocab_size + 1, bias=False) for _ in range(self.n_vq)] + ) + + # Per-codebook audio embeddings (emb_ext in upstream) + self.audio_embeddings = nn.ModuleList( + [nn.Embedding(self.audio_vocab_size + 1, hidden_size) for _ in range(self.n_vq)] + ) + + # GPU-resident per-request buffer keys (avoid CPU round-trips) + self.gpu_resident_buffer_keys: set[tuple[str, str]] = { + ("audio_codes", "current"), # last step's audio codes + ("audio_codes", "accumulated"), + ("hidden_states", "last"), + } + + # Static text-logit masks (lazy, built on first use to know device). + self._pre_exclude_text_ids = ( + self.pad_token_id, + self.audio_assistant_gen_slot_token_id, + self.audio_assistant_delay_slot_token_id, + self.audio_end_token_id, + ) + self._audio_keep_text_ids = ( + self.audio_assistant_gen_slot_token_id, + self.audio_assistant_delay_slot_token_id, + ) + + # Per-step state stash for `compute_logits` (populated by + # `make_omni_output`, which runs immediately before the sampler). + self._batch_state: list[dict[str, Any]] | None = None + + # ------------------------------------------------------------------ + # vLLM-Omni hooks + # ------------------------------------------------------------------ + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + return self.model.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **_: Any, + ) -> torch.Tensor | IntermediateTensors: + return self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor | OmniOutput, + sampling_metadata: SamplingMetadata | None = None, + ) -> torch.Tensor | None: + """Return text-head logits with delay-pattern constraints applied. + + The mask follows upstream MOSS-TTS' generate loop: + * Forced tokens override the sampler when delayed_lengths is in the + audio-emit window (delay_slot for [0, n_vq), audio_end at n_vq). + * Outside that window, mask audio control tokens unless the model + is currently emitting audio (is_audio). + * Mask delay_slot at step 0 and im_end during the first n_vq steps, + matching upstream's anti-collapse heuristics. + """ + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if hidden_states is None: + return None + logits = self.logits_processor(self.text_lm_head, hidden_states, sampling_metadata) + if logits is None or self._batch_state is None: + return logits + + states = self._batch_state + if not states: + return logits + + n_vq = self.n_vq + device = logits.device + vocab_size = logits.shape[-1] + rows_per_state = max(1, logits.shape[0] // max(1, len(states))) + + for i, state in enumerate(states): + if not isinstance(state, dict): + continue + row_start = i * rows_per_state + row_end = min(row_start + rows_per_state, logits.shape[0]) + if row_start >= row_end: + continue + row = logits[row_start:row_end] + + delayed_lengths = int(state.get("delayed_lengths", -1)) + is_audio = bool(state.get("is_audio", False)) + step = int(state.get("step", 0)) + audio_lengths_cur = int(state.get("audio_lengths", 0)) + max_new_frames = int(state.get("max_new_frames", -1)) + + # ---- max_new_frames cap (force im_end / EOS) ---- + # When the request specified ``max_new_frames`` and we've emitted + # that many audio-bearing tokens, force ``im_end`` so the AR loop + # stops. Without this, short prompts run to deploy-default + # ``max_tokens`` and produce 100+ s of audio per ~10-word input. + if ( + max_new_frames > 0 + and is_audio + and audio_lengths_cur >= max_new_frames + and 0 <= self.im_end_token_id < vocab_size + ): + neg_inf = torch.full_like(row, float("-inf")) + neg_inf[..., self.im_end_token_id] = 0.0 + logits[row_start:row_end] = neg_inf + continue + + # ---- Forced tokens (delay-slot run / audio-end) ---- + forced: int | None = None + if 0 <= delayed_lengths < n_vq: + forced = self.audio_assistant_delay_slot_token_id + elif delayed_lengths == n_vq: + forced = self.audio_end_token_id + if forced is not None and 0 <= forced < vocab_size: + neg_inf = torch.full_like(row, float("-inf")) + neg_inf[..., forced] = 0.0 + row = neg_inf + logits[row_start:row_end] = row + continue + + # ---- Pre-exclusion masks (delayed_lengths == -1 sentinel) ---- + if is_audio: + # Only delay_slot or gen_slot are valid in audio mode. + mask = torch.ones(vocab_size, dtype=torch.bool, device=device) + for tok in self._audio_keep_text_ids: + if 0 <= tok < vocab_size: + mask[tok] = False + row = row.masked_fill(mask, float("-inf")) + else: + # Mask audio control tokens during text generation. + for tok in self._pre_exclude_text_ids: + if 0 <= tok < vocab_size: + row[..., tok] = float("-inf") + + # ---- Step-conditional masks ---- + if step == 0: + tok = self.audio_assistant_delay_slot_token_id + if 0 <= tok < vocab_size: + row[..., tok] = float("-inf") + if step <= n_vq: + tok = self.im_end_token_id + if 0 <= tok < vocab_size: + row[..., tok] = float("-inf") + + logits[row_start:row_end] = row + + return logits + + def _initial_state(self, prompt_ids: torch.Tensor) -> dict[str, Any]: + """Initialise the AR generation state machine from a prompt. + + Mirrors upstream's ``generate`` head: + * is_audio = True iff the last prompt token is audio_start or gen_slot + and the prompt contains a prior audio_start. + * audio_lengths counts tokens since the last audio_start in that case. + * delayed_lengths starts at the int64-max sentinel (-1 here). + """ + if prompt_ids.numel() == 0: + return { + "audio_lengths": 0, + "delayed_lengths": -1, + "is_audio": False, + "step": 0, + } + prompt_cpu = prompt_ids.detach().to("cpu", dtype=torch.long).reshape(-1) + seq_len = int(prompt_cpu.shape[0]) + last_id = int(prompt_cpu[-1]) + is_continuation = last_id in ( + self.audio_start_token_id, + self.audio_assistant_gen_slot_token_id, + ) + last_audio_start = -1 + if is_continuation: + matches = (prompt_cpu == self.audio_start_token_id).nonzero(as_tuple=True)[0] + if matches.numel() > 0: + last_audio_start = int(matches[-1]) + is_audio = is_continuation and last_audio_start != -1 + audio_lengths = (seq_len - last_audio_start) if is_audio else 0 + return { + "audio_lengths": int(audio_lengths), + "delayed_lengths": -1, + "is_audio": bool(is_audio), + "step": 0, + } + + def _advance_state(self, state: dict[str, Any], sampled_id: int) -> dict[str, Any]: + """Update state given the text token sampled at the previous step.""" + audio_lengths = int(state.get("audio_lengths", 0)) + delayed_lengths = int(state.get("delayed_lengths", -1)) + is_audio = bool(state.get("is_audio", False)) + + if sampled_id in ( + self.audio_start_token_id, + self.audio_assistant_gen_slot_token_id, + self.audio_assistant_delay_slot_token_id, + ): + audio_lengths += 1 + if sampled_id == self.audio_end_token_id: + audio_lengths = 0 + + if delayed_lengths == -1: + if sampled_id == self.audio_assistant_delay_slot_token_id: + delayed_lengths = 0 + else: + delayed_lengths += 1 + if delayed_lengths > self.n_vq: + delayed_lengths = -1 + + if sampled_id == self.audio_start_token_id: + is_audio = True + if sampled_id == self.audio_end_token_id: + is_audio = False + + state = dict(state) + state["audio_lengths"] = audio_lengths + state["delayed_lengths"] = delayed_lengths + state["is_audio"] = is_audio + state["step"] = int(state.get("step", 0)) + 1 + return state + + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None, + **info_dict: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Build per-step input embeddings (text + audio additive fusion). + + Prefill: initialise the per-request state machine from the prompt. + Decode: update the state with the just-sampled text token, then build + the combined text+audio embedding using the previous step's codes. + """ + device = input_ids.device + span_len = int(input_ids.shape[0]) + audio_state = info_dict.get("audio_state") + is_first_call = not isinstance(audio_state, dict) + + if span_len > 1 or is_first_call: + # Prefill (or first call). Initialise the per-request state and + # build text embeddings. If the request carries reference-audio + # codes (``codes.ref`` from the upstream MossTTSDelayProcessor's + # delay-pattern grid, shape ``(L_full, n_vq)``), additively embed + # them at the matching prefill positions so the talker can attend + # to the speaker's timbre when generating its response. + audio_state = self._initial_state(input_ids) + embeds = self.model.embed_tokens(input_ids) + + ref_codes = (info_dict.get("codes", {}) or {}).get("ref") + last_ref_row = None + ref_offset = int(info_dict.get("ref_offset", 0)) + if isinstance(ref_codes, torch.Tensor) and ref_codes.numel() > 0: + if ref_codes.dim() == 1: + if ref_codes.numel() % self.n_vq == 0: + ref_codes = ref_codes.view(-1, self.n_vq) + else: + ref_codes = None # malformed, skip silently + if isinstance(ref_codes, torch.Tensor) and ref_codes.dim() == 2: + # Slice the window matching this prefill chunk (chunked + # prefill calls preprocess multiple times for one request). + end_off = ref_offset + span_len + chunk = ref_codes[ref_offset:end_off] + if chunk.numel() > 0 and chunk.shape[0] == span_len: + codes = chunk.to(device=device, dtype=torch.long).clamp_(0, self.audio_vocab_size) + audio_embed = torch.zeros_like(embeds) + for i, emb_layer in enumerate(self.audio_embeddings): + audio_embed = audio_embed + emb_layer(codes[:, i]) + embeds = embeds + audio_embed + last_ref_row = codes[-1].detach() + + current_codes = ( + last_ref_row + if last_ref_row is not None + else torch.full( + (self.n_vq,), + self.audio_pad_code, + dtype=torch.long, + device=device, + ) + ) + # Lift max_new_frames from request additional_information into + # audio_state so compute_logits can force im_end once we've + # emitted that many audio frames (mirrors the moss_tts_nano cap). + max_new_frames_req = info_dict.get("max_new_frames") + if isinstance(max_new_frames_req, (list, tuple)) and max_new_frames_req: + max_new_frames_req = max_new_frames_req[0] + try: + audio_state["max_new_frames"] = int(max_new_frames_req) if max_new_frames_req is not None else -1 + except (TypeError, ValueError): + audio_state["max_new_frames"] = -1 + info_update: dict[str, Any] = { + "audio_state": audio_state, + "audio_codes": {"current": current_codes}, + "ref_offset": ref_offset + span_len, + } + return input_ids, embeds, info_update + + # Decode step: input_ids is the text token sampled at step n-1. + sampled_id = int(input_ids.reshape(-1)[0].item()) + audio_state = self._advance_state(audio_state, sampled_id) + + # Combined embedding: text(t) + Σᵢ audio_emb_i(code_i_{t-1}). + text_embed = self.model.embed_tokens(input_ids.reshape(1)) # (1, H) + audio_codes_buf = (info_dict.get("audio_codes", {}) or {}).get("current") + if isinstance(audio_codes_buf, torch.Tensor) and audio_codes_buf.numel() == self.n_vq: + codes = audio_codes_buf.to(device=device, dtype=torch.long) + audio_embed = torch.zeros_like(text_embed) + for i, emb_layer in enumerate(self.audio_embeddings): + code_i = codes[i].clamp(0, self.audio_vocab_size).reshape(1) + audio_embed = audio_embed + emb_layer(code_i) + combined = text_embed + audio_embed + else: + combined = text_embed + + return input_ids, combined, {"audio_state": audio_state} + + def postprocess(self, hidden_states: torch.Tensor, **_: Any) -> dict[str, Any]: + if hidden_states.numel() == 0: + return {} + return {"hidden_states": {"last": hidden_states[-1].detach()}} + + @staticmethod + def _sample_with_top_k( + logits: torch.Tensor, + top_k: int, + temperature: float, + ) -> torch.Tensor: + """Top-k sampling on a (..., V) logits tensor returning (...,) ids.""" + if temperature > 0: + 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) + kth = top_vals[..., -1:].expand_as(logits) + logits = torch.where(logits < kth, torch.full_like(logits, float("-inf")), logits) + if temperature <= 0: + return logits.argmax(dim=-1) + probs = torch.softmax(logits, dim=-1) + flat = probs.reshape(-1, probs.shape[-1]) + sampled = torch.multinomial(flat, num_samples=1).reshape(probs.shape[:-1]) + return sampled + + def _sample_audio_codes( + self, + last_h: torch.Tensor, # (1, H) + state: dict[str, Any], + ) -> torch.Tensor: # (n_vq,) + """Sample one row of n_vq audio codes given current state. + + Mirrors upstream's pre/post audio masks: + pre_audio_mask = audio_lengths > arange(n_vq) + post_audio_mask = arange(n_vq) > delayed_lengths - 1 + (or all-True when delayed_lengths is sentinel) + sampling_audio_mask = pre & post → heads to sample; rest = pad_code + """ + device = last_h.device + n_vq = self.n_vq + audio_lengths = int(state.get("audio_lengths", 0)) + delayed_lengths = int(state.get("delayed_lengths", -1)) + + idx = torch.arange(n_vq, device=device) + pre_mask = idx < audio_lengths + if delayed_lengths == -1: + post_mask = torch.ones(n_vq, dtype=torch.bool, device=device) + else: + post_mask = idx > (delayed_lengths - 1) + sampling_mask = pre_mask & post_mask + + codes = torch.full((n_vq,), self.audio_pad_code, dtype=torch.long, device=device) + if not bool(sampling_mask.any()): + return codes + + # Use deploy YAML defaults. The audio sampler isn't routed through + # vLLM's sampler so we apply the upstream's audio_top_k locally. + audio_top_k = 25 + audio_temp = 1.7 + + for i in range(n_vq): + if not bool(sampling_mask[i]): + continue + head = self.audio_heads[i] + logits = head(last_h).reshape(-1) # (V,) + logits[..., -1] = float("-inf") # invalid sentinel + logits[..., self.audio_pad_code] = float("-inf") + sampled = self._sample_with_top_k(logits, audio_top_k, audio_temp) + codes[i] = sampled.long() + return codes + + def make_omni_output( + self, + model_outputs: torch.Tensor | OmniOutput, + **kwargs: Any, + ) -> OmniOutput: + """Sample audio codes per request and stash text-mask state. + + Per-request state lives in ``info["audio_state"]``. Audio codes are + accumulated in ``info["audio_codes"]["accumulated"]`` (T_acc, NQ) and + the most recent row is stored in ``info["audio_codes"]["current"]`` + for the next preprocess step. + """ + if isinstance(model_outputs, OmniOutput): + self._batch_state = None + return model_outputs + + hidden = model_outputs # (S, H) + info_dicts: list[dict[str, Any]] = ( + kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") or [] + ) + + # Ensure each request has an initialised state (defensive: typically + # populated by preprocess on the first call). + for info in info_dicts: + if isinstance(info, dict) and not isinstance(info.get("audio_state"), dict): + info["audio_state"] = { + "audio_lengths": 0, + "delayed_lengths": -1, + "is_audio": False, + "step": 0, + } + + # Stash the per-row state for compute_logits to apply masks. logits + # rows align with hidden rows, which align with info_dicts in order. + self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts] + + audio_codes_list: list[torch.Tensor] = [] + + if hidden.numel() > 0: + num_rows = hidden.shape[0] + rows_per_req = max(1, num_rows // max(1, len(info_dicts) or 1)) + + for i, info in enumerate(info_dicts): + if not isinstance(info, dict): + continue + row_start = i * rows_per_req + row_end = min(row_start + rows_per_req, num_rows) + if row_start >= row_end: + continue + + # Sample new audio codes from the last hidden state of this request. + last_h = hidden[row_end - 1].unsqueeze(0) # (1, H) + state = info.get("audio_state", {}) or {} + new_codes = self._sample_audio_codes(last_h, state) # (n_vq,) + + acc = (info.get("audio_codes", {}) or {}).get("accumulated") + if isinstance(acc, torch.Tensor) and acc.numel() > 0: + updated_acc = torch.cat([acc.to(new_codes.device), new_codes.unsqueeze(0)], dim=0) + else: + updated_acc = new_codes.unsqueeze(0) + + info["audio_codes"] = { + "current": new_codes, + "accumulated": updated_acc, + } + audio_codes_list.append(updated_acc) + + if not audio_codes_list: + return OmniOutput( + text_hidden_states=hidden, + multimodal_outputs={}, + ) + + # Forward the accumulated audio codes (T, NQ) of the first request to + # Stage 1. Multi-request batching for the codec is not yet supported + # by this adapter; the existing code already takes the first stage + # output anyway. + return OmniOutput( + text_hidden_states=hidden, + multimodal_outputs={"codes": {"audio": audio_codes_list[0]}}, + ) + + # ------------------------------------------------------------------ + # Weight loading + # ------------------------------------------------------------------ + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Map HF weight names to vLLM-Omni module names. + + HF layout (MossTTSDelayModel): + language_model.model.* → model.* + language_model.lm_head.weight (if present) + emb_ext.{i}.weight → audio_embeddings.{i}.weight + lm_heads.0.weight → text_lm_head.weight + lm_heads.{i+1}.weight → audio_heads.{i}.weight (i ≥ 0) + """ + loaded: set[str] = set() + params_dict = dict(self.named_parameters()) + + # Buffer Qwen3 backbone weights (stripped of `language_model.` prefix) + # and delegate to self.model.load_weights() so q/k/v and gate/up are + # correctly fused into qkv_proj/gate_up_proj. + backbone_weights: list[tuple[str, torch.Tensor]] = [] + + for name, tensor in weights: + # Qwen3 backbone — checkpoint stores keys as language_model.. + if name.startswith("language_model.") and not name.startswith("language_model.lm_head"): + backbone_weights.append((name[len("language_model.") :], tensor)) + continue + + # Text LM head (lm_heads.0 in upstream == language_model.lm_head) + if name in ("lm_heads.0.weight", "language_model.lm_head.weight"): + tgt = "text_lm_head.weight" + if tgt in params_dict: + default_weight_loader(params_dict[tgt], tensor) + loaded.add(tgt) + continue + + # Audio heads: lm_heads.{k}.weight for k >= 1 → audio_heads.{k-1}.weight + if name.startswith("lm_heads."): + parts = name.split(".") + try: + k = int(parts[1]) + except (IndexError, ValueError): + continue + if k >= 1: + tgt = f"audio_heads.{k - 1}.weight" + if tgt in params_dict: + default_weight_loader(params_dict[tgt], tensor) + loaded.add(tgt) + continue + + # Audio embeddings: emb_ext.{i}.weight → audio_embeddings.{i}.weight + if name.startswith("emb_ext."): + tgt = name.replace("emb_ext.", "audio_embeddings.", 1) + if tgt in params_dict: + default_weight_loader(params_dict[tgt], tensor) + loaded.add(tgt) + continue + + backbone_loaded = self.model.load_weights(iter(backbone_weights)) + for n in backbone_loaded: + loaded.add(f"model.{n}") + + return loaded + + +# --------------------------------------------------------------------------- +# MossTTSRealtimeTalkerForGeneration +# --------------------------------------------------------------------------- + + +class MossTTSRealtimeTalkerForGeneration(nn.Module): + """Stage-0 talker for MossTTSRealtime (1.7B, TTFB ~180 ms). + + Architecture differs from the delay model: + * Backbone (Qwen3) consumes ``embed_tokens[0](text) + Σᵢ embed_tokens[i+1](audio_i)``. + * The model does NOT have a text LM head — the text column at every + decode step is forced to ``text_pad`` (or ``eos`` when the audio EOS + token has been emitted), so we synthesise a deterministic logit row + to feed the vLLM sampler. + * Per-step audio generation runs the small ``local_transformer`` (4-layer + Qwen3-style decoder, ``rvq=16`` codebooks) inside ``make_omni_output``. + Stop condition: codebook-0 token equals ``audio_eos_token`` (1026). + """ + + have_multimodal_outputs: bool = True + has_preprocess: bool = True + has_postprocess: bool = True + + AUDIO_BOS = 1025 + AUDIO_EOS = 1026 + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.vllm_config = vllm_config + self.config: MossTTSRealtimeConfig = vllm_config.model_config.hf_config + lang_cfg = self.config.language_config + local_cfg = self.config.local_config + + self.hidden_size: int = int(lang_cfg.hidden_size) + self.n_vq: int = int(self.config.rvq) + self.audio_vocab_size: int = int(self.config.audio_vocab_size) + self.audio_pad_token: int = int(self.config.audio_pad_token) + self.text_pad_id: int = int(self.config.text_pad) + self.audio_eos_id: int = int(getattr(self.config, "eos_token_id", 151645)) + self.text_vocab_size: int = int(lang_cfg.vocab_size) + + # Qwen3 backbone (uses the inner language_config). vLLM exposes its + # internal weights under ``model.*`` so we keep that prefix and remap + # the upstream ``language_model.*`` keys at load time. + backbone_vllm_config = copy.copy(vllm_config) + # Swap in the inner Qwen3 config so vLLM picks the right num_layers, + # heads etc.; KV cache sizing already uses get_text_config(). + backbone_model_config = copy.copy(vllm_config.model_config) + backbone_model_config.hf_config = lang_cfg + backbone_model_config.hf_text_config = lang_cfg + backbone_vllm_config.model_config = backbone_model_config + self.model = Qwen3Model( + vllm_config=backbone_vllm_config, + prefix=_maybe_prefix(prefix, "model"), + ) + + # Outer per-channel embeddings: index 0 is text (vocab=text_vocab_size), + # indices 1..rvq are audio codebooks (vocab=audio_vocab_size). Match + # upstream's ``MossTTSRealtime.embed_tokens`` ModuleList exactly. + self.embed_tokens = nn.ModuleList() + self.embed_tokens.append( + nn.Embedding( + self.text_vocab_size, + self.hidden_size, + padding_idx=int(getattr(lang_cfg, "pad_token_id", 151643) or 151643), + ) + ) + for _ in range(self.n_vq): + self.embed_tokens.append( + nn.Embedding(self.audio_vocab_size, self.hidden_size, padding_idx=self.audio_pad_token) + ) + + # Local depth transformer + per-codebook LM heads. + self.local_transformer = MossTTSRealtimeLocalTransformer(local_cfg) + self.local_lm_heads = nn.ModuleList( + [nn.Linear(int(local_cfg.hidden_size), self.audio_vocab_size, bias=False) for _ in range(self.n_vq)] + ) + + # No real text LM head — compute_logits builds a one-hot row directly. + # We still need a logits_processor so vLLM downstream stays happy with + # something callable; a minimal pass-through suffices. + self.logits_processor = LogitsProcessor(self.text_vocab_size) + self._batch_state: list[dict[str, Any]] | None = None + + self.gpu_resident_buffer_keys: set[tuple[str, str]] = { + ("audio_codes", "current"), + ("audio_codes", "accumulated"), + ("hidden_states", "last"), + } + + # ------------------------------------------------------------------ + # vLLM hooks + # ------------------------------------------------------------------ + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + # Used by vllm dummy profiling; real prefill goes through preprocess(). + return self.embed_tokens[0](input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **_: Any, + ) -> torch.Tensor | IntermediateTensors: + return self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor | OmniOutput, + sampling_metadata: SamplingMetadata | None = None, + ) -> torch.Tensor | None: + """Synthesise a one-hot text logit row per request. + + The realtime model has no text LM head — text is always either + ``text_pad`` (continue) or ``eos`` (stop because audio EOS just fired). + """ + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if hidden_states is None or hidden_states.numel() == 0: + return None + + B = hidden_states.shape[0] + V = self.text_vocab_size + logits = hidden_states.new_full((B, V), float("-inf")) + + states = self._batch_state or [] + rows_per_state = max(1, B // max(1, len(states) or 1)) + for i, state in enumerate(states): + r0 = i * rows_per_state + r1 = min(r0 + rows_per_state, B) + if r0 >= r1: + continue + if not isinstance(state, dict): + logits[r0:r1, self.text_pad_id] = 0.0 + continue + if state.get("is_stopping"): + logits[r0:r1, self.audio_eos_id] = 0.0 + continue + # Streaming text: emit the next remaining text token, then advance + # the cursor. Once exhausted, fall back to ``text_pad``. + cursor = int(state.get("text_cursor", 0)) + remaining = state.get("remaining_text") or [] + if 0 <= cursor < len(remaining): + tok = int(remaining[cursor]) + state["text_cursor"] = cursor + 1 + else: + tok = self.text_pad_id + logits[r0:r1, tok] = 0.0 + # Defensive: if no states recorded yet (very first call before + # make_omni_output ran), default to text_pad everywhere. + if not states: + logits[:, self.text_pad_id] = 0.0 + return logits + + # ------------------------------------------------------------------ + # Embedding (text + audio codebooks, additive) + # ------------------------------------------------------------------ + + def _build_input_embeds( + self, + text_ids: torch.Tensor, # (T,) + audio_codes: torch.Tensor | None, # (T, n_vq) or None + ) -> torch.Tensor: + embeds = self.embed_tokens[0](text_ids) + if audio_codes is None: + return embeds + for i in range(self.n_vq): + col = audio_codes[:, i].clamp_(0, self.audio_vocab_size - 1) + embeds = embeds + self.embed_tokens[i + 1](col) + return embeds + + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None, + **info_dict: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + device = input_ids.device + span_len = int(input_ids.shape[0]) + audio_state = info_dict.get("audio_state") + is_first_call = not isinstance(audio_state, dict) + + if span_len > 1 or is_first_call: + # Prefill: read the full reference-audio code grid (channels x T) + # from ``info_dict["codes"]["ref"]``; slice the chunk that aligns + # with this prefill window. + ref_codes = (info_dict.get("codes", {}) or {}).get("ref") + ref_offset = int(info_dict.get("ref_offset", 0)) + chunk_audio = None + if isinstance(ref_codes, torch.Tensor) and ref_codes.numel() > 0: + if ref_codes.dim() == 1 and ref_codes.numel() % self.n_vq == 0: + ref_codes = ref_codes.view(-1, self.n_vq) + if isinstance(ref_codes, torch.Tensor) and ref_codes.dim() == 2: + end_off = ref_offset + span_len + sliced = ref_codes[ref_offset:end_off] + if sliced.shape[0] == span_len: + chunk_audio = sliced.to(device=device, dtype=torch.long) + embeds = self._build_input_embeds(input_ids, chunk_audio) + current_codes = ( + chunk_audio[-1].detach() + if chunk_audio is not None + else torch.full((self.n_vq,), self.audio_pad_token, dtype=torch.long, device=device) + ) + # Capture the streaming-text list from the request (set by the + # realtime end2end builder; empty for the delay variants which + # don't take this code path anyway). + remaining = (info_dict.get("ids", {}) or {}).get("all") + if not isinstance(remaining, list): + remaining = [] + info_update: dict[str, Any] = { + "audio_state": { + "is_stopping": False, + "step": 0, + "text_cursor": 0, + "remaining_text": list(remaining), + }, + "audio_codes": {"current": current_codes}, + "ref_offset": ref_offset + span_len, + } + return input_ids, embeds, info_update + + # Decode step: text_token from the just-sampled vLLM logit, plus the + # audio codes the local transformer produced last step. + prev_codes = (info_dict.get("audio_codes", {}) or {}).get("current") + if prev_codes is None: + prev_codes = torch.full((self.n_vq,), self.audio_pad_token, dtype=torch.long, device=device) + embeds = self._build_input_embeds( + input_ids.reshape(-1), + prev_codes.to(device=device, dtype=torch.long).unsqueeze(0), + ) + return input_ids, embeds, {} + + def postprocess(self, hidden_states: torch.Tensor, **_: Any) -> dict[str, Any]: + if hidden_states.numel() == 0: + return {} + return {"hidden_states": {"last": hidden_states[-1].detach()}} + + # ------------------------------------------------------------------ + # Per-step audio generation via local transformer + # ------------------------------------------------------------------ + + def make_omni_output( + self, + model_outputs: torch.Tensor | OmniOutput, + **kwargs: Any, + ) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + self._batch_state = None + return model_outputs + + hidden = model_outputs # (S, H) + info_dicts: list[dict[str, Any]] = ( + kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") or [] + ) + + # Defensive state init. + for info in info_dicts: + if isinstance(info, dict) and not isinstance(info.get("audio_state"), dict): + info["audio_state"] = {"is_stopping": False, "step": 0} + + self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts] + + audio_codes_list: list[torch.Tensor] = [] + if hidden.numel() > 0 and info_dicts: + num_rows = hidden.shape[0] + rows_per_req = max(1, num_rows // max(1, len(info_dicts) or 1)) + for i, info in enumerate(info_dicts): + if not isinstance(info, dict): + continue + row_start = i * rows_per_req + row_end = min(row_start + rows_per_req, num_rows) + if row_start >= row_end: + continue + + state = info.get("audio_state", {}) or {} + if state.get("is_stopping"): + continue # already stopped — no more audio frames + + last_h = hidden[row_end - 1].unsqueeze(0) # (1, H) + # Sampling parameters mirror upstream ``MossTTSRealtimeInference.generate``: + # 0.8 / 0.6 / 30 + 1.1 repetition penalty over a 50-frame window. + rep_window = 50 + hist_per_cb: list[list[int]] = [] + acc_for_hist = (info.get("audio_codes", {}) or {}).get("accumulated") + if isinstance(acc_for_hist, torch.Tensor) and acc_for_hist.numel() > 0: + tail = acc_for_hist[-rep_window:].long().cpu().tolist() + for cb in range(self.n_vq): + hist_per_cb.append([row[cb] for row in tail]) + else: + hist_per_cb = [[] for _ in range(self.n_vq)] + new_codes = self.local_transformer.generate_frame( + last_h, + self.local_lm_heads, + temperature=0.8, + top_p=0.6, + top_k=30, + do_sample=True, + repetition_penalty=1.1, + history_per_codebook=hist_per_cb, + ).squeeze(0) # (n_vq,) + if int(state.get("step", 0)) < 5 or int(state.get("step", 0)) % 50 == 0: + logger.info( + "[MossTTSRealtime make_omni] step=%d ch0=%d cursor=%d/%d", + int(state.get("step", 0)), + int(new_codes[0].item()), + int(state.get("text_cursor", 0)), + len(state.get("remaining_text") or []), + ) + + ch0 = int(new_codes[0].item()) + # Stop condition mirrors upstream: codebook 0 == eos_audio_id. + if ch0 == self.AUDIO_EOS: + state["is_stopping"] = True + state["step"] = int(state.get("step", 0)) + 1 + info["audio_codes"] = { + "current": new_codes, + "accumulated": (info.get("audio_codes", {}) or {}).get("accumulated"), + } + continue # don't append the eos frame to accumulated + + if ch0 in (self.AUDIO_BOS, self.audio_pad_token): + # Skip the bos / pad frames — they don't decode to real audio. + state["step"] = int(state.get("step", 0)) + 1 + info["audio_codes"] = { + "current": new_codes, + "accumulated": (info.get("audio_codes", {}) or {}).get("accumulated"), + } + continue + + acc = (info.get("audio_codes", {}) or {}).get("accumulated") + if isinstance(acc, torch.Tensor) and acc.numel() > 0: + updated_acc = torch.cat([acc.to(new_codes.device), new_codes.unsqueeze(0)], dim=0) + else: + updated_acc = new_codes.unsqueeze(0) + + info["audio_codes"] = {"current": new_codes, "accumulated": updated_acc} + state["step"] = int(state.get("step", 0)) + 1 + audio_codes_list.append(updated_acc) + + if not audio_codes_list: + return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) + + # Forward the first request's accumulated codes (multi-request + # batching for the codec is not yet wired up — same as delay). + # The realtime variant emits raw codes (no delay pattern), so we + # signal that to the chunk processor via a 1-D bool tensor in + # ``meta.finished`` (already a tensor field on the schema). The + # processor reads its truth value as "skip de-delay" only for + # realtime; the delay path doesn't populate ``meta`` at all. + device = audio_codes_list[0].device + return OmniOutput( + text_hidden_states=hidden, + multimodal_outputs={ + "codes": {"audio": audio_codes_list[0]}, + "meta": {"finished": torch.tensor([True], dtype=torch.bool, device=device)}, + }, + ) + + # ------------------------------------------------------------------ + # Weight loading + # ------------------------------------------------------------------ + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Remap upstream MossTTSRealtime checkpoint names → vendored layout. + + Mapping: + embed_tokens.{i}.* → embed_tokens.{i}.* (kept) + language_model.embed_tokens.* → model.embed_tokens.* (Qwen3 inner) + language_model.layers.* → model.layers.* + language_model.norm.* → model.norm.* + local_transformer.model.embed_tokens.* → local_transformer.model.codec_embedding.* + local_transformer.model.* → local_transformer.model.* (shared body, kept) + local_transformer.local_lm_heads.* → local_lm_heads.* (top-level) + """ + loaded: set[str] = set() + params_dict = dict(self.named_parameters()) + + # Qwen3 backbone weights need vLLM's stacked-params loader to fuse + # ``q_proj``/``k_proj``/``v_proj`` → ``qkv_proj`` and + # ``gate_proj``/``up_proj`` → ``gate_up_proj``. ``default_weight_loader`` + # alone leaves those fused params un-initialised, which silently turns + # the backbone into a slightly-corrupted model that never emits EOS. + backbone_weights: list[tuple[str, torch.Tensor]] = [] + skipped: list[str] = [] + for name, tensor in weights: + if name.startswith("language_model."): + backbone_weights.append((name[len("language_model.") :], tensor)) + continue + if name.startswith("local_transformer.model."): + # Shared CodePredictorBaseModel keeps the ``.model.`` nesting and + # the per-layer/norm names verbatim; only the codebook embedding + # is renamed (upstream ``embed_tokens`` → ``codec_embedding``). + sub = name[len("local_transformer.model.") :] + if sub.startswith("rotary_emb.inv_freq"): + continue # non-persistent buffer, recomputed at runtime + if sub.startswith("embed_tokens."): + sub = "codec_embedding." + sub[len("embed_tokens.") :] + tgt = "local_transformer.model." + sub + elif name.startswith("local_transformer.local_lm_heads."): + tgt = "local_lm_heads." + name[len("local_transformer.local_lm_heads.") :] + else: + tgt = name + if tgt in params_dict: + default_weight_loader(params_dict[tgt], tensor) + loaded.add(tgt) + else: + skipped.append(f"{name}->{tgt}") + + # Delegate Qwen3 weights to its own loader (handles fused params). + backbone_loaded = self.model.load_weights(iter(backbone_weights)) + for n in backbone_loaded: + loaded.add(f"model.{n}") + + logger.info( + "[MossTTSRealtime] loaded %d/%d params; skipped=%d (first 5: %s)", + len(loaded), + len(params_dict), + len(skipped), + skipped[:5], + ) + not_loaded = [n for n in params_dict if n not in loaded] + if not_loaded: + logger.warning("[MossTTSRealtime] %d params NOT loaded (first 5: %s)", len(not_loaded), not_loaded[:5]) + return loaded + + +__all__ = [ + "MossTTSDelayTalkerForGeneration", + "MossTTSRealtimeTalkerForGeneration", +] diff --git a/vllm_omni/model_executor/models/moss_tts/pipeline.py b/vllm_omni/model_executor/models/moss_tts/pipeline.py new file mode 100644 index 000000000..fbbdd8897 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/pipeline.py @@ -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"] diff --git a/vllm_omni/model_executor/models/moss_tts/reference_encoder.py b/vllm_omni/model_executor/models/moss_tts/reference_encoder.py new file mode 100644 index 000000000..701f673a5 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/reference_encoder.py @@ -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 diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index d642c6efb..efcc42da6 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -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", diff --git a/vllm_omni/model_executor/stage_input_processors/moss_tts.py b/vllm_omni/model_executor/stage_input_processors/moss_tts.py new file mode 100644 index 000000000..562318d73 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/moss_tts.py @@ -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"]