Upload 2 files

utilities/data/add_on.py

@@ -0,0 +1,508 @@
+import os
+import torch
+import numpy as np
+import torchaudio
+import matplotlib.pyplot as plt
+
+CACHE = {
+    "get_vits_phoneme_ids": {
+        "PAD_LENGTH": 310,
+        "_pad": "_",
+        "_punctuation": ';:,.!?¡¿—…"«»“” ',
+        "_letters": "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz",
+        "_letters_ipa": "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ",
+        "_special": "♪☎☒☝⚠",
+    }
+}
+
+CACHE["get_vits_phoneme_ids"]["symbols"] = (
+    [CACHE["get_vits_phoneme_ids"]["_pad"]]
+    + list(CACHE["get_vits_phoneme_ids"]["_punctuation"])
+    + list(CACHE["get_vits_phoneme_ids"]["_letters"])
+    + list(CACHE["get_vits_phoneme_ids"]["_letters_ipa"])
+    + list(CACHE["get_vits_phoneme_ids"]["_special"])
+)
+CACHE["get_vits_phoneme_ids"]["_symbol_to_id"] = {
+    s: i for i, s in enumerate(CACHE["get_vits_phoneme_ids"]["symbols"])
+}
+
+
+def get_vits_phoneme_ids(config, dl_output, metadata):
+    pad_token_id = 0
+    pad_length = CACHE["get_vits_phoneme_ids"]["PAD_LENGTH"]
+    _symbol_to_id = CACHE["get_vits_phoneme_ids"]["_symbol_to_id"]
+
+    assert (
+        "phonemes" in metadata.keys()
+    ), "You must provide vits phonemes on using addon get_vits_phoneme_ids"
+    clean_text = metadata["phonemes"]
+    sequence = []
+
+    for symbol in clean_text:
+        symbol_id = _symbol_to_id[symbol]
+        sequence += [symbol_id]
+
+    inserted_zero_sequence = [0] * (len(sequence) * 2)
+    inserted_zero_sequence[1::2] = sequence
+    inserted_zero_sequence = inserted_zero_sequence + [0]
+
+    def _pad_phonemes(phonemes_list):
+        return phonemes_list + [pad_token_id] * (pad_length - len(phonemes_list))
+
+    return {"phoneme_idx": torch.LongTensor(_pad_phonemes(inserted_zero_sequence))}
+
+
+def get_vits_phoneme_ids_no_padding(config, dl_output, metadata):
+    pad_token_id = 0
+    pad_length = CACHE["get_vits_phoneme_ids"]["PAD_LENGTH"]
+    _symbol_to_id = CACHE["get_vits_phoneme_ids"]["_symbol_to_id"]
+
+    assert (
+        "phonemes" in metadata.keys()
+    ), "You must provide vits phonemes on using addon get_vits_phoneme_ids"
+    clean_text = metadata["phonemes"] + "⚠"
+    sequence = []
+
+    for symbol in clean_text:
+        if symbol not in _symbol_to_id.keys():
+            print("%s is not in the vocabulary. %s" % (symbol, clean_text))
+            symbol = "_"
+        symbol_id = _symbol_to_id[symbol]
+        sequence += [symbol_id]
+
+    def _pad_phonemes(phonemes_list):
+        return phonemes_list + [pad_token_id] * (pad_length - len(phonemes_list))
+
+    sequence = sequence[:pad_length]
+
+    return {"phoneme_idx": torch.LongTensor(_pad_phonemes(sequence))}
+
+
+def calculate_relative_bandwidth(config, dl_output, metadata):
+    assert "stft" in dl_output.keys()
+
+    # The last dimension of the stft feature is the frequency dimension
+    freq_dimensions = dl_output["stft"].size(-1)
+
+    freq_energy_dist = torch.sum(dl_output["stft"], dim=0)
+    freq_energy_dist = torch.cumsum(freq_energy_dist, dim=0)
+    total_energy = freq_energy_dist[-1]
+
+    percentile_5th = total_energy * 0.05
+    percentile_95th = total_energy * 0.95
+
+    lower_idx = torch.argmin(torch.abs(percentile_5th - freq_energy_dist))
+    higher_idx = torch.argmin(torch.abs(percentile_95th - freq_energy_dist))
+
+    lower_idx = int((lower_idx / freq_dimensions) * 1000)
+    higher_idx = int((higher_idx / freq_dimensions) * 1000)
+
+    return {"freq_energy_percentile": torch.LongTensor([lower_idx, higher_idx])}
+
+
+def calculate_mel_spec_relative_bandwidth_as_extra_channel(config, dl_output, metadata):
+    assert "stft" in dl_output.keys()
+    linear_mel_spec = torch.exp(torch.clip(dl_output["log_mel_spec"], max=10))
+
+    # The last dimension of the stft feature is the frequency dimension
+    freq_dimensions = linear_mel_spec.size(-1)
+    freq_energy_dist = torch.sum(linear_mel_spec, dim=0)
+    freq_energy_dist = torch.cumsum(freq_energy_dist, dim=0)
+    total_energy = freq_energy_dist[-1]
+
+    percentile_5th = total_energy * 0.05
+    percentile_95th = total_energy * 0.95
+
+    lower_idx = torch.argmin(torch.abs(percentile_5th - freq_energy_dist))
+    higher_idx = torch.argmin(torch.abs(percentile_95th - freq_energy_dist))
+
+    latent_t_size = config["model"]["params"]["latent_t_size"]
+    latent_f_size = config["model"]["params"]["latent_f_size"]
+
+    lower_idx = int(latent_f_size * float((lower_idx / freq_dimensions)))
+    higher_idx = int(latent_f_size * float((higher_idx / freq_dimensions)))
+
+    bandwidth_condition = torch.zeros((latent_t_size, latent_f_size))
+    bandwidth_condition[:, lower_idx:higher_idx] += 1.0
+
+    return {
+        "mel_spec_bandwidth_cond_extra_channel": bandwidth_condition,
+        "freq_energy_percentile": torch.LongTensor([lower_idx, higher_idx]),
+    }
+
+
+def waveform_rs_48k(config, dl_output, metadata):
+    waveform = dl_output["waveform"]  # [1, samples]
+    sampling_rate = dl_output["sampling_rate"]
+
+    if sampling_rate != 48000:
+        waveform_48k = torchaudio.functional.resample(
+            waveform, orig_freq=sampling_rate, new_freq=48000
+        )
+    else:
+        waveform_48k = waveform
+
+    return {"waveform_48k": waveform_48k}
+
+
+def extract_vits_phoneme_and_flant5_text(config, dl_output, metadata):
+    assert (
+        "phoneme" not in metadata.keys()
+    ), "The metadata of speech you use seems belong to fastspeech. Please check dataset_root.json"
+
+    if "phonemes" in metadata.keys():
+        new_item = get_vits_phoneme_ids_no_padding(config, dl_output, metadata)
+        new_item["text"] = ""  # We assume TTS data does not have text description
+    else:
+        fake_metadata = {"phonemes": ""}  # Add empty phoneme sequence
+        new_item = get_vits_phoneme_ids_no_padding(config, dl_output, fake_metadata)
+
+    return new_item
+
+
+def extract_fs2_phoneme_and_flant5_text(config, dl_output, metadata):
+    if "phoneme" in metadata.keys():
+        new_item = extract_fs2_phoneme_g2p_en_feature(config, dl_output, metadata)
+        new_item["text"] = ""
+    else:
+        fake_metadata = {"phoneme": []}
+        new_item = extract_fs2_phoneme_g2p_en_feature(config, dl_output, fake_metadata)
+    return new_item
+
+
+def extract_fs2_phoneme_g2p_en_feature(config, dl_output, metadata):
+    PAD_LENGTH = 135
+
+    phonemes_lookup_dict = {
+        "K": 0,
+        "IH2": 1,
+        "NG": 2,
+        "OW2": 3,
+        "AH2": 4,
+        "F": 5,
+        "AE0": 6,
+        "IY0": 7,
+        "SH": 8,
+        "G": 9,
+        "W": 10,
+        "UW1": 11,
+        "AO2": 12,
+        "AW2": 13,
+        "UW0": 14,
+        "EY2": 15,
+        "UW2": 16,
+        "AE2": 17,
+        "IH0": 18,
+        "P": 19,
+        "D": 20,
+        "ER1": 21,
+        "AA1": 22,
+        "EH0": 23,
+        "UH1": 24,
+        "N": 25,
+        "V": 26,
+        "AY1": 27,
+        "EY1": 28,
+        "UH2": 29,
+        "EH1": 30,
+        "L": 31,
+        "AA2": 32,
+        "R": 33,
+        "OY1": 34,
+        "Y": 35,
+        "ER2": 36,
+        "S": 37,
+        "AE1": 38,
+        "AH1": 39,
+        "JH": 40,
+        "ER0": 41,
+        "EH2": 42,
+        "IY2": 43,
+        "OY2": 44,
+        "AW1": 45,
+        "IH1": 46,
+        "IY1": 47,
+        "OW0": 48,
+        "AO0": 49,
+        "AY0": 50,
+        "EY0": 51,
+        "AY2": 52,
+        "UH0": 53,
+        "M": 54,
+        "TH": 55,
+        "T": 56,
+        "OY0": 57,
+        "AW0": 58,
+        "DH": 59,
+        "Z": 60,
+        "spn": 61,
+        "AH0": 62,
+        "sp": 63,
+        "AO1": 64,
+        "OW1": 65,
+        "ZH": 66,
+        "B": 67,
+        "AA0": 68,
+        "CH": 69,
+        "HH": 70,
+    }
+    pad_token_id = len(phonemes_lookup_dict.keys())
+
+    assert (
+        "phoneme" in metadata.keys()
+    ), "The dataloader add-on extract_phoneme_g2p_en_feature will output phoneme id, which is not specified in your dataset"
+
+    phonemes = [
+        phonemes_lookup_dict[x]
+        for x in metadata["phoneme"]
+        if (x in phonemes_lookup_dict.keys())
+    ]
+
+    if (len(phonemes) / PAD_LENGTH) > 5:
+        print(
+            "Warning: Phonemes length is too long and is truncated too much! %s"
+            % metadata
+        )
+
+    phonemes = phonemes[:PAD_LENGTH]
+
+    def _pad_phonemes(phonemes_list):
+        return phonemes_list + [pad_token_id] * (PAD_LENGTH - len(phonemes_list))
+
+    return {"phoneme_idx": torch.LongTensor(_pad_phonemes(phonemes))}
+
+
+def extract_phoneme_g2p_en_feature(config, dl_output, metadata):
+    PAD_LENGTH = 250
+
+    phonemes_lookup_dict = {
+        " ": 0,
+        "AA": 1,
+        "AE": 2,
+        "AH": 3,
+        "AO": 4,
+        "AW": 5,
+        "AY": 6,
+        "B": 7,
+        "CH": 8,
+        "D": 9,
+        "DH": 10,
+        "EH": 11,
+        "ER": 12,
+        "EY": 13,
+        "F": 14,
+        "G": 15,
+        "HH": 16,
+        "IH": 17,
+        "IY": 18,
+        "JH": 19,
+        "K": 20,
+        "L": 21,
+        "M": 22,
+        "N": 23,
+        "NG": 24,
+        "OW": 25,
+        "OY": 26,
+        "P": 27,
+        "R": 28,
+        "S": 29,
+        "SH": 30,
+        "T": 31,
+        "TH": 32,
+        "UH": 33,
+        "UW": 34,
+        "V": 35,
+        "W": 36,
+        "Y": 37,
+        "Z": 38,
+        "ZH": 39,
+    }
+    pad_token_id = len(phonemes_lookup_dict.keys())
+
+    assert (
+        "phoneme" in metadata.keys()
+    ), "The dataloader add-on extract_phoneme_g2p_en_feature will output phoneme id, which is not specified in your dataset"
+    phonemes = [
+        phonemes_lookup_dict[x]
+        for x in metadata["phoneme"]
+        if (x in phonemes_lookup_dict.keys())
+    ]
+
+    if (len(phonemes) / PAD_LENGTH) > 5:
+        print(
+            "Warning: Phonemes length is too long and is truncated too much! %s"
+            % metadata
+        )
+
+    phonemes = phonemes[:PAD_LENGTH]
+
+    def _pad_phonemes(phonemes_list):
+        return phonemes_list + [pad_token_id] * (PAD_LENGTH - len(phonemes_list))
+
+    return {"phoneme_idx": torch.LongTensor(_pad_phonemes(phonemes))}
+
+
+def extract_kaldi_fbank_feature(config, dl_output, metadata):
+    norm_mean = -4.2677393
+    norm_std = 4.5689974
+
+    waveform = dl_output["waveform"]  # [1, samples]
+    sampling_rate = dl_output["sampling_rate"]
+    log_mel_spec_hifigan = dl_output["log_mel_spec"]
+
+    if sampling_rate != 16000:
+        waveform_16k = torchaudio.functional.resample(
+            waveform, orig_freq=sampling_rate, new_freq=16000
+        )
+    else:
+        waveform_16k = waveform
+
+    waveform_16k = waveform_16k - waveform_16k.mean()
+    fbank = torchaudio.compliance.kaldi.fbank(
+        waveform_16k,
+        htk_compat=True,
+        sample_frequency=16000,
+        use_energy=False,
+        window_type="hanning",
+        num_mel_bins=128,
+        dither=0.0,
+        frame_shift=10,
+    )
+
+    TARGET_LEN = log_mel_spec_hifigan.size(0)
+
+    # cut and pad
+    n_frames = fbank.shape[0]
+    p = TARGET_LEN - n_frames
+    if p > 0:
+        m = torch.nn.ZeroPad2d((0, 0, 0, p))
+        fbank = m(fbank)
+    elif p < 0:
+        fbank = fbank[:TARGET_LEN, :]
+
+    fbank = (fbank - norm_mean) / (norm_std * 2)
+
+    return {"ta_kaldi_fbank": fbank}  # [1024, 128]
+
+
+def extract_kaldi_fbank_feature_32k(config, dl_output, metadata):
+    norm_mean = -4.2677393
+    norm_std = 4.5689974
+
+    waveform = dl_output["waveform"]  # [1, samples]
+    sampling_rate = dl_output["sampling_rate"]
+    log_mel_spec_hifigan = dl_output["log_mel_spec"]
+
+    if sampling_rate != 32000:
+        waveform_32k = torchaudio.functional.resample(
+            waveform, orig_freq=sampling_rate, new_freq=32000
+        )
+    else:
+        waveform_32k = waveform
+
+    waveform_32k = waveform_32k - waveform_32k.mean()
+    fbank = torchaudio.compliance.kaldi.fbank(
+        waveform_32k,
+        htk_compat=True,
+        sample_frequency=32000,
+        use_energy=False,
+        window_type="hanning",
+        num_mel_bins=128,
+        dither=0.0,
+        frame_shift=10,
+    )
+
+    TARGET_LEN = log_mel_spec_hifigan.size(0)
+
+    # cut and pad
+    n_frames = fbank.shape[0]
+    p = TARGET_LEN - n_frames
+    if p > 0:
+        m = torch.nn.ZeroPad2d((0, 0, 0, p))
+        fbank = m(fbank)
+    elif p < 0:
+        fbank = fbank[:TARGET_LEN, :]
+
+    fbank = (fbank - norm_mean) / (norm_std * 2)
+
+    return {"ta_kaldi_fbank": fbank}  # [1024, 128]
+
+
+# Use the beat and downbeat information as music conditions
+def extract_drum_beat(config, dl_output, metadata):
+    def visualization(conditional_signal, mel_spectrogram, filename):
+        import soundfile as sf
+
+        sf.write(
+            os.path.basename(dl_output["fname"]),
+            np.array(dl_output["waveform"])[0],
+            dl_output["sampling_rate"],
+        )
+        plt.figure(figsize=(10, 10))
+
+        plt.subplot(211)
+        plt.imshow(np.array(conditional_signal).T, aspect="auto")
+        plt.title("Conditional Signal")
+
+        plt.subplot(212)
+        plt.imshow(np.array(mel_spectrogram).T, aspect="auto")
+        plt.title("Mel Spectrogram")
+
+        plt.savefig(filename)
+        plt.close()
+
+    assert "sample_rate" in metadata and "beat" in metadata and "downbeat" in metadata
+
+    sampling_rate = metadata["sample_rate"]
+    duration = dl_output["duration"]
+    # The dataloader segment length before performing torch resampling
+    original_segment_length_before_resample = int(sampling_rate * duration)
+
+    random_start_sample = int(dl_output["random_start_sample_in_original_audio_file"])
+
+    # The sample idx for beat and downbeat, relatively to the segmented audio
+    beat = [
+        x - random_start_sample
+        for x in metadata["beat"]
+        if (
+            x - random_start_sample >= 0
+            and x - random_start_sample <= original_segment_length_before_resample
+        )
+    ]
+    downbeat = [
+        x - random_start_sample
+        for x in metadata["downbeat"]
+        if (
+            x - random_start_sample >= 0
+            and x - random_start_sample <= original_segment_length_before_resample
+        )
+    ]
+
+    latent_shape = (
+        config["model"]["params"]["latent_t_size"],
+        config["model"]["params"]["latent_f_size"],
+    )
+    conditional_signal = torch.zeros(latent_shape)
+
+    # beat: -0.5
+    # downbeat: +1.0
+    # 0: none; -0.5: beat; 1.0: downbeat; 0.5: downbeat+beat
+    for each in beat:
+        beat_index = int(
+            (each / original_segment_length_before_resample) * latent_shape[0]
+        )
+        beat_index = min(beat_index, conditional_signal.size(0) - 1)
+
+        conditional_signal[beat_index, :] -= 0.5
+
+    for each in downbeat:
+        beat_index = int(
+            (each / original_segment_length_before_resample) * latent_shape[0]
+        )
+        beat_index = min(beat_index, conditional_signal.size(0) - 1)
+
+        conditional_signal[beat_index, :] += 1.0
+
+    # visualization(conditional_signal, dl_output["log_mel_spec"], filename = os.path.basename(dl_output["fname"])+".png")
+
+    return {"cond_beat_downbeat": conditional_signal}

