[Model] Add IndexTTS2 text-to-speech support (#3838)

Signed-off-by: BeatSeat <wendavid552@gmail.com>
Signed-off-by: David <david@podcast.ai>
Signed-off-by: wendavid552 <wendavid552@gmail.com>
Signed-off-by: WenDavid <davidzhangcf@foxmail.com>

.buildkite/test-merge.yml

@@ -363,6 +363,42 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: "TTS · IndexTTS2 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_indextts2.py
+          - tests/e2e/offline_inference/test_indextts2.py
+          - tests/model_executor/models/indextts2/
+          - tests/model_executor/stage_input_processors/test_indextts2.py
+          - tests/assets/indextts2/
+          - vllm_omni/model_executor/models/indextts2/
+          - vllm_omni/model_executor/stage_input_processors/indextts2.py
+          - vllm_omni/deploy/indextts2.yaml
+          - vllm_omni/deploy/indextts2_low_latency.yaml
+          - vllm_omni/entrypoints/openai/serving_speech.py
+          - vllm_omni/entrypoints/openai/tts_adapters/indextts2.py
+          - vllm_omni/config/pipeline_registry.py
+          - vllm_omni/model_executor/models/registry.py
+        timeout_in_minutes: 20
+        commands:
+          - |
+            timeout 20m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              pytest -s -v tests/e2e/online_serving/test_indextts2.py tests/e2e/offline_inference/test_indextts2.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+            "
+        agents:
+          queue: "gpu_4_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              gpus: all
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
       - label: "Omni · Qwen3-Omni Test"
         source_file_dependencies:
           - tests/e2e/offline_inference/test_qwen3_omni.py

docs/models/supported_models.md

@@ -74,6 +74,7 @@ th {
 | `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | |
 | `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | |
 | `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | |
+| `IndexTTS2TalkerForConditionalGeneration` | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✅︎ | | | |
 | `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` | ✅︎ | ✅︎ | | |

examples/offline_inference/text_to_speech/README.md

@@ -23,6 +23,7 @@ list of supported architectures across all modalities, see
 | OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz |
 | 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 |
 | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz |
 
 ## Common Quick Start
@@ -436,6 +437,50 @@ Streaming is exposed through the online OpenAI Speech API (`stream=true`). See [
 
 ---
 
+## IndexTTS-2
+
+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.
+
+### Quick start
+```bash
+python examples/offline_inference/text_to_speech/indextts2/end2end.py \
+    --model IndexTeam/IndexTTS-2 \
+    --text "你好，这是一个语音合成测试。" \
+    --ref-audio /path/to/reference.wav
+```
+
+### Emotion control
+```bash
+# Emotion from reference audio
+python examples/offline_inference/text_to_speech/indextts2/end2end.py \
+    --model IndexTeam/IndexTTS-2 \
+    --text "今天天气真好！" \
+    --ref-audio /path/to/ref.wav \
+    --emo-audio /path/to/happy.wav
+
+# Emotion from 8-dim vector (happy angry sad afraid disgusted melancholy surprised calm)
+python examples/offline_inference/text_to_speech/indextts2/end2end.py \
+    --model IndexTeam/IndexTTS-2 \
+    --text "今天天气真好！" \
+    --ref-audio /path/to/ref.wav \
+    --emo-vector 1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
+
+# Emotion from text description
+python examples/offline_inference/text_to_speech/indextts2/end2end.py \
+    --model IndexTeam/IndexTTS-2 \
+    --text "今天天气真好！" \
+    --ref-audio /path/to/ref.wav \
+    --emo-text "happy and excited"
+```
+
+### 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).
+
+---
+
 ## Voxtral TTS
 
 Voxtral-4B-TTS (Mistral). Has its own argparse surface; uses voice presets and the `mistral_common` `SpeechRequest` protocol.

examples/offline_inference/text_to_speech/indextts2/end2end.py

@@ -0,0 +1,187 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline inference example for IndexTTS2 via vLLM-Omni.
+
+Two-stage pipeline: GPT AR (Stage 0) → S2Mel + BigVGAN (Stage 1).
+Output is 22050 Hz mono WAV.
+
+Usage:
+  python end2end.py \
+    --model /path/to/IndexTeam/IndexTTS-2 \
+    --text "你好，这是一个语音合成测试。" \
+    --ref-audio /path/to/ref.wav
+
+  # With emotion audio:
+  python end2end.py \
+    --model /path/to/IndexTeam/IndexTTS-2 \
+    --text "今天天气真好！" \
+    --ref-audio /path/to/ref.wav \
+    --emo-audio /path/to/happy.wav
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import soundfile as sf
+import torch
+from vllm import SamplingParams
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+
+from vllm_omni import Omni  # noqa: E402
+from vllm_omni.model_executor.models.indextts2.prompt_utils import (  # noqa: E402
+    build_indextts2_prefill_prompt_ids,
+)
+
+
+def build_request(
+    model: str,
+    text: str,
+    ref_audio_path: str | None = None,
+    emo_audio_path: str | None = None,
+    emo_text: str | None = None,
+    emo_vector: list[float] | None = None,
+    emo_alpha: float | None = None,
+    use_emo_text: bool = False,
+    use_random: bool = False,
+) -> dict:
+    additional: dict = {"text": [text]}
+    if ref_audio_path:
+        additional["voice"] = [str(ref_audio_path)]
+    if emo_audio_path:
+        additional["emo_audio"] = [str(emo_audio_path)]
+    if emo_text:
+        additional["emo_text"] = [emo_text]
+    if emo_vector is not None:
+        additional["emo_vector"] = [emo_vector]
+    if emo_alpha is not None:
+        additional["emo_alpha"] = [emo_alpha]
+    if use_emo_text:
+        additional["use_emo_text"] = [True]
+    if use_random:
+        additional["use_random"] = [True]
+    return {
+        "prompt_token_ids": build_indextts2_prefill_prompt_ids(model, text),
+        "additional_information": additional,
+    }
+
+
+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)
+    print(f"  Saved {path} ({audio_np.shape}, {sample_rate} Hz)")
+
+
+def extract_audio(mm: dict) -> tuple[torch.Tensor | None, int]:
+    audio = mm.get("audio")
+    if audio is None:
+        audio = mm.get("model_outputs")
+    if isinstance(audio, list):
+        chunks = [chunk.reshape(-1) for chunk in audio if isinstance(chunk, torch.Tensor) and chunk.numel() > 0]
+        audio = torch.cat(chunks, dim=0) if chunks else None
+
+    sr_val = mm.get("sr")
+    if isinstance(sr_val, list):
+        sr_val = sr_val[-1] if sr_val else None
+    if hasattr(sr_val, "item"):
+        sample_rate = int(sr_val.item())
+    else:
+        sample_rate = int(sr_val) if sr_val is not None else 22050
+
+    return audio if isinstance(audio, torch.Tensor) else None, sample_rate
+
+
+def main(args) -> None:
+    omni = Omni(
+        model=args.model,
+        deploy_config=args.deploy_config,
+        stage_init_timeout=args.stage_init_timeout,
+    )
+
+    # Stage 0: GPT AR (text → mel codes)
+    gpt_sampling = SamplingParams(
+        temperature=0.8,
+        top_p=0.8,
+        top_k=30,
+        max_tokens=1500,
+        repetition_penalty=10.0,
+        stop_token_ids=[8193],
+        seed=args.seed if args.seed is not None else 42,
+        detokenize=False,
+    )
+    # Stage 1: S2Mel + BigVGAN (non-AR, params mostly ignored)
+    s2mel_sampling = SamplingParams(
+        temperature=0.0,
+        max_tokens=65536,
+        detokenize=True,
+    )
+    sampling_params = [gpt_sampling, s2mel_sampling]
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    print(f"Synthesizing: {args.text!r}")
+    if args.ref_audio:
+        print(f"  ref_audio: {args.ref_audio}")
+    inputs = build_request(
+        model=args.model,
+        text=args.text,
+        ref_audio_path=args.ref_audio,
+        emo_audio_path=args.emo_audio,
+        emo_text=args.emo_text,
+        emo_vector=args.emo_vector,
+        emo_alpha=args.emo_alpha,
+        use_emo_text=args.use_emo_text,
+        use_random=args.use_random,
+    )
+
+    for i, omni_out in enumerate(omni.generate(inputs, sampling_params_list=sampling_params)):
+        mm = omni_out.multimodal_output
+        if not mm:
+            print(f"  [req {i}] No multimodal output.")
+            continue
+        audio, sr = extract_audio(mm)
+        if audio is None:
+            print(f"  [req {i}] No waveform in multimodal_output.")
+            continue
+        out_path = str(output_dir / f"output_{i}.wav")
+        save_audio(audio.cpu(), out_path, sr)
+
+    print("Done.")
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="IndexTTS2 offline inference")
+    parser.add_argument("--model", required=True, help="HF model path for IndexTTS2.")
+    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.")
+    parser.add_argument("--emo-text", default=None, help="Emotion description text.")
+    parser.add_argument(
+        "--emo-vector",
+        type=float,
+        nargs=8,
+        default=None,
+        help="8-dim emotion vector: happy angry sad afraid disgusted melancholic surprised calm.",
+    )
+    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(
+        "--output-dir",
+        default=os.path.join(
+            os.environ.get("XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache")),
+            "indextts2_output",
+        ),
+    )
+    parser.add_argument("--deploy-config", default=None)
+    parser.add_argument("--stage-init-timeout", type=int, default=600)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    main(parse_args())

examples/online_serving/text_to_speech/README.md

@@ -16,6 +16,7 @@ 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 | — |
 | 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) | — | ✓ |
@@ -34,7 +35,7 @@ Launch the server (defaults shown — adjust `--port`, `--gpu-memory-utilization
 vllm serve <hf-repo-or-local-path> --omni --port 8091
 ```
 
-Send a TTS request via curl:
+Send a TTS request via curl. These generic snippets assume a model with a preset/default voice; voice-cloning-only models such as IndexTTS-2 require `ref_audio` or an uploaded audio `voice` (see model-specific sections below).
 
 ```bash
 curl -X POST http://localhost:8091/v1/audio/speech \
@@ -90,7 +91,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
     }' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
 ```
 
-Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 24 kHz for the others).
+Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 22.05 kHz for IndexTTS-2, and 24 kHz for many others).
 
 For full request-shape documentation (all parameters, response formats, error codes), see the [Speech API reference](../../../docs/serving/speech_api.md).
 
@@ -135,6 +136,40 @@ bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
 
 ---
 
+## IndexTTS-2
+
+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`.
+
+### 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
+```
+
+### Sending requests
+```bash
+# Voice cloning (ref_audio required)
+python examples/online_serving/text_to_speech/indextts2/speech_client.py \
+    --text "你好，世界！" \
+    --ref-audio /path/to/reference.wav
+
+# With emotion audio
+python examples/online_serving/text_to_speech/indextts2/speech_client.py \
+    --text "今天心情很好！" \
+    --ref-audio /path/to/ref.wav \
+    --emo-audio /path/to/happy.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.
+- 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).
+
+---
+
 ## Fish Speech S2 Pro
 
 4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec.

examples/online_serving/text_to_speech/indextts2/run_server.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+# Launch vLLM-Omni server for IndexTTS2
+#
+# 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
+
+set -e
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+ROOT_DIR="$(cd -- "$SCRIPT_DIR/../../../.." && pwd)"
+
+MODEL="${MODEL:-IndexTeam/IndexTTS-2}"
+PORT="${PORT:-8092}"
+DEPLOY_CONFIG="${DEPLOY_CONFIG:-$ROOT_DIR/vllm_omni/deploy/indextts2.yaml}"
+
+echo "Starting IndexTTS2 server with model: $MODEL"
+
+FLASHINFER_DISABLE_VERSION_CHECK=1 \
+vllm serve "$MODEL" \
+    --host 0.0.0.0 \
+    --port "$PORT" \
+    --omni \
+    --trust-remote-code \
+    --deploy-config "$DEPLOY_CONFIG"

examples/online_serving/text_to_speech/indextts2/speech_client.py

@@ -0,0 +1,126 @@
+# 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.
+
+Examples:
+    # With reference audio for voice cloning
+    python speech_client.py --text "你好，世界！" \
+        --ref-audio /path/to/reference.wav
+
+    # With emotion audio
+    python speech_client.py --text "今天心情很好！" \
+        --ref-audio /path/to/ref.wav \
+        --emo-audio /path/to/happy.wav
+
+Server setup:
+    vllm serve IndexTeam/IndexTTS-2 --omni --host 0.0.0.0 --port 8092
+"""
+
+from __future__ import annotations
+
+import argparse
+import base64
+import os
+
+import httpx
+
+DEFAULT_API_BASE = "http://localhost:8092"
+DEFAULT_API_KEY = "sk-empty"
+
+
+def encode_audio_to_base64(audio_path: str) -> str:
+    if not os.path.exists(audio_path):
+        raise FileNotFoundError(f"Audio file not found: {audio_path}")
+    ext = audio_path.lower().rsplit(".", 1)[-1]
+    mime = {"wav": "audio/wav", "mp3": "audio/mpeg", "flac": "audio/flac"}.get(ext, "audio/wav")
+    with open(audio_path, "rb") as f:
+        b64 = base64.b64encode(f.read()).decode("utf-8")
+    return f"data:{mime};base64,{b64}"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="IndexTTS2 OpenAI speech client")
+    parser.add_argument("--text", type=str, required=True)
+    parser.add_argument("--ref-audio", type=str, default=None, help="Reference audio for voice cloning")
+    parser.add_argument("--emo-audio", type=str, default=None, help="Emotion reference audio")
+    parser.add_argument("--emo-text", type=str, default=None, help="Emotion description text")
+    parser.add_argument(
+        "--emo-vector",
+        type=float,
+        nargs=8,
+        default=None,
+        help="8-dim emotion vector: happy angry sad afraid disgusted melancholic surprised calm",
+    )
+    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("--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 not args.ref_audio and not args.voice:
+        parser.error("IndexTTS2 requires --ref-audio or --voice for voice cloning")
+
+    payload: dict = {
+        "model": args.model,
+        "input": args.text,
+        "response_format": args.response_format,
+    }
+    if args.voice:
+        payload["voice"] = args.voice
+
+    if args.ref_audio:
+        ref = args.ref_audio
+        if ref.startswith(("http://", "https://", "data:")):
+            payload["ref_audio"] = ref
+        else:
+            payload["ref_audio"] = encode_audio_to_base64(ref)
+
+    extra_params = {}
+    if args.emo_audio:
+        emo = args.emo_audio
+        if emo.startswith(("http://", "https://", "data:")):
+            extra_params["emo_audio"] = emo
+        else:
+            extra_params["emo_audio"] = encode_audio_to_base64(emo)
+    if args.emo_text:
+        extra_params["emo_text"] = args.emo_text
+    if args.emo_vector is not None:
+        extra_params["emo_vector"] = args.emo_vector
+    if args.emo_alpha is not None:
+        extra_params["emo_alpha"] = args.emo_alpha
+    if args.use_emo_text:
+        extra_params["use_emo_text"] = True
+    if args.use_random:
+        extra_params["use_random"] = True
+    if extra_params:
+        payload["extra_params"] = extra_params
+
+    url = f"{args.api_base}/v1/audio/speech"
+    print(f"POST {url}")
+    print(f"  text: {args.text}")
+    if args.ref_audio:
+        print(f"  ref_audio: {args.ref_audio[:80]}...")
+
+    with httpx.Client(timeout=300) as client:
+        resp = client.post(
+            url,
+            json=payload,
+            headers={"Authorization": f"Bearer {args.api_key}"},
+        )
+
+    if resp.status_code != 200:
+        print(f"Error {resp.status_code}: {resp.text[:500]}")
+        return
+
+    with open(args.output, "wb") as f:
+        f.write(resp.content)
+    print(f"Saved: {args.output} ({len(resp.content):,} bytes)")
+
+
+if __name__ == "__main__":
+    main()

recipes/IndexTeam/IndexTTS-2.md

@@ -0,0 +1,139 @@
+# IndexTTS-2 for voice-cloned TTS on 1x GPU
+
+## Summary
+
+- Vendor: IndexTeam
+- Model: `IndexTeam/IndexTTS-2`
+- Task: Text-to-speech with voice cloning and optional emotion control
+- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to serve IndexTTS-2 as a two-stage TTS system for
+voice-cloned speech synthesis. Each request provides synthesis text plus either
+a reference audio clip or an uploaded audio voice; the model returns 22.05 kHz
+mono speech. IndexTTS-2 also supports optional emotion conditioning through
+model-specific `extra_params`.
+
+IndexTTS-2 is **not an async-chunk streaming model** in the bundled deployment
+configuration: the S2Mel flow-matching stage consumes the full mel-code sequence
+from the AR talker before it produces audio. The OpenAI-compatible endpoint can
+still be called with `stream=true` as an HTTP compatibility path, but that should
+not be interpreted as incremental model streaming.
+
+## References
+
+- Upstream model:
+  [IndexTeam/IndexTTS-2 on Hugging Face](https://huggingface.co/IndexTeam/IndexTTS-2)
+- Online serving example:
+  [`examples/online_serving/text_to_speech/README.md#indextts-2`](../../examples/online_serving/text_to_speech/README.md#indextts-2)
+- Offline inference example:
+  [`examples/offline_inference/text_to_speech/README.md#indextts-2`](../../examples/offline_inference/text_to_speech/README.md#indextts-2)
+- OpenAI-compatible client:
+  [`examples/online_serving/text_to_speech/indextts2/speech_client.py`](../../examples/online_serving/text_to_speech/indextts2/speech_client.py)
+- Deploy config:
+  [`vllm_omni/deploy/indextts2.yaml`](../../vllm_omni/deploy/indextts2.yaml)
+
+## Hardware Support
+
+### GPU
+
+### 1x L4 24GB or larger CUDA GPU
+
+#### 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
+
+#### Command
+
+Start the server from the repository root:
+
+```bash
+vllm serve IndexTeam/IndexTTS-2 --omni --trust-remote-code --port 8092
+```
+
+The bundled deployment config is selected automatically for the `indextts2`
+model type. To make the config choice explicit, pass it directly:
+
+```bash
+vllm serve IndexTeam/IndexTTS-2 --omni --trust-remote-code --port 8092 \
+  --deploy-config vllm_omni/deploy/indextts2.yaml
+```
+
+#### Verification
+
+Run the bundled client. Local audio paths passed to this client are converted to
+base64 data URLs before they are sent to the HTTP API:
+
+```bash
+python examples/online_serving/text_to_speech/indextts2/speech_client.py \
+  --api-base http://localhost:8092 \
+  --text "你好，这是 IndexTTS-2 语音合成测试。" \
+  --ref-audio /path/to/reference.wav \
+  --output indextts2.wav
+```
+
+For a raw HTTP request, pass `ref_audio` as an HTTP(S) URL, a base64 data URL,
+or a `file://` URI allowed by the server's `--allowed-local-media-path`. Do not
+pass an arbitrary local filesystem path directly in JSON unless your client
+converts it first:
+
+```bash
+curl http://localhost:8092/v1/audio/speech \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "IndexTeam/IndexTTS-2",
+    "input": "你好，这是 IndexTTS-2 语音合成测试。",
+    "response_format": "wav",
+    "ref_audio": "data:audio/wav;base64,<BASE64_ENCODED_AUDIO>"
+  }' \
+  --output indextts2.wav
+```
+
+Optional emotion conditioning is passed through `extra_params`. Use one emotion
+source mode at a time:
+
+```bash
+curl http://localhost:8092/v1/audio/speech \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "IndexTeam/IndexTTS-2",
+    "input": "今天心情很好！",
+    "response_format": "wav",
+    "ref_audio": "data:audio/wav;base64,<BASE64_ENCODED_REFERENCE_AUDIO>",
+    "extra_params": {
+      "emo_vector": [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+      "emo_alpha": 0.8
+    }
+  }' \
+  --output indextts2_happy.wav
+```
+
+#### Notes
+
+- Hardware scope: the default config is a single-GPU CUDA deployment. Stage 0
+  and Stage 1 share GPU 0 by default.
+- Audio output: 22.05 kHz mono WAV.
+- Voice cloning: `ref_audio` is required for the documented raw request path.
+  Alternatively, pass `voice` only when it names an uploaded audio voice from
+  `/v1/audio/voices`; IndexTTS-2 does not provide a built-in text-only preset
+  voice in this serving recipe.
+- Reference audio format: use an HTTP(S) URL, a base64 data URL, or a `file://`
+  URI allowed by the server's `--allowed-local-media-path` for raw API calls.
+  The bundled Python client accepts local paths and converts them to base64 data
+  URLs.
+- Emotion control: `extra_params.emo_audio`, `extra_params.emo_vector`, and
+  `extra_params.use_emo_text` are alternative conditioning modes. Official
+  precedence is `use_emo_text` > `emo_vector` > `emo_audio` > same emotion as
+  the speaker reference. The `emo_vector` order is `[happy, angry, sad, afraid,
+  disgusted, melancholic, surprised, calm]`.
+- Streaming limitation: `async_chunk` is disabled in
+  `vllm_omni/deploy/indextts2.yaml`; S2Mel requires the full mel-code sequence
+  before audio generation.
+- First request can be slower because reference encoders and auxiliary models
+  are lazily loaded and cached.

recipes/README.md

@@ -43,6 +43,7 @@ recipes/
 | [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
 | [`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 |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
 | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
 | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |

tests/core/sched/test_omni_scheduling_coordinator.py

@@ -115,6 +115,7 @@ class TestFullPayloadCoordinatorSelection(unittest.TestCase):
             ("MiMoAudioModel", "code2wav"),
             ("Qwen3TTSCode2Wav", "code2wav"),
             ("CosyVoice3Model", "cosyvoice3_code2wav"),
+            ("IndexTTS2S2MelDecoder", "indextts2_s2mel_decoder"),
             ("DyninOmniForConditionalGeneration", "token2image"),
             ("DyninOmniForConditionalGeneration", "token2audio"),
         }

tests/e2e/offline_inference/test_indextts2.py

@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E offline inference tests for IndexTTS2 two-stage pipeline.
+
+Stage 0 (GPT AR) → Stage 1 (S2Mel + BigVGAN). Output is 22050 Hz mono WAV.
+Covers: basic TTS.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni import Omni
+from vllm_omni.model_executor.models.indextts2.prompt_utils import (
+    build_indextts2_prefill_prompt_ids,
+)
+
+MODEL_NAME = "IndexTeam/IndexTTS-2"
+STAGE_CONFIG = get_deploy_config_path("indextts2.yaml")
+
+_OMNI_RUNNER_PARAM = (
+    MODEL_NAME,
+    STAGE_CONFIG,
+)
+
+ASSETS_DIR = Path(__file__).resolve().parents[2] / "assets" / "indextts2"
+REF_AUDIO = str(ASSETS_DIR / "ref_audio.wav")
+
+pytestmark = [
+    pytest.mark.advanced_model,
+    pytest.mark.tts,
+    pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
+]
+
+SAMPLE_RATE = 22050
+
+
+def _audio_from_mm(mm: dict) -> torch.Tensor | None:
+    audio = mm.get("audio")
+    if audio is None:
+        audio = mm.get("model_outputs")
+    if isinstance(audio, list):
+        chunks = [chunk.reshape(-1) for chunk in audio if isinstance(chunk, torch.Tensor) and chunk.numel() > 0]
+        return torch.cat(chunks, dim=0) if chunks else None
+    return audio if isinstance(audio, torch.Tensor) else None
+
+
+def _sample_rate_from_mm(mm: dict) -> int:
+    sr = mm.get("sr")
+    if isinstance(sr, list):
+        sr = sr[-1] if sr else None
+    if hasattr(sr, "item"):
+        return int(sr.item())
+    return int(sr) if sr is not None else SAMPLE_RATE
+
+
+def _build_request(text: str, ref_audio: str, **extra) -> dict:
+    info = {
+        "text": [text],
+        "voice": [ref_audio],
+    }
+    info.update(extra)
+    return {
+        "prompt_token_ids": build_indextts2_prefill_prompt_ids(MODEL_NAME, text),
+        "additional_information": info,
+    }
+
+
+def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]:
+    for omni_out in omni.generate(request):
+        mm = omni_out.multimodal_output
+        assert mm is not None, "Expected multimodal_output"
+        audio = _audio_from_mm(mm)
+        assert audio is not None, "Expected audio output"
+        return audio.cpu(), _sample_rate_from_mm(mm)
+    raise AssertionError("No outputs received")
+
+
+def _assert_valid_audio(audio: torch.Tensor, sr: int) -> None:
+    assert sr == SAMPLE_RATE
+    assert audio.numel() > 0
+    assert torch.any(audio != 0).item()
+
+
+@hardware_test(res={"cuda": "L4"})
+def test_basic_english(omni_runner: OmniRunner) -> None:
+    """English TTS: text + ref_audio, no emotion conditioning."""
+    req = _build_request("Hello, this is a voice synthesis test.", REF_AUDIO)
+    audio, sr = _collect_audio(omni_runner.omni, req)
+    _assert_valid_audio(audio, sr)

tests/e2e/online_serving/test_indextts2.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E online serving tests for IndexTTS2 via /v1/audio/speech endpoint.
+
+Two-stage pipeline: GPT AR → S2Mel + BigVGAN. Output is 22050 Hz mono WAV.
+Covers: basic TTS and the speech endpoint's stream=True compatibility path.
+"""
+
+from __future__ import annotations
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import load_test_audio_data_url
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+pytestmark = [pytest.mark.advanced_model, pytest.mark.tts]
+
+MODEL = "IndexTeam/IndexTTS-2"
+
+REF_AUDIO_URL = load_test_audio_data_url("indextts2/ref_audio.wav")
+
+tts_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_deploy_config_path("indextts2.yaml"),
+            server_args=["--disable-log-stats"],
+        ),
+        id="indextts2",
+    )
+]
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_basic_english(omni_server, openai_client) -> None:
+    """
+    Basic English TTS: text + ref_audio, no emotion.
+    Deploy Setting: default yaml
+    Input Modal: text + reference audio
+    Output Modal: audio (22050 Hz, WAV)
+    Input Setting: stream=False
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": "Hello, this is a voice cloning demo.",
+        "stream": False,
+        "response_format": "wav",
+        "ref_audio": REF_AUDIO_URL,
+        "min_audio_bytes": 1024,
+    }
+    openai_client.send_audio_speech_request(request_config)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_streaming(omni_server, openai_client) -> None:
+    """
+    stream=True compatibility: IndexTTS2 itself is non-async-chunking, but the
+    OpenAI speech endpoint should still accept the streaming request path and
+    return a valid PCM response.
+
+    Deploy Setting: default yaml
+    Input Modal: text + reference audio
+    Output Modal: audio (22050 Hz, PCM)
+    Input Setting: stream=True
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": "Hello, this is a voice cloning demo.",
+        "stream": True,
+        "response_format": "pcm",
+        "ref_audio": REF_AUDIO_URL,
+        "min_audio_bytes": 1024,
+    }
+    openai_client.send_audio_speech_request(request_config)

tests/model_executor/stage_input_processors/test_indextts2.py

