localTTS.py

import torch
import numpy as np
from scipy.signal import firwin, lfilter
from scipy.io.wavfile import write
import soxr
import json
import os
import sys

# Ensure TTS-TT2 and HiFi-GAN are in your Python path
tacotron_lib_path = os.path.join(os.path.dirname(__file__), 'libs', 'Local-Tacotron-2')
hifigan_lib_path = os.path.join(os.path.dirname(__file__), 'libs', 'Local-Hifi-Gan')

sys.path.append(tacotron_lib_path)
sys.path.append(hifigan_lib_path)

from hparams import create_hparams
from model import Tacotron2
from text import text_to_sequence
from env import AttrDict
from meldataset import mel_spectrogram, MAX_WAV_VALUE
from models import Generator
from denoiser import Denoiser

from .pronunciation import PronunciationProcessor

"""LocalTTS: A class for local text-to-speech synthesis using Tacotron2 and HiFi-GAN models."""
class LocalTTS:
    def __init__(self, deviceType: str, tacotron_dir: str = '1_TACOTRON_MODELS', hifigan_dir: str = '0_HIFIGAN_MODELS', cmu_dict_dir: str = 'CMU_DICTIONARY', superres_dir: str = 'SR_hifigan'):
        """
        Initializes the LocalTTS class by loading Tacotron2 and HiFi-GAN models.
        Args:
            deviceType (str): The device type to use ('cpu' or 'cuda').
            tacotron_dir (str): Directory containing Tacotron2 model files.
            hifigan_dir (str): Directory containing HiFi-GAN model files.
            cmu_dict_dir (str): Directory containing the CMU pronunciation dictionary.
            superres_dir (str): Directory containing the super-resolution HiFi-GAN model files.
        """

        self.deviceType = deviceType
        self.device = torch.device(self.deviceType)
        self.is_cuda = self.device.type == 'cuda'
        self.super_res = 3

        self.pronunciation = PronunciationProcessor(cmu_dict_dir)

        self.tacotron2_models = {}
        self.hifigan_models = {}

        self.__load_all_tacotron2(tacotron_dir)
        self.__load_all_hifigan(hifigan_dir)

        self.available_tacotron_models = set(self.tacotron2_models.keys())
        self.available_hifigan_models = set(self.hifigan_models.keys())
        
        hifigan2, h2, denoiser2 = self.__load_hifigan(os.path.join(superres_dir, 'Superres_Twilight_33000'), 'config_32k')
        self.hifigan_models['Superres_Twilight_33000'] = (hifigan2, h2, denoiser2)
        self.high_pass_filter = firwin(101, cutoff=10500, fs=h2.sampling_rate, pass_zero=False)
    
    def _apply_pronounciation_dictionary(self, text : str, punctuation=r"!?,.;:'\"", EOS_Token=True) -> str:
        """
        Applies the pronunciation dictionary to the input text.
        Args:
            text (str): The input text to process.
            punctuation (str): Punctuation characters to consider.
            EOS_Token (bool): Whether to append an end-of-sentence token.
            Returns:
            str: The processed text with pronunciation dictionary applied.
        """
        return self.pronunciation.apply(
            text,
            punctuation=punctuation,
            eos_token=EOS_Token
        )
    
    def __load_all_tacotron2(self, tacotron2_dir : str):
        """
        Loads all Tacotron2 models from the specified directory.
        Args:
            tacotron2_dir (str): Directory containing Tacotron2 model files.
        """
        for file in os.listdir(tacotron2_dir):
            if file.startswith('.'): continue
            model_name = file.split('.')[0]
            model_path = os.path.join(tacotron2_dir, file)
            model, hparams = self.__load_tacotron2(model_path)
            self.tacotron2_models[model_name] = (model, hparams)
            print(f"Loaded Tacotron2 model: {model_name}")

    def __load_tacotron2(self, model_path : str) -> tuple:
        """
        Loads a Tacotron2 model from the specified path.
        Args:
            model_path (str): Path to the Tacotron2 model file.
        Returns:
            tuple: A tuple containing the Tacotron2 model and its hyperparameters.
        """
        hparams = create_hparams()
        hparams.sampling_rate = 22050
        hparams.max_decoder_steps = 3000
        hparams.gate_threshold = 0.25
        model = Tacotron2(hparams)
        if self.is_cuda:
            state_dict = torch.load(model_path, weights_only=True)['state_dict']
            model.load_state_dict(state_dict)
            _ = model.cuda().eval().half()
        else:
            state_dict = torch.load(model_path, map_location=self.device, weights_only=True)['state_dict']
            model.load_state_dict(state_dict)
            _ = model.to(self.device).eval()
        return model, hparams
    
    def __load_all_hifigan(self, hifigan_dir : str):
        """
        Loads all HiFi-GAN models from the specified directory.
        Args:
            hifigan_dir (str): Directory containing HiFi-GAN model files.
        """
        for file in os.listdir(hifigan_dir):
            model_name = file.split('.')[0]
            model_path = os.path.join(hifigan_dir, file)
            hifigan, h, denoiser = self.__load_hifigan(model_path, 'config_v1')
            self.hifigan_models[model_name] = (hifigan, h, denoiser)
            print(f"Loaded HiFi-GAN model: {model_name}")

    def __load_hifigan(self, model_path : str, conf_name : str) -> tuple:
        """
        Loads a HiFi-GAN model from the specified path and configuration name.
        Args:
            model_path (str): Path to the HiFi-GAN model file.
            conf_name (str): Configuration name for the HiFi-GAN model.
        Returns:
            tuple: A tuple containing the HiFi-GAN model, its hyperparameters, and the denoiser.
        """
        with open(os.path.join(hifigan_lib_path, f'{conf_name}.json')) as f:
            json_config = json.loads(f.read())
        h = AttrDict(json_config)
        torch.manual_seed(h.seed)
        hifigan = Generator(h).to(self.device)
        state_dict_g = torch.load(model_path, map_location=self.device, weights_only=True)
        hifigan.load_state_dict(state_dict_g['generator'])
        hifigan.eval()
        hifigan.remove_weight_norm()
        denoiser = Denoiser(hifigan, mode='normal')
        return hifigan, h, denoiser
    
    def infer(self, text : str, model_name : str, hifigan_model_name : str, output_file : str, pronounciation_dictionary : bool = True, EOS_Token : bool = True):
        """
        Performs inference to generate speech audio from the given text using specified models.
        Args:
            text (str): Input text to be synthesized.
            model_name (str): Name of the Tacotron2 model to use.
            hifigan_model_name (str): Name of the HiFi-GAN model to use.
            output_file (str): Path to save the generated audio file.
            pronounciation_dictionary (bool, optional): Whether to apply a pronunciation dictionary. Defaults to True.
            EOS_Token (bool, optional): Whether to append an end-of-sentence token. Defaults to True.
        """
        if model_name not in self.available_tacotron_models:
            raise ValueError(f"Tacotron2 model not loaded: {model_name}")

        if hifigan_model_name not in self.available_hifigan_models:
            hifigan_model_name = "universal"

        if hifigan_model_name not in self.available_hifigan_models:
            raise ValueError(f"HiFi-GAN model not loaded: {hifigan_model_name}")

        if pronounciation_dictionary:
            # Apply pronunciation dictionary and ensure text ends with a period
            text = self._apply_pronounciation_dictionary(text, EOS_Token=EOS_Token)
            if EOS_Token: text=text+"."
        else:
            # Ensure text ends with a semicolon and period
            if text[-1] != ";" and EOS_Token: text=text+";" + "."
        # ---
        # Load models - Get the Tacotron2 and HiFi-GAN models based on model names, plus the super-resolution model
        # ---
        model, _ = self.tacotron2_models[model_name]
        hifigan, h, denoiser = self.hifigan_models[hifigan_model_name]
        hifigan_sr, h2, _ = self.hifigan_models.get("Superres_Twilight_33000", (None, None, None))
        # ---

        # ---
        # Inference - With inference_mode for efficiency
        # ---
        with torch.inference_mode():
            # Text to sequence conversion
            sequence = np.array(text_to_sequence(text, ['english_cleaners']))[None, :]

            # Move to device
            sequence = torch.as_tensor(sequence, device=self.device).long()

            mel_outputs_postnet = model.inference_mel_only(sequence)

            # HiFi-GAN Vocoder - Convert mel spectrogram to audio waveform
            y_g_hat = hifigan(mel_outputs_postnet.float())
            audio = y_g_hat.squeeze()
            audio = audio * MAX_WAV_VALUE
            audio_denoised = denoiser(audio.view(1, -1), strength=35)[:, 0]
            audio_denoised = audio_denoised.cpu().numpy().reshape(-1)

            # Resampling and Super-Resolution - Upsample and apply super-resolution
            normalize = (MAX_WAV_VALUE / np.max(np.abs(audio_denoised))) ** 0.9
            audio_denoised = audio_denoised * normalize
            wave = soxr.resample(
                audio_denoised.astype(np.float32, copy=False),
                h.sampling_rate,
                h2.sampling_rate,
                quality="HQ"
            ).astype(np.float32, copy=False)

            wave_out = wave.astype(np.int16)

            # Super-Resolution HiFi-GAN
            wave = wave / MAX_WAV_VALUE
            wave = torch.FloatTensor(wave).to(self.device)
            new_mel = mel_spectrogram(wave.unsqueeze(0), h2.n_fft, h2.num_mels,
                                    h2.sampling_rate, h2.hop_size, h2.win_size,
                                    h2.fmin, h2.fmax)
            y_g_hat2 = hifigan_sr(new_mel)
            audio2 = y_g_hat2.squeeze()
            audio2 = audio2 * MAX_WAV_VALUE
            audio2_denoised = denoiser(audio2.view(1, -1), strength=35)[:, 0]
            audio2_denoised = audio2_denoised.cpu().numpy().reshape(-1)

            # High-Frequency Filtering and Mixing
            y = lfilter(
                self.high_pass_filter.astype(np.float32, copy=False),
                np.array([1.0], dtype=np.float32),
                audio2_denoised.astype(np.float32, copy=False)
            )
            y *= self.super_res  # superres_strength
            y_out = y.astype(np.int16)
            y_padded = np.zeros(wave_out.shape)
            y_padded[: y_out.shape[0]] = y_out
            sr_mix = wave_out + y_padded
            sr_mix = sr_mix / normalize

            # Save output audio file
            write(output_file, h2.sampling_rate, sr_mix.astype(np.int16))
            print(f"Audio saved to {output_file}")