Upload 6 files

config.json

@@ -0,0 +1,60 @@
+{
+  "_name_or_path": "AstraMindAI/xtts2",
+  "architectures": [
+    "Xtts"
+  ],
+  "auto_map": {
+    "AutoConfig": "AstraMindAI/xtts2--xtts2_config.XTTSConfig",
+    "AutoModelForCausalLM": "AstraMindAI/xtts2--xtts2_modeling.Xtts"
+  },
+  "cond_d_vector_in_each_upsampling_layer": true,
+  "d_vector_dim": 512,
+  "decoder_input_dim": 1024,
+  "input_sample_rate": 22050,
+  "model_type": "xtts_hifigan",
+  "output_hop_length": 256,
+  "output_sample_rate": 24000,
+  "resblock_dilation_sizes": [
+    [
+      1,
+      3,
+      5
+    ],
+    [
+      1,
+      3,
+      5
+    ],
+    [
+      1,
+      3,
+      5
+    ]
+  ],
+  "resblock_kernel_sizes": [
+    3,
+    7,
+    11
+  ],
+  "speaker_encoder_config": {
+    "model_config": null,
+    "model_name": "speaker_encoder",
+    "preprocess_config": null,
+    "speaker_embedding_dim": 512,
+    "use_torch_spec": true
+  },
+  "transformers_version": "4.45.1",
+  "upsample_initial_channel": 512,
+  "upsample_kernel_sizes": [
+    16,
+    16,
+    4,
+    4
+  ],
+  "upsample_rates": [
+    8,
+    8,
+    2,
+    2
+  ]
+}

gpt_config.py

@@ -0,0 +1,172 @@
+from dataclasses import asdict, dataclass, field
+from typing import Dict, Optional, List
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+@dataclass
+class XTTSAudioConfig:
+    """Configuration for audio processing parameters"""
+    sample_rate: int = 22050
+    output_sample_rate: int = 24000
+    mel_channels: int = 80
+    hop_length: int = 256
+    win_length: int = 1024
+    n_fft: int = 1024
+    fmin: int = 0
+    fmax: int = 8000
+    power: float = 1.0
+    mel_norms_file: Optional[str] = None
+
+
+class XTTSGPTConfig(PretrainedConfig):
+    """Configuration class for the GPT component of XTTS"""
+    model_type = "xtts_gpt"
+
+    def __init__(
+            self,
+            # Model architecture
+            vocab_size: int = 256,
+            num_chars: int = 255,
+
+            # GPT parameters
+            gpt_batch_size: int = 1,
+            gpt_max_audio_tokens: int = 605,
+            gpt_max_text_tokens: int = 402,
+            gpt_max_prompt_tokens: int = 70,
+            gpt_layers: int = 30,
+            gpt_n_model_channels: int = 1024,
+            gpt_n_heads: int = 16,
+            gpt_number_text_tokens: int = 6681,
+            gpt_start_text_token: Optional[int] = None,
+            gpt_stop_text_token: Optional[int] = None,
+            gpt_num_audio_tokens: int = 1026,
+            gpt_start_audio_token: int = 1024,
+            gpt_stop_audio_token: int = 1025,
+            gpt_code_stride_len: int = 1024,
+            gpt_use_masking_gt_prompt_approach: bool = True,
+            gpt_use_perceiver_resampler: bool = True,
+            gpt_checkpointing: bool = False,
+            gpt_train_solo_embeddings: bool = False,
+
+            # Training parameters
+            enable_redaction: bool = False,
+            kv_cache: bool = True,
+            perceiver_cond_length_compression: int = 256,
+            label_smoothing: float = 0.0,
+
+            # Generation parameters
+            temperature: float = 0.75,
+            length_penalty: float = 1.0,
+            repetition_penalty: float = 5.0,
+            top_k: int = 50,
+            top_p: float = 0.85,
+            gpt_cond_len: int = 30,
+            gpt_cond_chunk_len: int = 4,
+            max_ref_len: int = 30,
+            sound_norm_refs: bool = False,
+
+            # Audio processing
+            audio_config: Optional[XTTSAudioConfig] = None,
+
+            # Constants and limits
+            duration_const: int = 102400,
+            char_limits: Optional[Dict[str, int]] = None,
+            languages: Optional[List[str]] = None,
+            pad_token_id: Optional[int] = None,
+            bos_token_id: Optional[int] = None,
+            eos_token_id: Optional[int] = None,
+            **kwargs,
+    ):
+        if char_limits is None:
+            char_limits = {
+                "en": 250, "de": 253, "fr": 273, "es": 239,
+                "it": 213, "pt": 203, "pl": 224, "zh": 82,
+                "ar": 166, "cs": 186, "ru": 182, "nl": 251,
+                "tr": 226, "ja": 71, "hu": 224, "ko": 95,
+            }
+
+        if languages is None:
+            languages = [
+                "en", "es", "fr", "de", "it", "pt", "pl", "tr", "ru", "nl",
+                "cs", "ar", "zh-cn", "hu", "ko", "ja", "hi"
+            ]
+
+        if audio_config is None:
+            audio_config = XTTSAudioConfig()
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            **kwargs
+        )
+
+        self.vocab_size = vocab_size
+        self.num_chars = num_chars
+
+        # GPT parameters
+        self.gpt_batch_size = gpt_batch_size
+        self.gpt_max_audio_tokens = gpt_max_audio_tokens
+        self.gpt_max_text_tokens = gpt_max_text_tokens
+        self.gpt_max_prompt_tokens = gpt_max_prompt_tokens
+        self.gpt_layers = gpt_layers
+        self.gpt_n_model_channels = gpt_n_model_channels
+        self.gpt_n_heads = gpt_n_heads
+        self.gpt_number_text_tokens = gpt_number_text_tokens
+        self.gpt_start_text_token = gpt_start_text_token
+        self.gpt_stop_text_token = gpt_stop_text_token
+        self.gpt_num_audio_tokens = gpt_num_audio_tokens
+        self.gpt_start_audio_token = gpt_start_audio_token
+        self.gpt_stop_audio_token = gpt_stop_audio_token
+        self.gpt_code_stride_len = gpt_code_stride_len
+        self.gpt_use_masking_gt_prompt_approach = gpt_use_masking_gt_prompt_approach
+        self.gpt_use_perceiver_resampler = gpt_use_perceiver_resampler
+        self.gpt_checkpointing = gpt_checkpointing
+        self.gpt_train_solo_embeddings = gpt_train_solo_embeddings
+
+        # Training parameters
+        self.enable_redaction = enable_redaction
+        self.kv_cache = kv_cache
+        self.perceiver_cond_length_compression = perceiver_cond_length_compression
+        self.label_smoothing = label_smoothing
+
+        # Generation parameters
+        self.temperature = temperature
+        self.length_penalty = length_penalty
+        self.repetition_penalty = repetition_penalty
+        self.top_k = top_k
+        self.top_p = top_p
+        self.gpt_cond_len = gpt_cond_len
+        self.gpt_cond_chunk_len = gpt_cond_chunk_len
+        self.max_ref_len = max_ref_len
+        self.sound_norm_refs = sound_norm_refs
+
+        # Audio processing
+        self.audio_config = audio_config
+
+        # Constants and limits
+        self.duration_const = duration_const
+        self.char_limits = char_limits
+        self.languages = languages
+
+    def to_dict(self):
+        """Convert config to dictionary"""
+        config_dict = super().to_dict()
+        config_dict["audio_config"] = asdict(self.audio_config)
+        return config_dict
+
+    @classmethod
+    def from_dict(cls, config_dict):
+        """Create config from dictionary"""
+        audio_config = XTTSAudioConfig(**config_dict.pop("audio_config", {}))
+        return cls(audio_config=audio_config, **config_dict)
+
+    def update_with_tokenizer(self, tokenizer=None):
+        """Update configuration values based on tokenizer"""
+        if tokenizer is not None:
+            self.gpt_number_text_tokens = tokenizer.get_vocab_size()
+            self.gpt_start_text_token = tokenizer.bos_token_id
+            self.gpt_stop_text_token = tokenizer.eos_token_id

