[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:
23 files changed
+2362
-145
No files matched your search
@@ -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():
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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"]
|
||||
@@ -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",
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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),
|
||||
)
|
||||
|
||||
@@ -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]
|
||||
@@ -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."""
|
||||
|
||||
|
||||
Reference in new issue
Block a user