inference.py

import nltk

nltk.download("punkt")
import phonemizer
import torch

torch.manual_seed(0)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
import random

random.seed(0)
import numpy as np

np.random.seed(0)
import time
import random
import yaml
from munch import Munch
from pydub import AudioSegment
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
import torchaudio
import librosa
import soundfile as sf
from nltk.tokenize import word_tokenize
from models import *
from utils import *
from text_utils import TextCleaner
import phonemizer
from Utils.PLBERT.util import load_plbert
from Modules.diffusion.sampler import DiffusionSampler, ADPM2Sampler, KarrasSchedule


textclenaer = TextCleaner()


to_mel = torchaudio.transforms.MelSpectrogram(
    n_mels=80, n_fft=2048, win_length=1200, hop_length=300
)
mean, std = -4, 4

SPEECH_SPEEDS = {"slowest": 0.4, "slow": 0.5, "medium": 0.8, "default": 1.0}
SPEECH_SPEEDS_LIST = list(SPEECH_SPEEDS.keys())
STYLE_TTS2_SAMPLE_RATE = 24000


class StyleTTS2Inference:
    def __init__(self, model_dir, language="hi", reference_audio_path=None):
        self.model_dir = model_dir
        self.config_path = f"{model_dir}/config.yml"
        self.reference_audio_path = reference_audio_path
        self.check_if_model_dir_has_pth_files_and_config()
        self.config = yaml.safe_load(open(self.config_path))
        self.language = language
        self.global_phonemizer = phonemizer.backend.EspeakBackend(
            language=self.language, preserve_punctuation=True, with_stress=True
        )
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"Using {self.device} device")
        self.load_model()
        self.compute_style()

    def check_if_model_dir_has_pth_files_and_config(self):
        # Check first if the model directory exists
        if not os.path.exists(self.model_dir):
            raise FileNotFoundError(f"Model directory {self.model_dir} not found")
        # Check for pth files
        files = [f for f in os.listdir(f"{self.model_dir}/") if f.endswith(".pth")]
        if len(files) == 0:
            raise FileNotFoundError(f"No .pth files found in {self.model_dir}")
        # Check for config file
        if not os.path.exists(self.config_path):
            raise FileNotFoundError(
                f"Config file {self.config_path} not found in {self.model_dir}"
            )
        # Check for reference audio file
        if self.reference_audio_path and not os.path.exists(self.reference_audio_path):
            raise FileNotFoundError(
                f"Reference audio file {self.reference_audio_path} not found"
            )

    def load_model(self):
        # load pretrained ASR model
        ASR_config = self.config.get("ASR_config", False)
        ASR_path = self.config.get("ASR_path", False)
        text_aligner = load_ASR_models(ASR_path, ASR_config)

        # load pretrained F0 model
        F0_path = self.config.get("F0_path", False)
        pitch_extractor = load_F0_models(F0_path)

        # load BERT model
        BERT_path = self.config.get("PLBERT_dir", False)
        plbert = load_plbert(BERT_path)

        self.model_params = recursive_munch(self.config["model_params"])
        self.model = build_model(
            self.model_params, text_aligner, pitch_extractor, plbert
        )
        _ = [self.model[key].eval() for key in self.model]
        _ = [self.model[key].to(self.device) for key in self.model]
        files = [f for f in os.listdir(f"{self.model_dir}/") if f.endswith(".pth")]
        params_whole = torch.load(f"{self.model_dir}/" + files[0], map_location="cpu")
        params = params_whole["net"]
        for key in self.model:
            if key in params:
                print("%s loaded" % key)
                try:
                    self.model[key].load_state_dict(params[key])
                except:
                    from collections import OrderedDict

                    state_dict = params[key]
                    new_state_dict = OrderedDict()
                    for k, v in state_dict.items():
                        name = k[7:]  # remove `module.`
                        new_state_dict[name] = v
                    # load params
                    self.model[key].load_state_dict(new_state_dict, strict=False)
        #             except:
        #                 _load(params[key], model[key])
        _ = [self.model[key].eval() for key in self.model]
        self.sampler = DiffusionSampler(
            self.model.diffusion.diffusion,
            sampler=ADPM2Sampler(),
            sigma_schedule=KarrasSchedule(
                sigma_min=0.0001, sigma_max=3.0, rho=9.0
            ),  # empirical parameters
            clamp=False,
        )

    def length_to_mask(self, lengths):
        mask = (
            torch.arange(lengths.max())
            .unsqueeze(0)
            .expand(lengths.shape[0], -1)
            .type_as(lengths)
        )
        mask = torch.gt(mask + 1, lengths.unsqueeze(1))
        return mask

    def preprocess(self, wave):
        wave_tensor = torch.from_numpy(wave).float()
        mel_tensor = to_mel(wave_tensor)
        mel_tensor = (torch.log(1e-5 + mel_tensor.unsqueeze(0)) - mean) / std
        return mel_tensor

    def compute_style(self):
        path = self.reference_audio_path
        wave, sr = librosa.load(path, sr=STYLE_TTS2_SAMPLE_RATE)
        audio, index = librosa.effects.trim(wave, top_db=30)
        if sr != STYLE_TTS2_SAMPLE_RATE:
            audio = librosa.resample(
                audio, orig_sr=sr, target_sr=STYLE_TTS2_SAMPLE_RATE
            )
        mel_tensor = self.preprocess(audio).to(self.device)

        with torch.no_grad():
            ref_s = self.model.style_encoder(mel_tensor.unsqueeze(1))
            ref_p = self.model.predictor_encoder(mel_tensor.unsqueeze(1))

        self.computed_style = torch.cat([ref_s, ref_p], dim=1)

    def generate_output_path(self, extension="wav"):
        return f"output_{time.time()}.{extension}"

    def inference(
        self,
        text,
        alpha=0.3,
        beta=0.7,
        diffusion_steps=5,
        embedding_scale=1,
        speech_speed="default",
        output_path=None,
        sample_rate=STYLE_TTS2_SAMPLE_RATE,
        extention="wav",
    ):
        if speech_speed not in SPEECH_SPEEDS:
            raise ValueError(
                f"Speech speed {speech_speed} not supported. Supported values are {', '.join(SPEECH_SPEEDS_LIST)}."
            )
        if extention not in ["wav", "mp3"]:
            raise ValueError(
                f"Extension {extention} not supported. Supported values are 'wav', 'mp3'."
            )
        ref_s = self.computed_style
        text = text.strip()
        ps = self.global_phonemizer.phonemize([text])
        ps = word_tokenize(ps[0])
        ps = " ".join(ps)
        tokens = textclenaer(ps)
        tokens.insert(0, 0)
        tokens = torch.LongTensor(tokens).to(self.device).unsqueeze(0)

        with torch.no_grad():
            input_lengths = torch.LongTensor([tokens.shape[-1]]).to(self.device)
            text_mask = length_to_mask(input_lengths).to(self.device)

            t_en = self.model.text_encoder(tokens, input_lengths, text_mask)
            bert_dur = self.model.bert(tokens, attention_mask=(~text_mask).int())
            d_en = self.model.bert_encoder(bert_dur).transpose(-1, -2)

            s_pred = self.sampler(
                noise=torch.randn((1, 256)).unsqueeze(1).to(self.device),
                embedding=bert_dur,
                embedding_scale=embedding_scale,
                features=ref_s,  # reference from the same speaker as the embedding
                num_steps=diffusion_steps,
            ).squeeze(1)

            s = s_pred[:, 128:]
            ref = s_pred[:, :128]

            ref = alpha * ref + (1 - alpha) * ref_s[:, :128]
            s = beta * s + (1 - beta) * ref_s[:, 128:]

            d = self.model.predictor.text_encoder(d_en, s, input_lengths, text_mask)

            x, _ = self.model.predictor.lstm(d)
            duration = (
                self.model.predictor.duration_proj(x) * SPEECH_SPEEDS[speech_speed]
            )
            duration = torch.sigmoid(duration).sum(axis=-1)
            pred_dur = torch.round(duration.squeeze()).clamp(min=1)
            pred_aln_trg = torch.zeros(input_lengths, int(pred_dur.sum().data))
            c_frame = 0
            for i in range(pred_aln_trg.size(0)):
                pred_aln_trg[i, c_frame : c_frame + int(pred_dur[i].data)] = 1
                c_frame += int(pred_dur[i].data)

            # encode prosody
            en = d.transpose(-1, -2) @ pred_aln_trg.unsqueeze(0).to(self.device)
            if self.model_params.decoder.type == "hifigan":
                asr_new = torch.zeros_like(en)
                asr_new[:, :, 0] = en[:, :, 0]
                asr_new[:, :, 1:] = en[:, :, 0:-1]
                en = asr_new

            F0_pred, N_pred = self.model.predictor.F0Ntrain(en, s)

            asr = t_en @ pred_aln_trg.unsqueeze(0).to(self.device)
            if self.model_params.decoder.type == "hifigan":
                asr_new = torch.zeros_like(asr)
                asr_new[:, :, 0] = asr[:, :, 0]
                asr_new[:, :, 1:] = asr[:, :, 0:-1]
                asr = asr_new

            out = self.model.decoder(asr, F0_pred, N_pred, ref.squeeze().unsqueeze(0))

        np_arr = (
            out.squeeze().cpu().numpy()[..., :-50]
        )  # weird pulse at the end of the model, need to be fixed later
        if not output_path:
            output_path = self.generate_output_path("wav")

        sf.write(output_path, np_arr, STYLE_TTS2_SAMPLE_RATE)

        # Handle sample rate conversion
        if sample_rate != STYLE_TTS2_SAMPLE_RATE:
            wave, sr = librosa.load(output_path, sr=STYLE_TTS2_SAMPLE_RATE)
            audio = librosa.resample(wave, orig_sr=sr, target_sr=sample_rate)
            sf.write(output_path, audio, sample_rate)

        # Handle extention conversion
        if extention == "mp3":
            sound = AudioSegment.from_wav(output_path)
            output_path_mp3 = output_path.replace("wav", "mp3")
            sound.export(output_path_mp3, format="mp3")
            # Delete the wav file
            os.remove(output_path)
            output_path = output_path_mp3

        return output_path


# Example Inference Usage
# model_dir = "path/to/model/dir"
# tts = StyleTTS2Inference(model_dir, language="hi", reference_audio_path="path/to/reference/audio.wav")
# text = "भारत एक बहुत बड़ा देश है।"
# audio_bytes = tts.inference(text)