xtts2_modeling.py

import asyncio
import functools
import logging
import random
import time
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import Optional, List, Tuple, Union, AsyncGenerator, Dict, Any
from concurrent.futures import ThreadPoolExecutor

import librosa
import torch
import numpy as np
import torchaudio
import sounddevice as sd
import io
from torch import nn
from IPython.display import Audio, display

from vllm import AsyncLLMEngine, AsyncEngineArgs, SamplingParams, TokensPrompt, RequestOutput
from vllm.multimodal import MultiModalDataDict
from vllm.utils import Counter

from TTS.TTS.tts.layers.xtts.hifigan_decoder import HifiDecoder

from TTS.tts.layers.xtts.latent_encoder import ConditioningEncoder  # noqa
from TTS.tts.layers.xtts.perceiver_encoder import PerceiverResampler  # noqa

from .xtts2_config import XTTSConfig, XTTSGPTConfig
from .tokenizer import XTTSTokenizerFast

from ..xtts2_gpt.xtts2_gpt_modeling import LearnedPositionEmbeddings


def wav_to_mel_cloning(
        wav,
        mel_norms_file="../experiments/clips_mel_norms.pth",
        mel_norms=None,
        device=torch.device("cpu"),
        n_fft=4096,
        hop_length=1024,
        win_length=4096,
        power=2,
        normalized=False,
        sample_rate=22050,
        f_min=0,
        f_max=8000,
        n_mels=80,
):
    mel_stft = torchaudio.transforms.MelSpectrogram(
        n_fft=n_fft,
        hop_length=hop_length,
        win_length=win_length,
        power=power,
        normalized=normalized,
        sample_rate=sample_rate,
        f_min=f_min,
        f_max=f_max,
        n_mels=n_mels,
        norm="slaney",
    ).to(device)
    wav = wav.to(device)
    mel = mel_stft(wav)
    mel = torch.log(torch.clamp(mel, min=1e-5))
    if mel_norms is None:
        mel_norms = torch.load(mel_norms_file, map_location=device)
    mel = mel / mel_norms.unsqueeze(0).unsqueeze(-1)
    return mel


def load_audio(audiopath, sampling_rate):
    audio, lsr = torchaudio.load(audiopath)

    # Stereo to mono if needed
    if audio.size(0) != 1:
        audio = torch.mean(audio, dim=0, keepdim=True)

    if lsr != sampling_rate:
        audio = torchaudio.functional.resample(audio, lsr, sampling_rate)

    # Clip audio invalid values
    audio.clip_(-1, 1)
    return audio


@dataclass
class XTTSRequest:
    """Container for XTTS inference request data"""
    request_id: str
    text: Union[AsyncGenerator[str, None], str]
    language: str
    speaker_file: str  # Path to the speaker audio file
    generate_every_n_chars: Optional[int] = None
    temperature: float = 0.75
    top_p: float = 0.85
    top_k: int = 50
    repetition_penalty: float = 5.0
    length_penalty: float = 1.0
    do_sample: bool = True
    max_ref_length: int = 60
    gpt_cond_len: int = 30
    gpt_cond_chunk_len: int = 4


import threading

class HiddenStatesCollector:
    def __init__(self):
        self.outputs = {}
        self.lock = threading.Lock()

    def __call__(self, outputs: Optional[torch.Tensor], request_id: str):
        """Save outputs for a specific request"""
        with self.lock:
            if request_id not in self.outputs:
                self.outputs[request_id] = []
            self.outputs[request_id].append(outputs)

    def get_hidden_states(self, request_id) -> Optional[torch.Tensor]:
        with self.lock:
            outputs = self.outputs.pop(request_id, None)
        if outputs is not None:
            outputs = torch.cat(outputs, dim=0)
        return outputs

    def bind_to_request(self, request_id: str):
        def bound_collector(outputs: Optional[torch.Tensor], _request_id: str = None):
            self(outputs, request_id)
        return bound_collector

class ExtendedSamplingParams(SamplingParams, kw_only=True):
    """Extended sampling parameters that allows additional fields while maintaining compatibility with SamplingParams.

    This class inherits from SamplingParams and allows adding new required fields
    without conflicting with the base class's optional fields ordering.
    """
    hidden_state_collector: HiddenStatesCollector  # New required field