@@ -0,0 +1,336 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import torch
+
+from vllm_omni.model_executor.stage_input_processors.indextts2 import (
+    _strip_stop_token,
+    talker2s2mel,
+    talker2s2mel_full_payload,
+    talker2s2mel_token_only,
+)
+
+STOP_MEL_TOKEN = 8193
+LATENT_DIM = 16
+
+
+def _make_talker_output(
+    *,
+    mel_codes: torch.Tensor,
+    latent: torch.Tensor,
+    s_ref: torch.Tensor | None = None,
+    ref_mel: torch.Tensor | None = None,
+    style: torch.Tensor | None = None,
+    finished: bool = True,
+):
+    meta = {}
+    if s_ref is not None:
+        meta["S_ref"] = s_ref
+    if ref_mel is not None:
+        meta["ref_mel"] = ref_mel
+    if style is not None:
+        meta["style"] = style
+
+    mm = {
+        "codes": {"mel": mel_codes},
+        "hidden_states": {"latent": latent},
+        "meta": meta,
+    }
+    return SimpleNamespace(
+        finished=finished,
+        outputs=[SimpleNamespace(multimodal_output=mm)],
+    )
+
+
+# ── _strip_stop_token ──
+
+
+def test_strip_stop_token_basic():
+    codes = torch.tensor([10, 20, 30, STOP_MEL_TOKEN, 99])
+    latent = torch.randn(5, LATENT_DIM)
+
+    c, lat, lens = _strip_stop_token(codes, latent)
+
+    assert c.shape == (1, 3)
+    assert lat.shape == (1, 3, LATENT_DIM)
+    assert lens.tolist() == [3]
+    assert (c[0] == torch.tensor([10, 20, 30])).all()
+
+
+def test_strip_stop_token_no_stop():
+    codes = torch.tensor([10, 20, 30])
+    latent = torch.randn(3, LATENT_DIM)
+
+    c, lat, lens = _strip_stop_token(codes, latent)
+
+    assert c.shape == (1, 3)
+    assert lens.tolist() == [3]
+
+
+def test_strip_stop_token_at_start():
+    codes = torch.tensor([STOP_MEL_TOKEN, 10, 20])
+    latent = torch.randn(3, LATENT_DIM)
+
+    c, lat, lens = _strip_stop_token(codes, latent)
+
+    assert c.shape == (1, 0)
+    assert lat.shape == (1, 0, LATENT_DIM)
+    assert lens.tolist() == [0]
+
+
+def test_strip_stop_token_batch_with_padding():
+    codes = torch.tensor(
+        [
+            [10, 20, STOP_MEL_TOKEN, 0],
+            [30, 40, 50, STOP_MEL_TOKEN],
+        ]
+    )
+    latent = torch.randn(2, 4, LATENT_DIM)
+
+    c, lat, lens = _strip_stop_token(codes, latent)
+
+    assert c.shape == (2, 3)
+    assert lat.shape == (2, 3, LATENT_DIM)
+    assert lens.tolist() == [2, 3]
+    assert c[0, 2].item() == 0  # padded with a valid semantic code
+
+
+def test_strip_stop_token_all_stop():
+    codes = torch.tensor([STOP_MEL_TOKEN])
+    latent = torch.randn(1, LATENT_DIM)
+
+    c, lat, lens = _strip_stop_token(codes, latent)
+
+    assert lens.tolist() == [0]
+    assert c.shape[1] == 0
+    assert lat.shape[1] == 0
+
+
+# ── talker2s2mel ──
+
+
+def test_talker2s2mel_basic():
+    mel_codes = torch.tensor([10, 20, 30, STOP_MEL_TOKEN])
+    latent = torch.randn(4, LATENT_DIM)
+    s_ref = torch.randn(1, 128)
+    ref_mel = torch.randn(1, 80, 200)
+    style = torch.randn(1, 256)
+
+    outputs = talker2s2mel(
+        source_outputs=[
+            _make_talker_output(
+                mel_codes=mel_codes,
+                latent=latent,
+                s_ref=s_ref,
+                ref_mel=ref_mel,
+                style=style,
+            )
+        ]
+    )
+
+    assert len(outputs) == 1
+    info = outputs[0]["additional_information"]
+    assert info["mel_codes"].shape == (1, 3)
+    assert info["latent"].shape == (1, 3, LATENT_DIM)
+    assert info["code_lens"].tolist() == [3]
+    assert "S_ref" in info
+    assert "ref_mel" in info
+    assert "style" in info
+    assert outputs[0]["prompt_token_ids"] == [0]
+
+
+def test_talker2s2mel_skips_unfinished():
+    mel_codes = torch.tensor([10, 20])
+    latent = torch.randn(2, LATENT_DIM)
+
+    outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent, finished=False)])
+
+    assert len(outputs) == 0
+
+
+def test_talker2s2mel_skips_empty_mel():
+    outputs = talker2s2mel(
+        source_outputs=[
+            _make_talker_output(
+                mel_codes=torch.tensor([]),
+                latent=torch.randn(0, LATENT_DIM),
+            )
+        ]
+    )
+
+    assert len(outputs) == 0
+
+
+def test_talker2s2mel_missing_optional_meta():
+    mel_codes = torch.tensor([10, 20])
+    latent = torch.randn(2, LATENT_DIM)
+
+    outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent)])
+
+    assert len(outputs) == 1
+    info = outputs[0]["additional_information"]
+    assert "S_ref" not in info
+    assert "ref_mel" not in info
+    assert "style" not in info
+
+
+def test_talker2s2mel_batch():
+    out0 = _make_talker_output(
+        mel_codes=torch.tensor([1, 2, STOP_MEL_TOKEN]),
+        latent=torch.randn(3, LATENT_DIM),
+        s_ref=torch.randn(1, 64),
+    )
+    out1 = _make_talker_output(
+        mel_codes=torch.tensor([5, 6, 7, 8]),
+        latent=torch.randn(4, LATENT_DIM),
+        ref_mel=torch.randn(1, 80, 100),
+    )
+
+    outputs = talker2s2mel(source_outputs=[out0, out1])
+
+    assert len(outputs) == 2
+    assert outputs[0]["additional_information"]["code_lens"].tolist() == [2]
+    assert outputs[1]["additional_information"]["code_lens"].tolist() == [4]
+
+
+def test_talker2s2mel_tensors_on_cpu():
+    device = torch.device("cpu")
+    mel_codes = torch.tensor([10, 20], device=device)
+    latent = torch.randn(2, LATENT_DIM, device=device)
+    s_ref = torch.randn(1, 64, device=device)
+
+    outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent, s_ref=s_ref)])
+
+    info = outputs[0]["additional_information"]
+    assert info["latent"].device.type == "cpu"
+    assert info["mel_codes"].device.type == "cpu"
+    assert info["S_ref"].device.type == "cpu"
+
+
+def test_talker2s2mel_no_multimodal_dict():
+    output = SimpleNamespace(
+        finished=True,
+        outputs=[SimpleNamespace(multimodal_output="not_a_dict")],
+    )
+
+    results = talker2s2mel(source_outputs=[output])
+
+    assert len(results) == 0
+
+
+# ── full-payload connector path ──
+
+
+def _conditioning_tensors():
+    return {
+        "S_ref": torch.randn(1, 3, LATENT_DIM),
+        "ref_mel": torch.randn(1, 80, 5),
+        "style": torch.randn(1, 4),
+    }
+
+
+def test_talker2s2mel_token_only_placeholder_finished_only():
+    finished = _make_talker_output(mel_codes=torch.tensor([1]), latent=torch.randn(1, LATENT_DIM), finished=True)
+    unfinished = _make_talker_output(mel_codes=torch.tensor([2]), latent=torch.randn(1, LATENT_DIM), finished=False)
+
+    outputs = talker2s2mel_token_only([unfinished, finished])
+
+    assert len(outputs) == 1
+    assert outputs[0]["prompt_token_ids"] == [0]
+    assert outputs[0]["additional_information"] is None
+    assert outputs[0]["multi_modal_data"] is None
+
+
+def test_talker2s2mel_full_payload_flat_keys_builds_s2mel_contract():
+    cond = _conditioning_tensors()
+    latent = torch.randn(3, LATENT_DIM)
+    payload = {
+        "codes.mel": torch.tensor([[10], [20], [30]]),
+        "hidden_states.latent": latent,
+        "meta.S_ref": cond["S_ref"],
+        "meta.ref_mel": cond["ref_mel"],
+        "meta.style": cond["style"],
+    }
+
+    result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-flat"))
+
+    assert result is not None
+    assert result["mel_codes"].shape == (1, 3)
+    assert result["mel_codes"].tolist() == [[10, 20, 30]]
+    assert result["latent"].shape == (1, 3, LATENT_DIM)
+    assert result["code_lens"].tolist() == [3]
+    assert result["S_ref"].device.type == "cpu"
+    assert result["ref_mel"].device.type == "cpu"
+    assert result["style"].device.type == "cpu"
+    assert result["S_ref"].data_ptr() != cond["S_ref"].data_ptr()
+
+
+def test_talker2s2mel_full_payload_nested_fallback_input():
+    cond = _conditioning_tensors()
+    payload = {
+        "codes": {"mel": torch.tensor([[4], [5]])},
+        "hidden_states": {"latent": torch.randn(2, LATENT_DIM)},
+        "meta": cond,
+    }
+
+    result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-nested"))
+
+    assert result is not None
+    assert result["mel_codes"].tolist() == [[4, 5]]
+    assert result["latent"].shape == (1, 2, LATENT_DIM)
+
+
+def test_talker2s2mel_full_payload_normalizes_one_by_t_mel_row():
+    cond = _conditioning_tensors()
+    payload = {
+        "codes.mel": torch.tensor([[7, 8, 9]]),
+        "hidden_states.latent": torch.randn(3, LATENT_DIM),
+        "meta": cond,
+    }
+
+    result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-row"))
+
+    assert result is not None
+    assert result["mel_codes"].tolist() == [[7, 8, 9]]
+
+
+def test_talker2s2mel_full_payload_trims_stop_token_and_matching_latent():
+    cond = _conditioning_tensors()
+    latent = torch.arange(4 * LATENT_DIM, dtype=torch.float32).reshape(4, LATENT_DIM)
+    payload = {
+        "codes.mel": torch.tensor([[1], [2], [STOP_MEL_TOKEN], [99]]),
+        "hidden_states.latent": latent,
+        "meta": cond,
+    }
+
+    result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-stop"))
+
+    assert result is not None
+    assert result["mel_codes"].tolist() == [[1, 2]]
+    assert result["latent"].shape == (1, 2, LATENT_DIM)
+    assert result["latent"][0, 0].tolist() == latent[0].tolist()
+    assert result["latent"][0, 1].tolist() == latent[1].tolist()
+    assert result["code_lens"].tolist() == [2]
+
+
+def test_talker2s2mel_full_payload_crops_mel_latent_to_common_length():
+    cond = _conditioning_tensors()
+    payload = {
+        "codes.mel": torch.tensor([[1], [2], [3], [4]]),
+        "hidden_states.latent": torch.randn(2, LATENT_DIM),
+        "meta": cond,
+    }
+
+    result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-crop"))
+
+    assert result is not None
+    assert result["mel_codes"].tolist() == [[1, 2]]
+    assert result["latent"].shape == (1, 2, LATENT_DIM)
+    assert result["code_lens"].tolist() == [2]
+
+
+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

vllm_omni/config/pipeline_registry.py

@@ -169,4 +169,8 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
         "vllm_omni.model_executor.models.dynin_omni.pipeline",
         "DYNIN_OMNI_PIPELINE",
     ),
+    "indextts2": (
+        "vllm_omni.model_executor.models.indextts2.pipeline",
+        "INDEXTTS2_PIPELINE",
+    ),
 }

vllm_omni/core/sched/omni_scheduling_coordinator.py

