NancyFyongandNancyFyong
e74fe3c6e8
Signed-off-by: NancyFyong <NancyFyong@users.noreply.github.com> Co-authored-by: NancyFyong <NancyFyong@users.noreply.github.com>
113 lines
4.6 KiB
Python
113 lines
4.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""E2E offline tests for Audio8 TTS Preview (DualAR, 44.1 kHz codec).
|
|
|
|
The prompt is pre-tokenized here (not built from ``mm_processor_kwargs``)
|
|
because Audio8 TTS shares Fish Speech's protocol: text-only prompts are plain
|
|
token ids, and voice-clone prompts are a placeholder of the exact final length
|
|
whose real embeddings are built model-side from the encoded reference audio.
|
|
"""
|
|
|
|
import os
|
|
|
|
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
|
|
|
import pytest
|
|
|
|
from tests.helpers.mark import hardware_test
|
|
from tests.helpers.runtime import OmniRunner, OmniRunnerHandler
|
|
from tests.helpers.stage_config import get_deploy_config_path
|
|
|
|
MODEL = os.environ.get("AUDIO8_TTS_MODEL_PATH", "Audio8/Audio8-TTS-Preview-0.6b")
|
|
DEPLOY_CONFIG = get_deploy_config_path("audio8_tts.yaml")
|
|
SAMPLE_RATE = 44100
|
|
TEXT = "The weather is nice today, perfect for a walk in the park."
|
|
|
|
# Four distinct inputs for the concurrency test: distinct text is required so a
|
|
# leaked per-request codec buffer (which would make two requests decode to the
|
|
# same audio) actually fails the pairwise-difference assertion.
|
|
CONCURRENT_TEXTS = [
|
|
"The weather is nice today, perfect for a walk in the park.",
|
|
"She sold seashells by the seashore all through the summer.",
|
|
"Our train departs at a quarter past nine tomorrow morning.",
|
|
"He planted rows of tomatoes and basil behind the old house.",
|
|
]
|
|
|
|
# The checkpoint registers its own `arktts` config in vllm-omni, so remote code
|
|
# must stay disabled (transformers would otherwise prefer the checkpoint's
|
|
# auto_map and bypass the Qwen2-shaped view Qwen2Model needs).
|
|
_OMNI_RUNNER_PARAM = (MODEL, DEPLOY_CONFIG, {"trust_remote_code": False})
|
|
|
|
pytestmark = pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True)
|
|
|
|
|
|
def _text_only_prompt(text: str = TEXT) -> dict:
|
|
from transformers import AutoTokenizer
|
|
|
|
from vllm_omni.model_executor.models.audio8_tts.prompt_utils import build_text_only_prompt_ids
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL)
|
|
prompt_ids, normalized_text = build_text_only_prompt_ids(tokenizer, text)
|
|
return {"prompt_token_ids": prompt_ids, "additional_information": {"text": [normalized_text]}}
|
|
|
|
|
|
# At ``core_model`` (L1/L2) the harness patches every stage to
|
|
# ``load_format: dummy``, so the AR runs on random weights: it never emits EOS
|
|
# and always generates to ``max_tokens`` (1024 frames ~= 47.5 s). Duration
|
|
# bounds are therefore only meaningful once real weights load, at
|
|
# ``advanced_model`` / ``full_model``.
|
|
_REAL_WEIGHT_RUN_LEVELS = frozenset({"advanced_model", "full_model"})
|
|
|
|
|
|
def _duration_bounds(run_level: str) -> dict[str, float]:
|
|
if run_level not in _REAL_WEIGHT_RUN_LEVELS:
|
|
return {}
|
|
return {"min_duration_s": 0.5, "max_duration_s": 30.0}
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.advanced_model
|
|
@pytest.mark.tts
|
|
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
|
def test_text_to_audio_001(omni_runner: OmniRunner, run_level: str) -> None:
|
|
"""Default deploy, single request, text-only synthesis.
|
|
|
|
Deploy Setting: audio8_tts.yaml
|
|
Input Modal: text
|
|
Output Modal: audio
|
|
"""
|
|
prompt = _text_only_prompt()
|
|
request_config = {
|
|
"input": TEXT,
|
|
"prompt_token_ids": prompt["prompt_token_ids"],
|
|
"additional_information": prompt["additional_information"],
|
|
"response_format": "wav",
|
|
"expected_sample_rate": SAMPLE_RATE,
|
|
**_duration_bounds(run_level),
|
|
}
|
|
OmniRunnerHandler(omni_runner).send_tokenized_tts_request(request_config)
|
|
|
|
|
|
@pytest.mark.advanced_model
|
|
@pytest.mark.tts
|
|
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
|
def test_text_to_audio_concurrent_002(omni_runner: OmniRunner, run_level: str) -> None:
|
|
"""Four concurrent requests with distinct text: stage 0 runs ``max_num_seqs: 4``.
|
|
|
|
Per-request codec state is keyed by request id. The four inputs are
|
|
deliberately distinct and the decoded outputs are asserted pairwise
|
|
different, so leaked codec state (which would make two requests decode to
|
|
identical audio) fails here, and only here.
|
|
"""
|
|
prompts = [_text_only_prompt(text) for text in CONCURRENT_TEXTS]
|
|
request_config = {
|
|
"input": CONCURRENT_TEXTS[0],
|
|
"prompt_token_ids": [p["prompt_token_ids"] for p in prompts],
|
|
"additional_information": [p["additional_information"] for p in prompts],
|
|
"response_format": "wav",
|
|
"expected_sample_rate": SAMPLE_RATE,
|
|
"assert_distinct_outputs": True,
|
|
**_duration_bounds(run_level),
|
|
}
|
|
OmniRunnerHandler(omni_runner).send_tokenized_tts_request(request_config)
|