class LogitsRepetitionPenalizer:
    """A logits processor that applies repetition penalty to prevent repetitive text generation."""

    def __init__(self, repetition_penalty: float):
        if repetition_penalty < 0:
            raise ValueError("Repetition penalty must be non-negative")
        self.repetition_penalty = repetition_penalty

    def __call__(self, token_ids: List[int], logits: torch.Tensor) -> torch.Tensor:
        """Apply repetition penalty to the logits based on previous tokens."""
        # If no repetition penalty or no tokens to check, return original logits
        if self.repetition_penalty == 1.0 or not token_ids:
            return logits

        # Create a mask for the repeated tokens
        repeated_tokens = torch.tensor(token_ids,
                                       device=logits.device,
                                       dtype=torch.long)

        # Get logits of repeated tokens
        repeated_logits = logits[repeated_tokens]

        # Apply penalty: divide positive logits by penalty, multiply negative logits by penalty
        repeated_logits = torch.where(
            repeated_logits > 0,
            repeated_logits / self.repetition_penalty,
            repeated_logits * self.repetition_penalty
        )

        # Update only the logits for repeated tokens
        logits[repeated_tokens] = repeated_logits

        return logits


@dataclass
class XTTSOutput:
    """Container for XTTS inference output with integrated audio utilities"""
    request_id: str
    wav: np.ndarray
    sample_rate: int = 24000

    def to_tensor(self) -> torch.Tensor:
        """Convert numpy array to torch tensor"""
        if isinstance(self.wav, np.ndarray):
            return torch.from_numpy(self.wav)
        return self.wav

    def to_bytes(self, format: str = 'wav', sample_width: int = 2) -> bytes:
        """Convert audio to bytes format.

        Args:
            format: Output format ('wav' or 'raw')
            sample_width: Bit depth (1, 2, or 4 bytes per sample)

        Returns:
            Audio data as bytes
        """
        # Convert to tensor if needed
        wav_tensor = self.to_tensor()

        # Ensure correct shape (1, N) for torchaudio
        if wav_tensor.dim() == 1:
            wav_tensor = wav_tensor.unsqueeze(0)

        # Normalize to [-1, 1]
        wav_tensor = torch.clamp(wav_tensor, -1.0, 1.0)

        if format == 'wav':
            buffer = io.BytesIO()
            torchaudio.save(
                buffer,
                wav_tensor,
                self.sample_rate,
                format="wav",
                encoding="PCM_S" if sample_width == 2 else "PCM_F",
                bits_per_sample=sample_width * 8
            )
            return buffer.getvalue()

        elif format == 'raw':
            # Scale to appropriate range based on sample width
            if sample_width == 2:  # 16-bit
                wav_tensor = (wav_tensor * 32767).to(torch.int16)
            elif sample_width == 4:  # 32-bit
                wav_tensor = (wav_tensor * 2147483647).to(torch.int32)
            else:  # 8-bit
                wav_tensor = (wav_tensor * 127).to(torch.int8)
            return wav_tensor.cpu().numpy().tobytes()

        else:
            raise ValueError(f"Unsupported format: {format}")

    def save(self,
             filename: Union[str, Path],
             sample_rate: Optional[int] = None,
             format: Optional[str] = None) -> None:
        """Save audio to file.

        Args:
            filename: Output filename
            sample_rate: Optional new sample rate for resampling
            format: Optional format override (default: inferred from extension)
        """
        wav_tensor = self.to_tensor()
        if wav_tensor.dim() == 1:
            wav_tensor = wav_tensor.unsqueeze(0)

        # Resample if needed
        if sample_rate and sample_rate != self.sample_rate:
            wav_tensor = torchaudio.functional.resample(
                wav_tensor,
                orig_freq=self.sample_rate,
                new_freq=sample_rate
            )
        else:
            sample_rate = self.sample_rate

        torchaudio.save(
            filename,
            wav_tensor,
            sample_rate,
            format=format
        )

    def resample(self, new_sample_rate: int) -> 'XTTSOutput':
        """Create new XTTSOutput with resampled audio.

        Args:
            new_sample_rate: Target sample rate

        Returns:
            New XTTSOutput instance with resampled audio
        """
        wav_tensor = self.to_tensor()
        if wav_tensor.dim() == 1:
            wav_tensor = wav_tensor.unsqueeze(0)

        resampled = torchaudio.functional.resample(
            wav_tensor,
            orig_freq=self.sample_rate,
            new_freq=new_sample_rate
        )

        return XTTSOutput(
            request_id=self.request_id,
            wav=resampled.squeeze().numpy(),
            sample_rate=new_sample_rate
        )

    def get_info(self) -> Tuple[int, int, float]:
        """Get audio information.

        Returns:
            Tuple of (number of samples, sample rate, duration in seconds)
        """
        n_samples = len(self.wav)
        duration = n_samples / self.sample_rate
        return n_samples, self.sample_rate, duration

    @classmethod
    def from_tensor(cls, request_id: str, tensor: torch.Tensor, sample_rate: int = 24000) -> 'XTTSOutput':
        """Create XTTSOutput from torch tensor.

        Args:
            request_id: Request identifier
            tensor: Audio tensor
            sample_rate: Sample rate of the audio

        Returns:
            New XTTSOutput instance
        """
        return cls(
            request_id=request_id,
            wav=tensor.squeeze().cpu().numpy(),
            sample_rate=sample_rate
        )

    @classmethod
    def from_file(cls, request_id: str, filename: Union[str, Path]) -> 'XTTSOutput':
        """Create XTTSOutput from audio file.

        Args:
            request_id: Request identifier
            filename: Path to audio file

        Returns:
            New XTTSOutput instance
        """
        wav_tensor, sample_rate = torchaudio.load(filename)
        return cls.from_tensor(request_id, wav_tensor, sample_rate)

    def play(self) -> None:
        """Play the audio through the default sound device.
        For use in regular Python scripts/applications."""
        # Ensure the audio is in the correct format
        if isinstance(self.wav, torch.Tensor):
            audio_data = self.wav.cpu().numpy()
        else:
            audio_data = self.wav

        # Ensure float32 and normalize
        if audio_data.dtype != np.float32:
            audio_data = audio_data.astype(np.float32)
        audio_data = np.clip(audio_data, -1.0, 1.0)

        # Play the audio
        sd.play(audio_data, self.sample_rate)
        sd.wait()  # Wait until the audio is finished playing

    def display(self) -> Optional[Audio]:
        """Display audio player in Jupyter notebook.
        Returns Audio widget if in notebook, None otherwise."""
        try:
            # Convert to bytes
            audio_bytes = self.to_bytes(format='wav')

            # Create and display audio widget
            audio_widget = Audio(audio_bytes, rate=self.sample_rate, autoplay=False)
            display(audio_widget)
            return audio_widget
        except Exception as e:
            print(f"Could not display audio widget: {str(e)}")
            print("Try using .play() method instead")
            return None

    def preview(self) -> None:
        """Smart play method that chooses appropriate playback method."""
        try:
            # Try notebook display first
            if self.display() is None:
                # Fall back to sounddevice if not in notebook
                self.play()
        except Exception as e:
            print(f"Error playing audio: {str(e)}")