@@ -51,6 +51,9 @@ _FULL_PAYLOAD_INPUT_STAGES: frozenset[tuple[str, str]] = frozenset(
         ("Qwen3TTSCode2Wav", "code2wav"),
         # cosyvoice3: cosyvoice3_talker (Stage 0) -> cosyvoice3_code2wav (Stage 1).
         ("CosyVoice3Model", "cosyvoice3_code2wav"),
+        # indextts2: indextts2_talker (Stage 0) -> indextts2_s2mel_decoder
+        # (Stage 1). Stage 1 consumes the complete mel/latent payload.
+        ("IndexTTS2S2MelDecoder", "indextts2_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.

vllm_omni/deploy/indextts2.yaml

@@ -0,0 +1,57 @@
+async_chunk: false
+trust_remote_code: true
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      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
+    gpu_memory_utilization: 0.4
+    enable_prefix_caching: false
+    skip_tokenizer_init: true
+    tokenizer: gpt2
+    async_scheduling: true
+    max_num_batched_tokens: 512
+    max_model_len: 2560
+    devices: "0"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    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
+    tokenizer: gpt2
+    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:
+      diffusion_steps: 12
+      s2mel_dit_bf16: true
+      s2mel_vocoder_bf16: true
+      s2mel_vocoder_cuda_graph: true
+      s2mel_dit_cuda_graph: true
+      s2mel_dit_cuda_graph_max_graphs: 4
+      s2mel_vocoder_capture_sizes: [128, 256, 384, 512]
+      s2mel_vocoder_compile_shapes: [256, 384, 512]
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 32767

vllm_omni/deploy/indextts2_low_latency.yaml

@@ -0,0 +1,76 @@
+# IndexTTS2 low-latency deploy profile.
+#
+# Stage 0 (GPT2 AR talker) uses vLLM FULL_DECODE_ONLY CUDA graph + FlashInfer
+# attention backend + async_scheduling=false, mirroring the Higgs low-latency
+# profile. This reduces decode launch overhead and per-step scheduling
+# fragmentation for c1-c4 workloads.
+#
+# Usage:
+#   vllm serve IndexTeam/IndexTTS-2 --omni --trust-remote-code \
+#     --deploy-config vllm_omni/deploy/indextts2_low_latency.yaml
+#
+# For 2-GPU testing, copy this file and set stage 1 devices: "1".
+
+async_chunk: false
+trust_remote_code: true
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      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: 16
+    gpu_memory_utilization: 0.6
+    enable_prefix_caching: false
+    skip_tokenizer_init: true
+    tokenizer: gpt2
+    enforce_eager: false
+    attention_backend: FLASHINFER
+    async_scheduling: false
+    max_num_batched_tokens: 4096
+    max_model_len: 2560
+    devices: "0"
+    compilation_config:
+      cudagraph_mode: FULL_DECODE_ONLY
+      cudagraph_capture_sizes: [1, 2, 4, 8, 16]
+      cudagraph_num_of_warmups: 1
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    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
+    tokenizer: gpt2
+    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:
+      diffusion_steps: 12
+      s2mel_dit_bf16: true
+      s2mel_vocoder_bf16: true
+      s2mel_vocoder_cuda_graph: true
+      s2mel_dit_cuda_graph: true
+      s2mel_dit_cuda_graph_max_graphs: 4
+      s2mel_vocoder_capture_sizes: [128, 256, 384, 512]
+      s2mel_vocoder_compile_shapes: [256, 384, 512]
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 32767

vllm_omni/engine/arg_utils.py

@@ -20,6 +20,8 @@ logger = init_logger(__name__)
 _ARCH_TO_MODEL_TYPE: dict[str, str] = {
     "CosyVoice3Model": "cosyvoice3",
     "GLMTTSForConditionalGeneration": "glm_tts",
+    "IndexTTS2S2MelDecoder": "indextts2",
+    "IndexTTS2TalkerForConditionalGeneration": "indextts2",
     "OmniVoiceModel": "omnivoice",
     "VoxCPM2TalkerForConditionalGeneration": "voxcpm2",
 }
@@ -35,6 +37,9 @@ def _register_omni_hf_configs() -> None:
     try:
         from transformers import AutoConfig
 
+        from vllm_omni.model_executor.models.indextts2.configuration_indextts2 import (
+            IndexTTS2Config,
+        )
         from vllm_omni.model_executor.models.ming_tts.config_ming_tts import MingDenseConfig
         from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
             Qwen3TTSConfig,
@@ -57,6 +62,7 @@ def _register_omni_hf_configs() -> None:
 
     for model_type, config_cls in [
         ("dense", MingDenseConfig),
+        ("indextts2", IndexTTS2Config),
         ("qwen3_tts", Qwen3TTSConfig),
         ("cosyvoice3", CosyVoice3Config),
         ("glm_tts", GLMTTSConfig),

vllm_omni/entrypoints/openai/serving_speech.py

@@ -85,6 +85,7 @@ _HIGGS_AUDIO_V2_TTS_MODEL_STAGES = {"higgs_audio_v2"}
 _HIGGS_V3_TTS_MODEL_STAGES = {"higgs_audio_v3"}
 _GLM_TTS_MODEL_STAGES = {"glm_tts"}
 _STEP_AUDIO2_TTS_MODEL_STAGES = {"step_audio2_thinker"}
+_INDEXTTS2_TTS_MODEL_STAGES = {"indextts2_talker"}
 _TTS_MODEL_STAGES: set[str] = (
     _VOXTRAL_TTS_MODEL_STAGES
     | _QWEN3_TTS_MODEL_STAGES
@@ -100,6 +101,7 @@ _TTS_MODEL_STAGES: set[str] = (
     | _MOSS_TTS_FULL_MODEL_STAGES
     | _GLM_TTS_MODEL_STAGES
     | _STEP_AUDIO2_TTS_MODEL_STAGES
+    | _INDEXTTS2_TTS_MODEL_STAGES
 )
 _SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = {
     "fish_tts",
@@ -109,6 +111,7 @@ _SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = {
     "voxcpm2",
     "higgs_audio_v2",
     "higgs_audio_v3",
+    "indextts2",
 }
 _TTS_LANGUAGES = frozenset(
     {
@@ -693,6 +696,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             return "glm_tts"
         if model_stage in _STEP_AUDIO2_TTS_MODEL_STAGES:
             return "step_audio2"
+        if model_stage in _INDEXTTS2_TTS_MODEL_STAGES:
+            return "indextts2"
         return None
 
     def _get_custom_voice_dir(self) -> str | None:
@@ -3429,6 +3434,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             model_type = "higgs_audio_v2"
         elif self._tts_model_type == "glm_tts":
             model_type = "glm_tts"
+        elif self._tts_model_type == "indextts2":
+            model_type = "indextts2"
         elif self._is_tts:
             model_type = tts_params.get("task_type", ["unknown"])[0]
         else:

vllm_omni/entrypoints/openai/tts_adapters/__init__.py

@@ -57,6 +57,7 @@ from vllm_omni.entrypoints.openai.tts_adapters import (  # noqa: E402,F401
     glm_tts,
     higgs_audio_v2,
     higgs_audio_v3,
+    indextts2,
     ming_tts,
     moss_tts,
     omnivoice,

vllm_omni/entrypoints/openai/tts_adapters/indextts2.py

@@ -0,0 +1,243 @@
+# SPDX-License-Identifier: Apache-2.0
+"""IndexTTS2 serving adapter."""
+
+from __future__ import annotations
+
+import hashlib
+import math
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import torch
+from vllm.inputs import tokens_input
+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
+
+if TYPE_CHECKING:
+    from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+
+_TTS_MAX_NEW_TOKENS_MIN = 1
+_TTS_MAX_NEW_TOKENS_MAX = 4096
+_INDEXTTS2_EMOTION_KEYS = ("emo_audio", "emo_vector", "emo_alpha", "emo_text", "use_emo_text", "use_random")
+
+
+def _update_conditioning_hash(h: hashlib._Hash, value: Any) -> None:
+    """Hash request conditioning values without serializing huge repr strings."""
+    if value is None:
+        h.update(b"none")
+        return
+    if isinstance(value, torch.Tensor):
+        arr = value.detach().cpu().numpy()
+        h.update(b"tensor")
+        h.update(repr((arr.shape, str(arr.dtype))).encode("utf-8"))
+        h.update(np.ascontiguousarray(arr).tobytes())
+        return
+    if isinstance(value, np.ndarray):
+        arr = np.ascontiguousarray(value)
+        h.update(b"ndarray")
+        h.update(repr((arr.shape, str(arr.dtype))).encode("utf-8"))
+        h.update(arr.tobytes())
+        return
+    if isinstance(value, (list, tuple)) and len(value) == 2 and isinstance(value[1], (int, np.integer)):
+        samples = value[0]
+        if isinstance(samples, (list, tuple, np.ndarray, torch.Tensor)):
+            if isinstance(samples, torch.Tensor):
+                samples_arr = samples.detach().cpu().numpy()
+            else:
+                samples_arr = np.asarray(samples, dtype=np.float32)
+            samples_arr = np.ascontiguousarray(samples_arr, dtype=np.float32)
+            h.update(b"audio")
+            h.update(int(value[1]).to_bytes(8, byteorder="little", signed=True))
+            h.update(int(samples_arr.size).to_bytes(8, byteorder="little", signed=False))
+            h.update(samples_arr.tobytes())
+            return
+    if isinstance(value, Mapping):
+        h.update(b"mapping")
+        for key in sorted(value):
+            h.update(b"\x01")
+            h.update(repr(key).encode("utf-8"))
+            h.update(b"\x02")
+            _update_conditioning_hash(h, value[key])
+        return
+    if isinstance(value, (list, tuple)):
+        h.update(f"sequence:{len(value)}".encode())
+        for item in value:
+            h.update(b"\x03")
+            _update_conditioning_hash(h, item)
+        return
+    h.update(repr(value).encode("utf-8"))
+
+
+def _first_conditioning_value(value: Any) -> Any:
+    if isinstance(value, list) and len(value) == 1:
+        return value[0]
+    return value
+
+
+def indextts2_conditioning_cache_salt(
+    request: OpenAICreateSpeechRequest,
+    tts_params: dict[str, Any],
+    *,
+    request_id: str | None = None,
+) -> str:
+    """Stable prefix-cache salt for IndexTTS2 placeholder prompts.
+
+    IndexTTS2 builds prefill embeddings from ``additional_information`` (speaker
+    audio plus emotion controls), while the scheduler only sees dummy token ids.
+    Include all supported control inputs so prefix-cache hits are only possible
+    when text, speaker, and emotion conditioning are identical.
+    """
+    h = hashlib.sha256()
+    for key in (
+        "text",
+        "voice",
+        "voice_name",
+        "voice_created_at",
+        "emo_audio",
+        "emo_vector",
+        "emo_alpha",
+        "emo_text",
+        "use_emo_text",
+        "use_random",
+    ):
+        h.update(b"\x00")
+        h.update(key.encode("utf-8"))
+        h.update(b"\x00")
+        _update_conditioning_hash(h, tts_params.get(key))
+
+    for key, value in (
+        ("input", request.input),
+        ("ref_audio", request.ref_audio),
+    ):
+        h.update(b"\x00")
+        h.update(key.encode("utf-8"))
+        h.update(b"\x00")
+        _update_conditioning_hash(h, value)
+
+    use_random = bool(_first_conditioning_value(tts_params.get("use_random")))
+    if use_random:
+        h.update(b"\x00indextts2-use-random-request\x00")
+        h.update((request_id or random_uuid()).encode("utf-8"))
+    return h.hexdigest()[:32]
+
+
+@register_tts_adapter
+class IndexTTS2Adapter(ARTTSAdapter):
+    stage_keys = frozenset({"indextts2_talker"})
+    name = "indextts2"
+
+    def validate(self, request: OpenAICreateSpeechRequest) -> str | None:
+        server = self.ctx.server
+        if not request.input or not request.input.strip():
+            return "Input text cannot be empty"
+
+        voice_lower = request.voice.lower() if isinstance(request.voice, str) else None
+        has_uploaded_voice = bool(
+            voice_lower
+            and voice_lower in server.uploaded_speakers
+            and server._get_uploaded_audio_data(voice_lower) is not None
+        )
+        if request.ref_audio is None and not has_uploaded_voice:
+            return "IndexTTS2 requires 'ref_audio' or an uploaded audio voice for voice cloning"
+        if request.ref_audio is not None:
+            fmt_err = server._validate_ref_audio_format(request.ref_audio)
+            if fmt_err:
+                return fmt_err
+        if request.max_new_tokens is not None:
+            if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN:
+                return f"max_new_tokens must be at least {_TTS_MAX_NEW_TOKENS_MIN}"
+            if request.max_new_tokens > _TTS_MAX_NEW_TOKENS_MAX:
+                return f"max_new_tokens cannot exceed {_TTS_MAX_NEW_TOKENS_MAX}"
+        if request.extra_params is not None:
+            if not isinstance(request.extra_params, dict):
+                return "extra_params must be a JSON object/dict."
+            extra_err = self._validate_extra_params(request.extra_params)
+            if extra_err:
+                return extra_err
+        return None
+
+    def _validate_extra_params(self, extras: Mapping[str, Any]) -> str | None:
+        server = self.ctx.server
+        if "emo_audio" in extras:
+            value = extras["emo_audio"]
+            if not isinstance(value, str):
+                return "extra_params.emo_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)"
+            fmt_err = server._validate_ref_audio_format(value)
+            if fmt_err:
+                return fmt_err.replace("ref_audio", "extra_params.emo_audio", 1)
+
+        if "emo_vector" in extras:
+            value = extras["emo_vector"]
+            if not isinstance(value, list) or len(value) != 8:
+                return "extra_params.emo_vector must be a list of 8 numbers"
+            for item in value:
+                if isinstance(item, bool) or not isinstance(item, (int, float)) or not math.isfinite(float(item)):
+                    return "extra_params.emo_vector must be a list of 8 finite numbers"
+                if float(item) < 0.0 or float(item) > 1.2:
+                    return "extra_params.emo_vector values must be in [0, 1.2]"
+
+        if "emo_alpha" in extras:
+            value = extras["emo_alpha"]
+            if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(float(value)):
+                return "extra_params.emo_alpha must be a finite number"
+
+        if "emo_text" in extras and extras["emo_text"] is not None and not isinstance(extras["emo_text"], str):
+            return "extra_params.emo_text must be a string"
+
+        for key in ("use_emo_text", "use_random"):
+            if key in extras and not isinstance(extras[key], bool):
+                return f"extra_params.{key} must be a boolean"
+        return None
+
+    async def build(
+        self,
+        request: OpenAICreateSpeechRequest,
+        sampling_params_list: list,
+        has_inline_ref_audio: bool,
+    ) -> PreparedRequest:
+        del sampling_params_list, has_inline_ref_audio
+
+        server = self.ctx.server
+        tts_params = await self._build_params(request)
+        from vllm_omni.model_executor.models.indextts2.prompt_utils import (
+            estimate_indextts2_prefill_prompt_len,
+        )
+
+        ph_len = estimate_indextts2_prefill_prompt_len(
+            server.engine_client.model_config.model,
+            request.input,
+        )
+        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")
+
+    async def _build_params(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]:
+        server = self.ctx.server
+        params: dict[str, Any] = {"text": [request.input]}
+        voice_lower = request.voice.lower() if isinstance(request.voice, str) else None
+        ref_audio_source = request.ref_audio
+        using_uploaded_voice = False
+        if ref_audio_source is None and voice_lower and voice_lower in server.uploaded_speakers:
+            ref_audio_source = server._get_uploaded_audio_data(voice_lower)
+            using_uploaded_voice = ref_audio_source is not None
+        if ref_audio_source is not None and isinstance(ref_audio_source, str):
+            wav_list, sr = await server._resolve_ref_audio(ref_audio_source)
+            params["voice"] = [[wav_list, sr]]
+        if using_uploaded_voice and voice_lower:
+            params["voice_name"] = [voice_lower]
+            params["voice_created_at"] = [server._voice_created_at(voice_lower)]
+
+        extras = request.extra_params if isinstance(request.extra_params, dict) else {}
+        for key in _INDEXTTS2_EMOTION_KEYS:
+            if key not in extras:
+                continue
+            if key == "emo_audio":
+                wav_list, sr = await server._resolve_ref_audio(extras[key])
+                params[key] = [[wav_list, sr]]
+            else:
+                params[key] = [extras[key]]
+        return params

vllm_omni/model_executor/models/indextts2/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/model_executor/models/indextts2/bigvgan_cuda_graph.py

@@ -0,0 +1,284 @@
+"""CUDA graph wrapper for BigVGAN vocoding in IndexTTS2 Stage 1.
+
+BigVGAN is heavily CPU launch-bound (Snake/AliasFree small-op storm:
+~3500 copy_ + ~2000 mul + ~340 cudnn_convolution launches per request,
+GPU busy <15%). Capturing the forward for fixed mel-frame buckets and
+replaying collapses thousands of kernel launches into one graph launch.
+
+Dynamic mel inputs are right zero-padded to the nearest captured bucket
+and the waveform is trimmed back to ``actual_frames * upsample_factor``
+after replay. Padding leakage into the kept region is bounded by the
+conv receptive field at the tail (~20-38 mel frames, ~0.2-0.4s audio)
+and is validated via WER/SIM A/B.
+
+The compile layer (``compile_shapes``) additionally applies
+``torch.compile`` to the BigVGAN forward before CUDA-graph capture,
+letting Inductor fuse the Snake/alias-free small-op chain. The
+three-tier dispatch at runtime is: compiled-graph > plain-graph > eager.
+
+Modeled on ``CUDAGraphDecoderWrapper``
+(vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py)
+and ``CUDAGraphGLMTTSDiTWrapper``
+(vllm_omni/model_executor/models/glm_tts/glm_tts_dit_wrapper.py).
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Callable
+
+import torch
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+logger = init_logger(__name__)
+
+# Mel-frame buckets at 22050 Hz / hop 256 (~86 fps).
+# Keep buckets modest so that CUDA-graph private pools do not exhaust
+# GPU memory when stage0 and stage1 share the same device (~22 GiB).
+# Inputs longer than the largest bucket fall back to eager BigVGAN.
+DEFAULT_CAPTURE_SIZES = [128, 256, 384, 512]
+
+
+class CUDAGraphBigVGANWrapper:
+    """Capture BigVGAN vocoding for fixed mel-frame buckets.
+
+    Supports an optional torch.compile layer on top of plain CUDA graphs.
+    When ``compile_shapes`` is provided, those buckets are compiled with
+    ``torch.compile(mode="default", fullgraph=False, dynamic=False)`` and
+    then captured into CUDA graphs — giving both Inductor operator fusion
+    and zero-launch-overhead replay.
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        *,
+        enabled: bool = True,
+        capture_sizes: list[int] | None = None,
+        compile_shapes: list[int] | None = None,
+    ) -> None:
+        self.model = model
+        self.enabled = bool(enabled)
+        self.capture_sizes = sorted(set(capture_sizes or DEFAULT_CAPTURE_SIZES))
+        self.compile_shapes = sorted(set(compile_shapes or []))
+        self.num_mels = int(model.h.num_mels)
+        self.upsample_factor = int(math.prod(model.h.upsample_rates))
+
+        # Plain CUDA graph tier
+        self.graphs: dict[int, torch.cuda.CUDAGraph] = {}
+        self.static_mel: dict[int, torch.Tensor] = {}
+        self.static_wav: dict[int, torch.Tensor] = {}
+
+        # Compiled CUDA graph tier
+        self._compiled_decoder: Callable[[torch.Tensor], torch.Tensor] | None = None
+        self._compiled_graphs: dict[int, torch.cuda.CUDAGraph] = {}
+        self._compiled_static_mel: dict[int, torch.Tensor] = {}
+        self._compiled_static_wav: dict[int, torch.Tensor] = {}
+        self._compiled_shapes: set[int] = set()
+
+        self._warmed_up = False
+        self.last_call_info: dict[str, int | str | None] = {}
+
+    def warmup(self, device: torch.device, dtype: torch.dtype) -> None:
+        if not self.enabled or self._warmed_up or device.type != "cuda" or torch.cuda.is_current_stream_capturing():
+            return
+        self.model.eval()
+
+        # Precompute Snake/SnakeBeta caches so torch.compile sees stable
+        # buffers instead of lazy init inside the traced graph.
+        self._precompute_snake_caches()
+
+        logger.info("Starting BigVGAN CUDA graph warmup for mel-frame buckets: %s", self.capture_sizes)
+        for size in self.capture_sizes:
+            try:
+                self._capture(size, device, dtype)
+                logger.info("Captured BigVGAN CUDA graph for mel frames=%d", size)
+            except Exception:
+                logger.warning("Failed to capture BigVGAN CUDA graph for size=%d", size, exc_info=True)
+
+        if self.compile_shapes:
+            self._warmup_compile_shapes(device, dtype)
+
+        self._warmed_up = True
+
+    def _precompute_snake_caches(self) -> None:
+        from vllm_omni.model_executor.models.common.snake_activation import Snake, SnakeBeta
+
+        for module in self.model.modules():
+            if isinstance(module, (Snake, SnakeBeta)):
+                module.precompute_exp_cache()
+
+    # ------------------------------------------------------------------
+    # Plain CUDA graph capture
+    # ------------------------------------------------------------------
+
+    def _capture(self, size: int, device: torch.device, dtype: torch.dtype) -> None:
+        mel = torch.zeros(1, self.num_mels, size, device=device, dtype=dtype)
+        self.static_mel[size] = mel
+        # Eager warm run so cuDNN autotuning happens outside capture.
+        with torch.no_grad():
+            _ = self.model(mel)
+        torch.accelerator.synchronize(device)
+
+        graph = torch.cuda.CUDAGraph()
+        with torch.no_grad(), torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()):
+            self.static_wav[size] = self.model(mel)
+        self.graphs[size] = graph
+
+    # ------------------------------------------------------------------
+    # Compiled CUDA graph capture (torch.compile + capture)
+    # ------------------------------------------------------------------
+
+    def _warmup_compile_shapes(self, device: torch.device, dtype: torch.dtype) -> None:
+        logger.info(
+            "Starting torch.compile + CUDA Graph warmup for BigVGAN shapes: %s",
+            self.compile_shapes,
+        )
+        try:
+            self._compiled_decoder = torch.compile(
+                self.model.forward,
+                mode="default",
+                fullgraph=False,
+                dynamic=False,
+            )
+        except Exception:
+            logger.warning("Failed to create torch.compile BigVGAN wrapper", exc_info=True)
+            self._compiled_decoder = None
+            return
+
+        for size in self.compile_shapes:
+            try:
+                self._capture_compiled(size, device, dtype)
+                self._compiled_shapes.add(size)
+                logger.info(
+                    "  torch.compile + CUDA Graph ready for mel frames=%d",
+                    size,
+                )
+            except Exception:
+                logger.warning(
+                    "  Failed to capture torch.compile CUDA Graph for size=%d; "
+                    "this shape will use plain CUDA graph/eager fallback",
+                    size,
+                    exc_info=True,
+                )
+        logger.info(
+            "torch.compile + CUDA Graph warmup complete: %d/%d shapes ready",
+            len(self._compiled_shapes),
+            len(self.compile_shapes),
+        )
+
+    def _capture_compiled(self, size: int, device: torch.device, dtype: torch.dtype) -> None:
+        if self._compiled_decoder is None:
+            raise RuntimeError("Compiled decoder is not initialized")
+
+        mel = torch.zeros(1, self.num_mels, size, device=device, dtype=dtype)
+        # Warmup iterations let Inductor finish tracing/autotuning before capture.
+        with torch.inference_mode():
+            for _ in range(5):
+                _ = self._compiled_decoder(mel)
+        torch.accelerator.synchronize(device)
+
+        graph = torch.cuda.CUDAGraph()
+        with torch.inference_mode():
+            with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()):
+                static_wav = self._compiled_decoder(mel)
+
+        self._compiled_static_mel[size] = mel
+        self._compiled_static_wav[size] = static_wav
+        self._compiled_graphs[size] = graph
+
+    # ------------------------------------------------------------------
+    # Dispatch
+    # ------------------------------------------------------------------
+
+    def _get_padded_size(self, actual_frames: int) -> int | None:
+        for size in self.capture_sizes:
+            if actual_frames <= size:
+                return size
+        return None
+
+    def __call__(self, mel: torch.Tensor) -> torch.Tensor:
+        """Vocode ``mel`` [1, num_mels, T]; falls back to eager when ineligible."""
+        actual_frames = int(mel.shape[-1]) if mel.ndim >= 3 else 0
+
+        # Eligibility guard
+        if (
+            not self.enabled
+            or not self._warmed_up
+            or mel.shape[0] != 1
+            or mel.device.type != "cuda"
+            or torch.cuda.is_current_stream_capturing()
+        ):
+            self.last_call_info = {
+                "mode": "eager",
+                "reason": "ineligible",
+                "batch": int(mel.shape[0]) if mel.ndim > 0 else 0,
+                "frames": actual_frames,
+                "bucket": None,
+                "padding": None,
+            }
+            return self.model(mel)
+
+        # Tier 1: compiled CUDA graph (exact shape match)
+        if actual_frames in self._compiled_shapes:
+            return self._replay_compiled(actual_frames, actual_frames, mel, "exact_hit")
+
+        # Tier 1b: compiled CUDA graph (padded shape)
+        padded = self._get_padded_size(actual_frames)
+        if padded is not None and padded in self._compiled_shapes:
+            return self._replay_compiled(actual_frames, padded, mel, "padded_hit")
+
+        # Tier 2: plain CUDA graph
+        size = padded
+        if size is None or size not in self.graphs:
+            self.last_call_info = {
+                "mode": "eager",
+                "reason": "no_bucket",
+                "batch": int(mel.shape[0]),
+                "frames": actual_frames,
+                "bucket": size,
+                "padding": None if size is None else int(size - actual_frames),
+            }
+            return self.model(mel)
+
+        buf = self.static_mel[size]
+        if actual_frames == size:
+            buf.copy_(mel.to(buf.dtype))
+        else:
+            buf.zero_()
+            buf[..., :actual_frames].copy_(mel.to(buf.dtype))
+        self.graphs[size].replay()
+        self.last_call_info = {
+            "mode": "graph",
+            "reason": "hit",
+            "batch": int(mel.shape[0]),
+            "frames": actual_frames,
+            "bucket": int(size),
+            "padding": int(size - actual_frames),
+        }
+        return self.static_wav[size][..., : actual_frames * self.upsample_factor].clone()
+
+    def _replay_compiled(
+        self,
+        actual_frames: int,
+        bucket: int,
+        mel: torch.Tensor,
+        reason: str,
+    ) -> torch.Tensor:
+        buf = self._compiled_static_mel[bucket]
+        if actual_frames == bucket:
+            buf.copy_(mel.to(buf.dtype))
+        else:
+            buf.zero_()
+            buf[..., :actual_frames].copy_(mel.to(buf.dtype))
+        self._compiled_graphs[bucket].replay()
+        self.last_call_info = {
+            "mode": "compile_graph",
+            "reason": reason,
+            "batch": int(mel.shape[0]),
+            "frames": actual_frames,
+            "bucket": int(bucket),
+            "padding": int(bucket - actual_frames),
+        }
+        return self._compiled_static_wav[bucket][..., : actual_frames * self.upsample_factor].clone()

vllm_omni/model_executor/models/indextts2/configuration_indextts2.py

@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import copy
+
+from transformers import PretrainedConfig
+
+
+class IndexTTS2Config(PretrainedConfig):
+    model_type = "indextts2"
+
+    # Defaults for the official IndexTTS2 checkpoint.  Values from HF
+    # config.json (passed via ``kwargs``) take precedence over these.
+    _DEFAULTS = {
+        "num_attention_heads": 20,
+        "hidden_size": 1280,
+        "dataset": {
+            "bpe_model": "bpe.model",
+            "sample_rate": 24000,
+            "squeeze": False,
+            "mel": {
+                "sample_rate": 24000,
+                "n_fft": 1024,
+                "hop_length": 256,
+                "win_length": 1024,
+                "n_mels": 100,
+                "mel_fmin": 0,
+                "normalize": False,
+            },
+        },
+        "gpt": {
+            "model_dim": 1280,
+            "max_mel_tokens": 1815,
+            "max_text_tokens": 600,
+            "heads": 20,
+            "use_mel_codes_as_input": True,
+            "mel_length_compression": 1024,
+            "layers": 24,
+            "number_text_tokens": 12000,
+            "number_mel_codes": 8194,
+            "start_mel_token": 8192,
+            "stop_mel_token": 8193,
+            "start_text_token": 0,
+            "stop_text_token": 1,
+            "train_solo_embeddings": False,
+            "condition_num_latent": 32,
+            "condition_type": "conformer_perceiver",
+            "condition_module": {
+                "output_size": 512,
+                "linear_units": 2048,
+                "attention_heads": 8,
+                "num_blocks": 6,
+                "input_layer": "conv2d2",
+                "perceiver_mult": 2,
+            },
+            "emo_condition_module": {
+                "output_size": 512,
+                "linear_units": 1024,
+                "attention_heads": 4,
+                "num_blocks": 4,
+                "input_layer": "conv2d2",
+                "perceiver_mult": 2,
+            },
+        },
+        "semantic_codec": {
+            "codebook_size": 8192,
+            "hidden_size": 1024,
+            "codebook_dim": 8,
+            "vocos_dim": 384,
+            "vocos_intermediate_dim": 2048,
+            "vocos_num_layers": 12,
+        },
+        "s2mel": {
+            "preprocess_params": {
+                "sr": 22050,
+                "spect_params": {
+                    "n_fft": 1024,
+                    "win_length": 1024,
+                    "hop_length": 256,
+                    "n_mels": 80,
+                    "fmin": 0,
+                    "fmax": "None",
+                },
+            },
+            "dit_type": "DiT",
+            "reg_loss_type": "l1",
+            "style_encoder": {"dim": 192},
+            "length_regulator": {
+                "channels": 512,
+                "is_discrete": False,
+                "in_channels": 1024,
+                "content_codebook_size": 2048,
+                "sampling_ratios": [1, 1, 1, 1],
+                "vector_quantize": False,
+                "n_codebooks": 1,
+                "quantizer_dropout": 0.0,
+                "f0_condition": False,
+                "n_f0_bins": 512,
+            },
+            "DiT": {
+                "hidden_dim": 512,
+                "num_heads": 8,
+                "depth": 13,
+                "class_dropout_prob": 0.1,
+                "block_size": 8192,
+                "in_channels": 80,
+                "style_condition": True,
+                "final_layer_type": "wavenet",
+                "target": "mel",
+                "content_dim": 512,
+                "content_codebook_size": 1024,
+                "content_type": "discrete",
+                "f0_condition": False,
+                "n_f0_bins": 512,
+                "content_codebooks": 1,
+                "is_causal": False,
+                "long_skip_connection": True,
+                "zero_prompt_speech_token": False,
+                "time_as_token": False,
+                "style_as_token": False,
+                "uvit_skip_connection": True,
+                "add_resblock_in_transformer": False,
+            },
+            "wavenet": {
+                "hidden_dim": 512,
+                "num_layers": 8,
+                "kernel_size": 5,
+                "dilation_rate": 1,
+                "p_dropout": 0.2,
+                "style_condition": True,
+            },
+        },
+        "gpt_checkpoint": "gpt.pth",
+        "w2v_stat": "wav2vec2bert_stats.pt",
+        "s2mel_checkpoint": "s2mel.pth",
+        "emo_matrix": "feat2.pt",
+        "spk_matrix": "feat1.pt",
+        "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"},
+        "version": 2.0,
+        # Default-off performance tracing for IndexTTS2 overhead analysis.
+    }
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        defaults = copy.deepcopy(self._DEFAULTS)
+        for attr, default in defaults.items():
+            if not hasattr(self, attr):
+                setattr(self, attr, default)
+        # 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
+        )

vllm_omni/model_executor/models/indextts2/dit_cuda_graph.py

@@ -0,0 +1,164 @@
+"""Per-shape CUDA graph capture for the S2Mel DiT transformer core.
+
+Profiling showed the DiT estimator forward is CPU launch-bound (the GPU
+is idle ~85% of CFM wall time). Within one request the CFM runs 12 Euler
+steps with *identical* tensor shapes (T_in never changes), so the graph
+captured on the first step is replayed by every subsequent step. Across
+requests T_in varies, so graphs are cached per shape with a small LRU;
+a cache miss costs roughly one extra eager step (warm run + capture).
+
+Capture safety relies on the IndexTTS2 S2Mel decoder only enabling this
+runner for full-mask shapes:
+
+- ``x_lens == T`` for every sample (the decoder feeds full-length sequences), so
+  the attention padding mask is all-True and dense flash attention is
+  mathematically identical. The decoder sets ``_assume_full_mask`` on
+  the gpt_fast ``Attention`` modules so the backend skips its
+  ``torch.any(~mask)`` D2H sync, which would abort stream capture.
+
+Modeled on ``CUDAGraphGLMTTSDiTWrapper``
+(vllm_omni/model_executor/models/glm_tts/glm_tts_dit_wrapper.py).
+"""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from dataclasses import dataclass
+
+import torch
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class _GraphEntry:
+    graph: torch.cuda.CUDAGraph
+    static_x: torch.Tensor
+    static_c: torch.Tensor
+    static_pos: torch.Tensor
+    static_mask: torch.Tensor | None
+    static_out: torch.Tensor
+
+
+class CUDAGraphDiTRunner:
+    """Replay the gpt_fast Transformer forward via per-shape CUDA graphs."""
+
+    def __init__(self, transformer: torch.nn.Module, *, max_graphs: int = 4, enabled: bool = True) -> None:
+        self.transformer = transformer
+        self.max_graphs = int(max_graphs)
+        self.enabled = bool(enabled)
+        self._cache: OrderedDict[tuple, _GraphEntry] = OrderedDict()
+        self.last_call_info: dict[str, object] = {}
+
+    def __call__(
+        self,
+        x_in: torch.Tensor,
+        c: torch.Tensor,
+        input_pos: torch.Tensor,
+        mask: torch.Tensor | None,
+    ) -> torch.Tensor:
+        if not self.enabled or x_in.device.type != "cuda" or torch.cuda.is_current_stream_capturing():
+            self.last_call_info = {
+                "mode": "eager",
+                "reason": "ineligible",
+                "shape": tuple(x_in.shape),
+                "cache_size": len(self._cache),
+            }
+            return self.transformer(x_in, c, input_pos, mask)
+
+        mask_key = None if mask is None else (tuple(mask.shape), mask.dtype)
+        device_key = (x_in.device.type, x_in.device.index)
+        key = (
+            device_key,
+            tuple(x_in.shape),
+            x_in.dtype,
+            tuple(c.shape),
+            c.dtype,
+            tuple(input_pos.shape),
+            input_pos.dtype,
+            mask_key,
+        )
+        entry = self._cache.get(key)
+        if entry is None:
+            entry = self._capture(x_in, c, input_pos, mask)
+            if entry is None:
+                self.last_call_info = {
+                    "mode": "eager",
+                    "reason": "capture_failed",
+                    "shape": tuple(x_in.shape),
+                    "cache_size": len(self._cache),
+                }
+                return self.transformer(x_in, c, input_pos, mask)
+            self._cache[key] = entry
+            while len(self._cache) > self.max_graphs:
+                _, evicted = self._cache.popitem(last=False)
+                del evicted
+            # Kernels are only recorded (not executed) during capture, so
+            # replay once to produce the first real output.
+            entry.graph.replay()
+            self.last_call_info = {
+                "mode": "graph",
+                "reason": "capture",
+                "shape": tuple(x_in.shape),
+                "cache_size": len(self._cache),
+            }
+            return entry.static_out.clone()
+
+        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.last_call_info = {
+            "mode": "graph",
+            "reason": "hit",
+            "shape": tuple(x_in.shape),
+            "cache_size": len(self._cache),
+        }
+        return entry.static_out.clone()
+
+    def _capture(
+        self,
+        x_in: torch.Tensor,
+        c: torch.Tensor,
+        input_pos: torch.Tensor,
+        mask: torch.Tensor | None,
+    ) -> _GraphEntry | None:
+        static_x = x_in.clone()
+        static_c = c.clone()
+        static_pos = input_pos.clone()
+        # The mask content is constant (all-True; see module docstring) and
+        # ignored by Attention via `_assume_full_mask`, but keep a static
+        # clone so the captured graph never dereferences a freed tensor.
+        static_mask = mask.contiguous().clone() if mask is not None else None
+        try:
+            # Eager warm run: cuDNN/flash-attn autotuning and lazy inits
+            # must happen outside stream capture.
+            with torch.no_grad():
+                _ = self.transformer(static_x, static_c, static_pos, static_mask)
+            torch.accelerator.synchronize(x_in.device)
+
+            graph = torch.cuda.CUDAGraph()
+            with torch.no_grad(), torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()):
+                static_out = self.transformer(static_x, static_c, static_pos, static_mask)
+        except Exception:
+            logger.warning(
+                "Disabling S2Mel DiT CUDA graphs after capture failure for shape=%s",
+                tuple(x_in.shape),
+                exc_info=True,
+            )
+            self.enabled = False
+            return None
+        logger.debug("Captured S2Mel DiT CUDA graph for shape=%s", tuple(x_in.shape))
+        return _GraphEntry(
+            graph=graph,
+            static_x=static_x,
+            static_c=static_c,
+            static_pos=static_pos,
+            static_mask=static_mask,
+            static_out=static_out,
+        )

vllm_omni/model_executor/models/indextts2/gpt/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/model_executor/models/indextts2/gpt/conformer/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/model_executor/models/indextts2/gpt/conformer/attention.py

@@ -0,0 +1,288 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright (c) 2019 Shigeki Karita
+#               2020 Mobvoi Inc (Binbin Zhang)
+#               2022 Xingchen Song (sxc19@mails.tsinghua.edu.cn)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Multi-Head Attention layer definition."""
+
+import math
+
+import torch
+from torch import nn
+
+
+class MultiHeadedAttention(nn.Module):
+    """Multi-Head Attention layer.
+
+    Args:
+        n_head (int): The number of heads.
+        n_feat (int): The number of features.
+        dropout_rate (float): Dropout rate.
+
+    """
+
+    def __init__(self, n_head: int, n_feat: int, dropout_rate: float):
+        """Construct an MultiHeadedAttention object."""
+        super().__init__()
+        assert n_feat % n_head == 0
+        # We assume d_v always equals d_k
+        self.d_k = n_feat // n_head
+        self.h = n_head
+        self.linear_q = nn.Linear(n_feat, n_feat)
+        self.linear_k = nn.Linear(n_feat, n_feat)
+        self.linear_v = nn.Linear(n_feat, n_feat)
+        self.linear_out = nn.Linear(n_feat, n_feat)
+        self.dropout = nn.Dropout(p=dropout_rate)
+
+    def forward_qkv(
+        self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Transform query, key and value.
+
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+
+        Returns:
+            torch.Tensor: Transformed query tensor, size
+                (#batch, n_head, time1, d_k).
+            torch.Tensor: Transformed key tensor, size
+                (#batch, n_head, time2, d_k).
+            torch.Tensor: Transformed value tensor, size
+                (#batch, n_head, time2, d_k).
+
+        """
+        n_batch = query.size(0)
+        q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k)
+        k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k)
+        v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k)
+        q = q.transpose(1, 2)  # (batch, head, time1, d_k)
+        k = k.transpose(1, 2)  # (batch, head, time2, d_k)
+        v = v.transpose(1, 2)  # (batch, head, time2, d_k)
+
+        return q, k, v
+
+    def forward_attention(
+        self, value: torch.Tensor, scores: torch.Tensor, mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool)
+    ) -> torch.Tensor:
+        """Compute attention context vector.
+
+        Args:
+            value (torch.Tensor): Transformed value, size
+                (#batch, n_head, time2, d_k).
+            scores (torch.Tensor): Attention score, size
+                (#batch, n_head, time1, time2).
+            mask (torch.Tensor): Mask, size (#batch, 1, time2) or
+                (#batch, time1, time2), (0, 0, 0) means fake mask.
+
+        Returns:
+            torch.Tensor: Transformed value (#batch, time1, d_model)
+                weighted by the attention score (#batch, time1, time2).
+
+        """
+        n_batch = value.size(0)
+        # NOTE(xcsong): When will `if mask.size(2) > 0` be True?
+        #   1. onnx(16/4) [WHY? Because we feed real cache & real mask for the
+        #           1st chunk to ease the onnx export.]
+        #   2. pytorch training
+        if mask.size(2) > 0:  # time2 > 0
+            mask = mask.unsqueeze(1).eq(0)  # (batch, 1, *, time2)
+            # For last chunk, time2 might be larger than scores.size(-1)
+            mask = mask[:, :, :, : scores.size(-1)]  # (batch, 1, *, time2)
+            scores = scores.masked_fill(mask, -float("inf"))
+            attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0)  # (batch, head, time1, time2)
+        # NOTE(xcsong): When will `if mask.size(2) > 0` be False?
+        #   1. onnx(16/-1, -1/-1, 16/0)
+        #   2. jit (16/-1, -1/-1, 16/0, 16/4)
+        else:
+            attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)
+
+        p_attn = self.dropout(attn)
+        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
+        x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)  # (batch, time1, d_model)
+
+        return self.linear_out(x)  # (batch, time1, d_model)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
+        pos_emb: torch.Tensor = torch.empty(0),
+        cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Compute scaled dot product attention.
+
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+            mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
+                (#batch, time1, time2).
+                1.When applying cross attention between decoder and encoder,
+                the batch padding mask for input is in (#batch, 1, T) shape.
+                2.When applying self attention of encoder,
+                the mask is in (#batch, T, T)  shape.
+                3.When applying self attention of decoder,
+                the mask is in (#batch, L, L)  shape.
+                4.If the different position in decoder see different block
+                of the encoder, such as Mocha, the passed in mask could be
+                in (#batch, L, T) shape. But there is no such case in current
+                Wenet.
+            cache (torch.Tensor): Cache tensor (1, head, cache_t, d_k * 2),
+                where `cache_t == chunk_size * num_decoding_left_chunks`
+                and `head * d_k == size`
+
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time1, d_model).
+            torch.Tensor: Cache tensor (1, head, cache_t + time1, d_k * 2)
+                where `cache_t == chunk_size * num_decoding_left_chunks`
+                and `head * d_k == size`
+
+        """
+        q, k, v = self.forward_qkv(query, key, value)
+
+        if cache.size(0) > 0:
+            key_cache, value_cache = torch.split(cache, cache.size(-1) // 2, dim=-1)
+            k = torch.cat([key_cache, k], dim=2)
+            v = torch.cat([value_cache, v], dim=2)
+        new_cache = torch.cat((k, v), dim=-1)
+
+        # Build boolean attention mask for SDPA.
+        # Input mask convention: True = attend, False = ignore.
+        # SDPA boolean mask convention: True = attend, False = masked to -inf.
+        if mask.size(2) > 0:
+            attn_mask = mask.unsqueeze(1)  # (batch, 1, *, time2)
+            attn_mask = attn_mask[:, :, :, : k.size(2)]
+        else:
+            attn_mask = None
+
+        # q/k/v: (batch, head, time, d_k)
+        x = torch.nn.functional.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=attn_mask,
+            dropout_p=self.dropout.p if self.training else 0.0,
+            scale=1.0 / math.sqrt(self.d_k),
+        )
+        n_batch = query.size(0)
+        x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)
+
+        return self.linear_out(x), new_cache
+
+
+class RelPositionMultiHeadedAttention(MultiHeadedAttention):
+    """Multi-Head Attention layer with relative position encoding.
+    Paper: https://arxiv.org/abs/1901.02860
+    Args:
+        n_head (int): The number of heads.
+        n_feat (int): The number of features.
+        dropout_rate (float): Dropout rate.
+    """
+
+    def __init__(self, n_head, n_feat, dropout_rate):
+        """Construct an RelPositionMultiHeadedAttention object."""
+        super().__init__(n_head, n_feat, dropout_rate)
+        # linear transformation for positional encoding
+        self.linear_pos = nn.Linear(n_feat, n_feat, bias=False)
+        # these two learnable bias are used in matrix c and matrix d
+        # as described in https://arxiv.org/abs/1901.02860 Section 3.3
+        self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k))
+        self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k))
+        torch.nn.init.xavier_uniform_(self.pos_bias_u)
+        torch.nn.init.xavier_uniform_(self.pos_bias_v)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
+        pos_emb: torch.Tensor = torch.empty(0),
+        cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Compute 'Scaled Dot Product Attention' with rel. positional encoding.
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+            mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
+                (#batch, time1, time2), (0, 0, 0) means fake mask.
+            pos_emb (torch.Tensor): Positional embedding tensor
+                (#batch, time2, size).
+            cache (torch.Tensor): Cache tensor (1, head, cache_t, d_k * 2),
+                where `cache_t == chunk_size * num_decoding_left_chunks`
+                and `head * d_k == size`
+        Returns:
+            torch.Tensor: Output tensor (#batch, time1, d_model).
+            torch.Tensor: Cache tensor (1, head, cache_t + time1, d_k * 2)
+                where `cache_t == chunk_size * num_decoding_left_chunks`
+                and `head * d_k == size`
+        """
+        q, k, v = self.forward_qkv(query, key, value)
+        q = q.transpose(1, 2)  # (batch, time1, head, d_k)
+
+        # NOTE(xcsong):
+        #   when export onnx model, for 1st chunk, we feed
+        #       cache(1, head, 0, d_k * 2) (16/-1, -1/-1, 16/0 mode)
+        #       or cache(1, head, real_cache_t, d_k * 2) (16/4 mode).
+        #       In all modes, `if cache.size(0) > 0` will always be `True`
+        #       and we will always do splitting and
+        #       concatnation(this will simplify onnx export). Note that
+        #       it's OK to concat & split zero-shaped tensors(see code below).
+        #   when export jit  model, for 1st chunk, we always feed
+        #       cache(0, 0, 0, 0) since jit supports dynamic if-branch.
+        # >>> a = torch.ones((1, 2, 0, 4))
+        # >>> b = torch.ones((1, 2, 3, 4))
+        # >>> c = torch.cat((a, b), dim=2)
+        # >>> torch.equal(b, c)        # True
+        # >>> d = torch.split(a, 2, dim=-1)
+        # >>> torch.equal(d[0], d[1])  # True
+        if cache.size(0) > 0:
+            key_cache, value_cache = torch.split(cache, cache.size(-1) // 2, dim=-1)
+            k = torch.cat([key_cache, k], dim=2)
+            v = torch.cat([value_cache, v], dim=2)
+        # NOTE(xcsong): We do cache slicing in encoder.forward_chunk, since it's
+        #   non-trivial to calculate `next_cache_start` here.
+        new_cache = torch.cat((k, v), dim=-1)
+
+        n_batch_pos = pos_emb.size(0)
+        p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k)
+        p = p.transpose(1, 2)  # (batch, head, time1, d_k)
+
+        # (batch, head, time1, d_k)
+        q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2)
+        # (batch, head, time1, d_k)
+        q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2)
+
+        # compute attention score
+        # first compute matrix a and matrix c
+        # as described in https://arxiv.org/abs/1901.02860 Section 3.3
+        # (batch, head, time1, time2)
+        matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
+
+        # compute matrix b and matrix d
+        # (batch, head, time1, time2)
+        matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
+
+        scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k)  # (batch, head, time1, time2)
+
+        return self.forward_attention(v, scores, mask), new_cache

vllm_omni/model_executor/models/indextts2/gpt/conformer/embedding.py

@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright (c) 2020 Mobvoi Inc. (authors: Binbin Zhang, Di Wu)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Modified from ESPnet(https://github.com/espnet/espnet)
+
+"""Positional Encoding Module."""
+
+import math
+
+import torch
+import torch.nn.functional as F
+
+
+class PositionalEncoding(torch.nn.Module):
+    """Positional encoding.
+
+    :param int d_model: embedding dim
+    :param float dropout_rate: dropout rate
+    :param int max_len: maximum input length
+
+    PE(pos, 2i)   = sin(pos/(10000^(2i/dmodel)))
+    PE(pos, 2i+1) = cos(pos/(10000^(2i/dmodel)))
+    """
+
+    def __init__(self, d_model: int, dropout_rate: float, max_len: int = 5000, reverse: bool = False):
+        """Construct an PositionalEncoding object."""
+        super().__init__()
+        self.d_model = d_model
+        self.xscale = math.sqrt(self.d_model)
+        self.dropout = torch.nn.Dropout(p=dropout_rate)
+        self.max_len = max_len
+
+        pe = torch.zeros(self.max_len, self.d_model)
+        position = torch.arange(0, self.max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, self.d_model, 2) * -(math.log(10000.0) / self.d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x: torch.Tensor, offset: int | torch.Tensor = 0) -> tuple[torch.Tensor, torch.Tensor]:
+        """Add positional encoding.
+
+        Args:
+            x (torch.Tensor): Input. Its shape is (batch, time, ...)
+            offset (int, torch.tensor): position offset
+
+        Returns:
+            torch.Tensor: Encoded tensor. Its shape is (batch, time, ...)
+            torch.Tensor: for compatibility to RelPositionalEncoding
+        """
+
+        self.pe = self.pe.to(x.device)
+        pos_emb = self.position_encoding(offset, x.size(1), False)
+        x = x * self.xscale + pos_emb
+        return self.dropout(x), self.dropout(pos_emb)
+
+    def position_encoding(self, offset: int | torch.Tensor, size: int, apply_dropout: bool = True) -> torch.Tensor:
+        """For getting encoding in a streaming fashion
+
+        Attention!!!!!
+        we apply dropout only once at the whole utterance level in a none
+        streaming way, but will call this function several times with
+        increasing input size in a streaming scenario, so the dropout will
+        be applied several times.
+
+        Args:
+            offset (int or torch.tensor): start offset
+            size (int): required size of position encoding
+
+        Returns:
+            torch.Tensor: Corresponding encoding
+        """
+        # How to subscript a Union type:
+        #   https://github.com/pytorch/pytorch/issues/69434
+        if isinstance(offset, int):
+            assert offset + size < self.max_len
+            pos_emb = self.pe[:, offset : offset + size]
+        elif isinstance(offset, torch.Tensor) and offset.dim() == 0:  # scalar
+            assert offset + size < self.max_len
+            pos_emb = self.pe[:, offset : offset + size]
+        else:  # for batched streaming decoding on GPU
+            assert torch.max(offset) + size < self.max_len
+            index = offset.unsqueeze(1) + torch.arange(0, size).to(offset.device)  # B X T
+            flag = index > 0
+            # remove negative offset
+            index = index * flag
+            pos_emb = F.embedding(index, self.pe[0])  # B X T X d_model
+
+        if apply_dropout:
+            pos_emb = self.dropout(pos_emb)
+        return pos_emb
+
+
+class RelPositionalEncoding(PositionalEncoding):
+    """Relative positional encoding module.
+    See : Appendix B in https://arxiv.org/abs/1901.02860
+    Args:
+        d_model (int): Embedding dimension.
+        dropout_rate (float): Dropout rate.
+        max_len (int): Maximum input length.
+    """
+
+    def __init__(self, d_model: int, dropout_rate: float, max_len: int = 5000):
+        """Initialize class."""
+        super().__init__(d_model, dropout_rate, max_len, reverse=True)
+
+    def forward(self, x: torch.Tensor, offset: int | torch.Tensor = 0) -> tuple[torch.Tensor, torch.Tensor]:
+        """Compute positional encoding.
+        Args:
+            x (torch.Tensor): Input tensor (batch, time, `*`).
+        Returns:
+            torch.Tensor: Encoded tensor (batch, time, `*`).
+            torch.Tensor: Positional embedding tensor (1, time, `*`).
+        """
+        self.pe = self.pe.to(x.device)
+        x = x * self.xscale
+        pos_emb = self.position_encoding(offset, x.size(1), False)
+        return self.dropout(x), self.dropout(pos_emb)

vllm_omni/model_executor/models/indextts2/gpt/conformer/subsampling.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright (c) 2021 Mobvoi Inc (Binbin Zhang, Di Wu)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Modified from ESPnet(https://github.com/espnet/espnet)
+
+
+"""Subsampling layer definition."""
+
+import torch
+
+
+class BaseSubsampling(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.right_context = 0
+        self.subsampling_rate = 1
+
+    def position_encoding(self, offset: int | torch.Tensor, size: int) -> torch.Tensor:
+        return self.pos_enc.position_encoding(offset, size)
+
+
+class Conv2dSubsampling2(BaseSubsampling):
+    """Convolutional 2D subsampling (to 1/2 length).
+
+    Args:
+        idim (int): Input dimension.
+        odim (int): Output dimension.
+        dropout_rate (float): Dropout rate.
+
+    """
+
+    def __init__(self, idim: int, odim: int, dropout_rate: float, pos_enc_class: torch.nn.Module):
+        super().__init__()
+        self.conv = torch.nn.Sequential(
+            torch.nn.Conv2d(1, odim, 3, 2),
+            torch.nn.ReLU(),
+        )
+        self.out = torch.nn.Sequential(torch.nn.Linear(odim * ((idim - 1) // 2), odim))
+        self.pos_enc = pos_enc_class
+        self.subsampling_rate = 2
+        self.right_context = 2
+
+    def forward(
+        self, x: torch.Tensor, x_mask: torch.Tensor, offset: int | torch.Tensor = 0
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        x = x.unsqueeze(1)  # (b, c=1, t, f)
+        x = self.conv(x)
+        b, c, t, f = x.size()
+        x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f))
+        x, pos_emb = self.pos_enc(x, offset)
+        return x, pos_emb, x_mask[:, :, 2::2]

vllm_omni/model_executor/models/indextts2/gpt/conformer_encoder.py

@@ -0,0 +1,422 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+import torch.nn as nn
+
+from vllm_omni.model_executor.models.indextts2.gpt.conformer.attention import (
+    MultiHeadedAttention,
+    RelPositionMultiHeadedAttention,
+)
+from vllm_omni.model_executor.models.indextts2.gpt.conformer.embedding import (
+    PositionalEncoding,
+    RelPositionalEncoding,
+)
+from vllm_omni.model_executor.models.indextts2.gpt.conformer.subsampling import (
+    Conv2dSubsampling2,
+)
+from vllm_omni.model_executor.models.indextts2.utils.common import make_pad_mask
+
+
+class PositionwiseFeedForward(torch.nn.Module):
+    """Positionwise feed forward layer.
+
+    FeedForward are applied on each position of the sequence.
+    The output dim is same with the input dim.
+
+    Args:
+        idim (int): Input dimension.
+        hidden_units (int): The number of hidden units.
+        dropout_rate (float): Dropout rate.
+        activation (torch.nn.Module): Activation function
+    """
+
+    def __init__(
+        self, idim: int, hidden_units: int, dropout_rate: float, activation: torch.nn.Module = torch.nn.ReLU()
+    ):
+        """Construct a PositionwiseFeedForward object."""
+        super().__init__()
+        self.w_1 = torch.nn.Linear(idim, hidden_units)
+        self.activation = activation
+        self.dropout = torch.nn.Dropout(dropout_rate)
+        self.w_2 = torch.nn.Linear(hidden_units, idim)
+
+    def forward(self, xs: torch.Tensor) -> torch.Tensor:
+        """Forward function.
+
+        Args:
+            xs: input tensor (B, L, D)
+        Returns:
+            output tensor, (B, L, D)
+        """
+        return self.w_2(self.dropout(self.activation(self.w_1(xs))))
+
+
+class ConvolutionModule(nn.Module):
+    """ConvolutionModule in Conformer model."""
+
+    def __init__(self, channels: int, kernel_size: int = 15, activation: nn.Module = nn.ReLU(), bias: bool = True):
+        """Construct an ConvolutionModule object.
+        Args:
+            channels (int): The number of channels of conv layers.
+            kernel_size (int): Kernel size of conv layers.
+            causal (int): Whether use causal convolution or not
+        """
+        super().__init__()
+
+        self.pointwise_conv1 = nn.Conv1d(
+            channels,
+            2 * channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=bias,
+        )
+        # self.lorder is used to distinguish if it's a causal convolution,
+        # if self.lorder > 0: it's a causal convolution, the input will be
+        #    padded with self.lorder frames on the left in forward.
+        # else: it's a symmetrical convolution
+        # kernel_size should be an odd number for none causal convolution
+        assert (kernel_size - 1) % 2 == 0
+        padding = (kernel_size - 1) // 2
+        self.lorder = 0
+
+        self.depthwise_conv = nn.Conv1d(
+            channels,
+            channels,
+            kernel_size,
+            stride=1,
+            padding=padding,
+            groups=channels,
+            bias=bias,
+        )
+
+        self.use_layer_norm = True
+        self.norm = nn.LayerNorm(channels)
+
+        self.pointwise_conv2 = nn.Conv1d(
+            channels,
+            channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=bias,
+        )
+        self.activation = activation
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask_pad: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
+        cache: torch.Tensor = torch.zeros((0, 0, 0)),
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Compute convolution module.
+        Args:
+            x (torch.Tensor): Input tensor (#batch, time, channels).
+            mask_pad (torch.Tensor): used for batch padding (#batch, 1, time),
+                (0, 0, 0) means fake mask.
+            cache (torch.Tensor): left context cache, it is only
+                used in causal convolution (#batch, channels, cache_t),
+                (0, 0, 0) meas fake cache.
+        Returns:
+            torch.Tensor: Output tensor (#batch, time, channels).
+        """
+        # exchange the temporal dimension and the feature dimension
+        x = x.transpose(1, 2)  # (#batch, channels, time)
+
+        # mask batch padding
+        if mask_pad.size(2) > 0:  # time > 0
+            x.masked_fill_(~mask_pad, 0.0)
+
+        if self.lorder > 0:
+            if cache.size(2) == 0:  # cache_t == 0
+                x = nn.functional.pad(x, (self.lorder, 0), "constant", 0.0)
+            else:
+                assert cache.size(0) == x.size(0)  # equal batch
+                assert cache.size(1) == x.size(1)  # equal channel
+                x = torch.cat((cache, x), dim=2)
+            assert x.size(2) > self.lorder
+            new_cache = x[:, :, -self.lorder :]
+        else:
+            # It's better we just return None if no cache is required,
+            # However, for JIT export, here we just fake one tensor instead of
+            # None.
+            new_cache = torch.zeros((0, 0, 0), dtype=x.dtype, device=x.device)
+
+        # GLU mechanism
+        x = self.pointwise_conv1(x)  # (batch, 2*channel, dim)
+        x = nn.functional.glu(x, dim=1)  # (batch, channel, dim)
+
+        # 1D Depthwise Conv
+        x = self.depthwise_conv(x)
+        if self.use_layer_norm:
+            x = x.transpose(1, 2)
+        x = self.activation(self.norm(x))
+        if self.use_layer_norm:
+            x = x.transpose(1, 2)
+        x = self.pointwise_conv2(x)
+        # mask batch padding
+        if mask_pad.size(2) > 0:  # time > 0
+            x.masked_fill_(~mask_pad, 0.0)
+
+        return x.transpose(1, 2), new_cache
+
+
+class ConformerEncoderLayer(nn.Module):
+    """Encoder layer module.
+    Args:
+        size (int): Input dimension.
+        self_attn (torch.nn.Module): Self-attention module instance.
+            `MultiHeadedAttention` or `RelPositionMultiHeadedAttention`
+            instance can be used as the argument.
+        feed_forward (torch.nn.Module): Feed-forward module instance.
+            `PositionwiseFeedForward` instance can be used as the argument.
+        feed_forward_macaron (torch.nn.Module): Additional feed-forward module
+             instance.
+            `PositionwiseFeedForward` instance can be used as the argument.
+        conv_module (torch.nn.Module): Convolution module instance.
+            `ConvlutionModule` instance can be used as the argument.
+        dropout_rate (float): Dropout rate.
+        normalize_before (bool):
+            True: use layer_norm before each sub-block.
+            False: use layer_norm after each sub-block.
+    """
+
+    def __init__(
+        self,
+        size: int,
+        self_attn: torch.nn.Module,
+        feed_forward: nn.Module | None = None,
+        feed_forward_macaron: nn.Module | None = None,
+        conv_module: nn.Module | None = None,
+        dropout_rate: float = 0.1,
+        normalize_before: bool = True,
+    ):
+        super().__init__()
+        self.self_attn = self_attn
+        self.feed_forward = feed_forward
+        self.feed_forward_macaron = feed_forward_macaron
+        self.conv_module = conv_module
+        self.norm_ff = nn.LayerNorm(size, eps=1e-5)
+        self.norm_mha = nn.LayerNorm(size, eps=1e-5)
+        if feed_forward_macaron is not None:
+            self.norm_ff_macaron = nn.LayerNorm(size, eps=1e-5)
+            self.ff_scale = 0.5
+        else:
+            self.ff_scale = 1.0
+        if self.conv_module is not None:
+            self.norm_conv = nn.LayerNorm(size, eps=1e-5)
+            self.norm_final = nn.LayerNorm(size, eps=1e-5)
+        self.dropout = nn.Dropout(dropout_rate)
+        self.size = size
+        self.normalize_before = normalize_before
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: torch.Tensor,
+        pos_emb: torch.Tensor,
+        mask_pad: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
+        att_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
+        cnn_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Compute encoded features.
+
+        Args:
+            x (torch.Tensor): (#batch, time, size)
+            mask (torch.Tensor): Mask tensor for the input (#batch, time，time),
+                (0, 0, 0) means fake mask.
+            pos_emb (torch.Tensor): positional encoding, must not be None
+                for ConformerEncoderLayer.
+            mask_pad (torch.Tensor): batch padding mask used for conv module.
+                (#batch, 1，time), (0, 0, 0) means fake mask.
+            att_cache (torch.Tensor): Cache tensor of the KEY & VALUE
+                (#batch=1, head, cache_t1, d_k * 2), head * d_k == size.
+            cnn_cache (torch.Tensor): Convolution cache in conformer layer
+                (#batch=1, size, cache_t2)
+        Returns:
+            torch.Tensor: Output tensor (#batch, time, size).
+            torch.Tensor: Mask tensor (#batch, time, time).
+            torch.Tensor: att_cache tensor,
+                (#batch=1, head, cache_t1 + time, d_k * 2).
+            torch.Tensor: cnn_cache tensor (#batch, size, cache_t2).
+        """
+
+        if self.feed_forward_macaron is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.norm_ff_macaron(x)
+            x = residual + self.ff_scale * self.dropout(self.feed_forward_macaron(x))
+
+        residual = x
+        if self.normalize_before:
+            x = self.norm_mha(x)
+        x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb, att_cache)
+        x = residual + self.dropout(x_att)
+
+        new_cnn_cache = torch.zeros((0, 0, 0), dtype=x.dtype, device=x.device)
+        if self.conv_module is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.norm_conv(x)
+            x, new_cnn_cache = self.conv_module(x, mask_pad, cnn_cache)
+            x = residual + self.dropout(x)
+
+        residual = x
+        if self.normalize_before:
+            x = self.norm_ff(x)
+        x = residual + self.ff_scale * self.dropout(self.feed_forward(x))
+
+        if self.conv_module is not None:
+            x = self.norm_final(x)
+
+        return x, mask, new_att_cache, new_cnn_cache
+
+
+class BaseEncoder(torch.nn.Module):
+    def __init__(
+        self,
+        input_size: int,
+        output_size: int = 256,
+        attention_heads: int = 4,
+        linear_units: int = 2048,
+        num_blocks: int = 6,
+        dropout_rate: float = 0.0,
+        input_layer: str = "conv2d",
+        pos_enc_layer_type: str = "abs_pos",
+        normalize_before: bool = True,
+        concat_after: bool = False,
+    ):
+        super().__init__()
+        self._output_size = output_size
+
+        pos_enc_map = {"abs_pos": PositionalEncoding, "rel_pos": RelPositionalEncoding}
+        if pos_enc_layer_type not in pos_enc_map:
+            raise ValueError("unknown pos_enc_layer: " + pos_enc_layer_type)
+        pos_enc_class = pos_enc_map[pos_enc_layer_type]
+
+        subsampling_map = {"conv2d2": Conv2dSubsampling2}
+        if input_layer not in subsampling_map:
+            raise ValueError("unknown input_layer: " + input_layer)
+        subsampling_class = subsampling_map[input_layer]
+
+        self.embed = subsampling_class(
+            input_size,
+            output_size,
+            dropout_rate,
+            pos_enc_class(output_size, dropout_rate),
+        )
+
+        self.normalize_before = normalize_before
+        self.after_norm = torch.nn.LayerNorm(output_size, eps=1e-5)
+
+    def output_size(self) -> int:
+        return self._output_size
+
+    def forward(
+        self,
+        xs: torch.Tensor,
+        xs_lens: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Embed positions in tensor.
+
+        Args:
+            xs: padded input tensor (B, T, D)
+            xs_lens: input length (B)
+
+        Returns:
+            encoder output tensor xs, and subsampled masks
+            xs: padded output tensor (B, T' ~= T/subsample_rate, D)
+            masks: torch.Tensor batch padding mask after subsample
+                (B, 1, T' ~= T/subsample_rate)
+        """
+        T = xs.size(1)
+        masks = ~make_pad_mask(xs_lens, T).unsqueeze(1)  # (B, 1, T)
+        xs, pos_emb, masks = self.embed(xs, masks)
+        chunk_masks = masks
+        mask_pad = masks  # (B, 1, T/subsample_rate)
+        for layer in self.encoders:
+            xs, chunk_masks, _, _ = layer(xs, chunk_masks, pos_emb, mask_pad)
+        if self.normalize_before:
+            xs = self.after_norm(xs)
+        # Here we assume the mask is not changed in encoder layers, so just
+        # return the masks before encoder layers, and the masks will be used
+        # for cross attention with decoder later
+        return xs, masks
+
+
+class ConformerEncoder(BaseEncoder):
+    """Conformer encoder module."""
+
+    def __init__(
+        self,
+        input_size: int,
+        output_size: int = 256,
+        attention_heads: int = 4,
+        linear_units: int = 2048,
+        num_blocks: int = 6,
+        dropout_rate: float = 0.0,
+        input_layer: str = "conv2d",
+        pos_enc_layer_type: str = "rel_pos",
+        normalize_before: bool = True,
+        concat_after: bool = False,
+        macaron_style: bool = False,
+        use_cnn_module: bool = True,
+        cnn_module_kernel: int = 15,
+    ):
+        super().__init__(
+            input_size,
+            output_size,
+            attention_heads,
+            linear_units,
+            num_blocks,
+            dropout_rate,
+            input_layer,
+            pos_enc_layer_type,
+            normalize_before,
+            concat_after,
+        )
+
+        activation = torch.nn.SiLU()
+
+        # self-attention module definition
+        if pos_enc_layer_type != "rel_pos":
+            encoder_selfattn_layer = MultiHeadedAttention
+        else:
+            encoder_selfattn_layer = RelPositionMultiHeadedAttention
+        encoder_selfattn_layer_args = (
+            attention_heads,
+            output_size,
+            dropout_rate,
+        )
+
+        # feed-forward module definition
+        positionwise_layer = PositionwiseFeedForward
+        positionwise_layer_args = (
+            output_size,
+            linear_units,
+            dropout_rate,
+            activation,
+        )
+        # convolution module definition
+        convolution_layer = ConvolutionModule
+        convolution_layer_args = (
+            output_size,
+            cnn_module_kernel,
+            activation,
+        )
+
+        self.encoders = torch.nn.ModuleList(
+            [
+                ConformerEncoderLayer(
+                    output_size,
+                    encoder_selfattn_layer(*encoder_selfattn_layer_args),
+                    positionwise_layer(*positionwise_layer_args),
+                    positionwise_layer(*positionwise_layer_args) if macaron_style else None,
+                    convolution_layer(*convolution_layer_args) if use_cnn_module else None,
+                    dropout_rate,
+                    normalize_before,
+                )
+                for _ in range(num_blocks)
+            ]
+        )

vllm_omni/model_executor/models/indextts2/gpt/embeddings.py

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Learned position embeddings for IndexTTS2 GPT model."""
+
+import torch
+import torch.nn as nn
+
+
+class LearnedPositionEmbeddings(nn.Module):
+    def __init__(self, seq_len: int, model_dim: int, init: float = 0.02):
+        super().__init__()
+        self.emb = nn.Embedding(seq_len, model_dim)
+        self.emb.weight.data.normal_(mean=0.0, std=init)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        sl = x.shape[1]
+        return self.emb(torch.arange(0, sl, device=x.device))
+
+    def get_fixed_embedding(self, ind: int, dev: torch.device) -> torch.Tensor:
+        return self.emb(torch.tensor([ind], device=dev)).unsqueeze(0)

vllm_omni/model_executor/models/indextts2/gpt/perceiver.py

@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Adapted from lucidrains/naturalspeech2-pytorch  (PerceiverResampler)
+# https://github.com/lucidrains/naturalspeech2-pytorch/blob/659bec7f/
+#   naturalspeech2_pytorch/naturalspeech2_pytorch.py#L532
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from torch import nn
+
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if callable(d) else d
+
+
+class Attend(nn.Module):
+    def __init__(self, dropout=0.0):
+        super().__init__()
+        self.dropout = dropout
+
+    def forward(self, q, k, v, mask=None):
+        attn_mask = None
+        if exists(mask):
+            attn_mask = rearrange(mask, "b j -> b 1 1 j")
+
+        return F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=attn_mask,
+            dropout_p=self.dropout if self.training else 0.0,
+        )
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, scale=True):
+        super().__init__()
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(dim)) if scale else None
+
+    def forward(self, x):
+        gamma = default(self.gamma, 1)
+        return F.normalize(x, dim=-1) * self.scale * gamma
+
+
+class GEGLU(nn.Module):
+    def forward(self, x):
+        x, gate = x.chunk(2, dim=-1)
+        return F.gelu(gate) * x
+
+
+def FeedForward(dim, mult=4):
+    dim_inner = int(dim * mult * 2 / 3)
+    return nn.Sequential(nn.Linear(dim, dim_inner * 2), GEGLU(), nn.Linear(dim_inner, dim))
+
+
+class PerceiverResampler(nn.Module):
+    def __init__(
+        self,
+        dim,
+        depth=2,
+        dim_context=None,
+        num_latents=32,
+        dim_head=64,
+        heads=8,
+        ff_mult=4,
+    ):
+        super().__init__()
+        dim_context = default(dim_context, dim)
+
+        self.proj_context = nn.Linear(dim_context, dim) if dim_context != dim else nn.Identity()
+
+        self.latents = nn.Parameter(torch.randn(num_latents, dim))
+        nn.init.normal_(self.latents, std=0.02)
+
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        Attention(dim=dim, dim_head=dim_head, heads=heads, cross_attn_include_queries=True),
+                        FeedForward(dim=dim, mult=ff_mult),
+                    ]
+                )
+            )
+
+        self.norm = RMSNorm(dim)
+
+    def forward(self, x, mask=None):
+        batch = x.shape[0]
+
+        x = self.proj_context(x)
+
+        latents = repeat(self.latents, "n d -> b n d", b=batch)
+
+        for attn, ff in self.layers:
+            latents = attn(latents, x, mask=mask) + latents
+            latents = ff(latents) + latents
+
+        return self.norm(latents)
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        *,
+        dim_context=None,
+        dim_head=64,
+        heads=8,
+        dropout=0.0,
+        cross_attn_include_queries=False,
+    ):
+        super().__init__()
+        self.scale = dim_head**-0.5
+        self.heads = heads
+        self.cross_attn_include_queries = cross_attn_include_queries
+
+        dim_inner = dim_head * heads
+        dim_context = default(dim_context, dim)
+
+        self.attend = Attend(dropout=dropout)
+        self.to_q = nn.Linear(dim, dim_inner, bias=False)
+        self.to_kv = nn.Linear(dim_context, dim_inner * 2, bias=False)
+        self.to_out = nn.Linear(dim_inner, dim, bias=False)
+
+    def forward(self, x, context=None, mask=None):
+        h, has_context = self.heads, exists(context)
+
+        context = default(context, x)
+
+        if has_context and self.cross_attn_include_queries:
+            context = torch.cat((x, context), dim=-2)
+
+        q, k, v = (self.to_q(x), *self.to_kv(context).chunk(2, dim=-1))
+        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v))
+
+        out = self.attend(q, k, v, mask=mask)
+
+        out = rearrange(out, "b h n d -> b n (h d)")
+        return self.to_out(out)

vllm_omni/model_executor/models/indextts2/indextts2_s2mel_decoder.py

@@ -0,0 +1,894 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""IndexTTS2 Stage 1: S2Mel decoder + BigVGAN vocoder.
+
+Receives mel_codes + latent from Stage 0 (GPT AR talker), runs flow matching
+to synthesize mel spectrogram, then BigVGAN to produce waveform audio.
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+import os
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+import torch.nn as nn
+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_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 .s2mel.modules.commons import AttrDict, MyModel
+
+logger = init_logger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Lazy loaders for external models
+# ---------------------------------------------------------------------------
+
+_bigvgan_models: dict[tuple[str, str], nn.Module] = {}
+
+
+def _patch_bigvgan_compat(cls):
+    """Monkey-patch BigVGAN._from_pretrained for huggingface_hub>=1.0 compat."""
+    if getattr(cls._from_pretrained, "_hf1_patched", False):
+        return
+    orig = cls._from_pretrained.__func__
+
+    @classmethod
+    def _compat(klass, *, proxies=None, resume_download=False, **kw):
+        return orig(klass, proxies=proxies, resume_download=resume_download, **kw)
+
+    _compat._hf1_patched = True
+    cls._from_pretrained = _compat
+
+
+def _strip_weight_norm(root: nn.Module) -> list[str]:
+    """Bake legacy weight_norm (weight_g/weight_v) into plain weights.
+
+    The DiT final-layer WaveNet (and any other S2Mel submodule) keeps
+    torch.nn.utils.weight_norm forward pre-hooks after loading, so every
+    CFM estimator step recomputes the normalization (~217 calls/request
+    in profiling). Inference never updates these weights, so we fold the
+    parametrization once at load time. Hooks live on the innermost conv
+    (e.g. SConv1d.conv.conv), hence the generic module walk instead of
+    the per-class remove_weight_norm() helpers.
+
+    Returns the qualified names (relative to ``root``) of the plain
+    weight parameters created by the fold, so callers can mark them as
+    checkpoint-initialized for vLLM's loader validation.
+    """
+    from torch.nn.utils import remove_weight_norm
+
+    folded: list[str] = []
+    for mod_name, module in root.named_modules():
+        for hook in list(module._forward_pre_hooks.values()):
+            if hook.__class__.__name__ == "WeightNorm":
+                remove_weight_norm(module, name=hook.name)
+                folded.append(f"{mod_name}.{hook.name}" if mod_name else hook.name)
+    return folded
+
+
+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:
+        return _bigvgan_models[cache_key]
+
+    lock = _repo_prefetch_lock(vocoder_name) if not os.path.isdir(vocoder_name) else None
+    if lock is not None:
+        lock.__enter__()
+    try:
+        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.BigVGAN)
+            bigvgan_model = bigvgan.BigVGAN.from_pretrained(vocoder_name)
+    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
+    return bigvgan_model
+
+
+# ---------------------------------------------------------------------------
+# Stage 1 Decoder
+# ---------------------------------------------------------------------------
+
+
+class IndexTTS2S2MelDecoder(nn.Module):
+    """S2Mel + BigVGAN decoder for IndexTTS2 Stage 1.
+
+    Receives from Stage 0:
+      - mel_codes: [T] mel code sequence
+      - latent: [T, 1280] hidden states from GPT AR
+      - S_ref: speaker semantic embeddings (from Wav2Vec2 → RepCodec)
+      - ref_mel: [80, T_ref] reference mel spectrogram
+      - style: [192] CAMPPlus style vector
+    """
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        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]
+
+        # --- Flags for vLLM-Omni framework ---
+        self.have_multimodal_outputs = True
+        self.has_preprocess = False
+        self.has_postprocess = False
+        self.enable_update_additional_information = True
+        self.requires_raw_input_tokens = True
+
+        # --- Build S2Mel model (CFM + LengthRegulator + gpt_layer) ---
+        s2mel_cfg = self.config.s2mel
+        s2mel_args = AttrDict(s2mel_cfg)
+        # Ensure nested dicts are also AttrDicts
+        for key in ["DiT", "wavenet", "length_regulator", "style_encoder", "preprocess_params"]:
+            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)
+
+        # 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
+
+        # Run the DiT estimator under bf16 autocast (flash attention + faster
+        # GEMMs); the CFM Euler solver stays float32. Disable via deploy YAML
+        # hf_overrides `s2mel_dit_bf16: false` if quality regresses.
+        self.s2mel_dit_bf16: bool = getattr(self.config, "s2mel_dit_bf16", True)
+        # Run BigVGAN vocoding under bf16 autocast (conv-heavy, ~2x faster).
+        # Disable via deploy YAML hf_overrides `s2mel_vocoder_bf16: false`.
+        self.s2mel_vocoder_bf16: bool = getattr(self.config, "s2mel_vocoder_bf16", True)
+        # 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._vocoder_graph: 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_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_compile_shapes=%s",
+            self.s2mel_dit_bf16,
+            self.s2mel_vocoder_bf16,
+            self.s2mel_vocoder_cuda_graph,
+            self.s2mel_dit_cuda_graph,
+            self.s2mel_dit_cuda_graph_max_graphs,
+            self.s2mel_vocoder_capture_sizes,
+            self.s2mel_vocoder_compile_shapes,
+        )
+
+    # ------------------------------------------------------------------
+    # vLLM hooks
+    # ------------------------------------------------------------------
+
+    def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
+        if input_ids.numel() == 0:
+            return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32)
+        return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32)
+
+    @torch.no_grad()
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor | None = None,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        runtime_additional_information: list[dict[str, Any]] | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor | OmniOutput:
+        """Run S2Mel flow matching + BigVGAN vocoding."""
+        model_intermediate_buffer = kwargs.get("model_intermediate_buffer")
+        if model_intermediate_buffer is None:
+            model_intermediate_buffer = runtime_additional_information
+        if runtime_additional_information is not None and "model_intermediate_buffer" not in kwargs:
+            logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer")
+
+        request_infos: list[dict[str, Any]] = []
+        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 {}
+
+        device = input_ids.device
+        model_dtype = self.s2mel.models["gpt_layer"][0].weight.dtype
+
+        logger.debug(
+            "[S2Mel forward] model_intermediate_buffer len=%d, keys=%s",
+            len(model_intermediate_buffer) if model_intermediate_buffer else 0,
+            list(additional_information.keys()) if additional_information else [],
+        )
+        for k, v in additional_information.items():
+            if isinstance(v, torch.Tensor):
+                logger.debug("[S2Mel forward] %s: shape=%s, dtype=%s", k, v.shape, v.dtype)
+
+        # Extract Stage 0 outputs and pad multiple ready Stage-1 requests into
+        # 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(
+                request_infos,
+                device=device,
+                model_dtype=model_dtype,
+            )
+            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)
+            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)
+
+        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},
+            )
+
+        logger.debug(
+            "[S2Mel] extracted tensors — mel_codes=%s latent=%s code_lens=%s s_ref=%s ref_mel=%s style=%s",
+            mel_codes.shape if mel_codes is not None else None,
+            latent.shape if latent is not None else None,
+            code_lens_raw,
+            s_ref.shape if isinstance(s_ref, torch.Tensor) else None,
+            ref_mel.shape if isinstance(ref_mel, torch.Tensor) else None,
+            style.shape if isinstance(style, torch.Tensor) else None,
+        )
+
+        # Ensure batch dimension
+        if mel_codes.ndim == 1:
+            mel_codes = mel_codes.unsqueeze(0)
+        if latent.ndim == 2:
+            latent = latent.unsqueeze(0)
+
+        # Compute code_lens
+        if code_lens_raw is not None:
+            if isinstance(code_lens_raw, torch.Tensor):
+                if "code_lens_list" not in locals():
+                    code_lens_list = [int(x) for x in code_lens_raw.detach().cpu().reshape(-1).tolist()]
+                code_lens = code_lens_raw.to(device=device, dtype=torch.long)
+            else:
+                code_lens_list = [int(x) for x in code_lens_raw]
+                code_lens = torch.tensor(code_lens_list, device=device, dtype=torch.long)
+        else:
+            # Strip stop token (8193) and compute actual length
+            stop_token = self.config.gpt.get("stop_mel_token", 8193)
+            code_lens_list = []
+            for i in range(mel_codes.shape[0]):
+                stop_mask = (mel_codes[i] == stop_token).nonzero(as_tuple=False)
+                if stop_mask.numel() > 0:
+                    code_lens_list.append(stop_mask[0].item())
+                else:
+                    code_lens_list.append(mel_codes.shape[1])
+            code_lens = torch.tensor(code_lens_list, device=device, dtype=torch.long)
+
+        # Trim to actual length
+        max_len = max(code_lens_list)
+        mel_codes = mel_codes[:, :max_len]
+        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,
+                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)
+        codebook_size = self.config.semantic_codec.get("codebook_size", 8192)
+        if logger.isEnabledFor(logging.DEBUG):
+            logger.debug(
+                "mel_codes debug: shape=%s, min=%d, max=%d",
+                mel_codes.shape,
+                mel_codes.min().item(),
+                mel_codes.max().item(),
+            )
+        if getattr(self.config, "s2mel_validate_code_range", False) and torch.any(
+            (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)
+
+        # 3. Length regulate: codes → mel frames
+        target_lengths_list = [int(length * self.mel_code_to_frame_ratio) for length in code_lens_list]
+        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)
+        cond = self.s2mel.models["length_regulator"](S_infer, ylens=target_lengths, n_quantizers=3, f0=None)[
+            0
+        ]  # [B, T_mel, 512]
+        logger.debug("[S2Mel] step3 length_regulator → cond=%s", cond.shape)
+
+        # 4. Reference prompt conditioning
+        if s_ref is not None and ref_mel is not None:
+            if ref_mel.ndim == 2:
+                ref_mel = ref_mel.unsqueeze(0)
+
+            if s_ref.ndim == 3:
+                # Official infer_v2.py uses the quantized embedding returned by
+                # semantic_codec.quantize() directly as S_ref.
+                s_ref_emb = s_ref.to(device=device, dtype=model_dtype)
+            else:
+                # Backward-compatible path for older payloads that carried
+                # codebook indices instead of quantized embeddings.
+                s_ref_codes = s_ref.long()
+                if s_ref_codes.ndim == 1:
+                    s_ref_codes = s_ref_codes.unsqueeze(0)  # [1, T]
+                codebook_size = self.config.semantic_codec.get("codebook_size", 8192)
+                if logger.isEnabledFor(logging.DEBUG):
+                    logger.debug(
+                        "S_ref debug: shape=%s, min=%d, max=%d, codebook_size=%d",
+                        s_ref_codes.shape,
+                        s_ref_codes.min().item(),
+                        s_ref_codes.max().item(),
+                        codebook_size,
+                    )
+                s_ref_codes = s_ref_codes.clamp(0, codebook_size - 1)
+                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]
+
+            ref_target_lengths = self._infer_ref_lengths(ref_mel)
+            prompt_condition = self.s2mel.models["length_regulator"](
+                s_ref_emb, ylens=ref_target_lengths, n_quantizers=3, f0=None
+            )[0]
+            cat_condition = torch.cat([prompt_condition, cond], dim=1)
+            logger.debug(
+                "[S2Mel] step4 ref conditioning — s_ref_emb=%s prompt_condition=%s cat_condition=%s",
+                s_ref_emb.shape,
+                prompt_condition.shape,
+                cat_condition.shape,
+            )
+        else:
+            prompt_condition = None
+            cat_condition = cond
+            ref_mel = torch.zeros(cond.shape[0], 80, 0, device=device, dtype=model_dtype)
+
+        if style is None:
+            style = torch.zeros(cat_condition.shape[0], 192, device=device, dtype=model_dtype)
+        if style.ndim == 1:
+            style = style.unsqueeze(0)
+
+        logger.debug(
+            "[S2Mel] step5 CFM input — cat_condition=%s ref_mel=%s style=%s steps=%d cfg_rate=%.2f",
+            cat_condition.shape,
+            ref_mel.shape,
+            style.shape,
+            self.diffusion_steps,
+            self.inference_cfg_rate,
+        )
+
+        # 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).
+        cfm = self.s2mel.models["cfm"]
+        cfm.estimator_autocast_dtype = torch.bfloat16 if self.s2mel_dit_bf16 else None
+        cond = cond.float()
+        if prompt_condition is not None:
+            prompt_condition = prompt_condition.float()
+        cat_condition = cat_condition.float()
+        ref_mel = ref_mel.float()
+        style = style.float()
+
+        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
+
+        estimator = cfm.estimator
+        required_cache_batch = max(2, int(cat_condition.shape[0]) * (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)
+
+        # 6. BigVGAN vocoding
+        vocoder_cfg = self.config.vocoder
+        vocoder_name = vocoder_cfg.get("name", "nvidia/bigvgan_v2_22khz_80band_256x")
+        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))
+        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))
+        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
+        # numerically equivalent to official infer_v2.py before its int16 save
+        # step, while avoiding double scaling in OpenAI/soundfile encoders.
+        if isinstance(wav, torch.Tensor) and wav.ndim == 0:
+            wav = wav.unsqueeze(0)
+
+        if logger.isEnabledFor(logging.DEBUG):
+            logger.debug(
+                "[S2Mel] step6 BigVGAN done → wav=%s min=%.1f max=%.1f sr=22050",
+                wav.shape if isinstance(wav, torch.Tensor) else [w.shape for w in wav],
+                wav.min().item() if isinstance(wav, torch.Tensor) else min(w.min().item() for w in wav),
+                wav.max().item() if isinstance(wav, torch.Tensor) else max(w.max().item() for w in wav),
+            )
+
+        audio_cpu = [w.cpu() for w in wav] if isinstance(wav, list) else wav.cpu()
+
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={
+                "audio": audio_cpu,
+                "sr": torch.tensor(22050, dtype=torch.int32),
+            },
+        )
+
+    def compute_logits(self, hidden_states: Any, sampling_metadata: Any = None) -> None:
+        return None
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def _build_batched_inputs(
+        self,
+        infos: list[dict[str, Any]],
+        *,
+        device: torch.device,
+        model_dtype: torch.dtype,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        list[int],
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+    ]:
+        stop_token = self.config.gpt.get("stop_mel_token", 8193)
+        mel_items: list[torch.Tensor] = []
+        latent_items: list[torch.Tensor] = []
+        code_lens: list[int] = []
+        s_ref_items: list[torch.Tensor | None] = []
+        ref_mel_items: list[torch.Tensor | None] = []
+        style_items: list[torch.Tensor | None] = []
+
+        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")
+
+            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])
+
+            raw_len = info.get("code_lens")
+            if isinstance(raw_len, torch.Tensor) and raw_len.numel() > 0:
+                code_len = int(raw_len.reshape(-1)[0].item())
+            elif raw_len is not None:
+                code_len = int(raw_len[0] if isinstance(raw_len, list) else raw_len)
+            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])))
+            mel_items.append(mel[:code_len])
+            latent_items.append(latent[:code_len])
+            code_lens.append(code_len)
+
+            s_ref = self._get_tensor(info, "S_ref", device)
+            if isinstance(s_ref, torch.Tensor) and s_ref.ndim == 3 and s_ref.shape[0] == 1:
+                s_ref = s_ref[0]
+            s_ref_items.append(s_ref)
+
+            ref_mel = self._get_tensor(info, "ref_mel", device, model_dtype)
+            if isinstance(ref_mel, torch.Tensor) and ref_mel.ndim == 3 and ref_mel.shape[0] == 1:
+                ref_mel = ref_mel[0]
+            ref_mel_items.append(ref_mel)
+
+            style = self._get_tensor(info, "style", device, model_dtype)
+            if isinstance(style, torch.Tensor) and style.ndim == 2 and style.shape[0] == 1:
+                style = style[0]
+            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)):
+            mel_batch[i, : mel.shape[0]] = mel
+            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_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 (
+            mel_batch,
+            latent_batch,
+            torch.tensor(code_lens, device=device, dtype=torch.long),
+            code_lens,
+            s_ref_batch,
+            ref_mel_batch,
+            style_batch,
+        )
+
+    @staticmethod
+    def _stack_optional_sequence(
+        items: list[torch.Tensor | None],
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor | None:
+        if not items or any(item is None for item in items):
+            return None
+        tensors = [item.to(device=device) for item in items if item is not None]
+        max_len = max(int(t.shape[0]) for t in tensors)
+        if tensors[0].ndim == 1:
+            out = torch.zeros(len(tensors), max_len, device=device, dtype=tensors[0].dtype)
+            for i, t in enumerate(tensors):
+                out[i, : t.shape[0]] = t
+            return out
+        out = torch.zeros(len(tensors), max_len, tensors[0].shape[-1], device=device, dtype=dtype)
+        for i, t in enumerate(tensors):
+            out[i, : t.shape[0], :] = t.to(dtype=dtype)
+        return out
+
+    @staticmethod
+    def _stack_optional_ref_mel(
+        items: list[torch.Tensor | None],
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor | None:
+        if not items or any(item is None for item in items):
+            return None
+        tensors = [item.to(device=device, dtype=dtype) for item in items if item is not None]
+        max_len = max(int(t.shape[-1]) for t in tensors)
+        out = torch.zeros(len(tensors), tensors[0].shape[-2], max_len, device=device, dtype=dtype)
+        for i, t in enumerate(tensors):
+            out[i, :, : t.shape[-1]] = t
+        return out
+
+    @staticmethod
+    def _stack_optional_style(
+        items: list[torch.Tensor | None],
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor | None:
+        if not items or any(item is None for item in items):
+            return None
+        return torch.stack(
+            [item.to(device=device, dtype=dtype).reshape(-1) for item in items if item is not None],
+            dim=0,
+        )
+
+    @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)
+
+    @staticmethod
+    def _get_tensor(
+        info: dict,
+        key: str,
+        device: torch.device,
+        dtype: torch.dtype | None = None,
+    ) -> torch.Tensor | None:
+        """Extract tensor from additional_information, handling nested dicts."""
+        val = info.get(key)
+        if val is None:
+            for sub_key in ["codes", "hidden_states", "meta"]:
+                sub = info.get(sub_key)
+                if isinstance(sub, dict) and key in sub:
+                    val = sub[key]
+                    break
+        if val is None:
+            return None
+        if isinstance(val, torch.Tensor):
+            return val.to(device=device, dtype=dtype) if dtype else val.to(device=device)
+        if isinstance(val, list):
+            if val and isinstance(val[0], torch.Tensor):
+                return val[0].to(device=device, dtype=dtype) if dtype else val[0].to(device=device)
+        return None
+
+    # ------------------------------------------------------------------
+    # Weight loading
+    # ------------------------------------------------------------------
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load weights from s2mel.pth checkpoint.
+
+        Bypasses vllm's default .pt iterator (same reason as talker: raw
+        tensor .pt files in model dir). Loads s2mel.pth directly and
+        flattens the nested ``state["net"]`` structure:
+
+        - ``net["cfm"][key]`` → ``s2mel.models.cfm.{key}``
+        - ``net["length_regulator"][key]`` → ``s2mel.models.length_regulator.{key}``
+        - ``net["gpt_layer"][key]`` → ``s2mel.models.gpt_layer.{key}``
+        """
+        _ = weights
+
+        ckpt_path = resolve_model_file(self.model_path, self.config.s2mel_checkpoint)
+        if ckpt_path is None:
+            raise FileNotFoundError(f"IndexTTS2 S2Mel checkpoint {self.config.s2mel_checkpoint!r} not found")
+        state = torch.load(ckpt_path, map_location="cpu", weights_only=True)
+        net = state["net"] if "net" in state else state
+
+        params_dict = dict(self.named_parameters(remove_duplicate=False))
+        loaded_params: set[str] = set()
+
+        prefix_map = {
+            "cfm.": "s2mel.models.cfm.",
+            "length_regulator.": "s2mel.models.length_regulator.",
+            "gpt_layer.": "s2mel.models.gpt_layer.",
+        }
+
+        # Flatten nested ModuleDict state: {module: {param: tensor}} → flat iter
+        flat_items: Iterable[tuple[str, torch.Tensor]]
+        if net and isinstance(next(iter(net.values())), dict):
+            flat_items = (
+                (f"{mod}.{param}", tensor) for mod, mod_state in net.items() for param, tensor in mod_state.items()
+            )
+        else:
+            flat_items = net.items()
+
+        for name, loaded_weight in flat_items:
+            mapped_name = name
+            for old_prefix, new_prefix in prefix_map.items():
+                if name.startswith(old_prefix):
+                    mapped_name = new_prefix + name[len(old_prefix) :]
+                    break
+
+            if mapped_name not in params_dict:
+                logger.debug("Skipping unrecognized S2Mel weight: %s → %s", name, mapped_name)
+                continue
+
+            param = params_dict[mapped_name]
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, loaded_weight)
+            loaded_params.add(mapped_name)
+
+        logger.info("Loaded %d weights for IndexTTS2S2MelDecoder", len(loaded_params))
+        required_prefixes = [
+            "s2mel.models.cfm.",
+            "s2mel.models.length_regulator.",
+            "s2mel.models.gpt_layer.",
+        ]
+        missing_prefixes = [
+            prefix for prefix in required_prefixes if not any(name.startswith(prefix) for name in loaded_params)
+        ]
+        if missing_prefixes:
+            raise RuntimeError(
+                "IndexTTS2 S2Mel checkpoint did not load required parameter groups: " + ", ".join(missing_prefixes)
+            )
+
+        stripped = _strip_weight_norm(self.s2mel)
+        if stripped:
+            # The fold replaces weight_g/weight_v with a plain `weight`
+            # parameter; register it as loaded so vLLM's initialization
+            # check does not flag it as missing from the checkpoint.
+            for name in stripped:
+                loaded_params.add(f"s2mel.{name}")
+            logger.info("Folded %d weight_norm parametrizations in S2Mel", len(stripped))
+
+        # Keep the CFM ODE solver in float32 once after loading instead of
+        # recursively re-casting the module on every request.
+        cfm = self.s2mel.models["cfm"]
+        cfm.float()
+        cfm.estimator_autocast_dtype = torch.bfloat16 if self.s2mel_dit_bf16 else None
+
+        self._warmup_external_models()
+        return loaded_params
+
+    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")
+            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)
+            logger.info("BigVGAN + semantic codec preloaded")
+        except Exception as e:
+            logger.warning("Failed to preload Stage 1 models: %s", e)
+
+    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.
+
+        Warmup is lazy (first request) following the GLM-TTS DiT wrapper
+        pattern; on warmup failure the wrapper is disabled and all
+        subsequent calls fall back to the eager model.
+        """
+        if not self.s2mel_vocoder_cuda_graph or device.type != "cuda":
+            return bigvgan
+        if self._vocoder_graph is None:
+            from vllm_omni.model_executor.models.indextts2.bigvgan_cuda_graph import (
+                CUDAGraphBigVGANWrapper,
+            )
+
+            self._vocoder_graph = CUDAGraphBigVGANWrapper(
+                bigvgan,
+                capture_sizes=self.s2mel_vocoder_capture_sizes,
+                compile_shapes=self.s2mel_vocoder_compile_shapes,
+            )
+            try:
+                self._vocoder_graph.warmup(device, dtype)
+            except Exception:
+                logger.warning("Disabling BigVGAN CUDA graphs after warmup failure", exc_info=True)
+                self._vocoder_graph.enabled = False
+        return self._vocoder_graph
+
+    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.
+        """
+        if not self.s2mel_dit_cuda_graph:
+            return
+        if getattr(estimator, "_cuda_graph_runner", None) is not None:
+            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,
+        )
+        logger.info(
+            "S2Mel DiT CUDA graph runner enabled (per-shape capture, LRU=%d)",
+            self.s2mel_dit_cuda_graph_max_graphs,
+        )

vllm_omni/model_executor/models/indextts2/pipeline.py

@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""IndexTTS2 pipeline: GPT AR talker (text → mel codes) → S2Mel + BigVGAN (mel → audio).
+
+Two-stage non-streaming pipeline. S2Mel flow matching (25 Euler steps) requires
+the full mel code sequence, so async_chunk is not applicable.
+"""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+_PROC = "vllm_omni.model_executor.stage_input_processors.indextts2"
+
+INDEXTTS2_PIPELINE = PipelineConfig(
+    model_type="indextts2",
+    model_arch="IndexTTS2TalkerForConditionalGeneration",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="indextts2_talker",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            owns_tokenizer=True,
+            engine_output_type="latent",
+            extras={"skip_tokenizer_init": True, "tokenizer": "gpt2"},
+            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_s2mel_decoder",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="audio",
+            engine_output_type="audio",
+            model_arch="IndexTTS2S2MelDecoder",
+            sync_process_input_func=f"{_PROC}.talker2s2mel_token_only",
+            extras={"skip_tokenizer_init": True, "tokenizer": "gpt2"},
+            sampling_constraints={"detokenize": True},
+        ),
+    ),
+)

vllm_omni/model_executor/models/indextts2/preprocess_utils.py

@@ -0,0 +1,324 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""External model loading, audio I/O, and emotion conditioning for IndexTTS2."""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from transformers.utils.hub import cached_file
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.model_loader.hub_prefetch import _repo_prefetch_lock
+
+logger = init_logger(__name__)
+
+# ---------------------------------------------------------------------------
+# Lazy external model caches — keyed by (model_path, device) for multi-device
+# ---------------------------------------------------------------------------
+
+_wav2vec2_cache: dict[tuple[str, str], tuple] = {}
+_semantic_codec_cache: dict[tuple[str, str], Any] = {}
+_campplus_cache: dict[tuple[str, str], Any] = {}
+_qwen_emotion_cache: dict[tuple[str, str], tuple] = {}
+
+
+def _freeze(model):
+    model.eval()
+    for p in model.parameters():
+        p.requires_grad_(False)
+    return model
+
+
+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
+    try:
+        return cached_file(model_path, filename)
+    except Exception as e:
+        logger.warning("Could not resolve IndexTTS2 asset %s from %s: %s", filename, model_path, e)
+        return None
+
+
+def load_wav2vec2(model_path: str, device: torch.device):
+    cache_key = (model_path, str(device))
+    if cache_key in _wav2vec2_cache:
+        return _wav2vec2_cache[cache_key]
+    from transformers import AutoFeatureExtractor, Wav2Vec2BertModel
+
+    model, processor = None, None
+    # Try local subfolder first (no lock needed), then HF hub with lock.
+    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)
+    except Exception:
+        try:
+            with _repo_prefetch_lock("facebook/w2v-bert-2.0"):
+                src_args = dict(pretrained_model_name_or_path="facebook/w2v-bert-2.0")
+                processor = AutoFeatureExtractor.from_pretrained(**src_args)
+                model = Wav2Vec2BertModel.from_pretrained(**src_args)
+        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")
+    _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))
+    if cache_key in _semantic_codec_cache:
+        return _semantic_codec_cache[cache_key]
+    from .utils.maskgct.repcodec_model import RepCodec
+
+    codec = RepCodec(
+        codebook_size=config.get("codebook_size", 8192),
+        hidden_size=config.get("hidden_size", 1024),
+        codebook_dim=config.get("codebook_dim", 8),
+    )
+    ckpt_path = resolve_model_file(model_path, "semantic_codec.pth")
+    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:
+        import safetensors.torch
+        from huggingface_hub import hf_hub_download
+
+        ckpt_path = hf_hub_download("amphion/MaskGCT", filename="semantic_codec/model.safetensors")
+        safetensors.torch.load_model(codec, ckpt_path)
+    _freeze(codec.to(device=device, dtype=torch.float32))
+    _semantic_codec_cache[cache_key] = codec
+    return codec
+
+
+def load_campplus(model_path: str, device: torch.device):
+    cache_key = (model_path, str(device))
+    if cache_key in _campplus_cache:
+        return _campplus_cache[cache_key]
+    from .utils.campplus.dtdnn import CAMPPlus
+
+    campplus = CAMPPlus(feat_dim=80, embedding_size=192)
+    ckpt_path = resolve_model_file(model_path, "campplus.pth")
+    if ckpt_path is None:
+        from huggingface_hub import hf_hub_download
+
+        ckpt_path = hf_hub_download("funasr/campplus", filename="campplus_cn_common.bin")
+    state = torch.load(ckpt_path, map_location="cpu", weights_only=True)
+    campplus.load_state_dict(state, strict=False)
+    _freeze(campplus.to(device=device, dtype=torch.float32))
+    _campplus_cache[cache_key] = campplus
+    return campplus
+
+
+def load_qwen_emotion(model_path: str, device: torch.device, *, trust_remote_code: bool = True):
+    cache_key = (model_path, str(device))
+    if cache_key in _qwen_emotion_cache:
+        return _qwen_emotion_cache[cache_key]
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    subfolder = "qwen0.6bemo4-merge"
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_path,
+            subfolder=subfolder,
+            trust_remote_code=trust_remote_code,
+        )
+        model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            subfolder=subfolder,
+            torch_dtype="float16",
+            device_map={"": str(device)},
+            trust_remote_code=trust_remote_code,
+        )
+    except Exception as e:
+        logger.warning(
+            "Could not load QwenEmotion from %s/%s: %s. use_emo_text will fall back to default emotion.",
+            model_path,
+            subfolder,
+            e,
+        )
+        return None, None
+    _freeze(model)
+    _qwen_emotion_cache[cache_key] = (model, tokenizer)
+    return model, tokenizer
+
+
+# ---------------------------------------------------------------------------
+# Audio utilities
+# ---------------------------------------------------------------------------
+
+
+def compute_fbank(wav_16k: torch.Tensor, device: torch.device) -> torch.Tensor:
+    """Compute 80-dim fbank features for CAMPPlus. Input: [T] at 16kHz."""
+    try:
+        import torchaudio.compliance.kaldi as kaldi
+
+        # kaldi.fbank is pure torch — keep on the wav's device (GPU when
+        # available) instead of forcing a CPU round-trip (~100ms for 15s).
+        wav = wav_16k.unsqueeze(0).float()
+        fbank = kaldi.fbank(
+            wav,
+            num_mel_bins=80,
+            dither=0,
+            sample_frequency=16000,
+        )
+        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)
+
+
+def wav2vec_extract(
+    wav_16k: torch.Tensor,
+    model: Any,
+    processor: Any,
+    device: torch.device,
+    w2v_stat: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Extract Wav2Vec2-BERT features. Returns [1, T, 1024]."""
+    wav_np = wav_16k.cpu().numpy()
+    inputs = processor(wav_np, sampling_rate=16000, return_tensors="pt")
+    input_features = inputs["input_features"].to(device=device, dtype=torch.float32)
+    attention_mask = inputs.get("attention_mask")
+    if attention_mask is not None:
+        attention_mask = attention_mask.to(device=device)
+    with torch.no_grad():
+        # bf16 autocast roughly halves the w2v-bert forward (~600M params,
+        # fp32 weights). Features are cast back to fp32 before stat
+        # normalization; downstream consumers are fp32.
+        if input_features.is_cuda:
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                outputs = model(
+                    input_features=input_features,
+                    attention_mask=attention_mask,
+                    output_hidden_states=True,
+                )
+        else:
+            outputs = model(
+                input_features=input_features,
+                attention_mask=attention_mask,
+                output_hidden_states=True,
+            )
+    feat = outputs.hidden_states[17].float()
+    if w2v_stat is not None:
+        mean = w2v_stat[0:1, :].to(device=device)
+        std = torch.sqrt(w2v_stat[1:2, :].to(device=device))
+        feat = (feat - mean) / std
+    return feat
+
+
+def load_reference_audio(
+    audio_path: str | tuple | list,
+    device: torch.device,
+    max_audio_length_seconds: float | None = 15,
+    mode: str = "speaker",
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Load reference audio and resample to 16kHz and 22.05kHz.
+
+    Accepts either a file path (str) or a pre-loaded (wav_list, sr) tuple
+    from the serving layer.
+
+    ``mode`` mirrors official IndexTTS2 v2:
+    - speaker: librosa default path first normalizes to 22.05kHz, truncates,
+      then derives the 16kHz wav2vec/CAMPPlus input from that 22.05kHz signal.
+    - emotion: librosa loads directly at 16kHz, then truncates.
+    """
+    if mode not in {"speaker", "emotion"}:
+        raise ValueError("IndexTTS2 audio mode must be 'speaker' or 'emotion'")
+
+    wav, sr = _load_audio_1d(audio_path)
+    # Resample on GPU when available: CPU sinc resampling for awkward ratios
+    # (e.g. 16000->22050, gcd=50) costs hundreds of ms; on GPU it is <5ms.
+    # Downstream consumers move tensors to CPU explicitly where needed.
+    if device.type == "cuda":
+        wav = wav.to(device)
+    if mode == "speaker":
+        wav_22k = _resample(wav, sr, 22050)
+        wav_22k = _truncate_audio(wav_22k, 22050, max_audio_length_seconds)
+        wav_16k = _resample(wav_22k, 22050, 16000)
+        return wav_16k, wav_22k
+
+    wav_16k = _resample(wav, sr, 16000)
+    wav_16k = _truncate_audio(wav_16k, 16000, max_audio_length_seconds)
+    wav_22k = _resample(wav_16k, 16000, 22050)
+    return wav_16k, wav_22k
+
+
+def _load_audio_1d(audio_path: str | tuple | list) -> tuple[torch.Tensor, int]:
+    if isinstance(audio_path, (list, tuple)) and len(audio_path) == 2:
+        wav_data, sr = audio_path
+        if isinstance(wav_data, np.ndarray):
+            wav = torch.from_numpy(wav_data).float()
+        elif isinstance(wav_data, list):
+            # np.asarray + from_numpy is ~10x faster than torch.tensor(list)
+            # for the long float lists the serving layer passes through msgspec.
+            wav = torch.from_numpy(np.asarray(wav_data, dtype=np.float32))
+        elif isinstance(wav_data, torch.Tensor):
+            wav = wav_data.float()
+        else:
+            raise TypeError(f"Unsupported audio data type: {type(wav_data)}")
+        return _mono_1d(wav), int(sr)
+
+    try:
+        import torchaudio
+
+        wav, sr = torchaudio.load(audio_path)
+        return _mono_1d(wav, channels_first=True), int(sr)
+    except (ImportError, RuntimeError, OSError):
+        pass
+
+    import soundfile as sf
+
+    audio_np, sr = sf.read(audio_path, dtype="float32", always_2d=False)
+    return _mono_1d(torch.from_numpy(audio_np).float(), channels_first=False), int(sr)
+
+
+def _mono_1d(wav: torch.Tensor, channels_first: bool | None = None) -> torch.Tensor:
+    if wav.ndim == 1:
+        return wav.contiguous()
+    if channels_first is True:
+        return wav.mean(dim=0).contiguous()
+    if channels_first is False:
+        return wav.mean(dim=-1).contiguous()
+    if wav.shape[0] <= 8 and wav.shape[-1] > wav.shape[0]:
+        return wav.mean(dim=0).contiguous()
+    return wav.mean(dim=-1).contiguous()
+
+
+def _truncate_audio(
+    wav: torch.Tensor,
+    sample_rate: int,
+    max_audio_length_seconds: float | None,
+) -> torch.Tensor:
+    if max_audio_length_seconds is None:
+        return wav
+    max_audio_samples = int(max_audio_length_seconds * sample_rate)
+    if max_audio_samples > 0 and wav.shape[-1] > max_audio_samples:
+        return wav[..., :max_audio_samples]
+    return wav
+
+
+def _resample(wav: torch.Tensor, orig_sr: int, target_sr: int) -> torch.Tensor:
+    if orig_sr == target_sr:
+        return wav
+    try:
+        import torchaudio
+
+        return torchaudio.functional.resample(wav, orig_sr, target_sr)
+    except (ImportError, RuntimeError, OSError):
+        return (
+            F.interpolate(
+                wav.unsqueeze(0).unsqueeze(0), scale_factor=target_sr / orig_sr, mode="linear", align_corners=False
+            )
+            .squeeze(0)
+            .squeeze(0)
+        )

vllm_omni/model_executor/models/indextts2/prompt_utils.py

@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Prompt helpers for IndexTTS2 talker prefill."""
+
+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.tokenizer import IndexTTS2Tokenizer
+
+_CONDITIONING_PREFIX_TOKENS = 34
+_TEXT_WRAPPER_TOKENS = 2
+_START_MEL_TOKENS = 1
+
+
+def _resolve_bpe_model_path(model_id_or_path: str) -> str:
+    path = resolve_model_file(model_id_or_path, "bpe.model")
+    if path is None:
+        raise FileNotFoundError(f"Could not resolve bpe.model for {model_id_or_path!r}")
+    return path
+
+
+@lru_cache(maxsize=16)
+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:
+    """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.
+    """
+    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
+
+
+def build_indextts2_prefill_prompt_ids(
+    model_id_or_path: str,
+    text: str,
+    *,
+    placeholder_token_id: int = 1,
+) -> list[int]:
+    return [placeholder_token_id] * estimate_indextts2_prefill_prompt_len(model_id_or_path, text)

vllm_omni/model_executor/models/indextts2/s2mel/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/model_executor/models/indextts2/s2mel/modules/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/model_executor/models/indextts2/s2mel/modules/audio.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import numpy as np
+import torch
+
+
+def dynamic_range_compression_torch(x, compression=1, clip_val=1e-5):
+    return torch.log(torch.clamp(x, min=clip_val) * compression)
+
+
+def spectral_normalize_torch(magnitudes):
+    return dynamic_range_compression_torch(magnitudes)
+
+
+def _hz_to_mel(freqs: np.ndarray | float) -> np.ndarray | float:
+    """librosa's default Slaney mel scale (htk=False)."""
+    scalar = not hasattr(freqs, "__len__")
+    freqs = np.atleast_1d(np.asarray(freqs, dtype=np.float64))
+    f_min = 0.0
+    f_sp = 200.0 / 3
+    mels = (freqs - f_min) / f_sp
+
+    min_log_hz = 1000.0
+    min_log_mel = (min_log_hz - f_min) / f_sp
+    logstep = np.log(6.4) / 27.0
+    log_t = freqs >= min_log_hz
+    mels[log_t] = min_log_mel + np.log(freqs[log_t] / min_log_hz) / logstep
+    return float(mels[0]) if scalar else mels
+
+
+def _mel_to_hz(mels: np.ndarray | float) -> np.ndarray | float:
+    """Inverse of librosa's default Slaney mel scale (htk=False)."""
+    scalar = not hasattr(mels, "__len__")
+    mels = np.atleast_1d(np.asarray(mels, dtype=np.float64))
+    f_min = 0.0
+    f_sp = 200.0 / 3
+    freqs = f_min + f_sp * mels
+
+    min_log_hz = 1000.0
+    min_log_mel = (min_log_hz - f_min) / f_sp
+    logstep = np.log(6.4) / 27.0
+    log_t = mels >= min_log_mel
+    freqs[log_t] = min_log_hz * np.exp(logstep * (mels[log_t] - min_log_mel))
+    return float(freqs[0]) if scalar else freqs
+
+
+def _mel_filterbank(sr: int, n_fft: int, n_mels: int, fmin: float, fmax: float | None) -> np.ndarray:
+    """Create a librosa.filters.mel-compatible Slaney filterbank."""
+    if fmax is None:
+        fmax = float(sr) / 2.0
+    mels = np.linspace(_hz_to_mel(fmin), _hz_to_mel(fmax), n_mels + 2)
+    freqs = _mel_to_hz(mels)
+    fft_freqs = np.fft.rfftfreq(n_fft, d=1.0 / sr)
+
+    fdiff = np.diff(freqs)
+    ramps = np.subtract.outer(freqs, fft_freqs)
+    weights = np.zeros((n_mels, len(fft_freqs)), dtype=np.float32)
+    for i in range(n_mels):
+        lower = -ramps[i] / fdiff[i]
+        upper = ramps[i + 2] / fdiff[i + 1]
+        weights[i] = np.maximum(0, np.minimum(lower, upper))
+
+    enorm = 2.0 / (freqs[2 : n_mels + 2] - freqs[:n_mels])
+    weights *= enorm[:, np.newaxis]
+    return weights
+
+
+mel_basis = {}
+hann_window = {}
+
+
+def mel_spectrogram(y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False):
+    global mel_basis, hann_window  # pylint: disable=global-statement
+    if f"{sampling_rate}_{fmax}_{y.device}" not in mel_basis:
+        mel = _mel_filterbank(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax)
+        mel_basis[str(sampling_rate) + "_" + str(fmax) + "_" + str(y.device)] = (
+            torch.from_numpy(mel).float().to(y.device)
+        )
+        hann_window[str(sampling_rate) + "_" + str(y.device)] = torch.hann_window(win_size).to(y.device)
+
+    y = torch.nn.functional.pad(
+        y.unsqueeze(1), (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)), mode="reflect"
+    )
+    y = y.squeeze(1)
+
+    spec = torch.view_as_real(
+        torch.stft(
+            y,
+            n_fft,
+            hop_length=hop_size,
+            win_length=win_size,
+            window=hann_window[str(sampling_rate) + "_" + str(y.device)],
+            center=center,
+            pad_mode="reflect",
+            normalized=False,
+            onesided=True,
+            return_complex=True,
+        )
+    )
+
+    spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9))
+
+    spec = torch.matmul(mel_basis[str(sampling_rate) + "_" + str(fmax) + "_" + str(y.device)], spec)
+    spec = spectral_normalize_torch(spec)
+
+    return spec

vllm_omni/model_executor/models/indextts2/s2mel/modules/bigvgan.py

@@ -0,0 +1,359 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.
+#   Licensed under the MIT license.
+
+# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
+#   LICENSE is in incl_licenses directory.
+
+import json
+import os
+
+import torch
+import torch.nn as nn
+from huggingface_hub import PyTorchModelHubMixin, hf_hub_download
+from torch.nn import Conv1d, ConvTranspose1d
+from torch.nn.utils import remove_weight_norm, weight_norm
+
+from vllm_omni.model_executor.models.common.alias_free_activation import AliasFreeActivation1d
+from vllm_omni.model_executor.models.common.snake_activation import Snake, SnakeBeta
+from vllm_omni.model_executor.models.indextts2.s2mel.modules.commons import AttrDict
+
+# ---------------------------------------------------------------------------
+# Helpers (inlined from env.py / utils.py)
+# ---------------------------------------------------------------------------
+
+
+def load_hparams_from_json(path) -> AttrDict:
+    with open(path) as f:
+        data = f.read()
+    return AttrDict(json.loads(data))
+
+
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+# ---------------------------------------------------------------------------
+# AMP Residual Blocks
+# ---------------------------------------------------------------------------
+
+
+class AMPBlock1(torch.nn.Module):
+    """AMPBlock with Snake/SnakeBeta activations and extra convs2 (dilation=1)."""
+
+    def __init__(
+        self,
+        h: AttrDict,
+        channels: int,
+        kernel_size: int = 3,
+        dilation: tuple = (1, 3, 5),
+        activation: str = None,
+    ):
+        super().__init__()
+
+        self.h = h
+
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        stride=1,
+                        dilation=d,
+                        padding=get_padding(kernel_size, d),
+                    )
+                )
+                for d in dilation
+            ]
+        )
+        self.convs1.apply(init_weights)
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        stride=1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                )
+                for _ in range(len(dilation))
+            ]
+        )
+        self.convs2.apply(init_weights)
+
+        self.num_layers = len(self.convs1) + len(self.convs2)
+
+        if activation == "snake":
+            self.activations = nn.ModuleList(
+                [
+                    AliasFreeActivation1d(activation=Snake(channels, alpha_logscale=h.snake_logscale))
+                    for _ in range(self.num_layers)
+                ]
+            )
+        elif activation == "snakebeta":
+            self.activations = nn.ModuleList(
+                [
+                    AliasFreeActivation1d(activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale))
+                    for _ in range(self.num_layers)
+                ]
+            )
+        else:
+            raise NotImplementedError(
+                "activation incorrectly specified. check the config file and look for 'activation'."
+            )
+
+    def forward(self, x):
+        acts1, acts2 = self.activations[::2], self.activations[1::2]
+        for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2):
+            xt = a1(x)
+            xt = c1(xt)
+            xt = a2(xt)
+            xt = c2(xt)
+            x = xt + x
+
+        return x
+
+    def remove_weight_norm(self):
+        for layer in self.convs1:
+            remove_weight_norm(layer)
+        for layer in self.convs2:
+            remove_weight_norm(layer)
+
+
+class AMPBlock2(torch.nn.Module):
+    """AMPBlock with Snake/SnakeBeta activations, no extra convs2."""
+
+    def __init__(
+        self,
+        h: AttrDict,
+        channels: int,
+        kernel_size: int = 3,
+        dilation: tuple = (1, 3, 5),
+        activation: str = None,
+    ):
+        super().__init__()
+
+        self.h = h
+
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        stride=1,
+                        dilation=d,
+                        padding=get_padding(kernel_size, d),
+                    )
+                )
+                for d in dilation
+            ]
+        )
+        self.convs.apply(init_weights)
+
+        self.num_layers = len(self.convs)
+
+        if activation == "snake":
+            self.activations = nn.ModuleList(
+                [
+                    AliasFreeActivation1d(activation=Snake(channels, alpha_logscale=h.snake_logscale))
+                    for _ in range(self.num_layers)
+                ]
+            )
+        elif activation == "snakebeta":
+            self.activations = nn.ModuleList(
+                [
+                    AliasFreeActivation1d(activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale))
+                    for _ in range(self.num_layers)
+                ]
+            )
+        else:
+            raise NotImplementedError(
+                "activation incorrectly specified. check the config file and look for 'activation'."
+            )
+
+    def forward(self, x):
+        for c, a in zip(self.convs, self.activations):
+            xt = a(x)
+            xt = c(xt)
+            x = xt + x
+
+        return x
+
+    def remove_weight_norm(self):
+        for layer in self.convs:
+            remove_weight_norm(layer)
+
+
+# ---------------------------------------------------------------------------
+# BigVGAN Generator
+# ---------------------------------------------------------------------------
+
+
+class BigVGAN(
+    torch.nn.Module,
+    PyTorchModelHubMixin,
+    library_name="bigvgan",
+    repo_url="https://github.com/NVIDIA/BigVGAN",
+    docs_url="https://github.com/NVIDIA/BigVGAN/blob/main/README.md",
+    pipeline_tag="audio-to-audio",
+    license="mit",
+    tags=["neural-vocoder", "audio-generation", "arxiv:2206.04658"],
+):
+    """BigVGAN neural vocoder with anti-aliased periodic activation."""
+
+    def __init__(self, h: AttrDict, use_cuda_kernel: bool = False):
+        super().__init__()
+        self.h = h
+
+        self.num_kernels = len(h.resblock_kernel_sizes)
+        self.num_upsamples = len(h.upsample_rates)
+
+        # Pre-conv
+        self.conv_pre = weight_norm(Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3))
+
+        # Define which AMPBlock to use. BigVGAN uses AMPBlock1 as default
+        if h.resblock == "1":
+            resblock_class = AMPBlock1
+        elif h.resblock == "2":
+            resblock_class = AMPBlock2
+        else:
+            raise ValueError(f"Incorrect resblock class specified in hyperparameters. Got {h.resblock}")
+
+        # Transposed conv-based upsamplers. does not apply anti-aliasing
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
+            self.ups.append(
+                nn.ModuleList(
+                    [
+                        weight_norm(
+                            ConvTranspose1d(
+                                h.upsample_initial_channel // (2**i),
+                                h.upsample_initial_channel // (2 ** (i + 1)),
+                                k,
+                                u,
+                                padding=(k - u) // 2,
+                            )
+                        )
+                    ]
+                )
+            )
+
+        # Residual blocks using anti-aliased multi-periodicity composition modules (AMP)
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = h.upsample_initial_channel // (2 ** (i + 1))
+            for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)):
+                self.resblocks.append(resblock_class(h, ch, k, d, activation=h.activation))
+
+        # Post-conv
+        activation_post = (
+            Snake(ch, alpha_logscale=h.snake_logscale)
+            if h.activation == "snake"
+            else (SnakeBeta(ch, alpha_logscale=h.snake_logscale) if h.activation == "snakebeta" else None)
+        )
+        if activation_post is None:
+            raise NotImplementedError(
+                "activation incorrectly specified. check the config file and look for 'activation'."
+            )
+
+        self.activation_post = AliasFreeActivation1d(activation=activation_post)
+
+        # Whether to use bias for the final conv_post
+        self.use_bias_at_final = h.get("use_bias_at_final", True)
+        self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3, bias=self.use_bias_at_final))
+
+        # Weight initialization
+        for i in range(len(self.ups)):
+            self.ups[i].apply(init_weights)
+        self.conv_post.apply(init_weights)
+
+        # Final tanh activation
+        self.use_tanh_at_final = h.get("use_tanh_at_final", True)
+
+    def forward(self, x):
+        # Pre-conv
+        x = self.conv_pre(x)
+
+        for i in range(self.num_upsamples):
+            # Upsampling
+            for i_up in range(len(self.ups[i])):
+                x = self.ups[i][i_up](x)
+            # AMP blocks
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+
+        # Post-conv
+        x = self.activation_post(x)
+        x = self.conv_post(x)
+        # Final tanh activation
+        if self.use_tanh_at_final:
+            x = torch.tanh(x)
+        else:
+            x = torch.clamp(x, min=-1.0, max=1.0)
+
+        return x
+
+    def remove_weight_norm(self):
+        try:
+            for up_layer in self.ups:
+                for sub in up_layer:
+                    remove_weight_norm(sub)
+            for block in self.resblocks:
+                block.remove_weight_norm()
+            remove_weight_norm(self.conv_pre)
+            remove_weight_norm(self.conv_post)
+        except ValueError:
+            pass
+
+    @classmethod
+    def _from_pretrained(cls, *, model_id: str, map_location: str = "cpu", **kwargs):
+        def _resolve(filename):
+            if os.path.isdir(model_id):
+                return os.path.join(model_id, filename)
+            return hf_hub_download(
+                repo_id=model_id,
+                filename=filename,
+                **{
+                    k: kwargs[k]
+                    for k in (
+                        "revision",
+                        "cache_dir",
+                        "force_download",
+                        "proxies",
+                        "resume_download",
+                        "token",
+                        "local_files_only",
+                    )
+                    if k in kwargs
+                },
+            )
+
+        h = load_hparams_from_json(_resolve("config.json"))
+        model = cls(h)
+
+        checkpoint_dict = torch.load(_resolve("bigvgan_generator.pt"), map_location=map_location)
+        try:
+            model.load_state_dict(checkpoint_dict["generator"], strict=False)
+        except RuntimeError:
+            model.remove_weight_norm()
+            model.load_state_dict(checkpoint_dict["generator"], strict=False)
+
+        return model

