msinference.py

from Modules.vits.models import VitsModel, VitsTokenizer
import sys
import tempfile
import re
import os
from collections import OrderedDict
from Modules.hifigan import Decoder
from Utils.PLBERT.util import load_plbert
import phonemizer
import torch
from cached_path import cached_path
import nltk
import audresample
nltk.download('punkt', download_dir='./')  # comment if downloaded once
nltk.download('punkt_tab', download_dir='./')
nltk.data.path.append('.')
import numpy as np
import yaml
import librosa
from models import ProsodyPredictor, TextEncoder, StyleEncoder, MelSpec
from nltk.tokenize import word_tokenize
from Utils.text_utils import transliterate_number
import textwrap

device = 'cpu'
if torch.cuda.is_available():
    device = 'cuda'

_pad = "$"
_punctuation = ';:,.!?¡¿—…"«»“” '
_letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
_letters_ipa = "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"

# Export all symbols:
symbols = [_pad] + list(_punctuation) + list(_letters) + list(_letters_ipa)

dicts = {}
for i in range(len((symbols))):
    dicts[symbols[i]] = i





class TextCleaner:
    def __init__(self, dummy=None):
        self.word_index_dictionary = dicts
        print(len(dicts))

    def __call__(self, text):
        indexes = []
        for char in text:
            try:
                indexes.append(self.word_index_dictionary[char])
            except KeyError:
                print('CLEAN', text)
        return indexes


textclenaer = TextCleaner()

def alpha_num(f):
    f = re.sub(' +', ' ', f)              # delete spaces
    f = re.sub(r'[^A-Z a-z0-9 ]+', '', f)  # del non alpha num
    return f

mel_spec = MelSpec().to(device)

def compute_style(path):
    x, sr = librosa.load(path, sr=24000)
    x, _ = librosa.effects.trim(x, top_db=30)
    if sr != 24000:
        x = librosa.resample(x, sr, 24000)
        
    with torch.no_grad():
        x = torch.from_numpy(x[None, :]).to(device=device, dtype=torch.float)
        
        mel_tensor = (torch.log(1e-5 + mel_spec(x)) + 4) / 4
        
        #mel_tensor = preprocess(audio).to(device)

        ref_s = style_encoder(mel_tensor)
        ref_p = predictor_encoder(mel_tensor)  # [bs, 11, 1, 128]

        s = torch.cat([ref_s, ref_p], dim=3)  # [bs, 11, 1, 256]

        s = s[:, :, 0, :].transpose(1, 2)  # [1, 128, 11]
    return s  # [1, 128, 11]

global_phonemizer = phonemizer.backend.EspeakBackend(
    language='en-us', preserve_punctuation=True,  with_stress=True)
# phonemizer = Phonemizer.from_checkpoint(str(cached_path('https://public-asai-dl-models.s3.eu-central-1.amazonaws.com/DeepPhonemizer/en_us_cmudict_ipa_forward.pt')))


args = yaml.safe_load(open(str('Utils/config.yml')))
ASR_config = args['ASR_config']

bert = load_plbert(args['PLBERT_dir']).eval().to(device)

decoder = Decoder(dim_in=512,
                  style_dim=128,
                  dim_out=80,  # n_mels
                  resblock_kernel_sizes=[3, 7, 11],
                  upsample_rates=[10, 5, 3, 2],
                  upsample_initial_channel=512,
                  resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
                  upsample_kernel_sizes=[20, 10, 6, 4]).eval().to(device)

text_encoder = TextEncoder(channels=512,
                           kernel_size=5,
                           depth=3,  # args['model_params']['n_layer'],
                           n_symbols=178,  # args['model_params']['n_token']
                           ).eval().to(device)

predictor = ProsodyPredictor(style_dim=128,
                             d_hid=512,
                             nlayers=3,  # OFFICIAL config.nlayers=5;
                             max_dur=50).eval().to(device)

style_encoder = StyleEncoder(dim_in=64,
                             style_dim=128,
                             max_conv_dim=512).eval().to(device)  # acoustic style encoder
predictor_encoder = StyleEncoder(dim_in=64,
                                 style_dim=128,
                                 max_conv_dim=512).eval().to(device)  # prosodic style encoder
bert_encoder = torch.nn.Linear(bert.config.hidden_size, 512).eval().to(device)

