diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index 23f89bad5..61de36938 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -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 | diff --git a/docs/user_guide/examples/online_serving/text_to_speech.md b/docs/user_guide/examples/online_serving/text_to_speech.md index f971cced1..bc1499b2d 100644 --- a/docs/user_guide/examples/online_serving/text_to_speech.md +++ b/docs/user_guide/examples/online_serving/text_to_speech.md @@ -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- `. --- @@ -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 diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index 12dafde41..a3b27b722 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -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- `. --- @@ -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 diff --git a/examples/online_serving/text_to_speech/qwen3_tts/precompute_custom_voice.py b/examples/online_serving/text_to_speech/qwen3_tts/precompute_custom_voice.py new file mode 100644 index 000000000..0fd765c6a --- /dev/null +++ b/examples/online_serving/text_to_speech/qwen3_tts/precompute_custom_voice.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=""`` 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() diff --git a/examples/online_serving/text_to_speech/voxcpm2/precompute_custom_voice.py b/examples/online_serving/text_to_speech/voxcpm2/precompute_custom_voice.py new file mode 100644 index 000000000..c9ffe617e --- /dev/null +++ b/examples/online_serving/text_to_speech/voxcpm2/precompute_custom_voice.py @@ -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=""`` 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() diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 872410aaa..4fcf3418a 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -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(" 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, diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py index 7da61fb33..5fd961b4f 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py @@ -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"], + }, + ) diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_prompt_len.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_prompt_len.py index 9493a00f2..10bc015e0 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_prompt_len.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_prompt_len.py @@ -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 diff --git a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py index 8f74ed094..9c4270c1e 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py @@ -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(): diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index e96f79f89..01facfba1 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -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) diff --git a/tests/test_tts_custom_voice_profiles.py b/tests/test_tts_custom_voice_profiles.py new file mode 100644 index 000000000..2c7792a8b --- /dev/null +++ b/tests/test_tts_custom_voice_profiles.py @@ -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"] diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index f6a88546f..cd0fe5dd5 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -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", ) diff --git a/vllm_omni/data_entry_keys.py b/vllm_omni/data_entry_keys.py index 09316b55e..aa0e1005c 100644 --- a/vllm_omni/data_entry_keys.py +++ b/vllm_omni/data_entry_keys.py @@ -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 diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index e650373a9..50e750c68 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -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 diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index e2a4b5af8..0ca4cd319 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -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, diff --git a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py index c9a0f266d..b227be19d 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py +++ b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py @@ -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`." diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py index b7c3c9d6e..7991d85a1 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py @@ -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}, diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index ca1c4b395..d97f21233 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py @@ -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 -------------------- diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py index 69b2a29ea..8b60457b0 100644 --- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py +++ b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py @@ -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( diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py index 1ffbdb931..158963c6e 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py @@ -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), ) diff --git a/vllm_omni/utils/custom_voice_io.py b/vllm_omni/utils/custom_voice_io.py new file mode 100644 index 000000000..9356c7c9a --- /dev/null +++ b/vllm_omni/utils/custom_voice_io.py @@ -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] diff --git a/vllm_omni/utils/speaker_cache.py b/vllm_omni/utils/speaker_cache.py index 80623d894..d29141489 100644 --- a/vllm_omni/utils/speaker_cache.py +++ b/vllm_omni/utils/speaker_cache.py @@ -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."""