Spaces:
Running
on
Zero
Running
on
Zero
| # 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 random | |
| import torch | |
| from torch.nn.utils.rnn import pad_sequence | |
| import json | |
| import os | |
| import numpy as np | |
| import librosa | |
| from utils.data_utils import * | |
| from processors.acoustic_extractor import cal_normalized_mel, load_mel_extrema | |
| from processors.content_extractor import ( | |
| ContentvecExtractor, | |
| WhisperExtractor, | |
| WenetExtractor, | |
| ) | |
| from models.base.base_dataset import ( | |
| BaseOfflineDataset, | |
| BaseOfflineCollator, | |
| BaseOnlineDataset, | |
| BaseOnlineCollator, | |
| ) | |
| from models.base.new_dataset import BaseTestDataset | |
| EPS = 1.0e-12 | |
| class SVCOfflineDataset(BaseOfflineDataset): | |
| def __init__(self, cfg, dataset, is_valid=False): | |
| BaseOfflineDataset.__init__(self, cfg, dataset, is_valid=is_valid) | |
| cfg = self.cfg | |
| if cfg.model.condition_encoder.use_whisper: | |
| self.whisper_aligner = WhisperExtractor(self.cfg) | |
| self.utt2whisper_path = load_content_feature_path( | |
| self.metadata, cfg.preprocess.processed_dir, cfg.preprocess.whisper_dir | |
| ) | |
| if cfg.model.condition_encoder.use_contentvec: | |
| self.contentvec_aligner = ContentvecExtractor(self.cfg) | |
| self.utt2contentVec_path = load_content_feature_path( | |
| self.metadata, | |
| cfg.preprocess.processed_dir, | |
| cfg.preprocess.contentvec_dir, | |
| ) | |
| if cfg.model.condition_encoder.use_mert: | |
| self.utt2mert_path = load_content_feature_path( | |
| self.metadata, cfg.preprocess.processed_dir, cfg.preprocess.mert_dir | |
| ) | |
| if cfg.model.condition_encoder.use_wenet: | |
| self.wenet_aligner = WenetExtractor(self.cfg) | |
| self.utt2wenet_path = load_content_feature_path( | |
| self.metadata, cfg.preprocess.processed_dir, cfg.preprocess.wenet_dir | |
| ) | |
| def __getitem__(self, index): | |
| single_feature = BaseOfflineDataset.__getitem__(self, index) | |
| utt_info = self.metadata[index] | |
| dataset = utt_info["Dataset"] | |
| uid = utt_info["Uid"] | |
| utt = "{}_{}".format(dataset, uid) | |
| if self.cfg.model.condition_encoder.use_whisper: | |
| assert "target_len" in single_feature.keys() | |
| aligned_whisper_feat = ( | |
| self.whisper_aligner.offline_resolution_transformation( | |
| np.load(self.utt2whisper_path[utt]), single_feature["target_len"] | |
| ) | |
| ) | |
| single_feature["whisper_feat"] = aligned_whisper_feat | |
| if self.cfg.model.condition_encoder.use_contentvec: | |
| assert "target_len" in single_feature.keys() | |
| aligned_contentvec = ( | |
| self.contentvec_aligner.offline_resolution_transformation( | |
| np.load(self.utt2contentVec_path[utt]), single_feature["target_len"] | |
| ) | |
| ) | |
| single_feature["contentvec_feat"] = aligned_contentvec | |
| if self.cfg.model.condition_encoder.use_mert: | |
| assert "target_len" in single_feature.keys() | |
| aligned_mert_feat = align_content_feature_length( | |
| np.load(self.utt2mert_path[utt]), | |
| single_feature["target_len"], | |
| source_hop=self.cfg.preprocess.mert_hop_size, | |
| ) | |
| single_feature["mert_feat"] = aligned_mert_feat | |
| if self.cfg.model.condition_encoder.use_wenet: | |
| assert "target_len" in single_feature.keys() | |
| aligned_wenet_feat = self.wenet_aligner.offline_resolution_transformation( | |
| np.load(self.utt2wenet_path[utt]), single_feature["target_len"] | |
| ) | |
| single_feature["wenet_feat"] = aligned_wenet_feat | |
| # print(single_feature.keys()) | |
| # for k, v in single_feature.items(): | |
| # if type(v) in [torch.Tensor, np.ndarray]: | |
| # print(k, v.shape) | |
| # else: | |
| # print(k, v) | |
| # exit() | |
| return self.clip_if_too_long(single_feature) | |
| def __len__(self): | |
| return len(self.metadata) | |
| def random_select(self, feature_seq_len, max_seq_len, ending_ts=2812): | |
| """ | |
| ending_ts: to avoid invalid whisper features for over 30s audios | |
| 2812 = 30 * 24000 // 256 | |
| """ | |
| ts = max(feature_seq_len - max_seq_len, 0) | |
| ts = min(ts, ending_ts - max_seq_len) | |
| start = random.randint(0, ts) | |
| end = start + max_seq_len | |
| return start, end | |
| def clip_if_too_long(self, sample, max_seq_len=512): | |
| """ | |
| sample : | |
| { | |
| 'spk_id': (1,), | |
| 'target_len': int | |
| 'mel': (seq_len, dim), | |
| 'frame_pitch': (seq_len,) | |
| 'frame_energy': (seq_len,) | |
| 'content_vector_feat': (seq_len, dim) | |
| } | |
| """ | |
| if sample["target_len"] <= max_seq_len: | |
| return sample | |
| start, end = self.random_select(sample["target_len"], max_seq_len) | |
| sample["target_len"] = end - start | |
| for k in sample.keys(): | |
| if k == "audio": | |
| # audio should be clipped in hop_size scale | |
| sample[k] = sample[k][ | |
| start | |
| * self.cfg.preprocess.hop_size : end | |
| * self.cfg.preprocess.hop_size | |
| ] | |
| elif k == "audio_len": | |
| sample[k] = (end - start) * self.cfg.preprocess.hop_size | |
| elif k not in ["spk_id", "target_len"]: | |
| sample[k] = sample[k][start:end] | |
| return sample | |
| class SVCOnlineDataset(BaseOnlineDataset): | |
| def __init__(self, cfg, dataset, is_valid=False): | |
| super().__init__(cfg, dataset, is_valid=is_valid) | |
| # Audio pretrained models' sample rates | |
| self.all_sample_rates = {self.sample_rate} | |
| if self.cfg.model.condition_encoder.use_whisper: | |
| self.all_sample_rates.add(self.cfg.preprocess.whisper_sample_rate) | |
| if self.cfg.model.condition_encoder.use_contentvec: | |
| self.all_sample_rates.add(self.cfg.preprocess.contentvec_sample_rate) | |
| if self.cfg.model.condition_encoder.use_wenet: | |
| self.all_sample_rates.add(self.cfg.preprocess.wenet_sample_rate) | |
| self.highest_sample_rate = max(list(self.all_sample_rates)) | |
| # The maximum duration (seconds) for one training sample | |
| self.max_duration = 6.0 | |
| self.max_n_frames = int(self.max_duration * self.highest_sample_rate) | |
| def random_select(self, wav, duration, wav_path): | |
| """ | |
| wav: (T,) | |
| """ | |
| if duration <= self.max_duration: | |
| return wav | |
| ts_frame = int((duration - self.max_duration) * self.highest_sample_rate) | |
| start = random.randint(0, ts_frame) | |
| end = start + self.max_n_frames | |
| if (wav[start:end] == 0).all(): | |
| print("*" * 20) | |
| print("Warning! The wav file {} has a lot of silience.".format(wav_path)) | |
| # There should be at least some frames that are not silience. Then we select them. | |
| assert (wav != 0).any() | |
| start = np.where(wav != 0)[0][0] | |
| end = start + self.max_n_frames | |
| return wav[start:end] | |
| def __getitem__(self, index): | |
| """ | |
| single_feature: dict, | |
| wav: (T,) | |
| wav_len: int | |
| target_len: int | |
| mask: (n_frames, 1) | |
| spk_id | |
| wav_{sr}: (T,) | |
| wav_{sr}_len: int | |
| """ | |
| single_feature = dict() | |
| utt_item = self.metadata[index] | |
| wav_path = utt_item["Path"] | |
| ### Use the highest sampling rate to load and randomly select ### | |
| highest_sr_wav, _ = librosa.load(wav_path, sr=self.highest_sample_rate) | |
| highest_sr_wav = self.random_select( | |
| highest_sr_wav, utt_item["Duration"], wav_path | |
| ) | |
| ### Waveforms under all the sample rates ### | |
| for sr in self.all_sample_rates: | |
| # Resample to the required sample rate | |
| if sr != self.highest_sample_rate: | |
| wav_sr = librosa.resample( | |
| highest_sr_wav, orig_sr=self.highest_sample_rate, target_sr=sr | |
| ) | |
| else: | |
| wav_sr = highest_sr_wav | |
| wav_sr = torch.as_tensor(wav_sr, dtype=torch.float32) | |
| single_feature["wav_{}".format(sr)] = wav_sr | |
| single_feature["wav_{}_len".format(sr)] = len(wav_sr) | |
| # For target sample rate | |
| if sr == self.sample_rate: | |
| wav_len = len(wav_sr) | |
| frame_len = wav_len // self.hop_size | |
| single_feature["wav"] = wav_sr | |
| single_feature["wav_len"] = wav_len | |
| single_feature["target_len"] = frame_len | |
| single_feature["mask"] = torch.ones(frame_len, 1, dtype=torch.long) | |
| ### Speaker ID ### | |
| if self.cfg.preprocess.use_spkid: | |
| utt = "{}_{}".format(utt_item["Dataset"], utt_item["Uid"]) | |
| single_feature["spk_id"] = torch.tensor( | |
| [self.spk2id[self.utt2spk[utt]]], dtype=torch.int32 | |
| ) | |
| return single_feature | |
| def __len__(self): | |
| return len(self.metadata) | |
| class SVCOfflineCollator(BaseOfflineCollator): | |
| def __init__(self, cfg): | |
| super().__init__(cfg) | |
| def __call__(self, batch): | |
| parsed_batch_features = super().__call__(batch) | |
| return parsed_batch_features | |
| class SVCOnlineCollator(BaseOnlineCollator): | |
| def __init__(self, cfg): | |
| super().__init__(cfg) | |
| def __call__(self, batch): | |
| """ | |
| SVCOnlineDataset.__getitem__: | |
| wav: (T,) | |
| wav_len: int | |
| target_len: int | |
| mask: (n_frames, 1) | |
| spk_id: (1) | |
| wav_{sr}: (T,) | |
| wav_{sr}_len: int | |
| Returns: | |
| wav: (B, T), torch.float32 | |
| wav_len: (B), torch.long | |
| target_len: (B), torch.long | |
| mask: (B, n_frames, 1), torch.long | |
| spk_id: (B, 1), torch.int32 | |
| wav_{sr}: (B, T) | |
| wav_{sr}_len: (B), torch.long | |
| """ | |
| packed_batch_features = dict() | |
| for key in batch[0].keys(): | |
| if "_len" in key: | |
| packed_batch_features[key] = torch.LongTensor([b[key] for b in batch]) | |
| else: | |
| packed_batch_features[key] = pad_sequence( | |
| [b[key] for b in batch], batch_first=True, padding_value=0 | |
| ) | |
| return packed_batch_features | |
| class SVCTestDataset(BaseTestDataset): | |
| def __init__(self, args, cfg, infer_type): | |
| BaseTestDataset.__init__(self, args, cfg, infer_type) | |
| self.metadata = self.get_metadata() | |
| target_singer = args.target_singer | |
| self.cfg = cfg | |
| self.trans_key = args.trans_key | |
| assert type(target_singer) == str | |
| self.target_singer = target_singer.split("_")[-1] | |
| self.target_dataset = target_singer.replace( | |
| "_{}".format(self.target_singer), "" | |
| ) | |
| if cfg.preprocess.mel_min_max_norm: | |
| if self.cfg.preprocess.features_extraction_mode == "online": | |
| # TODO: Change the hard code | |
| # Using an empirical mel extrema to normalize | |
| self.target_mel_extrema = load_mel_extrema(cfg.preprocess, "vctk") | |
| else: | |
| self.target_mel_extrema = load_mel_extrema( | |
| cfg.preprocess, self.target_dataset | |
| ) | |
| self.target_mel_extrema = torch.as_tensor( | |
| self.target_mel_extrema[0] | |
| ), torch.as_tensor(self.target_mel_extrema[1]) | |
| ######### Load source acoustic features ######### | |
| if cfg.preprocess.use_spkid: | |
| spk2id_path = os.path.join(args.acoustics_dir, cfg.preprocess.spk2id) | |
| # utt2sp_path = os.path.join(self.data_root, cfg.preprocess.utt2spk) | |
| with open(spk2id_path, "r", encoding="utf-8") as f: | |
| self.spk2id = json.load(f) | |
| # print("self.spk2id", self.spk2id) | |
| if cfg.preprocess.use_uv: | |
| self.utt2uv_path = { | |
| f'{utt_info["Dataset"]}_{utt_info["Uid"]}': os.path.join( | |
| cfg.preprocess.processed_dir, | |
| utt_info["Dataset"], | |
| cfg.preprocess.uv_dir, | |
| utt_info["Uid"] + ".npy", | |
| ) | |
| for utt_info in self.metadata | |
| } | |
| if cfg.preprocess.use_frame_pitch: | |
| self.utt2frame_pitch_path = { | |
| f'{utt_info["Dataset"]}_{utt_info["Uid"]}': os.path.join( | |
| cfg.preprocess.processed_dir, | |
| utt_info["Dataset"], | |
| cfg.preprocess.pitch_dir, | |
| utt_info["Uid"] + ".npy", | |
| ) | |
| for utt_info in self.metadata | |
| } | |
| # Target F0 median | |
| target_f0_statistics_path = os.path.join( | |
| cfg.preprocess.processed_dir, | |
| self.target_dataset, | |
| cfg.preprocess.pitch_dir, | |
| "statistics.json", | |
| ) | |
| self.target_pitch_median = json.load( | |
| open(target_f0_statistics_path, "r", encoding="utf-8") | |
| )[f"{self.target_dataset}_{self.target_singer}"]["voiced_positions"][ | |
| "median" | |
| ] | |
| # Source F0 median (if infer from file) | |
| if infer_type == "from_file": | |
| source_audio_name = cfg.inference.source_audio_name | |
| source_f0_statistics_path = os.path.join( | |
| cfg.preprocess.processed_dir, | |
| source_audio_name, | |
| cfg.preprocess.pitch_dir, | |
| "statistics.json", | |
| ) | |
| self.source_pitch_median = json.load( | |
| open(source_f0_statistics_path, "r", encoding="utf-8") | |
| )[f"{source_audio_name}_{source_audio_name}"]["voiced_positions"][ | |
| "median" | |
| ] | |
| else: | |
| self.source_pitch_median = None | |
| if cfg.preprocess.use_frame_energy: | |
| self.utt2frame_energy_path = { | |
| f'{utt_info["Dataset"]}_{utt_info["Uid"]}': os.path.join( | |
| cfg.preprocess.processed_dir, | |
| utt_info["Dataset"], | |
| cfg.preprocess.energy_dir, | |
| utt_info["Uid"] + ".npy", | |
| ) | |
| for utt_info in self.metadata | |
| } | |
| if cfg.preprocess.use_mel: | |
| self.utt2mel_path = { | |
| f'{utt_info["Dataset"]}_{utt_info["Uid"]}': os.path.join( | |
| cfg.preprocess.processed_dir, | |
| utt_info["Dataset"], | |
| cfg.preprocess.mel_dir, | |
| utt_info["Uid"] + ".npy", | |
| ) | |
| for utt_info in self.metadata | |
| } | |
| ######### Load source content features' path ######### | |
| if cfg.model.condition_encoder.use_whisper: | |
| self.whisper_aligner = WhisperExtractor(cfg) | |
| self.utt2whisper_path = load_content_feature_path( | |
| self.metadata, cfg.preprocess.processed_dir, cfg.preprocess.whisper_dir | |
| ) | |
| if cfg.model.condition_encoder.use_contentvec: | |
| self.contentvec_aligner = ContentvecExtractor(cfg) | |
| self.utt2contentVec_path = load_content_feature_path( | |
| self.metadata, | |
| cfg.preprocess.processed_dir, | |
| cfg.preprocess.contentvec_dir, | |
| ) | |
| if cfg.model.condition_encoder.use_mert: | |
| self.utt2mert_path = load_content_feature_path( | |
| self.metadata, cfg.preprocess.processed_dir, cfg.preprocess.mert_dir | |
| ) | |
| if cfg.model.condition_encoder.use_wenet: | |
| self.wenet_aligner = WenetExtractor(cfg) | |
| self.utt2wenet_path = load_content_feature_path( | |
| self.metadata, cfg.preprocess.processed_dir, cfg.preprocess.wenet_dir | |
| ) | |
| def __getitem__(self, index): | |
| single_feature = {} | |
| utt_info = self.metadata[index] | |
| dataset = utt_info["Dataset"] | |
| uid = utt_info["Uid"] | |
| utt = "{}_{}".format(dataset, uid) | |
| source_dataset = self.metadata[index]["Dataset"] | |
| if self.cfg.preprocess.use_spkid: | |
| single_feature["spk_id"] = np.array( | |
| [self.spk2id[f"{self.target_dataset}_{self.target_singer}"]], | |
| dtype=np.int32, | |
| ) | |
| ######### Get Acoustic Features Item ######### | |
| 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: | |
| # mel norm | |
| mel = cal_normalized_mel(mel, source_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_frame_pitch: | |
| frame_pitch_path = self.utt2frame_pitch_path[utt] | |
| frame_pitch = np.load(frame_pitch_path) | |
| if self.trans_key: | |
| try: | |
| self.trans_key = int(self.trans_key) | |
| except: | |
| pass | |
| if type(self.trans_key) == int: | |
| frame_pitch = transpose_key(frame_pitch, self.trans_key) | |
| elif self.trans_key: | |
| assert self.target_singer | |
| frame_pitch = pitch_shift_to_target( | |
| frame_pitch, self.target_pitch_median, self.source_pitch_median | |
| ) | |
| 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 | |
| ######### Get Content Features Item ######### | |
| if self.cfg.model.condition_encoder.use_whisper: | |
| assert "target_len" in single_feature.keys() | |
| aligned_whisper_feat = ( | |
| self.whisper_aligner.offline_resolution_transformation( | |
| np.load(self.utt2whisper_path[utt]), single_feature["target_len"] | |
| ) | |
| ) | |
| single_feature["whisper_feat"] = aligned_whisper_feat | |
| if self.cfg.model.condition_encoder.use_contentvec: | |
| assert "target_len" in single_feature.keys() | |
| aligned_contentvec = ( | |
| self.contentvec_aligner.offline_resolution_transformation( | |
| np.load(self.utt2contentVec_path[utt]), single_feature["target_len"] | |
| ) | |
| ) | |
| single_feature["contentvec_feat"] = aligned_contentvec | |
| if self.cfg.model.condition_encoder.use_mert: | |
| assert "target_len" in single_feature.keys() | |
| aligned_mert_feat = align_content_feature_length( | |
| np.load(self.utt2mert_path[utt]), | |
| single_feature["target_len"], | |
| source_hop=self.cfg.preprocess.mert_hop_size, | |
| ) | |
| single_feature["mert_feat"] = aligned_mert_feat | |
| if self.cfg.model.condition_encoder.use_wenet: | |
| assert "target_len" in single_feature.keys() | |
| aligned_wenet_feat = self.wenet_aligner.offline_resolution_transformation( | |
| np.load(self.utt2wenet_path[utt]), single_feature["target_len"] | |
| ) | |
| single_feature["wenet_feat"] = aligned_wenet_feat | |
| return single_feature | |
| def __len__(self): | |
| return len(self.metadata) | |
| class SVCTestCollator: | |
| """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 | |
| ) | |
| 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 | |