vllm-omni/tests/e2e/offline_inference/test_ming_tts.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

"""End-to-end offline inference tests for Ming-omni-tts."""

import asyncio
import uuid

import numpy as np
import pytest
import torch
from transformers import AutoTokenizer
from vllm import SamplingParams
from vllm.inputs import tokens_input

from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni import AsyncOmni
from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
    KEY_MAX_DECODE_STEPS,
    SAMPLE_RATE,
    TEXT_EOS_TOKEN_ID,
)
from vllm_omni.model_executor.models.ming_tts.prompt_assembly import DEFAULT_PROMPT, build_ming_dense_prompt

MODEL = "inclusionAI/Ming-omni-tts-0.5B"
DEPLOY_CONFIG = get_deploy_config_path("ming_tts.yaml")
TEST_TEXT = "我会一直在这里陪着你，直到你慢慢地沉入那个最温柔的梦里。"
TEST_INSTRUCTION = "轻柔的ASMR耳语，慢速，贴近麦克风"
MIN_AUDIO_SAMPLES = 1000


@pytest.fixture(scope="module")
def ming_tokenizer():
    return AutoTokenizer.from_pretrained(MODEL, trust_remote_code=False)


@pytest.fixture
def ming_engine():
    with OmniRunner(
        MODEL,
        deploy_config=DEPLOY_CONFIG,
        stage_init_timeout=300,
        enforce_eager=True,
    ) as runner:
        yield runner.omni


@pytest.fixture
def async_omni_engine():
    engine = AsyncOmni(
        model=MODEL,
        deploy_config=DEPLOY_CONFIG,
        stage_init_timeout=300,
        enforce_eager=True,
    )
    yield engine
    engine.shutdown()


def _build_prompt(
    tokenizer,
    *,
    text: str = TEST_TEXT,
    instruction=TEST_INSTRUCTION,
    use_zero_spk_emb: bool = True,
) -> dict:
    prompt_dict = build_ming_dense_prompt(
        tokenizer,
        prompt=DEFAULT_PROMPT,
        text=text,
        instruction=instruction,
        runtime_controls={KEY_MAX_DECODE_STEPS: 200},
        use_zero_spk_emb=use_zero_spk_emb,
    )
    prompt = tokens_input(prompt_token_ids=prompt_dict["prompt_token_ids"])
    prompt["prompt"] = prompt_dict["prompt"]
    prompt["text"] = prompt_dict["text"]
    prompt["additional_information"] = prompt_dict["additional_information"]
    return prompt


def _sampling_params_list() -> list[SamplingParams]:
    return [
        SamplingParams(
            temperature=0.0,
            max_tokens=201,
            stop_token_ids=[int(TEXT_EOS_TOKEN_ID)],
        ),
        SamplingParams(temperature=0.0, max_tokens=1),
    ]


def _flatten_audio(audio) -> torch.Tensor:
    if isinstance(audio, list):
        parts = [torch.as_tensor(item, dtype=torch.float32).reshape(-1).cpu() for item in audio]
        parts = [item for item in parts if item.numel() > 0]
        if not parts:
            return torch.zeros((0,), dtype=torch.float32)
        return torch.cat(parts, dim=0)
    return torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu()


def _extract_audio(multimodal_output: dict) -> torch.Tensor:
    assert isinstance(multimodal_output, dict), f"Expected dict, got {type(multimodal_output)}"
    audio = multimodal_output.get("audio", multimodal_output.get("model_outputs"))
    assert audio is not None, f"No audio output found, keys={list(multimodal_output.keys())}"
    waveform = _flatten_audio(audio)
    if waveform.numel() == 0:
        raise RuntimeError("Generated audio waveform is empty")
    return waveform


def _extract_sample_rate(multimodal_output: dict) -> int:
    sample_rate = multimodal_output.get("sr")
    if sample_rate is None:
        raise RuntimeError("Expected multimodal_output['sr']")
    if isinstance(sample_rate, list):
        sample_rate = sample_rate[-1]
    if hasattr(sample_rate, "item"):
        sample_rate = sample_rate.item()
    return int(sample_rate)


def _extract_final_audio_outputs(outputs):
    final_outputs = []
    for item in outputs:
        if getattr(item, "final_output_type", None) == "audio":
            final_outputs.append(item)
            continue
        request_output = getattr(item, "request_output", None)
        if request_output is None:
            continue
        multimodal_output = getattr(request_output, "multimodal_output", None)
        if isinstance(multimodal_output, dict):
            final_outputs.append(item)
            continue
        completions = getattr(request_output, "outputs", None) or []
        if any(isinstance(getattr(completion, "multimodal_output", None), dict) for completion in completions):
            final_outputs.append(item)
    return final_outputs


def _extract_multimodal_output(output) -> dict:
    multimodal_output = getattr(output, "multimodal_output", None)
    if isinstance(multimodal_output, dict):
        return multimodal_output

    request_output = getattr(output, "request_output", None)
    if request_output is not None:
        multimodal_output = getattr(request_output, "multimodal_output", None)
        if isinstance(multimodal_output, dict):
            return multimodal_output
        completions = getattr(request_output, "outputs", None) or []
        for completion in completions:
            multimodal_output = getattr(completion, "multimodal_output", None)
            if isinstance(multimodal_output, dict):
                return multimodal_output

    raise AssertionError("No multimodal audio output found in Ming generate results")


