107 lines
3.9 KiB
Python
107 lines
3.9 KiB
Python
|
# SPDX-License-Identifier: Apache-2.0
|
||
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||
|
|
"""E2E: codec-token streaming from a talker-only Qwen3-TTS deployment.
|
||
|
|
|
||
|
|
Deploy Setting: qwen3_tts_talker_only.yaml (single LLM_AR stage)
|
||
|
|
Input Modal: text (WebSocket /v1/audio/speech/stream, output_mode=codec_tokens)
|
||
|
|
Output Modal: codec frames decoded on the client with Qwen3TTSTokenizerV2Decoder
|
||
|
|
"""
|
||
|
|
|
||
|
|
import os
|
||
|
|
|
||
|
|
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||
|
|
|
||
|
|
import json
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from tests.helpers.mark import hardware_test
|
||
|
|
from tests.helpers.runtime import OmniServerParams
|
||
|
|
from tests.helpers.stage_config import get_deploy_config_path
|
||
|
|
from vllm_omni.entrypoints.openai.codec_stream import FLAG_EOS, unpack_codec_frame
|
||
|
|
|
||
|
|
MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice"
|
||
|
|
_STAGE_CONFIG = get_deploy_config_path("qwen3_tts_talker_only.yaml")
|
||
|
|
|
||
|
|
server_params = [
|
||
|
|
pytest.param(
|
||
|
|
OmniServerParams(model=MODEL, stage_config_path=_STAGE_CONFIG, server_args=["--trust-remote-code"]),
|
||
|
|
id="talker_only",
|
||
|
|
)
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def _stream_codec(port: int, text: str) -> tuple[dict, list[tuple[int, int, int, np.ndarray]], dict]:
|
||
|
|
from websockets.sync.client import connect
|
||
|
|
|
||
|
|
frames = []
|
||
|
|
with connect(f"ws://localhost:{port}/v1/audio/speech/stream", max_size=64 * 1024 * 1024) as ws:
|
||
|
|
ws.send(
|
||
|
|
json.dumps(
|
||
|
|
{
|
||
|
|
"type": "session.config",
|
||
|
|
"model": MODEL,
|
||
|
|
"output_mode": "codec_tokens",
|
||
|
|
"stream_audio": True,
|
||
|
|
"response_format": "pcm",
|
||
|
|
"voice": "vivian",
|
||
|
|
"task_type": "CustomVoice",
|
||
|
|
}
|
||
|
|
)
|
||
|
|
)
|
||
|
|
ws.send(json.dumps({"type": "input.text", "text": text}))
|
||
|
|
ws.send(json.dumps({"type": "input.done"}))
|
||
|
|
start = json.loads(ws.recv())
|
||
|
|
assert start["type"] == "codec.start", start
|
||
|
|
done = None
|
||
|
|
while True:
|
||
|
|
raw = ws.recv()
|
||
|
|
if isinstance(raw, (bytes, bytearray)):
|
||
|
|
frames.append(unpack_codec_frame(bytes(raw), start["codebooks"]))
|
||
|
|
continue
|
||
|
|
msg = json.loads(raw)
|
||
|
|
if msg["type"] == "codec.done":
|
||
|
|
done = msg
|
||
|
|
elif msg["type"] == "session.done":
|
||
|
|
break
|
||
|
|
elif msg["type"] == "error":
|
||
|
|
raise AssertionError(msg)
|
||
|
|
ws.send(json.dumps({"type": "session.close"}))
|
||
|
|
return start, frames, done
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.advanced_model
|
||
|
|
@pytest.mark.tts
|
||
|
|
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||
|
|
@pytest.mark.parametrize("omni_server", server_params, indirect=True)
|
||
|
|
def test_codec_token_stream_and_client_decode(omni_server) -> None:
|
||
|
|
start, frames, done = _stream_codec(omni_server.port, "Codec token streaming end to end test.")
|
||
|
|
assert start["codebooks"] == 16 and start["frame_rate_hz"] == 12.5 and start["sample_rate"] == 24000
|
||
|
|
assert frames, "no codec frames received"
|
||
|
|
seqs = [f[0] for f in frames]
|
||
|
|
assert seqs == list(range(len(seqs)))
|
||
|
|
assert frames[-1][2] & FLAG_EOS
|
||
|
|
total = sum(f[1] for f in frames)
|
||
|
|
assert total >= 10, f"too few frames: {total}"
|
||
|
|
assert done is not None and done["total_frames"] == total and done["error"] is False
|
||
|
|
codes = np.concatenate([f[3].astype(np.int64) for f in frames if f[1]], axis=0)
|
||
|
|
assert codes.min() >= 0 and codes.max() < 2048
|
||
|
|
|
||
|
|
import sys
|
||
|
|
|
||
|
|
sys.path.insert(
|
||
|
|
0,
|
||
|
|
str(
|
||
|
|
os.path.join(
|
||
|
|
os.path.dirname(__file__), "..", "..", "..", "examples", "online_serving", "text_to_speech", "qwen3_tts"
|
||
|
|
)
|
||
|
|
),
|
||
|
|
)
|
||
|
|
from edge_decoder_client import LocalDecoder
|
||
|
|
|
||
|
|
dec = LocalDecoder(MODEL, device="cpu")
|
||
|
|
pcm = np.concatenate([dec.decode(codes[i : i + 25]) for i in range(0, codes.shape[0], 25)])
|
||
|
|
expected_s = total / 12.5
|
||
|
|
assert abs(pcm.size / dec.sample_rate - expected_s) < 0.1 * expected_s + 0.2
|
||
|
|
assert float(np.abs(pcm).max()) > 1e-3
|