import os import statistics import torch from torch.utils.data import Dataset from tqdm import tqdm from Modules.Aligner.Aligner import Aligner from Modules.Aligner.CodecAlignerDataset import CodecAlignerDataset from Modules.ToucanTTS.DurationCalculator import DurationCalculator from Modules.ToucanTTS.EnergyCalculator import EnergyCalculator from Modules.ToucanTTS.PitchCalculator import Parselmouth from Preprocessing.AudioPreprocessor import AudioPreprocessor from Preprocessing.EnCodecAudioPreprocessor import CodecAudioPreprocessor from Preprocessing.TextFrontend import get_language_id from Preprocessing.articulatory_features import get_feature_to_index_lookup class TTSDataset(Dataset): def __init__(self, path_to_transcript_dict, acoustic_checkpoint_path, cache_dir, lang, loading_processes=os.cpu_count() if os.cpu_count() is not None else 10, min_len_in_seconds=1, max_len_in_seconds=15, device=torch.device("cpu"), rebuild_cache=False, ctc_selection=True, save_imgs=False, gpu_count=1, rank=0, annotate_silences=False): self.cache_dir = cache_dir self.device = device self.pttd = path_to_transcript_dict os.makedirs(cache_dir, exist_ok=True) if not os.path.exists(os.path.join(cache_dir, "tts_train_cache.pt")) or rebuild_cache: self._build_dataset_cache(path_to_transcript_dict=path_to_transcript_dict, acoustic_checkpoint_path=acoustic_checkpoint_path, cache_dir=cache_dir, lang=lang, loading_processes=loading_processes, min_len_in_seconds=min_len_in_seconds, max_len_in_seconds=max_len_in_seconds, device=device, rebuild_cache=rebuild_cache, ctc_selection=ctc_selection, save_imgs=save_imgs, gpu_count=gpu_count, rank=rank, annotate_silences=annotate_silences) self.cache_dir = cache_dir self.gpu_count = gpu_count self.rank = rank self.language_id = get_language_id(lang) self.datapoints = torch.load(os.path.join(self.cache_dir, "tts_train_cache.pt"), map_location='cpu') if self.gpu_count > 1: # we only keep a chunk of the dataset in memory to avoid redundancy. Which chunk, we figure out using the rank. while len(self.datapoints) % self.gpu_count != 0: self.datapoints.pop(-1) # a bit unfortunate, but if you're using multiple GPUs, you probably have a ton of datapoints anyway. chunksize = int(len(self.datapoints) / self.gpu_count) self.datapoints = self.datapoints[chunksize * self.rank:chunksize * (self.rank + 1)] print(f"Loaded a TTS dataset with {len(self.datapoints)} datapoints from {cache_dir}.") def _build_dataset_cache(self, path_to_transcript_dict, acoustic_checkpoint_path, cache_dir, lang, loading_processes=os.cpu_count() if os.cpu_count() is not None else 10, min_len_in_seconds=1, max_len_in_seconds=15, device=torch.device("cpu"), rebuild_cache=False, ctc_selection=True, save_imgs=False, gpu_count=1, rank=0, annotate_silences=False): if gpu_count != 1: import sys print("Please run the feature extraction using only a single GPU. Multi-GPU is only supported for training.") sys.exit() if not os.path.exists(os.path.join(cache_dir, "aligner_train_cache.pt")) or rebuild_cache: CodecAlignerDataset(path_to_transcript_dict=path_to_transcript_dict, cache_dir=cache_dir, lang=lang, loading_processes=loading_processes, min_len_in_seconds=min_len_in_seconds, max_len_in_seconds=max_len_in_seconds, rebuild_cache=rebuild_cache, device=device) datapoints = torch.load(os.path.join(cache_dir, "aligner_train_cache.pt"), map_location='cpu') # we use the aligner dataset as basis and augment it to contain the additional information we need for tts. self.dataset, _, speaker_embeddings, filepaths = datapoints print("... building dataset cache ...") self.codec_wrapper = CodecAudioPreprocessor(input_sr=-1, device=device) self.spec_extractor_for_features = AudioPreprocessor(input_sr=16000, output_sr=16000, device=device) self.datapoints = list() self.ctc_losses = list() self.acoustic_model = Aligner() self.acoustic_model.load_state_dict(torch.load(acoustic_checkpoint_path, map_location="cpu")["asr_model"]) self.acoustic_model = self.acoustic_model.to(device) self.acoustic_model.eval() # ========================================== # actual creation of datapoints starts here # ========================================== parsel = Parselmouth(fs=16000) energy_calc = EnergyCalculator(fs=16000).to(device) self.dc = DurationCalculator() vis_dir = os.path.join(cache_dir, "duration_vis") if save_imgs: os.makedirs(os.path.join(vis_dir, "post_clean"), exist_ok=True) if annotate_silences: os.makedirs(os.path.join(vis_dir, "pre_clean"), exist_ok=True) for index in tqdm(range(len(self.dataset))): codes = self.dataset[index][1] if codes.size()[0] != 24: # no clue why this is sometimes the case codes = codes.transpose(0, 1) decoded_wave = self.codec_wrapper.indexes_to_audio(codes.int().to(device)) decoded_wave_length = torch.LongTensor([len(decoded_wave)]) features = self.spec_extractor_for_features.audio_to_mel_spec_tensor(decoded_wave, explicit_sampling_rate=16000) feature_lengths = torch.LongTensor([len(features[0])]) text = self.dataset[index][0] cached_duration, ctc_loss = self._calculate_durations(text, index, os.path.join(vis_dir, "post_clean"), features, save_imgs) cached_energy = energy_calc(input_waves=torch.tensor(decoded_wave).unsqueeze(0).to(device), input_waves_lengths=decoded_wave_length, feats_lengths=feature_lengths, text=text, durations=cached_duration.unsqueeze(0), durations_lengths=torch.LongTensor([len(cached_duration)]))[0].squeeze(0).cpu() cached_pitch = parsel(input_waves=torch.tensor(decoded_wave).unsqueeze(0), input_waves_lengths=decoded_wave_length, feats_lengths=feature_lengths, text=text, durations=cached_duration.unsqueeze(0), durations_lengths=torch.LongTensor([len(cached_duration)]))[0].squeeze(0).cpu() self.datapoints.append([text, # text tensor torch.LongTensor([len(text)]), # length of text tensor codes, # codec tensor (in index form) feature_lengths, # length of spectrogram cached_duration.cpu(), # duration cached_energy.float(), # energy cached_pitch.float(), # pitch speaker_embeddings[index], # speaker embedding, filepaths[index] # path to the associated original raw audio file ]) self.ctc_losses.append(ctc_loss) # ============================= # done with datapoint creation # ============================= if ctc_selection and len(self.datapoints) > 300: # for less than 300 datapoints, we should not throw away anything. # now we can filter out some bad datapoints based on the CTC scores we collected mean_ctc = sum(self.ctc_losses) / len(self.ctc_losses) std_dev = statistics.stdev(self.ctc_losses) threshold = mean_ctc + (std_dev * 3.5) for index in range(len(self.ctc_losses), 0, -1): if self.ctc_losses[index - 1] > threshold: self.datapoints.pop(index - 1) print(f"Removing datapoint {index - 1}, because the CTC loss is 3.5 standard deviations higher than the mean. \n ctc: {round(self.ctc_losses[index - 1], 4)} vs. mean: {round(mean_ctc, 4)}") # save to cache if len(self.datapoints) > 0: torch.save(self.datapoints, os.path.join(cache_dir, "tts_train_cache.pt")) else: import sys print("No datapoints were prepared! Exiting...") sys.exit() del self.dataset def _calculate_durations(self, text, index, vis_dir, features, save_imgs): # We deal with the word boundaries by having 2 versions of text: with and without word boundaries. # We note the index of word boundaries and insert durations of 0 afterwards text_without_word_boundaries = list() indexes_of_word_boundaries = list() for phoneme_index, vector in enumerate(text): if vector[get_feature_to_index_lookup()["word-boundary"]] == 0: text_without_word_boundaries.append(vector.numpy().tolist()) else: indexes_of_word_boundaries.append(phoneme_index) matrix_without_word_boundaries = torch.Tensor(text_without_word_boundaries) alignment_path, ctc_loss = self.acoustic_model.inference(features=features.transpose(0, 1), tokens=matrix_without_word_boundaries.to(self.device), save_img_for_debug=os.path.join(vis_dir, f"{index}.png") if save_imgs else None, return_ctc=True) cached_duration = self.dc(torch.LongTensor(alignment_path), vis=None).cpu() for index_of_word_boundary in indexes_of_word_boundaries: cached_duration = torch.cat([cached_duration[:index_of_word_boundary], torch.LongTensor([0]), # insert a 0 duration wherever there is a word boundary cached_duration[index_of_word_boundary:]]) return cached_duration, ctc_loss def __getitem__(self, index): return self.datapoints[index][0], \ self.datapoints[index][1], \ self.datapoints[index][2], \ self.datapoints[index][3], \ self.datapoints[index][4], \ self.datapoints[index][5], \ self.datapoints[index][6], \ None, \ self.language_id, \ self.datapoints[index][7] def __len__(self): return len(self.datapoints) def remove_samples(self, list_of_samples_to_remove): for remove_id in sorted(list_of_samples_to_remove, reverse=True): self.datapoints.pop(remove_id) torch.save(self.datapoints, os.path.join(self.cache_dir, "tts_train_cache.pt")) print("Dataset updated!")