models/TTS_utils.py

import torch
from TTS.api import TTS
import time
import torchaudio
from TTS.tts.configs.xtts_config import XttsConfig
from TTS.tts.models.xtts import Xtts
import sounddevice as sd


def xtts_v2():
    """
    Load and return the XTTS v2 model.

    This function initializes the XTTS v2 model from the 🐸TTS library.
    The model is configured to use a GPU if available, otherwise it defaults to CPU.

    Returns:
        TTS: The initialized XTTS v2 model.
        
    Example usage:
        tts = xtts_v2()
    """
    # Get device
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # List available 🐸TTS models
    # print(TTS().list_models())

    # Init TTS
    tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to(device)


    return tts

def load_manual_xtts_v2(config_path, checkpoint_path):
    """
    Load the XTTS v2 model manually with configuration and checkpoint files.

    Args:
        config_path (str): Path to the configuration file.
            Example: "path/to/config.json"
        checkpoint_path (str): Path to the checkpoint directory.
            Example: "path/to/checkpoint/"

    Returns:
        Xtts: The loaded XTTS v2 model.
        
    Example usage:
        model = load_manual_xtts_v2("config.json", "checkpoint/")
    """
    print("Loading model...")
    config = XttsConfig()
    config.load_json(config_path)
    model = Xtts.init_from_config(config)
    model.load_checkpoint(config, checkpoint_dir=checkpoint_path, use_deepspeed=True)
    model.cuda()
    
    return model

import json
import concurrent.futures

# ----------------- StreamXTTSV2 -----------------
def get_text_order(json_path, num_elements, ):
    """
    Retrieve a specified number of text elements from a JSON file and update the file.

    Args:
        json_path (str): Path to the JSON file.
            Example: "path/to/data.json"
        num_elements (int): Number of elements to retrieve.
            Example: 3

    Returns:
        list: A list of tuples containing text, order, original_path, path_to_save, and language.
        
    Example usage:
        text_order = get_text_order("data.json", 3)
    """
    with open(json_path) as f:
        data = json.load(f)
    # check if the data is empty
    if not data['text']:
        return "No more text to process"
    if len(data['text']) < num_elements:
        num_elements = len(data['text'])
    text = data['text'][:num_elements]
    order = data['order'][:num_elements]
    original_path = data['original_path'][:num_elements]
    path_to_save = data['path_to_save'][:num_elements]
    language = data['language'][:num_elements]
    # remove the first elements
    data['text'] = data['text'][num_elements:]
    data['order'] = data['order'][num_elements:]
    data['original_path'] = data['original_path'][num_elements:]
    data['path_to_save'] = data['path_to_save'][num_elements:]
    data['language'] = data['language'][num_elements:]
    data['original_text'] = data['original_text'][num_elements:]
    # write the data back to the file
    with open(json_path, 'w') as f:
        json.dump(data, f)
    # make it return an array of arrays of text and order
    result = [i for i in zip(text, order, original_path, path_to_save, language)]
    return result

def append_text_order(json_path, text, order, original_path, path_to_save, language, original_text=None):
    """
    Append a text order to a JSON file.

    Args:
        json_path (str): Path to the JSON file.
            Example: "path/to/data.json"
        text (str): The text to append.
            Example: "Hello, world!"
        order (int): The order index.
            Example: 1
        original_path (str): Path to the original file.
            Example: "path/to/original.wav"
        path_to_save (str): Path to save the processed file.
            Example: "path/to/save.wav"
        language (str): Language of the text.
            Example: "en"
        original_text (str, optional): The original text if available.
            Example: "Hola, mundo!"

    Example usage:
        append_text_order("data.json", "Hello, world!", 1, "original.wav", "save.wav", "en", "Hola, mundo!")
    """
    with open(json_path) as f:
        data = json.load(f)
    data['text'].append(text)
    data['order'].append(order)
    data['original_path'].append(original_path)
    data['path_to_save'].append(path_to_save)
    data['language'].append(language)
    data['original_text'].append(original_text)
    with open(json_path, 'w') as f:
        json.dump(data, f)