tokenizer.py

@@ -0,0 +1,233 @@
+from typing import List, Optional, Union, Dict, Tuple, Any
+import os
+from functools import cached_property
+
+from transformers import PreTrainedTokenizerFast
+from transformers.tokenization_utils_base import TruncationStrategy, PaddingStrategy
+from tokenizers import Tokenizer, processors
+from tokenizers.pre_tokenizers import WhitespaceSplit
+from tokenizers.processors import TemplateProcessing
+import torch
+from hangul_romanize import Transliter
+from hangul_romanize.rule import academic
+import cutlet
+
+from TTS.tts.layers.xtts.tokenizer import (multilingual_cleaners, basic_cleaners,
+                                          chinese_transliterate, korean_transliterate,
+                                          japanese_cleaners)
+
+class XTTSTokenizerFast(PreTrainedTokenizerFast):
+    """
+    Fast Tokenizer implementation for XTTS model using HuggingFace's PreTrainedTokenizerFast
+    """
+    def __init__(
+            self,
+            vocab_file: str = None,
+            tokenizer_object: Optional[Tokenizer] = None,
+            unk_token: str = "[UNK]",
+            pad_token: str = "[PAD]",
+            bos_token: str = "[START]",
+            eos_token: str = "[STOP]",
+            clean_up_tokenization_spaces: bool = True,
+            **kwargs
+    ):
+        if tokenizer_object is None and vocab_file is not None:
+            tokenizer_object = Tokenizer.from_file(vocab_file)
+
+        if tokenizer_object is not None:
+            # Configure the tokenizer
+            tokenizer_object.pre_tokenizer = WhitespaceSplit()
+            tokenizer_object.enable_padding(
+                direction='right',
+                pad_id=tokenizer_object.token_to_id(pad_token) or 0,
+                pad_token=pad_token
+            )
+            tokenizer_object.post_processor = TemplateProcessing(
+                single=f"{bos_token} $A {eos_token}",
+                special_tokens=[
+                    (bos_token, tokenizer_object.token_to_id(bos_token)),
+                    (eos_token, tokenizer_object.token_to_id(eos_token)),
+                ],
+            )
+
+        super().__init__(
+            tokenizer_object=tokenizer_object,
+            unk_token=unk_token,
+            pad_token=pad_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            **kwargs
+        )
+
+        # Character limits per language
+        self.char_limits = {
+            "en": 250, "de": 253, "fr": 273, "es": 239,
+            "it": 213, "pt": 203, "pl": 224, "zh": 82,
+            "ar": 166, "cs": 186, "ru": 182, "nl": 251,
+            "tr": 226, "ja": 71, "hu": 224, "ko": 95,
+        }
+
+        # Initialize language tools
+        self._katsu = None
+        self._korean_transliter = Transliter(academic)
+
+    @cached_property
+    def katsu(self):
+        if self._katsu is None:
+            self._katsu = cutlet.Cutlet()
+        return self._katsu
+
+    def check_input_length(self, text: str, lang: str):
+        """Check if input text length is within limits for language"""
+        lang = lang.split("-")[0]  # remove region
+        limit = self.char_limits.get(lang, 250)
+        if len(text) > limit:
+            print(f"Warning: Text length exceeds {limit} char limit for '{lang}', may cause truncation.")
+
+    def preprocess_text(self, text: str, lang: str) -> str:
+        """Apply text preprocessing for language"""
+        if lang in {"ar", "cs", "de", "en", "es", "fr", "hu", "it",
+                   "nl", "pl", "pt", "ru", "tr", "zh", "ko"}:
+            text = multilingual_cleaners(text, lang)
+            if lang == "zh":
+                text = chinese_transliterate(text)
+            if lang == "ko":
+                text = korean_transliterate(text)
+        elif lang == "ja":
+            text = japanese_cleaners(text, self.katsu)
+        else:
+            text = basic_cleaners(text)
+        return text
+
+    def _batch_encode_plus(
+            self,
+            batch_text_or_text_pairs,
+            add_special_tokens: bool = True,
+            padding_strategy = PaddingStrategy.DO_NOT_PAD,
+            truncation_strategy = TruncationStrategy.DO_NOT_TRUNCATE,
+            max_length: Optional[int] = 402,
+            stride: int = 0,
+            is_split_into_words: bool = False,
+            pad_to_multiple_of: Optional[int] = None,
+            return_tensors: Optional[str] = None,
+            return_token_type_ids: Optional[bool] = None,
+            return_attention_mask: Optional[bool] = None,
+            return_overflowing_tokens: bool = False,
+            return_special_tokens_mask: bool = False,
+            return_offsets_mapping: bool = False,
+            return_length: bool = False,
+            verbose: bool = True,
+            **kwargs
+    ) -> Dict[str, Any]:
+        """
+        Override batch encoding to handle language-specific preprocessing
+        """
+        lang = kwargs.pop("lang", ["en"] * len(batch_text_or_text_pairs))
+        if isinstance(lang, str):
+            lang = [lang] * len(batch_text_or_text_pairs)
+
+        # Preprocess each text in the batch with its corresponding language
+        processed_texts = []
+        for text, text_lang in zip(batch_text_or_text_pairs, lang):
+            if isinstance(text, str):
+                # Check length and preprocess
+                self.check_input_length(text, text_lang)
+                processed_text = self.preprocess_text(text, text_lang)
+
+                # Format text with language tag and spaces
+                lang_code = "zh-cn" if text_lang == "zh" else text_lang
+                processed_text = f"[{lang_code}]{processed_text}"
+                processed_text = processed_text.replace(" ", "[SPACE]")
+
+                processed_texts.append(processed_text)
+            else:
+                processed_texts.append(text)
+
+        # Call the parent class's encoding method with processed texts
+        return super()._batch_encode_plus(
+            processed_texts,
+            add_special_tokens=add_special_tokens,
+            padding_strategy=padding_strategy,
+            truncation_strategy=truncation_strategy,
+            max_length=max_length,
+            stride=stride,
+            is_split_into_words=is_split_into_words,
+            pad_to_multiple_of=pad_to_multiple_of,
+            return_tensors=return_tensors,
+            return_token_type_ids=return_token_type_ids,
+            return_attention_mask=return_attention_mask,
+            return_overflowing_tokens=return_overflowing_tokens,
+            return_special_tokens_mask=return_special_tokens_mask,
+            return_offsets_mapping=return_offsets_mapping,
+            return_length=return_length,
+            verbose=verbose,
+            **kwargs
+        )
+
+    def __call__(
+            self,
+            text: Union[str, List[str]],
+            lang: Union[str, List[str]] = "en",
+            add_special_tokens: bool = True,
+            padding: Union[bool, str, PaddingStrategy] = True,  # Changed default to True
+            truncation: Union[bool, str, TruncationStrategy] = True,  # Changed default to True
+            max_length: Optional[int] = 402,
+            stride: int = 0,
+            return_tensors: Optional[str] = None,
+            return_token_type_ids: Optional[bool] = None,
+            return_attention_mask: Optional[bool] = True,  # Changed default to True
+            **kwargs
+    ):
+        """
+        Main tokenization method
+        Args:
+            text: Text or list of texts to tokenize
+            lang: Language code or list of language codes corresponding to each text
+            add_special_tokens: Whether to add special tokens
+            padding: Padding strategy (default True)
+            truncation: Truncation strategy (default True)
+            max_length: Maximum length
+            stride: Stride for truncation
+            return_tensors: Format of output tensors ("pt" for PyTorch)
+            return_token_type_ids: Whether to return token type IDs
+            return_attention_mask: Whether to return attention mask (default True)
+        """
+        # Convert single string to list for batch processing
+        if isinstance(text, str):
+            text = [text]
+            if isinstance(lang, str):
+                lang = [lang]
+
+        # Ensure text and lang lists have same length
+        if len(text) != len(lang):
+            raise ValueError(f"Number of texts ({len(text)}) must match number of language codes ({len(lang)})")
+
+        # Convert padding strategy
+        if isinstance(padding, bool):
+            padding_strategy = PaddingStrategy.MAX_LENGTH if padding else PaddingStrategy.DO_NOT_PAD
+        else:
+            padding_strategy = PaddingStrategy(padding)
+
+        # Convert truncation strategy
+        if isinstance(truncation, bool):
+            truncation_strategy = TruncationStrategy.LONGEST_FIRST if truncation else TruncationStrategy.DO_NOT_TRUNCATE
+        else:
+            truncation_strategy = TruncationStrategy(truncation)
+
+        # Use the batch encoding method
+        encoded = self._batch_encode_plus(
+            text,
+            add_special_tokens=add_special_tokens,
+            padding_strategy=padding_strategy,
+            truncation_strategy=truncation_strategy,
+            max_length=max_length,
+            stride=stride,
+            return_tensors=return_tensors,
+            return_token_type_ids=return_token_type_ids,
+            return_attention_mask=return_attention_mask,
+            lang=lang,
+            **kwargs
+        )
+
+        return encoded

