update models

Models/config.yaml

@@ -0,0 +1,80 @@
+log_dir: ./Models/Finetune
+save_freq: 1
+log_interval: 10
+device: cuda
+epochs: 50 
+batch_size: 2
+max_len: 310 # maximum number of frames
+pretrained_model: ./Models/Finetune/base_model.pth
+load_only_params: false # set to true if do not want to load epoch numbers and optimizer parameters
+debug: true
+
+data_params:
+  train_data: "../../Data_Speech/viVoice/train.txt" 
+  val_data: "../../Data_Speech/combine/combine_val.txt"
+  root_path: "../../Data_Speech/"
+
+symbol: #Total 189 symbols
+  pad: "$"
+  punctuation: ';:,.!?¡¿—…"«»“” '
+  letters: "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
+  letters_ipa: "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"
+  extend: "∫̆ăη͡123456" #ADD MORE SYMBOLS HERE
+
+preprocess_params:
+  sr: 24000
+  spect_params:
+    n_fft: 2048
+    win_length: 1200
+    hop_length: 300
+
+training_strats:
+  #All modules: 'decoder', 'predictor', 'text_encoder', 'style_encoder', 'text_aligner', 'pitch_extractor', 'mpd', 'msd'
+  freeze_modules: [''] # Not updated when training.
+  ignore_modules: [''] # Not loading => fresh start. IMPORTANT: 'text_aligner' and 'pitch_extractor' are util pretraineds DO NOT ignore them.
+
+model_params:
+  dim_in: 64 
+  hidden_dim: 512
+  max_conv_dim: 512
+  n_layer: 3
+  n_mels: 80
+  max_dur: 50 # maximum duration of a single phoneme
+  style_dim: 128 # style vector size
+  
+  dropout: 0.2
+
+  ASR_params:
+    input_dim: 80
+    hidden_dim: 256
+    n_layers: 6
+    token_embedding_dim: 512
+
+  JDC_params:
+    num_class: 1
+    seq_len: 192
+
+  # config for decoder
+  decoder: 
+      type: 'hifigan' # either hifigan or istftnet
+      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]
+  
+loss_params:
+    lambda_mel: 5. # mel reconstruction loss
+    lambda_gen: 1. # generator loss
+    
+    lambda_mono: 1. # monotonic alignment loss (TMA)
+    lambda_s2s: 1. # sequence-to-sequence loss (TMA)
+
+    lambda_F0: 1. # F0 reconstruction loss
+    lambda_norm: 1. # norm reconstruction loss
+    lambda_dur: 1. # duration loss
+    lambda_ce: 20. # duration predictor probability output CE loss
+
+optimizer_params:
+  lr: 0.0001 # general learning rate
+  ft_lr: 0.00001 # learning rate for acoustic modules

inference.py