# ----------------- StreamXTTSV2 -----------------
class StreamXTTSV2:
    """
    A class to handle streaming TTS using XTTS v2 model.

    Args:
        model (Xtts): The XTTS v2 model.
        sample_rate (int, optional): The sample rate for audio playback. Default is 24000.
        buffer_size (int, optional): The buffer size for audio playback. Default is 2.
    """
    def __init__(self, model, sample_rate=24000, buffer_size=2):
        self.model = model
        #self.gpt_cond_latent = gpt_cond_latent
        #self.speaker_embedding = speaker_embedding
        self.sample_rate = sample_rate
        self.buffer_size = buffer_size
        self.speed = 0.95
        self.stream_chunk_size = 40
        self.buffer = torch.Tensor().to('cpu')
        self.chunk_save = torch.Tensor().to('cpu')
        self.is_playing = False
        self.tasks_order = []
        self.order = 0
        self.initial = True

    def chunk_callback(self, chunk, i, output_dir, order):
        """
        Callback function to handle each chunk of audio during streaming.

        Args:
            chunk (torch.Tensor): The audio chunk.
                Example: tensor([0.1, 0.2, 0.3])
            i (int): The chunk index.
                Example: 1
            output_dir (str): Directory to save the chunk.
                Example: "output/"
            order (int): The order index.
                Example: 1
        """
        # Accumulate chunk into buffer
        self.buffer = torch.cat((self.buffer, chunk.squeeze().to('cpu')), dim=-1)
        self.chunk_save = torch.cat((self.chunk_save, chunk.squeeze().to('cpu')), dim=-1)
        chunk_filename = output_dir + f"chunk_{i}_{order}.wav"
        print(self.sample_rate)
        torchaudio.save(chunk_filename, self.chunk_save.unsqueeze(0), self.sample_rate)
        print(f"Chunk saved as {chunk_filename}")
        self.chunk_save = torch.Tensor().to('cpu')
        
        # Check if buffer has enough chunks to start playing
        if not self.is_playing and len(self.buffer) >= self.buffer_size:
            self.start_playback()

    def start_playback(self):
        """Start audio playback."""
        self.is_playing = True
        sd.play(self.buffer.numpy(), self.sample_rate, blocking=False)
        self.buffer = torch.Tensor().to('cpu')  # Reset buffer after starting playback

    def play(self, chunks, output_dir, path_to_save, order):
        """
        Play the audio chunks and save the complete audio.

        Args:
            chunks (list): List of audio chunks.
                Example: [tensor([0.1, 0.2, 0.3]), tensor([0.4, 0.5, 0.6])]
            output_dir (str): Directory to save the chunks.
                Example: "output/"
            path_to_save (str): Path to save the complete audio file.
                Example: "output/complete.wav"
            order (int): The order index.
                Example: 1
        """
        t0 = time.time()
        

        for i, chunk in enumerate(chunks):
            #print(chunk)
            if i == 0:
                print(f"Time to first chunk: {time.time() - t0}")
            print(f"Received chunk {i} of audio length {chunk.shape[-1]}")
            self.chunk_callback(chunk, i, output_dir, order)
        
        # Ensure all remaining audio is played
        while sd.get_stream().active:
            time.sleep(0.1)
        if len(self.buffer) > 0:
            sd.play(self.buffer.numpy(), self.sample_rate, blocking=True)
        
        # Save the complete audio to a file
        torchaudio.save(path_to_save, self.buffer.unsqueeze(0), self.sample_rate)
        print(f"Total audio length: {self.buffer.shape[-1]}")
        print("Audio playback finished.")
        #self.order += 1
        

    def inference_and_play(self, json_path, output_dir):
        """
        Perform inference and play the generated audio.

        Args:
            json_path (str): Path to the JSON file containing text orders.
                Example: "path/to/data.json"
            output_dir (str): Directory to save the chunks.
                Example: "output/"
        """
        print("Inference...")
        

        
        self.texts = get_text_order(json_path, 3)
        
        if self.texts == "No more text to process":
            print("No more text to process")
            return
        if self.texts == "Not enough text to process":
            print("Not enough text to process")
            return
        # is it returns a list of text and order
        if self.texts is not None:
            #print(self.texts)
            self.gpt_cond_latent, self.speaker_embedding = self.model.get_conditioning_latents(audio_path=[self.texts[0][2]])
            path_to_save = self.texts[0][3]
            #print(self.gpt_cond_latent, self.speaker_embedding)
            #print(self.texts)
            
            with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
                #text, order = get_text_order(texts)
                #print(text, order)
                futures = []
                print(self.texts)
                
                for text, i, path_a, path_s, lang in self.texts:
                    #print(text, i, path)
                    print(f"Processing text {i}: {text}")
                    print(f"Processing text {i}: {lang}")
                    future = executor.submit(self.model.inference_stream, text, lang, self.gpt_cond_latent, self.speaker_embedding, stream_chunk_size=self.stream_chunk_size, speed=self.speed)
                    #print(future.result())
                    futures.append(future)
                    
                
                for future, text in zip(futures, self.texts):
                    #print(text)
                    chunks = future.result()
                    print(text[1])
                    self.play(chunks, output_dir, path_to_save, text[1]) 
                    self.buffer = torch.Tensor().to('cpu')
            
            self.inference_and_play(json_path, output_dir )


