From 58007c0084b2a4411da57811e74290365ac80f04 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Wed, 12 Aug 2026 07:00:35 +0800 Subject: [PATCH] [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> --- .buildkite/intel/pipeline-intel.yml | 2 +- .buildkite/intel/scripts/run-xpu-test.sh | 17 +- benchmarks/tts/README.md | 28 + benchmarks/tts/bench_tts.py | 33 +- benchmarks/tts/model_configs.yaml | 11 + docker/Dockerfile.xpu | 9 +- docs/models/supported_models.md | 1 + .../text_to_speech/README.md | 37 +- .../text_to_speech/indextts2/end2end.py | 114 +- .../online_serving/text_to_speech/README.md | 48 +- .../text_to_speech/indextts2/run_server.sh | 27 +- .../text_to_speech/indextts2/speech_client.py | 39 +- pyproject.toml | 11 + recipes/IndexTeam/IndexTTS-2_5.md | 308 +++++ recipes/README.md | 1 + tests/benchmarks/test_bench_tts_cli.py | 36 + tests/config/test_config_factory.py | 10 + .../sched/test_omni_scheduling_coordinator.py | 1 + tests/engine/test_arg_utils.py | 152 +++ .../test_async_omni_engine_stage_init.py | 50 +- tests/engine/test_omni_request.py | 1 + .../openai_api/test_serving_speech.py | 110 +- .../openai_api/test_tts_adapter.py | 122 ++ .../openai_api/test_tts_detection.py | 20 +- .../test_indextts2_run_server.py | 108 ++ .../models/indextts2/__init__.py | 0 .../test_configuration_and_pipeline.py | 89 ++ .../models/indextts2/test_dit_cuda_graph.py | 75 ++ .../models/indextts2/test_enhanced_codec.py | 84 ++ .../models/indextts2/test_prompt_utils.py | 202 ++++ .../models/indextts2/test_recipe_defaults.py | 14 + .../models/indextts2/test_s2mel_v2_5.py | 1066 +++++++++++++++++ .../indextts2/test_stage_input_processor.py | 53 + .../models/indextts2/test_talker_v2_5.py | 476 ++++++++ .../models/indextts2/test_tokenizer_v2_5.py | 244 ++++ .../models/indextts2/test_wavenet_ragged.py | 196 +++ .../qwen3_tts/test_code_predictor_dtype.py | 23 + .../stage_input_processors/test_indextts2.py | 86 +- tests/test_indextts2_optional_dependencies.py | 155 +++ tests/utils/test_mm_outputs_partition.py | 31 + tests/worker/test_omni_gpu_model_runner.py | 44 + vllm_omni/config/pipeline_registry.py | 6 +- .../core/sched/omni_scheduling_coordinator.py | 5 +- vllm_omni/deploy/indextts2_5.yaml | 64 + vllm_omni/deploy/qwen2_5_omni.yaml | 12 +- vllm_omni/engine/arg_utils.py | 14 +- vllm_omni/engine/stage_init_utils.py | 22 +- vllm_omni/entrypoints/openai/api_server.py | 8 +- .../entrypoints/openai/protocol/audio.py | 13 +- vllm_omni/entrypoints/openai/serving_chat.py | 9 +- .../entrypoints/openai/serving_speech.py | 16 +- .../entrypoints/openai/tts_adapters/base.py | 2 + .../openai/tts_adapters/indextts2.py | 86 +- .../models/common/qwen3_code_predictor.py | 5 +- .../models/indextts2/codec/__init__.py | 7 + .../models/indextts2/codec/models.py | 138 +++ .../models/indextts2/codec/quantizer.py | 91 ++ .../models/indextts2/codec/vocos.py | 82 ++ .../indextts2/configuration_indextts2.py | 47 + .../models/indextts2/dit_cuda_graph.py | 55 +- .../indextts2/indextts2_s2mel_decoder.py | 901 +++++++++++--- .../models/indextts2/indextts2_talker.py | 471 ++++++-- .../models/indextts2/pipeline.py | 38 +- .../models/indextts2/preprocess_utils.py | 145 ++- .../models/indextts2/prompt_utils.py | 62 +- .../models/indextts2/s2mel/modules/commons.py | 4 +- .../s2mel/modules/diffusion_transformer.py | 23 +- .../indextts2/s2mel/modules/flow_matching.py | 136 ++- .../indextts2/s2mel/modules/gpt_fast/model.py | 67 +- .../models/indextts2/s2mel/modules/wavenet.py | 71 +- .../models/indextts2/text_processing_v2_5.py | 328 +++++ .../models/indextts2/tokenizer_v2_5.py | 286 +++++ vllm_omni/model_executor/models/registry.py | 11 + .../stage_input_processors/indextts2.py | 170 ++- vllm_omni/request.py | 4 + vllm_omni/utils/mm_outputs.py | 54 + vllm_omni/worker/gpu_ar_model_runner.py | 55 +- 77 files changed, 7123 insertions(+), 519 deletions(-) create mode 100644 recipes/IndexTeam/IndexTTS-2_5.md create mode 100644 tests/examples/online_serving/test_indextts2_run_server.py create mode 100644 tests/model_executor/models/indextts2/__init__.py create mode 100644 tests/model_executor/models/indextts2/test_configuration_and_pipeline.py create mode 100644 tests/model_executor/models/indextts2/test_dit_cuda_graph.py create mode 100644 tests/model_executor/models/indextts2/test_enhanced_codec.py create mode 100644 tests/model_executor/models/indextts2/test_prompt_utils.py create mode 100644 tests/model_executor/models/indextts2/test_recipe_defaults.py create mode 100644 tests/model_executor/models/indextts2/test_s2mel_v2_5.py create mode 100644 tests/model_executor/models/indextts2/test_stage_input_processor.py create mode 100644 tests/model_executor/models/indextts2/test_talker_v2_5.py create mode 100644 tests/model_executor/models/indextts2/test_tokenizer_v2_5.py create mode 100644 tests/model_executor/models/indextts2/test_wavenet_ragged.py create mode 100644 tests/test_indextts2_optional_dependencies.py create mode 100644 vllm_omni/deploy/indextts2_5.yaml create mode 100644 vllm_omni/model_executor/models/indextts2/codec/__init__.py create mode 100644 vllm_omni/model_executor/models/indextts2/codec/models.py create mode 100644 vllm_omni/model_executor/models/indextts2/codec/quantizer.py create mode 100644 vllm_omni/model_executor/models/indextts2/codec/vocos.py create mode 100644 vllm_omni/model_executor/models/indextts2/text_processing_v2_5.py create mode 100644 vllm_omni/model_executor/models/indextts2/tokenizer_v2_5.py diff --git a/.buildkite/intel/pipeline-intel.yml b/.buildkite/intel/pipeline-intel.yml index acb135101..2a1507456 100644 --- a/.buildkite/intel/pipeline-intel.yml +++ b/.buildkite/intel/pipeline-intel.yml @@ -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 diff --git a/.buildkite/intel/scripts/run-xpu-test.sh b/.buildkite/intel/scripts/run-xpu-test.sh index 6b40d58fb..742e39b78 100755 --- a/.buildkite/intel/scripts/run-xpu-test.sh +++ b/.buildkite/intel/scripts/run-xpu-test.sh @@ -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 diff --git a/benchmarks/tts/README.md b/benchmarks/tts/README.md index f6930a866..d3a0f9676 100644 --- a/benchmarks/tts/README.md +++ b/benchmarks/tts/README.md @@ -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 diff --git a/benchmarks/tts/bench_tts.py b/benchmarks/tts/bench_tts.py index c79d74a31..f8c1676be 100644 --- a/benchmarks/tts/bench_tts.py +++ b/benchmarks/tts/bench_tts.py @@ -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: diff --git a/benchmarks/tts/model_configs.yaml b/benchmarks/tts/model_configs.yaml index 2ec7addca..5450e8a7c 100644 --- a/benchmarks/tts/model_configs.yaml +++ b/benchmarks/tts/model_configs.yaml @@ -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. diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index 5d298040c..65ae05dd3 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -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 diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index add8b86b3..b9686047a 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -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` | ✅︎ | ✅︎ | | | — | diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index 180ab9954..56135b55e 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -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. --- diff --git a/examples/offline_inference/text_to_speech/indextts2/end2end.py b/examples/offline_inference/text_to_speech/indextts2/end2end.py index 858487ae1..084e2cc72 100644 --- a/examples/offline_inference/text_to_speech/indextts2/end2end.py +++ b/examples/offline_inference/text_to_speech/indextts2/end2end.py @@ -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__": diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index b331966b3..2f9d84bf3 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -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. --- diff --git a/examples/online_serving/text_to_speech/indextts2/run_server.sh b/examples/online_serving/text_to_speech/indextts2/run_server.sh index 4a17b3b04..c904fb19f 100644 --- a/examples/online_serving/text_to_speech/indextts2/run_server.sh +++ b/examples/online_serving/text_to_speech/indextts2/run_server.sh @@ -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" \ diff --git a/examples/online_serving/text_to_speech/indextts2/speech_client.py b/examples/online_serving/text_to_speech/indextts2/speech_client.py index b938f0bd3..65342a335 100644 --- a/examples/online_serving/text_to_speech/indextts2/speech_client.py +++ b/examples/online_serving/text_to_speech/indextts2/speech_client.py @@ -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 diff --git a/pyproject.toml b/pyproject.toml index bbb8433f9..f9fbdfb2a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/recipes/IndexTeam/IndexTTS-2_5.md b/recipes/IndexTeam/IndexTTS-2_5.md new file mode 100644 index 000000000..b5b741c7e --- /dev/null +++ b/recipes/IndexTeam/IndexTTS-2_5.md @@ -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,", + "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. diff --git a/recipes/README.md b/recipes/README.md index eb4016482..ca287bf70 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -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 | diff --git a/tests/benchmarks/test_bench_tts_cli.py b/tests/benchmarks/test_bench_tts_cli.py index b8a487f80..e32d94698 100644 --- a/tests/benchmarks/test_bench_tts_cli.py +++ b/tests/benchmarks/test_bench_tts_cli.py @@ -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( diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index eb774bb18..c3c3edbd1 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -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(): diff --git a/tests/core/sched/test_omni_scheduling_coordinator.py b/tests/core/sched/test_omni_scheduling_coordinator.py index a36a94e7d..a93929306 100644 --- a/tests/core/sched/test_omni_scheduling_coordinator.py +++ b/tests/core/sched/test_omni_scheduling_coordinator.py @@ -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"), } diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 2e2f75194..055866bd5 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -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() diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 67765aba6..27aee8706 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -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 diff --git a/tests/engine/test_omni_request.py b/tests/engine/test_omni_request.py index 4f4718f09..521c92b94 100644 --- a/tests/engine/test_omni_request.py +++ b/tests/engine/test_omni_request.py @@ -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(): diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index e1efd9286..32de8ac62 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -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.""" diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py index 89106900b..9ab744c83 100644 --- a/tests/entrypoints/openai_api/test_tts_adapter.py +++ b/tests/entrypoints/openai_api/test_tts_adapter.py @@ -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" diff --git a/tests/entrypoints/openai_api/test_tts_detection.py b/tests/entrypoints/openai_api/test_tts_detection.py index 27c6082c1..c63e5c8e4 100644 --- a/tests/entrypoints/openai_api/test_tts_detection.py +++ b/tests/entrypoints/openai_api/test_tts_detection.py @@ -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(): diff --git a/tests/examples/online_serving/test_indextts2_run_server.py b/tests/examples/online_serving/test_indextts2_run_server.py new file mode 100644 index 000000000..021706a7b --- /dev/null +++ b/tests/examples/online_serving/test_indextts2_run_server.py @@ -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 diff --git a/tests/model_executor/models/indextts2/__init__.py b/tests/model_executor/models/indextts2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/model_executor/models/indextts2/test_configuration_and_pipeline.py b/tests/model_executor/models/indextts2/test_configuration_and_pipeline.py new file mode 100644 index 000000000..f01a725d2 --- /dev/null +++ b/tests/model_executor/models/indextts2/test_configuration_and_pipeline.py @@ -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, + ] diff --git a/tests/model_executor/models/indextts2/test_dit_cuda_graph.py b/tests/model_executor/models/indextts2/test_dit_cuda_graph.py new file mode 100644 index 000000000..6bb476fb9 --- /dev/null +++ b/tests/model_executor/models/indextts2/test_dit_cuda_graph.py @@ -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, + } diff --git a/tests/model_executor/models/indextts2/test_enhanced_codec.py b/tests/model_executor/models/indextts2/test_enhanced_codec.py new file mode 100644 index 000000000..2f82eaf91 --- /dev/null +++ b/tests/model_executor/models/indextts2/test_enhanced_codec.py @@ -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) diff --git a/tests/model_executor/models/indextts2/test_prompt_utils.py b/tests/model_executor/models/indextts2/test_prompt_utils.py new file mode 100644 index 000000000..fe73e69a5 --- /dev/null +++ b/tests/model_executor/models/indextts2/test_prompt_utils.py @@ -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" diff --git a/tests/model_executor/models/indextts2/test_recipe_defaults.py b/tests/model_executor/models/indextts2/test_recipe_defaults.py new file mode 100644 index 000000000..2ed0ba977 --- /dev/null +++ b/tests/model_executor/models/indextts2/test_recipe_defaults.py @@ -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 diff --git a/tests/model_executor/models/indextts2/test_s2mel_v2_5.py b/tests/model_executor/models/indextts2/test_s2mel_v2_5.py new file mode 100644 index 000000000..7d7e6e8d2 --- /dev/null +++ b/tests/model_executor/models/indextts2/test_s2mel_v2_5.py @@ -0,0 +1,1066 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +import yaml + +from vllm_omni.model_executor.models.indextts2 import indextts2_s2mel_decoder +from vllm_omni.model_executor.models.indextts2.indextts2_s2mel_decoder import ( + IndexTTS2S2MelDecoder, +) +from vllm_omni.model_executor.models.indextts2.s2mel.modules import ( + flow_matching, +) +from vllm_omni.model_executor.models.indextts2.s2mel.modules.flow_matching import ( + BASECFM, + CFM, +) +from vllm_omni.model_executor.models.indextts2.s2mel.modules.gpt_fast import ( + model as gpt_fast_model, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _FakeEnhancedCodec: + downsample_scale = 2 + + def __init__(self): + self.decode_calls = 0 + + def decode(self, codes): + self.decode_calls += 1 + return torch.ones(codes.shape[0], codes.shape[1] * 2, 4) + + +class _FakeS2Mel: + def forward_gpt(self, latent): + return torch.full((*latent.shape[:-1], 4), 2.0) + + +class _FakeBigVGAN: + def __init__(self): + self.events = [] + + def remove_weight_norm(self): + self.events.append("remove_weight_norm") + + def to(self, *, device, dtype): + self.events.append(("to", device, dtype)) + return self + + def eval(self): + self.events.append("eval") + return self + + def parameters(self): + return [] + + +class _DefaultDtypeBigVGAN(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(1)) + self.fold_weight_dtype = None + self.fold_default_dtype = None + + def remove_weight_norm(self): + self.fold_weight_dtype = self.weight.dtype + self.fold_default_dtype = torch.get_default_dtype() + + +class _LengthRecordingBigVGAN(torch.nn.Module): + def __init__(self): + super().__init__() + self.input_lengths = [] + + def forward(self, mel): + self.input_lengths.append(mel.shape[-1]) + return torch.zeros( + mel.shape[0], + 1, + mel.shape[-1] * 256, + dtype=mel.dtype, + device=mel.device, + ) + + +class _NoiseEchoCFM(BASECFM): + def __init__(self): + torch.nn.Module.__init__(self) + self.in_channels = 2 + + def solve_euler( + self, + x, + x_lens, + prompt, + mu, + style, + f0, + t_span, + inference_cfg_rate=0.5, + prompt_lens=None, + unpad_data=None, + ): + del unpad_data + return x + + +class _PromptRecordingEstimator(torch.nn.Module): + style_as_token = False + time_as_token = False + is_causal = False + + def __init__(self): + super().__init__() + self.calls = [] + + def forward( + self, + x, + prompt, + x_lens, + _t, + _style, + mu, + *, + pre_mask, + unpad_data=None, + ): + self.calls.append( + { + "x": x.clone(), + "prompt": prompt.clone(), + "x_lens": x_lens.clone(), + "mu": mu.clone(), + "pre_mask": tuple(mask.clone() for mask in pre_mask), + "unpad_data": unpad_data, + } + ) + return torch.ones_like(x) + + +class _VariablePromptCFM(BASECFM): + def __init__(self): + torch.nn.Module.__init__(self) + self.in_channels = 2 + self.zero_prompt_speech_token = True + self.estimator = _PromptRecordingEstimator() + + +class _GroupingEstimator: + def __init__(self): + layers = [SimpleNamespace(attention=SimpleNamespace()) for _ in range(2)] + self.transformer = SimpleNamespace( + layers=layers, + freqs_cis=torch.ones(1), + max_batch_size=8, + ) + + +class _GroupingCFM: + in_channels = 2 + + def __init__(self): + self.estimator = _GroupingEstimator() + self.calls = [] + + 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=None, + ): + self.calls.append( + { + "mu": mu.clone(), + "x_lens": x_lens.clone(), + "prompt": prompt.clone(), + "prompt_lens": prompt_lens.clone(), + "style": style.clone(), + "f0": f0, + "n_timesteps": n_timesteps, + "temperature": temperature, + "inference_cfg_rate": inference_cfg_rate, + "initial_noise": initial_noise.clone(), + "unpad_data": unpad_data, + "full_mask": all(layer.attention._assume_full_mask for layer in self.estimator.transformer.layers), + } + ) + return mu[..., 0].unsqueeze(1).expand(-1, self.in_channels, -1).clone() + + +class _OrderRecordingBigVGAN(torch.nn.Module): + def __init__(self): + super().__init__() + self.calls = [] + + def forward(self, mel): + self.calls.append((int(mel.shape[-1]), float(mel[0, 0, 0]))) + return mel[:, :1].repeat_interleave(4, dim=-1) + + +class _MetadataRecordingDiffusionAttention(torch.nn.Module): + def __init__(self): + super().__init__() + self.metadata = [] + + def forward(self, query, _key, _value, *, attn_metadata): + self.metadata.append(attn_metadata) + return query + + +class _IdentityQKV(torch.nn.Module): + def forward(self, x): + return torch.cat([x, x, x], dim=-1) + + +def _make_decoder(*, use_gpt_latent: bool): + decoder = object.__new__(IndexTTS2S2MelDecoder) + decoder.use_gpt_latent = use_gpt_latent + decoder.semantic_codec_type = "enhanced" + decoder.s2mel = _FakeS2Mel() + return decoder + + +def test_v25_code_only_semantic_embedding_uses_enhanced_decode(): + decoder = _make_decoder(use_gpt_latent=False) + codec = _FakeEnhancedCodec() + + result = decoder._build_semantic_embedding( + codec=codec, + mel_codes=torch.tensor([[1, 2, 3]]), + latent=None, + dtype=torch.float32, + ) + + assert codec.decode_calls == 1 + assert result.shape == (1, 6, 4) + assert torch.equal(result, torch.ones(1, 6, 4)) + + +def test_v25_latent_mode_upsamples_projected_latent_before_addition(): + decoder = _make_decoder(use_gpt_latent=True) + codec = _FakeEnhancedCodec() + + result = decoder._build_semantic_embedding( + codec=codec, + mel_codes=torch.tensor([[1, 2, 3]]), + latent=torch.ones(1, 3, 8), + dtype=torch.float32, + ) + + assert result.shape == (1, 6, 4) + assert torch.equal(result, torch.full((1, 6, 4), 3.0)) + + +def test_s2mel_payload_latent_policy_must_match_stage_config(): + decoder = _make_decoder(use_gpt_latent=False) + + decoder._validate_payload_policy([{"use_gpt_latent": False}]) + with pytest.raises(ValueError, match="use_gpt_latent"): + decoder._validate_payload_policy([{"use_gpt_latent": True}]) + + +@pytest.mark.parametrize("missing_key", ["mel_codes", "S_ref", "ref_mel", "style"]) +def test_s2mel_payload_contract_requires_non_empty_conditioning(missing_key): + payload = { + "mel_codes": torch.ones(2, dtype=torch.long), + "S_ref": torch.ones(2, 4), + "ref_mel": torch.ones(80, 2), + "style": torch.ones(192), + } + payload[missing_key] = torch.empty(0) + + with pytest.raises(ValueError, match=f"non-empty {missing_key}"): + IndexTTS2S2MelDecoder._validate_payload_contract([payload]) + + +def test_s2mel_target_lengths_apply_request_local_duration_factors(): + assert IndexTTS2S2MelDecoder._target_lengths( + [100, 80], + semantic_time_scale=2, + mel_code_to_frame_ratio=1.72, + duration_factors=[0.5, 2.0], + ) == [172, 550] + + +def test_s2mel_duration_factors_default_to_one_without_request_infos(): + assert IndexTTS2S2MelDecoder._duration_factors([], batch_size=2) == [1.0, 1.0] + + +def test_s2mel_duration_factors_default_missing_request_value_to_one(): + assert IndexTTS2S2MelDecoder._duration_factors( + [{}, {"duration_factor": 0.5}], + batch_size=2, + ) == [1.0, 0.5] + + +def test_s2mel_duration_factors_accept_supported_bounds(): + assert IndexTTS2S2MelDecoder._duration_factors( + [{"duration_factor": 0.5}, {"duration_factor": 2.0}], + batch_size=2, + ) == [0.5, 2.0] + + +def test_s2mel_duration_factor_count_must_match_batch_size(): + with pytest.raises(ValueError) as exc_info: + IndexTTS2S2MelDecoder._duration_factors( + [{"duration_factor": 1.0}], + batch_size=2, + ) + + assert str(exc_info.value) == "IndexTTS duration-factor batch mismatch: factors=1 batch=2" + + +def test_s2mel_target_length_count_must_match_duration_factors(): + with pytest.raises(ValueError) as exc_info: + IndexTTS2S2MelDecoder._target_lengths( + [100, 80], + semantic_time_scale=2, + mel_code_to_frame_ratio=1.72, + duration_factors=[1.0], + ) + + assert str(exc_info.value) == "IndexTTS target-length batch mismatch: codes=2 factors=1" + + +@pytest.mark.parametrize("duration_factor", [float("nan"), float("inf"), float("-inf"), 0.0, -0.5, 0.49, 2.01]) +def test_s2mel_duration_factors_must_be_finite_and_in_supported_range(duration_factor): + with pytest.raises(ValueError) as exc_info: + IndexTTS2S2MelDecoder._duration_factors( + [{"duration_factor": duration_factor}], + batch_size=1, + ) + + assert str(exc_info.value) == "IndexTTS duration_factor values must be finite and between 0.5 and 2.0" + + +def test_cfm_unpad_data_builder_accounts_for_token_offsets_and_cfg_row_order(): + indices, cu_seqlens = flow_matching.build_cfm_unpad_data( + [3, 5], + padded_length=5, + token_offset=2, + repeat_for_cfg=True, + device=torch.device("cpu"), + ) + + # Effective lengths are [5, 7] and CFG order is + # [cond row 0, cond row 1, uncond row 0, uncond row 1]. The second + # copy must use the actual padded row offsets (14 and 21), not copied + # indices from the conditional half. + assert indices.dtype == torch.int64 + assert torch.equal( + indices, + torch.tensor( + [ + 0, + 1, + 2, + 3, + 4, + 7, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15, + 16, + 17, + 18, + 21, + 22, + 23, + 24, + 25, + 26, + 27, + ], + dtype=torch.int64, + ), + ) + assert cu_seqlens.dtype == torch.int32 + assert torch.equal(cu_seqlens, torch.tensor([0, 5, 12, 17, 24], dtype=torch.int32)) + + +def test_cfm_inference_uses_explicit_initial_noise(): + cfm = _NoiseEchoCFM() + initial_noise = torch.arange(12, dtype=torch.float32).reshape(1, 2, 6) + + result = cfm.inference( + mu=torch.zeros(1, 6, 4), + x_lens=torch.tensor([6]), + prompt=torch.zeros(1, 80, 0), + style=torch.zeros(1, 192), + f0=None, + n_timesteps=1, + temperature=0.5, + initial_noise=initial_noise, + ) + + assert torch.equal(result, initial_noise * 0.5) + + +def test_cfm_variable_prompt_lens_are_applied_per_row_under_cfg_each_step(): + cfm = _VariablePromptCFM() + initial_noise = torch.arange(20, dtype=torch.float32).reshape(2, 2, 5) + prompt = torch.tensor( + [ + [[11.0, 91.0, 92.0], [12.0, 93.0, 94.0]], + [[21.0, 22.0, 23.0], [24.0, 25.0, 26.0]], + ] + ) + mu = torch.arange(20, dtype=torch.float32).reshape(2, 5, 2) + 1 + prompt_lens = torch.tensor([1, 3], dtype=torch.long) + unpad_data = ( + torch.tensor([0, 1, 2, 3, 5, 6, 7, 8, 9, 10, 11, 12, 13, 15, 16, 17, 18, 19]), + torch.tensor([0, 4, 9, 13, 18], dtype=torch.int32), + ) + + result = cfm.inference( + mu=mu, + x_lens=torch.tensor([4, 5], dtype=torch.long), + prompt=prompt, + prompt_lens=prompt_lens, + style=torch.ones(2, 3), + f0=None, + n_timesteps=2, + inference_cfg_rate=0.7, + initial_noise=initial_noise, + unpad_data=unpad_data, + ) + + assert len(cfm.estimator.calls) == 2 + for call in cfm.estimator.calls: + # CFG packs conditional rows first and unconditional rows second. + assert torch.count_nonzero(call["x"][0, :, :1]) == 0 + assert torch.count_nonzero(call["x"][1, :, :3]) == 0 + assert torch.count_nonzero(call["x"][2, :, :1]) == 0 + assert torch.count_nonzero(call["x"][3, :, :3]) == 0 + assert torch.equal(call["prompt"][0, :, :1], prompt[0, :, :1]) + assert torch.count_nonzero(call["prompt"][0, :, 1:]) == 0 + assert torch.equal(call["prompt"][1, :, :3], prompt[1]) + assert torch.count_nonzero(call["prompt"][2:]) == 0 + assert torch.count_nonzero(call["mu"][0, :1]) == 0 + assert torch.equal(call["mu"][0, 1:], mu[0, 1:]) + assert torch.count_nonzero(call["mu"][1, :3]) == 0 + assert torch.equal(call["mu"][1, 3:], mu[1, 3:]) + assert torch.count_nonzero(call["mu"][2:]) == 0 + assert torch.equal(call["x_lens"], torch.tensor([4, 5])) + assert call["unpad_data"] is unpad_data + + assert torch.count_nonzero(result[0, :, :1]) == 0 + assert torch.count_nonzero(result[1, :, :3]) == 0 + + +def test_seeded_cfm_noise_is_request_local_and_batch_order_independent(): + kwargs = { + "channels": 2, + "length": 6, + "device": torch.device("cpu"), + "dtype": torch.float32, + } + + noise_40_42 = IndexTTS2S2MelDecoder._sample_cfm_noise( + seeds=[40, 42], + **kwargs, + ) + torch.randn(257) + noise_42_40 = IndexTTS2S2MelDecoder._sample_cfm_noise( + seeds=[42, 40], + **kwargs, + ) + + assert noise_40_42 is not None + assert noise_42_40 is not None + assert torch.equal(noise_40_42[0], noise_42_40[1]) + assert torch.equal(noise_40_42[1], noise_42_40[0]) + assert not torch.equal(noise_40_42[0], noise_40_42[1]) + + +def test_unseeded_cfm_noise_preserves_upstream_global_rng_behavior(): + assert ( + IndexTTS2S2MelDecoder._sample_cfm_noise( + seeds=[None, None], + channels=2, + length=6, + device=torch.device("cpu"), + dtype=torch.float32, + ) + is None + ) + + +def test_variable_length_cfm_groups_pack_batch_crop_and_restore_request_order(): + decoder = object.__new__(IndexTTS2S2MelDecoder) + decoder.diffusion_steps = 2 + decoder.inference_cfg_rate = 0.7 + decoder.s2mel_dit_cuda_graph = False + cfm = _GroupingCFM() + + # Original request order is deliberately different from total-length order. + # Stable sort by ref+target must produce groups [1, 3] and [2, 0]. + ref_lengths = [4, 1, 2, 2] + target_lengths = [3, 2, 3, 1] + seeds = [101, 202, 303, 404] + prompt_markers = [110.0, 120.0, 130.0, 140.0] + target_markers = [10.0, 20.0, 30.0, 40.0] + + prompt_condition = torch.full((4, 4, 2), -777.0) + cond = torch.full((4, 3, 2), -999.0) + ref_mel = torch.full((4, 2, 4), -555.0) + style = torch.zeros(4, 3) + for index, (ref_len, target_len) in enumerate(zip(ref_lengths, target_lengths)): + prompt_condition[index, :ref_len] = prompt_markers[index] + cond[index, :target_len] = target_markers[index] + ref_mel[index, :, :ref_len] = prompt_markers[index] + 1 + style[index] = index + 1 + + mels = decoder._run_cfm_groups( + cfm=cfm, + cond=cond, + prompt_condition=prompt_condition, + ref_mel=ref_mel, + style=style, + target_lengths=target_lengths, + ref_lengths=ref_lengths, + seeds=seeds, + max_group_size=2, + ) + + assert len(cfm.calls) == 2 + assert [call["x_lens"].tolist() for call in cfm.calls] == [[3, 3], [5, 7]] + assert [call["prompt_lens"].tolist() for call in cfm.calls] == [[1, 2], [2, 4]] + assert [call["style"][:, 0].tolist() for call in cfm.calls] == [[2.0, 4.0], [3.0, 1.0]] + + first, second = cfm.calls + assert torch.equal(first["mu"][0, :1], torch.full((1, 2), 120.0)) + assert torch.equal(first["mu"][0, 1:3], torch.full((2, 2), 20.0)) + assert torch.equal(first["mu"][1, :2], torch.full((2, 2), 140.0)) + assert torch.equal(first["mu"][1, 2:3], torch.full((1, 2), 40.0)) + assert torch.equal(second["mu"][0, :2], torch.full((2, 2), 130.0)) + assert torch.equal(second["mu"][0, 2:5], torch.full((3, 2), 30.0)) + assert torch.count_nonzero(second["mu"][0, 5:]) == 0 + assert torch.equal(second["mu"][1, :4], torch.full((4, 2), 110.0)) + assert torch.equal(second["mu"][1, 4:7], torch.full((3, 2), 10.0)) + + assert torch.equal(first["prompt"][0, :, :1], torch.full((2, 1), 121.0)) + assert torch.count_nonzero(first["prompt"][0, :, 1:]) == 0 + assert torch.equal(second["prompt"][0, :, :2], torch.full((2, 2), 131.0)) + assert torch.count_nonzero(second["prompt"][0, :, 2:]) == 0 + + # Equal total lengths have no tail padding and may use the full-mask path; + # the mixed-length group must explicitly disable it. + assert [call["full_mask"] for call in cfm.calls] == [True, False] + assert first["unpad_data"] is None + ragged_indices, ragged_cu_seqlens = second["unpad_data"] + assert torch.equal( + ragged_indices, + torch.tensor( + [ + 0, + 1, + 2, + 3, + 4, + 7, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15, + 16, + 17, + 18, + 21, + 22, + 23, + 24, + 25, + 26, + 27, + ] + ), + ) + assert torch.equal(ragged_cu_seqlens, torch.tensor([0, 5, 12, 17, 24], dtype=torch.int32)) + + for call, sorted_indices in zip(cfm.calls, ([1, 3], [2, 0])): + for row, request_index in enumerate(sorted_indices): + x_len = ref_lengths[request_index] + target_lengths[request_index] + generator = torch.Generator(device="cpu") + generator.manual_seed(seeds[request_index]) + expected = torch.randn((2, x_len), generator=generator) + assert torch.equal(call["initial_noise"][row, :, :x_len], expected) + assert torch.count_nonzero(call["initial_noise"][row, :, x_len:]) == 0 + + assert [tuple(mel.shape) for mel in mels] == [ + (2, 3), + (2, 2), + (2, 3), + (2, 1), + ] + for mel, marker in zip(mels, target_markers): + assert torch.equal(mel, torch.full_like(mel, marker)) + + +def test_bigvgan_folds_weight_norm_before_low_precision_cast(monkeypatch): + from vllm_omni.model_executor.models.indextts2.s2mel.modules import bigvgan + + model = _FakeBigVGAN() + monkeypatch.setattr( + bigvgan.BigVGAN, + "from_pretrained", + lambda *_args, **_kwargs: model, + ) + monkeypatch.setattr( + indextts2_s2mel_decoder, + "_patch_bigvgan_compat", + lambda _cls: None, + ) + monkeypatch.setattr( + indextts2_s2mel_decoder, + "_repo_prefetch_lock", + lambda _name: None, + ) + indextts2_s2mel_decoder._bigvgan_models.clear() + + loaded = indextts2_s2mel_decoder._load_bigvgan( + "fake-bigvgan", + torch.device("cpu"), + torch.bfloat16, + ) + + assert loaded is model + assert model.events == [ + "remove_weight_norm", + ("to", torch.device("cpu"), torch.bfloat16), + "eval", + ] + + +def test_bigvgan_folds_weight_norm_in_fp32_under_low_precision_loader_context( + monkeypatch, +): + from vllm_omni.model_executor.models.indextts2.s2mel.modules import bigvgan + + monkeypatch.setattr( + bigvgan.BigVGAN, + "from_pretrained", + lambda *_args, **_kwargs: _DefaultDtypeBigVGAN(), + ) + monkeypatch.setattr( + indextts2_s2mel_decoder, + "_patch_bigvgan_compat", + lambda _cls: None, + ) + monkeypatch.setattr( + indextts2_s2mel_decoder, + "_repo_prefetch_lock", + lambda _name: None, + ) + indextts2_s2mel_decoder._bigvgan_models.clear() + + previous_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.bfloat16) + try: + loaded = indextts2_s2mel_decoder._load_bigvgan( + "fake-bigvgan", + torch.device("cpu"), + torch.bfloat16, + ) + + assert loaded.fold_weight_dtype == torch.float32 + assert loaded.fold_default_dtype == torch.float32 + assert loaded.weight.dtype == torch.bfloat16 + assert torch.get_default_dtype() == torch.bfloat16 + finally: + torch.set_default_dtype(previous_dtype) + + +def _cfm_estimator_with_mixed_parameter_types(): + return torch.nn.ModuleDict( + { + "linear": torch.nn.Linear(3, 2), + "conv": torch.nn.Conv1d(2, 4, kernel_size=3, bias=False), + "norm": torch.nn.LayerNorm(4), + } + ) + + +def test_cfm_estimator_bf16_precast_only_casts_linear_and_conv1d(): + estimator = _cfm_estimator_with_mixed_parameter_types() + + parameter_count = indextts2_s2mel_decoder._precast_cfm_estimator_bf16( + estimator, + enabled=True, + ) + + assert parameter_count == 32 + assert estimator["linear"].weight.dtype == torch.bfloat16 + assert estimator["linear"].bias.dtype == torch.bfloat16 + assert estimator["conv"].weight.dtype == torch.bfloat16 + assert estimator["norm"].weight.dtype == torch.float32 + assert estimator["norm"].bias.dtype == torch.float32 + + +def test_cfm_estimator_bf16_precast_disabled_preserves_fp32(): + estimator = _cfm_estimator_with_mixed_parameter_types() + + parameter_count = indextts2_s2mel_decoder._precast_cfm_estimator_bf16( + estimator, + enabled=False, + ) + + assert parameter_count == 0 + assert all(parameter.dtype == torch.float32 for parameter in estimator.parameters()) + + +def test_v25_full_mask_fast_path_is_guarded_independently_of_cuda_graph(): + decoder = object.__new__(IndexTTS2S2MelDecoder) + decoder.s2mel_dit_cuda_graph = False + layers = [SimpleNamespace(attention=SimpleNamespace()) for _ in range(13)] + estimator = SimpleNamespace( + transformer=SimpleNamespace(layers=layers), + ) + + decoder._set_dit_full_mask_fast_path(estimator, enabled=True) + decoder._enable_dit_cuda_graph(estimator) + + assert all(layer.attention._assume_full_mask is True for layer in layers) + assert not hasattr(estimator, "_cuda_graph_runner") + + decoder._set_dit_full_mask_fast_path(estimator, enabled=False) + + assert all(layer.attention._assume_full_mask is False for layer in layers) + + +def test_indextts_attention_keeps_padding_metadata_model_local(monkeypatch): + attention = object.__new__(gpt_fast_model.Attention) + torch.nn.Module.__init__(attention) + attention.n_head = 1 + attention.n_local_heads = 1 + attention.head_dim = 2 + attention.dim = 2 + attention.wqkv = _IdentityQKV() + attention.wo = torch.nn.Identity() + attention.diff_attn = _MetadataRecordingDiffusionAttention() + monkeypatch.setattr(gpt_fast_model, "apply_rotary_emb", lambda tensor, _freqs: tensor) + + x = torch.ones(2, 3, 2) + mask = torch.tensor( + [ + [[[True, True, True], [True, True, True], [True, True, True]]], + [[[True, True, False], [True, True, False], [True, True, False]]], + ] + ) + attention._assume_full_mask = False + unpad_data = ( + torch.tensor([0, 1, 2, 3, 4], dtype=torch.int64), + torch.tensor([0, 3, 5], dtype=torch.int32), + ) + + attention(x, torch.empty(0), mask, unpad_data=unpad_data) + + padded_metadata = attention.diff_attn.metadata[-1] + assert padded_metadata.extra == {} + assert torch.equal(padded_metadata.attn_mask, mask[:, 0, 0, :]) + + attention._assume_full_mask = True + attention(x, torch.empty(0), mask) + + assert attention.diff_attn.metadata[-1] is None + + all_true_mask = torch.ones_like(mask) + for unknown_state in ("missing", None): + if unknown_state == "missing": + del attention._assume_full_mask + else: + attention._assume_full_mask = unknown_state + + attention(x, torch.empty(0), all_true_mask) + + fallback_metadata = attention.diff_attn.metadata[-1] + assert fallback_metadata.extra == {} + assert torch.equal(fallback_metadata.attn_mask, all_true_mask[:, 0, 0, :]) + + +def test_v25_ref_mel_lengths_come_from_unpadded_shapes(): + ref_mels = [ + torch.zeros(80, 5), + torch.ones(80, 3), + ] + + assert IndexTTS2S2MelDecoder._ref_mel_lengths(ref_mels) == [5, 3] + assert IndexTTS2S2MelDecoder._ref_mel_lengths([None, ref_mels[1]]) == [] + + +def test_v25_resolved_vocoder_source_is_cached(monkeypatch): + decoder = object.__new__(IndexTTS2S2MelDecoder) + decoder.model_path = "/models/indextts25" + decoder.config = SimpleNamespace(vocoder={"name": "fake-bigvgan"}) + decoder._resolved_vocoder_source = None + calls = [] + + def fake_resolve(model_path, vocoder_name): + calls.append((model_path, vocoder_name)) + return "/models/indextts25/hf_cache/bigvgan" + + monkeypatch.setattr( + indextts2_s2mel_decoder, + "_resolve_bigvgan_source", + fake_resolve, + ) + + first = decoder._get_resolved_vocoder_source() + second = decoder._get_resolved_vocoder_source() + + assert first == "/models/indextts25/hf_cache/bigvgan" + assert second == first + assert calls == [("/models/indextts25", "fake-bigvgan")] + + +def test_stage1_preload_does_not_swallow_cuda_runtime_failure(monkeypatch): + decoder = object.__new__(IndexTTS2S2MelDecoder) + torch.nn.Module.__init__(decoder) + decoder.s2mel = torch.nn.Linear(1, 1) + decoder.model_path = "/models/indextts25" + decoder.semantic_codec_type = "enhanced" + decoder.s2mel_vocoder_bf16 = False + decoder.s2mel_vocoder_cuda_graph = False + decoder._resolved_vocoder_source = "/models/bigvgan" + decoder.config = SimpleNamespace( + semantic_codec={}, + semantic_codec_checkpoint="codec.pth", + ) + + def fail_load(*_args, **_kwargs): + raise RuntimeError("CUDA error: illegal memory access") + + monkeypatch.setattr( + indextts2_s2mel_decoder, + "_load_bigvgan", + fail_load, + ) + + with pytest.raises(RuntimeError, match="CUDA error"): + decoder._warmup_external_models() + + +def _runtime_decoder(*, compile_dit: bool): + decoder = object.__new__(IndexTTS2S2MelDecoder) + calls = [] + decoder.s2mel = SimpleNamespace( + enable_torch_compile=lambda *, mode="default": calls.append("compile"), + ) + decoder.s2mel_dit_torch_compile = compile_dit + decoder.s2mel_dit_torch_compile_mode = "default" + decoder.s2mel_dit_cuda_graph = False + decoder._dit_runtime_configured = False + return decoder, calls + + +def test_configure_dit_runtime_compiles_only_on_cuda(): + cpu_decoder, cpu_calls = _runtime_decoder(compile_dit=True) + cuda_decoder, cuda_calls = _runtime_decoder(compile_dit=True) + + cpu_decoder._configure_dit_runtime(torch.device("cpu")) + cuda_decoder._configure_dit_runtime(torch.device("cuda")) + cuda_decoder._configure_dit_runtime(torch.device("cuda")) + + assert cpu_calls == [] + assert cuda_calls == ["compile"] + + +def test_cfm_torch_compile_forwards_reduce_overhead_mode(monkeypatch): + cfm = object.__new__(CFM) + torch.nn.Module.__init__(cfm) + eager_estimator = torch.nn.Identity() + cfm.estimator = eager_estimator + object.__setattr__(cfm, "_eager_estimator", eager_estimator) + calls = [] + + def fake_compile(model, **kwargs): + calls.append((model, kwargs)) + return model + + monkeypatch.setattr(flow_matching.torch, "compile", fake_compile) + + cfm.enable_torch_compile(mode="reduce-overhead") + + assert calls == [ + ( + eager_estimator, + { + "mode": "reduce-overhead", + "fullgraph": True, + "dynamic": True, + }, + ) + ] + + +def test_configure_dit_runtime_forwards_compile_mode_once(): + decoder = object.__new__(IndexTTS2S2MelDecoder) + calls = [] + decoder.s2mel = SimpleNamespace( + enable_torch_compile=lambda *, mode="default": calls.append(mode), + ) + decoder.s2mel_dit_torch_compile = True + decoder.s2mel_dit_torch_compile_mode = "reduce-overhead" + decoder.s2mel_dit_cuda_graph = False + decoder._dit_runtime_configured = False + + decoder._configure_dit_runtime(torch.device("cuda")) + decoder._configure_dit_runtime(torch.device("cuda")) + + assert calls == ["reduce-overhead"] + + +def _vocoder_decoder(*, compile_vocoder: bool, graph_vocoder: bool = False): + decoder = object.__new__(IndexTTS2S2MelDecoder) + decoder.s2mel_vocoder_torch_compile = compile_vocoder + decoder.s2mel_vocoder_cuda_graph = graph_vocoder + decoder._compiled_vocoder = None + return decoder + + +def test_vocoder_torch_compile_disabled_returns_original_model(monkeypatch): + decoder = _vocoder_decoder(compile_vocoder=False) + bigvgan = _LengthRecordingBigVGAN() + runner_factory = getattr(decoder, "_get_vocoder_runner", None) + + assert runner_factory is not None + + def unexpected_compile(*_args, **_kwargs): + pytest.fail("torch.compile must stay disabled") + + monkeypatch.setattr(indextts2_s2mel_decoder.torch, "compile", unexpected_compile) + + runner = runner_factory( + bigvgan, + torch.device("cuda"), + torch.bfloat16, + ) + + assert runner is bigvgan + + +def test_vocoder_torch_compile_is_cached_and_preserves_exact_lengths(monkeypatch): + decoder = _vocoder_decoder(compile_vocoder=True) + bigvgan = _LengthRecordingBigVGAN() + runner_factory = getattr(decoder, "_get_vocoder_runner", None) + compile_calls = [] + + assert runner_factory is not None + + def fake_compile(fn, **kwargs): + compile_calls.append((fn, kwargs)) + return fn + + monkeypatch.setattr(indextts2_s2mel_decoder.torch, "compile", fake_compile) + + first = runner_factory( + bigvgan, + torch.device("cuda"), + torch.bfloat16, + ) + second = runner_factory( + bigvgan, + torch.device("cuda"), + torch.bfloat16, + ) + first(torch.zeros(1, 80, 137)) + second(torch.zeros(1, 80, 221)) + + assert first is second + assert bigvgan.input_lengths == [137, 221] + assert len(compile_calls) == 1 + assert compile_calls[0][0] == bigvgan.forward + assert compile_calls[0][1] == { + "mode": "default", + "fullgraph": False, + "dynamic": True, + } + + +def test_bigvgan_vocodes_cropped_mels_one_request_at_a_time_in_order(): + decoder = object.__new__(IndexTTS2S2MelDecoder) + vocoder = _OrderRecordingBigVGAN() + mels = [ + torch.full((2, 3), 0.1), + torch.full((2, 2), 0.2), + torch.full((2, 4), 0.3), + ] + + wavs = decoder._vocode_mels( + mels=mels, + vocode=vocoder, + voc_dtype=torch.float32, + ) + + assert [length for length, _ in vocoder.calls] == [3, 2, 4] + assert [marker for _, marker in vocoder.calls] == pytest.approx([0.1, 0.2, 0.3]) + assert [tuple(wav.shape) for wav in wavs] == [(12,), (8,), (16,)] + assert [float(wav[0]) for wav in wavs] == pytest.approx([0.1, 0.2, 0.3]) + + +def test_vocoder_torch_compile_rejects_cuda_graph_combination(): + decoder = _vocoder_decoder( + compile_vocoder=True, + graph_vocoder=True, + ) + runner_factory = getattr(decoder, "_get_vocoder_runner", None) + + assert runner_factory is not None + with pytest.raises(ValueError, match="mutually exclusive"): + runner_factory( + _LengthRecordingBigVGAN(), + torch.device("cuda"), + torch.bfloat16, + ) + + +@pytest.mark.parametrize("recipe_name", ["indextts2_5.yaml"]) +def test_v25_recipe_enables_dynamic_dit_compile(recipe_name): + repo_root = Path(__file__).parents[4] + recipe_path = repo_root / "vllm_omni" / "deploy" / recipe_name + recipe = yaml.safe_load(recipe_path.read_text(encoding="utf-8")) + stage1 = next(stage for stage in recipe["stages"] if stage["stage_id"] == 1) + overrides = stage1["hf_overrides"] + + assert overrides["s2mel_dit_cuda_graph"] is False + assert overrides["s2mel_dit_bf16"] is True + assert overrides["s2mel_dit_torch_compile"] is True + assert overrides.get("s2mel_dit_torch_compile_mode") == "default" + assert overrides.get("s2mel_vocoder_torch_compile") is False + assert overrides.get("s2mel_vocoder_cuda_graph", False) is False + + +def test_v25_default_recipe_selects_validated_triton_backend_for_stage0(): + repo_root = Path(__file__).parents[4] + recipe_path = repo_root / "vllm_omni" / "deploy" / "indextts2_5.yaml" + recipe = yaml.safe_load(recipe_path.read_text(encoding="utf-8")) + stage0 = next(stage for stage in recipe["stages"] if stage["stage_id"] == 0) + + assert stage0["attention_backend"] == "TRITON_ATTN" + assert stage0["enable_chunked_prefill"] is False + assert stage0["max_num_batched_tokens"] == stage0["max_model_len"] diff --git a/tests/model_executor/models/indextts2/test_stage_input_processor.py b/tests/model_executor/models/indextts2/test_stage_input_processor.py new file mode 100644 index 000000000..f216ef2b0 --- /dev/null +++ b/tests/model_executor/models/indextts2/test_stage_input_processor.py @@ -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) diff --git a/tests/model_executor/models/indextts2/test_talker_v2_5.py b/tests/model_executor/models/indextts2/test_talker_v2_5.py new file mode 100644 index 000000000..c916354e3 --- /dev/null +++ b/tests/model_executor/models/indextts2/test_talker_v2_5.py @@ -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, + } diff --git a/tests/model_executor/models/indextts2/test_tokenizer_v2_5.py b/tests/model_executor/models/indextts2/test_tokenizer_v2_5.py new file mode 100644 index 000000000..d0b42be08 --- /dev/null +++ b/tests/model_executor/models/indextts2/test_tokenizer_v2_5.py @@ -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("") == ( + "<|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("", " ").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] diff --git a/tests/model_executor/models/indextts2/test_wavenet_ragged.py b/tests/model_executor/models/indextts2/test_wavenet_ragged.py new file mode 100644 index 000000000..d3269989b --- /dev/null +++ b/tests/model_executor/models/indextts2/test_wavenet_ragged.py @@ -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 diff --git a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py index 5e31c2cbd..20f8eb2f0 100644 --- a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py +++ b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py @@ -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 diff --git a/tests/model_executor/stage_input_processors/test_indextts2.py b/tests/model_executor/stage_input_processors/test_indextts2.py index 1380b99ff..5aa9ae4ab 100644 --- a/tests/model_executor/stage_input_processors/test_indextts2.py +++ b/tests/model_executor/stage_input_processors/test_indextts2.py @@ -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) diff --git a/tests/test_indextts2_optional_dependencies.py b/tests/test_indextts2_optional_dependencies.py new file mode 100644 index 000000000..cc3c46b63 --- /dev/null +++ b/tests/test_indextts2_optional_dependencies.py @@ -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() diff --git a/tests/utils/test_mm_outputs_partition.py b/tests/utils/test_mm_outputs_partition.py index d70b95a79..09cd3960a 100644 --- a/tests/utils/test_mm_outputs_partition.py +++ b/tests/utils/test_mm_outputs_partition.py @@ -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() diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index 949a2c7a8..ca282994a 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/tests/worker/test_omni_gpu_model_runner.py @@ -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]}} diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index d702f1ace..ac2e9d38d 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -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, } diff --git a/vllm_omni/core/sched/omni_scheduling_coordinator.py b/vllm_omni/core/sched/omni_scheduling_coordinator.py index 90820e51a..424fed2d9 100644 --- a/vllm_omni/core/sched/omni_scheduling_coordinator.py +++ b/vllm_omni/core/sched/omni_scheduling_coordinator.py @@ -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. diff --git a/vllm_omni/deploy/indextts2_5.yaml b/vllm_omni/deploy/indextts2_5.yaml new file mode 100644 index 000000000..3ea84e972 --- /dev/null +++ b/vllm_omni/deploy/indextts2_5.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 diff --git a/vllm_omni/deploy/qwen2_5_omni.yaml b/vllm_omni/deploy/qwen2_5_omni.yaml index 487ceefdd..1c1bed624 100644 --- a/vllm_omni/deploy/qwen2_5_omni.yaml +++ b/vllm_omni/deploy/qwen2_5_omni.yaml @@ -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. diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index c1512de99..a66f39ce3 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -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_") diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index e8284ba68..b0a95d22c 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -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, diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 100eb69ee..940f62f2e 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -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( diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index 92b2efd14..1564c4523 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -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 diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index ebc417615..ce019a970 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -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 diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 288625dfe..6b7ca1318 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -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." diff --git a/vllm_omni/entrypoints/openai/tts_adapters/base.py b/vllm_omni/entrypoints/openai/tts_adapters/base.py index 69d6c5c6e..b8679041c 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/base.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/base.py @@ -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 diff --git a/vllm_omni/entrypoints/openai/tts_adapters/indextts2.py b/vllm_omni/entrypoints/openai/tts_adapters/indextts2.py index 641e380da..d1cb2a469 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/indextts2.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/indextts2.py @@ -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 diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index 1bea2a124..8c9c134cf 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -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 diff --git a/vllm_omni/model_executor/models/indextts2/codec/__init__.py b/vllm_omni/model_executor/models/indextts2/codec/__init__.py new file mode 100644 index 000000000..55c98dcb9 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/codec/__init__.py @@ -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"] diff --git a/vllm_omni/model_executor/models/indextts2/codec/models.py b/vllm_omni/model_executor/models/indextts2/codec/models.py new file mode 100644 index 000000000..d88534ab8 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/codec/models.py @@ -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, + ) diff --git a/vllm_omni/model_executor/models/indextts2/codec/quantizer.py b/vllm_omni/model_executor/models/indextts2/codec/quantizer.py new file mode 100644 index 000000000..f5c44ee94 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/codec/quantizer.py @@ -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 diff --git a/vllm_omni/model_executor/models/indextts2/codec/vocos.py b/vllm_omni/model_executor/models/indextts2/codec/vocos.py new file mode 100644 index 000000000..7100cb9e6 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/codec/vocos.py @@ -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)) diff --git a/vllm_omni/model_executor/models/indextts2/configuration_indextts2.py b/vllm_omni/model_executor/models/indextts2/configuration_indextts2.py index 504923b1a..e39a34cc3 100644 --- a/vllm_omni/model_executor/models/indextts2/configuration_indextts2.py +++ b/vllm_omni/model_executor/models/indextts2/configuration_indextts2.py @@ -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) diff --git a/vllm_omni/model_executor/models/indextts2/dit_cuda_graph.py b/vllm_omni/model_executor/models/indextts2/dit_cuda_graph.py index 4d32b8c1b..8b75ebfa8 100644 --- a/vllm_omni/model_executor/models/indextts2/dit_cuda_graph.py +++ b/vllm_omni/model_executor/models/indextts2/dit_cuda_graph.py @@ -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, diff --git a/vllm_omni/model_executor/models/indextts2/indextts2_s2mel_decoder.py b/vllm_omni/model_executor/models/indextts2/indextts2_s2mel_decoder.py index a374f1b45..a6fb66329 100644 --- a/vllm_omni/model_executor/models/indextts2/indextts2_s2mel_decoder.py +++ b/vllm_omni/model_executor/models/indextts2/indextts2_s2mel_decoder.py @@ -16,17 +16,28 @@ from typing import Any import torch import torch.nn as nn +import torch.nn.functional as F from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.sequence import IntermediateTensors +from vllm.utils.torch_utils import set_default_torch_dtype from vllm_omni.diffusion.model_loader.hub_prefetch import _repo_prefetch_lock from vllm_omni.model_executor.models.output_templates import OmniOutput -from .configuration_indextts2 import IndexTTS2Config -from .preprocess_utils import load_semantic_codec, resolve_model_file +from .configuration_indextts2 import ( + INDEXTTS25_MAX_DURATION_FACTOR, + INDEXTTS25_MIN_DURATION_FACTOR, + IndexTTS2Config, +) +from .preprocess_utils import ( + load_semantic_codec, + resolve_model_directory, + resolve_model_file, +) from .s2mel.modules.commons import AttrDict, MyModel +from .s2mel.modules.flow_matching import build_cfm_unpad_data logger = init_logger(__name__) @@ -38,6 +49,40 @@ logger = init_logger(__name__) _bigvgan_models: dict[tuple[str, str], nn.Module] = {} +def _raise_if_cuda_runtime_error(error: BaseException) -> None: + if isinstance(error, torch.cuda.OutOfMemoryError): + raise error + if not isinstance(error, RuntimeError): + return + message = str(error).lower() + cuda_markers = ( + "cuda error", + "device-side assert", + "illegal memory access", + "out of memory", + "cublas", + "cudnn", + "cufft", + ) + if any(marker in message for marker in cuda_markers): + raise error + + +def _resolve_bigvgan_source(model_path: str, vocoder_name: str) -> str: + if os.path.isdir(vocoder_name): + return vocoder_name + local_dir = resolve_model_directory( + model_path, + "bigvgan", + os.path.join("hf_cache", "bigvgan"), + ) + if local_dir is not None: + return local_dir + if os.path.isdir(model_path): + raise FileNotFoundError(f"IndexTTS BigVGAN assets are missing from local bundle {model_path!r}") + return vocoder_name + + def _patch_bigvgan_compat(cls): """Monkey-patch BigVGAN._from_pretrained for huggingface_hub>=1.0 compat.""" if getattr(cls._from_pretrained, "_hf1_patched", False): @@ -78,6 +123,25 @@ def _strip_weight_norm(root: nn.Module) -> list[str]: return folded +def _precast_cfm_estimator_bf16( + estimator: nn.Module, + *, + enabled: bool, +) -> int: + """Pre-cast CFM GEMM/conv parameters while keeping other layers in fp32.""" + if not enabled: + return 0 + + parameter_count = 0 + for module in estimator.modules(): + if not isinstance(module, (nn.Linear, nn.Conv1d)): + continue + for parameter in module.parameters(recurse=False): + parameter.data = parameter.data.to(dtype=torch.bfloat16) + parameter_count += parameter.numel() + return parameter_count + + def _load_bigvgan(vocoder_name: str, device: torch.device, dtype: torch.dtype = torch.float32): cache_key = (vocoder_name, str(device), str(dtype)) if cache_key in _bigvgan_models: @@ -87,23 +151,25 @@ def _load_bigvgan(vocoder_name: str, device: torch.device, dtype: torch.dtype = if lock is not None: lock.__enter__() try: - try: - from .s2mel.modules import bigvgan as bigvgan_mod + with set_default_torch_dtype(torch.float32): + try: + from .s2mel.modules import bigvgan as bigvgan_mod - _patch_bigvgan_compat(bigvgan_mod.BigVGAN) - bigvgan_model = bigvgan_mod.BigVGAN.from_pretrained(vocoder_name) - except (ImportError, ModuleNotFoundError): - import bigvgan + _patch_bigvgan_compat(bigvgan_mod.BigVGAN) + bigvgan_model = bigvgan_mod.BigVGAN.from_pretrained(vocoder_name) + except (ImportError, ModuleNotFoundError): + import bigvgan - _patch_bigvgan_compat(bigvgan.BigVGAN) - bigvgan_model = bigvgan.BigVGAN.from_pretrained(vocoder_name) + _patch_bigvgan_compat(bigvgan.BigVGAN) + bigvgan_model = bigvgan.BigVGAN.from_pretrained(vocoder_name) + # Fold weight norm while its normalization is still computed in + # fp32. The vLLM model loader may otherwise leave the process-wide + # default dtype set to the stage's low-precision model dtype here. + bigvgan_model.remove_weight_norm() finally: if lock is not None: lock.__exit__(None, None, None) - # One-time weight cast (vs per-call autocast, which re-casts every conv's - # fp32 weights on each forward and is slower than fp32 on BigVGAN). bigvgan_model = bigvgan_model.to(device=device, dtype=dtype).eval() - bigvgan_model.remove_weight_norm() for p in bigvgan_model.parameters(): p.requires_grad_(False) _bigvgan_models[cache_key] = bigvgan_model @@ -131,6 +197,9 @@ class IndexTTS2S2MelDecoder(nn.Module): self.vllm_config = vllm_config self.model_path = vllm_config.model_config.model self.config: IndexTTS2Config = vllm_config.model_config.hf_config # type: ignore[assignment] + self._resolved_vocoder_source: str | None = None + self.use_gpt_latent = bool(getattr(self.config, "use_gpt_latent", True)) + self.semantic_codec_type = str(getattr(self.config, "semantic_codec_type", "repcodec")) # --- Flags for vLLM-Omni framework --- self.have_multimodal_outputs = True @@ -147,12 +216,19 @@ class IndexTTS2S2MelDecoder(nn.Module): if key in s2mel_args and isinstance(s2mel_args[key], dict): s2mel_args[key] = AttrDict(s2mel_args[key]) - self.s2mel = MyModel(s2mel_args, use_emovec=False, use_gpt_latent=True) + self.s2mel = MyModel( + s2mel_args, + use_emovec=False, + use_gpt_latent=self.use_gpt_latent, + ) # Diffusion config — configurable via deploy YAML hf_overrides self.diffusion_steps: int = getattr(self.config, "diffusion_steps", 25) self.inference_cfg_rate: float = getattr(self.config, "inference_cfg_rate", 0.7) self.mel_code_to_frame_ratio: float = 1.72 + self.s2mel_cfm_batch_size: int = int(getattr(self.config, "s2mel_cfm_batch_size", 1)) + if self.s2mel_cfm_batch_size < 1: + raise ValueError("s2mel_cfm_batch_size must be at least 1") # Run the DiT estimator under bf16 autocast (flash attention + faster # GEMMs); the CFM Euler solver stays float32. Disable via deploy YAML @@ -164,22 +240,47 @@ class IndexTTS2S2MelDecoder(nn.Module): # Capture BigVGAN into bucketed CUDA graphs (collapses the Snake/conv # small-op launch storm). Enable via `s2mel_vocoder_cuda_graph: true`. self.s2mel_vocoder_cuda_graph: bool = getattr(self.config, "s2mel_vocoder_cuda_graph", False) + self.s2mel_vocoder_torch_compile: bool = getattr( + self.config, + "s2mel_vocoder_torch_compile", + False, + ) self._vocoder_graph: Any = None + self._compiled_vocoder: Any = None # Capture the DiT transformer core into per-shape CUDA graphs. # Enable via `s2mel_dit_cuda_graph: true`. self.s2mel_dit_cuda_graph: bool = getattr(self.config, "s2mel_dit_cuda_graph", False) + self.s2mel_dit_torch_compile: bool = getattr( + self.config, + "s2mel_dit_torch_compile", + False, + ) + self.s2mel_dit_torch_compile_mode: str = str( + getattr( + self.config, + "s2mel_dit_torch_compile_mode", + "default", + ) + ) + self._dit_runtime_configured = False self.s2mel_dit_cuda_graph_max_graphs: int = int(getattr(self.config, "s2mel_dit_cuda_graph_max_graphs", 8)) self.s2mel_vocoder_capture_sizes: list[int] | None = getattr(self.config, "s2mel_vocoder_capture_sizes", None) self.s2mel_vocoder_compile_shapes: list[int] | None = getattr(self.config, "s2mel_vocoder_compile_shapes", None) logger.info( "[S2Mel] dit_bf16=%s, vocoder_bf16=%s, " - "vocoder_cuda_graph=%s, dit_cuda_graph=%s, dit_graph_lru=%d, vocoder_buckets=%s, " + "vocoder_cuda_graph=%s, vocoder_torch_compile=%s, " + "dit_cuda_graph=%s, dit_torch_compile=%s, dit_compile_mode=%s, " + "dit_graph_lru=%d, cfm_batch_size=%d, vocoder_buckets=%s, " "vocoder_compile_shapes=%s", self.s2mel_dit_bf16, self.s2mel_vocoder_bf16, self.s2mel_vocoder_cuda_graph, + self.s2mel_vocoder_torch_compile, self.s2mel_dit_cuda_graph, + self.s2mel_dit_torch_compile, + self.s2mel_dit_torch_compile_mode, self.s2mel_dit_cuda_graph_max_graphs, + self.s2mel_cfm_batch_size, self.s2mel_vocoder_capture_sizes, self.s2mel_vocoder_compile_shapes, ) @@ -214,9 +315,11 @@ class IndexTTS2S2MelDecoder(nn.Module): if model_intermediate_buffer: request_infos = [info for info in model_intermediate_buffer if isinstance(info, dict)] additional_information: dict[str, Any] = request_infos[0] if request_infos else {} + self._validate_payload_policy(request_infos) + self._validate_payload_contract(request_infos) device = input_ids.device - model_dtype = self.s2mel.models["gpt_layer"][0].weight.dtype + model_dtype = self._s2mel_model_dtype() logger.debug( "[S2Mel forward] model_intermediate_buffer len=%d, keys=%s", @@ -231,7 +334,16 @@ class IndexTTS2S2MelDecoder(nn.Module): # one S2Mel batch. The runner can split list-valued multimodal outputs # back to per-request payloads, so never silently drop batched items. if len(request_infos) > 1: - mel_codes, latent, code_lens, code_lens_list, s_ref, ref_mel, style = self._build_batched_inputs( + ( + mel_codes, + latent, + code_lens, + code_lens_list, + s_ref, + ref_mel, + style, + ref_lengths_list, + ) = self._build_batched_inputs( request_infos, device=device, model_dtype=model_dtype, @@ -239,19 +351,25 @@ class IndexTTS2S2MelDecoder(nn.Module): code_lens_raw = code_lens else: # Extract Stage 0 outputs and cast to model dtype - latent = self._get_tensor(additional_information, "latent", device, model_dtype) + latent = ( + self._get_tensor(additional_information, "latent", device, model_dtype) if self.use_gpt_latent else None + ) mel_codes = self._get_tensor(additional_information, "mel_codes", device) code_lens_raw = additional_information.get("code_lens") s_ref = self._get_tensor(additional_information, "S_ref", device) ref_mel = self._get_tensor(additional_information, "ref_mel", device, model_dtype) style = self._get_tensor(additional_information, "style", device, model_dtype) + ref_lengths_list = self._ref_mel_lengths([ref_mel]) - if mel_codes is None or latent is None: - logger.warning("S2Mel decoder received empty mel_codes or latent") - return OmniOutput( - text_hidden_states=torch.zeros(1, device=device), - multimodal_outputs={"audio": torch.zeros(1, device=device), "sr": 22050}, - ) + if mel_codes is None or (self.use_gpt_latent and latent is None): + if not request_infos: + # vLLM profiles generation stages once with synthetic token + # inputs and no upstream payload during engine initialization. + return OmniOutput( + text_hidden_states=torch.zeros(1, device=device), + multimodal_outputs={"audio": torch.zeros(1, device=device), "sr": 22050}, + ) + raise ValueError("S2Mel decoder requires mel_codes" + (" and latent" if self.use_gpt_latent else "")) logger.debug( "[S2Mel] extracted tensors — mel_codes=%s latent=%s code_lens=%s s_ref=%s ref_mel=%s style=%s", @@ -266,7 +384,7 @@ class IndexTTS2S2MelDecoder(nn.Module): # Ensure batch dimension if mel_codes.ndim == 1: mel_codes = mel_codes.unsqueeze(0) - if latent.ndim == 2: + if latent is not None and latent.ndim == 2: latent = latent.unsqueeze(0) # Compute code_lens @@ -289,28 +407,34 @@ class IndexTTS2S2MelDecoder(nn.Module): else: code_lens_list.append(mel_codes.shape[1]) code_lens = torch.tensor(code_lens_list, device=device, dtype=torch.long) + duration_factors = self._duration_factors(request_infos, len(code_lens_list)) # Trim to actual length max_len = max(code_lens_list) mel_codes = mel_codes[:, :max_len] - latent = latent[:, :max_len, :] + if latent is not None: + latent = latent[:, :max_len, :] if logger.isEnabledFor(logging.DEBUG): logger.debug( "[S2Mel] after trim — mel_codes=%s latent=%s code_lens=%s max_len=%d", mel_codes.shape, - latent.shape, + latent.shape if latent is not None else None, code_lens.tolist(), max_len, ) # --- S2Mel pipeline --- - # 1. Project GPT latent: 1280 → 1024 - latent = latent.to(device=device, dtype=model_dtype) - latent = self.s2mel.forward_gpt(latent).to(dtype=model_dtype) # [B, T, 1024] - logger.debug("[S2Mel] step1 forward_gpt → latent=%s dtype=%s", latent.shape, latent.dtype) - - # 2. Embed mel codes via semantic codec vq2emb - semantic_codec = load_semantic_codec(self.model_path, self.config.semantic_codec, device) + semantic_codec = load_semantic_codec( + self.model_path, + self.config.semantic_codec, + device, + codec_type=self.semantic_codec_type, + checkpoint_name=getattr( + self.config, + "semantic_codec_checkpoint", + None, + ), + ) codebook_size = self.config.semantic_codec.get("codebook_size", 8192) if logger.isEnabledFor(logging.DEBUG): logger.debug( @@ -323,14 +447,30 @@ class IndexTTS2S2MelDecoder(nn.Module): (mel_codes < 0) | (mel_codes >= codebook_size) ): raise ValueError(f"IndexTTS2 generated mel code outside semantic codebook range [0, {codebook_size - 1}]") - with torch.no_grad(): - S_infer = semantic_codec.quantizer.vq2emb(mel_codes.unsqueeze(1)) # [B, T, 1024] - S_infer = S_infer.transpose(1, 2).to(dtype=model_dtype) # [B, T, 1024] - S_infer = S_infer + latent - logger.debug("[S2Mel] step2 vq2emb → S_infer=%s dtype=%s", S_infer.shape, S_infer.dtype) + S_infer = self._build_semantic_embedding( + codec=semantic_codec, + mel_codes=mel_codes, + latent=latent, + dtype=model_dtype, + ) + logger.debug( + "[S2Mel] semantic decode → S_infer=%s dtype=%s codec=%s latent=%s", + S_infer.shape, + S_infer.dtype, + self.semantic_codec_type, + self.use_gpt_latent, + ) # 3. Length regulate: codes → mel frames - target_lengths_list = [int(length * self.mel_code_to_frame_ratio) for length in code_lens_list] + semantic_time_scale = ( + int(getattr(semantic_codec, "downsample_scale", 1) or 1) if self.semantic_codec_type == "enhanced" else 1 + ) + target_lengths_list = self._target_lengths( + code_lens_list, + semantic_time_scale=semantic_time_scale, + mel_code_to_frame_ratio=self.mel_code_to_frame_ratio, + duration_factors=duration_factors, + ) target_lengths = torch.tensor(target_lengths_list, device=device, dtype=torch.long) if logger.isEnabledFor(logging.DEBUG): logger.debug("[S2Mel] step3 length_regulator target_lengths=%s", target_lengths_list) @@ -364,11 +504,25 @@ class IndexTTS2S2MelDecoder(nn.Module): codebook_size, ) s_ref_codes = s_ref_codes.clamp(0, codebook_size - 1) + if self.semantic_codec_type != "repcodec": + raise ValueError("IndexTTS 2.5 expects W2V-BERT reference embeddings, not semantic-code indices") with torch.no_grad(): - s_ref_emb = semantic_codec.quantizer.vq2emb(s_ref_codes.unsqueeze(1)) # [B, 1024, T] - s_ref_emb = s_ref_emb.transpose(1, 2).to(dtype=model_dtype) # [B, T, 1024] + s_ref_emb = semantic_codec.quantizer.vq2emb(s_ref_codes.unsqueeze(1)) + s_ref_emb = s_ref_emb.transpose(1, 2).to(dtype=model_dtype) - ref_target_lengths = self._infer_ref_lengths(ref_mel) + ref_batch_size = int(ref_mel.shape[0]) + if len(ref_lengths_list) == 1 and ref_batch_size > 1: + ref_lengths_list = ref_lengths_list * ref_batch_size + if len(ref_lengths_list) != ref_batch_size: + raise ValueError( + "IndexTTS reference-mel length/batch mismatch: " + f"lengths={len(ref_lengths_list)} batch={ref_batch_size}" + ) + ref_target_lengths = torch.tensor( + ref_lengths_list, + device=device, + dtype=torch.long, + ) prompt_condition = self.s2mel.models["length_regulator"]( s_ref_emb, ylens=ref_target_lengths, n_quantizers=3, f0=None )[0] @@ -380,12 +534,10 @@ class IndexTTS2S2MelDecoder(nn.Module): cat_condition.shape, ) else: - prompt_condition = None - cat_condition = cond - ref_mel = torch.zeros(cond.shape[0], 80, 0, device=device, dtype=model_dtype) + raise RuntimeError("IndexTTS Stage 1 payload validation allowed missing S_ref or ref_mel") if style is None: - style = torch.zeros(cat_condition.shape[0], 192, device=device, dtype=model_dtype) + raise RuntimeError("IndexTTS Stage 1 payload validation allowed missing style") if style.ndim == 1: style = style.unsqueeze(0) @@ -401,6 +553,7 @@ class IndexTTS2S2MelDecoder(nn.Module): # 5. Flow matching inference (Euler ODE steps, CFG rate configurable) # CFM ODE solver runs in float32 — fp16 Euler steps diverge to noise. # The DiT estimator itself may run under bf16 autocast (see below). + self._configure_dit_runtime(device) cfm = self.s2mel.models["cfm"] cfm.estimator_autocast_dtype = torch.bfloat16 if self.s2mel_dit_bf16 else None cond = cond.float() @@ -410,126 +563,51 @@ class IndexTTS2S2MelDecoder(nn.Module): ref_mel = ref_mel.float() style = style.float() + cfm_batch_size = int(cat_condition.shape[0]) if prompt_condition is None: - ref_lengths = torch.zeros(cat_condition.shape[0], device=device, dtype=torch.long) - else: - ref_lengths = self._infer_ref_lengths(ref_mel) - ref_lengths_list = [int(length) for length in ref_lengths.tolist()] - if len(ref_lengths_list) == 1 and cat_condition.shape[0] > 1: - ref_lengths_list = ref_lengths_list * int(cat_condition.shape[0]) - ref_lengths = ref_lengths.expand(int(cat_condition.shape[0])) - max_ref_len = int(ref_mel.size(-1)) - ref_len = ref_lengths_list[0] if ref_lengths_list else max_ref_len - x_lens = target_lengths + ref_lengths + ref_lengths_list = [0] * cfm_batch_size + elif len(ref_lengths_list) != cfm_batch_size: + raise ValueError( + f"IndexTTS CFM reference-length/batch mismatch: lengths={len(ref_lengths_list)} batch={cfm_batch_size}" + ) estimator = cfm.estimator - required_cache_batch = max(2, int(cat_condition.shape[0]) * (2 if self.inference_cfg_rate > 0 else 1)) + max_cfm_group_size = min(cfm_batch_size, self.s2mel_cfm_batch_size) + required_cache_batch = max(2, max_cfm_group_size * (2 if self.inference_cfg_rate > 0 else 1)) if ( estimator.transformer.freqs_cis is None or getattr(estimator.transformer, "max_batch_size", -1) < required_cache_batch ): estimator.setup_caches(max_batch_size=required_cache_batch, max_seq_length=16384) logger.info("[S2Mel] DiT caches initialized (freqs_cis + causal_mask)") - - cat_len = int(cat_condition.size(1)) - x_lens_list = [target_len + ref_len_i for target_len, ref_len_i in zip(target_lengths_list, ref_lengths_list)] - same_ref_len = len(set(ref_lengths_list)) <= 1 - same_target_len = len(set(target_lengths_list)) <= 1 - dit_graph_full_mask = same_ref_len and same_target_len and all(length == cat_len for length in x_lens_list) - if dit_graph_full_mask: - self._enable_dit_cuda_graph(estimator) - with torch.no_grad(): - mel = cfm.inference( - cat_condition, - x_lens, - ref_mel, - style, - None, # f0 - self.diffusion_steps, - inference_cfg_rate=self.inference_cfg_rate, - ) - # Strip reference portion - mel = mel[:, :, ref_len:] - else: - # Batched Stage-1 requests are padded to a shared length. The DiT - # CUDA graph fast path assumes a full mask, so rebuild each request - # with its exact reference/target lengths before enabling the graph. - mels_list: list[torch.Tensor] = [] - for i in range(cat_condition.shape[0]): - self._enable_dit_cuda_graph(estimator) - ref_len_i = ref_lengths_list[i] - target_len_i = target_lengths_list[i] - if prompt_condition is not None: - prompt_i = prompt_condition[i : i + 1, :ref_len_i] - target_i = cond[i : i + 1, :target_len_i] - cond_i = torch.cat([prompt_i, target_i], dim=1) - else: - cond_i = cat_condition[i : i + 1, :target_len_i] - x_lens_i = torch.tensor([ref_len_i + target_len_i], device=device, dtype=torch.long) - ref_mel_i = ref_mel[i : i + 1, :, :ref_len_i] if ref_mel.ndim == 3 else ref_mel[:, :ref_len_i] - style_i = style[i : i + 1] if style.ndim == 2 else style - with torch.no_grad(): - mel_i = cfm.inference( - cond_i, - x_lens_i, - ref_mel_i, - style_i, - None, - self.diffusion_steps, - inference_cfg_rate=self.inference_cfg_rate, - ) - mels_list.append(mel_i[0, :, ref_len_i:]) - mel = mels_list - dit_graph_info = getattr(getattr(estimator, "_cuda_graph_runner", None), "last_call_info", None) - mel_shape = [tuple(m.shape) for m in mel] if isinstance(mel, list) else tuple(mel.shape) - logger.debug("[S2Mel] step5 CFM done → mel=%s (stripped ref %d frames)", mel_shape, ref_len) + cfm_seeds = self._cfm_request_seeds( + request_infos, + batch_size=int(cat_condition.shape[0]), + ) + mel = self._run_cfm_groups( + cfm=cfm, + cond=cond, + prompt_condition=prompt_condition, + ref_mel=ref_mel, + style=style, + target_lengths=target_lengths_list, + ref_lengths=ref_lengths_list, + seeds=cfm_seeds, + max_group_size=self.s2mel_cfm_batch_size, + ) + logger.debug("[S2Mel] step5 CFM done → mel=%s", [tuple(item.shape) for item in mel]) # 6. BigVGAN vocoding - vocoder_cfg = self.config.vocoder - vocoder_name = vocoder_cfg.get("name", "nvidia/bigvgan_v2_22khz_80band_256x") + vocoder_name = self._get_resolved_vocoder_source() voc_dtype = torch.bfloat16 if (self.s2mel_vocoder_bf16 and device.type == "cuda") else torch.float32 bigvgan = _load_bigvgan(vocoder_name, device, voc_dtype) - vocode = self._get_vocoder_graph(bigvgan, device, voc_dtype) - upsample_factor = int(math.prod(bigvgan.h.upsample_rates)) + vocode = self._get_vocoder_runner(bigvgan, device, voc_dtype) with torch.no_grad(): - wavs: list[torch.Tensor] = [] - vocoder_infos: list[dict[str, Any]] = [] - if isinstance(mel, list): - # Per-request CFM output: variable-length mels. Batch them - # for BigVGAN by padding to the max length. - batch_size = len(mel) - if batch_size == 1: - mel_i = mel[0:1].to(voc_dtype) if mel[0].ndim == 3 else mel[0].unsqueeze(0).to(voc_dtype) - wav_i = vocode(mel_i).float().squeeze(0).squeeze(0) - wavs.append(torch.clamp(wav_i, -1.0, 1.0)) - last_call_info = getattr(vocode, "last_call_info", None) - if isinstance(last_call_info, dict): - vocoder_infos.append(dict(last_call_info)) - else: - for i in range(batch_size): - mel_i = mel[i].unsqueeze(0).to(voc_dtype) if mel[i].ndim == 2 else mel[i : i + 1].to(voc_dtype) - wav_i = vocode(mel_i).float().squeeze(0).squeeze(0) - wavs.append(torch.clamp(wav_i, -1.0, 1.0)) - last_call_info = getattr(vocode, "last_call_info", None) - if isinstance(last_call_info, dict): - vocoder_infos.append(dict(last_call_info)) - else: - batch_size = mel.shape[0] - if batch_size == 1: - mel_i = mel[:, :, : target_lengths_list[0]].to(voc_dtype) - wav_i = vocode(mel_i).float().squeeze(0).squeeze(0) - wavs.append(torch.clamp(wav_i, -1.0, 1.0)) - last_call_info = getattr(vocode, "last_call_info", None) - if isinstance(last_call_info, dict): - vocoder_infos.append(dict(last_call_info)) - else: - for i in range(batch_size): - mel_i = mel[i : i + 1, :, : target_lengths_list[i]].to(voc_dtype) - wav_i = vocode(mel_i).float().squeeze(0).squeeze(0) - wavs.append(torch.clamp(wav_i, -1.0, 1.0)) - last_call_info = getattr(vocode, "last_call_info", None) - if isinstance(last_call_info, dict): - vocoder_infos.append(dict(last_call_info)) + wavs = self._vocode_mels( + mels=mel, + vocode=vocode, + voc_dtype=voc_dtype, + ) wav: torch.Tensor | list[torch.Tensor] = wavs[0] if len(wavs) == 1 else wavs # Keep the public vLLM-Omni audio contract as normalized float. This is @@ -563,6 +641,307 @@ class IndexTTS2S2MelDecoder(nn.Module): # Helpers # ------------------------------------------------------------------ + def _s2mel_model_dtype(self) -> torch.dtype: + if self.use_gpt_latent: + return self.s2mel.models["gpt_layer"][0].weight.dtype + return next(self.s2mel.models["length_regulator"].parameters()).dtype + + def _validate_payload_policy(self, infos: list[dict[str, Any]]) -> None: + for index, info in enumerate(infos): + if "use_gpt_latent" not in info: + raise ValueError(f"IndexTTS Stage 1 input {index} is missing use_gpt_latent") + payload_value = info["use_gpt_latent"] + if bool(payload_value) != self.use_gpt_latent: + raise ValueError( + "IndexTTS Stage 1 use_gpt_latent mismatch for input " + f"{index}: payload={bool(payload_value)} " + f"stage_config={self.use_gpt_latent}" + ) + + @staticmethod + def _validate_payload_contract(infos: list[dict[str, Any]]) -> None: + for index, info in enumerate(infos): + for key in ("mel_codes", "S_ref", "ref_mel", "style"): + value = info.get(key) + if not isinstance(value, torch.Tensor) or value.numel() == 0: + raise ValueError(f"IndexTTS Stage 1 input {index} is missing non-empty {key}") + + @staticmethod + def _duration_factors(infos: list[dict[str, Any]], batch_size: int) -> list[float]: + factors = [float(info.get("duration_factor", 1.0)) for info in infos] if infos else [1.0] * batch_size + if len(factors) != batch_size: + raise ValueError(f"IndexTTS duration-factor batch mismatch: factors={len(factors)} batch={batch_size}") + if any( + not math.isfinite(factor) or not INDEXTTS25_MIN_DURATION_FACTOR <= factor <= INDEXTTS25_MAX_DURATION_FACTOR + for factor in factors + ): + raise ValueError("IndexTTS duration_factor values must be finite and between 0.5 and 2.0") + return factors + + @staticmethod + def _target_lengths( + code_lens: list[int], + *, + semantic_time_scale: int, + mel_code_to_frame_ratio: float, + duration_factors: list[float], + ) -> list[int]: + if len(code_lens) != len(duration_factors): + raise ValueError( + f"IndexTTS target-length batch mismatch: codes={len(code_lens)} factors={len(duration_factors)}" + ) + return [ + int(length * semantic_time_scale * mel_code_to_frame_ratio * factor) + for length, factor in zip(code_lens, duration_factors) + ] + + def _run_cfm_groups( + self, + *, + cfm: Any, + cond: torch.Tensor, + prompt_condition: torch.Tensor | None, + ref_mel: torch.Tensor, + style: torch.Tensor, + target_lengths: list[int], + ref_lengths: list[int], + seeds: list[int | None], + max_group_size: int, + ) -> list[torch.Tensor]: + """Run tightly packed CFM groups and restore original request order.""" + batch_size = int(cond.shape[0]) + if max_group_size < 1: + raise ValueError("IndexTTS CFM max_group_size must be at least 1") + if not (len(target_lengths) == len(ref_lengths) == len(seeds) == batch_size): + raise ValueError( + "IndexTTS CFM input batch mismatch: " + f"batch={batch_size} target={len(target_lengths)} " + f"ref={len(ref_lengths)} seeds={len(seeds)}" + ) + if batch_size == 0: + return [] + + total_lengths = [ref_len + target_len for ref_len, target_len in zip(ref_lengths, target_lengths)] + # Python's sort is stable, so equal total lengths preserve arrival order. + sorted_indices = sorted( + range(batch_size), + key=total_lengths.__getitem__, + ) + outputs: list[torch.Tensor | None] = [None] * batch_size + estimator = cfm.estimator + runtime_estimator = getattr(estimator, "_orig_mod", estimator) + wavenet = getattr(runtime_estimator, "wavenet", None) + min_ragged_length = int(getattr(wavenet, "ragged_min_length", 1)) + groups = self._cfm_group_indices( + sorted_indices=sorted_indices, + total_lengths=total_lengths, + max_group_size=max_group_size, + min_ragged_length=min_ragged_length, + ) + + for group_indices in groups: + group_ref_lengths = [ref_lengths[index] for index in group_indices] + group_target_lengths = [target_lengths[index] for index in group_indices] + group_x_lengths = [ + ref_len + target_len for ref_len, target_len in zip(group_ref_lengths, group_target_lengths) + ] + group_length = max(group_x_lengths) + group_prompt_length = max(group_ref_lengths) + + packed_condition = cond.new_zeros( + len(group_indices), + group_length, + cond.shape[-1], + ) + packed_prompt = ref_mel.new_zeros( + len(group_indices), + ref_mel.shape[1], + group_prompt_length, + ) + for row, (index, ref_len, target_len) in enumerate( + zip(group_indices, group_ref_lengths, group_target_lengths) + ): + if ref_len: + if prompt_condition is None: + raise ValueError("IndexTTS CFM reference length requires prompt conditioning") + packed_condition[row, :ref_len].copy_(prompt_condition[index, :ref_len]) + packed_prompt[row, :, :ref_len].copy_(ref_mel[index, :, :ref_len]) + packed_condition[row, ref_len : ref_len + target_len].copy_(cond[index, :target_len]) + + group_index_tensor = torch.tensor( + group_indices, + device=style.device, + dtype=torch.long, + ) + packed_style = style.index_select(0, group_index_tensor) + x_lens = torch.tensor( + group_x_lengths, + device=cond.device, + dtype=torch.long, + ) + prompt_lens = torch.tensor( + group_ref_lengths, + device=cond.device, + dtype=torch.long, + ) + full_mask = all(length == group_length for length in group_x_lengths) + self._set_dit_full_mask_fast_path(estimator, enabled=full_mask) + if full_mask: + self._enable_dit_cuda_graph(estimator) + unpad_data = None + if not full_mask: + token_offset = int(bool(getattr(runtime_estimator, "style_as_token", False))) + int( + bool(getattr(runtime_estimator, "time_as_token", False)) + ) + unpad_data = build_cfm_unpad_data( + group_x_lengths, + padded_length=group_length, + token_offset=token_offset, + repeat_for_cfg=self.inference_cfg_rate > 0, + device=cond.device, + ) + + initial_noise = self._sample_cfm_noise( + seeds=[seeds[index] for index in group_indices], + channels=cfm.in_channels, + length=group_length, + lengths=group_x_lengths, + device=cond.device, + dtype=cond.dtype, + ) + with torch.no_grad(): + group_mel = cfm.inference( + packed_condition, + x_lens, + packed_prompt, + packed_style, + None, + self.diffusion_steps, + inference_cfg_rate=self.inference_cfg_rate, + initial_noise=initial_noise, + prompt_lens=prompt_lens, + unpad_data=unpad_data, + ) + + for row, (index, ref_len, target_len) in enumerate( + zip(group_indices, group_ref_lengths, group_target_lengths) + ): + outputs[index] = group_mel[row, :, ref_len : ref_len + target_len] + + if any(output is None for output in outputs): + raise RuntimeError("IndexTTS CFM failed to produce every request output") + return [output for output in outputs if output is not None] + + @staticmethod + def _cfm_group_indices( + *, + sorted_indices: list[int], + total_lengths: list[int], + max_group_size: int, + min_ragged_length: int, + ) -> list[list[int]]: + """Build packed groups while preserving exact short-input semantics.""" + groups: list[list[int]] = [] + start = 0 + while start < len(sorted_indices): + index = sorted_indices[start] + group_size = 1 if total_lengths[index] < min_ragged_length else max_group_size + group = sorted_indices[start : start + group_size] + groups.append(group) + start += len(group) + return groups + + @staticmethod + def _cfm_request_seeds( + infos: list[dict[str, Any]], + *, + batch_size: int, + ) -> list[int | None]: + seeds: list[int | None] = [] + for info in infos: + seed = info.get("seed") + if isinstance(seed, torch.Tensor): + seed = seed.reshape(-1)[0].item() if seed.numel() else None + elif isinstance(seed, (list, tuple)): + seed = seed[0] if seed else None + seeds.append(int(seed) if seed is not None else None) + + if not seeds: + return [None] * batch_size + if len(seeds) != batch_size: + raise ValueError(f"IndexTTS CFM seed/request batch mismatch: seeds={len(seeds)} batch={batch_size}") + return seeds + + @staticmethod + def _sample_cfm_noise( + *, + seeds: list[int | None], + channels: int, + length: int, + lengths: list[int] | None = None, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor | None: + if not seeds or all(seed is None for seed in seeds): + return None + if lengths is not None and len(lengths) != len(seeds): + raise ValueError(f"IndexTTS CFM noise length/seed mismatch: lengths={len(lengths)} seeds={len(seeds)}") + + rows: list[torch.Tensor] = [] + for index, seed in enumerate(seeds): + row_length = lengths[index] if lengths is not None else length + if row_length < 0 or row_length > length: + raise ValueError(f"IndexTTS CFM noise length out of range: row={row_length} padded={length}") + generator = None + if seed is not None: + generator = torch.Generator(device=device) + generator.manual_seed(seed) + row = torch.randn( + (1, channels, row_length), + device=device, + dtype=dtype, + generator=generator, + ) + rows.append(F.pad(row, (0, length - row_length))) + return torch.cat(rows, dim=0) + + def _build_semantic_embedding( + self, + *, + codec: Any, + mel_codes: torch.Tensor, + latent: torch.Tensor | None, + dtype: torch.dtype, + ) -> torch.Tensor: + if self.semantic_codec_type == "enhanced": + semantic = codec.decode(mel_codes) + else: + semantic = codec.quantizer.vq2emb(mel_codes.unsqueeze(1)) + semantic = semantic.transpose(1, 2) + semantic = semantic.to(device=mel_codes.device, dtype=dtype) + + if not self.use_gpt_latent: + return semantic + if latent is None: + raise ValueError("IndexTTS GPT latent mode requires a latent tensor") + + projected = self.s2mel.forward_gpt(latent.to(device=mel_codes.device, dtype=dtype)).to(dtype=dtype) + if projected.shape[1] != semantic.shape[1]: + # EnhancedCodec restores the encoder's temporal downsampling with + # nearest-neighbor upsampling. Mirror that time mapping for the + # optional GPT residual before addition. + projected = F.interpolate( + projected.transpose(1, 2), + size=semantic.shape[1], + mode="nearest", + ).transpose(1, 2) + if projected.shape != semantic.shape: + raise ValueError( + "IndexTTS semantic/GPT latent shape mismatch after alignment: " + f"semantic={tuple(semantic.shape)} latent={tuple(projected.shape)}" + ) + return semantic + projected + def _build_batched_inputs( self, infos: list[dict[str, Any]], @@ -571,12 +950,13 @@ class IndexTTS2S2MelDecoder(nn.Module): model_dtype: torch.dtype, ) -> tuple[ torch.Tensor, - torch.Tensor, + torch.Tensor | None, torch.Tensor, list[int], torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, + list[int], ]: stop_token = self.config.gpt.get("stop_mel_token", 8193) mel_items: list[torch.Tensor] = [] @@ -588,15 +968,18 @@ class IndexTTS2S2MelDecoder(nn.Module): for idx, info in enumerate(infos): mel = self._get_tensor(info, "mel_codes", device) - latent = self._get_tensor(info, "latent", device, model_dtype) - if mel is None or latent is None: - raise ValueError(f"S2Mel batched input {idx} is missing mel_codes or latent") + latent = self._get_tensor(info, "latent", device, model_dtype) if self.use_gpt_latent else None + if mel is None or (self.use_gpt_latent and latent is None): + raise ValueError( + f"S2Mel batched input {idx} is missing mel_codes" + (" or latent" if self.use_gpt_latent else "") + ) mel = mel.to(device=device, dtype=torch.long).reshape(-1) - if latent.ndim == 3 and latent.shape[0] == 1: - latent = latent[0] - elif latent.ndim != 2: - latent = latent.reshape(-1, latent.shape[-1]) + if latent is not None: + if latent.ndim == 3 and latent.shape[0] == 1: + latent = latent[0] + elif latent.ndim != 2: + latent = latent.reshape(-1, latent.shape[-1]) raw_len = info.get("code_lens") if isinstance(raw_len, torch.Tensor) and raw_len.numel() > 0: @@ -606,9 +989,16 @@ class IndexTTS2S2MelDecoder(nn.Module): else: stop_mask = (mel == stop_token).nonzero(as_tuple=False) code_len = int(stop_mask[0].item()) if stop_mask.numel() > 0 else int(mel.shape[0]) - code_len = max(0, min(code_len, int(mel.shape[0]), int(latent.shape[0]))) + code_len = max(0, min(code_len, int(mel.shape[0]))) + if latent is not None: + if int(latent.shape[0]) != int(mel.shape[0]): + raise ValueError( + f"S2Mel batched input {idx} mel/latent length mismatch: " + f"mel={mel.shape[0]} latent={latent.shape[0]}" + ) mel_items.append(mel[:code_len]) - latent_items.append(latent[:code_len]) + if latent is not None: + latent_items.append(latent[:code_len]) code_lens.append(code_len) s_ref = self._get_tensor(info, "S_ref", device) @@ -627,17 +1017,28 @@ class IndexTTS2S2MelDecoder(nn.Module): style_items.append(style) max_code_len = max(code_lens) if code_lens else 0 - latent_dim = latent_items[0].shape[-1] # Padding reaches vq2emb before length masking, so it must be a valid # semantic-code index. Actual sequence lengths are still carried by # code_lens/target_lengths; stop_token (8193) is outside the codebook. mel_batch = torch.zeros((len(mel_items), max_code_len), dtype=torch.long, device=device) - latent_batch = torch.zeros(len(latent_items), max_code_len, latent_dim, device=device, dtype=model_dtype) - for i, (mel, latent) in enumerate(zip(mel_items, latent_items)): + latent_batch: torch.Tensor | None = None + if self.use_gpt_latent: + latent_dim = latent_items[0].shape[-1] + latent_batch = torch.zeros( + len(latent_items), + max_code_len, + latent_dim, + device=device, + dtype=model_dtype, + ) + for i, mel in enumerate(mel_items): mel_batch[i, : mel.shape[0]] = mel - latent_batch[i, : latent.shape[0]] = latent.to(dtype=model_dtype) + if latent_batch is not None: + latent = latent_items[i] + latent_batch[i, : latent.shape[0]] = latent.to(dtype=model_dtype) s_ref_batch = self._stack_optional_sequence(s_ref_items, device=device, dtype=model_dtype) + ref_lengths = self._ref_mel_lengths(ref_mel_items) ref_mel_batch = self._stack_optional_ref_mel(ref_mel_items, device=device, dtype=model_dtype) style_batch = self._stack_optional_style(style_items, device=device, dtype=model_dtype) return ( @@ -648,6 +1049,7 @@ class IndexTTS2S2MelDecoder(nn.Module): s_ref_batch, ref_mel_batch, style_batch, + ref_lengths, ) @staticmethod @@ -702,12 +1104,11 @@ class IndexTTS2S2MelDecoder(nn.Module): ) @staticmethod - def _infer_ref_lengths(ref_mel: torch.Tensor) -> torch.Tensor: - if ref_mel.ndim == 2: - return torch.tensor([ref_mel.size(-1)], device=ref_mel.device, dtype=torch.long) - nonzero = ref_mel.abs().sum(dim=1) > 0 - lengths = nonzero.long().sum(dim=1) - return torch.clamp(lengths, min=1).to(device=ref_mel.device, dtype=torch.long) + def _ref_mel_lengths(items: list[torch.Tensor | None]) -> list[int]: + """Return unpadded reference lengths without inspecting CUDA values.""" + if not items or any(item is None for item in items): + return [] + return [int(item.shape[-1]) for item in items if item is not None] @staticmethod def _get_tensor( @@ -794,8 +1195,9 @@ class IndexTTS2S2MelDecoder(nn.Module): required_prefixes = [ "s2mel.models.cfm.", "s2mel.models.length_regulator.", - "s2mel.models.gpt_layer.", ] + if self.use_gpt_latent: + required_prefixes.append("s2mel.models.gpt_layer.") missing_prefixes = [ prefix for prefix in required_prefixes if not any(name.startswith(prefix) for name in loaded_params) ] @@ -817,29 +1219,126 @@ class IndexTTS2S2MelDecoder(nn.Module): # recursively re-casting the module on every request. cfm = self.s2mel.models["cfm"] cfm.float() + precast_parameter_count = _precast_cfm_estimator_bf16( + cfm.estimator, + enabled=self.s2mel_dit_bf16, + ) + if self.s2mel_dit_bf16: + logger.info( + "Pre-cast %d CFM estimator Linear/Conv1d parameters to bfloat16", + precast_parameter_count, + ) cfm.estimator_autocast_dtype = torch.bfloat16 if self.s2mel_dit_bf16 else None self._warmup_external_models() return loaded_params + def _configure_dit_runtime(self, device: torch.device) -> None: + if self._dit_runtime_configured: + return + if self.s2mel_dit_torch_compile and self.s2mel_dit_cuda_graph: + raise ValueError( + "s2mel_dit_torch_compile and s2mel_dit_cuda_graph are mutually exclusive", + ) + if self.s2mel_dit_torch_compile: + if device.type == "cuda": + self.s2mel.enable_torch_compile( + mode=self.s2mel_dit_torch_compile_mode, + ) + logger.info( + "S2Mel DiT torch.compile enabled with mode=%s", + self.s2mel_dit_torch_compile_mode, + ) + else: + logger.info( + "S2Mel DiT torch.compile skipped on device type %s", + device.type, + ) + self._dit_runtime_configured = True + def _warmup_external_models(self) -> None: """Preload vocoder + semantic codec to eliminate first-request latency.""" device = next((p.device for p in self.s2mel.parameters()), torch.device("cpu")) try: - vocoder_name = self.config.vocoder.get("name", "nvidia/bigvgan_v2_22khz_80band_256x") + vocoder_name = self._get_resolved_vocoder_source() voc_dtype = torch.bfloat16 if (self.s2mel_vocoder_bf16 and device.type == "cuda") else torch.float32 bigvgan = _load_bigvgan(vocoder_name, device, voc_dtype) - load_semantic_codec(self.model_path, self.config.semantic_codec, device) - # Warm up the CUDA graph + torch.compile wrapper at load time so the - # ~95s compile cost is paid during startup rather than on the first - # request. Mirrors the Qwen3-TTS pattern (enable_cudagraph in - # load_weights). _get_vocoder_graph constructs the wrapper and calls - # warmup() on first access. - if self.s2mel_vocoder_cuda_graph and device.type == "cuda": - self._get_vocoder_graph(bigvgan, device, voc_dtype) + load_semantic_codec( + self.model_path, + self.config.semantic_codec, + device, + codec_type=self.semantic_codec_type, + checkpoint_name=getattr( + self.config, + "semantic_codec_checkpoint", + None, + ), + ) + if (self.s2mel_vocoder_cuda_graph or self.s2mel_vocoder_torch_compile) and device.type == "cuda": + self._get_vocoder_runner(bigvgan, device, voc_dtype) logger.info("BigVGAN + semantic codec preloaded") except Exception as e: - logger.warning("Failed to preload Stage 1 models: %s", e) + _raise_if_cuda_runtime_error(e) + raise RuntimeError("Failed to preload required IndexTTS Stage 1 models") from e + + def _get_resolved_vocoder_source(self) -> str: + if self._resolved_vocoder_source is None: + vocoder_name = self.config.vocoder.get( + "name", + "nvidia/bigvgan_v2_22khz_80band_256x", + ) + self._resolved_vocoder_source = _resolve_bigvgan_source( + self.model_path, + vocoder_name, + ) + return self._resolved_vocoder_source + + def _vocode_mels( + self, + *, + mels: list[torch.Tensor], + vocode: Any, + voc_dtype: torch.dtype, + ) -> list[torch.Tensor]: + """Vocode each already-cropped request at its exact mel length.""" + wavs: list[torch.Tensor] = [] + for mel in mels: + if mel.ndim == 2: + mel_input = mel.unsqueeze(0) + elif mel.ndim == 3 and mel.shape[0] == 1: + mel_input = mel + else: + raise ValueError(f"IndexTTS BigVGAN expects one mel per request, got {tuple(mel.shape)}") + wav = vocode(mel_input.to(voc_dtype)).float().squeeze(0).squeeze(0) + wavs.append(torch.clamp(wav, -1.0, 1.0)) + return wavs + + def _get_vocoder_runner( + self, + bigvgan: nn.Module, + device: torch.device, + dtype: torch.dtype, + ): + """Return the selected exact-length BigVGAN execution path.""" + if self.s2mel_vocoder_torch_compile and self.s2mel_vocoder_cuda_graph: + raise ValueError( + "s2mel_vocoder_torch_compile and s2mel_vocoder_cuda_graph are mutually exclusive", + ) + if self.s2mel_vocoder_cuda_graph: + return self._get_vocoder_graph(bigvgan, device, dtype) + if not self.s2mel_vocoder_torch_compile or device.type != "cuda": + return bigvgan + if self._compiled_vocoder is None: + self._compiled_vocoder = torch.compile( + bigvgan.forward, + mode="default", + fullgraph=False, + dynamic=True, + ) + logger.info( + "BigVGAN dynamic torch.compile enabled with exact-length inputs", + ) + return self._compiled_vocoder def _get_vocoder_graph(self, bigvgan: nn.Module, device: torch.device, dtype: torch.dtype): """Return the CUDA-graph vocoder wrapper, or the eager model when disabled. @@ -867,23 +1366,41 @@ class IndexTTS2S2MelDecoder(nn.Module): self._vocoder_graph.enabled = False return self._vocoder_graph + @staticmethod + def _set_dit_full_mask_fast_path( + estimator: nn.Module, + *, + enabled: bool, + ) -> None: + """Select dense attention and physical WaveNet reflect boundaries.""" + runtime_estimator = getattr(estimator, "_orig_mod", estimator) + for layer in runtime_estimator.transformer.layers: + layer.attention._assume_full_mask = enabled + if hasattr(runtime_estimator, "wavenet"): + runtime_estimator.wavenet._assume_full_mask = enabled + + # A graph captured with the dense-attention branch must never replay + # for a padded input. Keep its cache resident, but route this call + # through eager execution until a later full-mask call re-enables it. + graph_runner = getattr(runtime_estimator, "_cuda_graph_runner", None) + if graph_runner is not None and not enabled: + graph_runner.enabled = False + def _enable_dit_cuda_graph(self, estimator: nn.Module) -> None: """Attach a per-shape CUDA graph runner to the DiT transformer core. - Sets `_assume_full_mask` on the attention modules: this decoder - always passes x_lens == T, so the padding mask is all-True and - dense attention is identical — the flag removes the backend's - per-step ``torch.any(~mask)`` D2H sync, which is also a - stream-capture blocker. + Full-mask semantics are enabled independently by + ``_set_dit_full_mask_fast_path`` so eager execution also avoids + the per-layer ``torch.any(~mask)`` host synchronization. """ if not self.s2mel_dit_cuda_graph: return - if getattr(estimator, "_cuda_graph_runner", None) is not None: + graph_runner = getattr(estimator, "_cuda_graph_runner", None) + if graph_runner is not None: + graph_runner.enabled = True return from vllm_omni.model_executor.models.indextts2.dit_cuda_graph import CUDAGraphDiTRunner - for layer in estimator.transformer.layers: - layer.attention._assume_full_mask = True estimator._cuda_graph_runner = CUDAGraphDiTRunner( estimator.transformer, max_graphs=self.s2mel_dit_cuda_graph_max_graphs, diff --git a/vllm_omni/model_executor/models/indextts2/indextts2_talker.py b/vllm_omni/model_executor/models/indextts2/indextts2_talker.py index 2b1244ccc..f9eae0c6f 100644 --- a/vllm_omni/model_executor/models/indextts2/indextts2_talker.py +++ b/vllm_omni/model_executor/models/indextts2/indextts2_talker.py @@ -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 diff --git a/vllm_omni/model_executor/models/indextts2/pipeline.py b/vllm_omni/model_executor/models/indextts2/pipeline.py index 50f866ffd..e29d61563 100644 --- a/vllm_omni/model_executor/models/indextts2/pipeline.py +++ b/vllm_omni/model_executor/models/indextts2/pipeline.py @@ -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}, ), ), diff --git a/vllm_omni/model_executor/models/indextts2/preprocess_utils.py b/vllm_omni/model_executor/models/indextts2/preprocess_utils.py index 452bce9bb..1b6fa9485 100644 --- a/vllm_omni/model_executor/models/indextts2/preprocess_utils.py +++ b/vllm_omni/model_executor/models/indextts2/preprocess_utils.py @@ -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( diff --git a/vllm_omni/model_executor/models/indextts2/prompt_utils.py b/vllm_omni/model_executor/models/indextts2/prompt_utils.py index c33202785..5a950036a 100644 --- a/vllm_omni/model_executor/models/indextts2/prompt_utils.py +++ b/vllm_omni/model_executor/models/indextts2/prompt_utils.py @@ -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 diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/commons.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/commons.py index dd86f5847..8aed5f93b 100644 --- a/vllm_omni/model_executor/models/indextts2/s2mel/modules/commons.py +++ b/vllm_omni/model_executor/models/indextts2/s2mel/modules/commons.py @@ -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) diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/diffusion_transformer.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/diffusion_transformer.py index f73fdc800..a76d998a4 100644 --- a/vllm_omni/model_executor/models/indextts2/s2mel/modules/diffusion_transformer.py +++ b/vllm_omni/model_executor/models/indextts2/s2mel/modules/diffusion_transformer.py @@ -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 diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/flow_matching.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/flow_matching.py index 5192867af..32c2f641c 100644 --- a/vllm_omni/model_executor/models/indextts2/s2mel/modules/flow_matching.py +++ b/vllm_omni/model_executor/models/indextts2/s2mel/modules/flow_matching.py @@ -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, ) diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/gpt_fast/model.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/gpt_fast/model.py index c410b6387..0e5b071d1 100644 --- a/vllm_omni/model_executor/models/indextts2/s2mel/modules/gpt_fast/model.py +++ b/vllm_omni/model_executor/models/indextts2/s2mel/modules/gpt_fast/model.py @@ -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) diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/wavenet.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/wavenet.py index 4c5381766..f7910055e 100644 --- a/vllm_omni/model_executor/models/indextts2/s2mel/modules/wavenet.py +++ b/vllm_omni/model_executor/models/indextts2/s2mel/modules/wavenet.py @@ -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, :] diff --git a/vllm_omni/model_executor/models/indextts2/text_processing_v2_5.py b/vllm_omni/model_executor/models/indextts2/text_processing_v2_5.py new file mode 100644 index 000000000..8eae3045c --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/text_processing_v2_5.py @@ -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"(? bool: + return bool(re.fullmatch(r"[\u3040-\u309f]+", text) or re.fullmatch(r"[\u30a0-\u30ff]+", text)) + + +def apply_pronunciation_annotations(text: str) -> str: + """Convert ```` 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("-", "")) + + 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*\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) diff --git a/vllm_omni/model_executor/models/indextts2/tokenizer_v2_5.py b/vllm_omni/model_executor/models/indextts2/tokenizer_v2_5.py new file mode 100644 index 000000000..6eae97703 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/tokenizer_v2_5.py @@ -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") diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index fd76dadd8..3c79ba34c 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -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", diff --git a/vllm_omni/model_executor/stage_input_processors/indextts2.py b/vllm_omni/model_executor/stage_input_processors/indextts2.py index 39d9d011c..3cf4e74c5 100644 --- a/vllm_omni/model_executor/stage_input_processors/indextts2.py +++ b/vllm_omni/model_executor/stage_input_processors/indextts2.py @@ -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 diff --git a/vllm_omni/request.py b/vllm_omni/request.py index 2f1fd5379..0606bcfc8 100644 --- a/vllm_omni/request.py +++ b/vllm_omni/request.py @@ -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 diff --git a/vllm_omni/utils/mm_outputs.py b/vllm_omni/utils/mm_outputs.py index af8f9ac14..1a380932b 100644 --- a/vllm_omni/utils/mm_outputs.py +++ b/vllm_omni/utils/mm_outputs.py @@ -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): diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index cb6e57579..b68c722ec 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -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"):