import glob import re import librosa import torch import yaml from sklearn.preprocessing import StandardScaler from torch import nn from modules.FastDiff.module.FastDiff_model import FastDiff as FastDiff_model from utils.hparams import hparams from modules.parallel_wavegan.utils import read_hdf5 from vocoders.base_vocoder import BaseVocoder, register_vocoder import numpy as np from modules.FastDiff.module.util import theta_timestep_loss, compute_hyperparams_given_schedule, sampling_given_noise_schedule def load_fastdiff_model(config_path, checkpoint_path): # load config with open(config_path) as f: config = yaml.load(f, Loader=yaml.Loader) # setup if torch.cuda.is_available(): device = torch.device("cuda") else: device = torch.device("cpu") model = FastDiff_model(audio_channels=config['audio_channels'], inner_channels=config['inner_channels'], cond_channels=config['cond_channels'], upsample_ratios=config['upsample_ratios'], lvc_layers_each_block=config['lvc_layers_each_block'], lvc_kernel_size=config['lvc_kernel_size'], kpnet_hidden_channels=config['kpnet_hidden_channels'], kpnet_conv_size=config['kpnet_conv_size'], dropout=config['dropout'], diffusion_step_embed_dim_in=config['diffusion_step_embed_dim_in'], diffusion_step_embed_dim_mid=config['diffusion_step_embed_dim_mid'], diffusion_step_embed_dim_out=config['diffusion_step_embed_dim_out'], use_weight_norm=config['use_weight_norm']) model.load_state_dict(torch.load(checkpoint_path, map_location="cpu")["state_dict"]["model"], strict=True) # Init hyperparameters by linear schedule noise_schedule = torch.linspace(float(config["beta_0"]), float(config["beta_T"]), int(config["T"])).cuda() diffusion_hyperparams = compute_hyperparams_given_schedule(noise_schedule) # map diffusion hyperparameters to gpu for key in diffusion_hyperparams: if key in ["beta", "alpha", "sigma"]: diffusion_hyperparams[key] = diffusion_hyperparams[key].cuda() diffusion_hyperparams = diffusion_hyperparams if config['noise_schedule'] != '': noise_schedule = config['noise_schedule'] if isinstance(noise_schedule, list): noise_schedule = torch.FloatTensor(noise_schedule).cuda() else: # Select Schedule try: reverse_step = int(hparams.get('N')) except: print('Please specify $N (the number of revere iterations) in config file. Now denoise with 4 iterations.') reverse_step = 4 if reverse_step == 1000: noise_schedule = torch.linspace(0.000001, 0.01, 1000).cuda() elif reverse_step == 200: noise_schedule = torch.linspace(0.0001, 0.02, 200).cuda() # Below are schedules derived by Noise Predictor elif reverse_step == 8: noise_schedule = [6.689325005027058e-07, 1.0033881153503899e-05, 0.00015496854030061513, 0.002387222135439515, 0.035597629845142365, 0.3681158423423767, 0.4735414385795593, 0.5] elif reverse_step == 6: noise_schedule = [1.7838445955931093e-06, 2.7984189728158526e-05, 0.00043231004383414984, 0.006634317338466644, 0.09357017278671265, 0.6000000238418579] elif reverse_step == 4: noise_schedule = [3.2176e-04, 2.5743e-03, 2.5376e-02, 7.0414e-01] elif reverse_step == 3: noise_schedule = [9.0000e-05, 9.0000e-03, 6.0000e-01] else: raise NotImplementedError if isinstance(noise_schedule, list): noise_schedule = torch.FloatTensor(noise_schedule).cuda() model.remove_weight_norm() model = model.eval().to(device) print(f"| Loaded model parameters from {checkpoint_path}.") print(f"| FastDiff device: {device}.") return model, diffusion_hyperparams, noise_schedule, config, device @register_vocoder class FastDiff(BaseVocoder): def __init__(self): if hparams['vocoder_ckpt'] == '': # load LJSpeech FastDiff pretrained model base_dir = 'checkpoint/FastDiff' config_path = f'{base_dir}/config.yaml' ckpt = sorted(glob.glob(f'{base_dir}/model_ckpt_steps_*.ckpt'), key= lambda x: int(re.findall(f'{base_dir}/model_ckpt_steps_(\d+).ckpt', x)[0]))[-1] print('| load FastDiff: ', ckpt) self.scaler = None self.model, self.dh, self.noise_schedule, self.config, self.device = load_fastdiff_model( config_path=config_path, checkpoint_path=ckpt, ) else: base_dir = hparams['vocoder_ckpt'] print(base_dir) config_path = f'{base_dir}/config.yaml' ckpt = sorted(glob.glob(f'{base_dir}/model_ckpt_steps_*.ckpt'), key= lambda x: int(re.findall(f'{base_dir}/model_ckpt_steps_(\d+).ckpt', x)[0]))[-1] print('| load FastDiff: ', ckpt) self.scaler = None self.model, self.dh, self.noise_schedule, self.config, self.device = load_fastdiff_model( config_path=config_path, checkpoint_path=ckpt, ) def spec2wav(self, mel, **kwargs): # start generation device = self.device with torch.no_grad(): c = torch.FloatTensor(mel).unsqueeze(0).transpose(2, 1).to(device) audio_length = c.shape[-1] * hparams["hop_size"] y = sampling_given_noise_schedule( self.model, (1, 1, audio_length), self.dh, self.noise_schedule, condition=c, ddim=False, return_sequence=False) wav_out = y.cpu().numpy() return wav_out @staticmethod def wav2spec(wav_fn, return_linear=False): from data_gen.tts.data_gen_utils import process_utterance res = process_utterance( wav_fn, fft_size=hparams['fft_size'], hop_size=hparams['hop_size'], win_length=hparams['win_size'], num_mels=hparams['audio_num_mel_bins'], fmin=hparams['fmin'], fmax=hparams['fmax'], sample_rate=hparams['audio_sample_rate'], loud_norm=hparams['loud_norm'], min_level_db=hparams['min_level_db'], return_linear=return_linear, vocoder='fastdiff', eps=float(hparams.get('wav2spec_eps', 1e-10))) if return_linear: return res[0], res[1].T, res[2].T # [T, 80], [T, n_fft] else: return res[0], res[1].T @staticmethod def wav2mfcc(wav_fn): fft_size = hparams['fft_size'] hop_size = hparams['hop_size'] win_length = hparams['win_size'] sample_rate = hparams['audio_sample_rate'] wav, _ = librosa.core.load(wav_fn, sr=sample_rate) mfcc = librosa.feature.mfcc(y=wav, sr=sample_rate, n_mfcc=13, n_fft=fft_size, hop_length=hop_size, win_length=win_length, pad_mode="constant", power=1.0) mfcc_delta = librosa.feature.delta(mfcc, order=1) mfcc_delta_delta = librosa.feature.delta(mfcc, order=2) mfcc = np.concatenate([mfcc, mfcc_delta, mfcc_delta_delta]).T return mfcc