def stream_prod(model, json_path, directory_path):
    """
    Stream production function for XTTS v2.

    Args:
        model (Xtts): The XTTS v2 model.
            Example: model = load_manual_xtts_v2("config.json", "checkpoint/")
        json_path (str): Path to the JSON file containing text orders.
            Example: "path/to/data.json"
        directory_path (str): Directory to save the chunks.
            Example: "output/"
    """
    streamer = StreamXTTSV2(model, buffer_size=2)
    results = streamer.inference_and_play(json_path, directory_path)
    if results is  None:
        time.sleep(3)
        stream_prod(model, json_path, directory_path)
    return "Streaming finished"


def just_inference(model, original_path, output_dir, text, lang, order):
    """
    Perform inference and save the generated audio.

    Args:
        model (Xtts): The XTTS v2 model.
            Example: model = load_manual_xtts_v2("config.json", "checkpoint/")
        original_path (str): Path to the original audio file.
            Example: "path/to/original.wav"
        output_dir (str): Directory to save the generated audio file.
            Example: "output/"
        text (str): The text to be synthesized.
            Example: "Hello, world!"
        lang (str): The language of the text.
            Example: "en"
        order (int): The order index.
            Example: 1

    Returns:
        tuple: A tuple containing the path to the saved audio file and the time to first chunk.
            Example: ("output/complete.wav", 1.23)
    """
    print("Inference...")
    path_to_save = output_dir
    t0 = time.time()
    gpt_cond_latent, speaker_embedding = model.get_conditioning_latents(audio_path=[original_path])

    chunks = model.inference_stream(
        text,
        lang,
        gpt_cond_latent,
        speaker_embedding,
        stream_chunk_size= 15 ,
        speed=0.95
        #temperature=0.1,
        #enable_text_splitting=True,
    )
    full_audio = torch.Tensor().to('cpu') 
    wav_chuncks = []
    for i, chunk in enumerate(chunks):
        if i == 1:
            time_to_first_chunk = time.time() - t0
            print(f"Time to first chunck: {time_to_first_chunk}")
        print(f"Received chunk {i} of audio length {chunk.shape[-1]}")
        wav_chuncks.append(chunk)
        full_audio = torch.cat((full_audio, chunk.squeeze().to('cpu')), dim=-1)
    
    
    
    # Save the complete audio to a file
    torchaudio.save(path_to_save, full_audio.unsqueeze(0), 24000)
    
    print("Inference finished")
    return path_to_save, time_to_first_chunk