From a8434a60da1f40a0d118a326908fd20c9ebe75c7 Mon Sep 17 00:00:00 2001 From: akshatvishu Date: Thu, 17 Sep 2026 00:09:51 +0530 Subject: [PATCH] [BugFix][TTS] Fix MOSS Realtime serving and generation (#5661) Signed-off-by: akshatvishu --- .../test_moss_tts_expansion.py | 9 +- .../test_moss_tts_realtime.py | 5 +- .../test_moss_tts_realtime_expansion.py | 5 +- .../online_serving/test_moss_tts_expansion.py | 28 +- .../test_moss_tts_realtime_serving.py | 302 ++++++++++++++++++ .../moss_tts/test_audio_tokenizer_v2.py | 54 ++++ .../moss_tts/test_realtime_chunked_prefill.py | 86 +++++ .../test_realtime_generation_limit.py | 53 +++ .../models/moss_tts/test_realtime_prompt.py | 78 +++++ .../openai/tts_adapters/moss_tts.py | 71 +++- .../moss_tts/modeling_moss_tts_talker.py | 57 ++-- .../models/moss_tts/realtime_prompt.py | 60 ++++ .../models/moss_tts/reference_encoder.py | 27 +- 13 files changed, 771 insertions(+), 64 deletions(-) create mode 100644 tests/entrypoints/openai_api/test_moss_tts_realtime_serving.py create mode 100644 tests/model_executor/models/moss_tts/test_audio_tokenizer_v2.py create mode 100644 tests/model_executor/models/moss_tts/test_realtime_chunked_prefill.py create mode 100644 tests/model_executor/models/moss_tts/test_realtime_generation_limit.py create mode 100644 tests/model_executor/models/moss_tts/test_realtime_prompt.py create mode 100644 vllm_omni/model_executor/models/moss_tts/realtime_prompt.py diff --git a/tests/e2e/offline_inference/test_moss_tts_expansion.py b/tests/e2e/offline_inference/test_moss_tts_expansion.py index 81a64fcd9..b899252b4 100644 --- a/tests/e2e/offline_inference/test_moss_tts_expansion.py +++ b/tests/e2e/offline_inference/test_moss_tts_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E offline inference tests for MOSS-VoiceGenerator (MossTTSDelayModel, 1.7B). Weekly expansion coverage (``pytest.mark.slow``) for the delay-model path. @@ -74,7 +74,6 @@ def _get_test_config() -> str: pytestmark = [ pytest.mark.slow, - pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"), pytest.mark.tts, pytest.mark.parametrize( "omni_runner", @@ -137,7 +136,7 @@ def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor return torch.cat(chunks, dim=0), sr_final -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) def test_moss_tts_delay_english(omni_runner: OmniRunner) -> None: """VoiceGenerator: English instruction produces non-empty 24 kHz audio.""" req = _build_request( @@ -151,7 +150,7 @@ def test_moss_tts_delay_english(omni_runner: OmniRunner) -> None: assert not torch.all(audio == 0), "Audio is silence" -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) def test_moss_tts_delay_chinese(omni_runner: OmniRunner) -> None: """VoiceGenerator: Chinese input produces non-empty audio.""" req = _build_request( @@ -165,7 +164,7 @@ def test_moss_tts_delay_chinese(omni_runner: OmniRunner) -> None: assert not torch.all(audio == 0) -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) def test_moss_tts_delay_batch(omni_runner: OmniRunner) -> None: """VoiceGenerator: batch of two requests each returns non-empty audio.""" requests = [ diff --git a/tests/e2e/offline_inference/test_moss_tts_realtime.py b/tests/e2e/offline_inference/test_moss_tts_realtime.py index 81d97cd74..0f0bf59c7 100644 --- a/tests/e2e/offline_inference/test_moss_tts_realtime.py +++ b/tests/e2e/offline_inference/test_moss_tts_realtime.py @@ -95,7 +95,6 @@ def _get_test_config() -> str: # --------------------------------------------------------------------------- pytestmark = [ - pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"), pytest.mark.slow, pytest.mark.tts, pytest.mark.parametrize( @@ -272,7 +271,7 @@ def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor # --------------------------------------------------------------------------- -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, 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.") @@ -283,7 +282,7 @@ def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str) assert not torch.all(audio == 0), "Audio is silence" -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, 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, "你好，这是语音合成测试。") diff --git a/tests/e2e/offline_inference/test_moss_tts_realtime_expansion.py b/tests/e2e/offline_inference/test_moss_tts_realtime_expansion.py index b84e5a0ff..678fc7e55 100644 --- a/tests/e2e/offline_inference/test_moss_tts_realtime_expansion.py +++ b/tests/e2e/offline_inference/test_moss_tts_realtime_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E offline inference tests for MOSS-TTS-Realtime (MossTTSRealtime, 1.7B). Exercises the local-transformer generation path on L4. Delay-model coverage is in @@ -33,7 +33,6 @@ _OMNI_RUNNER_PARAM = ( pytestmark = [ pytest.mark.slow, pytest.mark.tts, - pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"), pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), ] @@ -118,7 +117,7 @@ def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor raise AssertionError("No stage outputs received") -@hardware_test(res={"cuda": "L4"}) +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path) -> None: """MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio.""" req = _build_request("This is a real-time TTS streaming test.", ref_audio_path) diff --git a/tests/e2e/online_serving/test_moss_tts_expansion.py b/tests/e2e/online_serving/test_moss_tts_expansion.py index 598326a1e..bd38a7fc7 100644 --- a/tests/e2e/online_serving/test_moss_tts_expansion.py +++ b/tests/e2e/online_serving/test_moss_tts_expansion.py @@ -30,20 +30,16 @@ from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path -_SKIP_ISSUE_6417 = pytest.mark.skip( - reason="https://github.com/vllm-project/vllm-omni/issues/6417", -) - pytestmark = [ + pytest.mark.advanced_model, pytest.mark.slow, pytest.mark.tts, - _SKIP_ISSUE_6417, ] MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime" REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/HEAD/assets/audio/reference_zh_1.wav" -# Voice-clone output is not reliably transcribed by the Whisper check used at -# full_model run_level; assert non-trivial WAV payload size instead. +# Reject a response that contains only a header or a very short audio payload. +# Advanced and full runs also validate the transcript. _MIN_AUDIO_BYTES = 10_000 @@ -77,7 +73,9 @@ def get_prompt(prompt_type: str = "text") -> str: """ prompts = { "text": "Hello, this is a real-time voice cloning demo for testing.", - "chinese": "你好，这是一段实时语音合成测试。", + # Match the offline test prompt. Whisper-small can transcribe "一段" + # as "一本", which makes this short sample fail the similarity check. + "chinese": "你好，这是语音合成测试。", } return prompts.get(prompt_type, prompts["text"]) @@ -94,7 +92,7 @@ tts_server_params = [ ] -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_text_to_audio_001(omni_server, online_client, ref_audio_data_url) -> None: """ @@ -105,9 +103,8 @@ def test_text_to_audio_001(omni_server, online_client, ref_audio_data_url) -> No Input Setting: stream=False Datasets: single request - NOTE: ``min_audio_bytes`` skips Whisper transcript similarity — cloned - voice output is often empty or mismatched under ASR without indicating a - real serving regression. + ``min_audio_bytes`` adds a payload-size check. It does not disable the + transcript check used by advanced and full runs. """ request_config = { "model": omni_server.model, @@ -121,7 +118,7 @@ def test_text_to_audio_001(omni_server, online_client, ref_audio_data_url) -> No online_client.send_audio_speech_request(request_config) -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_text_to_audio_002_streaming(omni_server, online_client, ref_audio_data_url) -> None: """ @@ -150,7 +147,7 @@ def test_text_to_audio_002_streaming(omni_server, online_client, ref_audio_data_ online_client.send_audio_speech_request(request_config) -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_text_to_audio_003_chinese(omni_server, online_client, ref_audio_data_url) -> None: """ @@ -161,7 +158,8 @@ def test_text_to_audio_003_chinese(omni_server, online_client, ref_audio_data_ur Input Setting: stream=False Datasets: single request - NOTE: same ``min_audio_bytes`` rationale as test_text_to_audio_001. + ``min_audio_bytes`` adds a payload-size check. It does not disable the + transcript check used by advanced and full runs. """ request_config = { "model": omni_server.model, diff --git a/tests/entrypoints/openai_api/test_moss_tts_realtime_serving.py b/tests/entrypoints/openai_api/test_moss_tts_realtime_serving.py new file mode 100644 index 000000000..d69597634 --- /dev/null +++ b/tests/entrypoints/openai_api/test_moss_tts_realtime_serving.py @@ -0,0 +1,302 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import asyncio +import threading +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +import torch + +from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest +from vllm_omni.entrypoints.openai.tts_adapters import moss_tts as adapter_module +from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext +from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import MossTTSAdapter +from vllm_omni.model_executor.models.moss_tts import reference_encoder +from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize( + ("hf_config", "expected_codec_path"), + [ + (SimpleNamespace(), "OpenMOSS-Team/MOSS-Audio-Tokenizer"), + ( + SimpleNamespace(codec_model_name_or_path="/models/custom-moss-codec"), + "/models/custom-moss-codec", + ), + ], + ids=["default-codec", "configured-codec"], +) +def test_realtime_components_use_the_realtime_model_and_codec( + monkeypatch: pytest.MonkeyPatch, + hf_config: SimpleNamespace, + expected_codec_path: str, +) -> None: + engine_client = SimpleNamespace( + model_config=SimpleNamespace( + model="OpenMOSS-Team/MOSS-TTS-Realtime", + hf_config=hf_config, + ) + ) + + server = MossTTSAdapter(SpeechServingContext(server=object(), engine_client=engine_client)) + + tokenizer = object() + codec = type("Codec", (), {"to": lambda self, device: self, "eval": lambda self: self})() + processor_calls = [] + + class Processor: + def __init__(self, *, tokenizer): + processor_calls.append(tokenizer) + + class_calls = [] + tokenizer_calls = [] + codec_calls = [] + + def fake_get_class(class_reference, model_id): + class_calls.append((class_reference, model_id)) + return Processor + + def fake_load_tokenizer(model_id, *, trust_remote_code): + tokenizer_calls.append((model_id, trust_remote_code)) + return tokenizer + + def fake_load_codec(model_id, *, trust_remote_code): + codec_calls.append((model_id, trust_remote_code)) + return codec + + monkeypatch.setattr(adapter_module, "get_class_from_dynamic_module", fake_get_class) + monkeypatch.setattr(adapter_module.AutoTokenizer, "from_pretrained", fake_load_tokenizer) + monkeypatch.setattr(adapter_module.AutoModel, "from_pretrained", fake_load_codec) + + components = server._get_moss_realtime_components() + + assert components[0] is tokenizer + assert components[2] is codec + assert server._get_moss_realtime_components() is components + assert class_calls == [ + ( + "processing_mossttsrealtime.MossTTSRealtimeProcessor", + "OpenMOSS-Team/MOSS-TTS-Realtime", + ) + ] + assert tokenizer_calls == [("OpenMOSS-Team/MOSS-TTS-Realtime", True)] + assert codec_calls == [(expected_codec_path, True)] + assert processor_calls == [tokenizer] + + +def test_realtime_serving_builds_the_talker_prompt(monkeypatch: pytest.MonkeyPatch) -> None: + server = MossTTSAdapter(SpeechServingContext(server=SimpleNamespace(uploaded_speakers={}))) + server._moss_variant = "realtime" + tokenizer = object() + processor = object() + codec = object() + request_thread = threading.get_ident() + + def load_components(): + assert threading.get_ident() != request_thread + return tokenizer, processor, codec + + server._get_moss_realtime_components = load_components + + reference_codes = torch.arange(64, dtype=torch.int64).reshape(4, 16) + encode = AsyncMock(return_value=([reference_codes], {0: "reference-cache-key"})) + monkeypatch.setattr(server, "_encode_moss_references", encode) + + build_call = None + + def fake_build(actual_tokenizer, actual_processor, text, actual_codes): + nonlocal build_call + build_call = (actual_tokenizer, actual_processor, text, actual_codes) + return { + "prompt_token_ids": [10, 11], + "codes": {"ref": actual_codes}, + "ids": {"all": [12]}, + } + + monkeypatch.setattr(adapter_module, "build_realtime_prompt", fake_build) + + params = asyncio.run( + server._build_moss_tts_params( + OpenAICreateSpeechRequest( + input="speak this text", + ref_audio="data:audio/wav;base64,AAAA", + max_new_tokens=50, + ) + ) + ) + + encode.assert_awaited_once() + assert encode.call_args.kwargs == {"has_inline_ref_audio": False, "two_speaker": False} + assert encode.call_args.args[0].ref_audio == "data:audio/wav;base64,AAAA" + assert build_call == (tokenizer, processor, "speak this text", reference_codes) + assert params == { + "prompt_token_ids": [10, 11], + "codes": {"ref": reference_codes}, + "ids": {"all": [12]}, + "max_new_frames": [50], + "ref_audio_cache_key": "reference-cache-key", + } + assert "prompt_audio_array" not in params + + +def test_cancelled_component_waiter_does_not_duplicate_codec_load(monkeypatch): + adapter = MossTTSAdapter( + SpeechServingContext( + server=object(), + engine_client=SimpleNamespace( + model_config=SimpleNamespace(model="OpenMOSS-Team/MOSS-TTS-Realtime", hf_config=SimpleNamespace()) + ), + ) + ) + started = threading.Event() + release = threading.Event() + second_started = threading.Event() + loads = [] + codec = type("Codec", (), {"to": lambda self, device: self, "eval": lambda self: self})() + + def load_codec(*args, **kwargs): + loads.append(args) + started.set() + assert release.wait(5), "codec loader was not released" + return codec + + monkeypatch.setattr(adapter_module, "get_class_from_dynamic_module", lambda *args: lambda **kwargs: object()) + monkeypatch.setattr(adapter_module.AutoTokenizer, "from_pretrained", lambda *args, **kwargs: object()) + monkeypatch.setattr(adapter_module.AutoModel, "from_pretrained", load_codec) + + def second_load(): + second_started.set() + return adapter._get_moss_realtime_components() + + async def run(): + first = asyncio.create_task(asyncio.to_thread(adapter._get_moss_realtime_components)) + try: + assert await asyncio.to_thread(started.wait, 5) + first.cancel() + with pytest.raises(asyncio.CancelledError): + await first + second = asyncio.create_task(asyncio.to_thread(second_load)) + try: + assert await asyncio.to_thread(second_started.wait, 5) + finally: + release.set() + components = await asyncio.wait_for(second, 5) + assert components[2] is codec + assert adapter._get_moss_realtime_components() is components + assert len(loads) == 1 + finally: + release.set() + + asyncio.run(run()) + + +@pytest.mark.parametrize("voice_name", [None, "speaker"]) +def test_realtime_reference_cache_reuses_codes(voice_name): + cache = SpeakerEmbeddingCache() + resolve = AsyncMock(return_value=([0.0], 24000, "audio-key")) + request_thread = threading.get_ident() + encode_calls = [] + expected = torch.arange(64).reshape(4, 16) + + def batch_encode(wavs, *, num_quantizers): + assert threading.get_ident() != request_thread + assert num_quantizers == 16 + encode_calls.append(wavs) + return SimpleNamespace(audio_codes=expected.T.unsqueeze(1), audio_codes_lengths=torch.tensor([4])) + + encoder = reference_encoder.build_reference_encoder( + SimpleNamespace(batch_encode=batch_encode), variant="realtime", speaker_cache=cache + ) + + async def run(): + kwargs = dict( + resolve_ref_audio=resolve, + get_artifact_key=lambda key: "content-key", + voice_name=voice_name, + voice_created_at=1, + ) + try: + first, first_key = await encoder.encode("reference", **kwargs) + second, second_key = await encoder.encode("reference", **kwargs) + assert torch.equal(first, expected) + assert torch.equal(second, expected) + assert first.dtype == torch.int64 + assert first_key == "audio-key" + assert second_key == (None if voice_name else "audio-key") + second.zero_() + third, _ = await encoder.encode("another-reference", **kwargs) + assert torch.equal(third, expected) + assert len(encode_calls) == 1 + assert resolve.await_count == (1 if voice_name else 3) + key = cache.make_cache_key( + voice_name or "ref:content-key", "moss_tts_realtime_nq16", 1 if voice_name else 0 + ) + assert cache.get(key)["codes"].dtype == torch.int32 + if voice_name: + kwargs["voice_created_at"] = 2 + await encoder.encode("reference", **kwargs) + assert len(encode_calls) == 2 + assert resolve.await_count == 2 + finally: + await encoder.aclose() + + asyncio.run(run()) + + +def test_realtime_reference_batch_trims_each_item(): + cache = SpeakerEmbeddingCache() + codes = torch.arange(16 * 2 * 5).reshape(16, 2, 5) + calls = [] + + def batch_encode(wavs, *, num_quantizers): + calls.append(wavs) + assert num_quantizers == 16 + assert [wav.shape for wav in wavs] == [torch.Size([8]), torch.Size([12])] + return SimpleNamespace(audio_codes=codes, audio_codes_lengths=torch.tensor([3, 5])) + + encoder = reference_encoder.build_reference_encoder( + SimpleNamespace(batch_encode=batch_encode), variant="realtime", speaker_cache=cache + ) + + async def resolve(ref): + return [0.0] * (8 if ref == "short" else 12), 24000, ref + + async def run(): + try: + short, long = await asyncio.gather( + encoder.encode("short", resolve_ref_audio=resolve, get_artifact_key=lambda key: key), + encoder.encode("long", resolve_ref_audio=resolve, get_artifact_key=lambda key: key), + ) + torch.testing.assert_close(short[0], codes[:, 0, :3].T) + torch.testing.assert_close(long[0], codes[:, 1, :5].T) + assert len(calls) == 1 + finally: + await encoder.aclose() + + asyncio.run(run()) + + +def test_failed_realtime_reference_encoding_does_not_populate_cache(): + cache = SpeakerEmbeddingCache() + resolve = AsyncMock(return_value=([0.0], 24000, "audio-key")) + + def fail(*args, **kwargs): + raise ValueError("invalid reference") + + encoder = reference_encoder.build_reference_encoder( + SimpleNamespace(batch_encode=fail), variant="realtime", speaker_cache=cache + ) + + async def run(): + try: + with pytest.raises(ValueError, match="invalid reference"): + await encoder.encode("reference", resolve_ref_audio=resolve, get_artifact_key=lambda key: None) + assert cache.memory_bytes() == 0 + finally: + await encoder.aclose() + + asyncio.run(run()) diff --git a/tests/model_executor/models/moss_tts/test_audio_tokenizer_v2.py b/tests/model_executor/models/moss_tts/test_audio_tokenizer_v2.py new file mode 100644 index 000000000..f5d799dd7 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_audio_tokenizer_v2.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Regression tests for MOSS audio tokenizer v2 projection modules.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import ( + MossAudioTokenizerProjectedTransformer, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_projected_transformer_keeps_learned_projections_when_dimensions_match() -> None: + module = MossAudioTokenizerProjectedTransformer( + input_dimension=8, + output_dimension=8, + d_model=8, + module_type="transformer", + num_heads=1, + num_layers=0, + positional_embedding="rope", + ) + + assert isinstance(module.input_proj, nn.Linear) + assert isinstance(module.output_proj, nn.Linear) + assert module.input_proj.weight.shape == (8, 8) + assert module.output_proj.weight.shape == (8, 8) + + x = torch.randn(2, 8, 5) + lengths = torch.tensor([5, 3], dtype=torch.long) + output, output_lengths = module(x, lengths) + + assert output.shape == x.shape + torch.testing.assert_close(output_lengths, lengths) + + +def test_projected_transformer_keeps_linear_projections_when_dimensions_differ() -> None: + module = MossAudioTokenizerProjectedTransformer( + input_dimension=8, + output_dimension=6, + d_model=10, + module_type="transformer", + num_heads=1, + num_layers=0, + positional_embedding="rope", + ) + + assert isinstance(module.input_proj, nn.Linear) + assert isinstance(module.output_proj, nn.Linear) diff --git a/tests/model_executor/models/moss_tts/test_realtime_chunked_prefill.py b/tests/model_executor/models/moss_tts/test_realtime_chunked_prefill.py new file mode 100644 index 000000000..916213318 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_realtime_chunked_prefill.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_realtime_talker_preserves_state_across_chunked_prefill() -> None: + from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import ( + MossTTSRealtimeTalkerForGeneration, + ) + + talker = MossTTSRealtimeTalkerForGeneration.__new__(MossTTSRealtimeTalkerForGeneration) + talker.n_vq = 2 + talker.audio_pad_token = 1024 + talker.text_vocab_size = 16 + talker.text_pad_id = 2 + talker._batch_state = None + talker._batch_state_spans = None + + embedded_audio: list[torch.Tensor | None] = [] + + def build_input_embeds(input_ids, audio_codes): + embedded_audio.append(audio_codes.clone() if audio_codes is not None else None) + return torch.zeros((input_ids.shape[0], 1)) + + generated: list[None] = [] + + def generate_frame(*args, **kwargs): + generated.append(None) + return torch.ones((1, talker.n_vq), dtype=torch.long) + + talker._build_input_embeds = build_input_embeds + talker.local_transformer = SimpleNamespace(generate_frame=generate_frame) + talker.local_lm_heads = [] + + ref_codes = torch.arange(10, dtype=torch.long).reshape(5, 2) + info: dict[str, Any] = { + "codes": {"ref": ref_codes}, + "ids": {"all": [7, 8]}, + "max_new_frames": [3], + } + chunks = [(0, 2, False), (2, 2, False), (4, 1, True)] + + state: dict[str, Any] | None = None + for computed, length, eligible in chunks: + _, _, update = talker.preprocess( + torch.ones(length, dtype=torch.long), + None, + _omni_is_prefill=True, + _omni_num_computed_tokens=computed, + _omni_prompt_len=5, + **info, + ) + info.update(update) + if state is None: + state = info["audio_state"] + else: + assert info["audio_state"] is state + + output = talker.make_omni_output( + torch.ones((length, 1)), + model_intermediate_buffer=[info], + request_token_spans=[(0, length)], + request_sample_eligible=[eligible], + ) + if not eligible: + assert output.multimodal_outputs == {} + logits = talker.compute_logits(torch.ones((1, 1))) + assert logits.argmax().item() == (7 if eligible else talker.text_pad_id) + + assert state is not None + assert info["ref_offset"] == 5 + assert state["step"] == 1 + assert state["text_cursor"] == 1 + assert state["remaining_text"] == [7, 8] + assert len(generated) == 1 + assert output.multimodal_outputs["codes"]["audio"][0].shape == (1, 2) + for actual, expected in zip(embedded_audio, ref_codes.split((2, 2, 1)), strict=True): + assert actual is not None + torch.testing.assert_close(actual, expected) diff --git a/tests/model_executor/models/moss_tts/test_realtime_generation_limit.py b/tests/model_executor/models/moss_tts/test_realtime_generation_limit.py new file mode 100644 index 000000000..1bbfaeeef --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_realtime_generation_limit.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_realtime_talker_stops_at_max_new_frames() -> None: + from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import ( + MossTTSRealtimeTalkerForGeneration, + ) + + talker = MossTTSRealtimeTalkerForGeneration.__new__(MossTTSRealtimeTalkerForGeneration) + talker.n_vq = 4 + talker.audio_pad_token = 1024 + talker.text_vocab_size = 8 + talker.text_pad_id = 2 + talker.audio_eos_id = 3 + talker._batch_state = None + talker._build_input_embeds = lambda input_ids, audio_codes: torch.zeros((input_ids.shape[0], 1)) + + generated: list[None] = [] + + def generate_frame(*args, **kwargs): + generated.append(None) + return torch.ones((1, talker.n_vq), dtype=torch.long) + + talker.local_transformer = SimpleNamespace(generate_frame=generate_frame) + talker.local_lm_heads = [] + + _, _, info = talker.preprocess( + torch.tensor([1, 2]), + None, + max_new_frames=[2], + ) + assert info["audio_state"]["max_new_frames"] == 2 + + for _ in range(3): + output = talker.make_omni_output( + torch.ones((1, 1)), + runtime_additional_information=[info], + request_token_spans=[(0, 1)], + ) + + assert len(generated) == 2 + assert info["audio_state"]["step"] == 2 + assert info["audio_state"]["is_stopping"] is True + assert output.multimodal_outputs == {} + assert talker.compute_logits(torch.ones((1, 1))).argmax().item() == talker.audio_eos_id diff --git a/tests/model_executor/models/moss_tts/test_realtime_prompt.py b/tests/model_executor/models/moss_tts/test_realtime_prompt.py new file mode 100644 index 000000000..cb36f6dd8 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_realtime_prompt.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import numpy as np +import pytest +import torch + +from vllm_omni.model_executor.models.moss_tts.realtime_prompt import build_realtime_prompt + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Tokenizer: + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + del add_special_tokens + if text == "<|im_start|>assistant\n": + return [20, 21] + return list(range(100, 100 + len(text.split()))) + + +class _Processor: + audio_channel_pad = 1024 + audio_bos_token = 1025 + + def __init__(self) -> None: + self.reference_codes = None + + def make_ensemble(self, prompt_audio_tokens): + self.reference_codes = prompt_audio_tokens + return np.full((3, 17), self.audio_channel_pad, dtype=np.int64) + + +def test_build_realtime_prompt_aligns_reference_and_prefill_text() -> None: + processor = _Processor() + reference_codes = torch.arange(80, dtype=torch.int64).reshape(5, 16) + + params = build_realtime_prompt( + _Tokenizer(), + processor, + "one two three", + reference_codes, + ) + + assert np.array_equal(processor.reference_codes, reference_codes.numpy()) + assert params["prompt_token_ids"] == [1024, 1024, 1024, 20, 21, 100, 101, 102] + audio_grid = params["codes"]["ref"] + assert audio_grid.shape == (len(params["prompt_token_ids"]), 16) + assert audio_grid[-1, 0].item() == processor.audio_bos_token + assert "ids" not in params + + +def test_build_realtime_prompt_streams_text_after_first_twelve_tokens() -> None: + params = build_realtime_prompt( + _Tokenizer(), + _Processor(), + " ".join(f"word{i}" for i in range(15)), + torch.zeros((2, 16), dtype=torch.int64), + ) + + assert params["prompt_token_ids"][-12:] == list(range(100, 112)) + assert params["ids"] == {"all": [112, 113, 114]} + + +def test_build_realtime_prompt_rejects_wrong_processor_layout(monkeypatch: pytest.MonkeyPatch) -> None: + processor = _Processor() + monkeypatch.setattr( + processor, + "make_ensemble", + lambda prompt_audio_tokens: np.zeros((3, 16), dtype=np.int64), + ) + + with pytest.raises(ValueError, match=r"expected \(sequence_length, 17\)"): + build_realtime_prompt( + _Tokenizer(), + processor, + "one", + torch.zeros((2, 16), dtype=torch.int64), + ) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py index bed20efc2..555fec42d 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py @@ -6,8 +6,12 @@ Both variants share the same build/validate flow (``_build_moss_tts_params`` handles each); they are registered under distinct model-type names. """ +import asyncio +import threading from typing import TYPE_CHECKING, Any, cast +from transformers import AutoModel, AutoTokenizer +from transformers.dynamic_module_utils import get_class_from_dynamic_module from vllm.inputs import tokens_input from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter @@ -18,6 +22,7 @@ from vllm_omni.entrypoints.openai.tts_adapters.base import ( apply_max_new_tokens, conditioning_cache_salt, ) +from vllm_omni.model_executor.models.moss_tts.realtime_prompt import build_realtime_prompt if TYPE_CHECKING: from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest @@ -30,6 +35,7 @@ class _MossTTSAdapterBase(ARTTSAdapter): super().__init__(ctx) self._moss_variant = None if self.name == "moss_tts_nano" else self._detect_moss_variant() self._moss_processor_cache = None + self._moss_realtime_components_lock = threading.Lock() @property def engine_client(self): @@ -63,7 +69,7 @@ class _MossTTSAdapterBase(ARTTSAdapter): # inside the model package; this layer only supplies the processor and # the process-wide speaker cache. encoder = build_reference_encoder( - self._get_moss_processor(), + self._get_moss_realtime_components()[2] if self._moss_variant == "realtime" else self._get_moss_processor(), variant=cast(str, self._moss_variant), speaker_cache=self._speaker_cache, ) @@ -157,6 +163,45 @@ class _MossTTSAdapterBase(ARTTSAdapter): self._moss_processor_cache = proc return proc + def _get_moss_realtime_components(self): + with self._moss_realtime_components_lock: + cached = getattr(self, "_moss_realtime_components", None) + if cached is not None: + return cached + + model_id = self.engine_client.model_config.model + processor_cls = get_class_from_dynamic_module( + "processing_mossttsrealtime.MossTTSRealtimeProcessor", + model_id, + ) + tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) + processor = processor_cls(tokenizer=tokenizer) + hf_config = self.engine_client.model_config.hf_config + codec_path = str( + getattr( + hf_config, + "codec_model_name_or_path", + getattr( + hf_config, + "audio_tokenizer_name_or_path", + "OpenMOSS-Team/MOSS-Audio-Tokenizer", + ), + ) + ) + codec = ( + AutoModel.from_pretrained( + codec_path, + trust_remote_code=True, + ) + # Reference encoding runs in the API process. Keep the codec on CPU + # so it does not reserve accelerator memory outside the model workers. + .to("cpu") + .eval() + ) + cached = (tokenizer, processor, codec) + self._moss_realtime_components = cached + return cached + async def _build_moss_tts_params( self, request: "OpenAICreateSpeechRequest", @@ -196,24 +241,18 @@ class _MossTTSAdapterBase(ARTTSAdapter): params["ref_audio_cache_key"] = cache_key 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 = { - "text": [request.input or ""], - "mode": ["voice_clone"], - } + tokenizer, processor, _ = await asyncio.to_thread(self._get_moss_realtime_components) + references, realtime_resolve_keys = await self._encode_moss_references( + request, + has_inline_ref_audio=has_inline_ref_audio, + two_speaker=False, + ) + params = build_realtime_prompt(tokenizer, processor, request.input or "", references[0]) if request.max_new_tokens is not None: params["max_new_frames"] = [request.max_new_tokens] - wav_list, sr, cache_key = await self._resolve_ref_audio(cast(str, request.ref_audio)) - params["prompt_audio_array"] = [[wav_list, sr]] - params["ref_audio_cache_key"] = cache_key + if 0 in realtime_resolve_keys: + params["ref_audio_cache_key"] = realtime_resolve_keys[0] return params # ---- MossTTSDelay family (tts/ttsd/sound_effect/voice_generator) 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 9d2d2d171..8d36086ef 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 @@ -44,10 +44,10 @@ def _maybe_prefix(prefix: str, name: str) -> str: def _iter_state_row_spans( - states: Sequence[dict[str, Any]], + states: Sequence[dict[str, Any] | None], spans: Sequence[tuple[int, int]] | None, num_rows: int, -) -> Iterable[tuple[dict[str, Any], int, int]]: +) -> Iterable[tuple[dict[str, Any] | None, int, int]]: """Yield state/logit-row ranges without assuming equal request row counts.""" if num_rows <= 0 or not states: return @@ -798,6 +798,9 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): have_multimodal_outputs: bool = True has_preprocess: bool = True has_postprocess: bool = True + # Intermediate prefill chunks reach make_omni_output but must not produce + # audio until the scheduler has completed the prompt. + requires_request_sample_eligibility: bool = True AUDIO_BOS = 1025 AUDIO_EOS = 1026 @@ -858,7 +861,7 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): # 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._batch_state: list[dict[str, Any] | None] | None = None # Stacked weight cache built after load_weights(). # shape: (n_vq, audio_vocab_size, hidden_size) @@ -972,13 +975,14 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): span_len = int(input_ids.shape[0]) audio_state = info_dict.get("audio_state") is_first_call = not isinstance(audio_state, dict) + is_prefill = bool(info_dict.get("_omni_is_prefill", span_len > 1)) - if span_len > 1 or is_first_call: + if is_prefill 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)) + ref_offset = int(info_dict.get("_omni_num_computed_tokens", 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: @@ -994,22 +998,24 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): 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": { + "audio_codes": {"current": current_codes}, + "ref_offset": ref_offset + span_len, + } + if is_first_call: + remaining = (info_dict.get("ids", {}) or {}).get("all") + if not isinstance(remaining, list): + remaining = [] + max_new_frames = info_dict.get("max_new_frames") + if isinstance(max_new_frames, (list, tuple)): + max_new_frames = max_new_frames[0] if max_new_frames else None + info_update["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, - } + "max_new_frames": int(max_new_frames) if max_new_frames is not None else -1, + } return input_ids, embeds, info_update # Decode step: text_token from the just-sampled vLLM logit, plus the @@ -1046,13 +1052,21 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): info_dicts: list[dict[str, Any]] = ( kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") or [] ) + sample_eligible = kwargs.get("request_sample_eligible") + eligible = [ + sample_eligible is None or i >= len(sample_eligible) or bool(sample_eligible[i]) + for i in range(len(info_dicts)) + ] # 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] + self._batch_state = [ + (info["audio_state"] if isinstance(info, dict) and eligible[i] else None) + for i, info in enumerate(info_dicts) + ] self._batch_state_spans = kwargs.get("request_token_spans") # See delay talker: real per-request row spans from the runner are @@ -1074,6 +1088,8 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): for i, info in enumerate(info_dicts): if not isinstance(info, dict): continue + if not eligible[i]: + continue row_start, row_end = spans[i] row_end = min(int(row_end), num_rows) if row_start >= row_end: @@ -1115,10 +1131,13 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): ) ch0 = int(new_codes[0].item()) + state["step"] = int(state.get("step", 0)) + 1 + max_new_frames = int(state.get("max_new_frames", -1)) + if max_new_frames > 0 and state["step"] >= max_new_frames: + state["is_stopping"] = True # 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"), @@ -1127,7 +1146,6 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): 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"), @@ -1141,7 +1159,6 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): updated_acc = new_codes.unsqueeze(0) info["audio_codes"] = {"current": new_codes, "accumulated": updated_acc} - state["step"] = int(state.get("step", 0)) + 1 per_req_codes[i] = new_codes.unsqueeze(0) have_codes = True diff --git a/vllm_omni/model_executor/models/moss_tts/realtime_prompt.py b/vllm_omni/model_executor/models/moss_tts/realtime_prompt.py new file mode 100644 index 000000000..985f76205 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/realtime_prompt.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Prompt construction for MOSS-TTS-Realtime.""" + +from typing import Any + +import numpy as np +import torch + +REALTIME_AUDIO_CHANNELS = 16 +REALTIME_PREFILL_TEXT_TOKENS = 12 + + +def build_realtime_prompt( + tokenizer: Any, + processor: Any, + text: str, + reference_codes: torch.Tensor, +) -> dict[str, Any]: + """Build the mixed text/audio prompt consumed by the Realtime talker.""" + audio_tokens = reference_codes.detach().cpu().numpy() + system_grid = np.asarray(processor.make_ensemble(prompt_audio_tokens=audio_tokens), dtype=np.int64) + expected_width = REALTIME_AUDIO_CHANNELS + 1 + if system_grid.ndim != 2 or system_grid.shape[1] != expected_width: + raise ValueError( + f"MOSS-TTS-Realtime processor returned shape {system_grid.shape}; " + f"expected (sequence_length, {expected_width})" + ) + + assistant_ids = tokenizer.encode("<|im_start|>assistant\n", add_special_tokens=False) + assistant_grid = np.full( + (len(assistant_ids), expected_width), + processor.audio_channel_pad, + dtype=np.int64, + ) + assistant_grid[:, 0] = assistant_ids + + text_ids = tokenizer.encode(text, add_special_tokens=False) + if not text_ids: + raise ValueError("MOSS-TTS-Realtime input is empty after tokenization") + + prefill_len = min(len(text_ids), REALTIME_PREFILL_TEXT_TOKENS) + text_grid = np.full( + (prefill_len, expected_width), + processor.audio_channel_pad, + dtype=np.int64, + ) + text_grid[:, 0] = text_ids[:prefill_len] + text_grid[-1, 1] = processor.audio_bos_token + + prompt_grid = np.concatenate((system_grid, assistant_grid, text_grid), axis=0) + params: dict[str, Any] = { + "prompt_token_ids": prompt_grid[:, 0].tolist(), + "codes": {"ref": torch.from_numpy(prompt_grid[:, 1:].copy())}, + } + remaining_text_ids = text_ids[prefill_len:] + if remaining_text_ids: + params["ids"] = {"all": remaining_text_ids} + return params diff --git a/vllm_omni/model_executor/models/moss_tts/reference_encoder.py b/vllm_omni/model_executor/models/moss_tts/reference_encoder.py index 36b2616e7..a828d59e0 100644 --- a/vllm_omni/model_executor/models/moss_tts/reference_encoder.py +++ b/vllm_omni/model_executor/models/moss_tts/reference_encoder.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Reference-audio encoding + speaker cache for the MOSS-TTS-family talker. This lives in the model package (not the shared serving layer) so all @@ -360,6 +363,18 @@ class MossReferenceEncoder: await self._batcher.aclose() +class MossRealtimeReferenceEncoder(MossReferenceEncoder): + """Use the realtime codec with the shared reference cache and batcher.""" + + def _encode_prepared(self, prepared: list[torch.Tensor]) -> list[torch.Tensor]: + with torch.no_grad(): + encoded = self._processor.batch_encode([wav.squeeze(0) for wav in prepared], num_quantizers=self._n_vq) + return [ + encoded.audio_codes[:, i, : int(length.item())].transpose(0, 1).contiguous() + for i, length in enumerate(encoded.audio_codes_lengths) + ] + + def build_reference_encoder( processor: Any, *, @@ -369,9 +384,17 @@ def build_reference_encoder( """Build the per-server encoder for a MOSS-TTS ``variant``. Derives the encode geometry (``n_vq`` and the working sample rate) from the - upstream processor's ``model_config`` so the variant knowledge stays in the - model package rather than in ``serving_speech.py``. + upstream processor's ``model_config``; Realtime takes its codec directly + and uses 16 codebooks at 24 kHz. Variant knowledge stays in the model package. """ + if variant == "realtime": + return MossRealtimeReferenceEncoder( + processor, + variant=variant, + n_vq=16, + sr_target=24000, + speaker_cache=speaker_cache, + ) n_vq = int(getattr(processor.model_config, "n_vq", 32)) # Local-v1.5 encodes reference audio at a fixed 24 kHz working rate # regardless of its 48 kHz stereo *output* codec -- mirrors the offline