Upload 3 files

config.json

@@ -1,14 +1,6 @@
 {
- "_name_or_path": "AstraMindAI/xtts2-gpt",
-  "architectures": [
-    "XttsGPT"
-  ],
-  "torch_dtype": "float32",
-  "auto_map": {
-    "AutoConfig": "AstraMindAI/xtts2-gpt--gpt_config.XTTSGPTConfig",
-    "AutoModelForCausalLM": "AstraMindAI/xtts2-gpt--xtts2_gpt_modeling.XttsGPT",
-    "AutoTokenizer": "AstraMindAI/xtts2-gpt--tokenizer.XTTSTokenizerFast"
-  },
+  "activation_function": "gelu",
+  "attn_pdrop": 0.1,
   "audio_config": {
     "fmax": 8000,
     "fmin": 0,
@@ -21,6 +13,7 @@
     "sample_rate": 22050,
     "win_length": 1024
   },
+  "batch_size": 1,
   "char_limits": {
     "ar": 166,
     "cs": 186,
@@ -39,28 +32,14 @@
     "tr": 226,
     "zh": 82
   },
+  "checkpointing": false,
+  "code_stride_len": 1024,
+  "cond_chunk_len": 4,
+  "cond_len": 30,
   "duration_const": 102400,
+  "embd_pdrop": 0.1,
   "enable_redaction": false,
-  "gpt_batch_size": 1,
-  "gpt_checkpointing": false,
-  "gpt_code_stride_len": 1024,
-  "gpt_cond_chunk_len": 4,
-  "gpt_cond_len": 30,
-  "gpt_layers": 30,
-  "gpt_max_audio_tokens": 605,
-  "gpt_max_prompt_tokens": 70,
-  "gpt_max_text_tokens": 402,
-  "gpt_n_heads": 16,
-  "gpt_n_model_channels": 1024,
-  "gpt_num_audio_tokens": 1026,
-  "gpt_number_text_tokens": 6681,
-  "gpt_start_audio_token": 1024,
-  "gpt_start_text_token": null,
-  "gpt_stop_audio_token": 1025,
-  "gpt_stop_text_token": null,
-  "gpt_train_solo_embeddings": false,
-  "gpt_use_masking_gt_prompt_approach": true,
-  "gpt_use_perceiver_resampler": true,
+  "hidden_size": 1024,
   "kv_cache": true,
   "label_smoothing": 0.0,
   "languages": [
@@ -82,60 +61,34 @@
     "ja",
     "hi"
   ],
+  "layer_norm_epsilon": 1e-05,
+  "max_audio_tokens": 605,
+  "max_position_embeddings": 2048,
+  "max_prompt_tokens": 70,
   "max_ref_len": 30,
+  "max_text_tokens": 402,
   "model_type": "xtts_gpt",
+  "n_inner": null,
+  "num_attention_heads": 16,
+  "num_audio_tokens": 1026,
   "num_chars": 255,
+  "num_hidden_layers": 30,
+  "number_text_tokens": 6681,
   "perceiver_cond_length_compression": 256,
+  "reorder_and_upcast_attn": false,
+  "repetition_penalty": 5.0,
+  "resid_pdrop": 0.1,
+  "scale_attn_by_inverse_layer_idx": false,
   "sound_norm_refs": false,
-  "transformers_version": "4.45.1",
-  "vocab_size": 256,
-  "cond_d_vector_in_each_upsampling_layer": true,
-  "d_vector_dim": 512,
-  "decoder_input_dim": 1024,
-  "input_sample_rate": 22050,
-  "hifi_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
-  },
-  "upsample_initial_channel": 512,
-  "upsample_kernel_sizes": [
-    16,
-    16,
-    4,
-    4
-  ],
-  "upsample_rates": [
-    8,
-    8,
-    2,
-    2
-  ]
+  "start_audio_token": 1024,
+  "start_text_token": null,
+  "stop_audio_token": 1025,
+  "stop_text_token": null,
+  "temperature": 0.75,
+  "top_p": 0.85,
+  "train_solo_embeddings": false,
+  "transformers_version": "4.46.0",
+  "use_masking_gt_prompt_approach": true,
+  "use_perceiver_resampler": true,
+  "vocab_size": 256
 }

