Update inference.py and meldataset,py

inference.py

@@ -65,9 +65,31 @@ class StyleTTS2(torch.nn.Module):
         super().__init__()
         self.register_buffer("get_device", torch.empty(0))
         self.preprocess = Preprocess()
-
-        config = yaml.safe_load(open(config_path))
+        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'
 
@@ -186,7 +208,7 @@ class StyleTTS2(torch.nn.Module):
         speed = min(max(speed, 0.0001), 2) #speed range [0, 2]
         
         phonem = ' '.join(word_tokenize(phonem))
-        tokens = TextCleaner()(phonem)
+        tokens = self.cleaner(phonem)
         tokens.insert(0, 0)
         tokens.append(0)
         tokens = torch.LongTensor(tokens).to(device).unsqueeze(0)
@@ -231,17 +253,34 @@ class StyleTTS2(torch.nn.Module):
         
         return out.squeeze().cpu().numpy(), duration.mean()
     
-    def get_styles(self, speaker, denoise=0.3, avg_style=True):
-        if avg_style:   split_dur = 3
-        else:           split_dur = 0
-        style = {}
-        ref_s = self.__compute_style(speaker['path'], denoise=denoise, split_dur=split_dur)
+    def get_styles(self, speaker, denoise=0.3, avg_style=True, load_styles=False):
+        if not load_styles:
+            if avg_style:   split_dur = 3
+            else:           split_dur = 0
+            self.ref_s = self.__compute_style(speaker['path'], denoise=denoise, split_dur=split_dur)
+        else:
+            if self.ref_s is None:
+                raise Exception("Have to compute or load the styles first!")
         style = {
-            'style': ref_s,
+            'style': self.ref_s,
             'path': speaker['path'],
             'speed': speaker['speed'],
         }
         return style
+    
+    def save_styles(self, save_dir):
+        if self.ref_s is not None:
+            torch.save(self.ref_s, save_dir)
+            print("Saved styles!")
+        else:
+            raise Exception("Have to compute the styles before saving it.")
+
+    def load_styles(self, save_dir):
+        try:
+            self.ref_s = torch.load(save_dir)
+            print("Loaded styles!")
+        except Exception as e:
+            print(e)
 
     def generate(self, phonem, style, stabilize=True, n_merge=16):
         if stabilize:   smooth_value=0.2

meldataset.py

@@ -1,7 +1,5 @@
 #coding: utf-8
-import os
 import os.path as osp
-import time
 import random
 import numpy as np
 import random
@@ -9,10 +7,10 @@ import soundfile as sf
 import librosa
 
 import torch
-from torch import nn
-import torch.nn.functional as F
 import torchaudio
-from torch.utils.data import DataLoader
+import torch.utils.data
+import torch.distributed as dist
+from multiprocessing import Pool
 
 import logging
 logger = logging.getLogger(__name__)
@@ -20,33 +18,19 @@ logger.setLevel(logging.DEBUG)
 
 import pandas as pd
 
-##########################################################
-_pad = "$"
-_punctuation = ';:,.!?¡¿—…"«»“” '
-_letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
-_letters_ipa = "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"
-_extend = "" #ADD MORE SYMBOLS HERE
-
-# Export all symbols:
-symbols = [_pad] + list(_punctuation) + list(_letters) + list(_letters_ipa) + list(_extend)
-
-dicts = {}
-for i in range(len((symbols))):
-    dicts[symbols[i]] = i
-
-# Copy this code somewhere else then run with print(len(dicts) + 1) to check total symbols
-##########################################################
-
 class TextCleaner:
-    def __init__(self, dummy=None):
-        self.word_index_dictionary = dicts
+    def __init__(self, symbol_dict, debug=True):
+        self.word_index_dictionary = symbol_dict
+        self.debug = debug
     def __call__(self, text):
         indexes = []
         for char in text:
             try:
                 indexes.append(self.word_index_dictionary[char])
             except KeyError as e:
-                #print(char)
+                if self.debug:
+                    print("\nWARNING UNKNOWN IPA CHARACTERS/LETTERS: ", char)
+                    print("To ignore set 'debug' to false in the config")
                 continue
         return indexes
 
@@ -75,17 +59,16 @@ class FilePathDataset(torch.utils.data.Dataset):
     def __init__(self,
                  data_list,
                  root_path,
+                 symbol_dict,
                  sr=24000,
                  data_augmentation=False,
-                 validation=False
+                 validation=False,
+                 debug=True
                  ):
 
-        spect_params = SPECT_PARAMS
-        mel_params = MEL_PARAMS
-
         _data_list = [l.strip().split('|') for l in data_list]
         self.data_list = _data_list #[data if len(data) == 3 else (*data, 0) for data in _data_list] #append speakerid=0 for all
