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

Signed-off-by: Sy03 <1370724210@qq.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Dan250124 <416947747@qq.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

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-<field> <value>`.
 
 ---
@@ -474,6 +500,23 @@ python voxcpm2/openai_speech_client.py \
 ```
 The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs.
 
+### Precomputed custom voices
+For repeated VoxCPM2 speakers, precompute the prompt cache and load it through `custom_voice_dir`:
+```bash
+python voxcpm2/precompute_custom_voice.py \
+    --model openbmb/VoxCPM2 \
+    --voice-name alice \
+    --ref-audio /path/to/reference.wav \
+    --mode ref_continuation \
+    --prompt-text "Original transcript of the reference audio" \
+    --output-dir /path/to/custom_voices
+```
+Add the output directory to the deploy config:
+```yaml
+custom_voice_dir: /path/to/custom_voices
+```
+After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use `voice="alice"` without sending `ref_audio`.
+
 ### Gradio demo (gapless streaming via AudioWorklet)
 ```bash
 python voxcpm2/gradio_demo.py

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="<name>"`` and no per-request ref_audio.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+REPO_ROOT = Path(__file__).resolve().parents[4]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from vllm_omni.utils.custom_voice_io import safe_voice_stem  # noqa: E402
+
+MANIFEST_NAME = "custom_voice_manifest.json"
+_SPEAKER_ENCODER_SAMPLE_RATE = 24000
+_SPEAKER_ENCODER_NUM_MELS = 128
+_SPEAKER_ENCODER_N_FFT = 1024
+_SPEAKER_ENCODER_HOP_SIZE = 256
+_SPEAKER_ENCODER_WIN_SIZE = 1024
+_SPEAKER_ENCODER_FMIN = 0
+_SPEAKER_ENCODER_FMAX = 12000
+
+
+def _resolve_model_dir(model: str) -> str:
+    if os.path.isdir(model):
+        return model
+    from transformers.utils.hub import cached_file
+
+    config_path = cached_file(model, "config.json")
+    if config_path is None:
+        raise FileNotFoundError(f"Could not resolve config.json for {model}")
+    return os.path.dirname(config_path)
+
+
+def _read_audio_mono(path: str) -> tuple[np.ndarray, int]:
+    import soundfile as sf
+
+    wav, sr = sf.read(path, dtype="float32")
+    if wav.ndim > 1:
+        wav = wav.mean(axis=-1)
+    return np.asarray(wav, dtype=np.float32), int(sr)
+
+
+def _resample(wav: np.ndarray, sr: int, target_sr: int) -> np.ndarray:
+    if sr == target_sr:
+        return wav.astype(np.float32)
+    from vllm.multimodal.audio import AudioResampler
+
+    return AudioResampler(target_sr=target_sr).resample(wav.astype(np.float32), orig_sr=int(sr))
+
+
+def _load_speaker_encoder(model_dir: str, device: torch.device):
+    from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSSpeakerEncoder
+
+    config = Qwen3TTSConfig.from_pretrained(model_dir)
+    encoder = Qwen3TTSSpeakerEncoder(config.speaker_encoder_config)
+    index_path = Path(model_dir) / "model.safetensors.index.json"
+    state: dict[str, torch.Tensor] = {}
+    if index_path.exists():
+        index = json.loads(index_path.read_text(encoding="utf-8"))
+        weight_map = index.get("weight_map", {})
+        shard_names = sorted({v for k, v in weight_map.items() if k.startswith("speaker_encoder.")})
+        shard_paths = [Path(model_dir) / name for name in shard_names]
+    else:
+        shard_paths = sorted(Path(model_dir).glob("model*.safetensors"))
+    for shard in shard_paths:
+        with safe_open(str(shard), framework="pt", device="cpu") as f:
+            for key in f.keys():
+                if key.startswith("speaker_encoder."):
+                    state[key.removeprefix("speaker_encoder.")] = f.get_tensor(key)
+    if not state:
+        raise RuntimeError(
+            f"No speaker_encoder.* tensors found under {model_dir}. "
+            "Use a Qwen3-TTS Base checkpoint for pre-computing custom voice profiles."
+        )
+    encoder.load_state_dict(state)
+    encoder.to(device=device, dtype=torch.bfloat16).eval()
+    return config, encoder
+
+
+def _speaker_embedding(
+    encoder: torch.nn.Module,
+    speaker_encoder_config: object,
+    wav: np.ndarray,
+    sr: int,
+    device: torch.device,
+) -> torch.Tensor:
+    from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import mel_spectrogram
+
+    target_sr = int(getattr(speaker_encoder_config, "sample_rate", _SPEAKER_ENCODER_SAMPLE_RATE))
+    mel_dim = int(getattr(speaker_encoder_config, "mel_dim", _SPEAKER_ENCODER_NUM_MELS))
+    if target_sr != _SPEAKER_ENCODER_SAMPLE_RATE or mel_dim != _SPEAKER_ENCODER_NUM_MELS:
+        raise ValueError(
+            "Unsupported Qwen3-TTS speaker encoder mel config for precompute: "
+            f"sample_rate={target_sr}, mel_dim={mel_dim}. "
+            f"Expected sample_rate={_SPEAKER_ENCODER_SAMPLE_RATE}, mel_dim={_SPEAKER_ENCODER_NUM_MELS} "
+            "to match the serving-time speaker embedding path."
+        )
+    wav = _resample(wav, sr, target_sr)
+    wav_t = torch.from_numpy(wav).to(device=device, dtype=torch.float32)
+    mels = mel_spectrogram(
+        wav_t.unsqueeze(0),
+        n_fft=_SPEAKER_ENCODER_N_FFT,
+        num_mels=_SPEAKER_ENCODER_NUM_MELS,
+        sampling_rate=_SPEAKER_ENCODER_SAMPLE_RATE,
+        hop_size=_SPEAKER_ENCODER_HOP_SIZE,
+        win_size=_SPEAKER_ENCODER_WIN_SIZE,
+        fmin=_SPEAKER_ENCODER_FMIN,
+        fmax=_SPEAKER_ENCODER_FMAX,
+    ).transpose(1, 2)
+    with torch.inference_mode():
+        return encoder(mels.to(device=device, dtype=torch.bfloat16))[0].float().cpu().contiguous()
+
+
+def _ref_code(model_dir: str, wav: np.ndarray, sr: int, device: torch.device) -> torch.Tensor:
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import Qwen3TTSTokenizer
+
+    tok = Qwen3TTSTokenizer.from_pretrained(
+        str(Path(model_dir) / "speech_tokenizer"),
+        torch_dtype=torch.bfloat16,
+    )
+    try:
+        del tok.model.decoder
+        tok.model.decoder = None
+        tok.model.encoder.to(device)
+        tok.device = device
+    except Exception:
+        tok.device = device
+    with torch.inference_mode():
+        enc = tok.encode(wav, sr=int(sr), return_dict=True)
+    codes = getattr(enc, "audio_codes", None)
+    if isinstance(codes, list):
+        codes = codes[0] if codes else None
+    if not isinstance(codes, torch.Tensor):
+        raise RuntimeError("Speech tokenizer did not return audio_codes")
+    if codes.ndim == 3:
+        codes = codes[0]
+    return codes.to(dtype=torch.int32, device="cpu").contiguous()
+
+
+def _load_manifest(output_dir: Path, model: str, hidden_size: int) -> dict[str, Any]:
+    path = output_dir / MANIFEST_NAME
+    if path.exists():
+        return json.loads(path.read_text(encoding="utf-8"))
+    return {
+        "schema_version": 1,
+        "model_type": "qwen3_tts",
+        "model": model,
+        "hidden_size": hidden_size,
+        "voices": {},
+    }
+
+
+def _write_voice(
+    *,
+    model: str,
+    output_dir: Path,
+    voice_name: str,
+    ref_audio: str,
+    ref_text: str | None,
+    mode: str,
+    speaker_description: str | None,
+    device: torch.device,
+) -> None:
+    model_dir = _resolve_model_dir(model)
+    config, encoder = _load_speaker_encoder(model_dir, device)
+    wav, sr = _read_audio_mono(ref_audio)
+    if wav.size < 1024:
+        raise ValueError(f"Reference audio too short: {wav.size} samples")
+
+    tensors: dict[str, torch.Tensor] = {
+        "speaker_embedding": _speaker_embedding(encoder, config.speaker_encoder_config, wav, sr, device),
+    }
+    if mode == "icl":
+        if not ref_text or not ref_text.strip():
+            raise ValueError("--ref-text is required for --mode icl")
+        tensors["ref_code"] = _ref_code(model_dir, wav, sr, device)
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    filename = f"{safe_voice_stem(voice_name)}.safetensors"
+    save_file(tensors, str(output_dir / filename))
+
+    hidden_size = int(getattr(config.talker_config, "hidden_size", tensors["speaker_embedding"].numel()))
+    manifest = _load_manifest(output_dir, model, hidden_size)
+    entry: dict[str, Any] = {
+        "name": voice_name,
+        "file": filename,
+        "mode": mode,
+        "embedding_dim": int(tensors["speaker_embedding"].numel()),
+    }
+    if "ref_code" in tensors:
+        entry["ref_code_length"] = int(tensors["ref_code"].shape[0])
+    if ref_text:
+        entry["ref_text"] = ref_text
+    if speaker_description:
+        entry["speaker_description"] = speaker_description
+    manifest.setdefault("voices", {})[voice_name] = entry
+    (output_dir / MANIFEST_NAME).write_text(json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8")
+    print(f"Wrote {output_dir / filename}")
+    print(f"Updated {output_dir / MANIFEST_NAME}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Pre-compute Qwen3-TTS custom voice profile")
+    parser.add_argument("--model", required=True, help="Qwen3-TTS Base model path or Hugging Face ID")
+    parser.add_argument("--voice-name", required=True)
+    parser.add_argument("--ref-audio", required=True)
+    parser.add_argument("--ref-text", default=None)
+    parser.add_argument("--mode", choices=["xvec", "icl"], default="xvec")
+    parser.add_argument("--speaker-description", default=None)
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    args = parser.parse_args()
+
+    _write_voice(
+        model=args.model,
+        output_dir=Path(args.output_dir),
+        voice_name=args.voice_name,
+        ref_audio=args.ref_audio,
+        ref_text=args.ref_text,
+        mode=args.mode,
+        speaker_description=args.speaker_description,
+        device=torch.device(args.device),
+    )
+
+
+if __name__ == "__main__":
+    main()

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="<name>"`` and no per-request ref_audio.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors.torch import save_file
+
+REPO_ROOT = Path(__file__).resolve().parents[4]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from vllm_omni.utils.custom_voice_io import safe_voice_stem  # noqa: E402
+
+MANIFEST_NAME = "custom_voice_manifest.json"
+
+
+def _load_tts(model: str, device: torch.device):
+    from vllm_omni.model_executor.models.voxcpm2.voxcpm2_import_utils import import_voxcpm2_core
+
+    VoxCPM = import_voxcpm2_core()
+    native = VoxCPM.from_pretrained(model, load_denoiser=False, optimize=False)
+    return native.tts_model.to(device).eval()
+
+
+def _load_manifest(output_dir: Path, model: str) -> dict[str, Any]:
+    path = output_dir / MANIFEST_NAME
+    if path.exists():
+        return json.loads(path.read_text(encoding="utf-8"))
+    return {
+        "schema_version": 1,
+        "model_type": "voxcpm2",
+        "model": model,
+        "voices": {},
+    }
+
+
+def _write_voice(
+    *,
+    model: str,
+    output_dir: Path,
+    voice_name: str,
+    ref_audio: str,
+    prompt_text: str | None,
+    mode: str,
+    speaker_description: str | None,
+    device: torch.device,
+) -> None:
+    if mode in ("continuation", "ref_continuation") and not prompt_text:
+        raise ValueError("--prompt-text is required for continuation/ref_continuation modes")
+
+    tts = _load_tts(model, device)
+    tensors: dict[str, torch.Tensor] = {}
+    with torch.inference_mode():
+        if mode in ("reference", "ref_continuation"):
+            tensors["ref_audio_feat"] = tts._encode_wav(ref_audio, padding_mode="right").float().cpu().contiguous()
+        if mode in ("continuation", "ref_continuation"):
+            tensors["audio_feat"] = tts._encode_wav(ref_audio, padding_mode="left").float().cpu().contiguous()
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    filename = f"{safe_voice_stem(voice_name)}.safetensors"
+    save_file(tensors, str(output_dir / filename))
+
+    manifest = _load_manifest(output_dir, model)
+    entry: dict[str, Any] = {
+        "name": voice_name,
+        "file": filename,
+        "mode": mode,
+    }
+    if "ref_audio_feat" in tensors:
+        entry["ref_audio_feat_len"] = int(tensors["ref_audio_feat"].shape[0])
+    if "audio_feat" in tensors:
+        entry["audio_feat_len"] = int(tensors["audio_feat"].shape[0])
+    if prompt_text:
+        entry["prompt_text"] = prompt_text
+    if speaker_description:
+        entry["speaker_description"] = speaker_description
+
+    manifest.setdefault("voices", {})[voice_name] = entry
+    (output_dir / MANIFEST_NAME).write_text(json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8")
+    print(f"Wrote {output_dir / filename}")
+    print(f"Updated {output_dir / MANIFEST_NAME}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Pre-compute VoxCPM2 custom voice profile")
+    parser.add_argument("--model", default="openbmb/VoxCPM2", help="VoxCPM2 model path or Hugging Face ID")
+    parser.add_argument("--voice-name", required=True)
+    parser.add_argument("--ref-audio", required=True)
+    parser.add_argument(
+        "--prompt-text",
+        default=None,
+        help="Transcript of ref audio for continuation/ref_continuation modes",
+    )
+    parser.add_argument(
+        "--mode",
+        choices=["reference", "continuation", "ref_continuation"],
+        default="reference",
+    )
+    parser.add_argument("--speaker-description", default=None)
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    args = parser.parse_args()
+
+    _write_voice(
+        model=args.model,
+        output_dir=Path(args.output_dir),
+        voice_name=args.voice_name,
+        ref_audio=args.ref_audio,
+        prompt_text=args.prompt_text,
+        mode=args.mode,
+        speaker_description=args.speaker_description,
+        device=torch.device(args.device),
+    )
+
+
+if __name__ == "__main__":
+    main()