vllm_omni/model_executor/models/indextts2/s2mel/modules/commons.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from torch import nn
+
+
+class AttrDict(dict):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.__dict__ = self
+
+
+def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
+    n_channels_int = n_channels if isinstance(n_channels, int) else n_channels[0]
+    in_act = input_a + input_b
+    t_act_part, s_act_part = torch.split(in_act, n_channels_int, dim=1)
+    t_act = torch.tanh(t_act_part)
+    s_act = torch.sigmoid(s_act_part)
+    acts = t_act * s_act
+    return acts
+
+
+def sequence_mask(length, max_length=None):
+    if max_length is None:
+        max_length = length.max()
+    x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+    return x.unsqueeze(0) < length.unsqueeze(1)
+
+
+class MyModel(nn.Module):
+    def __init__(self, args, use_emovec=False, use_gpt_latent=False):
+        super().__init__()
+        from vllm_omni.model_executor.models.indextts2.s2mel.modules.flow_matching import CFM
+        from vllm_omni.model_executor.models.indextts2.s2mel.modules.length_regulator import (
+            InterpolateRegulator,
+        )
+
+        length_regulator = InterpolateRegulator(
+            channels=args.length_regulator.channels,
+            sampling_ratios=args.length_regulator.sampling_ratios,
+            is_discrete=args.length_regulator.is_discrete,
+            in_channels=getattr(args.length_regulator, "in_channels", None),
+            vector_quantize=getattr(args.length_regulator, "vector_quantize", False),
+            codebook_size=args.length_regulator.content_codebook_size,
+            n_codebooks=getattr(args.length_regulator, "n_codebooks", 1),
+            quantizer_dropout=getattr(args.length_regulator, "quantizer_dropout", 0.0),
+            f0_condition=getattr(args.length_regulator, "f0_condition", False),
+            n_f0_bins=getattr(args.length_regulator, "n_f0_bins", 512),
+        )
+
+        if use_gpt_latent:
+            self.models = nn.ModuleDict(
+                {
+                    "cfm": CFM(args),
+                    "length_regulator": length_regulator,
+                    "gpt_layer": torch.nn.Sequential(
+                        torch.nn.Linear(1280, 256),
+                        torch.nn.Linear(256, 128),
+                        torch.nn.Linear(128, 1024),
+                    ),
+                }
+            )
+        else:
+            self.models = nn.ModuleDict({"cfm": CFM(args), "length_regulator": length_regulator})
+
+    def forward_gpt(self, x):
+        x = self.models["gpt_layer"](x)
+        return x
+
+    def enable_torch_compile(self):
+        if "cfm" in self.models:
+            self.models["cfm"].enable_torch_compile()