-        self.text_cleaner = TextCleaner()
+        self.text_cleaner = TextCleaner(symbol_dict, debug)
         self.sr = sr
 
         self.df = pd.DataFrame(self.data_list)
@@ -195,9 +178,13 @@ class Collater(object):
         return waves, texts, input_lengths, mels, output_lengths
 
 
+def get_length(wave_path, root_path):
+    info = sf.info(osp.join(root_path, wave_path))
+    return info.frames * (24000 / info.samplerate)
 
 def build_dataloader(path_list,
                      root_path,
+                     symbol_dict,
                      validation=False,
                      batch_size=4,
                      num_workers=1,
@@ -205,14 +192,116 @@ def build_dataloader(path_list,
                      collate_config={},
                      dataset_config={}):
     
-    dataset = FilePathDataset(path_list, root_path, validation=validation, **dataset_config)
+    dataset = FilePathDataset(path_list, root_path, symbol_dict, validation=validation, **dataset_config)
     collate_fn = Collater(**collate_config)
-    data_loader = DataLoader(dataset,
-                             batch_size=batch_size,
-                             shuffle=(not validation),
-                             num_workers=num_workers,
-                             drop_last=(not validation),
-                             collate_fn=collate_fn,
-                             pin_memory=(device != 'cpu'))
-
-    return data_loader
+    
+    print("Getting sample lengths...")
+    
+    num_processes = num_workers * 2
+    if num_processes != 0:
+        list_of_tuples = [(d[0], root_path) for d in dataset.data_list]
+        with Pool(processes=num_processes) as pool:
+            sample_lengths = pool.starmap(get_length, list_of_tuples, chunksize=16)
+    else:
+        sample_lengths = []
+        for d in dataset.data_list:
+            sample_lengths.append(get_length(d[0], root_path))
+
+    data_loader = torch.utils.data.DataLoader(
+        dataset,
+        num_workers=num_workers,
+        batch_sampler=BatchSampler(
+            sample_lengths,
+            batch_size,
+            shuffle=(not validation),
+            drop_last=(not validation),
+            num_replicas=1,
+            rank=0,
+        ),
+        collate_fn=collate_fn,
+        pin_memory=(device != "cpu"),
+    )
+
+    return data_loader
+
+#https://github.com/duerig/StyleTTS2/
+class BatchSampler(torch.utils.data.Sampler):
+    def __init__(
+        self,
+        sample_lengths,
+        batch_sizes,
+        num_replicas=None,
+        rank=None,
+        shuffle=True,
+        drop_last=False,
+    ):
+        self.batch_sizes = batch_sizes
+        if num_replicas is None:
+            self.num_replicas = dist.get_world_size()
+        else:
+            self.num_replicas = num_replicas
+        if rank is None:
+            self.rank = dist.get_rank()
+        else:
+            self.rank = rank
+        self.shuffle = shuffle
+        self.drop_last = drop_last
+
+        self.time_bins = {}
+        self.epoch = 0
+        self.total_len = 0
+        self.last_bin = None
+
+        for i in range(len(sample_lengths)):
+            bin_num = self.get_time_bin(sample_lengths[i])
+            if bin_num != -1:
+                if bin_num not in self.time_bins:
+                    self.time_bins[bin_num] = []
+                self.time_bins[bin_num].append(i)
+
+        for key in self.time_bins.keys():
+            val = self.time_bins[key]
+            total_batch = self.batch_sizes * num_replicas
+            self.total_len += len(val) // total_batch
+            if not self.drop_last and len(val) % total_batch != 0:
+                self.total_len += 1
+
+    def __iter__(self):
+        sampler_order = list(self.time_bins.keys())
+        sampler_indices = []
+
+        if self.shuffle:
+            sampler_indices = torch.randperm(len(sampler_order)).tolist()
+        else:
+            sampler_indices = list(range(len(sampler_order)))
+
+        for index in sampler_indices:
+            key = sampler_order[index]
+            current_bin = self.time_bins[key]
+            dist = torch.utils.data.distributed.DistributedSampler(
+                current_bin,
+                num_replicas=self.num_replicas,
+                rank=self.rank,
+                shuffle=self.shuffle,
+                drop_last=self.drop_last,
+            )
+            dist.set_epoch(self.epoch)
+            sampler = torch.utils.data.sampler.BatchSampler(
+                dist, self.batch_sizes, self.drop_last
+            )
+            for item_list in sampler:
+                self.last_bin = key
+                yield [current_bin[i] for i in item_list]
+
+    def __len__(self):
+        return self.total_len
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch
+
+    def get_time_bin(self, sample_count):
+        result = -1
+        frames = sample_count // 300
+        if frames >= 20:
+            result = (frames - 20) // 20
+        return result