[New Model] Support IndexTTS 2.5 (#5957)

Signed-off-by: Sysy <1370724210@qq.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

vllm_omni/model_executor/models/indextts2/codec/quantizer.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: MIT
+# SPDX-FileCopyrightText: Copyright contributors to the Amphion project
+"""Decode-only factorized residual vector quantizer.
+
+Adapted from the Amphion codec implementation vendored by IndexTTS 2.5.
+Training-only paths are intentionally omitted.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils import weight_norm
+
+
+def _weight_norm_conv1d(*args, **kwargs) -> nn.Conv1d:
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+class FactorizedVectorQuantizeDecoder(nn.Module):
+    def __init__(
+        self,
+        *,
+        input_dim: int,
+        codebook_size: int,
+        codebook_dim: int,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.codebook_dim = codebook_dim
+        if input_dim != codebook_dim:
+            # Keep upstream module names and weight-normalized parameter layout
+            # so codec.pth keys load without conversion.
+            self.in_project = _weight_norm_conv1d(
+                input_dim,
+                codebook_dim,
+                kernel_size=1,
+            )
+            self.out_project = _weight_norm_conv1d(
+                codebook_dim,
+                input_dim,
+                kernel_size=1,
+            )
+        else:
+            self.in_project = nn.Identity()
+            self.out_project = nn.Identity()
+        self.codebook = nn.Embedding(codebook_size, codebook_dim)
+
+    def vq2emb(self, codes: torch.Tensor) -> torch.Tensor:
+        embedding = F.embedding(codes, self.codebook.weight).transpose(1, 2)
+        return self.out_project(embedding)
+
+
+class ResidualVQDecoder(nn.Module):
+    def __init__(
+        self,
+        *,
+        input_dim: int,
+        num_quantizers: int,
+        codebook_size: int,
+        codebook_dim: int,
+    ):
+        super().__init__()
+        self.num_quantizers = num_quantizers
+        self.quantizers = nn.ModuleList(
+            [
+                FactorizedVectorQuantizeDecoder(
+                    input_dim=input_dim,
+                    codebook_size=codebook_size,
+                    codebook_dim=codebook_dim,
+                )
+                for _ in range(num_quantizers)
+            ]
+        )
+
+    def vq2emb(
+        self,
+        codes: torch.Tensor,
+        n_quantizers: int | None = None,
+    ) -> torch.Tensor:
+        count = self.num_quantizers if n_quantizers is None else n_quantizers
+        quantized: torch.Tensor | None = None
+        for index, quantizer in enumerate(self.quantizers):
+            if index >= count:
+                break
+            value = quantizer.vq2emb(codes[index])
+            quantized = value if quantized is None else quantized + value
+        if quantized is None:
+            raise ValueError("EnhancedCodec requires at least one quantizer")
+        return quantized