diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index 84c93a2f3..f2d5e972b 100644 --- a/.buildkite/test-merge.yml +++ b/.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 diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 7c5db00d6..79c44d740 100644 --- a/docs/models/supported_models.md +++ b/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` | ✅︎ | ✅︎ | | | diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index 3da75bc0f..ea1d36051 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/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. diff --git a/examples/offline_inference/text_to_speech/indextts2/end2end.py b/examples/offline_inference/text_to_speech/indextts2/end2end.py new file mode 100644 index 000000000..858487ae1 --- /dev/null +++ b/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()) diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index 1c180440f..1a3096fab 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/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 --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. diff --git a/examples/online_serving/text_to_speech/indextts2/run_server.sh b/examples/online_serving/text_to_speech/indextts2/run_server.sh new file mode 100644 index 000000000..4a17b3b04 --- /dev/null +++ b/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" diff --git a/examples/online_serving/text_to_speech/indextts2/speech_client.py b/examples/online_serving/text_to_speech/indextts2/speech_client.py new file mode 100644 index 000000000..b938f0bd3 --- /dev/null +++ b/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() diff --git a/recipes/IndexTeam/IndexTTS-2.md b/recipes/IndexTeam/IndexTTS-2.md new file mode 100644 index 000000000..54c1119dd --- /dev/null +++ b/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," + }' \ + --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,", + "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. diff --git a/recipes/README.md b/recipes/README.md index 30f1577be..4da489c59 100644 --- a/recipes/README.md +++ b/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 | diff --git a/tests/assets/indextts2/ref_audio.wav b/tests/assets/indextts2/ref_audio.wav new file mode 100644 index 000000000..d13a4c847 Binary files /dev/null and b/tests/assets/indextts2/ref_audio.wav differ diff --git a/tests/core/sched/test_omni_scheduling_coordinator.py b/tests/core/sched/test_omni_scheduling_coordinator.py index 1b36cd784..068cbd362 100644 --- a/tests/core/sched/test_omni_scheduling_coordinator.py +++ b/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"), } diff --git a/tests/e2e/offline_inference/test_indextts2.py b/tests/e2e/offline_inference/test_indextts2.py new file mode 100644 index 000000000..6c2786efa --- /dev/null +++ b/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) diff --git a/tests/e2e/online_serving/test_indextts2.py b/tests/e2e/online_serving/test_indextts2.py new file mode 100644 index 000000000..c0fceb1ce --- /dev/null +++ b/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) diff --git a/tests/model_executor/stage_input_processors/test_indextts2.py b/tests/model_executor/stage_input_processors/test_indextts2.py new file mode 100644 index 000000000..59f14b414 --- /dev/null +++ b/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 diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index f7109f395..63548ef0f 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/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", + ), } diff --git a/vllm_omni/core/sched/omni_scheduling_coordinator.py b/vllm_omni/core/sched/omni_scheduling_coordinator.py index 4056fd938..af6c73d6d 100644 --- a/vllm_omni/core/sched/omni_scheduling_coordinator.py +++ b/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. diff --git a/vllm_omni/deploy/indextts2.yaml b/vllm_omni/deploy/indextts2.yaml new file mode 100644 index 000000000..75eb60188 --- /dev/null +++ b/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 diff --git a/vllm_omni/deploy/indextts2_low_latency.yaml b/vllm_omni/deploy/indextts2_low_latency.yaml new file mode 100644 index 000000000..621116ab2 --- /dev/null +++ b/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 diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 148412012..83bb5b1c1 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/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), diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 0ae947efb..b5b98a359 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/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: diff --git a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py index ca7b9612a..a2b9cce09 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py +++ b/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, diff --git a/vllm_omni/entrypoints/openai/tts_adapters/indextts2.py b/vllm_omni/entrypoints/openai/tts_adapters/indextts2.py new file mode 100644 index 000000000..641e380da --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/indextts2/__init__.py b/vllm_omni/model_executor/models/indextts2/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/indextts2/bigvgan_cuda_graph.py b/vllm_omni/model_executor/models/indextts2/bigvgan_cuda_graph.py new file mode 100644 index 000000000..c64a46404 --- /dev/null +++ b/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() diff --git a/vllm_omni/model_executor/models/indextts2/configuration_indextts2.py b/vllm_omni/model_executor/models/indextts2/configuration_indextts2.py new file mode 100644 index 000000000..504923b1a --- /dev/null +++ b/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 + ) diff --git a/vllm_omni/model_executor/models/indextts2/dit_cuda_graph.py b/vllm_omni/model_executor/models/indextts2/dit_cuda_graph.py new file mode 100644 index 000000000..4d32b8c1b --- /dev/null +++ b/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, + ) diff --git a/vllm_omni/model_executor/models/indextts2/gpt/__init__.py b/vllm_omni/model_executor/models/indextts2/gpt/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/indextts2/gpt/conformer/__init__.py b/vllm_omni/model_executor/models/indextts2/gpt/conformer/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/indextts2/gpt/conformer/attention.py b/vllm_omni/model_executor/models/indextts2/gpt/conformer/attention.py new file mode 100644 index 000000000..7b34ab381 --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/indextts2/gpt/conformer/embedding.py b/vllm_omni/model_executor/models/indextts2/gpt/conformer/embedding.py new file mode 100644 index 000000000..9c52df586 --- /dev/null +++ b/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) diff --git a/vllm_omni/model_executor/models/indextts2/gpt/conformer/subsampling.py b/vllm_omni/model_executor/models/indextts2/gpt/conformer/subsampling.py new file mode 100644 index 000000000..f8e0b7aa0 --- /dev/null +++ b/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] diff --git a/vllm_omni/model_executor/models/indextts2/gpt/conformer_encoder.py b/vllm_omni/model_executor/models/indextts2/gpt/conformer_encoder.py new file mode 100644 index 000000000..67eebbab3 --- /dev/null +++ b/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) + ] + ) diff --git a/vllm_omni/model_executor/models/indextts2/gpt/embeddings.py b/vllm_omni/model_executor/models/indextts2/gpt/embeddings.py new file mode 100644 index 000000000..bc9e8e466 --- /dev/null +++ b/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) diff --git a/vllm_omni/model_executor/models/indextts2/gpt/perceiver.py b/vllm_omni/model_executor/models/indextts2/gpt/perceiver.py new file mode 100644 index 000000000..98c7becfa --- /dev/null +++ b/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) diff --git a/vllm_omni/model_executor/models/indextts2/indextts2_s2mel_decoder.py b/vllm_omni/model_executor/models/indextts2/indextts2_s2mel_decoder.py new file mode 100644 index 000000000..a374f1b45 --- /dev/null +++ b/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, + ) diff --git a/vllm_omni/model_executor/models/indextts2/indextts2_talker.py b/vllm_omni/model_executor/models/indextts2/indextts2_talker.py new file mode 100644 index 000000000..2b1244ccc --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/indextts2_talker.py @@ -0,0 +1,1062 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""IndexTTS2 Stage 0: GPT-2 AR Talker with vLLM-native PagedAttention. + +Predicts mel codes autoregressively and collects hidden_states as latent +for Stage 1 (S2Mel + BigVGAN). +""" + +from __future__ import annotations + +import json +import math +import re +from collections.abc import Iterable +from typing import Any + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig +from vllm.distributed import get_pp_group +from vllm.logger import init_logger +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.gpt2 import GPT2Block +from vllm.model_executor.models.utils import ( + is_pp_missing_parameter, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_prefix, +) +from vllm.sequence import IntermediateTensors + +from vllm_omni.data_entry_keys import OmniPayload +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.utils.speaker_cache import get_speaker_cache + +from .configuration_indextts2 import IndexTTS2Config +from .gpt.conformer_encoder import ConformerEncoder +from .gpt.embeddings import LearnedPositionEmbeddings +from .gpt.perceiver import PerceiverResampler +from .preprocess_utils import ( + compute_fbank, + load_campplus, + load_qwen_emotion, + load_reference_audio, + load_semantic_codec, + load_wav2vec2, + resolve_model_file, + wav2vec_extract, +) + +logger = init_logger(__name__) + + +def _find_most_similar_cosine(query: torch.Tensor, matrix: torch.Tensor) -> int: + sims = F.cosine_similarity(query.float(), matrix.float(), dim=1) + return int(torch.argmax(sims).item()) + + +@support_torch_compile( + dynamic_arg_dims={ + "input_ids": 0, + "positions": 0, + "intermediate_tensors": 0, + "inputs_embeds": 0, + } +) +class IndexTTS2TalkerForConditionalGeneration(nn.Module): + """vLLM-native GPT-2 AR talker for IndexTTS2. + + Stage 0 of the two-stage pipeline. Predicts mel codes (8194 vocab) + and accumulates hidden_states as latent for Stage 1 S2Mel decoder. + """ + + 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] + gpt_cfg = self.config.gpt + + self.model_dim = gpt_cfg["model_dim"] + self.num_layers = gpt_cfg["layers"] + self.num_heads = gpt_cfg["heads"] + self.max_mel_tokens = gpt_cfg["max_mel_tokens"] + self.max_text_tokens = gpt_cfg["max_text_tokens"] + self.number_mel_codes = gpt_cfg["number_mel_codes"] + self.start_mel_token = gpt_cfg["start_mel_token"] + self.stop_mel_token = gpt_cfg["stop_mel_token"] + self.number_text_tokens = gpt_cfg.get("number_text_tokens", 12000) + self.condition_num_latent = gpt_cfg.get("condition_num_latent", 32) + + # --- Flags for vLLM-Omni framework --- + self.have_multimodal_outputs = True + self.has_preprocess = True + self.has_postprocess = True + self.requires_raw_input_tokens = True + self.enable_update_additional_information = True + # S2Mel consumes the complete code sequence only at request end, but + # the connector accumulator needs to observe each appendable AR delta. + # Keep runner request-end gating disabled and emit per-step rows from + # make_omni_output; the full-payload builder reconstructs the sequence. + self.omni_payload_at_request_end = False + self.omni_request_end_token_ids = () + self._speaker_cache = get_speaker_cache() + self.gpu_resident_buffer_keys: set[tuple[str, str]] = { + ("codes", "mel"), + ("hidden_states", "latent"), + ("meta", "mel_start_offset"), + ("meta", "latent_acc"), + ("meta", "mel_code_count"), + } + + # --- GPT-2 transformer (vLLM-native with PagedAttention) --- + # Build a GPT2Config-like object for GPT2Block + from transformers import GPT2Config as HFGpt2Config + + max_seq_len = self.max_mel_tokens + self.max_text_tokens + self.condition_num_latent + 4 + gpt2_config = HFGpt2Config( + vocab_size=self.number_mel_codes, + n_positions=max_seq_len, + n_ctx=max_seq_len, + n_embd=self.model_dim, + n_layer=self.num_layers, + n_head=self.num_heads, + activation_function="gelu_new", + layer_norm_epsilon=1e-5, + ) + # Store for GPT2Block construction + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + self.start_layer, self.end_layer, self.h = make_layers( + gpt2_config.n_layer, + lambda prefix: GPT2Block(gpt2_config, cache_config, quant_config, prefix=prefix), + prefix=f"{prefix}.h", + ) + self.ln_f = nn.LayerNorm(self.model_dim, eps=gpt2_config.layer_norm_epsilon) + self.final_norm = nn.LayerNorm(self.model_dim, eps=gpt2_config.layer_norm_epsilon) + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states"], self.model_dim + ) + + # --- Embeddings --- + self.text_embedding = nn.Embedding(self.number_text_tokens + 1, self.model_dim) + self.mel_embedding = nn.Embedding(self.number_mel_codes, self.model_dim) + self.text_pos_embedding = LearnedPositionEmbeddings(self.max_text_tokens + 2, self.model_dim) + self.mel_pos_embedding = LearnedPositionEmbeddings(self.max_mel_tokens + 2 + 1, self.model_dim) + self.speed_emb = nn.Embedding(2, self.model_dim) + self.emo_layer = nn.Linear(self.model_dim, self.model_dim) + self.emovec_layer = nn.Linear(1024, self.model_dim) + + # --- Mel head (logits) --- + if get_pp_group().is_last_rank: + self.mel_head = ParallelLMHead( + self.number_mel_codes, + self.model_dim, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "mel_head"), + ) + else: + self.mel_head = None # type: ignore[assignment] + self.logits_processor = LogitsProcessor(self.number_mel_codes) + + # --- Conditioning modules (speaker / emotion) --- + cond_cfg = gpt_cfg.get("condition_module", {}) + emo_cond_cfg = gpt_cfg.get("emo_condition_module", {}) + + self.conditioning_encoder = ConformerEncoder( + input_size=1024, + output_size=cond_cfg.get("output_size", 512), + linear_units=cond_cfg.get("linear_units", 2048), + attention_heads=cond_cfg.get("attention_heads", 8), + num_blocks=cond_cfg.get("num_blocks", 6), + input_layer=cond_cfg.get("input_layer", "conv2d2"), + ) + self.perceiver_encoder = PerceiverResampler( + self.model_dim, + dim_context=cond_cfg.get("output_size", 512), + ff_mult=cond_cfg.get("perceiver_mult", 2), + heads=cond_cfg.get("attention_heads", 8), + num_latents=self.condition_num_latent, + ) + + self.emo_conditioning_encoder = ConformerEncoder( + input_size=1024, + output_size=emo_cond_cfg.get("output_size", 512), + linear_units=emo_cond_cfg.get("linear_units", 1024), + attention_heads=emo_cond_cfg.get("attention_heads", 4), + num_blocks=emo_cond_cfg.get("num_blocks", 4), + input_layer=emo_cond_cfg.get("input_layer", "conv2d2"), + ) + self.emo_perceiver_encoder = PerceiverResampler( + 1024, + dim_context=emo_cond_cfg.get("output_size", 512), + ff_mult=emo_cond_cfg.get("perceiver_mult", 2), + heads=emo_cond_cfg.get("attention_heads", 4), + num_latents=1, + ) + + # Padding masks for perceiver cross-attention + self.cond_mask_pad = nn.ConstantPad1d((self.condition_num_latent, 0), True) + self.emo_cond_mask_pad = nn.ConstantPad1d((1, 0), True) + + # --- Lazy-loaded external models --- + self._w2v_stat: torch.Tensor | None = None + self._emo_matrix: torch.Tensor | None = None + self._spk_matrix: torch.Tensor | None = None + self._text_tokenizer: Any = None + + # Initialize embeddings per GPT-2 convention + for emb in [self.text_embedding, self.mel_embedding]: + emb.weight.data.normal_(mean=0.0, std=0.02) + self.speed_emb.weight.data.normal_(mean=0.0, std=0.0) + + # ------------------------------------------------------------------ + # vLLM required hooks + # ------------------------------------------------------------------ + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + return self.mel_embedding(input_ids) + + def _extract_mel_offsets(self, positions: torch.Tensor, kwargs: dict[str, Any]) -> torch.Tensor: + """Build per-token mel_start_offset from model_intermediate_buffer.""" + info_dicts = kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") or [] + if "runtime_additional_information" in kwargs and "model_intermediate_buffer" not in kwargs: + logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer") + if not info_dicts: + logger.warning_once("[OFFSET] no model_intermediate_buffer available; defaulting mel offsets to 0") + return torch.zeros_like(positions) + offsets = [] + for info in info_dicts: + if isinstance(info, dict): + meta = info.get("meta", {}) + val = meta.get("mel_start_offset", 0) + offsets.append(int(val) if not isinstance(val, int) else val) + else: + offsets.append(0) + return torch.tensor(offsets, dtype=positions.dtype, device=positions.device) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor | IntermediateTensors: + """AR transformer forward. + + During prefill: inputs_embeds is set by preprocess(), input_ids ignored. + During decode: inputs_embeds is mel_embedding(input_ids) + mel_pos. + """ + if get_pp_group().is_first_rank: + if inputs_embeds is None: + logger.warning("[FORWARD] inputs_embeds=None — fallback decode path") + mel_start_offset = self._extract_mel_offsets(positions, kwargs) + mel_pos = positions - mel_start_offset + mel_pos = torch.clamp(mel_pos, min=0) + inputs_embeds = self.mel_embedding(input_ids) + self.mel_pos_embedding.emb(mel_pos) + hidden_states = inputs_embeds + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + + for layer in self.h[self.start_layer : self.end_layer]: + hidden_states = layer(hidden_states) + + if not get_pp_group().is_last_rank: + return IntermediateTensors({"hidden_states": hidden_states}) + + return self.final_norm(self.ln_f(hidden_states)) + + def compute_logits( + self, + hidden_states: torch.Tensor | OmniOutput, + sampling_metadata: Any = None, + ) -> torch.Tensor | None: + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if hidden_states is None or self.mel_head is None: + return None + return self.logits_processor(self.mel_head, hidden_states) + + # ------------------------------------------------------------------ + # Omni multimodal output plumbing + # ------------------------------------------------------------------ + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: + """Collect mel_codes and hidden_states from intermediate buffer.""" + if isinstance(model_outputs, OmniOutput): + return model_outputs + + hidden = model_outputs + info_dicts = kwargs.get("model_intermediate_buffer") + if info_dicts is None: + info_dicts = kwargs.get("runtime_additional_information") or [] + if "runtime_additional_information" in kwargs and "model_intermediate_buffer" not in kwargs: + logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer") + + device = hidden.device + latent_dim = hidden.shape[-1] + _zero_codes = torch.zeros(0, 1, dtype=torch.long, device=device) + _zero_latent = torch.zeros(0, latent_dim, dtype=hidden.dtype, device=device) + + # Per-request lists: to_payload_element dispatches lists via + # element[idx], giving correct per-request mapping even during + # mixed prefill+decode steps. Zero-length placeholders for prefill + # requests are harmless — torch.cat ignores them in accumulation. + # + # Export semantics are appendable deltas for the existing full-payload + # accumulator: codes.mel is a 2-D row [1, 1], hidden_states.latent is + # the aligned one-row [1, D] slice from meta.latent_acc. + mel_codes_list: list[torch.Tensor] = [] + latent_list: list[torch.Tensor] = [] + meta_list: list[dict[str, Any]] = [] + + for info in info_dicts: + if not isinstance(info, dict): + mel_codes_list.append(_zero_codes) + latent_list.append(_zero_latent) + meta_list.append({}) + continue + codes = info.get("codes", {}) + mc = codes.get("mel") + info_meta = info.get("meta", {}) + # mel_len now comes from the int counter in meta, not from + # code_seq.shape[0] (codes.mel is a single-token delta). + mel_len = int(info_meta.get("mel_code_count", 0)) + if isinstance(mc, torch.Tensor) and mc.numel() > 0 and mel_len > 0: + code_seq = mc.reshape(-1) + mel_codes_list.append(code_seq[-1:].reshape(1, 1).contiguous()) + else: + mel_codes_list.append(_zero_codes) + + lat_acc = info_meta.get("latent_acc") + if isinstance(lat_acc, torch.Tensor) and lat_acc.numel() > 0 and mel_len > 0: + lat_seq = lat_acc.reshape(-1, lat_acc.shape[-1]) if lat_acc.ndim != 2 else lat_acc + idx = min(max(mel_len - 1, 0), int(lat_seq.shape[0]) - 1) + latent_list.append(lat_seq[idx : idx + 1].contiguous()) + else: + hs = info.get("hidden_states", {}) + lat = hs.get("latent") + if isinstance(lat, torch.Tensor) and lat.numel() > 0: + lat_seq = lat.reshape(-1, lat.shape[-1]) if lat.ndim != 2 else lat + latent_list.append(lat_seq[-1:].contiguous()) + else: + latent_list.append(_zero_latent) + + req_meta: dict[str, Any] = {} + if mel_len == 1: + s_ref = info_meta.get("S_ref") + ref_mel = info_meta.get("ref_mel") + style = info_meta.get("style") + if isinstance(s_ref, torch.Tensor): + req_meta["S_ref"] = s_ref + if isinstance(ref_mel, torch.Tensor): + req_meta["ref_mel"] = ref_mel + if isinstance(style, torch.Tensor): + req_meta["style"] = style + meta_list.append(req_meta) + + if not any(t.numel() > 0 for t in mel_codes_list): + return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) + + mm: OmniPayload = { + "codes": {"mel": mel_codes_list}, + "hidden_states": {"latent": latent_list}, + } + if any(meta_list): + mm["meta"] = meta_list + + return OmniOutput(text_hidden_states=hidden, multimodal_outputs=mm) + + # ------------------------------------------------------------------ + # preprocess / postprocess + # ------------------------------------------------------------------ + + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None, + **info_dict: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Build prompt embeddings for prefill; compute mel embeddings for decode. + + Prefill layout:: + + [conds(32) + emo_vec(1) + duration(2)] [text_emb + text_pos] [start_mel + mel_pos(0)] + + Decode: ``mel_embedding(token) + mel_pos_embedding(step)``. + """ + span_len = int(input_ids.shape[0]) + meta = info_dict.get("meta", {}) + + # --- Decode path: span_len==1 and mel_start_offset already set --- + if span_len == 1 and isinstance(meta, dict) and meta.get("mel_start_offset") is not None: + return self._preprocess_decode(input_ids, info_dict, meta) + + # --- Prefill path --- + return self._preprocess_prefill(input_ids, input_embeds, info_dict, span_len) + + def _preprocess_decode( + self, + input_ids: torch.Tensor, + info_dict: dict[str, Any], + meta: dict[str, Any], + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Lightweight decode step: embed token + track code count.""" + # Track code count via int counter instead of accumulating a full tensor + # (the connector reconstructs the full sequence from per-step deltas). + prev_count = meta.get("mel_code_count", 0) + if not isinstance(prev_count, int): + prev_count = int(prev_count) + decode_step = prev_count + 1 + + # Only keep the current token; no O(T) history accumulation. + new_codes = input_ids.detach() + + mel_pos = input_ids.new_full((1,), decode_step) + embeds = self.mel_embedding(input_ids) + self.mel_pos_embedding.emb(mel_pos) + + return ( + input_ids, + embeds, + { + "codes": {"mel": new_codes}, + "meta": {"mel_code_count": decode_step}, + }, + ) + + def _preprocess_prefill( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None, + info_dict: dict[str, Any], + span_len: int, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Full conditioning pipeline — runs once during prefill.""" + additional_information = info_dict.get("additional_information") + if isinstance(additional_information, dict): + merged: dict[str, Any] = {k: v for k, v in info_dict.items() if k != "additional_information"} + for k, v in additional_information.items(): + merged.setdefault(k, v) + info_dict = merged + + device = input_ids.device + + def _first(key, default=None): + v = info_dict.get(key) + return v[0] if isinstance(v, list) and v else default + + text = _first("text") + if not text: + raise ValueError("Missing additional_information.text for IndexTTS2 talker.") + + voice_path = _first("voice") + if voice_path is None: + raise ValueError("IndexTTS2 requires voice (reference audio path) in additional_information.") + + emo_audio_path = _first("emo_audio") + emo_text = _first("emo_text") + use_emo_text = bool(_first("use_emo_text", False)) + emo_vector = _first("emo_vector") + emo_alpha = float(_first("emo_alpha", 1.0)) + use_random = bool(_first("use_random", False)) + _raw_emo_voice = _first("emo_voice_name") + emo_voice_name = str(_raw_emo_voice).strip().lower() if _raw_emo_voice else None + + # --- Speaker cache lookup --- + _voice_name = _first("voice_name") + _voice_created_at = int(_first("voice_created_at", 0)) + _speaker_cache_key = None + _cached = None + if _voice_name: + _speaker_cache_key = self._speaker_cache.make_cache_key(_voice_name, "indextts2", _voice_created_at) + _cached = self._speaker_cache.get(_speaker_cache_key) + + # --- Load audio and extract features --- + # Keep the speaker-cache hit path truly lazy: cached speaker artifacts are + # sufficient for conditioning and same-speaker emotion reuse, so avoid + # loading/resampling reference audio and touching the Wav2Vec2 singleton + # unless a miss (or an uncached separate emotion audio) actually needs it. + wav_16k: torch.Tensor | None = None + w2v_model = None + w2v_proc = None + + if _cached is not None: + s_ref = _cached["S_ref"].to(device) + style = _cached["style"].to(device) + ref_mel = _cached["ref_mel"].to(device) + spk_cond_emb = _cached["spk_cond_emb"].to(device) + else: + wav_16k, wav_22k = load_reference_audio(voice_path, device, mode="speaker") + w2v_model, w2v_proc = load_wav2vec2(self.model_path, device) + self._ensure_w2v_stat_loaded(device) + spk_cond_emb = wav2vec_extract(wav_16k, w2v_model, w2v_proc, device, self._w2v_stat) + + semantic_codec = load_semantic_codec(self.model_path, self.config.semantic_codec, device) + with torch.no_grad(): + _, s_ref = semantic_codec.quantize(spk_cond_emb) # [B, T, 1024] quantized embeddings + + campplus = load_campplus(self.model_path, device) + fbank = compute_fbank(wav_16k, device) + with torch.no_grad(): + style = campplus(fbank) # [1, 192] + + ref_mel = self._compute_mel_22k(wav_22k, device) # [1, 80, T_ref] + + if _speaker_cache_key is not None: + self._speaker_cache.put( + _speaker_cache_key, + { + "S_ref": s_ref.detach().cpu().contiguous(), + "style": style.detach().cpu().contiguous(), + "ref_mel": ref_mel.detach().cpu().contiguous(), + "spk_cond_emb": spk_cond_emb.detach().cpu().contiguous(), + }, + ) + + model_dtype = next(self.conditioning_encoder.parameters()).dtype + spk_cond_emb = spk_cond_emb.to(device=device, dtype=model_dtype) + spk_lens = torch.tensor([spk_cond_emb.shape[1]], device=device, dtype=torch.long) + speech_cond, mask = self.conditioning_encoder(spk_cond_emb, spk_lens) + conds_mask = self.cond_mask_pad(mask.squeeze(1)) + conds = self.perceiver_encoder(speech_cond, conds_mask) # [1, 32, D] + + emo_vec = self._compute_emotion_vector( + wav_16k=wav_16k, + speaker_audio_path=voice_path, + emo_audio_path=emo_audio_path, + emo_voice_name=emo_voice_name, + main_text=text, + use_emo_text=use_emo_text, + emo_text=emo_text, + emo_vector=emo_vector, + emo_alpha=emo_alpha, + use_random=use_random, + style=style, + spk_cond_emb=spk_cond_emb, + w2v_model=w2v_model, + w2v_proc=w2v_proc, + device=device, + ) + + speed_zero = torch.zeros(1, device=device, dtype=torch.long) + speed_one = torch.ones(1, device=device, dtype=torch.long) + duration_emb = self.speed_emb(speed_zero) + duration_emb_half = self.speed_emb(speed_one) + + conds_with_emo = conds + emo_vec.unsqueeze(1) # [1, 32, D] + conds_prefix = torch.cat( + [conds_with_emo, duration_emb_half.unsqueeze(1), duration_emb.unsqueeze(1)], + dim=1, + ) # [1, 34, D] + + text_tokens = self._tokenize_text(text, device) # [1, L] + text_emb = self.text_embedding(text_tokens) + self.text_pos_embedding(text_tokens) + + start_mel = self.mel_embedding(torch.tensor([[self.start_mel_token]], device=device, dtype=torch.long)) + start_mel_pos = self.mel_pos_embedding.get_fixed_embedding(0, device) + start_mel = start_mel + start_mel_pos # [1, 1, D] + + inputs_embeds = torch.cat([conds_prefix, text_emb, start_mel], dim=1).squeeze(0) + mel_start_offset = int(conds_prefix.shape[1]) + int(text_emb.shape[1]) + prefill_mel_index = int(inputs_embeds.shape[0]) - 1 + + if int(inputs_embeds.shape[0]) < span_len: + pad_n = span_len - int(inputs_embeds.shape[0]) + pad_emb = torch.zeros(pad_n, self.model_dim, device=device, dtype=inputs_embeds.dtype) + inputs_embeds = torch.cat([inputs_embeds, pad_emb], dim=0) + elif int(inputs_embeds.shape[0]) > span_len: + inputs_embeds = inputs_embeds[:span_len] + + input_ids_out = torch.full_like(input_ids, self.start_mel_token) + + info_update: dict[str, Any] = { + "meta": { + "mel_start_offset": mel_start_offset, + "prefill_mel_index": prefill_mel_index, + "mel_code_count": 0, + "S_ref": s_ref.cpu().contiguous(), + "ref_mel": ref_mel.cpu().contiguous(), + "style": style.cpu().contiguous(), + "latent_acc": torch.zeros(0, self.model_dim, dtype=inputs_embeds.dtype, device=device), + }, + "codes": {"mel": torch.zeros(0, dtype=torch.long, device=device)}, + "hidden_states": {"latent": torch.zeros(0, self.model_dim, dtype=inputs_embeds.dtype, device=device)}, + } + + return input_ids_out, inputs_embeds, info_update + + def postprocess( + self, + hidden_states: torch.Tensor, + multimodal_outputs: Any = None, + **kwargs: Any, + ) -> dict[str, Any]: + """Store current decode-step hidden state for Stage 1. + + Only the current row is kept; the connector's full-payload accumulator + reconstructs the complete latent sequence from per-step deltas emitted + by make_omni_output. + """ + if hidden_states.shape[0] != 1: + meta = kwargs.get("meta", {}) + prefill_mel_index = int(meta.get("prefill_mel_index", hidden_states.shape[0] - 1)) + prefill_mel_index = min(max(prefill_mel_index, 0), hidden_states.shape[0] - 1) + return {"meta": {"latent_acc": hidden_states[prefill_mel_index : prefill_mel_index + 1].detach()}} + + return {"meta": {"latent_acc": hidden_states.detach()}} + + def _compute_mel_22k(self, wav_22k: torch.Tensor, device: torch.device) -> torch.Tensor: + """Compute 80-band mel spectrogram at 22.05kHz for Stage 1.""" + s2mel_cfg = self.config.s2mel.get("preprocess_params", {}) + sr = s2mel_cfg.get("sr", 22050) + spect = s2mel_cfg.get("spect_params", {}) + n_fft = spect.get("n_fft", 1024) + hop_length = spect.get("hop_length", 256) + win_length = spect.get("win_length", 1024) + n_mels = spect.get("n_mels", 80) + + # mel_spectrogram is pure torch with per-device caches — run on the + # wav's device (GPU after load_reference_audio) instead of CPU. + wav = wav_22k.float() + if wav.ndim == 1: + wav = wav.unsqueeze(0) + + from .s2mel.modules.audio import mel_spectrogram as mel_fn + + mel = mel_fn(wav, n_fft, n_mels, sr, hop_length, win_length, 0, None) + return mel.to(device=device) # [1, 80, T] + + # ------------------------------------------------------------------ + # Emotion conditioning + # ------------------------------------------------------------------ + + def _compute_emotion_vector( + self, + *, + wav_16k: torch.Tensor | None, + speaker_audio_path: Any, + emo_audio_path: str | None, + emo_voice_name: str | None, + main_text: str, + use_emo_text: bool, + emo_text: str | None, + emo_vector: Any, + emo_alpha: float, + use_random: bool, + style: torch.Tensor, + spk_cond_emb: torch.Tensor, + w2v_model: Any, + w2v_proc: Any, + device: torch.device, + ) -> torch.Tensor: + """Full emotion control logic aligned with official infer_v2.py. Returns [1, D].""" + # --- Mutual exclusion: text/vector guidance clears emotion reference audio --- + # Mirrors official infer_v2.py: use_emo_text is evaluated before explicit + # vectors, so text guidance wins if both are supplied by an internal caller. + if use_emo_text: + if emo_text is None: + emo_text = main_text + emo_vector = self._predict_emotion_from_text(emo_text, device) + if use_emo_text or emo_vector is not None: + emo_audio_path = None + + # --- Validate emo_vector (user-provided or text-predicted) --- + if emo_vector is not None: + if isinstance(emo_vector, torch.Tensor): + emo_vector = emo_vector.tolist() + elif isinstance(emo_vector, np.ndarray): + emo_vector = emo_vector.tolist() + if not isinstance(emo_vector, list) or len(emo_vector) != len(self._DESIRED_ORDER): + raise ValueError( + "IndexTTS2 emo_vector must contain 8 values in order: " + "happy, angry, sad, afraid, disgusted, melancholic, surprised, calm" + ) + if any(isinstance(v, bool) for v in emo_vector): + raise ValueError("IndexTTS2 emo_vector values must be numeric") + try: + emo_vector = [float(v) for v in emo_vector] + except (TypeError, ValueError) as exc: + raise ValueError("IndexTTS2 emo_vector values must be numeric") from exc + if not all(math.isfinite(v) for v in emo_vector): + raise ValueError("IndexTTS2 emo_vector values must be finite") + if not all(0.0 <= v <= 1.2 for v in emo_vector): + raise ValueError("IndexTTS2 emo_vector values must be between 0.0 and 1.2") + + # Alpha pre-scaling (official infer_v2.py:414-418): + # when emo_alpha != 1.0, scale emo_vector down before overlay. + emo_vector_scale = max(0.0, min(1.0, emo_alpha)) + if emo_vector_scale != 1.0: + emo_vector = [int(x * emo_vector_scale * 10000) / 10000 for x in emo_vector] + + # --- Path B: emo_vector (8-dim) → spk_matrix cosine lookup --- + emovec_mat = None + weight_sum = 0.0 + if emo_vector is not None: + weight_sum = sum(emo_vector) + vec = torch.tensor(emo_vector, device=device, dtype=torch.float32) + emovec_mat = self._compute_emo_from_distribution(vec, style, device, use_random=use_random) + + # --- Audio-based emotion with alpha blending --- + if emo_audio_path is None: + # No separate emotion audio: reuse speaker w2v-bert features instead of + # reloading + re-extracting the same audio (saves ~0.5s per prefill). + # Official code resamples orig→16k directly for emotion vs orig→22k→16k + # for speaker; the resulting w2v features are near-identical. + effective_alpha = 1.0 + emo_cond_emb = spk_cond_emb + else: + effective_alpha = emo_alpha + # emotion audio cache (enabled when emo_voice_name is provided) + _emo_cache_key = None + _emo_cached = None + if emo_voice_name: + _emo_cache_key = self._speaker_cache.make_cache_key(emo_voice_name, "indextts2_emo", 0) + _emo_cached = self._speaker_cache.get(_emo_cache_key) + + if _emo_cached is not None: + emo_cond_emb = _emo_cached["spk_cond_emb"].to(device) + logger.debug("[EMO] cache HIT for %s", emo_voice_name) + else: + emo_wav_16k, _ = load_reference_audio(emo_audio_path, device, mode="emotion") + if w2v_model is None or w2v_proc is None: + w2v_model, w2v_proc = load_wav2vec2(self.model_path, device) + self._ensure_w2v_stat_loaded(device) + emo_cond_emb = wav2vec_extract(emo_wav_16k, w2v_model, w2v_proc, device, self._w2v_stat) + if _emo_cache_key is not None: + self._speaker_cache.put( + _emo_cache_key, + { + "spk_cond_emb": emo_cond_emb.detach().cpu().contiguous(), + }, + ) + logger.debug("[EMO] cache MISS, cached for %s", emo_voice_name) + + emovec = self._merge_emovec(spk_cond_emb, emo_cond_emb, effective_alpha) + + # --- emo_vector overlay (official infer_v2.py:560-561) --- + # emovec_mat is raw feat2.pt weighted sum [1, model_dim], NOT projected. + # Mix: emovec_mat contributes the emotion prototype; (1-sum)*emovec + # preserves the audio-based emotion as a "base color". + if emovec_mat is not None: + emovec = emovec_mat.to(dtype=emovec.dtype) + (1.0 - weight_sum) * emovec + + return emovec + + _MELANCHOLIC_WORDS = frozenset( + { + "低落", + "melancholy", + "melancholic", + "depression", + "depressed", + "gloomy", + } + ) + + _DESIRED_ORDER = ["高兴", "愤怒", "悲伤", "恐惧", "反感", "低落", "惊讶", "自然"] + + def _predict_emotion_from_text(self, text: str, device: torch.device) -> list[float] | None: + """Use QwenEmotion CausalLM to predict 8-dim emotion vector from text (aligned with official).""" + model, tokenizer = load_qwen_emotion( + self.model_path, + device, + trust_remote_code=self.vllm_config.model_config.trust_remote_code, + ) + if model is None or tokenizer is None: + return None + + messages = [ + {"role": "system", "content": "文本情感分类"}, + {"role": "user", "content": text}, + ] + prompt_text = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + model_inputs = tokenizer([prompt_text], return_tensors="pt").to(model.device) + + with torch.no_grad(): + generated_ids = model.generate( + **model_inputs, + max_new_tokens=512, + pad_token_id=tokenizer.eos_token_id, + ) + output_ids = generated_ids[0][len(model_inputs.input_ids[0]) :].tolist() + + try: + index = len(output_ids) - output_ids[::-1].index(151668) + except ValueError: + index = 0 + content_str = tokenizer.decode(output_ids[index:], skip_special_tokens=True) + + try: + content = json.loads(content_str) + except json.JSONDecodeError: + content = { + m.group(1): float(m.group(2)) for m in re.finditer(r'([^\s":.,]+?)"?\s*:\s*([\d.]+)', content_str) + } + + text_lower = text.lower() + if any(w in text_lower for w in self._MELANCHOLIC_WORDS): + content["悲伤"], content["低落"] = content.get("低落", 0.0), content.get("悲伤", 0.0) + + max_score, min_score = 1.2, 0.0 + emo_vector = [] + for cn_key in self._DESIRED_ORDER: + val = content.get(cn_key, 0.0) + if isinstance(val, str): + try: + val = float(val) + except ValueError: + val = 0.0 + emo_vector.append(max(min_score, min(max_score, val))) + + if all(v <= 0.0 for v in emo_vector): + emo_vector[7] = 1.0 # default calm + + return emo_vector + + def _compute_emo_from_distribution( + self, + emo_dist: torch.Tensor, + style: torch.Tensor, + device: torch.device, + use_random: bool = False, + ) -> torch.Tensor: + """From 8-dim emotion distribution + CAMPPlus style, build emovec_mat via spk_matrix cosine lookup. + + Returns [1, model_dim] (feat2.pt rows are already model_dim). + """ + self._ensure_matrix_loaded("_emo_matrix", "emo_matrix", "Emotion matrix") + self._ensure_matrix_loaded("_spk_matrix", "spk_matrix", "Speaker matrix") + if self._emo_matrix is None or self._spk_matrix is None: + return torch.zeros(1, self.model_dim, device=device) + + weight_vector = emo_dist.to(device=device, dtype=torch.float32) + if weight_vector.ndim == 2: + weight_vector = weight_vector.squeeze(0) # [8] + + if use_random: + indices = [int(torch.randint(0, sub.shape[0], (1,), device=device).item()) for sub in self._spk_matrix] + else: + indices = [_find_most_similar_cosine(style, sub.to(device)) for sub in self._spk_matrix] + emo_rows = torch.cat( + [sub.to(device)[idx].unsqueeze(0) for idx, sub in zip(indices, self._emo_matrix)], dim=0 + ) # [8, model_dim] + + return (weight_vector.unsqueeze(1) * emo_rows).sum(0).unsqueeze(0) # [1, model_dim] + + def _compute_audio_emo_vec(self, cond_emb: torch.Tensor) -> torch.Tensor: + """Conformer → Perceiver → emovec_layer → emo_layer. Input: [1, T, 1024], output: [1, D].""" + emo_dtype = next(self.emo_conditioning_encoder.parameters()).dtype + cond_emb = cond_emb.to(device=cond_emb.device, dtype=emo_dtype) + cond_lens = torch.tensor([cond_emb.shape[1]], device=cond_emb.device, dtype=torch.long) + emo_cond, emo_mask = self.emo_conditioning_encoder(cond_emb, cond_lens) + emo_conds_mask = self.emo_cond_mask_pad(emo_mask.squeeze(1)) + emo_percept = self.emo_perceiver_encoder(emo_cond, emo_conds_mask) # [1, 1, 1024] + emo_vec_syn = self.emovec_layer(emo_percept.squeeze(1)) # [1, D] + return self.emo_layer(emo_vec_syn) # [1, D] + + def _merge_emovec( + self, + spk_cond_emb: torch.Tensor, + emo_cond_emb: torch.Tensor, + alpha: float, + ) -> torch.Tensor: + """Alpha-blend speaker and emotion audio vectors: base + alpha * (emo - base).""" + if emo_cond_emb is spk_cond_emb or alpha == 1.0: + return self._compute_audio_emo_vec(emo_cond_emb) + if alpha == 0.0: + return self._compute_audio_emo_vec(spk_cond_emb) + base_vec = self._compute_audio_emo_vec(spk_cond_emb) + emo_vec = self._compute_audio_emo_vec(emo_cond_emb) + return base_vec + alpha * (emo_vec - base_vec) # [1, D] + + # ------------------------------------------------------------------ + # Text tokenizer + # ------------------------------------------------------------------ + + def _tokenize_text(self, text: str, device: torch.device) -> torch.Tensor: + """Tokenize text using the BPE tokenizer, add start/stop tokens.""" + if self._text_tokenizer is None: + from .tokenizer import IndexTTS2Tokenizer + + bpe_path = resolve_model_file(self.model_path, "bpe.model") + if bpe_path is None: + raise FileNotFoundError(f"BPE model not found in {self.model_path}") + self._text_tokenizer = IndexTTS2Tokenizer(bpe_path, model_dir=self.model_path) + + token_ids = self._text_tokenizer.encode(text, add_special_tokens=False) + + # Wrap with start/stop text tokens + start_text = 0 + stop_text = 1 + token_ids = [start_text] + token_ids + [stop_text] + return torch.tensor([token_ids], device=device, dtype=torch.long) + + # ------------------------------------------------------------------ + # Lazy loading helpers + # ------------------------------------------------------------------ + + def _ensure_w2v_stat_loaded(self, device: torch.device) -> None: + if self._w2v_stat is not None: + return + stat_path = resolve_model_file(self.model_path, self.config.w2v_stat) + if stat_path is not None: + raw = torch.load(stat_path, map_location="cpu", weights_only=True) + if isinstance(raw, dict): + self._w2v_stat = torch.stack([raw["mean"], raw["var"]]) # [2, 1024] + else: + self._w2v_stat = raw + else: + logger.warning("Wav2Vec2-BERT stat file not found in %s", self.model_path) + + def _ensure_matrix_loaded(self, attr: str, cfg_key: str, label: str) -> None: + if getattr(self, attr) is not None: + return + path = resolve_model_file(self.model_path, getattr(self.config, cfg_key)) + if path is not None: + raw = torch.load(path, map_location="cpu", weights_only=True) + setattr(self, attr, torch.split(raw, list(self.config.emo_num))) + else: + logger.warning("%s not found in %s", label, self.model_path) + + # ------------------------------------------------------------------ + # Weight loading + # ------------------------------------------------------------------ + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights from IndexTTS2 checkpoint (gpt.pth). + + Bypasses vllm's default .pt iterator because the model directory + contains raw-tensor files (feat1.pt, feat2.pt) that are not state + dicts and would crash ``pt_weights_iterator``. + + Weight mapping from gpt.pth checkpoint → vLLM model params: + gpt.h.{i}.* → h.{i}.* (strip ``gpt.`` prefix) + gpt.ln_f.* → ln_f.* + final_norm.* → final_norm.* + text_head.* → (skipped) + → identity + """ + _ = weights # don't iterate — raw .pt files crash the default iterator + + ckpt_path = resolve_model_file(self.model_path, self.config.gpt_checkpoint) + if ckpt_path is None: + raise FileNotFoundError(f"IndexTTS2 GPT checkpoint {self.config.gpt_checkpoint!r} not found") + state = torch.load(ckpt_path, map_location="cpu", weights_only=True) + + prefix_map = { + "gpt.h.": "h.", + "gpt.ln_f.": "ln_f.", + "final_norm.": "final_norm.", + "text_head.": "_skip_text_head.", + } + + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded_params: set[str] = set() + + for name, loaded_weight in state.items(): + # Skip attention masks + if ".attn.bias" in name or ".attn.masked_bias" in name: + continue + # Skip wpe (null position embeddings) + if ".wpe." in name or ".wte." in name: + continue + + # Remap checkpoint name to model name + 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 + + # Skip unmapped / intentionally skipped weights + if mapped_name.startswith("_skip_"): + continue + + if is_pp_missing_parameter(mapped_name, self): + continue + + if mapped_name not in params_dict: + logger.debug("Skipping unrecognized weight: %s → %s", name, mapped_name) + continue + + param = params_dict[mapped_name] + + # GPT-2 Conv1D → Linear transpose + for conv1d_name in ["c_attn", "c_proj", "c_fc"]: + if conv1d_name not in mapped_name: + continue + if not mapped_name.endswith(".weight"): + continue + loaded_weight = loaded_weight.t() + + 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 IndexTTS2TalkerForConditionalGeneration", + len(loaded_params), + ) + required_prefixes = [ + "text_embedding.", + "mel_embedding.", + "text_pos_embedding.", + "mel_pos_embedding.", + "speed_emb.", + "emo_layer.", + "emovec_layer.", + "conditioning_encoder.", + "perceiver_encoder.", + "emo_conditioning_encoder.", + "emo_perceiver_encoder.", + "h.", + "ln_f.", + "final_norm.", + ] + if self.mel_head is not None: + required_prefixes.append("mel_head.") + 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 GPT checkpoint did not load required parameter groups: " + ", ".join(missing_prefixes) + ) + + # Ensure all sub-modules on correct device (some nn.Parameter created + # with torch.Tensor() end up on CPU even when model is on CUDA). + target_device = next((p.device for p in self.h.parameters()), torch.device("cpu")) + for mod in [ + self.conditioning_encoder, + self.emo_conditioning_encoder, + self.perceiver_encoder, + self.emo_perceiver_encoder, + ]: + mod.to(target_device) + + self._warmup_external_models(target_device) + return loaded_params + + def _warmup_external_models(self, device: torch.device) -> None: + """Preload external models to eliminate first-request latency.""" + try: + load_wav2vec2(self.model_path, device) + load_semantic_codec(self.model_path, self.config.semantic_codec, device) + load_campplus(self.model_path, device) + self._ensure_w2v_stat_loaded(device) + logger.info("External models preloaded: Wav2Vec2, RepCodec, CAMPPlus") + except Exception as e: + logger.warning("Failed to preload external models: %s", e) diff --git a/vllm_omni/model_executor/models/indextts2/pipeline.py b/vllm_omni/model_executor/models/indextts2/pipeline.py new file mode 100644 index 000000000..2ce02cee8 --- /dev/null +++ b/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}, + ), + ), +) diff --git a/vllm_omni/model_executor/models/indextts2/preprocess_utils.py b/vllm_omni/model_executor/models/indextts2/preprocess_utils.py new file mode 100644 index 000000000..452bce9bb --- /dev/null +++ b/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) + ) diff --git a/vllm_omni/model_executor/models/indextts2/prompt_utils.py b/vllm_omni/model_executor/models/indextts2/prompt_utils.py new file mode 100644 index 000000000..c33202785 --- /dev/null +++ b/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) diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/__init__.py b/vllm_omni/model_executor/models/indextts2/s2mel/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/__init__.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/audio.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/audio.py new file mode 100644 index 000000000..21bb45d73 --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/bigvgan.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/bigvgan.py new file mode 100644 index 000000000..5876bb147 --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/commons.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/commons.py new file mode 100644 index 000000000..dd86f5847 --- /dev/null +++ b/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() diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/diffusion_transformer.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/diffusion_transformer.py new file mode 100644 index 000000000..f73fdc800 --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/encodec.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/encodec.py new file mode 100644 index 000000000..a41141dcb --- /dev/null +++ b/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) diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/flow_matching.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/flow_matching.py new file mode 100644 index 000000000..5192867af --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/gpt_fast/__init__.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/gpt_fast/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/gpt_fast/model.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/gpt_fast/model.py new file mode 100644 index 000000000..c410b6387 --- /dev/null +++ b/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) diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/length_regulator.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/length_regulator.py new file mode 100644 index 000000000..1f2757781 --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/wavenet.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/wavenet.py new file mode 100644 index 000000000..4c5381766 --- /dev/null +++ b/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) diff --git a/vllm_omni/model_executor/models/indextts2/tokenizer.py b/vllm_omni/model_executor/models/indextts2/tokenizer.py new file mode 100644 index 000000000..be20d02be --- /dev/null +++ b/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, + ) diff --git a/vllm_omni/model_executor/models/indextts2/utils/__init__.py b/vllm_omni/model_executor/models/indextts2/utils/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/indextts2/utils/campplus/__init__.py b/vllm_omni/model_executor/models/indextts2/utils/campplus/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/indextts2/utils/campplus/dtdnn.py b/vllm_omni/model_executor/models/indextts2/utils/campplus/dtdnn.py new file mode 100644 index 000000000..71766d04a --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/indextts2/utils/campplus/layers.py b/vllm_omni/model_executor/models/indextts2/utils/campplus/layers.py new file mode 100644 index 000000000..4919cf16e --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/indextts2/utils/common.py b/vllm_omni/model_executor/models/indextts2/utils/common.py new file mode 100644 index 000000000..c969dc159 --- /dev/null +++ b/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"") + + words = line.split() + # restore english sentences + sent_placeholder_pattern = re.compile(r"()") + 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 diff --git a/vllm_omni/model_executor/models/indextts2/utils/front.py b/vllm_omni/model_executor/models/indextts2/utils/front.py new file mode 100644 index 000000000..3c3bd4cb7 --- /dev/null +++ b/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 "" + + @property + def pad_token(self): + return None + + @property + def bos_token(self): + return "" + + @property + def eos_token(self): + return "" + + @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, + ) diff --git a/vllm_omni/model_executor/models/indextts2/utils/maskgct/__init__.py b/vllm_omni/model_executor/models/indextts2/utils/maskgct/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/indextts2/utils/maskgct/quantize.py b/vllm_omni/model_executor/models/indextts2/utils/maskgct/quantize.py new file mode 100644 index 000000000..e4fcc6c79 --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/indextts2/utils/maskgct/repcodec_model.py b/vllm_omni/model_executor/models/indextts2/utils/maskgct/repcodec_model.py new file mode 100644 index 000000000..fbccdb363 --- /dev/null +++ b/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) diff --git a/vllm_omni/model_executor/models/indextts2/utils/maskgct/vocos.py b/vllm_omni/model_executor/models/indextts2/utils/maskgct/vocos.py new file mode 100644 index 000000000..ce735a9ad --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index d057502c2..279de283d 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/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", diff --git a/vllm_omni/model_executor/stage_input_processors/indextts2.py b/vllm_omni/model_executor/stage_input_processors/indextts2.py new file mode 100644 index 000000000..9d38c8f2a --- /dev/null +++ b/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") diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index a23d6f4e2..23e6a4aac 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/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( diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index bc3d6ea2d..07f394842 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/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(