gpt_config.py

@@ -29,27 +29,34 @@ class XTTSGPTConfig(PretrainedConfig):
             self,
             # Model architecture
             vocab_size: int = 256,
+            hidden_size: int = 1024,  # Changed from gpt_n_model_channels
+            num_hidden_layers: int = 30,  # Changed from gpt_layers
+            num_attention_heads: int = 16,  # Changed from gpt_n_heads
+            n_inner: Optional[int] = None,  # Added for GPT-2 compatibility
+            max_position_embeddings: int = 2048,  # Added for positional embeddings
+            layer_norm_epsilon: float = 1e-5,  # Added for layer norm
+            activation_function: str = "gelu",  # Added activation function
+            resid_pdrop: float = 0.1,  # Added dropout rates
+            embd_pdrop: float = 0.1,
+            attn_pdrop: float = 0.1,
+
+            # Specific XTTS parameters
             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,
+            batch_size: int = 1,  # Changed from gpt_batch_size
+            max_audio_tokens: int = 605,  # Changed from gpt_max_audio_tokens
+            max_text_tokens: int = 402,  # Changed from gpt_max_text_tokens
+            max_prompt_tokens: int = 70,  # Changed from gpt_max_prompt_tokens
+            number_text_tokens: int = 6681,  # Changed from gpt_number_text_tokens
+            start_text_token: Optional[int] = None,  # Changed from gpt_start_text_token
+            stop_text_token: Optional[int] = None,  # Changed from gpt_stop_text_token
+            num_audio_tokens: int = 1026,  # Changed from gpt_num_audio_tokens
+            start_audio_token: int = 1024,  # Changed from gpt_start_audio_token
+            stop_audio_token: int = 1025,  # Changed from gpt_stop_audio_token
+            code_stride_len: int = 1024,  # Changed from gpt_code_stride_len
+            use_masking_gt_prompt_approach: bool = True,  # Changed from gpt_use_masking_gt_prompt_approach
+            use_perceiver_resampler: bool = True,  # Changed from gpt_use_perceiver_resampler
+            checkpointing: bool = False,  # Changed from gpt_checkpointing
+            train_solo_embeddings: bool = False,  # Changed from gpt_train_solo_embeddings
 
             # Training parameters
             enable_redaction: bool = False,
@@ -58,13 +65,13 @@ class XTTSGPTConfig(PretrainedConfig):
             label_smoothing: float = 0.0,
 
             # Generation parameters
-            #temperature: float = 0.75, will trow a warning
-            #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,
+            temperature: float = 0.75,
+            length_penalty: float = 1.0,
+            repetition_penalty: float = 5.0,
+            top_k: int = 50,
+            top_p: float = 0.85,
+            cond_len: int = 30,  # Changed from gpt_cond_len
+            cond_chunk_len: int = 4,  # Changed from gpt_cond_chunk_len
             max_ref_len: int = 30,
             sound_norm_refs: bool = False,
 
@@ -78,6 +85,12 @@ class XTTSGPTConfig(PretrainedConfig):
             pad_token_id: Optional[int] = None,
             bos_token_id: Optional[int] = None,
             eos_token_id: Optional[int] = None,
+
+            # GPT-2 compatibility flags
+            scale_attn_by_inverse_layer_idx: bool = False,
+            reorder_and_upcast_attn: bool = False,
+            add_cross_attention: bool = False,
+            tie_word_embeddings: bool = True,
             **kwargs,
     ):
         if char_limits is None:
@@ -105,27 +118,34 @@ class XTTSGPTConfig(PretrainedConfig):
         )
 
         self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.n_inner = n_inner
+        self.max_position_embeddings = max_position_embeddings
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.activation_function = activation_function
+        self.resid_pdrop = resid_pdrop
+        self.embd_pdrop = embd_pdrop
+        self.attn_pdrop = attn_pdrop
+
+        # XTTS specific parameters
         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