xtts2_config.py

@@ -0,0 +1,418 @@
+from dataclasses import asdict, dataclass
+from typing import Dict, List, Optional
+from transformers.configuration_utils import PretrainedConfig
+
+
+@dataclass
+class SpeakerEncoderConfig:
+    """Configuration for the speaker encoder component"""
+    model_name: str = "speaker_encoder"
+    preprocess_config: Optional[Dict] = None
+    model_config: Optional[Dict] = None
+    speaker_embedding_dim: int = 512
+    use_torch_spec: bool = True
+
+
+@dataclass
+class XTTSAudioConfig:
+    """Configuration for audio processing parameters"""
+    sample_rate: int = 22050
+    output_sample_rate: int = 24000
+    mel_channels: int = 80
+    hop_length: int = 256
+    win_length: int = 1024
+    n_fft: int = 1024
+    fmin: int = 0
+    fmax: int = 8000
+    power: float = 1.0
+    mel_norms_file: Optional[str] = None
+
+
+class XTTSConfig(PretrainedConfig):
+    """Combined configuration class for XTTS including both HifiGAN and GPT components"""
+    model_type = "xtts"
+
+    def __init__(
+            self,
+            # HifiGAN Audio parameters
+            input_sample_rate: int = 22050,
+            output_sample_rate: int = 24000,
+            output_hop_length: int = 256,
+
+            # HifiGAN Model architecture
+            decoder_input_dim: int = 1024,
+            d_vector_dim: int = 512,
+            cond_d_vector_in_each_upsampling_layer: bool = True,
+
+            # HifiGAN Upsampling parameters
+            upsample_rates: List[int] = None,
+            upsample_kernel_sizes: List[int] = None,
+            upsample_initial_channel: int = 512,
+
+            # HifiGAN Resblock parameters
+            resblock_kernel_sizes: List[int] = None,
+            resblock_dilation_sizes: List[List[int]] = None,
+
+            # HifiGAN Speaker encoder
+            speaker_encoder_config: Optional[Dict] = None,
+
+            # GPT Model architecture
+            vocab_size: int = 256,
+            num_chars: int = 255,
+
+            # GPT parameters
+            gpt_batch_size: int = 1,
+            gpt_max_audio_tokens: int = 605,
+            gpt_max_text_tokens: int = 402,
+            gpt_max_prompt_tokens: int = 70,
+            gpt_layers: int = 30,
+            gpt_n_model_channels: int = 1024,
+            gpt_n_heads: int = 16,
+            gpt_number_text_tokens: int = 6681,
+            gpt_start_text_token: Optional[int] = None,
+            gpt_stop_text_token: Optional[int] = None,
+            gpt_num_audio_tokens: int = 1026,
+            gpt_start_audio_token: int = 1024,
+            gpt_stop_audio_token: int = 1025,
+            gpt_code_stride_len: int = 1024,
+            gpt_use_masking_gt_prompt_approach: bool = True,
+            gpt_use_perceiver_resampler: bool = True,
+            gpt_checkpointing: bool = False,
+            gpt_train_solo_embeddings: bool = False,
+
+            # GPT Training parameters
+            enable_redaction: bool = False,
+            kv_cache: bool = True,
+            perceiver_cond_length_compression: int = 256,
+            label_smoothing: float = 0.0,
+
+            # GPT Generation parameters
+            temperature: float = 0.75,
+            length_penalty: float = 1.0,
+            repetition_penalty: float = 5.0,
+            top_k: int = 50,
+            top_p: float = 0.85,
+            gpt_cond_len: int = 30,
+            gpt_cond_chunk_len: int = 4,
+            max_ref_len: int = 30,
+            sound_norm_refs: bool = False,
+
+            # GPT Audio processing
+            audio_config: Optional[XTTSAudioConfig] = None,
+
+            # GPT Constants and limits
+            duration_const: int = 102400,
+            char_limits: Optional[Dict[str, int]] = None,
+            languages: Optional[List[str]] = None,
+
+            # Base config parameters
+            pad_token_id: Optional[int] = None,
+            bos_token_id: Optional[int] = None,
+            eos_token_id: Optional[int] = None,
+            **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            **kwargs
+        )
+
+        # Set default lists for HifiGAN
+        if upsample_rates is None:
+            upsample_rates = [8, 8, 2, 2]
+        if upsample_kernel_sizes is None:
+            upsample_kernel_sizes = [16, 16, 4, 4]
+        if resblock_kernel_sizes is None:
+            resblock_kernel_sizes = [3, 7, 11]
+        if resblock_dilation_sizes is None:
+            resblock_dilation_sizes = [[1, 3, 5], [1, 3, 5], [1, 3, 5]]
+
+        # Set default dicts for GPT
+        if char_limits is None:
+            char_limits = {
+                "en": 250, "de": 253, "fr": 273, "es": 239,
+                "it": 213, "pt": 203, "pl": 224, "zh": 82,
+                "ar": 166, "cs": 186, "ru": 182, "nl": 251,
+                "tr": 226, "ja": 71, "hu": 224, "ko": 95,
+            }
+
+        if languages is None:
+            languages = [
+                "en", "es", "fr", "de", "it", "pt", "pl", "tr", "ru", "nl",
+                "cs", "ar", "zh-cn", "hu", "ko", "ja", "hi"
+            ]
+
+        # Initialize HifiGAN parameters
+        # Audio parameters
+        self.input_sample_rate = input_sample_rate
+        self.output_sample_rate = output_sample_rate
+        self.output_hop_length = output_hop_length
+
+        # Model architecture
+        self.decoder_input_dim = decoder_input_dim
+        self.d_vector_dim = d_vector_dim
+        self.cond_d_vector_in_each_upsampling_layer = cond_d_vector_in_each_upsampling_layer
+
+        # Upsampling parameters
+        self.upsample_rates = upsample_rates
+        self.upsample_kernel_sizes = upsample_kernel_sizes
+        self.upsample_initial_channel = upsample_initial_channel
+
+        # Resblock parameters
+        self.resblock_kernel_sizes = resblock_kernel_sizes
+        self.resblock_dilation_sizes = resblock_dilation_sizes
+
+        # Speaker encoder - store as dictionary
+        if speaker_encoder_config is None:
+            self.speaker_encoder_config = asdict(SpeakerEncoderConfig())
+        elif isinstance(speaker_encoder_config, dict):
+            default_config = asdict(SpeakerEncoderConfig())
+            default_config.update(speaker_encoder_config)
+            self.speaker_encoder_config = default_config
+        elif isinstance(speaker_encoder_config, SpeakerEncoderConfig):
+            self.speaker_encoder_config = asdict(speaker_encoder_config)
+        else:
+            raise ValueError("speaker_encoder_config must be either a dictionary or SpeakerEncoderConfig instance")
+
+        # Initialize GPT parameters
+        self.vocab_size = vocab_size
+        self.num_chars = num_chars
+
+        # GPT model parameters
+        self.gpt_batch_size = gpt_batch_size
+        self.gpt_max_audio_tokens = gpt_max_audio_tokens
+        self.gpt_max_text_tokens = gpt_max_text_tokens
+        self.gpt_max_prompt_tokens = gpt_max_prompt_tokens
+        self.gpt_layers = gpt_layers
+        self.gpt_n_model_channels = gpt_n_model_channels
+        self.gpt_n_heads = gpt_n_heads
+        self.gpt_number_text_tokens = gpt_number_text_tokens
+        self.gpt_start_text_token = gpt_start_text_token
+        self.gpt_stop_text_token = gpt_stop_text_token
+        self.gpt_num_audio_tokens = gpt_num_audio_tokens
+        self.gpt_start_audio_token = gpt_start_audio_token
+        self.gpt_stop_audio_token = gpt_stop_audio_token
+        self.gpt_code_stride_len = gpt_code_stride_len
+        self.gpt_use_masking_gt_prompt_approach = gpt_use_masking_gt_prompt_approach
+        self.gpt_use_perceiver_resampler = gpt_use_perceiver_resampler
+        self.gpt_checkpointing = gpt_checkpointing
+        self.gpt_train_solo_embeddings = gpt_train_solo_embeddings
+
+        # Training parameters
+        self.enable_redaction = enable_redaction
+        self.kv_cache = kv_cache
+        self.perceiver_cond_length_compression = perceiver_cond_length_compression
+        self.label_smoothing = label_smoothing
+
+        # Generation parameters
+        self.temperature = temperature
+        self.length_penalty = length_penalty
+        self.repetition_penalty = repetition_penalty
+        self.top_k = top_k
+        self.top_p = top_p
+        self.gpt_cond_len = gpt_cond_len
+        self.gpt_cond_chunk_len = gpt_cond_chunk_len
+        self.max_ref_len = max_ref_len
+        self.sound_norm_refs = sound_norm_refs
+
+        # Audio processing
+        if audio_config is None:
+            audio_config = XTTSAudioConfig()
+        elif isinstance(audio_config, dict):
+            audio_config = XTTSAudioConfig(**audio_config)
+        self.audio_config = audio_config
+
+        # Constants and limits
+        self.duration_const = duration_const
+        self.char_limits = char_limits
+        self.languages = languages
+
+    def to_dict(self) -> Dict:
+        """Convert the config to a dictionary format."""
+        # Get parent class dict
+        output = super().to_dict()
+
+        # Add all attributes
+        output.update({
+            # HifiGAN parameters
+            "input_sample_rate": self.input_sample_rate,
+            "output_sample_rate": self.output_sample_rate,
+            "output_hop_length": self.output_hop_length,
+            "decoder_input_dim": self.decoder_input_dim,
+            "d_vector_dim": self.d_vector_dim,
+            "cond_d_vector_in_each_upsampling_layer": self.cond_d_vector_in_each_upsampling_layer,
+            "upsample_rates": self.upsample_rates,
+            "upsample_kernel_sizes": self.upsample_kernel_sizes,
+            "upsample_initial_channel": self.upsample_initial_channel,
+            "resblock_kernel_sizes": self.resblock_kernel_sizes,
+            "resblock_dilation_sizes": self.resblock_dilation_sizes,
+            "speaker_encoder_config": self.speaker_encoder_config,
+
+            # GPT parameters
+            "vocab_size": self.vocab_size,
+            "num_chars": self.num_chars,
+            "gpt_batch_size": self.gpt_batch_size,
+            "gpt_max_audio_tokens": self.gpt_max_audio_tokens,
+            "gpt_max_text_tokens": self.gpt_max_text_tokens,
+            "gpt_max_prompt_tokens": self.gpt_max_prompt_tokens,
+            "gpt_layers": self.gpt_layers,
+            "gpt_n_model_channels": self.gpt_n_model_channels,
+            "gpt_n_heads": self.gpt_n_heads,
+            "gpt_number_text_tokens": self.gpt_number_text_tokens,
+            "gpt_start_text_token": self.gpt_start_text_token,
+            "gpt_stop_text_token": self.gpt_stop_text_token,
+            "gpt_num_audio_tokens": self.gpt_num_audio_tokens,
+            "gpt_start_audio_token": self.gpt_start_audio_token,
+            "gpt_stop_audio_token": self.gpt_stop_audio_token,
+            "gpt_code_stride_len": self.gpt_code_stride_len,
+            "gpt_use_masking_gt_prompt_approach": self.gpt_use_masking_gt_prompt_approach,
+            "gpt_use_perceiver_resampler": self.gpt_use_perceiver_resampler,
+            "gpt_checkpointing": self.gpt_checkpointing,
+            "gpt_train_solo_embeddings": self.gpt_train_solo_embeddings,
+            "enable_redaction": self.enable_redaction,
+            "kv_cache": self.kv_cache,
+            "perceiver_cond_length_compression": self.perceiver_cond_length_compression,
+            "label_smoothing": self.label_smoothing,
+            "temperature": self.temperature,
+            "length_penalty": self.length_penalty,
+            "repetition_penalty": self.repetition_penalty,
+            "top_k": self.top_k,
+            "top_p": self.top_p,
+            "gpt_cond_len": self.gpt_cond_len,
+            "gpt_cond_chunk_len": self.gpt_cond_chunk_len,
+            "max_ref_len": self.max_ref_len,
+            "sound_norm_refs": self.sound_norm_refs,
+            "audio_config": asdict(self.audio_config),
+            "duration_const": self.duration_const,
+            "char_limits": self.char_limits,
+            "languages": self.languages,
+        })
+
+        return output
+
+    @classmethod
+    def from_dict(cls, config_dict: Dict) -> "XTTSConfig":
+        """Create a config instance from a dictionary."""
+        config_copy = config_dict.copy()
+
+        # Handle special nested configs
+        if "audio_config" in config_copy:
+            config_copy["audio_config"] = XTTSAudioConfig(**config_copy["audio_config"])
+
+        return cls(**config_copy)
+
+    def get_speaker_encoder_config(self) -> SpeakerEncoderConfig:
+        """Get speaker encoder config as a SpeakerEncoderConfig instance"""
+        return SpeakerEncoderConfig(**self.speaker_encoder_config)
+
+    def update_with_tokenizer(self, tokenizer=None):
+        """Update configuration values based on tokenizer"""
+        if tokenizer is not None:
+            self.gpt_number_text_tokens = tokenizer.get_vocab_size()
+            self.gpt_start_text_token = tokenizer.bos_token_id
+            self.gpt_stop_text_token = tokenizer.eos_token_id
+            self.vocab_size = tokenizer.get_vocab_size()
+            self.pad_token_id = tokenizer.pad_token_id
+            self.bos_token_id = tokenizer.bos_token_id
+            self.eos_token_id = tokenizer.eos_token_id
+
+    def get_hifigan_config(self) -> Dict:
+        """Extract HiFiGAN-specific configuration"""
+        return {
+            "input_sample_rate": self.input_sample_rate,
+            "output_sample_rate": self.output_sample_rate,
+            "output_hop_length": self.output_hop_length,
+            "decoder_input_dim": self.decoder_input_dim,
+            "d_vector_dim": self.d_vector_dim,
+            "cond_d_vector_in_each_upsampling_layer": self.cond_d_vector_in_each_upsampling_layer,
+            "upsample_rates": self.upsample_rates,
+            "upsample_kernel_sizes": self.upsample_kernel_sizes,
+            "upsample_initial_channel": self.upsample_initial_channel,
+            "resblock_kernel_sizes": self.resblock_kernel_sizes,
+            "resblock_dilation_sizes": self.resblock_dilation_sizes,
+            "speaker_encoder_config": self.speaker_encoder_config
+        }
+
+    def get_gpt_config(self) -> Dict:
+        """Extract GPT-specific configuration"""
+        return {
+            "vocab_size": self.vocab_size,
+            "num_chars": self.num_chars,
+            "gpt_batch_size": self.gpt_batch_size,
+            "gpt_max_audio_tokens": self.gpt_max_audio_tokens,
+            "gpt_max_text_tokens": self.gpt_max_text_tokens,
+            "gpt_max_prompt_tokens": self.gpt_max_prompt_tokens,
+            "gpt_layers": self.gpt_layers,
+            "gpt_n_model_channels": self.gpt_n_model_channels,
+            "gpt_n_heads": self.gpt_n_heads,
+            "gpt_number_text_tokens": self.gpt_number_text_tokens,
+            "gpt_start_text_token": self.gpt_start_text_token,
+            "gpt_stop_text_token": self.gpt_stop_text_token,
+            "gpt_num_audio_tokens": self.gpt_num_audio_tokens,
+            "gpt_start_audio_token": self.gpt_start_audio_token,
+            "gpt_stop_audio_token": self.gpt_stop_audio_token,
+            "gpt_code_stride_len": self.gpt_code_stride_len,
+            "gpt_use_masking_gt_prompt_approach": self.gpt_use_masking_gt_prompt_approach,
+            "gpt_use_perceiver_resampler": self.gpt_use_perceiver_resampler,
+            "gpt_checkpointing": self.gpt_checkpointing,
+            "gpt_train_solo_embeddings": self.gpt_train_solo_embeddings,
+            "enable_redaction": self.enable_redaction,
+            "kv_cache": self.kv_cache,
+            "perceiver_cond_length_compression": self.perceiver_cond_length_compression,
+            "label_smoothing": self.label_smoothing,
+            "audio_config": self.audio_config,
+            "pad_token_id": self.pad_token_id,
+            "bos_token_id": self.bos_token_id,
+            "eos_token_id": self.eos_token_id
+        }
+
+    def get_generation_config(self) -> Dict:
+        """Extract generation-specific configuration"""
+        return {
+            "temperature": self.temperature,
+            "length_penalty": self.length_penalty,
+            "repetition_penalty": self.repetition_penalty,
+            "top_k": self.top_k,
+            "top_p": self.top_p,
+            "gpt_cond_len": self.gpt_cond_len,
+            "gpt_cond_chunk_len": self.gpt_cond_chunk_len,
+            "max_ref_len": self.max_ref_len,
+            "sound_norm_refs": self.sound_norm_refs
+        }
+
+    def validate(self):
+        """Validate configuration values"""
+        if self.gpt_max_text_tokens <= 0:
+            raise ValueError("gpt_max_text_tokens must be positive")
+        if self.gpt_max_audio_tokens <= 0:
+            raise ValueError("gpt_max_audio_tokens must be positive")
+        if self.gpt_layers <= 0:
+            raise ValueError("gpt_layers must be positive")
+        if self.gpt_n_heads <= 0:
+            raise ValueError("gpt_n_heads must be positive")
+        if self.gpt_n_model_channels <= 0:
+            raise ValueError("gpt_n_model_channels must be positive")
+        if len(self.upsample_rates) != len(self.upsample_kernel_sizes):
+            raise ValueError("upsample_rates and upsample_kernel_sizes must have same length")
+        if not all(isinstance(x, int) and x > 0 for x in self.upsample_rates):
+            raise ValueError("all upsample_rates must be positive integers")
+
+    def get_audio_config(self) -> XTTSAudioConfig:
+        """Get the audio configuration"""
+        return self.audio_config
+
+    @property
+    def num_hidden_layers(self) -> int:
+        """Get number of hidden layers (alias for gpt_layers)"""
+        return self.gpt_layers
+
+    @property
+    def hidden_size(self) -> int:
+        """Get hidden size (alias for gpt_n_model_channels)"""
+        return self.gpt_n_model_channels
+
+    @property
+    def num_attention_heads(self) -> int:
+        """Get number of attention heads (alias for gpt_n_heads)"""
+        return self.gpt_n_heads