utilities/data/dataset.py

@@ -0,0 +1,518 @@
+import os
+import pandas as pd
+
+import audiosr.utilities.audio as Audio
+from audiosr.utilities.tools import load_json
+
+import random
+from torch.utils.data import Dataset
+import torch.nn.functional
+import torch
+import numpy as np
+import torchaudio
+
+
+class AudioDataset(Dataset):
+    def __init__(
+        self,
+        config=None,
+        split="train",
+        waveform_only=False,
+        add_ons=[],
+        dataset_json_path=None,  #
+    ):
+        """
+        Dataset that manages audio recordings
+        :param audio_conf: Dictionary containing the audio loading and preprocessing settings
+        :param dataset_json_file
+        """
+        self.config = config
+        self.split = split
+        self.pad_wav_start_sample = 0  # If none, random choose
+        self.trim_wav = False
+        self.waveform_only = waveform_only
+        self.add_ons = [eval(x) for x in add_ons]
+        print("Add-ons:", self.add_ons)
+
+        self.build_setting_parameters()
+
+        # For an external dataset
+        if dataset_json_path is not None:
+            assert type(dataset_json_path) == str
+            print("Load metadata from %s" % dataset_json_path)
+            self.data = load_json(dataset_json_path)["data"]
+            self.id2label, self.index_dict, self.num2label = {}, {}, {}
+        else:
+            self.metadata_root = load_json(self.config["metadata_root"])
+            self.dataset_name = self.config["data"][self.split]
+            assert split in self.config["data"].keys(), (
+                "The dataset split %s you specified is not present in the config. You can choose from %s"
+                % (split, self.config["data"].keys())
+            )
+            self.build_dataset()
+            self.build_id_to_label()
+
+        self.build_dsp()
+        self.label_num = len(self.index_dict)
+        print("Dataset initialize finished")
+
+    def __getitem__(self, index):
+        (
+            fname,
+            waveform,
+            stft,
+            log_mel_spec,
+            label_vector,  # the one-hot representation of the audio class
+            # the metadata of the sampled audio file and the mixup audio file (if exist)
+            (datum, mix_datum),
+            random_start,
+        ) = self.feature_extraction(index)
+        text = self.get_sample_text_caption(datum, mix_datum, label_vector)
+
+        data = {
+            "text": text,  # list
+            "fname": self.text_to_filename(text)
+            if (len(fname) == 0)
+            else fname,  # list
+            # tensor, [batchsize, class_num]
+            "label_vector": "" if (label_vector is None) else label_vector.float(),
+            # tensor, [batchsize, 1, samples_num]
+            "waveform": "" if (waveform is None) else waveform.float(),
+            # tensor, [batchsize, t-steps, f-bins]
+            "stft": "" if (stft is None) else stft.float(),
+            # tensor, [batchsize, t-steps, mel-bins]
+            "log_mel_spec": "" if (log_mel_spec is None) else log_mel_spec.float(),
+            "duration": self.duration,
+            "sampling_rate": self.sampling_rate,
+            "random_start_sample_in_original_audio_file": random_start,
+        }
+
+        for add_on in self.add_ons:
+            data.update(add_on(self.config, data, self.data[index]))
+
+        if data["text"] is None:
+            print("Warning: The model return None on key text", fname)
+            data["text"] = ""
+
+        return data
+
+    def text_to_filename(self, text):
+        return text.replace(" ", "_").replace("'", "_").replace('"', "_")
+
+    def get_dataset_root_path(self, dataset):
+        assert dataset in self.metadata_root.keys()
+        return self.metadata_root[dataset]
+
+    def get_dataset_metadata_path(self, dataset, key):
+        # key: train, test, val, class_label_indices
+        try:
+            if dataset in self.metadata_root["metadata"]["path"].keys():
+                return self.metadata_root["metadata"]["path"][dataset][key]
+        except:
+            raise ValueError(
+                'Dataset %s does not metadata "%s" specified' % (dataset, key)
+            )
+            # return None
+
+    def __len__(self):
+        return len(self.data)
+
+    def feature_extraction(self, index):
+        if index > len(self.data) - 1:
+            print(
+                "The index of the dataloader is out of range: %s/%s"
+                % (index, len(self.data))
+            )
+            index = random.randint(0, len(self.data) - 1)
+
+        # Read wave file and extract feature
+        while True:
+            try:
+                label_indices = np.zeros(self.label_num, dtype=np.float32)
+                datum = self.data[index]
+                (
+                    log_mel_spec,
+                    stft,
+                    mix_lambda,
+                    waveform,
+                    random_start,
+                ) = self.read_audio_file(datum["wav"])
+                mix_datum = None
+                if self.label_num > 0 and "labels" in datum.keys():
+                    for label_str in datum["labels"].split(","):
+                        label_indices[int(self.index_dict[label_str])] = 1.0
+
+                # If the key "label" is not in the metadata, return all zero vector
+                label_indices = torch.FloatTensor(label_indices)
+                break
+            except Exception as e:
+                index = (index + 1) % len(self.data)
+                print(
+                    "Error encounter during audio feature extraction: ", e, datum["wav"]
+                )
+                continue
+
+        # The filename of the wav file
+        fname = datum["wav"]
+        # t_step = log_mel_spec.size(0)
+        # waveform = torch.FloatTensor(waveform[..., : int(self.hopsize * t_step)])
+        waveform = torch.FloatTensor(waveform)
+
+        return (
+            fname,
+            waveform,
+            stft,
+            log_mel_spec,
+            label_indices,
+            (datum, mix_datum),
+            random_start,
+        )
+
+    # def augmentation(self, log_mel_spec):
+    #     assert torch.min(log_mel_spec) < 0
+    #     log_mel_spec = log_mel_spec.exp()
+
+    #     log_mel_spec = torch.transpose(log_mel_spec, 0, 1)
+    #     # this is just to satisfy new torchaudio version.
+    #     log_mel_spec = log_mel_spec.unsqueeze(0)
+    #     if self.freqm != 0:
+    #         log_mel_spec = self.frequency_masking(log_mel_spec, self.freqm)
+    #     if self.timem != 0:
+    #         log_mel_spec = self.time_masking(
+    #             log_mel_spec, self.timem)  # self.timem=0
+
+    #     log_mel_spec = (log_mel_spec + 1e-7).log()
+    #     # squeeze back
+    #     log_mel_spec = log_mel_spec.squeeze(0)
+    #     log_mel_spec = torch.transpose(log_mel_spec, 0, 1)
+    #     return log_mel_spec
+
+    def build_setting_parameters(self):
+        # Read from the json config
+        self.melbins = self.config["preprocessing"]["mel"]["n_mel_channels"]
+        # self.freqm = self.config["preprocessing"]["mel"]["freqm"]
+        # self.timem = self.config["preprocessing"]["mel"]["timem"]
+        self.sampling_rate = self.config["preprocessing"]["audio"]["sampling_rate"]
+        self.hopsize = self.config["preprocessing"]["stft"]["hop_length"]
+        self.duration = self.config["preprocessing"]["audio"]["duration"]
+        self.target_length = int(self.duration * self.sampling_rate / self.hopsize)
+
+        self.mixup = self.config["augmentation"]["mixup"]
+
+        # Calculate parameter derivations
+        # self.waveform_sample_length = int(self.target_length * self.hopsize)
+
+        # if (self.config["balance_sampling_weight"]):
+        #     self.samples_weight = np.loadtxt(
+        #         self.config["balance_sampling_weight"], delimiter=","
+        #     )
+
+        if "train" not in self.split:
+            self.mixup = 0.0
+            # self.freqm = 0
+            # self.timem = 0
+
+    def _relative_path_to_absolute_path(self, metadata, dataset_name):
+        root_path = self.get_dataset_root_path(dataset_name)
+        for i in range(len(metadata["data"])):
+            assert "wav" in metadata["data"][i].keys(), metadata["data"][i]
+            assert metadata["data"][i]["wav"][0] != "/", (
+                "The dataset metadata should only contain relative path to the audio file: "
+                + str(metadata["data"][i]["wav"])
+            )
+            metadata["data"][i]["wav"] = os.path.join(
+                root_path, metadata["data"][i]["wav"]
+            )
+        return metadata
+
+    def build_dataset(self):
+        self.data = []
+        print("Build dataset split %s from %s" % (self.split, self.dataset_name))
+        if type(self.dataset_name) is str:
+            data_json = load_json(
+                self.get_dataset_metadata_path(self.dataset_name, key=self.split)
+            )
+            data_json = self._relative_path_to_absolute_path(
+                data_json, self.dataset_name
+            )
+            self.data = data_json["data"]
+        elif type(self.dataset_name) is list:
+            for dataset_name in self.dataset_name:
+                data_json = load_json(
+                    self.get_dataset_metadata_path(dataset_name, key=self.split)
+                )
+                data_json = self._relative_path_to_absolute_path(
+                    data_json, dataset_name
+                )
+                self.data += data_json["data"]
+        else:
+            raise Exception("Invalid data format")
+        print("Data size: {}".format(len(self.data)))
+
+    def build_dsp(self):
+        self.STFT = Audio.stft.TacotronSTFT(
+            self.config["preprocessing"]["stft"]["filter_length"],
+            self.config["preprocessing"]["stft"]["hop_length"],
+            self.config["preprocessing"]["stft"]["win_length"],
+            self.config["preprocessing"]["mel"]["n_mel_channels"],
+            self.config["preprocessing"]["audio"]["sampling_rate"],
+            self.config["preprocessing"]["mel"]["mel_fmin"],
+            self.config["preprocessing"]["mel"]["mel_fmax"],
+        )
+        # self.stft_transform = torchaudio.transforms.Spectrogram(
+        #     n_fft=1024, hop_length=160
+        # )
+        # self.melscale_transform = torchaudio.transforms.MelScale(
+        #     sample_rate=16000, n_stft=1024 // 2 + 1, n_mels=64
+        # )
+
+    def build_id_to_label(self):
+        id2label = {}
+        id2num = {}
+        num2label = {}
+        class_label_indices_path = self.get_dataset_metadata_path(
+            dataset=self.config["data"]["class_label_indices"],
+            key="class_label_indices",
+        )
+        if class_label_indices_path is not None:
+            df = pd.read_csv(class_label_indices_path)
+            for _, row in df.iterrows():
+                index, mid, display_name = row["index"], row["mid"], row["display_name"]
+                id2label[mid] = display_name
+                id2num[mid] = index
+                num2label[index] = display_name
+            self.id2label, self.index_dict, self.num2label = id2label, id2num, num2label
+        else:
+            self.id2label, self.index_dict, self.num2label = {}, {}, {}
+
+    def resample(self, waveform, sr):
+        waveform = torchaudio.functional.resample(waveform, sr, self.sampling_rate)
+        # waveform = librosa.resample(waveform, sr, self.sampling_rate)
+        return waveform
+
+        # if sr == 16000:
+        #     return waveform
+        # if sr == 32000 and self.sampling_rate == 16000:
+        #     waveform = waveform[::2]
+        #     return waveform
+        # if sr == 48000 and self.sampling_rate == 16000:
+        #     waveform = waveform[::3]
+        #     return waveform
+        # else:
+        #     raise ValueError(
+        #         "We currently only support 16k audio generation. You need to resample you audio file to 16k, 32k, or 48k: %s, %s"
+        #         % (sr, self.sampling_rate)
+        #     )
+
+    def normalize_wav(self, waveform):
+        waveform = waveform - np.mean(waveform)
+        waveform = waveform / (np.max(np.abs(waveform)) + 1e-8)
+        return waveform * 0.5  # Manually limit the maximum amplitude into 0.5
+
+    def random_segment_wav(self, waveform, target_length):
+        waveform_length = waveform.shape[-1]
+        assert waveform_length > 100, "Waveform is too short, %s" % waveform_length
+
+        # Too short
+        if (waveform_length - target_length) <= 0:
+            return waveform, 0
+
+        random_start = int(self.random_uniform(0, waveform_length - target_length))
+        return waveform[:, random_start : random_start + target_length], random_start
+
+    def pad_wav(self, waveform, target_length):
+        waveform_length = waveform.shape[-1]
+        assert waveform_length > 100, "Waveform is too short, %s" % waveform_length
+
+        if waveform_length == target_length:
+            return waveform
+
+        # Pad
+        temp_wav = np.zeros((1, target_length), dtype=np.float32)
+        if self.pad_wav_start_sample is None:
+            rand_start = int(self.random_uniform(0, target_length - waveform_length))
+        else:
+            rand_start = 0
+
+        temp_wav[:, rand_start : rand_start + waveform_length] = waveform
+        return temp_wav
+
+    def trim_wav(self, waveform):
+        if np.max(np.abs(waveform)) < 0.0001:
+            return waveform
+
+        def detect_leading_silence(waveform, threshold=0.0001):
+            chunk_size = 1000
+            waveform_length = waveform.shape[0]
+            start = 0
+            while start + chunk_size < waveform_length:
+                if np.max(np.abs(waveform[start : start + chunk_size])) < threshold:
+                    start += chunk_size
+                else:
+                    break
+            return start
+
+        def detect_ending_silence(waveform, threshold=0.0001):
+            chunk_size = 1000
+            waveform_length = waveform.shape[0]
+            start = waveform_length
+            while start - chunk_size > 0:
+                if np.max(np.abs(waveform[start - chunk_size : start])) < threshold:
+                    start -= chunk_size
+                else:
+                    break
+            if start == waveform_length:
+                return start
+            else:
+                return start + chunk_size
+
+        start = detect_leading_silence(waveform)
+        end = detect_ending_silence(waveform)
+
+        return waveform[start:end]
+
+    def read_wav_file(self, filename):
+        # waveform, sr = librosa.load(filename, sr=None, mono=True) # 4 times slower
+        waveform, sr = torchaudio.load(filename)
+
+        waveform, random_start = self.random_segment_wav(
+            waveform, target_length=int(sr * self.duration)
+        )
+
+        waveform = self.resample(waveform, sr)
+        # random_start = int(random_start * (self.sampling_rate / sr))
+
+        waveform = waveform.numpy()[0, ...]
+
+        waveform = self.normalize_wav(waveform)
+
+        if self.trim_wav:
+            waveform = self.trim_wav(waveform)
+
+        waveform = waveform[None, ...]
+        waveform = self.pad_wav(
+            waveform, target_length=int(self.sampling_rate * self.duration)
+        )
+        return waveform, random_start
+
+    def mix_two_waveforms(self, waveform1, waveform2):
+        mix_lambda = np.random.beta(5, 5)
+        mix_waveform = mix_lambda * waveform1 + (1 - mix_lambda) * waveform2
+        return self.normalize_wav(mix_waveform), mix_lambda
+
+    def read_audio_file(self, filename, filename2=None):
+        if os.path.exists(filename):
+            waveform, random_start = self.read_wav_file(filename)
+        else:
+            print(
+                'Warning [dataset.py]: The wav path "',
+                filename,
+                '" is not find in the metadata. Use empty waveform instead.',
+            )
+            target_length = int(self.sampling_rate * self.duration)
+            waveform = torch.zeros((1, target_length))
+            random_start = 0
+
+        mix_lambda = 0.0
+        # log_mel_spec, stft = self.wav_feature_extraction_torchaudio(waveform) # this line is faster, but this implementation is not aligned with HiFi-GAN
+        if not self.waveform_only:
+            log_mel_spec, stft = self.wav_feature_extraction(waveform)
+        else:
+            # Load waveform data only
+            # Use zero array to keep the format unified
+            log_mel_spec, stft = None, None
+
+        return log_mel_spec, stft, mix_lambda, waveform, random_start
+
+    def get_sample_text_caption(self, datum, mix_datum, label_indices):
+        text = self.label_indices_to_text(datum, label_indices)
+        if mix_datum is not None:
+            text += " " + self.label_indices_to_text(mix_datum, label_indices)
+        return text
+
+    # This one is significantly slower than "wav_feature_extraction_torchaudio" if num_worker > 1
+    def wav_feature_extraction(self, waveform):
+        waveform = waveform[0, ...]
+        waveform = torch.FloatTensor(waveform)
+
+        log_mel_spec, stft, energy = Audio.tools.get_mel_from_wav(waveform, self.STFT)
+
+        log_mel_spec = torch.FloatTensor(log_mel_spec.T)
+        stft = torch.FloatTensor(stft.T)
+
+        log_mel_spec, stft = self.pad_spec(log_mel_spec), self.pad_spec(stft)
+        return log_mel_spec, stft
+
+    # @profile
+    # def wav_feature_extraction_torchaudio(self, waveform):
+    #     waveform = waveform[0, ...]
+    #     waveform = torch.FloatTensor(waveform)
+
+    #     stft = self.stft_transform(waveform)
+    #     mel_spec = self.melscale_transform(stft)
+    #     log_mel_spec = torch.log(mel_spec + 1e-7)
+
+    #     log_mel_spec = torch.FloatTensor(log_mel_spec.T)
+    #     stft = torch.FloatTensor(stft.T)
+
+    #     log_mel_spec, stft = self.pad_spec(log_mel_spec), self.pad_spec(stft)
+    #     return log_mel_spec, stft
+
+    def pad_spec(self, log_mel_spec):
+        n_frames = log_mel_spec.shape[0]
+        p = self.target_length - n_frames
+        # cut and pad
+        if p > 0:
+            m = torch.nn.ZeroPad2d((0, 0, 0, p))
+            log_mel_spec = m(log_mel_spec)
+        elif p < 0:
+            log_mel_spec = log_mel_spec[0 : self.target_length, :]
+
+        if log_mel_spec.size(-1) % 2 != 0:
+            log_mel_spec = log_mel_spec[..., :-1]
+
+        return log_mel_spec
+
+    def _read_datum_caption(self, datum):
+        caption_keys = [x for x in datum.keys() if ("caption" in x)]
+        random_index = torch.randint(0, len(caption_keys), (1,))[0].item()
+        return datum[caption_keys[random_index]]
+
+    def _is_contain_caption(self, datum):
+        caption_keys = [x for x in datum.keys() if ("caption" in x)]
+        return len(caption_keys) > 0
+
+    def label_indices_to_text(self, datum, label_indices):
+        if self._is_contain_caption(datum):
+            return self._read_datum_caption(datum)
+        elif "label" in datum.keys():
+            name_indices = torch.where(label_indices > 0.1)[0]
+            # description_header = "This audio contains the sound of "
+            description_header = ""
+            labels = ""
+            for id, each in enumerate(name_indices):
+                if id == len(name_indices) - 1:
+                    labels += "%s." % self.num2label[int(each)]
+                else:
+                    labels += "%s, " % self.num2label[int(each)]
+            return description_header + labels
+        else:
+            return ""  # TODO, if both label and caption are not provided, return empty string
+
+    def random_uniform(self, start, end):
+        val = torch.rand(1).item()
+        return start + (end - start) * val
+
+    def frequency_masking(self, log_mel_spec, freqm):
+        bs, freq, tsteps = log_mel_spec.size()
+        mask_len = int(self.random_uniform(freqm // 8, freqm))
+        mask_start = int(self.random_uniform(start=0, end=freq - mask_len))
+        log_mel_spec[:, mask_start : mask_start + mask_len, :] *= 0.0
+        return log_mel_spec
+
+    def time_masking(self, log_mel_spec, timem):
+        bs, freq, tsteps = log_mel_spec.size()
+        mask_len = int(self.random_uniform(timem // 8, timem))
+        mask_start = int(self.random_uniform(start=0, end=tsteps - mask_len))
+        log_mel_spec[:, :, mask_start : mask_start + mask_len] *= 0.0
+        return log_mel_spec