Architectures/ToucanTTS/TTSDataset.py

import os
import statistics

import torch
from torch.utils.data import Dataset
from tqdm import tqdm

from Architectures.Aligner.Aligner import Aligner
from Architectures.Aligner.CodecAlignerDataset import CodecAlignerDataset
from Architectures.ToucanTTS.DurationCalculator import DurationCalculator
from Architectures.ToucanTTS.EnergyCalculator import EnergyCalculator
from Architectures.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
from Utility.utils import remove_elements


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)

        torch.hub._validate_not_a_forked_repo = lambda a, b, c: True  # torch 1.9 has a bug in the hub loading, this is a workaround
        # careful: assumes 16kHz or 8kHz audio
        silero_model, utils = torch.hub.load(repo_or_dir='snakers4/silero-vad', model='silero_vad', force_reload=False, onnx=False, verbose=False)
        (get_speech_timestamps, save_audio, read_audio, VADIterator, collect_chunks) = utils
        torch.set_grad_enabled(True)  # finding this issue was very infuriating: silero sets
        # this to false globally during model loading rather than using inference_mode or no_grad
        silero_model = silero_model.to(device)

        # ==========================================
        # 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]

            if annotate_silences:
                text = self._annotate_silences(text, get_speech_timestamps, index, vis_dir, decoded_wave, device, features, silero_model, save_imgs, decoded_wave_length)
            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 _annotate_silences(self, text, get_speech_timestamps, index, vis_dir, decoded_wave, device, features, silero_model, save_imgs, decoded_wave_length):
        """
        Takes in a text tensor and returns a text tensor with pauses added in all locations, where there are actually pauses in the speech signal. Unfortunately, this tends to make mistakes and not work quite as intended yet. I might revisit it in the future, if I see the need for extremely accurate labels for a small dataset of e.g. special data.
        """
        text_with_pauses = list()
        for phoneme_index, vector in enumerate(text):
            # We add pauses before every word boundary, and later we remove the ones that were added too much
            if vector[get_feature_to_index_lookup()["word-boundary"]] == 1:
                if text[phoneme_index - 1][get_feature_to_index_lookup()["silence"]] != 1:
                    text_with_pauses.append([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0.,
                                             0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                                             0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                                             0., 0., 0., 0., 0., 0., 0., 0.])
                text_with_pauses.append(vector.numpy().tolist())
            else:
                text_with_pauses.append(vector.numpy().tolist())
        text = torch.Tensor(text_with_pauses)

        cached_duration, _ = self._calculate_durations(text, index, os.path.join(vis_dir, "pre_clean"), features, save_imgs)

        cumsum = 0
        potential_silences = list()
        phoneme_indexes_of_silences = list()
        for phoneme_index, phone in enumerate(text):
            if phone[get_feature_to_index_lookup()["silence"]] == 1 or phone[get_feature_to_index_lookup()["end of sentence"]] == 1 or phone[get_feature_to_index_lookup()["questionmark"]] == 1 or phone[get_feature_to_index_lookup()["exclamationmark"]] == 1 or phone[get_feature_to_index_lookup()["fullstop"]] == 1:
                potential_silences.append([cumsum, cumsum + cached_duration[phoneme_index]])
                phoneme_indexes_of_silences.append(phoneme_index)
            cumsum = cumsum + cached_duration[phoneme_index]
        with torch.inference_mode():
            speech_timestamps = get_speech_timestamps(torch.Tensor(decoded_wave).to(device), silero_model, sampling_rate=16000)
        vad_silences = list()
        prev_end = 0
        for speech_segment in speech_timestamps:
            if prev_end != 0:
                vad_silences.append([prev_end, speech_segment["start"]])
            prev_end = speech_segment["end"]
        # at this point we know all the silences and we know the legal silences.
        # We have to transform them both into ratios, so we can compare them.
        # If a silence overlaps with a legal silence, it can stay.
        illegal_silences = list()
        for silence_index, silence in enumerate(potential_silences):
            illegal = True
            start = silence[0] / len(features)
            end = silence[1] / len(features)
            for legal_silence in vad_silences:
                legal_start = legal_silence[0] / decoded_wave_length
                legal_end = legal_silence[1] / decoded_wave_length
                if legal_start < start < legal_end or legal_start < end < legal_end:
                    illegal = False
                    break
            if illegal:
                # If it is an illegal silence, it is marked for removal in the original wave according to ration with real samplingrate.
                illegal_silences.append(phoneme_indexes_of_silences[silence_index])

        text = remove_elements(text, illegal_silences)  # now we have all the silences where there should be silences and none where there shouldn't be any. We have to run the aligner again with this new information.
        return text

    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!")