+        self.batch_size = batch_size
+        self.max_audio_tokens = max_audio_tokens
+        self.max_text_tokens = max_text_tokens
+        self.max_prompt_tokens = max_prompt_tokens
+        self.number_text_tokens = number_text_tokens
+        self.start_text_token = start_text_token
+        self.stop_text_token = stop_text_token
+        self.num_audio_tokens = num_audio_tokens
+        self.start_audio_token = start_audio_token
+        self.stop_audio_token = stop_audio_token
+        self.code_stride_len = code_stride_len
+        self.use_masking_gt_prompt_approach = use_masking_gt_prompt_approach
+        self.use_perceiver_resampler = use_perceiver_resampler
+        self.checkpointing = checkpointing
+        self.train_solo_embeddings = train_solo_embeddings
 
         # Training parameters
         self.enable_redaction = enable_redaction
@@ -134,8 +154,13 @@ class XTTSGPTConfig(PretrainedConfig):
         self.label_smoothing = label_smoothing
 
         # Generation parameters
-        self.gpt_cond_len = gpt_cond_len
-        self.gpt_cond_chunk_len = gpt_cond_chunk_len
+        self.temperature = temperature
+        self.length_penalty = length_penalty
+        self.repetition_penalty = repetition_penalty
+        self.top_k = top_k
+        self.top_p = top_p
+        self.cond_len = cond_len
+        self.cond_chunk_len = cond_chunk_len
         self.max_ref_len = max_ref_len
         self.sound_norm_refs = sound_norm_refs
 
@@ -147,6 +172,12 @@ class XTTSGPTConfig(PretrainedConfig):
         self.char_limits = char_limits
         self.languages = languages
 
+        # GPT-2 compatibility flags
+        self.scale_attn_by_inverse_layer_idx = scale_attn_by_inverse_layer_idx
+        self.reorder_and_upcast_attn = reorder_and_upcast_attn
+        self.add_cross_attention = add_cross_attention
+        self.tie_word_embeddings = tie_word_embeddings
+
     def to_dict(self):
         """Convert config to dictionary"""
         config_dict = super().to_dict()
@@ -154,14 +185,14 @@ class XTTSGPTConfig(PretrainedConfig):
         return config_dict
 
     @classmethod
-    def from_dict(cls, config_dict, *arg, **kwargs):
+    def from_dict(cls, config_dict, *args, **kwargs):
         """Create config from dictionary"""
         audio_config = XTTSAudioConfig(**config_dict.pop("audio_config", {}))
-        return cls(audio_config=audio_config, **config_dict)
+        return cls(audio_config=audio_config, **config_dict, **kwargs)
 
     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.number_text_tokens = tokenizer.get_vocab_size()
+            self.start_text_token = tokenizer.bos_token_id
+            self.stop_text_token = tokenizer.eos_token_id

xtts2_gpt_modeling.py

@@ -8,19 +8,20 @@ from torch.nn import functional as F
 from typing import List, Optional, Union, Iterable, Tuple, Mapping
 
 from transformers import PretrainedConfig
-from vllm.attention import AttentionMetadata
-from vllm.config import CacheConfig
-from vllm.distributed import get_pp_group
+from vllm.attention import AttentionMetadata, Attention
+from vllm.config import CacheConfig, MultiModalConfig
+from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
 from vllm.inputs import InputContext, INPUT_REGISTRY
-from vllm.model_executor.layers.linear import ColumnParallelLinear
+from vllm.model_executor.layers.activation import get_act_fn
+from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear
+from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.sampler import Sampler, SamplerOutput
 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
-from vllm.model_executor.models.gpt2 import GPT2Block
 from vllm.model_executor.sampling_metadata import SamplingMetadata
 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalInputs
 from vllm.sequence import IntermediateTensors, SequenceData, VLLM_TOKEN_ID_ARRAY_TYPE
-from vllm.model_executor.models.interfaces import SupportsMultiModal
+from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
 
 from TTS.tts.layers.xtts.latent_encoder import ConditioningEncoder # noqa
 from TTS.tts.layers.xtts.perceiver_encoder import PerceiverResampler # noqa
