[BugFix][TTS] Fix MOSS Realtime serving and generation (#5661)

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>

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)

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)

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

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),
+        )