Spaces:
Running
on
T4
Running
on
T4
| # Copyright (c) 2023 Amphion. | |
| # | |
| # This source code is licensed under the MIT license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| import json | |
| import os | |
| import torchaudio | |
| import numpy as np | |
| import torch | |
| from utils.data_utils import * | |
| from torch.nn.utils.rnn import pad_sequence | |
| from text import text_to_sequence | |
| from text.text_token_collation import phoneIDCollation | |
| from processors.acoustic_extractor import cal_normalized_mel | |
| from models.base.base_dataset import ( | |
| BaseDataset, | |
| BaseCollator, | |
| BaseTestDataset, | |
| BaseTestCollator, | |
| ) | |
| from processors.content_extractor import ( | |
| ContentvecExtractor, | |
| WenetExtractor, | |
| WhisperExtractor, | |
| ) | |
| class TTSDataset(BaseDataset): | |
| def __init__(self, cfg, dataset, is_valid=False): | |
| """ | |
| Args: | |
| cfg: config | |
| dataset: dataset name | |
| is_valid: whether to use train or valid dataset | |
| """ | |
| assert isinstance(dataset, str) | |
| self.cfg = cfg | |
| processed_data_dir = os.path.join(cfg.preprocess.processed_dir, dataset) | |
| meta_file = cfg.preprocess.valid_file if is_valid else cfg.preprocess.train_file | |
| self.metafile_path = os.path.join(processed_data_dir, meta_file) | |
| self.metadata = self.get_metadata() | |
| """ | |
| load spk2id and utt2spk from json file | |
| spk2id: {spk1: 0, spk2: 1, ...} | |
| utt2spk: {dataset_uid: spk1, ...} | |
| """ | |
| if cfg.preprocess.use_spkid: | |
| dataset = self.metadata[0]["Dataset"] | |
| spk2id_path = os.path.join(processed_data_dir, cfg.preprocess.spk2id) | |
| with open(spk2id_path, "r") as f: | |
| self.spk2id = json.load(f) | |
| utt2spk_path = os.path.join(processed_data_dir, cfg.preprocess.utt2spk) | |
| self.utt2spk = dict() | |
| with open(utt2spk_path, "r") as f: | |
| for line in f.readlines(): | |
| utt, spk = line.strip().split("\t") | |
| self.utt2spk[utt] = spk | |
| if cfg.preprocess.use_uv: | |
| self.utt2uv_path = {} | |
| for utt_info in self.metadata: | |
| dataset = utt_info["Dataset"] | |
| uid = utt_info["Uid"] | |
| utt = "{}_{}".format(dataset, uid) | |
| self.utt2uv_path[utt] = os.path.join( | |
| cfg.preprocess.processed_dir, | |
| dataset, | |
| cfg.preprocess.uv_dir, | |
| uid + ".npy", | |
| ) | |
| if cfg.preprocess.use_frame_pitch: | |
| self.utt2frame_pitch_path = {} | |
| for utt_info in self.metadata: | |
| dataset = utt_info["Dataset"] | |
| uid = utt_info["Uid"] | |
| utt = "{}_{}".format(dataset, uid) | |
| self.utt2frame_pitch_path[utt] = os.path.join( | |
| cfg.preprocess.processed_dir, | |
| dataset, | |
| cfg.preprocess.pitch_dir, | |
| uid + ".npy", | |
| ) | |
| if cfg.preprocess.use_frame_energy: | |
| self.utt2frame_energy_path = {} | |
| for utt_info in self.metadata: | |
| dataset = utt_info["Dataset"] | |
| uid = utt_info["Uid"] | |
| utt = "{}_{}".format(dataset, uid) | |
| self.utt2frame_energy_path[utt] = os.path.join( | |
| cfg.preprocess.processed_dir, | |
| dataset, | |
| cfg.preprocess.energy_dir, | |
| uid + ".npy", | |
| ) | |
| if cfg.preprocess.use_mel: | |
| self.utt2mel_path = {} | |
| for utt_info in self.metadata: | |
| dataset = utt_info["Dataset"] | |
| uid = utt_info["Uid"] | |
| utt = "{}_{}".format(dataset, uid) | |
| self.utt2mel_path[utt] = os.path.join( | |
| cfg.preprocess.processed_dir, | |
| dataset, | |
| cfg.preprocess.mel_dir, | |
| uid + ".npy", | |
| ) | |
| if cfg.preprocess.use_linear: | |
| self.utt2linear_path = {} | |
| for utt_info in self.metadata: | |
| dataset = utt_info["Dataset"] | |
| uid = utt_info["Uid"] | |
| utt = "{}_{}".format(dataset, uid) | |
| self.utt2linear_path[utt] = os.path.join( | |
| cfg.preprocess.processed_dir, | |
| dataset, | |
| cfg.preprocess.linear_dir, | |
| uid + ".npy", | |
| ) | |
| if cfg.preprocess.use_audio: | |
| self.utt2audio_path = {} | |
| for utt_info in self.metadata: | |
| dataset = utt_info["Dataset"] | |
| uid = utt_info["Uid"] | |
| utt = "{}_{}".format(dataset, uid) | |
| if cfg.preprocess.extract_audio: | |
| self.utt2audio_path[utt] = os.path.join( | |
| cfg.preprocess.processed_dir, | |
| dataset, | |
| cfg.preprocess.audio_dir, | |
| uid + ".wav", | |
| ) | |
| else: | |
| self.utt2audio_path[utt] = utt_info["Path"] | |
| # self.utt2audio_path[utt] = os.path.join( | |
| # cfg.preprocess.processed_dir, | |
| # dataset, | |
| # cfg.preprocess.audio_dir, | |
| # uid + ".numpy", | |
| # ) | |
| elif cfg.preprocess.use_label: | |
| self.utt2label_path = {} | |
| for utt_info in self.metadata: | |
| dataset = utt_info["Dataset"] | |
| uid = utt_info["Uid"] | |
| utt = "{}_{}".format(dataset, uid) | |
| self.utt2label_path[utt] = os.path.join( | |
| cfg.preprocess.processed_dir, | |
| dataset, | |
| cfg.preprocess.label_dir, | |
| uid + ".npy", | |
| ) | |
| elif cfg.preprocess.use_one_hot: | |
| self.utt2one_hot_path = {} | |
| for utt_info in self.metadata: | |
| dataset = utt_info["Dataset"] | |
| uid = utt_info["Uid"] | |
| utt = "{}_{}".format(dataset, uid) | |
| self.utt2one_hot_path[utt] = os.path.join( | |
| cfg.preprocess.processed_dir, | |
| dataset, | |
| cfg.preprocess.one_hot_dir, | |
| uid + ".npy", | |
| ) | |
| if cfg.preprocess.use_text or cfg.preprocess.use_phone: | |
| self.utt2seq = {} | |
| for utt_info in self.metadata: | |
| dataset = utt_info["Dataset"] | |
| uid = utt_info["Uid"] | |
| utt = "{}_{}".format(dataset, uid) | |
| if cfg.preprocess.use_text: | |
| text = utt_info["Text"] | |
| sequence = text_to_sequence(text, cfg.preprocess.text_cleaners) | |
| elif cfg.preprocess.use_phone: | |
| # load phoneme squence from phone file | |
| phone_path = os.path.join( | |
| processed_data_dir, cfg.preprocess.phone_dir, uid + ".phone" | |
| ) | |
| with open(phone_path, "r") as fin: | |
| phones = fin.readlines() | |
| assert len(phones) == 1 | |
| phones = phones[0].strip() | |
| phones_seq = phones.split(" ") | |
| phon_id_collator = phoneIDCollation(cfg, dataset=dataset) | |
| sequence = phon_id_collator.get_phone_id_sequence(cfg, phones_seq) | |
| self.utt2seq[utt] = sequence | |
| def __getitem__(self, index): | |
| utt_info = self.metadata[index] | |
| dataset = utt_info["Dataset"] | |
| uid = utt_info["Uid"] | |
| utt = "{}_{}".format(dataset, uid) | |
| single_feature = dict() | |
| if self.cfg.preprocess.use_spkid: | |
| single_feature["spk_id"] = np.array( | |
| [self.spk2id[self.utt2spk[utt]]], dtype=np.int32 | |
| ) | |
| if self.cfg.preprocess.use_mel: | |
| mel = np.load(self.utt2mel_path[utt]) | |
| assert mel.shape[0] == self.cfg.preprocess.n_mel # [n_mels, T] | |
| if self.cfg.preprocess.use_min_max_norm_mel: | |
| # do mel norm | |
| mel = cal_normalized_mel(mel, utt_info["Dataset"], self.cfg.preprocess) | |
| if "target_len" not in single_feature.keys(): | |
| single_feature["target_len"] = mel.shape[1] | |
| single_feature["mel"] = mel.T # [T, n_mels] | |
| if self.cfg.preprocess.use_linear: | |
| linear = np.load(self.utt2linear_path[utt]) | |
| if "target_len" not in single_feature.keys(): | |
| single_feature["target_len"] = linear.shape[1] | |
| single_feature["linear"] = linear.T # [T, n_linear] | |
| if self.cfg.preprocess.use_frame_pitch: | |
| frame_pitch_path = self.utt2frame_pitch_path[utt] | |
| frame_pitch = np.load(frame_pitch_path) | |
| if "target_len" not in single_feature.keys(): | |
| single_feature["target_len"] = len(frame_pitch) | |
| aligned_frame_pitch = align_length( | |
| frame_pitch, single_feature["target_len"] | |
| ) | |
| single_feature["frame_pitch"] = aligned_frame_pitch | |
| if self.cfg.preprocess.use_uv: | |
| frame_uv_path = self.utt2uv_path[utt] | |
| frame_uv = np.load(frame_uv_path) | |
| aligned_frame_uv = align_length(frame_uv, single_feature["target_len"]) | |
| aligned_frame_uv = [ | |
| 0 if frame_uv else 1 for frame_uv in aligned_frame_uv | |
| ] | |
| aligned_frame_uv = np.array(aligned_frame_uv) | |
| single_feature["frame_uv"] = aligned_frame_uv | |
| if self.cfg.preprocess.use_frame_energy: | |
| frame_energy_path = self.utt2frame_energy_path[utt] | |
| frame_energy = np.load(frame_energy_path) | |
| if "target_len" not in single_feature.keys(): | |
| single_feature["target_len"] = len(frame_energy) | |
| aligned_frame_energy = align_length( | |
| frame_energy, single_feature["target_len"] | |
| ) | |
| single_feature["frame_energy"] = aligned_frame_energy | |
| if self.cfg.preprocess.use_audio: | |
| audio, sr = torchaudio.load(self.utt2audio_path[utt]) | |
| audio = audio.cpu().numpy().squeeze() | |
| single_feature["audio"] = audio | |
| single_feature["audio_len"] = audio.shape[0] | |
| if self.cfg.preprocess.use_phone or self.cfg.preprocess.use_text: | |
| single_feature["phone_seq"] = np.array(self.utt2seq[utt]) | |
| single_feature["phone_len"] = len(self.utt2seq[utt]) | |
| return single_feature | |
| def __len__(self): | |
| return super().__len__() | |
| def get_metadata(self): | |
| return super().get_metadata() | |
| class TTSCollator(BaseCollator): | |
| """Zero-pads model inputs and targets based on number of frames per step""" | |
| def __init__(self, cfg): | |
| super().__init__(cfg) | |
| def __call__(self, batch): | |
| parsed_batch_features = super().__call__(batch) | |
| return parsed_batch_features | |
| class TTSTestDataset(BaseTestDataset): | |
| def __init__(self, args, cfg): | |
| self.cfg = cfg | |
| # inference from test list file | |
| if args.test_list_file is not None: | |
| # construst metadata | |
| self.metadata = [] | |
| with open(args.test_list_file, "r") as fin: | |
| for idx, line in enumerate(fin.readlines()): | |
| utt_info = {} | |
| utt_info["Dataset"] = "test" | |
| utt_info["Text"] = line.strip() | |
| utt_info["Uid"] = str(idx) | |
| self.metadata.append(utt_info) | |
| else: | |
| assert args.testing_set | |
| self.metafile_path = os.path.join( | |
| cfg.preprocess.processed_dir, | |
| args.dataset, | |
| "{}.json".format(args.testing_set), | |
| ) | |
| self.metadata = self.get_metadata() | |
| def __getitem__(self, index): | |
| single_feature = {} | |
| return single_feature | |
| def __len__(self): | |
| return len(self.metadata) | |
| class TTSTestCollator(BaseTestCollator): | |
| """Zero-pads model inputs and targets based on number of frames per step""" | |
| def __init__(self, cfg): | |
| self.cfg = cfg | |
| def __call__(self, batch): | |
| packed_batch_features = dict() | |
| # mel: [b, T, n_mels] | |
| # frame_pitch, frame_energy: [1, T] | |
| # target_len: [1] | |
| # spk_id: [b, 1] | |
| # mask: [b, T, 1] | |
| for key in batch[0].keys(): | |
| if key == "target_len": | |
| packed_batch_features["target_len"] = torch.LongTensor( | |
| [b["target_len"] for b in batch] | |
| ) | |
| masks = [ | |
| torch.ones((b["target_len"], 1), dtype=torch.long) for b in batch | |
| ] | |
| packed_batch_features["mask"] = pad_sequence( | |
| masks, batch_first=True, padding_value=0 | |
| ) | |
| elif key == "phone_len": | |
| packed_batch_features["phone_len"] = torch.LongTensor( | |
| [b["phone_len"] for b in batch] | |
| ) | |
| masks = [ | |
| torch.ones((b["phone_len"], 1), dtype=torch.long) for b in batch | |
| ] | |
| packed_batch_features["phn_mask"] = pad_sequence( | |
| masks, batch_first=True, padding_value=0 | |
| ) | |
| elif key == "audio_len": | |
| packed_batch_features["audio_len"] = torch.LongTensor( | |
| [b["audio_len"] for b in batch] | |
| ) | |
| masks = [ | |
| torch.ones((b["audio_len"], 1), dtype=torch.long) for b in batch | |
| ] | |
| else: | |
| values = [torch.from_numpy(b[key]) for b in batch] | |
| packed_batch_features[key] = pad_sequence( | |
| values, batch_first=True, padding_value=0 | |
| ) | |
| return packed_batch_features | |