Spaces:
Build error
Build error
File size: 9,029 Bytes
e882f51 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 |
import json
import logging
import os
import random
from copy import deepcopy
import numpy as np
import yaml
from resemblyzer import VoiceEncoder
from tqdm import tqdm
from infer_tools.f0_static import static_f0_time
from modules.vocoders.nsf_hifigan import NsfHifiGAN
from preprocessing.hubertinfer import HubertEncoder
from preprocessing.process_pipeline import File2Batch
from preprocessing.process_pipeline import get_pitch_parselmouth, get_pitch_crepe
from utils.hparams import hparams
from utils.hparams import set_hparams
from utils.indexed_datasets import IndexedDatasetBuilder
os.environ["OMP_NUM_THREADS"] = "1"
BASE_ITEM_ATTRIBUTES = ['wav_fn', 'spk_id']
class SvcBinarizer:
'''
Base class for data processing.
1. *process* and *process_data_split*:
process entire data, generate the train-test split (support parallel processing);
2. *process_item*:
process singe piece of data;
3. *get_pitch*:
infer the pitch using some algorithm;
4. *get_align*:
get the alignment using 'mel2ph' format (see https://arxiv.org/abs/1905.09263).
5. phoneme encoder, voice encoder, etc.
Subclasses should define:
1. *load_metadata*:
how to read multiple datasets from files;
2. *train_item_names*, *valid_item_names*, *test_item_names*:
how to split the dataset;
3. load_ph_set:
the phoneme set.
'''
def __init__(self, data_dir=None, item_attributes=None):
self.spk_map = None
self.vocoder = NsfHifiGAN()
self.phone_encoder = HubertEncoder(pt_path=hparams['hubert_path'])
if item_attributes is None:
item_attributes = BASE_ITEM_ATTRIBUTES
if data_dir is None:
data_dir = hparams['raw_data_dir']
if 'speakers' not in hparams:
speakers = hparams['datasets']
hparams['speakers'] = hparams['datasets']
else:
speakers = hparams['speakers']
assert isinstance(speakers, list), 'Speakers must be a list'
assert len(speakers) == len(set(speakers)), 'Speakers cannot contain duplicate names'
self.raw_data_dirs = data_dir if isinstance(data_dir, list) else [data_dir]
assert len(speakers) == len(self.raw_data_dirs), \
'Number of raw data dirs must equal number of speaker names!'
self.speakers = speakers
self.binarization_args = hparams['binarization_args']
self.items = {}
# every item in self.items has some attributes
self.item_attributes = item_attributes
# load each dataset
for ds_id, data_dir in enumerate(self.raw_data_dirs):
self.load_meta_data(data_dir, ds_id)
if ds_id == 0:
# check program correctness
assert all([attr in self.item_attributes for attr in list(self.items.values())[0].keys()])
self.item_names = sorted(list(self.items.keys()))
if self.binarization_args['shuffle']:
random.seed(hparams['seed'])
random.shuffle(self.item_names)
# set default get_pitch algorithm
if hparams['use_crepe']:
self.get_pitch_algorithm = get_pitch_crepe
else:
self.get_pitch_algorithm = get_pitch_parselmouth
print('spkers: ', set(self.speakers))
self._train_item_names, self._test_item_names = self.split_train_test_set(self.item_names)
@staticmethod
def split_train_test_set(item_names):
auto_test = item_names[-5:]
item_names = set(deepcopy(item_names))
if hparams['choose_test_manually']:
prefixes = set([str(pr) for pr in hparams['test_prefixes']])
test_item_names = set()
# Add prefixes that specified speaker index and matches exactly item name to test set
for prefix in deepcopy(prefixes):
if prefix in item_names:
test_item_names.add(prefix)
prefixes.remove(prefix)
# Add prefixes that exactly matches item name without speaker id to test set
for prefix in deepcopy(prefixes):
for name in item_names:
if name.split(':')[-1] == prefix:
test_item_names.add(name)
prefixes.remove(prefix)
# Add names with one of the remaining prefixes to test set
for prefix in deepcopy(prefixes):
for name in item_names:
if name.startswith(prefix):
test_item_names.add(name)
prefixes.remove(prefix)
for prefix in prefixes:
for name in item_names:
if name.split(':')[-1].startswith(prefix):
test_item_names.add(name)
test_item_names = sorted(list(test_item_names))
else:
test_item_names = auto_test
train_item_names = [x for x in item_names if x not in set(test_item_names)]
logging.info("train {}".format(len(train_item_names)))
logging.info("test {}".format(len(test_item_names)))
return train_item_names, test_item_names
@property
def train_item_names(self):
return self._train_item_names
@property
def valid_item_names(self):
return self._test_item_names
@property
def test_item_names(self):
return self._test_item_names
def load_meta_data(self, raw_data_dir, ds_id):
self.items.update(File2Batch.file2temporary_dict(raw_data_dir, ds_id))
@staticmethod
def build_spk_map():
spk_map = {x: i for i, x in enumerate(hparams['speakers'])}
assert len(spk_map) <= hparams['num_spk'], 'Actual number of speakers should be smaller than num_spk!'
return spk_map
def item_name2spk_id(self, item_name):
return self.spk_map[self.items[item_name]['spk_id']]
def meta_data_iterator(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:
meta_data = self.items[item_name]
yield item_name, meta_data
def process(self):
os.makedirs(hparams['binary_data_dir'], exist_ok=True)
self.spk_map = self.build_spk_map()
print("| spk_map: ", self.spk_map)
spk_map_fn = f"{hparams['binary_data_dir']}/spk_map.json"
json.dump(self.spk_map, open(spk_map_fn, 'w', encoding='utf-8'))
self.process_data_split('valid')
self.process_data_split('test')
self.process_data_split('train')
def process_data_split(self, prefix):
data_dir = hparams['binary_data_dir']
args = []
builder = IndexedDatasetBuilder(f'{data_dir}/{prefix}')
lengths = []
total_sec = 0
if self.binarization_args['with_spk_embed']:
voice_encoder = VoiceEncoder().cuda()
for item_name, meta_data in self.meta_data_iterator(prefix):
args.append([item_name, meta_data, self.binarization_args])
spec_min = []
spec_max = []
f0_dict = {}
# code for single cpu processing
for i in tqdm(reversed(range(len(args))), total=len(args)):
a = args[i]
item = self.process_item(*a)
if item is None:
continue
item['spk_embed'] = voice_encoder.embed_utterance(item['wav']) \
if self.binarization_args['with_spk_embed'] else None
spec_min.append(item['spec_min'])
spec_max.append(item['spec_max'])
f0_dict[item['wav_fn']] = item['f0']
builder.add_item(item)
lengths.append(item['len'])
total_sec += item['sec']
if prefix == 'train':
spec_max = np.max(spec_max, 0)
spec_min = np.min(spec_min, 0)
pitch_time = static_f0_time(f0_dict)
with open(hparams['config_path'], encoding='utf-8') as f:
_hparams = yaml.safe_load(f)
_hparams['spec_max'] = spec_max.tolist()
_hparams['spec_min'] = spec_min.tolist()
if self.speakers == 1:
_hparams['f0_static'] = json.dumps(pitch_time)
with open(hparams['config_path'], 'w', encoding='utf-8') as f:
yaml.safe_dump(_hparams, f)
builder.finalize()
np.save(f'{data_dir}/{prefix}_lengths.npy', lengths)
print(f"| {prefix} total duration: {total_sec:.3f}s")
def process_item(self, item_name, meta_data, binarization_args):
from preprocessing.process_pipeline import File2Batch
return File2Batch.temporary_dict2processed_input(item_name, meta_data, self.phone_encoder)
if __name__ == "__main__":
set_hparams()
SvcBinarizer().process()
|