@@ -32,17 +33,147 @@ _AUDIO_PLACEHOLDER_TOKEN = 8192  # Using XTTS start_audio_token as placeholder
 _AUDIO_TOKENS_PER_SECOND = 6.25
 _CODE_STRIDE_LEN = 1024
 
+class GPT2Attention(nn.Module):
+    def __init__(
+            self,
+            config: PretrainedConfig,
+            cache_config: Optional[CacheConfig] = None,
+            quant_config: Optional[QuantizationConfig] = None,
+            prefix: str = "",
+    ):
+        super().__init__()
+        total_num_heads = config.num_attention_heads
+        self.hidden_size = config.hidden_size
+        tensor_model_parallel_world_size = get_tensor_model_parallel_world_size()
+        assert total_num_heads % tensor_model_parallel_world_size == 0
+        self.num_heads = total_num_heads // tensor_model_parallel_world_size
+        self.head_dim = self.hidden_size // total_num_heads
+        self.scale = self.head_dim**-0.5
+
+        self.c_attn = QKVParallelLinear(
+            self.hidden_size,
+            self.head_dim,
+            total_num_heads,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.c_attn",
+        )
+        self.c_proj = RowParallelLinear(
+            self.hidden_size,
+            self.hidden_size,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.c_proj",
+        )
+        self.attn = Attention(
+            self.num_heads,
+            self.head_dim,
+            scale=self.scale,
+            cache_config=cache_config,
+            quant_config=quant_config
+        )
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            kv_cache: torch.Tensor,
+            attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        qkv, _ = self.c_attn(hidden_states)
+        q, k, v = qkv.chunk(chunks=3, dim=-1)
+        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
+        attn_output, _ = self.c_proj(attn_output)
+        return attn_output
+
+
+class GPT2MLP(nn.Module):
+    def __init__(
+            self,
+            intermediate_size: int,
+            config: PretrainedConfig,
+            quant_config: Optional[QuantizationConfig] = None,
+            prefix: str = "",
+    ):
+        super().__init__()
+        hidden_size = config.hidden_size
+
+        self.c_fc = ColumnParallelLinear(
+            hidden_size,
+            intermediate_size,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.c_fc",
+        )
+        self.c_proj = RowParallelLinear(
+            intermediate_size,
+            hidden_size,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.c_proj",
+        )
+        self.act = get_act_fn(config.activation_function, quant_config, intermediate_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states, _ = self.c_fc(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states, _ = self.c_proj(hidden_states)
+        return hidden_states
+
+
+class GPT2Block(nn.Module):
+    def __init__(
+            self,
+            config: PretrainedConfig,
+            cache_config: Optional[CacheConfig] = None,
+            quant_config: Optional[QuantizationConfig] = None,
+            prefix: str = "",
+    ):
+        super().__init__()
+        hidden_size = config.hidden_size
+        inner_dim = config.n_inner if config.n_inner is not None else 4 * hidden_size
+
+        self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+        self.attn = GPT2Attention(
+            config,
+            cache_config,
+            quant_config,
+            prefix=f"{prefix}.attn"
+        )
+        self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+        self.mlp = GPT2MLP(
+            inner_dim,
+            config,
+            quant_config,
+            prefix=f"{prefix}.mlp"
+        )
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            kv_cache: torch.Tensor,
+            attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.ln_1(hidden_states)
+        attn_output = self.attn(
+            hidden_states=hidden_states,
+            kv_cache=kv_cache,
+            attn_metadata=attn_metadata,
+        )
+        hidden_states = attn_output + residual
+
+        residual = hidden_states
+        hidden_states = self.ln_2(hidden_states)
+        feed_forward_hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + feed_forward_hidden_states
+        return hidden_states
+
+
 
 def get_xtts_max_audio_tokens(ctx: InputContext) -> int:
     """Calculate maximum audio tokens based on text context and audio duration."""
-    # Based on GPT config and common XTTS settings
-    text_context = ctx.model_config.max_seq_len - 100  # Reserve space for text
-    # Allow for ~30 seconds of audio (similar to whisper chunks)
-    max_audio_duration = 30.0
-    audio_tokens = math.ceil(max_audio_duration * _AUDIO_TOKENS_PER_SECOND)
-    total_tokens = text_context + audio_tokens + 4  # +4 for special tokens
-
-    return min(total_tokens, 1000)  # Cap at 1000 tokens as specified
+    # Based on GPT config and XTTSv2 settings
+    return 608
 
 
 def dummy_seq_data_for_xtts(
@@ -73,7 +204,7 @@ def dummy_conditioning_for_xtts(
 ) -> dict:
     """Create dummy conditioning data for XTTS."""
     return {
-        "cond_latents": [(torch.zeros(80, 1024), 22050) for _ in range(audio_count)]
+        "audio": [(torch.zeros(80, 1024), 22050) for _ in range(audio_count)]
     }
 
 
@@ -106,10 +237,11 @@ def input_mapper_for_xtts(ctx: InputContext, data: object) -> MultiModalInputs:
 @MULTIMODAL_REGISTRY.register_input_mapper("audio", input_mapper_for_xtts)
 @MULTIMODAL_REGISTRY.register_max_multimodal_tokens("audio", get_xtts_max_audio_tokens)
 @INPUT_REGISTRY.register_dummy_data(dummy_data_for_xtts)
-class XttsGPT(nn.Module, SupportsMultiModal):
+class XttsGPT(nn.Module, SupportsMultiModal, SupportsPP):
     def __init__(
             self,
             config: PretrainedConfig,
+            multimodal_config: MultiModalConfig,
             cache_config: Optional[CacheConfig] = None,
             quant_config: Optional["QuantizationConfig"] = None,
     ):
@@ -119,14 +251,16 @@ class XttsGPT(nn.Module, SupportsMultiModal):
 
         # XTTS specific components
         self.conditioning_encoder = ConditioningEncoder(
-            80, config.n_embd, num_attn_heads=config.n_head
+            config.audio_config.mel_channels,
+            config.hidden_size,
+            num_attn_heads=config.num_attention_heads
         )
 
         if config.use_perceiver_resampler:
             self.conditioning_perceiver = PerceiverResampler(
-                dim=config.n_embd,
+                dim=config.hidden_size,
                 depth=2,
-                dim_context=config.n_embd,
+                dim_context=config.hidden_size,
                 num_latents=32,
                 dim_head=64,
                 heads=8,
@@ -144,7 +278,7 @@ class XttsGPT(nn.Module, SupportsMultiModal):
 
         # Prediction heads
         self.text_head = ColumnParallelLinear(
-            config.n_embd,
+            config.hidden_size,
             config.vocab_size,
             bias=False,
             quant_config=quant_config,
@@ -152,7 +286,7 @@ class XttsGPT(nn.Module, SupportsMultiModal):
         )
 
         self.mel_head = ColumnParallelLinear(
-            config.n_embd,
+            config.hidden_size,
             config.num_audio_tokens,
             bias=False,
             quant_config=quant_config,
@@ -176,15 +310,9 @@ class XttsGPT(nn.Module, SupportsMultiModal):
             conds = cond_input.unsqueeze(1)
         return conds
 
-    def forward(
-            self,
-            input_ids: torch.Tensor,
-            positions: torch.Tensor,
-            kv_caches: List[torch.Tensor],
-            attn_metadata: AttentionMetadata,
-            intermediate_tensors: Optional[IntermediateTensors] = None,
-            cond_latents: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
+    def forward(self, input_ids: torch.Tensor, positions: torch.Tensor, kv_caches: List[torch.Tensor],
+            attn_metadata: AttentionMetadata, intermediate_tensors: Optional[IntermediateTensors] = None,
+            cond_latents: Optional[torch.Tensor] = None ) -> torch.Tensor:
         """Forward pass following VLLM pattern."""
         if cond_latents is not None:
             # Combine conditioning with input embeddings
@@ -250,25 +378,39 @@ class XttsGPT2Model(nn.Module):
             self,
             config: PretrainedConfig,
             cache_config: Optional[CacheConfig] = None,
-            quant_config: Optional["QuantizationConfig"] = None,
+            quant_config: Optional[QuantizationConfig] = None,
             prefix: str = "",
     ):
         super().__init__()
         self.config = config
-        self.text_embedding = VocabParallelEmbedding(config.number_text_tokens, config.n_embd)
-        self.mel_embedding = VocabParallelEmbedding(config.num_audio_tokens, config.n_embd)
+
+        self.text_embedding = VocabParallelEmbedding(
+            config.number_text_tokens,
+            config.hidden_size
+        )
+        self.mel_embedding = VocabParallelEmbedding(
+            config.num_audio_tokens,
+            config.hidden_size
+        )
 
         self.text_pos_embedding = (
-            LearnedPositionEmbeddings(config.max_text_seq_len, config.n_embd)
-            if config.max_mel_seq_len != -1
-            else functools.partial(config.null_position_embeddings, dim=config.n_embd)
+            LearnedPositionEmbeddings(
+                config.max_text_tokens + 2,
+                config.hidden_size
+            )
+            if config.max_audio_tokens != -1
+            else functools.partial(config.null_position_embeddings, dim=config.hidden_size)
         )
+
         self.mel_pos_embedding = (
-            LearnedPositionEmbeddings(config.max_mel_seq_len, config.n_embd)
-            if config.max_mel_seq_len != -1
-            else functools.partial(config.null_position_embeddings, dim=config.n_embd)
+            LearnedPositionEmbeddings(
+                config.max_audio_tokens + 3,
+                config.hidden_size
+            )
+            if config.max_audio_tokens != -1
+            else functools.partial(config.null_position_embeddings, dim=config.hidden_size)
         )
-        # Build gpt blocks
+
         self.h = nn.ModuleList([
             GPT2Block(
                 config,
@@ -278,32 +420,38 @@ class XttsGPT2Model(nn.Module):
             ) for i in range(config.num_hidden_layers)
         ])
 
-        self.final_norm = nn.LayerNorm(
-            config.n_embd,
-            eps=config.layer_norm_epsilon
-        )
+        self.ln_f = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)
+
+    def get_input_embeddings(self):
+        return self.text_embedding
 
-    def forward( # TODO: this is not correct, allieeate it with the correct implementation
-        self,
-        input_ids: torch.Tensor,
-        position_ids: torch.Tensor,
-        kv_caches: List[torch.Tensor],
-        attn_metadata: AttentionMetadata,
-        intermediate_tensors: Optional[IntermediateTensors],
+    def forward(
+            self,
+            input_ids: Optional[torch.Tensor] = None,
+            positions: Optional[torch.Tensor] = None,
+            inputs_embeds: Optional[torch.Tensor] = None,
+            kv_caches: List[torch.Tensor] = None,
+            attn_metadata: AttentionMetadata = None,
+            intermediate_tensors: Optional[IntermediateTensors] = None,
     ) -> Union[torch.Tensor, IntermediateTensors]:
         if get_pp_group().is_first_rank:
-            inputs_embeds = self.wte(input_ids)
-            position_embeds = self.wpe(position_ids)
-            hidden_states = inputs_embeds + position_embeds
+            if inputs_embeds is None:
+                inputs_embeds = self.text_embedding(input_ids)
+            hidden_states = inputs_embeds
+
+            if positions is not None:
+                position_embeds = self.text_pos_embedding(positions)
+                hidden_states = hidden_states + position_embeds
         else:
             assert intermediate_tensors is not None
             hidden_states = intermediate_tensors["hidden_states"]
 
-        for i in range(self.start_layer, self.end_layer):
-            layer = self.h[i]
-            hidden_states = layer(hidden_states,
-                                  kv_caches[i - self.start_layer],
-                                  attn_metadata)
+        for i, block in enumerate(self.h):
+            hidden_states = block(
+                hidden_states,
+                kv_caches[i],
+                attn_metadata
+            )
 
         if not get_pp_group().is_last_rank:
             return IntermediateTensors({"hidden_states": hidden_states})