[Feature] Add precomputed custom voices and Qwen3-TTS ref-context cache (#3492)

Signed-off-by: Sy03 <1370724210@qq.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Dan250124 <416947747@qq.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
This commit is contained in:
authored and GitHub committed 2026-05-31 17:51:24 +00:00
1 parent 57e5f0fe69
commit d4c139508f
23 files changed
+2362 -145

No files matched your search

+38
View File
@@ -381,6 +381,44 @@ are cached in-process with a shared LRU so repeated requests with the same
all TTS model types; deleting a voice invalidates every model-type slot at
once.
### Precomputed Custom Voices
Qwen3-TTS Base and VoxCPM2 can load offline-precomputed voices at startup.
Generate a directory containing `custom_voice_manifest.json` plus one
`.safetensors` file per voice, then set the pipeline-wide deploy config field:
```yaml
custom_voice_dir: /path/to/custom_voices
```
Qwen3-TTS profiles are created with:
```bash
python examples/online_serving/text_to_speech/qwen3_tts/precompute_custom_voice.py \
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
--voice-name alice \
--ref-audio /path/to/reference.wav \
--ref-text "Original transcript of the reference audio" \
--mode icl \
--output-dir /path/to/custom_voices
```
VoxCPM2 profiles are created with:
```bash
python examples/online_serving/text_to_speech/voxcpm2/precompute_custom_voice.py \
--model openbmb/VoxCPM2 \
--voice-name alice \
--ref-audio /path/to/reference.wav \
--mode ref_continuation \
--prompt-text "Original transcript of the reference audio" \
--output-dir /path/to/custom_voices
```
Only profiles whose safetensors payload can be loaded and validated are exposed
by `GET /v1/audio/voices`. Valid precomputed voices can be used in
`POST /v1/audio/speech` by passing `voice="alice"` without `ref_audio`.
**Configuration (environment variables):**
| Variable | Default | Description |
@@ -380,6 +380,31 @@ curl -X POST http://localhost:8091/v1/audio/voices \
```
Uploaded voices are then usable as `voice="custom_voice_1"` on subsequent requests.
### Precomputed custom voices
For reused Base voice-cloning speakers, precompute the reference artifacts once and load them at server startup:
```bash
python qwen3_tts/precompute_custom_voice.py \
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
--voice-name alice \
--ref-audio /path/to/reference.wav \
--ref-text "Original transcript of the reference audio" \
--mode icl \
--output-dir /path/to/custom_voices
```
`--mode icl` stores both `speaker_embedding` and `ref_code`; `--mode xvec` stores only the speaker embedding. Add the output directory to a deploy config:
```yaml
custom_voice_dir: /path/to/custom_voices
```
Then start the server with that config and call the Speech API with only the voice name:
```bash
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --deploy-config /path/to/qwen3_tts_custom_voice.yaml
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{"input":"Hello from a precomputed voice.","voice":"alice","task_type":"Base"}' \
--output alice.wav
```
### Streaming PCM
```bash
curl -X POST http://localhost:8091/v1/audio/speech \
@@ -419,7 +444,8 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000
`qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
### Notes
- Base voice cloning has per-request reference-audio cost; the uniproc default keeps IPC overhead off the critical path. See the executor-backend section above for background.
- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
- With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request.
- `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
---
@@ -448,6 +474,23 @@ python voxcpm2/openai_speech_client.py \
```
The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs.
### Precomputed custom voices
For repeated VoxCPM2 speakers, precompute the prompt cache and load it through `custom_voice_dir`:
```bash
python voxcpm2/precompute_custom_voice.py \
--model openbmb/VoxCPM2 \
--voice-name alice \
--ref-audio /path/to/reference.wav \
--mode ref_continuation \
--prompt-text "Original transcript of the reference audio" \
--output-dir /path/to/custom_voices
```
Add the output directory to the deploy config:
```yaml
custom_voice_dir: /path/to/custom_voices
```
After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use `voice="alice"` without sending `ref_audio`.
### Gradio demo (gapless streaming via AudioWorklet)
```bash
python voxcpm2/gradio_demo.py
@@ -412,6 +412,31 @@ curl -X POST http://localhost:8091/v1/audio/voices \
```
Uploaded voices are then usable as `voice="custom_voice_1"` on subsequent requests.
### Precomputed custom voices
For reused Base voice-cloning speakers, precompute the reference artifacts once and load them at server startup:
```bash
python qwen3_tts/precompute_custom_voice.py \
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
--voice-name alice \
--ref-audio /path/to/reference.wav \
--ref-text "Original transcript of the reference audio" \
--mode icl \
--output-dir /path/to/custom_voices
```
`--mode icl` stores both `speaker_embedding` and `ref_code`; `--mode xvec` stores only the speaker embedding. Add the output directory to a deploy config:
```yaml
custom_voice_dir: /path/to/custom_voices
```
Then start the server with that config and call the Speech API with only the voice name:
```bash
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --deploy-config /path/to/qwen3_tts_custom_voice.yaml
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{"input":"Hello from a precomputed voice.","voice":"alice","task_type":"Base"}' \
--output alice.wav
```
### Streaming PCM
```bash
curl -X POST http://localhost:8091/v1/audio/speech \
@@ -447,7 +472,8 @@ python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-de
`qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
### Notes
- Base voice cloning has per-request reference-audio cost; the uniproc default keeps IPC overhead off the critical path. See the executor-backend section above for background.
- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
- With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request.
- `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
---
@@ -474,6 +500,23 @@ python voxcpm2/openai_speech_client.py \
```
The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs.
### Precomputed custom voices
For repeated VoxCPM2 speakers, precompute the prompt cache and load it through `custom_voice_dir`:
```bash
python voxcpm2/precompute_custom_voice.py \
--model openbmb/VoxCPM2 \
--voice-name alice \
--ref-audio /path/to/reference.wav \
--mode ref_continuation \
--prompt-text "Original transcript of the reference audio" \
--output-dir /path/to/custom_voices
```
Add the output directory to the deploy config:
```yaml
custom_voice_dir: /path/to/custom_voices
```
After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use `voice="alice"` without sending `ref_audio`.
### Gradio demo (gapless streaming via AudioWorklet)
```bash
python voxcpm2/gradio_demo.py
@@ -0,0 +1,243 @@
"""Pre-compute Qwen3-TTS custom voice profiles.
The generated directory can be passed to the server via
``custom_voice_dir`` in ``vllm_omni/deploy/qwen3_tts.yaml``. Requests can then
use ``/v1/audio/speech`` with ``voice="<name>"`` and no per-request ref_audio.
"""
from __future__ import annotations
import argparse
import json
import os
import sys
from pathlib import Path
from typing import Any
import numpy as np
import torch
from safetensors import safe_open
from safetensors.torch import save_file
REPO_ROOT = Path(__file__).resolve().parents[4]
if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
from vllm_omni.utils.custom_voice_io import safe_voice_stem # noqa: E402
MANIFEST_NAME = "custom_voice_manifest.json"
_SPEAKER_ENCODER_SAMPLE_RATE = 24000
_SPEAKER_ENCODER_NUM_MELS = 128
_SPEAKER_ENCODER_N_FFT = 1024
_SPEAKER_ENCODER_HOP_SIZE = 256
_SPEAKER_ENCODER_WIN_SIZE = 1024
_SPEAKER_ENCODER_FMIN = 0
_SPEAKER_ENCODER_FMAX = 12000
def _resolve_model_dir(model: str) -> str:
if os.path.isdir(model):
return model
from transformers.utils.hub import cached_file
config_path = cached_file(model, "config.json")
if config_path is None:
raise FileNotFoundError(f"Could not resolve config.json for {model}")
return os.path.dirname(config_path)
def _read_audio_mono(path: str) -> tuple[np.ndarray, int]:
import soundfile as sf
wav, sr = sf.read(path, dtype="float32")
if wav.ndim > 1:
wav = wav.mean(axis=-1)
return np.asarray(wav, dtype=np.float32), int(sr)
def _resample(wav: np.ndarray, sr: int, target_sr: int) -> np.ndarray:
if sr == target_sr:
return wav.astype(np.float32)
from vllm.multimodal.audio import AudioResampler
return AudioResampler(target_sr=target_sr).resample(wav.astype(np.float32), orig_sr=int(sr))
def _load_speaker_encoder(model_dir: str, device: torch.device):
from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSSpeakerEncoder
config = Qwen3TTSConfig.from_pretrained(model_dir)
encoder = Qwen3TTSSpeakerEncoder(config.speaker_encoder_config)
index_path = Path(model_dir) / "model.safetensors.index.json"
state: dict[str, torch.Tensor] = {}
if index_path.exists():
index = json.loads(index_path.read_text(encoding="utf-8"))
weight_map = index.get("weight_map", {})
shard_names = sorted({v for k, v in weight_map.items() if k.startswith("speaker_encoder.")})
shard_paths = [Path(model_dir) / name for name in shard_names]
else:
shard_paths = sorted(Path(model_dir).glob("model*.safetensors"))
for shard in shard_paths:
with safe_open(str(shard), framework="pt", device="cpu") as f:
for key in f.keys():
if key.startswith("speaker_encoder."):
state[key.removeprefix("speaker_encoder.")] = f.get_tensor(key)
if not state:
raise RuntimeError(
f"No speaker_encoder.* tensors found under {model_dir}. "
"Use a Qwen3-TTS Base checkpoint for pre-computing custom voice profiles."
)
encoder.load_state_dict(state)
encoder.to(device=device, dtype=torch.bfloat16).eval()
return config, encoder
def _speaker_embedding(
encoder: torch.nn.Module,
speaker_encoder_config: object,
wav: np.ndarray,
sr: int,
device: torch.device,
) -> torch.Tensor:
from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import mel_spectrogram
target_sr = int(getattr(speaker_encoder_config, "sample_rate", _SPEAKER_ENCODER_SAMPLE_RATE))
mel_dim = int(getattr(speaker_encoder_config, "mel_dim", _SPEAKER_ENCODER_NUM_MELS))
if target_sr != _SPEAKER_ENCODER_SAMPLE_RATE or mel_dim != _SPEAKER_ENCODER_NUM_MELS:
raise ValueError(
"Unsupported Qwen3-TTS speaker encoder mel config for precompute: "
f"sample_rate={target_sr}, mel_dim={mel_dim}. "
f"Expected sample_rate={_SPEAKER_ENCODER_SAMPLE_RATE}, mel_dim={_SPEAKER_ENCODER_NUM_MELS} "
"to match the serving-time speaker embedding path."
)
wav = _resample(wav, sr, target_sr)
wav_t = torch.from_numpy(wav).to(device=device, dtype=torch.float32)
mels = mel_spectrogram(
wav_t.unsqueeze(0),
n_fft=_SPEAKER_ENCODER_N_FFT,
num_mels=_SPEAKER_ENCODER_NUM_MELS,
sampling_rate=_SPEAKER_ENCODER_SAMPLE_RATE,
hop_size=_SPEAKER_ENCODER_HOP_SIZE,
win_size=_SPEAKER_ENCODER_WIN_SIZE,
fmin=_SPEAKER_ENCODER_FMIN,
fmax=_SPEAKER_ENCODER_FMAX,
).transpose(1, 2)
with torch.inference_mode():
return encoder(mels.to(device=device, dtype=torch.bfloat16))[0].float().cpu().contiguous()
def _ref_code(model_dir: str, wav: np.ndarray, sr: int, device: torch.device) -> torch.Tensor:
from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import Qwen3TTSTokenizer
tok = Qwen3TTSTokenizer.from_pretrained(
str(Path(model_dir) / "speech_tokenizer"),
torch_dtype=torch.bfloat16,
)
try:
del tok.model.decoder
tok.model.decoder = None
tok.model.encoder.to(device)
tok.device = device
except Exception:
tok.device = device
with torch.inference_mode():
enc = tok.encode(wav, sr=int(sr), return_dict=True)
codes = getattr(enc, "audio_codes", None)
if isinstance(codes, list):
codes = codes[0] if codes else None
if not isinstance(codes, torch.Tensor):
raise RuntimeError("Speech tokenizer did not return audio_codes")
if codes.ndim == 3:
codes = codes[0]
return codes.to(dtype=torch.int32, device="cpu").contiguous()
def _load_manifest(output_dir: Path, model: str, hidden_size: int) -> dict[str, Any]:
path = output_dir / MANIFEST_NAME
if path.exists():
return json.loads(path.read_text(encoding="utf-8"))
return {
"schema_version": 1,
"model_type": "qwen3_tts",
"model": model,
"hidden_size": hidden_size,
"voices": {},
}
def _write_voice(
*,
model: str,
output_dir: Path,
voice_name: str,
ref_audio: str,
ref_text: str | None,
mode: str,
speaker_description: str | None,
device: torch.device,
) -> None:
model_dir = _resolve_model_dir(model)
config, encoder = _load_speaker_encoder(model_dir, device)
wav, sr = _read_audio_mono(ref_audio)
if wav.size < 1024:
raise ValueError(f"Reference audio too short: {wav.size} samples")
tensors: dict[str, torch.Tensor] = {
"speaker_embedding": _speaker_embedding(encoder, config.speaker_encoder_config, wav, sr, device),
}
if mode == "icl":
if not ref_text or not ref_text.strip():
raise ValueError("--ref-text is required for --mode icl")
tensors["ref_code"] = _ref_code(model_dir, wav, sr, device)
output_dir.mkdir(parents=True, exist_ok=True)
filename = f"{safe_voice_stem(voice_name)}.safetensors"
save_file(tensors, str(output_dir / filename))
hidden_size = int(getattr(config.talker_config, "hidden_size", tensors["speaker_embedding"].numel()))
manifest = _load_manifest(output_dir, model, hidden_size)
entry: dict[str, Any] = {
"name": voice_name,
"file": filename,
"mode": mode,
"embedding_dim": int(tensors["speaker_embedding"].numel()),
}
if "ref_code" in tensors:
entry["ref_code_length"] = int(tensors["ref_code"].shape[0])
if ref_text:
entry["ref_text"] = ref_text
if speaker_description:
entry["speaker_description"] = speaker_description
manifest.setdefault("voices", {})[voice_name] = entry
(output_dir / MANIFEST_NAME).write_text(json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8")
print(f"Wrote {output_dir / filename}")
print(f"Updated {output_dir / MANIFEST_NAME}")
def main() -> None:
parser = argparse.ArgumentParser(description="Pre-compute Qwen3-TTS custom voice profile")
parser.add_argument("--model", required=True, help="Qwen3-TTS Base model path or Hugging Face ID")
parser.add_argument("--voice-name", required=True)
parser.add_argument("--ref-audio", required=True)
parser.add_argument("--ref-text", default=None)
parser.add_argument("--mode", choices=["xvec", "icl"], default="xvec")
parser.add_argument("--speaker-description", default=None)
parser.add_argument("--output-dir", required=True)
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
args = parser.parse_args()
_write_voice(
model=args.model,
output_dir=Path(args.output_dir),
voice_name=args.voice_name,
ref_audio=args.ref_audio,
ref_text=args.ref_text,
mode=args.mode,
speaker_description=args.speaker_description,
device=torch.device(args.device),
)
if __name__ == "__main__":
main()
@@ -0,0 +1,128 @@
"""Pre-compute VoxCPM2 custom voice profiles.
The generated directory can be passed to the server via
``custom_voice_dir`` in ``vllm_omni/deploy/voxcpm2.yaml``. Requests can then
use ``/v1/audio/speech`` with ``voice="<name>"`` and no per-request ref_audio.
"""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
from typing import Any
import torch
from safetensors.torch import save_file
REPO_ROOT = Path(__file__).resolve().parents[4]
if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
from vllm_omni.utils.custom_voice_io import safe_voice_stem # noqa: E402
MANIFEST_NAME = "custom_voice_manifest.json"
def _load_tts(model: str, device: torch.device):
from vllm_omni.model_executor.models.voxcpm2.voxcpm2_import_utils import import_voxcpm2_core
VoxCPM = import_voxcpm2_core()
native = VoxCPM.from_pretrained(model, load_denoiser=False, optimize=False)
return native.tts_model.to(device).eval()
def _load_manifest(output_dir: Path, model: str) -> dict[str, Any]:
path = output_dir / MANIFEST_NAME
if path.exists():
return json.loads(path.read_text(encoding="utf-8"))
return {
"schema_version": 1,
"model_type": "voxcpm2",
"model": model,
"voices": {},
}
def _write_voice(
*,
model: str,
output_dir: Path,
voice_name: str,
ref_audio: str,
prompt_text: str | None,
mode: str,
speaker_description: str | None,
device: torch.device,
) -> None:
if mode in ("continuation", "ref_continuation") and not prompt_text:
raise ValueError("--prompt-text is required for continuation/ref_continuation modes")
tts = _load_tts(model, device)
tensors: dict[str, torch.Tensor] = {}
with torch.inference_mode():
if mode in ("reference", "ref_continuation"):
tensors["ref_audio_feat"] = tts._encode_wav(ref_audio, padding_mode="right").float().cpu().contiguous()
if mode in ("continuation", "ref_continuation"):
tensors["audio_feat"] = tts._encode_wav(ref_audio, padding_mode="left").float().cpu().contiguous()
output_dir.mkdir(parents=True, exist_ok=True)
filename = f"{safe_voice_stem(voice_name)}.safetensors"
save_file(tensors, str(output_dir / filename))
manifest = _load_manifest(output_dir, model)
entry: dict[str, Any] = {
"name": voice_name,
"file": filename,
"mode": mode,
}
if "ref_audio_feat" in tensors:
entry["ref_audio_feat_len"] = int(tensors["ref_audio_feat"].shape[0])
if "audio_feat" in tensors:
entry["audio_feat_len"] = int(tensors["audio_feat"].shape[0])
if prompt_text:
entry["prompt_text"] = prompt_text
if speaker_description:
entry["speaker_description"] = speaker_description
manifest.setdefault("voices", {})[voice_name] = entry
(output_dir / MANIFEST_NAME).write_text(json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8")
print(f"Wrote {output_dir / filename}")
print(f"Updated {output_dir / MANIFEST_NAME}")
def main() -> None:
parser = argparse.ArgumentParser(description="Pre-compute VoxCPM2 custom voice profile")
parser.add_argument("--model", default="openbmb/VoxCPM2", help="VoxCPM2 model path or Hugging Face ID")
parser.add_argument("--voice-name", required=True)
parser.add_argument("--ref-audio", required=True)
parser.add_argument(
"--prompt-text",
default=None,
help="Transcript of ref audio for continuation/ref_continuation modes",
)
parser.add_argument(
"--mode",
choices=["reference", "continuation", "ref_continuation"],
default="reference",
)
parser.add_argument("--speaker-description", default=None)
parser.add_argument("--output-dir", required=True)
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
args = parser.parse_args()
_write_voice(
model=args.model,
output_dir=Path(args.output_dir),
voice_name=args.voice_name,
ref_audio=args.ref_audio,
prompt_text=args.prompt_text,
mode=args.mode,
speaker_description=args.speaker_description,
device=torch.device(args.device),
)
if __name__ == "__main__":
main()
@@ -1,9 +1,13 @@
# tests/entrypoints/openai/test_serving_speech.py
import asyncio
import base64
import hashlib
import io
import json
import logging
import os
import struct
import wave
from inspect import Signature, signature
from pathlib import Path
from types import SimpleNamespace
@@ -170,6 +174,27 @@ def create_mock_audio_output_for_test(
)
def _write_custom_voice_manifest(root: Path, *, model_type: str, voices: dict) -> None:
payload = {
"schema_version": 1,
"model_type": model_type,
"voices": voices,
}
(root / "custom_voice_manifest.json").write_text(json.dumps(payload), encoding="utf-8")
def _wav_data_url(samples: np.ndarray, sample_rate: int) -> str:
pcm = (np.clip(samples, -1.0, 1.0) * 32767).astype("<i2")
buf = io.BytesIO()
with wave.open(buf, "wb") as wav:
wav.setnchannels(1)
wav.setsampwidth(2)
wav.setframerate(sample_rate)
wav.writeframes(pcm.tobytes())
encoded = base64.b64encode(buf.getvalue()).decode("ascii")
return f"data:audio/wav;base64,{encoded}"
@pytest.fixture
def test_app(mocker: MockerFixture):
# Mock the engine client
@@ -972,6 +997,233 @@ class TestTTSMethods:
params = speech_server._build_tts_params(req)
assert "voice_clone_prompt" not in params
@pytest.mark.asyncio
async def test_resolve_ref_audio_reuses_decoded_audio_for_same_source(self, speech_server):
wav = np.linspace(-0.5, 0.5, 48000, dtype=np.float32)
ref_audio = _wav_data_url(wav, 24000)
speech_server.model_config.allowed_local_media_path = ""
speech_server.model_config.allowed_media_domains = None
first = await speech_server._resolve_ref_audio(ref_audio)
second = await speech_server._resolve_ref_audio(ref_audio)
assert first[1] == 24000
assert second[1] == 24000
assert first[0] is second[0]
assert first[0][0] == pytest.approx(float(wav[0]), abs=1e-4)
assert speech_server._get_resolved_ref_audio_artifact_key(
ref_audio
) == speech_server._make_ref_audio_artifact_cache_key(np.asarray(first[0], dtype=np.float32), 24000)
def test_precomputed_qwen3_voice_infers_base_without_ref_audio(self, speech_server):
"""Precomputed Qwen3 voices are reusable by name without per-request ref_audio."""
speech_server._tts_model_type = "qwen3_tts"
speech_server.precomputed_speakers = {
"alice": {
"name": "Alice",
"model_type": "qwen3_tts",
"mode": "icl",
"ref_text": "reference transcript",
"ref_code_length": 3,
}
}
speech_server.supported_speakers = {"alice"}
req = OpenAICreateSpeechRequest(input="Hello", voice="Alice")
assert speech_server._validate_tts_request(req) is None
assert req.task_type == "Base"
params = speech_server._build_tts_params(req)
assert params["task_type"] == ["Base"]
assert params["speaker"] == ["alice"]
assert params["x_vector_only_mode"] == [False]
assert params["ref_text"] == ["reference transcript"]
assert params["ref_code_length"] == [3]
assert "ref_audio" not in params
def test_uploaded_qwen3_voice_wins_over_same_named_precomputed_voice(self, speech_server, tmp_path):
"""Qwen3 uploaded voices should take precedence over same-name precomputed voices."""
from safetensors.torch import save_file
uploaded_path = tmp_path / "alice.safetensors"
save_file({"speaker_embedding": torch.tensor([0.1] * 4)}, str(uploaded_path))
speech_server._tts_model_type = "qwen3_tts"
speech_server.uploaded_speakers_dir = tmp_path
speech_server.uploaded_speakers = {
"alice": {
"name": "alice",
"created_at": 123,
"file_path": str(uploaded_path),
"embedding_source": "direct",
}
}
speech_server.precomputed_speakers = {
"alice": {
"name": "Alice",
"model_type": "qwen3_tts",
"mode": "icl",
"ref_text": "precomputed transcript",
"ref_code_length": 3,
}
}
req = OpenAICreateSpeechRequest(input="Hello", voice="Alice")
assert speech_server._validate_tts_request(req) is None
params = speech_server._build_tts_params(req)
assert params["speaker"] == ["alice"]
assert params["voice_created_at"] == [123]
assert params["task_type"] == ["Base"]
assert "voice_clone_prompt" in params
assert "ref_code_length" not in params
assert params.get("ref_text") != ["precomputed transcript"]
def test_precomputed_qwen3_missing_safetensors_is_not_registered(self, speech_server, tmp_path):
"""Manifest entries are not supported speakers unless their safetensors load."""
_write_custom_voice_manifest(
tmp_path,
model_type="qwen3_tts",
voices={"Alice": {"file": "missing.safetensors", "mode": "xvec", "embedding_dim": 4}},
)
speech_server._tts_model_type = "qwen3_tts"
speech_server.engine_client.model_config = SimpleNamespace(
hf_config=SimpleNamespace(
custom_voice_dir=str(tmp_path),
talker_config=SimpleNamespace(hidden_size=4),
)
)
profiles = speech_server._load_precomputed_speakers()
assert profiles == {}
speech_server.precomputed_speakers = profiles
speech_server.supported_speakers = set(profiles)
assert "alice" not in speech_server.supported_speakers
req = OpenAICreateSpeechRequest(input="Hello", voice="Alice")
assert speech_server._validate_tts_request(req) is not None
def test_precomputed_qwen3_icl_without_ref_code_is_not_registered(self, speech_server, tmp_path):
"""Qwen3 ICL profiles without ref_code cannot be exposed by the API layer."""
from safetensors.torch import save_file
save_file({"speaker_embedding": torch.arange(4, dtype=torch.float32)}, str(tmp_path / "alice.safetensors"))
_write_custom_voice_manifest(
tmp_path,
model_type="qwen3_tts",
voices={
"Alice": {
"file": "alice.safetensors",
"mode": "icl",
"ref_text": "reference transcript",
"embedding_dim": 4,
}
},
)
speech_server._tts_model_type = "qwen3_tts"
speech_server.engine_client.model_config = SimpleNamespace(
hf_config=SimpleNamespace(
custom_voice_dir=str(tmp_path),
talker_config=SimpleNamespace(hidden_size=4),
)
)
profiles = speech_server._load_precomputed_speakers()
assert profiles == {}
speech_server.precomputed_speakers = profiles
speech_server.supported_speakers = set(profiles)
assert "alice" not in speech_server.supported_speakers
req = OpenAICreateSpeechRequest(input="Hello", voice="Alice")
assert speech_server._validate_tts_request(req) is not None
def test_precomputed_voxcpm2_missing_safetensors_is_not_registered(self, speech_server, tmp_path):
"""VoxCPM2 must not advertise a manifest-only voice that cannot hit prompt cache."""
_write_custom_voice_manifest(
tmp_path,
model_type="voxcpm2",
voices={"Bob": {"file": "missing.safetensors", "mode": "reference", "ref_audio_feat_len": 2}},
)
speech_server._tts_model_type = "voxcpm2"
speech_server.engine_client.model_config = SimpleNamespace(
hf_config=SimpleNamespace(custom_voice_dir=str(tmp_path))
)
profiles = speech_server._load_precomputed_speakers()
assert profiles == {}
speech_server.precomputed_speakers = profiles
speech_server.supported_speakers = set(profiles)
assert "bob" not in speech_server.supported_speakers
req = OpenAICreateSpeechRequest(input="Hello", voice="Bob")
assert speech_server._validate_tts_request(req) is not None
assert speech_server._validate_tts_request(OpenAICreateSpeechRequest(input="Hello")) is None
def test_prepare_voxcpm2_rejects_supported_speaker_without_voice_profile(self, speech_server, mocker):
"""VoxCPM2 named voices must be uploaded or precomputed on the real request path."""
speech_server._tts_model_type = "voxcpm2"
speech_server.supported_speakers = {"bob", "default"}
speech_server.uploaded_speakers = {}
speech_server.precomputed_speakers = {}
speech_server.engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)]
speech_server.engine_client.generate = mocker.MagicMock(return_value="generator")
speech_server._build_voxcpm2_prompt = mocker.AsyncMock(
return_value={"prompt_token_ids": [1], "additional_information": {}}
)
with pytest.raises(ValueError, match="Invalid voice 'bob'"):
asyncio.run(speech_server._prepare_speech_generation(OpenAICreateSpeechRequest(input="Hello", voice="Bob")))
speech_server._build_voxcpm2_prompt.assert_not_awaited()
speech_server.engine_client.generate.assert_not_called()
def test_prepare_voxcpm2_accepts_default_voice(self, speech_server, mocker):
"""VoxCPM2 default voice preserves the built-in zero-shot request path."""
speech_server._tts_model_type = "voxcpm2"
speech_server.supported_speakers = {"default"}
speech_server.uploaded_speakers = {}
speech_server.precomputed_speakers = {}
speech_server.engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)]
speech_server.engine_client.generate = mocker.MagicMock(return_value=iter(()))
speech_server._build_voxcpm2_prompt = mocker.AsyncMock(
return_value={"prompt_token_ids": [1], "additional_information": {}}
)
asyncio.run(speech_server._prepare_speech_generation(OpenAICreateSpeechRequest(input="Hello", voice="default")))
speech_server._build_voxcpm2_prompt.assert_awaited_once()
speech_server.engine_client.generate.assert_called_once()
def test_prepare_voxcpm2_precomputed_voice_sets_model_cache_key(self, speech_server, mocker):
"""VoxCPM2 precomputed voices must carry voice metadata to the model cache lookup."""
speech_server._tts_model_type = "voxcpm2"
speech_server.supported_speakers = {"alice"}
speech_server.uploaded_speakers = {}
speech_server.precomputed_speakers = {
"alice": {
"name": "Alice",
"model_type": "voxcpm2",
"mode": "reference",
"ref_audio_feat_len": 2,
}
}
speech_server.engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)]
speech_server.engine_client.generate = mocker.MagicMock(return_value=iter(()))
speech_server._build_voxcpm2_prompt = mocker.AsyncMock(
return_value={
"prompt_token_ids": [1],
"additional_information": {
"voice_profile": speech_server.precomputed_speakers["alice"],
},
}
)
asyncio.run(speech_server._prepare_speech_generation(OpenAICreateSpeechRequest(input="Hello", voice="Alice")))
prompt = speech_server.engine_client.generate.call_args.kwargs["prompt"]
additional = prompt["additional_information"]
assert additional["voice_name"] == "alice"
assert additional["voice_created_at"] == 0
def test_build_tts_params(self, speech_server):
"""Test TTS parameter building."""
req = OpenAICreateSpeechRequest(input="Hello", voice="Ryan", language="English")
@@ -2497,6 +2749,95 @@ class TestTTSAsyncOffloading:
qwen3_tts_server._build_tts_params.assert_called_once()
qwen3_tts_server._estimate_prompt_len_async.assert_awaited_once()
def test_qwen3_repeated_ref_audio_hot_path_sends_cache_key_without_waveform(self, qwen3_tts_server):
"""After a ref artifact is marked ready, repeated requests avoid ref_audio payload IPC."""
wav_list = [0.0] * 48000
artifact_key = "a" * 40
ref_audio = "data:audio/wav;base64,same"
qwen3_tts_server._put_resolved_ref_audio(
hashlib.sha1(ref_audio.encode("utf-8")).hexdigest(),
wav_list,
24000,
artifact_key,
)
qwen3_tts_server._ref_audio_model_artifact_ready.add(artifact_key)
qwen3_tts_server._codec_frame_rate = 25.0
qwen3_tts_server._tts_tokenizer = lambda _text, padding=False: {"input_ids": list(range(10))}
qwen3_tts_server.engine_client.model_config.hf_config.talker_config = SimpleNamespace(
codec_language_id={},
spk_is_dialect={},
)
request = OpenAICreateSpeechRequest(
input="hello",
task_type="Base",
ref_audio=ref_audio,
ref_text="reference",
)
request_id, _generator, tts_params = asyncio.run(
qwen3_tts_server._prepare_speech_generation(request, request_id="req-hot")
)
assert request_id == "req-hot"
assert "ref_audio" not in tts_params
assert tts_params["_qwen3_tts_ref_audio_cache_key"] == [artifact_key]
assert tts_params["ref_code_length"] == [50]
prompt = qwen3_tts_server.engine_client.generate.call_args.kwargs["prompt"]
assert prompt["additional_information"] is tts_params
def test_qwen3_ref_audio_artifact_ready_is_evicted_with_resolve_cache(self, qwen3_tts_server):
qwen3_tts_server._ref_audio_resolve_cache_max_entries = 1
qwen3_tts_server._ref_audio_resolve_cache_max_bytes = 1_000_000
qwen3_tts_server._put_resolved_ref_audio("ref-a", [0.0] * 8, 24000, "artifact-a")
qwen3_tts_server._ref_audio_model_artifact_ready.add("artifact-a")
qwen3_tts_server._put_resolved_ref_audio("ref-b", [0.0] * 8, 24000, "artifact-b")
assert "artifact-a" not in qwen3_tts_server._ref_audio_model_artifact_ready
assert "artifact-b" in {entry[3] for entry in qwen3_tts_server._ref_audio_resolve_cache.values()}
@pytest.mark.asyncio
async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_error(self, qwen3_tts_server):
async def failing_generator():
raise ValueError("boom")
yield # pragma: no cover
qwen3_tts_server._request_ref_audio_artifact_keys["req-fail"] = "artifact-fail"
with pytest.raises(ValueError, match="boom"):
await anext(qwen3_tts_server._generate_audio_chunks(failing_generator(), "req-fail"))
assert "req-fail" not in qwen3_tts_server._request_ref_audio_artifact_keys
assert "artifact-fail" not in qwen3_tts_server._ref_audio_model_artifact_ready
@pytest.mark.asyncio
async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_close(self, qwen3_tts_server):
async def pcm_generator():
yield SimpleNamespace(
multimodal_output={
"audio": torch.zeros(16, dtype=torch.float32),
"sr": 24000,
}
)
await asyncio.sleep(0)
qwen3_tts_server._request_ref_audio_artifact_keys["req-close"] = "artifact-close"
stream = qwen3_tts_server._generate_audio_chunks(pcm_generator(), "req-close")
assert await anext(stream)
await stream.aclose()
assert "req-close" not in qwen3_tts_server._request_ref_audio_artifact_keys
assert "artifact-close" not in qwen3_tts_server._ref_audio_model_artifact_ready
def test_qwen3_ref_audio_artifact_ready_requires_live_resolve_cache_entry(self, qwen3_tts_server):
qwen3_tts_server._request_ref_audio_artifact_keys["req-evicted"] = "artifact-evicted"
qwen3_tts_server._mark_ref_audio_artifact_ready_for_request("req-evicted")
assert "req-evicted" not in qwen3_tts_server._request_ref_audio_artifact_keys
assert "artifact-evicted" not in qwen3_tts_server._ref_audio_model_artifact_ready
def test_shutdown_is_idempotent(self, mocker: MockerFixture):
"""Calling shutdown() twice should not raise."""
mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())
@@ -25,6 +25,7 @@ class _FakeDecoder(nn.Module):
self.total_upsample = total_upsample
self.decode_calls: list[dict[str, int]] = []
self.batched_decode_calls: list[dict[str, int]] = []
self.decode_codes: list[torch.Tensor] = []
self.cudagraph_calls: list[dict[str, int | torch.device]] = []
def to(self, *args, **kwargs):
@@ -37,6 +38,7 @@ class _FakeDecoder(nn.Module):
chunk_size: int = 300,
left_context_size: int = 25,
) -> torch.Tensor:
self.decode_codes.append(codes.detach().cpu().clone())
self.decode_calls.append(
{
"chunk_size": chunk_size,
@@ -181,6 +183,136 @@ def test_forward_trims_trailing_padding_without_context():
torch.testing.assert_close(audio, expected)
def test_forward_reuses_cached_ref_context_for_followup_chunk():
model = _make_model()
first_codes = torch.tensor(
[
9,
8,
1,
2,
9,
8,
3,
4,
],
dtype=torch.long,
)
model.forward(
input_ids=first_codes,
runtime_additional_information=[
{
"meta": {
"left_context_size": 2,
"ref_context_size": 2,
"ref_context_request_id": "rid",
"ref_context_included": True,
}
}
],
)
followup_codes = torch.tensor([5, 6, 7, 15, 16, 17], dtype=torch.long)
model.forward(
input_ids=followup_codes,
runtime_additional_information=[
{
"meta": {
"left_context_size": 3,
"ref_context_size": 2,
"ref_context_request_id": "rid",
"ref_context_included": False,
"finished": torch.tensor(True),
}
}
],
)
torch.testing.assert_close(
model.decoder.decode_codes[-1],
torch.tensor(
[
[
[9, 8, 5, 6, 7],
[9, 8, 15, 16, 17],
]
],
dtype=torch.long,
),
)
assert "rid" not in model._ref_context_cache
def test_forward_fails_fast_when_ref_context_cache_is_missing():
model = _make_model()
followup_codes = torch.tensor([5, 6, 7, 15, 16, 17], dtype=torch.long)
with pytest.raises(ValueError, match="Missing Qwen3-TTS ref context cache"):
model.forward(
input_ids=followup_codes,
runtime_additional_information=[
{
"meta": {
"left_context_size": 3,
"ref_context_size": 2,
"ref_context_request_id": "rid",
"ref_context_included": False,
}
}
],
)
def test_ref_context_cache_evicts_lru_entries_when_requests_abort():
model = _make_model()
model._ref_context_cache_max_entries = 2
first_codes = torch.tensor([9, 8, 1, 2, 9, 8, 3, 4], dtype=torch.long)
for rid in ("a", "b", "c"):
model.forward(
input_ids=first_codes,
runtime_additional_information=[
{
"meta": {
"left_context_size": 2,
"ref_context_size": 2,
"ref_context_request_id": rid,
"ref_context_included": True,
}
}
],
)
assert list(model._ref_context_cache) == ["b", "c"]
assert model._ref_context_cache_bytes == sum(model._tensor_nbytes(t) for t in model._ref_context_cache.values())
def test_ref_context_cache_evicts_to_byte_cap():
model = _make_model()
model._ref_context_cache_max_entries = 100
model._ref_context_cache_max_bytes = 33
first_codes = torch.tensor([9, 8, 1, 2, 9, 8, 3, 4], dtype=torch.long)
for rid in ("a", "b"):
model.forward(
input_ids=first_codes,
runtime_additional_information=[
{
"meta": {
"left_context_size": 2,
"ref_context_size": 2,
"ref_context_request_id": rid,
"ref_context_included": True,
}
}
],
)
assert list(model._ref_context_cache) == ["b"]
assert model._ref_context_cache_bytes == 32
def test_connector_codec_chunking_does_not_override_decode_chunking():
model = _make_model(
async_chunk=True,
@@ -2,6 +2,7 @@ from collections import OrderedDict
from types import SimpleNamespace
import numpy as np
import pytest
import torch
from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
@@ -9,6 +10,7 @@ from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
PRECOMPUTED_REF_CODE_KEY,
PRECOMPUTED_REF_IDS_KEY,
PRECOMPUTED_TEXT_IDS_KEY,
REF_AUDIO_CACHE_KEY,
Qwen3TTSPromptEmbedsBuilder,
)
from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
@@ -493,6 +495,41 @@ def test_base_voice_clone_batch_preprocess_skips_after_initial_prefill_state_exi
assert NORMALIZED_REF_AUDIO_KEY not in buf["r1"]
def test_base_voice_clone_batch_preprocess_uses_serving_artifact_cache_key_without_normalize():
builder = _make_minimal_builder()
ref_code = torch.full((2, 2), 5, dtype=torch.long)
builder.put_ref_audio_artifacts("same-ref", ref_code=ref_code)
builder.normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("normalize not expected"))
builder._encode_ref_audio_batch_fn = lambda *a, **kw: (_ for _ in ()).throw(
AssertionError("speech tokenizer not expected")
)
class FakeTextTokenizer:
def __call__(self, texts, *, padding=False):
return {"input_ids": [[7, 8, 9] for _ in texts]}
builder._text_tokenizer = FakeTextTokenizer()
buf = {
"r1": {
"task_type": ["Base"],
"text": ["one"],
"ref_audio": ["a.wav"],
"ref_text": ["hello"],
"x_vector_only_mode": [False],
REF_AUDIO_CACHE_KEY: ["same-ref"],
}
}
builder.preprocess_batch(
req_ids=["r1"],
model_intermediate_buffer=buf,
device=torch.device("cpu"),
)
assert torch.equal(buf["r1"]["codes"][PRECOMPUTED_REF_CODE_KEY], ref_code)
assert NORMALIZED_REF_AUDIO_KEY not in buf["r1"]
def test_base_voice_clone_uses_batched_ref_code_without_serial_encode():
builder = _make_minimal_builder()
device_param = torch.nn.Parameter(torch.empty(0))
@@ -536,3 +573,112 @@ def test_base_voice_clone_uses_batched_ref_code_without_serial_encode():
assert speaker_wav_ids == [id(ref_audio)]
assert ref_code_len == 2
assert torch.equal(out_ref_code, ref_code)
def test_voice_clone_prompt_artifacts_do_not_resolve_ref_audio_for_cache_fill():
builder = _make_minimal_builder()
device_param = torch.nn.Parameter(torch.empty(0))
builder._text_embedding = _stub_text_embedding(device_param)
builder._text_projection = lambda embeds: embeds
builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
class FakeTokenizer:
def __call__(self, *_args, **_kwargs):
return {"input_ids": torch.arange(8, dtype=torch.long).reshape(1, -1)}
builder._text_tokenizer = FakeTokenizer()
builder._generate_icl_prompt = lambda **kwargs: (
torch.ones((1, 2, 4), device=kwargs["ref_code"].device),
torch.ones((1, 4), device=kwargs["ref_code"].device),
)
builder.normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("normalize not expected"))
builder.encode_ref_audio_to_code = lambda _wav, _sr: (_ for _ in ()).throw(AssertionError("encode not expected"))
builder.extract_speaker_embedding = lambda _wav, _sr: (_ for _ in ()).throw(
AssertionError("speaker embedding not expected")
)
ref_code = torch.arange(4, dtype=torch.long).reshape(2, 2)
_prompt, _trailing, ref_code_len, out_ref_code = builder.build_prompt_embeds(
task_type="Base",
info_dict={
"text": ["hello"],
"ref_audio": ["ref.wav"],
"ref_ids": torch.arange(8, dtype=torch.long).reshape(1, -1),
"non_streaming_mode": [False],
"voice_clone_prompt": {
"ref_code": ref_code,
"ref_spk_embedding": torch.ones(4, dtype=torch.bfloat16),
"icl_mode": True,
},
},
)
assert ref_code_len == 2
assert torch.equal(out_ref_code, ref_code)
def test_ref_audio_artifact_only_uses_cache_without_ref_audio_payload():
builder = _make_minimal_builder()
device_param = torch.nn.Parameter(torch.empty(0))
builder._text_embedding = _stub_text_embedding(device_param)
builder._text_projection = lambda embeds: embeds
builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
class FakeTokenizer:
def __call__(self, *_args, **_kwargs):
return {"input_ids": torch.arange(8, dtype=torch.long).reshape(1, -1)}
builder._text_tokenizer = FakeTokenizer()
builder._generate_icl_prompt = lambda **kwargs: (
torch.ones((1, 2, 4), device=kwargs["ref_code"].device),
torch.ones((1, 4), device=kwargs["ref_code"].device),
)
builder.normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("normalize not expected"))
builder.encode_ref_audio_to_code = lambda _wav, _sr: (_ for _ in ()).throw(AssertionError("encode not expected"))
builder.extract_speaker_embedding = lambda _wav, _sr: (_ for _ in ()).throw(
AssertionError("speaker embedding not expected")
)
ref_code = torch.arange(4, dtype=torch.long).reshape(2, 2)
builder.put_ref_audio_artifacts(
"same-ref",
ref_code=ref_code,
ref_spk_embedding=torch.ones(4, dtype=torch.bfloat16),
)
_prompt, _trailing, ref_code_len, out_ref_code = builder.build_prompt_embeds(
task_type="Base",
info_dict={
"text": ["hello"],
"ref_ids": torch.arange(8, dtype=torch.long).reshape(1, -1),
"non_streaming_mode": [False],
REF_AUDIO_CACHE_KEY: ["same-ref"],
},
)
assert ref_code_len == 2
assert torch.equal(out_ref_code, ref_code)
def test_ref_audio_artifact_only_cache_miss_fails_fast():
builder = _make_minimal_builder()
device_param = torch.nn.Parameter(torch.empty(0))
builder._text_embedding = _stub_text_embedding(device_param)
builder._text_projection = lambda embeds: embeds
builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
class FakeTokenizer:
def __call__(self, *_args, **_kwargs):
return {"input_ids": torch.arange(8, dtype=torch.long).reshape(1, -1)}
builder._text_tokenizer = FakeTokenizer()
with pytest.raises(RuntimeError, match="artifact cache miss"):
builder.build_prompt_embeds(
task_type="Base",
info_dict={
"text": ["hello"],
"ref_ids": torch.arange(8, dtype=torch.long).reshape(1, -1),
"non_streaming_mode": [False],
REF_AUDIO_CACHE_KEY: ["missing-ref"],
},
)
@@ -79,6 +79,27 @@ def test_estimate_prompt_len_handles_1d_ref_code() -> None:
assert est > 50, f"got {est}; 1D ref_code must contribute its own length"
def test_estimate_prompt_len_uses_ref_code_length_without_ref_audio() -> None:
est = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
additional_information={
"task_type": ["Base"],
"text": ["hello"],
"ref_text": ["reference transcript"],
"x_vector_only_mode": [False],
"ref_code_length": [6],
"_qwen3_tts_ref_audio_cache_key": ["same-ref"],
"language": ["English"],
},
task_type="Base",
tokenize_prompt=lambda _text: list(range(10)),
codec_language_id=None,
spk_is_dialect=None,
estimate_ref_code_len=lambda _ref_audio: None,
)
assert est == 15
def test_estimate_prompt_len_uses_list_of_int_ref_ids_from_voice_clone_prompt() -> None:
"""`voice_clone_prompt.ref_ids` as a list-of-int (the natural pre-tokenized shape)
must be read as a sequence, not unwrapped to its first element. Without the fix
@@ -269,11 +269,14 @@ def test_first_streaming_chunk_prepends_ref_code_context():
assert payload is not None
assert payload.meta.left_context_size == 2
assert payload.meta.ref_context_size == 2
assert payload.meta.ref_context_request_id == rid
assert payload.meta.ref_context_included is True
assert len(payload.codes.audio) == _Q * 12
def test_ref_code_context_applies_to_all_streaming_chunks():
"""ref_code is prepended as decoder context on every chunk, not just the first."""
def test_followup_ref_code_context_is_sent_as_metadata_handle():
"""Follow-up chunks keep full ref context semantically without resending it."""
tm = _tm()
rid = "r-ref2"
tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(35)]
@@ -289,9 +292,13 @@ def test_ref_code_context_applies_to_all_streaming_chunks():
)
assert payload is not None
# ref_code (2 frames) prepended as left context on second chunk too
# ref_code (2 frames) is represented in metadata so Code2Wav can restore it
# from its request-local cache. It must not be resent in codes.audio.
assert payload.meta.left_context_size == 10 + 2
assert len(payload.codes.audio) == _Q * (35 + 2)
assert payload.meta.ref_context_size == 2
assert payload.meta.ref_context_request_id == rid
assert payload.meta.ref_context_included is False
assert len(payload.codes.audio) == _Q * 35
def test_streaming_ref_code_context_is_bounded_for_batchable_shapes():
+29
View File
@@ -1027,6 +1027,9 @@ class TestDeployConfigLoading:
"async_scheduling",
"compilation_config",
"config_format",
"custom_voice_dir",
"data_parallel_size",
"devices",
"disable_hybrid_kv_cache_manager",
"enable_flashinfer_autotune",
"enforce_eager",
@@ -1068,6 +1071,32 @@ class TestDeployConfigLoading:
f"added={actual_fields - expected_fields}, removed={expected_fields - actual_fields}"
)
def test_custom_voice_dir_is_pipeline_wide_engine_arg(self, tmp_path):
import vllm_omni.model_executor.models.qwen3_tts.pipeline # noqa: F401
from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
deploy_path = tmp_path / "qwen3_tts_custom_voice.yaml"
custom_voice_dir = tmp_path / "voices"
deploy_path.write_text(
f"""
async_chunk: true
custom_voice_dir: {custom_voice_dir}
stages:
- stage_id: 0
devices: "0"
- stage_id: 1
devices: "0"
""",
encoding="utf-8",
)
deploy = load_deploy_config(deploy_path)
pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
stages = merge_pipeline_deploy(pipeline, deploy)
assert deploy.custom_voice_dir == str(custom_voice_dir)
assert {s.yaml_engine_args.get("custom_voice_dir") for s in stages} == {str(custom_voice_dir)}
def test_load_qwen3_omni_moe_deploy_config(self):
deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
deploy = load_deploy_config(deploy_path)
+283
View File
@@ -0,0 +1,283 @@
import json
from types import SimpleNamespace
import pytest
import torch
from safetensors.torch import save_file
from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _write_manifest(root, *, model_type: str, voices: dict) -> None:
payload = {
"schema_version": 1,
"model_type": model_type,
"voices": voices,
}
(root / "custom_voice_manifest.json").write_text(json.dumps(payload), encoding="utf-8")
def test_qwen3_custom_voice_profiles_warm_speaker_cache(tmp_path):
from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
Qwen3TTSTalkerForConditionalGeneration,
)
save_file(
{
"speaker_embedding": torch.arange(4, dtype=torch.float32),
"ref_code": torch.arange(6, dtype=torch.int32).reshape(3, 2),
},
str(tmp_path / "alice.safetensors"),
)
_write_manifest(
tmp_path,
model_type="qwen3_tts",
voices={
"Alice": {
"file": "alice.safetensors",
"mode": "icl",
"ref_text": "reference transcript",
"embedding_dim": 4,
}
},
)
model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
model.config = SimpleNamespace(
custom_voice_dir=str(tmp_path),
speaker_encoder_config=SimpleNamespace(enc_dim=4),
)
model._speaker_cache = SpeakerEmbeddingCache()
model._load_custom_voice_profiles()
key = model._speaker_cache.make_cache_key("alice", model_type="qwen3_tts_icl", created_at=0)
cached = model._speaker_cache.get(key)
assert cached is not None
assert cached["icl_mode"] is True
assert cached["ref_text"] == "reference transcript"
torch.testing.assert_close(cached["ref_spk_embedding"], torch.arange(4, dtype=torch.float32))
torch.testing.assert_close(cached["ref_code"], torch.arange(6, dtype=torch.int32).reshape(3, 2))
def test_qwen3_icl_profile_without_ref_code_is_not_downgraded(tmp_path):
from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
Qwen3TTSTalkerForConditionalGeneration,
)
save_file({"speaker_embedding": torch.arange(4, dtype=torch.float32)}, str(tmp_path / "alice.safetensors"))
_write_manifest(
tmp_path,
model_type="qwen3_tts",
voices={
"Alice": {
"file": "alice.safetensors",
"mode": "icl",
"ref_text": "reference transcript",
"embedding_dim": 4,
}
},
)
model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
model.config = SimpleNamespace(
custom_voice_dir=str(tmp_path),
speaker_encoder_config=SimpleNamespace(enc_dim=4),
)
model._speaker_cache = SpeakerEmbeddingCache()
model._load_custom_voice_profiles()
icl_key = model._speaker_cache.make_cache_key("alice", model_type="qwen3_tts_icl", created_at=0)
xvec_key = model._speaker_cache.make_cache_key("alice", model_type="qwen3_tts_xvec", created_at=0)
assert model._speaker_cache.get(icl_key) is None
assert model._speaker_cache.get(xvec_key) is None
def test_qwen3_prompt_len_uses_precomputed_ref_code_length():
from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
Qwen3TTSPromptEmbedsBuilder,
)
prompt_len = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
additional_information={
"text": ["hello"],
"task_type": ["Base"],
"speaker": ["alice"],
"x_vector_only_mode": [False],
"ref_code_length": [6],
"ref_text": ["reference transcript"],
},
task_type="Base",
tokenize_prompt=lambda _text: list(range(10)),
codec_language_id=None,
spk_is_dialect=None,
estimate_ref_code_len=lambda _ref_audio: None,
)
assert prompt_len == 15
def test_custom_voice_profile_file_must_stay_under_manifest_dir(tmp_path):
from vllm_omni.utils.speaker_cache import iter_custom_voice_profiles
_write_manifest(
tmp_path,
model_type="qwen3_tts",
voices={"Alice": {"file": "../alice.safetensors", "mode": "xvec"}},
)
assert iter_custom_voice_profiles(tmp_path, expected_model_type="qwen3_tts") == []
def test_custom_voice_profiles_drop_stale_provenance_metadata(tmp_path):
from vllm_omni.utils.speaker_cache import iter_custom_voice_profiles
save_file({"speaker_embedding": torch.arange(4, dtype=torch.float32)}, str(tmp_path / "alice.safetensors"))
_write_manifest(
tmp_path,
model_type="qwen3_tts",
voices={
"Alice": {
"file": "alice.safetensors",
"mode": "xvec",
"audio_codec_config_hash": "old-hash",
"audio_codec_version": "old-codec",
"model_revision": "old-revision",
}
},
)
profiles = iter_custom_voice_profiles(tmp_path, expected_model_type="qwen3_tts")
assert len(profiles) == 1
assert "audio_codec_config_hash" not in profiles[0]
assert "audio_codec_version" not in profiles[0]
assert "model_revision" not in profiles[0]
def test_voxcpm2_custom_voice_profiles_warm_full_prompt_cache(tmp_path):
from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
VoxCPM2TalkerForConditionalGeneration,
)
save_file(
{
"ref_audio_feat": torch.ones(2, 3, 4, dtype=torch.float32),
"audio_feat": torch.full((5, 3, 4), 2.0, dtype=torch.float32),
},
str(tmp_path / "bob.safetensors"),
)
_write_manifest(
tmp_path,
model_type="voxcpm2",
voices={
"Bob": {
"file": "bob.safetensors",
"mode": "ref_continuation",
"prompt_text": "prompt transcript",
"ref_audio_feat_len": 2,
"audio_feat_len": 5,
}
},
)
model = VoxCPM2TalkerForConditionalGeneration.__new__(VoxCPM2TalkerForConditionalGeneration)
model.config = SimpleNamespace(custom_voice_dir=str(tmp_path))
model._speaker_cache = SpeakerEmbeddingCache()
model._load_custom_voice_profiles()
key = model._speaker_cache.make_cache_key("bob", model_type="voxcpm2", created_at=0)
cached = model._speaker_cache.get(key)
assert cached is not None
assert cached["mode"] == "ref_continuation"
assert cached["prompt_text"] == "prompt transcript"
torch.testing.assert_close(cached["ref_audio_feat"], torch.ones(2, 3, 4))
torch.testing.assert_close(cached["audio_feat"], torch.full((5, 3, 4), 2.0))
def test_voxcpm2_custom_voice_profiles_skip_unloadable_manifest_entry(tmp_path):
from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
VoxCPM2TalkerForConditionalGeneration,
)
_write_manifest(
tmp_path,
model_type="voxcpm2",
voices={
"Bob": {
"file": "missing.safetensors",
"mode": "reference",
"ref_audio_feat_len": 2,
}
},
)
model = VoxCPM2TalkerForConditionalGeneration.__new__(VoxCPM2TalkerForConditionalGeneration)
model.config = SimpleNamespace(custom_voice_dir=str(tmp_path))
model._speaker_cache = SpeakerEmbeddingCache()
model._load_custom_voice_profiles()
key = model._speaker_cache.make_cache_key("bob", model_type="voxcpm2", created_at=0)
assert model._speaker_cache.get(key) is None
def test_voxcpm2_precomputed_voice_cache_miss_raises_before_zero_shot():
from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
VoxCPM2TalkerForConditionalGeneration,
)
model = VoxCPM2TalkerForConditionalGeneration.__new__(VoxCPM2TalkerForConditionalGeneration)
model.config = SimpleNamespace(bos_token_id=1)
model._speaker_cache = SpeakerEmbeddingCache()
model._get_multichar_zh_split = lambda: {}
model._get_or_create_state = lambda _req_id: SimpleNamespace()
with pytest.raises(ValueError, match="not loaded in the model cache"):
model.preprocess(
torch.tensor([1, 2], dtype=torch.long),
None,
request_id="req-1",
text_token_ids=[[2]],
voice_name=["bob"],
voice_created_at=0,
voice_profile={"mode": "reference"},
)
def test_voxcpm2_prompt_can_account_for_precomputed_profile_lengths():
from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import build_voxcpm2_prompt
class _Tokenizer:
bos_token_id = 1
unk_token_id = 0
def get_vocab(self):
return {}
def encode(self, text, add_special_tokens=True):
ids = [10 + i for i, _ in enumerate(text.split())]
return [self.bos_token_id, *ids] if add_special_tokens else ids
prompt = build_voxcpm2_prompt(
hf_config=SimpleNamespace(audio_vae_config={"sample_rate": 48000, "encoder_rates": [2]}, patch_size=3),
tokenizer=_Tokenizer(),
split_map={},
text="hello world",
voice_profile={
"mode": "ref_continuation",
"ref_audio_feat_len": 2,
"audio_feat_len": 5,
"prompt_text": "prompt words",
},
)
# text(2) + audio_start(1) + ref prefix(2 + ref_len) + prompt_text(2) + prompt_audio_len(5)
assert len(prompt["prompt_token_ids"]) == 14
additional = prompt["additional_information"]
assert additional["voice_profile"]["mode"] == "ref_continuation"
assert additional["prompt_text"] == ["prompt words"]
+3
View File
@@ -516,6 +516,7 @@ class DeployConfig:
enable_chunked_prefill: bool | None = None
data_parallel_size: int | None = None
pipeline_parallel_size: int | None = None
custom_voice_dir: str | None = None
_STAGE_RESERVED_KEYS = frozenset(
@@ -693,6 +694,7 @@ def load_deploy_config(path: str | Path) -> DeployConfig:
"enable_chunked_prefill",
"data_parallel_size",
"pipeline_parallel_size",
"custom_voice_dir",
):
if name in raw_dict:
kwargs[name] = raw_dict[name]
@@ -814,6 +816,7 @@ _PIPELINE_WIDE_ENGINE_FIELDS: tuple[str, ...] = (
"enable_chunked_prefill",
"data_parallel_size",
"pipeline_parallel_size",
"custom_voice_dir",
)
+6
View File
@@ -77,6 +77,9 @@ class OmniPayloadMeta(TypedDict, total=False):
decode_flag: bool
codec_streaming: bool
ref_code_len: int
ref_context_size: int
ref_context_request_id: str
ref_context_included: bool
talker_prefill_offset: int
@@ -161,6 +164,9 @@ class MetaStruct(_StructBase):
decode_flag: bool | None = None
codec_streaming: bool | None = None
ref_code_len: int | None = None
ref_context_size: int | None = None
ref_context_request_id: str | None = None
ref_context_included: bool | None = None
talker_prefill_offset: int | None = None
codec_chunk_frames: int | None = None
codec_left_context_frames: int | None = None
+4
View File
@@ -824,6 +824,10 @@ def build_vllm_config(
if upgraded is not vllm_config.quant_config:
vllm_config = replace(vllm_config, quant_config=upgraded)
custom_voice_dir = engine_args_dict.get("custom_voice_dir")
if custom_voice_dir:
setattr(vllm_config.model_config.hf_config, "custom_voice_dir", custom_voice_dir)
return vllm_config, executor_class
+331 -106
View File
@@ -8,6 +8,7 @@ import os
import re
import struct
import time
from collections import OrderedDict
from concurrent.futures import ThreadPoolExecutor
from http import HTTPStatus
from pathlib import Path
@@ -55,7 +56,11 @@ from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import (
create_instruction as ming_create_instruction,
)
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.utils.speaker_cache import get_speaker_cache
from vllm_omni.utils.speaker_cache import (
get_speaker_cache,
iter_custom_voice_profiles,
load_validated_profile_tensors,
)
logger = init_logger(__name__)
@@ -107,6 +112,9 @@ _TTS_LANGUAGES: set[str] = {
}
_REF_AUDIO_MIN_DURATION = 1.0 # seconds
_REF_AUDIO_MAX_DURATION = 30.0 # seconds
_REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES = 256
_REF_AUDIO_RESOLVE_CACHE_MAX_BYTES = 256 * 1024 * 1024
_QWEN3_TTS_REF_AUDIO_CACHE_KEY = "_qwen3_tts_ref_audio_cache_key"
_TTS_MAX_INSTRUCTIONS_LENGTH = 500
_TTS_MAX_NEW_TOKENS_MIN = 1
_TTS_MAX_NEW_TOKENS_MAX = 4096
@@ -261,8 +269,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
logger.warning("Invalid SPEAKER_MAX_UPLOADED=%r; using default 1000", _raw_cap)
self._max_uploaded_speakers = 1000
self.uploaded_speakers: dict[str, dict] = {}
self.precomputed_speakers: dict[str, dict[str, Any]] = {}
self.supported_speakers: set[str] = set()
self._ref_audio_data_url_cache: dict[str, str] = {}
self._ref_audio_resolve_cache: OrderedDict[str, tuple[list[float], int, int, str]] = OrderedDict()
self._ref_audio_resolve_cache_bytes = 0
self._ref_audio_resolve_cache_max_entries = _REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES
self._ref_audio_resolve_cache_max_bytes = _REF_AUDIO_RESOLVE_CACHE_MAX_BYTES
self._ref_audio_model_artifact_ready: set[str] = set()
self._request_ref_audio_artifact_keys: dict[str, str] = {}
self._speaker_cache = get_speaker_cache()
self._last_upload_ts = 0
self._upload_lock = asyncio.Lock()
@@ -395,6 +410,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
)
# Determine TTS model type or None
self._tts_model_type = self._detect_tts_model_type()
self.precomputed_speakers = self._load_precomputed_speakers()
# GLM-TTS lazy-cached resources (populated on first GLM-TTS request)
self._glm_tts_text_tokenizer: object | None = None
@@ -405,6 +421,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
# Merge built-in speakers into the set initialized by _init_speaker_storage.
self.supported_speakers |= self._load_supported_speakers()
self.supported_speakers |= set(self.precomputed_speakers)
self._tts_tokenizer = None
self._voxcpm2_tokenizer = None
self._voxcpm2_split_map: dict[int, list[int]] = {}
@@ -570,6 +587,45 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
return "glm_tts"
return None
def _get_custom_voice_dir(self) -> str | None:
try:
value = getattr(self.engine_client.model_config.hf_config, "custom_voice_dir", None)
except AttributeError:
return None
if isinstance(value, os.PathLike):
return os.fspath(value)
if isinstance(value, str) and value:
return value
return None
def _load_precomputed_speakers(self) -> dict[str, dict[str, Any]]:
"""Load precomputed voice names from ``custom_voice_dir`` for API validation."""
if self._tts_model_type not in ("qwen3_tts", "voxcpm2"):
return {}
custom_voice_dir = self._get_custom_voice_dir()
if not custom_voice_dir:
return {}
profiles: dict[str, dict[str, Any]] = {}
qwen3_embedding_dim = self._get_qwen_tts_expected_speaker_embedding_dim()
for profile in iter_custom_voice_profiles(custom_voice_dir, expected_model_type=self._tts_model_type):
tensors = load_validated_profile_tensors(
profile,
expected_model_type=self._tts_model_type,
qwen3_embedding_dim=qwen3_embedding_dim,
)
if tensors is None:
continue
profiles[profile["voice_name_lower"]] = profile
if profiles:
logger.info(
"Loaded %d precomputed %s voice profile(s) from %s",
len(profiles),
self._tts_model_type,
custom_voice_dir,
)
return profiles
def _compute_max_instructions_length(self) -> int:
"""Compute max instructions length with precedence: CLI > stage config > default.
@@ -748,11 +804,14 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
self._voxcpm2_encode("") # lazy-init tokenizer + split_map
ref_audio = None
ref_sr = None
voice_profile = None
if request.ref_audio is not None:
ref_audio, ref_sr = await self._resolve_ref_audio(request.ref_audio)
elif uploaded_ref is not None:
wav_np, ref_sr = uploaded_ref
ref_audio = wav_np.tolist()
elif request.voice is not None:
voice_profile = self.precomputed_speakers.get(request.voice.lower())
return build_voxcpm2_prompt(
hf_config=self.engine_client.model_config.hf_config,
tokenizer=self._voxcpm2_tokenizer,
@@ -761,6 +820,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
ref_audio=ref_audio,
ref_sr=ref_sr,
ref_text=request.ref_text,
voice_profile=voice_profile,
)
def _load_uploaded_audio(self, voice_name: str) -> tuple[np.ndarray, int] | None:
@@ -1205,12 +1265,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
if self._tts_model_type == "cosyvoice3":
return self._validate_cosyvoice3_request(request)
if self._tts_model_type == "voxcpm2":
if request.max_new_tokens is not None:
if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN:
return f"max_new_tokens must be at least {_TTS_MAX_NEW_TOKENS_MIN}"
if request.max_new_tokens > _TTS_MAX_NEW_TOKENS_MAX:
return f"max_new_tokens cannot exceed {_TTS_MAX_NEW_TOKENS_MAX}"
return None # VoxCPM2 accepts any text input
return self._validate_voxcpm2_request(request)
if self._tts_model_type == "ming_flash_omni_tts":
return self._validate_ming_tts_request(request)
if self._tts_model_type == "moss_tts_nano":
@@ -1219,6 +1274,26 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
return self._validate_glm_tts_request(request)
return self._validate_qwen_tts_request(request)
def _validate_voxcpm2_request(self, request: OpenAICreateSpeechRequest) -> str | None:
"""Validate VoxCPM2 request parameters. Returns error message or None."""
if not request.input or not request.input.strip():
return "Input text cannot be empty"
if request.voice is not None:
request.voice = request.voice.lower()
available_voices = set(self.uploaded_speakers) | set(self.precomputed_speakers) | {"default"}
if request.voice not in available_voices:
supported = ", ".join(sorted(available_voices)) or "none"
return f"Invalid voice '{request.voice}'. Supported: {supported}"
if request.max_new_tokens is not None:
if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN:
return f"max_new_tokens must be at least {_TTS_MAX_NEW_TOKENS_MIN}"
if request.max_new_tokens > _TTS_MAX_NEW_TOKENS_MAX:
return f"max_new_tokens cannot exceed {_TTS_MAX_NEW_TOKENS_MAX}"
return None
def _voxcpm2_encode(self, text: str) -> list[int]:
"""Tokenize text for VoxCPM2, splitting multichar Chinese tokens."""
from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
@@ -1315,11 +1390,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
# Infer Base task when ref_audio or ref_text is provided without explicit task_type.
if request.task_type is None and (request.ref_audio is not None or request.ref_text is not None):
request.task_type = "Base"
task_type = request.task_type or "CustomVoice"
# Normalize voice to lowercase for case-insensitive matching
if request.voice is not None:
request.voice = request.voice.lower()
if request.task_type is None and request.voice in self.precomputed_speakers:
request.task_type = "Base"
task_type = request.task_type or "CustomVoice"
# Validate input is not empty
if not request.input or not request.input.strip():
@@ -1379,6 +1456,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
file_path = Path(speaker_info["file_path"])
if not file_path.exists():
return f"Data file for uploaded speaker '{request.voice}' not found on disk"
elif voice_lower in self.precomputed_speakers:
profile = self.precomputed_speakers[voice_lower]
mode = str(profile.get("mode") or "xvec").lower()
ref_text = request.ref_text or profile.get("ref_text")
if mode == "icl" and (not isinstance(ref_text, str) or not ref_text.strip()):
return (
f"Precomputed voice '{request.voice}' uses ICL mode but has no ref_text in "
"the request or custom voice manifest"
)
else:
# need ref_audio for built-in speaker
if request.ref_audio is None:
@@ -1677,6 +1763,19 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
URLs, ``data:`` base64 URIs, and ``file:`` local paths (the latter
gated by ``--allowed-local-media-path``).
"""
cache_key = hashlib.sha1(ref_audio_str.encode("utf-8")).hexdigest()
cached = self._ref_audio_resolve_cache.get(cache_key)
if cached is not None:
self._ref_audio_resolve_cache.move_to_end(cache_key)
wav_list, sr, _, _ = cached
logger.debug(
"Resolved ref_audio from cache: samples=%d sr=%d duration_s=%.3f",
len(wav_list),
sr,
len(wav_list) / sr if sr > 0 else 0.0,
)
return wav_list, sr
# In diffusion mode, model_config may not be available
if self._diffusion_mode:
connector = MediaConnector()
@@ -1686,11 +1785,14 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
allowed_local_media_path=model_config.allowed_local_media_path,
allowed_media_domains=model_config.allowed_media_domains,
)
fetch_start_s = time.perf_counter()
wav_np, sr = await connector.fetch_audio_async(ref_audio_str)
fetch_decode_ms = (time.perf_counter() - fetch_start_s) * 1000.0
wav_np = np.asarray(wav_np, dtype=np.float32)
if wav_np.ndim > 1:
wav_np = np.mean(wav_np, axis=-1)
sr = int(sr)
artifact_key = self._make_ref_audio_artifact_cache_key(wav_np, sr)
duration = len(wav_np) / sr if sr > 0 else 0.0
if duration < _REF_AUDIO_MIN_DURATION:
raise ValueError(
@@ -1702,7 +1804,83 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
f"Reference audio too long ({duration:.1f}s). "
f"Maximum {_REF_AUDIO_MAX_DURATION:.0f}s supported — use a shorter clip."
)
return wav_np.tolist(), sr
tolist_start_s = time.perf_counter()
wav_list = wav_np.tolist()
tolist_ms = (time.perf_counter() - tolist_start_s) * 1000.0
logger.debug(
"Resolved ref_audio: fetch_decode_ms=%.3f tolist_ms=%.3f samples=%d sr=%d duration_s=%.3f",
fetch_decode_ms,
tolist_ms,
len(wav_np),
sr,
duration,
)
self._put_resolved_ref_audio(cache_key, wav_list, sr, artifact_key)
return wav_list, sr
@staticmethod
def _make_ref_audio_artifact_cache_key(wav: np.ndarray, sr: int) -> str:
wav_f32 = wav.astype(np.float32, copy=False).reshape(-1)
h = hashlib.sha1()
h.update(int(sr).to_bytes(4, byteorder="little", signed=False))
h.update(int(wav_f32.size).to_bytes(8, byteorder="little", signed=False))
h.update(wav_f32.tobytes(order="C"))
return h.hexdigest()
def _get_resolved_ref_audio_artifact_key(self, ref_audio_str: str) -> str | None:
source_key = hashlib.sha1(ref_audio_str.encode("utf-8")).hexdigest()
cached = self._ref_audio_resolve_cache.get(source_key)
if cached is None:
return None
self._ref_audio_resolve_cache.move_to_end(source_key)
return cached[3]
def _put_resolved_ref_audio(self, cache_key: str, wav_list: list[float], sr: int, artifact_key: str) -> None:
if self._ref_audio_resolve_cache_max_entries <= 0 or self._ref_audio_resolve_cache_max_bytes <= 0:
return
# Approximate list[float] storage. CPython float objects add per-element
# overhead, so max_entries remains the hard cache cap.
size = len(wav_list) * 40
if size > self._ref_audio_resolve_cache_max_bytes:
return
previous = self._ref_audio_resolve_cache.pop(cache_key, None)
if previous is not None:
self._ref_audio_resolve_cache_bytes -= previous[2]
if previous[3] != artifact_key:
self._discard_ref_audio_artifact_ready_if_unreferenced(previous[3])
self._ref_audio_resolve_cache[cache_key] = (wav_list, int(sr), size, artifact_key)
self._ref_audio_resolve_cache_bytes += size
while len(self._ref_audio_resolve_cache) > self._ref_audio_resolve_cache_max_entries:
_, (_, _, old_size, old_artifact_key) = self._ref_audio_resolve_cache.popitem(last=False)
self._ref_audio_resolve_cache_bytes -= old_size
self._discard_ref_audio_artifact_ready_if_unreferenced(old_artifact_key)
while self._ref_audio_resolve_cache_bytes > self._ref_audio_resolve_cache_max_bytes:
_, (_, _, old_size, old_artifact_key) = self._ref_audio_resolve_cache.popitem(last=False)
self._ref_audio_resolve_cache_bytes -= old_size
self._discard_ref_audio_artifact_ready_if_unreferenced(old_artifact_key)
def _discard_ref_audio_artifact_ready_if_unreferenced(self, artifact_key: str) -> None:
if artifact_key and all(entry[3] != artifact_key for entry in self._ref_audio_resolve_cache.values()):
self._ref_audio_model_artifact_ready.discard(artifact_key)
def _qwen3_tts_can_use_ref_audio_artifact_only(self, tts_params: dict[str, Any], artifact_key: str | None) -> bool:
if self._tts_model_type != "qwen3_tts":
return False
if not artifact_key or artifact_key not in self._ref_audio_model_artifact_ready:
return False
return (tts_params.get("task_type") or ["CustomVoice"])[0] == "Base"
def _track_ref_audio_artifact_warmup(self, request_id: str, artifact_key: str | None) -> None:
if artifact_key:
self._request_ref_audio_artifact_keys[request_id] = artifact_key
def _mark_ref_audio_artifact_ready_for_request(self, request_id: str) -> None:
artifact_key = self._request_ref_audio_artifact_keys.pop(request_id, None)
if artifact_key and any(entry[3] == artifact_key for entry in self._ref_audio_resolve_cache.values()):
self._ref_audio_model_artifact_ready.add(artifact_key)
def _discard_ref_audio_artifact_warmup(self, request_id: str) -> None:
self._request_ref_audio_artifact_keys.pop(request_id, None)
async def _generate_audio_chunks(
self,
@@ -1733,6 +1911,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
first_chunk = True
first_audio_chunk_s: float | None = None
stream_start_s = request_start_s if request_start_s is not None else time.perf_counter()
artifact_ready = False
try:
async for res in generator:
@@ -1787,6 +1966,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
if first_audio_chunk_s is None:
first_audio_chunk_s = time.perf_counter()
yield self.create_audio(audio_obj).audio_data
self._mark_ref_audio_artifact_ready_for_request(request_id)
artifact_ready = True
total_ms = (time.perf_counter() - stream_start_s) * 1000.0
if first_audio_chunk_s is not None:
first_chunk_ms = (first_audio_chunk_s - stream_start_s) * 1000.0
@@ -1840,6 +2021,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
)
logger.exception("Streaming speech generation failed for %s: %s", request_id, e)
raise
finally:
if not artifact_ready:
self._discard_ref_audio_artifact_warmup(request_id)
@staticmethod
def _extract_audio_output(res) -> tuple[dict | None, str | None]:
@@ -1893,13 +2077,14 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
# Speaker (voice)
if request.voice is not None:
voice_lower = request.voice.lower()
params["speaker"] = [request.voice]
params["voice_created_at"] = [self._voice_created_at(request.voice.lower())]
params["voice_created_at"] = [self._voice_created_at(voice_lower)]
# Uploaded voices use task_type="Base" (CustomVoice requires built-in spk_id).
# If ref_text was provided at upload time, use in-context cloning; otherwise x_vector only.
if request.voice.lower() in self.uploaded_speakers and request.ref_audio is None:
speaker_info = self.uploaded_speakers[request.voice.lower()]
if voice_lower in self.uploaded_speakers and request.ref_audio is None:
speaker_info = self.uploaded_speakers[voice_lower]
# Check if this voice was uploaded with a pre-computed embedding.
# Populate request.speaker_embedding so the existing code path
@@ -1907,6 +2092,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
embedding = self._get_uploaded_speaker_embedding(request.voice)
if embedding is not None:
request.speaker_embedding = embedding
params["speaker"] = [voice_lower]
params["task_type"] = ["Base"]
logger.info("Auto-set speaker_embedding for uploaded voice: %s", request.voice)
else:
@@ -1914,6 +2100,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
if not audio_data:
raise ValueError(f"Audio file for uploaded voice '{request.voice}' is missing or corrupted")
stored_ref_text = speaker_info.get("ref_text")
params["speaker"] = [voice_lower]
params["ref_audio"] = [audio_data]
params["task_type"] = ["Base"]
if stored_ref_text:
@@ -1924,6 +2111,19 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
logger.info(
"Auto-set ref_audio for uploaded voice: %s (icl=%s)", request.voice, bool(stored_ref_text)
)
elif voice_lower in self.precomputed_speakers and request.ref_audio is None:
profile = self.precomputed_speakers[voice_lower]
mode = str(profile.get("mode") or "xvec").lower()
params["speaker"] = [voice_lower]
params["task_type"] = ["Base"]
params["x_vector_only_mode"] = [mode != "icl"]
ref_text = request.ref_text or profile.get("ref_text")
if isinstance(ref_text, str) and ref_text.strip():
params["ref_text"] = [ref_text]
ref_code_length = profile.get("ref_code_length")
if mode == "icl" and ref_code_length:
params["ref_code_length"] = [int(ref_code_length)]
logger.info("Using precomputed Qwen3-TTS custom voice profile: %s (mode=%s)", voice_lower, mode)
elif params["task_type"][0] == "CustomVoice":
params["speaker"] = ["Vivian"] # Default for CustomVoice
@@ -2334,6 +2534,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
if self.engine_client.errored:
raise self.engine_client.dead_error
request_id = request_id or f"speech-{random_uuid()}"
qwen3_ref_audio_warmup_artifact_key: str | None = None
# If this is a streaming request, we need to coerce
# cumulative outputs to delta outputs; this ensures
# we don't emit redundant MM data & drain after emitting.
@@ -2412,12 +2615,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
elif self._tts_model_type == "voxcpm2":
# voxcpm2 doesn't use `_apply_uploaded_speaker` because the prompt builder needs the
# raw waveform tuple for prefill-length accounting, not a base64 data URL.
validation_error = self._validate_voxcpm2_request(request)
if validation_error:
raise ValueError(validation_error)
uploaded_ref: tuple[np.ndarray, int] | None = None
if request.voice:
voice_lower = request.voice.lower()
if voice_lower not in self.uploaded_speakers and voice_lower not in self.supported_speakers:
all_voices = sorted(self.uploaded_speakers.keys() | self.supported_speakers)
raise ValueError(f"Invalid voice '{request.voice}'. Supported: {', '.join(all_voices) or 'none'}")
if voice_lower in self.uploaded_speakers and not has_inline_ref_audio:
if self.uploaded_speakers[voice_lower].get("embedding_source") == "direct":
raise ValueError(
@@ -2430,7 +2634,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
tts_params = {}
if request.voice:
voice_lower = request.voice.lower()
if voice_lower in self.uploaded_speakers and not has_inline_ref_audio:
if voice_lower in self.uploaded_speakers or voice_lower in self.precomputed_speakers:
additional = prompt.setdefault("additional_information", {})
additional["voice_name"] = voice_lower
additional["voice_created_at"] = self._voice_created_at(voice_lower)
@@ -2478,7 +2682,18 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
ref_audio_source = tts_params["ref_audio"][0]
if ref_audio_source is not None and isinstance(ref_audio_source, str):
wav_list, sr = await self._resolve_ref_audio(ref_audio_source)
tts_params["ref_audio"] = [[wav_list, sr]]
artifact_key = self._get_resolved_ref_audio_artifact_key(ref_audio_source)
if self._tts_model_type == "qwen3_tts" and artifact_key:
tts_params[_QWEN3_TTS_REF_AUDIO_CACHE_KEY] = [artifact_key]
ref_code_length = self._estimate_ref_code_len([wav_list, sr])
if self._tts_model_type == "qwen3_tts" and ref_code_length is not None:
tts_params["ref_code_length"] = [int(ref_code_length)]
if self._qwen3_tts_can_use_ref_audio_artifact_only(tts_params, artifact_key):
logger.debug("Using Qwen3-TTS ref_audio artifact-only path: %s", artifact_key)
else:
tts_params["ref_audio"] = [[wav_list, sr]]
if self._tts_model_type == "qwen3_tts":
qwen3_ref_audio_warmup_artifact_key = artifact_key
ph_len = await self._estimate_prompt_len_async(tts_params)
prompt = tokens_input(prompt_token_ids=[1] * ph_len)
@@ -2506,7 +2721,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
tts_params = {}
prompt = {"prompt": request.input}
request_id = request_id or f"speech-{random_uuid()}"
if self._is_fish_speech:
model_type = "fish_speech"
elif self._tts_model_type == "covo_audio":
@@ -2643,6 +2857,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
sampling_params_list=sampling_params_list,
output_modalities=["audio"],
)
self._track_ref_audio_artifact_warmup(request_id, qwen3_ref_audio_warmup_artifact_key)
return request_id, generator, tts_params
async def _generate_pcm_chunks(self, generator, request_id: str):
@@ -2657,8 +2872,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
async def _iter_pcm_audio_bytes(self, request: OpenAICreateSpeechRequest):
"""Yield raw PCM bytes for a speech request as soon as chunks are decoded."""
request_id, generator, _ = await self._prepare_speech_generation(request)
async for chunk in self._generate_pcm_chunks(generator, request_id):
yield chunk
try:
async for chunk in self._generate_pcm_chunks(generator, request_id):
yield chunk
finally:
self._discard_ref_audio_artifact_warmup(request_id)
async def _generate_audio_bytes(
self,
@@ -2667,101 +2885,108 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
request_id: str | None = None,
) -> tuple[bytes | str, str]:
request_id, generator, _ = await self._prepare_speech_generation(request, request_id=request_id)
artifact_ready = False
# MOSS-TTS-Nano emits delta chunks per yield (single-stage,
# async_chunk=false). The engine surfaces each yield as its own
# RequestOutput, so we need to accumulate across the async-for loop —
# final_output alone only carries the last (often empty) sentinel.
is_moss = self._tts_model_type == "moss_tts_nano"
moss_chunks: list[Any] = []
moss_sample_rate: int | None = None
try:
# MOSS-TTS-Nano emits delta chunks per yield (single-stage,
# async_chunk=false). The engine surfaces each yield as its own
# RequestOutput, so we need to accumulate across the async-for loop —
# final_output alone only carries the last (often empty) sentinel.
is_moss = self._tts_model_type == "moss_tts_nano"
moss_chunks: list[Any] = []
moss_sample_rate: int | None = None
final_output: OmniRequestOutput | None = None
async for res in generator:
final_output = res
if not is_moss:
continue
try:
step_audio, step_key = self._extract_audio_output(res)
except Exception:
continue
if step_key is None:
continue
chunk = step_audio[step_key]
candidates = chunk if isinstance(chunk, list) else [chunk]
for cand in candidates:
if hasattr(cand, "numel") and cand.numel() > 0:
moss_chunks.append(cand)
sr_step = step_audio.get("sr")
if sr_step is not None:
sr_val_step = sr_step[-1] if isinstance(sr_step, list) and sr_step else sr_step
moss_sample_rate = int(sr_val_step.item()) if hasattr(sr_val_step, "item") else int(sr_val_step)
final_output: OmniRequestOutput | None = None
async for res in generator:
final_output = res
if not is_moss:
continue
try:
step_audio, step_key = self._extract_audio_output(res)
except Exception:
continue
if step_key is None:
continue
chunk = step_audio[step_key]
candidates = chunk if isinstance(chunk, list) else [chunk]
for cand in candidates:
if hasattr(cand, "numel") and cand.numel() > 0:
moss_chunks.append(cand)
sr_step = step_audio.get("sr")
if sr_step is not None:
sr_val_step = sr_step[-1] if isinstance(sr_step, list) and sr_step else sr_step
moss_sample_rate = int(sr_val_step.item()) if hasattr(sr_val_step, "item") else int(sr_val_step)
if final_output is None:
raise ValueError("No output generated from the model.")
if final_output is None:
raise ValueError("No output generated from the model.")
audio_output, audio_key = self._extract_audio_output(final_output)
if audio_key is None:
raise ValueError("TTS model did not produce audio output.")
audio_output, audio_key = self._extract_audio_output(final_output)
if audio_key is None:
raise ValueError("TTS model did not produce audio output.")
audio_tensor = audio_output[audio_key]
sr_raw = audio_output.get("sr", 24000)
sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
sample_rate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
audio_tensor = audio_output[audio_key]
sr_raw = audio_output.get("sr", 24000)
sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
sample_rate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
if is_moss:
# Prefer the engine's own consolidated audio when present. After the
# vllm 0.20 rebase non-stream requests resolve to FINAL_ONLY, so
# final_output already carries the full concatenated waveform; the
# delta-accumulator below is kept as a fallback for DELTA-style
# engines that surface chunks one yield at a time.
if isinstance(audio_tensor, list):
non_empty_final = [c for c in audio_tensor if hasattr(c, "numel") and c.numel() > 0]
final_audio = torch.cat(non_empty_final, dim=-1) if non_empty_final else None
elif hasattr(audio_tensor, "numel") and audio_tensor.numel() > 0:
final_audio = audio_tensor
else:
final_audio = None
if is_moss:
# Prefer the engine's own consolidated audio when present. After the
# vllm 0.20 rebase non-stream requests resolve to FINAL_ONLY, so
# final_output already carries the full concatenated waveform; the
# delta-accumulator below is kept as a fallback for DELTA-style
# engines that surface chunks one yield at a time.
if isinstance(audio_tensor, list):
non_empty_final = [c for c in audio_tensor if hasattr(c, "numel") and c.numel() > 0]
final_audio = torch.cat(non_empty_final, dim=-1) if non_empty_final else None
elif hasattr(audio_tensor, "numel") and audio_tensor.numel() > 0:
final_audio = audio_tensor
else:
final_audio = None
if final_audio is not None:
audio_tensor = final_audio
elif moss_chunks:
audio_tensor = torch.cat(moss_chunks, dim=-1)
else:
audio_tensor = np.zeros((0,), dtype=np.float32)
if moss_sample_rate is not None:
sample_rate = moss_sample_rate
elif isinstance(audio_tensor, list):
async_chunk = bool(getattr(self.engine_client.model_config, "async_chunk", False))
if async_chunk:
non_empty_chunks = [candidate for candidate in audio_tensor if candidate.numel() > 0]
audio_tensor = (
torch.cat(non_empty_chunks, dim=-1) if non_empty_chunks else np.zeros((0,), dtype=np.float32)
)
else:
audio_history = audio_tensor
audio_tensor = np.zeros((0,), dtype=np.float32)
# Non-async Qwen3-TTS returns cumulative history snapshots, so keep the latest non-empty tensor.
for candidate in reversed(audio_history):
if candidate.numel() > 0:
audio_tensor = candidate
break
if hasattr(audio_tensor, "float"):
audio_tensor = audio_tensor.float().detach().cpu().numpy()
if final_audio is not None:
audio_tensor = final_audio
elif moss_chunks:
audio_tensor = torch.cat(moss_chunks, dim=-1)
else:
audio_tensor = np.zeros((0,), dtype=np.float32)
if moss_sample_rate is not None:
sample_rate = moss_sample_rate
elif isinstance(audio_tensor, list):
async_chunk = bool(getattr(self.engine_client.model_config, "async_chunk", False))
if async_chunk:
non_empty_chunks = [candidate for candidate in audio_tensor if candidate.numel() > 0]
audio_tensor = (
torch.cat(non_empty_chunks, dim=-1) if non_empty_chunks else np.zeros((0,), dtype=np.float32)
)
else:
audio_history = audio_tensor
audio_tensor = np.zeros((0,), dtype=np.float32)
# Non-async Qwen3-TTS returns cumulative history snapshots, so keep the latest non-empty tensor.
for candidate in reversed(audio_history):
if candidate.numel() > 0:
audio_tensor = candidate
break
if hasattr(audio_tensor, "float"):
audio_tensor = audio_tensor.float().detach().cpu().numpy()
if audio_tensor.ndim > 1:
audio_tensor = audio_tensor.squeeze()
if audio_tensor.ndim > 1:
audio_tensor = audio_tensor.squeeze()
audio_obj = CreateAudio(
audio_tensor=audio_tensor,
sample_rate=sample_rate,
response_format=request.response_format or "wav",
speed=request.speed or 1.0,
stream_format=request.stream_format,
base64_encode=base64_encode,
)
audio_response: AudioResponse = self.create_audio(audio_obj)
return audio_response.audio_data, audio_response.media_type
audio_obj = CreateAudio(
audio_tensor=audio_tensor,
sample_rate=sample_rate,
response_format=request.response_format or "wav",
speed=request.speed or 1.0,
stream_format=request.stream_format,
base64_encode=base64_encode,
)
audio_response: AudioResponse = self.create_audio(audio_obj)
self._mark_ref_audio_artifact_ready_for_request(request_id)
artifact_ready = True
return audio_response.audio_data, audio_response.media_type
finally:
if not artifact_ready:
self._discard_ref_audio_artifact_warmup(request_id)
async def _create_diffusion_speech(
self,
@@ -757,6 +757,14 @@ class Qwen3TTSPromptEmbedsBuilder:
ref_audio_list = info_dict.get("ref_audio")
if not isinstance(ref_audio_list, list) or not ref_audio_list:
continue
cache_key_from_serving = first_value(info_dict.get(REF_AUDIO_CACHE_KEY))
if isinstance(cache_key_from_serving, str) and cache_key_from_serving:
cached = self.get_ref_audio_artifacts(cache_key_from_serving)
if cached is not None:
cached_ref_code = coerce_ref_code_tensor(cached.get("ref_code"), device=device)
if isinstance(cached_ref_code, torch.Tensor):
info_dict.setdefault("codes", {})[PRECOMPUTED_REF_CODE_KEY] = cached_ref_code
continue
try:
wav, sr = self.normalize_ref_audio(ref_audio_list[0])
except Exception:
@@ -1029,6 +1037,19 @@ class Qwen3TTSPromptEmbedsBuilder:
xvec_only = bool((info_dict.get("x_vector_only_mode") or [False])[0])
in_context_mode = not xvec_only
voice_clone_prompt = _normalize_voice_clone_prompt(info_dict.get("voice_clone_prompt"))
if voice_clone_prompt is not None and "icl_mode" in voice_clone_prompt:
icl_flag = _as_singleton(voice_clone_prompt.get("icl_mode"))
if isinstance(icl_flag, bool):
in_context_mode = icl_flag
xvec_only = not in_context_mode
prompt_has_ref_code = voice_clone_prompt is not None and self._has_ref_code_like(
voice_clone_prompt.get("ref_code")
)
prompt_spk = voice_clone_prompt.get("ref_spk_embedding") if voice_clone_prompt is not None else None
prompt_has_speaker = isinstance(prompt_spk, torch.Tensor) or (
isinstance(prompt_spk, (list, np.ndarray)) and len(prompt_spk) > 0
)
needs_ref_audio_cache_key = not (prompt_has_speaker and (not in_context_mode or prompt_has_ref_code))
ref_audio_wav: np.ndarray | None = None
ref_audio_sr: int | None = None
@@ -1056,7 +1077,8 @@ class Qwen3TTSPromptEmbedsBuilder:
# from ref_audio: x_vector_only_mode, precomputed ref_code
# with no ref_spk_embedding, voice_clone_prompt.icl_mode=False.
ref_audio_cache_key = info_dict.pop(REF_AUDIO_CACHE_KEY, None)
if not (isinstance(ref_audio_cache_key, str) and ref_audio_cache_key):
ref_audio_cache_key = first_value(ref_audio_cache_key)
if needs_ref_audio_cache_key and not (isinstance(ref_audio_cache_key, str) and ref_audio_cache_key):
normalized = info_dict.get(NORMALIZED_REF_AUDIO_KEY)
wav_peek: np.ndarray | None = None
sr_peek: int | None = None
@@ -1075,13 +1097,25 @@ class Qwen3TTSPromptEmbedsBuilder:
cached_artifacts = None
if isinstance(ref_audio_cache_key, str) and ref_audio_cache_key:
cached_artifacts = self.get_ref_audio_artifacts(ref_audio_cache_key)
ref_audio_list = info_dict.get("ref_audio")
has_ref_audio_payload = isinstance(ref_audio_list, list) and bool(ref_audio_list)
artifact_only = (
isinstance(ref_audio_cache_key, str) and bool(ref_audio_cache_key) and not has_ref_audio_payload
)
if artifact_only and cached_artifacts is None:
raise RuntimeError(
"Qwen3-TTS ref_audio artifact cache miss for artifact-only request; "
"retry with ref_audio or precompute a custom voice profile."
)
# Speaker cache: only for uploaded (named) speakers
_speaker_cache_key = None
_cache_lookup_voice = None
if voice_clone_prompt is None and self._speaker_cache is not None:
_speaker_list = info_dict.get("speaker")
if isinstance(_speaker_list, list) and _speaker_list:
_voice_name = str(_speaker_list[0]).lower()
_cache_lookup_voice = _voice_name
# Per-mode namespace — xvec and icl produce different artifacts
# for the same voice, so they must not share a cache slot.
_mode = "xvec" if xvec_only else "icl"
@@ -1103,15 +1137,18 @@ class Qwen3TTSPromptEmbedsBuilder:
"ref_code": ref_code_cached,
"ref_spk_embedding": ref_spk_embed_cached,
"icl_mode": _cached.get("icl_mode"),
"ref_text": _cached.get("ref_text"),
}
_speaker_cache_key = None # hit → don't store again
# Official implementation may pass `voice_clone_prompt.icl_mode`.
if voice_clone_prompt is not None and "icl_mode" in voice_clone_prompt:
icl_flag = _as_singleton(voice_clone_prompt.get("icl_mode"))
if isinstance(icl_flag, bool):
in_context_mode = icl_flag
xvec_only = not in_context_mode
if voice_clone_prompt is None and _speaker_cache_key is not None:
ref_audio_list = info_dict.get("ref_audio")
if not isinstance(ref_audio_list, list) or not ref_audio_list:
raise ValueError(
f"Qwen3-TTS speaker '{_cache_lookup_voice}' was requested without ref_audio, "
"but no precomputed cache entry was loaded"
)
ref_code = None
if voice_clone_prompt is not None:
ref_code = _as_singleton(voice_clone_prompt.get("ref_code"))
@@ -1136,6 +1173,8 @@ class Qwen3TTSPromptEmbedsBuilder:
if isinstance(cached_ref_code, torch.Tensor):
ref_code_t = cached_ref_code
ref_code_len = int(ref_code_t.shape[0])
if ref_code_t is None and in_context_mode and artifact_only:
raise RuntimeError("Qwen3-TTS ref_audio artifact cache entry is missing ref_code.")
if ref_code_t is None and in_context_mode:
# Compute ref_code from ref_audio if not provided.
wav_np, sr = _get_ref_audio()
@@ -1161,6 +1200,8 @@ class Qwen3TTSPromptEmbedsBuilder:
.to(device=input_ids.device, dtype=torch.bfloat16)
.view(1, 1, -1)
)
elif artifact_only:
raise RuntimeError("Qwen3-TTS ref_audio artifact cache entry is missing ref_spk_embedding.")
else:
wav_np, sr = _get_ref_audio()
speaker_embed = self.extract_speaker_embedding(wav_np, sr).view(1, 1, -1)
@@ -1207,6 +1248,8 @@ class Qwen3TTSPromptEmbedsBuilder:
)
if ref_ids is None:
ref_text = _as_singleton(info_dict.get("ref_text"))
if (not isinstance(ref_text, str) or not ref_text.strip()) and voice_clone_prompt is not None:
ref_text = _as_singleton(voice_clone_prompt.get("ref_text"))
if isinstance(ref_text, str) and ref_text.strip():
ref_ids = tok(
build_ref_text(ref_text),
@@ -1489,6 +1532,13 @@ class Qwen3TTSPromptEmbedsBuilder:
except Exception:
ref_code_len = None
if ref_code_len is None:
ref_code_length = _first(info.get("ref_code_length"), None)
try:
if ref_code_length is not None:
ref_code_len = int(ref_code_length) # type: ignore[arg-type]
except (TypeError, ValueError):
ref_code_len = None
if ref_code_len is None and estimate_ref_code_len is not None:
ref_code_len = estimate_ref_code_len(info.get("ref_audio"))
if ref_code_len is None:
@@ -1511,6 +1561,10 @@ class Qwen3TTSPromptEmbedsBuilder:
if ref_ids is None:
ref_text = _first(info.get("ref_text"), "")
if (not isinstance(ref_text, str) or not ref_text.strip()) and isinstance(
voice_clone_prompt, dict
):
ref_text = _first(voice_clone_prompt.get("ref_text"), "")
if not isinstance(ref_text, str) or not ref_text.strip():
raise ValueError(
"Base in-context non-streaming requires `ref_text` or tokenized `ref_ids`."
@@ -1,7 +1,7 @@
from __future__ import annotations
import os
from collections import Counter
from collections import Counter, OrderedDict
from collections.abc import Iterable
from typing import Any
@@ -24,6 +24,9 @@ from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
logger = init_logger(__name__)
_REF_CONTEXT_CACHE_MAX_ENTRIES = 4096
_REF_CONTEXT_CACHE_MAX_BYTES = 64 * 1024 * 1024
def _codec_ids_from_payload_or_input(
input_ids: torch.Tensor,
@@ -100,6 +103,10 @@ class Qwen3TTSCode2Wav(nn.Module):
self._output_sample_rate = int(tok_config.output_sample_rate)
self._total_upsample = int(self.decoder.total_upsample)
self._decoder_sliding_window = int(getattr(dec_config, "sliding_window", 0) or 0)
self._ref_context_cache: OrderedDict[str, torch.Tensor] = OrderedDict()
self._ref_context_cache_bytes = 0
self._ref_context_cache_max_entries = _REF_CONTEXT_CACHE_MAX_ENTRIES
self._ref_context_cache_max_bytes = _REF_CONTEXT_CACHE_MAX_BYTES
def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
# This stage ignores token embeddings. Keep a stable dummy embedding for vLLM runner.
@@ -205,6 +212,49 @@ class Qwen3TTSCode2Wav(nn.Module):
self._batch_stats_actual_frames.update(actual_frames)
self._batch_stats_bucket_groups[(group_size, bucket_frames)] += 1
@staticmethod
def _tensor_nbytes(tensor: torch.Tensor) -> int:
return int(tensor.numel() * tensor.element_size())
def _evict_ref_context_cache_if_needed(self) -> None:
evicted = 0
while len(self._ref_context_cache) > self._ref_context_cache_max_entries:
_, cached = self._ref_context_cache.popitem(last=False)
self._ref_context_cache_bytes -= self._tensor_nbytes(cached)
evicted += 1
while self._ref_context_cache_bytes > self._ref_context_cache_max_bytes and len(self._ref_context_cache) > 1:
_, cached = self._ref_context_cache.popitem(last=False)
self._ref_context_cache_bytes -= self._tensor_nbytes(cached)
evicted += 1
if evicted:
logger.debug(
"Evicted %d Qwen3-TTS ref context cache entries; entries=%d bytes=%d",
evicted,
len(self._ref_context_cache),
self._ref_context_cache_bytes,
)
def _cache_ref_context(self, request_id: str, tensor: torch.Tensor) -> None:
previous = self._ref_context_cache.pop(request_id, None)
if previous is not None:
self._ref_context_cache_bytes -= self._tensor_nbytes(previous)
cached = tensor.detach().contiguous()
self._ref_context_cache[request_id] = cached
self._ref_context_cache.move_to_end(request_id)
self._ref_context_cache_bytes += self._tensor_nbytes(cached)
self._evict_ref_context_cache_if_needed()
def _get_ref_context(self, request_id: str) -> torch.Tensor | None:
cached = self._ref_context_cache.get(request_id)
if cached is not None:
self._ref_context_cache.move_to_end(request_id)
return cached
def _pop_ref_context(self, request_id: str) -> None:
cached = self._ref_context_cache.pop(request_id, None)
if cached is not None:
self._ref_context_cache_bytes -= self._tensor_nbytes(cached)
def log_decode_batch_stats(self) -> None:
if not self._batch_stats_enabled or self._batch_stats_requests == 0:
return
@@ -268,6 +318,32 @@ class Qwen3TTSCode2Wav(nn.Module):
valid_codes_qf: list[torch.Tensor] = []
valid_indices: list[int] = []
left_context_size = [0] * len(request_ids_list)
ref_context_size = [0] * len(request_ids_list)
ref_context_request_ids: list[str | None] = [None] * len(request_ids_list)
ref_context_included = [False] * len(request_ids_list)
finished_flags = [False] * len(request_ids_list)
def _meta_int(value: Any) -> int:
if isinstance(value, list):
value = value[0] if value else 0
if isinstance(value, torch.Tensor):
value = value.reshape(-1)[0].item() if value.numel() > 0 else 0
return int(value or 0)
def _meta_str(value: Any) -> str | None:
if isinstance(value, list):
value = value[0] if value else None
if value is None:
return None
return str(value)
def _meta_bool(value: Any) -> bool:
if isinstance(value, list):
value = value[0] if value else False
if isinstance(value, torch.Tensor):
return bool(value.reshape(-1)[0].item()) if value.numel() > 0 else False
return bool(value)
if runtime_infos:
for i, info in enumerate(runtime_infos):
if i >= len(left_context_size):
@@ -276,7 +352,15 @@ class Qwen3TTSCode2Wav(nn.Module):
continue
meta = info.get("meta", {})
if "left_context_size" in meta:
left_context_size[i] = meta["left_context_size"]
left_context_size[i] = _meta_int(meta["left_context_size"])
if "ref_context_size" in meta:
ref_context_size[i] = _meta_int(meta["ref_context_size"])
if "ref_context_request_id" in meta:
ref_context_request_ids[i] = _meta_str(meta["ref_context_request_id"])
if "ref_context_included" in meta:
ref_context_included[i] = _meta_bool(meta["ref_context_included"])
if "finished" in meta:
finished_flags[i] = _meta_bool(meta["finished"])
for i, req_ids in enumerate(request_ids_list):
runtime_info = runtime_infos[i] if i < len(runtime_infos) else None
req_ids = _codec_ids_from_payload_or_input(req_ids, runtime_info)
@@ -284,6 +368,7 @@ class Qwen3TTSCode2Wav(nn.Module):
parsed.append((0, 0))
continue
ctx_frames = left_context_size[i]
ref_ctx_frames = ref_context_size[i]
flat = req_ids
n = flat.numel()
if n == 0 or n % q != 0:
@@ -302,6 +387,25 @@ class Qwen3TTSCode2Wav(nn.Module):
frames = n // q
# [q*F] -> [Q, F] for direct decoder call (decoder expects [B, Q, F])
codes_qf = flat.reshape(q, frames)
ref_req_id = ref_context_request_ids[i]
if ref_req_id is not None and ref_ctx_frames > 0:
if ref_context_included[i]:
if frames < ref_ctx_frames:
raise ValueError(
"Qwen3-TTS ref context metadata says ref prefix is included, "
f"but frames={frames} < ref_context_size={ref_ctx_frames}"
)
self._cache_ref_context(ref_req_id, codes_qf[:, :ref_ctx_frames])
else:
cached_ref = self._get_ref_context(ref_req_id)
if cached_ref is None:
raise ValueError(
"Missing Qwen3-TTS ref context cache for "
f"request {ref_req_id!r}; first chunk must include ref_code"
)
cached_ref = cached_ref.to(device=codes_qf.device, dtype=codes_qf.dtype)
codes_qf = torch.cat((cached_ref, codes_qf), dim=1)
frames = int(codes_qf.shape[1])
parsed.append((ctx_frames, frames))
valid_codes_qf.append(codes_qf)
valid_indices.append(i)
@@ -443,6 +547,10 @@ class Qwen3TTSCode2Wav(nn.Module):
if wav.shape[0] > 0:
audios[idx] = wav.to(dtype=torch.float32).reshape(-1)
for req_id, finished in zip(ref_context_request_ids, finished_flags, strict=False):
if req_id is not None and finished:
self._pop_ref_context(req_id)
return OmniOutput(
text_hidden_states=None,
multimodal_outputs={"model_outputs": audios, "sr": srs},
@@ -23,7 +23,11 @@ from vllm.sequence import IntermediateTensors
from vllm_omni.data_entry_keys import OmniPayload
from vllm_omni.model_executor.models.output_templates import OmniOutput
from vllm_omni.utils.speaker_cache import get_speaker_cache
from vllm_omni.utils.speaker_cache import (
get_speaker_cache,
iter_custom_voice_profiles,
load_validated_profile_tensors,
)
from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSSpeakerEncoderConfig, Qwen3TTSTalkerConfig
from .prompt_embeds_builder import Qwen3TTSPromptEmbedsBuilder
@@ -441,6 +445,46 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
encode_ref_audio_batch=self._encode_ref_audio_batch,
speaker_cache=self._speaker_cache,
)
self._load_custom_voice_profiles()
# -------------------- custom voice profiles --------------------
def _load_custom_voice_profiles(self) -> None:
"""Preload offline Qwen3-TTS custom voice profiles into speaker cache."""
custom_voice_dir = getattr(self.config, "custom_voice_dir", None)
if not custom_voice_dir:
return
expected_dim = int(getattr(self.config.speaker_encoder_config, "enc_dim", 0) or 0)
loaded = 0
for profile in iter_custom_voice_profiles(custom_voice_dir, expected_model_type="qwen3_tts"):
tensors = load_validated_profile_tensors(
profile,
expected_model_type="qwen3_tts",
qwen3_embedding_dim=expected_dim,
)
if tensors is None:
continue
speaker_embedding = tensors["speaker_embedding"].reshape(-1).contiguous().cpu()
mode = str(profile.get("mode") or "xvec").lower()
ref_code = tensors.get("ref_code")
artifacts: dict[str, Any] = {
"ref_spk_embedding": speaker_embedding,
"ref_code": ref_code.contiguous().cpu() if isinstance(ref_code, torch.Tensor) else None,
"icl_mode": mode == "icl",
"ref_text": profile.get("ref_text"),
}
key = self._speaker_cache.make_cache_key(
profile["voice_name_lower"],
model_type=f"qwen3_tts_{mode}",
created_at=0,
)
self._speaker_cache.put(key, artifacts)
loaded += 1
if loaded:
logger.info("Loaded %d precomputed Qwen3-TTS custom voice profile(s) from %s", loaded, custom_voice_dir)
# -------------------- vLLM required hooks --------------------
@@ -35,7 +35,11 @@ from vllm.platforms import current_platform
from vllm.sequence import IntermediateTensors
from vllm_omni.model_executor.models.output_templates import OmniOutput
from vllm_omni.utils.speaker_cache import get_speaker_cache
from vllm_omni.utils.speaker_cache import (
get_speaker_cache,
iter_custom_voice_profiles,
load_validated_profile_tensors,
)
from .minicpm4_paged import MiniCPM4PagedForVoxCPM2, MiniCPM4PagedResidualLM
from .voxcpm2_import_utils import import_voxcpm2_core
@@ -97,6 +101,7 @@ def build_voxcpm2_prompt(
ref_audio: Any | None = None,
ref_sr: int | None = None,
ref_text: str | None = None,
voice_profile: dict[str, Any] | None = None,
) -> dict[str, Any]:
"""Build a VoxCPM2 prefill prompt whose ``prompt_token_ids`` length matches
the talker-side prefill length.
@@ -110,7 +115,24 @@ def build_voxcpm2_prompt(
ids = ids[1:]
prefill_len = len(ids) + 1 # + audio_start
additional: dict[str, Any] = {"text_token_ids": [ids]}
if ref_audio is not None:
if voice_profile is not None and ref_audio is None:
mode = str(voice_profile.get("mode") or "reference").lower()
ref_audio_feat_len = int(voice_profile.get("ref_audio_feat_len") or 0)
audio_feat_len = int(voice_profile.get("audio_feat_len") or 0)
prompt_text = voice_profile.get("prompt_text") or voice_profile.get("ref_text")
additional["voice_profile"] = dict(voice_profile)
if mode in ("reference", "ref_continuation"):
prefill_len += ref_audio_feat_len + 2 # ref_start / ref_end
if mode in ("continuation", "ref_continuation"):
if isinstance(prompt_text, str) and prompt_text:
additional["prompt_text"] = [prompt_text]
prompt_ids = split_multichar_chinese(tokenizer.encode(prompt_text, add_special_tokens=True), split_map)
if prompt_ids and prompt_ids[0] == bos:
prompt_ids = prompt_ids[1:]
prefill_len += len(prompt_ids)
prefill_len += audio_feat_len
elif ref_audio is not None:
vae = hf_config.audio_vae_config
patch_samples = hf_config.patch_size * math.prod(vae["encoder_rates"])
ref_len = math.ceil(math.ceil(len(ref_audio) * vae["sample_rate"] / ref_sr) / patch_samples)
@@ -450,6 +472,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
# Speaker cache for ref_audio_feat across requests
self._speaker_cache = get_speaker_cache()
self._load_custom_voice_profiles()
self._perf = _PerfTimer(enabled=_ENABLE_PROFILING)
self._cfm_buffers: _CFMBufferManager | None = None
@@ -472,6 +495,51 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
# one-shot by design: fires at most once per process to avoid log spam.
self._active_state_warned = False
# -------------------- custom voice profiles --------------------
@staticmethod
def _clone_prompt_cache(cache: dict[str, Any]) -> dict[str, Any]:
cloned: dict[str, Any] = {}
for key, value in cache.items():
cloned[key] = value.clone() if isinstance(value, torch.Tensor) else value
return cloned
def _load_custom_voice_profiles(self) -> None:
"""Preload offline VoxCPM2 prompt caches into the shared speaker cache."""
custom_voice_dir = getattr(self.config, "custom_voice_dir", None)
if not custom_voice_dir:
return
loaded = 0
for profile in iter_custom_voice_profiles(custom_voice_dir, expected_model_type="voxcpm2"):
tensors = load_validated_profile_tensors(profile, expected_model_type="voxcpm2")
if tensors is None:
continue
ref_audio_feat = tensors.get("ref_audio_feat")
audio_feat = tensors.get("audio_feat")
mode = str(profile.get("mode") or "").lower()
prompt_cache: dict[str, Any] = {"mode": mode}
if ref_audio_feat is not None:
prompt_cache["ref_audio_feat"] = ref_audio_feat.contiguous().cpu()
if audio_feat is not None:
prompt_cache["audio_feat"] = audio_feat.contiguous().cpu()
prompt_text = profile.get("prompt_text") or profile.get("ref_text")
if isinstance(prompt_text, str) and prompt_text:
prompt_cache["prompt_text"] = prompt_text
key = self._speaker_cache.make_cache_key(
profile["voice_name_lower"],
model_type="voxcpm2",
created_at=0,
)
self._speaker_cache.put(key, prompt_cache)
loaded += 1
if loaded:
logger.info("Loaded %d precomputed VoxCPM2 custom voice profile(s) from %s", loaded, custom_voice_dir)
@property
def tts(self) -> nn.Module:
return self._tts
@@ -1161,6 +1229,10 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
prompt_audio = prompt_audio[0] if prompt_audio else None
if isinstance(prompt_text, list):
prompt_text = prompt_text[0] if prompt_text else None
voice_profile = info_dict.get("voice_profile")
if isinstance(voice_profile, list):
voice_profile = voice_profile[0] if voice_profile else None
requires_precomputed_cache = isinstance(voice_profile, dict)
state.prompt_cache = None
voice_name = info_dict.get("voice_name")
@@ -1168,20 +1240,29 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
voice_name = voice_name[0] if voice_name else None
_created_at = int(info_dict.get("voice_created_at") or 0)
if ref_audio or (prompt_audio and prompt_text):
# Check speaker cache for reference-only mode
if voice_name and ref_audio and not prompt_audio:
_cache_key = self._speaker_cache.make_cache_key(
voice_name, model_type="voxcpm2", created_at=_created_at
)
cached = self._speaker_cache.get(_cache_key)
if cached is not None:
if voice_name:
_cache_key = self._speaker_cache.make_cache_key(
voice_name, model_type="voxcpm2", created_at=_created_at
)
cached = self._speaker_cache.get(_cache_key)
if cached is not None:
if "mode" in cached:
state.prompt_cache = self._clone_prompt_cache(cached)
elif "ref_audio_feat" in cached:
state.prompt_cache = {
"mode": "reference",
"ref_audio_feat": cached["ref_audio_feat"].clone(),
}
logger.debug("Speaker cache HIT for VoxCPM2 speaker '%s'", voice_name)
logger.debug("Speaker cache HIT for VoxCPM2 speaker '%s'", voice_name)
if state.prompt_cache is None and requires_precomputed_cache:
requested_mode = voice_profile.get("mode") if isinstance(voice_profile, dict) else None
raise ValueError(
f"Precomputed VoxCPM2 voice '{voice_name}' was accepted by serving but is not loaded "
f"in the model cache (voice_created_at={_created_at}, mode={requested_mode})"
)
if state.prompt_cache is None and (ref_audio or (prompt_audio and prompt_text)):
if state.prompt_cache is None:
try:
state.prompt_cache = self._build_prompt_cache(
@@ -248,12 +248,29 @@ def talker2code2wav_async_chunk(
left_context_size = max(0, end_index - context_length)
window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:]
# Prepend a bounded ref_code tail as decoder context for every chunk so the
# vocoder keeps voice-clone speaker identity without making Stage1 shapes
# depend on full reference-audio length. The decoder is causal with sliding
# attention, so frames older than this context window cannot affect the
# emitted chunk. Use `.get()` (not `.pop()`) to keep ref_code for later chunks.
# Prepend the bounded ref_code tail to the first emitted chunk so Code2Wav
# can cache the same reference context that mainline would otherwise send on
# every chunk. Follow-up chunks send only a metadata handle; Code2Wav restores
# this context locally before decode.
ref_code = request_payload.get(request_id)
ref_context_size = 0
ref_context_request_id: str | None = None
ref_context_included = False
if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
if ref_code.ndim == 1:
num_quantizers = len(window_frames[0])
if ref_code.numel() % num_quantizers == 0:
ref_code = ref_code.reshape(-1, num_quantizers)
else:
logger.warning(
"Ignoring malformed ref_code with %d elements not divisible by num_quantizers=%d",
ref_code.numel(),
num_quantizers,
)
ref_code = None
elif ref_code.ndim != 2:
logger.warning("Ignoring malformed ref_code shape %s", tuple(ref_code.shape))
ref_code = None
if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
ref_context = ref_code
if ref_code_context_frames > 0 and int(ref_context.shape[0]) > ref_code_context_frames:
@@ -263,9 +280,15 @@ def talker2code2wav_async_chunk(
int(ref_context.shape[0]),
)
ref_context = ref_context[-ref_code_context_frames:]
ref_frames = ref_context.tolist()
window_frames = ref_frames + window_frames
left_context_size += len(ref_frames)
ref_context_size = int(ref_context.shape[0]) if ref_context.ndim > 1 else 0
if ref_context_size > 0:
ref_context_request_id = request_id
emitted_chunks = int(transfer_manager.put_req_chunk.get(request_id, 0))
if emitted_chunks <= 0:
ref_frames = ref_context.tolist()
window_frames = ref_frames + window_frames
ref_context_included = True
left_context_size += ref_context_size
num_quantizers = len(window_frames[0])
num_frames = len(window_frames)
@@ -274,12 +297,18 @@ def talker2code2wav_async_chunk(
dtype=torch.long,
)
meta = MetaStruct(
left_context_size=left_context_size,
finished=torch.tensor(finished, dtype=torch.bool),
)
if ref_context_size > 0 and ref_context_request_id is not None:
meta.ref_context_size = ref_context_size
meta.ref_context_request_id = ref_context_request_id
meta.ref_context_included = ref_context_included
return OmniPayloadStruct(
codes=CodesStruct(audio=code_predictor_codes),
meta=MetaStruct(
left_context_size=left_context_size,
finished=torch.tensor(finished, dtype=torch.bool),
),
meta=meta,
speaker=extract_speaker_from_request(request),
language=extract_language_from_request(request),
)
+12
View File
@@ -0,0 +1,12 @@
"""Shared helpers for custom voice profile tools."""
from __future__ import annotations
import re
def safe_voice_stem(name: str) -> str:
stem = re.sub(r"[^a-zA-Z0-9_.-]+", "_", name.strip())
if not stem or stem in (".", ".."):
raise ValueError(f"Invalid voice name: {name!r}")
return stem[:200]
+197
View File
@@ -6,8 +6,11 @@ has its own slot. Access via :func:`get_speaker_cache`.
from __future__ import annotations
import json
import os
import threading
from collections import OrderedDict
from pathlib import Path
from typing import Any
import torch
@@ -16,6 +19,8 @@ from vllm.logger import init_logger
logger = init_logger(__name__)
_MAX_BYTES = 512 * 1024**2 # 512 MiB
_CUSTOM_VOICE_MANIFEST = "custom_voice_manifest.json"
_CUSTOM_VOICE_SCHEMA_VERSION = 1
_SINGLETON: SpeakerEmbeddingCache | None = None
_SINGLETON_LOCK = threading.Lock()
@@ -31,6 +36,198 @@ def _estimate_tensor_bytes(obj: object) -> int:
return 0
def _as_voice_map(voices: object) -> dict[str, dict[str, Any]]:
if not isinstance(voices, dict):
return {}
result: dict[str, dict[str, Any]] = {}
for name, info in voices.items():
entry = dict(info) if isinstance(info, dict) else {}
entry.setdefault("name", str(name))
result[str(name)] = entry
return result
def _load_custom_voice_manifest(custom_voice_dir: str | os.PathLike[str] | None) -> dict[str, Any] | None:
if not custom_voice_dir:
return None
manifest_path = Path(custom_voice_dir).expanduser() / _CUSTOM_VOICE_MANIFEST
if not manifest_path.exists():
return None
try:
return json.loads(manifest_path.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError) as exc:
logger.warning("Failed to load custom voice manifest %s: %s", manifest_path, exc)
return None
def iter_custom_voice_profiles(
custom_voice_dir: str | os.PathLike[str] | None,
*,
expected_model_type: str | None = None,
) -> list[dict[str, Any]]:
manifest = _load_custom_voice_manifest(custom_voice_dir)
if manifest is None:
return []
manifest_model_type = manifest.get("model_type")
if expected_model_type and manifest_model_type and manifest_model_type != expected_model_type:
logger.warning(
"Skipping custom voice manifest for model_type=%s while loading %s",
manifest_model_type,
expected_model_type,
)
return []
root = Path(custom_voice_dir).expanduser().resolve() # type: ignore[arg-type]
profiles: list[dict[str, Any]] = []
for raw_name, raw_info in _as_voice_map(manifest.get("voices", {})).items():
name = str(raw_info.get("name") or raw_name).strip()
if not name:
continue
lower = name.lower()
info = dict(raw_info)
for stale_key in ("model_revision", "audio_codec_version", "audio_codec_config_hash"):
info.pop(stale_key, None)
filename = info.get("file") or info.get("filename") or f"{lower}.safetensors"
file_path = (root / str(filename)).resolve()
if not file_path.is_relative_to(root):
logger.warning("Skipping custom voice %s: file path escapes custom_voice_dir", name)
continue
info.update(
{
"name": name,
"voice_name_lower": lower,
"file": str(filename),
"file_path": str(file_path),
"custom_voice_dir": str(root),
"model_type": manifest_model_type or expected_model_type,
"schema_version": manifest.get("schema_version", _CUSTOM_VOICE_SCHEMA_VERSION),
}
)
profiles.append(info)
return profiles
def _load_profile_tensors(profile: dict[str, Any]) -> dict[str, torch.Tensor] | None:
try:
from safetensors.torch import load_file
except ImportError:
logger.warning("safetensors is required to load custom voice profiles")
return None
file_path = profile.get("file_path")
if not file_path:
return None
try:
return load_file(str(file_path))
except Exception as exc:
logger.warning("Failed to load custom voice profile %s: %s", file_path, exc)
return None
def _first_dim(tensor: torch.Tensor | None) -> int:
if tensor is None or tensor.ndim == 0:
return 0
return int(tensor.shape[0])
def _validate_qwen3_tts_profile(
profile: dict[str, Any],
tensors: dict[str, torch.Tensor],
*,
expected_embedding_dim: int | None,
) -> str | None:
# Successful validation also normalizes/augments `profile` for serving metadata.
speaker_embedding = tensors.get("speaker_embedding")
if not isinstance(speaker_embedding, torch.Tensor):
return "missing required tensor `speaker_embedding`"
embedding_dim = int(speaker_embedding.reshape(-1).numel())
if expected_embedding_dim and embedding_dim != expected_embedding_dim:
return f"speaker_embedding dim={embedding_dim}, expected={expected_embedding_dim}"
mode = str(profile.get("mode") or "xvec").lower()
if mode not in ("xvec", "icl"):
return f"invalid mode={mode!r}; expected 'xvec' or 'icl'"
profile["mode"] = mode
profile["embedding_dim"] = embedding_dim
if mode == "icl":
ref_code = tensors.get("ref_code")
if not isinstance(ref_code, torch.Tensor):
return "mode='icl' requires tensor `ref_code`"
ref_code_length = _first_dim(ref_code)
if ref_code_length <= 0:
return "mode='icl' requires non-empty tensor `ref_code`"
profile["ref_code_length"] = ref_code_length
return None
def _validate_voxcpm2_profile(profile: dict[str, Any], tensors: dict[str, torch.Tensor]) -> str | None:
# Successful validation also normalizes/augments `profile` for serving metadata.
ref_audio_feat = tensors.get("ref_audio_feat")
audio_feat = tensors.get("audio_feat")
has_ref_audio_feat = isinstance(ref_audio_feat, torch.Tensor)
has_audio_feat = isinstance(audio_feat, torch.Tensor)
if not has_ref_audio_feat and not has_audio_feat:
return "missing required tensor `ref_audio_feat` or `audio_feat`"
mode = str(profile.get("mode") or "").lower()
if not mode:
if has_ref_audio_feat and has_audio_feat:
mode = "ref_continuation"
elif has_ref_audio_feat:
mode = "reference"
else:
mode = "continuation"
if mode not in ("reference", "continuation", "ref_continuation"):
return f"invalid mode={mode!r}; expected 'reference', 'continuation', or 'ref_continuation'"
if mode in ("reference", "ref_continuation") and not has_ref_audio_feat:
return f"mode={mode!r} requires tensor `ref_audio_feat`"
if mode in ("continuation", "ref_continuation") and not has_audio_feat:
return f"mode={mode!r} requires tensor `audio_feat`"
ref_audio_feat_len = _first_dim(ref_audio_feat)
audio_feat_len = _first_dim(audio_feat)
if mode in ("reference", "ref_continuation") and ref_audio_feat_len <= 0:
return f"mode={mode!r} requires non-empty tensor `ref_audio_feat`"
if mode in ("continuation", "ref_continuation") and audio_feat_len <= 0:
return f"mode={mode!r} requires non-empty tensor `audio_feat`"
profile["mode"] = mode
if has_ref_audio_feat:
profile["ref_audio_feat_len"] = ref_audio_feat_len
if has_audio_feat:
profile["audio_feat_len"] = audio_feat_len
return None
def load_validated_profile_tensors(
profile: dict[str, Any],
*,
expected_model_type: str,
qwen3_embedding_dim: int | None = None,
) -> dict[str, torch.Tensor] | None:
tensors = _load_profile_tensors(profile)
error = None
if tensors is None:
error = "safetensors file could not be loaded"
elif str(profile.get("model_type") or expected_model_type or "") != expected_model_type:
error = f"model_type={profile.get('model_type')!r}, expected={expected_model_type!r}"
elif expected_model_type == "qwen3_tts":
error = _validate_qwen3_tts_profile(profile, tensors, expected_embedding_dim=qwen3_embedding_dim)
elif expected_model_type == "voxcpm2":
error = _validate_voxcpm2_profile(profile, tensors)
else:
error = f"unsupported custom voice model_type={expected_model_type!r}"
if error is not None:
logger.warning("Skipping custom %s voice %s: %s", expected_model_type, profile.get("name"), error)
return None
return tensors
class SpeakerEmbeddingCache:
"""Thread-safe in-memory LRU cache for speaker extraction artifacts."""