add data_gen

data_gen/tts/base_binarizer.py

@@ -0,0 +1,324 @@
+import json
+import os
+import random
+from re import L
+import traceback
+from functools import partial
+
+import numpy as np
+from resemblyzer import VoiceEncoder
+from tqdm import tqdm
+
+import utils.commons.single_thread_env  # NOQA
+from utils.audio import librosa_wav2spec
+from utils.audio.align import get_mel2ph, mel2token_to_dur
+from utils.audio.cwt import get_lf0_cwt, get_cont_lf0
+from utils.audio.pitch.utils import f0_to_coarse
+from utils.audio.pitch_extractors import extract_pitch_simple
+from utils.commons.hparams import hparams
+from utils.commons.indexed_datasets import IndexedDatasetBuilder
+from utils.commons.multiprocess_utils import multiprocess_run_tqdm
+from utils.os_utils import remove_file, copy_file
+
+np.seterr(divide='ignore', invalid='ignore')
+
+
+class BinarizationError(Exception):
+    pass
+
+sentence2graph_parser = None
+
+
+class BaseBinarizer:
+    def __init__(self, processed_data_dir=None):
+        if processed_data_dir is None:
+            processed_data_dir = hparams['processed_data_dir']
+        self.processed_data_dir = processed_data_dir
+        self.binarization_args = hparams['binarization_args']
+        self.items = {}
+        self.item_names = []
+
+        global sentence2graph_parser
+        from modules.tts.syntaspeech.syntactic_graph_buider import Sentence2GraphParser
+
+        if hparams['ds_name'] == 'libritts':
+            # Unfortunately, we found when processing libritts with multi-processing will incur pytorch.multiprocessing
+            # so we use single thread with cuda graph builder 
+            # it take about 20 hours in a PC with 24-cores-cpu and a RTX2080Ti to process the whole LibriTTS
+            # so run the binarization and take a break!
+            sentence2graph_parser = Sentence2GraphParser("en", use_gpu=True)
+        elif hparams['ds_name'] == 'ljspeech':
+            # use multi-processing, thus gpu is disabled
+            # it takes about 30 minutes for binarization
+            sentence2graph_parser = Sentence2GraphParser("en", use_gpu=False)
+        elif hparams['preprocess_args']['txt_processor'] == 'zh':
+            # use multi-processing, thus gpu is disabled
+            # it takes about 30 minutes for binarization
+            sentence2graph_parser = Sentence2GraphParser("zh", use_gpu=False)
+        else:
+             raise NotImplementedError
+
+    def load_meta_data(self):
+        processed_data_dir = self.processed_data_dir
+        items_list = json.load(open(f"{processed_data_dir}/metadata.json"))
+        for r in tqdm(items_list, desc='Loading meta data.'):
+            item_name = r['item_name']
+            self.items[item_name] = r
+            self.item_names.append(item_name)
+        if self.binarization_args['shuffle']:
+            random.seed(1234)
+            random.shuffle(self.item_names)
+
+    @property
+    def train_item_names(self):
+        range_ = self._convert_range(self.binarization_args['train_range'])
+        return self.item_names[range_[0]:range_[1]]
+
+    @property
+    def valid_item_names(self):
+        range_ = self._convert_range(self.binarization_args['valid_range'])
+        return self.item_names[range_[0]:range_[1]]
+
+    @property
+    def test_item_names(self):
+        range_ = self._convert_range(self.binarization_args['test_range'])
+        return self.item_names[range_[0]:range_[1]]
+
+    def _convert_range(self, range_):
+        if range_[1] == -1:
+            range_[1] = len(self.item_names)
+        return range_
+
+    def meta_data(self, prefix):
+        if prefix == 'valid':
+            item_names = self.valid_item_names
+        elif prefix == 'test':
+            item_names = self.test_item_names
+        else:
+            item_names = self.train_item_names
+        for item_name in item_names:
+            yield self.items[item_name]
+
+    def process(self):
+        self.load_meta_data()
+        os.makedirs(hparams['binary_data_dir'], exist_ok=True)
+        for fn in ['phone_set.json', 'word_set.json', 'spk_map.json']:
+            remove_file(f"{hparams['binary_data_dir']}/{fn}")
+            copy_file(f"{hparams['processed_data_dir']}/{fn}", f"{hparams['binary_data_dir']}/{fn}")
+        if hparams['ds_name'] in ['ljspeech', 'biaobei']:
+            self.process_data('valid')
+            self.process_data('test')
+            self.process_data('train')
+        elif hparams['ds_name'] in ['libritts']:
+            self.process_data_single_processing('valid')
+            self.process_data_single_processing('test')
+            self.process_data_single_processing('train')
+        else:
+            raise NotImplementedError
+
+    def process_data(self, prefix):
+        data_dir = hparams['binary_data_dir']
+        builder = IndexedDatasetBuilder(f'{data_dir}/{prefix}')
+        meta_data = list(self.meta_data(prefix))
+        process_item = partial(self.process_item, binarization_args=self.binarization_args)
+        ph_lengths = []
+        mel_lengths = []
+        total_sec = 0
+        items = []
+        args = [{'item': item} for item in meta_data]
+
+        for item_id, item in multiprocess_run_tqdm(process_item, args, desc='Processing data'):
+            if item is not None:
+                items.append(item)
+        if self.binarization_args['with_spk_embed']:
+            args = [{'wav': item['wav']} for item in items]
+            for item_id, spk_embed in multiprocess_run_tqdm(
+                    self.get_spk_embed, args,
+                    init_ctx_func=lambda wid: {'voice_encoder': VoiceEncoder().cuda()}, num_workers=4,
+                    desc='Extracting spk embed'):
+                items[item_id]['spk_embed'] = spk_embed
+
+        for item in items:
+            if not self.binarization_args['with_wav'] and 'wav' in item:
+                del item['wav']
+            builder.add_item(item)
+            mel_lengths.append(item['len'])
+            assert item['len'] > 0, (item['item_name'], item['txt'], item['mel2ph'])
+            if 'ph_len' in item:
+                ph_lengths.append(item['ph_len'])
+            total_sec += item['sec']
+        builder.finalize()
+        np.save(f'{data_dir}/{prefix}_lengths.npy', mel_lengths)
+        if len(ph_lengths) > 0:
+            np.save(f'{data_dir}/{prefix}_ph_lengths.npy', ph_lengths)
+        print(f"| {prefix} total duration: {total_sec:.3f}s")
+
+
+    def process_data_single_processing(self, prefix):
+        data_dir = hparams['binary_data_dir']
+        builder = IndexedDatasetBuilder(f'{data_dir}/{prefix}')
+        meta_data = list(self.meta_data(prefix))
+        ph_lengths = []
+        mel_lengths = []
+        total_sec = 0
+        items = []
+        args = [{'item': item} for item in meta_data]
+
+        for raw_item in tqdm(meta_data):
+            item = self.process_item(raw_item, self.binarization_args)
+            if item is not None:
+                if item['dgl_graph'].num_nodes() != np.array(item['ph2word']).max():
+                    print(f"Skip Item: {item['item_name']} word nodes number incorrect!")
+                    continue
+
+                items.append(item)
+
+        if self.binarization_args['with_spk_embed']:
+            args = [{'wav': item['wav']} for item in items]
+            for item_id, spk_embed in multiprocess_run_tqdm(
+                    self.get_spk_embed, args,
+                    init_ctx_func=lambda wid: {'voice_encoder': VoiceEncoder().cuda()}, num_workers=4,
+                    desc='Extracting spk embed'):
+                items[item_id]['spk_embed'] = spk_embed
+
+        for item in items:
+            if not self.binarization_args['with_wav'] and 'wav' in item:
+                del item['wav']
+            builder.add_item(item)
+            mel_lengths.append(item['len'])
+            assert item['len'] > 0, (item['item_name'], item['txt'], item['mel2ph'])
+            if 'ph_len' in item:
+                ph_lengths.append(item['ph_len'])
+            total_sec += item['sec']
+        builder.finalize()
+        np.save(f'{data_dir}/{prefix}_lengths.npy', mel_lengths)
+        if len(ph_lengths) > 0:
+            np.save(f'{data_dir}/{prefix}_ph_lengths.npy', ph_lengths)
+        print(f"| {prefix} total duration: {total_sec:.3f}s")
+
+    @classmethod
+    def process_item(cls, item, binarization_args):
+        try:
+            item['ph_len'] = len(item['ph_token'])
+            item_name = item['item_name']
+            wav_fn = item['wav_fn']
+            wav, mel = cls.process_audio(wav_fn, item, binarization_args)
+        except Exception as e:
+            print(f"| Skip item ({e}) for index error. item_name: {item_name}, wav_fn: {wav_fn}")
+            return None
+        try:
+            n_bos_frames, n_eos_frames = 0, 0
+            if binarization_args['with_align']:
+                tg_fn = f"{hparams['processed_data_dir']}/mfa_outputs/{item_name}.TextGrid"
+                item['tg_fn'] = tg_fn
+                cls.process_align(tg_fn, item)
+                if binarization_args['trim_eos_bos']:
+                    n_bos_frames = item['dur'][0]
+                    n_eos_frames = item['dur'][-1]
+                    T = len(mel)
+                    item['mel'] = mel[n_bos_frames:T - n_eos_frames]
+
+                    item['mel2ph'] = item['mel2ph'][n_bos_frames:T - n_eos_frames]
+                    item['mel2word'] = item['mel2word'][n_bos_frames:T - n_eos_frames]
+                    item['dur'] = item['dur'][1:-1]
+                    item['dur_word'] = item['dur_word'][1:-1]
+                    item['len'] = item['mel'].shape[0]
+                    item['wav'] = wav[n_bos_frames * hparams['hop_size']:len(wav) - n_eos_frames * hparams['hop_size']]
+            if binarization_args['with_f0']:
+                cls.process_pitch(item, n_bos_frames, n_eos_frames)
+        except BinarizationError as e:
+            print(f"| Skip item ({e}). item_name: {item_name}, wav_fn: {wav_fn}")
+            return None
+        except Exception as e:
+            traceback.print_exc()
+            print(f"| Skip item. item_name: {item_name}, wav_fn: {wav_fn}")
+            return None
+
+        if item['mel'].shape[0] < 128:
+            print(f"Skip Item: {item['item_name']} Mel-spectrogram is shorter than 128!")
+            return None
+        # fix one bad case of stanza
+        if item['txt'].endswith('yn .'):
+            item['txt'] = item['txt'][:-4]+'y .'
+        try:
+            language = sentence2graph_parser.language
+            if language == 'en':
+                dgl_graph, etypes = sentence2graph_parser.parse(item['txt'])
+            elif language == 'zh':
+                dgl_graph, etypes = sentence2graph_parser.parse(item['txt'], item['word'].split(" "), item['ph_gb_word'].split(" "))
+            else:
+                raise NotImplementedError
+            item['dgl_graph'] = dgl_graph
+            item['edge_types'] = etypes
+        except:
+            print(f"| Dependency Parsing Error! Skip item. item_name: {item_name}, wav_fn: {wav_fn}")
+            return None
+        return item
+
+    @classmethod
+    def process_audio(cls, wav_fn, res, binarization_args):
+        wav2spec_dict = librosa_wav2spec(
+            wav_fn,
+            fft_size=hparams['fft_size'],
+            hop_size=hparams['hop_size'],
+            win_length=hparams['win_size'],
+            num_mels=hparams['audio_num_mel_bins'],
+            fmin=hparams['fmin'],
+            fmax=hparams['fmax'],
+            sample_rate=hparams['audio_sample_rate'],
+            loud_norm=hparams['loud_norm'])
+        mel = wav2spec_dict['mel']
+        wav = wav2spec_dict['wav'].astype(np.float16)
+        if binarization_args['with_linear']:
+            res['linear'] = wav2spec_dict['linear']
+        res.update({'mel': mel, 'wav': wav, 'sec': len(wav) / hparams['audio_sample_rate'], 'len': mel.shape[0]})
+        return wav, mel
+
+    @staticmethod
+    def process_align(tg_fn, item):
+        ph = item['ph']
+        mel = item['mel']
+        ph_token = item['ph_token']
+        if tg_fn is not None and os.path.exists(tg_fn):
+            mel2ph, dur = get_mel2ph(tg_fn, ph, mel, hparams['hop_size'], hparams['audio_sample_rate'],
+                                     hparams['binarization_args']['min_sil_duration'])
+        else:
+            raise BinarizationError(f"Align not found")
+        if np.array(mel2ph).max() - 1 >= len(ph_token):
+            raise BinarizationError(
+                f"Align does not match: mel2ph.max() - 1: {np.array(mel2ph).max() - 1}, len(phone_encoded): {len(ph_token)}")
+        item['mel2ph'] = mel2ph
+        item['dur'] = dur
+
+        ph2word = item['ph2word']
+        mel2word = [ph2word[p - 1] for p in item['mel2ph']]
+        item['mel2word'] = mel2word  # [T_mel]
+        dur_word = mel2token_to_dur(mel2word, len(item['word_token']))
+        item['dur_word'] = dur_word.tolist()  # [T_word]
+
+    @staticmethod
+    def process_pitch(item, n_bos_frames, n_eos_frames):
+        wav, mel = item['wav'], item['mel']
+        f0 = extract_pitch_simple(item['wav'])
+        if sum(f0) == 0:
+            raise BinarizationError("Empty f0")
+        assert len(mel) == len(f0), (len(mel), len(f0))
+        pitch_coarse = f0_to_coarse(f0)
+        item['f0'] = f0
+        item['pitch'] = pitch_coarse
+        if hparams['binarization_args']['with_f0cwt']:
+            uv, cont_lf0_lpf = get_cont_lf0(f0)
+            logf0s_mean_org, logf0s_std_org = np.mean(cont_lf0_lpf), np.std(cont_lf0_lpf)
+            cont_lf0_lpf_norm = (cont_lf0_lpf - logf0s_mean_org) / logf0s_std_org
+            cwt_spec, scales = get_lf0_cwt(cont_lf0_lpf_norm)
+            item['cwt_spec'] = cwt_spec
+            item['cwt_mean'] = logf0s_mean_org
+            item['cwt_std'] = logf0s_std_org
+
+    @staticmethod
+    def get_spk_embed(wav, ctx):
+        return ctx['voice_encoder'].embed_utterance(wav.astype(float))
+
+    @property
+    def num_workers(self):
+        return int(os.getenv('N_PROC', hparams.get('N_PROC', os.cpu_count())))

