vllm-omni/vllm_omni/model_executor/models/indextts2/configuration_indextts2.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import copy

from transformers import PretrainedConfig

INDEXTTS25_MIN_DURATION_FACTOR = 0.5
INDEXTTS25_MAX_DURATION_FACTOR = 2.0


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"},
        "tokenizer_file": "bpe.model",
        "semantic_codec_checkpoint": "semantic_codec.pth",
        "semantic_codec_type": "repcodec",
        "use_gpt_latent": True,
        "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)
        # vLLM kernel warmup reads the top-level vocabulary size even when
        # tokenizer initialization is skipped for the audio-code model.
        self.vocab_size = int(getattr(self, "vocab_size", 0) or self.gpt["number_mel_codes"])
        # 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
        )


class IndexTTS25Config(IndexTTS2Config):
    """Configuration for the official IndexTTS 2.5 checkpoint.

    IndexTTS 2.5 keeps the GPT/S2Mel backbone shapes used by IndexTTS 2,
    but switches text tokenization and speaker conditioning, and decodes
    semantic codes with EnhancedCodec. GPT latent conditioning is disabled by
    default upstream; the opt-in path is an experimental vLLM-Omni-specific
    latent variant with no official runnable reference output.
    """

    model_type = "indextts2_5"
    _DEFAULTS = copy.deepcopy(IndexTTS2Config._DEFAULTS)
    _DEFAULTS["gpt"] = {
        **_DEFAULTS["gpt"],
        "number_text_tokens": 60509,
    }
    _DEFAULTS.update(
        {
            "tokenizer_file": "multilingual_zh_ja_yue_char_del.tiktoken",
            "semantic_codec_checkpoint": "codec.pth",
            "semantic_codec_type": "enhanced",
            "use_gpt_latent": False,
            "version": 2.5,
        }
    )

    def __init__(self, **kwargs):
        # HF overrides are sometimes intentionally partial. Merge nested GPT
        # values with the 2.5 defaults so dimensions needed by the shared
        # implementation remain present while explicit values still win.
        explicit_gpt = kwargs.pop("gpt", None)
        merged_gpt = copy.deepcopy(self._DEFAULTS["gpt"])
        if explicit_gpt is not None:
            merged_gpt.update(explicit_gpt)
        super().__init__(gpt=merged_gpt, **kwargs)