vllm_omni/model_executor/models/indextts2/s2mel/modules/diffusion_transformer.py

@@ -0,0 +1,251 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import math
+
+import torch
+from torch import nn
+from torch.nn.utils import weight_norm
+
+from vllm_omni.model_executor.models.indextts2.s2mel.modules.commons import sequence_mask
+from vllm_omni.model_executor.models.indextts2.s2mel.modules.gpt_fast.model import ModelArgs, Transformer
+from vllm_omni.model_executor.models.indextts2.s2mel.modules.wavenet import WN
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+#################################################################################
+#               Embedding Layers for Timesteps and Class Labels                 #
+#################################################################################
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+        self.max_period = 10000
+        self.scale = 1000
+
+        half = frequency_embedding_size // 2
+        freqs = torch.exp(-math.log(self.max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half)
+        self.register_buffer("freqs", freqs)
+
+    def timestep_embedding(self, t):
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element.
+                          These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+
+        args = self.scale * t[:, None].float() * self.freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if self.frequency_embedding_size % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t)
+        t_emb = self.mlp(t_freq.to(self.mlp[0].weight.dtype))
+        return t_emb
+
+
+class FinalLayer(nn.Module):
+    """
+    The final layer of DiT.
+    """
+
+    def __init__(self, hidden_size, patch_size, out_channels):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = weight_norm(nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True))
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
+
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+
+class DiT(torch.nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.time_as_token = getattr(args.DiT, "time_as_token", False)
+        self.style_as_token = getattr(args.DiT, "style_as_token", False)
+        self.uvit_skip_connection = getattr(args.DiT, "uvit_skip_connection", False)
+        model_args = ModelArgs(
+            block_size=16384,  # args.DiT.block_size,
+            n_layer=args.DiT.depth,
+            n_head=args.DiT.num_heads,
+            dim=args.DiT.hidden_dim,
+            head_dim=args.DiT.hidden_dim // args.DiT.num_heads,
+            vocab_size=1024,
+            uvit_skip_connection=self.uvit_skip_connection,
+            time_as_token=self.time_as_token,
+        )
+        self.transformer = Transformer(model_args)
+        self.in_channels = args.DiT.in_channels
+        self.out_channels = args.DiT.in_channels
+        self.num_heads = args.DiT.num_heads
+
+        self.x_embedder = weight_norm(nn.Linear(args.DiT.in_channels, args.DiT.hidden_dim, bias=True))
+
+        self.content_type = args.DiT.content_type  # 'discrete' or 'continuous'
+        self.content_codebook_size = args.DiT.content_codebook_size  # for discrete content
+        self.content_dim = args.DiT.content_dim  # for continuous content
+        self.cond_projection = nn.Linear(args.DiT.content_dim, args.DiT.hidden_dim, bias=True)  # continuous content
+
+        self.is_causal = args.DiT.is_causal
+
+        self.t_embedder = TimestepEmbedder(args.DiT.hidden_dim)
+
+        input_pos = torch.arange(16384)
+        self.register_buffer("input_pos", input_pos)
+
+        self.final_layer_type = args.DiT.final_layer_type  # mlp or wavenet
+        if self.final_layer_type == "wavenet":
+            self.t_embedder2 = TimestepEmbedder(args.wavenet.hidden_dim)
+            self.conv1 = nn.Linear(args.DiT.hidden_dim, args.wavenet.hidden_dim)
+            self.conv2 = nn.Conv1d(args.wavenet.hidden_dim, args.DiT.in_channels, 1)
+            self.wavenet = WN(
+                hidden_channels=args.wavenet.hidden_dim,
+                kernel_size=args.wavenet.kernel_size,
+                dilation_rate=args.wavenet.dilation_rate,
+                n_layers=args.wavenet.num_layers,
+                gin_channels=args.wavenet.hidden_dim,
+                p_dropout=args.wavenet.p_dropout,
+                causal=False,
+            )
+            self.final_layer = FinalLayer(args.wavenet.hidden_dim, 1, args.wavenet.hidden_dim)
+            self.res_projection = nn.Linear(
+                args.DiT.hidden_dim, args.wavenet.hidden_dim
+            )  # residual connection from transformer output to final output
+            self.wavenet_style_condition = args.wavenet.style_condition
+            assert args.DiT.style_condition == args.wavenet.style_condition
+        else:
+            self.final_mlp = nn.Sequential(
+                nn.Linear(args.DiT.hidden_dim, args.DiT.hidden_dim),
+                nn.SiLU(),
+                nn.Linear(args.DiT.hidden_dim, args.DiT.in_channels),
+            )
+        self.transformer_style_condition = args.DiT.style_condition
+
+        self.class_dropout_prob = args.DiT.class_dropout_prob
+        self.content_mask_embedder = nn.Embedding(1, args.DiT.hidden_dim)
+
+        self.long_skip_connection = args.DiT.long_skip_connection
+        self.skip_linear = nn.Linear(args.DiT.hidden_dim + args.DiT.in_channels, args.DiT.hidden_dim)
+
+        self.cond_x_merge_linear = nn.Linear(
+            args.DiT.hidden_dim
+            + args.DiT.in_channels * 2
+            + args.style_encoder.dim * self.transformer_style_condition * (not self.style_as_token),
+            args.DiT.hidden_dim,
+        )
+        if self.style_as_token:
+            self.style_in = nn.Linear(args.style_encoder.dim, args.DiT.hidden_dim)
+
+    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):
+        """
+        x (torch.Tensor): random noise
+        prompt_x (torch.Tensor): reference mel + zero mel
+            shape: (batch_size, 80, 795+1068)
+        x_lens (torch.Tensor): mel frames output
+            shape: (batch_size, mel_timesteps)
+        t (torch.Tensor): radshape:
+            shape: (batch_size)
+        style (torch.Tensor): reference global style
+            shape: (batch_size, 192)
+        cond (torch.Tensor): semantic info of reference audio and altered audio
+            shape: (batch_size, mel_timesteps(795+1069), 512)
+        pre_mask (torch.Tensor | None): precomputed (x_mask, x_mask_expanded) to
+            avoid recomputation across ODE steps. Tuple of (x_mask, x_mask_expanded).
+        """
+        class_dropout = False
+        if self.training and torch.rand(1) < self.class_dropout_prob:
+            class_dropout = True
+        if not self.training and mask_content:
+            class_dropout = True
+        cond_in_module = self.cond_projection
+
+        B, _, T = x.size()
+
+        t1 = self.t_embedder(t)  # (N, D) # t1 [2, 512]
+        cond = cond_in_module(cond)  # cond [2,1863,512]->[2,1863,512]
+
+        x = x.transpose(1, 2)  # [2,1863,80]
+        prompt_x = prompt_x.transpose(1, 2)  # [2,1863,80]
+
+        x_in = torch.cat([x, prompt_x, cond], dim=-1)  # 80+80+512=672 [2, 1863, 672]
+
+        if self.transformer_style_condition and not self.style_as_token:  # True and True
+            x_in = torch.cat([x_in, style[:, None, :].expand(-1, T, -1)], dim=-1)  # [2, 1863, 864]
+
+        if class_dropout:  # False
+            x_in[..., self.in_channels :] = x_in[..., self.in_channels :] * 0  # 80维后全置为0
+
+        x_in = self.cond_x_merge_linear(x_in)  # (N, T, D) [2, 1863, 512]
+
+        if self.style_as_token:  # False
+            style = self.style_in(style)
+            style = torch.zeros_like(style) if class_dropout else style
+            x_in = torch.cat([style.unsqueeze(1), x_in], dim=1)
+
+        if self.time_as_token:  # False
+            x_in = torch.cat([t1.unsqueeze(1), x_in], dim=1)
+
+        if pre_mask is not None:
+            x_mask, x_mask_expanded = pre_mask
+        else:
+            x_mask = (
+                sequence_mask(x_lens + self.style_as_token + self.time_as_token, max_length=x_in.size(1))
+                .to(x.device)
+                .unsqueeze(1)
+            )
+            input_pos = self.input_pos[: x_in.size(1)]
+            x_mask_expanded = x_mask[:, None, :].expand(-1, -1, x_in.size(1), -1) if not self.is_causal else None
+
+        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:
+            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 = x_res[:, 1:] if self.time_as_token else x_res
+        x_res = x_res[:, 1:] if self.style_as_token else x_res
+
+        if self.long_skip_connection:  # True
+            x_res = self.skip_linear(torch.cat([x_res, x], dim=-1))
+        if self.final_layer_type == "wavenet":
+            x = self.conv1(x_res)
+            x = x.transpose(1, 2)
+            t2 = self.t_embedder2(t)
+            x = self.wavenet(x, x_mask, g=t2.unsqueeze(2)).transpose(1, 2) + self.res_projection(
+                x_res
+            )  # long residual connection
+            x = self.final_layer(x, t1).transpose(1, 2)
+            x = self.conv2(x)
+        else:
+            x = self.final_mlp(x_res)
+            x = x.transpose(1, 2)
+        # x [2,80,1863]
+        return x

