Youyuan Liandprincepride
423f34326e
Signed-off-by: Youyuan Li <105433145+Darcy-Lee@users.noreply.github.com> Signed-off-by: princepride <wangzhipeng628@gmail.com> Co-authored-by: princepride <wangzhipeng628@gmail.com>
83 lines
2.8 KiB
Python
83 lines
2.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""End-to-end regressions for YuE2 offline WAV export."""
|
|
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import soundfile as sf
|
|
|
|
from tests.helpers.mark import hardware_marks
|
|
from tests.helpers.runtime import get_model_prefix
|
|
from vllm_omni.transformers_utils.repo_utils import hf_api
|
|
|
|
pytestmark = [
|
|
pytest.mark.full_model,
|
|
pytest.mark.tts,
|
|
*hardware_marks(res={"cuda": ["L4", "H100", "B200"]}, num_cards=1),
|
|
]
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[3]
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def vae_path() -> Path:
|
|
path = Path(os.environ.get("YUE2_VAE", get_model_prefix() + "m-a-p/YuE2-Vae"))
|
|
if not path.is_dir():
|
|
path = Path(hf_api().snapshot_download(str(path), allow_patterns=["*.json", "*.safetensors"]))
|
|
return path
|
|
|
|
|
|
@pytest.mark.parametrize("max_frames,truncated", [(200, True), (9000, False)], ids=["budget", "natural-eos"])
|
|
def test_offline_example_saves_audio(vae_path: Path, tmp_path: Path, max_frames: int, truncated: bool) -> None:
|
|
model = Path(os.environ.get("YUE2_MODEL_DIR", get_model_prefix() + "m-a-p/YuE2-3B"))
|
|
if not model.is_dir():
|
|
model = Path(
|
|
hf_api().snapshot_download(str(model), allow_patterns=["*.json", "*.safetensors", "qwen.tiktoken"])
|
|
)
|
|
output = tmp_path / "song.wav"
|
|
env = os.environ.copy()
|
|
env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
|
env["PYTHONPATH"] = str(REPO_ROOT) + os.pathsep + env.get("PYTHONPATH", "")
|
|
completed = subprocess.run(
|
|
[
|
|
sys.executable,
|
|
str(REPO_ROOT / "examples/offline_inference/yue2/end2end.py"),
|
|
"--model",
|
|
str(model),
|
|
"--vae",
|
|
str(vae_path),
|
|
"--style",
|
|
"gentle acoustic folk, clear English vocals, 90 BPM",
|
|
"--lyrics",
|
|
"[Verse]\nMorning light across the bay\nCarry all my dreams away\n[Outro]",
|
|
"--cot",
|
|
"off",
|
|
"--seed",
|
|
"831001",
|
|
"--max-frames",
|
|
str(max_frames),
|
|
"--output",
|
|
str(output),
|
|
],
|
|
cwd=REPO_ROOT,
|
|
env=env,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=900,
|
|
)
|
|
assert completed.returncode == 0, completed.stdout + completed.stderr
|
|
assert f"truncated={truncated}" in completed.stdout
|
|
audio, sample_rate = sf.read(output, always_2d=True)
|
|
assert sample_rate == 48000 and audio.shape[1] == 2
|
|
assert np.isfinite(audio).all() and np.abs(audio).max() > 0
|
|
duration = len(audio) / sample_rate
|
|
if truncated:
|
|
assert duration == pytest.approx(max_frames / 25, abs=0.1)
|
|
else:
|
|
assert 8 <= duration < max_frames / 25
|