OzoneAsai/Style-Bert-VITS2
0
1import os2import random3import torch4import torch.utils.data5from tqdm import tqdm6import numpy as np7from tools.log import logger8import commons9from mel_processing import spectrogram_torch, mel_spectrogram_torch10from utils import load_wav_to_torch, load_filepaths_and_text11from text import cleaned_text_to_sequence12from config import config13 14"""Multi speaker version"""15 16 17class TextAudioSpeakerLoader(torch.utils.data.Dataset):18 """19 1) loads audio, speaker_id, text pairs20 2) normalizes text and converts them to sequences of integers21 3) computes spectrograms from audio files.22 """23 24 def __init__(self, audiopaths_sid_text, hparams):25 self.audiopaths_sid_text = load_filepaths_and_text(audiopaths_sid_text)26 self.max_wav_value = hparams.max_wav_value27 self.sampling_rate = hparams.sampling_rate28 self.filter_length = hparams.filter_length29 self.hop_length = hparams.hop_length30 self.win_length = hparams.win_length31 self.sampling_rate = hparams.sampling_rate32 self.spk_map = hparams.spk2id33 self.hparams = hparams34 35 self.use_mel_spec_posterior = getattr(36 hparams, "use_mel_posterior_encoder", False37 )38 if self.use_mel_spec_posterior:39 self.n_mel_channels = getattr(hparams, "n_mel_channels", 80)40 41 self.cleaned_text = getattr(hparams, "cleaned_text", False)42 43 self.add_blank = hparams.add_blank44 self.min_text_len = getattr(hparams, "min_text_len", 1)45 self.max_text_len = getattr(hparams, "max_text_len", 384)46 47 random.seed(1234)48 random.shuffle(self.audiopaths_sid_text)49 self._filter()50 51 def _filter(self):52 """53 Filter text & store spec lengths54 """55 # Store spectrogram lengths for Bucketing56 # wav_length ~= file_size / (wav_channels * Bytes per dim) = file_size / (1 * 2)57 # spec_length = wav_length // hop_length58 59 audiopaths_sid_text_new = []60 lengths = []61 skipped = 062 logger.info("Init dataset...")63 for _id, spk, language, text, phones, tone, word2ph in tqdm(64 self.audiopaths_sid_text65 ):66 audiopath = f"{_id}"67 if self.min_text_len <= len(phones) and len(phones) <= self.max_text_len:68 phones = phones.split(" ")69 tone = [int(i) for i in tone.split(" ")]70 word2ph = [int(i) for i in word2ph.split(" ")]71 audiopaths_sid_text_new.append(72 [audiopath, spk, language, text, phones, tone, word2ph]73 )74 lengths.append(os.path.getsize(audiopath) // (2 * self.hop_length))75 else:76 skipped += 177 logger.info(78 "skipped: "79 + str(skipped)80 + ", total: "81 + str(len(self.audiopaths_sid_text))82 )83 self.audiopaths_sid_text = audiopaths_sid_text_new84 self.lengths = lengths85 86 def get_audio_text_speaker_pair(self, audiopath_sid_text):87 # separate filename, speaker_id and text88 audiopath, sid, language, text, phones, tone, word2ph = audiopath_sid_text89 90 bert, ja_bert, en_bert, phones, tone, language = self.get_text(91 text, word2ph, phones, tone, language, audiopath92 )93 94 spec, wav = self.get_audio(audiopath)95 sid = torch.LongTensor([int(self.spk_map[sid])])96 style_vec = torch.FloatTensor(np.load(f"{audiopath}.npy"))97 return (98 phones,99 spec,100 wav,101 sid,102 tone,103 language,104 bert,105 ja_bert,106 en_bert,107 style_vec,108 )109 110 def get_audio(self, filename):111 audio, sampling_rate = load_wav_to_torch(filename)112 if sampling_rate != self.sampling_rate:113 raise ValueError(114 "{} {} SR doesn't match target {} SR".format(115 filename, sampling_rate, self.sampling_rate116 )117 )118 audio_norm = audio / self.max_wav_value119 audio_norm = audio_norm.unsqueeze(0)120 spec_filename = filename.replace(".wav", ".spec.pt")121 if self.use_mel_spec_posterior:122 spec_filename = spec_filename.replace(".spec.pt", ".mel.pt")123 try:124 spec = torch.load(spec_filename)125 except:126 if self.use_mel_spec_posterior:127 spec = mel_spectrogram_torch(128 audio_norm,129 self.filter_length,130 self.n_mel_channels,131 self.sampling_rate,132 self.hop_length,133 self.win_length,134 self.hparams.mel_fmin,135 self.hparams.mel_fmax,136 center=False,137 )138 else:139 spec = spectrogram_torch(140 audio_norm,141 self.filter_length,142 self.sampling_rate,143 self.hop_length,144 self.win_length,145 center=False,146 )147 spec = torch.squeeze(spec, 0)148 if config.train_ms_config.spec_cache:149 torch.save(spec, spec_filename)150 return spec, audio_norm151 152 def get_text(self, text, word2ph, phone, tone, language_str, wav_path):153 phone, tone, language = cleaned_text_to_sequence(phone, tone, language_str)154 if self.add_blank:155 phone = commons.intersperse(phone, 0)156 tone = commons.intersperse(tone, 0)157 language = commons.intersperse(language, 0)158 for i in range(len(word2ph)):159 word2ph[i] = word2ph[i] * 2160 word2ph[0] += 1161 bert_path = wav_path.replace(".wav", ".bert.pt")162 try:163 bert_ori = torch.load(bert_path)164 assert bert_ori.shape[-1] == len(phone)165 except Exception as e:166 logger.warning("Bert load Failed")167 logger.warning(e)168 169 if language_str == "ZH":170 bert = bert_ori171 ja_bert = torch.zeros(1024, len(phone))172 en_bert = torch.zeros(1024, len(phone))173 elif language_str == "JP":174 bert = torch.zeros(1024, len(phone))175 ja_bert = bert_ori176 en_bert = torch.zeros(1024, len(phone))177 elif language_str == "EN":178 bert = torch.zeros(1024, len(phone))179 ja_bert = torch.zeros(1024, len(phone))180 en_bert = bert_ori181 phone = torch.LongTensor(phone)182 tone = torch.LongTensor(tone)183 language = torch.LongTensor(language)184 return bert, ja_bert, en_bert, phone, tone, language185 186 def get_sid(self, sid):187 sid = torch.LongTensor([int(sid)])188 return sid189 190 def __getitem__(self, index):191 return self.get_audio_text_speaker_pair(self.audiopaths_sid_text[index])192 193 def __len__(self):194 return len(self.audiopaths_sid_text)195 196 197class TextAudioSpeakerCollate:198 """Zero-pads model inputs and targets"""199 200 def __init__(self, return_ids=False):201 self.return_ids = return_ids202 203 def __call__(self, batch):204 """Collate's training batch from normalized text, audio and speaker identities205 PARAMS206 ------207 batch: [text_normalized, spec_normalized, wav_normalized, sid]208 """209 # Right zero-pad all one-hot text sequences to max input length210 _, ids_sorted_decreasing = torch.sort(211 torch.LongTensor([x[1].size(1) for x in batch]), dim=0, descending=True212 )213 214 max_text_len = max([len(x[0]) for x in batch])215 max_spec_len = max([x[1].size(1) for x in batch])216 max_wav_len = max([x[2].size(1) for x in batch])217 218 text_lengths = torch.LongTensor(len(batch))219 spec_lengths = torch.LongTensor(len(batch))220 wav_lengths = torch.LongTensor(len(batch))221 sid = torch.LongTensor(len(batch))222 223 text_padded = torch.LongTensor(len(batch), max_text_len)224 tone_padded = torch.LongTensor(len(batch), max_text_len)225 language_padded = torch.LongTensor(len(batch), max_text_len)226 bert_padded = torch.FloatTensor(len(batch), 1024, max_text_len)227 ja_bert_padded = torch.FloatTensor(len(batch), 1024, max_text_len)228 en_bert_padded = torch.FloatTensor(len(batch), 1024, max_text_len)229 style_vec = torch.FloatTensor(len(batch), 256)230 231 spec_padded = torch.FloatTensor(len(batch), batch[0][1].size(0), max_spec_len)232 wav_padded = torch.FloatTensor(len(batch), 1, max_wav_len)233 text_padded.zero_()234 tone_padded.zero_()235 language_padded.zero_()236 spec_padded.zero_()237 wav_padded.zero_()238 bert_padded.zero_()239 ja_bert_padded.zero_()240 en_bert_padded.zero_()241 style_vec.zero_()242 243 for i in range(len(ids_sorted_decreasing)):244 row = batch[ids_sorted_decreasing[i]]245 246 text = row[0]247 text_padded[i, : text.size(0)] = text248 text_lengths[i] = text.size(0)249 250 spec = row[1]251 spec_padded[i, :, : spec.size(1)] = spec252 spec_lengths[i] = spec.size(1)253 254 wav = row[2]255 wav_padded[i, :, : wav.size(1)] = wav256 wav_lengths[i] = wav.size(1)257 258 sid[i] = row[3]259 260 tone = row[4]261 tone_padded[i, : tone.size(0)] = tone262 263 language = row[5]264 language_padded[i, : language.size(0)] = language265 266 bert = row[6]267 bert_padded[i, :, : bert.size(1)] = bert268 269 ja_bert = row[7]270 ja_bert_padded[i, :, : ja_bert.size(1)] = ja_bert271 272 en_bert = row[8]273 en_bert_padded[i, :, : en_bert.size(1)] = en_bert274 275 style_vec[i, :] = row[9]276 277 return (278 text_padded,279 text_lengths,280 spec_padded,281 spec_lengths,282 wav_padded,283 wav_lengths,284 sid,285 tone_padded,286 language_padded,287 bert_padded,288 ja_bert_padded,289 en_bert_padded,290 style_vec,291 )292 293 294class DistributedBucketSampler(torch.utils.data.distributed.DistributedSampler):295 """296 Maintain similar input lengths in a batch.297 Length groups are specified by boundaries.298 Ex) boundaries = [b1, b2, b3] -> any batch is included either {x | b1 < length(x) <=b2} or {x | b2 < length(x) <= b3}.299 300 It removes samples which are not included in the boundaries.301 Ex) boundaries = [b1, b2, b3] -> any x s.t. length(x) <= b1 or length(x) > b3 are discarded.302 """303 304 def __init__(305 self,306 dataset,307 batch_size,308 boundaries,309 num_replicas=None,310 rank=None,311 shuffle=True,312 ):313 super().__init__(dataset, num_replicas=num_replicas, rank=rank, shuffle=shuffle)314 self.lengths = dataset.lengths315 self.batch_size = batch_size316 self.boundaries = boundaries317 318 self.buckets, self.num_samples_per_bucket = self._create_buckets()319 logger.info(f"Bucket info: {self.num_samples_per_bucket}")320 logger.info(321 f"Unused samples: {len(self.lengths) - sum(self.num_samples_per_bucket)}"322 )323 self.total_size = sum(self.num_samples_per_bucket)324 self.num_samples = self.total_size // self.num_replicas325 326 def _create_buckets(self):327 buckets = [[] for _ in range(len(self.boundaries) - 1)]328 for i in range(len(self.lengths)):329 length = self.lengths[i]330 idx_bucket = self._bisect(length)331 if idx_bucket != -1:332 buckets[idx_bucket].append(i)333 334 try:335 for i in range(len(buckets) - 1, 0, -1):336 if len(buckets[i]) == 0:337 buckets.pop(i)338 self.boundaries.pop(i + 1)339 assert all(len(bucket) > 0 for bucket in buckets)340 # When one bucket is not traversed341 except Exception as e:342 print("Bucket warning ", e)343 for i in range(len(buckets) - 1, -1, -1):344 if len(buckets[i]) == 0:345 buckets.pop(i)346 self.boundaries.pop(i + 1)347 348 num_samples_per_bucket = []349 for i in range(len(buckets)):350 len_bucket = len(buckets[i])351 total_batch_size = self.num_replicas * self.batch_size352 rem = (353 total_batch_size - (len_bucket % total_batch_size)354 ) % total_batch_size355 num_samples_per_bucket.append(len_bucket + rem)356 return buckets, num_samples_per_bucket357 358 def __iter__(self):359 # deterministically shuffle based on epoch360 g = torch.Generator()361 g.manual_seed(self.epoch)362 363 indices = []364 if self.shuffle:365 for bucket in self.buckets:366 indices.append(torch.randperm(len(bucket), generator=g).tolist())367 else:368 for bucket in self.buckets:369 indices.append(list(range(len(bucket))))370 371 batches = []372 for i in range(len(self.buckets)):373 bucket = self.buckets[i]374 len_bucket = len(bucket)375 if len_bucket == 0:376 continue377 ids_bucket = indices[i]378 num_samples_bucket = self.num_samples_per_bucket[i]379 380 # add extra samples to make it evenly divisible381 rem = num_samples_bucket - len_bucket382 ids_bucket = (383 ids_bucket384 + ids_bucket * (rem // len_bucket)385 + ids_bucket[: (rem % len_bucket)]386 )387 388 # subsample389 ids_bucket = ids_bucket[self.rank :: self.num_replicas]390 391 # batching392 for j in range(len(ids_bucket) // self.batch_size):393 batch = [394 bucket[idx]395 for idx in ids_bucket[396 j * self.batch_size : (j + 1) * self.batch_size397 ]398 ]399 batches.append(batch)400 401 if self.shuffle:402 batch_ids = torch.randperm(len(batches), generator=g).tolist()403 batches = [batches[i] for i in batch_ids]404 self.batches = batches405 406 assert len(self.batches) * self.batch_size == self.num_samples407 return iter(self.batches)408 409 def _bisect(self, x, lo=0, hi=None):410 if hi is None:411 hi = len(self.boundaries) - 1412 413 if hi > lo:414 mid = (hi + lo) // 2415 if self.boundaries[mid] < x and x <= self.boundaries[mid + 1]:416 return mid417 elif x <= self.boundaries[mid]:418 return self._bisect(x, lo, mid)419 else:420 return self._bisect(x, mid + 1, hi)421 else:422 return -1423 424 def __len__(self):425 return self.num_samples // self.batch_size426