[Perf] Stacked audio ops: batched matmul and embedding gather in MOSS-TTS talker (#4230)
Signed-off-by: leo.yang <leo.yang.engineer@gmail.com>
This commit is contained in:
yangyonggit authored and GitHub committed
1 parent
37a442447e
commit
63ba11b8b1
3 files changed
+534
-166
No files matched your search
@@ -1,77 +1,97 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""E2E offline inference tests for MOSS-TTS two-stage pipeline.
|
||||
"""E2E offline inference tests for MOSS-VoiceGenerator (MossTTSDelayModel, 1.7B).
|
||||
|
||||
Tests the MossTTSDelayModel path using MOSS-VoiceGenerator (1.7B) — the
|
||||
smallest MossTTSDelayModel variant — so the test runs on a single A10G or L4
|
||||
without requiring an 80 GB GPU.
|
||||
Uses the standard omni_runner + pytestmark pattern (one module-scoped engine
|
||||
per file, skill invariant I4). The realtime variant lives in the sibling file
|
||||
test_moss_tts_realtime.py.
|
||||
|
||||
A second test class covers MOSS-TTS-Realtime (1.7B, MossTTSRealtime) to
|
||||
exercise the local-transformer generation path.
|
||||
MOSS-VoiceGenerator synthesizes speech from a text *instruction* that describes
|
||||
the desired voice (e.g. "a warm female voice with an American accent"). It does
|
||||
NOT accept a reference audio clip. The request format requires running
|
||||
AutoProcessor.from_pretrained once per call to encode the (text, instruction)
|
||||
pair into the (prompt_token_ids, codes.ref) grid that the MossTTSDelayModel
|
||||
talker expects — same as examples/offline_inference/text_to_speech/moss_tts/
|
||||
end2end.py:_build_unified_codes.
|
||||
|
||||
Both classes use a module-scoped engine so the model is loaded once per test
|
||||
module execution. Do NOT add a second engine fixture in this file — that
|
||||
triggers mid-module teardown races (see skill invariant I4).
|
||||
No determinism test: VoiceGenerator produces variable-length output even with
|
||||
a fixed seed; waveform length reproducibility is not guaranteed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gc
|
||||
import os
|
||||
import urllib.request
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm import SamplingParams
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from vllm_omni import Omni
|
||||
from tests.helpers.runtime import OmniRunner
|
||||
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Shared constants
|
||||
# Constants
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Selected by the nightly "TTS · Function Test" (-m "full_model and L4 and tts").
|
||||
# Without the tts + run-level marks these tests are collected but never run.
|
||||
# TODO: Fix this test
|
||||
# pytestmark = [pytest.mark.full_model, pytest.mark.tts]
|
||||
MODEL = "OpenMOSS-Team/MOSS-VoiceGenerator"
|
||||
SAMPLE_RATE = 24_000
|
||||
|
||||
SAMPLE_RATE = 24_000 # All MOSS-TTS full variants output 24 kHz
|
||||
|
||||
REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/reference_zh_1.wav"
|
||||
|
||||
_DEFAULT_SAMPLING = SamplingParams(
|
||||
# Stage 0 = AR talker; max_tokens=512 keeps tests fast (~20 s of audio).
|
||||
# Stage 1 = codec decoder; greedy, large context to hold all codec tokens.
|
||||
_STAGE0_PARAMS = SamplingParams(
|
||||
temperature=1.7,
|
||||
top_p=0.8,
|
||||
top_k=25,
|
||||
max_tokens=512, # short for tests (~20 s at 24 kHz, ~12.5 tokens/s)
|
||||
max_tokens=512,
|
||||
seed=42,
|
||||
detokenize=False,
|
||||
)
|
||||
_STAGE1_PARAMS = SamplingParams(
|
||||
temperature=0.0,
|
||||
top_p=1.0,
|
||||
top_k=-1,
|
||||
max_tokens=65536,
|
||||
seed=42,
|
||||
detokenize=False,
|
||||
)
|
||||
_SAMPLING = [_STAGE0_PARAMS, _STAGE1_PARAMS]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Session fixtures
|
||||
# Deploy config
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def ref_audio_path(tmp_path_factory) -> str:
|
||||
"""Download the upstream reference clip once per session.
|
||||
def _get_test_config() -> str:
|
||||
"""Derive a CI-friendly config from moss_voice_generator.yaml.
|
||||
|
||||
Set ``MOSS_TTS_SKIP_ON_NET_FAIL=1`` to skip in air-gapped environments.
|
||||
Reduces Stage 0 gpu_memory_utilization from 0.60 → 0.45 to leave headroom
|
||||
for Stage 1 (0.12) on a shared L4/A10G. max_num_seqs=1 keeps per-test
|
||||
peak memory predictable.
|
||||
"""
|
||||
cache_dir = tmp_path_factory.mktemp("moss_tts_ref")
|
||||
target = cache_dir / "zh_1.wav"
|
||||
try:
|
||||
with urllib.request.urlopen(REF_AUDIO_URL, timeout=30) as resp:
|
||||
target.write_bytes(resp.read())
|
||||
except Exception as exc:
|
||||
msg = f"Cannot fetch reference clip {REF_AUDIO_URL}: {exc}"
|
||||
if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"):
|
||||
pytest.skip(msg)
|
||||
pytest.fail(msg)
|
||||
if not target.exists() or target.stat().st_size == 0:
|
||||
pytest.fail(f"Reference clip empty after download: {target}")
|
||||
return str(target)
|
||||
return modify_stage_config(
|
||||
get_deploy_config_path("moss_voice_generator.yaml"),
|
||||
updates={
|
||||
"stages": {
|
||||
0: {
|
||||
"max_num_seqs": 1,
|
||||
"gpu_memory_utilization": 0.45,
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# pytestmark — one engine for the whole module
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.full_model,
|
||||
pytest.mark.tts,
|
||||
pytest.mark.parametrize("omni_runner", [(MODEL, _get_test_config())], indirect=True),
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -79,147 +99,142 @@ def ref_audio_path(tmp_path_factory) -> str:
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _build_request(text: str, ref_audio_path: str, seed: int = 42) -> dict:
|
||||
def _build_request(text: str, instruction: str) -> dict:
|
||||
"""Build a VoiceGenerator request using AutoProcessor.
|
||||
|
||||
Calls proc.build_user_message + proc() to produce the unified-codes grid
|
||||
(L, 1+n_vq) and splits it into prompt_token_ids + codes.ref, exactly as
|
||||
end2end.py:_build_unified_codes does for the voice_generator variant.
|
||||
The processor (including its CPU-resident audio_tokenizer) is freed before
|
||||
returning so it does not compete with the running engine.
|
||||
|
||||
Set MOSS_TTS_SKIP_ON_NET_FAIL=1 to skip in air-gapped environments where
|
||||
the HF snapshot is not cached.
|
||||
"""
|
||||
from transformers import AutoProcessor
|
||||
|
||||
try:
|
||||
proc = AutoProcessor.from_pretrained(MODEL, trust_remote_code=True)
|
||||
except Exception as exc:
|
||||
msg = f"Cannot load AutoProcessor for {MODEL}: {exc}"
|
||||
if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"):
|
||||
pytest.skip(msg)
|
||||
pytest.fail(msg)
|
||||
|
||||
user_msg = proc.build_user_message(text=text, instruction=instruction)
|
||||
batch = proc(conversations=[[user_msg]], mode="generation")
|
||||
unified = batch["input_ids"][0] # (L, 1+n_vq)
|
||||
text_ids = unified[:, 0].tolist()
|
||||
audio_codes = unified[:, 1:].contiguous().to(torch.int64) # (L, n_vq)
|
||||
del proc
|
||||
gc.collect()
|
||||
|
||||
return {
|
||||
"prompt": "<|im_start|>assistant\n",
|
||||
"additional_information": {
|
||||
"task_type": ["voice_clone"],
|
||||
"text": [text],
|
||||
"mode": ["voice_clone"],
|
||||
"prompt_audio_path": [ref_audio_path],
|
||||
"seed": [seed],
|
||||
},
|
||||
"prompt_token_ids": text_ids,
|
||||
"additional_information": {"codes": {"ref": audio_codes}},
|
||||
}
|
||||
|
||||
|
||||
def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]:
|
||||
"""Run one request; return (waveform_cpu, sample_rate)."""
|
||||
for stage_outputs in omni.generate(request, _DEFAULT_SAMPLING):
|
||||
for req_output in stage_outputs.request_output:
|
||||
mm = req_output.outputs[0].multimodal_output
|
||||
assert mm is not None, "Expected non-None multimodal_output"
|
||||
# Never use ``a or b`` on tensor values — a multi-element tensor
|
||||
# raises on truthiness. Check ``is None`` explicitly and handle the
|
||||
# pre-consolidation list form.
|
||||
audio = mm.get("audio")
|
||||
if audio is None:
|
||||
audio = mm.get("model_outputs")
|
||||
sr = mm.get("sr")
|
||||
assert audio is not None, "Expected 'audio' or 'model_outputs' in multimodal_output"
|
||||
if isinstance(audio, list):
|
||||
audio = torch.cat(
|
||||
[t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0],
|
||||
dim=0,
|
||||
)
|
||||
assert isinstance(audio, torch.Tensor), f"Expected Tensor, got {type(audio)}"
|
||||
return audio.reshape(-1).cpu(), int(sr.item()) if sr is not None else SAMPLE_RATE
|
||||
raise AssertionError("No stage outputs received")
|
||||
def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor, int]:
|
||||
"""Run one request and return (waveform_cpu, sample_rate).
|
||||
|
||||
Iterates all OmniRequestOutputs from generate() via the top-level
|
||||
multimodal_output property (which aggregates from completion outputs).
|
||||
AR-stage outputs have empty multimodal_output and are silently skipped.
|
||||
With async_chunk=True and FINAL_ONLY, the codec stage consolidates its
|
||||
chunks into a single tensor before yielding — but we handle list fallback
|
||||
for robustness.
|
||||
"""
|
||||
chunks: list[torch.Tensor] = []
|
||||
sr_final = SAMPLE_RATE
|
||||
for out in omni_runner.omni.generate(request, _SAMPLING):
|
||||
mm = out.multimodal_output
|
||||
if not mm:
|
||||
continue
|
||||
audio = mm.get("audio")
|
||||
if audio is None:
|
||||
audio = mm.get("model_outputs")
|
||||
if audio is None:
|
||||
continue
|
||||
sr = mm.get("sr")
|
||||
if sr is not None:
|
||||
sr_final = int(sr.item() if hasattr(sr, "item") else sr)
|
||||
if isinstance(audio, list):
|
||||
audio = torch.cat(
|
||||
[t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0],
|
||||
dim=0,
|
||||
)
|
||||
if isinstance(audio, torch.Tensor) and audio.numel() > 0:
|
||||
chunks.append(audio.reshape(-1).cpu())
|
||||
if not chunks:
|
||||
raise AssertionError("No audio output received from generate()")
|
||||
return torch.cat(chunks, dim=0), sr_final
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# MossTTSDelayModel — MOSS-VoiceGenerator 1.7B
|
||||
# (tests the delay-pattern generation path on a small GPU)
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_DELAY_MODEL = "OpenMOSS-Team/MOSS-VoiceGenerator"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def delay_engine():
|
||||
return Omni(model=_DELAY_MODEL, stage_init_timeout=300)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"})
|
||||
def test_moss_tts_delay_english(delay_engine, ref_audio_path):
|
||||
"""MossTTSDelayModel: English voice_clone produces non-empty 24 kHz audio."""
|
||||
req = _build_request("Hello, this is a MOSS-TTS voice cloning test.", ref_audio_path)
|
||||
audio, sr = _collect_audio(delay_engine, req)
|
||||
@pytest.mark.advanced_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_moss_tts_english(omni_runner: OmniRunner) -> None:
|
||||
"""VoiceGenerator: English instruction produces non-empty 24 kHz audio."""
|
||||
req = _build_request(
|
||||
"Hello, this is a MOSS voice design test.",
|
||||
"a warm female voice with an American accent",
|
||||
)
|
||||
audio, sr = _collect_audio(omni_runner, req)
|
||||
|
||||
assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}"
|
||||
assert audio.numel() > 0, "Audio tensor is empty"
|
||||
assert not torch.all(audio == 0), "Audio is silence"
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"})
|
||||
def test_moss_tts_delay_chinese(delay_engine, ref_audio_path):
|
||||
"""MossTTSDelayModel: Chinese input produces non-empty audio."""
|
||||
req = _build_request("你好，这是语音合成测试。", ref_audio_path)
|
||||
audio, sr = _collect_audio(delay_engine, req)
|
||||
@pytest.mark.advanced_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_moss_tts_chinese(omni_runner: OmniRunner) -> None:
|
||||
"""VoiceGenerator: Chinese input produces non-empty audio."""
|
||||
req = _build_request(
|
||||
"你好，这是语音合成测试。",
|
||||
"一个清晰温暖的女声",
|
||||
)
|
||||
audio, sr = _collect_audio(omni_runner, req)
|
||||
|
||||
assert sr == SAMPLE_RATE
|
||||
assert audio.numel() > 0
|
||||
assert not torch.all(audio == 0)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"})
|
||||
def test_moss_tts_delay_deterministic(delay_engine, ref_audio_path):
|
||||
"""MossTTSDelayModel: same seed yields identical waveforms."""
|
||||
req = _build_request("Reproducibility check.", ref_audio_path, seed=99)
|
||||
audio1, _ = _collect_audio(delay_engine, req)
|
||||
audio2, _ = _collect_audio(delay_engine, req)
|
||||
|
||||
assert audio1.shape == audio2.shape, "Shapes differ across identical seeds"
|
||||
assert torch.allclose(audio1, audio2, atol=1e-4), "Waveforms differ across identical seeds"
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"})
|
||||
def test_moss_tts_delay_batch(delay_engine, ref_audio_path):
|
||||
"""MossTTSDelayModel: batch of two requests each returns non-empty audio."""
|
||||
@pytest.mark.advanced_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_moss_tts_batch(omni_runner: OmniRunner) -> None:
|
||||
"""VoiceGenerator: batch of two requests each returns non-empty audio."""
|
||||
requests = [
|
||||
_build_request("First sentence.", ref_audio_path),
|
||||
_build_request("Second sentence.", ref_audio_path),
|
||||
_build_request("First sentence.", "a warm female voice"),
|
||||
_build_request("Second sentence.", "a young male voice"),
|
||||
]
|
||||
results: list[torch.Tensor] = []
|
||||
for stage_outputs in delay_engine.generate(requests, [_DEFAULT_SAMPLING] * 2):
|
||||
for req_output in stage_outputs.request_output:
|
||||
mm = req_output.outputs[0].multimodal_output
|
||||
assert mm is not None
|
||||
audio = mm.get("audio")
|
||||
if audio is None:
|
||||
audio = mm.get("model_outputs")
|
||||
assert audio is not None
|
||||
if isinstance(audio, list):
|
||||
audio = torch.cat(
|
||||
[t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0],
|
||||
dim=0,
|
||||
)
|
||||
for out in omni_runner.omni.generate(requests, _SAMPLING):
|
||||
mm = out.multimodal_output
|
||||
if not mm:
|
||||
continue
|
||||
audio = mm.get("audio")
|
||||
if audio is None:
|
||||
audio = mm.get("model_outputs")
|
||||
if audio is None:
|
||||
continue
|
||||
if isinstance(audio, list):
|
||||
audio = torch.cat(
|
||||
[t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0],
|
||||
dim=0,
|
||||
)
|
||||
if isinstance(audio, torch.Tensor) and audio.numel() > 0:
|
||||
results.append(audio.reshape(-1).cpu())
|
||||
|
||||
assert len(results) == 2, f"Expected 2 outputs, got {len(results)}"
|
||||
# With async_chunk streaming + FINAL_ONLY, each request produces one
|
||||
# consolidated output. The >= 2 guard handles backward-compatibility if
|
||||
# consolidation behavior changes.
|
||||
assert len(results) >= 2, f"Expected at least 2 audio outputs, got {len(results)}"
|
||||
for i, audio in enumerate(results):
|
||||
assert audio.numel() > 0, f"Audio[{i}] is empty"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# MossTTSRealtime — MOSS-TTS-Realtime 1.7B
|
||||
# (tests the local-transformer generation path)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_REALTIME_MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def realtime_engine():
|
||||
return Omni(model=_REALTIME_MODEL, stage_init_timeout=300)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"})
|
||||
def test_moss_tts_realtime_english(realtime_engine, ref_audio_path):
|
||||
"""MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio."""
|
||||
req = _build_request("This is a real-time TTS streaming test.", ref_audio_path)
|
||||
audio, sr = _collect_audio(realtime_engine, req)
|
||||
|
||||
assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}"
|
||||
assert audio.numel() > 0, "Audio tensor is empty"
|
||||
assert not torch.all(audio == 0), "Audio is silence"
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"})
|
||||
def test_moss_tts_realtime_deterministic(realtime_engine, ref_audio_path):
|
||||
"""MossTTSRealtime: same seed yields identical waveforms."""
|
||||
req = _build_request("Determinism check for realtime model.", ref_audio_path, seed=7)
|
||||
audio1, _ = _collect_audio(realtime_engine, req)
|
||||
audio2, _ = _collect_audio(realtime_engine, req)
|
||||
|
||||
assert audio1.shape == audio2.shape
|
||||
assert torch.allclose(audio1, audio2, atol=1e-4)
|
||||
assert audio.numel() > 0, f"Audio chunk[{i}] is empty"
|
||||
@@ -0,0 +1,290 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""E2E offline inference tests for MOSS-TTS-Realtime (MossTTSRealtime, 1.7B).
|
||||
|
||||
Uses the standard omni_runner + pytestmark pattern (one module-scoped engine
|
||||
per file, skill invariant I4). The delay-model variant lives in the sibling
|
||||
file test_moss_tts.py.
|
||||
|
||||
The request format mirrors end2end.py:_build_realtime_prompt exactly:
|
||||
1. snapshot_download to locate the MossTTSRealtimeProcessor module.
|
||||
2. MOSS-Audio-Tokenizer (separate HF repo) to encode the reference clip.
|
||||
3. Build the (L, 17) int grid; split into prompt_token_ids + codes.ref +
|
||||
ids.all (remaining text tokens fed one-per-step during decode).
|
||||
|
||||
Set MOSS_TTS_SKIP_ON_NET_FAIL=1 to skip in air-gapped environments.
|
||||
|
||||
No determinism test: MossTTSRealtime uses a local depth transformer that is
|
||||
sensitive to async scheduling timing; waveform reproducibility is not
|
||||
guaranteed across back-to-back calls.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gc
|
||||
import importlib.util
|
||||
import os
|
||||
import sys
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import soundfile as sf
|
||||
import torch
|
||||
from vllm import SamplingParams
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniRunner
|
||||
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Constants
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime"
|
||||
SAMPLE_RATE = 24_000
|
||||
|
||||
REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/reference_zh_1.wav"
|
||||
|
||||
_STAGE0_PARAMS = SamplingParams(
|
||||
temperature=1.0,
|
||||
top_p=0.95,
|
||||
top_k=50,
|
||||
max_tokens=512,
|
||||
seed=42,
|
||||
detokenize=False,
|
||||
)
|
||||
_STAGE1_PARAMS = SamplingParams(
|
||||
temperature=0.0,
|
||||
top_p=1.0,
|
||||
top_k=-1,
|
||||
max_tokens=65536,
|
||||
seed=42,
|
||||
detokenize=False,
|
||||
)
|
||||
_SAMPLING = [_STAGE0_PARAMS, _STAGE1_PARAMS]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Deploy config
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _get_test_config() -> str:
|
||||
"""Derive a CI-friendly config from moss_tts_realtime.yaml.
|
||||
|
||||
Reduces Stage 0 gpu_memory_utilization from 0.60 → 0.45 to leave headroom
|
||||
for Stage 1 (0.12) on a shared L4/A10G. max_num_seqs=1 keeps per-test
|
||||
peak memory predictable.
|
||||
"""
|
||||
return modify_stage_config(
|
||||
get_deploy_config_path("moss_tts_realtime.yaml"),
|
||||
updates={
|
||||
"stages": {
|
||||
0: {
|
||||
"max_num_seqs": 1,
|
||||
"gpu_memory_utilization": 0.45,
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# pytestmark — one engine for the whole module
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.full_model,
|
||||
pytest.mark.tts,
|
||||
pytest.mark.parametrize("omni_runner", [(MODEL, _get_test_config())], indirect=True),
|
||||
]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def ref_audio_path(tmp_path_factory) -> str:
|
||||
"""Download the upstream reference clip once per test module.
|
||||
|
||||
Set MOSS_TTS_SKIP_ON_NET_FAIL=1 to skip in air-gapped environments.
|
||||
"""
|
||||
target = tmp_path_factory.mktemp("moss_tts_realtime_ref") / "zh_1.wav"
|
||||
try:
|
||||
with urllib.request.urlopen(REF_AUDIO_URL, timeout=30) as resp:
|
||||
target.write_bytes(resp.read())
|
||||
except Exception as exc:
|
||||
msg = f"Cannot fetch reference clip {REF_AUDIO_URL}: {exc}"
|
||||
if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"):
|
||||
pytest.skip(msg)
|
||||
pytest.fail(msg)
|
||||
if not target.exists() or target.stat().st_size == 0:
|
||||
pytest.fail(f"Reference clip empty after download: {target}")
|
||||
return str(target)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _load_ref_audio(path: str, target_sr: int = 24000) -> torch.Tensor:
|
||||
"""Load and resample reference audio to (1, T) float32 at target_sr."""
|
||||
data, sr = sf.read(path, always_2d=True) # (T, C)
|
||||
wav = torch.from_numpy(data.T.astype("float32")) # (C, T)
|
||||
if wav.shape[0] > 1:
|
||||
wav = wav.mean(dim=0, keepdim=True)
|
||||
if sr != target_sr:
|
||||
import torchaudio
|
||||
|
||||
wav = torchaudio.functional.resample(wav, sr, target_sr)
|
||||
return wav # (1, T)
|
||||
|
||||
|
||||
def _build_request(ref_audio_path: str, text: str) -> dict:
|
||||
"""Build a MOSS-TTS-Realtime request.
|
||||
|
||||
Mirrors end2end.py:_build_realtime_prompt exactly:
|
||||
1. Locate MossTTSRealtimeProcessor in the HF snapshot.
|
||||
2. Encode ref audio via MOSS-Audio-Tokenizer (separate repo).
|
||||
3. Build the (L, 17) int grid and return prompt_token_ids +
|
||||
additional_information: {codes.ref, ids.all}.
|
||||
|
||||
Frees the codec model before returning to avoid competing with the
|
||||
running vllm engine for GPU/CPU memory.
|
||||
"""
|
||||
from huggingface_hub import snapshot_download
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
|
||||
# Step 1: locate the realtime processor module in the snapshot.
|
||||
try:
|
||||
snap_dir = Path(snapshot_download(repo_id=MODEL))
|
||||
except Exception as exc:
|
||||
msg = f"Cannot locate snapshot for {MODEL}: {exc}"
|
||||
if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"):
|
||||
pytest.skip(msg)
|
||||
pytest.fail(msg)
|
||||
|
||||
proc_module_path = snap_dir / "processing_mossttsrealtime.py"
|
||||
if not proc_module_path.exists():
|
||||
pytest.fail(f"Realtime processor module missing at {proc_module_path}")
|
||||
|
||||
mod_name = "_moss_tts_realtime_proc"
|
||||
if mod_name not in sys.modules:
|
||||
spec = importlib.util.spec_from_file_location(mod_name, proc_module_path)
|
||||
mod = importlib.util.module_from_spec(spec) # type: ignore[arg-type]
|
||||
sys.modules[mod_name] = mod
|
||||
spec.loader.exec_module(mod) # type: ignore[union-attr]
|
||||
else:
|
||||
mod = sys.modules[mod_name]
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
|
||||
realtime_processor = mod.MossTTSRealtimeProcessor(tokenizer=tokenizer)
|
||||
|
||||
# Step 2: encode the reference audio clip via MOSS-Audio-Tokenizer.
|
||||
try:
|
||||
codec = AutoModel.from_pretrained("OpenMOSS-Team/MOSS-Audio-Tokenizer", trust_remote_code=True).to("cpu").eval()
|
||||
except Exception as exc:
|
||||
msg = f"Cannot load MOSS-Audio-Tokenizer: {exc}"
|
||||
if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"):
|
||||
pytest.skip(msg)
|
||||
pytest.fail(msg)
|
||||
|
||||
ref_wav = _load_ref_audio(ref_audio_path)
|
||||
wav_2d = ref_wav if ref_wav.dim() == 2 else ref_wav.unsqueeze(0)
|
||||
with torch.no_grad():
|
||||
enc = codec.batch_encode([wav_2d.squeeze(0)], num_quantizers=16)
|
||||
codes = enc.audio_codes[:, 0, : int(enc.audio_codes_lengths[0].item())]
|
||||
audio_tokens = codes.transpose(0, 1).contiguous().cpu().numpy() # (T, 16)
|
||||
del codec
|
||||
gc.collect()
|
||||
|
||||
# Step 3: build the (L, 17) prefill grid.
|
||||
system_grid = realtime_processor.make_ensemble(prompt_audio_tokens=audio_tokens)
|
||||
|
||||
assistant_header_ids = tokenizer.encode("<|im_start|>assistant\n", add_special_tokens=False)
|
||||
assistant_grid = np.full((len(assistant_header_ids), 17), 1024, dtype=np.int64)
|
||||
assistant_grid[:, 0] = assistant_header_ids
|
||||
|
||||
text_ids_only = tokenizer.encode(text, add_special_tokens=False)
|
||||
if not text_ids_only:
|
||||
pytest.fail(f"Empty text after tokenization: {text!r}")
|
||||
|
||||
PREFILL_MAX_TEXT = 12
|
||||
cur_len = min(len(text_ids_only), PREFILL_MAX_TEXT)
|
||||
text_grid = np.full((cur_len, 17), 1024, dtype=np.int64)
|
||||
text_grid[:, 0] = text_ids_only[:cur_len]
|
||||
text_grid[-1, 1] = 1025 # audio_bos at the last prefilled text token
|
||||
|
||||
grid = np.concatenate([system_grid, assistant_grid, text_grid], axis=0)
|
||||
prompt_token_ids = grid[:, 0].tolist()
|
||||
audio_codes_tensor = torch.from_numpy(grid[:, 1:].astype(np.int64).copy()) # (L, 16)
|
||||
remaining_text_ids = list(text_ids_only[cur_len:])
|
||||
|
||||
additional_info: dict = {"codes": {"ref": audio_codes_tensor}}
|
||||
if remaining_text_ids:
|
||||
additional_info["ids"] = {"all": remaining_text_ids}
|
||||
|
||||
return {
|
||||
"prompt_token_ids": prompt_token_ids,
|
||||
"additional_information": additional_info,
|
||||
}
|
||||
|
||||
|
||||
def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor, int]:
|
||||
"""Run one request and return (waveform_cpu, sample_rate)."""
|
||||
chunks: list[torch.Tensor] = []
|
||||
sr_final = SAMPLE_RATE
|
||||
for out in omni_runner.omni.generate(request, _SAMPLING):
|
||||
mm = out.multimodal_output
|
||||
if not mm:
|
||||
continue
|
||||
audio = mm.get("audio")
|
||||
if audio is None:
|
||||
audio = mm.get("model_outputs")
|
||||
if audio is None:
|
||||
continue
|
||||
sr = mm.get("sr")
|
||||
if sr is not None:
|
||||
sr_final = int(sr.item() if hasattr(sr, "item") else sr)
|
||||
if isinstance(audio, list):
|
||||
audio = torch.cat(
|
||||
[t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0],
|
||||
dim=0,
|
||||
)
|
||||
if isinstance(audio, torch.Tensor) and audio.numel() > 0:
|
||||
chunks.append(audio.reshape(-1).cpu())
|
||||
if not chunks:
|
||||
raise AssertionError("No audio output received from generate()")
|
||||
return torch.cat(chunks, dim=0), sr_final
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@hardware_test(res={"cuda": "L4"}, 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.")
|
||||
audio, sr = _collect_audio(omni_runner, req)
|
||||
|
||||
assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}"
|
||||
assert audio.numel() > 0, "Audio tensor is empty"
|
||||
assert not torch.all(audio == 0), "Audio is silence"
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@hardware_test(res={"cuda": "L4"}, 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, "你好，这是语音合成测试。")
|
||||
audio, sr = _collect_audio(omni_runner, req)
|
||||
|
||||
assert sr == SAMPLE_RATE
|
||||
assert audio.numel() > 0
|
||||
assert not torch.all(audio == 0)
|
||||
@@ -153,6 +153,12 @@ class MossTTSDelayTalkerForGeneration(nn.Module):
|
||||
# `make_omni_output`, which runs immediately before the sampler).
|
||||
self._batch_state: list[dict[str, Any]] | None = None
|
||||
|
||||
# Stacked weight caches built after load_weights(). Avoids a Python
|
||||
# loop over n_vq heads/embeddings at every decode step.
|
||||
# shapes: (n_vq, audio_vocab+1, hidden_size)
|
||||
self._stacked_audio_head_w: torch.Tensor | None = None
|
||||
self._stacked_audio_emb_w: torch.Tensor | None = None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# vLLM-Omni hooks
|
||||
# ------------------------------------------------------------------
|
||||
@@ -394,9 +400,15 @@ class MossTTSDelayTalkerForGeneration(nn.Module):
|
||||
chunk = ref_codes[ref_offset:end_off]
|
||||
if chunk.numel() > 0 and chunk.shape[0] == span_len:
|
||||
codes = chunk.to(device=device, dtype=torch.long).clamp_(0, self.audio_vocab_size)
|
||||
audio_embed = torch.zeros_like(embeds)
|
||||
for i, emb_layer in enumerate(self.audio_embeddings):
|
||||
audio_embed = audio_embed + emb_layer(codes[:, i])
|
||||
if self._stacked_audio_emb_w is not None:
|
||||
codes_nq = codes.t() # (n_vq, L)
|
||||
audio_embed = self._stacked_audio_emb_w[
|
||||
torch.arange(self.n_vq, device=device)[:, None], codes_nq
|
||||
].sum(0)
|
||||
else:
|
||||
audio_embed = torch.zeros_like(embeds)
|
||||
for i, emb_layer in enumerate(self.audio_embeddings):
|
||||
audio_embed = audio_embed + emb_layer(codes[:, i])
|
||||
embeds = embeds + audio_embed
|
||||
last_ref_row = codes[-1].detach()
|
||||
|
||||
@@ -436,10 +448,17 @@ class MossTTSDelayTalkerForGeneration(nn.Module):
|
||||
audio_codes_buf = (info_dict.get("audio_codes", {}) or {}).get("current")
|
||||
if isinstance(audio_codes_buf, torch.Tensor) and audio_codes_buf.numel() == self.n_vq:
|
||||
codes = audio_codes_buf.to(device=device, dtype=torch.long)
|
||||
audio_embed = torch.zeros_like(text_embed)
|
||||
for i, emb_layer in enumerate(self.audio_embeddings):
|
||||
code_i = codes[i].clamp(0, self.audio_vocab_size).reshape(1)
|
||||
audio_embed = audio_embed + emb_layer(code_i)
|
||||
if self._stacked_audio_emb_w is not None:
|
||||
# Single gather: stacked_emb[i, codes[i]] for each i → (n_vq, H) → sum → (1, H)
|
||||
n_vq_idx = torch.arange(self.n_vq, device=device)
|
||||
audio_embed = self._stacked_audio_emb_w[n_vq_idx, codes.clamp(0, self.audio_vocab_size)].sum(
|
||||
0, keepdim=True
|
||||
)
|
||||
else:
|
||||
audio_embed = torch.zeros_like(text_embed)
|
||||
for i, emb_layer in enumerate(self.audio_embeddings):
|
||||
code_i = codes[i].clamp(0, self.audio_vocab_size).reshape(1)
|
||||
audio_embed = audio_embed + emb_layer(code_i)
|
||||
combined = text_embed + audio_embed
|
||||
else:
|
||||
combined = text_embed
|
||||
@@ -501,17 +520,38 @@ class MossTTSDelayTalkerForGeneration(nn.Module):
|
||||
if not bool(sampling_mask.any()):
|
||||
return codes
|
||||
|
||||
# Use deploy YAML defaults. The audio sampler isn't routed through
|
||||
# vLLM's sampler so we apply the upstream's audio_top_k locally.
|
||||
audio_top_k = 25
|
||||
audio_temp = 1.7
|
||||
|
||||
if self._stacked_audio_head_w is not None:
|
||||
# Single batched matmul: (n_vq, V+1, H) @ (H,) → (n_vq, V+1).
|
||||
# Replaces n_vq serial nn.Linear calls.
|
||||
all_logits = self._stacked_audio_head_w @ last_h.reshape(-1) # (n_vq, V+1)
|
||||
all_logits[:, -1] = float("-inf") # invalid sentinel
|
||||
all_logits[:, self.audio_pad_code] = float("-inf")
|
||||
|
||||
active_idx = sampling_mask.nonzero(as_tuple=True)[0] # (n_active,)
|
||||
active_logits = all_logits[active_idx] # (n_active, V+1)
|
||||
active_logits = active_logits / max(audio_temp, 1e-6)
|
||||
if audio_top_k > 0 and audio_top_k < active_logits.shape[-1]:
|
||||
top_vals, _ = torch.topk(active_logits, audio_top_k, dim=-1)
|
||||
kth = top_vals[:, -1:].expand_as(active_logits)
|
||||
active_logits = torch.where(
|
||||
active_logits < kth,
|
||||
torch.full_like(active_logits, float("-inf")),
|
||||
active_logits,
|
||||
)
|
||||
probs = torch.softmax(active_logits, dim=-1)
|
||||
codes[active_idx] = torch.multinomial(probs, num_samples=1).squeeze(-1).long()
|
||||
return codes
|
||||
|
||||
# Fallback: per-head loop (used before load_weights() completes).
|
||||
for i in range(n_vq):
|
||||
if not bool(sampling_mask[i]):
|
||||
continue
|
||||
head = self.audio_heads[i]
|
||||
logits = head(last_h).reshape(-1) # (V,)
|
||||
logits[..., -1] = float("-inf") # invalid sentinel
|
||||
logits[..., -1] = float("-inf")
|
||||
logits[..., self.audio_pad_code] = float("-inf")
|
||||
sampled = self._sample_with_top_k(logits, audio_top_k, audio_temp)
|
||||
codes[i] = sampled.long()
|
||||
@@ -678,6 +718,14 @@ class MossTTSDelayTalkerForGeneration(nn.Module):
|
||||
for n in backbone_loaded:
|
||||
loaded.add(f"model.{n}")
|
||||
|
||||
# Build stacked weight caches for vectorised decode ops.
|
||||
self._stacked_audio_head_w = torch.stack(
|
||||
[h.weight.detach() for h in self.audio_heads], dim=0
|
||||
) # (n_vq, audio_vocab+1, hidden_size)
|
||||
self._stacked_audio_emb_w = torch.stack(
|
||||
[e.weight.detach() for e in self.audio_embeddings], dim=0
|
||||
) # (n_vq, audio_vocab+1, hidden_size)
|
||||
|
||||
return loaded
|
||||
|
||||
|
||||
@@ -765,6 +813,10 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
|
||||
self.logits_processor = LogitsProcessor(self.text_vocab_size)
|
||||
self._batch_state: list[dict[str, Any]] | None = None
|
||||
|
||||
# Stacked weight cache built after load_weights().
|
||||
# shape: (n_vq, audio_vocab_size, hidden_size)
|
||||
self._stacked_audio_emb_w: torch.Tensor | None = None
|
||||
|
||||
self.gpu_resident_buffer_keys: set[tuple[str, str]] = {
|
||||
("audio_codes", "current"),
|
||||
("audio_codes", "accumulated"),
|
||||
@@ -854,6 +906,11 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
|
||||
embeds = self.embed_tokens[0](text_ids)
|
||||
if audio_codes is None:
|
||||
return embeds
|
||||
if self._stacked_audio_emb_w is not None:
|
||||
# Single gather: (n_vq, vocab, H)[arange[:,None], codes.T] → (n_vq, T, H) → sum → (T, H)
|
||||
n_vq_idx = torch.arange(self.n_vq, device=audio_codes.device)
|
||||
codes_t = audio_codes.t().clamp(0, self.audio_vocab_size - 1) # (n_vq, T)
|
||||
return embeds + self._stacked_audio_emb_w[n_vq_idx.unsqueeze(1), codes_t].sum(0)
|
||||
for i in range(self.n_vq):
|
||||
col = audio_codes[:, i].clamp_(0, self.audio_vocab_size - 1)
|
||||
embeds = embeds + self.embed_tokens[i + 1](col)
|
||||
@@ -1115,6 +1172,12 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
|
||||
for n in backbone_loaded:
|
||||
loaded.add(f"model.{n}")
|
||||
|
||||
# Build stacked embedding cache for vectorised decode ops.
|
||||
# embed_tokens[0] is the text embedding; [1..n_vq] are audio codebooks.
|
||||
self._stacked_audio_emb_w = torch.stack(
|
||||
[self.embed_tokens[i + 1].weight.detach() for i in range(self.n_vq)], dim=0
|
||||
) # (n_vq, audio_vocab_size, hidden_size)
|
||||
|
||||
logger.info(
|
||||
"[MossTTSRealtime] loaded %d/%d params; skipped=%d (first 5: %s)",
|
||||
len(loaded),
|
||||
|
||||
Reference in new issue
Block a user