data_gen/tts/base_preprocess.py

@@ -0,0 +1,251 @@
+import json
+import os
+import random
+import re
+import traceback
+from collections import Counter
+from functools import partial
+
+import librosa
+from tqdm import tqdm
+from data_gen.tts.txt_processors.base_text_processor import get_txt_processor_cls
+from data_gen.tts.wav_processors.base_processor import get_wav_processor_cls
+from utils.commons.hparams import hparams
+from utils.commons.multiprocess_utils import multiprocess_run_tqdm
+from utils.os_utils import link_file, move_file, remove_file
+from utils.text.text_encoder import is_sil_phoneme, build_token_encoder
+
+
+class BasePreprocessor:
+    def __init__(self):
+        self.preprocess_args = hparams['preprocess_args']
+        txt_processor = self.preprocess_args['txt_processor']
+        self.txt_processor = get_txt_processor_cls(txt_processor)
+        self.raw_data_dir = hparams['raw_data_dir']
+        self.processed_dir = hparams['processed_data_dir']
+        self.spk_map_fn = f"{self.processed_dir}/spk_map.json"
+
+    def meta_data(self):
+        """
+
+        :return: {'item_name': Str, 'wav_fn': Str, 'txt': Str, 'spk_name': Str, 'txt_loader': None or Func}
+        """
+        raise NotImplementedError
+
+    def process(self):
+        processed_dir = self.processed_dir
+        wav_processed_tmp_dir = f'{processed_dir}/processed_tmp'
+        remove_file(wav_processed_tmp_dir)
+        os.makedirs(wav_processed_tmp_dir, exist_ok=True)
+        wav_processed_dir = f'{processed_dir}/{self.wav_processed_dirname}'
+        remove_file(wav_processed_dir)
+        os.makedirs(wav_processed_dir, exist_ok=True)
+
+        meta_data = list(tqdm(self.meta_data(), desc='Load meta data'))
+        item_names = [d['item_name'] for d in meta_data]
+        assert len(item_names) == len(set(item_names)), 'Key `item_name` should be Unique.'
+
+        # preprocess data
+        phone_list = []
+        word_list = []
+        spk_names = set()
+        process_item = partial(self.preprocess_first_pass,
+                               txt_processor=self.txt_processor,
+                               wav_processed_dir=wav_processed_dir,
+                               wav_processed_tmp=wav_processed_tmp_dir,
+                               preprocess_args=self.preprocess_args)
+        items = []
+        args = [{
+            'item_name': item_raw['item_name'],
+            'txt_raw': item_raw['txt'],
+            'wav_fn': item_raw['wav_fn'],
+            'txt_loader': item_raw.get('txt_loader'),
+            'others': item_raw.get('others', None)
+        } for item_raw in meta_data]
+        for item_, (item_id, item) in zip(meta_data, multiprocess_run_tqdm(process_item, args, desc='Preprocess')):
+            if item is not None:
+                item_.update(item)
+                item = item_
+                if 'txt_loader' in item:
+                    del item['txt_loader']
+                item['id'] = item_id
+                item['spk_name'] = item.get('spk_name', '<SINGLE_SPK>')
+                item['others'] = item.get('others', None)
+                phone_list += item['ph'].split(" ")
+                word_list += item['word'].split(" ")
+                spk_names.add(item['spk_name'])
+                items.append(item)
+
+        # add encoded tokens
+        ph_encoder, word_encoder = self._phone_encoder(phone_list), self._word_encoder(word_list)
+        spk_map = self.build_spk_map(spk_names)
+        args = [{
+            'ph': item['ph'], 'word': item['word'], 'spk_name': item['spk_name'],
+            'word_encoder': word_encoder, 'ph_encoder': ph_encoder, 'spk_map': spk_map
+        } for item in items]
+        for idx, item_new_kv in multiprocess_run_tqdm(self.preprocess_second_pass, args, desc='Add encoded tokens'):
+            items[idx].update(item_new_kv)
+
+        # build mfa data
+        if self.preprocess_args['use_mfa']:
+            mfa_dict = set()
+            mfa_input_dir = f'{processed_dir}/mfa_inputs'
+            remove_file(mfa_input_dir)
+            # group MFA inputs for better parallelism
+            mfa_groups = [i // self.preprocess_args['nsample_per_mfa_group'] for i in range(len(items))]
+            if self.preprocess_args['mfa_group_shuffle']:
+                random.seed(hparams['seed'])
+                random.shuffle(mfa_groups)
+            args = [{
+                'item': item, 'mfa_input_dir': mfa_input_dir,
+                'mfa_group': mfa_group, 'wav_processed_tmp': wav_processed_tmp_dir,
+                'preprocess_args': self.preprocess_args
+            } for item, mfa_group in zip(items, mfa_groups)]
+            for i, (ph_gb_word_nosil, new_wav_align_fn) in multiprocess_run_tqdm(
+                    self.build_mfa_inputs, args, desc='Build MFA data'):
+                items[i]['wav_align_fn'] = new_wav_align_fn
+                for w in ph_gb_word_nosil.split(" "):
+                    mfa_dict.add(f"{w} {w.replace('_', ' ')}")
+            mfa_dict = sorted(mfa_dict)
+            with open(f'{processed_dir}/mfa_dict.txt', 'w') as f:
+                f.writelines([f'{l}\n' for l in mfa_dict])
+        with open(f"{processed_dir}/{self.meta_csv_filename}.json", 'w') as f:
+            f.write(re.sub(r'\n\s+([\d+\]])', r'\1', json.dumps(items, ensure_ascii=False, sort_keys=False, indent=1)))
+        remove_file(wav_processed_tmp_dir)
+
+    @classmethod
+    def preprocess_first_pass(cls, item_name, txt_raw, txt_processor,
+                              wav_fn, wav_processed_dir, wav_processed_tmp,
+                              preprocess_args, txt_loader=None, others=None):
+        try:
+            if txt_loader is not None:
+                txt_raw = txt_loader(txt_raw)
+            ph, txt, word, ph2word, ph_gb_word = cls.txt_to_ph(txt_processor, txt_raw, preprocess_args)
+            wav_fn, wav_align_fn = cls.process_wav(
+                item_name, wav_fn,
+                hparams['processed_data_dir'],
+                wav_processed_tmp, preprocess_args)
+
+            # wav for binarization
+            ext = os.path.splitext(wav_fn)[1]
+            os.makedirs(wav_processed_dir, exist_ok=True)
+            new_wav_fn = f"{wav_processed_dir}/{item_name}{ext}"
+            move_link_func = move_file if os.path.dirname(wav_fn) == wav_processed_tmp else link_file
+            move_link_func(wav_fn, new_wav_fn)
+            return {
+                'txt': txt, 'txt_raw': txt_raw, 'ph': ph,
+                'word': word, 'ph2word': ph2word, 'ph_gb_word': ph_gb_word,
+                'wav_fn': new_wav_fn, 'wav_align_fn': wav_align_fn,
+                'others': others
+            }
+        except:
+            traceback.print_exc()
+            print(f"| Error is caught. item_name: {item_name}.")
+            return None
+
+    @staticmethod
+    def txt_to_ph(txt_processor, txt_raw, preprocess_args):
+        txt_struct, txt = txt_processor.process(txt_raw, preprocess_args)
+        ph = [p for w in txt_struct for p in w[1]]
+        ph_gb_word = ["_".join(w[1]) for w in txt_struct]
+        words = [w[0] for w in txt_struct]
+        # word_id=0 is reserved for padding
+        ph2word = [w_id + 1 for w_id, w in enumerate(txt_struct) for _ in range(len(w[1]))]
+        return " ".join(ph), txt, " ".join(words), ph2word, " ".join(ph_gb_word)
+
+    @staticmethod
+    def process_wav(item_name, wav_fn, processed_dir, wav_processed_tmp, preprocess_args):
+        processors = [get_wav_processor_cls(v) for v in preprocess_args['wav_processors']]
+        processors = [k() for k in processors if k is not None]
+        if len(processors) >= 1:
+            sr_file = librosa.core.get_samplerate(wav_fn)
+            output_fn_for_align = None
+            ext = os.path.splitext(wav_fn)[1]
+            input_fn = f"{wav_processed_tmp}/{item_name}{ext}"
+            link_file(wav_fn, input_fn)
+            for p in processors:
+                outputs = p.process(input_fn, sr_file, wav_processed_tmp, processed_dir, item_name, preprocess_args)
+                if len(outputs) == 3:
+                    input_fn, sr, output_fn_for_align = outputs
+                else:
+                    input_fn, sr = outputs
+            return input_fn, output_fn_for_align
+        else:
+            return wav_fn, wav_fn
+
+    def _phone_encoder(self, ph_set):
+        ph_set_fn = f"{self.processed_dir}/phone_set.json"
+        if self.preprocess_args['reset_phone_dict'] or not os.path.exists(ph_set_fn):
+            ph_set = sorted(set(ph_set))
+            json.dump(ph_set, open(ph_set_fn, 'w'), ensure_ascii=False)
+            print("| Build phone set: ", ph_set)
+        else:
+            ph_set = json.load(open(ph_set_fn, 'r'))
+            print("| Load phone set: ", ph_set)
+        return build_token_encoder(ph_set_fn)
+
+    def _word_encoder(self, word_set):
+        word_set_fn = f"{self.processed_dir}/word_set.json"
+        if self.preprocess_args['reset_word_dict']:
+            word_set = Counter(word_set)
+            total_words = sum(word_set.values())
+            word_set = word_set.most_common(hparams['word_dict_size'])
+            num_unk_words = total_words - sum([x[1] for x in word_set])
+            word_set = ['<BOS>', '<EOS>'] + [x[0] for x in word_set]
+            word_set = sorted(set(word_set))
+            json.dump(word_set, open(word_set_fn, 'w'), ensure_ascii=False)
+            print(f"| Build word set. Size: {len(word_set)}, #total words: {total_words},"
+                  f" #unk_words: {num_unk_words}, word_set[:10]:, {word_set[:10]}.")
+        else:
+            word_set = json.load(open(word_set_fn, 'r'))
+            print("| Load word set. Size: ", len(word_set), word_set[:10])
+        return build_token_encoder(word_set_fn)
+
+    @classmethod
+    def preprocess_second_pass(cls, word, ph, spk_name, word_encoder, ph_encoder, spk_map):
+        word_token = word_encoder.encode(word)
+        ph_token = ph_encoder.encode(ph)
+        spk_id = spk_map[spk_name]
+        return {'word_token': word_token, 'ph_token': ph_token, 'spk_id': spk_id}
+
+    def build_spk_map(self, spk_names):
+        spk_map = {x: i for i, x in enumerate(sorted(list(spk_names)))}
+        assert len(spk_map) == 0 or len(spk_map) <= hparams['num_spk'], len(spk_map)
+        print(f"| Number of spks: {len(spk_map)}, spk_map: {spk_map}")
+        json.dump(spk_map, open(self.spk_map_fn, 'w'), ensure_ascii=False)
+        return spk_map
+
+    @classmethod
+    def build_mfa_inputs(cls, item, mfa_input_dir, mfa_group, wav_processed_tmp, preprocess_args):
+        item_name = item['item_name']
+        wav_align_fn = item['wav_align_fn']
+        ph_gb_word = item['ph_gb_word']
+        ext = os.path.splitext(wav_align_fn)[1]
+        mfa_input_group_dir = f'{mfa_input_dir}/{mfa_group}'
+        os.makedirs(mfa_input_group_dir, exist_ok=True)
+        new_wav_align_fn = f"{mfa_input_group_dir}/{item_name}{ext}"
+        move_link_func = move_file if os.path.dirname(wav_align_fn) == wav_processed_tmp else link_file
+        move_link_func(wav_align_fn, new_wav_align_fn)
+        ph_gb_word_nosil = " ".join(["_".join([p for p in w.split("_") if not is_sil_phoneme(p)])
+                                     for w in ph_gb_word.split(" ") if not is_sil_phoneme(w)])
+        with open(f'{mfa_input_group_dir}/{item_name}.lab', 'w') as f_txt:
+            f_txt.write(ph_gb_word_nosil)
+        return ph_gb_word_nosil, new_wav_align_fn
+
+    def load_spk_map(self, base_dir):
+        spk_map_fn = f"{base_dir}/spk_map.json"
+        spk_map = json.load(open(spk_map_fn, 'r'))
+        return spk_map
+
+    def load_dict(self, base_dir):
+        ph_encoder = build_token_encoder(f'{base_dir}/phone_set.json')
+        word_encoder = build_token_encoder(f'{base_dir}/word_set.json')
+        return ph_encoder, word_encoder
+
+    @property
+    def meta_csv_filename(self):
+        return 'metadata'
+
+    @property
+    def wav_processed_dirname(self):
+        return 'wav_processed'

data_gen/tts/runs/align_and_binarize.py

@@ -0,0 +1,12 @@
+import utils.commons.single_thread_env  # NOQA
+from utils.commons.hparams import set_hparams, hparams
+from data_gen.tts.runs.binarize import binarize
+from data_gen.tts.runs.preprocess import preprocess
+from data_gen.tts.runs.train_mfa_align import train_mfa_align
+
+if __name__ == '__main__':
+    set_hparams()
+    preprocess()
+    if hparams['preprocess_args']['use_mfa']:
+        train_mfa_align()
+    binarize()

data_gen/tts/runs/binarize.py

@@ -0,0 +1,17 @@
+import utils.commons.single_thread_env  # NOQA
+from utils.commons.hparams import hparams, set_hparams
+import importlib
+
+
+def binarize():
+    binarizer_cls = hparams.get("binarizer_cls", 'data_gen.tts.base_binarizer.BaseBinarizer')
+    pkg = ".".join(binarizer_cls.split(".")[:-1])
+    cls_name = binarizer_cls.split(".")[-1]
+    binarizer_cls = getattr(importlib.import_module(pkg), cls_name)
+    print("| Binarizer: ", binarizer_cls)
+    binarizer_cls().process()
+
+
+if __name__ == '__main__':
+    set_hparams()
+    binarize()

data_gen/tts/runs/preprocess.py

@@ -0,0 +1,17 @@
+import utils.commons.single_thread_env  # NOQA
+from utils.commons.hparams import hparams, set_hparams
+import importlib
+
+
+def preprocess():
+    assert hparams['preprocess_cls'] != ''
+
+    pkg = ".".join(hparams["preprocess_cls"].split(".")[:-1])
+    cls_name = hparams["preprocess_cls"].split(".")[-1]
+    process_cls = getattr(importlib.import_module(pkg), cls_name)
+    process_cls().process()
+
+
+if __name__ == '__main__':
+    set_hparams()
+    preprocess()

data_gen/tts/runs/train_mfa_align.py

@@ -0,0 +1,46 @@
+import utils.commons.single_thread_env  # NOQA
+import glob
+import subprocess
+from textgrid import TextGrid
+import os
+from utils.commons.hparams import hparams, set_hparams
+
+
+def train_mfa_align(mfa_outputs="mfa_outputs",
+                    mfa_inputs="mfa_inputs",
+                    model_name=None, pretrain_model_name=None,
+                    mfa_cmd='train'):
+    CORPUS = hparams['processed_data_dir'].split("/")[-1]
+    NUM_JOB = int(os.getenv('N_PROC', os.cpu_count()))
+    env_vars = [f'CORPUS={CORPUS}', f'NUM_JOB={NUM_JOB}']
+    if mfa_outputs is not None:
+        env_vars.append(f'MFA_OUTPUTS={mfa_outputs}')
+    if mfa_inputs is not None:
+        env_vars.append(f'MFA_INPUTS={mfa_inputs}')
+    if model_name is not None:
+        env_vars.append(f'MODEL_NAME={model_name}')
+    if pretrain_model_name is not None:
+        env_vars.append(f'PRETRAIN_MODEL_NAME={pretrain_model_name}')
+    if mfa_cmd is not None:
+        env_vars.append(f'MFA_CMD={mfa_cmd}')
+    env_str = ' '.join(env_vars)
+    print(f"| Run MFA for {CORPUS}. Env vars: {env_str}")
+    subprocess.check_call(f'{env_str} bash mfa_usr/run_mfa_train_align.sh', shell=True)
+    mfa_offset = hparams['preprocess_args']['mfa_offset']
+    if mfa_offset > 0:
+        for tg_fn in glob.glob(f'{hparams["processed_data_dir"]}/{mfa_outputs}/*.TextGrid'):
+            tg = TextGrid.fromFile(tg_fn)
+            max_time = tg.maxTime
+            for tier in tg.tiers:
+                for interval in tier.intervals:
+                    interval.maxTime = min(interval.maxTime + mfa_offset, max_time)
+                    interval.minTime = min(interval.minTime + mfa_offset, max_time)
+                tier.intervals[0].minTime = 0
+                tier.maxTime = min(tier.maxTime + mfa_offset, max_time)
+            tg.write(tg_fn)
+            TextGrid.fromFile(tg_fn)
+
+
+if __name__ == '__main__':
+    set_hparams(print_hparams=False)
+    train_mfa_align()

data_gen/tts/txt_processors/__init__.py

@@ -0,0 +1 @@
+from . import en, zh

data_gen/tts/txt_processors/base_text_processor.py

@@ -0,0 +1,50 @@
+from utils.text.text_encoder import is_sil_phoneme
+
+REGISTERED_TEXT_PROCESSORS = {}
+
+
+def register_txt_processors(name):
+    def _f(cls):
+        REGISTERED_TEXT_PROCESSORS[name] = cls
+        return cls
+
+    return _f
+
+
+def get_txt_processor_cls(name):
+    return REGISTERED_TEXT_PROCESSORS.get(name, None)
+
+
+class BaseTxtProcessor:
+    @staticmethod
+    def sp_phonemes():
+        return ['|']
+
+    @classmethod
+    def process(cls, txt, preprocess_args):
+        raise NotImplementedError
+
+    @classmethod
+    def postprocess(cls, txt_struct, preprocess_args):
+        # remove sil phoneme in head and tail
+        while len(txt_struct) > 0 and is_sil_phoneme(txt_struct[0][0]):
+            txt_struct = txt_struct[1:]
+        while len(txt_struct) > 0 and is_sil_phoneme(txt_struct[-1][0]):
+            txt_struct = txt_struct[:-1]
+        if preprocess_args['with_phsep']:
+            txt_struct = cls.add_bdr(txt_struct)
+        if preprocess_args['add_eos_bos']:
+            txt_struct = [["<BOS>", ["<BOS>"]]] + txt_struct + [["<EOS>", ["<EOS>"]]]
+        return txt_struct
+
+    @classmethod
+    def add_bdr(cls, txt_struct):
+        txt_struct_ = []
+        for i, ts in enumerate(txt_struct):
+            txt_struct_.append(ts)
+            if i != len(txt_struct) - 1 and \
+                    not is_sil_phoneme(txt_struct[i][0]) and not is_sil_phoneme(txt_struct[i + 1][0]):
+                # txt_struct_.append(['|', ['|']])
+                # We disbale the sep token because it is imcompatible with syntactic graph.
+                pass
+        return txt_struct_

data_gen/tts/txt_processors/en.py

@@ -0,0 +1,78 @@
+import re
+import unicodedata
+
+from g2p_en import G2p
+from g2p_en.expand import normalize_numbers
+from nltk import pos_tag
+from nltk.tokenize import TweetTokenizer
+
+from data_gen.tts.txt_processors.base_text_processor import BaseTxtProcessor, register_txt_processors
+from utils.text.text_encoder import PUNCS, is_sil_phoneme
+
+
+class EnG2p(G2p):
+    word_tokenize = TweetTokenizer().tokenize
+
+    def __call__(self, text):
+        # preprocessing
+        words = EnG2p.word_tokenize(text)
+        tokens = pos_tag(words)  # tuples of (word, tag)
+
+        # steps
+        prons = []
+        for word, pos in tokens:
+            if re.search("[a-z]", word) is None:
+                pron = [word]
+
+            elif word in self.homograph2features:  # Check homograph
+                pron1, pron2, pos1 = self.homograph2features[word]
+                if pos.startswith(pos1):
+                    pron = pron1
+                else:
+                    pron = pron2
+            elif word in self.cmu:  # lookup CMU dict
+                pron = self.cmu[word][0]
+            else:  # predict for oov
+                pron = self.predict(word)
+
+            prons.extend(pron)
+            prons.extend([" "])
+
+        return prons[:-1]
+
+
+@register_txt_processors('en')
+class TxtProcessor(BaseTxtProcessor):
+    g2p = EnG2p()
+
+    @staticmethod
+    def preprocess_text(text):
+        text = normalize_numbers(text)
+        text = ''.join(char for char in unicodedata.normalize('NFD', text)
+                       if unicodedata.category(char) != 'Mn')  # Strip accents
+        text = text.lower()
+        text = re.sub("[\'\"()]+", "", text)
+        text = re.sub("[-]+", " ", text)
+        text = re.sub(f"[^ a-z{PUNCS}]", "", text)
+        text = re.sub(f" ?([{PUNCS}]) ?", r"\1", text)  # !! -> !
+        text = re.sub(f"([{PUNCS}])+", r"\1", text)  # !! -> !
+        text = text.replace("i.e.", "that is")
+        text = text.replace("i.e.", "that is")
+        text = text.replace("etc.", "etc")
+        text = re.sub(f"([{PUNCS}])", r" \1 ", text)
+        text = re.sub(rf"\s+", r" ", text)
+        return text
+
+    @classmethod
+    def process(cls, txt, preprocess_args):
+        txt = cls.preprocess_text(txt).strip()
+        phs = cls.g2p(txt)
+        txt_struct = [[w, []] for w in txt.split(" ")]
+        i_word = 0
+        for p in phs:
+            if p == ' ':
+                i_word += 1
+            else:
+                txt_struct[i_word][1].append(p)
+        txt_struct = cls.postprocess(txt_struct, preprocess_args)
+        return txt_struct, txt

data_gen/tts/txt_processors/zh.py

@@ -0,0 +1,110 @@
+import re
+import jieba
+from pypinyin import pinyin, Style
+from utils.text.text_norm import NSWNormalizer
+from data_gen.tts.txt_processors.base_text_processor import BaseTxtProcessor, register_txt_processors
+from utils.text.text_encoder import PUNCS, is_sil_phoneme
+
+ALL_SHENMU = ['zh', 'ch', 'sh', 'b', 'p', 'm', 'f', 'd', 't', 'n', 'l', 'g', 'k', 'h', 'j',
+              'q', 'x', 'r', 'z', 'c', 's', 'y', 'w']
+
+
+@register_txt_processors('zh')
+class TxtProcessor(BaseTxtProcessor):
+    table = {ord(f): ord(t) for f, t in zip(
+        u'：，。！？【】（）％＃＠＆１２３４５６７８９０',
+        u':,.!?[]()%#@&1234567890')}
+
+    @staticmethod
+    def sp_phonemes():
+        return ['|', '#']
+
+    @staticmethod
+    def preprocess_text(text):
+        text = text.translate(TxtProcessor.table)
+        text = NSWNormalizer(text).normalize(remove_punc=False).lower()
+        text = re.sub("[\'\"()]+", "", text)
+        text = re.sub("[-]+", " ", text)
+        text = re.sub(f"[^ A-Za-z\u4e00-\u9fff{PUNCS}]", "", text)
+        text = re.sub(f"([{PUNCS}])+", r"\1", text)  # !! -> !
+        text = re.sub(f"([{PUNCS}])", r" \1 ", text)
+        text = re.sub(rf"\s+", r"", text)
+        text = re.sub(rf"[A-Za-z]+", r"$", text)
+        return text
+
+    @classmethod
+    def pinyin_with_en(cls, txt, style):
+        x = pinyin(txt, style)
+        x = [t[0] for t in x]
+        x_ = []
+        for t in x:
+            if '$' not in t:
+                x_.append(t)
+            else:
+                x_ += list(t)
+        x_ = [t if t != '$' else 'ENG' for t in x_]
+        return x_
+
+    @classmethod
+    def process(cls, txt, pre_align_args):
+        txt = cls.preprocess_text(txt)
+        # https://blog.csdn.net/zhoulei124/article/details/89055403
+        pre_align_args['use_tone'] = True
+        shengmu = cls.pinyin_with_en(txt, style=Style.INITIALS)
+        yunmu = cls.pinyin_with_en(txt, style=
+        Style.FINALS_TONE3 if pre_align_args['use_tone'] else Style.FINALS)
+        assert len(shengmu) == len(yunmu)
+        ph_list = []
+        for a, b in zip(shengmu, yunmu):
+            if a == b:
+                ph_list += [a]
+            else:
+                ph_list += [a + "%" + b]
+        seg_list = '#'.join(jieba.cut(txt))
+        assert len(ph_list) == len([s for s in seg_list if s != '#']), (ph_list, seg_list)
+
+        # 加入词边界'#'
+        ph_list_ = []
+        seg_idx = 0
+        for p in ph_list:
+            if seg_list[seg_idx] == '#':
+                ph_list_.append('#')
+                seg_idx += 1
+            elif len(ph_list_) > 0:
+                ph_list_.append("|")
+            seg_idx += 1
+            finished = False
+            if not finished:
+                ph_list_ += [x for x in p.split("%") if x != '']
+
+        ph_list = ph_list_
+
+        # 去除静音符号周围的词边界标记 [..., '#', ',', '#', ...]
+        sil_phonemes = list(PUNCS) + TxtProcessor.sp_phonemes()
+        ph_list_ = []
+        for i in range(0, len(ph_list), 1):
+            if ph_list[i] != '#' or (ph_list[i - 1] not in sil_phonemes and ph_list[i + 1] not in sil_phonemes):
+                ph_list_.append(ph_list[i])
+        ph_list = ph_list_
+        txt_struct = [[w, []] for w in txt]
+        i = 0
+        for ph in ph_list:
+            if ph == '|' or ph == '#':
+                i += 1
+                continue
+            elif ph in [',', '.']:
+                i += 1
+                txt_struct[i][1].append(ph)
+                i += 1
+                continue
+            txt_struct[i][1].append(ph)
+        # return ph_list, txt
+        txt_struct.insert(0, ['<BOS>', ['<BOS>']])
+        txt_struct.append(['<EOS>', ['<EOS>']])
+        return txt_struct, txt
+
+
+if __name__ == '__main__':
+    t = 'simon演唱过后，simon还进行了simon精彩的文艺演出simon.'
+    phs, txt = TxtProcessor.process(t, {'use_tone': True})
+    print(phs, txt)

data_gen/tts/wav_processors/__init__.py

@@ -0,0 +1,2 @@
+from . import base_processor
+from . import common_processors

data_gen/tts/wav_processors/base_processor.py

@@ -0,0 +1,25 @@
+REGISTERED_WAV_PROCESSORS = {}
+
+
+def register_wav_processors(name):
+    def _f(cls):
+        REGISTERED_WAV_PROCESSORS[name] = cls
+        return cls
+
+    return _f
+
+
+def get_wav_processor_cls(name):
+    return REGISTERED_WAV_PROCESSORS.get(name, None)
+
+
+class BaseWavProcessor:
+    @property
+    def name(self):
+        raise NotImplementedError
+
+    def output_fn(self, input_fn):
+        return f'{input_fn[:-4]}_{self.name}.wav'
+
+    def process(self, input_fn, sr, tmp_dir, processed_dir, item_name, preprocess_args):
+        raise NotImplementedError

data_gen/tts/wav_processors/common_processors.py

@@ -0,0 +1,86 @@
+import os
+import subprocess
+import librosa
+import numpy as np
+from data_gen.tts.wav_processors.base_processor import BaseWavProcessor, register_wav_processors
+from utils.audio import trim_long_silences
+from utils.audio.io import save_wav
+from utils.audio.rnnoise import rnnoise
+from utils.commons.hparams import hparams
+
+
+@register_wav_processors(name='sox_to_wav')
+class ConvertToWavProcessor(BaseWavProcessor):
+    @property
+    def name(self):
+        return 'ToWav'
+
+    def process(self, input_fn, sr, tmp_dir, processed_dir, item_name, preprocess_args):
+        if input_fn[-4:] == '.wav':
+            return input_fn, sr
+        else:
+            output_fn = self.output_fn(input_fn)
+            subprocess.check_call(f'sox -v 0.95 "{input_fn}" -t wav "{output_fn}"', shell=True)
+            return output_fn, sr
+
+
+@register_wav_processors(name='sox_resample')
+class ResampleProcessor(BaseWavProcessor):
+    @property
+    def name(self):
+        return 'Resample'
+
+    def process(self, input_fn, sr, tmp_dir, processed_dir, item_name, preprocess_args):
+        output_fn = self.output_fn(input_fn)
+        sr_file = librosa.core.get_samplerate(input_fn)
+        if sr != sr_file:
+            subprocess.check_call(f'sox -v 0.95 "{input_fn}" -r{sr} "{output_fn}"', shell=True)
+            y, _ = librosa.core.load(input_fn, sr=sr)
+            y, _ = librosa.effects.trim(y)
+            save_wav(y, output_fn, sr)
+            return output_fn, sr
+        else:
+            return input_fn, sr
+
+
+@register_wav_processors(name='trim_sil')
+class TrimSILProcessor(BaseWavProcessor):
+    @property
+    def name(self):
+        return 'TrimSIL'
+
+    def process(self, input_fn, sr, tmp_dir, processed_dir, item_name, preprocess_args):
+        output_fn = self.output_fn(input_fn)
+        y, _ = librosa.core.load(input_fn, sr=sr)
+        y, _ = librosa.effects.trim(y)
+        save_wav(y, output_fn, sr)
+        return output_fn
+
+
+@register_wav_processors(name='trim_all_sil')
+class TrimAllSILProcessor(BaseWavProcessor):
+    @property
+    def name(self):
+        return 'TrimSIL'
+
+    def process(self, input_fn, sr, tmp_dir, processed_dir, item_name, preprocess_args):
+        output_fn = self.output_fn(input_fn)
+        y, audio_mask, _ = trim_long_silences(
+            input_fn, vad_max_silence_length=preprocess_args.get('vad_max_silence_length', 12))
+        save_wav(y, output_fn, sr)
+        if preprocess_args['save_sil_mask']:
+            os.makedirs(f'{processed_dir}/sil_mask', exist_ok=True)
+            np.save(f'{processed_dir}/sil_mask/{item_name}.npy', audio_mask)
+        return output_fn, sr
+
+
+@register_wav_processors(name='denoise')
+class DenoiseProcessor(BaseWavProcessor):
+    @property
+    def name(self):
+        return 'Denoise'
+
+    def process(self, input_fn, sr, tmp_dir, processed_dir, item_name, preprocess_args):
+        output_fn = self.output_fn(input_fn)
+        rnnoise(input_fn, output_fn, out_sample_rate=sr)
+        return output_fn, sr