xtts2_modeling.py

@@ -0,0 +1,259 @@
+import asyncio
+from dataclasses import dataclass
+from typing import Optional, List, Tuple
+from concurrent.futures import ThreadPoolExecutor
+import torch
+import numpy as np
+from transformers import PreTrainedModel
+
+from vllm import AsyncLLMEngine, AsyncEngineArgs, SamplingParams, TokensPrompt
+from vllm.multimodal import MultiModalDataDict
+from vllm.utils import Counter
+
+from TTS.TTS.tts.layers.xtts.hifigan_decoder import HifiDecoder
+from gpt_config import XTTSGPTConfig
+from xtts2_config import XTTSConfig
+from tokenizer import XTTSTokenizerFast
+
+
+@dataclass
+class XTTSRequest:
+    """Container for XTTS inference request data"""
+    request_id: str
+    text: str
+    language: str
+    gpt_cond_latent: torch.Tensor
+    speaker_embedding: torch.Tensor
+    temperature: float = 0.75
+    top_p: float = 0.85
+    top_k: int = 50
+    repetition_penalty: float = 10.0
+    length_penalty: float = 1.0
+    do_sample: bool = True
+
+
+@dataclass
+class XTTSOutput:
+    """Container for XTTS inference output"""
+    request_id: str
+    wav: np.ndarray
+    gpt_latents: np.ndarray
+    speaker_embedding: torch.Tensor
+
+
+class Xtts(PreTrainedModel):
+    """Async XTTS model implementation using VLLM's AsyncEngine."""
+
+    def __init__(self, hifi_config: XTTSConfig, gpt_config: XTTSGPTConfig, tensor_parallel_size: int = 1, **kwargs):
+        self.hifi_config = hifi_config
+        self.gpt_config = gpt_config
+        self.tp = tensor_parallel_size
+        self.tokenizer = XTTSTokenizerFast.from_pretrained("AstraMindAI/xtts2-gpt")
+        self.request_counter = Counter()
+        self.executor = ThreadPoolExecutor(max_workers=4)  # For CPU-bound tasks
+        self.init_models()
+        self.register_buffer("mel_stats", torch.ones(80))
+
+    @staticmethod
+    def get_memory_percentage(memory: int) -> float:
+        """Get memory percentage."""
+        return memory / torch.cuda.get_device_properties(0).total_memory
+
+    async def init_models(self):
+        """Initialize models with AsyncVLLMEngine."""
+        # Initialize VLLM engine
+        engine_args = AsyncEngineArgs(
+            model=self.gpt_config.model_dir,
+            tensor_parallel_size=self.tp,
+            dtype="auto ",
+            max_model_len=self.gpt_config.gpt_max_text_tokens + self.gpt_config.gpt_max_audio_tokens,
+            gpu_memory_utilization=self.get_memory_percentage(2),# since the model neds 2 gb we need to calc the bare minimum memory
+            trust_remote_code=True,
+            skip_tokenizer_init=True, # no need to initialize tokenizer, we use our own
+            max_num_batched_tokens=4096,
+            max_num_seqs=256,
+        )
+
+        self.llm_engine = AsyncLLMEngine.from_engine_args(engine_args)
+        self.llm_engine = AsyncLLMEngine
+        # Initialize HiFi-GAN decoder
+        self.hifigan_decoder = HifiDecoder(
+            input_sample_rate=self.hifi_config.input_sample_rate,
+            output_sample_rate=self.hifi_config.output_sample_rate,
+            output_hop_length=self.hifi_config.output_hop_length,
+            ar_mel_length_compression=self.hifi_config.gpt_code_stride_len,
+            decoder_input_dim=self.hifi_config.decoder_input_dim,
+            d_vector_dim=self.hifi_config.d_vector_dim,
+            cond_d_vector_in_each_upsampling_layer=self.hifi_config.cond_d_vector_in_each_upsampling_layer,
+        )
+
+    @classmethod
+    def from_pretrained(
+            cls,
+            pretrained_model_name_or_path: str,
+            torch_dtype: torch.dtype = torch.float16,
+            device_map: Optional[str] = "auto",
+            tensor_parallel_size: int = 1,
+            **kwargs,
+    ) -> "Xtts":
+        """Load pretrained XTTS model from HuggingFace Hub.
+
+        Args:
+            pretrained_model_name_or_path (str): Path to pretrained weights or HF Hub model id
+            torch_dtype (torch.dtype, optional): Type to load the model as. Defaults to float16.
+            device_map (str, optional): Device mapping strategy. Defaults to "auto".
+            **kwargs: Additional arguments passed to the model.
+
+        Returns:
+            Xtts: Loaded model instance
+        """
+        from huggingface_hub import hf_hub_download
+        import json
+        import os
+
+        # Download and load configs
+        if not os.path.exists(pretrained_model_name_or_path):
+            config_file = hf_hub_download(
+                repo_id=pretrained_model_name_or_path,
+                filename="../xtts2_gpt/config.json"
+            )
+            with open(config_file, 'r') as f:
+                config = json.load(f)
+
+            gpt_config_file = hf_hub_download(
+                repo_id=pretrained_model_name_or_path,
+                filename="gpt_config.py"
+            )
+            with open(gpt_config_file, 'r') as f:
+                gpt_config = json.loads(f.read())
+
+            hifigan_config_file = hf_hub_download(
+                repo_id=pretrained_model_name_or_path,
+                filename="xtts2_config.py"
+            )
+            with open(hifigan_config_file, 'r') as f:
+                hifigan_config = json.loads(f.read())
+        else:
+            # Load from local path
+            with open(os.path.join(pretrained_model_name_or_path, "config.json"), 'r') as f:
+                config = json.load(f)
+
+
+        # Initialize configs
+        gpt_config = XTTSGPTConfig(**config)
+        hifi_config = XTTSConfig(**config)
+
+        # Initialize model
+        model = cls(
+            hifi_config=hifi_config,
+            gpt_config=gpt_config,
+            tensor_parallel_size=tensor_parallel_size,
+            **kwargs
+        )
+
+        # Load model weights
+        if not os.path.exists(pretrained_model_name_or_path):
+            gpt_weights = hf_hub_download(
+                repo_id=pretrained_model_name_or_path,
+                filename="../xtts2_gpt/xttsv2-gpt.safetensors"
+            )
+            hifigan_weights = hf_hub_download(
+                repo_id=pretrained_model_name_or_path,
+                filename="xttsv2-hifigan-mel.safetensors"
+            )
+        else:
+            gpt_weights = os.path.join(pretrained_model_name_or_path, "xttsv2-gpt.safetensors")
+            hifigan_weights = os.path.join(pretrained_model_name_or_path, "xttsv2-hifigan-mel.safetensors")
+
+        # Load GPT weights
+        import safetensors.torch
+        state_dict = safetensors.torch.load_file(gpt_weights)
+        model.gpt.load_state_dict(state_dict)
+
+        # Load HiFi-GAN weights
+        hifigan_state = safetensors.torch.load_file(hifigan_weights)
+        model.hifigan_decoder.load_state_dict(hifigan_state)
+
+        # Set model properties
+        model.config = config
+
+        # Cast model to specified dtype
+        model = model.to(torch_dtype)
+
+        # Handle device mapping
+        if device_map:
+            from accelerate import dispatch_model
+            model = dispatch_model(model, device_map=device_map)
+
+        return model
+
+    def prepare_inputs(self, text: str, language: str, gpt_cond_latent: torch.Tensor) -> Tuple[List[int], torch.Tensor]:
+        """Prepare input text with conditioning tokens."""
+        # Add special tokens and conditioning format
+        # Format: <|condition|>latent_data<|endofcondition|>text<|endoftext|>
+        text_tokens = self.tokenizer.encode(text, lang=language)
+        return text_tokens, gpt_cond_latent
+
+
+
+    async def generate_speech_async(self, request: XTTSRequest) -> XTTSOutput:
+        """Generate speech for a single request asynchronously."""
+        # Prepare input with conditioning
+        tokens, gpt_cond_latent = self.prepare_inputs(
+            request.text,
+            request.language,
+            request.gpt_cond_latent
+        )
+
+        # Setup sampling parameters
+        sampling_params = SamplingParams(
+            temperature=request.temperature,
+            top_p=request.top_p,
+            top_k=request.top_k,
+            repetition_penalty=request.repetition_penalty,
+            max_tokens=self.gpt_config.gpt_max_audio_tokens,
+            stop=['</s>', '<|endoftext|>']
+        )
+        engine_inputs = TokensPrompt( prompt_token_ids = tokens )
+        if gpt_cond_latent is not None:
+            engine_inputs["multi_modal_data"] = MultiModalDataDict({"audio":gpt_cond_latent})
+        # Generate tokens using VLLM
+        output_generator = self.llm_engine.generate(
+            inputs=engine_inputs,
+            sampling_params=sampling_params,
+            request_id=request.request_id
+        )
+
+        async for outputs in output_generator:
+            # Extract generated tokens
+            generated_tokens = outputs.outputs[0].token_ids
+
+            # Convert to hidden states (this step depends on your model architecture)
+            hidden_states = await self._tokens_to_hidden_states(generated_tokens)
+
+            # Generate audio using HiFi-GAN (run in thread pool to avoid blocking)
+            wav = await asyncio.get_event_loop().run_in_executor(
+                self.executor,
+                lambda: self.hifigan_decoder(
+                    hidden_states,
+                    g=request.speaker_embedding
+                ).cpu().numpy().squeeze()
+            )
+
+            return XTTSOutput(
+                request_id=request.request_id,
+                wav=wav,
+                gpt_latents=hidden_states.cpu().numpy(),
+                speaker_embedding=request.speaker_embedding
+            )
+
+
+    async def _tokens_to_hidden_states(self, tokens: List[int]) -> torch.Tensor:
+        """Convert generated tokens to hidden states."""
+        # This implementation depends on your specific model architecture
+        # You'll need to adapt this based on how your model processes tokens
+        # This is a placeholder implementation
+        token_tensor = torch.tensor(tokens, device=self.device)
+        # Use VLLM's engine to get hidden states
+        hidden_states = await self.llm_engine.encode(token_tensor)
+        return hidden_states

xttsv2-hifigan-mel.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6eaf6c236291478363be6da06c0869551c51bf0c8983fd2dd70561a4a1f1ace3
+size 103599512