# params_whole = torch.load('freevc2/yl4579_styletts2.pth' map_location='cpu')
params_whole = torch.load(str(cached_path(
    "hf://yl4579/StyleTTS2-LibriTTS/Models/LibriTTS/epochs_2nd_00020.pth")), map_location='cpu', weights_only=True)
params = params_whole['net']
#params['decoder'].pop('module.generator.m_source.l_linear.weight')
#params['decoder'].pop('module.generator.m_source.l_linear.bias')  # SourceHNSf

def _del_prefix(d):
    # del ".module"
    out = OrderedDict()
    for k, v in d.items():
        out[k[7:]] = v
    return out


bert.load_state_dict(_del_prefix(params['bert']), strict=True)
bert_encoder.load_state_dict(_del_prefix(params['bert_encoder']), strict=True)
# XTRA non-ckpt LSTMs nlayers add slowiness to voice
predictor.load_state_dict(_del_prefix(params['predictor']), strict=True)
decoder.load_state_dict(_del_prefix(params['decoder']), strict=True)
text_encoder.load_state_dict(_del_prefix(params['text_encoder']), strict=True)
predictor_encoder.load_state_dict(_del_prefix(
    params['predictor_encoder']), strict=True)
style_encoder.load_state_dict(_del_prefix(
    params['style_encoder']), strict=True)

def inference(text,
              ref_s):
    # text = transliterate_number(text, lang='en').strip()  # Transliteration only used for foreign()  # perhaps add xtra . after ? ;
    ps = global_phonemizer.phonemize([text])
    ps = word_tokenize(ps[0])
    ps = ' '.join(ps)
    tokens = textclenaer(ps)
    tokens.insert(0, 0)
    tokens = torch.LongTensor(tokens).to(device).unsqueeze(0)
    with torch.no_grad():
        hidden_states = text_encoder(tokens)
        bert_dur = bert(tokens, attention_mask=torch.ones_like(tokens))
        d_en = bert_encoder(bert_dur).transpose(-1, -2)

        aln_trg, F0_pred, N_pred = predictor(d_en=d_en, s=ref_s[:, 128:, :])

        asr = torch.bmm(aln_trg, hidden_states)
        asr = asr.transpose(1, 2)
        asr = torch.cat([asr[:, :, 0:1], asr[:, :, 0:-1]], 2)
        x = decoder(asr=asr,              # [1, 512, 201]
                    F0_curve=F0_pred,     # [1, 1, 402] 2x time
                    N=N_pred,             # [1, 1, 402] 2x time
                    s=ref_s[:, :128, :])  # [1, 256, 1]

    x = x.cpu().numpy()[0, 0, :]
    x[-400:] = 0  # noisy pulse produced for unterminated sentences, in absence of punctuation, (not sure if same behaviour for all voices)

    # StyleTTS2 is 24kHz -> Resample to 16kHz as is AudioGen / MMS

    if x.shape[0] > 10:

        x = audresample.resample(signal=x.astype(np.float32),
                                 original_rate=24000,
                                 target_rate=16000)[0, :]  # audresample reshapes (64,) -> (1,64) | Volume Normalisation applies in api.py:tts_multi_sentence()

    else:
        print('\n\n\n\n\nEMPTY TTS\n\n\n\n\n\nn', x.shape)
        x = np.zeros(0)
    return x


# ___________________________________________________________

# https://huggingface.co/spaces/mms-meta/MMS/blob/main/tts.py
# ___________________________________________________________
# -*- coding: utf-8 -*-
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

TTS_LANGUAGES = {}
# with open('_d.csv', 'w') as f2:
with open(f"Utils/all_langs.csv") as f:
    for line in f:
        iso, name = line.split(",", 1)
        TTS_LANGUAGES[iso.strip()] = name.strip()
        # f2.write(iso + ',' + name.replace("a S","")+'\n')


# ==============================================================================================
# LOAD hun / ron / serbian - rmc-script_latin / cyrillic-Carpathian (not Vlax)
# ==============================================================================================

PHONEME_MAP = {
    'služ': 'sloooozz',  # 'službeno'
    'suver': 'siuveeerra',  # 'suverena'
    'država': 'dirrezav',  # 'država'
    'iči': 'ici',  # 'Graniči'
    's ': 'se',  # a s with space
    'q': 'ku',
    'w': 'aou',
    'z': 's',
    "š": "s",
    'th': 'ta',
    'v': 'vv',
    # "ć": "č",
    # "đ": "ď",
    # "lj": "ľ",
    # "nj": "ň",
    "ž": "z",
    # "c": "č"
}