class Xtts(nn.Module):
    """Async XTTS model implementation using VLLM's AsyncEngine."""

    def __init__(self, hifi_config: XTTSConfig, gpt_config: XTTSGPTConfig, tensor_parallel_size: int = 1, **kwargs):
        super().__init__()

        self.hifi_config = hifi_config
        self.gpt_config = gpt_config
        self.mel_bos_token_id = gpt_config.start_audio_token
        self.mel_eos_token_id = gpt_config.stop_audio_token
        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.hidden_states_collector = HiddenStatesCollector()

        # Register buffer before creating modules
        self.register_buffer("mel_stats", torch.ones(80))

        # Initialize all nn.Module components
        self.conditioning_encoder = ConditioningEncoder(
            gpt_config.audio_config.mel_channels,
            gpt_config.hidden_size,
            num_attn_heads=gpt_config.num_attention_heads
        )

        self.text_embedding = nn.Embedding(
            gpt_config.number_text_tokens,
            gpt_config.hidden_size
        )

        self.text_pos_embedding = (
            LearnedPositionEmbeddings(
                gpt_config.max_text_tokens + 2,
                gpt_config.hidden_size,
                supports_pp=False
            )
            if gpt_config.max_audio_tokens != -1
            else functools.partial(gpt_config.null_position_embeddings, dim=gpt_config.hidden_size)
        )

        if gpt_config.use_perceiver_resampler:
            self.conditioning_perceiver = PerceiverResampler(
                dim=gpt_config.hidden_size,
                depth=2,
                dim_context=gpt_config.hidden_size,
                num_latents=32,
                dim_head=64,
                heads=8,
                ff_mult=4,
                use_flash_attn=False,
            )

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

        # Kept for model loading purposes
        self.text_head = nn.Linear(gpt_config.hidden_size, gpt_config.number_text_tokens, bias=True)
        self.final_norm = nn.LayerNorm(gpt_config.hidden_size, eps=1e-5, bias=True)

        # Initialize VLLM engine at the end
        self.init_vllm_engine()

        # Semaphore for concurrency control
        self.max_concurrency = 10
        self.semaphore = asyncio.BoundedSemaphore(self.max_concurrency)

    def half(self):
        # We cannot permit downcasting since it will throw an error while padding
        return

    def to(self, *args, **kwargs):
        # Block downcasting
        dtype = kwargs.get('dtype', None)
        if dtype == torch.float16 or dtype == torch.bfloat16:
            kwargs['dtype'] = torch.float32
        elif len(args) > 0 and (args[0] == torch.float16 or args[0] == torch.bfloat16):
            args = list(args)
            args[0] = torch.float32
            args = tuple(args)
        return super().to(*args, **kwargs)

    @property
    def device(self):
        """Get the current device of the model."""
        return next(self.parameters()).device

    @property
    def dtype(self):
        """Get the current dtype of the model."""
        return next(self.parameters()).dtype

    @staticmethod
    def get_memory_percentage(memory: int) -> float:
        """Get memory percentage."""
        total_memory = torch.cuda.get_device_properties(0).total_memory
        reserved_memory = torch.cuda.memory_reserved(0)
        allocated_memory = torch.cuda.memory_allocated(0)
        available_memory = total_memory - reserved_memory - allocated_memory
        return memory / available_memory

    def init_vllm_engine(self):
        """Initialize models with AsyncVLLMEngine."""
        engine_args = AsyncEngineArgs(
            model="AstraMindAI/xtts2-gpt",
            tensor_parallel_size=self.tp,
            dtype="auto",
            disable_log_stats=True,
            max_model_len=self.gpt_config.max_text_tokens + self.gpt_config.max_audio_tokens,
            gpu_memory_utilization=self.get_memory_percentage(3 * 1024 ** 3),
            trust_remote_code=True,
            enforce_eager=True,
            limit_mm_per_prompt={"audio": 1},
            max_num_batched_tokens=7296,
        )

        self.llm_engine = AsyncLLMEngine.from_engine_args(engine_args)

    @classmethod
    def from_pretrained(
            cls,
            pretrained_model_name_or_path: str,
            torch_dtype: torch.dtype = torch.float32,
            device_map: Optional[str] = "auto",
            tensor_parallel_size: int = 1,
            **kwargs,
    ) -> "Xtts":
        """Load pretrained XTTS model from HuggingFace Hub."""
        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="config.json"
            )
            with open(config_file, 'r') as f:
                config = json.load(f)

        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['gpt_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):
            hifigan_weights = hf_hub_download(
                repo_id=pretrained_model_name_or_path,
                filename="xtts-v2.safetensors"
            )
        else:
            hifigan_weights = os.path.join(pretrained_model_name_or_path, "xtts-v2.safetensors")

        import safetensors.torch

        # Load HiFi-GAN weights
        hifigan_state = safetensors.torch.load_file(hifigan_weights)
        model.load_state_dict(hifigan_state)

        # Set model properties
        model.config = config

        # Cast model to specified dtype
        model = model.to(torch_dtype)
        model = model.to('cuda')

        return model

    @staticmethod
    def load_audio(audio_path: Union[str, Path], sampling_rate: int = 22050) -> torch.Tensor:
        audio, lsr = torchaudio.load(audio_path)

        # Stereo to mono if needed
        if audio.size(0) != 1:
            audio = torch.mean(audio, dim=0, keepdim=True)

        if lsr != sampling_rate:
            audio = torchaudio.functional.resample(audio, lsr, sampling_rate)

        # Clip audio invalid values
        audio.clip_(-1, 1)
        return audio

    @torch.inference_mode()
    def get_speaker_embedding(self, audio, sr):
        audio_16k = torchaudio.functional.resample(audio, sr, 16000)
        return (
            self.hifigan_decoder.speaker_encoder.forward(audio_16k.to(self.device), l2_norm=True)
            .unsqueeze(-1)
            .to(self.device)
        )

    @torch.inference_mode()
    def get_gpt_cond_latents(self, audio, sr, length: int = 30, chunk_length: int = 6):
        """Compute the conditioning latents for the GPT model from the given audio."""
        if sr != 22050:
            audio = torchaudio.functional.resample(audio, sr, 22050)
        if length > 0:
            audio = audio[:, : 22050 * length]
        if self.gpt_config.use_perceiver_resampler:
            style_embs = []
            for i in range(0, audio.shape[1], 22050 * chunk_length):
                audio_chunk = audio[:, i: i + 22050 * chunk_length]

                # if the chunk is too short ignore it
                if audio_chunk.size(-1) < 22050 * 0.33:
                    continue

                mel_chunk = wav_to_mel_cloning(
                    audio_chunk,
                    mel_norms=self.mel_stats.cpu(),
                    n_fft=2048,
                    hop_length=256,
                    win_length=1024,
                    power=2,
                    normalized=False,
                    sample_rate=22050,
                    f_min=0,
                    f_max=8000,
                    n_mels=80,
                )
                style_emb = self.get_style_emb(mel_chunk.to(self.device), None)
                style_embs.append(style_emb)

            # mean style embedding
            cond_latent = torch.stack(style_embs).mean(dim=0)
        else:
            mel = wav_to_mel_cloning(
                audio,
                mel_norms=self.mel_stats.cpu(),
                n_fft=4096,
                hop_length=1024,
                win_length=4096,
                power=2,
                normalized=False,
                sample_rate=22050,
                f_min=0,
                f_max=8000,
                n_mels=80,
            )
            cond_latent = self.get_style_emb(mel.to(self.device))
        return cond_latent.transpose(1, 2)

    @torch.inference_mode()
    def get_conditioning_latents(
            self,
            audio_path,
            max_ref_length=30,
            gpt_cond_len=6,
            gpt_cond_chunk_len=6,
            librosa_trim_db=None,
            sound_norm_refs=False,
            load_sr=22050,
    ):
        """Get the conditioning latents for the GPT model from the given audio."""
        # Deal with multiple references
        assert isinstance(audio_path, str) or isinstance(audio_path, list), "audio_path must be a string or a list."

        if not isinstance(audio_path, list):
            audio_paths = [audio_path]
        else:
            audio_paths = audio_path

        speaker_embeddings = []
        audios = []
        for file_path in audio_paths:
            audio = load_audio(file_path, load_sr)
            audio = audio[:, : load_sr * max_ref_length].to(self.device).to(self.dtype)
            if sound_norm_refs:
                audio = (audio / torch.abs(audio).max()) * 0.75
            if librosa_trim_db is not None:
                audio = librosa.effects.trim(audio, top_db=librosa_trim_db)[0]

            # Compute latents for the decoder
            speaker_embedding = self.get_speaker_embedding(audio, load_sr)
            speaker_embeddings.append(speaker_embedding)

            audios.append(audio)

        # Merge all the audios and compute the latents for the GPT
        full_audio = torch.cat(audios, dim=-1)
        gpt_cond_latents = self.get_gpt_cond_latents(
            full_audio, load_sr, length=gpt_cond_len, chunk_length=gpt_cond_chunk_len
        )  # [1, 1024, T]

        speaker_embedding = torch.stack(speaker_embeddings)
        speaker_embedding = speaker_embedding.mean(dim=0)

        return gpt_cond_latents, speaker_embedding

    def get_style_emb(self, cond_input: torch.Tensor, return_latent: bool = False) -> torch.Tensor:
        """Get conditioning embeddings from mel spectrograms."""
        if not return_latent:
            if cond_input.ndim == 4:
                cond_input = cond_input.squeeze(1)
            conds = self.conditioning_encoder(cond_input)

            if hasattr(self, 'conditioning_perceiver'):
                conds = self.conditioning_perceiver(
                    conds.permute(0, 2, 1)
                ).transpose(1, 2)
        else:
            conds = cond_input.unsqueeze(1)
        return conds

    async def prepare_text_tokens_async(self, text: str, language: str, split_text=False) \
            -> Tuple[List[Union[int, List[int]]], List[torch.Tensor]]:
        """Prepare text tokens for the given text and language."""

        async def elaborate_tokens(text_tokens: List[int]) -> torch.Tensor:
            text_tokens.insert(0, self.tokenizer.bos_token_id)
            text_tokens.append(self.tokenizer.eos_token_id)
            return torch.tensor(text_tokens).unsqueeze(0).to(self.text_embedding.weight.device)

        async def embed_tokens(text_tokens: Union[torch.Tensor, List[torch.Tensor]]) -> List[torch.Tensor]:
            embeds = []
            if isinstance(text_tokens, list):
                for list_element in text_tokens:
                    embeds.append(self.text_embedding(list_element) + self.text_pos_embedding(list_element))
            else:
                embeds.append(self.text_embedding(text_tokens) + self.text_pos_embedding(text_tokens))
            return embeds

        fake_tokens_for_audio_generation = []
        if split_text:
            text_tokens = self.tokenizer.batch_encode_with_split(text, lang=[language])
            for idx, text_token in enumerate(text_tokens):
                text_tokens[idx] = await elaborate_tokens(text_token)
                fake_tokens_for_audio_generation.append([1] * len(text_token))
        else:
            text_tokens = self.tokenizer.batch_encode(text, lang=[language])
            text_tokens = await elaborate_tokens(text_tokens)
            fake_tokens_for_audio_generation = [1] * len(text_tokens)
        return fake_tokens_for_audio_generation, await embed_tokens(text_tokens)

    async def prepare_inputs_async(self, text: str, language: str, speaker_file: Union[str, Path],
                                   max_ref_length: int, gpt_cond_len: int, gpt_cond_chunk_len: int, split_text: bool) \
            -> Tuple[List[List[int]], List[torch.Tensor], torch.Tensor]:
        """Prepare input text with conditioning tokens. Return combined conditioning latents"""
        # Tokenize text based on the language
        text_tokens, text_embeddings = await self.prepare_text_tokens_async(text, language, split_text)

        # Load the speaker file and convert it to a tensor
        gpt_cond_latent, speaker_embeddings = await self.get_conditioning_latents_async(
            speaker_file,
            max_ref_length,
            gpt_cond_len,
            gpt_cond_chunk_len
        )

        cond_latents = []
        for text_embedding in text_embeddings:
            # Concatenate along sequence dimension
            cond_latents.append((torch.cat([gpt_cond_latent, text_embedding], dim=1).squeeze(0)
                                 .to(self.llm_engine.engine.model_config.dtype)))

        return text_tokens, cond_latents, speaker_embeddings

    async def get_conditioning_latents_async(
            self,
            audio_path,
            max_ref_length=30,
            gpt_cond_len=6,
            gpt_cond_chunk_len=6,
            librosa_trim_db=None,
            sound_norm_refs=False,
            load_sr=22050,
    ):
        """Async version of get_conditioning_latents with concurrency control."""
        async with self.semaphore:
            # Run the original get_conditioning_latents in executor
            result = await asyncio.get_event_loop().run_in_executor(
                None,
                functools.partial(self.get_conditioning_latents,
                                  audio_path,
                                  max_ref_length,
                                  gpt_cond_len,
                                  gpt_cond_chunk_len,
                                  librosa_trim_db,
                                  sound_norm_refs,
                                  load_sr)
            )
        return result

    async def get_model_logits(self, token_ids: List[int], conditioning: MultiModalDataDict) -> torch.Tensor:
        """Get model logits for a specific request"""
        request_id = uuid.uuid4().hex

        # Add start and end tokens
        token_ids = [self.mel_bos_token_id] + token_ids + [self.mel_eos_token_id] * 5

        engine_inputs = TokensPrompt(prompt_token_ids=token_ids)
        engine_inputs["multi_modal_data"] = conditioning

        # Bind the collector to this request
        bound_collector = self.hidden_states_collector.bind_to_request(request_id)

        # Set up sampling parameters with the bound collector
        sampling_params = ExtendedSamplingParams(
            detokenize=False,
            max_tokens=1,
            hidden_state_collector=bound_collector,
        )

        # Generate with unique request ID
        generator = self.llm_engine.generate(
            prompt=engine_inputs,
            sampling_params=sampling_params,
            request_id=request_id
        )

        # Consume the generator with a timeout
        try:
            async def consume_generator():
                async for _ in generator:
                    pass

            await asyncio.wait_for(consume_generator(), timeout=300)
        except asyncio.TimeoutError:
            raise RuntimeError("Timeout while generating logits")

        # Get the collected hidden states
        hidden_states = self.hidden_states_collector.get_hidden_states(request_id)

        if hidden_states is None:
            raise RuntimeError(f"No hidden states collected for request {request_id}")

        return hidden_states[-len(token_ids):, ...].unsqueeze(0).to(self.device).to(self.dtype)


    async def process_tokens_to_speech(
            self,
            generators: List[AsyncGenerator[RequestOutput, None]],
            speaker_embeddings: torch.Tensor,
            multimodal_data: List[torch.Tensor],
            chunk_size: int = 20,
    ) -> AsyncGenerator[XTTSOutput, None]:
        """
        Process multiple token generators concurrently and emit results sequentially.
        Uses a queue-based approach to handle multiple generators reliably.
        """
        # Create a queue for each generator to store its results
        queues = [asyncio.Queue() for _ in generators]

        # Create tasks for processing each generator
        tasks = []
        for i, generator in enumerate(generators):
            task = asyncio.create_task(
                self._process_single_generator(
                    generator,
                    queues[i],
                    speaker_embeddings,
                    multimodal_data[i],
                    chunk_size
                )
            )
            tasks.append(task)

        try:
            # Process queues in sequence
            for i, queue in enumerate(queues):
                while True:
                    result = await queue.get()
                    if result is None:
                        # This generator has finished
                        break
                    else:
                        yield result

        finally:
            # Ensure all tasks are properly cleaned up
            for task in tasks:
                if not task.done():
                    task.cancel()
            await asyncio.gather(*tasks, return_exceptions=True)

    async def _process_single_generator(
            self,
            generator: AsyncGenerator[RequestOutput, None],
            queue: asyncio.Queue,
            speaker_embeddings: torch.Tensor,
            gpt_embed_input: torch.Tensor,
            chunk_size: int
    ) -> None:
        """Process a single generator and put results in its queue."""
        try:
            last_decoded_token = 0
            accumulated_tokens = []

            async for output in generator:
                # Get new tokens
                new_tokens = output.outputs[0].token_ids[last_decoded_token:]
                accumulated_tokens.extend(new_tokens)
                last_decoded_token = len(accumulated_tokens)

                # Process tokens when we have enough or it's the final output
                if output.finished:# or len(accumulated_tokens) >= chunk_size: se lascio con acculated token mi ripete gli stesis toke, why??
                    # Process the accumulated tokens
                    hidden_states = await self.get_model_logits(
                        accumulated_tokens,
                        {
                            "audio": {
                                'embeds': gpt_embed_input,
                                "is_logits_only_mode": True
                            }
                        }
                    )

                    # Generate audio segment
                    wav = await asyncio.get_event_loop().run_in_executor(
                        self.executor,
                        lambda: self.hifigan_decoder.inference(
                            hidden_states,
                            g=speaker_embeddings
                        ).cpu().numpy().squeeze()
                    )

                    # Put result in queue
                    await queue.put(XTTSOutput(
                        request_id=output.request_id,
                        wav=wav
                    ))

                    # Reset accumulated tokens
                    accumulated_tokens = []

                if output.finished:
                    break

        except Exception as e:
            logging.error(f"Error in generator processing: {e}")
        finally:
            # Signal completion
            await queue.put(None)

    async def generate_speech_async_from_streaming_source(self, request: XTTSRequest) -> AsyncGenerator[XTTSOutput, None]:
        """Generate speech for streaming source of text, making a streaming source of audio tokens and then decoding
        and returning a streaming audio response."""
        assert isinstance(request.text, AsyncGenerator), "Text must be an AsyncGenerator for streaming source."
        # Prepare input with conditioning
        gpt_cond_latent, speaker_embeddings = await self.get_conditioning_latents_async(
            request.speaker_file,
            request.max_ref_length,
            request.gpt_cond_len,
            request.gpt_cond_chunk_len
        )
        sampling_params = SamplingParams(
            temperature=request.temperature,
            top_p=request.top_p,
            detokenize=False,
            top_k=request.top_k,
            logits_processors=[LogitsRepetitionPenalizer(request.repetition_penalty)],
            repetition_penalty=1.0,  # Since we're handling repetition penalty manually
            max_tokens=self.gpt_config.gpt_max_audio_tokens,
            ignore_eos=True,  # Ignore the tokenizer eos token since it is for textual generation
            stop_token_ids=[self.mel_eos_token_id],
        )

        accumulated_text = ""
        async for text in request.text:
            text = text.strip()
            accumulated_text += text

            if len(accumulated_text) > request.generate_every_n_chars:
                tokens, embeddings = await self.prepare_text_tokens_async(accumulated_text, request.language)
                gpt_embed_input = [torch.cat([gpt_cond_latent, embeddings[0]], dim=0)]

                engine_inputs = TokensPrompt(prompt_token_ids=tokens)
                if gpt_embed_input is not None:
                    engine_inputs["multi_modal_data"] = {"audio": {"embeds": gpt_embed_input, "is_logits_only_mode": False}}
                token_generator = [self.llm_engine.generate(
                    prompt=engine_inputs,
                    sampling_params=sampling_params,
                    request_id=request.request_id,
                )]
                # Process tokens to speech
                async for output in self.process_tokens_to_speech(
                        token_generator,
                        speaker_embeddings,
                        gpt_embed_input,
                        chunk_size=50
                ):
                    yield output

                accumulated_text = ""

    async def generate_speech_from_text_async(self, request: XTTSRequest) -> AsyncGenerator[XTTSOutput, None]:
        """Generate speech for a single request asynchronously."""
        # Prepare input with conditioning
        tokens_list, gpt_embed_inputs, speaker_embeddings = await self.prepare_inputs_async(
            request.text,
            request.language,
            request.speaker_file,
            request.max_ref_length,
            request.gpt_cond_len,
            request.gpt_cond_chunk_len,
            split_text=True  # Split text to avoid OOM on big texts
        )

        # Start all requests in parallel
        generators = []
        for seq_index, sequence in enumerate(tokens_list):
            sampling_params = SamplingParams(
                temperature=request.temperature,
                top_p=request.top_p,
                detokenize=False,
                top_k=request.top_k,
                logits_processors=[LogitsRepetitionPenalizer(request.repetition_penalty)],
                repetition_penalty=1.0,  # Since we're handling repetition penalty manually
                max_tokens=self.gpt_config.gpt_max_audio_tokens,
                ignore_eos=True,  # Ignore the tokenizer eos token since it is for textual generation
                stop_token_ids=[self.mel_eos_token_id],
            )

            engine_inputs = TokensPrompt(prompt_token_ids=sequence)
            if gpt_embed_inputs is not None:
                engine_inputs["multi_modal_data"] = {"audio": {"embeds": gpt_embed_inputs[seq_index], "is_logits_only_mode": False}}

            # Get audio token generator from VLLM
            token_generator = self.llm_engine.generate(
                prompt=engine_inputs,
                sampling_params=sampling_params,
                request_id=f"{request.request_id}_{seq_index}",
            )
            generators.append(token_generator)

        # Process tokens to speech
        async for output in self.process_tokens_to_speech(
                generators,
                speaker_embeddings,
                gpt_embed_inputs,
                chunk_size=50
        ):
            yield output

    def generate_speech_from_text(self, request: XTTSRequest) -> List[XTTSOutput]:
        """
        Synchronous wrapper for generate_speech_from_text_async.

        Args:
            request: XTTSRequest object containing generation parameters

        Returns:
            List of XTTSOutput containing the generated speech segments
        """

        async def _collect_outputs():
            outputs = []
            async for output in self.generate_speech_from_text_async(request):
                outputs.append(output)
            return outputs

        # Run the async code in an event loop
        import asyncio

        # Get or create an event loop
        try:
            loop = asyncio.get_event_loop()
        except RuntimeError:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)

        if loop.is_running():
            # Create a new loop if the current one is running
            new_loop = asyncio.new_event_loop()
            results = new_loop.run_until_complete(_collect_outputs())
            new_loop.close()
        else:
            results = loop.run_until_complete(_collect_outputs())

        return results