[New Model] Support IndexTTS 2.5 (#5957)
Signed-off-by: Sysy <1370724210@qq.com> Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
This commit is contained in:
1 parent
1ff2753469
commit
58007c0084
77 files changed
+7123
-519
No files matched your search
@@ -10,7 +10,7 @@ steps:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
# Buildkite will automatically replace this with the actual commit hash
|
||||
VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
|
||||
VLLM_VERSION: "v0.26.0"
|
||||
VLLM_VERSION: "v0.27.0"
|
||||
priority: 100
|
||||
timeout_in_minutes: 60
|
||||
soft_fail: false
|
||||
@@ -34,8 +34,23 @@ docker_build() {
|
||||
docker build "${out[@]}" "$@" -f docker/Dockerfile.xpu .
|
||||
}
|
||||
|
||||
build_vllm_base() (
|
||||
local vllm_source_dir
|
||||
vllm_source_dir=$(mktemp -d)
|
||||
trap 'rm -rf "${vllm_source_dir}"' EXIT
|
||||
|
||||
git clone --depth 1 --branch "${VLLM_VERSION}" \
|
||||
https://github.com/vllm-project/vllm "${vllm_source_dir}"
|
||||
|
||||
local out=("${export_args[@]/'{{IMAGE}}'/${base_image_name}}")
|
||||
docker build "${out[@]}" \
|
||||
--target vllm-openai \
|
||||
-f "${vllm_source_dir}/docker/Dockerfile.xpu" \
|
||||
"${vllm_source_dir}"
|
||||
)
|
||||
|
||||
if [ -z "$(docker images -q "${base_image_name}")" ]; then
|
||||
docker_build "${base_image_name}" --target vllm-base --build-arg "VLLM_VERSION=${VLLM_VERSION}"
|
||||
build_vllm_base
|
||||
fi
|
||||
|
||||
# Try building the docker image
|
||||
|
||||
@@ -139,6 +139,34 @@ python benchmarks/tts/bench_tts.py \
|
||||
--output-dir ./results
|
||||
```
|
||||
|
||||
#### Fixed IndexTTS 2.5 performance sweep
|
||||
|
||||
IndexTTS 2.5 is registered in `model_configs.yaml` like the other TTS models.
|
||||
Use `--served-model-name` when the server was launched from a local native
|
||||
bundle instead of the registry key. The command below reproduces the acceptance
|
||||
workload: Seed-TTS Eval EN, `n=500`, concurrency `4/8/16/32`, five warmups, and
|
||||
dataset seed 0.
|
||||
|
||||
```bash
|
||||
python benchmarks/tts/bench_tts.py \
|
||||
--model IndexTeam/IndexTTS-2.5 \
|
||||
--served-model-name /path/to/indextts-2.5 \
|
||||
--task voice_clone --locale en \
|
||||
--dataset-path /path/to/seedtts_testset \
|
||||
--host 127.0.0.1 --port 8092 \
|
||||
--concurrency 4 8 16 32 \
|
||||
--num-prompts 500 --num-warmups 5 \
|
||||
--output-dir ./results/indextts2_5 \
|
||||
-- \
|
||||
--tokenizer /path/to/indextts-2.5/qwen0.6bemo4-merge \
|
||||
--seed 0 --metric-percentiles 50,95,99 --disable-tqdm --save-detailed
|
||||
```
|
||||
|
||||
The wrapper benchmarks an already-running server. For reproducible quality
|
||||
comparisons, add `--request-seed 42`; production-performance runs should omit
|
||||
it. The trailing `--seed 0` controls Seed-TTS row selection rather than model
|
||||
sampling.
|
||||
|
||||
### 4. Plot a sweep
|
||||
|
||||
```bash
|
||||
|
||||
@@ -84,12 +84,17 @@ def build_bench_args(
|
||||
result_filename: str | None,
|
||||
extra_cli_args: list[str],
|
||||
output_len: int | None = None,
|
||||
served_model_name: str | None = None,
|
||||
num_warmups: int = 2,
|
||||
request_seed: int | None = None,
|
||||
) -> list[str]:
|
||||
"""Build the ``vllm bench serve --omni`` command for one (task, concurrency) run."""
|
||||
dataset_name = _TASK_TO_DATASET[task]
|
||||
backend: str = model_cfg["backend"]
|
||||
endpoint: str = model_cfg["endpoint"]
|
||||
task_extra_body: dict[str, Any] = (model_cfg.get("task_extra_body") or {}).get(task) or {}
|
||||
task_extra_body: dict[str, Any] = dict((model_cfg.get("task_extra_body") or {}).get(task) or {})
|
||||
if request_seed is not None:
|
||||
task_extra_body["seed"] = request_seed
|
||||
|
||||
# Resolve dataset path
|
||||
if dataset_path:
|
||||
@@ -109,7 +114,7 @@ def build_bench_args(
|
||||
"--port",
|
||||
str(port),
|
||||
"--model",
|
||||
model,
|
||||
served_model_name or model,
|
||||
"--backend",
|
||||
backend,
|
||||
"--endpoint",
|
||||
@@ -119,7 +124,7 @@ def build_bench_args(
|
||||
"--num-prompts",
|
||||
str(num_prompts),
|
||||
"--num-warmups",
|
||||
"2",
|
||||
str(num_warmups),
|
||||
"--percentile-metrics",
|
||||
"ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun",
|
||||
]
|
||||
@@ -152,7 +157,8 @@ def build_bench_args(
|
||||
if result_filename:
|
||||
cmd += ["--result-filename", result_filename]
|
||||
|
||||
cmd += extra_cli_args
|
||||
passthrough = extra_cli_args[1:] if extra_cli_args[:1] == ["--"] else extra_cli_args
|
||||
cmd += passthrough
|
||||
return cmd
|
||||
|
||||
|
||||
@@ -225,6 +231,12 @@ def main() -> None:
|
||||
parser.add_argument(
|
||||
"--model", required=True, help="HuggingFace model ID (e.g. Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--served-model-name",
|
||||
default=None,
|
||||
help="Model name/path accepted by the running server. Use this when --model selects a registry entry "
|
||||
"but the server was launched from a local bundle path.",
|
||||
)
|
||||
parser.add_argument("--task", default="all", help="Task type: voice_clone | default_voice | voice_design | all")
|
||||
parser.add_argument("--locale", default="en", choices=["en", "zh"])
|
||||
parser.add_argument("--concurrency", type=int, nargs="+", default=[1, 4], metavar="N")
|
||||
@@ -236,6 +248,13 @@ def main() -> None:
|
||||
metavar="N",
|
||||
help="Number of prompts per run. If one value, applied to all concurrency levels.",
|
||||
)
|
||||
parser.add_argument("--num-warmups", type=int, default=2, help="Warmup requests before each measured run")
|
||||
parser.add_argument(
|
||||
"--request-seed",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Optional model sampling seed added to every request body for reproducible quality runs",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dataset-path", default=None, help="Root of seed-tts-eval dataset (required for voice_clone/default_voice)"
|
||||
)
|
||||
@@ -256,6 +275,9 @@ def main() -> None:
|
||||
parser.add_argument("extra", nargs=argparse.REMAINDER, help="Extra args passed directly to vllm bench serve")
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.num_warmups < 0:
|
||||
parser.error("--num-warmups cannot be negative")
|
||||
|
||||
model_configs = load_model_configs(Path(args.model_configs))
|
||||
if args.model not in model_configs:
|
||||
known = "\n ".join(model_configs.keys())
|
||||
@@ -308,6 +330,9 @@ def main() -> None:
|
||||
result_filename=result_filename,
|
||||
extra_cli_args=args.extra or [],
|
||||
output_len=args.output_len,
|
||||
served_model_name=args.served_model_name,
|
||||
num_warmups=args.num_warmups,
|
||||
request_seed=args.request_seed,
|
||||
)
|
||||
result = run_one_benchmark(cmd)
|
||||
if result is not None:
|
||||
|
||||
@@ -7,6 +7,17 @@
|
||||
# is tracked here.
|
||||
|
||||
models:
|
||||
IndexTeam/IndexTTS-2.5:
|
||||
supported_tasks: [voice_clone]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
trust_remote_code: true
|
||||
task_extra_body:
|
||||
voice_clone:
|
||||
extra_params:
|
||||
lang: en
|
||||
text_normalization: true
|
||||
|
||||
# -CustomVoice checkpoints lack speaker_encoder weights, so voice_clone is
|
||||
# NOT supported (an attempt raises ValueError from _extract_speaker_embedding
|
||||
# at model runtime). Use -Base for voice_clone.
|
||||
|
||||
@@ -88,7 +88,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
|
||||
ENV UV_INDEX_STRATEGY="unsafe-best-match"
|
||||
ENV UV_LINK_MODE="copy"
|
||||
|
||||
ARG VLLM_VERSION=v0.26.0
|
||||
ARG VLLM_VERSION=v0.27.0
|
||||
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
|
||||
WORKDIR /workspace/vllm
|
||||
|
||||
@@ -108,7 +108,7 @@ ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip uninstall triton triton-xpu && \
|
||||
uv pip install triton-xpu==3.7.1 && \
|
||||
uv pip install triton-xpu==3.7.2 && \
|
||||
uv pip install --no-build-isolation --no-deps .
|
||||
|
||||
CMD ["/bin/bash"]
|
||||
@@ -144,12 +144,13 @@ WORKDIR /workspace/vllm-omni
|
||||
COPY . .
|
||||
|
||||
ENV VLLM_OMNI_TARGET_DEVICE=xpu
|
||||
# Keep the upstream oneCCL packages paired with the XPU torch runtime.
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --no-build-isolation \
|
||||
--extra-index-url https://mirrors.ustc.edu.cn/pypi/simple/ \
|
||||
".[dev]" && \
|
||||
uv pip uninstall triton oneccl oneccl-devel && \
|
||||
uv pip install triton-xpu==3.7.1 --reinstall
|
||||
uv pip uninstall triton && \
|
||||
uv pip install triton-xpu==3.7.2 --reinstall
|
||||
|
||||
FROM vllm-omni AS vllm-omni-openai
|
||||
|
||||
|
||||
@@ -80,6 +80,7 @@ th {
|
||||
| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | | — |
|
||||
| `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | | [Published](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b) |
|
||||
| `IndexTTS2TalkerForConditionalGeneration` | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✅︎ | | | | — |
|
||||
| `IndexTTS25TalkerForConditionalGeneration` | IndexTTS-2.5 | native `checkpoints/` bundle | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/IndexTeam/IndexTTS-2_5.md) |
|
||||
| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | | — |
|
||||
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — |
|
||||
|
||||
@@ -24,6 +24,7 @@ list of supported architectures across all modalities, see
|
||||
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants | 24 kHz |
|
||||
| VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | ✓ (online) | continuation | 48 kHz |
|
||||
| IndexTTS-2 | `IndexTeam/IndexTTS-2` | 2 (AR talker + S2Mel DiT + BigVGAN) | ✓ (required) | — | emotion control (`--emo-audio`, `--emo-text`, `--emo-vector`) | 22.05 kHz |
|
||||
| IndexTTS-2.5 | native `checkpoints/` bundle | 2 (AR talker + EnhancedCodec + S2Mel DiT + BigVGAN) | ✓ (required) | — | multilingual (`--lang`) + emotion control | 22.05 kHz |
|
||||
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz |
|
||||
|
||||
## Common Quick Start
|
||||
@@ -437,11 +438,14 @@ Streaming is exposed through the online OpenAI Speech API (`stream=true`). See [
|
||||
|
||||
---
|
||||
|
||||
## IndexTTS-2
|
||||
## IndexTTS-2 and IndexTTS-2.5
|
||||
|
||||
2-stage TTS pipeline (GPT AR talker + S2Mel CFM DiT + BigVGAN vocoder) at 22.05 kHz. Every request requires reference audio for zero-shot voice cloning. Supports emotion conditioning via audio, text, or 8-dim vector.
|
||||
Both versions use a 2-stage TTS pipeline and produce 22.05 kHz mono audio. Stage 0 is a GPT AR talker. IndexTTS-2 Stage 1 uses RepCodec semantic embeddings, GPT latent, S2Mel CFM/DiT, and BigVGAN. IndexTTS-2.5 replaces RepCodec with EnhancedCodec and uses the official code-only path (`use_gpt_latent=false`). Every request requires reference audio for zero-shot voice cloning.
|
||||
|
||||
### Quick start
|
||||
|
||||
IndexTTS-2:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/indextts2/end2end.py \
|
||||
--model IndexTeam/IndexTTS-2 \
|
||||
@@ -449,6 +453,24 @@ python examples/offline_inference/text_to_speech/indextts2/end2end.py \
|
||||
--ref-audio /path/to/reference.wav
|
||||
```
|
||||
|
||||
IndexTTS-2 uses its existing frontend and does not require the IndexTTS-2.5 text dependencies. Before running IndexTTS-2.5, install its optional frontend dependencies:
|
||||
|
||||
```bash
|
||||
pip install 'vllm-omni[indextts2]'
|
||||
```
|
||||
|
||||
IndexTTS-2.5 with a native repository/lab bundle whose model root contains `checkpoints/`:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/indextts2/end2end.py \
|
||||
--model /path/to/indextts-2.5 \
|
||||
--model-version 2.5 \
|
||||
--lang ja \
|
||||
--text "こんにちは、音声合成のテストです。" \
|
||||
--ref-audio /path/to/reference.wav
|
||||
|
||||
```
|
||||
|
||||
### Emotion control
|
||||
```bash
|
||||
# Emotion from reference audio
|
||||
@@ -474,10 +496,13 @@ python examples/offline_inference/text_to_speech/indextts2/end2end.py \
|
||||
```
|
||||
|
||||
### Notes
|
||||
- `--ref-audio` is **required** — IndexTTS-2 does not support text-only synthesis.
|
||||
- Stage 0 (AR Talker): GPT-2 generates mel codes from text + reference audio.
|
||||
- Stage 1 (S2Mel + BigVGAN): CFM DiT converts mel codes to waveform at 22.05 kHz.
|
||||
- Deploy config: `vllm_omni/deploy/indextts2.yaml`. Stage 1 runs with `enforce_eager: true` (DiT has dynamic shapes).
|
||||
- `--ref-audio` is **required** — neither version provides text-only synthesis.
|
||||
- IndexTTS-2.5 Stage 0 (AR Talker) generates mel codes from text + reference audio using plain vLLM sampling. It does not reproduce the official default `num_beams=3` beam search. For parity comparisons, run upstream with `num_beams=1`; output quality can differ from the official beam-search result.
|
||||
- IndexTTS-2.5 accepts language codes such as `zh`, `en`, `zhen` (mixed Chinese/English), `ja`, and `yue`. `Mandarin` is a vLLM-Omni convenience alias for `zh`. Japanese (`ja`) uses `fugashi` tokenization and produces audio, but it does not automatically expand numbers, dates, or percentages; write those inputs as readable Japanese text before calling the model. Use `--no-text-normalization` only when the input is already normalized.
|
||||
- A seed controls Stage 0 AR sampling and per-request CFM noise. Different concurrent batch compositions do not guarantee a bit-identical waveform.
|
||||
- Stage 1 (semantic codec + S2Mel + BigVGAN): CFM/DiT converts the complete semantic-code sequence to waveform.
|
||||
- Deploy configs: `vllm_omni/deploy/indextts2.yaml` and `indextts2_5.yaml`.
|
||||
- `async_chunk=false` is intentional for IndexTTS-2.5 correctness: S2Mel consumes the completed semantic sequence.
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -1,13 +1,16 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline inference example for IndexTTS2 via vLLM-Omni.
|
||||
"""Offline inference example for IndexTTS 2.0 and 2.5 via vLLM-Omni.
|
||||
|
||||
Two-stage pipeline: GPT AR (Stage 0) → S2Mel + BigVGAN (Stage 1).
|
||||
Two-stage pipeline: GPT AR (Stage 0) → semantic codec + S2Mel + BigVGAN
|
||||
(Stage 1).
|
||||
Output is 22050 Hz mono WAV.
|
||||
|
||||
Usage:
|
||||
python end2end.py \
|
||||
--model /path/to/IndexTeam/IndexTTS-2 \
|
||||
--model /path/to/indextts-2.5 \
|
||||
--model-version 2.5 \
|
||||
--lang zh \
|
||||
--text "你好,这是一个语音合成测试。" \
|
||||
--ref-audio /path/to/ref.wav
|
||||
|
||||
@@ -47,6 +50,10 @@ def build_request(
|
||||
emo_alpha: float | None = None,
|
||||
use_emo_text: bool = False,
|
||||
use_random: bool = False,
|
||||
model_type: str = "indextts2",
|
||||
lang: str = "zh",
|
||||
text_normalization: bool = True,
|
||||
tokenizer_file: str | None = None,
|
||||
) -> dict:
|
||||
additional: dict = {"text": [text]}
|
||||
if ref_audio_path:
|
||||
@@ -63,12 +70,44 @@ def build_request(
|
||||
additional["use_emo_text"] = [True]
|
||||
if use_random:
|
||||
additional["use_random"] = [True]
|
||||
if model_type == "indextts2_5":
|
||||
additional["lang"] = [lang]
|
||||
additional["text_normalization"] = [text_normalization]
|
||||
prompt_kwargs = {
|
||||
"model_type": model_type,
|
||||
"lang": lang,
|
||||
"text_normalization": text_normalization,
|
||||
}
|
||||
if tokenizer_file is not None:
|
||||
prompt_kwargs["tokenizer_file"] = tokenizer_file
|
||||
return {
|
||||
"prompt_token_ids": build_indextts2_prefill_prompt_ids(model, text),
|
||||
"prompt_token_ids": build_indextts2_prefill_prompt_ids(
|
||||
model,
|
||||
text,
|
||||
**prompt_kwargs,
|
||||
),
|
||||
"additional_information": additional,
|
||||
}
|
||||
|
||||
|
||||
def _get_stage0_tokenizer_file(omni: Omni) -> str | None:
|
||||
"""Read the effective tokenizer override from the resolved Stage 0 config."""
|
||||
stage_configs = omni.stage_configs
|
||||
if not stage_configs:
|
||||
return None
|
||||
engine_args = getattr(stage_configs[0], "engine_args", None)
|
||||
if engine_args is None:
|
||||
return None
|
||||
hf_overrides = getattr(engine_args, "hf_overrides", None)
|
||||
if hf_overrides is None and isinstance(engine_args, dict):
|
||||
hf_overrides = engine_args.get("hf_overrides")
|
||||
if hf_overrides is None:
|
||||
return None
|
||||
if hasattr(hf_overrides, "get"):
|
||||
return hf_overrides.get("tokenizer_file")
|
||||
return getattr(hf_overrides, "tokenizer_file", None)
|
||||
|
||||
|
||||
def save_audio(waveform: torch.Tensor, path: str, sample_rate: int = 22050) -> None:
|
||||
audio_np = waveform.float().numpy()
|
||||
sf.write(path, audio_np, sample_rate)
|
||||
@@ -95,9 +134,14 @@ def extract_audio(mm: dict) -> tuple[torch.Tensor | None, int]:
|
||||
|
||||
|
||||
def main(args) -> None:
|
||||
model_type = "indextts2_5" if args.model_version == "2.5" else "indextts2"
|
||||
deploy_config = args.deploy_config
|
||||
if deploy_config is None and args.model_version == "2.5":
|
||||
repo_root = Path(__file__).resolve().parents[4]
|
||||
deploy_config = str(repo_root / "vllm_omni" / "deploy" / "indextts2_5.yaml")
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
deploy_config=args.deploy_config,
|
||||
deploy_config=deploy_config,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
|
||||
@@ -126,7 +170,7 @@ def main(args) -> None:
|
||||
print(f"Synthesizing: {args.text!r}")
|
||||
if args.ref_audio:
|
||||
print(f" ref_audio: {args.ref_audio}")
|
||||
inputs = build_request(
|
||||
request_kwargs = dict(
|
||||
model=args.model,
|
||||
text=args.text,
|
||||
ref_audio_path=args.ref_audio,
|
||||
@@ -136,7 +180,17 @@ def main(args) -> None:
|
||||
emo_alpha=args.emo_alpha,
|
||||
use_emo_text=args.use_emo_text,
|
||||
use_random=args.use_random,
|
||||
model_type=model_type,
|
||||
lang=args.lang,
|
||||
text_normalization=args.text_normalization,
|
||||
)
|
||||
if model_type == "indextts2_5":
|
||||
tokenizer_file = _get_stage0_tokenizer_file(omni)
|
||||
if tokenizer_file is not None:
|
||||
request_kwargs["tokenizer_file"] = tokenizer_file
|
||||
inputs = build_request(**request_kwargs)
|
||||
if model_type == "indextts2_5":
|
||||
inputs["additional_information"]["duration_factor"] = [1.0 / getattr(args, "speed", 1.0)]
|
||||
|
||||
for i, omni_out in enumerate(omni.generate(inputs, sampling_params_list=sampling_params)):
|
||||
mm = omni_out.multimodal_output
|
||||
@@ -154,8 +208,37 @@ def main(args) -> None:
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="IndexTTS2 offline inference")
|
||||
parser.add_argument("--model", required=True, help="HF model path for IndexTTS2.")
|
||||
parser = FlexibleArgumentParser(description="IndexTTS 2.0/2.5 offline inference")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
required=True,
|
||||
help="HF model path or native bundle containing checkpoints/.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model-version",
|
||||
choices=("2.0", "2.5"),
|
||||
default="2.0",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lang",
|
||||
default="zh",
|
||||
help=(
|
||||
"IndexTTS 2.5 language code, for example zh/en/zhen/ja/yue; "
|
||||
"zhen is mixed Chinese/English and Mandarin is a vLLM-Omni alias for zh."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-text-normalization",
|
||||
action="store_false",
|
||||
dest="text_normalization",
|
||||
help="Disable IndexTTS 2.5 text normalization.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--speed",
|
||||
type=float,
|
||||
default=1.0,
|
||||
help="IndexTTS 2.5 native synthesis speed in [0.5, 2.0]; 2.0 is faster.",
|
||||
)
|
||||
parser.add_argument("--text", default="你好,这是IndexTTS2语音合成测试。")
|
||||
parser.add_argument("--ref-audio", required=True, help="Reference audio for voice cloning.")
|
||||
parser.add_argument("--emo-audio", default=None, help="Emotion reference audio.")
|
||||
@@ -170,7 +253,15 @@ def parse_args():
|
||||
parser.add_argument("--emo-alpha", type=float, default=None, help="Emotion weight in [0, 1].")
|
||||
parser.add_argument("--use-emo-text", action="store_true", help="Infer emotion vector from emo-text or text.")
|
||||
parser.add_argument("--use-random", action="store_true", help="Use random emotion prototypes.")
|
||||
parser.add_argument("--seed", type=int, default=None)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=None,
|
||||
help=(
|
||||
"Seed Stage 0 AR sampling and per-request CFM noise; changing the "
|
||||
"concurrent batch composition may still change the final waveform."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
default=os.path.join(
|
||||
@@ -180,7 +271,10 @@ def parse_args():
|
||||
)
|
||||
parser.add_argument("--deploy-config", default=None)
|
||||
parser.add_argument("--stage-init-timeout", type=int, default=600)
|
||||
return parser.parse_args()
|
||||
args = parser.parse_args()
|
||||
if args.model_version == "2.5" and not 0.5 <= args.speed <= 2.0:
|
||||
parser.error("IndexTTS 2.5 --speed must be between 0.5 and 2.0")
|
||||
return args
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -16,7 +16,8 @@ For the full list of supported architectures across all modalities, see
|
||||
|---|---|---|---|---|---|
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
|
||||
| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
|
||||
| IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✓ (`ref_audio` or uploaded `voice`) | compat only, non-chunk | uploaded audio voice only; no presets | — |
|
||||
| IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✓ (`ref_audio` or uploaded `voice`) | `stream=true` response, non-chunk | uploaded audio voice only; no presets | — |
|
||||
| IndexTTS-2.5 | native `checkpoints/` bundle | ✓ (`ref_audio` or uploaded `voice`) | `stream=true` response, non-chunk | uploaded audio voice only; no presets | — |
|
||||
| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | ✓ (`ref_audio` / `speaker_embedding`) | ✓ (PCM stream) | IP labels + structured `instructions` | — |
|
||||
| Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | — (caption-controlled) | — | caption fields (`instructions`) | — |
|
||||
| MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✓ (`ref_audio` required) | ✓ (PCM stream) | — | ✓ |
|
||||
@@ -138,15 +139,29 @@ bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
|
||||
|
||||
---
|
||||
|
||||
## IndexTTS-2
|
||||
## IndexTTS-2 and IndexTTS-2.5
|
||||
|
||||
2-stage TTS (GPT AR + S2Mel CFM DiT + BigVGAN) at 22.05 kHz. Requests use `ref_audio` for voice cloning, or an uploaded audio `voice` from `/v1/audio/voices`. Supports emotion conditioning via `emo_audio`, `emo_text`, or `emo_vector` passed in `extra_params`.
|
||||
2-stage TTS at 22.05 kHz. Requests use `ref_audio` for voice cloning, or an uploaded audio `voice` from `/v1/audio/voices`. IndexTTS-2.5 uses a multilingual tokenizer, CAMPPlus speaker projection, EnhancedCodec, and code-only Stage 0→1 transfer by default. Both versions support emotion conditioning via `emo_audio`, `emo_text`, or `emo_vector` passed in `extra_params`.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
IndexTTS-2 uses its existing frontend and does not require the IndexTTS-2.5 text dependencies. Before launching IndexTTS-2.5, install its optional frontend dependencies:
|
||||
|
||||
```bash
|
||||
pip install 'vllm-omni[indextts2]'
|
||||
```
|
||||
|
||||
### Launch
|
||||
```bash
|
||||
vllm serve IndexTeam/IndexTTS-2 --omni --trust-remote-code --port 8092
|
||||
# or, to pass the bundled deploy config explicitly:
|
||||
bash examples/online_serving/text_to_speech/indextts2/run_server.sh
|
||||
|
||||
# IndexTTS-2.5 default: use_gpt_latent=false
|
||||
MODEL_VERSION=2.5 \
|
||||
MODEL=/path/to/indextts-2.5 \
|
||||
bash examples/online_serving/text_to_speech/indextts2/run_server.sh
|
||||
|
||||
```
|
||||
|
||||
### Sending requests
|
||||
@@ -161,14 +176,35 @@ python examples/online_serving/text_to_speech/indextts2/speech_client.py \
|
||||
--text "今天心情很好!" \
|
||||
--ref-audio /path/to/ref.wav \
|
||||
--emo-audio /path/to/happy.wav
|
||||
|
||||
# IndexTTS-2.5 Japanese request
|
||||
python examples/online_serving/text_to_speech/indextts2/speech_client.py \
|
||||
--model-version 2.5 \
|
||||
--model /path/to/indextts-2.5 \
|
||||
--lang ja \
|
||||
--text "こんにちは、音声合成のテストです。" \
|
||||
--ref-audio /path/to/ref.wav
|
||||
|
||||
# IndexTTS-2.5 native speed control
|
||||
python examples/online_serving/text_to_speech/indextts2/speech_client.py \
|
||||
--model-version 2.5 \
|
||||
--model /path/to/indextts-2.5 \
|
||||
--speed 2.0 \
|
||||
--text "这是两倍速度的语音合成测试。" \
|
||||
--ref-audio /path/to/ref.wav
|
||||
```
|
||||
|
||||
### Notes
|
||||
- Output: 22.05 kHz mono WAV.
|
||||
- Provide `ref_audio` on the documented raw request path, or pass `voice` only when it names an uploaded audio voice; IndexTTS-2 does not provide a built-in text-only preset voice.
|
||||
- Provide `ref_audio` on the documented raw request path, or pass `voice` only when it names an uploaded audio voice; neither version provides a built-in text-only preset voice.
|
||||
- IndexTTS-2.5 request controls `lang` and `text_normalization` are carried in `extra_params`; the client does this for `--model-version 2.5`.
|
||||
- IndexTTS-2.5 accepts the public `speed` request field in `[0.5, 2.0]`: `2.0` generates shorter, faster speech and `0.5` generates longer, slower speech. This is native Stage 1 duration control, so serving does not apply the generic playback-speed adjustment a second time. IndexTTS-2 does not use this control.
|
||||
- IndexTTS-2.5 accepts language codes such as `zh`, `en`, `zhen` (mixed Chinese/English), `ja`, and `yue`. `Mandarin` is a vLLM-Omni convenience alias for `zh`. Japanese (`ja`) uses `fugashi` tokenization and produces audio, but it does not automatically expand numbers, dates, or percentages; callers should first write those inputs as readable Japanese text.
|
||||
- A request `seed` controls Stage 0 AR sampling and per-request CFM noise. Different concurrent batch compositions do not guarantee a bit-identical waveform.
|
||||
- IndexTTS-2.5 Stage 0 uses plain vLLM sampling and does not reproduce the official default `num_beams=3` beam search. For parity comparisons, run upstream with `num_beams=1`; output quality can differ from the official beam-search result.
|
||||
- Emotion params (`emo_audio`, `emo_text`, `emo_vector`, `emo_alpha`, `use_emo_text`, `use_random`) are passed via the `extra_params` field. Official precedence is `use_emo_text` > `emo_vector` > `emo_audio` > same emotion as the speaker reference.
|
||||
- `stream=true` is accepted as an OpenAI-compatible response path, but IndexTTS-2 is not async-chunk streaming; audio is produced after S2Mel receives the full mel-code sequence.
|
||||
- Deploy config: `vllm_omni/deploy/indextts2.yaml` (auto-loaded).
|
||||
- IndexTTS-2.5 uses `vllm_omni/deploy/indextts2_5.yaml` with the official code-only Stage 0 to Stage 1 contract.
|
||||
- For IndexTTS-2.5, set `MODEL` to the local native bundle; asset discovery accepts its nested `checkpoints/` layout.
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -1,20 +1,37 @@
|
||||
#!/bin/bash
|
||||
# Launch vLLM-Omni server for IndexTTS2
|
||||
# Launch vLLM-Omni server for IndexTTS 2.0 or 2.5.
|
||||
#
|
||||
# Usage from repository root:
|
||||
# examples/online_serving/text_to_speech/indextts2/run_server.sh
|
||||
# CUDA_VISIBLE_DEVICES=0 PORT=8092 MODEL=/path/to/IndexTeam/IndexTTS-2 examples/online_serving/text_to_speech/indextts2/run_server.sh
|
||||
# MODEL_VERSION=2.5 MODEL=/path/to/native/bundle \
|
||||
# examples/online_serving/text_to_speech/indextts2/run_server.sh
|
||||
|
||||
set -e
|
||||
|
||||
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
|
||||
ROOT_DIR="$(cd -- "$SCRIPT_DIR/../../../.." && pwd)"
|
||||
|
||||
MODEL="${MODEL:-IndexTeam/IndexTTS-2}"
|
||||
MODEL_VERSION="${MODEL_VERSION:-2.0}"
|
||||
PORT="${PORT:-8092}"
|
||||
DEPLOY_CONFIG="${DEPLOY_CONFIG:-$ROOT_DIR/vllm_omni/deploy/indextts2.yaml}"
|
||||
|
||||
echo "Starting IndexTTS2 server with model: $MODEL"
|
||||
if [[ "$MODEL_VERSION" == "2.5" ]]; then
|
||||
if [[ -z "${MODEL:-}" ]]; then
|
||||
echo "Usage: MODEL_VERSION=2.5 MODEL=/path/to/native/bundle bash $0" >&2
|
||||
exit 2
|
||||
fi
|
||||
DEFAULT_DEPLOY_CONFIG="$ROOT_DIR/vllm_omni/deploy/indextts2_5.yaml"
|
||||
elif [[ "$MODEL_VERSION" == "2.0" ]]; then
|
||||
MODEL="${MODEL:-IndexTeam/IndexTTS-2}"
|
||||
DEFAULT_DEPLOY_CONFIG="$ROOT_DIR/vllm_omni/deploy/indextts2.yaml"
|
||||
else
|
||||
echo "MODEL_VERSION must be 2.0 or 2.5" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
DEPLOY_CONFIG="${DEPLOY_CONFIG:-$DEFAULT_DEPLOY_CONFIG}"
|
||||
|
||||
echo "Starting IndexTTS $MODEL_VERSION server with model: $MODEL"
|
||||
echo "Deploy config: $DEPLOY_CONFIG"
|
||||
|
||||
FLASHINFER_DISABLE_VERSION_CHECK=1 \
|
||||
vllm serve "$MODEL" \
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""OpenAI-compatible client for IndexTTS2 TTS via /v1/audio/speech endpoint.
|
||||
"""OpenAI-compatible client for IndexTTS 2.0/2.5 via /v1/audio/speech.
|
||||
|
||||
Examples:
|
||||
# With reference audio for voice cloning
|
||||
@@ -54,13 +54,44 @@ def main() -> None:
|
||||
parser.add_argument("--emo-alpha", type=float, default=None, help="Emotion weight in [0, 1]")
|
||||
parser.add_argument("--use-emo-text", action="store_true", help="Infer emotion vector from emo-text or text")
|
||||
parser.add_argument("--use-random", action="store_true", help="Use random emotion prototypes")
|
||||
parser.add_argument("--model", type=str, default="IndexTeam/IndexTTS-2")
|
||||
parser.add_argument(
|
||||
"--model-version",
|
||||
choices=("2.0", "2.5"),
|
||||
default="2.0",
|
||||
)
|
||||
parser.add_argument("--model", type=str, default=None)
|
||||
parser.add_argument(
|
||||
"--lang",
|
||||
default="zh",
|
||||
help=(
|
||||
"IndexTTS 2.5 language code, for example zh/en/zhen/ja/yue; "
|
||||
"zhen is mixed Chinese/English and Mandarin is a vLLM-Omni alias for zh"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-text-normalization",
|
||||
action="store_false",
|
||||
dest="text_normalization",
|
||||
help="Disable IndexTTS 2.5 text normalization",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--speed",
|
||||
type=float,
|
||||
default=1.0,
|
||||
help="IndexTTS 2.5 native synthesis speed in [0.5, 2.0]; 2.0 is faster",
|
||||
)
|
||||
parser.add_argument("--voice", type=str, default=None, help="Uploaded voice name to use instead of --ref-audio")
|
||||
parser.add_argument("--output", type=str, default="output.wav")
|
||||
parser.add_argument("--api-base", type=str, default=DEFAULT_API_BASE)
|
||||
parser.add_argument("--api-key", type=str, default=DEFAULT_API_KEY)
|
||||
parser.add_argument("--response-format", type=str, default="wav")
|
||||
args = parser.parse_args()
|
||||
if args.model is None:
|
||||
if args.model_version == "2.5":
|
||||
parser.error("--model is required for IndexTTS 2.5")
|
||||
args.model = "IndexTeam/IndexTTS-2"
|
||||
if args.model_version == "2.5" and not 0.5 <= args.speed <= 2.0:
|
||||
parser.error("IndexTTS 2.5 --speed must be between 0.5 and 2.0")
|
||||
|
||||
if not args.ref_audio and not args.voice:
|
||||
parser.error("IndexTTS2 requires --ref-audio or --voice for voice cloning")
|
||||
@@ -97,6 +128,10 @@ def main() -> None:
|
||||
extra_params["use_emo_text"] = True
|
||||
if args.use_random:
|
||||
extra_params["use_random"] = True
|
||||
if args.model_version == "2.5":
|
||||
payload["speed"] = args.speed
|
||||
extra_params["lang"] = args.lang
|
||||
extra_params["text_normalization"] = args.text_normalization
|
||||
if extra_params:
|
||||
payload["extra_params"] = extra_params
|
||||
|
||||
|
||||
@@ -84,6 +84,13 @@ forced-aligner = [
|
||||
"qwen-asr>=0.0.6", # faithful Japanese/Korean word segmentation for the TTS forced aligner
|
||||
]
|
||||
|
||||
indextts2 = [
|
||||
"fugashi>=1.2.0",
|
||||
"unidic-lite>=1.0.0",
|
||||
"wetext>=0.0.9; sys_platform != 'linux'",
|
||||
"WeTextProcessing; sys_platform == 'linux'",
|
||||
]
|
||||
|
||||
# Local-test-only extras
|
||||
local = [
|
||||
"FlagEmbedding",
|
||||
@@ -327,3 +334,7 @@ byt5 = "byt5"
|
||||
# "seeked" is the standard HTML media element DOM event name (fires when a
|
||||
# seek completes), used by the word-timestamps Gradio demo. Not a typo.
|
||||
seeked = "seeked"
|
||||
# IndexTTS 2.5 language codes and speech-emotion-recognition event token.
|
||||
fo = "fo"
|
||||
ba = "ba"
|
||||
SER = "SER"
|
||||
@@ -0,0 +1,308 @@
|
||||
# IndexTTS-2.5 for multilingual voice-cloned TTS on 1x GPU
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: IndexTeam
|
||||
- Model: `IndexTeam/IndexTTS-2.5`
|
||||
- Task: Multilingual text-to-speech with voice cloning, emotion control, and
|
||||
native speed control
|
||||
- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API, or
|
||||
offline inference through `Omni`
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe to serve IndexTTS-2.5 as a two-stage TTS system. Stage 0 is an
|
||||
autoregressive talker; Stage 1 uses EnhancedCodec, S2Mel CFM/DiT, and BigVGAN
|
||||
to produce 22.05 kHz mono speech. Each request supplies synthesis text and
|
||||
reference audio, or an uploaded audio voice, for zero-shot voice cloning.
|
||||
|
||||
IndexTTS-2.5 adds multilingual text processing and model-native speed control
|
||||
to the IndexTTS-2 serving contract.
|
||||
|
||||
## References
|
||||
|
||||
- Upstream model:
|
||||
[IndexTeam/IndexTTS-2.5 on Hugging Face](https://huggingface.co/IndexTeam/IndexTTS-2.5)
|
||||
- Online serving example:
|
||||
[`examples/online_serving/text_to_speech/README.md#indextts-2-and-indextts-25`](../../examples/online_serving/text_to_speech/README.md#indextts-2-and-indextts-25)
|
||||
- Offline inference example:
|
||||
[`examples/offline_inference/text_to_speech/README.md#indextts-2-and-indextts-25`](../../examples/offline_inference/text_to_speech/README.md#indextts-2-and-indextts-25)
|
||||
- OpenAI-compatible client:
|
||||
[`examples/online_serving/text_to_speech/indextts2/speech_client.py`](../../examples/online_serving/text_to_speech/indextts2/speech_client.py)
|
||||
- Standard deploy config:
|
||||
[`vllm_omni/deploy/indextts2_5.yaml`](../../vllm_omni/deploy/indextts2_5.yaml)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
### GPU
|
||||
|
||||
### 1x NVIDIA H20 96GB
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux
|
||||
- Python: 3.10+
|
||||
- Driver / runtime: NVIDIA CUDA environment
|
||||
- vLLM version: Match the repository requirements for your checkout
|
||||
- vLLM-Omni version or commit: Use the commit you are deploying from
|
||||
|
||||
Install vLLM-Omni with the IndexTTS text-processing dependencies:
|
||||
|
||||
```bash
|
||||
pip install 'vllm-omni[indextts2]'
|
||||
```
|
||||
|
||||
Obtain the native IndexTTS-2.5 bundle and point `MODEL` at its root. The model
|
||||
loader accepts the upstream nested `checkpoints/` layout.
|
||||
|
||||
#### Command
|
||||
|
||||
Start the standard code-only server from the repository root:
|
||||
|
||||
```bash
|
||||
MODEL_VERSION=2.5 \
|
||||
MODEL=/path/to/indextts-2.5 \
|
||||
bash examples/online_serving/text_to_speech/indextts2/run_server.sh
|
||||
```
|
||||
|
||||
This selects `vllm_omni/deploy/indextts2_5.yaml`, which uses
|
||||
`use_gpt_latent=false`. To launch directly:
|
||||
|
||||
```bash
|
||||
vllm serve /path/to/indextts-2.5 \
|
||||
--omni \
|
||||
--trust-remote-code \
|
||||
--port 8092 \
|
||||
--deploy-config vllm_omni/deploy/indextts2_5.yaml
|
||||
```
|
||||
|
||||
#### Optional: NVIDIA MPS for higher single-GPU throughput
|
||||
|
||||
Stage 0 and Stage 1 run in separate processes on the same GPU. Enabling
|
||||
NVIDIA Multi-Process Service (MPS) can improve steady-state throughput by
|
||||
allowing work from both stages to overlap more effectively. MPS is an opt-in
|
||||
deployment optimization: this recipe does not start it automatically, and the
|
||||
gain depends on the request mix and concurrency, so benchmark it with the
|
||||
intended workload before enabling it in production.
|
||||
|
||||
Pay attention to CUDA device renumbering. For example, if the MPS control
|
||||
daemon is started with physical `CUDA_VISIBLE_DEVICES=1`, that GPU is exposed
|
||||
to MPS clients as logical device `0`; launch the server with
|
||||
`CUDA_VISIBLE_DEVICES=0`, not `1`. Use a dedicated MPS pipe/log directory when
|
||||
the host is shared with other services, and stop only the MPS daemon owned by
|
||||
this deployment.
|
||||
|
||||
#### Verification
|
||||
|
||||
Send a multilingual voice-cloning request with the bundled client. Local audio
|
||||
paths are converted to base64 data URLs before transmission:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/indextts2/speech_client.py \
|
||||
--api-base http://localhost:8092 \
|
||||
--model-version 2.5 \
|
||||
--model /path/to/indextts-2.5 \
|
||||
--lang zh \
|
||||
--text "你好,这是 IndexTTS-2.5 语音合成测试。" \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--output indextts2_5.wav
|
||||
```
|
||||
|
||||
##### Native speed control
|
||||
|
||||
The public `speed` field is handled natively. Its accepted range is
|
||||
`[0.5, 2.0]`; values above `1.0` produce shorter, faster speech. The serving
|
||||
adapter converts the API convention to the model-native control with
|
||||
`duration_factor = 1.0 / speed`, so `speed=1.25` uses
|
||||
`duration_factor=0.8`:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/indextts2/speech_client.py \
|
||||
--api-base http://localhost:8092 \
|
||||
--model-version 2.5 \
|
||||
--model /path/to/indextts-2.5 \
|
||||
--speed 1.25 \
|
||||
--text "这是使用模型原生语速控制的测试。" \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--output indextts2_5_speed.wav
|
||||
```
|
||||
|
||||
##### Upload and reuse a named voice
|
||||
|
||||
Upload a reference recording once when callers should reuse the same voice
|
||||
without sending `ref_audio` on every request. `consent` is required by the
|
||||
voice-storage API:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8092/v1/audio/voices \
|
||||
-F "audio_sample=@/path/to/reference.wav" \
|
||||
-F "consent=user-consent-id" \
|
||||
-F "name=indextts_demo_voice" \
|
||||
-F "speaker_description=IndexTTS-2.5 demonstration voice"
|
||||
```
|
||||
|
||||
Confirm that the voice was stored, then synthesize with its name instead of a
|
||||
reference-audio payload:
|
||||
|
||||
```bash
|
||||
curl http://localhost:8092/v1/audio/voices
|
||||
|
||||
python examples/online_serving/text_to_speech/indextts2/speech_client.py \
|
||||
--api-base http://localhost:8092 \
|
||||
--model-version 2.5 \
|
||||
--model /path/to/indextts-2.5 \
|
||||
--voice indextts_demo_voice \
|
||||
--lang zh \
|
||||
--text "你好,这是复用已上传音色的语音合成测试。" \
|
||||
--output indextts2_5_named_voice.wav
|
||||
```
|
||||
|
||||
There are no built-in text-only preset voices. Each named IndexTTS voice must
|
||||
first be created from an uploaded reference recording.
|
||||
|
||||
##### Multilingual synthesis and text normalization
|
||||
|
||||
The client accepts `zh` (Mandarin), `en` (English), `zhen` (mixed Chinese and
|
||||
English), `ja` (Japanese), and `yue` (Cantonese). The following commands use
|
||||
the same reference voice across all supported language modes:
|
||||
|
||||
```bash
|
||||
CLIENT=examples/online_serving/text_to_speech/indextts2/speech_client.py
|
||||
MODEL=/path/to/indextts-2.5
|
||||
REF_AUDIO=/path/to/reference.wav
|
||||
|
||||
python "$CLIENT" --model-version 2.5 --model "$MODEL" \
|
||||
--ref-audio "$REF_AUDIO" --lang zh \
|
||||
--text "你好,这是中文语音合成测试。" --output indextts2_5_zh.wav
|
||||
|
||||
python "$CLIENT" --model-version 2.5 --model "$MODEL" \
|
||||
--ref-audio "$REF_AUDIO" --lang en \
|
||||
--text "Hello, this is an English speech synthesis test." \
|
||||
--output indextts2_5_en.wav
|
||||
|
||||
python "$CLIENT" --model-version 2.5 --model "$MODEL" \
|
||||
--ref-audio "$REF_AUDIO" --lang zhen \
|
||||
--text "Hello,欢迎使用 IndexTTS 二点五。" --output indextts2_5_zhen.wav
|
||||
|
||||
python "$CLIENT" --model-version 2.5 --model "$MODEL" \
|
||||
--ref-audio "$REF_AUDIO" --lang ja \
|
||||
--text "こんにちは、これは日本語の音声合成テストです。" \
|
||||
--output indextts2_5_ja.wav
|
||||
|
||||
python "$CLIENT" --model-version 2.5 --model "$MODEL" \
|
||||
--ref-audio "$REF_AUDIO" --lang yue \
|
||||
--text "你好,呢段係粵語語音合成測試。" --output indextts2_5_yue.wav
|
||||
```
|
||||
|
||||
Text normalization is enabled by default. Disable it only when the input is
|
||||
already written exactly as it should be spoken:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/indextts2/speech_client.py \
|
||||
--model-version 2.5 --model /path/to/indextts-2.5 \
|
||||
--ref-audio /path/to/reference.wav --lang en --no-text-normalization \
|
||||
--text "The temperature is twenty five degrees." \
|
||||
--output indextts2_5_pre_normalized.wav
|
||||
```
|
||||
|
||||
For raw HTTP requests, put the same language and text-normalization controls
|
||||
under `extra_params`:
|
||||
|
||||
```bash
|
||||
curl http://localhost:8092/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "/path/to/indextts-2.5",
|
||||
"input": "This is an IndexTTS-2.5 test.",
|
||||
"response_format": "wav",
|
||||
"speed": 1.0,
|
||||
"ref_audio": "data:audio/wav;base64,<BASE64_ENCODED_AUDIO>",
|
||||
"extra_params": {
|
||||
"lang": "en",
|
||||
"text_normalization": true
|
||||
}
|
||||
}' \
|
||||
--output indextts2_5_en.wav
|
||||
```
|
||||
|
||||
##### Emotion control
|
||||
|
||||
IndexTTS-2.5 supports an explicit eight-value emotion vector, an emotion text
|
||||
description, or a separate emotion-reference recording. Use one source at a
|
||||
time. Vector values are ordered as `happy`, `angry`, `sad`, `afraid`,
|
||||
`disgusted`, `melancholic`, `surprised`, and `calm`; each value must be in
|
||||
`[0, 1.2]`. `emo_alpha` controls the strength in `[0, 1]`.
|
||||
|
||||
Use an explicit emotion vector for deterministic control:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/indextts2/speech_client.py \
|
||||
--model-version 2.5 --model /path/to/indextts-2.5 \
|
||||
--ref-audio /path/to/reference.wav --lang zh \
|
||||
--text "今天真是令人开心的一天!" \
|
||||
--emo-vector 1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 \
|
||||
--emo-alpha 0.8 --output indextts2_5_happy_vector.wav
|
||||
```
|
||||
|
||||
Infer the emotion from a natural-language description with `use_emo_text`:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/indextts2/speech_client.py \
|
||||
--model-version 2.5 --model /path/to/indextts-2.5 \
|
||||
--ref-audio /path/to/reference.wav --lang zh \
|
||||
--text "我们终于完成了这项工作!" \
|
||||
--use-emo-text --emo-text "开心、兴奋而且充满活力" \
|
||||
--emo-alpha 0.8 --output indextts2_5_happy_text.wav
|
||||
```
|
||||
|
||||
Transfer emotion from a recording while retaining the speaker identity from
|
||||
`ref_audio`:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/indextts2/speech_client.py \
|
||||
--model-version 2.5 --model /path/to/indextts-2.5 \
|
||||
--ref-audio /path/to/reference.wav --lang zh \
|
||||
--text "请用情绪参考音频中的表达方式朗读这句话。" \
|
||||
--emo-audio /path/to/emotion_reference.wav \
|
||||
--emo-alpha 0.8 --output indextts2_5_emotion_audio.wav
|
||||
```
|
||||
|
||||
If multiple emotion sources are supplied, the official precedence is
|
||||
`use_emo_text` > `emo_vector` > `emo_audio` > the emotion in the speaker
|
||||
reference. `--use-random` is also available for exploratory random emotion
|
||||
prototypes, but should not be used when reproducible conditioning is required.
|
||||
|
||||
For offline inference:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/indextts2/end2end.py \
|
||||
--model /path/to/indextts-2.5 \
|
||||
--model-version 2.5 \
|
||||
--lang zh \
|
||||
--speed 1.0 \
|
||||
--text "你好,这是离线语音合成测试。" \
|
||||
--ref-audio /path/to/reference.wav
|
||||
```
|
||||
|
||||
#### Notes
|
||||
|
||||
- Hardware scope: the standard recipe has been exercised on one NVIDIA H200.
|
||||
It does not claim an unverified throughput or quality result.
|
||||
- Audio output: 22.05 kHz mono WAV.
|
||||
- Voice cloning: reference audio is required on the documented raw request
|
||||
path. Alternatively, `voice` may name an uploaded audio voice; there is no
|
||||
built-in text-only preset voice.
|
||||
- Native speed: serving maps `speed` to the model's duration factor and skips
|
||||
generic waveform speed adjustment, so speed is applied exactly once.
|
||||
- Languages: common codes include `zh`, `en`, `zhen` (mixed Chinese/English),
|
||||
`ja`, and `yue`; `Mandarin` is accepted as an alias for `zh`.
|
||||
- Japanese normalization: `ja` tokenization does not automatically expand
|
||||
numbers, dates, or percentages. Write these inputs as readable Japanese text
|
||||
before inference.
|
||||
- Emotion controls: `use_emo_text`, `emo_vector`, and `emo_audio` are
|
||||
alternative conditioning modes. Their precedence is `use_emo_text` >
|
||||
`emo_vector` > `emo_audio` > the speaker-reference emotion.
|
||||
- Sampling difference: Stage 0 uses plain vLLM sampling, not the upstream
|
||||
default `num_beams=3` beam search. Use upstream `num_beams=1` for parity
|
||||
comparisons; output may differ from the official beam-search result.
|
||||
@@ -50,6 +50,7 @@ recipes/
|
||||
| [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
|
||||
| [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
|
||||
| [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU |
|
||||
| [`IndexTeam/IndexTTS-2_5.md`](./IndexTeam/IndexTTS-2_5.md) | Online + offline multilingual voice-cloned TTS with native speed and emotion control | 1x NVIDIA H20 96GB |
|
||||
| [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 |
|
||||
| [`MiniMaxAI/MiniMax-H3-MUSA.md`](./MiniMaxAI/MiniMax-H3-MUSA.md) | MiniMax H3 T2VA, FL2VA, and Ref2VA on MUSA | 4x MTT S5000 (TP4 Ref2VA profile) |
|
||||
| [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
|
||||
|
||||
@@ -51,6 +51,14 @@ def test_load_model_configs(model_configs_path: Path) -> None:
|
||||
assert configs["test/ModelA"]["supported_tasks"] == ["voice_clone", "default_voice"]
|
||||
|
||||
|
||||
def test_indextts25_is_registered_in_shared_model_configs() -> None:
|
||||
config_path = Path(bench_tts.__file__).with_name("model_configs.yaml")
|
||||
config = bench_tts.load_model_configs(config_path)["IndexTeam/IndexTTS-2.5"]
|
||||
|
||||
assert config["supported_tasks"] == ["voice_clone"]
|
||||
assert config["task_extra_body"]["voice_clone"]["extra_params"]["lang"] == "en"
|
||||
|
||||
|
||||
def test_build_bench_args_voice_clone(model_configs_path: Path) -> None:
|
||||
configs = bench_tts.load_model_configs(model_configs_path)
|
||||
cmd = bench_tts.build_bench_args(
|
||||
@@ -120,6 +128,34 @@ def test_build_bench_args_wer_eval_adds_flag(model_configs_path: Path) -> None:
|
||||
assert "--seed-tts-wer-eval" in cmd
|
||||
|
||||
|
||||
def test_build_bench_args_supports_local_model_and_shared_sweep_options(model_configs_path: Path) -> None:
|
||||
configs = bench_tts.load_model_configs(model_configs_path)
|
||||
cmd = bench_tts.build_bench_args(
|
||||
host="localhost",
|
||||
port=8092,
|
||||
model="test/ModelA",
|
||||
served_model_name="/models/indextts25",
|
||||
task="voice_clone",
|
||||
model_cfg=configs["test/ModelA"],
|
||||
locale="en",
|
||||
num_prompts=500,
|
||||
num_warmups=5,
|
||||
request_seed=42,
|
||||
concurrency=8,
|
||||
dataset_path="/data/seed-tts",
|
||||
wer_eval=False,
|
||||
output_dir=None,
|
||||
result_filename=None,
|
||||
extra_cli_args=["--", "--tokenizer", "/models/indextts25/qwen0.6bemo4-merge"],
|
||||
)
|
||||
|
||||
assert cmd[cmd.index("--model") + 1] == "/models/indextts25"
|
||||
assert cmd[cmd.index("--num-warmups") + 1] == "5"
|
||||
assert cmd[-2:] == ["--tokenizer", "/models/indextts25/qwen0.6bemo4-merge"]
|
||||
extra_body = json.loads(cmd[cmd.index("--extra-body") + 1])
|
||||
assert extra_body == {"task_type": "Base", "seed": 42}
|
||||
|
||||
|
||||
def test_unsupported_task_exits(model_configs_path: Path, capsys: pytest.CaptureFixture, mocker) -> None:
|
||||
# ModelB does not support voice_design
|
||||
mocker.patch.object(
|
||||
|
||||
@@ -2276,6 +2276,16 @@ class TestPlatformOverrides:
|
||||
# Stage 2 unaffected fields stay at base
|
||||
assert deploy.stages[2].enforce_eager is False
|
||||
|
||||
def test_qwen2_5_omni_xpu_uses_eager_ar_stages(self):
|
||||
deploy_path = Path(get_deploy_config_path("qwen2_5_omni.yaml"))
|
||||
|
||||
deploy = load_deploy_config(deploy_path)
|
||||
deploy = _apply_platform_overrides(deploy, platform="xpu")
|
||||
|
||||
assert deploy.stages[0].enforce_eager is True
|
||||
assert deploy.stages[1].enforce_eager is True
|
||||
assert deploy.stages[2].enforce_eager is True
|
||||
|
||||
def test_xpu_overrides(self):
|
||||
deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
|
||||
if not deploy_path.exists():
|
||||
|
||||
@@ -126,6 +126,7 @@ class TestFullPayloadCoordinatorSelection(unittest.TestCase):
|
||||
("AudexCode2Wav", "audex_code2wav"),
|
||||
("AudexXCodec1", "audex_xcodec"),
|
||||
("IndexTTS2S2MelDecoder", "indextts2_s2mel_decoder"),
|
||||
("IndexTTS25S2MelDecoder", "indextts2_5_s2mel_decoder"),
|
||||
("DyninOmniForConditionalGeneration", "token2image"),
|
||||
("DyninOmniForConditionalGeneration", "token2audio"),
|
||||
}
|
||||
|
||||
@@ -6,6 +6,9 @@ explicitly patch values that differ from vLLM.
|
||||
|
||||
import argparse
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
@@ -194,6 +197,155 @@ def test_qwen3_tts_code2wav_injects_max_position_embeddings(monkeypatch):
|
||||
}
|
||||
|
||||
|
||||
def test_patch_missing_local_hf_config(tmp_path):
|
||||
"""Native model bundles may only contain a checkpoints/ directory."""
|
||||
(tmp_path / "checkpoints").mkdir()
|
||||
args = object.__new__(OmniEngineArgs)
|
||||
args.model = str(tmp_path)
|
||||
args.model_arch = "IndexTTS25TalkerForConditionalGeneration"
|
||||
args.hf_config_path = None
|
||||
|
||||
try:
|
||||
args._patch_empty_hf_config("indextts2_5")
|
||||
|
||||
assert args.hf_config_path is not None
|
||||
config_path = args.hf_config_path + "/config.json"
|
||||
with open(config_path, encoding="utf-8") as config_file:
|
||||
config = json.load(config_file)
|
||||
assert config == {
|
||||
"model_type": "indextts2_5",
|
||||
"architectures": ["IndexTTS25TalkerForConditionalGeneration"],
|
||||
}
|
||||
finally:
|
||||
if hasattr(args, "_temp_config_dir"):
|
||||
shutil.rmtree(args._temp_config_dir, ignore_errors=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"config_entry",
|
||||
["malformed", "directory", "broken_symlink", "permission_error"],
|
||||
)
|
||||
def test_non_missing_local_hf_config_error_reaches_parent_loader(tmp_path, monkeypatch, config_entry):
|
||||
"""Non-missing config errors must reach vLLM's normal loader."""
|
||||
config_path = tmp_path / "config.json"
|
||||
if config_entry == "malformed":
|
||||
config_path.write_text("{not valid json", encoding="utf-8")
|
||||
loader_error = json.JSONDecodeError("invalid config", "{not valid json", 1)
|
||||
elif config_entry == "directory":
|
||||
config_path.mkdir()
|
||||
loader_error = IsADirectoryError(str(config_path))
|
||||
elif config_entry == "broken_symlink":
|
||||
config_path.symlink_to(tmp_path / "missing-target.json")
|
||||
loader_error = FileNotFoundError(str(config_path))
|
||||
else:
|
||||
config_path.write_text("{}", encoding="utf-8")
|
||||
loader_error = PermissionError(str(config_path))
|
||||
|
||||
parent_inputs = {}
|
||||
|
||||
def controlled_get_config_dict(cls, model, **kwargs):
|
||||
if model == str(tmp_path):
|
||||
raise loader_error
|
||||
with open(os.path.join(model, "config.json"), encoding="utf-8") as config_file:
|
||||
return json.load(config_file), {}
|
||||
|
||||
def fake_parent_create_model_config(self):
|
||||
parent_inputs["called"] = True
|
||||
parent_inputs["model"] = self.model
|
||||
parent_inputs["hf_config_path"] = self.hf_config_path
|
||||
config_root = self.hf_config_path or self.model
|
||||
parent_inputs["config_root"] = config_root
|
||||
return PretrainedConfig.get_config_dict(config_root)[0]
|
||||
|
||||
monkeypatch.setattr(PretrainedConfig, "get_config_dict", classmethod(controlled_get_config_dict))
|
||||
monkeypatch.setattr(OmniEngineArgs, "_ensure_omni_models_registered", lambda self: True)
|
||||
monkeypatch.setattr(EngineArgs, "create_model_config", fake_parent_create_model_config)
|
||||
monkeypatch.setattr(
|
||||
OmniModelConfig,
|
||||
"from_vllm_model_config",
|
||||
classmethod(lambda cls, model_config, **omni_kwargs: model_config),
|
||||
)
|
||||
|
||||
args = OmniEngineArgs(
|
||||
model=str(tmp_path),
|
||||
model_arch="IndexTTS25TalkerForConditionalGeneration",
|
||||
)
|
||||
|
||||
with pytest.raises(type(loader_error)) as exc_info:
|
||||
args.create_model_config()
|
||||
|
||||
assert str(exc_info.value) == str(loader_error)
|
||||
assert parent_inputs["called"] is True
|
||||
assert parent_inputs["model"] == str(tmp_path)
|
||||
assert parent_inputs["config_root"] == str(tmp_path)
|
||||
assert parent_inputs["hf_config_path"] is None
|
||||
assert args.model == str(tmp_path)
|
||||
assert args.hf_config_path is None
|
||||
assert not hasattr(args, "_temp_config_dir")
|
||||
|
||||
|
||||
def test_patch_valid_hf_config_without_model_type_preserves_keys(tmp_path):
|
||||
"""A valid partial config should be copied before required keys are added."""
|
||||
original_config = {
|
||||
"hidden_size": 1024,
|
||||
"custom_key": {"nested": True},
|
||||
}
|
||||
(tmp_path / "config.json").write_text(json.dumps(original_config), encoding="utf-8")
|
||||
args = object.__new__(OmniEngineArgs)
|
||||
args.model = str(tmp_path)
|
||||
args.model_arch = "IndexTTS25TalkerForConditionalGeneration"
|
||||
args.hf_config_path = None
|
||||
|
||||
try:
|
||||
args._patch_empty_hf_config("indextts2_5")
|
||||
|
||||
assert args.hf_config_path is not None
|
||||
with open(os.path.join(args.hf_config_path, "config.json"), encoding="utf-8") as config_file:
|
||||
patched_config = json.load(config_file)
|
||||
assert patched_config["hidden_size"] == original_config["hidden_size"]
|
||||
assert patched_config["custom_key"] == original_config["custom_key"]
|
||||
assert patched_config["model_type"] == "indextts2_5"
|
||||
assert patched_config["architectures"] == ["IndexTTS25TalkerForConditionalGeneration"]
|
||||
finally:
|
||||
if hasattr(args, "_temp_config_dir"):
|
||||
shutil.rmtree(args._temp_config_dir, ignore_errors=True)
|
||||
|
||||
|
||||
def test_remote_hf_config_error_reaches_parent_loader(monkeypatch):
|
||||
"""Remote resolution failures must stay on vLLM's normal error path."""
|
||||
loader_error = OSError("remote config resolution failed")
|
||||
parent_inputs = {}
|
||||
|
||||
def controlled_get_config_dict(cls, model, **kwargs):
|
||||
raise loader_error
|
||||
|
||||
def fake_parent_create_model_config(self):
|
||||
parent_inputs["called"] = True
|
||||
parent_inputs["model"] = self.model
|
||||
parent_inputs["hf_config_path"] = self.hf_config_path
|
||||
return PretrainedConfig.get_config_dict(self.model)[0]
|
||||
|
||||
monkeypatch.setattr(PretrainedConfig, "get_config_dict", classmethod(controlled_get_config_dict))
|
||||
monkeypatch.setattr(OmniEngineArgs, "_ensure_omni_models_registered", lambda self: True)
|
||||
monkeypatch.setattr(EngineArgs, "create_model_config", fake_parent_create_model_config)
|
||||
|
||||
args = OmniEngineArgs(
|
||||
model="remote/model",
|
||||
model_arch="IndexTTS25TalkerForConditionalGeneration",
|
||||
)
|
||||
|
||||
with pytest.raises(OSError) as exc_info:
|
||||
args.create_model_config()
|
||||
|
||||
assert str(exc_info.value) == str(loader_error)
|
||||
assert parent_inputs["called"] is True
|
||||
assert parent_inputs["model"] == "remote/model"
|
||||
assert parent_inputs["hf_config_path"] is None
|
||||
assert args.model == "remote/model"
|
||||
assert args.hf_config_path is None
|
||||
assert not hasattr(args, "_temp_config_dir")
|
||||
|
||||
|
||||
def test_stage_specific_text_config_override():
|
||||
"""Stage swap must refresh hf_text_config, dependent attrs, and model_arch_config."""
|
||||
vllm_config = EngineArgs().create_model_config()
|
||||
|
||||
@@ -412,7 +412,7 @@ def test_stage_runtime_passes_log_stats_to_llm_replica_launch(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
runtime_mod.StageEngineCoreClientBase,
|
||||
"make_async_mp_client",
|
||||
lambda **kwargs: (captured.__setitem__("client_log_stats", kwargs["log_stats"]) or stage_client),
|
||||
lambda **kwargs: captured.__setitem__("client_log_stats", kwargs["log_stats"]) or stage_client,
|
||||
)
|
||||
|
||||
assert runtime._initialize_local_llm_replica(plan, stage_init_timeout=1) is stage_client
|
||||
@@ -807,9 +807,9 @@ def test_build_stage0_input_processor_uses_omni_input_preprocessor(monkeypatch):
|
||||
import vllm_omni.engine.stage_init_utils as init_mod
|
||||
|
||||
class DummyInputProcessor:
|
||||
def __init__(self, vllm_config):
|
||||
def __init__(self, vllm_config, renderer=None):
|
||||
self.vllm_config = vllm_config
|
||||
self.renderer = object()
|
||||
self.renderer = renderer or object()
|
||||
self.input_preprocessor = None
|
||||
|
||||
class DummyOmniInputPreprocessor:
|
||||
@@ -828,6 +828,50 @@ def test_build_stage0_input_processor_uses_omni_input_preprocessor(monkeypatch):
|
||||
assert input_processor.input_preprocessor.renderer is input_processor.renderer
|
||||
|
||||
|
||||
def test_build_stage0_input_processor_does_not_resolve_tokenizer_when_skipped(
|
||||
monkeypatch,
|
||||
):
|
||||
import vllm_omni.engine.stage_init_utils as init_mod
|
||||
|
||||
token_only_renderer = object()
|
||||
seen = {}
|
||||
|
||||
class DummyInputProcessor:
|
||||
def __init__(self, vllm_config, renderer=None):
|
||||
if renderer is None:
|
||||
raise AssertionError("skip_tokenizer_init must supply a renderer")
|
||||
seen["renderer"] = renderer
|
||||
self.renderer = renderer
|
||||
self.input_preprocessor = None
|
||||
|
||||
class DummyOmniInputPreprocessor:
|
||||
def __init__(self, vllm_config, renderer=None):
|
||||
seen["preprocessor_renderer"] = renderer
|
||||
|
||||
monkeypatch.setattr(init_mod, "InputProcessor", DummyInputProcessor)
|
||||
monkeypatch.setattr(
|
||||
init_mod,
|
||||
"_build_token_only_renderer",
|
||||
lambda _config: token_only_renderer,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
init_mod,
|
||||
"OmniInputPreprocessor",
|
||||
DummyOmniInputPreprocessor,
|
||||
)
|
||||
|
||||
config = types.SimpleNamespace(
|
||||
model_config=types.SimpleNamespace(
|
||||
skip_tokenizer_init=True,
|
||||
try_get_generation_config=lambda: {},
|
||||
)
|
||||
)
|
||||
build_stage0_input_processor(config)
|
||||
|
||||
assert seen["renderer"] is token_only_renderer
|
||||
assert seen["preprocessor_renderer"] is token_only_renderer
|
||||
|
||||
|
||||
def test_inject_kv_stage_info_infers_sender_tp_topology():
|
||||
from vllm_omni.engine.stage_init_utils import inject_kv_stage_info
|
||||
|
||||
|
||||
@@ -52,6 +52,7 @@ def test_positional_construction_matches_base_request():
|
||||
assert req.external_req_id is None
|
||||
assert req.additional_information is None
|
||||
assert req.prompt_embeds_payload is None
|
||||
assert req.num_stale_output_tokens == 0
|
||||
|
||||
|
||||
def test_positional_prompt_embeds_does_not_collide():
|
||||
|
||||
@@ -743,6 +743,48 @@ class TestTTSMethods:
|
||||
assert speech_server._is_tts is False
|
||||
assert speech_server._tts_stage is None
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model_type", "expected_speed"),
|
||||
[("indextts2_5", 1.0), ("qwen3_tts", 2.0)],
|
||||
)
|
||||
def test_audio_encode_speed_respects_adapter_native_control(
|
||||
self,
|
||||
speech_server,
|
||||
model_type,
|
||||
expected_speed,
|
||||
):
|
||||
speech_server._tts_model_type = model_type
|
||||
request = OpenAICreateSpeechRequest(input="Hello", speed=2.0)
|
||||
|
||||
assert speech_server._audio_encode_speed(request) == expected_speed
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_streaming_native_speed_skips_generic_audio_adjustment(
|
||||
self,
|
||||
speech_server,
|
||||
mocker: MockerFixture,
|
||||
):
|
||||
async def mock_generate():
|
||||
yield create_mock_audio_output_for_test()
|
||||
|
||||
speech_server._tts_model_type = "indextts2_5"
|
||||
mocker.patch.object(
|
||||
speech_server,
|
||||
"_prepare_speech_generation",
|
||||
new=mocker.AsyncMock(return_value=("speech-native-speed", mock_generate(), {})),
|
||||
)
|
||||
create_audio = mocker.patch.object(
|
||||
speech_server,
|
||||
"create_audio",
|
||||
return_value=SimpleNamespace(audio_data=b"dummy", media_type="audio/wav"),
|
||||
)
|
||||
|
||||
await speech_server._generate_audio_bytes(OpenAICreateSpeechRequest(input="Hello", speed=2.0))
|
||||
|
||||
audio_obj = create_audio.call_args.args[0]
|
||||
assert isinstance(audio_obj, CreateAudio)
|
||||
assert audio_obj.speed == 1.0
|
||||
|
||||
def test_is_tts_detection_with_tts_stage(self, mocker: MockerFixture):
|
||||
"""Test TTS model detection when TTS stage exists."""
|
||||
mock_engine_client = mocker.MagicMock()
|
||||
@@ -2178,17 +2220,17 @@ class TestStreamingProtocolValidation:
|
||||
assert req.is_streaming() is False
|
||||
|
||||
def test_stream_validation_errors(self):
|
||||
"""stream=True requires response_format in ('pcm', 'wav') and speed=1.0."""
|
||||
"""The request schema validates formats; model-aware speed checks happen in serving."""
|
||||
with pytest.raises(ValidationError, match="requires response_format='pcm' or 'wav'"):
|
||||
OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="mp3")
|
||||
with pytest.raises(ValidationError, match="Speed adjustment is not supported"):
|
||||
OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="pcm", speed=2.0)
|
||||
request = OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="pcm", speed=2.0)
|
||||
assert request.speed == 2.0
|
||||
|
||||
def test_stream_format_audio_validation_errors(self):
|
||||
with pytest.raises(ValidationError, match="requires response_format='pcm' or 'wav'"):
|
||||
OpenAICreateSpeechRequest(input="Hello", stream_format="audio", response_format="mp3")
|
||||
with pytest.raises(ValidationError, match="Speed adjustment is not supported"):
|
||||
OpenAICreateSpeechRequest(input="Hello", stream_format="audio", response_format="pcm", speed=2.0)
|
||||
request = OpenAICreateSpeechRequest(input="Hello", stream_format="audio", response_format="pcm", speed=2.0)
|
||||
assert request.speed == 2.0
|
||||
|
||||
def test_stream_valid(self):
|
||||
"""stream=True + response_format in ('pcm', 'wav') + speed=1.0 is accepted as SSE."""
|
||||
@@ -2271,6 +2313,18 @@ class TestStreamingResponse:
|
||||
models=mock_models,
|
||||
request_logger=mocker.MagicMock(),
|
||||
)
|
||||
speech_server._tts_model_type = None
|
||||
speech_server._adapter = None
|
||||
speech_server._diffusion_mode = False
|
||||
mocker.patch.object(speech_server, "_uses_native_speed_control", return_value=False)
|
||||
mocker.patch.object(
|
||||
speech_server,
|
||||
"create_error_response",
|
||||
side_effect=lambda message, **_: JSONResponse(
|
||||
status_code=400,
|
||||
content={"error": {"message": message}},
|
||||
),
|
||||
)
|
||||
|
||||
original_create_speech = speech_server.create_speech
|
||||
sig = signature(original_create_speech)
|
||||
@@ -2285,6 +2339,7 @@ class TestStreamingResponse:
|
||||
|
||||
app = FastAPI()
|
||||
app.add_api_route("/v1/audio/speech", speech_server.create_speech, methods=["POST"], response_model=None)
|
||||
app.state.speech_server = speech_server
|
||||
return app
|
||||
|
||||
@staticmethod
|
||||
@@ -2393,6 +2448,39 @@ class TestStreamingResponse:
|
||||
assert response.status_code in (400, 422)
|
||||
assert "audio/" not in response.headers.get("content-type", "")
|
||||
|
||||
@pytest.mark.parametrize("stream_format", ["sse", "audio"])
|
||||
def test_native_speed_control_accepts_streaming_speed(
|
||||
self,
|
||||
streaming_app,
|
||||
mocker: MockerFixture,
|
||||
stream_format: str,
|
||||
):
|
||||
speech_server = streaming_app.state.speech_server
|
||||
speech_server._tts_model_type = "indextts2_5"
|
||||
speech_server._uses_native_speed_control.return_value = True
|
||||
|
||||
async def prepare(request, request_id=None, **kwargs):
|
||||
del kwargs
|
||||
generator = speech_server.engine_client.generate(prompt={}, request_id=request_id)
|
||||
return request_id, generator, {"duration_factor": [1.0 / request.speed]}
|
||||
|
||||
mocker.patch.object(speech_server, "_prepare_speech_generation", side_effect=prepare)
|
||||
|
||||
client = TestClient(streaming_app)
|
||||
response = client.post(
|
||||
"/v1/audio/speech",
|
||||
json={
|
||||
"input": "Hello",
|
||||
"stream_format": stream_format,
|
||||
"response_format": "pcm",
|
||||
"speed": 2.0,
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
expected_media_type = "text/event-stream" if stream_format == "sse" else "audio/pcm"
|
||||
assert expected_media_type in response.headers["content-type"]
|
||||
|
||||
@pytest.fixture
|
||||
def erroring_streaming_app(self, mocker: MockerFixture):
|
||||
"""Test app whose mock engine raises mid-stream, to exercise the SSE error event."""
|
||||
@@ -2616,6 +2704,18 @@ def test_streaming_speech_session_config_accepts_non_streaming_mode():
|
||||
assert config.non_streaming_mode is True
|
||||
|
||||
|
||||
def test_streaming_speech_session_config_scopes_native_speed_to_http():
|
||||
with pytest.raises(
|
||||
ValidationError,
|
||||
match="native speed control is only available through HTTP streaming",
|
||||
):
|
||||
StreamingSpeechSessionConfig(
|
||||
stream_audio=True,
|
||||
response_format="pcm",
|
||||
speed=1.5,
|
||||
)
|
||||
|
||||
|
||||
class TestAsyncOmniSupportedTasks:
|
||||
"""Test that AsyncOmni reports correct supported tasks based on output modalities."""
|
||||
|
||||
|
||||
@@ -4,6 +4,9 @@
|
||||
Pure-Python registry/resolution logic; no model or GPU resources are loaded.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.entrypoints.openai.tts_adapters import (
|
||||
@@ -11,9 +14,19 @@ from vllm_omni.entrypoints.openai.tts_adapters import (
|
||||
ARTTSAdapter,
|
||||
DiffusionTTSAdapter,
|
||||
all_tts_model_types,
|
||||
detect_tts_model_type,
|
||||
resolve_adapter,
|
||||
)
|
||||
from vllm_omni.entrypoints.openai.tts_adapters.indextts2 import (
|
||||
IndexTTS2Adapter,
|
||||
IndexTTS25Adapter,
|
||||
indextts2_conditioning_cache_salt,
|
||||
)
|
||||
from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSAdapter
|
||||
from vllm_omni.model_executor.models.indextts2 import prompt_utils
|
||||
from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import (
|
||||
INDEXTTS25_TOKENIZER_FILE,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -35,6 +48,8 @@ EXPECTED_MODEL_TYPES = {
|
||||
"higgs_audio_v3",
|
||||
"glm_tts",
|
||||
"step_audio2",
|
||||
"indextts2",
|
||||
"indextts2_5",
|
||||
}
|
||||
|
||||
|
||||
@@ -84,6 +99,113 @@ def test_qwen3_tts_metadata():
|
||||
assert issubclass(Qwen3TTSAdapter, ARTTSAdapter)
|
||||
|
||||
|
||||
def test_indextts_adapters_are_versioned():
|
||||
assert resolve_adapter("indextts2") is IndexTTS2Adapter
|
||||
assert resolve_adapter("indextts2_5") is IndexTTS25Adapter
|
||||
assert IndexTTS25Adapter.stage_keys == frozenset({"indextts2_5_talker"})
|
||||
assert detect_tts_model_type("indextts2_5_talker", None) == "indextts2_5"
|
||||
|
||||
|
||||
def test_indextts25_validates_explicit_language():
|
||||
adapter = IndexTTS25Adapter(type("Context", (), {"server": object()})())
|
||||
|
||||
assert adapter._validate_extra_params({"lang": "ja"}) is None
|
||||
assert "Unsupported IndexTTS 2.5 language" in adapter._validate_extra_params({"lang": "xx-invalid"})
|
||||
|
||||
|
||||
def _indextts25_adapter_and_request(*, speed: float):
|
||||
server = SimpleNamespace(
|
||||
uploaded_speakers={},
|
||||
_validate_ref_audio_format=lambda ref_audio: None,
|
||||
)
|
||||
adapter = IndexTTS25Adapter(SimpleNamespace(server=server))
|
||||
request = SimpleNamespace(
|
||||
input="hello",
|
||||
voice="alloy",
|
||||
ref_audio=object(),
|
||||
max_new_tokens=None,
|
||||
extra_params=None,
|
||||
speed=speed,
|
||||
)
|
||||
return adapter, request
|
||||
|
||||
|
||||
def test_indextts25_uses_native_speed_control_duration_factor():
|
||||
adapter, fast_request = _indextts25_adapter_and_request(speed=2.0)
|
||||
_, slow_request = _indextts25_adapter_and_request(speed=0.5)
|
||||
|
||||
assert adapter.native_speed_control is True
|
||||
assert adapter.validate(fast_request) is None
|
||||
assert adapter.validate(slow_request) is None
|
||||
assert asyncio.run(adapter._build_params(fast_request))["duration_factor"] == [0.5]
|
||||
assert asyncio.run(adapter._build_params(slow_request))["duration_factor"] == [2.0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("speed", [0.49, 2.01])
|
||||
def test_indextts25_rejects_out_of_range_native_speed(speed):
|
||||
adapter, request = _indextts25_adapter_and_request(speed=speed)
|
||||
|
||||
assert adapter.validate(request) == "IndexTTS 2.5 speed must be between 0.5 and 2.0"
|
||||
|
||||
|
||||
def test_indextts25_speed_does_not_change_conditioning_cache_salt():
|
||||
adapter, fast_request = _indextts25_adapter_and_request(speed=2.0)
|
||||
_, slow_request = _indextts25_adapter_and_request(speed=0.5)
|
||||
slow_request.ref_audio = fast_request.ref_audio
|
||||
|
||||
fast_params = asyncio.run(adapter._build_params(fast_request))
|
||||
slow_params = asyncio.run(adapter._build_params(slow_request))
|
||||
|
||||
assert indextts2_conditioning_cache_salt(
|
||||
fast_request,
|
||||
fast_params,
|
||||
) == indextts2_conditioning_cache_salt(slow_request, slow_params)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("hf_config", "expected_tokenizer_file"),
|
||||
[
|
||||
(SimpleNamespace(tokenizer_file="custom-tokenizer.tiktoken"), "custom-tokenizer.tiktoken"),
|
||||
(SimpleNamespace(), INDEXTTS25_TOKENIZER_FILE),
|
||||
],
|
||||
)
|
||||
def test_indextts25_build_uses_configured_tokenizer_file(
|
||||
monkeypatch,
|
||||
hf_config,
|
||||
expected_tokenizer_file,
|
||||
):
|
||||
captured = {}
|
||||
|
||||
def fake_estimate(*args, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return 4
|
||||
|
||||
async def fake_build_params(request):
|
||||
return {"lang": ["en"], "text_normalization": [True]}
|
||||
|
||||
monkeypatch.setattr(
|
||||
prompt_utils,
|
||||
"estimate_indextts2_prefill_prompt_len",
|
||||
fake_estimate,
|
||||
)
|
||||
server = SimpleNamespace(
|
||||
engine_client=SimpleNamespace(
|
||||
model_config=SimpleNamespace(
|
||||
model="/model",
|
||||
hf_config=hf_config,
|
||||
)
|
||||
)
|
||||
)
|
||||
adapter = IndexTTS25Adapter(SimpleNamespace(server=server))
|
||||
monkeypatch.setattr(adapter, "_build_params", fake_build_params)
|
||||
request = SimpleNamespace(input="hello", ref_audio=None)
|
||||
|
||||
prepared = asyncio.run(adapter.build(request, [], False))
|
||||
|
||||
assert prepared.prompt["prompt_token_ids"] == [1] * 4
|
||||
assert captured["tokenizer_file"] == expected_tokenizer_file
|
||||
|
||||
|
||||
def test_diffusion_adapter_extra_body_params_fallback():
|
||||
class _DiffAdapter(DiffusionTTSAdapter):
|
||||
name = "diff_probe"
|
||||
|
||||
@@ -149,6 +149,7 @@ _PIPELINE_STAGES = [
|
||||
"glm_tts_dit",
|
||||
"higgs_audio_v2",
|
||||
"higgs_audio_v3",
|
||||
"indextts2_5_talker",
|
||||
"indextts2_s2mel_decoder",
|
||||
"indextts2_talker",
|
||||
"latent_generator",
|
||||
@@ -213,18 +214,23 @@ _CLAIMED_ARCH_CASES = [
|
||||
("fused_thinker_talker", "Qwen3TTSForConditionalGeneration"),
|
||||
]
|
||||
|
||||
_REACHABLE = list(itertools.product(_STAGES, _UNCLAIMED_ARCHS)) + _CLAIMED_ARCH_CASES
|
||||
# The frozen ladder can only be compared over stages it knew about. New
|
||||
# adapter-owned stages are covered by their adapter tests and by the registry
|
||||
# completeness checks below.
|
||||
_NEW_ADAPTER_STAGE_KEYS = all_tts_stage_keys() - _LEGACY_TTS_MODEL_STAGES
|
||||
_LEGACY_REACHABLE_STAGES = [stage for stage in _STAGES if stage not in _NEW_ADAPTER_STAGE_KEYS]
|
||||
_REACHABLE = list(itertools.product(_LEGACY_REACHABLE_STAGES, _UNCLAIMED_ARCHS)) + _CLAIMED_ARCH_CASES
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("model_stage", "model_arch"), _REACHABLE)
|
||||
def test_detection_matches_legacy_ladder(model_stage, model_arch):
|
||||
"""Registry detection is identical to the ladder it replaced.
|
||||
|
||||
Covers every ``model_stage`` any in-tree pipeline emits (``_PIPELINE_STAGES``,
|
||||
plus negatives) against every architecture no adapter claims, and each
|
||||
claimed architecture against the stage keys it really ships with. Pairing a
|
||||
claimed architecture with an unrelated model's stage key is excluded because
|
||||
no pipeline can produce it, and is covered separately by
|
||||
Covers every stage the frozen ladder knew about, plus non-TTS pipeline
|
||||
stages and negatives, against every architecture no adapter claims. New
|
||||
adapter-owned stages are tested by their adapter suites. Pairing a claimed
|
||||
architecture with an unrelated model's stage key is excluded because no
|
||||
pipeline can produce it, and is covered separately by
|
||||
``test_arch_matching_is_a_fallback_not_an_override``.
|
||||
"""
|
||||
assert detect_tts_model_type(model_stage, model_arch) == _legacy_detect(model_stage, model_arch)
|
||||
@@ -251,7 +257,7 @@ def test_arch_matching_is_a_fallback_not_an_override():
|
||||
|
||||
def test_stage_keys_cover_legacy_stage_set():
|
||||
"""No stage key was dropped when the module constants were deleted."""
|
||||
assert all_tts_stage_keys() == frozenset(_LEGACY_TTS_MODEL_STAGES)
|
||||
assert _LEGACY_TTS_MODEL_STAGES <= all_tts_stage_keys()
|
||||
|
||||
|
||||
def test_pipeline_stage_list_is_complete():
|
||||
|
||||
@@ -0,0 +1,108 @@
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[3]
|
||||
SCRIPT = ROOT / "examples/online_serving/text_to_speech/indextts2/run_server.sh"
|
||||
README = ROOT / "examples/online_serving/text_to_speech/README.md"
|
||||
MACHINE_SPECIFIC_HOME_PATH = re.compile(r"(?:/home|/Users)/[^/\s`]+/")
|
||||
|
||||
|
||||
def run_server(
|
||||
tmp_path: Path,
|
||||
*,
|
||||
model_version: str | None = None,
|
||||
model: str | None = None,
|
||||
) -> tuple[subprocess.CompletedProcess[str], Path]:
|
||||
argv_file = tmp_path / "vllm.argv"
|
||||
stub = tmp_path / "vllm"
|
||||
stub.write_text(
|
||||
'#!/bin/bash\nprintf "%s\\0" "$@" > "$VLLM_ARGV_FILE"\n',
|
||||
encoding="utf-8",
|
||||
)
|
||||
stub.chmod(0o755)
|
||||
|
||||
env = os.environ.copy()
|
||||
for name in ("MODEL_VERSION", "MODEL", "DEPLOY_CONFIG", "PORT"):
|
||||
env.pop(name, None)
|
||||
env["PATH"] = f"{tmp_path}{os.pathsep}{env['PATH']}"
|
||||
env["VLLM_ARGV_FILE"] = str(argv_file)
|
||||
|
||||
if model_version is not None:
|
||||
env["MODEL_VERSION"] = model_version
|
||||
if model is not None:
|
||||
env["MODEL"] = model
|
||||
result = subprocess.run(
|
||||
["bash", str(SCRIPT)],
|
||||
cwd=ROOT,
|
||||
env=env,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
return result, argv_file
|
||||
|
||||
|
||||
def read_argv(argv_file: Path) -> list[str]:
|
||||
return argv_file.read_bytes().decode().rstrip("\0").split("\0")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", [None, ""])
|
||||
def test_version_2_5_requires_model_before_invoking_vllm(
|
||||
tmp_path: Path,
|
||||
model: str | None,
|
||||
) -> None:
|
||||
result, argv_file = run_server(tmp_path, model_version="2.5", model=model)
|
||||
|
||||
assert result.returncode == 2
|
||||
assert "MODEL=/path/to/native/bundle" in result.stderr
|
||||
assert not argv_file.exists()
|
||||
|
||||
|
||||
def test_version_2_5_selects_deploy_config(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
model = "/models/indextts-2.5"
|
||||
result, argv_file = run_server(
|
||||
tmp_path,
|
||||
model_version="2.5",
|
||||
model=model,
|
||||
)
|
||||
|
||||
assert result.returncode == 0, result.stderr
|
||||
argv = read_argv(argv_file)
|
||||
assert argv[:2] == ["serve", model]
|
||||
deploy_config_index = argv.index("--deploy-config") + 1
|
||||
assert Path(argv[deploy_config_index]) == ROOT / "vllm_omni/deploy/indextts2_5.yaml"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_version", [None, "2.0"])
|
||||
def test_default_and_version_2_0_use_huggingface_model(
|
||||
tmp_path: Path,
|
||||
model_version: str | None,
|
||||
) -> None:
|
||||
result, argv_file = run_server(tmp_path, model_version=model_version)
|
||||
|
||||
assert result.returncode == 0, result.stderr
|
||||
argv = read_argv(argv_file)
|
||||
assert argv[:2] == ["serve", "IndexTeam/IndexTTS-2"]
|
||||
deploy_config_index = argv.index("--deploy-config") + 1
|
||||
assert Path(argv[deploy_config_index]) == ROOT / "vllm_omni/deploy/indextts2.yaml"
|
||||
|
||||
|
||||
def test_invalid_model_version_exits_before_invoking_vllm(tmp_path: Path) -> None:
|
||||
result, argv_file = run_server(tmp_path, model_version="3.0")
|
||||
|
||||
assert result.returncode == 2
|
||||
assert "MODEL_VERSION must be 2.0 or 2.5" in result.stderr
|
||||
assert not argv_file.exists()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", [SCRIPT, README])
|
||||
def test_public_example_does_not_contain_machine_specific_home_path(path: Path) -> None:
|
||||
assert MACHINE_SPECIFIC_HOME_PATH.search(path.read_text(encoding="utf-8")) is None
|
||||
Whitespace-only changes.
@@ -0,0 +1,89 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
|
||||
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
|
||||
from vllm_omni.model_executor.models.indextts2.configuration_indextts2 import (
|
||||
IndexTTS2Config,
|
||||
IndexTTS25Config,
|
||||
)
|
||||
from vllm_omni.model_executor.models.indextts2.pipeline import INDEXTTS25_PIPELINE
|
||||
|
||||
|
||||
def test_indextts25_defaults_are_distinct_from_v2():
|
||||
v2 = IndexTTS2Config()
|
||||
v25 = IndexTTS25Config()
|
||||
|
||||
assert v2.model_type == "indextts2"
|
||||
assert v2.gpt["number_text_tokens"] == 12000
|
||||
assert v2.gpt["condition_num_latent"] == 32
|
||||
assert v2.use_gpt_latent is True
|
||||
assert v2.semantic_codec_type == "repcodec"
|
||||
assert v2.vocab_size == v2.gpt["number_mel_codes"]
|
||||
|
||||
assert v25.model_type == "indextts2_5"
|
||||
assert v25.gpt["number_text_tokens"] == 60509
|
||||
assert v25.gpt["condition_num_latent"] == 32
|
||||
assert v25.use_gpt_latent is False
|
||||
assert v25.semantic_codec_type == "enhanced"
|
||||
assert v25.semantic_codec_checkpoint == "codec.pth"
|
||||
assert v25.tokenizer_file == "multilingual_zh_ja_yue_char_del.tiktoken"
|
||||
assert v25.vocab_size == v25.gpt["number_mel_codes"]
|
||||
|
||||
|
||||
def test_indextts25_explicit_overrides_win_over_defaults():
|
||||
config = IndexTTS25Config(
|
||||
use_gpt_latent=True,
|
||||
gpt={
|
||||
"number_text_tokens": 70000,
|
||||
"condition_num_latent": 2,
|
||||
},
|
||||
vocab_size=9000,
|
||||
)
|
||||
|
||||
assert config.use_gpt_latent is True
|
||||
assert config.gpt["number_text_tokens"] == 70000
|
||||
assert config.gpt["condition_num_latent"] == 2
|
||||
assert config.gpt["model_dim"] == config.hidden_size
|
||||
assert config.vocab_size == 9000
|
||||
|
||||
|
||||
def test_indextts25_pipeline_is_registered_with_two_distinct_stages():
|
||||
assert OMNI_PIPELINES["indextts2_5"] is INDEXTTS25_PIPELINE
|
||||
assert INDEXTTS25_PIPELINE.default_deploy_config_name == "indextts2_5.yaml"
|
||||
assert [stage.model_stage for stage in INDEXTTS25_PIPELINE.stages] == [
|
||||
"indextts2_5_talker",
|
||||
"indextts2_5_s2mel_decoder",
|
||||
]
|
||||
assert [stage.model_arch for stage in INDEXTTS25_PIPELINE.stages] == [
|
||||
None,
|
||||
"IndexTTS25S2MelDecoder",
|
||||
]
|
||||
assert all("tokenizer" not in stage.extras for stage in INDEXTTS25_PIPELINE.stages)
|
||||
|
||||
|
||||
def test_indextts25_default_deploy_selects_validated_triton_backend():
|
||||
deploy_path = Path(__file__).parents[4] / "vllm_omni" / "deploy" / "indextts2_5.yaml"
|
||||
config = yaml.safe_load(deploy_path.read_text(encoding="utf-8"))
|
||||
stage0 = next(stage for stage in config["stages"] if stage["stage_id"] == 0)
|
||||
|
||||
assert stage0["attention_backend"] == "TRITON_ATTN"
|
||||
assert stage0["enable_chunked_prefill"] is False
|
||||
|
||||
|
||||
def test_indextts25_deploy_configs_keep_latent_policy_in_sync():
|
||||
deploy_dir = Path(__file__).parents[4] / "vllm_omni" / "deploy"
|
||||
expected = {"indextts2_5.yaml": False}
|
||||
for filename, use_gpt_latent in expected.items():
|
||||
with open(deploy_dir / filename, encoding="utf-8") as deploy_file:
|
||||
config = yaml.safe_load(deploy_file)
|
||||
assert config["pipeline"] == "indextts2_5"
|
||||
assert config["async_chunk"] is False
|
||||
assert all("tokenizer" not in stage for stage in config["stages"])
|
||||
assert [stage["hf_overrides"]["use_gpt_latent"] for stage in config["stages"]] == [
|
||||
use_gpt_latent,
|
||||
use_gpt_latent,
|
||||
]
|
||||
@@ -0,0 +1,75 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.model_executor.models.indextts2.dit_cuda_graph import (
|
||||
CUDAGraphDiTRunner,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _CopySpy:
|
||||
def __init__(self):
|
||||
self.sources = []
|
||||
|
||||
def copy_(self, source):
|
||||
self.sources.append(source)
|
||||
return self
|
||||
|
||||
|
||||
class _GraphSpy:
|
||||
def __init__(self):
|
||||
self.replays = 0
|
||||
|
||||
def replay(self):
|
||||
self.replays += 1
|
||||
|
||||
|
||||
def test_replay_only_updates_dynamic_inputs_and_reuses_static_output():
|
||||
runner = CUDAGraphDiTRunner(transformer=object())
|
||||
graph = _GraphSpy()
|
||||
static_x = _CopySpy()
|
||||
static_c = _CopySpy()
|
||||
static_pos = _CopySpy()
|
||||
static_mask = _CopySpy()
|
||||
static_out = object()
|
||||
entry = SimpleNamespace(
|
||||
graph=graph,
|
||||
static_x=static_x,
|
||||
static_c=static_c,
|
||||
static_pos=static_pos,
|
||||
static_mask=static_mask,
|
||||
static_out=static_out,
|
||||
)
|
||||
x_in = object()
|
||||
c = object()
|
||||
|
||||
output = runner._replay(entry, x_in=x_in, c=c)
|
||||
|
||||
assert static_x.sources == [x_in]
|
||||
assert static_c.sources == [c]
|
||||
assert static_pos.sources == []
|
||||
assert static_mask.sources == []
|
||||
assert graph.replays == 1
|
||||
assert output is static_out
|
||||
|
||||
|
||||
def test_graph_stats_start_empty_and_remain_bounded():
|
||||
runner = CUDAGraphDiTRunner(transformer=object(), max_graphs=4)
|
||||
|
||||
stats = runner.stats_snapshot()
|
||||
stats["hits"] = 99
|
||||
|
||||
assert runner.stats_snapshot() == {
|
||||
"calls": 0,
|
||||
"hits": 0,
|
||||
"captures": 0,
|
||||
"capture_failures": 0,
|
||||
"evictions": 0,
|
||||
"eager": 0,
|
||||
"cache_size": 0,
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import torch
|
||||
|
||||
from vllm_omni.model_executor.models.indextts2 import codec as codec_module
|
||||
from vllm_omni.model_executor.models.indextts2 import preprocess_utils
|
||||
from vllm_omni.model_executor.models.indextts2.codec import EnhancedCodec
|
||||
|
||||
|
||||
def test_enhanced_codec_decode_restores_hidden_size_and_doubles_time():
|
||||
codec = EnhancedCodec(
|
||||
codebook_size=8,
|
||||
hidden_size=4,
|
||||
codebook_dim=2,
|
||||
vocos_dim=4,
|
||||
vocos_intermediate_dim=8,
|
||||
vocos_num_layers=1,
|
||||
downsample_scale=2,
|
||||
).eval()
|
||||
|
||||
decoded = codec.decode(torch.tensor([[0, 1, 2]]))
|
||||
|
||||
assert decoded.shape == (1, 6, 4)
|
||||
|
||||
|
||||
def test_enhanced_codec_inference_module_does_not_allocate_unused_encoder():
|
||||
codec = EnhancedCodec(
|
||||
codebook_size=8,
|
||||
hidden_size=4,
|
||||
codebook_dim=2,
|
||||
vocos_dim=4,
|
||||
vocos_intermediate_dim=8,
|
||||
vocos_num_layers=1,
|
||||
)
|
||||
|
||||
parameter_names = set(dict(codec.named_parameters()))
|
||||
assert not any(name.startswith("encoder.") for name in parameter_names)
|
||||
assert not any(name.startswith("down.") for name in parameter_names)
|
||||
assert any(name.startswith("decoder.") for name in parameter_names)
|
||||
assert any(name.startswith("quantizer.") for name in parameter_names)
|
||||
|
||||
|
||||
def test_semantic_codec_loader_selects_enhanced_checkpoint(tmp_path, monkeypatch):
|
||||
checkpoint = tmp_path / "codec.pth"
|
||||
checkpoint.touch()
|
||||
seen = {}
|
||||
|
||||
class FakeCodec:
|
||||
def load_checkpoint(self, path):
|
||||
seen["path"] = path
|
||||
|
||||
def to(self, **kwargs):
|
||||
seen["to"] = kwargs
|
||||
return self
|
||||
|
||||
def eval(self):
|
||||
return self
|
||||
|
||||
def parameters(self):
|
||||
return ()
|
||||
|
||||
monkeypatch.setattr(codec_module, "EnhancedCodec", lambda **kwargs: FakeCodec())
|
||||
preprocess_utils._semantic_codec_cache.clear()
|
||||
|
||||
codec = preprocess_utils.load_semantic_codec(
|
||||
str(tmp_path),
|
||||
{"codebook_size": 8192},
|
||||
torch.device("cpu"),
|
||||
codec_type="enhanced",
|
||||
checkpoint_name="codec.pth",
|
||||
)
|
||||
|
||||
assert isinstance(codec, FakeCodec)
|
||||
assert seen["path"] == str(checkpoint)
|
||||
assert seen["to"] == {"device": torch.device("cpu"), "dtype": torch.float32}
|
||||
|
||||
|
||||
def test_model_asset_resolver_supports_native_checkpoints_layout(tmp_path):
|
||||
checkpoint = tmp_path / "checkpoints" / "codec.pth"
|
||||
checkpoint.parent.mkdir()
|
||||
checkpoint.touch()
|
||||
|
||||
assert preprocess_utils.resolve_model_file(str(tmp_path), "codec.pth") == str(checkpoint)
|
||||
@@ -0,0 +1,202 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import importlib.util
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.model_executor.models.indextts2 import prompt_utils
|
||||
from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import (
|
||||
INDEXTTS25_TOKENIZER_FILE,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _load_indextts2_offline_example():
|
||||
example_path = (
|
||||
Path(__file__).resolve().parents[4]
|
||||
/ "examples"
|
||||
/ "offline_inference"
|
||||
/ "text_to_speech"
|
||||
/ "indextts2"
|
||||
/ "end2end.py"
|
||||
)
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"indextts2_offline_example_test",
|
||||
example_path,
|
||||
)
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def test_offline_main_uses_resolved_stage0_tokenizer_file(monkeypatch, tmp_path):
|
||||
monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
example = _load_indextts2_offline_example()
|
||||
stage0 = SimpleNamespace(engine_args=SimpleNamespace(hf_overrides={"tokenizer_file": "custom-tokenizer.tiktoken"}))
|
||||
fake_omni = SimpleNamespace(
|
||||
stage_configs=[stage0],
|
||||
generate=lambda *args, **kwargs: [],
|
||||
)
|
||||
captured = {}
|
||||
|
||||
monkeypatch.setattr(example, "Omni", lambda **kwargs: fake_omni)
|
||||
monkeypatch.setattr(example, "SamplingParams", lambda **kwargs: kwargs)
|
||||
monkeypatch.setattr(
|
||||
example,
|
||||
"build_request",
|
||||
lambda **kwargs: captured.update(kwargs) or {"additional_information": {}},
|
||||
)
|
||||
|
||||
example.main(
|
||||
SimpleNamespace(
|
||||
model="/model",
|
||||
model_version="2.5",
|
||||
deploy_config="/deploy/custom.yaml",
|
||||
use_gpt_latent=False,
|
||||
stage_init_timeout=300,
|
||||
seed=42,
|
||||
output_dir=str(tmp_path),
|
||||
text="hello",
|
||||
ref_audio=None,
|
||||
emo_audio=None,
|
||||
emo_text=None,
|
||||
emo_vector=None,
|
||||
emo_alpha=None,
|
||||
use_emo_text=False,
|
||||
use_random=False,
|
||||
lang="en",
|
||||
text_normalization=True,
|
||||
)
|
||||
)
|
||||
|
||||
assert captured["tokenizer_file"] == "custom-tokenizer.tiktoken"
|
||||
|
||||
|
||||
def test_indextts25_prompt_len_uses_three_conditioning_tokens(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
def fake_prepare(*args, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return [11, 12, 13, 14], 0
|
||||
|
||||
monkeypatch.setattr(
|
||||
prompt_utils,
|
||||
"prepare_indextts25_text",
|
||||
fake_prepare,
|
||||
)
|
||||
|
||||
prompt_len = prompt_utils.estimate_indextts2_prefill_prompt_len(
|
||||
"/model",
|
||||
"hello",
|
||||
model_type="indextts2_5",
|
||||
lang="en",
|
||||
text_normalization=False,
|
||||
)
|
||||
|
||||
# Official prepare_gpt_inputs wraps the language-prefixed tokenizer IDs
|
||||
# with start_text=0 and stop_text=1 before appending start_mel.
|
||||
assert prompt_len == 3 + (4 + 2) + 1
|
||||
assert captured["text_normalization"] is False
|
||||
assert captured["tokenizer_file"] == INDEXTTS25_TOKENIZER_FILE
|
||||
|
||||
|
||||
def test_indextts25_prompt_len_uses_custom_tokenizer_file(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
def fake_prepare(*args, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return [11, 12], 0
|
||||
|
||||
monkeypatch.setattr(
|
||||
prompt_utils,
|
||||
"prepare_indextts25_text",
|
||||
fake_prepare,
|
||||
)
|
||||
|
||||
prompt_utils.estimate_indextts2_prefill_prompt_len(
|
||||
"/model",
|
||||
"hello",
|
||||
model_type="indextts2_5",
|
||||
tokenizer_file="custom-tokenizer.tiktoken",
|
||||
)
|
||||
|
||||
assert captured["tokenizer_file"] == "custom-tokenizer.tiktoken"
|
||||
|
||||
|
||||
def test_indextts25_prompt_len_filters_existing_text_wrapper_ids(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
prompt_utils,
|
||||
"prepare_indextts25_text",
|
||||
lambda *args, **kwargs: ([0, 58838, 1, 42, 0], 0),
|
||||
)
|
||||
|
||||
prompt_len = prompt_utils.estimate_indextts2_prefill_prompt_len(
|
||||
"/model",
|
||||
"hello",
|
||||
model_type="indextts2_5",
|
||||
lang="en",
|
||||
)
|
||||
|
||||
assert prompt_len == 3 + (2 + 2) + 1
|
||||
|
||||
|
||||
def test_indextts2_prompt_len_ignores_v25_tokenizer_file(monkeypatch):
|
||||
class FakeTokenizer:
|
||||
def encode(self, text, *, add_special_tokens):
|
||||
assert text == "hello"
|
||||
assert add_special_tokens is False
|
||||
return [11, 12]
|
||||
|
||||
def fail_prepare(*args, **kwargs):
|
||||
raise AssertionError("IndexTTS 2 must not use the IndexTTS 2.5 tokenizer")
|
||||
|
||||
monkeypatch.setattr(
|
||||
prompt_utils,
|
||||
"_get_text_tokenizer",
|
||||
lambda model_id_or_path: FakeTokenizer(),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
prompt_utils,
|
||||
"prepare_indextts25_text",
|
||||
fail_prepare,
|
||||
)
|
||||
|
||||
prompt_len = prompt_utils.estimate_indextts2_prefill_prompt_len(
|
||||
"/model",
|
||||
"hello",
|
||||
model_type="indextts2",
|
||||
tokenizer_file="custom.tiktoken",
|
||||
)
|
||||
|
||||
assert prompt_len == 34 + (2 + 2) + 1
|
||||
|
||||
|
||||
def test_indextts25_prompt_ids_preserve_placeholder_value(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
def fake_estimate(*args, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return 10
|
||||
|
||||
monkeypatch.setattr(
|
||||
prompt_utils,
|
||||
"estimate_indextts2_prefill_prompt_len",
|
||||
fake_estimate,
|
||||
)
|
||||
|
||||
prompt_ids = prompt_utils.build_indextts2_prefill_prompt_ids(
|
||||
"/model",
|
||||
"hello",
|
||||
model_type="indextts2_5",
|
||||
lang="en",
|
||||
tokenizer_file="custom-tokenizer.tiktoken",
|
||||
placeholder_token_id=17,
|
||||
)
|
||||
|
||||
assert prompt_ids == [17] * 10
|
||||
assert captured["tokenizer_file"] == "custom-tokenizer.tiktoken"
|
||||
@@ -0,0 +1,14 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
|
||||
|
||||
def test_indextts25_default_recipe_enables_validated_cfm_batching():
|
||||
recipe_path = Path(__file__).parents[4] / "vllm_omni" / "deploy" / "indextts2_5.yaml"
|
||||
recipe = yaml.safe_load(recipe_path.read_text(encoding="utf-8"))
|
||||
stage1 = next(stage for stage in recipe["stages"] if stage["stage_id"] == 1)
|
||||
|
||||
assert stage1["hf_overrides"]["s2mel_cfm_batch_size"] == 4
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,53 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.model_executor.stage_input_processors.indextts2 import (
|
||||
STOP_MEL_TOKEN,
|
||||
talker2s2mel_full_payload,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_no_latent_full_payload_uses_sampled_output_ids_and_strips_stop():
|
||||
request = SimpleNamespace(
|
||||
request_id="req-1",
|
||||
output_token_ids=[17, 23, STOP_MEL_TOKEN],
|
||||
sampling_params=SimpleNamespace(seed=42),
|
||||
)
|
||||
payload = {
|
||||
"meta": {
|
||||
"use_gpt_latent": False,
|
||||
"S_ref": torch.ones(1, 4),
|
||||
"ref_mel": torch.ones(1, 80, 3),
|
||||
"style": torch.ones(1, 192),
|
||||
}
|
||||
}
|
||||
|
||||
result = talker2s2mel_full_payload(None, payload, request)
|
||||
|
||||
assert result is not None
|
||||
assert result["mel_codes"].tolist() == [[17, 23]]
|
||||
assert result["code_lens"].tolist() == [2]
|
||||
assert result["use_gpt_latent"] is False
|
||||
assert result["seed"] == 42
|
||||
|
||||
|
||||
def test_no_latent_full_payload_rejects_legacy_payload_codes():
|
||||
request = SimpleNamespace(
|
||||
request_id="legacy",
|
||||
output_token_ids=[],
|
||||
sampling_params=SimpleNamespace(seed=None),
|
||||
)
|
||||
payload = {
|
||||
"codes": {"mel": torch.tensor([31, 37, STOP_MEL_TOKEN])},
|
||||
"meta": {"use_gpt_latent": False},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="no completed output_token_ids"):
|
||||
talker2s2mel_full_payload(None, payload, request)
|
||||
@@ -0,0 +1,476 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.config import CompilationMode
|
||||
|
||||
from vllm_omni.model_executor.models.indextts2.indextts2_talker import (
|
||||
IndexTTS2TalkerForConditionalGeneration,
|
||||
_normalize_indextts_checkpoint_key,
|
||||
_require_exact_indextts_prefill_span,
|
||||
resolve_indextts_conditioning_policy,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_v25_conditioning_policy():
|
||||
policy = resolve_indextts_conditioning_policy("indextts2_5")
|
||||
|
||||
assert policy.prefix_tokens == 3
|
||||
assert policy.uses_conformer_perceiver is False
|
||||
assert policy.uses_campplus_projection is True
|
||||
assert policy.uses_language_embedding is True
|
||||
|
||||
|
||||
def test_v2_conditioning_policy_is_unchanged():
|
||||
policy = resolve_indextts_conditioning_policy("indextts2")
|
||||
|
||||
assert policy.prefix_tokens == 34
|
||||
assert policy.uses_conformer_perceiver is True
|
||||
assert policy.uses_campplus_projection is False
|
||||
assert policy.uses_language_embedding is False
|
||||
|
||||
|
||||
def test_v25_text_tokens_wrap_language_prefixed_ids_with_start_and_stop(
|
||||
monkeypatch,
|
||||
):
|
||||
import vllm_omni.model_executor.models.indextts2.text_processing_v2_5 as text_processing
|
||||
|
||||
monkeypatch.setattr(
|
||||
text_processing,
|
||||
"prepare_indextts25_text",
|
||||
lambda *args, **kwargs: ([58838, 42], 7),
|
||||
)
|
||||
talker = object.__new__(IndexTTS2TalkerForConditionalGeneration)
|
||||
talker.conditioning_policy = resolve_indextts_conditioning_policy("indextts2_5")
|
||||
talker.model_path = "/model"
|
||||
talker.config = SimpleNamespace(tokenizer_file="tokenizer.tiktoken")
|
||||
|
||||
token_ids, language_id = talker._tokenize_text(
|
||||
"hello",
|
||||
torch.device("cpu"),
|
||||
lang="ja",
|
||||
)
|
||||
|
||||
assert token_ids.tolist() == [[0, 58838, 42, 1]]
|
||||
assert language_id.tolist() == [7]
|
||||
|
||||
|
||||
def test_v25_text_tokens_filter_existing_start_and_stop_ids(monkeypatch):
|
||||
import vllm_omni.model_executor.models.indextts2.text_processing_v2_5 as text_processing
|
||||
|
||||
monkeypatch.setattr(
|
||||
text_processing,
|
||||
"prepare_indextts25_text",
|
||||
lambda *args, **kwargs: ([0, 58838, 1, 42, 0], 7),
|
||||
)
|
||||
talker = object.__new__(IndexTTS2TalkerForConditionalGeneration)
|
||||
talker.conditioning_policy = resolve_indextts_conditioning_policy("indextts2_5")
|
||||
talker.model_path = "/model"
|
||||
talker.config = SimpleNamespace(tokenizer_file="tokenizer.tiktoken")
|
||||
|
||||
token_ids, language_id = talker._tokenize_text(
|
||||
"hello",
|
||||
torch.device("cpu"),
|
||||
lang="en",
|
||||
)
|
||||
|
||||
assert token_ids.tolist() == [[0, 58838, 42, 1]]
|
||||
assert language_id.tolist() == [7]
|
||||
|
||||
|
||||
def test_native_fsdp_checkpoint_keys_are_normalized_at_every_nesting_level():
|
||||
assert _normalize_indextts_checkpoint_key("_fsdp_wrapped_module.spk_emb_proj.weight") == "spk_emb_proj.weight"
|
||||
assert (
|
||||
_normalize_indextts_checkpoint_key("_fsdp_wrapped_module.mel_embedding._fsdp_wrapped_module.weight")
|
||||
== "mel_embedding.weight"
|
||||
)
|
||||
|
||||
|
||||
def test_prefill_span_mismatch_fails_instead_of_padding_or_truncating():
|
||||
for actual_span_len in (9, 11):
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="model_type=indextts2_5.*lang=es.*text_token_count=6",
|
||||
):
|
||||
_require_exact_indextts_prefill_span(
|
||||
actual_span_len=actual_span_len,
|
||||
scheduled_span_len=10,
|
||||
model_type="indextts2_5",
|
||||
lang="es",
|
||||
text_token_count=6,
|
||||
)
|
||||
|
||||
|
||||
def test_matching_prefill_span_is_accepted():
|
||||
_require_exact_indextts_prefill_span(
|
||||
actual_span_len=10,
|
||||
scheduled_span_len=10,
|
||||
model_type="indextts2_5",
|
||||
lang="en",
|
||||
text_token_count=6,
|
||||
)
|
||||
|
||||
|
||||
def _make_minimal_talker(
|
||||
monkeypatch,
|
||||
*,
|
||||
model_type,
|
||||
use_gpt_latent,
|
||||
):
|
||||
import vllm_omni.model_executor.models.indextts2.indextts2_talker as talker_module
|
||||
|
||||
seen = {}
|
||||
|
||||
def fake_parallel_lm_head(vocab_size, hidden_size, *, bias=False, **_kwargs):
|
||||
seen["bias"] = bias
|
||||
return torch.nn.Linear(hidden_size, vocab_size, bias=bias)
|
||||
|
||||
monkeypatch.setattr(talker_module, "ParallelLMHead", fake_parallel_lm_head)
|
||||
monkeypatch.setattr(
|
||||
talker_module,
|
||||
"make_layers",
|
||||
lambda *_args, **_kwargs: (
|
||||
0,
|
||||
1,
|
||||
torch.nn.ModuleList([torch.nn.Linear(8, 8)]),
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
talker_module,
|
||||
"get_pp_group",
|
||||
lambda: SimpleNamespace(is_last_rank=lambda: True),
|
||||
)
|
||||
monkeypatch.setattr(talker_module, "get_speaker_cache", object)
|
||||
monkeypatch.setattr(
|
||||
talker_module,
|
||||
"LogitsProcessor",
|
||||
lambda *_args, **_kwargs: object(),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
talker_module,
|
||||
"ConformerEncoder",
|
||||
lambda **_kwargs: torch.nn.Linear(1, 1),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
talker_module,
|
||||
"PerceiverResampler",
|
||||
lambda *_args, **_kwargs: torch.nn.Linear(1, 1),
|
||||
)
|
||||
|
||||
gpt = {
|
||||
"model_dim": 8,
|
||||
"layers": 0,
|
||||
"heads": 1,
|
||||
"max_mel_tokens": 4,
|
||||
"max_text_tokens": 4,
|
||||
"number_mel_codes": 16,
|
||||
"start_mel_token": 14,
|
||||
"stop_mel_token": 15,
|
||||
"number_text_tokens": 32,
|
||||
"condition_num_latent": 1,
|
||||
"condition_module": {},
|
||||
"emo_condition_module": {},
|
||||
}
|
||||
config = SimpleNamespace(
|
||||
gpt=gpt,
|
||||
model_type=model_type,
|
||||
use_gpt_latent=use_gpt_latent,
|
||||
gpt_checkpoint="gpt.pth",
|
||||
)
|
||||
vllm_config = SimpleNamespace(
|
||||
model_config=SimpleNamespace(model="/model", hf_config=config),
|
||||
cache_config=None,
|
||||
quant_config=None,
|
||||
compilation_config=SimpleNamespace(mode=CompilationMode.NONE),
|
||||
)
|
||||
|
||||
return (
|
||||
IndexTTS2TalkerForConditionalGeneration(vllm_config=vllm_config),
|
||||
seen,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model_type", "use_gpt_latent"),
|
||||
[
|
||||
("indextts2", True),
|
||||
("indextts2_5", False),
|
||||
],
|
||||
)
|
||||
def test_mel_head_preserves_official_checkpoint_bias_and_hidden_payload_policy(
|
||||
monkeypatch,
|
||||
model_type,
|
||||
use_gpt_latent,
|
||||
):
|
||||
talker, seen = _make_minimal_talker(
|
||||
monkeypatch,
|
||||
model_type=model_type,
|
||||
use_gpt_latent=use_gpt_latent,
|
||||
)
|
||||
|
||||
assert seen["bias"] is True
|
||||
assert talker.mel_head.bias is not None
|
||||
assert talker.omni_pooler_payload_include_hidden is use_gpt_latent
|
||||
assert talker.omni_payload_at_request_end is (model_type == "indextts2_5")
|
||||
|
||||
|
||||
def test_v2_load_weights_loads_exact_mel_head_bias_parameter(monkeypatch):
|
||||
import vllm_omni.model_executor.models.indextts2.indextts2_talker as talker_module
|
||||
|
||||
talker, _ = _make_minimal_talker(
|
||||
monkeypatch,
|
||||
model_type="indextts2",
|
||||
use_gpt_latent=True,
|
||||
)
|
||||
checkpoint = {
|
||||
name: parameter.detach().clone() for name, parameter in talker.named_parameters(remove_duplicate=False)
|
||||
}
|
||||
expected_bias = torch.arange(
|
||||
talker.number_mel_codes,
|
||||
dtype=talker.mel_head.bias.dtype,
|
||||
)
|
||||
checkpoint["mel_head.bias"] = expected_bias
|
||||
with torch.no_grad():
|
||||
talker.mel_head.bias.zero_()
|
||||
|
||||
monkeypatch.setattr(
|
||||
talker_module,
|
||||
"resolve_model_file",
|
||||
lambda *_args, **_kwargs: "/model/gpt.pth",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
talker_module.torch,
|
||||
"load",
|
||||
lambda *_args, **_kwargs: checkpoint,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
talker_module,
|
||||
"is_pp_missing_parameter",
|
||||
lambda *_args, **_kwargs: False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
talker,
|
||||
"_warmup_external_models",
|
||||
lambda _device: None,
|
||||
)
|
||||
|
||||
loaded = talker.load_weights([])
|
||||
|
||||
assert "mel_head.bias" in loaded
|
||||
torch.testing.assert_close(talker.mel_head.bias, expected_bias)
|
||||
|
||||
|
||||
def test_compute_logits_passes_mel_head_bias_to_logits_processor():
|
||||
talker = object.__new__(IndexTTS2TalkerForConditionalGeneration)
|
||||
torch.nn.Module.__init__(talker)
|
||||
talker.mel_head = torch.nn.Linear(3, 4, bias=True)
|
||||
hidden_states = torch.zeros(1, 3)
|
||||
seen = {}
|
||||
|
||||
def fake_logits_processor(head, hidden, **kwargs):
|
||||
seen["head"] = head
|
||||
seen["hidden"] = hidden
|
||||
seen.update(kwargs)
|
||||
return hidden
|
||||
|
||||
talker.logits_processor = fake_logits_processor
|
||||
|
||||
result = talker.compute_logits(hidden_states)
|
||||
|
||||
assert result is hidden_states
|
||||
assert seen["head"] is talker.mel_head
|
||||
assert seen["hidden"] is hidden_states
|
||||
assert seen["embedding_bias"] is talker.mel_head.bias
|
||||
|
||||
|
||||
def _make_cached_v25_prefill_talker(monkeypatch):
|
||||
talker, _ = _make_minimal_talker(
|
||||
monkeypatch,
|
||||
model_type="indextts2_5",
|
||||
use_gpt_latent=False,
|
||||
)
|
||||
talker._speaker_cache = SimpleNamespace(
|
||||
make_cache_key=lambda *_args: "cached-speaker",
|
||||
get=lambda _key: {
|
||||
"S_ref": torch.ones(1, 3, 8),
|
||||
"ref_mel": torch.ones(1, 80, 3),
|
||||
"style": torch.ones(1, 192),
|
||||
"spk_cond_emb": torch.ones(1, 3, 8),
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
talker,
|
||||
"_compute_emotion_vector",
|
||||
lambda **_kwargs: torch.zeros(1, talker.model_dim),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
talker,
|
||||
"_tokenize_text",
|
||||
lambda *_args, **_kwargs: (torch.tensor([[1]]), None),
|
||||
)
|
||||
return talker
|
||||
|
||||
|
||||
def test_prefill_preserves_duration_factor_in_metadata(monkeypatch):
|
||||
talker = _make_cached_v25_prefill_talker(monkeypatch)
|
||||
|
||||
_, _, info_update = talker._preprocess_prefill(
|
||||
torch.zeros(5, dtype=torch.long),
|
||||
None,
|
||||
{
|
||||
"text": ["hello"],
|
||||
"voice": ["/speaker.wav"],
|
||||
"voice_name": ["cached"],
|
||||
"duration_factor": [0.5],
|
||||
},
|
||||
span_len=5,
|
||||
)
|
||||
|
||||
assert info_update["meta"]["duration_factor"] == 0.5
|
||||
|
||||
|
||||
@pytest.mark.parametrize("duration_factor", [0.49, 2.01])
|
||||
def test_prefill_rejects_duration_factor_outside_supported_range(monkeypatch, duration_factor):
|
||||
talker = _make_cached_v25_prefill_talker(monkeypatch)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="^IndexTTS 2.5 duration_factor must be between 0.5 and 2.0$",
|
||||
):
|
||||
talker._preprocess_prefill(
|
||||
torch.zeros(5, dtype=torch.long),
|
||||
None,
|
||||
{
|
||||
"text": ["hello"],
|
||||
"voice": ["/speaker.wav"],
|
||||
"voice_name": ["cached"],
|
||||
"duration_factor": [duration_factor],
|
||||
},
|
||||
span_len=5,
|
||||
)
|
||||
|
||||
|
||||
def test_no_latent_sampled_token_payload_emits_only_first_step_metadata():
|
||||
talker = object.__new__(IndexTTS2TalkerForConditionalGeneration)
|
||||
talker.use_gpt_latent = False
|
||||
hidden = torch.zeros(1, 8)
|
||||
info = {
|
||||
"codes": {"mel": torch.tensor([17])},
|
||||
"meta": {
|
||||
"mel_code_count": 1,
|
||||
"S_ref": torch.ones(1, 4),
|
||||
"ref_mel": torch.ones(1, 80, 3),
|
||||
"style": torch.ones(1, 192),
|
||||
"duration_factor": 0.5,
|
||||
},
|
||||
}
|
||||
|
||||
first_decode_output = talker.make_omni_output(
|
||||
hidden,
|
||||
model_intermediate_buffer=[info],
|
||||
)
|
||||
|
||||
assert "codes" not in first_decode_output.multimodal_outputs
|
||||
assert "hidden_states" not in first_decode_output.multimodal_outputs
|
||||
assert first_decode_output.multimodal_outputs["meta"][0]["S_ref"] is info["meta"]["S_ref"]
|
||||
assert first_decode_output.multimodal_outputs["meta"][0]["duration_factor"] == 0.5
|
||||
|
||||
later = talker.make_omni_output(
|
||||
hidden,
|
||||
model_intermediate_buffer=[
|
||||
{
|
||||
"codes": {"mel": torch.tensor([23])},
|
||||
"meta": {"mel_code_count": 2},
|
||||
}
|
||||
],
|
||||
)
|
||||
assert later.multimodal_outputs == {"meta": [{"use_gpt_latent": False}]}
|
||||
|
||||
|
||||
def test_first_decode_payload_defaults_missing_duration_factor_to_one():
|
||||
talker = object.__new__(IndexTTS2TalkerForConditionalGeneration)
|
||||
talker.use_gpt_latent = False
|
||||
hidden = torch.zeros(1, 8)
|
||||
info = {
|
||||
"codes": {"mel": torch.tensor([17])},
|
||||
"meta": {
|
||||
"mel_code_count": 1,
|
||||
"S_ref": torch.ones(1, 4),
|
||||
"ref_mel": torch.ones(1, 80, 3),
|
||||
"style": torch.ones(1, 192),
|
||||
},
|
||||
}
|
||||
|
||||
output = talker.make_omni_output(
|
||||
hidden,
|
||||
model_intermediate_buffer=[info],
|
||||
)
|
||||
|
||||
assert output.multimodal_outputs["meta"][0]["duration_factor"] == 1.0
|
||||
|
||||
|
||||
def test_latent_sampled_token_payload_keeps_aligned_gpu_deltas():
|
||||
talker = object.__new__(IndexTTS2TalkerForConditionalGeneration)
|
||||
talker.use_gpt_latent = True
|
||||
hidden = torch.zeros(1, 8)
|
||||
info = {
|
||||
"codes": {"mel": torch.tensor([17])},
|
||||
"meta": {
|
||||
"mel_code_count": 1,
|
||||
"latent_acc": torch.arange(8, dtype=torch.float32).reshape(1, 8),
|
||||
"duration_factor": 1.0,
|
||||
},
|
||||
}
|
||||
|
||||
output = talker.make_omni_output(
|
||||
hidden,
|
||||
model_intermediate_buffer=[info],
|
||||
)
|
||||
|
||||
assert output.multimodal_outputs["codes"]["mel"][0].tolist() == [[17]]
|
||||
assert output.multimodal_outputs["hidden_states"]["latent"][0].tolist() == [list(range(8))]
|
||||
|
||||
|
||||
def test_v25_default_emotion_reuses_speaker_w2v_features(monkeypatch):
|
||||
talker = object.__new__(IndexTTS2TalkerForConditionalGeneration)
|
||||
speaker_features = torch.zeros(1, 4, 1024)
|
||||
expected = torch.ones(1, 1280)
|
||||
seen = {}
|
||||
|
||||
def fake_merge(spk_cond_emb, emo_cond_emb, alpha):
|
||||
seen["spk"] = spk_cond_emb
|
||||
seen["emo"] = emo_cond_emb
|
||||
seen["alpha"] = alpha
|
||||
return expected
|
||||
|
||||
monkeypatch.setattr(talker, "_merge_emovec", fake_merge)
|
||||
|
||||
result = talker._compute_emotion_vector(
|
||||
wav_16k=None,
|
||||
speaker_audio_path="/speaker.wav",
|
||||
emo_audio_path=None,
|
||||
emo_voice_name=None,
|
||||
main_text="hello",
|
||||
use_emo_text=False,
|
||||
emo_text=None,
|
||||
emo_vector=None,
|
||||
emo_alpha=0.25,
|
||||
use_random=False,
|
||||
style=torch.zeros(1, 192),
|
||||
spk_cond_emb=speaker_features,
|
||||
w2v_model=None,
|
||||
w2v_proc=None,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
|
||||
assert result is expected
|
||||
assert seen == {
|
||||
"spk": speaker_features,
|
||||
"emo": speaker_features,
|
||||
"alpha": 1.0,
|
||||
}
|
||||
@@ -0,0 +1,244 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import builtins
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.model_executor.models.indextts2 import text_processing_v2_5
|
||||
from vllm_omni.model_executor.models.indextts2.text_processing_v2_5 import (
|
||||
apply_pronunciation_annotations,
|
||||
clean_indextts25_text,
|
||||
prepare_indextts25_text,
|
||||
)
|
||||
from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import (
|
||||
INDEXTTS25_SPECIAL_TOKENS,
|
||||
INDEXTTS25_VOCAB_SIZE,
|
||||
LANGUAGE_DICT,
|
||||
lang_to_token,
|
||||
normalize_language_code,
|
||||
resolve_indextts25_tokenizer_file,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_special_token_layout_matches_official_checkpoint_vocab():
|
||||
assert len(INDEXTTS25_SPECIAL_TOKENS) == 1673
|
||||
assert len(set(INDEXTTS25_SPECIAL_TOKENS)) == 1673
|
||||
assert 58836 + len(INDEXTTS25_SPECIAL_TOKENS) == INDEXTTS25_VOCAB_SIZE
|
||||
assert INDEXTTS25_VOCAB_SIZE == 60509
|
||||
|
||||
|
||||
def test_pronunciation_annotations_select_english_and_chinese_markers():
|
||||
assert apply_pronunciation_annotations("<going|G OW1 . IH0 NG>") == (
|
||||
"<|SPECIAL_TOKEN_1|>G OW1 . IH0 NG<|SPECIAL_TOKEN_1|>"
|
||||
)
|
||||
assert apply_pronunciation_annotations("<行|xing2>") == ("<|SPECIAL_TOKEN_2|>XING2<|SPECIAL_TOKEN_2|>")
|
||||
|
||||
|
||||
def test_prepare_text_applies_language_prefix_case_and_annotation(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
def fake_encode(text, *, model_dir, tokenizer_file):
|
||||
captured["text"] = text
|
||||
captured["model_dir"] = model_dir
|
||||
captured["tokenizer_file"] = tokenizer_file
|
||||
return [7, 8, 9]
|
||||
|
||||
monkeypatch.setattr(text_processing_v2_5, "encode_indextts25_text", fake_encode)
|
||||
|
||||
token_ids, lang_id = prepare_indextts25_text(
|
||||
"HELLO <行|xing2>",
|
||||
lang="zh",
|
||||
model_dir="/model",
|
||||
text_normalization=False,
|
||||
)
|
||||
|
||||
assert token_ids == [7, 8, 9]
|
||||
assert lang_id == LANGUAGE_DICT["zh"]
|
||||
assert captured == {
|
||||
"text": "<|zh|> hello <|SPECIAL_TOKEN_2|>XING2<|SPECIAL_TOKEN_2|>",
|
||||
"model_dir": "/model",
|
||||
"tokenizer_file": "multilingual_zh_ja_yue_char_del.tiktoken",
|
||||
}
|
||||
|
||||
|
||||
def test_mandarin_alias_is_an_intentional_vllm_omni_convenience():
|
||||
assert normalize_language_code("Mandarin") == "zh"
|
||||
assert lang_to_token("mandarin") == LANGUAGE_DICT["zh"]
|
||||
|
||||
|
||||
def test_invalid_language_does_not_fall_back_to_common():
|
||||
with pytest.raises(ValueError, match="Unsupported IndexTTS 2.5 language"):
|
||||
lang_to_token("xx-invalid")
|
||||
|
||||
|
||||
def test_zhen_uses_mixed_normalization():
|
||||
calls = []
|
||||
|
||||
def fake_mixed_normalizer(text):
|
||||
calls.append(text)
|
||||
return "MIXED"
|
||||
|
||||
original = text_processing_v2_5._normalize_zh_or_en
|
||||
text_processing_v2_5._normalize_zh_or_en = fake_mixed_normalizer
|
||||
try:
|
||||
result = text_processing_v2_5._normalize_with_official_backend(
|
||||
"中文 and English",
|
||||
"zhen",
|
||||
)
|
||||
finally:
|
||||
text_processing_v2_5._normalize_zh_or_en = original
|
||||
|
||||
assert calls == ["中文 and English"]
|
||||
assert result == "MIXED"
|
||||
|
||||
|
||||
def test_spanish_normalization_does_not_hide_broken_dependencies(monkeypatch):
|
||||
original_import = builtins.__import__
|
||||
|
||||
def controlled_import(name, *args, **kwargs):
|
||||
if name == "nemo_text_processing.text_normalization.normalize":
|
||||
raise RuntimeError("broken NeMo grammar")
|
||||
return original_import(name, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(builtins, "__import__", controlled_import)
|
||||
|
||||
with pytest.raises(RuntimeError, match="broken NeMo grammar"):
|
||||
text_processing_v2_5._normalize_with_official_backend("hola", "es")
|
||||
|
||||
|
||||
def test_zhen_keeps_literal_prefix_and_uses_common_embedding(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
def fake_encode(text, *, model_dir, tokenizer_file):
|
||||
captured["text"] = text
|
||||
return [7, 8, 9]
|
||||
|
||||
monkeypatch.setattr(text_processing_v2_5, "encode_indextts25_text", fake_encode)
|
||||
|
||||
token_ids, lang_id = prepare_indextts25_text(
|
||||
"HELLO 中文",
|
||||
lang="zhen",
|
||||
model_dir="/model",
|
||||
text_normalization=False,
|
||||
)
|
||||
|
||||
assert token_ids == [7, 8, 9]
|
||||
assert captured["text"] == "<|zhen|> hello 中文"
|
||||
assert lang_id == LANGUAGE_DICT["common"]
|
||||
assert "<|zhen|>" not in INDEXTTS25_SPECIAL_TOKENS
|
||||
assert "<|common|>" not in INDEXTTS25_SPECIAL_TOKENS
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
(",,,", "…"),
|
||||
("测试,,,结束", "测试…结束"),
|
||||
],
|
||||
)
|
||||
def test_repeated_full_width_commas_use_longest_first_cleanup(text, expected):
|
||||
assert clean_indextts25_text(text) == expected
|
||||
|
||||
|
||||
def test_default_text_normalizer_is_shared_by_prompt_and_talker(monkeypatch):
|
||||
calls = []
|
||||
|
||||
def fake_normalizer(text, lang):
|
||||
calls.append((text, lang))
|
||||
return "twenty five percent"
|
||||
|
||||
monkeypatch.setattr(
|
||||
text_processing_v2_5,
|
||||
"_normalize_with_official_backend",
|
||||
fake_normalizer,
|
||||
)
|
||||
|
||||
result = text_processing_v2_5.normalize_indextts25_text(
|
||||
"25%",
|
||||
lang="en",
|
||||
text_normalization=True,
|
||||
)
|
||||
|
||||
assert calls == [("25%", "en")]
|
||||
assert result == "twenty five percent"
|
||||
|
||||
|
||||
def test_mixed_text_uses_official_content_routing_and_protection(monkeypatch):
|
||||
calls = []
|
||||
|
||||
class FakeNormalizer:
|
||||
def __init__(self, lang):
|
||||
self.lang = lang
|
||||
|
||||
def normalize(self, text):
|
||||
calls.append((self.lang, text))
|
||||
return text.replace("<H>", " <H> ").replace("5", "five")
|
||||
|
||||
monkeypatch.setattr(
|
||||
text_processing_v2_5,
|
||||
"_load_wetext_normalizer",
|
||||
lambda lang: FakeNormalizer(lang),
|
||||
)
|
||||
|
||||
result = text_processing_v2_5._normalize_with_official_backend(
|
||||
"GPT-5-nano 中文 xuan4",
|
||||
"en",
|
||||
)
|
||||
|
||||
assert calls[0][0] == "zh"
|
||||
assert result == "GPT-five-nano 中文 XVAN4"
|
||||
|
||||
|
||||
def test_tokenizer_resolves_native_checkpoints_subdirectory(tmp_path):
|
||||
tokenizer = tmp_path / "checkpoints" / "multilingual_zh_ja_yue_char_del.tiktoken"
|
||||
tokenizer.parent.mkdir()
|
||||
tokenizer.touch()
|
||||
|
||||
assert resolve_indextts25_tokenizer_file(str(tmp_path)) == str(tokenizer)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("lang", "text", "expected"),
|
||||
[
|
||||
(
|
||||
"es",
|
||||
"En la actualidad se continúa utilizando una corona abierta.",
|
||||
"<|es|> EN LA ACTUALIDAD SE CONTINÚA UTILIZANDO UNA CORONA ABIERTA.",
|
||||
),
|
||||
(
|
||||
"ar",
|
||||
"فكلُّها خياراتٌ اتَّخذتَها بنفسِك.",
|
||||
"<|ar|> فكلُّها خياراتٌ اتَّخذتَها بنفسِك.",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_official_spanish_and_arabic_cases_preserve_expected_text_contract(
|
||||
monkeypatch,
|
||||
lang,
|
||||
text,
|
||||
expected,
|
||||
):
|
||||
captured = {}
|
||||
|
||||
def fake_encode(value, *, model_dir, tokenizer_file):
|
||||
captured["text"] = value
|
||||
return [1]
|
||||
|
||||
monkeypatch.setattr(
|
||||
text_processing_v2_5,
|
||||
"encode_indextts25_text",
|
||||
fake_encode,
|
||||
)
|
||||
|
||||
_, lang_id = prepare_indextts25_text(
|
||||
text,
|
||||
lang=lang,
|
||||
model_dir="/model",
|
||||
text_normalization=False,
|
||||
)
|
||||
|
||||
assert captured["text"] == expected
|
||||
assert lang_id == LANGUAGE_DICT[lang]
|
||||
@@ -0,0 +1,196 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.model_executor.models.indextts2.indextts2_s2mel_decoder import (
|
||||
IndexTTS2S2MelDecoder,
|
||||
)
|
||||
from vllm_omni.model_executor.models.indextts2.s2mel.modules.wavenet import WN
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _make_wavenet() -> WN:
|
||||
torch.manual_seed(0)
|
||||
return WN(
|
||||
hidden_channels=2,
|
||||
kernel_size=5,
|
||||
dilation_rate=1,
|
||||
n_layers=8,
|
||||
gin_channels=2,
|
||||
p_dropout=0.0,
|
||||
causal=False,
|
||||
).eval()
|
||||
|
||||
|
||||
def _sequence_mask(lengths: torch.Tensor, max_length: int) -> torch.Tensor:
|
||||
positions = torch.arange(max_length, device=lengths.device)
|
||||
return (positions.unsqueeze(0) < lengths.unsqueeze(1)).unsqueeze(1)
|
||||
|
||||
|
||||
def _singleton_outputs(
|
||||
wavenet: WN,
|
||||
inputs: torch.Tensor,
|
||||
lengths: list[int],
|
||||
conditioning: torch.Tensor,
|
||||
) -> list[torch.Tensor]:
|
||||
# This is the independent oracle: every request ends at the physical tensor
|
||||
# boundary and therefore uses SConv1d's original reflect-padding behavior.
|
||||
wavenet._assume_full_mask = True
|
||||
return [
|
||||
wavenet(
|
||||
inputs[row : row + 1, :, :length],
|
||||
torch.ones(1, 1, length, dtype=torch.bool),
|
||||
g=conditioning[row : row + 1],
|
||||
)
|
||||
for row, length in enumerate(lengths)
|
||||
]
|
||||
|
||||
|
||||
def test_ragged_wavenet_matches_original_singleton_reflect_boundaries():
|
||||
wavenet = _make_wavenet()
|
||||
lengths = [6, 5, 4, 3]
|
||||
inputs = torch.randn(len(lengths), 2, max(lengths))
|
||||
conditioning = torch.randn(len(lengths), 2, 1)
|
||||
expected = _singleton_outputs(wavenet, inputs, lengths, conditioning)
|
||||
|
||||
mask = _sequence_mask(torch.tensor(lengths), inputs.shape[-1])
|
||||
wavenet._assume_full_mask = False
|
||||
# Keep deliberately non-zero invalid tails. The first WaveNet layer sees
|
||||
# transformer output before the recurrent mask has cleared those values.
|
||||
actual = wavenet(inputs, mask, g=conditioning)
|
||||
|
||||
for row, (length, singleton) in enumerate(zip(lengths, expected)):
|
||||
torch.testing.assert_close(
|
||||
actual[row : row + 1, :, :length],
|
||||
singleton,
|
||||
rtol=1e-5,
|
||||
atol=1e-6,
|
||||
)
|
||||
|
||||
|
||||
def test_wavenet_exposes_minimum_safe_ragged_length():
|
||||
# kernel=5, dilation=1 has two right-reflect ghost positions, so a
|
||||
# logical row must contain at least three frames before it can use the
|
||||
# vectorized ragged path. Shorter rows stay on original singleton SConv1d.
|
||||
assert _make_wavenet().ragged_min_length == 3
|
||||
|
||||
|
||||
def test_cfm_grouping_isolates_rows_too_short_for_logical_reflect():
|
||||
groups = IndexTTS2S2MelDecoder._cfm_group_indices(
|
||||
sorted_indices=[0, 1, 2, 3, 4],
|
||||
total_lengths=[0, 1, 2, 3, 6],
|
||||
max_group_size=2,
|
||||
min_ragged_length=3,
|
||||
)
|
||||
|
||||
assert groups == [[0], [1], [2], [3, 4]]
|
||||
|
||||
|
||||
def test_ragged_wavenet_preserves_cfg_cond_uncond_row_order():
|
||||
wavenet = _make_wavenet()
|
||||
base_lengths = [3, 6]
|
||||
conditional = torch.randn(2, 2, 6)
|
||||
unconditional = torch.randn(2, 2, 6)
|
||||
cfg_inputs = torch.cat([conditional, unconditional], dim=0)
|
||||
cfg_lengths = base_lengths + base_lengths
|
||||
cfg_conditioning = torch.randn(4, 2, 1)
|
||||
expected = _singleton_outputs(
|
||||
wavenet,
|
||||
cfg_inputs,
|
||||
cfg_lengths,
|
||||
cfg_conditioning,
|
||||
)
|
||||
|
||||
mask = _sequence_mask(torch.tensor(cfg_lengths), cfg_inputs.shape[-1])
|
||||
wavenet._assume_full_mask = False
|
||||
actual = wavenet(cfg_inputs, mask, g=cfg_conditioning)
|
||||
|
||||
for row, (length, singleton) in enumerate(zip(cfg_lengths, expected)):
|
||||
torch.testing.assert_close(
|
||||
actual[row : row + 1, :, :length],
|
||||
singleton,
|
||||
rtol=1e-5,
|
||||
atol=1e-6,
|
||||
)
|
||||
|
||||
|
||||
def test_full_mask_wavenet_path_is_bit_exact_with_original_path():
|
||||
wavenet = _make_wavenet()
|
||||
inputs = torch.randn(3, 2, 6)
|
||||
mask = torch.ones(3, 1, 6, dtype=torch.bool)
|
||||
conditioning = torch.randn(3, 2, 1)
|
||||
|
||||
if hasattr(wavenet, "_assume_full_mask"):
|
||||
del wavenet._assume_full_mask
|
||||
expected = wavenet(inputs, mask, g=conditioning)
|
||||
|
||||
wavenet._assume_full_mask = True
|
||||
actual = wavenet(inputs, mask, g=conditioning)
|
||||
|
||||
assert torch.equal(actual, expected)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(torch, "compile"), reason="torch.compile unavailable")
|
||||
def test_ragged_wavenet_is_one_dynamic_fullgraph():
|
||||
wavenet = _make_wavenet()
|
||||
lengths = [3, 6]
|
||||
inputs = torch.randn(2, 2, 6)
|
||||
conditioning = torch.randn(2, 2, 1)
|
||||
mask = _sequence_mask(torch.tensor(lengths), inputs.shape[-1])
|
||||
expected = _singleton_outputs(wavenet, inputs, lengths, conditioning)
|
||||
wavenet._assume_full_mask = False
|
||||
|
||||
compiled = torch.compile(
|
||||
wavenet,
|
||||
backend="eager",
|
||||
fullgraph=True,
|
||||
dynamic=True,
|
||||
)
|
||||
actual = compiled(inputs, mask, g=conditioning)
|
||||
|
||||
for row, (length, singleton) in enumerate(zip(lengths, expected)):
|
||||
torch.testing.assert_close(
|
||||
actual[row : row + 1, :, :length],
|
||||
singleton,
|
||||
rtol=1e-5,
|
||||
atol=1e-6,
|
||||
)
|
||||
|
||||
|
||||
def test_decoder_full_mask_switch_also_controls_wavenet_boundary_path():
|
||||
decoder = object.__new__(IndexTTS2S2MelDecoder)
|
||||
attention_layers = [SimpleNamespace(attention=SimpleNamespace()) for _ in range(2)]
|
||||
estimator = SimpleNamespace(
|
||||
transformer=SimpleNamespace(layers=attention_layers),
|
||||
wavenet=SimpleNamespace(),
|
||||
)
|
||||
|
||||
decoder._set_dit_full_mask_fast_path(estimator, enabled=False)
|
||||
|
||||
assert all(layer.attention._assume_full_mask is False for layer in attention_layers)
|
||||
assert estimator.wavenet._assume_full_mask is False
|
||||
|
||||
decoder._set_dit_full_mask_fast_path(estimator, enabled=True)
|
||||
|
||||
assert all(layer.attention._assume_full_mask is True for layer in attention_layers)
|
||||
assert estimator.wavenet._assume_full_mask is True
|
||||
|
||||
|
||||
def test_decoder_full_mask_switch_unwraps_compiled_estimator():
|
||||
decoder = object.__new__(IndexTTS2S2MelDecoder)
|
||||
attention_layers = [SimpleNamespace(attention=SimpleNamespace()) for _ in range(2)]
|
||||
eager_estimator = SimpleNamespace(
|
||||
transformer=SimpleNamespace(layers=attention_layers),
|
||||
wavenet=SimpleNamespace(),
|
||||
)
|
||||
compiled_estimator = SimpleNamespace(_orig_mod=eager_estimator)
|
||||
|
||||
decoder._set_dit_full_mask_fast_path(compiled_estimator, enabled=False)
|
||||
|
||||
assert all(layer.attention._assume_full_mask is False for layer in attention_layers)
|
||||
assert eager_estimator.wavenet._assume_full_mask is False
|
||||
@@ -57,6 +57,7 @@ def _build_mock_modules(mocker: MockerFixture) -> dict[str, object]:
|
||||
platforms_mock = mocker.MagicMock()
|
||||
platforms_mock.current_omni_platform.supports_torch_inductor.return_value = False
|
||||
platforms_mock.current_omni_platform.is_npu.return_value = False
|
||||
platforms_mock.current_omni_platform.is_xpu.return_value = False
|
||||
|
||||
logger_mock = mocker.MagicMock()
|
||||
logger_mock.init_logger = lambda name: mocker.MagicMock()
|
||||
@@ -284,6 +285,28 @@ class TestCodePredictorDtypeAlignment:
|
||||
predictor._setup_compile()
|
||||
assert predictor._model_dtype == torch.float16
|
||||
|
||||
def test_setup_compile_uses_eager_on_xpu(self, mocker: MockerFixture, loaded_target_classes) -> None:
|
||||
"""XPU should avoid the torch.compile path used by CUDA."""
|
||||
_, _, code_predictor_wrapper, _, _ = loaded_target_classes
|
||||
cp_config, talker_config = _make_tiny_config(loaded_target_classes)
|
||||
vllm_config = _make_vllm_config(mocker, max_num_seqs=2)
|
||||
|
||||
predictor = code_predictor_wrapper(
|
||||
vllm_config=vllm_config,
|
||||
config=cp_config,
|
||||
talker_config=talker_config,
|
||||
)
|
||||
|
||||
common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
|
||||
mocker.patch.object(common_mod.current_omni_platform, "supports_torch_inductor", return_value=True)
|
||||
mocker.patch.object(common_mod.current_omni_platform, "is_xpu", return_value=True)
|
||||
compile_mock = mocker.patch.object(common_mod.torch, "compile")
|
||||
|
||||
predictor._setup_compile()
|
||||
|
||||
compile_mock.assert_not_called()
|
||||
assert predictor._compiled_model_fwd == predictor.model.forward
|
||||
|
||||
def test_forward_with_mismatched_input_dtype(self, mocker: MockerFixture, loaded_target_classes) -> None:
|
||||
"""forward() should not crash when inputs are float32 but model is float16."""
|
||||
_, _, code_predictor_wrapper, _, _ = loaded_target_classes
|
||||
|
||||
@@ -116,6 +116,7 @@ def test_strip_stop_token_all_stop():
|
||||
|
||||
def _conditioning_tensors():
|
||||
return {
|
||||
"use_gpt_latent": True,
|
||||
"S_ref": torch.randn(1, 3, LATENT_DIM),
|
||||
"ref_mel": torch.randn(1, 80, 5),
|
||||
"style": torch.randn(1, 4),
|
||||
@@ -143,6 +144,8 @@ def test_talker2s2mel_full_payload_flat_keys_builds_s2mel_contract():
|
||||
"meta.S_ref": cond["S_ref"],
|
||||
"meta.ref_mel": cond["ref_mel"],
|
||||
"meta.style": cond["style"],
|
||||
"meta.use_gpt_latent": True,
|
||||
"meta.duration_factor": 0.5,
|
||||
}
|
||||
|
||||
result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-flat"))
|
||||
@@ -156,6 +159,45 @@ def test_talker2s2mel_full_payload_flat_keys_builds_s2mel_contract():
|
||||
assert result["ref_mel"].device.type == "cpu"
|
||||
assert result["style"].device.type == "cpu"
|
||||
assert result["S_ref"].data_ptr() != cond["S_ref"].data_ptr()
|
||||
assert result["duration_factor"] == 0.5
|
||||
|
||||
|
||||
def test_talker2s2mel_full_payload_defaults_missing_duration_factor_to_one():
|
||||
cond = _conditioning_tensors()
|
||||
payload = {
|
||||
"codes.mel": torch.tensor([[10], [20]]),
|
||||
"hidden_states.latent": torch.randn(2, LATENT_DIM),
|
||||
"meta.S_ref": cond["S_ref"],
|
||||
"meta.ref_mel": cond["ref_mel"],
|
||||
"meta.style": cond["style"],
|
||||
"meta.use_gpt_latent": True,
|
||||
}
|
||||
|
||||
result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-default-duration"))
|
||||
|
||||
assert result is not None
|
||||
assert result["duration_factor"] == 1.0
|
||||
|
||||
|
||||
def test_talker2s2mel_full_payload_preserves_explicit_request_seed():
|
||||
cond = _conditioning_tensors()
|
||||
payload = {
|
||||
"codes.mel": torch.tensor([[10], [20]]),
|
||||
"hidden_states.latent": torch.randn(2, LATENT_DIM),
|
||||
"meta.S_ref": cond["S_ref"],
|
||||
"meta.ref_mel": cond["ref_mel"],
|
||||
"meta.style": cond["style"],
|
||||
"meta.use_gpt_latent": True,
|
||||
}
|
||||
request = SimpleNamespace(
|
||||
request_id="r-seeded",
|
||||
sampling_params=SimpleNamespace(seed=0),
|
||||
)
|
||||
|
||||
result = talker2s2mel_full_payload(None, payload, request)
|
||||
|
||||
assert result is not None
|
||||
assert result["seed"] == 0
|
||||
|
||||
|
||||
def test_talker2s2mel_full_payload_nested_fallback_input():
|
||||
@@ -206,7 +248,7 @@ def test_talker2s2mel_full_payload_trims_stop_token_and_matching_latent():
|
||||
assert result["code_lens"].tolist() == [2]
|
||||
|
||||
|
||||
def test_talker2s2mel_full_payload_crops_mel_latent_to_common_length():
|
||||
def test_talker2s2mel_full_payload_rejects_mel_latent_length_mismatch():
|
||||
cond = _conditioning_tensors()
|
||||
payload = {
|
||||
"codes.mel": torch.tensor([[1], [2], [3], [4]]),
|
||||
@@ -214,14 +256,44 @@ def test_talker2s2mel_full_payload_crops_mel_latent_to_common_length():
|
||||
"meta": cond,
|
||||
}
|
||||
|
||||
result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-crop"))
|
||||
with pytest.raises(ValueError, match="mel/latent length mismatch"):
|
||||
talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-crop"))
|
||||
|
||||
|
||||
def test_talker2s2mel_full_payload_v25_code_only_does_not_invent_latent():
|
||||
cond = _conditioning_tensors()
|
||||
payload = {
|
||||
"meta.use_gpt_latent": False,
|
||||
"meta.S_ref": cond["S_ref"],
|
||||
"meta.ref_mel": cond["ref_mel"],
|
||||
"meta.style": cond["style"],
|
||||
}
|
||||
|
||||
request = SimpleNamespace(
|
||||
request_id="r-v25",
|
||||
output_token_ids=[3, 4, STOP_MEL_TOKEN],
|
||||
)
|
||||
result = talker2s2mel_full_payload(None, payload, request)
|
||||
|
||||
assert result is not None
|
||||
assert result["mel_codes"].tolist() == [[1, 2]]
|
||||
assert result["latent"].shape == (1, 2, LATENT_DIM)
|
||||
assert result["mel_codes"].tolist() == [[3, 4]]
|
||||
assert result["code_lens"].tolist() == [2]
|
||||
assert result["use_gpt_latent"] is False
|
||||
assert "latent" not in result
|
||||
|
||||
|
||||
def test_talker2s2mel_full_payload_missing_required_fields_returns_none():
|
||||
assert talker2s2mel_full_payload(None, {"hidden_states.latent": torch.randn(1, LATENT_DIM)}, None) is None
|
||||
assert talker2s2mel_full_payload(None, {"codes.mel": torch.tensor([[1]])}, None) is None
|
||||
def test_talker2s2mel_full_payload_latent_mode_rejects_missing_latent():
|
||||
payload = {
|
||||
"codes.mel": torch.tensor([[3], [4], [STOP_MEL_TOKEN]]),
|
||||
"meta.use_gpt_latent": True,
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="missing hidden_states.latent"):
|
||||
talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-latent"))
|
||||
|
||||
|
||||
def test_talker2s2mel_full_payload_missing_required_fields_raises():
|
||||
with pytest.raises(ValueError, match="missing meta.use_gpt_latent"):
|
||||
talker2s2mel_full_payload(None, {"hidden_states.latent": torch.randn(1, LATENT_DIM)}, None)
|
||||
with pytest.raises(ValueError, match="missing meta.use_gpt_latent"):
|
||||
talker2s2mel_full_payload(None, {"codes.mel": torch.tensor([[1]])}, None)
|
||||
@@ -0,0 +1,155 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import builtins
|
||||
import importlib.util
|
||||
import platform
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import tomllib
|
||||
from packaging.requirements import Requirement
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
ROOT = Path(__file__).parents[1]
|
||||
TEXT_PROCESSING_PATH = ROOT / "vllm_omni" / "model_executor" / "models" / "indextts2" / "text_processing_v2_5.py"
|
||||
INDEXTTS2_FRONTEND_PACKAGES = {
|
||||
"fugashi",
|
||||
"unidic-lite",
|
||||
"wetext",
|
||||
"wetextprocessing",
|
||||
}
|
||||
INSTALL_COMMAND = "pip install 'vllm-omni[indextts2]'"
|
||||
|
||||
|
||||
def _load_text_processing_module(monkeypatch: pytest.MonkeyPatch):
|
||||
package_name = "_indextts2_optional_dependencies_test"
|
||||
package = types.ModuleType(package_name)
|
||||
package.__path__ = [str(TEXT_PROCESSING_PATH.parent)]
|
||||
tokenizer = types.ModuleType(f"{package_name}.tokenizer_v2_5")
|
||||
tokenizer.INDEXTTS25_TOKENIZER_FILE = "tokenizer.tiktoken"
|
||||
tokenizer.encode_indextts25_text = lambda *args, **kwargs: []
|
||||
tokenizer.lang_to_token = lambda lang: 0
|
||||
tokenizer.normalize_language_code = lambda lang: lang
|
||||
monkeypatch.setitem(sys.modules, package_name, package)
|
||||
monkeypatch.setitem(sys.modules, tokenizer.__name__, tokenizer)
|
||||
|
||||
module_name = f"{package_name}.text_processing_v2_5"
|
||||
spec = importlib.util.spec_from_file_location(module_name, TEXT_PROCESSING_PATH)
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
monkeypatch.setitem(sys.modules, module_name, module)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def _requirements_from_file(path: Path) -> list[Requirement]:
|
||||
requirements = []
|
||||
for raw_line in path.read_text().splitlines():
|
||||
line = raw_line.split("#", 1)[0].strip()
|
||||
if not line:
|
||||
continue
|
||||
if line.startswith(("-r ", "--requirement ")):
|
||||
included = line.split(maxsplit=1)[1]
|
||||
requirements.extend(_requirements_from_file(path.parent / included))
|
||||
continue
|
||||
requirements.append(Requirement(line))
|
||||
return requirements
|
||||
|
||||
|
||||
def _block_import(monkeypatch: pytest.MonkeyPatch, package: str) -> None:
|
||||
original_import = builtins.__import__
|
||||
|
||||
def controlled_import(name, *args, **kwargs):
|
||||
if name == package or name.startswith(f"{package}."):
|
||||
raise ModuleNotFoundError(f"No module named {package!r}", name=package)
|
||||
return original_import(name, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(builtins, "__import__", controlled_import)
|
||||
|
||||
|
||||
def test_indextts2_extra_declares_platform_specific_frontend_dependencies():
|
||||
project = tomllib.loads((ROOT / "pyproject.toml").read_text())
|
||||
requirements = [Requirement(value) for value in project["project"]["optional-dependencies"]["indextts2"]]
|
||||
|
||||
assert {
|
||||
(requirement.name.lower(), str(requirement.specifier), str(requirement.marker) if requirement.marker else None)
|
||||
for requirement in requirements
|
||||
} == {
|
||||
("fugashi", ">=1.2.0", None),
|
||||
("unidic-lite", ">=1.0.0", None),
|
||||
("wetext", ">=0.0.9", 'sys_platform != "linux"'),
|
||||
("wetextprocessing", "", 'sys_platform == "linux"'),
|
||||
}
|
||||
|
||||
|
||||
def test_common_requirements_exclude_indextts2_frontend_dependencies():
|
||||
requirements = _requirements_from_file(ROOT / "requirements" / "common.txt")
|
||||
|
||||
assert not ({requirement.name.lower() for requirement in requirements} & INDEXTTS2_FRONTEND_PACKAGES)
|
||||
|
||||
|
||||
def test_wetext_normalizer_constructor_error_propagates(monkeypatch):
|
||||
text_processing_v2_5 = _load_text_processing_module(monkeypatch)
|
||||
sentinel = RuntimeError("wetext constructor sentinel")
|
||||
|
||||
class FailingNormalizer:
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise sentinel
|
||||
|
||||
wetext = types.ModuleType("wetext")
|
||||
wetext.Normalizer = FailingNormalizer
|
||||
monkeypatch.setitem(sys.modules, "wetext", wetext)
|
||||
monkeypatch.setattr(platform, "system", lambda: "Darwin")
|
||||
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
text_processing_v2_5._load_wetext_normalizer("zh")
|
||||
|
||||
assert exc_info.value is sentinel
|
||||
|
||||
|
||||
def test_japanese_tagger_constructor_error_propagates(monkeypatch):
|
||||
text_processing_v2_5 = _load_text_processing_module(monkeypatch)
|
||||
sentinel = RuntimeError("fugashi constructor sentinel")
|
||||
|
||||
fugashi = types.ModuleType("fugashi")
|
||||
|
||||
def failing_tagger():
|
||||
raise sentinel
|
||||
|
||||
fugashi.Tagger = failing_tagger
|
||||
monkeypatch.setitem(sys.modules, "fugashi", fugashi)
|
||||
monkeypatch.setitem(sys.modules, "unidic_lite", types.ModuleType("unidic_lite"))
|
||||
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
text_processing_v2_5._load_japanese_tagger()
|
||||
|
||||
assert exc_info.value is sentinel
|
||||
|
||||
|
||||
def test_missing_wetext_backend_reports_indextts2_extra(monkeypatch):
|
||||
text_processing_v2_5 = _load_text_processing_module(monkeypatch)
|
||||
text_processing_v2_5._load_wetext_normalizer.cache_clear()
|
||||
monkeypatch.setattr(platform, "system", lambda: "Darwin")
|
||||
_block_import(monkeypatch, "wetext")
|
||||
|
||||
with pytest.raises(RuntimeError, match="IndexTTS 2.5 text normalization") as exc_info:
|
||||
text_processing_v2_5._load_wetext_normalizer("zh")
|
||||
|
||||
assert INSTALL_COMMAND in str(exc_info.value)
|
||||
text_processing_v2_5._load_wetext_normalizer.cache_clear()
|
||||
|
||||
|
||||
def test_missing_japanese_tagger_reports_indextts2_extra(monkeypatch):
|
||||
text_processing_v2_5 = _load_text_processing_module(monkeypatch)
|
||||
text_processing_v2_5._load_japanese_tagger.cache_clear()
|
||||
_block_import(monkeypatch, "fugashi")
|
||||
|
||||
with pytest.raises(RuntimeError, match="IndexTTS 2.5 Japanese preprocessing") as exc_info:
|
||||
text_processing_v2_5._load_japanese_tagger()
|
||||
|
||||
assert INSTALL_COMMAND in str(exc_info.value)
|
||||
text_processing_v2_5._load_japanese_tagger.cache_clear()
|
||||
@@ -1,6 +1,7 @@
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import vllm_omni.utils.mm_outputs as mm_outputs
|
||||
from vllm_omni.utils.mm_outputs import partition_flat_payload, partition_payload_list
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -69,3 +70,33 @@ def test_partition_duplex_audio_transcript_metadata_to_client_mm():
|
||||
assert "meta.turn_end" in client
|
||||
assert "meta.native_duplex_segment_text" not in client
|
||||
assert "meta.native_duplex_segment_text" in inter
|
||||
|
||||
|
||||
def test_snapshot_mm_payload_coalesces_compatible_tensor_list(monkeypatch):
|
||||
chunks = [
|
||||
torch.tensor([[7]], dtype=torch.long),
|
||||
torch.empty(0, 1, dtype=torch.long),
|
||||
torch.tensor([[9], [10]], dtype=torch.long),
|
||||
]
|
||||
real_cat = torch.cat
|
||||
cat_shapes = []
|
||||
|
||||
def recording_cat(tensors, *args, **kwargs):
|
||||
cat_shapes.append([tuple(tensor.shape) for tensor in tensors])
|
||||
return real_cat(tensors, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(torch, "cat", recording_cat)
|
||||
|
||||
result = mm_outputs.snapshot_mm_payload({"codes.mel": chunks})
|
||||
|
||||
assert cat_shapes == [[(1, 1), (0, 1), (2, 1)]]
|
||||
output_chunks = result["codes.mel"]
|
||||
assert [tuple(chunk.shape) for chunk in output_chunks] == [
|
||||
(1, 1),
|
||||
(0, 1),
|
||||
(2, 1),
|
||||
]
|
||||
torch.testing.assert_close(output_chunks[0], chunks[0])
|
||||
torch.testing.assert_close(output_chunks[1], chunks[1])
|
||||
torch.testing.assert_close(output_chunks[2], chunks[2])
|
||||
assert output_chunks[0].untyped_storage().data_ptr() == output_chunks[2].untyped_storage().data_ptr()
|
||||
@@ -5,6 +5,7 @@ import pytest
|
||||
import torch
|
||||
from vllm.v1.cudagraph_dispatcher import CUDAGraphMode
|
||||
|
||||
from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
|
||||
from vllm_omni.worker.gpu_model_runner import (
|
||||
OmniGPUModelRunner,
|
||||
_filter_mrope_kwargs_for_model,
|
||||
@@ -691,6 +692,49 @@ def test_full_payload_output_accumulation_hook_matrix():
|
||||
assert not runner._should_accumulate_full_payload_output()
|
||||
|
||||
|
||||
def _make_request_end_payload_runner(*, enabled=True, prefix_cache=None):
|
||||
runner = object.__new__(GPUARModelRunner)
|
||||
runner.model = SimpleNamespace(omni_payload_at_request_end=enabled)
|
||||
runner.omni_prefix_cache = prefix_cache
|
||||
runner.model_config = SimpleNamespace(
|
||||
model_arch="IndexTTS25TalkerForConditionalGeneration",
|
||||
model_stage="indextts2_5_talker",
|
||||
async_chunk=False,
|
||||
final_output=False,
|
||||
custom_process_next_stage_input_func="module.full_payload",
|
||||
)
|
||||
runner._custom_process_func = object()
|
||||
runner._pending_full_payload_send = {}
|
||||
runner._stage_id = 0
|
||||
runner._omni_connector = object()
|
||||
return runner
|
||||
|
||||
|
||||
def test_request_end_payload_d2h_gate_requires_opt_in_and_no_prefix_cache():
|
||||
assert _make_request_end_payload_runner()._should_defer_full_payload_d2h()
|
||||
assert not _make_request_end_payload_runner(enabled=False)._should_defer_full_payload_d2h()
|
||||
assert not _make_request_end_payload_runner(prefix_cache=object())._should_defer_full_payload_d2h()
|
||||
|
||||
|
||||
def test_request_end_payload_suppresses_per_step_multimodal_outputs():
|
||||
runner = _make_request_end_payload_runner()
|
||||
|
||||
def unexpected_build(_payload):
|
||||
raise AssertionError("request-end payloads must stay inside the GPU accumulator")
|
||||
|
||||
runner._build_multimodal_outputs = unexpected_build
|
||||
pooler_inter = [{"codes.mel": torch.tensor([[7]])}]
|
||||
|
||||
inter_stage, client = runner._build_omni_step_outputs(
|
||||
pooler_inter,
|
||||
pooler_inter,
|
||||
defer_full_payload_d2h=True,
|
||||
)
|
||||
|
||||
assert inter_stage is None
|
||||
assert client is None
|
||||
|
||||
|
||||
def test_sync_local_stage_payloads_retains_payload_until_request_is_active():
|
||||
runner = object.__new__(OmniGPUModelRunner)
|
||||
payload = {"codes": {"audio": [1, 2, 3]}}
|
||||
|
||||
@@ -61,7 +61,10 @@ from vllm_omni.model_executor.models.hunyuan_image3.pipeline import (
|
||||
HUNYUAN_IMAGE3_PIPELINE,
|
||||
)
|
||||
from vllm_omni.model_executor.models.hunyuan_video.pipeline import HUNYUAN_VIDEO_15_PIPELINE
|
||||
from vllm_omni.model_executor.models.indextts2.pipeline import INDEXTTS2_PIPELINE
|
||||
from vllm_omni.model_executor.models.indextts2.pipeline import (
|
||||
INDEXTTS2_PIPELINE,
|
||||
INDEXTTS25_PIPELINE,
|
||||
)
|
||||
from vllm_omni.model_executor.models.lance.pipeline import LANCE_PIPELINE
|
||||
from vllm_omni.model_executor.models.mammoth_moda2.pipeline import (
|
||||
MAMMOTH_MODA2_AR_PIPELINE,
|
||||
@@ -164,6 +167,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
|
||||
"higgs_multimodal_qwen3": HIGGS_AUDIO_V3_PIPELINE,
|
||||
"dynin_omni": DYNIN_OMNI_PIPELINE,
|
||||
"indextts2": INDEXTTS2_PIPELINE,
|
||||
"indextts2_5": INDEXTTS25_PIPELINE,
|
||||
"soulxsinger_svc": SOULXSINGER_SVC_PIPELINE,
|
||||
"soulxsinger_svs": SOULXSINGER_SVS_PIPELINE,
|
||||
}
|
||||
|
||||
@@ -57,9 +57,10 @@ _FULL_PAYLOAD_INPUT_STAGES: frozenset[tuple[str, str]] = frozenset(
|
||||
# audex TTA: tta thinker (Stage 0) -> XCodec1 (Stage 1, always sync
|
||||
# full-payload — CNN codec decoded over the full sequence).
|
||||
("AudexXCodec1", "audex_xcodec"),
|
||||
# indextts2: indextts2_talker (Stage 0) -> indextts2_s2mel_decoder
|
||||
# (Stage 1). Stage 1 consumes the complete mel/latent payload.
|
||||
# indextts2 / indextts2_5: talker (Stage 0) -> s2mel decoder
|
||||
# (Stage 1). Stage 1 consumes the complete mel/optional-latent payload.
|
||||
("IndexTTS2S2MelDecoder", "indextts2_s2mel_decoder"),
|
||||
("IndexTTS25S2MelDecoder", "indextts2_5_s2mel_decoder"),
|
||||
# dynin: token2text (Stage 0) -> token2image (Stage 1) ->
|
||||
# token2audio (Stage 2). Producer wires via
|
||||
# custom_process_next_stage_input_func: *_full_payload in deploy yaml.
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
async_chunk: false
|
||||
trust_remote_code: true
|
||||
pipeline: indextts2_5
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
codec_streaming: false
|
||||
connector_get_sleep_s: 0.01
|
||||
connector_get_max_wait_first_chunk: 5000
|
||||
connector_get_max_wait: 500
|
||||
stages:
|
||||
- stage_id: 0
|
||||
max_num_seqs: 4
|
||||
# Triton matches the native GPT-2 logits for this non-RoPE model.
|
||||
# FlashInfer produces materially different semantic-token logits.
|
||||
attention_backend: TRITON_ATTN
|
||||
gpu_memory_utilization: 0.4
|
||||
enable_prefix_caching: false
|
||||
enable_chunked_prefill: false
|
||||
skip_tokenizer_init: true
|
||||
async_scheduling: true
|
||||
max_num_batched_tokens: 2560
|
||||
max_model_len: 2560
|
||||
devices: "0"
|
||||
output_connectors:
|
||||
to_stage_1: connector_of_shared_memory
|
||||
hf_overrides:
|
||||
use_gpt_latent: false
|
||||
# Stage 0 uses plain vLLM sampling, not the upstream default
|
||||
# num_beams=3 beam search. Use upstream num_beams=1 for parity comparisons.
|
||||
default_sampling_params:
|
||||
temperature: 0.8
|
||||
top_p: 0.8
|
||||
top_k: 30
|
||||
max_tokens: 1500
|
||||
repetition_penalty: 10.0
|
||||
stop_token_ids: [8193]
|
||||
detokenize: false
|
||||
- stage_id: 1
|
||||
max_num_seqs: 4
|
||||
gpu_memory_utilization: 0.4
|
||||
enable_prefix_caching: false
|
||||
skip_tokenizer_init: true
|
||||
enforce_eager: true
|
||||
async_scheduling: true
|
||||
max_num_batched_tokens: 8192
|
||||
max_model_len: 32768
|
||||
devices: "0"
|
||||
input_connectors:
|
||||
from_stage_0: connector_of_shared_memory
|
||||
hf_overrides:
|
||||
use_gpt_latent: false
|
||||
diffusion_steps: 25
|
||||
s2mel_dit_bf16: true
|
||||
s2mel_cfm_batch_size: 4
|
||||
inference_cfg_rate: 0.7
|
||||
s2mel_dit_cuda_graph: false
|
||||
s2mel_dit_torch_compile: true
|
||||
s2mel_dit_torch_compile_mode: default
|
||||
s2mel_vocoder_torch_compile: false
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
max_tokens: 32767
|
||||
@@ -13,8 +13,8 @@
|
||||
# the custom generation loop — set explicitly everywhere).
|
||||
# * AR stages (thinker, talker): model-dependent. Qwen2.5-Omni runs eager
|
||||
# on CUDA (thinker uses custom ops that don't trace cleanly); NPU / XPU
|
||||
# platform overrides flip back to false where cudagraph is verified.
|
||||
# Qwen3-Omni / Qwen3-TTS AR stages use the default (false = cudagraph on).
|
||||
# overrides are documented below. Qwen3-Omni / Qwen3-TTS AR stages use the
|
||||
# default (false = cudagraph on).
|
||||
async_chunk: false
|
||||
|
||||
stages:
|
||||
@@ -82,13 +82,15 @@ platforms:
|
||||
|
||||
xpu:
|
||||
stages:
|
||||
# Verified on 2x Intel Arc Pro B60. Both AR stages use cudagraphs.
|
||||
# PyTorch 2.13 XPU Dynamo raises duplicate-handler assertions while
|
||||
# compiling both Qwen2.5 AR stages. Keep them eager until that toolchain
|
||||
# issue is fixed; stage 2 already inherits enforce_eager: true above.
|
||||
- stage_id: 0
|
||||
gpu_memory_utilization: 0.9
|
||||
enforce_eager: false
|
||||
enforce_eager: true
|
||||
- stage_id: 1
|
||||
gpu_memory_utilization: 0.5
|
||||
enforce_eager: false
|
||||
enforce_eager: true
|
||||
- stage_id: 2
|
||||
gpu_memory_utilization: 0.3
|
||||
# Stage 2 colocates with stage 1's device on XPU.
|
||||
|
||||
@@ -22,6 +22,8 @@ _ARCH_TO_MODEL_TYPE: dict[str, str] = {
|
||||
"GLMTTSForConditionalGeneration": "glm_tts",
|
||||
"IndexTTS2S2MelDecoder": "indextts2",
|
||||
"IndexTTS2TalkerForConditionalGeneration": "indextts2",
|
||||
"IndexTTS25S2MelDecoder": "indextts2_5",
|
||||
"IndexTTS25TalkerForConditionalGeneration": "indextts2_5",
|
||||
"OmniVoiceModel": "omnivoice",
|
||||
# PersonaPlex ships an empty config.json, so create_model_config() must patch
|
||||
# model_type=personaplex from the arch name or the staged pipeline can't load
|
||||
@@ -44,6 +46,7 @@ def _register_omni_hf_configs() -> None:
|
||||
|
||||
from vllm_omni.model_executor.models.indextts2.configuration_indextts2 import (
|
||||
IndexTTS2Config,
|
||||
IndexTTS25Config,
|
||||
)
|
||||
from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
|
||||
MingDenseConfig,
|
||||
@@ -79,6 +82,7 @@ def _register_omni_hf_configs() -> None:
|
||||
("dense", MingDenseConfig),
|
||||
("bailingmm", MingMoeConfig),
|
||||
("indextts2", IndexTTS2Config),
|
||||
("indextts2_5", IndexTTS25Config),
|
||||
("moss_tts_local", MossTTSLocalConfig),
|
||||
("moss_tts_realtime", MossTTSRealtimeConfig),
|
||||
("qwen3_tts", Qwen3TTSConfig),
|
||||
@@ -250,7 +254,15 @@ class OmniEngineArgs(EngineArgs):
|
||||
if config_dict.get("model_type"):
|
||||
return # config.json already has model_type, no patching needed
|
||||
except Exception:
|
||||
return # can't load config, let vLLM handle the error
|
||||
# The official IndexTTS 2.5 bundle has no HuggingFace config.json.
|
||||
# Keep this exception model-scoped so other loader failures retain
|
||||
# vLLM's normal error path.
|
||||
if model_type != "indextts2_5" or not os.path.isdir(self.model):
|
||||
return
|
||||
config_path = os.path.join(self.model, "config.json")
|
||||
if os.path.lexists(config_path):
|
||||
return
|
||||
config_dict = {}
|
||||
|
||||
# Create a temp dir with a patched config.json
|
||||
temp_dir = tempfile.mkdtemp(prefix="omni_hf_config_")
|
||||
|
||||
@@ -20,6 +20,7 @@ from dataclasses import dataclass, fields, replace
|
||||
from typing import Any, Literal, cast
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.renderers import BaseRenderer
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.tokenizers import cached_tokenizer_from_config
|
||||
from vllm.usage.usage_lib import UsageContext
|
||||
@@ -1119,11 +1120,30 @@ def build_llm_stage_output_processor(
|
||||
)
|
||||
|
||||
|
||||
class _TokenOnlyRenderer(BaseRenderer):
|
||||
"""Renderer for stages that explicitly skip tokenizer initialization."""
|
||||
|
||||
def render_messages(self, messages, params):
|
||||
raise ValueError(
|
||||
"Chat messages are unavailable when skip_tokenizer_init=True; submit prompt_token_ids or prompt_embeds"
|
||||
)
|
||||
|
||||
|
||||
def _build_token_only_renderer(stage_vllm_config: Any) -> BaseRenderer:
|
||||
return _TokenOnlyRenderer(stage_vllm_config, tokenizer=None)
|
||||
|
||||
|
||||
def build_stage0_input_processor(stage_vllm_config: Any) -> InputProcessor:
|
||||
"""Build the shared stage-0 input processor."""
|
||||
|
||||
patch_generation_config_if_needed(stage_vllm_config.model_config)
|
||||
input_processor = InputProcessor(vllm_config=stage_vllm_config)
|
||||
if bool(getattr(stage_vllm_config.model_config, "skip_tokenizer_init", False)):
|
||||
input_processor = InputProcessor(
|
||||
vllm_config=stage_vllm_config,
|
||||
renderer=_build_token_only_renderer(stage_vllm_config),
|
||||
)
|
||||
else:
|
||||
input_processor = InputProcessor(vllm_config=stage_vllm_config)
|
||||
input_processor.input_preprocessor = OmniInputPreprocessor(
|
||||
vllm_config=stage_vllm_config,
|
||||
renderer=input_processor.renderer,
|
||||
|
||||
@@ -996,9 +996,11 @@ async def omni_init_app_state(
|
||||
)
|
||||
|
||||
# Warm up chat template processing to avoid first-request latency
|
||||
# Upstream f5ffc59b6a removed OpenAIServingChat.warmup() and moved the
|
||||
# warmup onto the renderer (OnlineRenderer.warmup()); mirror upstream.
|
||||
state.online_renderer.warmup()
|
||||
# Upstream f5ffc59b6a moved warmup onto OnlineRenderer. Accelerator
|
||||
# images can temporarily lag that renderer API, where warmup is optional.
|
||||
renderer_warmup = getattr(state.online_renderer, "warmup", None)
|
||||
if renderer_warmup is not None:
|
||||
renderer_warmup()
|
||||
|
||||
state.openai_serving_completion = (
|
||||
OpenAIServingCompletion(
|
||||
|
||||
@@ -82,7 +82,8 @@ class OpenAICreateSpeechRequest(BaseModel):
|
||||
description=(
|
||||
"Streaming switch; defaults to OpenAI speech.audio.* SSE events. "
|
||||
"Set stream_format='audio' to opt into raw pcm/wav byte streaming. "
|
||||
"Requires response_format='pcm' or 'wav'. Speed adjustment is not supported when streaming."
|
||||
"HTTP streaming requires response_format='pcm' or 'wav'. "
|
||||
"Models with native speed control may accept speed adjustment over HTTP."
|
||||
),
|
||||
)
|
||||
|
||||
@@ -317,8 +318,6 @@ class OpenAICreateSpeechRequest(BaseModel):
|
||||
)
|
||||
if self.speed is None:
|
||||
self.speed = 1.0
|
||||
elif self.speed != 1.0:
|
||||
raise ValueError("Speed adjustment is not supported when streaming. Set speed=1.0 or omit it.")
|
||||
return self
|
||||
|
||||
|
||||
@@ -532,7 +531,8 @@ class StreamingSpeechSessionConfig(BaseModel):
|
||||
default=False,
|
||||
description=(
|
||||
"If true, send raw PCM audio chunks progressively over WebSocket. "
|
||||
"Requires response_format='pcm'. Speed adjustment is not supported when streaming."
|
||||
"Requires response_format='pcm'. WebSocket streaming currently requires "
|
||||
"speed=1.0, including for models with native speed control."
|
||||
),
|
||||
)
|
||||
word_timestamps: bool = Field(
|
||||
@@ -556,5 +556,8 @@ class StreamingSpeechSessionConfig(BaseModel):
|
||||
if self.speed is None:
|
||||
self.speed = 1.0
|
||||
elif self.speed != 1.0:
|
||||
raise ValueError("Speed adjustment is not supported when stream_audio=true. Set speed=1.0 or omit it.")
|
||||
raise ValueError(
|
||||
"WebSocket stream_audio=true currently requires speed=1.0; "
|
||||
"native speed control is only available through HTTP streaming."
|
||||
)
|
||||
return self
|
||||
@@ -83,7 +83,14 @@ from vllm.inputs import PromptType
|
||||
from vllm.logger import init_logger
|
||||
from vllm.multimodal.media.connector import MediaConnector
|
||||
from vllm.outputs import RequestOutput
|
||||
from vllm.parser.mistral import MistralToolCall
|
||||
|
||||
try:
|
||||
from vllm.parser.mistral import MistralToolCall
|
||||
except ImportError:
|
||||
# Accelerator images can lag the core vLLM release by one version.
|
||||
from vllm.tool_parsers.mistral_tool_parser import MistralToolCall
|
||||
|
||||
|
||||
from vllm.reasoning import ReasoningParser
|
||||
from vllm.renderers import BaseRenderer, merge_kwargs
|
||||
from vllm.renderers.inputs import TokPrompt
|
||||
|
||||
@@ -101,6 +101,7 @@ _SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = {
|
||||
"higgs_audio_v2",
|
||||
"higgs_audio_v3",
|
||||
"indextts2",
|
||||
"indextts2_5",
|
||||
"audex",
|
||||
"audex_tta",
|
||||
}
|
||||
@@ -513,6 +514,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
self._adapter = adapter_cls(ctx)
|
||||
return self._adapter
|
||||
|
||||
def _uses_native_speed_control(self) -> bool:
|
||||
adapter = self._get_tts_adapter()
|
||||
return bool(adapter is not None and adapter.native_speed_control)
|
||||
|
||||
def _audio_encode_speed(self, request: OpenAICreateSpeechRequest) -> float:
|
||||
if self._uses_native_speed_control():
|
||||
return 1.0
|
||||
return float(request.speed or 1.0)
|
||||
|
||||
async def warmup(self) -> None:
|
||||
"""Run a synthetic speech request to trigger all first-request warmup.
|
||||
|
||||
@@ -3473,7 +3483,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
audio_tensor=audio_tensor,
|
||||
sample_rate=sample_rate,
|
||||
response_format=request.response_format or "wav",
|
||||
speed=request.speed or 1.0,
|
||||
speed=self._audio_encode_speed(request),
|
||||
base64_encode=base64_encode,
|
||||
)
|
||||
audio_response: AudioResponse = self.create_audio(audio_obj)
|
||||
@@ -3587,7 +3597,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
audio_tensor=audio_tensor,
|
||||
sample_rate=sample_rate,
|
||||
response_format=request.response_format or "wav",
|
||||
speed=request.speed or 1.0,
|
||||
speed=self._audio_encode_speed(request),
|
||||
base64_encode=False,
|
||||
)
|
||||
audio_response: AudioResponse = self.create_audio(audio_obj)
|
||||
@@ -3629,7 +3639,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
return response_format, self.create_error_response(
|
||||
f"{mode_label} is only supported for 'pcm' and 'wav' formats. Got '{response_format}'."
|
||||
)
|
||||
if request.speed is not None and request.speed != 1.0:
|
||||
if request.speed is not None and request.speed != 1.0 and not self._uses_native_speed_control():
|
||||
return response_format, self.create_error_response(
|
||||
f"{mode_label} is not supported with speed adjustment. "
|
||||
"Use a non-streaming request or remove the speed parameter."
|
||||
|
||||
@@ -115,6 +115,8 @@ class TTSModelAdapter(ABC):
|
||||
detect_priority: ClassVar[int] = 100
|
||||
#: Serving backend: ``"ar"`` (engine_client) or ``"diffusion"``.
|
||||
backend: ClassVar[str] = "ar"
|
||||
#: Whether the model consumes ``request.speed`` in its native parameters.
|
||||
native_speed_control: ClassVar[bool] = False
|
||||
|
||||
max_new_tokens_min = 1
|
||||
|
||||
|
||||
@@ -6,7 +6,7 @@ from __future__ import annotations
|
||||
import hashlib
|
||||
import math
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from typing import TYPE_CHECKING, Any, ClassVar
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
@@ -15,6 +15,10 @@ from vllm.utils import random_uuid
|
||||
|
||||
from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
|
||||
from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest
|
||||
from vllm_omni.model_executor.models.indextts2.configuration_indextts2 import (
|
||||
INDEXTTS25_MAX_DURATION_FACTOR,
|
||||
INDEXTTS25_MIN_DURATION_FACTOR,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
|
||||
@@ -102,6 +106,8 @@ def indextts2_conditioning_cache_salt(
|
||||
"emo_text",
|
||||
"use_emo_text",
|
||||
"use_random",
|
||||
"lang",
|
||||
"text_normalization",
|
||||
):
|
||||
h.update(b"\x00")
|
||||
h.update(key.encode("utf-8"))
|
||||
@@ -206,14 +212,34 @@ class IndexTTS2Adapter(ARTTSAdapter):
|
||||
estimate_indextts2_prefill_prompt_len,
|
||||
)
|
||||
|
||||
prompt_kwargs: dict[str, Any] = {}
|
||||
if self.name == "indextts2_5":
|
||||
from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import (
|
||||
INDEXTTS25_TOKENIZER_FILE,
|
||||
)
|
||||
|
||||
hf_config = getattr(server.engine_client.model_config, "hf_config", None)
|
||||
prompt_kwargs["tokenizer_file"] = getattr(
|
||||
hf_config,
|
||||
"tokenizer_file",
|
||||
INDEXTTS25_TOKENIZER_FILE,
|
||||
)
|
||||
ph_len = estimate_indextts2_prefill_prompt_len(
|
||||
server.engine_client.model_config.model,
|
||||
request.input,
|
||||
model_type=self.name,
|
||||
lang=str(_first_conditioning_value(tts_params.get("lang")) or "zh"),
|
||||
text_normalization=bool(
|
||||
_first_conditioning_value(tts_params.get("text_normalization"))
|
||||
if "text_normalization" in tts_params
|
||||
else True
|
||||
),
|
||||
**prompt_kwargs,
|
||||
)
|
||||
prompt = tokens_input(prompt_token_ids=[1] * ph_len)
|
||||
prompt["additional_information"] = tts_params
|
||||
prompt["cache_salt"] = indextts2_conditioning_cache_salt(request, tts_params)
|
||||
return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type="indextts2")
|
||||
return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type=self.name)
|
||||
|
||||
async def _build_params(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]:
|
||||
server = self.ctx.server
|
||||
@@ -241,3 +267,59 @@ class IndexTTS2Adapter(ARTTSAdapter):
|
||||
else:
|
||||
params[key] = [extras[key]]
|
||||
return params
|
||||
|
||||
|
||||
@register_tts_adapter
|
||||
class IndexTTS25Adapter(IndexTTS2Adapter):
|
||||
stage_keys = frozenset({"indextts2_5_talker"})
|
||||
name = "indextts2_5"
|
||||
native_speed_control: ClassVar[bool] = True
|
||||
|
||||
def validate(self, request: OpenAICreateSpeechRequest) -> str | None:
|
||||
error = super().validate(request)
|
||||
if error is not None:
|
||||
return error
|
||||
speed = request.speed if request.speed is not None else 1.0
|
||||
if speed < INDEXTTS25_MIN_DURATION_FACTOR or speed > INDEXTTS25_MAX_DURATION_FACTOR:
|
||||
return "IndexTTS 2.5 speed must be between 0.5 and 2.0"
|
||||
return None
|
||||
|
||||
def _validate_extra_params(self, extras: Mapping[str, Any]) -> str | None:
|
||||
error = super()._validate_extra_params(extras)
|
||||
if error is not None:
|
||||
return error
|
||||
|
||||
if "lang" in extras:
|
||||
value = extras["lang"]
|
||||
if not isinstance(value, str):
|
||||
return "extra_params.lang must be a string"
|
||||
from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import (
|
||||
normalize_language_code,
|
||||
)
|
||||
|
||||
try:
|
||||
normalize_language_code(value)
|
||||
except ValueError as exc:
|
||||
return str(exc)
|
||||
if "text_normalization" in extras and not isinstance(
|
||||
extras["text_normalization"],
|
||||
bool,
|
||||
):
|
||||
return "extra_params.text_normalization must be a boolean"
|
||||
return None
|
||||
|
||||
async def _build_params(
|
||||
self,
|
||||
request: OpenAICreateSpeechRequest,
|
||||
) -> dict[str, Any]:
|
||||
params = await super()._build_params(request)
|
||||
extras = request.extra_params if isinstance(request.extra_params, dict) else {}
|
||||
from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import (
|
||||
normalize_language_code,
|
||||
)
|
||||
|
||||
params["lang"] = [normalize_language_code(str(extras.get("lang", "zh")))]
|
||||
params["text_normalization"] = [bool(extras.get("text_normalization", True))]
|
||||
speed = request.speed if request.speed is not None else 1.0
|
||||
params["duration_factor"] = [1.0 / speed]
|
||||
return params
|
||||
@@ -571,7 +571,10 @@ class CodePredictorWrapper(nn.Module):
|
||||
self._lm_heads_list = list(self.lm_head)
|
||||
self._codec_embeds_list = list(self.model.codec_embedding)
|
||||
|
||||
if not current_omni_platform.supports_torch_inductor():
|
||||
# Torch 2.13 XPU Dynamo can double-register built-in handlers when
|
||||
# spawned workers compile this predictor. Keep this narrow path eager
|
||||
# until the upstream XPU compiler issue is resolved.
|
||||
if current_omni_platform.is_xpu() or not current_omni_platform.supports_torch_inductor():
|
||||
# NPU or other platforms without Inductor support
|
||||
self._compiled_model_fwd = self.model.forward
|
||||
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
# SPDX-License-Identifier: MIT
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the Amphion project
|
||||
"""Inference-only EnhancedCodec runtime for IndexTTS 2.5."""
|
||||
|
||||
from .models import EnhancedCodec
|
||||
|
||||
__all__ = ["EnhancedCodec"]
|
||||
@@ -0,0 +1,138 @@
|
||||
# SPDX-License-Identifier: MIT
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the Amphion project
|
||||
"""Inference-only EnhancedCodec decoder for IndexTTS 2.5.
|
||||
|
||||
The checkpoint originates from the Amphion EnhancedCodec implementation
|
||||
vendored by IndexTTS 2.5. The encoder and training paths are deliberately not
|
||||
constructed because vLLM-Omni only performs semantic-code decoding here.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .quantizer import ResidualVQDecoder
|
||||
from .vocos import VocosBackbone
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _config_value(
|
||||
cfg: Mapping[str, Any] | Any | None,
|
||||
name: str,
|
||||
default: Any,
|
||||
) -> Any:
|
||||
if cfg is None:
|
||||
return default
|
||||
if isinstance(cfg, Mapping):
|
||||
return cfg.get(name, default)
|
||||
return getattr(cfg, name, default)
|
||||
|
||||
|
||||
class EnhancedCodec(nn.Module):
|
||||
"""Subset of EnhancedCodec required for ``decode(codes)``."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
codebook_size: int = 8192,
|
||||
hidden_size: int = 1024,
|
||||
codebook_dim: int = 8,
|
||||
vocos_dim: int = 384,
|
||||
vocos_intermediate_dim: int = 2048,
|
||||
vocos_num_layers: int = 12,
|
||||
num_quantizers: int = 1,
|
||||
downsample_scale: int | None = 2,
|
||||
cfg: Mapping[str, Any] | Any | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
codebook_size = int(_config_value(cfg, "codebook_size", codebook_size))
|
||||
hidden_size = int(_config_value(cfg, "hidden_size", hidden_size))
|
||||
codebook_dim = int(_config_value(cfg, "codebook_dim", codebook_dim))
|
||||
vocos_dim = int(_config_value(cfg, "vocos_dim", vocos_dim))
|
||||
vocos_intermediate_dim = int(_config_value(cfg, "vocos_intermediate_dim", vocos_intermediate_dim))
|
||||
vocos_num_layers = int(_config_value(cfg, "vocos_num_layers", vocos_num_layers))
|
||||
num_quantizers = int(_config_value(cfg, "num_quantizers", num_quantizers))
|
||||
downsample_scale = _config_value(
|
||||
cfg,
|
||||
"downsample_scale",
|
||||
downsample_scale,
|
||||
)
|
||||
|
||||
self.hidden_size = hidden_size
|
||||
self.num_quantizers = num_quantizers
|
||||
self.downsample_scale = downsample_scale
|
||||
self.decoder = nn.Sequential(
|
||||
VocosBackbone(
|
||||
input_channels=hidden_size,
|
||||
dim=vocos_dim,
|
||||
intermediate_dim=vocos_intermediate_dim,
|
||||
num_layers=vocos_num_layers,
|
||||
),
|
||||
nn.Linear(vocos_dim, hidden_size),
|
||||
)
|
||||
self.quantizer = ResidualVQDecoder(
|
||||
input_dim=hidden_size,
|
||||
num_quantizers=num_quantizers,
|
||||
codebook_size=codebook_size,
|
||||
codebook_dim=codebook_dim,
|
||||
)
|
||||
if downsample_scale is not None and downsample_scale > 1:
|
||||
self.up = nn.Conv1d(
|
||||
hidden_size,
|
||||
hidden_size,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
)
|
||||
|
||||
def decode(self, codes: torch.Tensor) -> torch.Tensor:
|
||||
"""Decode `[B,T]` or `[N,B,T]` codes to `[B,T*scale,hidden]`."""
|
||||
if codes.ndim == 2:
|
||||
codes = codes.unsqueeze(0)
|
||||
if codes.ndim != 3:
|
||||
raise ValueError(f"EnhancedCodec codes must have shape [B,T] or [N,B,T], got {tuple(codes.shape)}")
|
||||
if codes.shape[0] != self.num_quantizers:
|
||||
raise ValueError(f"EnhancedCodec expected {self.num_quantizers} codebooks, got {codes.shape[0]}")
|
||||
|
||||
hidden = self.quantizer.vq2emb(codes)
|
||||
hidden = self.decoder(hidden)
|
||||
if self.downsample_scale is not None and self.downsample_scale > 1:
|
||||
hidden = hidden.transpose(1, 2)
|
||||
hidden = F.interpolate(
|
||||
hidden,
|
||||
scale_factor=self.downsample_scale,
|
||||
mode="nearest",
|
||||
)
|
||||
hidden = self.up(hidden).transpose(1, 2)
|
||||
return hidden
|
||||
|
||||
def load_checkpoint(self, checkpoint_path: str) -> None:
|
||||
path = Path(checkpoint_path)
|
||||
if not path.is_file():
|
||||
raise FileNotFoundError(f"EnhancedCodec checkpoint not found: {path}")
|
||||
checkpoint = torch.load(path, map_location="cpu", weights_only=True)
|
||||
saved_state = checkpoint.get("model", checkpoint)
|
||||
own_state = self.state_dict()
|
||||
missing = [
|
||||
name
|
||||
for name, value in own_state.items()
|
||||
if name not in saved_state or saved_state[name].shape != value.shape
|
||||
]
|
||||
if missing:
|
||||
raise RuntimeError("EnhancedCodec checkpoint is missing required decoder weights: " + ", ".join(missing))
|
||||
self.load_state_dict(
|
||||
{name: saved_state[name] for name in own_state},
|
||||
strict=True,
|
||||
)
|
||||
logger.info(
|
||||
"Loaded %d inference-only EnhancedCodec tensors from %s",
|
||||
len(own_state),
|
||||
path,
|
||||
)
|
||||
@@ -0,0 +1,91 @@
|
||||
# SPDX-License-Identifier: MIT
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the Amphion project
|
||||
"""Decode-only factorized residual vector quantizer.
|
||||
|
||||
Adapted from the Amphion codec implementation vendored by IndexTTS 2.5.
|
||||
Training-only paths are intentionally omitted.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.nn.utils import weight_norm
|
||||
|
||||
|
||||
def _weight_norm_conv1d(*args, **kwargs) -> nn.Conv1d:
|
||||
return weight_norm(nn.Conv1d(*args, **kwargs))
|
||||
|
||||
|
||||
class FactorizedVectorQuantizeDecoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
input_dim: int,
|
||||
codebook_size: int,
|
||||
codebook_dim: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.input_dim = input_dim
|
||||
self.codebook_dim = codebook_dim
|
||||
if input_dim != codebook_dim:
|
||||
# Keep upstream module names and weight-normalized parameter layout
|
||||
# so codec.pth keys load without conversion.
|
||||
self.in_project = _weight_norm_conv1d(
|
||||
input_dim,
|
||||
codebook_dim,
|
||||
kernel_size=1,
|
||||
)
|
||||
self.out_project = _weight_norm_conv1d(
|
||||
codebook_dim,
|
||||
input_dim,
|
||||
kernel_size=1,
|
||||
)
|
||||
else:
|
||||
self.in_project = nn.Identity()
|
||||
self.out_project = nn.Identity()
|
||||
self.codebook = nn.Embedding(codebook_size, codebook_dim)
|
||||
|
||||
def vq2emb(self, codes: torch.Tensor) -> torch.Tensor:
|
||||
embedding = F.embedding(codes, self.codebook.weight).transpose(1, 2)
|
||||
return self.out_project(embedding)
|
||||
|
||||
|
||||
class ResidualVQDecoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
input_dim: int,
|
||||
num_quantizers: int,
|
||||
codebook_size: int,
|
||||
codebook_dim: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.num_quantizers = num_quantizers
|
||||
self.quantizers = nn.ModuleList(
|
||||
[
|
||||
FactorizedVectorQuantizeDecoder(
|
||||
input_dim=input_dim,
|
||||
codebook_size=codebook_size,
|
||||
codebook_dim=codebook_dim,
|
||||
)
|
||||
for _ in range(num_quantizers)
|
||||
]
|
||||
)
|
||||
|
||||
def vq2emb(
|
||||
self,
|
||||
codes: torch.Tensor,
|
||||
n_quantizers: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
count = self.num_quantizers if n_quantizers is None else n_quantizers
|
||||
quantized: torch.Tensor | None = None
|
||||
for index, quantizer in enumerate(self.quantizers):
|
||||
if index >= count:
|
||||
break
|
||||
value = quantizer.vq2emb(codes[index])
|
||||
quantized = value if quantized is None else quantized + value
|
||||
if quantized is None:
|
||||
raise ValueError("EnhancedCodec requires at least one quantizer")
|
||||
return quantized
|
||||
@@ -0,0 +1,82 @@
|
||||
# SPDX-License-Identifier: MIT
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the Amphion project
|
||||
"""ConvNeXt Vocos backbone used by the IndexTTS 2.5 codec decoder."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class ConvNeXtBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
dim: int,
|
||||
intermediate_dim: int,
|
||||
layer_scale_init_value: float,
|
||||
):
|
||||
super().__init__()
|
||||
self.dwconv = nn.Conv1d(
|
||||
dim,
|
||||
dim,
|
||||
kernel_size=7,
|
||||
padding=3,
|
||||
groups=dim,
|
||||
)
|
||||
self.norm = nn.LayerNorm(dim, eps=1e-6)
|
||||
self.pwconv1 = nn.Linear(dim, intermediate_dim)
|
||||
self.act = nn.GELU()
|
||||
self.pwconv2 = nn.Linear(intermediate_dim, dim)
|
||||
self.gamma = nn.Parameter(
|
||||
layer_scale_init_value * torch.ones(dim),
|
||||
requires_grad=True,
|
||||
)
|
||||
|
||||
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
|
||||
residual = inputs
|
||||
hidden = self.dwconv(inputs).transpose(1, 2)
|
||||
hidden = self.norm(hidden)
|
||||
hidden = self.pwconv1(hidden)
|
||||
hidden = self.act(hidden)
|
||||
hidden = self.pwconv2(hidden)
|
||||
hidden = self.gamma * hidden
|
||||
return residual + hidden.transpose(1, 2)
|
||||
|
||||
|
||||
class VocosBackbone(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
input_channels: int,
|
||||
dim: int,
|
||||
intermediate_dim: int,
|
||||
num_layers: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.embed = nn.Conv1d(
|
||||
input_channels,
|
||||
dim,
|
||||
kernel_size=7,
|
||||
padding=3,
|
||||
)
|
||||
self.norm = nn.LayerNorm(dim, eps=1e-6)
|
||||
layer_scale = 1 / num_layers
|
||||
self.convnext = nn.ModuleList(
|
||||
[
|
||||
ConvNeXtBlock(
|
||||
dim=dim,
|
||||
intermediate_dim=intermediate_dim,
|
||||
layer_scale_init_value=layer_scale,
|
||||
)
|
||||
for _ in range(num_layers)
|
||||
]
|
||||
)
|
||||
self.final_layer_norm = nn.LayerNorm(dim, eps=1e-6)
|
||||
|
||||
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
|
||||
hidden = self.embed(inputs)
|
||||
hidden = self.norm(hidden.transpose(1, 2)).transpose(1, 2)
|
||||
for block in self.convnext:
|
||||
hidden = block(hidden)
|
||||
return self.final_layer_norm(hidden.transpose(1, 2))
|
||||
@@ -5,6 +5,9 @@ import copy
|
||||
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
INDEXTTS25_MIN_DURATION_FACTOR = 0.5
|
||||
INDEXTTS25_MAX_DURATION_FACTOR = 2.0
|
||||
|
||||
|
||||
class IndexTTS2Config(PretrainedConfig):
|
||||
model_type = "indextts2"
|
||||
@@ -138,6 +141,10 @@ class IndexTTS2Config(PretrainedConfig):
|
||||
"emo_num": [3, 17, 2, 8, 4, 5, 10, 24],
|
||||
"qwen_emo_path": "qwen0.6bemo4-merge/",
|
||||
"vocoder": {"type": "bigvgan", "name": "nvidia/bigvgan_v2_22khz_80band_256x"},
|
||||
"tokenizer_file": "bpe.model",
|
||||
"semantic_codec_checkpoint": "semantic_codec.pth",
|
||||
"semantic_codec_type": "repcodec",
|
||||
"use_gpt_latent": True,
|
||||
"version": 2.0,
|
||||
# Default-off performance tracing for IndexTTS2 overhead analysis.
|
||||
}
|
||||
@@ -148,9 +155,49 @@ class IndexTTS2Config(PretrainedConfig):
|
||||
for attr, default in defaults.items():
|
||||
if not hasattr(self, attr):
|
||||
setattr(self, attr, default)
|
||||
# vLLM kernel warmup reads the top-level vocabulary size even when
|
||||
# tokenizer initialization is skipped for the audio-code model.
|
||||
self.vocab_size = int(getattr(self, "vocab_size", 0) or self.gpt["number_mel_codes"])
|
||||
# Keep gpt.model_dim in sync with top-level hidden_size.
|
||||
if isinstance(self.gpt, dict) and self.gpt.get("model_dim") != self.hidden_size:
|
||||
self.gpt["model_dim"] = self.hidden_size
|
||||
self.output_sample_rate = int(
|
||||
self.s2mel["preprocess_params"].get("sr", 22050) if isinstance(self.s2mel, dict) else 22050
|
||||
)
|
||||
|
||||
|
||||
class IndexTTS25Config(IndexTTS2Config):
|
||||
"""Configuration for the official IndexTTS 2.5 checkpoint.
|
||||
|
||||
IndexTTS 2.5 keeps the GPT/S2Mel backbone shapes used by IndexTTS 2,
|
||||
but switches text tokenization and speaker conditioning, and decodes
|
||||
semantic codes with EnhancedCodec. GPT latent conditioning is disabled by
|
||||
default upstream; the opt-in path is an experimental vLLM-Omni-specific
|
||||
latent variant with no official runnable reference output.
|
||||
"""
|
||||
|
||||
model_type = "indextts2_5"
|
||||
_DEFAULTS = copy.deepcopy(IndexTTS2Config._DEFAULTS)
|
||||
_DEFAULTS["gpt"] = {
|
||||
**_DEFAULTS["gpt"],
|
||||
"number_text_tokens": 60509,
|
||||
}
|
||||
_DEFAULTS.update(
|
||||
{
|
||||
"tokenizer_file": "multilingual_zh_ja_yue_char_del.tiktoken",
|
||||
"semantic_codec_checkpoint": "codec.pth",
|
||||
"semantic_codec_type": "enhanced",
|
||||
"use_gpt_latent": False,
|
||||
"version": 2.5,
|
||||
}
|
||||
)
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
# HF overrides are sometimes intentionally partial. Merge nested GPT
|
||||
# values with the 2.5 defaults so dimensions needed by the shared
|
||||
# implementation remain present while explicit values still win.
|
||||
explicit_gpt = kwargs.pop("gpt", None)
|
||||
merged_gpt = copy.deepcopy(self._DEFAULTS["gpt"])
|
||||
if explicit_gpt is not None:
|
||||
merged_gpt.update(explicit_gpt)
|
||||
super().__init__(gpt=merged_gpt, **kwargs)
|
||||
@@ -51,6 +51,41 @@ class CUDAGraphDiTRunner:
|
||||
self.enabled = bool(enabled)
|
||||
self._cache: OrderedDict[tuple, _GraphEntry] = OrderedDict()
|
||||
self.last_call_info: dict[str, object] = {}
|
||||
self._stats = {
|
||||
"calls": 0,
|
||||
"hits": 0,
|
||||
"captures": 0,
|
||||
"capture_failures": 0,
|
||||
"evictions": 0,
|
||||
"eager": 0,
|
||||
}
|
||||
|
||||
def stats_snapshot(self) -> dict[str, object]:
|
||||
"""Return bounded cumulative graph-cache telemetry."""
|
||||
return {
|
||||
**self._stats,
|
||||
"cache_size": len(self._cache),
|
||||
}
|
||||
|
||||
def _replay(
|
||||
self,
|
||||
entry: _GraphEntry,
|
||||
*,
|
||||
x_in: torch.Tensor,
|
||||
c: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Replay with only the inputs that vary between Euler steps.
|
||||
|
||||
``input_pos`` is the same arange slice for a given cache key and the
|
||||
runner is enabled only for all-True masks, so their capture-time
|
||||
buffers remain valid. The returned output is also the graph's static
|
||||
buffer: all downstream consumers are enqueued on the same CUDA stream
|
||||
before the next replay can overwrite it.
|
||||
"""
|
||||
entry.static_x.copy_(x_in)
|
||||
entry.static_c.copy_(c)
|
||||
entry.graph.replay()
|
||||
return entry.static_out
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
@@ -59,7 +94,9 @@ class CUDAGraphDiTRunner:
|
||||
input_pos: torch.Tensor,
|
||||
mask: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
self._stats["calls"] += 1
|
||||
if not self.enabled or x_in.device.type != "cuda" or torch.cuda.is_current_stream_capturing():
|
||||
self._stats["eager"] += 1
|
||||
self.last_call_info = {
|
||||
"mode": "eager",
|
||||
"reason": "ineligible",
|
||||
@@ -84,6 +121,8 @@ class CUDAGraphDiTRunner:
|
||||
if entry is None:
|
||||
entry = self._capture(x_in, c, input_pos, mask)
|
||||
if entry is None:
|
||||
self._stats["capture_failures"] += 1
|
||||
self._stats["eager"] += 1
|
||||
self.last_call_info = {
|
||||
"mode": "eager",
|
||||
"reason": "capture_failed",
|
||||
@@ -91,10 +130,12 @@ class CUDAGraphDiTRunner:
|
||||
"cache_size": len(self._cache),
|
||||
}
|
||||
return self.transformer(x_in, c, input_pos, mask)
|
||||
self._stats["captures"] += 1
|
||||
self._cache[key] = entry
|
||||
while len(self._cache) > self.max_graphs:
|
||||
_, evicted = self._cache.popitem(last=False)
|
||||
del evicted
|
||||
self._stats["evictions"] += 1
|
||||
# Kernels are only recorded (not executed) during capture, so
|
||||
# replay once to produce the first real output.
|
||||
entry.graph.replay()
|
||||
@@ -104,22 +145,20 @@ class CUDAGraphDiTRunner:
|
||||
"shape": tuple(x_in.shape),
|
||||
"cache_size": len(self._cache),
|
||||
}
|
||||
return entry.static_out.clone()
|
||||
return entry.static_out
|
||||
|
||||
self._cache.move_to_end(key)
|
||||
entry.static_x.copy_(x_in)
|
||||
entry.static_c.copy_(c)
|
||||
entry.static_pos.copy_(input_pos)
|
||||
if entry.static_mask is not None and mask is not None:
|
||||
entry.static_mask.copy_(mask)
|
||||
entry.graph.replay()
|
||||
self._stats["hits"] += 1
|
||||
output = self._replay(entry, x_in=x_in, c=c)
|
||||
self.last_call_info = {
|
||||
"mode": "graph",
|
||||
"reason": "hit",
|
||||
"shape": tuple(x_in.shape),
|
||||
"cache_size": len(self._cache),
|
||||
}
|
||||
return entry.static_out.clone()
|
||||
if self._stats["calls"] % 1000 == 0:
|
||||
logger.info("S2Mel DiT CUDA graph stats: %s", self.stats_snapshot())
|
||||
return output
|
||||
|
||||
def _capture(
|
||||
self,
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -12,6 +12,7 @@ import json
|
||||
import math
|
||||
import re
|
||||
from collections.abc import Iterable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
@@ -38,7 +39,11 @@ 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 .configuration_indextts2 import IndexTTS2Config
|
||||
from .configuration_indextts2 import (
|
||||
INDEXTTS25_MAX_DURATION_FACTOR,
|
||||
INDEXTTS25_MIN_DURATION_FACTOR,
|
||||
IndexTTS2Config,
|
||||
)
|
||||
from .gpt.conformer_encoder import ConformerEncoder
|
||||
from .gpt.embeddings import LearnedPositionEmbeddings
|
||||
from .gpt.perceiver import PerceiverResampler
|
||||
@@ -56,6 +61,57 @@ from .preprocess_utils import (
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class IndexTTSConditioningPolicy:
|
||||
prefix_tokens: int
|
||||
uses_conformer_perceiver: bool
|
||||
uses_campplus_projection: bool
|
||||
uses_language_embedding: bool
|
||||
|
||||
|
||||
def resolve_indextts_conditioning_policy(model_type: str) -> IndexTTSConditioningPolicy:
|
||||
if model_type == "indextts2_5":
|
||||
return IndexTTSConditioningPolicy(
|
||||
prefix_tokens=3,
|
||||
uses_conformer_perceiver=False,
|
||||
uses_campplus_projection=True,
|
||||
uses_language_embedding=True,
|
||||
)
|
||||
return IndexTTSConditioningPolicy(
|
||||
prefix_tokens=34,
|
||||
uses_conformer_perceiver=True,
|
||||
uses_campplus_projection=False,
|
||||
uses_language_embedding=False,
|
||||
)
|
||||
|
||||
|
||||
def _normalize_indextts_checkpoint_key(name: str) -> str:
|
||||
"""Remove FSDP wrapper path components from native checkpoint keys."""
|
||||
return name.replace("_fsdp_wrapped_module.", "")
|
||||
|
||||
|
||||
def _require_exact_indextts_prefill_span(
|
||||
*,
|
||||
actual_span_len: int,
|
||||
scheduled_span_len: int,
|
||||
model_type: str,
|
||||
lang: str,
|
||||
text_token_count: int,
|
||||
) -> None:
|
||||
"""Reject prompt-sizing drift before it can move/remove ``start_mel``."""
|
||||
if actual_span_len == scheduled_span_len:
|
||||
return
|
||||
raise ValueError(
|
||||
"IndexTTS prefill embedding span mismatch: "
|
||||
f"scheduled_span_len={scheduled_span_len}, "
|
||||
f"actual_span_len={actual_span_len}, "
|
||||
f"model_type={model_type}, lang={lang}, "
|
||||
f"text_token_count={text_token_count}. "
|
||||
"IndexTTS requires exact placeholder sizing and does not support "
|
||||
"chunked prefill; keep enable_chunked_prefill=false."
|
||||
)
|
||||
|
||||
|
||||
def _find_most_similar_cosine(query: torch.Tensor, matrix: torch.Tensor) -> int:
|
||||
sims = F.cosine_similarity(query.float(), matrix.float(), dim=1)
|
||||
return int(torch.argmax(sims).item())
|
||||
@@ -82,6 +138,8 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
self.model_path = vllm_config.model_config.model
|
||||
self.config: IndexTTS2Config = vllm_config.model_config.hf_config # type: ignore[assignment]
|
||||
gpt_cfg = self.config.gpt
|
||||
self.conditioning_policy = resolve_indextts_conditioning_policy(self.config.model_type)
|
||||
self.use_gpt_latent = bool(getattr(self.config, "use_gpt_latent", True))
|
||||
|
||||
self.model_dim = gpt_cfg["model_dim"]
|
||||
self.num_layers = gpt_cfg["layers"]
|
||||
@@ -100,26 +158,37 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
self.has_postprocess = True
|
||||
self.requires_raw_input_tokens = True
|
||||
self.enable_update_additional_information = True
|
||||
# S2Mel consumes the complete code sequence only at request end, but
|
||||
# the connector accumulator needs to observe each appendable AR delta.
|
||||
# Keep runner request-end gating disabled and emit per-step rows from
|
||||
# make_omni_output; the full-payload builder reconstructs the sequence.
|
||||
self.omni_payload_at_request_end = False
|
||||
# Stage 1 consumes the hidden-state stream only in the legacy
|
||||
# ``use_gpt_latent`` mode. IndexTTS 2.5 transports mel codes plus
|
||||
# reference conditioning, so asking the generic runner to include
|
||||
# hidden states would add a synchronous [num_tokens, model_dim] D2H
|
||||
# copy on every decode step without adding anything to the payload.
|
||||
self.omni_pooler_payload_include_hidden = self.use_gpt_latent
|
||||
# S2Mel consumes the complete sequence only at request end. In latent
|
||||
# mode the runner retains aligned code/latent deltas on GPU; no-latent
|
||||
# mode reuses the CPU output-token history and retains only conditioning
|
||||
# metadata here.
|
||||
self.omni_payload_at_request_end = self.config.model_type == "indextts2_5"
|
||||
self.omni_request_end_token_ids = ()
|
||||
self._speaker_cache = get_speaker_cache()
|
||||
self.gpu_resident_buffer_keys: set[tuple[str, str]] = {
|
||||
("codes", "mel"),
|
||||
("hidden_states", "latent"),
|
||||
("meta", "mel_start_offset"),
|
||||
("meta", "latent_acc"),
|
||||
("meta", "mel_code_count"),
|
||||
}
|
||||
if self.use_gpt_latent:
|
||||
self.gpu_resident_buffer_keys.update(
|
||||
{
|
||||
("hidden_states", "latent"),
|
||||
("meta", "latent_acc"),
|
||||
}
|
||||
)
|
||||
|
||||
# --- GPT-2 transformer (vLLM-native with PagedAttention) ---
|
||||
# Build a GPT2Config-like object for GPT2Block
|
||||
from transformers import GPT2Config as HFGpt2Config
|
||||
|
||||
max_seq_len = self.max_mel_tokens + self.max_text_tokens + self.condition_num_latent + 4
|
||||
max_seq_len = self.max_mel_tokens + self.max_text_tokens + self.conditioning_policy.prefix_tokens + 3
|
||||
gpt2_config = HFGpt2Config(
|
||||
vocab_size=self.number_mel_codes,
|
||||
n_positions=max_seq_len,
|
||||
@@ -149,7 +218,13 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
self.mel_embedding = nn.Embedding(self.number_mel_codes, self.model_dim)
|
||||
self.text_pos_embedding = LearnedPositionEmbeddings(self.max_text_tokens + 2, self.model_dim)
|
||||
self.mel_pos_embedding = LearnedPositionEmbeddings(self.max_mel_tokens + 2 + 1, self.model_dim)
|
||||
self.speed_emb = nn.Embedding(2, self.model_dim)
|
||||
if self.conditioning_policy.uses_campplus_projection:
|
||||
from .tokenizer_v2_5 import LANGUAGE_DICT
|
||||
|
||||
self.spk_emb_proj = nn.Linear(192, self.model_dim)
|
||||
self.lang_embedding = nn.Embedding(len(LANGUAGE_DICT) + 1, self.model_dim)
|
||||
else:
|
||||
self.speed_emb = nn.Embedding(2, self.model_dim)
|
||||
self.emo_layer = nn.Linear(self.model_dim, self.model_dim)
|
||||
self.emovec_layer = nn.Linear(1024, self.model_dim)
|
||||
|
||||
@@ -158,6 +233,7 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
self.mel_head = ParallelLMHead(
|
||||
self.number_mel_codes,
|
||||
self.model_dim,
|
||||
bias=True,
|
||||
quant_config=quant_config,
|
||||
prefix=maybe_prefix(prefix, "mel_head"),
|
||||
)
|
||||
@@ -169,21 +245,22 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
cond_cfg = gpt_cfg.get("condition_module", {})
|
||||
emo_cond_cfg = gpt_cfg.get("emo_condition_module", {})
|
||||
|
||||
self.conditioning_encoder = ConformerEncoder(
|
||||
input_size=1024,
|
||||
output_size=cond_cfg.get("output_size", 512),
|
||||
linear_units=cond_cfg.get("linear_units", 2048),
|
||||
attention_heads=cond_cfg.get("attention_heads", 8),
|
||||
num_blocks=cond_cfg.get("num_blocks", 6),
|
||||
input_layer=cond_cfg.get("input_layer", "conv2d2"),
|
||||
)
|
||||
self.perceiver_encoder = PerceiverResampler(
|
||||
self.model_dim,
|
||||
dim_context=cond_cfg.get("output_size", 512),
|
||||
ff_mult=cond_cfg.get("perceiver_mult", 2),
|
||||
heads=cond_cfg.get("attention_heads", 8),
|
||||
num_latents=self.condition_num_latent,
|
||||
)
|
||||
if self.conditioning_policy.uses_conformer_perceiver:
|
||||
self.conditioning_encoder = ConformerEncoder(
|
||||
input_size=1024,
|
||||
output_size=cond_cfg.get("output_size", 512),
|
||||
linear_units=cond_cfg.get("linear_units", 2048),
|
||||
attention_heads=cond_cfg.get("attention_heads", 8),
|
||||
num_blocks=cond_cfg.get("num_blocks", 6),
|
||||
input_layer=cond_cfg.get("input_layer", "conv2d2"),
|
||||
)
|
||||
self.perceiver_encoder = PerceiverResampler(
|
||||
self.model_dim,
|
||||
dim_context=cond_cfg.get("output_size", 512),
|
||||
ff_mult=cond_cfg.get("perceiver_mult", 2),
|
||||
heads=cond_cfg.get("attention_heads", 8),
|
||||
num_latents=self.condition_num_latent,
|
||||
)
|
||||
|
||||
self.emo_conditioning_encoder = ConformerEncoder(
|
||||
input_size=1024,
|
||||
@@ -202,7 +279,8 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
)
|
||||
|
||||
# Padding masks for perceiver cross-attention
|
||||
self.cond_mask_pad = nn.ConstantPad1d((self.condition_num_latent, 0), True)
|
||||
if self.conditioning_policy.uses_conformer_perceiver:
|
||||
self.cond_mask_pad = nn.ConstantPad1d((self.condition_num_latent, 0), True)
|
||||
self.emo_cond_mask_pad = nn.ConstantPad1d((1, 0), True)
|
||||
|
||||
# --- Lazy-loaded external models ---
|
||||
@@ -214,7 +292,8 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
# Initialize embeddings per GPT-2 convention
|
||||
for emb in [self.text_embedding, self.mel_embedding]:
|
||||
emb.weight.data.normal_(mean=0.0, std=0.02)
|
||||
self.speed_emb.weight.data.normal_(mean=0.0, std=0.0)
|
||||
if hasattr(self, "speed_emb"):
|
||||
self.speed_emb.weight.data.normal_(mean=0.0, std=0.0)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# vLLM required hooks
|
||||
@@ -283,14 +362,23 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
hidden_states = hidden_states.text_hidden_states
|
||||
if hidden_states is None or self.mel_head is None:
|
||||
return None
|
||||
return self.logits_processor(self.mel_head, hidden_states)
|
||||
# `LogitsProcessor` never reads `lm_head.bias`: the unquantized apply()
|
||||
# path is `F.linear(x, layer.weight, bias)` where `bias` comes from the
|
||||
# `embedding_bias` argument. IndexTTS's mel head is `Linear(..., bias=True)`
|
||||
# upstream, so the bias has to be passed explicitly or logits silently
|
||||
# drop it.
|
||||
return self.logits_processor(
|
||||
self.mel_head,
|
||||
hidden_states,
|
||||
embedding_bias=getattr(self.mel_head, "bias", None),
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Omni multimodal output plumbing
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput:
|
||||
"""Collect mel_codes and hidden_states from intermediate buffer."""
|
||||
"""Build the per-step Stage 1 payload from the intermediate buffer."""
|
||||
if isinstance(model_outputs, OmniOutput):
|
||||
return model_outputs
|
||||
|
||||
@@ -303,58 +391,71 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
|
||||
device = hidden.device
|
||||
latent_dim = hidden.shape[-1]
|
||||
_zero_codes = torch.zeros(0, 1, dtype=torch.long, device=device)
|
||||
_zero_latent = torch.zeros(0, latent_dim, dtype=hidden.dtype, device=device)
|
||||
_zero_codes = torch.zeros(0, 1, dtype=torch.long, device=device) if self.use_gpt_latent else None
|
||||
_zero_latent = torch.zeros(0, latent_dim, dtype=hidden.dtype, device=device) if self.use_gpt_latent else None
|
||||
|
||||
# Per-request lists: to_payload_element dispatches lists via
|
||||
# element[idx], giving correct per-request mapping even during
|
||||
# mixed prefill+decode steps. Zero-length placeholders for prefill
|
||||
# requests are harmless — torch.cat ignores them in accumulation.
|
||||
#
|
||||
# Export semantics are appendable deltas for the existing full-payload
|
||||
# accumulator: codes.mel is a 2-D row [1, 1], hidden_states.latent is
|
||||
# the aligned one-row [1, D] slice from meta.latent_acc.
|
||||
# Latent-mode exports are appendable deltas for the existing
|
||||
# full-payload accumulator: codes.mel is a 2-D row [1, 1], while
|
||||
# hidden_states.latent is the aligned one-row [1, D] slice from
|
||||
# meta.latent_acc. No-latent mode only uses meta_list.
|
||||
mel_codes_list: list[torch.Tensor] = []
|
||||
latent_list: list[torch.Tensor] = []
|
||||
meta_list: list[dict[str, Any]] = []
|
||||
|
||||
for info in info_dicts:
|
||||
if not isinstance(info, dict):
|
||||
mel_codes_list.append(_zero_codes)
|
||||
latent_list.append(_zero_latent)
|
||||
if self.use_gpt_latent:
|
||||
assert _zero_codes is not None
|
||||
mel_codes_list.append(_zero_codes)
|
||||
assert _zero_latent is not None
|
||||
latent_list.append(_zero_latent)
|
||||
meta_list.append({})
|
||||
continue
|
||||
codes = info.get("codes", {})
|
||||
mc = codes.get("mel")
|
||||
info_meta = info.get("meta", {})
|
||||
# mel_len now comes from the int counter in meta, not from
|
||||
# code_seq.shape[0] (codes.mel is a single-token delta).
|
||||
mel_len = int(info_meta.get("mel_code_count", 0))
|
||||
if isinstance(mc, torch.Tensor) and mc.numel() > 0 and mel_len > 0:
|
||||
code_seq = mc.reshape(-1)
|
||||
mel_codes_list.append(code_seq[-1:].reshape(1, 1).contiguous())
|
||||
else:
|
||||
mel_codes_list.append(_zero_codes)
|
||||
|
||||
lat_acc = info_meta.get("latent_acc")
|
||||
if isinstance(lat_acc, torch.Tensor) and lat_acc.numel() > 0 and mel_len > 0:
|
||||
lat_seq = lat_acc.reshape(-1, lat_acc.shape[-1]) if lat_acc.ndim != 2 else lat_acc
|
||||
idx = min(max(mel_len - 1, 0), int(lat_seq.shape[0]) - 1)
|
||||
latent_list.append(lat_seq[idx : idx + 1].contiguous())
|
||||
else:
|
||||
hs = info.get("hidden_states", {})
|
||||
lat = hs.get("latent")
|
||||
if isinstance(lat, torch.Tensor) and lat.numel() > 0:
|
||||
lat_seq = lat.reshape(-1, lat.shape[-1]) if lat.ndim != 2 else lat
|
||||
latent_list.append(lat_seq[-1:].contiguous())
|
||||
if self.use_gpt_latent:
|
||||
codes = info.get("codes", {})
|
||||
mc = codes.get("mel")
|
||||
if isinstance(mc, torch.Tensor) and mc.numel() > 0 and mel_len > 0:
|
||||
code_seq = mc.reshape(-1)
|
||||
mel_codes_list.append(code_seq[-1:].reshape(1, 1).contiguous())
|
||||
else:
|
||||
latent_list.append(_zero_latent)
|
||||
assert _zero_codes is not None
|
||||
mel_codes_list.append(_zero_codes)
|
||||
|
||||
req_meta: dict[str, Any] = {}
|
||||
lat_acc = info_meta.get("latent_acc")
|
||||
if isinstance(lat_acc, torch.Tensor) and lat_acc.numel() > 0 and mel_len > 0:
|
||||
lat_seq = lat_acc.reshape(-1, lat_acc.shape[-1]) if lat_acc.ndim != 2 else lat_acc
|
||||
idx = min(max(mel_len - 1, 0), int(lat_seq.shape[0]) - 1)
|
||||
latent_list.append(lat_seq[idx : idx + 1].contiguous())
|
||||
else:
|
||||
hs = info.get("hidden_states", {})
|
||||
lat = hs.get("latent")
|
||||
if isinstance(lat, torch.Tensor) and lat.numel() > 0:
|
||||
lat_seq = lat.reshape(-1, lat.shape[-1]) if lat.ndim != 2 else lat
|
||||
latent_list.append(lat_seq[-1:].contiguous())
|
||||
else:
|
||||
assert _zero_latent is not None
|
||||
latent_list.append(_zero_latent)
|
||||
|
||||
# Scalar metadata is kept with REPLACE semantics by the generic
|
||||
# full-payload accumulator. Emit the policy on every step so the
|
||||
# final payload contract never depends on which decode row carried
|
||||
# the one-time conditioning tensors.
|
||||
req_meta: dict[str, Any] = {"use_gpt_latent": self.use_gpt_latent}
|
||||
if mel_len == 1:
|
||||
s_ref = info_meta.get("S_ref")
|
||||
ref_mel = info_meta.get("ref_mel")
|
||||
style = info_meta.get("style")
|
||||
req_meta["duration_factor"] = info_meta.get("duration_factor", 1.0)
|
||||
if isinstance(s_ref, torch.Tensor):
|
||||
req_meta["S_ref"] = s_ref
|
||||
if isinstance(ref_mel, torch.Tensor):
|
||||
@@ -363,13 +464,22 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
req_meta["style"] = style
|
||||
meta_list.append(req_meta)
|
||||
|
||||
if not any(t.numel() > 0 for t in mel_codes_list):
|
||||
return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
|
||||
|
||||
mm: OmniPayload = {
|
||||
"codes": {"mel": mel_codes_list},
|
||||
"hidden_states": {"latent": latent_list},
|
||||
}
|
||||
if self.use_gpt_latent:
|
||||
if not any(t.numel() > 0 for t in mel_codes_list):
|
||||
return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
|
||||
mm: OmniPayload = {
|
||||
"codes": {"mel": mel_codes_list},
|
||||
}
|
||||
mm["hidden_states"] = {"latent": latent_list}
|
||||
else:
|
||||
# vLLM already materializes sampled token IDs on CPU for scheduler
|
||||
# bookkeeping and stop-token handling. Stage 1 consumes the complete
|
||||
# sequence only after request completion, so retain the conditioning
|
||||
# metadata once and read the finished request.output_token_ids there
|
||||
# instead of snapshotting a duplicate GPU code row every step.
|
||||
if not any(meta_list):
|
||||
return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
|
||||
mm = {}
|
||||
if any(meta_list):
|
||||
mm["meta"] = meta_list
|
||||
|
||||
@@ -387,9 +497,8 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
|
||||
"""Build prompt embeddings for prefill; compute mel embeddings for decode.
|
||||
|
||||
Prefill layout::
|
||||
|
||||
[conds(32) + emo_vec(1) + duration(2)] [text_emb + text_pos] [start_mel + mel_pos(0)]
|
||||
IndexTTS 2 prefill uses 32 speaker latents plus two duration tokens.
|
||||
IndexTTS 2.5 uses one CAMPPlus speaker token plus two zero tokens.
|
||||
|
||||
Decode: ``mel_embedding(token) + mel_pos_embedding(step)``.
|
||||
"""
|
||||
@@ -453,6 +562,12 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
v = info_dict.get(key)
|
||||
return v[0] if isinstance(v, list) and v else default
|
||||
|
||||
duration_factor = float(_first("duration_factor", 1.0))
|
||||
# Direct Omni callers bypass the serving adapter, so retain the
|
||||
# model-side guard as the final fail-fast validation layer.
|
||||
if not (INDEXTTS25_MIN_DURATION_FACTOR <= duration_factor <= INDEXTTS25_MAX_DURATION_FACTOR):
|
||||
raise ValueError("IndexTTS 2.5 duration_factor must be between 0.5 and 2.0")
|
||||
|
||||
text = _first("text")
|
||||
if not text:
|
||||
raise ValueError("Missing additional_information.text for IndexTTS2 talker.")
|
||||
@@ -467,6 +582,8 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
emo_vector = _first("emo_vector")
|
||||
emo_alpha = float(_first("emo_alpha", 1.0))
|
||||
use_random = bool(_first("use_random", False))
|
||||
lang = str(_first("lang", "zh"))
|
||||
text_normalization = bool(_first("text_normalization", True))
|
||||
_raw_emo_voice = _first("emo_voice_name")
|
||||
emo_voice_name = str(_raw_emo_voice).strip().lower() if _raw_emo_voice else None
|
||||
|
||||
@@ -499,9 +616,19 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
self._ensure_w2v_stat_loaded(device)
|
||||
spk_cond_emb = wav2vec_extract(wav_16k, w2v_model, w2v_proc, device, self._w2v_stat)
|
||||
|
||||
semantic_codec = load_semantic_codec(self.model_path, self.config.semantic_codec, device)
|
||||
with torch.no_grad():
|
||||
_, s_ref = semantic_codec.quantize(spk_cond_emb) # [B, T, 1024] quantized embeddings
|
||||
if self.conditioning_policy.uses_campplus_projection:
|
||||
# Official 2.5 keeps W2V-BERT features as the reference
|
||||
# semantic sequence; EnhancedCodec is only used for generated
|
||||
# semantic codes in Stage 1.
|
||||
s_ref = spk_cond_emb
|
||||
else:
|
||||
semantic_codec = load_semantic_codec(
|
||||
self.model_path,
|
||||
self.config.semantic_codec,
|
||||
device,
|
||||
)
|
||||
with torch.no_grad():
|
||||
_, s_ref = semantic_codec.quantize(spk_cond_emb)
|
||||
|
||||
campplus = load_campplus(self.model_path, device)
|
||||
fbank = compute_fbank(wav_16k, device)
|
||||
@@ -521,12 +648,16 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
},
|
||||
)
|
||||
|
||||
model_dtype = next(self.conditioning_encoder.parameters()).dtype
|
||||
if self.conditioning_policy.uses_campplus_projection:
|
||||
model_dtype = self.spk_emb_proj.weight.dtype
|
||||
else:
|
||||
model_dtype = next(self.conditioning_encoder.parameters()).dtype
|
||||
spk_cond_emb = spk_cond_emb.to(device=device, dtype=model_dtype)
|
||||
spk_lens = torch.tensor([spk_cond_emb.shape[1]], device=device, dtype=torch.long)
|
||||
speech_cond, mask = self.conditioning_encoder(spk_cond_emb, spk_lens)
|
||||
conds_mask = self.cond_mask_pad(mask.squeeze(1))
|
||||
conds = self.perceiver_encoder(speech_cond, conds_mask) # [1, 32, D]
|
||||
if self.conditioning_policy.uses_conformer_perceiver:
|
||||
spk_lens = torch.tensor([spk_cond_emb.shape[1]], device=device, dtype=torch.long)
|
||||
speech_cond, mask = self.conditioning_encoder(spk_cond_emb, spk_lens)
|
||||
conds_mask = self.cond_mask_pad(mask.squeeze(1))
|
||||
conds = self.perceiver_encoder(speech_cond, conds_mask)
|
||||
|
||||
emo_vec = self._compute_emotion_vector(
|
||||
wav_16k=wav_16k,
|
||||
@@ -546,19 +677,39 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
device=device,
|
||||
)
|
||||
|
||||
speed_zero = torch.zeros(1, device=device, dtype=torch.long)
|
||||
speed_one = torch.ones(1, device=device, dtype=torch.long)
|
||||
duration_emb = self.speed_emb(speed_zero)
|
||||
duration_emb_half = self.speed_emb(speed_one)
|
||||
if self.conditioning_policy.uses_campplus_projection:
|
||||
speaker_token = self.spk_emb_proj(style.to(dtype=model_dtype)).unsqueeze(1)
|
||||
zero_tokens = torch.zeros(
|
||||
speaker_token.shape[0],
|
||||
2,
|
||||
self.model_dim,
|
||||
device=device,
|
||||
dtype=speaker_token.dtype,
|
||||
)
|
||||
conds_prefix = torch.cat(
|
||||
[speaker_token + emo_vec.unsqueeze(1), zero_tokens],
|
||||
dim=1,
|
||||
)
|
||||
else:
|
||||
speed_zero = torch.zeros(1, device=device, dtype=torch.long)
|
||||
speed_one = torch.ones(1, device=device, dtype=torch.long)
|
||||
duration_emb = self.speed_emb(speed_zero)
|
||||
duration_emb_half = self.speed_emb(speed_one)
|
||||
conds_with_emo = conds + emo_vec.unsqueeze(1)
|
||||
conds_prefix = torch.cat(
|
||||
[conds_with_emo, duration_emb_half.unsqueeze(1), duration_emb.unsqueeze(1)],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
conds_with_emo = conds + emo_vec.unsqueeze(1) # [1, 32, D]
|
||||
conds_prefix = torch.cat(
|
||||
[conds_with_emo, duration_emb_half.unsqueeze(1), duration_emb.unsqueeze(1)],
|
||||
dim=1,
|
||||
) # [1, 34, D]
|
||||
|
||||
text_tokens = self._tokenize_text(text, device) # [1, L]
|
||||
text_tokens, lang_id = self._tokenize_text(
|
||||
text,
|
||||
device,
|
||||
lang=lang,
|
||||
text_normalization=text_normalization,
|
||||
)
|
||||
text_emb = self.text_embedding(text_tokens) + self.text_pos_embedding(text_tokens)
|
||||
if lang_id is not None:
|
||||
text_emb = text_emb + self.lang_embedding(lang_id).unsqueeze(1)
|
||||
|
||||
start_mel = self.mel_embedding(torch.tensor([[self.start_mel_token]], device=device, dtype=torch.long))
|
||||
start_mel_pos = self.mel_pos_embedding.get_fixed_embedding(0, device)
|
||||
@@ -567,13 +718,13 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
inputs_embeds = torch.cat([conds_prefix, text_emb, start_mel], dim=1).squeeze(0)
|
||||
mel_start_offset = int(conds_prefix.shape[1]) + int(text_emb.shape[1])
|
||||
prefill_mel_index = int(inputs_embeds.shape[0]) - 1
|
||||
|
||||
if int(inputs_embeds.shape[0]) < span_len:
|
||||
pad_n = span_len - int(inputs_embeds.shape[0])
|
||||
pad_emb = torch.zeros(pad_n, self.model_dim, device=device, dtype=inputs_embeds.dtype)
|
||||
inputs_embeds = torch.cat([inputs_embeds, pad_emb], dim=0)
|
||||
elif int(inputs_embeds.shape[0]) > span_len:
|
||||
inputs_embeds = inputs_embeds[:span_len]
|
||||
_require_exact_indextts_prefill_span(
|
||||
actual_span_len=int(inputs_embeds.shape[0]),
|
||||
scheduled_span_len=span_len,
|
||||
model_type=str(getattr(self.config, "model_type", "indextts2")),
|
||||
lang=lang,
|
||||
text_token_count=int(text_tokens.shape[1]),
|
||||
)
|
||||
|
||||
input_ids_out = torch.full_like(input_ids, self.start_mel_token)
|
||||
|
||||
@@ -585,11 +736,26 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
"S_ref": s_ref.cpu().contiguous(),
|
||||
"ref_mel": ref_mel.cpu().contiguous(),
|
||||
"style": style.cpu().contiguous(),
|
||||
"latent_acc": torch.zeros(0, self.model_dim, dtype=inputs_embeds.dtype, device=device),
|
||||
"use_gpt_latent": self.use_gpt_latent,
|
||||
"duration_factor": duration_factor,
|
||||
},
|
||||
"codes": {"mel": torch.zeros(0, dtype=torch.long, device=device)},
|
||||
"hidden_states": {"latent": torch.zeros(0, self.model_dim, dtype=inputs_embeds.dtype, device=device)},
|
||||
}
|
||||
if self.use_gpt_latent:
|
||||
info_update["meta"]["latent_acc"] = torch.zeros(
|
||||
0,
|
||||
self.model_dim,
|
||||
dtype=inputs_embeds.dtype,
|
||||
device=device,
|
||||
)
|
||||
info_update["hidden_states"] = {
|
||||
"latent": torch.zeros(
|
||||
0,
|
||||
self.model_dim,
|
||||
dtype=inputs_embeds.dtype,
|
||||
device=device,
|
||||
)
|
||||
}
|
||||
|
||||
return input_ids_out, inputs_embeds, info_update
|
||||
|
||||
@@ -605,6 +771,8 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
reconstructs the complete latent sequence from per-step deltas emitted
|
||||
by make_omni_output.
|
||||
"""
|
||||
if not self.use_gpt_latent:
|
||||
return {}
|
||||
if hidden_states.shape[0] != 1:
|
||||
meta = kwargs.get("meta", {})
|
||||
prefill_mel_index = int(meta.get("prefill_mel_index", hidden_states.shape[0] - 1))
|
||||
@@ -665,6 +833,8 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
if emo_text is None:
|
||||
emo_text = main_text
|
||||
emo_vector = self._predict_emotion_from_text(emo_text, device)
|
||||
if emo_vector is None:
|
||||
raise RuntimeError("IndexTTS use_emo_text requested but QwenEmotion could not be loaded")
|
||||
if use_emo_text or emo_vector is not None:
|
||||
emo_audio_path = None
|
||||
|
||||
@@ -840,7 +1010,7 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
self._ensure_matrix_loaded("_emo_matrix", "emo_matrix", "Emotion matrix")
|
||||
self._ensure_matrix_loaded("_spk_matrix", "spk_matrix", "Speaker matrix")
|
||||
if self._emo_matrix is None or self._spk_matrix is None:
|
||||
return torch.zeros(1, self.model_dim, device=device)
|
||||
raise RuntimeError("IndexTTS emotion control requires both emotion and speaker matrices")
|
||||
|
||||
weight_vector = emo_dist.to(device=device, dtype=torch.float32)
|
||||
if weight_vector.ndim == 2:
|
||||
@@ -886,23 +1056,47 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
# Text tokenizer
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _tokenize_text(self, text: str, device: torch.device) -> torch.Tensor:
|
||||
"""Tokenize text using the BPE tokenizer, add start/stop tokens."""
|
||||
if self._text_tokenizer is None:
|
||||
from .tokenizer import IndexTTS2Tokenizer
|
||||
|
||||
bpe_path = resolve_model_file(self.model_path, "bpe.model")
|
||||
if bpe_path is None:
|
||||
raise FileNotFoundError(f"BPE model not found in {self.model_path}")
|
||||
self._text_tokenizer = IndexTTS2Tokenizer(bpe_path, model_dir=self.model_path)
|
||||
|
||||
token_ids = self._text_tokenizer.encode(text, add_special_tokens=False)
|
||||
|
||||
# Wrap with start/stop text tokens
|
||||
def _tokenize_text(
|
||||
self,
|
||||
text: str,
|
||||
device: torch.device,
|
||||
*,
|
||||
lang: str = "zh",
|
||||
text_normalization: bool = True,
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
"""Tokenize text and add the checkpoint start/stop text tokens."""
|
||||
start_text = 0
|
||||
stop_text = 1
|
||||
lang_id: int | None = None
|
||||
if self.conditioning_policy.uses_language_embedding:
|
||||
from .text_processing_v2_5 import prepare_indextts25_text
|
||||
|
||||
token_ids, lang_id = prepare_indextts25_text(
|
||||
text,
|
||||
lang=lang,
|
||||
model_dir=self.model_path,
|
||||
tokenizer_file=self.config.tokenizer_file,
|
||||
text_normalization=text_normalization,
|
||||
)
|
||||
token_ids = [token_id for token_id in token_ids if token_id not in {start_text, stop_text}]
|
||||
else:
|
||||
if self._text_tokenizer is None:
|
||||
from .tokenizer import IndexTTS2Tokenizer
|
||||
|
||||
bpe_path = resolve_model_file(self.model_path, "bpe.model")
|
||||
if bpe_path is None:
|
||||
raise FileNotFoundError(f"BPE model not found in {self.model_path}")
|
||||
self._text_tokenizer = IndexTTS2Tokenizer(bpe_path, model_dir=self.model_path)
|
||||
|
||||
token_ids = self._text_tokenizer.encode(text, add_special_tokens=False)
|
||||
|
||||
# The 2.5 tokenizer begins with the explicit ``<|lang|>`` token.
|
||||
# Official prepare_gpt_inputs filters any existing wrapper tokens,
|
||||
# then adds start_text=0 and stop_text=1 around that tokenizer output.
|
||||
token_ids = [start_text] + token_ids + [stop_text]
|
||||
return torch.tensor([token_ids], device=device, dtype=torch.long)
|
||||
token_tensor = torch.tensor([token_ids], device=device, dtype=torch.long)
|
||||
lang_tensor = torch.tensor([lang_id], device=device, dtype=torch.long) if lang_id is not None else None
|
||||
return token_tensor, lang_tensor
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Lazy loading helpers
|
||||
@@ -919,7 +1113,7 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
else:
|
||||
self._w2v_stat = raw
|
||||
else:
|
||||
logger.warning("Wav2Vec2-BERT stat file not found in %s", self.model_path)
|
||||
raise FileNotFoundError(f"Wav2Vec2-BERT stat file not found in {self.model_path!r}")
|
||||
|
||||
def _ensure_matrix_loaded(self, attr: str, cfg_key: str, label: str) -> None:
|
||||
if getattr(self, attr) is not None:
|
||||
@@ -929,7 +1123,7 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
raw = torch.load(path, map_location="cpu", weights_only=True)
|
||||
setattr(self, attr, torch.split(raw, list(self.config.emo_num)))
|
||||
else:
|
||||
logger.warning("%s not found in %s", label, self.model_path)
|
||||
raise FileNotFoundError(f"{label} not found in {self.model_path!r}")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Weight loading
|
||||
@@ -966,7 +1160,8 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
params_dict = dict(self.named_parameters(remove_duplicate=False))
|
||||
loaded_params: set[str] = set()
|
||||
|
||||
for name, loaded_weight in state.items():
|
||||
for checkpoint_name, loaded_weight in state.items():
|
||||
name = _normalize_indextts_checkpoint_key(checkpoint_name)
|
||||
# Skip attention masks
|
||||
if ".attn.bias" in name or ".attn.masked_bias" in name:
|
||||
continue
|
||||
@@ -989,7 +1184,11 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
continue
|
||||
|
||||
if mapped_name not in params_dict:
|
||||
logger.debug("Skipping unrecognized weight: %s → %s", name, mapped_name)
|
||||
logger.debug(
|
||||
"Skipping unrecognized weight: %s → %s",
|
||||
checkpoint_name,
|
||||
mapped_name,
|
||||
)
|
||||
continue
|
||||
|
||||
param = params_dict[mapped_name]
|
||||
@@ -1015,17 +1214,29 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
"mel_embedding.",
|
||||
"text_pos_embedding.",
|
||||
"mel_pos_embedding.",
|
||||
"speed_emb.",
|
||||
"emo_layer.",
|
||||
"emovec_layer.",
|
||||
"conditioning_encoder.",
|
||||
"perceiver_encoder.",
|
||||
"emo_conditioning_encoder.",
|
||||
"emo_perceiver_encoder.",
|
||||
"h.",
|
||||
"ln_f.",
|
||||
"final_norm.",
|
||||
]
|
||||
if self.conditioning_policy.uses_campplus_projection:
|
||||
required_prefixes.extend(
|
||||
[
|
||||
"spk_emb_proj.",
|
||||
"lang_embedding.",
|
||||
]
|
||||
)
|
||||
else:
|
||||
required_prefixes.extend(
|
||||
[
|
||||
"speed_emb.",
|
||||
"conditioning_encoder.",
|
||||
"perceiver_encoder.",
|
||||
]
|
||||
)
|
||||
if self.mel_head is not None:
|
||||
required_prefixes.append("mel_head.")
|
||||
missing_prefixes = [
|
||||
@@ -1033,18 +1244,22 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
]
|
||||
if missing_prefixes:
|
||||
raise RuntimeError(
|
||||
"IndexTTS2 GPT checkpoint did not load required parameter groups: " + ", ".join(missing_prefixes)
|
||||
f"{self.config.model_type} GPT checkpoint did not load required parameter groups: "
|
||||
+ ", ".join(missing_prefixes)
|
||||
)
|
||||
|
||||
# Ensure all sub-modules on correct device (some nn.Parameter created
|
||||
# with torch.Tensor() end up on CPU even when model is on CUDA).
|
||||
target_device = next((p.device for p in self.h.parameters()), torch.device("cpu"))
|
||||
for mod in [
|
||||
self.conditioning_encoder,
|
||||
modules = [
|
||||
self.emo_conditioning_encoder,
|
||||
self.perceiver_encoder,
|
||||
self.emo_perceiver_encoder,
|
||||
]:
|
||||
]
|
||||
if self.conditioning_policy.uses_campplus_projection:
|
||||
modules.extend([self.spk_emb_proj, self.lang_embedding])
|
||||
else:
|
||||
modules.extend([self.conditioning_encoder, self.perceiver_encoder])
|
||||
for mod in modules:
|
||||
mod.to(target_device)
|
||||
|
||||
self._warmup_external_models(target_device)
|
||||
@@ -1054,9 +1269,17 @@ class IndexTTS2TalkerForConditionalGeneration(nn.Module):
|
||||
"""Preload external models to eliminate first-request latency."""
|
||||
try:
|
||||
load_wav2vec2(self.model_path, device)
|
||||
load_semantic_codec(self.model_path, self.config.semantic_codec, device)
|
||||
if not self.conditioning_policy.uses_campplus_projection:
|
||||
load_semantic_codec(
|
||||
self.model_path,
|
||||
self.config.semantic_codec,
|
||||
device,
|
||||
)
|
||||
load_campplus(self.model_path, device)
|
||||
self._ensure_w2v_stat_loaded(device)
|
||||
logger.info("External models preloaded: Wav2Vec2, RepCodec, CAMPPlus")
|
||||
logger.info(
|
||||
"External models preloaded: Wav2Vec2, %sCAMPPlus",
|
||||
"RepCodec, " if not self.conditioning_policy.uses_campplus_projection else "",
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to preload external models: %s", e)
|
||||
raise RuntimeError("Failed to preload required IndexTTS Stage 0 models") from e
|
||||
@@ -43,7 +43,43 @@ INDEXTTS2_PIPELINE = PipelineConfig(
|
||||
engine_output_type="audio",
|
||||
model_arch="IndexTTS2S2MelDecoder",
|
||||
sync_process_input_func=f"{_PROC}.talker2s2mel_token_only",
|
||||
extras={"skip_tokenizer_init": True, "tokenizer": "gpt2"},
|
||||
extras={"skip_tokenizer_init": True},
|
||||
sampling_constraints={"detokenize": True},
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
INDEXTTS25_PIPELINE = PipelineConfig(
|
||||
model_type="indextts2_5",
|
||||
default_deploy_config_name="indextts2_5.yaml",
|
||||
model_arch="IndexTTS25TalkerForConditionalGeneration",
|
||||
stages=(
|
||||
StagePipelineConfig(
|
||||
stage_id=0,
|
||||
model_stage="indextts2_5_talker",
|
||||
execution_type=StageExecutionType.LLM_AR,
|
||||
input_sources=(),
|
||||
owns_tokenizer=True,
|
||||
engine_output_type="latent",
|
||||
extras={"skip_tokenizer_init": True},
|
||||
custom_process_next_stage_input_func=f"{_PROC}.talker2s2mel_full_payload",
|
||||
sampling_constraints={
|
||||
"detokenize": False,
|
||||
"stop_token_ids": [8193],
|
||||
},
|
||||
),
|
||||
StagePipelineConfig(
|
||||
stage_id=1,
|
||||
model_stage="indextts2_5_s2mel_decoder",
|
||||
execution_type=StageExecutionType.LLM_GENERATION,
|
||||
input_sources=(0,),
|
||||
final_output=True,
|
||||
final_output_type="audio",
|
||||
engine_output_type="audio",
|
||||
model_arch="IndexTTS25S2MelDecoder",
|
||||
sync_process_input_func=f"{_PROC}.talker2s2mel_token_only",
|
||||
extras={"skip_tokenizer_init": True},
|
||||
sampling_constraints={"detokenize": True},
|
||||
),
|
||||
),
|
||||
|
||||
@@ -22,7 +22,7 @@ logger = init_logger(__name__)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_wav2vec2_cache: dict[tuple[str, str], tuple] = {}
|
||||
_semantic_codec_cache: dict[tuple[str, str], Any] = {}
|
||||
_semantic_codec_cache: dict[tuple[str, str, str], Any] = {}
|
||||
_campplus_cache: dict[tuple[str, str], Any] = {}
|
||||
_qwen_emotion_cache: dict[tuple[str, str], tuple] = {}
|
||||
|
||||
@@ -34,11 +34,40 @@ def _freeze(model):
|
||||
return model
|
||||
|
||||
|
||||
def local_model_roots(model_path: str) -> tuple[str, ...]:
|
||||
"""Return native IndexTTS asset roots in lookup order.
|
||||
|
||||
HuggingFace-converted checkpoints place assets at the model root, while
|
||||
the official IndexTTS repository and lab-published bundle keep them under
|
||||
``checkpoints/``.
|
||||
"""
|
||||
roots = [model_path]
|
||||
checkpoints = os.path.join(model_path, "checkpoints")
|
||||
if os.path.isdir(checkpoints):
|
||||
roots.append(checkpoints)
|
||||
return tuple(roots)
|
||||
|
||||
|
||||
def resolve_model_directory(
|
||||
model_path: str,
|
||||
*relative_paths: str,
|
||||
) -> str | None:
|
||||
for root in local_model_roots(model_path):
|
||||
for relative_path in relative_paths:
|
||||
candidate = os.path.join(root, relative_path)
|
||||
if os.path.isdir(candidate):
|
||||
return candidate
|
||||
return None
|
||||
|
||||
|
||||
def resolve_model_file(model_path: str, filename: str) -> str | None:
|
||||
"""Resolve an IndexTTS2 asset from a local model dir or HF repo id."""
|
||||
local_path = os.path.join(model_path, filename)
|
||||
if os.path.isfile(local_path):
|
||||
return local_path
|
||||
for root in local_model_roots(model_path):
|
||||
local_path = os.path.join(root, filename)
|
||||
if os.path.isfile(local_path):
|
||||
return local_path
|
||||
if os.path.isdir(model_path):
|
||||
return None
|
||||
try:
|
||||
return cached_file(model_path, filename)
|
||||
except Exception as e:
|
||||
@@ -53,11 +82,38 @@ def load_wav2vec2(model_path: str, device: torch.device):
|
||||
from transformers import AutoFeatureExtractor, Wav2Vec2BertModel
|
||||
|
||||
model, processor = None, None
|
||||
# Try local subfolder first (no lock needed), then HF hub with lock.
|
||||
# Try the layouts shipped by IndexTTS 2 and 2.5 before the Hub fallback.
|
||||
local_dir = resolve_model_directory(
|
||||
model_path,
|
||||
"wav2vec2bert",
|
||||
"w2v-bert-2.0",
|
||||
os.path.join("hf_cache", "w2v-bert-2.0"),
|
||||
)
|
||||
if local_dir is not None:
|
||||
try:
|
||||
processor = AutoFeatureExtractor.from_pretrained(
|
||||
local_dir,
|
||||
local_files_only=True,
|
||||
)
|
||||
model = Wav2Vec2BertModel.from_pretrained(
|
||||
local_dir,
|
||||
local_files_only=True,
|
||||
)
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"Failed to load local IndexTTS Wav2Vec2-BERT assets from {local_dir!r}") from exc
|
||||
|
||||
if local_dir is None and os.path.isdir(model_path):
|
||||
raise FileNotFoundError(f"IndexTTS Wav2Vec2-BERT assets are missing from local bundle {model_path!r}")
|
||||
|
||||
# Try a remote model subfolder, then the public base model.
|
||||
try:
|
||||
src_args = dict(pretrained_model_name_or_path=model_path, subfolder="wav2vec2bert")
|
||||
processor = AutoFeatureExtractor.from_pretrained(**src_args)
|
||||
model = Wav2Vec2BertModel.from_pretrained(**src_args)
|
||||
if model is None or processor is None:
|
||||
src_args = dict(
|
||||
pretrained_model_name_or_path=model_path,
|
||||
subfolder="wav2vec2bert",
|
||||
)
|
||||
processor = AutoFeatureExtractor.from_pretrained(**src_args)
|
||||
model = Wav2Vec2BertModel.from_pretrained(**src_args)
|
||||
except Exception:
|
||||
try:
|
||||
with _repo_prefetch_lock("facebook/w2v-bert-2.0"):
|
||||
@@ -67,16 +123,43 @@ def load_wav2vec2(model_path: str, device: torch.device):
|
||||
except Exception:
|
||||
pass
|
||||
if model is None or processor is None:
|
||||
raise RuntimeError(f"Could not load Wav2Vec2-BERT from {model_path}/wav2vec2bert or facebook/w2v-bert-2.0")
|
||||
raise RuntimeError(
|
||||
f"Could not load Wav2Vec2-BERT from IndexTTS model assets under {model_path!r} or facebook/w2v-bert-2.0"
|
||||
)
|
||||
_freeze(model.to(device=device, dtype=torch.float32))
|
||||
_wav2vec2_cache[cache_key] = (model, processor)
|
||||
return model, processor
|
||||
|
||||
|
||||
def load_semantic_codec(model_path: str, config: dict, device: torch.device):
|
||||
cache_key = (model_path, str(device))
|
||||
def load_semantic_codec(
|
||||
model_path: str,
|
||||
config: dict,
|
||||
device: torch.device,
|
||||
*,
|
||||
codec_type: str = "repcodec",
|
||||
checkpoint_name: str | None = None,
|
||||
):
|
||||
cache_key = (model_path, str(device), codec_type)
|
||||
if cache_key in _semantic_codec_cache:
|
||||
return _semantic_codec_cache[cache_key]
|
||||
|
||||
if codec_type == "enhanced":
|
||||
from .codec import EnhancedCodec
|
||||
|
||||
codec = EnhancedCodec(**config, cfg=config)
|
||||
checkpoint_name = checkpoint_name or "codec.pth"
|
||||
ckpt_path = resolve_model_file(model_path, checkpoint_name)
|
||||
if ckpt_path is None:
|
||||
raise FileNotFoundError(
|
||||
f"IndexTTS 2.5 EnhancedCodec checkpoint {checkpoint_name!r} not found in {model_path!r}"
|
||||
)
|
||||
codec.load_checkpoint(ckpt_path)
|
||||
_freeze(codec.to(device=device, dtype=torch.float32))
|
||||
_semantic_codec_cache[cache_key] = codec
|
||||
return codec
|
||||
if codec_type != "repcodec":
|
||||
raise ValueError(f"Unsupported IndexTTS semantic codec type: {codec_type!r}")
|
||||
|
||||
from .utils.maskgct.repcodec_model import RepCodec
|
||||
|
||||
codec = RepCodec(
|
||||
@@ -84,11 +167,14 @@ def load_semantic_codec(model_path: str, config: dict, device: torch.device):
|
||||
hidden_size=config.get("hidden_size", 1024),
|
||||
codebook_dim=config.get("codebook_dim", 8),
|
||||
)
|
||||
ckpt_path = resolve_model_file(model_path, "semantic_codec.pth")
|
||||
checkpoint_name = checkpoint_name or "semantic_codec.pth"
|
||||
ckpt_path = resolve_model_file(model_path, checkpoint_name)
|
||||
if ckpt_path is not None:
|
||||
state = torch.load(ckpt_path, map_location="cpu", weights_only=True)
|
||||
codec.load_state_dict(state, strict=False)
|
||||
else:
|
||||
if os.path.isdir(model_path):
|
||||
raise FileNotFoundError(f"IndexTTS RepCodec checkpoint is missing from local bundle {model_path!r}")
|
||||
import safetensors.torch
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
@@ -106,8 +192,19 @@ def load_campplus(model_path: str, device: torch.device):
|
||||
from .utils.campplus.dtdnn import CAMPPlus
|
||||
|
||||
campplus = CAMPPlus(feat_dim=80, embedding_size=192)
|
||||
ckpt_path = resolve_model_file(model_path, "campplus.pth")
|
||||
ckpt_path = None
|
||||
for relative_path in (
|
||||
"campplus.pth",
|
||||
"campplus_cn_common.bin",
|
||||
os.path.join("hf_cache", "campplus_cn_common.bin"),
|
||||
os.path.join("campplus", "campplus_cn_common.bin"),
|
||||
):
|
||||
ckpt_path = resolve_model_file(model_path, relative_path)
|
||||
if ckpt_path is not None:
|
||||
break
|
||||
if ckpt_path is None:
|
||||
if os.path.isdir(model_path):
|
||||
raise FileNotFoundError(f"IndexTTS CAMPPlus checkpoint is missing from local bundle {model_path!r}")
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
ckpt_path = hf_hub_download("funasr/campplus", filename="campplus_cn_common.bin")
|
||||
@@ -125,20 +222,29 @@ def load_qwen_emotion(model_path: str, device: torch.device, *, trust_remote_cod
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
subfolder = "qwen0.6bemo4-merge"
|
||||
local_dir = resolve_model_directory(model_path, subfolder)
|
||||
source_args = (
|
||||
{"pretrained_model_name_or_path": local_dir, "local_files_only": True}
|
||||
if local_dir is not None
|
||||
else {
|
||||
"pretrained_model_name_or_path": model_path,
|
||||
"subfolder": subfolder,
|
||||
}
|
||||
)
|
||||
try:
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
model_path,
|
||||
subfolder=subfolder,
|
||||
**source_args,
|
||||
trust_remote_code=trust_remote_code,
|
||||
)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_path,
|
||||
subfolder=subfolder,
|
||||
**source_args,
|
||||
torch_dtype="float16",
|
||||
device_map={"": str(device)},
|
||||
trust_remote_code=trust_remote_code,
|
||||
)
|
||||
except Exception as e:
|
||||
if local_dir is not None:
|
||||
raise RuntimeError(f"Failed to load local IndexTTS QwenEmotion assets from {local_dir!r}") from e
|
||||
logger.warning(
|
||||
"Could not load QwenEmotion from %s/%s: %s. use_emo_text will fall back to default emotion.",
|
||||
model_path,
|
||||
@@ -172,9 +278,8 @@ def compute_fbank(wav_16k: torch.Tensor, device: torch.device) -> torch.Tensor:
|
||||
)
|
||||
fbank = fbank - fbank.mean(dim=0, keepdim=True)
|
||||
return fbank.unsqueeze(0).to(device=device) # [1, T, 80]
|
||||
except (ImportError, RuntimeError, OSError):
|
||||
logger.warning("torchaudio not available, returning zero fbank")
|
||||
return torch.zeros(1, 100, 80, device=device)
|
||||
except (ImportError, RuntimeError, OSError) as exc:
|
||||
raise RuntimeError("IndexTTS requires torchaudio Kaldi fbank support") from exc
|
||||
|
||||
|
||||
def wav2vec_extract(
|
||||
|
||||
@@ -7,10 +7,19 @@ from __future__ import annotations
|
||||
from functools import lru_cache
|
||||
|
||||
from vllm_omni.model_executor.models.indextts2.preprocess_utils import resolve_model_file
|
||||
from vllm_omni.model_executor.models.indextts2.text_processing_v2_5 import (
|
||||
prepare_indextts25_text,
|
||||
)
|
||||
from vllm_omni.model_executor.models.indextts2.tokenizer import IndexTTS2Tokenizer
|
||||
from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import (
|
||||
INDEXTTS25_TOKENIZER_FILE,
|
||||
)
|
||||
|
||||
_CONDITIONING_PREFIX_TOKENS = 34
|
||||
_TEXT_WRAPPER_TOKENS = 2
|
||||
_V2_CONDITIONING_PREFIX_TOKENS = 34
|
||||
_V25_CONDITIONING_PREFIX_TOKENS = 3
|
||||
_V2_TEXT_WRAPPER_TOKENS = 2
|
||||
_V25_TEXT_WRAPPER_TOKENS = 2
|
||||
_V25_TEXT_WRAPPER_TOKEN_IDS = frozenset({0, 1})
|
||||
_START_MEL_TOKENS = 1
|
||||
|
||||
|
||||
@@ -26,22 +35,55 @@ def _get_text_tokenizer(model_id_or_path: str) -> IndexTTS2Tokenizer:
|
||||
return IndexTTS2Tokenizer(_resolve_bpe_model_path(model_id_or_path), model_dir=model_id_or_path)
|
||||
|
||||
|
||||
def estimate_indextts2_prefill_prompt_len(model_id_or_path: str, text: str) -> int:
|
||||
def estimate_indextts2_prefill_prompt_len(
|
||||
model_id_or_path: str,
|
||||
text: str,
|
||||
*,
|
||||
model_type: str = "indextts2",
|
||||
lang: str = "zh",
|
||||
text_normalization: bool = True,
|
||||
tokenizer_file: str = INDEXTTS25_TOKENIZER_FILE,
|
||||
) -> int:
|
||||
"""Return the placeholder prompt length expected by the IndexTTS2 talker.
|
||||
|
||||
The model-side prefill replaces placeholder embeddings with:
|
||||
34 conditioning tokens, text tokens wrapped with start/stop markers, and
|
||||
one start-mel token.
|
||||
IndexTTS 2 uses 34 conditioning tokens. IndexTTS 2.5 uses one projected
|
||||
CAMPPlus speaker token followed by two zero tokens.
|
||||
"""
|
||||
tokenizer = _get_text_tokenizer(model_id_or_path)
|
||||
text_token_count = len(tokenizer.encode(text, add_special_tokens=False)) + _TEXT_WRAPPER_TOKENS
|
||||
return _CONDITIONING_PREFIX_TOKENS + text_token_count + _START_MEL_TOKENS
|
||||
if model_type == "indextts2_5":
|
||||
text_ids, _ = prepare_indextts25_text(
|
||||
text,
|
||||
lang=lang,
|
||||
model_dir=model_id_or_path,
|
||||
text_normalization=text_normalization,
|
||||
tokenizer_file=tokenizer_file,
|
||||
)
|
||||
conditioning_prefix_tokens = _V25_CONDITIONING_PREFIX_TOKENS
|
||||
text_token_count = (
|
||||
sum(token_id not in _V25_TEXT_WRAPPER_TOKEN_IDS for token_id in text_ids) + _V25_TEXT_WRAPPER_TOKENS
|
||||
)
|
||||
else:
|
||||
tokenizer = _get_text_tokenizer(model_id_or_path)
|
||||
text_token_count = len(tokenizer.encode(text, add_special_tokens=False)) + _V2_TEXT_WRAPPER_TOKENS
|
||||
conditioning_prefix_tokens = _V2_CONDITIONING_PREFIX_TOKENS
|
||||
return conditioning_prefix_tokens + text_token_count + _START_MEL_TOKENS
|
||||
|
||||
|
||||
def build_indextts2_prefill_prompt_ids(
|
||||
model_id_or_path: str,
|
||||
text: str,
|
||||
*,
|
||||
model_type: str = "indextts2",
|
||||
lang: str = "zh",
|
||||
text_normalization: bool = True,
|
||||
tokenizer_file: str = INDEXTTS25_TOKENIZER_FILE,
|
||||
placeholder_token_id: int = 1,
|
||||
) -> list[int]:
|
||||
return [placeholder_token_id] * estimate_indextts2_prefill_prompt_len(model_id_or_path, text)
|
||||
prompt_len = estimate_indextts2_prefill_prompt_len(
|
||||
model_id_or_path,
|
||||
text,
|
||||
model_type=model_type,
|
||||
lang=lang,
|
||||
text_normalization=text_normalization,
|
||||
tokenizer_file=tokenizer_file,
|
||||
)
|
||||
return [placeholder_token_id] * prompt_len
|
||||
@@ -68,6 +68,6 @@ class MyModel(nn.Module):
|
||||
x = self.models["gpt_layer"](x)
|
||||
return x
|
||||
|
||||
def enable_torch_compile(self):
|
||||
def enable_torch_compile(self, mode: str = "default"):
|
||||
if "cfm" in self.models:
|
||||
self.models["cfm"].enable_torch_compile()
|
||||
self.models["cfm"].enable_torch_compile(mode=mode)
|
||||
@@ -163,7 +163,18 @@ class DiT(torch.nn.Module):
|
||||
def setup_caches(self, max_batch_size, max_seq_length):
|
||||
self.transformer.setup_caches(max_batch_size, max_seq_length, use_kv_cache=False)
|
||||
|
||||
def forward(self, x, prompt_x, x_lens, t, style, cond, mask_content=False, pre_mask=None):
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
prompt_x,
|
||||
x_lens,
|
||||
t,
|
||||
style,
|
||||
cond,
|
||||
mask_content=False,
|
||||
pre_mask=None,
|
||||
unpad_data: tuple[torch.Tensor, torch.Tensor] | None = None,
|
||||
):
|
||||
"""
|
||||
x (torch.Tensor): random noise
|
||||
prompt_x (torch.Tensor): reference mel + zero mel
|
||||
@@ -226,10 +237,16 @@ class DiT(torch.nn.Module):
|
||||
input_pos = self.input_pos[: x_in.size(1)] # (T,) range(0,1863)
|
||||
# Route through per-shape CUDA graph runner if enabled.
|
||||
graph_runner = getattr(self, "_cuda_graph_runner", None)
|
||||
if graph_runner is not None:
|
||||
if graph_runner is not None and unpad_data is None:
|
||||
x_res = graph_runner(x_in, t1.unsqueeze(1), input_pos, x_mask_expanded)
|
||||
else:
|
||||
x_res = self.transformer(x_in, t1.unsqueeze(1), input_pos, x_mask_expanded) # [2, 1863, 512]
|
||||
x_res = self.transformer(
|
||||
x_in,
|
||||
t1.unsqueeze(1),
|
||||
input_pos,
|
||||
x_mask_expanded,
|
||||
unpad_data=unpad_data,
|
||||
) # [2, 1863, 512]
|
||||
x_res = x_res[:, 1:] if self.time_as_token else x_res
|
||||
x_res = x_res[:, 1:] if self.style_as_token else x_res
|
||||
|
||||
|
||||
@@ -8,6 +8,41 @@ import torch
|
||||
from vllm_omni.model_executor.models.indextts2.s2mel.modules.diffusion_transformer import DiT
|
||||
|
||||
|
||||
def build_cfm_unpad_data(
|
||||
lengths: list[int],
|
||||
*,
|
||||
padded_length: int,
|
||||
token_offset: int,
|
||||
repeat_for_cfg: bool,
|
||||
device: torch.device,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Build row-major FlashAttention unpadding metadata for one CFM group."""
|
||||
if not lengths:
|
||||
raise ValueError("CFM unpad data requires at least one sequence length")
|
||||
if token_offset < 0:
|
||||
raise ValueError(f"CFM unpad token offset must be non-negative: {token_offset}")
|
||||
if any(length < 0 or length > padded_length for length in lengths):
|
||||
raise ValueError(f"CFM unpad lengths out of range: lengths={lengths} padded={padded_length}")
|
||||
if max(lengths) != padded_length:
|
||||
raise ValueError(f"CFM unpad group must contain its padded max: lengths={lengths} padded={padded_length}")
|
||||
|
||||
valid_lengths = [length + token_offset for length in lengths]
|
||||
if repeat_for_cfg:
|
||||
valid_lengths = [*valid_lengths, *valid_lengths]
|
||||
padded_stride = padded_length + token_offset
|
||||
indices = torch.cat(
|
||||
[
|
||||
torch.arange(valid_length, device=device, dtype=torch.int64) + row * padded_stride
|
||||
for row, valid_length in enumerate(valid_lengths)
|
||||
]
|
||||
)
|
||||
cumulative_lengths = [0]
|
||||
for valid_length in valid_lengths:
|
||||
cumulative_lengths.append(cumulative_lengths[-1] + valid_length)
|
||||
cu_seqlens = torch.tensor(cumulative_lengths, device=device, dtype=torch.int32)
|
||||
return indices, cu_seqlens
|
||||
|
||||
|
||||
class BASECFM(torch.nn.Module, ABC):
|
||||
def __init__(
|
||||
self,
|
||||
@@ -26,7 +61,20 @@ class BASECFM(torch.nn.Module, ABC):
|
||||
self.zero_prompt_speech_token = False
|
||||
|
||||
@torch.inference_mode()
|
||||
def inference(self, mu, x_lens, prompt, style, f0, n_timesteps, temperature=1.0, inference_cfg_rate=0.5):
|
||||
def inference(
|
||||
self,
|
||||
mu,
|
||||
x_lens,
|
||||
prompt,
|
||||
style,
|
||||
f0,
|
||||
n_timesteps,
|
||||
temperature=1.0,
|
||||
inference_cfg_rate=0.5,
|
||||
initial_noise=None,
|
||||
prompt_lens=None,
|
||||
unpad_data: tuple[torch.Tensor, torch.Tensor] | None = None,
|
||||
):
|
||||
"""Forward diffusion
|
||||
|
||||
Args:
|
||||
@@ -47,11 +95,48 @@ class BASECFM(torch.nn.Module, ABC):
|
||||
shape: (batch_size, 80, mel_timesteps)
|
||||
"""
|
||||
B, T = mu.size(0), mu.size(1)
|
||||
z = torch.randn([B, self.in_channels, T], device=mu.device, dtype=mu.dtype) * temperature
|
||||
expected_noise_shape = (B, self.in_channels, T)
|
||||
if initial_noise is None:
|
||||
z = torch.randn(
|
||||
expected_noise_shape,
|
||||
device=mu.device,
|
||||
dtype=mu.dtype,
|
||||
)
|
||||
else:
|
||||
if tuple(initial_noise.shape) != expected_noise_shape:
|
||||
raise ValueError(
|
||||
"CFM initial noise shape mismatch: "
|
||||
f"expected={expected_noise_shape}, got={tuple(initial_noise.shape)}"
|
||||
)
|
||||
z = initial_noise.to(device=mu.device, dtype=mu.dtype)
|
||||
z = z * temperature
|
||||
t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype)
|
||||
return self.solve_euler(z, x_lens, prompt, mu, style, f0, t_span, inference_cfg_rate)
|
||||
return self.solve_euler(
|
||||
z,
|
||||
x_lens,
|
||||
prompt,
|
||||
mu,
|
||||
style,
|
||||
f0,
|
||||
t_span,
|
||||
inference_cfg_rate,
|
||||
prompt_lens=prompt_lens,
|
||||
unpad_data=unpad_data,
|
||||
)
|
||||
|
||||
def solve_euler(self, x, x_lens, prompt, mu, style, f0, t_span, inference_cfg_rate=0.5):
|
||||
def solve_euler(
|
||||
self,
|
||||
x,
|
||||
x_lens,
|
||||
prompt,
|
||||
mu,
|
||||
style,
|
||||
f0,
|
||||
t_span,
|
||||
inference_cfg_rate=0.5,
|
||||
prompt_lens=None,
|
||||
unpad_data: tuple[torch.Tensor, torch.Tensor] | None = None,
|
||||
):
|
||||
"""
|
||||
Fixed euler solver for ODEs.
|
||||
Args:
|
||||
@@ -82,19 +167,35 @@ class BASECFM(torch.nn.Module, ABC):
|
||||
out = self.estimator(*args, **kw)
|
||||
return out.float()
|
||||
|
||||
prompt_len = prompt.size(-1)
|
||||
batch = x.shape[0]
|
||||
if prompt_lens is None:
|
||||
prompt_lens = torch.full(
|
||||
(batch,),
|
||||
prompt.size(-1),
|
||||
device=x.device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
else:
|
||||
prompt_lens = prompt_lens.to(device=x.device, dtype=torch.long).reshape(-1)
|
||||
if prompt_lens.shape != (batch,):
|
||||
raise ValueError(f"CFM prompt length/batch mismatch: lengths={prompt_lens.numel()} batch={batch}")
|
||||
|
||||
prompt_positions = torch.arange(prompt.size(-1), device=x.device)
|
||||
prompt_mask = prompt_positions.unsqueeze(0) < prompt_lens.unsqueeze(1)
|
||||
x_positions = torch.arange(x.size(-1), device=x.device)
|
||||
noise_prefix_mask = x_positions.unsqueeze(0) < prompt_lens.unsqueeze(1)
|
||||
|
||||
prompt_x = torch.zeros_like(x)
|
||||
prompt_x[..., :prompt_len] = prompt[..., :prompt_len]
|
||||
x[..., :prompt_len] = 0
|
||||
prompt_x[..., : prompt.size(-1)] = prompt * prompt_mask.unsqueeze(1).to(prompt.dtype)
|
||||
x.masked_fill_(noise_prefix_mask.unsqueeze(1), 0)
|
||||
if self.zero_prompt_speech_token:
|
||||
mu[..., :prompt_len] = 0
|
||||
mu = mu.masked_fill(noise_prefix_mask.unsqueeze(-1), 0)
|
||||
|
||||
# Precompute mask once — T and x_lens are constant across all ODE steps.
|
||||
# x_mask [B, 1, T] broadcasts to [2B, 1, T] when CFG doubles batch.
|
||||
pre_mask = self._precompute_mask(x, x_lens)
|
||||
|
||||
if inference_cfg_rate > 0:
|
||||
batch = x.shape[0]
|
||||
cfg_prompt_x = prompt_x.new_zeros((2 * batch, *prompt_x.shape[1:]))
|
||||
cfg_prompt_x[:batch].copy_(prompt_x)
|
||||
cfg_style = style.new_zeros((2 * batch, *style.shape[1:]))
|
||||
@@ -120,16 +221,26 @@ class BASECFM(torch.nn.Module, ABC):
|
||||
cfg_style,
|
||||
cfg_mu,
|
||||
pre_mask=cfg_pre_mask,
|
||||
unpad_data=unpad_data,
|
||||
)
|
||||
|
||||
dphi_dt, cfg_dphi_dt = stacked_dphi_dt.chunk(2, dim=0)
|
||||
dphi_dt = (1.0 + inference_cfg_rate) * dphi_dt - inference_cfg_rate * cfg_dphi_dt
|
||||
else:
|
||||
dphi_dt = _run_estimator(x, prompt_x, x_lens, t.unsqueeze(0), style, mu, pre_mask=pre_mask)
|
||||
dphi_dt = _run_estimator(
|
||||
x,
|
||||
prompt_x,
|
||||
x_lens,
|
||||
t.unsqueeze(0),
|
||||
style,
|
||||
mu,
|
||||
pre_mask=pre_mask,
|
||||
unpad_data=unpad_data,
|
||||
)
|
||||
|
||||
x = x + dt * dphi_dt
|
||||
t = t + dt
|
||||
x[:, :, :prompt_len] = 0
|
||||
x.masked_fill_(noise_prefix_mask.unsqueeze(1), 0)
|
||||
|
||||
return x
|
||||
|
||||
@@ -180,7 +291,7 @@ class CFM(BASECFM):
|
||||
raise NotImplementedError(f"Unknown diffusion type {args.dit_type}")
|
||||
self._compiled = False
|
||||
|
||||
def enable_torch_compile(self):
|
||||
def enable_torch_compile(self, mode: str = "default"):
|
||||
"""Enable torch.compile optimization for the estimator model.
|
||||
|
||||
This method applies torch.compile to the estimator (DiT model) for significant
|
||||
@@ -191,6 +302,7 @@ class CFM(BASECFM):
|
||||
torch._inductor.config.reorder_for_compute_comm_overlap = True
|
||||
self.estimator = torch.compile(
|
||||
self._eager_estimator,
|
||||
mode=mode,
|
||||
fullgraph=True,
|
||||
dynamic=True,
|
||||
)
|
||||
|
||||
@@ -116,6 +116,7 @@ class Transformer(nn.Module):
|
||||
context: Tensor | None = None,
|
||||
context_input_pos: Tensor | None = None,
|
||||
cross_attention_mask: Tensor | None = None,
|
||||
unpad_data: tuple[Tensor, Tensor] | None = None,
|
||||
) -> Tensor:
|
||||
assert self.freqs_cis is not None, "Caches must be initialized first"
|
||||
if mask is None: # in case of non-causal model
|
||||
@@ -135,7 +136,18 @@ class Transformer(nn.Module):
|
||||
skip_in_x = skip_in_x_list.pop(-1)
|
||||
else:
|
||||
skip_in_x = None
|
||||
x = layer(x, c, input_pos, freqs_cis, mask, context, context_freqs_cis, cross_attention_mask, skip_in_x)
|
||||
x = layer(
|
||||
x,
|
||||
c,
|
||||
input_pos,
|
||||
freqs_cis,
|
||||
mask,
|
||||
context,
|
||||
context_freqs_cis,
|
||||
cross_attention_mask,
|
||||
skip_in_x,
|
||||
unpad_data=unpad_data,
|
||||
)
|
||||
if self.uvit_skip_connection and i in self.layers_emit_skip:
|
||||
skip_in_x_list.append(x)
|
||||
x = self.norm(x, c)
|
||||
@@ -169,11 +181,18 @@ class TransformerBlock(nn.Module):
|
||||
context_freqs_cis: Tensor | None = None,
|
||||
cross_attention_mask: Tensor | None = None,
|
||||
skip_in_x: Tensor | None = None,
|
||||
unpad_data: tuple[Tensor, Tensor] | None = None,
|
||||
) -> Tensor:
|
||||
c = None if self.time_as_token else c
|
||||
if self.uvit_skip_connection and skip_in_x is not None:
|
||||
x = self.skip_in_linear(torch.cat([x, skip_in_x], dim=-1))
|
||||
h = x + self.attention(self.attention_norm(x, c), freqs_cis, mask, input_pos)
|
||||
h = x + self.attention(
|
||||
self.attention_norm(x, c),
|
||||
freqs_cis,
|
||||
mask,
|
||||
input_pos,
|
||||
unpad_data=unpad_data,
|
||||
)
|
||||
out = h + self.feed_forward(self.ffn_norm(h, c))
|
||||
return out
|
||||
|
||||
@@ -218,6 +237,7 @@ class Attention(nn.Module):
|
||||
input_pos: Tensor | None = None,
|
||||
context: Tensor | None = None,
|
||||
context_freqs_cis: Tensor | None = None,
|
||||
unpad_data: tuple[Tensor, Tensor] | None = None,
|
||||
) -> Tensor:
|
||||
bsz, seqlen, _ = x.shape
|
||||
|
||||
@@ -252,7 +272,37 @@ class Attention(nn.Module):
|
||||
# building AttentionMetadata entirely — this avoids the backend's
|
||||
# `torch.any(~mask)` D2H sync per ODE step and keeps the forward
|
||||
# CUDA-graph capturable.
|
||||
if getattr(self, "_assume_full_mask", False):
|
||||
full_mask_state = getattr(self, "_assume_full_mask", None)
|
||||
packed_varlen = (
|
||||
unpad_data is not None
|
||||
and context is None
|
||||
and q.device.type == "cuda"
|
||||
and q.dtype != torch.float32
|
||||
and hasattr(getattr(self.diff_attn, "attention", None), "_forward_varlen_packed")
|
||||
)
|
||||
if packed_varlen:
|
||||
indices, cu_seqlens = unpad_data
|
||||
packed_q = q.flatten(0, 1).index_select(0, indices).unsqueeze(0)
|
||||
packed_k = k.flatten(0, 1).index_select(0, indices).unsqueeze(0)
|
||||
packed_v = v.flatten(0, 1).index_select(0, indices).unsqueeze(0)
|
||||
packed_metadata = AttentionMetadata(
|
||||
extra={
|
||||
"cu_seqlens_q": cu_seqlens,
|
||||
"cu_seqlens_k": cu_seqlens,
|
||||
"max_seqlen_q": seqlen,
|
||||
"max_seqlen_k": seqlen,
|
||||
}
|
||||
)
|
||||
packed_y = self.diff_attn(
|
||||
packed_q,
|
||||
packed_k,
|
||||
packed_v,
|
||||
attn_metadata=packed_metadata,
|
||||
)
|
||||
flat_y = torch.zeros_like(q.flatten(0, 1))
|
||||
flat_y.index_copy_(0, indices, packed_y.squeeze(0))
|
||||
y = flat_y.view_as(q)
|
||||
elif full_mask_state is True:
|
||||
attn_metadata = None
|
||||
else:
|
||||
attn_mask = mask
|
||||
@@ -260,8 +310,15 @@ class Attention(nn.Module):
|
||||
# (B, 1, T, T) → (B, T) — take the first query row;
|
||||
# all query positions share the same padding pattern.
|
||||
attn_mask = attn_mask[:, 0, 0, :]
|
||||
attn_metadata = AttentionMetadata(attn_mask=attn_mask) if attn_mask is not None else None
|
||||
y = self.diff_attn(q, k, v, attn_metadata=attn_metadata)
|
||||
attn_metadata = (
|
||||
AttentionMetadata(
|
||||
attn_mask=attn_mask,
|
||||
)
|
||||
if attn_mask is not None
|
||||
else None
|
||||
)
|
||||
if not packed_varlen:
|
||||
y = self.diff_attn(q, k, v, attn_metadata=attn_metadata)
|
||||
y = y.flatten(2, 3)
|
||||
|
||||
y = self.wo(y)
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from vllm_omni.model_executor.models.indextts2.s2mel.modules.encodec import SConv1d
|
||||
|
||||
@@ -53,14 +54,82 @@ class WN(torch.nn.Module):
|
||||
res_skip_layer = conv1d_type(hidden_channels, res_skip_channels, 1, norm="weight_norm", causal=causal)
|
||||
self.res_skip_layers.append(res_skip_layer)
|
||||
|
||||
# The vectorized ragged path rebuilds logical right-reflect values in
|
||||
# place and therefore requires each row to be strictly longer than the
|
||||
# widest right padding. The decoder keeps shorter rows on the original
|
||||
# exact-length singleton path, preserving SConv1d's existing behavior.
|
||||
self.ragged_min_length = 1 + max(
|
||||
((layer.conv.conv.kernel_size[0] - 1) * layer.conv.conv.dilation[0] + 1 - layer.conv.conv.stride[0]) // 2
|
||||
for layer in self.in_layers
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _logical_reflect_in_layer(
|
||||
in_layer: SConv1d,
|
||||
x: torch.Tensor,
|
||||
logical_lengths: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Run a WaveNet input convolution at each row's logical boundary.
|
||||
|
||||
``SConv1d`` reflects at the physical tensor boundary. In a padded
|
||||
request batch, shorter rows end before that boundary, so their right
|
||||
reflection values must be rebuilt before the convolution. This path is
|
||||
intentionally WaveNet-specific: its input layers are non-causal,
|
||||
reflect-padded, stride-one convolutions. The decoder keeps rows that
|
||||
are not strictly longer than the right reflection width on the
|
||||
original singleton path.
|
||||
"""
|
||||
conv = in_layer.conv.conv
|
||||
stride = conv.stride[0]
|
||||
dilation = conv.dilation[0]
|
||||
effective_kernel = (conv.kernel_size[0] - 1) * dilation + 1
|
||||
if in_layer.causal or in_layer.pad_mode != "reflect" or stride != 1:
|
||||
raise RuntimeError(
|
||||
"Ragged WaveNet batching requires non-causal reflect-padded stride-one input convolutions"
|
||||
)
|
||||
|
||||
padding_total = effective_kernel - stride
|
||||
padding_right = padding_total // 2
|
||||
padding_left = padding_total - padding_right
|
||||
padded = F.pad(x, (padding_left, padding_right), mode="reflect")
|
||||
|
||||
if padding_right:
|
||||
offsets = torch.arange(
|
||||
padding_right,
|
||||
device=x.device,
|
||||
dtype=torch.long,
|
||||
).unsqueeze(0)
|
||||
lengths = logical_lengths.to(device=x.device, dtype=torch.long).reshape(-1, 1)
|
||||
source_positions = padding_left + lengths - 2 - offsets
|
||||
target_positions = padding_left + lengths + offsets
|
||||
source_indices = source_positions.unsqueeze(1).expand(-1, x.shape[1], -1)
|
||||
target_indices = target_positions.unsqueeze(1).expand(-1, x.shape[1], -1)
|
||||
reflected = torch.gather(padded, dim=2, index=source_indices)
|
||||
padded.scatter_(dim=2, index=target_indices, src=reflected)
|
||||
|
||||
return in_layer.conv(padded)
|
||||
|
||||
def forward(self, x, x_mask, g=None, **kwargs):
|
||||
output = torch.zeros_like(x)
|
||||
|
||||
logical_lengths = None
|
||||
if getattr(self, "_assume_full_mask", None) is False:
|
||||
# Under CFG, x_mask has already been repeated from B to 2B in the
|
||||
# same [conditional rows, unconditional rows] order as x.
|
||||
logical_lengths = x_mask[:, 0, :].sum(dim=-1)
|
||||
|
||||
if g is not None:
|
||||
g = self.cond_layer(g)
|
||||
|
||||
for i in range(self.n_layers):
|
||||
x_in = self.in_layers[i](x)
|
||||
if logical_lengths is None:
|
||||
x_in = self.in_layers[i](x)
|
||||
else:
|
||||
x_in = self._logical_reflect_in_layer(
|
||||
self.in_layers[i],
|
||||
x,
|
||||
logical_lengths,
|
||||
)
|
||||
if g is not None:
|
||||
cond_offset = i * 2 * self.hidden_channels
|
||||
g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :]
|
||||
|
||||
@@ -0,0 +1,328 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Text preparation shared by IndexTTS 2.5 prompt sizing and model prefill."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Callable
|
||||
from functools import lru_cache
|
||||
|
||||
from .tokenizer_v2_5 import (
|
||||
INDEXTTS25_TOKENIZER_FILE,
|
||||
encode_indextts25_text,
|
||||
lang_to_token,
|
||||
normalize_language_code,
|
||||
)
|
||||
|
||||
PRONUNCIATION_ANNOTATION_PATTERN = re.compile(r"<([^|>\n]+)\|([^>\n]+)>")
|
||||
_SPECIAL_TOKEN_PATTERN = re.compile(r"<\|([^|]+)\|>")
|
||||
_PINYIN_TONE_PATTERN = re.compile(
|
||||
r"(?<![a-z])"
|
||||
r"((?:[bpmfdtnlgkhjqxzcsryw]|[zcs]h)?"
|
||||
r"(?:[aeiouüv]|[ae]i|u[aio]|ao|ou|i[aue]|[uüv]e|[uvü]ang?|"
|
||||
r"uai|[aeiuv]n|[aeio]ng|ia[no]|i[ao]ng)|ng|er)([1-5])",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
_NAME_PATTERN = re.compile(r"[\u4e00-\u9fff]+(?:[-·—][\u4e00-\u9fff]+){1,2}")
|
||||
_TECH_TERM_PATTERN = re.compile(r"[A-Za-z][A-Za-z0-9]*(?:-[A-Za-z0-9]+)+")
|
||||
_ENGLISH_CONTRACTION_PATTERN = re.compile(
|
||||
r"(what|where|who|which|how|t?here|it|s?he|that|this)'s",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
_EMAIL_PATTERN = re.compile(r"^[a-zA-Z0-9]+@[a-zA-Z0-9]+\.[a-zA-Z]+$")
|
||||
|
||||
_CHAR_REPLACEMENTS = {
|
||||
":": ",",
|
||||
";": ",",
|
||||
";": ",",
|
||||
",": ",",
|
||||
"。": ".",
|
||||
"!": "!",
|
||||
"?": "?",
|
||||
"\n": " ",
|
||||
"·": "-",
|
||||
"、": ",",
|
||||
"...": "…",
|
||||
",,,": "…",
|
||||
",,,": "…",
|
||||
"……": "…",
|
||||
"“": "'",
|
||||
"”": "'",
|
||||
'"': "'",
|
||||
"‘": "'",
|
||||
"’": "'",
|
||||
"(": "'",
|
||||
")": "'",
|
||||
"(": "'",
|
||||
")": "'",
|
||||
"《": "'",
|
||||
"》": "'",
|
||||
"【": "'",
|
||||
"】": "'",
|
||||
"[": "'",
|
||||
"]": "'",
|
||||
"—": "-",
|
||||
"~": "-",
|
||||
"~": "-",
|
||||
"「": "'",
|
||||
"」": "'",
|
||||
":": ",",
|
||||
}
|
||||
# Deliberately prefer the longest replacement. This diverges from the
|
||||
# insertion-ordered upstream regex so ``,,,`` becomes one ellipsis, not ``,,,``.
|
||||
_CLEAN_PATTERN = re.compile("|".join(re.escape(key) for key in sorted(_CHAR_REPLACEMENTS, key=len, reverse=True)))
|
||||
|
||||
|
||||
def _is_kana(text: str) -> bool:
|
||||
return bool(re.fullmatch(r"[\u3040-\u309f]+", text) or re.fullmatch(r"[\u30a0-\u30ff]+", text))
|
||||
|
||||
|
||||
def apply_pronunciation_annotations(text: str) -> str:
|
||||
"""Convert ``<word|pronunciation>`` to the official special-token form."""
|
||||
|
||||
def replace(match: re.Match[str]) -> str:
|
||||
word = match.group(1)
|
||||
pronunciation = match.group(2).upper()
|
||||
if _is_kana(pronunciation):
|
||||
return f" {pronunciation} "
|
||||
token = "SPECIAL_TOKEN_2" if re.search(r"[\u4e00-\u9fff]", word) else "SPECIAL_TOKEN_1"
|
||||
return f"<|{token}|>{pronunciation}<|{token}|>"
|
||||
|
||||
return PRONUNCIATION_ANNOTATION_PATTERN.sub(replace, text)
|
||||
|
||||
|
||||
def clean_indextts25_text(text: str) -> str:
|
||||
return _CLEAN_PATTERN.sub(lambda match: _CHAR_REPLACEMENTS[match.group()], text)
|
||||
|
||||
|
||||
@lru_cache(maxsize=2)
|
||||
def _load_wetext_normalizer(lang: str):
|
||||
import platform
|
||||
|
||||
if platform.system() == "Linux":
|
||||
try:
|
||||
if lang == "zh":
|
||||
from tn.chinese.normalizer import Normalizer
|
||||
else:
|
||||
from tn.english.normalizer import Normalizer
|
||||
except ModuleNotFoundError as exc:
|
||||
if exc.name != "tn":
|
||||
raise
|
||||
raise RuntimeError(
|
||||
f"IndexTTS 2.5 text normalization backend for {lang!r} is unavailable. "
|
||||
"Install it with pip install 'vllm-omni[indextts2]'"
|
||||
) from exc
|
||||
if lang == "zh":
|
||||
return Normalizer(
|
||||
remove_interjections=False,
|
||||
remove_erhua=False,
|
||||
overwrite_cache=False,
|
||||
)
|
||||
return Normalizer(overwrite_cache=False)
|
||||
|
||||
try:
|
||||
from wetext import Normalizer
|
||||
except ModuleNotFoundError as exc:
|
||||
if exc.name != "wetext":
|
||||
raise
|
||||
raise RuntimeError(
|
||||
f"IndexTTS 2.5 text normalization backend for {lang!r} is unavailable. "
|
||||
"Install it with pip install 'vllm-omni[indextts2]'"
|
||||
) from exc
|
||||
if lang == "zh":
|
||||
return Normalizer(remove_erhua=False, lang="zh", operator="tn")
|
||||
return Normalizer(lang="en", operator="tn")
|
||||
|
||||
|
||||
def _uses_chinese_normalizer(text: str) -> bool:
|
||||
has_chinese = bool(re.search(r"[\u4e00-\u9fff]", text))
|
||||
has_alpha = bool(re.search(r"[a-zA-Z]", text))
|
||||
if has_chinese or not has_alpha or _EMAIL_PATTERN.fullmatch(text):
|
||||
return True
|
||||
return bool(_PINYIN_TONE_PATTERN.search(text))
|
||||
|
||||
|
||||
def _protect_values(
|
||||
text: str,
|
||||
values: list[str],
|
||||
*,
|
||||
prefix: str,
|
||||
) -> tuple[str, list[str]]:
|
||||
unique_values = sorted(set(values), key=len, reverse=True)
|
||||
for index, value in enumerate(unique_values):
|
||||
marker = f"<{prefix}_{chr(ord('a') + index)}>"
|
||||
text = text.replace(value, marker)
|
||||
return text, unique_values
|
||||
|
||||
|
||||
def _restore_values(
|
||||
text: str,
|
||||
values: list[str],
|
||||
*,
|
||||
prefix: str,
|
||||
transform: Callable[[str], str] | None = None,
|
||||
) -> str:
|
||||
for index, value in enumerate(values):
|
||||
marker = f"<{prefix}_{chr(ord('a') + index)}>"
|
||||
text = text.replace(marker, transform(value) if transform else value)
|
||||
return text
|
||||
|
||||
|
||||
def _correct_pinyin(pinyin: str) -> str:
|
||||
if not pinyin or pinyin[0] not in "jqxJQX":
|
||||
return pinyin.upper()
|
||||
return re.sub(
|
||||
r"([jqx])[uü](n|e|an)*(\d)",
|
||||
r"\g<1>v\g<2>\g<3>",
|
||||
pinyin,
|
||||
flags=re.IGNORECASE,
|
||||
).upper()
|
||||
|
||||
|
||||
def _normalize_zh_or_en(text: str) -> str:
|
||||
normalizer_lang = "zh" if _uses_chinese_normalizer(text) else "en"
|
||||
text = _ENGLISH_CONTRACTION_PATTERN.sub(r"\1 is", text)
|
||||
|
||||
tech_terms = _TECH_TERM_PATTERN.findall(text)
|
||||
for term in sorted(set(tech_terms), key=len, reverse=True):
|
||||
text = text.replace(term, term.replace("-", "<H>"))
|
||||
|
||||
pinyin_values: list[str] = []
|
||||
name_values: list[str] = []
|
||||
if normalizer_lang == "zh":
|
||||
pinyin_values = [match.group(0) for match in _PINYIN_TONE_PATTERN.finditer(text)]
|
||||
text, pinyin_values = _protect_values(
|
||||
text,
|
||||
pinyin_values,
|
||||
prefix="pinyin",
|
||||
)
|
||||
text, name_values = _protect_values(
|
||||
text,
|
||||
_NAME_PATTERN.findall(text),
|
||||
prefix="n",
|
||||
)
|
||||
|
||||
result = _load_wetext_normalizer(normalizer_lang).normalize(text)
|
||||
result = _restore_values(result, name_values, prefix="n")
|
||||
result = _restore_values(
|
||||
result,
|
||||
pinyin_values,
|
||||
prefix="pinyin",
|
||||
transform=_correct_pinyin,
|
||||
)
|
||||
if tech_terms:
|
||||
result = re.sub(r"\s*<H>\s*", "-", result)
|
||||
if normalizer_lang == "zh":
|
||||
result = result.replace("$", ".")
|
||||
return clean_indextts25_text(result)
|
||||
|
||||
|
||||
def _normalize_with_official_backend(text: str, lang: str) -> str:
|
||||
if lang in {"zh", "en", "zhen"}:
|
||||
return _normalize_zh_or_en(text)
|
||||
if lang == "es":
|
||||
# NeMo is optional in the upstream implementation and falls back to
|
||||
# raw text when its grammar package is unavailable.
|
||||
try:
|
||||
from nemo_text_processing.text_normalization.normalize import (
|
||||
Normalizer,
|
||||
)
|
||||
|
||||
normalizer = Normalizer(input_case="cased", lang="es")
|
||||
return normalizer.normalize(text, verbose=False)
|
||||
except ModuleNotFoundError as exc:
|
||||
if exc.name != "nemo_text_processing":
|
||||
raise
|
||||
return text
|
||||
return text
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _load_japanese_tagger():
|
||||
try:
|
||||
import fugashi
|
||||
import unidic_lite
|
||||
except ModuleNotFoundError as exc:
|
||||
if exc.name not in {"fugashi", "unidic_lite"}:
|
||||
raise
|
||||
raise RuntimeError(
|
||||
"IndexTTS 2.5 Japanese preprocessing requires fugashi and unidic-lite. "
|
||||
"Install them with pip install 'vllm-omni[indextts2]'"
|
||||
) from exc
|
||||
_ = unidic_lite
|
||||
return fugashi.Tagger()
|
||||
|
||||
|
||||
def _process_japanese_text(text: str) -> str:
|
||||
"""Tokenize Japanese like the official g2p_ratio=0 frontend."""
|
||||
|
||||
def process_segment(segment: str) -> str:
|
||||
return " ".join(token.surface for token in _load_japanese_tagger()(segment))
|
||||
|
||||
parts = re.split(r"( +)", text)
|
||||
return "".join(process_segment(part) if part.strip() else part for part in parts)
|
||||
|
||||
|
||||
def normalize_indextts25_text(
|
||||
text: str,
|
||||
*,
|
||||
lang: str,
|
||||
text_normalization: bool = True,
|
||||
normalizer: Callable[[str, str], str] | None = None,
|
||||
) -> str:
|
||||
"""Apply the deterministic portion of the official preprocessing order.
|
||||
|
||||
A caller-provided normalizer is used for number/date expansion. Keeping it
|
||||
injectable lets prompt sizing and Talker prefill use the same normalizer
|
||||
instance and avoids constructing heavyweight TN grammars per request.
|
||||
"""
|
||||
lang = normalize_language_code(lang)
|
||||
text = clean_indextts25_text(text)
|
||||
if text_normalization:
|
||||
normalize = normalizer or _normalize_with_official_backend
|
||||
text = normalize(text, lang)
|
||||
if lang in {"ja", "zh", "en", "zhen"}:
|
||||
text = text.lower()
|
||||
elif lang == "es":
|
||||
text = text.upper()
|
||||
text = apply_pronunciation_annotations(text)
|
||||
if lang == "ja":
|
||||
text = _process_japanese_text(text)
|
||||
return _SPECIAL_TOKEN_PATTERN.sub(
|
||||
lambda match: f"<|{match.group(1).upper()}|>",
|
||||
text,
|
||||
)
|
||||
|
||||
|
||||
def prepare_indextts25_text(
|
||||
text: str,
|
||||
*,
|
||||
lang: str,
|
||||
model_dir: str,
|
||||
tokenizer_file: str = INDEXTTS25_TOKENIZER_FILE,
|
||||
text_normalization: bool = True,
|
||||
normalizer: Callable[[str, str], str] | None = None,
|
||||
) -> tuple[list[int], int]:
|
||||
if not text or not text.strip():
|
||||
raise ValueError("IndexTTS 2.5 text cannot be empty")
|
||||
lang_code = normalize_language_code(lang)
|
||||
# When no heavyweight TN callback is installed, punctuation/casing and
|
||||
# annotations are still deterministic. The Talker installs the official
|
||||
# language-specific normalizer before enabling text_normalization.
|
||||
normalized = normalize_indextts25_text(
|
||||
text,
|
||||
lang=lang_code,
|
||||
text_normalization=text_normalization,
|
||||
normalizer=normalizer,
|
||||
)
|
||||
# ``zhen`` is not a registered special token. Keep its official literal
|
||||
# prefix so tiktoken encodes it through existing mergeable ranks; remapping
|
||||
# it to ``common`` would change both the text ABI and preprocessing route.
|
||||
prefixed = f"<|{lang_code}|> {normalized}"
|
||||
token_ids = encode_indextts25_text(
|
||||
prefixed,
|
||||
model_dir=model_dir,
|
||||
tokenizer_file=tokenizer_file,
|
||||
)
|
||||
return token_ids, lang_to_token(lang_code)
|
||||
@@ -0,0 +1,286 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""IndexTTS 2.5 multilingual tiktoken vocabulary.
|
||||
|
||||
The vocabulary layout mirrors the official IndexTTS 2.5 tokenizer: 58,836
|
||||
mergeable ranks followed by 1,673 registered special tokens.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
INDEXTTS25_TOKENIZER_FILE = "multilingual_zh_ja_yue_char_del.tiktoken"
|
||||
INDEXTTS25_MERGEABLE_RANKS = 58836
|
||||
INDEXTTS25_VOCAB_SIZE = 60509
|
||||
INDEXTTS25_NUM_SPECIAL_LANGUAGES = 99
|
||||
|
||||
# Order is checkpoint ABI: lang_embedding rows were trained against these IDs.
|
||||
LANGUAGES = (
|
||||
"en",
|
||||
"zh",
|
||||
"de",
|
||||
"es",
|
||||
"ru",
|
||||
"ko",
|
||||
"fr",
|
||||
"ja",
|
||||
"pt",
|
||||
"tr",
|
||||
"pl",
|
||||
"ca",
|
||||
"nl",
|
||||
"ar",
|
||||
"sv",
|
||||
"it",
|
||||
"id",
|
||||
"hi",
|
||||
"fi",
|
||||
"vi",
|
||||
"he",
|
||||
"uk",
|
||||
"el",
|
||||
"ms",
|
||||
"cs",
|
||||
"ro",
|
||||
"da",
|
||||
"hu",
|
||||
"ta",
|
||||
"no",
|
||||
"th",
|
||||
"ur",
|
||||
"hr",
|
||||
"bg",
|
||||
"lt",
|
||||
"la",
|
||||
"mi",
|
||||
"ml",
|
||||
"cy",
|
||||
"sk",
|
||||
"te",
|
||||
"fa",
|
||||
"lv",
|
||||
"bn",
|
||||
"sr",
|
||||
"az",
|
||||
"sl",
|
||||
"kn",
|
||||
"et",
|
||||
"mk",
|
||||
"br",
|
||||
"eu",
|
||||
"is",
|
||||
"hy",
|
||||
"ne",
|
||||
"mn",
|
||||
"bs",
|
||||
"kk",
|
||||
"sq",
|
||||
"sw",
|
||||
"gl",
|
||||
"mr",
|
||||
"pa",
|
||||
"si",
|
||||
"km",
|
||||
"sn",
|
||||
"yo",
|
||||
"so",
|
||||
"af",
|
||||
"oc",
|
||||
"ka",
|
||||
"be",
|
||||
"tg",
|
||||
"sd",
|
||||
"gu",
|
||||
"am",
|
||||
"yi",
|
||||
"lo",
|
||||
"uz",
|
||||
"fo",
|
||||
"ht",
|
||||
"ps",
|
||||
"tk",
|
||||
"nn",
|
||||
"mt",
|
||||
"sa",
|
||||
"lb",
|
||||
"my",
|
||||
"bo",
|
||||
"tl",
|
||||
"mg",
|
||||
"as",
|
||||
"tt",
|
||||
"haw",
|
||||
"ln",
|
||||
"ha",
|
||||
"ba",
|
||||
"jw",
|
||||
"su",
|
||||
"yue",
|
||||
"minnan",
|
||||
"wuyu",
|
||||
"dialect",
|
||||
"zh/en",
|
||||
"en/zh",
|
||||
"common",
|
||||
)
|
||||
LANGUAGE_DICT = {lang: index for index, lang in enumerate(LANGUAGES)}
|
||||
|
||||
# ``zhen`` is an official mixed Chinese/English frontend mode, but it has no
|
||||
# dedicated checkpoint language-embedding row or registered special token.
|
||||
_FRONTEND_LANGUAGE_TO_EMBEDDING = {"zhen": "common"}
|
||||
|
||||
TO_LANGUAGE_CODE = {
|
||||
"english": "en",
|
||||
"chinese": "zh",
|
||||
"german": "de",
|
||||
"spanish": "es",
|
||||
"russian": "ru",
|
||||
"korean": "ko",
|
||||
"french": "fr",
|
||||
"japanese": "ja",
|
||||
"portuguese": "pt",
|
||||
"turkish": "tr",
|
||||
"cantonese": "yue",
|
||||
# Intentional vLLM-Omni convenience alias; upstream would route the name
|
||||
# ``mandarin`` to the ``common`` language-embedding row.
|
||||
"mandarin": "zh",
|
||||
"burmese": "my",
|
||||
"valencian": "ca",
|
||||
"flemish": "nl",
|
||||
"haitian": "ht",
|
||||
"letzeburgesch": "lb",
|
||||
"pushto": "ps",
|
||||
"panjabi": "pa",
|
||||
"moldavian": "ro",
|
||||
"moldovan": "ro",
|
||||
"sinhalese": "si",
|
||||
"castilian": "es",
|
||||
}
|
||||
|
||||
_AUDIO_EVENTS = (
|
||||
"ASR",
|
||||
"AED",
|
||||
"SER",
|
||||
"Speech",
|
||||
"/Speech",
|
||||
"BGM",
|
||||
"/BGM",
|
||||
"Laughter",
|
||||
"/Laughter",
|
||||
"Applause",
|
||||
"/Applause",
|
||||
)
|
||||
_EMOTIONS = ("HAPPY", "SAD", "ANGRY", "NEUTRAL")
|
||||
_TTS_VOCAL_TOKENS = (
|
||||
"TTS/B",
|
||||
"TTS/O",
|
||||
"TTS/Q",
|
||||
"TTS/A",
|
||||
"TTS/CO",
|
||||
"TTS/CL",
|
||||
"TTS/H",
|
||||
*(f"TTS/SP{i:02d}" for i in range(1, 14)),
|
||||
)
|
||||
|
||||
INDEXTTS25_SPECIAL_TOKENS = (
|
||||
"<|endoftext|>",
|
||||
"<|startoftranscript|>",
|
||||
*(f"<|{lang}|>" for lang in LANGUAGES[:INDEXTTS25_NUM_SPECIAL_LANGUAGES]),
|
||||
*(f"<|{event}|>" for event in _AUDIO_EVENTS),
|
||||
*(f"<|{emotion}|>" for emotion in _EMOTIONS),
|
||||
"<|translate|>",
|
||||
"<|transcribe|>",
|
||||
"<|startoflm|>",
|
||||
"<|startofprev|>",
|
||||
"<|nospeech|>",
|
||||
"<|notimestamps|>",
|
||||
*(f"<|SPECIAL_TOKEN_{i}|>" for i in range(1, 31)),
|
||||
*(f"<|{token}|>" for token in _TTS_VOCAL_TOKENS),
|
||||
*(f"<|{i * 0.02:.2f}|>" for i in range(1501)),
|
||||
)
|
||||
|
||||
_PATTERN = r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
|
||||
|
||||
|
||||
def normalize_language_code(lang: str) -> str:
|
||||
normalized = lang.strip().lower()
|
||||
normalized = TO_LANGUAGE_CODE.get(normalized, normalized)
|
||||
if normalized not in LANGUAGE_DICT and normalized not in _FRONTEND_LANGUAGE_TO_EMBEDDING:
|
||||
supported = (*LANGUAGES, *sorted(_FRONTEND_LANGUAGE_TO_EMBEDDING))
|
||||
raise ValueError(f"Unsupported IndexTTS 2.5 language {lang!r}; expected one of {', '.join(supported)}")
|
||||
return normalized
|
||||
|
||||
|
||||
def lang_to_token(lang: str) -> int:
|
||||
normalized = normalize_language_code(lang)
|
||||
embedding_language = _FRONTEND_LANGUAGE_TO_EMBEDDING.get(normalized, normalized)
|
||||
return LANGUAGE_DICT[embedding_language]
|
||||
|
||||
|
||||
def resolve_indextts25_tokenizer_file(
|
||||
model_dir: str,
|
||||
tokenizer_file: str = INDEXTTS25_TOKENIZER_FILE,
|
||||
) -> str:
|
||||
model_path = Path(model_dir).expanduser()
|
||||
for root in (model_path, model_path / "checkpoints"):
|
||||
direct = root / tokenizer_file
|
||||
if direct.is_file():
|
||||
return str(direct)
|
||||
if model_path.is_file() and model_path.name == tokenizer_file:
|
||||
return str(model_path)
|
||||
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
return hf_hub_download(repo_id=model_dir, filename=tokenizer_file)
|
||||
except Exception as exc:
|
||||
raise FileNotFoundError(
|
||||
f"Could not resolve IndexTTS 2.5 tokenizer {tokenizer_file!r} from {model_dir!r}"
|
||||
) from exc
|
||||
|
||||
|
||||
@lru_cache(maxsize=16)
|
||||
def load_indextts25_encoding(
|
||||
model_dir: str,
|
||||
tokenizer_file: str = INDEXTTS25_TOKENIZER_FILE,
|
||||
) -> Any:
|
||||
import tiktoken
|
||||
|
||||
vocab_path = resolve_indextts25_tokenizer_file(model_dir, tokenizer_file)
|
||||
with open(vocab_path, encoding="utf-8") as vocab:
|
||||
ranks = {
|
||||
base64.b64decode(token): int(rank) for line in vocab if line.strip() for token, rank in (line.split(),)
|
||||
}
|
||||
|
||||
if len(ranks) != INDEXTTS25_MERGEABLE_RANKS:
|
||||
raise ValueError(
|
||||
f"IndexTTS 2.5 tokenizer {vocab_path!r} has {len(ranks)} "
|
||||
f"mergeable ranks; expected {INDEXTTS25_MERGEABLE_RANKS}"
|
||||
)
|
||||
special_tokens = {token: len(ranks) + index for index, token in enumerate(INDEXTTS25_SPECIAL_TOKENS)}
|
||||
explicit_n_vocab = len(ranks) + len(special_tokens)
|
||||
if explicit_n_vocab != INDEXTTS25_VOCAB_SIZE:
|
||||
raise ValueError(
|
||||
f"IndexTTS 2.5 tokenizer builds vocab size {explicit_n_vocab}; expected {INDEXTTS25_VOCAB_SIZE}"
|
||||
)
|
||||
return tiktoken.Encoding(
|
||||
name=Path(vocab_path).name,
|
||||
explicit_n_vocab=explicit_n_vocab,
|
||||
pat_str=_PATTERN,
|
||||
mergeable_ranks=ranks,
|
||||
special_tokens=special_tokens,
|
||||
)
|
||||
|
||||
|
||||
def encode_indextts25_text(
|
||||
text: str,
|
||||
*,
|
||||
model_dir: str,
|
||||
tokenizer_file: str = INDEXTTS25_TOKENIZER_FILE,
|
||||
) -> list[int]:
|
||||
encoding = load_indextts25_encoding(model_dir, tokenizer_file)
|
||||
return encoding.encode(text, allowed_special="all")
|
||||
@@ -369,6 +369,17 @@ _OMNI_MODELS = {
|
||||
"indextts2_s2mel_decoder",
|
||||
"IndexTTS2S2MelDecoder",
|
||||
),
|
||||
## IndexTTS 2.5 (shared implementations, versioned HF config)
|
||||
"IndexTTS25TalkerForConditionalGeneration": (
|
||||
"indextts2",
|
||||
"indextts2_talker",
|
||||
"IndexTTS2TalkerForConditionalGeneration",
|
||||
),
|
||||
"IndexTTS25S2MelDecoder": (
|
||||
"indextts2",
|
||||
"indextts2_s2mel_decoder",
|
||||
"IndexTTS2S2MelDecoder",
|
||||
),
|
||||
## Ming-flash-omni-2.0
|
||||
"MingFlashOmniForConditionalGeneration": (
|
||||
"ming_flash_omni",
|
||||
|
||||
@@ -26,21 +26,21 @@ def _cpu_view(tensor: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
def _strip_stop_token(
|
||||
codes: torch.Tensor,
|
||||
latent: torch.Tensor,
|
||||
latent: torch.Tensor | None,
|
||||
stop_mel_token: int = STOP_MEL_TOKEN,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]:
|
||||
"""Strip at the first stop token, matching official IndexTTS2 v2.
|
||||
|
||||
Returns: (codes [B, T'], latent [B, T', D], code_lens [B]).
|
||||
Returns codes, optional aligned latent, and code lengths.
|
||||
"""
|
||||
if codes.ndim == 1:
|
||||
codes = codes.unsqueeze(0)
|
||||
if latent.ndim == 2:
|
||||
if latent is not None and latent.ndim == 2:
|
||||
latent = latent.unsqueeze(0)
|
||||
|
||||
if codes.device.type != "cpu":
|
||||
codes = codes.detach().cpu()
|
||||
if latent.device.type != "cpu":
|
||||
if latent is not None and latent.device.type != "cpu":
|
||||
latent = latent.detach().cpu()
|
||||
|
||||
device = codes.device
|
||||
@@ -50,8 +50,6 @@ def _strip_stop_token(
|
||||
|
||||
for i in range(codes.shape[0]):
|
||||
code = codes[i]
|
||||
lat = latent[i]
|
||||
|
||||
stop_mask = (code == stop_mel_token).nonzero(as_tuple=False)
|
||||
if stop_mask.numel() > 0:
|
||||
valid_len = int(stop_mask[0].item())
|
||||
@@ -59,27 +57,50 @@ def _strip_stop_token(
|
||||
valid_len = int(code.shape[0])
|
||||
code_lens.append(valid_len)
|
||||
codes_out.append(code[:valid_len])
|
||||
latent_out.append(lat[:valid_len])
|
||||
if latent is not None:
|
||||
latent_out.append(latent[i, :valid_len])
|
||||
|
||||
max_len = max(code_lens) if code_lens else 0
|
||||
if max_len == 0:
|
||||
empty_codes = torch.zeros(codes.shape[0], 0, dtype=torch.long, device=device)
|
||||
empty_latent = (
|
||||
torch.zeros(
|
||||
codes.shape[0],
|
||||
0,
|
||||
latent.shape[-1],
|
||||
device=device,
|
||||
dtype=latent.dtype,
|
||||
)
|
||||
if latent is not None
|
||||
else None
|
||||
)
|
||||
return (
|
||||
torch.zeros(codes.shape[0], 0, dtype=torch.long, device=device),
|
||||
torch.zeros(codes.shape[0], 0, latent.shape[-1], device=device, dtype=latent.dtype),
|
||||
torch.zeros(codes.shape[0], dtype=torch.long, device=device),
|
||||
empty_codes,
|
||||
empty_latent,
|
||||
torch.zeros(
|
||||
codes.shape[0],
|
||||
dtype=torch.long,
|
||||
device=device,
|
||||
),
|
||||
)
|
||||
|
||||
padded_codes = torch.zeros((len(codes_out), max_len), dtype=torch.long, device=device)
|
||||
padded_latent = torch.zeros(
|
||||
len(latent_out),
|
||||
max_len,
|
||||
latent_out[0].shape[-1],
|
||||
device=device,
|
||||
dtype=latent_out[0].dtype,
|
||||
padded_latent = (
|
||||
torch.zeros(
|
||||
len(latent_out),
|
||||
max_len,
|
||||
latent_out[0].shape[-1],
|
||||
device=device,
|
||||
dtype=latent_out[0].dtype,
|
||||
)
|
||||
if latent_out
|
||||
else None
|
||||
)
|
||||
for i, (c, lat) in enumerate(zip(codes_out, latent_out)):
|
||||
for i, c in enumerate(codes_out):
|
||||
padded_codes[i, : c.shape[0]] = c
|
||||
padded_latent[i, : lat.shape[0]] = lat
|
||||
if padded_latent is not None:
|
||||
lat = latent_out[i]
|
||||
padded_latent[i, : lat.shape[0]] = lat
|
||||
|
||||
return padded_codes, padded_latent, torch.tensor(code_lens, dtype=torch.long, device=device)
|
||||
|
||||
@@ -105,26 +126,33 @@ def _normalize_latent_sequence(latent: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
def _build_s2mel_additional_information(
|
||||
mel_codes: torch.Tensor,
|
||||
latent: torch.Tensor,
|
||||
latent: torch.Tensor | None,
|
||||
meta: dict[str, Any],
|
||||
*,
|
||||
use_gpt_latent: bool,
|
||||
context: str,
|
||||
) -> dict[str, Any]:
|
||||
"""Build the Stage-1 S2Mel tensor contract shared by legacy and connector paths."""
|
||||
mel_codes_clean, latent_clean, code_lens = _strip_stop_token(mel_codes, latent)
|
||||
|
||||
additional_information = {
|
||||
"latent": _cpu_view(latent_clean),
|
||||
"mel_codes": _cpu_view(mel_codes_clean),
|
||||
"code_lens": _cpu_view(code_lens),
|
||||
"use_gpt_latent": use_gpt_latent,
|
||||
}
|
||||
if use_gpt_latent:
|
||||
if latent_clean is None:
|
||||
raise ValueError(f"{context} requires GPT latent but none was provided")
|
||||
additional_information["latent"] = _cpu_view(latent_clean)
|
||||
|
||||
duration_factor = meta.get("duration_factor", 1.0)
|
||||
additional_information["duration_factor"] = 1.0 if duration_factor is None else float(duration_factor)
|
||||
|
||||
for key in ("S_ref", "ref_mel", "style"):
|
||||
val = meta.get(key)
|
||||
if isinstance(val, torch.Tensor):
|
||||
additional_information[key] = _cpu_view(val)
|
||||
else:
|
||||
logger.warning("[%s] %s MISSING — Stage 1 will use fallback", context, key)
|
||||
if not isinstance(val, torch.Tensor) or val.numel() == 0:
|
||||
raise ValueError(f"{context} requires non-empty IndexTTS conditioning field {key}")
|
||||
additional_information[key] = _cpu_view(val)
|
||||
|
||||
return additional_information
|
||||
|
||||
@@ -143,6 +171,12 @@ def _request_id(request: Any) -> str:
|
||||
return str(getattr(request, "external_req_id", None) or getattr(request, "request_id", "?"))
|
||||
|
||||
|
||||
def _request_seed(request: Any) -> int | None:
|
||||
sampling_params = getattr(request, "sampling_params", None)
|
||||
seed = getattr(sampling_params, "seed", None)
|
||||
return int(seed) if seed is not None else None
|
||||
|
||||
|
||||
def talker2s2mel_token_only(
|
||||
source_outputs: list[Any],
|
||||
prompt: Any = None,
|
||||
@@ -177,36 +211,78 @@ def talker2s2mel_full_payload(
|
||||
del transfer_manager
|
||||
rid = _request_id(request)
|
||||
if not isinstance(pooling_output, dict):
|
||||
logger.warning("talker2s2mel_full_payload: pooling_output not a dict for req=%s", rid)
|
||||
return None
|
||||
raise ValueError(f"IndexTTS Stage 0 payload must be a dict for req={rid}")
|
||||
|
||||
use_gpt_latent_value = _get_payload_value(
|
||||
pooling_output,
|
||||
"meta.use_gpt_latent",
|
||||
"meta",
|
||||
"use_gpt_latent",
|
||||
)
|
||||
if use_gpt_latent_value is None:
|
||||
raise ValueError(f"IndexTTS payload is missing meta.use_gpt_latent for req={rid}")
|
||||
use_gpt_latent = bool(use_gpt_latent_value)
|
||||
payload_latent = _get_payload_value(
|
||||
pooling_output,
|
||||
"hidden_states.latent",
|
||||
"hidden_states",
|
||||
"latent",
|
||||
)
|
||||
payload_mel_codes = _get_payload_value(
|
||||
pooling_output,
|
||||
"codes.mel",
|
||||
"codes",
|
||||
"mel",
|
||||
)
|
||||
if use_gpt_latent:
|
||||
mel_codes = payload_mel_codes
|
||||
latent = payload_latent
|
||||
else:
|
||||
output_token_ids = list(getattr(request, "output_token_ids", None) or [])
|
||||
if not output_token_ids:
|
||||
raise ValueError(f"IndexTTS code-only payload has no completed output_token_ids for req={rid}")
|
||||
mel_codes = torch.tensor(output_token_ids, dtype=torch.long)
|
||||
latent = None
|
||||
|
||||
mel_codes = _get_payload_value(pooling_output, "codes.mel", "codes", "mel")
|
||||
if not isinstance(mel_codes, torch.Tensor) or mel_codes.numel() == 0:
|
||||
logger.warning("talker2s2mel_full_payload: missing codes.mel for req=%s", rid)
|
||||
return None
|
||||
|
||||
latent = _get_payload_value(pooling_output, "hidden_states.latent", "hidden_states", "latent")
|
||||
if not isinstance(latent, torch.Tensor) or latent.numel() == 0:
|
||||
logger.warning("talker2s2mel_full_payload: missing hidden_states.latent for req=%s", rid)
|
||||
return None
|
||||
raise ValueError(f"IndexTTS payload is missing mel codes for req={rid} (use_gpt_latent={use_gpt_latent})")
|
||||
|
||||
mel_seq = _normalize_mel_sequence(mel_codes)
|
||||
latent_seq = _normalize_latent_sequence(latent)
|
||||
if mel_seq.numel() == 0 or latent_seq.numel() == 0:
|
||||
logger.warning("talker2s2mel_full_payload: empty normalized mel/latent for req=%s", rid)
|
||||
return None
|
||||
if mel_seq.numel() == 0:
|
||||
raise ValueError(f"IndexTTS payload has an empty mel sequence for req={rid}")
|
||||
|
||||
common_len = min(int(mel_seq.shape[0]), int(latent_seq.shape[0]))
|
||||
if common_len <= 0:
|
||||
logger.warning("talker2s2mel_full_payload: no common mel/latent length for req=%s", rid)
|
||||
return None
|
||||
|
||||
mel_seq = mel_seq[:common_len]
|
||||
latent_seq = latent_seq[:common_len]
|
||||
latent_seq: torch.Tensor | None = None
|
||||
if use_gpt_latent:
|
||||
if not isinstance(latent, torch.Tensor) or latent.numel() == 0:
|
||||
raise ValueError(f"IndexTTS latent payload is missing hidden_states.latent for req={rid}")
|
||||
latent_seq = _normalize_latent_sequence(latent)
|
||||
if latent_seq.numel() == 0:
|
||||
raise ValueError(f"IndexTTS payload has an empty latent sequence for req={rid}")
|
||||
if mel_seq.shape[0] != latent_seq.shape[0]:
|
||||
raise ValueError(
|
||||
f"IndexTTS mel/latent length mismatch for req={rid}: "
|
||||
f"mel={mel_seq.shape[0]} latent={latent_seq.shape[0]}"
|
||||
)
|
||||
|
||||
meta = {
|
||||
"S_ref": _get_payload_value(pooling_output, "meta.S_ref", "meta", "S_ref"),
|
||||
"ref_mel": _get_payload_value(pooling_output, "meta.ref_mel", "meta", "ref_mel"),
|
||||
"style": _get_payload_value(pooling_output, "meta.style", "meta", "style"),
|
||||
"duration_factor": _get_payload_value(
|
||||
pooling_output,
|
||||
"meta.duration_factor",
|
||||
"meta",
|
||||
"duration_factor",
|
||||
),
|
||||
}
|
||||
return _build_s2mel_additional_information(mel_seq, latent_seq, meta, context="full_payload")
|
||||
additional_information = _build_s2mel_additional_information(
|
||||
mel_seq,
|
||||
latent_seq,
|
||||
meta,
|
||||
use_gpt_latent=use_gpt_latent,
|
||||
context="full_payload",
|
||||
)
|
||||
seed = _request_seed(request)
|
||||
if seed is not None:
|
||||
additional_information["seed"] = seed
|
||||
return additional_information
|
||||
@@ -41,6 +41,10 @@ class OmniRequest(Request):
|
||||
if prompt_embeds is not None:
|
||||
kwargs["prompt_embeds"] = self._maybe_decode_prompt_embeds(prompt_embeds)
|
||||
super().__init__(*args, **kwargs)
|
||||
# vLLM 0.27 owns this counter; accelerator images still based on 0.26
|
||||
# do not. Keep the Omni scheduler's stale-output drain compatible with
|
||||
# both request layouts until those images move to 0.27.
|
||||
self.num_stale_output_tokens = int(getattr(self, "num_stale_output_tokens", 0) or 0)
|
||||
# Preserve serialized prompt embeddings payload (optional)
|
||||
self.prompt_embeds_payload: PromptEmbedsPayload | None = (
|
||||
prompt_embeds if isinstance(prompt_embeds, PromptEmbedsPayload) else None
|
||||
|
||||
@@ -106,6 +106,60 @@ def build_mm_cpu(multimodal_outputs: dict) -> dict[str, object]:
|
||||
return mm_cpu
|
||||
|
||||
|
||||
def snapshot_mm_payload(multimodal_outputs: dict) -> dict[str, object]:
|
||||
"""Snapshot a multimodal payload without moving it off its device.
|
||||
|
||||
Request-end full-payload producers retain these tensors across decode
|
||||
steps. CUDA graph outputs and runner input buffers may be reused on the
|
||||
next step, so retaining views is not sufficient: CUDA tensors must be
|
||||
copied first. Compatible per-request tensor lists are packed with one
|
||||
``torch.cat`` and restored as views, avoiding one tiny copy kernel per
|
||||
active request.
|
||||
"""
|
||||
if not multimodal_outputs:
|
||||
return {}
|
||||
return {key: _snapshot_payload_value(value) for key, value in multimodal_outputs.items()}
|
||||
|
||||
|
||||
def _snapshot_payload_value(value):
|
||||
if isinstance(value, torch.Tensor):
|
||||
tensor = value.detach()
|
||||
return tensor.clone() if tensor.device.type != "cpu" else tensor
|
||||
if isinstance(value, dict):
|
||||
return {key: _snapshot_payload_value(item) for key, item in value.items()}
|
||||
if isinstance(value, list):
|
||||
if not value:
|
||||
return value
|
||||
first = value[0]
|
||||
if (
|
||||
len(value) > 1
|
||||
and isinstance(first, torch.Tensor)
|
||||
and first.ndim > 0
|
||||
and first.layout == torch.strided
|
||||
and all(
|
||||
isinstance(item, torch.Tensor)
|
||||
and item.ndim == first.ndim
|
||||
and item.layout == first.layout
|
||||
and item.device == first.device
|
||||
and item.dtype == first.dtype
|
||||
and tuple(item.shape[1:]) == tuple(first.shape[1:])
|
||||
for item in value
|
||||
)
|
||||
):
|
||||
packed = torch.cat([item.detach() for item in value], dim=0)
|
||||
output = []
|
||||
offset = 0
|
||||
for item in value:
|
||||
length = item.shape[0]
|
||||
output.append(packed.narrow(0, offset, length))
|
||||
offset += length
|
||||
return output
|
||||
return [_snapshot_payload_value(item) for item in value]
|
||||
if isinstance(value, tuple):
|
||||
return tuple(_snapshot_payload_value(item) for item in value)
|
||||
return value
|
||||
|
||||
|
||||
def _to_cpu(value):
|
||||
"""Recursively detach + move tensors to CPU; preserve dict/list nesting."""
|
||||
if isinstance(value, torch.Tensor):
|
||||
|
||||
@@ -46,7 +46,12 @@ from vllm_omni.distributed.omni_connectors.utils.config import (
|
||||
stage_sends_async_output,
|
||||
)
|
||||
from vllm_omni.outputs import OmniModelRunnerOutput
|
||||
from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_payload_list, to_payload_element
|
||||
from vllm_omni.utils.mm_outputs import (
|
||||
build_mm_cpu,
|
||||
partition_payload_list,
|
||||
snapshot_mm_payload,
|
||||
to_payload_element,
|
||||
)
|
||||
from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
|
||||
from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
|
||||
from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest
|
||||
@@ -1715,6 +1720,35 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
|
||||
def _model_omni_pooler_payload_include_hidden(self) -> bool:
|
||||
return self._runner_model_omni_flag("omni_pooler_payload_include_hidden", default=True)
|
||||
|
||||
def _should_defer_full_payload_d2h(self) -> bool:
|
||||
"""Keep opted-in full payloads on device until request completion."""
|
||||
return (
|
||||
getattr(self, "omni_prefix_cache", None) is None
|
||||
and self._runner_model_omni_flag("omni_payload_at_request_end")
|
||||
and self._should_accumulate_full_payload_output()
|
||||
)
|
||||
|
||||
def _build_omni_step_outputs(
|
||||
self,
|
||||
pooler_inter: list[dict[str, object]],
|
||||
pooler_client: list[dict[str, object]],
|
||||
*,
|
||||
defer_full_payload_d2h: bool,
|
||||
) -> tuple[
|
||||
list[dict[str, torch.Tensor] | None] | None,
|
||||
list[dict[str, torch.Tensor] | None] | None,
|
||||
]:
|
||||
if defer_full_payload_d2h:
|
||||
# The connector-side full-payload accumulator owns the snapshots.
|
||||
# Nothing from an intermediate stage needs to cross the worker/core
|
||||
# boundary on each decode step.
|
||||
return None, None
|
||||
inter_stage_outputs = self._build_multimodal_outputs(pooler_inter)
|
||||
multimodal_outputs = (
|
||||
inter_stage_outputs if pooler_client is pooler_inter else self._build_multimodal_outputs(pooler_client)
|
||||
)
|
||||
return inter_stage_outputs, multimodal_outputs
|
||||
|
||||
def _should_use_async_omni_output(self) -> bool:
|
||||
if not self.use_async_scheduling:
|
||||
return False
|
||||
@@ -1879,6 +1913,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
|
||||
|
||||
needs_pooler_payload = len(downstream_req_ids) > 0
|
||||
downstream_req_id_set = set(downstream_req_ids)
|
||||
defer_full_payload_d2h = needs_pooler_payload and self._should_defer_full_payload_d2h()
|
||||
hidden_states_cpu = None
|
||||
req_hidden_states_cpu: dict[str, torch.Tensor] | None = None
|
||||
include_hidden_payload = self._model_omni_pooler_payload_include_hidden()
|
||||
@@ -1933,10 +1968,13 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
|
||||
needs_scheduled_hidden_payload=needs_scheduled_hidden_payload,
|
||||
)
|
||||
if combined_multimodal_outputs is None:
|
||||
with record_function_or_nullcontext("omni_output_builder:build_mm_cpu"):
|
||||
mm_cpu = build_mm_cpu(
|
||||
flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs
|
||||
)
|
||||
flat_mm = flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs
|
||||
if defer_full_payload_d2h:
|
||||
with record_function_or_nullcontext("omni_output_builder:snapshot_mm_payload"):
|
||||
mm_cpu = snapshot_mm_payload(flat_mm)
|
||||
else:
|
||||
with record_function_or_nullcontext("omni_output_builder:build_mm_cpu"):
|
||||
mm_cpu = build_mm_cpu(flat_mm)
|
||||
|
||||
with record_function_or_nullcontext("omni_output_builder:process_additional_information"):
|
||||
if not postprocess_already_applied:
|
||||
@@ -1995,9 +2033,10 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
|
||||
self.accumulate_full_payload_output(rid, pooler_inter[i], req_state)
|
||||
|
||||
with record_function_or_nullcontext("omni_output_builder:build_multimodal_outputs"):
|
||||
inter_stage_outputs = self._build_multimodal_outputs(pooler_inter)
|
||||
multimodal_outputs = (
|
||||
inter_stage_outputs if pooler_client is pooler_inter else self._build_multimodal_outputs(pooler_client)
|
||||
inter_stage_outputs, multimodal_outputs = self._build_omni_step_outputs(
|
||||
pooler_inter,
|
||||
pooler_client,
|
||||
defer_full_payload_d2h=defer_full_payload_d2h,
|
||||
)
|
||||
|
||||
with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):
|
||||
|
||||
Reference in new issue
Block a user