@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_ming_tts_offline_basic(ming_engine, ming_tokenizer) -> None:
    """Test blocking Ming generation through Omni."""
    outputs = ming_engine.generate(
        prompts=[_build_prompt(ming_tokenizer)],
        sampling_params_list=_sampling_params_list(),
        py_generator=False,
    )
    final_outputs = _extract_final_audio_outputs(outputs)
    assert len(final_outputs) == 1, f"Expected one final audio output, got {len(final_outputs)}"
    multimodal_output = _extract_multimodal_output(final_outputs[0])
    waveform = _extract_audio(multimodal_output)
    sample_rate = _extract_sample_rate(multimodal_output)
    assert waveform.ndim == 1
    assert waveform.shape[0] == waveform.numel()
    assert waveform.numel() > MIN_AUDIO_SAMPLES
    assert np.max(np.abs(waveform.numpy())) > 0.01, "Audio appears silent"
    assert sample_rate == SAMPLE_RATE, f"Expected Ming output sample rate {SAMPLE_RATE}, got {sample_rate}"


@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_ming_tts_speaker_conditioning_differs(ming_engine, ming_tokenizer) -> None:
    """Test that different Ming speaker controls produce different waveform outputs."""
    style_outputs = ming_engine.generate(
        prompts=[_build_prompt(ming_tokenizer)],
        sampling_params_list=_sampling_params_list(),
        py_generator=False,
    )
    ip_outputs = ming_engine.generate(
        prompts=[_build_prompt(ming_tokenizer, text=TEST_TEXT, instruction={"IP": "灵小甄"}, use_zero_spk_emb=True)],
        sampling_params_list=_sampling_params_list(),
        py_generator=False,
    )

    style_final_outputs = _extract_final_audio_outputs(style_outputs)
    ip_final_outputs = _extract_final_audio_outputs(ip_outputs)
    assert len(style_final_outputs) == 1, "No style audio output produced"
    assert len(ip_final_outputs) == 1, "No IP audio output produced"

    style_waveform = _extract_audio(_extract_multimodal_output(style_final_outputs[0]))
    ip_waveform = _extract_audio(_extract_multimodal_output(ip_final_outputs[0]))
    assert style_waveform.numel() > MIN_AUDIO_SAMPLES
    assert ip_waveform.numel() > MIN_AUDIO_SAMPLES
    assert np.max(np.abs(style_waveform.numpy())) > 0.01, "Style audio appears silent"
    assert np.max(np.abs(ip_waveform.numpy())) > 0.01, "IP audio appears silent"

    overlap = min(int(style_waveform.numel()), int(ip_waveform.numel()))
    mean_abs_diff = torch.mean(torch.abs(style_waveform[:overlap] - ip_waveform[:overlap])).item()
    assert style_waveform.shape != ip_waveform.shape or mean_abs_diff > 1e-4, (
        "Speaker-conditioned outputs should differ, but style and IP waveforms were effectively identical"
    )


@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_ming_tts_multiple_prompts_queued(ming_engine, ming_tokenizer) -> None:
    """Regression: supported max_num_seqs=1 config must still drain queued prompts."""
    prompts = [
        _build_prompt(
            ming_tokenizer,
            text="第一条语音用于验证 Ming dense 队列中的长请求可以完成。",
            instruction="平静自然的旁白，语速中等",
        ),
        _build_prompt(
            ming_tokenizer,
            text="第二条短语音也必须生成完成。",
            instruction="清晰明亮的女声",
        ),
    ]
    outputs = ming_engine.generate(
        prompts=prompts,
        sampling_params_list=_sampling_params_list(),
        py_generator=False,
    )
    final_outputs = _extract_final_audio_outputs(outputs)
    assert len(final_outputs) == len(prompts), f"Expected {len(prompts)} audio outputs, got {len(final_outputs)}"
    for i, output in enumerate(final_outputs):
        waveform = _extract_audio(_extract_multimodal_output(output))
        duration_s = waveform.numel() / SAMPLE_RATE
        assert 0.1 < duration_s < 30.0, f"Request {i} audio duration out of range: {duration_s:.2f}s"
        assert np.max(np.abs(waveform.numpy())) > 0.01, f"Request {i} audio appears silent"


@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_ming_tts_offline_streaming(async_omni_engine, ming_tokenizer) -> None:
    """Test async_chunk streaming Ming generation through AsyncOmni."""

    async def _run() -> None:
        all_audio_chunks = []
        accumulated_samples = 0
        chunk_idx = 0
        sample_rate = None
        async for stage_output in async_omni_engine.generate(
            prompt=_build_prompt(ming_tokenizer),
            request_id=str(uuid.uuid4()),
            sampling_params_list=_sampling_params_list(),
        ):
            multimodal_output = stage_output.multimodal_output or {}
            audio = multimodal_output.get("audio")
            if "sr" in multimodal_output:
                sample_rate = _extract_sample_rate(multimodal_output)
            if audio is None:
                continue
            finished = stage_output.finished
            if isinstance(audio, torch.Tensor):
                if finished:
                    audio_chunk = audio[accumulated_samples:].float().detach().cpu()
                else:
                    audio_chunk = audio.float().detach().cpu()
            elif isinstance(audio, list):
                audio_chunk = _flatten_audio(audio)
            else:
                audio_chunk = torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu()
            accumulated_samples += int(audio_chunk.numel())
            chunk_idx += 1
            if audio_chunk.numel() > 0:
                all_audio_chunks.append(audio_chunk)
        assert all_audio_chunks, "No streaming audio chunks received"
        waveform = torch.cat(all_audio_chunks, dim=0)
        assert waveform.numel() > MIN_AUDIO_SAMPLES
        assert np.max(np.abs(waveform.numpy())) > 0.01, "Audio appears silent"
        assert sample_rate is not None, "Streaming path did not return a sample rate"
        assert sample_rate == SAMPLE_RATE, f"Expected Ming output sample rate {SAMPLE_RATE}, got {sample_rate}"

    asyncio.run(_run())