vllm_omni/model_executor/models/indextts2/s2mel/modules/encodec.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Convolutional layers wrappers and utilities."""
+
+import math
+import typing as tp
+import warnings
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.utils import weight_norm
+
+
+def apply_parametrization_norm(module: nn.Module, norm: str = "none") -> nn.Module:
+    if norm == "weight_norm":
+        return weight_norm(module)
+    return module
+
+
+def get_norm_module(module: nn.Module, causal: bool = False, norm: str = "none", **norm_kwargs) -> nn.Module:
+    return nn.Identity()
+
+
+def get_extra_padding_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0) -> int:
+    length = x.shape[-1]
+    n_frames = (length - kernel_size + padding_total) / stride + 1
+    ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)
+    return ideal_length - length
+
+
+def pad1d(x: torch.Tensor, paddings: tuple[int, int], mode: str = "zero", value: float = 0.0):
+    length = x.shape[-1]
+    padding_left, padding_right = paddings
+    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+    if mode == "reflect":
+        max_pad = max(padding_left, padding_right)
+        extra_pad = 0
+        if length <= max_pad:
+            extra_pad = max_pad - length + 1
+            x = F.pad(x, (0, extra_pad))
+        padded = F.pad(x, paddings, mode, value)
+        end = padded.shape[-1] - extra_pad
+        return padded[..., :end]
+    else:
+        return F.pad(x, paddings, mode, value)
+
+
+class NormConv1d(nn.Module):
+    def __init__(self, *args, causal: bool = False, norm: str = "none", norm_kwargs: dict[str, tp.Any] = {}, **kwargs):
+        super().__init__()
+        self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.norm(x)
+        return x
+
+
+class SConv1d(nn.Module):
+    """Conv1d with builtin handling of asymmetric or causal padding and normalization."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        dilation: int = 1,
+        groups: int = 1,
+        bias: bool = True,
+        causal: bool = False,
+        norm: str = "none",
+        norm_kwargs: dict[str, tp.Any] = {},
+        pad_mode: str = "reflect",
+        **kwargs,
+    ):
+        super().__init__()
+        if stride > 1 and dilation > 1:
+            warnings.warn(
+                "SConv1d has been initialized with stride > 1 and dilation > 1"
+                f" (kernel_size={kernel_size} stride={stride}, dilation={dilation})."
+            )
+        self.conv = NormConv1d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+            causal=causal,
+            norm=norm,
+            norm_kwargs=norm_kwargs,
+        )
+        self.causal = causal
+        self.pad_mode = pad_mode
+
+    def forward(self, x):
+        B, C, T = x.shape
+        kernel_size = self.conv.conv.kernel_size[0]
+        stride = self.conv.conv.stride[0]
+        dilation = self.conv.conv.dilation[0]
+        kernel_size = (kernel_size - 1) * dilation + 1
+        padding_total = kernel_size - stride
+        extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
+        if self.causal:
+            x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode)
+        else:
+            padding_right = padding_total // 2
+            padding_left = padding_total - padding_right
+            x = pad1d(x, (padding_left, padding_right + extra_padding), mode=self.pad_mode)
+        return self.conv(x)

vllm_omni/model_executor/models/indextts2/s2mel/modules/flow_matching.py

