Spaces:
Build error
Build error
| import audioldm_train.utilities.audio as Audio | |
| import os | |
| from librosa.filters import mel as librosa_mel_fn | |
| import torchaudio | |
| import torch | |
| import numpy as np | |
| def dynamic_range_compression_torch(x, C=1, clip_val=1e-5): | |
| return torch.log(torch.clamp(x, min=clip_val) * C) | |
| def dynamic_range_decompression_torch(x, C=1): | |
| return torch.exp(x) / C | |
| def spectral_normalize_torch(magnitudes): | |
| output = dynamic_range_compression_torch(magnitudes) | |
| return output | |
| def spectral_de_normalize_torch(magnitudes): | |
| output = dynamic_range_decompression_torch(magnitudes) | |
| return output | |
| class Preprocessor: | |
| def __init__(self,config,waveform_only=False): | |
| self.pad_wav_start_sample = 0 | |
| self.trim_wavs = False | |
| self.waveform_only = waveform_only | |
| self.config = config | |
| # self.duration = duration | |
| self.build_setting_parameters() | |
| self.build_dsp() | |
| def build_setting_parameters(self): | |
| # Read from the json config | |
| self.melbins = self.config["preprocessing"]["mel"]["n_mel_channels"] | |
| 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"] | |
| def build_dsp(self): | |
| self.mel_basis = {} | |
| self.hann_window = {} | |
| self.filter_length = self.config["preprocessing"]["stft"]["filter_length"] | |
| self.hop_length = self.config["preprocessing"]["stft"]["hop_length"] | |
| self.win_length = self.config["preprocessing"]["stft"]["win_length"] | |
| self.n_mel = self.config["preprocessing"]["mel"]["n_mel_channels"] | |
| self.sampling_rate = self.config["preprocessing"]["audio"]["sampling_rate"] | |
| self.mel_fmin = self.config["preprocessing"]["mel"]["mel_fmin"] | |
| self.mel_fmax = self.config["preprocessing"]["mel"]["mel_fmax"] | |
| 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"],) | |
| def random_uniform(self, start, end): | |
| val = torch.rand(1).item() | |
| return start + (end - start) * val | |
| 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 | |
| for i in range(10): | |
| random_start = int(self.random_uniform(0, waveform_length - target_length)) | |
| if torch.max( | |
| torch.abs(waveform[:, random_start : random_start + target_length]) | |
| > 1e-4 | |
| ): | |
| break | |
| return waveform[:, random_start : random_start + target_length], random_start | |
| 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 trim_wav(self, waveform): | |
| if np.max(np.abs(waveform)) < 0.0001: | |
| return waveform | |
| 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 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 | |
| # # To avoid dim mismatch | |
| # if temp_wav.shape[1] < waveform_length: | |
| # waveform = waveform[:temp_wav.shape[1]] | |
| temp_wav[:, rand_start : rand_start + waveform_length] = waveform | |
| return temp_wav | |
| 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)[0] | |
| log_mel_spec, stft = self.mel_spectrogram_train(waveform.unsqueeze(0)) | |
| 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 mel_spectrogram_train(self, y): | |
| if torch.min(y) < -1.0: | |
| print("train min value is ", torch.min(y)) | |
| if torch.max(y) > 1.0: | |
| print("train max value is ", torch.max(y)) | |
| if self.mel_fmax not in self.mel_basis: | |
| mel = librosa_mel_fn( | |
| self.sampling_rate, | |
| self.filter_length, | |
| self.n_mel, | |
| self.mel_fmin, | |
| self.mel_fmax | |
| ) | |
| self.mel_basis[str(self.mel_fmax) + "_" + str(y.device)] = ( | |
| torch.from_numpy(mel).float().to(y.device) | |
| ) | |
| self.hann_window[str(y.device)] = torch.hann_window(self.win_length).to( | |
| y.device | |
| ) | |
| y = torch.nn.functional.pad( | |
| y.unsqueeze(1), | |
| ( | |
| int((self.filter_length - self.hop_length) / 2), | |
| int((self.filter_length - self.hop_length) / 2), | |
| ), | |
| mode="reflect", | |
| ) | |
| y = y.squeeze(1) | |
| stft_spec = torch.stft( | |
| y, | |
| self.filter_length, | |
| hop_length=self.hop_length, | |
| win_length=self.win_length, | |
| window=self.hann_window[str(y.device)], | |
| center=False, | |
| pad_mode="reflect", | |
| normalized=False, | |
| onesided=True, | |
| return_complex=True, | |
| ) | |
| stft_spec = torch.abs(stft_spec) | |
| mel = spectral_normalize_torch( | |
| torch.matmul( | |
| self.mel_basis[str(self.mel_fmax) + "_" + str(y.device)], stft_spec | |
| ) | |
| ) | |
| return mel[0], stft_spec[0] | |
| def trim_wav(self, waveform): | |
| if np.max(np.abs(waveform)) < 0.0001: | |
| return waveform | |
| def resample(self, waveform, sr): | |
| waveform = torchaudio.functional.resample(waveform, sr, self.sampling_rate) | |
| return waveform | |
| def read_wav_file(self, filename): | |
| # waveform, sr = librosa.load(filename, sr=None, mono=True) # 4 times slower | |
| waveform, sr = torchaudio.load(filename) | |
| # print(waveform.shape) | |
| # waveform, random_start = self.random_segment_wav( | |
| # waveform, target_length=int(sr * self.duration) | |
| # ) | |
| duration = waveform.shape[1]/sr | |
| # print(duration) | |
| waveform, random_start = self.random_segment_wav( | |
| waveform, target_length=int(sr * duration) | |
| ) | |
| waveform = self.resample(waveform, sr) | |
| # Compute audio duration | |
| # random_start = int(random_start * (self.sampling_rate / sr)) | |
| waveform = waveform.numpy()[0, ...] | |
| waveform = self.normalize_wav(waveform) | |
| if self.trim_wavs: | |
| waveform = self.trim_wav(waveform) | |
| waveform = waveform[None, ...] | |
| # waveform = self.pad_wav( | |
| # waveform, target_length=int(self.sampling_rate * self.duration) | |
| # ) | |
| # waveform, target_length=int(self.sampling_rate * duration) | |
| waveform = self.pad_wav( | |
| waveform, target_length=round(self.sampling_rate * duration) | |
| ) | |
| return waveform, random_start, duration | |
| def read_audio_file(self, filename, filename2=None): | |
| if os.path.exists(filename): | |
| waveform, random_start,duration = self.read_wav_file(filename) | |
| # Rewrite target_length and duration | |
| self.duration = duration | |
| self.target_length = int(self.duration * self.sampling_rate / self.hopsize) | |
| else: | |
| print( | |
| 'Non-fatal Warning [dataset.py]: The wav path "', | |
| filename, | |
| '" is not find in the metadata. Use empty waveform instead. This is normal in the inference process.', | |
| ) | |
| target_length = int(self.sampling_rate * self.duration) | |
| waveform = torch.zeros((1, target_length)) | |
| random_start = 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, waveform, random_start |