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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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