@@ -0,0 +1,201 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from abc import ABC
+
+import torch
+
+from vllm_omni.model_executor.models.indextts2.s2mel.modules.diffusion_transformer import DiT
+
+
+class BASECFM(torch.nn.Module, ABC):
+    def __init__(
+        self,
+        args,
+    ):
+        super().__init__()
+        self.sigma_min = 1e-6
+
+        self.estimator = None
+
+        self.in_channels = args.DiT.in_channels
+
+        if hasattr(args.DiT, "zero_prompt_speech_token"):
+            self.zero_prompt_speech_token = args.DiT.zero_prompt_speech_token
+        else:
+            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):
+        """Forward diffusion
+
+        Args:
+            mu (torch.Tensor): semantic info of reference audio and altered audio
+                shape: (batch_size, mel_timesteps(795+1069), 512)
+            x_lens (torch.Tensor): mel frames output
+                shape: (batch_size, mel_timesteps)
+            prompt (torch.Tensor): reference mel
+                shape: (batch_size, 80, 795)
+            style (torch.Tensor): reference global style
+                shape: (batch_size, 192)
+            f0 (None): unused, reserved for f0 conditioning
+            n_timesteps (int): number of diffusion steps
+            temperature (float, optional): temperature for scaling noise. Defaults to 1.0.
+
+        Returns:
+            sample (torch.Tensor): generated mel-spectrogram
+                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
+        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)
+
+    def solve_euler(self, x, x_lens, prompt, mu, style, f0, t_span, inference_cfg_rate=0.5):
+        """
+        Fixed euler solver for ODEs.
+        Args:
+            x (torch.Tensor): random noise
+            t_span (torch.Tensor): n_timesteps interpolated
+                shape: (n_timesteps + 1,)
+            mu (torch.Tensor): semantic info of reference audio and altered audio
+                shape: (batch_size, mel_timesteps(795+1069), 512)
+            x_lens (torch.Tensor): mel frames output
+                shape: (batch_size, mel_timesteps)
+            prompt (torch.Tensor): reference mel
+                shape: (batch_size, 80, 795)
+            style (torch.Tensor): reference global style
+                shape: (batch_size, 192)
+        """
+        t = t_span[0]
+
+        # Optional bf16 autocast for the DiT estimator (set by the decoder via
+        # `estimator_autocast_dtype`). The Euler solver state stays float32 —
+        # only the estimator forward runs in reduced precision, which restores
+        # flash-attention eligibility and halves GEMM cost.
+        autocast_dtype = getattr(self, "estimator_autocast_dtype", None)
+
+        def _run_estimator(*args, **kw):
+            if autocast_dtype is None:
+                return self.estimator(*args, **kw)
+            with torch.autocast(x.device.type, dtype=autocast_dtype):
+                out = self.estimator(*args, **kw)
+            return out.float()
+
+        prompt_len = prompt.size(-1)
+        prompt_x = torch.zeros_like(x)
+        prompt_x[..., :prompt_len] = prompt[..., :prompt_len]
+        x[..., :prompt_len] = 0
+        if self.zero_prompt_speech_token:
+            mu[..., :prompt_len] = 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:]))
+            cfg_style[:batch].copy_(style)
+            cfg_mu = mu.new_zeros((2 * batch, *mu.shape[1:]))
+            cfg_mu[:batch].copy_(mu)
+            cfg_x = x.new_empty((2 * batch, *x.shape[1:]))
+            cfg_t = x.new_empty((2 * batch,))
+            cfg_pre_mask = self._repeat_pre_mask_for_cfg(pre_mask)
+
+        for step in range(1, len(t_span)):
+            dt = t_span[step] - t_span[step - 1]
+            if inference_cfg_rate > 0:
+                cfg_x[:batch].copy_(x)
+                cfg_x[batch:].copy_(x)
+                cfg_t.fill_(t)
+
+                stacked_dphi_dt = _run_estimator(
+                    cfg_x,
+                    cfg_prompt_x,
+                    x_lens,
+                    cfg_t,
+                    cfg_style,
+                    cfg_mu,
+                    pre_mask=cfg_pre_mask,
+                )
+
+                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)
+
+            x = x + dt * dphi_dt
+            t = t + dt
+            x[:, :, :prompt_len] = 0
+
+        return x
+
+    @staticmethod
+    def _repeat_pre_mask_for_cfg(pre_mask):
+        """Duplicate precomputed masks when CFG doubles the batch.
+
+        The original IndexTTS2 inference path used B=1, where a two-element
+        timestep tensor happened to match CFG. Stage-1 batching makes the
+        estimator input batch 2B, so all batch-indexed masks must be repeated
+        consistently with ``torch.cat([cond, uncond], dim=0)``.
+        """
+        if not isinstance(pre_mask, tuple):
+            return pre_mask
+        repeated = []
+        for mask in pre_mask:
+            if isinstance(mask, torch.Tensor):
+                repeated.append(torch.cat([mask, mask], dim=0))
+            else:
+                repeated.append(mask)
+        return tuple(repeated)
+
+    def _precompute_mask(self, x: torch.Tensor, x_lens: torch.Tensor):
+        """Precompute padding mask for DiT — reused across all ODE steps."""
+        from .commons import sequence_mask
+
+        estimator = self._eager_estimator if hasattr(self, "_eager_estimator") else self.estimator
+        T = x.size(2)
+        style_offset = getattr(estimator, "style_as_token", False)
+        time_offset = getattr(estimator, "time_as_token", False)
+        T_in = T + int(style_offset) + int(time_offset)
+        is_causal = getattr(estimator, "is_causal", False)
+
+        x_mask = sequence_mask(x_lens + int(style_offset) + int(time_offset), max_length=T_in).to(x.device).unsqueeze(1)
+        if is_causal:
+            return (x_mask, None)
+        x_mask_expanded = x_mask[:, None, :].expand(-1, -1, T_in, -1)
+        return (x_mask, x_mask_expanded)
+
+
+class CFM(BASECFM):
+    def __init__(self, args):
+        super().__init__(args)
+        if args.dit_type == "DiT":
+            self.estimator = DiT(args)
+            object.__setattr__(self, "_eager_estimator", self.estimator)
+        else:
+            raise NotImplementedError(f"Unknown diffusion type {args.dit_type}")
+        self._compiled = False
+
+    def enable_torch_compile(self):
+        """Enable torch.compile optimization for the estimator model.
+
+        This method applies torch.compile to the estimator (DiT model) for significant
+        performance improvements during inference. It also configures distributed
+        training optimizations if applicable.
+        """
+        if torch.distributed.is_initialized():
+            torch._inductor.config.reorder_for_compute_comm_overlap = True
+        self.estimator = torch.compile(
+            self._eager_estimator,
+            fullgraph=True,
+            dynamic=True,
+        )
+        self._compiled = True
+
+    def disable_torch_compile(self):
+        self.estimator = self._eager_estimator
+        self._compiled = False

vllm_omni/model_executor/models/indextts2/s2mel/modules/gpt_fast/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/model_executor/models/indextts2/s2mel/modules/gpt_fast/model.py

@@ -0,0 +1,303 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import functional as F
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.layer import Attention as DiffusionAttention
+from vllm_omni.diffusion.layers.norm import RMSNorm
+
+
+def find_multiple(n: int, k: int) -> int:
+    if n % k == 0:
+        return n
+    return n + k - (n % k)
+
+
+class AdaptiveLayerNorm(nn.Module):
+    r"""Adaptive Layer Normalization"""
+
+    def __init__(self, d_model, norm) -> None:
+        super().__init__()
+        self.project_layer = nn.Linear(d_model, 2 * d_model)
+        self.norm = norm
+        self.d_model = d_model
+        self.eps = getattr(self.norm, "eps", getattr(self.norm, "variance_epsilon", 1e-5))
+
+    def forward(self, input: Tensor, embedding: Tensor = None) -> Tensor:
+        if embedding is None:
+            return self.norm(input)
+        weight, bias = torch.split(
+            self.project_layer(embedding),
+            split_size_or_sections=self.d_model,
+            dim=-1,
+        )
+        return weight * self.norm(input) + bias
+
+
+@dataclass
+class ModelArgs:
+    block_size: int = 2048
+    vocab_size: int = 32000
+    n_layer: int = 32
+    n_head: int = 32
+    dim: int = 4096
+    intermediate_size: int = None
+    n_local_heads: int = -1
+    head_dim: int = 64
+    rope_base: float = 10000
+    norm_eps: float = 1e-5
+    has_cross_attention: bool = False
+    context_dim: int = 0
+    uvit_skip_connection: bool = False
+    time_as_token: bool = False
+
+    def __post_init__(self):
+        if self.n_local_heads == -1:
+            self.n_local_heads = self.n_head
+        if self.intermediate_size is None:
+            hidden_dim = 4 * self.dim
+            n_hidden = int(2 * hidden_dim / 3)
+            self.intermediate_size = find_multiple(n_hidden, 256)
+
+
+class Transformer(nn.Module):
+    def __init__(self, config: ModelArgs) -> None:
+        super().__init__()
+        self.config = config
+
+        self.layers = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layer))
+        self.norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+
+        self.freqs_cis: Tensor | None = None
+        self.mask_cache: Tensor | None = None
+        self.max_batch_size = -1
+        self.max_seq_length = -1
+
+    def setup_caches(self, max_batch_size, max_seq_length, use_kv_cache=False):
+        if self.max_seq_length >= max_seq_length and self.max_batch_size >= max_batch_size:
+            return
+        max_seq_length = find_multiple(max_seq_length, 8)
+        self.max_seq_length = max_seq_length
+        self.max_batch_size = max_batch_size
+        dtype = self.norm.project_layer.weight.dtype
+        device = self.norm.project_layer.weight.device
+
+        self.freqs_cis = precompute_freqs_cis(
+            self.config.block_size, self.config.head_dim, self.config.rope_base, dtype
+        ).to(device)
+        self.causal_mask = torch.tril(torch.ones(self.max_seq_length, self.max_seq_length, dtype=torch.bool)).to(device)
+        self.use_kv_cache = use_kv_cache
+        self.uvit_skip_connection = self.config.uvit_skip_connection
+        if self.uvit_skip_connection:
+            self.layers_emit_skip = [i for i in range(self.config.n_layer) if i < self.config.n_layer // 2]
+            self.layers_receive_skip = [i for i in range(self.config.n_layer) if i > self.config.n_layer // 2]
+        else:
+            self.layers_emit_skip = []
+            self.layers_receive_skip = []
+
+    def forward(
+        self,
+        x: Tensor,
+        c: Tensor,
+        input_pos: Tensor | None = None,
+        mask: Tensor | None = None,
+        context: Tensor | None = None,
+        context_input_pos: Tensor | None = None,
+        cross_attention_mask: 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
+            if not self.training and self.use_kv_cache:
+                mask = self.causal_mask[None, None, input_pos]
+            else:
+                mask = self.causal_mask[None, None, input_pos]
+                mask = mask[..., input_pos]
+        freqs_cis = self.freqs_cis[input_pos]
+        if context is not None:
+            context_freqs_cis = self.freqs_cis[context_input_pos]
+        else:
+            context_freqs_cis = None
+        skip_in_x_list = []
+        for i, layer in enumerate(self.layers):
+            if self.uvit_skip_connection and i in self.layers_receive_skip:
+                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)
+            if self.uvit_skip_connection and i in self.layers_emit_skip:
+                skip_in_x_list.append(x)
+        x = self.norm(x, c)
+        return x
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config: ModelArgs) -> None:
+        super().__init__()
+        self.attention = Attention(config)
+        self.feed_forward = FeedForward(config)
+        self.ffn_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+        self.attention_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+
+        if config.uvit_skip_connection:
+            self.skip_in_linear = nn.Linear(config.dim * 2, config.dim)
+            self.uvit_skip_connection = True
+        else:
+            self.uvit_skip_connection = False
+
+        self.time_as_token = config.time_as_token
+
+    def forward(
+        self,
+        x: Tensor,
+        c: Tensor,
+        input_pos: Tensor,
+        freqs_cis: Tensor,
+        mask: Tensor,
+        context: Tensor | None = None,
+        context_freqs_cis: Tensor | None = None,
+        cross_attention_mask: Tensor | None = None,
+        skip_in_x: 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)
+        out = h + self.feed_forward(self.ffn_norm(h, c))
+        return out
+
+
+class Attention(nn.Module):
+    def __init__(self, config: ModelArgs, is_cross_attention: bool = False):
+        super().__init__()
+        assert config.dim % config.n_head == 0
+
+        total_head_dim = (config.n_head + 2 * config.n_local_heads) * config.head_dim
+        # key, query, value projections for all heads, but in a batch
+        if is_cross_attention:
+            self.wq = nn.Linear(config.dim, config.n_head * config.head_dim, bias=False)
+            self.wkv = nn.Linear(config.context_dim, 2 * config.n_local_heads * config.head_dim, bias=False)
+        else:
+            self.wqkv = nn.Linear(config.dim, total_head_dim, bias=False)
+        self.wo = nn.Linear(config.head_dim * config.n_head, config.dim, bias=False)
+        self.kv_cache = None
+
+        self.n_head = config.n_head
+        self.head_dim = config.head_dim
+        self.n_local_heads = config.n_local_heads
+        self.dim = config.dim
+
+        # vLLM-Omni diffusion attention backend.
+        # Auto-selects FlashAttention / SageAttention / SDPA based on
+        # DIFFUSION_ATTENTION_BACKEND env var or platform default.
+        self.diff_attn = DiffusionAttention(
+            num_heads=config.n_head,
+            head_size=config.head_dim,
+            softmax_scale=1.0 / math.sqrt(config.head_dim),
+            causal=config.time_as_token,
+            num_kv_heads=config.n_local_heads,
+            role="cross" if is_cross_attention else "self",
+        )
+
+    def forward(
+        self,
+        x: Tensor,
+        freqs_cis: Tensor,
+        mask: Tensor,
+        input_pos: Tensor | None = None,
+        context: Tensor | None = None,
+        context_freqs_cis: Tensor | None = None,
+    ) -> Tensor:
+        bsz, seqlen, _ = x.shape
+
+        kv_size = self.n_local_heads * self.head_dim
+        if context is None:
+            q, k, v = self.wqkv(x).split([kv_size, kv_size, kv_size], dim=-1)
+            context_seqlen = seqlen
+        else:
+            q = self.wq(x)
+            k, v = self.wkv(context).split([kv_size, kv_size], dim=-1)
+            context_seqlen = context.shape[1]
+
+        q = q.view(bsz, seqlen, self.n_head, self.head_dim)
+        k = k.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)
+        v = v.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)
+
+        q = apply_rotary_emb(q, freqs_cis)
+        k = apply_rotary_emb(k, context_freqs_cis if context_freqs_cis is not None else freqs_cis)
+
+        # vLLM-Omni diffusion attention backend.
+        # Expects (B, S, H, D) layout — skip the transpose.
+        # GQA handled internally via num_kv_heads init param.
+        # KV cache is not used by the S2Mel DiT, so skip it here.
+        #
+        # Mask format: diffusion_transformer.py passes a 4D padding mask
+        # (B, 1, T, T). The FlashAttention backend requires 2D (B, T).
+        # Squeeze to 2D for universal compatibility.
+        #
+        # Fast path: when the caller guarantees no padding (x_lens == T,
+        # the S2Mel CFM always runs full-length sequences), the mask is
+        # all-True and dense attention is mathematically identical. Skip
+        # 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):
+            attn_metadata = None
+        else:
+            attn_mask = mask
+            if attn_mask is not None and attn_mask.dim() == 4:
+                # (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)
+        y = y.flatten(2, 3)
+
+        y = self.wo(y)
+        return y
+
+
+class FeedForward(nn.Module):
+    def __init__(self, config: ModelArgs) -> None:
+        super().__init__()
+        self.w1 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+        self.w3 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+        self.w2 = nn.Linear(config.intermediate_size, config.dim, bias=False)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+def precompute_freqs_cis(seq_len: int, n_elem: int, base: int = 10000, dtype: torch.dtype = torch.bfloat16) -> Tensor:
+    freqs = 1.0 / (base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem))
+    t = torch.arange(seq_len, device=freqs.device)
+    freqs = torch.outer(t, freqs)
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+    cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1)
+    return cache.to(dtype=dtype)
+
+
+def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
+    xshaped = x.float().reshape(*x.shape[:-1], -1, 2)
+    freqs_cis = freqs_cis.view(1, xshaped.size(1), 1, xshaped.size(3), 2)
+    x_out2 = torch.stack(
+        [
+            xshaped[..., 0] * freqs_cis[..., 0] - xshaped[..., 1] * freqs_cis[..., 1],
+            xshaped[..., 1] * freqs_cis[..., 0] + xshaped[..., 0] * freqs_cis[..., 1],
+        ],
+        -1,
+    )
+
+    x_out2 = x_out2.flatten(3)
+    return x_out2.type_as(x)

vllm_omni/model_executor/models/indextts2/s2mel/modules/length_regulator.py

@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from vllm_omni.model_executor.models.indextts2.s2mel.modules.commons import sequence_mask
+
+# f0_bin = 256
+f0_max = 1100.0
+f0_min = 50.0
+f0_mel_min = 1127 * np.log(1 + f0_min / 700)
+f0_mel_max = 1127 * np.log(1 + f0_max / 700)
+
+
+def f0_to_coarse(f0, f0_bin):
+    f0_mel = 1127 * (1 + f0 / 700).log()
+    a = (f0_bin - 2) / (f0_mel_max - f0_mel_min)
+    b = f0_mel_min * a - 1.0
+    f0_mel = torch.where(f0_mel > 0, f0_mel * a - b, f0_mel)
+    f0_coarse = torch.round(f0_mel).long()
+    f0_coarse = f0_coarse * (f0_coarse > 0)
+    f0_coarse = f0_coarse + ((f0_coarse < 1) * 1)
+    f0_coarse = f0_coarse * (f0_coarse < f0_bin)
+    f0_coarse = f0_coarse + ((f0_coarse >= f0_bin) * (f0_bin - 1))
+    return f0_coarse
+
+
+class InterpolateRegulator(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        sampling_ratios: tuple,
+        is_discrete: bool = False,
+        in_channels: int = None,  # only applies to continuous input
+        vector_quantize: bool = False,  # whether to use vector quantization, only applies to continuous input
+        codebook_size: int = 1024,  # for discrete only
+        out_channels: int = None,
+        groups: int = 1,
+        n_codebooks: int = 1,  # number of codebooks
+        quantizer_dropout: float = 0.0,  # dropout for quantizer
+        f0_condition: bool = False,
+        n_f0_bins: int = 512,
+    ):
+        super().__init__()
+        self.sampling_ratios = sampling_ratios
+        out_channels = out_channels or channels
+        model = nn.ModuleList([])
+        if len(sampling_ratios) > 0:
+            self.interpolate = True
+            for _ in sampling_ratios:
+                module = nn.Conv1d(channels, channels, 3, 1, 1)
+                norm = nn.GroupNorm(groups, channels)
+                act = nn.Mish()
+                model.extend([module, norm, act])
+        else:
+            self.interpolate = False
+        model.append(nn.Conv1d(channels, out_channels, 1, 1))
+        self.model = nn.Sequential(*model)
+        self.embedding = nn.Embedding(codebook_size, channels)
+        self.is_discrete = is_discrete
+
+        self.mask_token = nn.Parameter(torch.zeros(1, channels))
+
+        self.n_codebooks = n_codebooks
+        if n_codebooks > 1:
+            self.extra_codebooks = nn.ModuleList(
+                [nn.Embedding(codebook_size, channels) for _ in range(n_codebooks - 1)]
+            )
+            self.extra_codebook_mask_tokens = nn.ParameterList(
+                [nn.Parameter(torch.zeros(1, channels)) for _ in range(n_codebooks - 1)]
+            )
+        self.quantizer_dropout = quantizer_dropout
+
+        if f0_condition:
+            self.f0_embedding = nn.Embedding(n_f0_bins, channels)
+            self.f0_condition = f0_condition
+            self.n_f0_bins = n_f0_bins
+            self.f0_bins = torch.arange(2, 1024, 1024 // n_f0_bins)
+            self.f0_mask = nn.Parameter(torch.zeros(1, channels))
+        else:
+            self.f0_condition = False
+
+        if not is_discrete:
+            self.content_in_proj = nn.Linear(in_channels, channels)
+            if vector_quantize:
+                from dac.nn.quantize import VectorQuantize
+
+                self.vq = VectorQuantize(channels, codebook_size, 8)
+
+    def forward(self, x, ylens=None, n_quantizers=None, f0=None):
+        # apply token drop
+        if self.training:
+            n_quantizers = torch.ones((x.shape[0],)) * self.n_codebooks
+            dropout = torch.randint(1, self.n_codebooks + 1, (x.shape[0],))
+            n_dropout = int(x.shape[0] * self.quantizer_dropout)
+            n_quantizers[:n_dropout] = dropout[:n_dropout]
+            n_quantizers = n_quantizers.to(x.device)
+            # decide whether to drop for each sample in batch
+        else:
+            n_quantizers = torch.ones((x.shape[0],), device=x.device) * (
+                self.n_codebooks if n_quantizers is None else n_quantizers
+            )
+        if self.is_discrete:
+            if self.n_codebooks > 1:
+                assert len(x.size()) == 3
+                x_emb = self.embedding(x[:, 0])
+                for i, emb in enumerate(self.extra_codebooks):
+                    x_emb = x_emb + (n_quantizers > i + 1)[..., None, None] * emb(x[:, i + 1])
+                    # add mask token if not using this codebook
+                    # x_emb = x_emb + (n_quantizers <= i+1)[..., None, None] * self.extra_codebook_mask_tokens[i]
+                x = x_emb
+            elif self.n_codebooks == 1:
+                if len(x.size()) == 2:
+                    x = self.embedding(x)
+                else:
+                    x = self.embedding(x[:, 0])
+        else:
+            x = self.content_in_proj(x)
+        # x in (B, T, D)
+        mask = sequence_mask(ylens).unsqueeze(-1)
+        if self.interpolate:
+            x = F.interpolate(x.transpose(1, 2).contiguous(), size=ylens.max(), mode="nearest")
+        else:
+            x = x.transpose(1, 2).contiguous()
+            mask = mask[:, : x.size(2), :]
+            ylens = ylens.clamp(max=x.size(2)).long()
+        if self.f0_condition:
+            if f0 is None:
+                x = x + self.f0_mask.unsqueeze(-1)
+            else:
+                # quantized_f0 = torch.bucketize(f0, self.f0_bins.to(f0.device))  # (N, T)
+                quantized_f0 = f0_to_coarse(f0, self.n_f0_bins)
+                quantized_f0 = quantized_f0.clamp(0, self.n_f0_bins - 1).long()
+                f0_emb = self.f0_embedding(quantized_f0)
+                f0_emb = F.interpolate(f0_emb.transpose(1, 2).contiguous(), size=ylens.max(), mode="nearest")
+                x = x + f0_emb
+        out = self.model(x).transpose(1, 2).contiguous()
+        if hasattr(self, "vq"):
+            (
+                out_q,
+                commitment_loss,
+                codebook_loss,
+                codes,
+                out,
+            ) = self.vq(out.transpose(1, 2))
+            out_q = out_q.transpose(1, 2)
+            return out_q * mask, ylens, codes, commitment_loss, codebook_loss
+        olens = ylens
+        return out * mask, olens, None, None, None

vllm_omni/model_executor/models/indextts2/s2mel/modules/wavenet.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from torch import nn
+
+from vllm_omni.model_executor.models.indextts2.s2mel.modules.encodec import SConv1d
+
+from . import commons
+
+
+class WN(torch.nn.Module):
+    def __init__(
+        self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0, causal=False
+    ):
+        super().__init__()
+        conv1d_type = SConv1d
+        assert kernel_size % 2 == 1
+        self.hidden_channels = hidden_channels
+        self.kernel_size = (kernel_size,)
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.gin_channels = gin_channels
+        self.p_dropout = p_dropout
+
+        self.in_layers = torch.nn.ModuleList()
+        self.res_skip_layers = torch.nn.ModuleList()
+        self.drop = nn.Dropout(p_dropout)
+
+        if gin_channels != 0:
+            self.cond_layer = conv1d_type(gin_channels, 2 * hidden_channels * n_layers, 1, norm="weight_norm")
+
+        for i in range(n_layers):
+            dilation = dilation_rate**i
+            padding = int((kernel_size * dilation - dilation) / 2)
+            in_layer = conv1d_type(
+                hidden_channels,
+                2 * hidden_channels,
+                kernel_size,
+                dilation=dilation,
+                padding=padding,
+                norm="weight_norm",
+                causal=causal,
+            )
+            self.in_layers.append(in_layer)
+
+            # last one is not necessary
+            if i < n_layers - 1:
+                res_skip_channels = 2 * hidden_channels
+            else:
+                res_skip_channels = hidden_channels
+
+            res_skip_layer = conv1d_type(hidden_channels, res_skip_channels, 1, norm="weight_norm", causal=causal)
+            self.res_skip_layers.append(res_skip_layer)
+
+    def forward(self, x, x_mask, g=None, **kwargs):
+        output = torch.zeros_like(x)
+
+        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 g is not None:
+                cond_offset = i * 2 * self.hidden_channels
+                g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :]
+            else:
+                g_l = torch.zeros_like(x_in)
+
+            acts = commons.fused_add_tanh_sigmoid_multiply(x_in, g_l, self.hidden_channels)
+            acts = self.drop(acts)
+
+            res_skip_acts = self.res_skip_layers[i](acts)
+            if i < self.n_layers - 1:
+                res_acts = res_skip_acts[:, : self.hidden_channels, :]
+                x = (x + res_acts) * x_mask
+                output = output + res_skip_acts[:, self.hidden_channels :, :]
+            else:
+                output = output + res_skip_acts
+        return output * x_mask
+
+    def remove_weight_norm(self):
+        if self.gin_channels != 0:
+            torch.nn.utils.remove_weight_norm(self.cond_layer)
+        for conv_layer in self.in_layers:
+            torch.nn.utils.remove_weight_norm(conv_layer)
+        for conv_layer in self.res_skip_layers:
+            torch.nn.utils.remove_weight_norm(conv_layer)

vllm_omni/model_executor/models/indextts2/tokenizer.py

@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+
+from transformers import PreTrainedTokenizer
+from vllm.logger import init_logger
+
+from vllm_omni.model_executor.models.indextts2.utils.front import TextTokenizer
+
+logger = init_logger(__name__)
+
+
+class IndexTTS2Tokenizer(PreTrainedTokenizer):
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        vocab_file = kwargs.pop("vocab_file", None)
+        if vocab_file is None:
+            vocab_file = os.path.join(pretrained_model_name_or_path, "bpe.model")
+        logger.info(f"IndexTTS2Tokenizer.from_pretrained: {pretrained_model_name_or_path}")
+        return cls(vocab_file, **kwargs)
+
+    def __init__(self, vocab_file: str, **kwargs):
+        logger.info(f"IndexTTS2Tokenizer.__init__: vocab_file={vocab_file}")
+        self.vocab_file = vocab_file
+        self._tok = TextTokenizer(vocab_file)
+        logger.info("IndexTTS2Tokenizer initialized successfully")
+        super().__init__(**kwargs)
+
+    @property
+    def vocab_size(self):
+        return self._tok.vocab_size
+
+    def get_vocab(self):
+        return self._tok.get_vocab()
+
+    def _tokenize(self, text):
+        return self._tok.tokenize(text)
+
+    def _convert_token_to_id(self, token):
+        return self._tok.convert_tokens_to_ids(token)[0]
+
+    def convert_tokens_to_string(self, tokens):
+        return self._tok.decode(self._tok.convert_tokens_to_ids(tokens))
+
+    def split_segments(
+        self,
+        tokenized: list[str],
+        max_text_tokens_per_segment: int = 120,
+        quick_streaming_tokens: int = 0,
+    ) -> list[list[str]]:
+        return self._tok.split_segments(
+            tokenized,
+            max_text_tokens_per_segment=max_text_tokens_per_segment,
+            quick_streaming_tokens=quick_streaming_tokens,
+        )

vllm_omni/model_executor/models/indextts2/utils/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/model_executor/models/indextts2/utils/campplus/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/model_executor/models/indextts2/utils/campplus/dtdnn.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
+
+from collections import OrderedDict
+
+import torch.nn.functional as F
+from torch import nn
+
+from .layers import (
+    BasicResBlock,
+    CAMDenseTDNNBlock,
+    DenseLayer,
+    StatsPool,
+    TDNNLayer,
+    TransitLayer,
+    get_nonlinear,
+)
+
+
+class FCM(nn.Module):
+    def __init__(self, block=BasicResBlock, num_blocks=[2, 2], m_channels=32, feat_dim=80):
+        super().__init__()
+        self.in_planes = m_channels
+        self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(m_channels)
+
+        self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=2)
+        self.layer2 = self._make_layer(block, m_channels, num_blocks[1], stride=2)
+
+        self.conv2 = nn.Conv2d(m_channels, m_channels, kernel_size=3, stride=(2, 1), padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(m_channels)
+        self.out_channels = m_channels * (feat_dim // 8)
+
+    def _make_layer(self, block, planes, num_blocks, stride):
+        strides = [stride] + [1] * (num_blocks - 1)
+        layers = []
+        for stride in strides:
+            layers.append(block(self.in_planes, planes, stride))
+            self.in_planes = planes * block.expansion
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = x.unsqueeze(1)
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.layer1(out)
+        out = self.layer2(out)
+        out = F.relu(self.bn2(self.conv2(out)))
+
+        shape = out.shape
+        out = out.reshape(shape[0], shape[1] * shape[2], shape[3])
+        return out
+
+
+class CAMPPlus(nn.Module):
+    def __init__(
+        self,
+        feat_dim=80,
+        embedding_size=512,
+        growth_rate=32,
+        bn_size=4,
+        init_channels=128,
+        config_str="batchnorm-relu",
+        memory_efficient=True,
+    ):
+        super().__init__()
+
+        self.head = FCM(feat_dim=feat_dim)
+        channels = self.head.out_channels
+
+        self.xvector = nn.Sequential(
+            OrderedDict(
+                [
+                    (
+                        "tdnn",
+                        TDNNLayer(channels, init_channels, 5, stride=2, dilation=1, padding=-1, config_str=config_str),
+                    ),
+                ]
+            )
+        )
+        channels = init_channels
+        for i, (num_layers, kernel_size, dilation) in enumerate(zip((12, 24, 16), (3, 3, 3), (1, 2, 2))):
+            block = CAMDenseTDNNBlock(
+                num_layers=num_layers,
+                in_channels=channels,
+                out_channels=growth_rate,
+                bn_channels=bn_size * growth_rate,
+                kernel_size=kernel_size,
+                dilation=dilation,
+                config_str=config_str,
+                memory_efficient=memory_efficient,
+            )
+            self.xvector.add_module(f"block{i + 1}", block)
+            channels = channels + num_layers * growth_rate
+            self.xvector.add_module(
+                f"transit{i + 1}", TransitLayer(channels, channels // 2, bias=False, config_str=config_str)
+            )
+            channels //= 2
+
+        self.xvector.add_module("out_nonlinear", get_nonlinear(config_str, channels))
+
+        self.xvector.add_module("stats", StatsPool())
+        self.xvector.add_module("dense", DenseLayer(channels * 2, embedding_size, config_str="batchnorm_"))
+
+        for m in self.modules():
+            if isinstance(m, (nn.Conv1d, nn.Linear)):
+                nn.init.kaiming_normal_(m.weight.data)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+    def forward(self, x):
+        x = x.permute(0, 2, 1)  # (B,T,F) => (B,F,T)
+        x = self.head(x)
+        x = self.xvector(x)
+        return x

vllm_omni/model_executor/models/indextts2/utils/campplus/layers.py

@@ -0,0 +1,221 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint as cp
+from torch import nn
+
+
+def get_nonlinear(config_str, channels):
+    nonlinear = nn.Sequential()
+    for name in config_str.split("-"):
+        if name == "relu":
+            nonlinear.add_module("relu", nn.ReLU(inplace=True))
+        elif name == "prelu":
+            nonlinear.add_module("prelu", nn.PReLU(channels))
+        elif name == "batchnorm":
+            nonlinear.add_module("batchnorm", nn.BatchNorm1d(channels))
+        elif name == "batchnorm_":
+            nonlinear.add_module("batchnorm", nn.BatchNorm1d(channels, affine=False))
+        else:
+            raise ValueError(f"Unexpected module ({name}).")
+    return nonlinear
+
+
+def statistics_pooling(x, dim=-1, keepdim=False, unbiased=True, eps=1e-2):
+    mean = x.mean(dim=dim)
+    std = x.std(dim=dim, unbiased=unbiased)
+    stats = torch.cat([mean, std], dim=-1)
+    if keepdim:
+        stats = stats.unsqueeze(dim=dim)
+    return stats
+
+
+class StatsPool(nn.Module):
+    def forward(self, x):
+        return statistics_pooling(x)
+
+
+class TDNNLayer(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        bias=False,
+        config_str="batchnorm-relu",
+    ):
+        super().__init__()
+        if padding < 0:
+            assert kernel_size % 2 == 1, f"Expect equal paddings, but got even kernel size ({kernel_size})"
+            padding = (kernel_size - 1) // 2 * dilation
+        self.linear = nn.Conv1d(
+            in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias
+        )
+        self.nonlinear = get_nonlinear(config_str, out_channels)
+
+    def forward(self, x):
+        x = self.linear(x)
+        x = self.nonlinear(x)
+        return x
+
+
+class CAMLayer(nn.Module):
+    def __init__(self, bn_channels, out_channels, kernel_size, stride, padding, dilation, bias, reduction=2):
+        super().__init__()
+        self.linear_local = nn.Conv1d(
+            bn_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias
+        )
+        self.linear1 = nn.Conv1d(bn_channels, bn_channels // reduction, 1)
+        self.relu = nn.ReLU(inplace=True)
+        self.linear2 = nn.Conv1d(bn_channels // reduction, out_channels, 1)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        y = self.linear_local(x)
+        context = x.mean(-1, keepdim=True) + self.seg_pooling(x)
+        context = self.relu(self.linear1(context))
+        m = self.sigmoid(self.linear2(context))
+        return y * m
+
+    def seg_pooling(self, x, seg_len=100, stype="avg"):
+        if stype == "avg":
+            seg = F.avg_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True)
+        elif stype == "max":
+            seg = F.max_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True)
+        else:
+            raise ValueError("Wrong segment pooling type.")
+        shape = seg.shape
+        seg = seg.unsqueeze(-1).expand(*shape, seg_len).reshape(*shape[:-1], -1)
+        seg = seg[..., : x.shape[-1]]
+        return seg
+
+
+class CAMDenseTDNNLayer(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        bn_channels,
+        kernel_size,
+        stride=1,
+        dilation=1,
+        bias=False,
+        config_str="batchnorm-relu",
+        memory_efficient=False,
+    ):
+        super().__init__()
+        assert kernel_size % 2 == 1, f"Expect equal paddings, but got even kernel size ({kernel_size})"
+        padding = (kernel_size - 1) // 2 * dilation
+        self.memory_efficient = memory_efficient
+        self.nonlinear1 = get_nonlinear(config_str, in_channels)
+        self.linear1 = nn.Conv1d(in_channels, bn_channels, 1, bias=False)
+        self.nonlinear2 = get_nonlinear(config_str, bn_channels)
+        self.cam_layer = CAMLayer(
+            bn_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias
+        )
+
+    def bn_function(self, x):
+        return self.linear1(self.nonlinear1(x))
+
+    def forward(self, x):
+        if self.training and self.memory_efficient:
+            x = cp.checkpoint(self.bn_function, x)
+        else:
+            x = self.bn_function(x)
+        x = self.cam_layer(self.nonlinear2(x))
+        return x
+
+
+class CAMDenseTDNNBlock(nn.ModuleList):
+    def __init__(
+        self,
+        num_layers,
+        in_channels,
+        out_channels,
+        bn_channels,
+        kernel_size,
+        stride=1,
+        dilation=1,
+        bias=False,
+        config_str="batchnorm-relu",
+        memory_efficient=False,
+    ):
+        super().__init__()
+        for i in range(num_layers):
+            layer = CAMDenseTDNNLayer(
+                in_channels=in_channels + i * out_channels,
+                out_channels=out_channels,
+                bn_channels=bn_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                dilation=dilation,
+                bias=bias,
+                config_str=config_str,
+                memory_efficient=memory_efficient,
+            )
+            self.add_module(f"tdnnd{i + 1}", layer)
+
+    def forward(self, x):
+        for layer in self:
+            x = torch.cat([x, layer(x)], dim=1)
+        return x
+
+
+class TransitLayer(nn.Module):
+    def __init__(self, in_channels, out_channels, bias=True, config_str="batchnorm-relu"):
+        super().__init__()
+        self.nonlinear = get_nonlinear(config_str, in_channels)
+        self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias)
+
+    def forward(self, x):
+        x = self.nonlinear(x)
+        x = self.linear(x)
+        return x
+
+
+class DenseLayer(nn.Module):
+    def __init__(self, in_channels, out_channels, bias=False, config_str="batchnorm-relu"):
+        super().__init__()
+        self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias)
+        self.nonlinear = get_nonlinear(config_str, out_channels)
+
+    def forward(self, x):
+        if len(x.shape) == 2:
+            x = self.linear(x.unsqueeze(dim=-1)).squeeze(dim=-1)
+        else:
+            x = self.linear(x)
+        x = self.nonlinear(x)
+        return x
+
+
+class BasicResBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, in_planes, planes, stride=1):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=(stride, 1), padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=(stride, 1), bias=False),
+                nn.BatchNorm2d(self.expansion * planes),
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out

vllm_omni/model_executor/models/indextts2/utils/common.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import re
+
+import torch
+
+
+def tokenize_by_CJK_char(line: str, do_upper_case=True) -> str:
+    """
+    Tokenize a line of text with CJK char.
+
+    Note: All return characters will be upper case.
+
+    Example:
+      input = "你好世界是 hello world 的中文"
+      output = "你 好 世 界 是 HELLO WORLD 的 中 文"
+
+    Args:
+      line:
+        The input text.
+
+    Return:
+      A new string tokenize by CJK char.
+    """
+    # CJK ranges from nltk/tokenize/util.py (alvations/nltk@79eed6d)
+    CJK_RANGE_PATTERN = (
+        r"([\u1100-\u11ff\u2e80-\ua4cf\ua840-\uD7AF\uF900-\uFAFF\uFE30-\uFE4F\uFF65-\uFFDC\U00020000-\U0002FFFF])"
+    )
+    chars = re.split(CJK_RANGE_PATTERN, line.strip())
+    return " ".join([w.strip().upper() if do_upper_case else w.strip() for w in chars if w.strip()])
+
+
+def de_tokenized_by_CJK_char(line: str, do_lower_case=False) -> str:
+    """
+    Example:
+      input = "你 好 世 界 是 HELLO WORLD 的 中 文"
+      output = "你好世界是 hello world 的中文"
+
+    do_lower_case:
+      input = "SEE YOU!"
+      output = "see you!"
+    """
+    # replace english words in the line with placeholders
+    english_word_pattern = re.compile(r"([A-Z]+(?:[\s'-][A-Z-]+)*)", re.IGNORECASE)
+    english_sents = english_word_pattern.findall(line)
+    for i, sent in enumerate(english_sents):
+        line = line.replace(sent, f"<sent_{i}>")
+
+    words = line.split()
+    # restore english sentences
+    sent_placeholder_pattern = re.compile(r"(<sent_(\d+)>)")
+    for i in range(len(words)):
+        all_matches = sent_placeholder_pattern.findall(words[i])
+        if all_matches:
+            # restore the english word
+            for h, j in all_matches:
+                placeholder_index = int(j)
+                words[i] = words[i].replace(h, english_sents[placeholder_index])
+                if do_lower_case:
+                    words[i] = words[i].lower()
+    return "".join(words)
+
+
+def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
+    """Make mask tensor containing indices of padded part.
+
+    See description of make_non_pad_mask.
+
+    Args:
+        lengths (torch.Tensor): Batch of lengths (B,).
+    Returns:
+        torch.Tensor: Mask tensor containing indices of padded part.
+
+    Examples:
+        >>> lengths = [5, 3, 2]
+        >>> make_pad_mask(lengths)
+        masks = [[0, 0, 0, 0 ,0],
+                 [0, 0, 0, 1, 1],
+                 [0, 0, 1, 1, 1]]
+    """
+    batch_size = lengths.size(0)
+    max_len = max_len if max_len > 0 else lengths.max().item()
+    seq_range = torch.arange(0, max_len, dtype=torch.int64, device=lengths.device)
+    seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)
+    seq_length_expand = lengths.unsqueeze(-1)
+    mask = seq_range_expand >= seq_length_expand
+    return mask

vllm_omni/model_executor/models/indextts2/utils/front.py

@@ -0,0 +1,208 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+import warnings
+from typing import overload
+
+from sentencepiece import SentencePieceProcessor
+
+from vllm_omni.model_executor.models.indextts2.utils.common import (
+    de_tokenized_by_CJK_char,
+    tokenize_by_CJK_char,
+)
+
+
+class TextTokenizer:
+    def __init__(self, vocab_file: str):
+        self.vocab_file = vocab_file
+
+        if self.vocab_file is None:
+            raise ValueError("vocab_file is None")
+        if not os.path.exists(self.vocab_file):
+            raise ValueError(f"vocab_file {self.vocab_file} does not exist")
+        self.sp_model = SentencePieceProcessor(model_file=self.vocab_file)
+        self.pre_tokenizers = [tokenize_by_CJK_char]
+
+    @property
+    def vocab_size(self):
+        return self.sp_model.GetPieceSize()
+
+    @property
+    def unk_token(self):
+        return "<unk>"
+
+    @property
+    def pad_token(self):
+        return None
+
+    @property
+    def bos_token(self):
+        return "<s>"
+
+    @property
+    def eos_token(self):
+        return "</s>"
+
+    @property
+    def pad_token_id(self):
+        return -1
+
+    @property
+    def bos_token_id(self):
+        return 0
+
+    @property
+    def eos_token_id(self):
+        return 1
+
+    @property
+    def unk_token_id(self):
+        return self.sp_model.unk_id()
+
+    @property
+    def special_tokens_map(self):
+        return {
+            "unk_token": self.unk_token,
+            "pad_token": self.pad_token,
+            "bos_token": self.bos_token,
+            "eos_token": self.eos_token,
+        }
+
+    def get_vocab(self):
+        vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
+        return vocab
+
+    @overload
+    def convert_ids_to_tokens(self, ids: int) -> str: ...
+
+    @overload
+    def convert_ids_to_tokens(self, ids: list[int]) -> list[str]: ...
+
+    def convert_ids_to_tokens(self, ids: list[int] | int):
+        return self.sp_model.IdToPiece(ids)
+
+    def convert_tokens_to_ids(self, tokens: list[str] | str) -> list[int]:
+        if isinstance(tokens, str):
+            tokens = [tokens]
+        return [self.sp_model.PieceToId(token) for token in tokens]
+
+    def tokenize(self, text: str) -> list[str]:
+        return self.encode(text, out_type=str)
+
+    def encode(self, text: str, **kwargs):
+        if len(text) == 0:
+            return []
+        if len(text.strip()) == 1:
+            return self.sp_model.Encode(text, out_type=kwargs.pop("out_type", int), **kwargs)
+        for pre_tokenizer in self.pre_tokenizers:
+            text = pre_tokenizer(text)
+        return self.sp_model.Encode(text, out_type=kwargs.pop("out_type", int), **kwargs)
+
+    def decode(self, ids: list[int] | int, do_lower_case=False, **kwargs):
+        if isinstance(ids, int):
+            ids = [ids]
+        decoded = self.sp_model.Decode(ids, out_type=kwargs.pop("out_type", str), **kwargs)
+        return de_tokenized_by_CJK_char(decoded, do_lower_case=do_lower_case)
+
+    @staticmethod
+    def split_segments_by_token(
+        tokenized_str: list[str],
+        split_tokens: list[str],
+        max_text_tokens_per_segment: int,
+        quick_streaming_tokens: int = 0,
+    ) -> list[list[str]]:
+        if len(tokenized_str) == 0:
+            return []
+        segments: list[list[str]] = []
+        current_segment = []
+        current_segment_tokens_len = 0
+        i = 0
+        while i < len(tokenized_str):
+            token = tokenized_str[i]
+            current_segment.append(token)
+            current_segment_tokens_len += 1
+            if not ("," in split_tokens or "▁," in split_tokens) and (
+                "," in current_segment or "▁," in current_segment
+            ):
+                sub_segments = TextTokenizer.split_segments_by_token(
+                    current_segment,
+                    [",", "▁,"],
+                    max_text_tokens_per_segment=max_text_tokens_per_segment,
+                    quick_streaming_tokens=quick_streaming_tokens,
+                )
+            elif "-" not in split_tokens and "-" in current_segment:
+                sub_segments = TextTokenizer.split_segments_by_token(
+                    current_segment,
+                    ["-"],
+                    max_text_tokens_per_segment=max_text_tokens_per_segment,
+                    quick_streaming_tokens=quick_streaming_tokens,
+                )
+            elif current_segment_tokens_len <= max_text_tokens_per_segment:
+                if token in split_tokens and current_segment_tokens_len > 2:
+                    if i < len(tokenized_str) - 1:
+                        if tokenized_str[i + 1] in ["'", "▁'"]:
+                            current_segment.append(tokenized_str[i + 1])
+                            i += 1
+                    segments.append(current_segment)
+                    current_segment = []
+                    current_segment_tokens_len = 0
+                i += 1
+                continue
+            else:
+                sub_segments = []
+                for j in range(0, len(current_segment), max_text_tokens_per_segment):
+                    if j + max_text_tokens_per_segment < len(current_segment):
+                        sub_segments.append(current_segment[j : j + max_text_tokens_per_segment])
+                    else:
+                        sub_segments.append(current_segment[j:])
+                warnings.warn(
+                    f"The tokens length of segment exceeds limit: {max_text_tokens_per_segment}, "
+                    f"Tokens in segment: {current_segment}."
+                    "Maybe unexpected behavior",
+                    RuntimeWarning,
+                )
+            segments.extend(sub_segments)
+            current_segment = []
+            current_segment_tokens_len = 0
+            i += 1
+        if current_segment_tokens_len > 0:
+            assert current_segment_tokens_len <= max_text_tokens_per_segment
+            segments.append(current_segment)
+        merged_segments = []
+        total_token = 0
+        for segment in segments:
+            total_token += len(segment)
+            if len(segment) == 0:
+                continue
+            if len(merged_segments) == 0:
+                merged_segments.append(segment)
+            elif (
+                len(merged_segments[-1]) + len(segment) <= max_text_tokens_per_segment
+                and total_token > quick_streaming_tokens
+            ):
+                merged_segments[-1] = merged_segments[-1] + segment
+            elif len(merged_segments[-1]) + len(segment) <= max_text_tokens_per_segment / 2:
+                merged_segments[-1] = merged_segments[-1] + segment
+            else:
+                merged_segments.append(segment)
+        return merged_segments
+
+    punctuation_marks_tokens = [
+        ".",
+        "!",
+        "?",
+        "▁.",
+        "▁?",
+        "▁...",
+    ]
+
+    def split_segments(
+        self, tokenized: list[str], max_text_tokens_per_segment=120, quick_streaming_tokens=0
+    ) -> list[list[str]]:
+        return TextTokenizer.split_segments_by_token(
+            tokenized,
+            self.punctuation_marks_tokens,
+            max_text_tokens_per_segment=max_text_tokens_per_segment,
+            quick_streaming_tokens=quick_streaming_tokens,
+        )

vllm_omni/model_executor/models/indextts2/utils/maskgct/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/model_executor/models/indextts2/utils/maskgct/quantize.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright (c) 2024 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Merged from factorized_vector_quantize.py and residual_vq.py.
+# Only FactorizedVectorQuantize ("fvq") quantizer type is retained
+# as it is the only one used by RepCodec in IndexTTS2.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch.nn.utils import weight_norm
+
+
+def WNConv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+class FactorizedVectorQuantize(nn.Module):
+    def __init__(
+        self,
+        input_dim,
+        codebook_size,
+        codebook_dim,
+        commitment=0.005,
+        codebook_loss_weight=1.0,
+        use_l2_normalize=True,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+        self.commitment = commitment
+        self.codebook_loss_weight = codebook_loss_weight
+        self.use_l2_normalize = use_l2_normalize
+
+        if self.input_dim != self.codebook_dim:
+            self.in_project = WNConv1d(self.input_dim, self.codebook_dim, kernel_size=1)
+            self.out_project = WNConv1d(self.codebook_dim, self.input_dim, kernel_size=1)
+        else:
+            self.in_project = nn.Identity()
+            self.out_project = nn.Identity()
+
+        self.codebook = nn.Embedding(self.codebook_size, self.codebook_dim)
+
+    def forward(self, z):
+        z_e = self.in_project(z)
+        z_q, indices = self.decode_latents(z_e)
+
+        if self.training:
+            commit_loss = F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2]) * self.commitment
+            codebook_loss = F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2]) * self.codebook_loss_weight
+        else:
+            commit_loss = torch.zeros(z.shape[0], device=z.device)
+            codebook_loss = torch.zeros(z.shape[0], device=z.device)
+
+        z_q = z_e + (z_q - z_e).detach()
+        z_q = self.out_project(z_q)
+
+        return z_q, commit_loss, codebook_loss, indices, z_e
+
+    def embed_code(self, embed_id):
+        return F.embedding(embed_id, self.codebook.weight)
+
+    def decode_code(self, embed_id):
+        return self.embed_code(embed_id).transpose(1, 2)
+
+    def decode_latents(self, latents):
+        encodings = rearrange(latents, "b d t -> (b t) d")
+        codebook = self.codebook.weight
+
+        if self.use_l2_normalize:
+            encodings = F.normalize(encodings)
+            codebook = F.normalize(codebook)
+
+        dist = (
+            encodings.pow(2).sum(1, keepdim=True)
+            - 2 * encodings @ codebook.t()
+            + codebook.pow(2).sum(1, keepdim=True).t()
+        )
+        indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))
+        z_q = self.decode_code(indices)
+        return z_q, indices
+
+    def vq2emb(self, vq, out_proj=True):
+        emb = self.decode_code(vq)
+        if out_proj:
+            emb = self.out_project(emb)
+        return emb
+
+
+class ResidualVQ(nn.Module):
+    """Residual Vector Quantization (SoundStream).
+
+    Only supports FactorizedVectorQuantize ("fvq") quantizer type.
+    """
+
+    def __init__(
+        self,
+        input_dim: int = 256,
+        num_quantizers: int = 8,
+        codebook_size: int = 1024,
+        codebook_dim: int = 256,
+        quantizer_type: str = "fvq",
+        quantizer_dropout: float = 0.5,
+        **kwargs,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.num_quantizers = num_quantizers
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+        self.quantizer_dropout = quantizer_dropout
+
+        self.quantizers = nn.ModuleList(
+            [
+                FactorizedVectorQuantize(
+                    input_dim=input_dim,
+                    codebook_size=codebook_size,
+                    codebook_dim=codebook_dim,
+                    **kwargs,
+                )
+                for _ in range(num_quantizers)
+            ]
+        )
+
+    def forward(self, z, n_quantizers: int = None):
+        quantized_out = 0.0
+        residual = z
+
+        all_commit_losses = []
+        all_codebook_losses = []
+        all_indices = []
+        all_quantized = []
+
+        if n_quantizers is None:
+            n_quantizers = self.num_quantizers
+
+        if self.training:
+            n_q = torch.ones((z.shape[0],)) * self.num_quantizers + 1
+            dropout = torch.randint(1, self.num_quantizers + 1, (z.shape[0],))
+            n_dropout = int(z.shape[0] * self.quantizer_dropout)
+            n_q[:n_dropout] = dropout[:n_dropout]
+            n_quantizers = n_q.to(z.device)
+
+        for i, quantizer in enumerate(self.quantizers):
+            if not self.training and i >= n_quantizers:
+                break
+
+            z_q_i, commit_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer(residual)
+
+            mask = torch.full((z.shape[0],), fill_value=i, device=z.device) < n_quantizers
+            quantized_out = quantized_out + z_q_i * mask[:, None, None]
+            residual = residual - z_q_i
+
+            commit_loss_i = (commit_loss_i * mask).mean()
+            codebook_loss_i = (codebook_loss_i * mask).mean()
+
+            all_commit_losses.append(commit_loss_i)
+            all_codebook_losses.append(codebook_loss_i)
+            all_indices.append(indices_i)
+            all_quantized.append(z_q_i)
+
+        (
+            all_commit_losses,
+            all_codebook_losses,
+            all_indices,
+            all_quantized,
+        ) = map(
+            torch.stack,
+            (
+                all_commit_losses,
+                all_codebook_losses,
+                all_indices,
+                all_quantized,
+            ),
+        )
+
+        return (
+            quantized_out,
+            all_indices,
+            all_commit_losses,
+            all_codebook_losses,
+            all_quantized,
+        )
+
+    def vq2emb(self, vq, n_quantizers=None):
+        quantized_out = 0.0
+        if n_quantizers is None:
+            n_quantizers = self.num_quantizers
+        for idx, quantizer in enumerate(self.quantizers):
+            if idx >= n_quantizers:
+                break
+            quantized_out += quantizer.vq2emb(vq[idx])
+        return quantized_out

vllm_omni/model_executor/models/indextts2/utils/maskgct/repcodec_model.py

@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright (c) 2024 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch.nn as nn
+from torch.nn import functional as F
+
+from vllm_omni.model_executor.models.indextts2.utils.maskgct.quantize import ResidualVQ
+from vllm_omni.model_executor.models.indextts2.utils.maskgct.vocos import VocosBackbone
+
+
+def init_weights(m):
+    if isinstance(m, nn.Conv1d):
+        nn.init.trunc_normal_(m.weight, std=0.02)
+        nn.init.constant_(m.bias, 0)
+    if isinstance(m, nn.Linear):
+        nn.init.trunc_normal_(m.weight, std=0.02)
+        nn.init.constant_(m.bias, 0)
+
+
+class RepCodec(nn.Module):
+    def __init__(
+        self,
+        codebook_size=8192,
+        hidden_size=1024,
+        codebook_dim=8,
+        vocos_dim=384,
+        vocos_intermediate_dim=2048,
+        vocos_num_layers=12,
+        num_quantizers=1,
+        downsample_scale=1,
+        cfg=None,
+    ):
+        super().__init__()
+        if cfg is not None:
+            _g = lambda attr, default: getattr(cfg, attr, default)  # noqa: E731
+            codebook_size = _g("codebook_size", codebook_size)
+            codebook_dim = _g("codebook_dim", codebook_dim)
+            hidden_size = _g("hidden_size", hidden_size)
+            vocos_dim = _g("vocos_dim", vocos_dim)
+            vocos_intermediate_dim = _g("vocos_intermediate_dim", vocos_intermediate_dim)
+            vocos_num_layers = _g("vocos_num_layers", vocos_num_layers)
+            num_quantizers = _g("num_quantizers", num_quantizers)
+            downsample_scale = _g("downsample_scale", downsample_scale)
+
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+        self.hidden_size = hidden_size
+        self.vocos_dim = vocos_dim
+        self.vocos_intermediate_dim = vocos_intermediate_dim
+        self.vocos_num_layers = vocos_num_layers
+        self.num_quantizers = num_quantizers
+        self.downsample_scale = downsample_scale
+
+        if self.downsample_scale is not None and self.downsample_scale > 1:
+            self.down = nn.Conv1d(self.hidden_size, self.hidden_size, kernel_size=3, stride=2, padding=1)
+            self.up = nn.Conv1d(self.hidden_size, self.hidden_size, kernel_size=3, stride=1, padding=1)
+
+        self.encoder = nn.Sequential(
+            VocosBackbone(
+                input_channels=self.hidden_size,
+                dim=self.vocos_dim,
+                intermediate_dim=self.vocos_intermediate_dim,
+                num_layers=self.vocos_num_layers,
+                adanorm_num_embeddings=None,
+            ),
+            nn.Linear(self.vocos_dim, self.hidden_size),
+        )
+        self.decoder = nn.Sequential(
+            VocosBackbone(
+                input_channels=self.hidden_size,
+                dim=self.vocos_dim,
+                intermediate_dim=self.vocos_intermediate_dim,
+                num_layers=self.vocos_num_layers,
+                adanorm_num_embeddings=None,
+            ),
+            nn.Linear(self.vocos_dim, self.hidden_size),
+        )
+
+        self.quantizer = ResidualVQ(
+            input_dim=hidden_size,
+            num_quantizers=num_quantizers,
+            codebook_size=codebook_size,
+            codebook_dim=codebook_dim,
+            quantizer_type="fvq",
+            quantizer_dropout=0.0,
+            commitment=0.15,
+            codebook_loss_weight=1.0,
+            use_l2_normalize=True,
+        )
+
+        self.reset_parameters()
+
+    def quantize(self, x):
+        if self.downsample_scale is not None and self.downsample_scale > 1:
+            x = x.transpose(1, 2)
+            x = self.down(x)
+            x = F.gelu(x)
+            x = x.transpose(1, 2)
+
+        x = self.encoder(x.transpose(1, 2)).transpose(1, 2)
+
+        (
+            quantized_out,
+            all_indices,
+            all_commit_losses,
+            all_codebook_losses,
+            _,
+        ) = self.quantizer(x)
+
+        if all_indices.shape[0] == 1:
+            return all_indices.squeeze(0), quantized_out.transpose(1, 2)
+        return all_indices, quantized_out.transpose(1, 2)
+
+    def reset_parameters(self):
+        self.apply(init_weights)

vllm_omni/model_executor/models/indextts2/utils/maskgct/vocos.py

@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright (c) 2024 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import torch
+from torch import nn
+
+
+class ConvNeXtBlock(nn.Module):
+    """ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal."""
+
+    def __init__(
+        self,
+        dim: int,
+        intermediate_dim: int,
+        layer_scale_init_value: float,
+        adanorm_num_embeddings: int | None = None,
+    ):
+        super().__init__()
+        self.dwconv = nn.Conv1d(dim, dim, kernel_size=7, padding=3, groups=dim)
+        self.adanorm = adanorm_num_embeddings is not None
+        if adanorm_num_embeddings:
+            raise ValueError("AdaLayerNorm is not supported in inference-only mode")
+        else:
+            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)
+            if layer_scale_init_value > 0
+            else None
+        )
+
+    def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor | None = None) -> torch.Tensor:
+        residual = x
+        x = self.dwconv(x)
+        x = x.transpose(1, 2)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.transpose(1, 2)
+        x = residual + x
+        return x
+
+
+class VocosBackbone(nn.Module):
+    """Vocos backbone module built with ConvNeXt blocks."""
+
+    def __init__(
+        self,
+        input_channels: int,
+        dim: int,
+        intermediate_dim: int,
+        num_layers: int,
+        layer_scale_init_value: float | None = None,
+        adanorm_num_embeddings: int | None = None,
+    ):
+        super().__init__()
+        self.input_channels = input_channels
+        self.embed = nn.Conv1d(input_channels, dim, kernel_size=7, padding=3)
+        self.adanorm = adanorm_num_embeddings is not None
+        if adanorm_num_embeddings:
+            raise ValueError("AdaLayerNorm is not supported in inference-only mode")
+        else:
+            self.norm = nn.LayerNorm(dim, eps=1e-6)
+        layer_scale_init_value = layer_scale_init_value or 1 / num_layers
+        self.convnext = nn.ModuleList(
+            [
+                ConvNeXtBlock(
+                    dim=dim,
+                    intermediate_dim=intermediate_dim,
+                    layer_scale_init_value=layer_scale_init_value,
+                    adanorm_num_embeddings=adanorm_num_embeddings,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.final_layer_norm = nn.LayerNorm(dim, eps=1e-6)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Conv1d, nn.Linear)):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            nn.init.constant_(m.bias, 0)
+
+    def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
+        x = self.embed(x)
+        x = self.norm(x.transpose(1, 2))
+        x = x.transpose(1, 2)
+        for conv_block in self.convnext:
+            x = conv_block(x)
+        x = self.final_layer_norm(x.transpose(1, 2))
+        return x

vllm_omni/model_executor/models/registry.py

@@ -305,6 +305,17 @@ _OMNI_MODELS = {
         "dynin_omni",
         "DyninOmniForConditionalGeneration",
     ),
+    ## IndexTTS2
+    "IndexTTS2TalkerForConditionalGeneration": (
+        "indextts2",
+        "indextts2_talker",
+        "IndexTTS2TalkerForConditionalGeneration",
+    ),
+    "IndexTTS2S2MelDecoder": (
+        "indextts2",
+        "indextts2_s2mel_decoder",
+        "IndexTTS2S2MelDecoder",
+    ),
     ## Ming-flash-omni-2.0
     "MingFlashOmniForConditionalGeneration": (
         "ming_flash_omni",

vllm_omni/model_executor/stage_input_processors/indextts2.py

@@ -0,0 +1,271 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Stage input processor for IndexTTS2: Talker (GPT AR) → S2Mel decoder."""
+
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+STOP_MEL_TOKEN = 8193
+
+
+def _cpu_view(tensor: torch.Tensor) -> torch.Tensor:
+    """Return a contiguous CPU tensor suitable for connector serialization.
+
+    Always returns an owning copy so downstream consumers cannot mutate the
+    source tensors that may still be referenced by the producer stage.
+    """
+    out = tensor.detach()
+    if out.device.type != "cpu":
+        return out.cpu().contiguous()
+    return out.clone().contiguous()
+
+
+def _strip_stop_token(
+    codes: torch.Tensor,
+    latent: torch.Tensor,
+    stop_mel_token: int = STOP_MEL_TOKEN,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Strip at the first stop token, matching official IndexTTS2 v2.
+
+    Returns: (codes [B, T'], latent [B, T', D], code_lens [B]).
+    """
+    if codes.ndim == 1:
+        codes = codes.unsqueeze(0)
+    if latent.ndim == 2:
+        latent = latent.unsqueeze(0)
+
+    if codes.device.type != "cpu":
+        codes = codes.detach().cpu()
+    if latent.device.type != "cpu":
+        latent = latent.detach().cpu()
+
+    device = codes.device
+    code_lens = []
+    codes_out = []
+    latent_out = []
+
+    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())
+        else:
+            valid_len = int(code.shape[0])
+        code_lens.append(valid_len)
+        codes_out.append(code[:valid_len])
+        latent_out.append(lat[:valid_len])
+
+    max_len = max(code_lens) if code_lens else 0
+    if max_len == 0:
+        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),
+        )
+
+    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,
+    )
+    for i, (c, lat) in enumerate(zip(codes_out, latent_out)):
+        padded_codes[i, : c.shape[0]] = c
+        padded_latent[i, : lat.shape[0]] = lat
+
+    return padded_codes, padded_latent, torch.tensor(code_lens, dtype=torch.long, device=device)
+
+
+def _normalize_mel_sequence(mel_codes: torch.Tensor) -> torch.Tensor:
+    mel_codes = mel_codes.to(torch.long)
+    if mel_codes.ndim <= 1:
+        return mel_codes.reshape(-1).contiguous() if mel_codes.ndim == 0 else mel_codes.contiguous()
+    if mel_codes.ndim == 2 and (mel_codes.shape[0] == 1 or mel_codes.shape[1] == 1):
+        return mel_codes.reshape(-1).contiguous()
+    return mel_codes.reshape(-1).contiguous()
+
+
+def _normalize_latent_sequence(latent: torch.Tensor) -> torch.Tensor:
+    if latent.ndim <= 1:
+        return latent.reshape(1, -1).contiguous()
+    if latent.ndim == 2:
+        return latent.contiguous()
+    if latent.ndim == 3 and latent.shape[0] == 1:
+        return latent[0].contiguous()
+    return latent.reshape(-1, latent.shape[-1]).contiguous()
+
+
+def _build_s2mel_additional_information(
+    mel_codes: torch.Tensor,
+    latent: torch.Tensor,
+    meta: dict[str, Any],
+    *,
+    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),
+    }
+
+    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)
+
+    return additional_information
+
+
+def _get_payload_value(pooling_output: dict[str, Any], dotted_key: str, nested_parent: str, nested_key: str) -> Any:
+    value = pooling_output.get(dotted_key)
+    if value is not None:
+        return value
+    nested = pooling_output.get(nested_parent)
+    if isinstance(nested, dict):
+        return nested.get(nested_key)
+    return None
+
+
+def _request_id(request: Any) -> str:
+    return str(getattr(request, "external_req_id", None) or getattr(request, "request_id", "?"))
+
+
+def talker2s2mel(
+    source_outputs: list[Any],
+    prompt: Any = None,
+    _requires_multimodal_data: bool = False,
+) -> list[Any]:
+    """Legacy orchestrator path: collect all Stage 0 output, format Stage 1 input."""
+    _old_nt = torch.get_num_threads()
+    torch.set_num_threads(1)
+    try:
+        return _talker2s2mel_impl(source_outputs)
+    finally:
+        torch.set_num_threads(_old_nt)
+
+
+def _talker2s2mel_impl(source_outputs: list[Any]) -> list[Any]:
+    from vllm_omni.inputs.data import OmniTokensPrompt
+
+    s2mel_inputs: list[OmniTokensPrompt] = []
+
+    for i, talker_output in enumerate(source_outputs):
+        if not talker_output.finished:
+            continue
+
+        output = talker_output.outputs[0]
+        mm = output.multimodal_output
+        if not isinstance(mm, dict):
+            logger.warning("Talker output %d has no multimodal_output dict", i)
+            continue
+
+        codes_dict = mm.get("codes", {})
+        mel_codes = codes_dict.get("mel")
+        if not isinstance(mel_codes, torch.Tensor) or mel_codes.numel() == 0:
+            logger.warning("Talker output %d has empty mel_codes", i)
+            continue
+        mel_codes = mel_codes.to(torch.long)
+
+        hs_dict = mm.get("hidden_states", {})
+        latent = hs_dict.get("latent")
+        if not isinstance(latent, torch.Tensor) or latent.numel() == 0:
+            logger.warning("Talker output %d has empty latent", i)
+            continue
+
+        meta = mm.get("meta", {})
+        if not isinstance(meta, dict):
+            meta = {}
+
+        additional_information = _build_s2mel_additional_information(mel_codes, latent, meta, context="talker2s2mel")
+        s2mel_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=[0],
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+                additional_information=additional_information,
+            )
+        )
+
+    return s2mel_inputs
+
+
+def talker2s2mel_token_only(
+    source_outputs: list[Any],
+    prompt: Any = None,
+    _requires_multimodal_data: bool = False,
+) -> list[Any]:
+    """Sync-side placeholder for Stage 1; tensors arrive via full-payload connector."""
+    from vllm_omni.inputs.data import OmniTokensPrompt
+
+    del prompt
+    s2mel_inputs: list[OmniTokensPrompt] = []
+    for talker_output in source_outputs:
+        if not talker_output.finished:
+            continue
+        s2mel_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=[0],
+                additional_information=None,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+            )
+        )
+    return s2mel_inputs
+
+
+def talker2s2mel_full_payload(
+    transfer_manager: Any,
+    pooling_output: dict[str, Any],
+    request: Any,
+    **_: Any,
+) -> dict[str, Any] | None:
+    """Build the complete S2Mel input from accumulated per-step talker deltas."""
+    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
+
+    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
+
+    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
+
+    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]
+
+    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"),
+    }
+    return _build_s2mel_additional_information(mel_seq, latent_seq, meta, context="full_payload")

vllm_omni/worker/gpu_ar_model_runner.py

@@ -303,6 +303,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             "Qwen3TTSCode2Wav",
             "CosyVoice3Model",
             "DyninOmniForConditionalGeneration",
+            "IndexTTS2TalkerForConditionalGeneration",
         }
         if getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS:
             self.init_omni_connectors(

vllm_omni/worker/gpu_generation_model_runner.py

@@ -64,6 +64,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
             "Qwen3TTSCode2Wav",
             "CosyVoice3Model",
             "DyninOmniForConditionalGeneration",
+            "IndexTTS2S2MelDecoder",
         }
         if getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS:
             self.init_omni_connectors(