Files
vllm-omni/vllm_omni/model_executor/models/indextts2/configuration_indextts2.py
T
BeatSeat 044240d8eb [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>
2026-06-22 19:24:01 +00:00

157 lines
5.4 KiB
Python

# 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
)