diff --git a/tests/e2e/offline_inference/test_moss_tts.py b/tests/e2e/offline_inference/test_moss_tts.py index 505db1af7..19a8ac6d6 100644 --- a/tests/e2e/offline_inference/test_moss_tts.py +++ b/tests/e2e/offline_inference/test_moss_tts.py @@ -1,77 +1,97 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""E2E offline inference tests for MOSS-TTS two-stage pipeline. +"""E2E offline inference tests for MOSS-VoiceGenerator (MossTTSDelayModel, 1.7B). -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. +Uses the standard omni_runner + pytestmark pattern (one module-scoped engine +per file, skill invariant I4). The realtime variant lives in the sibling file +test_moss_tts_realtime.py. -A second test class covers MOSS-TTS-Realtime (1.7B, MossTTSRealtime) to -exercise the local-transformer generation path. +MOSS-VoiceGenerator synthesizes speech from a text *instruction* that describes +the desired voice (e.g. "a warm female voice with an American accent"). It does +NOT accept a reference audio clip. The request format requires running +AutoProcessor.from_pretrained once per call to encode the (text, instruction) +pair into the (prompt_token_ids, codes.ref) grid that the MossTTSDelayModel +talker expects — same as examples/offline_inference/text_to_speech/moss_tts/ +end2end.py:_build_unified_codes. -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). +No determinism test: VoiceGenerator produces variable-length output even with +a fixed seed; waveform length reproducibility is not guaranteed. """ from __future__ import annotations +import gc 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 +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config # --------------------------------------------------------------------------- -# Shared constants +# 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. -# TODO: Fix this test -# pytestmark = [pytest.mark.full_model, pytest.mark.tts] +MODEL = "OpenMOSS-Team/MOSS-VoiceGenerator" +SAMPLE_RATE = 24_000 -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( +# Stage 0 = AR talker; max_tokens=512 keeps tests fast (~20 s of audio). +# Stage 1 = codec decoder; greedy, large context to hold all codec tokens. +_STAGE0_PARAMS = 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) + max_tokens=512, seed=42, detokenize=False, ) +_STAGE1_PARAMS = SamplingParams( + temperature=0.0, + top_p=1.0, + top_k=-1, + max_tokens=65536, + seed=42, + detokenize=False, +) +_SAMPLING = [_STAGE0_PARAMS, _STAGE1_PARAMS] # --------------------------------------------------------------------------- -# Session fixtures +# Deploy config # --------------------------------------------------------------------------- -@pytest.fixture(scope="session") -def ref_audio_path(tmp_path_factory) -> str: - """Download the upstream reference clip once per session. +def _get_test_config() -> str: + """Derive a CI-friendly config from moss_voice_generator.yaml. - Set ``MOSS_TTS_SKIP_ON_NET_FAIL=1`` to skip in air-gapped environments. + Reduces Stage 0 gpu_memory_utilization from 0.60 → 0.45 to leave headroom + for Stage 1 (0.12) on a shared L4/A10G. max_num_seqs=1 keeps per-test + peak memory predictable. """ - 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) + return modify_stage_config( + get_deploy_config_path("moss_voice_generator.yaml"), + updates={ + "stages": { + 0: { + "max_num_seqs": 1, + "gpu_memory_utilization": 0.45, + }, + }, + }, + ) + + +# --------------------------------------------------------------------------- +# pytestmark — one engine for the whole module +# --------------------------------------------------------------------------- + +pytestmark = [ + pytest.mark.full_model, + pytest.mark.tts, + pytest.mark.parametrize("omni_runner", [(MODEL, _get_test_config())], indirect=True), +] # --------------------------------------------------------------------------- @@ -79,147 +99,142 @@ def ref_audio_path(tmp_path_factory) -> str: # --------------------------------------------------------------------------- -def _build_request(text: str, ref_audio_path: str, seed: int = 42) -> dict: +def _build_request(text: str, instruction: str) -> dict: + """Build a VoiceGenerator request using AutoProcessor. + + Calls proc.build_user_message + proc() to produce the unified-codes grid + (L, 1+n_vq) and splits it into prompt_token_ids + codes.ref, exactly as + end2end.py:_build_unified_codes does for the voice_generator variant. + The processor (including its CPU-resident audio_tokenizer) is freed before + returning so it does not compete with the running engine. + + Set MOSS_TTS_SKIP_ON_NET_FAIL=1 to skip in air-gapped environments where + the HF snapshot is not cached. + """ + from transformers import AutoProcessor + + try: + proc = AutoProcessor.from_pretrained(MODEL, trust_remote_code=True) + except Exception as exc: + msg = f"Cannot load AutoProcessor for {MODEL}: {exc}" + if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"): + pytest.skip(msg) + pytest.fail(msg) + + user_msg = proc.build_user_message(text=text, instruction=instruction) + batch = proc(conversations=[[user_msg]], mode="generation") + unified = batch["input_ids"][0] # (L, 1+n_vq) + text_ids = unified[:, 0].tolist() + audio_codes = unified[:, 1:].contiguous().to(torch.int64) # (L, n_vq) + del proc + gc.collect() + 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], - }, + "prompt_token_ids": text_ids, + "additional_information": {"codes": {"ref": audio_codes}}, } -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") +def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor, int]: + """Run one request and return (waveform_cpu, sample_rate). + + Iterates all OmniRequestOutputs from generate() via the top-level + multimodal_output property (which aggregates from completion outputs). + AR-stage outputs have empty multimodal_output and are silently skipped. + With async_chunk=True and FINAL_ONLY, the codec stage consolidates its + chunks into a single tensor before yielding — but we handle list fallback + for robustness. + """ + chunks: list[torch.Tensor] = [] + sr_final = SAMPLE_RATE + for out in omni_runner.omni.generate(request, _SAMPLING): + mm = out.multimodal_output + if not mm: + continue + audio = mm.get("audio") + if audio is None: + audio = mm.get("model_outputs") + if audio is None: + continue + sr = mm.get("sr") + if sr is not None: + sr_final = int(sr.item() if hasattr(sr, "item") else sr) + 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, + ) + if isinstance(audio, torch.Tensor) and audio.numel() > 0: + chunks.append(audio.reshape(-1).cpu()) + if not chunks: + raise AssertionError("No audio output received from generate()") + return torch.cat(chunks, dim=0), sr_final # --------------------------------------------------------------------------- -# MossTTSDelayModel — MOSS-VoiceGenerator 1.7B -# (tests the delay-pattern generation path on a small GPU) +# Tests # --------------------------------------------------------------------------- -_DELAY_MODEL = "OpenMOSS-Team/MOSS-VoiceGenerator" - -@pytest.fixture(scope="module") -def delay_engine(): - return Omni(model=_DELAY_MODEL, stage_init_timeout=300) - - -@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) +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_moss_tts_english(omni_runner: OmniRunner) -> None: + """VoiceGenerator: English instruction produces non-empty 24 kHz audio.""" + req = _build_request( + "Hello, this is a MOSS voice design test.", + "a warm female voice with an American accent", + ) + audio, sr = _collect_audio(omni_runner, 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" -@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) +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_moss_tts_chinese(omni_runner: OmniRunner) -> None: + """VoiceGenerator: Chinese input produces non-empty audio.""" + req = _build_request( + "你好，这是语音合成测试。", + "一个清晰温暖的女声", + ) + audio, sr = _collect_audio(omni_runner, req) assert sr == SAMPLE_RATE assert audio.numel() > 0 assert not torch.all(audio == 0) -@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" - - -@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.""" +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_moss_tts_batch(omni_runner: OmniRunner) -> None: + """VoiceGenerator: batch of two requests each returns non-empty audio.""" requests = [ - _build_request("First sentence.", ref_audio_path), - _build_request("Second sentence.", ref_audio_path), + _build_request("First sentence.", "a warm female voice"), + _build_request("Second sentence.", "a young male voice"), ] 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, - ) + for out in omni_runner.omni.generate(requests, _SAMPLING): + mm = out.multimodal_output + if not mm: + continue + audio = mm.get("audio") + if audio is None: + audio = mm.get("model_outputs") + if audio is None: + continue + 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, + ) + if isinstance(audio, torch.Tensor) and audio.numel() > 0: results.append(audio.reshape(-1).cpu()) - assert len(results) == 2, f"Expected 2 outputs, got {len(results)}" + # With async_chunk streaming + FINAL_ONLY, each request produces one + # consolidated output. The >= 2 guard handles backward-compatibility if + # consolidation behavior changes. + assert len(results) >= 2, f"Expected at least 2 audio 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) - - -@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" - - -@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) + assert audio.numel() > 0, f"Audio chunk[{i}] is empty" diff --git a/tests/e2e/offline_inference/test_moss_tts_realtime.py b/tests/e2e/offline_inference/test_moss_tts_realtime.py new file mode 100644 index 000000000..540c8750e --- /dev/null +++ b/tests/e2e/offline_inference/test_moss_tts_realtime.py @@ -0,0 +1,290 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E offline inference tests for MOSS-TTS-Realtime (MossTTSRealtime, 1.7B). + +Uses the standard omni_runner + pytestmark pattern (one module-scoped engine +per file, skill invariant I4). The delay-model variant lives in the sibling +file test_moss_tts.py. + +The request format mirrors end2end.py:_build_realtime_prompt exactly: + 1. snapshot_download to locate the MossTTSRealtimeProcessor module. + 2. MOSS-Audio-Tokenizer (separate HF repo) to encode the reference clip. + 3. Build the (L, 17) int grid; split into prompt_token_ids + codes.ref + + ids.all (remaining text tokens fed one-per-step during decode). + +Set MOSS_TTS_SKIP_ON_NET_FAIL=1 to skip in air-gapped environments. + +No determinism test: MossTTSRealtime uses a local depth transformer that is +sensitive to async scheduling timing; waveform reproducibility is not +guaranteed across back-to-back calls. +""" + +from __future__ import annotations + +import gc +import importlib.util +import os +import sys +import urllib.request +from pathlib import Path + +import numpy as np +import pytest +import soundfile as sf +import torch +from vllm import SamplingParams + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime" +SAMPLE_RATE = 24_000 + +REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/reference_zh_1.wav" + +_STAGE0_PARAMS = SamplingParams( + temperature=1.0, + top_p=0.95, + top_k=50, + max_tokens=512, + seed=42, + detokenize=False, +) +_STAGE1_PARAMS = SamplingParams( + temperature=0.0, + top_p=1.0, + top_k=-1, + max_tokens=65536, + seed=42, + detokenize=False, +) +_SAMPLING = [_STAGE0_PARAMS, _STAGE1_PARAMS] + +# --------------------------------------------------------------------------- +# Deploy config +# --------------------------------------------------------------------------- + + +def _get_test_config() -> str: + """Derive a CI-friendly config from moss_tts_realtime.yaml. + + Reduces Stage 0 gpu_memory_utilization from 0.60 → 0.45 to leave headroom + for Stage 1 (0.12) on a shared L4/A10G. max_num_seqs=1 keeps per-test + peak memory predictable. + """ + return modify_stage_config( + get_deploy_config_path("moss_tts_realtime.yaml"), + updates={ + "stages": { + 0: { + "max_num_seqs": 1, + "gpu_memory_utilization": 0.45, + }, + }, + }, + ) + + +# --------------------------------------------------------------------------- +# pytestmark — one engine for the whole module +# --------------------------------------------------------------------------- + +pytestmark = [ + pytest.mark.full_model, + pytest.mark.tts, + pytest.mark.parametrize("omni_runner", [(MODEL, _get_test_config())], indirect=True), +] + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def ref_audio_path(tmp_path_factory) -> str: + """Download the upstream reference clip once per test module. + + Set MOSS_TTS_SKIP_ON_NET_FAIL=1 to skip in air-gapped environments. + """ + target = tmp_path_factory.mktemp("moss_tts_realtime_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: + 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 _load_ref_audio(path: str, target_sr: int = 24000) -> torch.Tensor: + """Load and resample reference audio to (1, T) float32 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_request(ref_audio_path: str, text: str) -> dict: + """Build a MOSS-TTS-Realtime request. + + Mirrors end2end.py:_build_realtime_prompt exactly: + 1. Locate MossTTSRealtimeProcessor in the HF snapshot. + 2. Encode ref audio via MOSS-Audio-Tokenizer (separate repo). + 3. Build the (L, 17) int grid and return prompt_token_ids + + additional_information: {codes.ref, ids.all}. + + Frees the codec model before returning to avoid competing with the + running vllm engine for GPU/CPU memory. + """ + from huggingface_hub import snapshot_download + from transformers import AutoModel, AutoTokenizer + + # Step 1: locate the realtime processor module in the snapshot. + try: + snap_dir = Path(snapshot_download(repo_id=MODEL)) + except Exception as exc: + msg = f"Cannot locate snapshot for {MODEL}: {exc}" + if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"): + pytest.skip(msg) + pytest.fail(msg) + + proc_module_path = snap_dir / "processing_mossttsrealtime.py" + if not proc_module_path.exists(): + pytest.fail(f"Realtime processor module missing at {proc_module_path}") + + mod_name = "_moss_tts_realtime_proc" + if mod_name not in sys.modules: + spec = importlib.util.spec_from_file_location(mod_name, proc_module_path) + mod = importlib.util.module_from_spec(spec) # type: ignore[arg-type] + sys.modules[mod_name] = mod + spec.loader.exec_module(mod) # type: ignore[union-attr] + else: + mod = sys.modules[mod_name] + + tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True) + realtime_processor = mod.MossTTSRealtimeProcessor(tokenizer=tokenizer) + + # Step 2: encode the reference audio clip via MOSS-Audio-Tokenizer. + try: + codec = AutoModel.from_pretrained("OpenMOSS-Team/MOSS-Audio-Tokenizer", trust_remote_code=True).to("cpu").eval() + except Exception as exc: + msg = f"Cannot load MOSS-Audio-Tokenizer: {exc}" + if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"): + pytest.skip(msg) + pytest.fail(msg) + + ref_wav = _load_ref_audio(ref_audio_path) + wav_2d = ref_wav if ref_wav.dim() == 2 else ref_wav.unsqueeze(0) + with torch.no_grad(): + enc = codec.batch_encode([wav_2d.squeeze(0)], num_quantizers=16) + 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() + + # Step 3: build the (L, 17) prefill grid. + system_grid = realtime_processor.make_ensemble(prompt_audio_tokens=audio_tokens) + + 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 + + text_ids_only = tokenizer.encode(text, add_special_tokens=False) + if not text_ids_only: + pytest.fail(f"Empty text after tokenization: {text!r}") + + 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 at the last prefilled text token + + grid = np.concatenate([system_grid, assistant_grid, text_grid], axis=0) + prompt_token_ids = grid[:, 0].tolist() + audio_codes_tensor = torch.from_numpy(grid[:, 1:].astype(np.int64).copy()) # (L, 16) + remaining_text_ids = list(text_ids_only[cur_len:]) + + additional_info: dict = {"codes": {"ref": audio_codes_tensor}} + if remaining_text_ids: + additional_info["ids"] = {"all": remaining_text_ids} + + return { + "prompt_token_ids": prompt_token_ids, + "additional_information": additional_info, + } + + +def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor, int]: + """Run one request and return (waveform_cpu, sample_rate).""" + chunks: list[torch.Tensor] = [] + sr_final = SAMPLE_RATE + for out in omni_runner.omni.generate(request, _SAMPLING): + mm = out.multimodal_output + if not mm: + continue + audio = mm.get("audio") + if audio is None: + audio = mm.get("model_outputs") + if audio is None: + continue + sr = mm.get("sr") + if sr is not None: + sr_final = int(sr.item() if hasattr(sr, "item") else sr) + 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, + ) + if isinstance(audio, torch.Tensor) and audio.numel() > 0: + chunks.append(audio.reshape(-1).cpu()) + if not chunks: + raise AssertionError("No audio output received from generate()") + return torch.cat(chunks, dim=0), sr_final + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str) -> None: + """MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio.""" + req = _build_request(ref_audio_path, "This is a real-time TTS streaming test.") + audio, sr = _collect_audio(omni_runner, 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.advanced_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_moss_tts_realtime_chinese(omni_runner: OmniRunner, ref_audio_path: str) -> None: + """MossTTSRealtime: Chinese input produces non-empty audio.""" + req = _build_request(ref_audio_path, "你好，这是语音合成测试。") + audio, sr = _collect_audio(omni_runner, req) + + assert sr == SAMPLE_RATE + assert audio.numel() > 0 + assert not torch.all(audio == 0) 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 index a03b4c1d1..d40e22014 100644 --- 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 @@ -153,6 +153,12 @@ class MossTTSDelayTalkerForGeneration(nn.Module): # `make_omni_output`, which runs immediately before the sampler). self._batch_state: list[dict[str, Any]] | None = None + # Stacked weight caches built after load_weights(). Avoids a Python + # loop over n_vq heads/embeddings at every decode step. + # shapes: (n_vq, audio_vocab+1, hidden_size) + self._stacked_audio_head_w: torch.Tensor | None = None + self._stacked_audio_emb_w: torch.Tensor | None = None + # ------------------------------------------------------------------ # vLLM-Omni hooks # ------------------------------------------------------------------ @@ -394,9 +400,15 @@ class MossTTSDelayTalkerForGeneration(nn.Module): 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]) + if self._stacked_audio_emb_w is not None: + codes_nq = codes.t() # (n_vq, L) + audio_embed = self._stacked_audio_emb_w[ + torch.arange(self.n_vq, device=device)[:, None], codes_nq + ].sum(0) + else: + 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() @@ -436,10 +448,17 @@ class MossTTSDelayTalkerForGeneration(nn.Module): 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) + if self._stacked_audio_emb_w is not None: + # Single gather: stacked_emb[i, codes[i]] for each i → (n_vq, H) → sum → (1, H) + n_vq_idx = torch.arange(self.n_vq, device=device) + audio_embed = self._stacked_audio_emb_w[n_vq_idx, codes.clamp(0, self.audio_vocab_size)].sum( + 0, keepdim=True + ) + else: + 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 @@ -501,17 +520,38 @@ class MossTTSDelayTalkerForGeneration(nn.Module): 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 + if self._stacked_audio_head_w is not None: + # Single batched matmul: (n_vq, V+1, H) @ (H,) → (n_vq, V+1). + # Replaces n_vq serial nn.Linear calls. + all_logits = self._stacked_audio_head_w @ last_h.reshape(-1) # (n_vq, V+1) + all_logits[:, -1] = float("-inf") # invalid sentinel + all_logits[:, self.audio_pad_code] = float("-inf") + + active_idx = sampling_mask.nonzero(as_tuple=True)[0] # (n_active,) + active_logits = all_logits[active_idx] # (n_active, V+1) + active_logits = active_logits / max(audio_temp, 1e-6) + if audio_top_k > 0 and audio_top_k < active_logits.shape[-1]: + top_vals, _ = torch.topk(active_logits, audio_top_k, dim=-1) + kth = top_vals[:, -1:].expand_as(active_logits) + active_logits = torch.where( + active_logits < kth, + torch.full_like(active_logits, float("-inf")), + active_logits, + ) + probs = torch.softmax(active_logits, dim=-1) + codes[active_idx] = torch.multinomial(probs, num_samples=1).squeeze(-1).long() + return codes + + # Fallback: per-head loop (used before load_weights() completes). 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[..., -1] = float("-inf") logits[..., self.audio_pad_code] = float("-inf") sampled = self._sample_with_top_k(logits, audio_top_k, audio_temp) codes[i] = sampled.long() @@ -678,6 +718,14 @@ class MossTTSDelayTalkerForGeneration(nn.Module): for n in backbone_loaded: loaded.add(f"model.{n}") + # Build stacked weight caches for vectorised decode ops. + self._stacked_audio_head_w = torch.stack( + [h.weight.detach() for h in self.audio_heads], dim=0 + ) # (n_vq, audio_vocab+1, hidden_size) + self._stacked_audio_emb_w = torch.stack( + [e.weight.detach() for e in self.audio_embeddings], dim=0 + ) # (n_vq, audio_vocab+1, hidden_size) + return loaded @@ -765,6 +813,10 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): self.logits_processor = LogitsProcessor(self.text_vocab_size) self._batch_state: list[dict[str, Any]] | None = None + # Stacked weight cache built after load_weights(). + # shape: (n_vq, audio_vocab_size, hidden_size) + self._stacked_audio_emb_w: torch.Tensor | None = None + self.gpu_resident_buffer_keys: set[tuple[str, str]] = { ("audio_codes", "current"), ("audio_codes", "accumulated"), @@ -854,6 +906,11 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): embeds = self.embed_tokens[0](text_ids) if audio_codes is None: return embeds + if self._stacked_audio_emb_w is not None: + # Single gather: (n_vq, vocab, H)[arange[:,None], codes.T] → (n_vq, T, H) → sum → (T, H) + n_vq_idx = torch.arange(self.n_vq, device=audio_codes.device) + codes_t = audio_codes.t().clamp(0, self.audio_vocab_size - 1) # (n_vq, T) + return embeds + self._stacked_audio_emb_w[n_vq_idx.unsqueeze(1), codes_t].sum(0) 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) @@ -1115,6 +1172,12 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): for n in backbone_loaded: loaded.add(f"model.{n}") + # Build stacked embedding cache for vectorised decode ops. + # embed_tokens[0] is the text embedding; [1..n_vq] are audio codebooks. + self._stacked_audio_emb_w = torch.stack( + [self.embed_tokens[i + 1].weight.detach() for i in range(self.n_vq)], dim=0 + ) # (n_vq, audio_vocab_size, hidden_size) + logger.info( "[MossTTSRealtime] loaded %d/%d params; skipped=%d (first 5: %s)", len(loaded),