@@ -0,0 +1,356 @@
+import re
+import yaml
+from munch import Munch
+import numpy as np
+import librosa
+import noisereduce as nr
+from meldataset import TextCleaner
+import torch
+import torchaudio
+from nltk.tokenize import word_tokenize
+import nltk
+nltk.download('punkt_tab')
+
+from models import ProsodyPredictor, TextEncoder, StyleEncoder
+from Modules.hifigan import Decoder
+
+import sys
+import phonemizer
+if sys.platform.startswith("win"):
+    try:
+        from phonemizer.backend.espeak.wrapper import EspeakWrapper
+        import espeakng_loader
+        EspeakWrapper.set_library(espeakng_loader.get_library_path())
+    except Exception as e:
+        print(e)
+
+def espeak_phn(text, lang):
+    try:
+        my_phonemizer = phonemizer.backend.EspeakBackend(language=lang, preserve_punctuation=True,  with_stress=True, language_switch='remove-flags')
+        return my_phonemizer.phonemize([text])[0]
+    except Exception as e:
+        print(e)
+
+class Preprocess:
+    def __text_normalize(self, text):
+        punctuation = ["，", "、", "،", ";", "(", "．", "。", "…", "!", "–", ":", "?"]
+        map_to = "."
+        punctuation_pattern = re.compile(f"[{''.join(re.escape(p) for p in punctuation)}]")
+        #replace punctuation that acts like a comma or period
+        text = punctuation_pattern.sub(map_to, text)
+        #replace consecutive whitespace chars with a single space and strip leading/trailing spaces
+        text = re.sub(r'\s+', ' ', text).strip()
+        return text
+    def __merge_fragments(self, texts, n):
+        merged = []
+        i = 0
+        while i < len(texts):
+            fragment = texts[i]
+            j = i + 1
+            while len(fragment.split()) < n and j < len(texts):
+                fragment += ", " + texts[j]
+                j += 1
+            merged.append(fragment)
+            i = j
+        if len(merged[-1].split()) < n and len(merged) > 1: #handle last sentence
+            merged[-2] = merged[-2] + ", " + merged[-1]
+            del merged[-1]
+        else:
+            merged[-1] = merged[-1]
+        return merged
+    def wave_preprocess(self, wave):
+        to_mel = torchaudio.transforms.MelSpectrogram(n_mels=80, n_fft=2048, win_length=1200, hop_length=300)
+        mean, std = -4, 4
+        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 text_preprocess(self, text, n_merge=12):
+        text_norm = self.__text_normalize(text).split(".")#split by sentences.
+        text_norm = [s.strip() for s in text_norm]
+        text_norm = list(filter(lambda x: x != '', text_norm)) #filter empty index
+        text_norm = self.__merge_fragments(text_norm, n=n_merge) #merge if a sentence has less that n 
+        return text_norm
+    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
+
+#For inference only
+class StyleTTS2(torch.nn.Module):
+    def __init__(self, config_path, models_path):
+        super().__init__()
+        self.register_buffer("get_device", torch.empty(0))
+        self.preprocess = Preprocess()
+        self.ref_s = None
+        config = yaml.safe_load(open(config_path, "r", encoding="utf-8"))
+        
+        try:
+            symbols = (
+                            list(config['symbol']['pad']) +
+                            list(config['symbol']['punctuation']) +
+                            list(config['symbol']['letters']) +
+                            list(config['symbol']['letters_ipa']) +
+                            list(config['symbol']['extend'])
+                        )
+            symbol_dict = {}
+            for i in range(len((symbols))):
+                symbol_dict[symbols[i]] = i
+            
+            n_token = len(symbol_dict) + 1
+            print("\nFound:", n_token, "symbols")
+        except Exception as e:
+            print(f"\nERROR: Cannot find {e} in config file!\nYour config file is likely outdated, please download updated version from the repository.")
+            raise SystemExit(1)
+        
+        args = self.__recursive_munch(config['model_params'])
+        args['n_token'] = n_token
+        
+        self.cleaner = TextCleaner(symbol_dict, debug=False)
+
+        assert args.decoder.type in ['hifigan'], 'Decoder type unknown'
+
+        self.decoder            = Decoder(dim_in=args.hidden_dim, style_dim=args.style_dim, dim_out=args.n_mels,
+                                        resblock_kernel_sizes = args.decoder.resblock_kernel_sizes,
+                                        upsample_rates = args.decoder.upsample_rates,
+                                        upsample_initial_channel=args.decoder.upsample_initial_channel,
+                                        resblock_dilation_sizes=args.decoder.resblock_dilation_sizes,
+                                        upsample_kernel_sizes=args.decoder.upsample_kernel_sizes)
+        self.predictor           = ProsodyPredictor(style_dim=args.style_dim, d_hid=args.hidden_dim, nlayers=args.n_layer, max_dur=args.max_dur, dropout=args.dropout)
+        self.text_encoder        = TextEncoder(channels=args.hidden_dim, kernel_size=5, depth=args.n_layer, n_symbols=args.n_token)
+        self.style_encoder       = StyleEncoder(dim_in=args.dim_in, style_dim=args.style_dim, max_conv_dim=args.hidden_dim)# acoustic style encoder
+
+        self.__load_models(models_path)
+    
+    def __recursive_munch(self, d):
+        if isinstance(d, dict):
+            return Munch((k, self.__recursive_munch(v)) for k, v in d.items())
+        elif isinstance(d, list):
+            return [self.__recursive_munch(v) for v in d]
+        else:
+            return d
+        
+    def __init_replacement_func(self, replacements):
+        replacement_iter = iter(replacements)
+        def replacement(match):
+            return next(replacement_iter)
+        return replacement
+    
+    def __replace_outliers_zscore(self, tensor, threshold=3.0, factor=0.95):
+        mean = tensor.mean()
+        std = tensor.std()
+        z = (tensor - mean) / std
+
+        # Identify outliers
+        outlier_mask = torch.abs(z) > threshold
+        # Compute replacement value, respecting sign
+        sign = torch.sign(tensor - mean)
+        replacement = mean + sign * (threshold * std * factor)
+
+        result = tensor.clone()
+        result[outlier_mask] = replacement[outlier_mask]
+
+        return result
+        
+    def __load_models(self, models_path):
+        module_params = []
+        model = {'decoder':self.decoder, 'predictor':self.predictor, 'text_encoder':self.text_encoder, 'style_encoder':self.style_encoder}
+
+        params_whole = torch.load(models_path, map_location='cpu')
+        params = params_whole['net']
+        params = {key: value for key, value in params.items() if key in model.keys()}
+
+        for key in model:
+            try:
+                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
+                model[key].load_state_dict(new_state_dict, strict=False)
+
+            total_params = sum(p.numel() for p in model[key].parameters())
+            print(key,":",total_params)
+            module_params.append(total_params)
+
+        print('\nTotal',":",sum(module_params))
+
+    def __compute_style(self, path, denoise, split_dur):
+        device = self.get_device.device
+        denoise = min(denoise, 1)
+        if split_dur != 0: split_dur = max(int(split_dur), 1)
+        max_samples = 24000*20 #max 20 seconds ref audio
+        print("Computing the style for:", path)
+        
+        wave, sr = librosa.load(path, sr=24000)
+        audio, index = librosa.effects.trim(wave, top_db=30)
+        if sr != 24000:
+            audio = librosa.resample(audio, sr, 24000)
+        if len(audio) > max_samples:
+            audio = audio[:max_samples]
+        
+        if denoise > 0.0:
+            audio_denoise = nr.reduce_noise(y=audio, sr=sr, n_fft=2048, win_length=1200, hop_length=300)
+            audio = audio*(1-denoise) + audio_denoise*denoise
+
+        with torch.no_grad():
+            if split_dur>0 and len(audio)/sr>=4: #Only effective if audio length is >= 4s
+                #This option will split the ref audio to multiple parts, calculate styles and average them
+                count = 0
+                ref_s = None
+                jump = sr*split_dur
+                total_len = len(audio)
+                
+                #Need to init before the loop
+                mel_tensor = self.preprocess.wave_preprocess(audio[0:jump]).to(device)
+                ref_s = self.style_encoder(mel_tensor.unsqueeze(1))
+                count += 1
+                for i in range(jump, total_len, jump):
+                    if i+jump >= total_len:
+                        left_dur = (total_len-i)/sr
+                        if left_dur >= 1: #Still count if left over dur is >= 1s
+                            mel_tensor = self.preprocess.wave_preprocess(audio[i:total_len]).to(device)
+                            ref_s += self.style_encoder(mel_tensor.unsqueeze(1))
+                            count += 1
+                        continue
+                    mel_tensor = self.preprocess.wave_preprocess(audio[i:i+jump]).to(device)
+                    ref_s += self.style_encoder(mel_tensor.unsqueeze(1))
+                    count += 1
+                ref_s /= count
+            else:
+                mel_tensor = self.preprocess.wave_preprocess(audio).to(device)
+                ref_s = self.style_encoder(mel_tensor.unsqueeze(1))
+
+        return ref_s
+        
+    def __inference(self, phonem, ref_s, speed=1, prev_d_mean=0, t=0.1):
+        device = self.get_device.device
+        speed = min(max(speed, 0.0001), 2) #speed range [0, 2]
+        
+        phonem = ' '.join(word_tokenize(phonem))
+        tokens = self.cleaner(phonem)
+        tokens.insert(0, 0)
+        tokens.append(0)
+        tokens = torch.LongTensor(tokens).to(device).unsqueeze(0)
+        
+        with torch.no_grad():
+            input_lengths = torch.LongTensor([tokens.shape[-1]]).to(device)
+            text_mask = self.preprocess.length_to_mask(input_lengths).to(device)
+
+            # encode
+            t_en = self.text_encoder(tokens, input_lengths, text_mask)
+            s = ref_s.to(device)
+        
+            # cal alignment
+            d = self.predictor.text_encoder(t_en, s, input_lengths, text_mask)
+            x, _ = self.predictor.lstm(d)
+            duration = self.predictor.duration_proj(x) 
+            duration = torch.sigmoid(duration).sum(axis=-1)
+
+            if prev_d_mean != 0:#Stabilize speaking speed between splits
+                dur_stats = torch.empty(duration.shape).normal_(mean=prev_d_mean, std=duration.std()).to(device)
+            else:
+                dur_stats = torch.empty(duration.shape).normal_(mean=duration.mean(), std=duration.std()).to(device)
+            duration = duration*(1-t) + dur_stats*t
+            duration[:,1:-2] = self.__replace_outliers_zscore(duration[:,1:-2]) #Normalize outlier
+
+            duration /= speed
+                
+            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)
+            alignment = pred_aln_trg.unsqueeze(0).to(device)
+
+            # encode prosody
+            en = (d.transpose(-1, -2) @ alignment)
+            F0_pred, N_pred = self.predictor.F0Ntrain(en, s)
+            asr = (t_en @ pred_aln_trg.unsqueeze(0).to(device))
+
+            out = self.decoder(asr, F0_pred, N_pred, s)
+        
+        return out.squeeze().cpu().numpy(), duration.mean()
+    
+    def get_styles(self, speakers, denoise=0.3, avg_style=True):
+        if avg_style:   split_dur = 2
+        else:           split_dur = 0
+        styles = {}
+        for id in speakers:
+            ref_s = self.__compute_style(speakers[id]['path'], denoise=denoise, split_dur=split_dur)
+            styles[id] = {
+                'style': ref_s,
+                'path': speakers[id]['path'],
+                'lang': speakers[id]['lang'],
+                'speed': speakers[id]['speed'],
+            }
+        return styles
+
+    def generate(self, text, styles, stabilize=True, n_merge=16, default_speaker= "[id_1]"):
+        if stabilize:   smooth_value=0.2
+        else:           smooth_value=0    
+        
+        list_wav        = []
+        prev_d_mean     = 0
+        lang_pattern    = r'\[([^\]]+)\]\{([^}]+)\}'
+
+        text = re.sub(r'[\n\r\t\f\v]', '', text)
+        #fix lang tokens span to multiple sents
+        find_lang_tokens = re.findall(lang_pattern, text)
+        if find_lang_tokens:
+            cus_text = []
+            for lang, t in find_lang_tokens:
+                parts = self.preprocess.text_preprocess(t, n_merge=0)
+                parts = ".".join([f"[{lang}]" + f"{{{p}}}"for p in parts])
+                cus_text.append(parts)
+            replacement_func = self.__init_replacement_func(cus_text)
+            text = re.sub(lang_pattern, replacement_func, text)
+
+        texts = re.split(r'(\[id_\d+\])', text) #split the text by speaker ids while keeping the ids.
+        if len(texts) <= 1 or bool(re.match(r'(\[id_\d+\])', texts[0]) == False): #Add a default speaker
+            texts.insert(0, default_speaker)
+        curr_id = None
+        for i in range(len(texts)): #remove consecutive ids
+            if bool(re.match(r'(\[id_\d+\])', texts[i])):
+                if texts[i]!=curr_id:
+                    curr_id = texts[i]
+                else:
+                    texts[i] = ''
+        del curr_id
+        texts = list(filter(lambda x: x != '', texts))
+
+        print("Generating Audio...")
+        for i in texts:
+            if bool(re.match(r'(\[id_\d+\])', i)):
+                #Set up env for matched speaker
+                speaker_id = i.strip('[]')
+                current_ref_s = styles[speaker_id]['style']
+                speed = styles[speaker_id]['speed']
+                continue
+            text_norm = self.preprocess.text_preprocess(i, n_merge=n_merge)
+            for sentence in text_norm:
+                cus_phonem = []
+                find_lang_tokens = re.findall(lang_pattern, sentence)
+                if find_lang_tokens:
+                    for lang, t in find_lang_tokens:
+                        try:
+                            phonem = espeak_phn(t, lang)
+                            cus_phonem.append(phonem)
+                        except Exception as e:
+                            print(e)
+                        
+                replacement_func = self.__init_replacement_func(cus_phonem)
+                phonem =  espeak_phn(sentence, styles[speaker_id]['lang'])
+                phonem = re.sub(lang_pattern, replacement_func, phonem)
+
+                wav, prev_d_mean = self.__inference(phonem, current_ref_s, speed=speed, prev_d_mean=prev_d_mean, t=smooth_value)
+                wav = wav[4000:-4000] #Remove weird pulse and silent tokens
+                list_wav.append(wav)
+        
+        final_wav = np.concatenate(list_wav)
+        final_wav = np.concatenate([np.zeros([4000]), final_wav, np.zeros([4000])], axis=0) # add padding
+        return final_wav