def fix_phones(text):
    for src, target in PHONEME_MAP.items():
        text = text.replace(src, target)
    # text = re.sub(r'\s+', '` `', text) #.strip() #.lower()
    # text = re.sub(r'\s+', '_     _', text)  # almost proper pausing

    return text.replace(',', '_     _').replace('.', '_    _')


def has_cyrillic(text):
    # https://stackoverflow.com/questions/48255244/python-check-if-a-string-contains-cyrillic-characters
    return bool(re.search('[\u0400-\u04FF]', text))


def foreign(text=None,   # split sentences here so we can prepend a txt for german to each sentence to
                         # fall on the male voice (Sink attn)
            lang='romanian',
            speed=None):

    # https://huggingface.co/dkounadis/artificial-styletts2/blob/main/Utils/all_langs.csv
    lang = lang.lower()

    # https://huggingface.co/spaces/mms-meta/MMS

    if 'hun' in lang:

        lang_code = 'hun'

    elif any([i in lang for i in ['ser', 'bosn', 'herzegov', 'montenegr', 'macedon']]):

        if has_cyrillic(text):  # check 0-th sentence if is cyrillic

            # romani carpathian (also has latin / cyrillic Vlax)
            lang_code = 'rmc-script_cyrillic'

        else:

            # romani carpathian (has also Vlax)
            lang_code = 'rmc-script_latin'

    elif 'rom' in lang:

        lang_code = 'ron'

    elif 'ger' in lang or 'deu' in lang or 'allem' in lang:

        lang_code = 'deu'

    elif 'alban' in lang:

        lang_code = 'sqi'

    else:

        lang_code = lang.split()[0].strip()

    # load VITS

    # net_g = VitsModel.from_pretrained(f'facebook/mms-tts-{lang_code}').eval().to(device)
    # tokenizer = VitsTokenizer.from_pretrained(f'facebook/mms-tts-{lang_code}')
    global cached_lang_code, cached_net_g, cached_tokenizer

    if 'cached_lang_code' not in globals() or cached_lang_code != lang_code:
        cached_lang_code = lang_code
        cached_net_g = VitsModel.from_pretrained(f'facebook/mms-tts-{lang_code}').eval().to(device)
        cached_tokenizer = VitsTokenizer.from_pretrained(f'facebook/mms-tts-{lang_code}')

    net_g = cached_net_g
    tokenizer = cached_tokenizer


    total_audio = []

    # Split long sentences if deu to control voice switch - for other languages let text no-split
    if not isinstance(text, list):
            # Split Very long sentences
            text = [sub_sent+' ' for sub_sent in textwrap.wrap(text, 440, break_long_words=0)]

    for _t in text:

        _t = _t.lower()

        # NUMBERS

        try:
          _t = transliterate_number(_t, lang=lang_code)
        except NotImplementedError:
          print('Transliterate Numbers - NotImplemented for {lang_code=}', _t,'\n____________________________________________')

        # PRONOUNC.

        if lang_code == 'rmc-script_latin':

            _t = fix_phones(_t)  # phonemes replace per language

        elif lang_code == 'ron':

            # tone
            _t = _t.replace("ţ", "ț"
                            ).replace('ț', 'ts').replace('î', 'u').replace('â', 'a').replace('ş', 's')

        # /data/dkounadis/.hf7/hub/models--facebook--mms-tts/snapshots/44cc7fb408064ef9ea6e7c59130d88cac1274671/models/rmc-script_latin/vocab.txt
        # input_ids / attention_mask
        inputs = tokenizer(_t, return_tensors="pt")

        with torch.no_grad():

            # MMS

            x = net_g(input_ids=inputs.input_ids.to(device),
                      attention_mask=inputs.attention_mask.to(device),
                      lang_code=lang_code,
                      )[0, :]

            # crop the 1st audio - is PREFIX text 156000 samples to chose deu voice / VitsAttention()

            total_audio.append(x)

        print(f'\n\n_______________________________ {_t} {x.shape=}')

    x = torch.cat(total_audio).cpu().numpy()

    # x /= np.abs(x).max() + 1e-7  ~ Volume normalisation @api.py:tts_multi_sentence() OR demo.py

    return x  # 16kHz - only resample  StyleTTS2 from 24Hkz -> 16kHz