diffusion

diffusion/data_loaders.py

@@ -0,0 +1,284 @@
+import os
+import random
+import re
+import numpy as np
+import librosa
+import torch
+import random
+from utils import repeat_expand_2d
+from tqdm import tqdm
+from torch.utils.data import Dataset
+
+def traverse_dir(
+        root_dir,
+        extensions,
+        amount=None,
+        str_include=None,
+        str_exclude=None,
+        is_pure=False,
+        is_sort=False,
+        is_ext=True):
+
+    file_list = []
+    cnt = 0
+    for root, _, files in os.walk(root_dir):
+        for file in files:
+            if any([file.endswith(f".{ext}") for ext in extensions]):
+                # path
+                mix_path = os.path.join(root, file)
+                pure_path = mix_path[len(root_dir)+1:] if is_pure else mix_path
+
+                # amount
+                if (amount is not None) and (cnt == amount):
+                    if is_sort:
+                        file_list.sort()
+                    return file_list
+                
+                # check string
+                if (str_include is not None) and (str_include not in pure_path):
+                    continue
+                if (str_exclude is not None) and (str_exclude in pure_path):
+                    continue
+                
+                if not is_ext:
+                    ext = pure_path.split('.')[-1]
+                    pure_path = pure_path[:-(len(ext)+1)]
+                file_list.append(pure_path)
+                cnt += 1
+    if is_sort:
+        file_list.sort()
+    return file_list
+
+
+def get_data_loaders(args, whole_audio=False):
+    data_train = AudioDataset(
+        filelists = args.data.training_files,
+        waveform_sec=args.data.duration,
+        hop_size=args.data.block_size,
+        sample_rate=args.data.sampling_rate,
+        load_all_data=args.train.cache_all_data,
+        whole_audio=whole_audio,
+        extensions=args.data.extensions,
+        n_spk=args.model.n_spk,
+        spk=args.spk,
+        device=args.train.cache_device,
+        fp16=args.train.cache_fp16,
+        use_aug=True)
+    loader_train = torch.utils.data.DataLoader(
+        data_train ,
+        batch_size=args.train.batch_size if not whole_audio else 1,
+        shuffle=True,
+        num_workers=args.train.num_workers if args.train.cache_device=='cpu' else 0,
+        persistent_workers=(args.train.num_workers > 0) if args.train.cache_device=='cpu' else False,
+        pin_memory=True if args.train.cache_device=='cpu' else False
+    )
+    data_valid = AudioDataset(
+        filelists = args.data.validation_files,
+        waveform_sec=args.data.duration,
+        hop_size=args.data.block_size,
+        sample_rate=args.data.sampling_rate,
+        load_all_data=args.train.cache_all_data,
+        whole_audio=True,
+        spk=args.spk,
+        extensions=args.data.extensions,
+        n_spk=args.model.n_spk)
+    loader_valid = torch.utils.data.DataLoader(
+        data_valid,
+        batch_size=1,
+        shuffle=False,
+        num_workers=0,
+        pin_memory=True
+    )
+    return loader_train, loader_valid 
+
+
+class AudioDataset(Dataset):
+    def __init__(
+        self,
+        filelists,
+        waveform_sec,
+        hop_size,
+        sample_rate,
+        spk,
+        load_all_data=True,
+        whole_audio=False,
+        extensions=['wav'],
+        n_spk=1,
+        device='cpu',
+        fp16=False,
+        use_aug=False,
+    ):
+        super().__init__()
+        
+        self.waveform_sec = waveform_sec
+        self.sample_rate = sample_rate
+        self.hop_size = hop_size
+        self.filelists = filelists
+        self.whole_audio = whole_audio
+        self.use_aug = use_aug
+        self.data_buffer={}
+        self.pitch_aug_dict = {}
+        # np.load(os.path.join(self.path_root, 'pitch_aug_dict.npy'), allow_pickle=True).item()
+        if load_all_data:
+            print('Load all the data filelists:', filelists)
+        else:
+            print('Load the f0, volume data filelists:', filelists)
+        with open(filelists,"r") as f:
+            self.paths = f.read().splitlines()
+        for name_ext in tqdm(self.paths, total=len(self.paths)):
+            name = os.path.splitext(name_ext)[0]
+            path_audio = name_ext
+            duration = librosa.get_duration(filename = path_audio, sr = self.sample_rate)
+            
+            path_f0 = name_ext + ".f0.npy"
+            f0,_ = np.load(path_f0,allow_pickle=True)
+            f0 = torch.from_numpy(np.array(f0,dtype=float)).float().unsqueeze(-1).to(device)
+                
+            path_volume = name_ext + ".vol.npy"
+            volume = np.load(path_volume)
+            volume = torch.from_numpy(volume).float().unsqueeze(-1).to(device)
+            
+            path_augvol = name_ext + ".aug_vol.npy"
+            aug_vol = np.load(path_augvol)
+            aug_vol = torch.from_numpy(aug_vol).float().unsqueeze(-1).to(device)
+                        
+            if n_spk is not None and n_spk > 1:
+                spk_name = name_ext.split("/")[-2]
+                spk_id = spk[spk_name] if spk_name in spk else 0
+                if spk_id < 0 or spk_id >= n_spk:
+                    raise ValueError(' [x] Muiti-speaker traing error : spk_id must be a positive integer from 0 to n_spk-1 ')
+            else:
+                spk_id = 0
+            spk_id = torch.LongTensor(np.array([spk_id])).to(device)
+
+            if load_all_data:
+                '''
+                audio, sr = librosa.load(path_audio, sr=self.sample_rate)
+                if len(audio.shape) > 1:
+                    audio = librosa.to_mono(audio)
+                audio = torch.from_numpy(audio).to(device)
+                '''
+                path_mel = name_ext + ".mel.npy"
+                mel = np.load(path_mel)
+                mel = torch.from_numpy(mel).to(device)
+                
+                path_augmel = name_ext + ".aug_mel.npy"
+                aug_mel,keyshift = np.load(path_augmel, allow_pickle=True)
+                aug_mel = np.array(aug_mel,dtype=float)
+                aug_mel = torch.from_numpy(aug_mel).to(device)
+                self.pitch_aug_dict[name_ext] = keyshift
+
+                path_units = name_ext + ".soft.pt"
+                units = torch.load(path_units).to(device)
+                units = units[0]  
+                units = repeat_expand_2d(units,f0.size(0)).transpose(0,1)
+                
+                if fp16:
+                    mel = mel.half()
+                    aug_mel = aug_mel.half()
+                    units = units.half()
+                    
+                self.data_buffer[name_ext] = {
+                        'duration': duration,
+                        'mel': mel,
+                        'aug_mel': aug_mel,
+                        'units': units,
+                        'f0': f0,
+                        'volume': volume,
+                        'aug_vol': aug_vol,
+                        'spk_id': spk_id
+                        }
+            else:
+                path_augmel = name_ext + ".aug_mel.npy"               
+                aug_mel,keyshift = np.load(path_augmel, allow_pickle=True)
+                self.pitch_aug_dict[name_ext] = keyshift
+                self.data_buffer[name_ext] = {
+                        'duration': duration,
+                        'f0': f0,
+                        'volume': volume,
+                        'aug_vol': aug_vol,
+                        'spk_id': spk_id
+                        }
+           
+
+    def __getitem__(self, file_idx):
+        name_ext = self.paths[file_idx]
+        data_buffer = self.data_buffer[name_ext]
+        # check duration. if too short, then skip
+        if data_buffer['duration'] < (self.waveform_sec + 0.1):
+            return self.__getitem__( (file_idx + 1) % len(self.paths))
+            
+        # get item
+        return self.get_data(name_ext, data_buffer)
+
+    def get_data(self, name_ext, data_buffer):
+        name = os.path.splitext(name_ext)[0]
+        frame_resolution = self.hop_size / self.sample_rate
+        duration = data_buffer['duration']
+        waveform_sec = duration if self.whole_audio else self.waveform_sec
+        
+        # load audio
+        idx_from = 0 if self.whole_audio else random.uniform(0, duration - waveform_sec - 0.1)
+        start_frame = int(idx_from / frame_resolution)
+        units_frame_len = int(waveform_sec / frame_resolution)
+        aug_flag = random.choice([True, False]) and self.use_aug
+        '''
+        audio = data_buffer.get('audio')
+        if audio is None:
+            path_audio = os.path.join(self.path_root, 'audio', name) + '.wav'
+            audio, sr = librosa.load(
+                    path_audio, 
+                    sr = self.sample_rate, 
+                    offset = start_frame * frame_resolution,
+                    duration = waveform_sec)
+            if len(audio.shape) > 1:
+                audio = librosa.to_mono(audio)
+            # clip audio into N seconds
+            audio = audio[ : audio.shape[-1] // self.hop_size * self.hop_size]       
+            audio = torch.from_numpy(audio).float()
+        else:
+            audio = audio[start_frame * self.hop_size : (start_frame + units_frame_len) * self.hop_size]
+        '''
+        # load mel
+        mel_key = 'aug_mel' if aug_flag else 'mel'
+        mel = data_buffer.get(mel_key)
+        if mel is None:
+            mel = name_ext + ".mel.npy"
+            mel = np.load(mel)
+            mel = mel[start_frame : start_frame + units_frame_len]
+            mel = torch.from_numpy(mel).float() 
+        else:
+            mel = mel[start_frame : start_frame + units_frame_len]
+            
+        # load f0
+        f0 = data_buffer.get('f0')
+        aug_shift = 0
+        if aug_flag:
+            aug_shift = self.pitch_aug_dict[name_ext]
+        f0_frames = 2 ** (aug_shift / 12) * f0[start_frame : start_frame + units_frame_len]
+        
+        # load units
+        units = data_buffer.get('units')
+        if units is None:
+            path_units = name_ext + ".soft.pt"
+            units = torch.load(path_units)
+            units = units[0]  
+            units = repeat_expand_2d(units,f0.size(0)).transpose(0,1)
+            
+        units = units[start_frame : start_frame + units_frame_len]
+
+        # load volume
+        vol_key = 'aug_vol' if aug_flag else 'volume'
+        volume = data_buffer.get(vol_key)
+        volume_frames = volume[start_frame : start_frame + units_frame_len]
+        
+        # load spk_id
+        spk_id = data_buffer.get('spk_id')
+        
+        # load shift
+        aug_shift = torch.from_numpy(np.array([[aug_shift]])).float()
+        
+        return dict(mel=mel, f0=f0_frames, volume=volume_frames, units=units, spk_id=spk_id, aug_shift=aug_shift, name=name, name_ext=name_ext)
+
+    def __len__(self):
+        return len(self.paths)

diffusion/diffusion.py

@@ -0,0 +1,317 @@
+from collections import deque
+from functools import partial
+from inspect import isfunction
+import torch.nn.functional as F
+import librosa.sequence
+import numpy as np
+import torch
+from torch import nn
+from tqdm import tqdm
+
+
+def exists(x):
+    return x is not None
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+def extract(a, t, x_shape):
+    b, *_ = t.shape
+    out = a.gather(-1, t)
+    return out.reshape(b, *((1,) * (len(x_shape) - 1)))
+
+
+def noise_like(shape, device, repeat=False):
+    repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
+    noise = lambda: torch.randn(shape, device=device)
+    return repeat_noise() if repeat else noise()
+
+
+def linear_beta_schedule(timesteps, max_beta=0.02):
+    """
+    linear schedule
+    """
+    betas = np.linspace(1e-4, max_beta, timesteps)
+    return betas
+
+
+def cosine_beta_schedule(timesteps, s=0.008):
+    """
+    cosine schedule
+    as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
+    """
+    steps = timesteps + 1
+    x = np.linspace(0, steps, steps)
+    alphas_cumprod = np.cos(((x / steps) + s) / (1 + s) * np.pi * 0.5) ** 2
+    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
+    return np.clip(betas, a_min=0, a_max=0.999)
+
+
+beta_schedule = {
+    "cosine": cosine_beta_schedule,
+    "linear": linear_beta_schedule,
+}
+
+
+class GaussianDiffusion(nn.Module):
+    def __init__(self, 
+                denoise_fn, 
+                out_dims=128,
+                timesteps=1000, 
+                k_step=1000,
+                max_beta=0.02,
+                spec_min=-12, 
+                spec_max=2):
+        super().__init__()
+        self.denoise_fn = denoise_fn
+        self.out_dims = out_dims
+        betas = beta_schedule['linear'](timesteps, max_beta=max_beta)
+
+        alphas = 1. - betas
+        alphas_cumprod = np.cumprod(alphas, axis=0)
+        alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
+
+        timesteps, = betas.shape
+        self.num_timesteps = int(timesteps)
+        self.k_step = k_step
+
+        self.noise_list = deque(maxlen=4)
+
+        to_torch = partial(torch.tensor, dtype=torch.float32)
+
+        self.register_buffer('betas', to_torch(betas))
+        self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
+        self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
+        self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
+        self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
+        self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
+        self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))
+
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+        posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
+        # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
+        self.register_buffer('posterior_variance', to_torch(posterior_variance))
+        # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
+        self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20))))
+        self.register_buffer('posterior_mean_coef1', to_torch(
+            betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)))
+        self.register_buffer('posterior_mean_coef2', to_torch(
+            (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod)))
+
+        self.register_buffer('spec_min', torch.FloatTensor([spec_min])[None, None, :out_dims])
+        self.register_buffer('spec_max', torch.FloatTensor([spec_max])[None, None, :out_dims])
+
+    def q_mean_variance(self, x_start, t):
+        mean = extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+        variance = extract(1. - self.alphas_cumprod, t, x_start.shape)
+        log_variance = extract(self.log_one_minus_alphas_cumprod, t, x_start.shape)
+        return mean, variance, log_variance
+
+    def predict_start_from_noise(self, x_t, t, noise):
+        return (
+                extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
+                extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
+        )
+
+    def q_posterior(self, x_start, x_t, t):
+        posterior_mean = (
+                extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +
+                extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        posterior_variance = extract(self.posterior_variance, t, x_t.shape)
+        posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+    def p_mean_variance(self, x, t, cond):
+        noise_pred = self.denoise_fn(x, t, cond=cond)
+        x_recon = self.predict_start_from_noise(x, t=t, noise=noise_pred)
+
+        x_recon.clamp_(-1., 1.)
+
+        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
+        return model_mean, posterior_variance, posterior_log_variance
+
+    @torch.no_grad()
+    def p_sample(self, x, t, cond, clip_denoised=True, repeat_noise=False):
+        b, *_, device = *x.shape, x.device
+        model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, cond=cond)
+        noise = noise_like(x.shape, device, repeat_noise)
+        # no noise when t == 0
+        nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
+        return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
+
+    @torch.no_grad()
+    def p_sample_plms(self, x, t, interval, cond, clip_denoised=True, repeat_noise=False):
+        """
+        Use the PLMS method from
+        [Pseudo Numerical Methods for Diffusion Models on Manifolds](https://arxiv.org/abs/2202.09778).
+        """
+
+        def get_x_pred(x, noise_t, t):
+            a_t = extract(self.alphas_cumprod, t, x.shape)
+            a_prev = extract(self.alphas_cumprod, torch.max(t - interval, torch.zeros_like(t)), x.shape)
+            a_t_sq, a_prev_sq = a_t.sqrt(), a_prev.sqrt()
+
+            x_delta = (a_prev - a_t) * ((1 / (a_t_sq * (a_t_sq + a_prev_sq))) * x - 1 / (
+                    a_t_sq * (((1 - a_prev) * a_t).sqrt() + ((1 - a_t) * a_prev).sqrt())) * noise_t)
+            x_pred = x + x_delta
+
+            return x_pred
+
+        noise_list = self.noise_list
+        noise_pred = self.denoise_fn(x, t, cond=cond)
+
+        if len(noise_list) == 0:
+            x_pred = get_x_pred(x, noise_pred, t)
+            noise_pred_prev = self.denoise_fn(x_pred, max(t - interval, 0), cond=cond)
+            noise_pred_prime = (noise_pred + noise_pred_prev) / 2
+        elif len(noise_list) == 1:
+            noise_pred_prime = (3 * noise_pred - noise_list[-1]) / 2
+        elif len(noise_list) == 2:
+            noise_pred_prime = (23 * noise_pred - 16 * noise_list[-1] + 5 * noise_list[-2]) / 12
+        else:
+            noise_pred_prime = (55 * noise_pred - 59 * noise_list[-1] + 37 * noise_list[-2] - 9 * noise_list[-3]) / 24
+
+        x_prev = get_x_pred(x, noise_pred_prime, t)
+        noise_list.append(noise_pred)
+
+        return x_prev
+
+    def q_sample(self, x_start, t, noise=None):
+        noise = default(noise, lambda: torch.randn_like(x_start))
+        return (
+                extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
+                extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
+        )
+
+    def p_losses(self, x_start, t, cond, noise=None, loss_type='l2'):
+        noise = default(noise, lambda: torch.randn_like(x_start))
+
+        x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
+        x_recon = self.denoise_fn(x_noisy, t, cond)
+
+        if loss_type == 'l1':
+            loss = (noise - x_recon).abs().mean()
+        elif loss_type == 'l2':
+            loss = F.mse_loss(noise, x_recon)
+        else:
+            raise NotImplementedError()
+
+        return loss
+
+    def forward(self, 
+                condition, 
+                gt_spec=None, 
+                infer=True, 
+                infer_speedup=10, 
+                method='dpm-solver',
+                k_step=300,
+                use_tqdm=True):
+        """
+            conditioning diffusion, use fastspeech2 encoder output as the condition
+        """
+        cond = condition.transpose(1, 2)
+        b, device = condition.shape[0], condition.device
+
+        if not infer:
+            spec = self.norm_spec(gt_spec)
+            t = torch.randint(0, self.k_step, (b,), device=device).long()
+            norm_spec = spec.transpose(1, 2)[:, None, :, :]  # [B, 1, M, T]
+            return self.p_losses(norm_spec, t, cond=cond)
+        else:
+            shape = (cond.shape[0], 1, self.out_dims, cond.shape[2])
+            
+            if gt_spec is None:
+                t = self.k_step
+                x = torch.randn(shape, device=device)
+            else:
+                t = k_step
+                norm_spec = self.norm_spec(gt_spec)
+                norm_spec = norm_spec.transpose(1, 2)[:, None, :, :]
+                x = self.q_sample(x_start=norm_spec, t=torch.tensor([t - 1], device=device).long())
+                        
+            if method is not None and infer_speedup > 1:
+                if method == 'dpm-solver':
+                    from .dpm_solver_pytorch import NoiseScheduleVP, model_wrapper, DPM_Solver
+                    # 1. Define the noise schedule.
+                    noise_schedule = NoiseScheduleVP(schedule='discrete', betas=self.betas[:t])
+
+                    # 2. Convert your discrete-time `model` to the continuous-time
+                    # noise prediction model. Here is an example for a diffusion model
+                    # `model` with the noise prediction type ("noise") .
+                    def my_wrapper(fn):
+                        def wrapped(x, t, **kwargs):
+                            ret = fn(x, t, **kwargs)
+                            if use_tqdm:
+                                self.bar.update(1)
+                            return ret
+
+                        return wrapped
+
+                    model_fn = model_wrapper(
+                        my_wrapper(self.denoise_fn),
+                        noise_schedule,
+                        model_type="noise",  # or "x_start" or "v" or "score"
+                        model_kwargs={"cond": cond}
+                    )
+
+                    # 3. Define dpm-solver and sample by singlestep DPM-Solver.
+                    # (We recommend singlestep DPM-Solver for unconditional sampling)
+                    # You can adjust the `steps` to balance the computation
+                    # costs and the sample quality.
+                    dpm_solver = DPM_Solver(model_fn, noise_schedule)
+
+                    steps = t // infer_speedup
+                    if use_tqdm:
+                        self.bar = tqdm(desc="sample time step", total=steps)
+                    x = dpm_solver.sample(
+                        x,
+                        steps=steps,
+                        order=3,
+                        skip_type="time_uniform",
+                        method="singlestep",
+                    )
+                    if use_tqdm:
+                        self.bar.close()
+                elif method == 'pndm':
+                    self.noise_list = deque(maxlen=4)
+                    if use_tqdm:
+                        for i in tqdm(
+                                reversed(range(0, t, infer_speedup)), desc='sample time step',
+                                total=t // infer_speedup,
+                        ):
+                            x = self.p_sample_plms(
+                                x, torch.full((b,), i, device=device, dtype=torch.long),
+                                infer_speedup, cond=cond
+                            )
+                    else:
+                        for i in reversed(range(0, t, infer_speedup)):
+                            x = self.p_sample_plms(
+                                x, torch.full((b,), i, device=device, dtype=torch.long),
+                                infer_speedup, cond=cond
+                            )
+                else:
+                    raise NotImplementedError(method)
+            else:
+                if use_tqdm:
+                    for i in tqdm(reversed(range(0, t)), desc='sample time step', total=t):
+                        x = self.p_sample(x, torch.full((b,), i, device=device, dtype=torch.long), cond)
+                else:
+                    for i in reversed(range(0, t)):
+                        x = self.p_sample(x, torch.full((b,), i, device=device, dtype=torch.long), cond)
+            x = x.squeeze(1).transpose(1, 2)  # [B, T, M]
+            return self.denorm_spec(x)
+
+    def norm_spec(self, x):
+        return (x - self.spec_min) / (self.spec_max - self.spec_min) * 2 - 1
+
+    def denorm_spec(self, x):
+        return (x + 1) / 2 * (self.spec_max - self.spec_min) + self.spec_min

diffusion/diffusion_onnx.py

@@ -0,0 +1,612 @@
+from collections import deque
+from functools import partial
+from inspect import isfunction
+import torch.nn.functional as F
+import librosa.sequence
+import numpy as np
+from torch.nn import Conv1d
+from torch.nn import Mish
+import torch
+from torch import nn
+from tqdm import tqdm
+import math
+
+
+def exists(x):
+    return x is not None
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+def extract(a, t):
+    return a[t].reshape((1, 1, 1, 1))
+
+
+def noise_like(shape, device, repeat=False):
+    repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
+    noise = lambda: torch.randn(shape, device=device)
+    return repeat_noise() if repeat else noise()
+
+
+def linear_beta_schedule(timesteps, max_beta=0.02):
+    """
+    linear schedule
+    """
+    betas = np.linspace(1e-4, max_beta, timesteps)
+    return betas
+
+
+def cosine_beta_schedule(timesteps, s=0.008):
+    """
+    cosine schedule
+    as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
+    """
+    steps = timesteps + 1
+    x = np.linspace(0, steps, steps)
+    alphas_cumprod = np.cos(((x / steps) + s) / (1 + s) * np.pi * 0.5) ** 2
+    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
+    return np.clip(betas, a_min=0, a_max=0.999)
+
+
+beta_schedule = {
+    "cosine": cosine_beta_schedule,
+    "linear": linear_beta_schedule,
+}
+
+
+def extract_1(a, t):
+    return a[t].reshape((1, 1, 1, 1))
+
+
+def predict_stage0(noise_pred, noise_pred_prev):
+    return (noise_pred + noise_pred_prev) / 2
+
+
+def predict_stage1(noise_pred, noise_list):
+    return (noise_pred * 3
+            - noise_list[-1]) / 2
+
+
+def predict_stage2(noise_pred, noise_list):
+    return (noise_pred * 23
+            - noise_list[-1] * 16
+            + noise_list[-2] * 5) / 12
+
+
+def predict_stage3(noise_pred, noise_list):
+    return (noise_pred * 55
+            - noise_list[-1] * 59
+            + noise_list[-2] * 37
+            - noise_list[-3] * 9) / 24
+
+
+class SinusoidalPosEmb(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+        self.half_dim = dim // 2
+        self.emb = 9.21034037 / (self.half_dim - 1)
+        self.emb = torch.exp(torch.arange(self.half_dim) * torch.tensor(-self.emb)).unsqueeze(0)
+        self.emb = self.emb.cpu()
+
+    def forward(self, x):
+        emb = self.emb * x
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, encoder_hidden, residual_channels, dilation):
+        super().__init__()
+        self.residual_channels = residual_channels
+        self.dilated_conv = Conv1d(residual_channels, 2 * residual_channels, 3, padding=dilation, dilation=dilation)
+        self.diffusion_projection = nn.Linear(residual_channels, residual_channels)
+        self.conditioner_projection = Conv1d(encoder_hidden, 2 * residual_channels, 1)
+        self.output_projection = Conv1d(residual_channels, 2 * residual_channels, 1)
+
+    def forward(self, x, conditioner, diffusion_step):
+        diffusion_step = self.diffusion_projection(diffusion_step).unsqueeze(-1)
+        conditioner = self.conditioner_projection(conditioner)
+        y = x + diffusion_step
+        y = self.dilated_conv(y) + conditioner
+
+        gate, filter_1 = torch.split(y, [self.residual_channels, self.residual_channels], dim=1)
+
+        y = torch.sigmoid(gate) * torch.tanh(filter_1)
+        y = self.output_projection(y)
+
+        residual, skip = torch.split(y, [self.residual_channels, self.residual_channels], dim=1)
+
+        return (x + residual) / 1.41421356, skip
+
+
+class DiffNet(nn.Module):
+    def __init__(self, in_dims, n_layers, n_chans, n_hidden):
+        super().__init__()
+        self.encoder_hidden = n_hidden
+        self.residual_layers = n_layers
+        self.residual_channels = n_chans
+        self.input_projection = Conv1d(in_dims, self.residual_channels, 1)
+        self.diffusion_embedding = SinusoidalPosEmb(self.residual_channels)
+        dim = self.residual_channels
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, dim * 4),
+            Mish(),
+            nn.Linear(dim * 4, dim)
+        )
+        self.residual_layers = nn.ModuleList([
+            ResidualBlock(self.encoder_hidden, self.residual_channels, 1)
+            for i in range(self.residual_layers)
+        ])
+        self.skip_projection = Conv1d(self.residual_channels, self.residual_channels, 1)
+        self.output_projection = Conv1d(self.residual_channels, in_dims, 1)
+        nn.init.zeros_(self.output_projection.weight)
+
+    def forward(self, spec, diffusion_step, cond):
+        x = spec.squeeze(0)
+        x = self.input_projection(x)  # x [B, residual_channel, T]
+        x = F.relu(x)
+        # skip = torch.randn_like(x)
+        diffusion_step = diffusion_step.float()
+        diffusion_step = self.diffusion_embedding(diffusion_step)
+        diffusion_step = self.mlp(diffusion_step)
+
+        x, skip = self.residual_layers[0](x, cond, diffusion_step)
+        # noinspection PyTypeChecker
+        for layer in self.residual_layers[1:]:
+            x, skip_connection = layer.forward(x, cond, diffusion_step)
+            skip = skip + skip_connection
+        x = skip / math.sqrt(len(self.residual_layers))
+        x = self.skip_projection(x)
+        x = F.relu(x)
+        x = self.output_projection(x)  # [B, 80, T]
+        return x.unsqueeze(1)
+
+
+class AfterDiffusion(nn.Module):
+    def __init__(self, spec_max, spec_min, v_type='a'):
+        super().__init__()
+        self.spec_max = spec_max
+        self.spec_min = spec_min
+        self.type = v_type
+
+    def forward(self, x):
+        x = x.squeeze(1).permute(0, 2, 1)
+        mel_out = (x + 1) / 2 * (self.spec_max - self.spec_min) + self.spec_min
+        if self.type == 'nsf-hifigan-log10':
+            mel_out = mel_out * 0.434294
+        return mel_out.transpose(2, 1)
+
+
+class Pred(nn.Module):
+    def __init__(self, alphas_cumprod):
+        super().__init__()
+        self.alphas_cumprod = alphas_cumprod
+
+    def forward(self, x_1, noise_t, t_1, t_prev):
+        a_t = extract(self.alphas_cumprod, t_1).cpu()
+        a_prev = extract(self.alphas_cumprod, t_prev).cpu()
+        a_t_sq, a_prev_sq = a_t.sqrt().cpu(), a_prev.sqrt().cpu()
+        x_delta = (a_prev - a_t) * ((1 / (a_t_sq * (a_t_sq + a_prev_sq))) * x_1 - 1 / (
+                a_t_sq * (((1 - a_prev) * a_t).sqrt() + ((1 - a_t) * a_prev).sqrt())) * noise_t)
+        x_pred = x_1 + x_delta.cpu()
+
+        return x_pred
+
+
+class GaussianDiffusion(nn.Module):
+    def __init__(self, 
+                out_dims=128,
+                n_layers=20,
+                n_chans=384,
+                n_hidden=256,
+                timesteps=1000, 
+                k_step=1000,
+                max_beta=0.02,
+                spec_min=-12, 
+                spec_max=2):
+        super().__init__()
+        self.denoise_fn = DiffNet(out_dims, n_layers, n_chans, n_hidden)
+        self.out_dims = out_dims
+        self.mel_bins = out_dims
+        self.n_hidden = n_hidden
+        betas = beta_schedule['linear'](timesteps, max_beta=max_beta)
+
+        alphas = 1. - betas
+        alphas_cumprod = np.cumprod(alphas, axis=0)
+        alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
+        timesteps, = betas.shape
+        self.num_timesteps = int(timesteps)
+        self.k_step = k_step
+
+        self.noise_list = deque(maxlen=4)
+
+        to_torch = partial(torch.tensor, dtype=torch.float32)
+
+        self.register_buffer('betas', to_torch(betas))
+        self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
+        self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
+        self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
+        self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
+        self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
+        self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))
+
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+        posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
+        # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
+        self.register_buffer('posterior_variance', to_torch(posterior_variance))
+        # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
+        self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20))))
+        self.register_buffer('posterior_mean_coef1', to_torch(
+            betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)))
+        self.register_buffer('posterior_mean_coef2', to_torch(
+            (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod)))
+
+        self.register_buffer('spec_min', torch.FloatTensor([spec_min])[None, None, :out_dims])
+        self.register_buffer('spec_max', torch.FloatTensor([spec_max])[None, None, :out_dims])
+        self.ad = AfterDiffusion(self.spec_max, self.spec_min)
+        self.xp = Pred(self.alphas_cumprod)
+
+    def q_mean_variance(self, x_start, t):
+        mean = extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+        variance = extract(1. - self.alphas_cumprod, t, x_start.shape)
+        log_variance = extract(self.log_one_minus_alphas_cumprod, t, x_start.shape)
+        return mean, variance, log_variance
+
+    def predict_start_from_noise(self, x_t, t, noise):
+        return (
+                extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
+                extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
+        )
+
+    def q_posterior(self, x_start, x_t, t):
+        posterior_mean = (
+                extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +
+                extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        posterior_variance = extract(self.posterior_variance, t, x_t.shape)
+        posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+    def p_mean_variance(self, x, t, cond):
+        noise_pred = self.denoise_fn(x, t, cond=cond)
+        x_recon = self.predict_start_from_noise(x, t=t, noise=noise_pred)
+
+        x_recon.clamp_(-1., 1.)
+
+        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
+        return model_mean, posterior_variance, posterior_log_variance
+
+    @torch.no_grad()
+    def p_sample(self, x, t, cond, clip_denoised=True, repeat_noise=False):
+        b, *_, device = *x.shape, x.device
+        model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, cond=cond)
+        noise = noise_like(x.shape, device, repeat_noise)
+        # no noise when t == 0
+        nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
+        return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
+
+    @torch.no_grad()
+    def p_sample_plms(self, x, t, interval, cond, clip_denoised=True, repeat_noise=False):
+        """
+        Use the PLMS method from
+        [Pseudo Numerical Methods for Diffusion Models on Manifolds](https://arxiv.org/abs/2202.09778).
+        """
+
+        def get_x_pred(x, noise_t, t):
+            a_t = extract(self.alphas_cumprod, t)
+            a_prev = extract(self.alphas_cumprod, torch.max(t - interval, torch.zeros_like(t)))
+            a_t_sq, a_prev_sq = a_t.sqrt(), a_prev.sqrt()
+
+            x_delta = (a_prev - a_t) * ((1 / (a_t_sq * (a_t_sq + a_prev_sq))) * x - 1 / (
+                    a_t_sq * (((1 - a_prev) * a_t).sqrt() + ((1 - a_t) * a_prev).sqrt())) * noise_t)
+            x_pred = x + x_delta
+
+            return x_pred
+
+        noise_list = self.noise_list
+        noise_pred = self.denoise_fn(x, t, cond=cond)
+
+        if len(noise_list) == 0:
+            x_pred = get_x_pred(x, noise_pred, t)
+            noise_pred_prev = self.denoise_fn(x_pred, max(t - interval, 0), cond=cond)
+            noise_pred_prime = (noise_pred + noise_pred_prev) / 2
+        elif len(noise_list) == 1:
+            noise_pred_prime = (3 * noise_pred - noise_list[-1]) / 2
+        elif len(noise_list) == 2:
+            noise_pred_prime = (23 * noise_pred - 16 * noise_list[-1] + 5 * noise_list[-2]) / 12
+        else:
+            noise_pred_prime = (55 * noise_pred - 59 * noise_list[-1] + 37 * noise_list[-2] - 9 * noise_list[-3]) / 24
+
+        x_prev = get_x_pred(x, noise_pred_prime, t)
+        noise_list.append(noise_pred)
+
+        return x_prev
+
+    def q_sample(self, x_start, t, noise=None):
+        noise = default(noise, lambda: torch.randn_like(x_start))
+        return (
+                extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
+                extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
+        )
+
+    def p_losses(self, x_start, t, cond, noise=None, loss_type='l2'):
+        noise = default(noise, lambda: torch.randn_like(x_start))
+
+        x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
+        x_recon = self.denoise_fn(x_noisy, t, cond)
+
+        if loss_type == 'l1':
+            loss = (noise - x_recon).abs().mean()
+        elif loss_type == 'l2':
+            loss = F.mse_loss(noise, x_recon)
+        else:
+            raise NotImplementedError()
+
+        return loss
+
+    def org_forward(self, 
+                condition, 
+                init_noise=None,
+                gt_spec=None, 
+                infer=True, 
+                infer_speedup=100, 
+                method='pndm',
+                k_step=1000,
+                use_tqdm=True):
+        """
+            conditioning diffusion, use fastspeech2 encoder output as the condition
+        """
+        cond = condition
+        b, device = condition.shape[0], condition.device
+        if not infer:
+            spec = self.norm_spec(gt_spec)
+            t = torch.randint(0, self.k_step, (b,), device=device).long()
+            norm_spec = spec.transpose(1, 2)[:, None, :, :]  # [B, 1, M, T]
+            return self.p_losses(norm_spec, t, cond=cond)
+        else:
+            shape = (cond.shape[0], 1, self.out_dims, cond.shape[2])
+            
+            if gt_spec is None:
+                t = self.k_step
+                if init_noise is None:
+                    x = torch.randn(shape, device=device)
+                else:
+                    x = init_noise
+            else:
+                t = k_step
+                norm_spec = self.norm_spec(gt_spec)
+                norm_spec = norm_spec.transpose(1, 2)[:, None, :, :]
+                x = self.q_sample(x_start=norm_spec, t=torch.tensor([t - 1], device=device).long())
+                        
+            if method is not None and infer_speedup > 1:
+                if method == 'dpm-solver':
+                    from .dpm_solver_pytorch import NoiseScheduleVP, model_wrapper, DPM_Solver
+                    # 1. Define the noise schedule.
+                    noise_schedule = NoiseScheduleVP(schedule='discrete', betas=self.betas[:t])
+
+                    # 2. Convert your discrete-time `model` to the continuous-time
+                    # noise prediction model. Here is an example for a diffusion model
+                    # `model` with the noise prediction type ("noise") .
+                    def my_wrapper(fn):
+                        def wrapped(x, t, **kwargs):
+                            ret = fn(x, t, **kwargs)
+                            if use_tqdm:
+                                self.bar.update(1)
+                            return ret
+
+                        return wrapped
+
+                    model_fn = model_wrapper(
+                        my_wrapper(self.denoise_fn),
+                        noise_schedule,
+                        model_type="noise",  # or "x_start" or "v" or "score"
+                        model_kwargs={"cond": cond}
+                    )
+
+                    # 3. Define dpm-solver and sample by singlestep DPM-Solver.
+                    # (We recommend singlestep DPM-Solver for unconditional sampling)
+                    # You can adjust the `steps` to balance the computation
+                    # costs and the sample quality.
+                    dpm_solver = DPM_Solver(model_fn, noise_schedule)
+
+                    steps = t // infer_speedup
+                    if use_tqdm:
+                        self.bar = tqdm(desc="sample time step", total=steps)
+                    x = dpm_solver.sample(
+                        x,
+                        steps=steps,
+                        order=3,
+                        skip_type="time_uniform",
+                        method="singlestep",
+                    )
+                    if use_tqdm:
+                        self.bar.close()
+                elif method == 'pndm':
+                    self.noise_list = deque(maxlen=4)
+                    if use_tqdm:
+                        for i in tqdm(
+                                reversed(range(0, t, infer_speedup)), desc='sample time step',
+                                total=t // infer_speedup,
+                        ):
+                            x = self.p_sample_plms(
+                                x, torch.full((b,), i, device=device, dtype=torch.long),
+                                infer_speedup, cond=cond
+                            )
+                    else:
+                        for i in reversed(range(0, t, infer_speedup)):
+                            x = self.p_sample_plms(
+                                x, torch.full((b,), i, device=device, dtype=torch.long),
+                                infer_speedup, cond=cond
+                            )
+                else:
+                    raise NotImplementedError(method)
+            else:
+                if use_tqdm:
+                    for i in tqdm(reversed(range(0, t)), desc='sample time step', total=t):
+                        x = self.p_sample(x, torch.full((b,), i, device=device, dtype=torch.long), cond)
+                else:
+                    for i in reversed(range(0, t)):
+                        x = self.p_sample(x, torch.full((b,), i, device=device, dtype=torch.long), cond)
+            x = x.squeeze(1).transpose(1, 2)  # [B, T, M]
+            return self.denorm_spec(x).transpose(2, 1)
+
+    def norm_spec(self, x):
+        return (x - self.spec_min) / (self.spec_max - self.spec_min) * 2 - 1
+
+    def denorm_spec(self, x):
+        return (x + 1) / 2 * (self.spec_max - self.spec_min) + self.spec_min
+
+    def get_x_pred(self, x_1, noise_t, t_1, t_prev):
+        a_t = extract(self.alphas_cumprod, t_1)
+        a_prev = extract(self.alphas_cumprod, t_prev)
+        a_t_sq, a_prev_sq = a_t.sqrt(), a_prev.sqrt()
+        x_delta = (a_prev - a_t) * ((1 / (a_t_sq * (a_t_sq + a_prev_sq))) * x_1 - 1 / (
+                a_t_sq * (((1 - a_prev) * a_t).sqrt() + ((1 - a_t) * a_prev).sqrt())) * noise_t)
+        x_pred = x_1 + x_delta
+        return x_pred
+
+    def OnnxExport(self, project_name=None, init_noise=None, hidden_channels=256, export_denoise=True, export_pred=True, export_after=True):
+        cond = torch.randn([1, self.n_hidden, 10]).cpu()
+        if init_noise is None:
+            x = torch.randn((1, 1, self.mel_bins, cond.shape[2]), dtype=torch.float32).cpu()
+        else:
+            x = init_noise
+        pndms = 100
+
+        org_y_x = self.org_forward(cond, init_noise=x)
+
+        device = cond.device
+        n_frames = cond.shape[2]
+        step_range = torch.arange(0, self.k_step, pndms, dtype=torch.long, device=device).flip(0)
+        plms_noise_stage = torch.tensor(0, dtype=torch.long, device=device)
+        noise_list = torch.zeros((0, 1, 1, self.mel_bins, n_frames), device=device)
+
+        ot = step_range[0]
+        ot_1 = torch.full((1,), ot, device=device, dtype=torch.long)
+        if export_denoise:
+            torch.onnx.export(
+                self.denoise_fn,
+                (x.cpu(), ot_1.cpu(), cond.cpu()),
+                f"{project_name}_denoise.onnx",
+                input_names=["noise", "time", "condition"],
+                output_names=["noise_pred"],
+                dynamic_axes={
+                    "noise": [3],
+                    "condition": [2]
+                },
+                opset_version=16
+            )
+
+        for t in step_range:
+            t_1 = torch.full((1,), t, device=device, dtype=torch.long)
+            noise_pred = self.denoise_fn(x, t_1, cond)
+            t_prev = t_1 - pndms
+            t_prev = t_prev * (t_prev > 0)
+            if plms_noise_stage == 0:
+                if export_pred:
+                    torch.onnx.export(
+                        self.xp,
+                        (x.cpu(), noise_pred.cpu(), t_1.cpu(), t_prev.cpu()),
+                        f"{project_name}_pred.onnx",
+                        input_names=["noise", "noise_pred", "time", "time_prev"],
+                        output_names=["noise_pred_o"],
+                        dynamic_axes={
+                            "noise": [3],
+                            "noise_pred": [3]
+                        },
+                        opset_version=16
+                    )
+
+                x_pred = self.get_x_pred(x, noise_pred, t_1, t_prev)
+                noise_pred_prev = self.denoise_fn(x_pred, t_prev, cond=cond)
+                noise_pred_prime = predict_stage0(noise_pred, noise_pred_prev)
+
+            elif plms_noise_stage == 1:
+                noise_pred_prime = predict_stage1(noise_pred, noise_list)
+
+            elif plms_noise_stage == 2:
+                noise_pred_prime = predict_stage2(noise_pred, noise_list)
+
+            else:
+                noise_pred_prime = predict_stage3(noise_pred, noise_list)
+
+            noise_pred = noise_pred.unsqueeze(0)
+
+            if plms_noise_stage < 3:
+                noise_list = torch.cat((noise_list, noise_pred), dim=0)
+                plms_noise_stage = plms_noise_stage + 1
+
+            else:
+                noise_list = torch.cat((noise_list[-2:], noise_pred), dim=0)
+
+            x = self.get_x_pred(x, noise_pred_prime, t_1, t_prev)
+        if export_after:
+            torch.onnx.export(
+                self.ad,
+                x.cpu(),
+                f"{project_name}_after.onnx",
+                input_names=["x"],
+                output_names=["mel_out"],
+                dynamic_axes={
+                    "x": [3]
+                },
+                opset_version=16
+            )
+        x = self.ad(x)
+
+        print((x == org_y_x).all())
+        return x
+
+    def forward(self, condition=None, init_noise=None, pndms=None, k_step=None):
+        cond = condition
+        x = init_noise
+
+        device = cond.device
+        n_frames = cond.shape[2]
+        step_range = torch.arange(0, k_step.item(), pndms.item(), dtype=torch.long, device=device).flip(0)
+        plms_noise_stage = torch.tensor(0, dtype=torch.long, device=device)
+        noise_list = torch.zeros((0, 1, 1, self.mel_bins, n_frames), device=device)
+
+        ot = step_range[0]
+        ot_1 = torch.full((1,), ot, device=device, dtype=torch.long)
+
+        for t in step_range:
+            t_1 = torch.full((1,), t, device=device, dtype=torch.long)
+            noise_pred = self.denoise_fn(x, t_1, cond)
+            t_prev = t_1 - pndms
+            t_prev = t_prev * (t_prev > 0)
+            if plms_noise_stage == 0:
+                x_pred = self.get_x_pred(x, noise_pred, t_1, t_prev)
+                noise_pred_prev = self.denoise_fn(x_pred, t_prev, cond=cond)
+                noise_pred_prime = predict_stage0(noise_pred, noise_pred_prev)
+
+            elif plms_noise_stage == 1:
+                noise_pred_prime = predict_stage1(noise_pred, noise_list)
+
+            elif plms_noise_stage == 2:
+                noise_pred_prime = predict_stage2(noise_pred, noise_list)
+
+            else:
+                noise_pred_prime = predict_stage3(noise_pred, noise_list)
+
+            noise_pred = noise_pred.unsqueeze(0)
+
+            if plms_noise_stage < 3:
+                noise_list = torch.cat((noise_list, noise_pred), dim=0)
+                plms_noise_stage = plms_noise_stage + 1
+
+            else:
+                noise_list = torch.cat((noise_list[-2:], noise_pred), dim=0)
+
+            x = self.get_x_pred(x, noise_pred_prime, t_1, t_prev)
+        x = self.ad(x)
+        return x

diffusion/dpm_solver_pytorch.py

@@ -0,0 +1,1201 @@
+import math
+
+import torch
+
+
+class NoiseScheduleVP:
+    def __init__(
+            self,
+            schedule='discrete',
+            betas=None,
+            alphas_cumprod=None,
+            continuous_beta_0=0.1,
+            continuous_beta_1=20.,
+    ):
+        """Create a wrapper class for the forward SDE (VP type).
+
+        ***
+        Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t.
+                We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for high-resolution images.
+        ***
+
+        The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ).
+        We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper).
+        Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have:
+
+            log_alpha_t = self.marginal_log_mean_coeff(t)
+            sigma_t = self.marginal_std(t)
+            lambda_t = self.marginal_lambda(t)
+
+        Moreover, as lambda(t) is an invertible function, we also support its inverse function:
+
+            t = self.inverse_lambda(lambda_t)
+
+        ===============================================================
+
+        We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and continuous-time DPMs (trained on t in [t_0, T]).
+
+        1. For discrete-time DPMs:
+
+            For discrete-time DPMs trained on n = 0, 1, ..., N-1, we convert the discrete steps to continuous time steps by:
+                t_i = (i + 1) / N
+            e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1.
+            We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3.
+
+            Args:
+                betas: A `torch.Tensor`. The beta array for the discrete-time DPM. (See the original DDPM paper for details)
+                alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. (See the original DDPM paper for details)
+
+            Note that we always have alphas_cumprod = cumprod(betas). Therefore, we only need to set one of `betas` and `alphas_cumprod`.
+
+            **Important**:  Please pay special attention for the args for `alphas_cumprod`:
+                The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. Specifically, DDPMs assume that
+                    q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ).
+                Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. In fact, we have
+                    alpha_{t_n} = \sqrt{\hat{alpha_n}},
+                and
+                    log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}).
+
+
+        2. For continuous-time DPMs:
+
+            We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). The hyperparameters for the noise
+            schedule are the default settings in DDPM and improved-DDPM:
+
+            Args:
+                beta_min: A `float` number. The smallest beta for the linear schedule.
+                beta_max: A `float` number. The largest beta for the linear schedule.
+                cosine_s: A `float` number. The hyperparameter in the cosine schedule.
+                cosine_beta_max: A `float` number. The hyperparameter in the cosine schedule.
+                T: A `float` number. The ending time of the forward process.
+
+        ===============================================================
+
+        Args:
+            schedule: A `str`. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs,
+                    'linear' or 'cosine' for continuous-time DPMs.
+        Returns:
+            A wrapper object of the forward SDE (VP type).
+        
+        ===============================================================
+
+        Example:
+
+        # For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1):
+        >>> ns = NoiseScheduleVP('discrete', betas=betas)
+
+        # For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array for n = 0, 1, ..., N - 1):
+        >>> ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod)
+
+        # For continuous-time DPMs (VPSDE), linear schedule:
+        >>> ns = NoiseScheduleVP('linear', continuous_beta_0=0.1, continuous_beta_1=20.)
+
+        """
+
+        if schedule not in ['discrete', 'linear', 'cosine']:
+            raise ValueError(
+                "Unsupported noise schedule {}. The schedule needs to be 'discrete' or 'linear' or 'cosine'".format(
+                    schedule))
+
+        self.schedule = schedule
+        if schedule == 'discrete':
+            if betas is not None:
+                log_alphas = 0.5 * torch.log(1 - betas).cumsum(dim=0)
+            else:
+                assert alphas_cumprod is not None
+                log_alphas = 0.5 * torch.log(alphas_cumprod)
+            self.total_N = len(log_alphas)
+            self.T = 1.
+            self.t_array = torch.linspace(0., 1., self.total_N + 1)[1:].reshape((1, -1))
+            self.log_alpha_array = log_alphas.reshape((1, -1,))
+        else:
+            self.total_N = 1000
+            self.beta_0 = continuous_beta_0
+            self.beta_1 = continuous_beta_1
+            self.cosine_s = 0.008
+            self.cosine_beta_max = 999.
+            self.cosine_t_max = math.atan(self.cosine_beta_max * (1. + self.cosine_s) / math.pi) * 2. * (
+                        1. + self.cosine_s) / math.pi - self.cosine_s
+            self.cosine_log_alpha_0 = math.log(math.cos(self.cosine_s / (1. + self.cosine_s) * math.pi / 2.))
+            self.schedule = schedule
+            if schedule == 'cosine':
+                # For the cosine schedule, T = 1 will have numerical issues. So we manually set the ending time T.
+                # Note that T = 0.9946 may be not the optimal setting. However, we find it works well.
+                self.T = 0.9946
+            else:
+                self.T = 1.
+
+    def marginal_log_mean_coeff(self, t):
+        """
+        Compute log(alpha_t) of a given continuous-time label t in [0, T].
+        """
+        if self.schedule == 'discrete':
+            return interpolate_fn(t.reshape((-1, 1)), self.t_array.to(t.device),
+                                  self.log_alpha_array.to(t.device)).reshape((-1))
+        elif self.schedule == 'linear':
+            return -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
+        elif self.schedule == 'cosine':
+            log_alpha_fn = lambda s: torch.log(torch.cos((s + self.cosine_s) / (1. + self.cosine_s) * math.pi / 2.))
+            log_alpha_t = log_alpha_fn(t) - self.cosine_log_alpha_0
+            return log_alpha_t
+
+    def marginal_alpha(self, t):
+        """
+        Compute alpha_t of a given continuous-time label t in [0, T].
+        """
+        return torch.exp(self.marginal_log_mean_coeff(t))
+
+    def marginal_std(self, t):
+        """
+        Compute sigma_t of a given continuous-time label t in [0, T].
+        """
+        return torch.sqrt(1. - torch.exp(2. * self.marginal_log_mean_coeff(t)))
+
+    def marginal_lambda(self, t):
+        """
+        Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T].
+        """
+        log_mean_coeff = self.marginal_log_mean_coeff(t)
+        log_std = 0.5 * torch.log(1. - torch.exp(2. * log_mean_coeff))
+        return log_mean_coeff - log_std
+
+    def inverse_lambda(self, lamb):
+        """
+        Compute the continuous-time label t in [0, T] of a given half-logSNR lambda_t.
+        """
+        if self.schedule == 'linear':
+            tmp = 2. * (self.beta_1 - self.beta_0) * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
+            Delta = self.beta_0 ** 2 + tmp
+            return tmp / (torch.sqrt(Delta) + self.beta_0) / (self.beta_1 - self.beta_0)
+        elif self.schedule == 'discrete':
+            log_alpha = -0.5 * torch.logaddexp(torch.zeros((1,)).to(lamb.device), -2. * lamb)
+            t = interpolate_fn(log_alpha.reshape((-1, 1)), torch.flip(self.log_alpha_array.to(lamb.device), [1]),
+                               torch.flip(self.t_array.to(lamb.device), [1]))
+            return t.reshape((-1,))
+        else:
+            log_alpha = -0.5 * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
+            t_fn = lambda log_alpha_t: torch.arccos(torch.exp(log_alpha_t + self.cosine_log_alpha_0)) * 2. * (
+                        1. + self.cosine_s) / math.pi - self.cosine_s
+            t = t_fn(log_alpha)
+            return t
+
+
+def model_wrapper(
+        model,
+        noise_schedule,
+        model_type="noise",
+        model_kwargs={},
+        guidance_type="uncond",
+        condition=None,
+        unconditional_condition=None,
+        guidance_scale=1.,
+        classifier_fn=None,
+        classifier_kwargs={},
+):
+    """Create a wrapper function for the noise prediction model.
+
+    DPM-Solver needs to solve the continuous-time diffusion ODEs. For DPMs trained on discrete-time labels, we need to
+    firstly wrap the model function to a noise prediction model that accepts the continuous time as the input.
+
+    We support four types of the diffusion model by setting `model_type`:
+
+        1. "noise": noise prediction model. (Trained by predicting noise).
+
+        2. "x_start": data prediction model. (Trained by predicting the data x_0 at time 0).
+
+        3. "v": velocity prediction model. (Trained by predicting the velocity).
+            The "v" prediction is derivation detailed in Appendix D of [1], and is used in Imagen-Video [2].
+
+            [1] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models."
+                arXiv preprint arXiv:2202.00512 (2022).
+            [2] Ho, Jonathan, et al. "Imagen Video: High Definition Video Generation with Diffusion Models."
+                arXiv preprint arXiv:2210.02303 (2022).
+    
+        4. "score": marginal score function. (Trained by denoising score matching).
+            Note that the score function and the noise prediction model follows a simple relationship:
+            ```
+                noise(x_t, t) = -sigma_t * score(x_t, t)
+            ```
+
+    We support three types of guided sampling by DPMs by setting `guidance_type`:
+        1. "uncond": unconditional sampling by DPMs.
+            The input `model` has the following format:
+            ``
+                model(x, t_input, **model_kwargs) -> noise | x_start | v | score
+            ``
+
+        2. "classifier": classifier guidance sampling [3] by DPMs and another classifier.
+            The input `model` has the following format:
+            ``
+                model(x, t_input, **model_kwargs) -> noise | x_start | v | score
+            `` 
+
+            The input `classifier_fn` has the following format:
+            ``
+                classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond)
+            ``
+
+            [3] P. Dhariwal and A. Q. Nichol, "Diffusion models beat GANs on image synthesis,"
+                in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 8780-8794.
+
+        3. "classifier-free": classifier-free guidance sampling by conditional DPMs.
+            The input `model` has the following format:
+            ``
+                model(x, t_input, cond, **model_kwargs) -> noise | x_start | v | score
+            `` 
+            And if cond == `unconditional_condition`, the model output is the unconditional DPM output.
+
+            [4] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance."
+                arXiv preprint arXiv:2207.12598 (2022).
+        
+
+    The `t_input` is the time label of the model, which may be discrete-time labels (i.e. 0 to 999)
+    or continuous-time labels (i.e. epsilon to T).
+
+    We wrap the model function to accept only `x` and `t_continuous` as inputs, and outputs the predicted noise:
+    ``
+        def model_fn(x, t_continuous) -> noise:
+            t_input = get_model_input_time(t_continuous)
+            return noise_pred(model, x, t_input, **model_kwargs)         
+    ``
+    where `t_continuous` is the continuous time labels (i.e. epsilon to T). And we use `model_fn` for DPM-Solver.
+
+    ===============================================================
+
+    Args:
+        model: A diffusion model with the corresponding format described above.
+        noise_schedule: A noise schedule object, such as NoiseScheduleVP.
+        model_type: A `str`. The parameterization type of the diffusion model.
+                    "noise" or "x_start" or "v" or "score".
+        model_kwargs: A `dict`. A dict for the other inputs of the model function.
+        guidance_type: A `str`. The type of the guidance for sampling.
+                    "uncond" or "classifier" or "classifier-free".
+        condition: A pytorch tensor. The condition for the guided sampling.
+                    Only used for "classifier" or "classifier-free" guidance type.
+        unconditional_condition: A pytorch tensor. The condition for the unconditional sampling.
+                    Only used for "classifier-free" guidance type.
+        guidance_scale: A `float`. The scale for the guided sampling.
+        classifier_fn: A classifier function. Only used for the classifier guidance.
+        classifier_kwargs: A `dict`. A dict for the other inputs of the classifier function.
+    Returns:
+        A noise prediction model that accepts the noised data and the continuous time as the inputs.
+    """
+
+    def get_model_input_time(t_continuous):
+        """
+        Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time.
+        For discrete-time DPMs, we convert `t_continuous` in [1 / N, 1] to `t_input` in [0, 1000 * (N - 1) / N].
+        For continuous-time DPMs, we just use `t_continuous`.
+        """
+        if noise_schedule.schedule == 'discrete':
+            return (t_continuous - 1. / noise_schedule.total_N) * noise_schedule.total_N
+        else:
+            return t_continuous
+
+    def noise_pred_fn(x, t_continuous, cond=None):
+        if t_continuous.reshape((-1,)).shape[0] == 1:
+            t_continuous = t_continuous.expand((x.shape[0]))
+        t_input = get_model_input_time(t_continuous)
+        if cond is None:
+            output = model(x, t_input, **model_kwargs)
+        else:
+            output = model(x, t_input, cond, **model_kwargs)
+        if model_type == "noise":
+            return output
+        elif model_type == "x_start":
+            alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
+            dims = x.dim()
+            return (x - expand_dims(alpha_t, dims) * output) / expand_dims(sigma_t, dims)
+        elif model_type == "v":
+            alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
+            dims = x.dim()
+            return expand_dims(alpha_t, dims) * output + expand_dims(sigma_t, dims) * x
+        elif model_type == "score":
+            sigma_t = noise_schedule.marginal_std(t_continuous)
+            dims = x.dim()
+            return -expand_dims(sigma_t, dims) * output
+
+    def cond_grad_fn(x, t_input):
+        """
+        Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t).
+        """
+        with torch.enable_grad():
+            x_in = x.detach().requires_grad_(True)
+            log_prob = classifier_fn(x_in, t_input, condition, **classifier_kwargs)
+            return torch.autograd.grad(log_prob.sum(), x_in)[0]
+
+    def model_fn(x, t_continuous):
+        """
+        The noise predicition model function that is used for DPM-Solver.
+        """
+        if t_continuous.reshape((-1,)).shape[0] == 1:
+            t_continuous = t_continuous.expand((x.shape[0]))
+        if guidance_type == "uncond":
+            return noise_pred_fn(x, t_continuous)
+        elif guidance_type == "classifier":
+            assert classifier_fn is not None
+            t_input = get_model_input_time(t_continuous)
+            cond_grad = cond_grad_fn(x, t_input)
+            sigma_t = noise_schedule.marginal_std(t_continuous)
+            noise = noise_pred_fn(x, t_continuous)
+            return noise - guidance_scale * expand_dims(sigma_t, dims=cond_grad.dim()) * cond_grad
+        elif guidance_type == "classifier-free":
+            if guidance_scale == 1. or unconditional_condition is None:
+                return noise_pred_fn(x, t_continuous, cond=condition)
+            else:
+                x_in = torch.cat([x] * 2)
+                t_in = torch.cat([t_continuous] * 2)
+                c_in = torch.cat([unconditional_condition, condition])
+                noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
+                return noise_uncond + guidance_scale * (noise - noise_uncond)
+
+    assert model_type in ["noise", "x_start", "v"]
+    assert guidance_type in ["uncond", "classifier", "classifier-free"]
+    return model_fn
+
+
+class DPM_Solver:
+    def __init__(self, model_fn, noise_schedule, predict_x0=False, thresholding=False, max_val=1.):
+        """Construct a DPM-Solver. 
+
+        We support both the noise prediction model ("predicting epsilon") and the data prediction model ("predicting x0").
+        If `predict_x0` is False, we use the solver for the noise prediction model (DPM-Solver).
+        If `predict_x0` is True, we use the solver for the data prediction model (DPM-Solver++).
+            In such case, we further support the "dynamic thresholding" in [1] when `thresholding` is True.
+            The "dynamic thresholding" can greatly improve the sample quality for pixel-space DPMs with large guidance scales.
+
+        Args:
+            model_fn: A noise prediction model function which accepts the continuous-time input (t in [epsilon, T]):
+                ``
+                def model_fn(x, t_continuous):
+                    return noise
+                ``
+            noise_schedule: A noise schedule object, such as NoiseScheduleVP.
+            predict_x0: A `bool`. If true, use the data prediction model; else, use the noise prediction model.
+            thresholding: A `bool`. Valid when `predict_x0` is True. Whether to use the "dynamic thresholding" in [1].
+            max_val: A `float`. Valid when both `predict_x0` and `thresholding` are True. The max value for thresholding.
+        
+        [1] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv:2205.11487, 2022b.
+        """
+        self.model = model_fn
+        self.noise_schedule = noise_schedule
+        self.predict_x0 = predict_x0
+        self.thresholding = thresholding
+        self.max_val = max_val
+
+    def noise_prediction_fn(self, x, t):
+        """
+        Return the noise prediction model.
+        """
+        return self.model(x, t)
+
+    def data_prediction_fn(self, x, t):
+        """
+        Return the data prediction model (with thresholding).
+        """
+        noise = self.noise_prediction_fn(x, t)
+        dims = x.dim()
+        alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t)
+        x0 = (x - expand_dims(sigma_t, dims) * noise) / expand_dims(alpha_t, dims)
+        if self.thresholding:
+            p = 0.995  # A hyperparameter in the paper of "Imagen" [1].
+            s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
+            s = expand_dims(torch.maximum(s, self.max_val * torch.ones_like(s).to(s.device)), dims)
+            x0 = torch.clamp(x0, -s, s) / s
+        return x0
+
+    def model_fn(self, x, t):
+        """
+        Convert the model to the noise prediction model or the data prediction model. 
+        """
+        if self.predict_x0:
+            return self.data_prediction_fn(x, t)
+        else:
+            return self.noise_prediction_fn(x, t)
+
+    def get_time_steps(self, skip_type, t_T, t_0, N, device):
+        """Compute the intermediate time steps for sampling.
+
+        Args:
+            skip_type: A `str`. The type for the spacing of the time steps. We support three types:
+                - 'logSNR': uniform logSNR for the time steps.
+                - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.)
+                - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)
+            t_T: A `float`. The starting time of the sampling (default is T).
+            t_0: A `float`. The ending time of the sampling (default is epsilon).
+            N: A `int`. The total number of the spacing of the time steps.
+            device: A torch device.
+        Returns:
+            A pytorch tensor of the time steps, with the shape (N + 1,).
+        """
+        if skip_type == 'logSNR':
+            lambda_T = self.noise_schedule.marginal_lambda(torch.tensor(t_T).to(device))
+            lambda_0 = self.noise_schedule.marginal_lambda(torch.tensor(t_0).to(device))
+            logSNR_steps = torch.linspace(lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1).to(device)
+            return self.noise_schedule.inverse_lambda(logSNR_steps)
+        elif skip_type == 'time_uniform':
+            return torch.linspace(t_T, t_0, N + 1).to(device)
+        elif skip_type == 'time_quadratic':
+            t_order = 2
+            t = torch.linspace(t_T ** (1. / t_order), t_0 ** (1. / t_order), N + 1).pow(t_order).to(device)
+            return t
+        else:
+            raise ValueError(
+                "Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(skip_type))
+
+    def get_orders_and_timesteps_for_singlestep_solver(self, steps, order, skip_type, t_T, t_0, device):
+        """
+        Get the order of each step for sampling by the singlestep DPM-Solver.
+
+        We combine both DPM-Solver-1,2,3 to use all the function evaluations, which is named as "DPM-Solver-fast".
+        Given a fixed number of function evaluations by `steps`, the sampling procedure by DPM-Solver-fast is:
+            - If order == 1:
+                We take `steps` of DPM-Solver-1 (i.e. DDIM).
+            - If order == 2:
+                - Denote K = (steps // 2). We take K or (K + 1) intermediate time steps for sampling.
+                - If steps % 2 == 0, we use K steps of DPM-Solver-2.
+                - If steps % 2 == 1, we use K steps of DPM-Solver-2 and 1 step of DPM-Solver-1.
+            - If order == 3:
+                - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling.
+                - If steps % 3 == 0, we use (K - 2) steps of DPM-Solver-3, and 1 step of DPM-Solver-2 and 1 step of DPM-Solver-1.
+                - If steps % 3 == 1, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-1.
+                - If steps % 3 == 2, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-2.
+
+        ============================================
+        Args:
+            order: A `int`. The max order for the solver (2 or 3).
+            steps: A `int`. The total number of function evaluations (NFE).
+            skip_type: A `str`. The type for the spacing of the time steps. We support three types:
+                - 'logSNR': uniform logSNR for the time steps.
+                - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.)
+                - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)
+            t_T: A `float`. The starting time of the sampling (default is T).
+            t_0: A `float`. The ending time of the sampling (default is epsilon).
+            device: A torch device.
+        Returns:
+            orders: A list of the solver order of each step.
+        """
+        if order == 3:
+            K = steps // 3 + 1
+            if steps % 3 == 0:
+                orders = [3, ] * (K - 2) + [2, 1]
+            elif steps % 3 == 1:
+                orders = [3, ] * (K - 1) + [1]
+            else:
+                orders = [3, ] * (K - 1) + [2]
+        elif order == 2:
+            if steps % 2 == 0:
+                K = steps // 2
+                orders = [2, ] * K
+            else:
+                K = steps // 2 + 1
+                orders = [2, ] * (K - 1) + [1]
+        elif order == 1:
+            K = 1
+            orders = [1, ] * steps
+        else:
+            raise ValueError("'order' must be '1' or '2' or '3'.")
+        if skip_type == 'logSNR':
+            # To reproduce the results in DPM-Solver paper
+            timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, K, device)
+        else:
+            timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, steps, device)[
+                torch.cumsum(torch.tensor([0, ] + orders), dim=0).to(device)]
+        return timesteps_outer, orders
+
+    def denoise_fn(self, x, s):
+        """
+        Denoise at the final step, which is equivalent to solve the ODE from lambda_s to infty by first-order discretization. 
+        """
+        return self.data_prediction_fn(x, s)
+
+    def dpm_solver_first_update(self, x, s, t, model_s=None, return_intermediate=False):
+        """
+        DPM-Solver-1 (equivalent to DDIM) from time `s` to time `t`.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `s`.
+            s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
+            t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
+            model_s: A pytorch tensor. The model function evaluated at time `s`.
+                If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
+            return_intermediate: A `bool`. If true, also return the model value at time `s`.
+        Returns:
+            x_t: A pytorch tensor. The approximated solution at time `t`.
+        """
+        ns = self.noise_schedule
+        dims = x.dim()
+        lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
+        h = lambda_t - lambda_s
+        log_alpha_s, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(t)
+        sigma_s, sigma_t = ns.marginal_std(s), ns.marginal_std(t)
+        alpha_t = torch.exp(log_alpha_t)
+
+        if self.predict_x0:
+            phi_1 = torch.expm1(-h)
+            if model_s is None:
+                model_s = self.model_fn(x, s)
+            x_t = (
+                    expand_dims(sigma_t / sigma_s, dims) * x
+                    - expand_dims(alpha_t * phi_1, dims) * model_s
+            )
+            if return_intermediate:
+                return x_t, {'model_s': model_s}
+            else:
+                return x_t
+        else:
+            phi_1 = torch.expm1(h)
+            if model_s is None:
+                model_s = self.model_fn(x, s)
+            x_t = (
+                    expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
+                    - expand_dims(sigma_t * phi_1, dims) * model_s
+            )
+            if return_intermediate:
+                return x_t, {'model_s': model_s}
+            else:
+                return x_t
+
+    def singlestep_dpm_solver_second_update(self, x, s, t, r1=0.5, model_s=None, return_intermediate=False,
+                                            solver_type='dpm_solver'):
+        """
+        Singlestep solver DPM-Solver-2 from time `s` to time `t`.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `s`.
+            s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
+            t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
+            r1: A `float`. The hyperparameter of the second-order solver.
+            model_s: A pytorch tensor. The model function evaluated at time `s`.
+                If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
+            return_intermediate: A `bool`. If true, also return the model value at time `s` and `s1` (the intermediate time).
+            solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
+                The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
+        Returns:
+            x_t: A pytorch tensor. The approximated solution at time `t`.
+        """
+        if solver_type not in ['dpm_solver', 'taylor']:
+            raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
+        if r1 is None:
+            r1 = 0.5
+        ns = self.noise_schedule
+        dims = x.dim()
+        lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
+        h = lambda_t - lambda_s
+        lambda_s1 = lambda_s + r1 * h
+        s1 = ns.inverse_lambda(lambda_s1)
+        log_alpha_s, log_alpha_s1, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(
+            s1), ns.marginal_log_mean_coeff(t)
+        sigma_s, sigma_s1, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(t)
+        alpha_s1, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_t)
+
+        if self.predict_x0:
+            phi_11 = torch.expm1(-r1 * h)
+            phi_1 = torch.expm1(-h)
+
+            if model_s is None:
+                model_s = self.model_fn(x, s)
+            x_s1 = (
+                    expand_dims(sigma_s1 / sigma_s, dims) * x
+                    - expand_dims(alpha_s1 * phi_11, dims) * model_s
+            )
+            model_s1 = self.model_fn(x_s1, s1)
+            if solver_type == 'dpm_solver':
+                x_t = (
+                        expand_dims(sigma_t / sigma_s, dims) * x
+                        - expand_dims(alpha_t * phi_1, dims) * model_s
+                        - (0.5 / r1) * expand_dims(alpha_t * phi_1, dims) * (model_s1 - model_s)
+                )
+            elif solver_type == 'taylor':
+                x_t = (
+                        expand_dims(sigma_t / sigma_s, dims) * x
+                        - expand_dims(alpha_t * phi_1, dims) * model_s
+                        + (1. / r1) * expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * (
+                                    model_s1 - model_s)
+                )
+        else:
+            phi_11 = torch.expm1(r1 * h)
+            phi_1 = torch.expm1(h)
+
+            if model_s is None:
+                model_s = self.model_fn(x, s)
+            x_s1 = (
+                    expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x
+                    - expand_dims(sigma_s1 * phi_11, dims) * model_s
+            )
+            model_s1 = self.model_fn(x_s1, s1)
+            if solver_type == 'dpm_solver':
+                x_t = (
+                        expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
+                        - expand_dims(sigma_t * phi_1, dims) * model_s
+                        - (0.5 / r1) * expand_dims(sigma_t * phi_1, dims) * (model_s1 - model_s)
+                )
+            elif solver_type == 'taylor':
+                x_t = (
+                        expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
+                        - expand_dims(sigma_t * phi_1, dims) * model_s
+                        - (1. / r1) * expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * (model_s1 - model_s)
+                )
+        if return_intermediate:
+            return x_t, {'model_s': model_s, 'model_s1': model_s1}
+        else:
+            return x_t
+
+    def singlestep_dpm_solver_third_update(self, x, s, t, r1=1. / 3., r2=2. / 3., model_s=None, model_s1=None,
+                                           return_intermediate=False, solver_type='dpm_solver'):
+        """
+        Singlestep solver DPM-Solver-3 from time `s` to time `t`.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `s`.
+            s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
+            t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
+            r1: A `float`. The hyperparameter of the third-order solver.
+            r2: A `float`. The hyperparameter of the third-order solver.
+            model_s: A pytorch tensor. The model function evaluated at time `s`.
+                If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
+            model_s1: A pytorch tensor. The model function evaluated at time `s1` (the intermediate time given by `r1`).
+                If `model_s1` is None, we evaluate the model at `s1`; otherwise we directly use it.
+            return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times).
+            solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
+                The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
+        Returns:
+            x_t: A pytorch tensor. The approximated solution at time `t`.
+        """
+        if solver_type not in ['dpm_solver', 'taylor']:
+            raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
+        if r1 is None:
+            r1 = 1. / 3.
+        if r2 is None:
+            r2 = 2. / 3.
+        ns = self.noise_schedule
+        dims = x.dim()
+        lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
+        h = lambda_t - lambda_s
+        lambda_s1 = lambda_s + r1 * h
+        lambda_s2 = lambda_s + r2 * h
+        s1 = ns.inverse_lambda(lambda_s1)
+        s2 = ns.inverse_lambda(lambda_s2)
+        log_alpha_s, log_alpha_s1, log_alpha_s2, log_alpha_t = ns.marginal_log_mean_coeff(
+            s), ns.marginal_log_mean_coeff(s1), ns.marginal_log_mean_coeff(s2), ns.marginal_log_mean_coeff(t)
+        sigma_s, sigma_s1, sigma_s2, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(
+            s2), ns.marginal_std(t)
+        alpha_s1, alpha_s2, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_s2), torch.exp(log_alpha_t)
+
+        if self.predict_x0:
+            phi_11 = torch.expm1(-r1 * h)
+            phi_12 = torch.expm1(-r2 * h)
+            phi_1 = torch.expm1(-h)
+            phi_22 = torch.expm1(-r2 * h) / (r2 * h) + 1.
+            phi_2 = phi_1 / h + 1.
+            phi_3 = phi_2 / h - 0.5
+
+            if model_s is None:
+                model_s = self.model_fn(x, s)
+            if model_s1 is None:
+                x_s1 = (
+                        expand_dims(sigma_s1 / sigma_s, dims) * x
+                        - expand_dims(alpha_s1 * phi_11, dims) * model_s
+                )
+                model_s1 = self.model_fn(x_s1, s1)
+            x_s2 = (
+                    expand_dims(sigma_s2 / sigma_s, dims) * x
+                    - expand_dims(alpha_s2 * phi_12, dims) * model_s
+                    + r2 / r1 * expand_dims(alpha_s2 * phi_22, dims) * (model_s1 - model_s)
+            )
+            model_s2 = self.model_fn(x_s2, s2)
+            if solver_type == 'dpm_solver':
+                x_t = (
+                        expand_dims(sigma_t / sigma_s, dims) * x
+                        - expand_dims(alpha_t * phi_1, dims) * model_s
+                        + (1. / r2) * expand_dims(alpha_t * phi_2, dims) * (model_s2 - model_s)
+                )
+            elif solver_type == 'taylor':
+                D1_0 = (1. / r1) * (model_s1 - model_s)
+                D1_1 = (1. / r2) * (model_s2 - model_s)
+                D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)
+                D2 = 2. * (D1_1 - D1_0) / (r2 - r1)
+                x_t = (
+                        expand_dims(sigma_t / sigma_s, dims) * x
+                        - expand_dims(alpha_t * phi_1, dims) * model_s
+                        + expand_dims(alpha_t * phi_2, dims) * D1
+                        - expand_dims(alpha_t * phi_3, dims) * D2
+                )
+        else:
+            phi_11 = torch.expm1(r1 * h)
+            phi_12 = torch.expm1(r2 * h)
+            phi_1 = torch.expm1(h)
+            phi_22 = torch.expm1(r2 * h) / (r2 * h) - 1.
+            phi_2 = phi_1 / h - 1.
+            phi_3 = phi_2 / h - 0.5
+
+            if model_s is None:
+                model_s = self.model_fn(x, s)
+            if model_s1 is None:
+                x_s1 = (
+                        expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x
+                        - expand_dims(sigma_s1 * phi_11, dims) * model_s
+                )
+                model_s1 = self.model_fn(x_s1, s1)
+            x_s2 = (
+                    expand_dims(torch.exp(log_alpha_s2 - log_alpha_s), dims) * x
+                    - expand_dims(sigma_s2 * phi_12, dims) * model_s
+                    - r2 / r1 * expand_dims(sigma_s2 * phi_22, dims) * (model_s1 - model_s)
+            )
+            model_s2 = self.model_fn(x_s2, s2)
+            if solver_type == 'dpm_solver':
+                x_t = (
+                        expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
+                        - expand_dims(sigma_t * phi_1, dims) * model_s
+                        - (1. / r2) * expand_dims(sigma_t * phi_2, dims) * (model_s2 - model_s)
+                )
+            elif solver_type == 'taylor':
+                D1_0 = (1. / r1) * (model_s1 - model_s)
+                D1_1 = (1. / r2) * (model_s2 - model_s)
+                D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)
+                D2 = 2. * (D1_1 - D1_0) / (r2 - r1)
+                x_t = (
+                        expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
+                        - expand_dims(sigma_t * phi_1, dims) * model_s
+                        - expand_dims(sigma_t * phi_2, dims) * D1
+                        - expand_dims(sigma_t * phi_3, dims) * D2
+                )
+
+        if return_intermediate:
+            return x_t, {'model_s': model_s, 'model_s1': model_s1, 'model_s2': model_s2}
+        else:
+            return x_t
+
+    def multistep_dpm_solver_second_update(self, x, model_prev_list, t_prev_list, t, solver_type="dpm_solver"):
+        """
+        Multistep solver DPM-Solver-2 from time `t_prev_list[-1]` to time `t`.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `s`.
+            model_prev_list: A list of pytorch tensor. The previous computed model values.
+            t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)
+            t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
+            solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
+                The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
+        Returns:
+            x_t: A pytorch tensor. The approximated solution at time `t`.
+        """
+        if solver_type not in ['dpm_solver', 'taylor']:
+            raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
+        ns = self.noise_schedule
+        dims = x.dim()
+        model_prev_1, model_prev_0 = model_prev_list
+        t_prev_1, t_prev_0 = t_prev_list
+        lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_1), ns.marginal_lambda(
+            t_prev_0), ns.marginal_lambda(t)
+        log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
+        sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
+        alpha_t = torch.exp(log_alpha_t)
+
+        h_0 = lambda_prev_0 - lambda_prev_1
+        h = lambda_t - lambda_prev_0
+        r0 = h_0 / h
+        D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1)
+        if self.predict_x0:
+            if solver_type == 'dpm_solver':
+                x_t = (
+                        expand_dims(sigma_t / sigma_prev_0, dims) * x
+                        - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
+                        - 0.5 * expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * D1_0
+                )
+            elif solver_type == 'taylor':
+                x_t = (
+                        expand_dims(sigma_t / sigma_prev_0, dims) * x
+                        - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
+                        + expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1_0
+                )
+        else:
+            if solver_type == 'dpm_solver':
+                x_t = (
+                        expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
+                        - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
+                        - 0.5 * expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * D1_0
+                )
+            elif solver_type == 'taylor':
+                x_t = (
+                        expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
+                        - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
+                        - expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1_0
+                )
+        return x_t
+
+    def multistep_dpm_solver_third_update(self, x, model_prev_list, t_prev_list, t, solver_type='dpm_solver'):
+        """
+        Multistep solver DPM-Solver-3 from time `t_prev_list[-1]` to time `t`.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `s`.
+            model_prev_list: A list of pytorch tensor. The previous computed model values.
+            t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)
+            t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
+            solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
+                The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
+        Returns:
+            x_t: A pytorch tensor. The approximated solution at time `t`.
+        """
+        ns = self.noise_schedule
+        dims = x.dim()
+        model_prev_2, model_prev_1, model_prev_0 = model_prev_list
+        t_prev_2, t_prev_1, t_prev_0 = t_prev_list
+        lambda_prev_2, lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_2), ns.marginal_lambda(
+            t_prev_1), ns.marginal_lambda(t_prev_0), ns.marginal_lambda(t)
+        log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
+        sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
+        alpha_t = torch.exp(log_alpha_t)
+
+        h_1 = lambda_prev_1 - lambda_prev_2
+        h_0 = lambda_prev_0 - lambda_prev_1
+        h = lambda_t - lambda_prev_0
+        r0, r1 = h_0 / h, h_1 / h
+        D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1)
+        D1_1 = expand_dims(1. / r1, dims) * (model_prev_1 - model_prev_2)
+        D1 = D1_0 + expand_dims(r0 / (r0 + r1), dims) * (D1_0 - D1_1)
+        D2 = expand_dims(1. / (r0 + r1), dims) * (D1_0 - D1_1)
+        if self.predict_x0:
+            x_t = (
+                    expand_dims(sigma_t / sigma_prev_0, dims) * x
+                    - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
+                    + expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1
+                    - expand_dims(alpha_t * ((torch.exp(-h) - 1. + h) / h ** 2 - 0.5), dims) * D2
+            )
+        else:
+            x_t = (
+                    expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
+                    - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
+                    - expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1
+                    - expand_dims(sigma_t * ((torch.exp(h) - 1. - h) / h ** 2 - 0.5), dims) * D2
+            )
+        return x_t
+
+    def singlestep_dpm_solver_update(self, x, s, t, order, return_intermediate=False, solver_type='dpm_solver', r1=None,
+                                     r2=None):
+        """
+        Singlestep DPM-Solver with the order `order` from time `s` to time `t`.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `s`.
+            s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
+            t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
+            order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3.
+            return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times).
+            solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
+                The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
+            r1: A `float`. The hyperparameter of the second-order or third-order solver.
+            r2: A `float`. The hyperparameter of the third-order solver.
+        Returns:
+            x_t: A pytorch tensor. The approximated solution at time `t`.
+        """
+        if order == 1:
+            return self.dpm_solver_first_update(x, s, t, return_intermediate=return_intermediate)
+        elif order == 2:
+            return self.singlestep_dpm_solver_second_update(x, s, t, return_intermediate=return_intermediate,
+                                                            solver_type=solver_type, r1=r1)
+        elif order == 3:
+            return self.singlestep_dpm_solver_third_update(x, s, t, return_intermediate=return_intermediate,
+                                                           solver_type=solver_type, r1=r1, r2=r2)
+        else:
+            raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order))
+
+    def multistep_dpm_solver_update(self, x, model_prev_list, t_prev_list, t, order, solver_type='dpm_solver'):
+        """
+        Multistep DPM-Solver with the order `order` from time `t_prev_list[-1]` to time `t`.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `s`.
+            model_prev_list: A list of pytorch tensor. The previous computed model values.
+            t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)
+            t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
+            order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3.
+            solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
+                The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
+        Returns:
+            x_t: A pytorch tensor. The approximated solution at time `t`.
+        """
+        if order == 1:
+            return self.dpm_solver_first_update(x, t_prev_list[-1], t, model_s=model_prev_list[-1])
+        elif order == 2:
+            return self.multistep_dpm_solver_second_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type)
+        elif order == 3:
+            return self.multistep_dpm_solver_third_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type)
+        else:
+            raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order))
+
+    def dpm_solver_adaptive(self, x, order, t_T, t_0, h_init=0.05, atol=0.0078, rtol=0.05, theta=0.9, t_err=1e-5,
+                            solver_type='dpm_solver'):
+        """
+        The adaptive step size solver based on singlestep DPM-Solver.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `t_T`.
+            order: A `int`. The (higher) order of the solver. We only support order == 2 or 3.
+            t_T: A `float`. The starting time of the sampling (default is T).
+            t_0: A `float`. The ending time of the sampling (default is epsilon).
+            h_init: A `float`. The initial step size (for logSNR).
+            atol: A `float`. The absolute tolerance of the solver. For image data, the default setting is 0.0078, followed [1].
+            rtol: A `float`. The relative tolerance of the solver. The default setting is 0.05.
+            theta: A `float`. The safety hyperparameter for adapting the step size. The default setting is 0.9, followed [1].
+            t_err: A `float`. The tolerance for the time. We solve the diffusion ODE until the absolute error between the 
+                current time and `t_0` is less than `t_err`. The default setting is 1e-5.
+            solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
+                The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
+        Returns:
+            x_0: A pytorch tensor. The approximated solution at time `t_0`.
+
+        [1] A. Jolicoeur-Martineau, K. Li, R. Piché-Taillefer, T. Kachman, and I. Mitliagkas, "Gotta go fast when generating data with score-based models," arXiv preprint arXiv:2105.14080, 2021.
+        """
+        ns = self.noise_schedule
+        s = t_T * torch.ones((x.shape[0],)).to(x)
+        lambda_s = ns.marginal_lambda(s)
+        lambda_0 = ns.marginal_lambda(t_0 * torch.ones_like(s).to(x))
+        h = h_init * torch.ones_like(s).to(x)
+        x_prev = x
+        nfe = 0
+        if order == 2:
+            r1 = 0.5
+            lower_update = lambda x, s, t: self.dpm_solver_first_update(x, s, t, return_intermediate=True)
+            higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1,
+                                                                                               solver_type=solver_type,
+                                                                                               **kwargs)
+        elif order == 3:
+            r1, r2 = 1. / 3., 2. / 3.
+            lower_update = lambda x, s, t: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1,
+                                                                                    return_intermediate=True,
+                                                                                    solver_type=solver_type)
+            higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_third_update(x, s, t, r1=r1, r2=r2,
+                                                                                              solver_type=solver_type,
+                                                                                              **kwargs)
+        else:
+            raise ValueError("For adaptive step size solver, order must be 2 or 3, got {}".format(order))
+        while torch.abs((s - t_0)).mean() > t_err:
+            t = ns.inverse_lambda(lambda_s + h)
+            x_lower, lower_noise_kwargs = lower_update(x, s, t)
+            x_higher = higher_update(x, s, t, **lower_noise_kwargs)
+            delta = torch.max(torch.ones_like(x).to(x) * atol, rtol * torch.max(torch.abs(x_lower), torch.abs(x_prev)))
+            norm_fn = lambda v: torch.sqrt(torch.square(v.reshape((v.shape[0], -1))).mean(dim=-1, keepdim=True))
+            E = norm_fn((x_higher - x_lower) / delta).max()
+            if torch.all(E <= 1.):
+                x = x_higher
+                s = t
+                x_prev = x_lower
+                lambda_s = ns.marginal_lambda(s)
+            h = torch.min(theta * h * torch.float_power(E, -1. / order).float(), lambda_0 - lambda_s)
+            nfe += order
+        print('adaptive solver nfe', nfe)
+        return x
+
+    def sample(self, x, steps=20, t_start=None, t_end=None, order=3, skip_type='time_uniform',
+               method='singlestep', denoise=False, solver_type='dpm_solver', atol=0.0078,
+               rtol=0.05,
+               ):
+        """
+        Compute the sample at time `t_end` by DPM-Solver, given the initial `x` at time `t_start`.
+
+        =====================================================
+
+        We support the following algorithms for both noise prediction model and data prediction model:
+            - 'singlestep':
+                Singlestep DPM-Solver (i.e. "DPM-Solver-fast" in the paper), which combines different orders of singlestep DPM-Solver. 
+                We combine all the singlestep solvers with order <= `order` to use up all the function evaluations (steps).
+                The total number of function evaluations (NFE) == `steps`.
+                Given a fixed NFE == `steps`, the sampling procedure is:
+                    - If `order` == 1:
+                        - Denote K = steps. We use K steps of DPM-Solver-1 (i.e. DDIM).
+                    - If `order` == 2:
+                        - Denote K = (steps // 2) + (steps % 2). We take K intermediate time steps for sampling.
+                        - If steps % 2 == 0, we use K steps of singlestep DPM-Solver-2.
+                        - If steps % 2 == 1, we use (K - 1) steps of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1.
+                    - If `order` == 3:
+                        - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling.
+                        - If steps % 3 == 0, we use (K - 2) steps of singlestep DPM-Solver-3, and 1 step of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1.
+                        - If steps % 3 == 1, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of DPM-Solver-1.
+                        - If steps % 3 == 2, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of singlestep DPM-Solver-2.
+            - 'multistep':
+                Multistep DPM-Solver with the order of `order`. The total number of function evaluations (NFE) == `steps`.
+                We initialize the first `order` values by lower order multistep solvers.
+                Given a fixed NFE == `steps`, the sampling procedure is:
+                    Denote K = steps.
+                    - If `order` == 1:
+                        - We use K steps of DPM-Solver-1 (i.e. DDIM).
+                    - If `order` == 2:
+                        - We firstly use 1 step of DPM-Solver-1, then use (K - 1) step of multistep DPM-Solver-2.
+                    - If `order` == 3:
+                        - We firstly use 1 step of DPM-Solver-1, then 1 step of multistep DPM-Solver-2, then (K - 2) step of multistep DPM-Solver-3.
+            - 'singlestep_fixed':
+                Fixed order singlestep DPM-Solver (i.e. DPM-Solver-1 or singlestep DPM-Solver-2 or singlestep DPM-Solver-3).
+                We use singlestep DPM-Solver-`order` for `order`=1 or 2 or 3, with total [`steps` // `order`] * `order` NFE.
+            - 'adaptive':
+                Adaptive step size DPM-Solver (i.e. "DPM-Solver-12" and "DPM-Solver-23" in the paper).
+                We ignore `steps` and use adaptive step size DPM-Solver with a higher order of `order`.
+                You can adjust the absolute tolerance `atol` and the relative tolerance `rtol` to balance the computatation costs
+                (NFE) and the sample quality.
+                    - If `order` == 2, we use DPM-Solver-12 which combines DPM-Solver-1 and singlestep DPM-Solver-2.
+                    - If `order` == 3, we use DPM-Solver-23 which combines singlestep DPM-Solver-2 and singlestep DPM-Solver-3.
+
+        =====================================================
+
+        Some advices for choosing the algorithm:
+            - For **unconditional sampling** or **guided sampling with small guidance scale** by DPMs:
+                Use singlestep DPM-Solver ("DPM-Solver-fast" in the paper) with `order = 3`.
+                e.g.
+                    >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=False)
+                    >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=3,
+                            skip_type='time_uniform', method='singlestep')
+            - For **guided sampling with large guidance scale** by DPMs:
+                Use multistep DPM-Solver with `predict_x0 = True` and `order = 2`.
+                e.g.
+                    >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=True)
+                    >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=2,
+                            skip_type='time_uniform', method='multistep')
+
+        We support three types of `skip_type`:
+            - 'logSNR': uniform logSNR for the time steps. **Recommended for low-resolutional images**
+            - 'time_uniform': uniform time for the time steps. **Recommended for high-resolutional images**.
+            - 'time_quadratic': quadratic time for the time steps.
+
+        =====================================================
+        Args:
+            x: A pytorch tensor. The initial value at time `t_start`
+                e.g. if `t_start` == T, then `x` is a sample from the standard normal distribution.
+            steps: A `int`. The total number of function evaluations (NFE).
+            t_start: A `float`. The starting time of the sampling.
+                If `T` is None, we use self.noise_schedule.T (default is 1.0).
+            t_end: A `float`. The ending time of the sampling.
+                If `t_end` is None, we use 1. / self.noise_schedule.total_N.
+                e.g. if total_N == 1000, we have `t_end` == 1e-3.
+                For discrete-time DPMs:
+                    - We recommend `t_end` == 1. / self.noise_schedule.total_N.
+                For continuous-time DPMs:
+                    - We recommend `t_end` == 1e-3 when `steps` <= 15; and `t_end` == 1e-4 when `steps` > 15.
+            order: A `int`. The order of DPM-Solver.
+            skip_type: A `str`. The type for the spacing of the time steps. 'time_uniform' or 'logSNR' or 'time_quadratic'.
+            method: A `str`. The method for sampling. 'singlestep' or 'multistep' or 'singlestep_fixed' or 'adaptive'.
+            denoise: A `bool`. Whether to denoise at the final step. Default is False.
+                If `denoise` is True, the total NFE is (`steps` + 1).
+            solver_type: A `str`. The taylor expansion type for the solver. `dpm_solver` or `taylor`. We recommend `dpm_solver`.
+            atol: A `float`. The absolute tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.
+            rtol: A `float`. The relative tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.
+        Returns:
+            x_end: A pytorch tensor. The approximated solution at time `t_end`.
+
+        """
+        t_0 = 1. / self.noise_schedule.total_N if t_end is None else t_end
+        t_T = self.noise_schedule.T if t_start is None else t_start
+        device = x.device
+        if method == 'adaptive':
+            with torch.no_grad():
+                x = self.dpm_solver_adaptive(x, order=order, t_T=t_T, t_0=t_0, atol=atol, rtol=rtol,
+                                             solver_type=solver_type)
+        elif method == 'multistep':
+            assert steps >= order
+            timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device)
+            assert timesteps.shape[0] - 1 == steps
+            with torch.no_grad():
+                vec_t = timesteps[0].expand((x.shape[0]))
+                model_prev_list = [self.model_fn(x, vec_t)]
+                t_prev_list = [vec_t]
+                # Init the first `order` values by lower order multistep DPM-Solver.
+                for init_order in range(1, order):
+                    vec_t = timesteps[init_order].expand(x.shape[0])
+                    x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, vec_t, init_order,
+                                                         solver_type=solver_type)
+                    model_prev_list.append(self.model_fn(x, vec_t))
+                    t_prev_list.append(vec_t)
+                # Compute the remaining values by `order`-th order multistep DPM-Solver.
+                for step in range(order, steps + 1):
+                    vec_t = timesteps[step].expand(x.shape[0])
+                    x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, vec_t, order,
+                                                         solver_type=solver_type)
+                    for i in range(order - 1):
+                        t_prev_list[i] = t_prev_list[i + 1]
+                        model_prev_list[i] = model_prev_list[i + 1]
+                    t_prev_list[-1] = vec_t
+                    # We do not need to evaluate the final model value.
+                    if step < steps:
+                        model_prev_list[-1] = self.model_fn(x, vec_t)
+        elif method in ['singlestep', 'singlestep_fixed']:
+            if method == 'singlestep':
+                timesteps_outer, orders = self.get_orders_and_timesteps_for_singlestep_solver(steps=steps, order=order,
+                                                                                              skip_type=skip_type,
+                                                                                              t_T=t_T, t_0=t_0,
+                                                                                              device=device)
+            elif method == 'singlestep_fixed':
+                K = steps // order
+                orders = [order, ] * K
+                timesteps_outer = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=K, device=device)
+            for i, order in enumerate(orders):
+                t_T_inner, t_0_inner = timesteps_outer[i], timesteps_outer[i + 1]
+                timesteps_inner = self.get_time_steps(skip_type=skip_type, t_T=t_T_inner.item(), t_0=t_0_inner.item(),
+                                                      N=order, device=device)
+                lambda_inner = self.noise_schedule.marginal_lambda(timesteps_inner)
+                vec_s, vec_t = t_T_inner.repeat(x.shape[0]), t_0_inner.repeat(x.shape[0])
+                h = lambda_inner[-1] - lambda_inner[0]
+                r1 = None if order <= 1 else (lambda_inner[1] - lambda_inner[0]) / h
+                r2 = None if order <= 2 else (lambda_inner[2] - lambda_inner[0]) / h
+                x = self.singlestep_dpm_solver_update(x, vec_s, vec_t, order, solver_type=solver_type, r1=r1, r2=r2)
+        if denoise:
+            x = self.denoise_fn(x, torch.ones((x.shape[0],)).to(device) * t_0)
+        return x
+
+
+#############################################################
+# other utility functions
+#############################################################
+
+def interpolate_fn(x, xp, yp):
+    """
+    A piecewise linear function y = f(x), using xp and yp as keypoints.
+    We implement f(x) in a differentiable way (i.e. applicable for autograd).
+    The function f(x) is well-defined for all x-axis. (For x beyond the bounds of xp, we use the outmost points of xp to define the linear function.)
+
+    Args:
+        x: PyTorch tensor with shape [N, C], where N is the batch size, C is the number of channels (we use C = 1 for DPM-Solver).
+        xp: PyTorch tensor with shape [C, K], where K is the number of keypoints.
+        yp: PyTorch tensor with shape [C, K].
+    Returns:
+        The function values f(x), with shape [N, C].
+    """
+    N, K = x.shape[0], xp.shape[1]
+    all_x = torch.cat([x.unsqueeze(2), xp.unsqueeze(0).repeat((N, 1, 1))], dim=2)
+    sorted_all_x, x_indices = torch.sort(all_x, dim=2)
+    x_idx = torch.argmin(x_indices, dim=2)
+    cand_start_idx = x_idx - 1
+    start_idx = torch.where(
+        torch.eq(x_idx, 0),
+        torch.tensor(1, device=x.device),
+        torch.where(
+            torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
+        ),
+    )
+    end_idx = torch.where(torch.eq(start_idx, cand_start_idx), start_idx + 2, start_idx + 1)
+    start_x = torch.gather(sorted_all_x, dim=2, index=start_idx.unsqueeze(2)).squeeze(2)
+    end_x = torch.gather(sorted_all_x, dim=2, index=end_idx.unsqueeze(2)).squeeze(2)
+    start_idx2 = torch.where(
+        torch.eq(x_idx, 0),
+        torch.tensor(0, device=x.device),
+        torch.where(
+            torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
+        ),
+    )
+    y_positions_expanded = yp.unsqueeze(0).expand(N, -1, -1)
+    start_y = torch.gather(y_positions_expanded, dim=2, index=start_idx2.unsqueeze(2)).squeeze(2)
+    end_y = torch.gather(y_positions_expanded, dim=2, index=(start_idx2 + 1).unsqueeze(2)).squeeze(2)
+    cand = start_y + (x - start_x) * (end_y - start_y) / (end_x - start_x)
+    return cand
+
+
+def expand_dims(v, dims):
+    """
+    Expand the tensor `v` to the dim `dims`.
+
+    Args:
+        `v`: a PyTorch tensor with shape [N].
+        `dim`: a `int`.
+    Returns:
+        a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`.
+    """
+    return v[(...,) + (None,) * (dims - 1)]

diffusion/how to export onnx.md

@@ -0,0 +1,4 @@
+- Open [onnx_export](onnx_export.py)
+- project_name = "dddsp" change "project_name" to your project name
+- model_path = f'{project_name}/model_500000.pt' change "model_path" to your model path
+- Run

diffusion/infer_gt_mel.py

@@ -0,0 +1,74 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+from diffusion.unit2mel import load_model_vocoder
+
+
+class DiffGtMel:
+    def __init__(self, project_path=None, device=None):
+        self.project_path = project_path
+        if device is not None:
+            self.device = device
+        else:
+            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.model = None
+        self.vocoder = None
+        self.args = None
+
+    def flush_model(self, project_path, ddsp_config=None):
+        if (self.model is None) or (project_path != self.project_path):
+            model, vocoder, args = load_model_vocoder(project_path, device=self.device)
+            if self.check_args(ddsp_config, args):
+                self.model = model
+                self.vocoder = vocoder
+                self.args = args
+
+    def check_args(self, args1, args2):
+        if args1.data.block_size != args2.data.block_size:
+            raise ValueError("DDSP与DIFF模型的block_size不一致")
+        if args1.data.sampling_rate != args2.data.sampling_rate:
+            raise ValueError("DDSP与DIFF模型的sampling_rate不一致")
+        if args1.data.encoder != args2.data.encoder:
+            raise ValueError("DDSP与DIFF模型的encoder不一致")
+        return True
+
+    def __call__(self, audio, f0, hubert, volume, acc=1, spk_id=1, k_step=0, method='pndm',
+                 spk_mix_dict=None, start_frame=0):
+        input_mel = self.vocoder.extract(audio, self.args.data.sampling_rate)
+        out_mel = self.model(
+            hubert,
+            f0,
+            volume,
+            spk_id=spk_id,
+            spk_mix_dict=spk_mix_dict,
+            gt_spec=input_mel,
+            infer=True,
+            infer_speedup=acc,
+            method=method,
+            k_step=k_step,
+            use_tqdm=False)
+        if start_frame > 0:
+            out_mel = out_mel[:, start_frame:, :]
+            f0 = f0[:, start_frame:, :]
+        output = self.vocoder.infer(out_mel, f0)
+        if start_frame > 0:
+            output = F.pad(output, (start_frame * self.vocoder.vocoder_hop_size, 0))
+        return output
+
+    def infer(self, audio, f0, hubert, volume, acc=1, spk_id=1, k_step=0, method='pndm', silence_front=0,
+              use_silence=False, spk_mix_dict=None):
+        start_frame = int(silence_front * self.vocoder.vocoder_sample_rate / self.vocoder.vocoder_hop_size)
+        if use_silence:
+            audio = audio[:, start_frame * self.vocoder.vocoder_hop_size:]
+            f0 = f0[:, start_frame:, :]
+            hubert = hubert[:, start_frame:, :]
+            volume = volume[:, start_frame:, :]
+            _start_frame = 0
+        else:
+            _start_frame = start_frame
+        audio = self.__call__(audio, f0, hubert, volume, acc=acc, spk_id=spk_id, k_step=k_step,
+                              method=method, spk_mix_dict=spk_mix_dict, start_frame=_start_frame)
+        if use_silence:
+            if start_frame > 0:
+                audio = F.pad(audio, (start_frame * self.vocoder.vocoder_hop_size, 0))
+        return audio

diffusion/logger/saver.py

@@ -0,0 +1,150 @@
+'''
+author: wayn391@mastertones
+'''
+
+import os
+import json
+import time
+import yaml
+import datetime
+import torch
+import matplotlib.pyplot as plt
+from . import utils
+from torch.utils.tensorboard import SummaryWriter
+
+class Saver(object):
+    def __init__(
+            self, 
+            args,
+            initial_global_step=-1):
+
+        self.expdir = args.env.expdir
+        self.sample_rate = args.data.sampling_rate
+        
+        # cold start
+        self.global_step = initial_global_step
+        self.init_time = time.time()
+        self.last_time = time.time()
+
+        # makedirs
+        os.makedirs(self.expdir, exist_ok=True)       
+
+        # path
+        self.path_log_info = os.path.join(self.expdir, 'log_info.txt')
+
+        # ckpt
+        os.makedirs(self.expdir, exist_ok=True)       
+
+        # writer
+        self.writer = SummaryWriter(os.path.join(self.expdir, 'logs'))
+        
+        # save config
+        path_config = os.path.join(self.expdir, 'config.yaml')
+        with open(path_config, "w") as out_config:
+            yaml.dump(dict(args), out_config)
+
+
+    def log_info(self, msg):
+        '''log method'''
+        if isinstance(msg, dict):
+            msg_list = []
+            for k, v in msg.items():
+                tmp_str = ''
+                if isinstance(v, int):
+                    tmp_str = '{}: {:,}'.format(k, v)
+                else:
+                    tmp_str = '{}: {}'.format(k, v)
+
+                msg_list.append(tmp_str)
+            msg_str = '\n'.join(msg_list)
+        else:
+            msg_str = msg
+        
+        # dsplay
+        print(msg_str)
+
+        # save
+        with open(self.path_log_info, 'a') as fp:
+            fp.write(msg_str+'\n')
+
+    def log_value(self, dict):
+        for k, v in dict.items():
+            self.writer.add_scalar(k, v, self.global_step)
+    
+    def log_spec(self, name, spec, spec_out, vmin=-14, vmax=3.5):  
+        spec_cat = torch.cat([(spec_out - spec).abs() + vmin, spec, spec_out], -1)
+        spec = spec_cat[0]
+        if isinstance(spec, torch.Tensor):
+            spec = spec.cpu().numpy()
+        fig = plt.figure(figsize=(12, 9))
+        plt.pcolor(spec.T, vmin=vmin, vmax=vmax)
+        plt.tight_layout()
+        self.writer.add_figure(name, fig, self.global_step)
+    
+    def log_audio(self, dict):
+        for k, v in dict.items():
+            self.writer.add_audio(k, v, global_step=self.global_step, sample_rate=self.sample_rate)
+    
+    def get_interval_time(self, update=True):
+        cur_time = time.time()
+        time_interval = cur_time - self.last_time
+        if update:
+            self.last_time = cur_time
+        return time_interval
+
+    def get_total_time(self, to_str=True):
+        total_time = time.time() - self.init_time
+        if to_str:
+            total_time = str(datetime.timedelta(
+                seconds=total_time))[:-5]
+        return total_time
+
+    def save_model(
+            self,
+            model, 
+            optimizer,
+            name='model',
+            postfix='',
+            to_json=False):
+        # path
+        if postfix:
+            postfix = '_' + postfix
+        path_pt = os.path.join(
+            self.expdir , name+postfix+'.pt')
+       
+        # check
+        print(' [*] model checkpoint saved: {}'.format(path_pt))
+
+        # save
+        if optimizer is not None:
+            torch.save({
+                'global_step': self.global_step,
+                'model': model.state_dict(),
+                'optimizer': optimizer.state_dict()}, path_pt)
+        else:
+            torch.save({
+                'global_step': self.global_step,
+                'model': model.state_dict()}, path_pt)
+            
+        # to json
+        if to_json:
+            path_json = os.path.join(
+                self.expdir , name+'.json')
+            utils.to_json(path_params, path_json)
+    
+    def delete_model(self, name='model', postfix=''):
+        # path
+        if postfix:
+            postfix = '_' + postfix
+        path_pt = os.path.join(
+            self.expdir , name+postfix+'.pt')
+       
+        # delete
+        if os.path.exists(path_pt):
+            os.remove(path_pt)
+            print(' [*] model checkpoint deleted: {}'.format(path_pt))
+        
+    def global_step_increment(self):
+        self.global_step += 1
+
+

diffusion/logger/utils.py

@@ -0,0 +1,126 @@
+import os
+import yaml
+import json
+import pickle
+import torch
+
+def traverse_dir(
+        root_dir,
+        extensions,
+        amount=None,
+        str_include=None,
+        str_exclude=None,
+        is_pure=False,
+        is_sort=False,
+        is_ext=True):
+
+    file_list = []
+    cnt = 0
+    for root, _, files in os.walk(root_dir):
+        for file in files:
+            if any([file.endswith(f".{ext}") for ext in extensions]):
+                # path
+                mix_path = os.path.join(root, file)
+                pure_path = mix_path[len(root_dir)+1:] if is_pure else mix_path
+
+                # amount
+                if (amount is not None) and (cnt == amount):
+                    if is_sort:
+                        file_list.sort()
+                    return file_list
+                
+                # check string
+                if (str_include is not None) and (str_include not in pure_path):
+                    continue
+                if (str_exclude is not None) and (str_exclude in pure_path):
+                    continue
+                
+                if not is_ext:
+                    ext = pure_path.split('.')[-1]
+                    pure_path = pure_path[:-(len(ext)+1)]
+                file_list.append(pure_path)
+                cnt += 1
+    if is_sort:
+        file_list.sort()
+    return file_list
+
+
+    
+class DotDict(dict):
+    def __getattr__(*args):         
+        val = dict.get(*args)         
+        return DotDict(val) if type(val) is dict else val   
+
+    __setattr__ = dict.__setitem__    
+    __delattr__ = dict.__delitem__
+
+
+def get_network_paras_amount(model_dict):
+    info = dict()
+    for model_name, model in model_dict.items():
+        # all_params = sum(p.numel() for p in model.parameters())
+        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+        info[model_name] = trainable_params
+    return info
+
+
+def load_config(path_config):
+    with open(path_config, "r") as config:
+        args = yaml.safe_load(config)
+    args = DotDict(args)
+    # print(args)
+    return args
+
+def save_config(path_config,config):
+    config = dict(config)
+    with open(path_config, "w") as f:
+        yaml.dump(config, f)
+
+def to_json(path_params, path_json):
+    params = torch.load(path_params, map_location=torch.device('cpu'))
+    raw_state_dict = {}
+    for k, v in params.items():
+        val = v.flatten().numpy().tolist()
+        raw_state_dict[k] = val
+
+    with open(path_json, 'w') as outfile:
+        json.dump(raw_state_dict, outfile,indent= "\t")
+
+
+def convert_tensor_to_numpy(tensor, is_squeeze=True):
+    if is_squeeze:
+        tensor = tensor.squeeze()
+    if tensor.requires_grad:
+        tensor = tensor.detach()
+    if tensor.is_cuda:
+        tensor = tensor.cpu()
+    return tensor.numpy()
+
+           
+def load_model(
+        expdir, 
+        model,
+        optimizer,
+        name='model',
+        postfix='',
+        device='cpu'):
+    if postfix == '':
+        postfix = '_' + postfix
+    path = os.path.join(expdir, name+postfix)
+    path_pt = traverse_dir(expdir, ['pt'], is_ext=False)
+    global_step = 0
+    if len(path_pt) > 0:
+        steps = [s[len(path):] for s in path_pt]
+        maxstep = max([int(s) if s.isdigit() else 0 for s in steps])
+        if maxstep >= 0:
+            path_pt = path+str(maxstep)+'.pt'
+        else:
+            path_pt = path+'best.pt'
+        print(' [*] restoring model from', path_pt)
+        ckpt = torch.load(path_pt, map_location=torch.device(device))
+        global_step = ckpt['global_step']
+        model.load_state_dict(ckpt['model'], strict=False)
+        if ckpt.get('optimizer') != None:
+            optimizer.load_state_dict(ckpt['optimizer'])
+    return global_step, model, optimizer

diffusion/onnx_export.py

@@ -0,0 +1,226 @@
+from diffusion_onnx import GaussianDiffusion
+import os
+import yaml
+import torch
+import torch.nn as nn
+import numpy as np
+from wavenet import WaveNet
+import torch.nn.functional as F
+import diffusion
+
+class DotDict(dict):
+    def __getattr__(*args):         
+        val = dict.get(*args)         
+        return DotDict(val) if type(val) is dict else val   
+
+    __setattr__ = dict.__setitem__    
+    __delattr__ = dict.__delitem__
+
+    
+def load_model_vocoder(
+        model_path,
+        device='cpu'):
+    config_file = os.path.join(os.path.split(model_path)[0], 'config.yaml')
+    with open(config_file, "r") as config:
+        args = yaml.safe_load(config)
+    args = DotDict(args)
+    
+    # load model
+    model = Unit2Mel(
+                args.data.encoder_out_channels, 
+                args.model.n_spk,
+                args.model.use_pitch_aug,
+                128,
+                args.model.n_layers,
+                args.model.n_chans,
+                args.model.n_hidden)
+    
+    print(' [Loading] ' + model_path)
+    ckpt = torch.load(model_path, map_location=torch.device(device))
+    model.to(device)
+    model.load_state_dict(ckpt['model'])
+    model.eval()
+    return model, args
+
+
+class Unit2Mel(nn.Module):
+    def __init__(
+            self,
+            input_channel,
+            n_spk,
+            use_pitch_aug=False,
+            out_dims=128,
+            n_layers=20, 
+            n_chans=384, 
+            n_hidden=256):
+        super().__init__()
+        self.unit_embed = nn.Linear(input_channel, n_hidden)
+        self.f0_embed = nn.Linear(1, n_hidden)
+        self.volume_embed = nn.Linear(1, n_hidden)
+        if use_pitch_aug:
+            self.aug_shift_embed = nn.Linear(1, n_hidden, bias=False)
+        else:
+            self.aug_shift_embed = None
+        self.n_spk = n_spk
+        if n_spk is not None and n_spk > 1:
+            self.spk_embed = nn.Embedding(n_spk, n_hidden)
+            
+        # diffusion
+        self.decoder = GaussianDiffusion(out_dims, n_layers, n_chans, n_hidden)
+        self.hidden_size = n_hidden
+        self.speaker_map = torch.zeros((self.n_spk,1,1,n_hidden))
+    
+        
+
+    def forward(self, units, mel2ph, f0, volume, g = None):
+        
+        '''
+        input: 
+            B x n_frames x n_unit
+        return: 
+            dict of B x n_frames x feat
+        '''
+
+        decoder_inp = F.pad(units, [0, 0, 1, 0])
+        mel2ph_ = mel2ph.unsqueeze(2).repeat([1, 1, units.shape[-1]])
+        units = torch.gather(decoder_inp, 1, mel2ph_)  # [B, T, H]
+
+        x = self.unit_embed(units) + self.f0_embed((1 + f0.unsqueeze(-1) / 700).log()) + self.volume_embed(volume.unsqueeze(-1))
+
+        if self.n_spk is not None and self.n_spk > 1:   # [N, S]  *  [S, B, 1, H]
+            g = g.reshape((g.shape[0], g.shape[1], 1, 1, 1))  # [N, S, B, 1, 1]
+            g = g * self.speaker_map  # [N, S, B, 1, H]
+            g = torch.sum(g, dim=1) # [N, 1, B, 1, H]
+            g = g.transpose(0, -1).transpose(0, -2).squeeze(0) # [B, H, N]
+            x = x.transpose(1, 2) + g
+            return x
+        else:
+            return x.transpose(1, 2)
+        
+
+    def init_spkembed(self, units, f0, volume, spk_id = None, spk_mix_dict = None, aug_shift = None,
+                gt_spec=None, infer=True, infer_speedup=10, method='dpm-solver', k_step=300, use_tqdm=True):
+        
+        '''
+        input: 
+            B x n_frames x n_unit
+        return: 
+            dict of B x n_frames x feat
+        '''
+        x = self.unit_embed(units) + self.f0_embed((1+ f0 / 700).log()) + self.volume_embed(volume)
+        if self.n_spk is not None and self.n_spk > 1:
+            if spk_mix_dict is not None:
+                spk_embed_mix = torch.zeros((1,1,self.hidden_size))
+                for k, v in spk_mix_dict.items():
+                    spk_id_torch = torch.LongTensor(np.array([[k]])).to(units.device)
+                    spk_embeddd = self.spk_embed(spk_id_torch)
+                    self.speaker_map[k] = spk_embeddd
+                    spk_embed_mix = spk_embed_mix + v * spk_embeddd
+                x = x + spk_embed_mix
+            else:
+                x = x + self.spk_embed(spk_id - 1)
+        self.speaker_map = self.speaker_map.unsqueeze(0)
+        self.speaker_map = self.speaker_map.detach()
+        return x.transpose(1, 2)
+
+    def OnnxExport(self, project_name=None, init_noise=None, export_encoder=True, export_denoise=True, export_pred=True, export_after=True):
+        hubert_hidden_size = 768
+        n_frames = 100
+        hubert = torch.randn((1, n_frames, hubert_hidden_size))
+        mel2ph = torch.arange(end=n_frames).unsqueeze(0).long()
+        f0 = torch.randn((1, n_frames))
+        volume = torch.randn((1, n_frames))
+        spk_mix = []
+        spks = {}
+        if self.n_spk is not None and self.n_spk > 1:
+            for i in range(self.n_spk):
+                spk_mix.append(1.0/float(self.n_spk))
+                spks.update({i:1.0/float(self.n_spk)})
+        spk_mix = torch.tensor(spk_mix)
+        spk_mix = spk_mix.repeat(n_frames, 1)
+        orgouttt = self.init_spkembed(hubert, f0.unsqueeze(-1), volume.unsqueeze(-1), spk_mix_dict=spks)
+        outtt = self.forward(hubert, mel2ph, f0, volume, spk_mix)
+        if export_encoder:
+            torch.onnx.export(
+                self,
+                (hubert, mel2ph, f0, volume, spk_mix),
+                f"{project_name}_encoder.onnx",
+                input_names=["hubert", "mel2ph", "f0", "volume", "spk_mix"],
+                output_names=["mel_pred"],
+                dynamic_axes={
+                    "hubert": [1],
+                    "f0": [1],
+                    "volume": [1],
+                    "mel2ph": [1],
+                    "spk_mix": [0],
+                },
+                opset_version=16
+            )
+        
+        self.decoder.OnnxExport(project_name, init_noise=init_noise, export_denoise=export_denoise, export_pred=export_pred, export_after=export_after)
+
+    def ExportOnnx(self, project_name=None):
+        hubert_hidden_size = 768
+        n_frames = 100
+        hubert = torch.randn((1, n_frames, hubert_hidden_size))
+        mel2ph = torch.arange(end=n_frames).unsqueeze(0).long()
+        f0 = torch.randn((1, n_frames))
+        volume = torch.randn((1, n_frames))
+        spk_mix = []
+        spks = {}
+        if self.n_spk is not None and self.n_spk > 1:
+            for i in range(self.n_spk):
+                spk_mix.append(1.0/float(self.n_spk))
+                spks.update({i:1.0/float(self.n_spk)})
+        spk_mix = torch.tensor(spk_mix)
+        orgouttt = self.orgforward(hubert, f0.unsqueeze(-1), volume.unsqueeze(-1), spk_mix_dict=spks)
+        outtt = self.forward(hubert, mel2ph, f0, volume, spk_mix)
+
+        torch.onnx.export(
+                self,
+                (hubert, mel2ph, f0, volume, spk_mix),
+                f"{project_name}_encoder.onnx",
+                input_names=["hubert", "mel2ph", "f0", "volume", "spk_mix"],
+                output_names=["mel_pred"],
+                dynamic_axes={
+                    "hubert": [1],
+                    "f0": [1],
+                    "volume": [1],
+                    "mel2ph": [1]
+                },
+                opset_version=16
+            )
+
+        condition = torch.randn(1,self.decoder.n_hidden,n_frames)
+        noise = torch.randn((1, 1, self.decoder.mel_bins, condition.shape[2]), dtype=torch.float32)
+        pndm_speedup = torch.LongTensor([100])
+        K_steps = torch.LongTensor([1000])
+        self.decoder = torch.jit.script(self.decoder)
+        self.decoder(condition, noise, pndm_speedup, K_steps)
+
+        torch.onnx.export(
+                self.decoder,
+                (condition, noise, pndm_speedup, K_steps),
+                f"{project_name}_diffusion.onnx",
+                input_names=["condition", "noise", "pndm_speedup", "K_steps"],
+                output_names=["mel"],
+                dynamic_axes={
+                    "condition": [2],
+                    "noise": [3],
+                },
+                opset_version=16
+            )
+
+
+if __name__ == "__main__":
+    project_name = "dddsp"
+    model_path = f'{project_name}/model_500000.pt'
+
+    model, _ = load_model_vocoder(model_path)
+
+    # 分开Diffusion导出（需要使用MoeSS/MoeVoiceStudio或者自己编写Pndm/Dpm采样）
+    model.OnnxExport(project_name, export_encoder=True, export_denoise=True, export_pred=True, export_after=True)
+
+    # 合并Diffusion导出（Encoder和Diffusion分开，直接将Encoder的结果和初始噪声输入Diffusion即可）
+    # model.ExportOnnx(project_name)
+

diffusion/solver.py

@@ -0,0 +1,195 @@
+import os
+import time
+import numpy as np
+import torch
+import librosa
+from diffusion.logger.saver import Saver
+from diffusion.logger import utils
+from torch import autocast
+from torch.cuda.amp import GradScaler
+
+def test(args, model, vocoder, loader_test, saver):
+    print(' [*] testing...')
+    model.eval()
+
+    # losses
+    test_loss = 0.
+    
+    # intialization
+    num_batches = len(loader_test)
+    rtf_all = []
+    
+    # run
+    with torch.no_grad():
+        for bidx, data in enumerate(loader_test):
+            fn = data['name'][0].split("/")[-1]
+            speaker = data['name'][0].split("/")[-2]
+            print('--------')
+            print('{}/{} - {}'.format(bidx, num_batches, fn))
+
+            # unpack data
+            for k in data.keys():
+                if not k.startswith('name'):
+                    data[k] = data[k].to(args.device)
+            print('>>', data['name'][0])
+
+            # forward
+            st_time = time.time()
+            mel = model(
+                    data['units'], 
+                    data['f0'], 
+                    data['volume'], 
+                    data['spk_id'],
+                    gt_spec=None,
+                    infer=True, 
+                    infer_speedup=args.infer.speedup, 
+                    method=args.infer.method)
+            signal = vocoder.infer(mel, data['f0'])
+            ed_time = time.time()
+                        
+            # RTF
+            run_time = ed_time - st_time
+            song_time = signal.shape[-1] / args.data.sampling_rate
+            rtf = run_time / song_time
+            print('RTF: {}  | {} / {}'.format(rtf, run_time, song_time))
+            rtf_all.append(rtf)
+           
+            # loss
+            for i in range(args.train.batch_size):
+                loss = model(
+                    data['units'], 
+                    data['f0'], 
+                    data['volume'], 
+                    data['spk_id'], 
+                    gt_spec=data['mel'],
+                    infer=False)
+                test_loss += loss.item()
+            
+            # log mel
+            saver.log_spec(f"{speaker}_{fn}.wav", data['mel'], mel)
+            
+            # log audi
+            path_audio = data['name_ext'][0]
+            audio, sr = librosa.load(path_audio, sr=args.data.sampling_rate)
+            if len(audio.shape) > 1:
+                audio = librosa.to_mono(audio)
+            audio = torch.from_numpy(audio).unsqueeze(0).to(signal)
+            saver.log_audio({f"{speaker}_{fn}_gt.wav": audio,f"{speaker}_{fn}_pred.wav": signal})
+    # report
+    test_loss /= args.train.batch_size
+    test_loss /= num_batches 
+    
+    # check
+    print(' [test_loss] test_loss:', test_loss)
+    print(' Real Time Factor', np.mean(rtf_all))
+    return test_loss
+
+
+def train(args, initial_global_step, model, optimizer, scheduler, vocoder, loader_train, loader_test):
+    # saver
+    saver = Saver(args, initial_global_step=initial_global_step)
+
+    # model size
+    params_count = utils.get_network_paras_amount({'model': model})
+    saver.log_info('--- model size ---')
+    saver.log_info(params_count)
+    
+    # run
+    num_batches = len(loader_train)
+    model.train()
+    saver.log_info('======= start training =======')
+    scaler = GradScaler()
+    if args.train.amp_dtype == 'fp32':
+        dtype = torch.float32
+    elif args.train.amp_dtype == 'fp16':
+        dtype = torch.float16
+    elif args.train.amp_dtype == 'bf16':
+        dtype = torch.bfloat16
+    else:
+        raise ValueError(' [x] Unknown amp_dtype: ' + args.train.amp_dtype)
+    saver.log_info("epoch|batch_idx/num_batches|output_dir|batch/s|lr|time|step")
+    for epoch in range(args.train.epochs):
+        for batch_idx, data in enumerate(loader_train):
+            saver.global_step_increment()
+            optimizer.zero_grad()
+
+            # unpack data
+            for k in data.keys():
+                if not k.startswith('name'):
+                    data[k] = data[k].to(args.device)
+            
+            # forward
+            if dtype == torch.float32:
+                loss = model(data['units'].float(), data['f0'], data['volume'], data['spk_id'], 
+                                aug_shift = data['aug_shift'], gt_spec=data['mel'].float(), infer=False)
+            else:
+                with autocast(device_type=args.device, dtype=dtype):
+                    loss = model(data['units'], data['f0'], data['volume'], data['spk_id'], 
+                                    aug_shift = data['aug_shift'], gt_spec=data['mel'], infer=False)
+            
+            # handle nan loss
+            if torch.isnan(loss):
+                raise ValueError(' [x] nan loss ')
+            else:
+                # backpropagate
+                if dtype == torch.float32:
+                    loss.backward()
+                    optimizer.step()
+                else:
+                    scaler.scale(loss).backward()
+                    scaler.step(optimizer)
+                    scaler.update()
+                scheduler.step()
+                
+            # log loss
+            if saver.global_step % args.train.interval_log == 0:
+                current_lr =  optimizer.param_groups[0]['lr']
+                saver.log_info(
+                    'epoch: {} | {:3d}/{:3d} | {} | batch/s: {:.2f} | lr: {:.6} | loss: {:.3f} | time: {} | step: {}'.format(
+                        epoch,
+                        batch_idx,
+                        num_batches,
+                        args.env.expdir,
+                        args.train.interval_log/saver.get_interval_time(),
+                        current_lr,
+                        loss.item(),
+                        saver.get_total_time(),
+                        saver.global_step
+                    )
+                )
+                
+                saver.log_value({
+                    'train/loss': loss.item()
+                })
+                
+                saver.log_value({
+                    'train/lr': current_lr
+                })
+            
+            # validation
+            if saver.global_step % args.train.interval_val == 0:
+                optimizer_save = optimizer if args.train.save_opt else None
+                
+                # save latest
+                saver.save_model(model, optimizer_save, postfix=f'{saver.global_step}')
+                last_val_step = saver.global_step - args.train.interval_val
+                if last_val_step % args.train.interval_force_save != 0:
+                    saver.delete_model(postfix=f'{last_val_step}')
+                
+                # run testing set
+                test_loss = test(args, model, vocoder, loader_test, saver)
+                
+                # log loss
+                saver.log_info(
+                    ' --- <validation> --- \nloss: {:.3f}. '.format(
+                        test_loss,
+                    )
+                )
+                
+                saver.log_value({
+                    'validation/loss': test_loss
+                })
+                
+                model.train()
+
+                          

diffusion/unit2mel.py

@@ -0,0 +1,147 @@
+import os
+import yaml
+import torch
+import torch.nn as nn
+import numpy as np
+from .diffusion import GaussianDiffusion
+from .wavenet import WaveNet
+from .vocoder import Vocoder
+
+class DotDict(dict):
+    def __getattr__(*args):         
+        val = dict.get(*args)         
+        return DotDict(val) if type(val) is dict else val   
+
+    __setattr__ = dict.__setitem__    
+    __delattr__ = dict.__delitem__
+
+    
+def load_model_vocoder(
+        model_path,
+        device='cpu',
+        config_path = None
+        ):
+    if config_path is None: config_file = os.path.join(os.path.split(model_path)[0], 'config.yaml')
+    else: config_file = config_path
+    
+    with open(config_file, "r") as config:
+        args = yaml.safe_load(config)
+    args = DotDict(args)
+    
+    # load vocoder
+    vocoder = Vocoder(args.vocoder.type, args.vocoder.ckpt, device=device)
+    
+    # load model
+    model = Unit2Mel(
+                args.data.encoder_out_channels, 
+                args.model.n_spk,
+                args.model.use_pitch_aug,
+                vocoder.dimension,
+                args.model.n_layers,
+                args.model.n_chans,
+                args.model.n_hidden)
+    
+    print(' [Loading] ' + model_path)
+    ckpt = torch.load(model_path, map_location=torch.device(device))
+    model.to(device)
+    model.load_state_dict(ckpt['model'])
+    model.eval()
+    return model, vocoder, args
+
+
+class Unit2Mel(nn.Module):
+    def __init__(
+            self,
+            input_channel,
+            n_spk,
+            use_pitch_aug=False,
+            out_dims=128,
+            n_layers=20, 
+            n_chans=384, 
+            n_hidden=256):
+        super().__init__()
+        self.unit_embed = nn.Linear(input_channel, n_hidden)
+        self.f0_embed = nn.Linear(1, n_hidden)
+        self.volume_embed = nn.Linear(1, n_hidden)
+        if use_pitch_aug:
+            self.aug_shift_embed = nn.Linear(1, n_hidden, bias=False)
+        else:
+            self.aug_shift_embed = None
+        self.n_spk = n_spk
+        if n_spk is not None and n_spk > 1:
+            self.spk_embed = nn.Embedding(n_spk, n_hidden)
+        
+        self.n_hidden = n_hidden
+        # diffusion
+        self.decoder = GaussianDiffusion(WaveNet(out_dims, n_layers, n_chans, n_hidden), out_dims=out_dims)
+        self.input_channel = input_channel
+    
+    def init_spkembed(self, units, f0, volume, spk_id = None, spk_mix_dict = None, aug_shift = None,
+                gt_spec=None, infer=True, infer_speedup=10, method='dpm-solver', k_step=300, use_tqdm=True):
+        
+        '''
+        input: 
+            B x n_frames x n_unit
+        return: 
+            dict of B x n_frames x feat
+        '''
+        x = self.unit_embed(units) + self.f0_embed((1+ f0 / 700).log()) + self.volume_embed(volume)
+        if self.n_spk is not None and self.n_spk > 1:
+            if spk_mix_dict is not None:
+                spk_embed_mix = torch.zeros((1,1,self.hidden_size))
+                for k, v in spk_mix_dict.items():
+                    spk_id_torch = torch.LongTensor(np.array([[k]])).to(units.device)
+                    spk_embeddd = self.spk_embed(spk_id_torch)
+                    self.speaker_map[k] = spk_embeddd
+                    spk_embed_mix = spk_embed_mix + v * spk_embeddd
+                x = x + spk_embed_mix
+            else:
+                x = x + self.spk_embed(spk_id - 1)
+        self.speaker_map = self.speaker_map.unsqueeze(0)
+        self.speaker_map = self.speaker_map.detach()
+        return x.transpose(1, 2)
+
+    def init_spkmix(self, n_spk):
+        self.speaker_map = torch.zeros((n_spk,1,1,self.n_hidden))
+        hubert_hidden_size = self.input_channel
+        n_frames = 10
+        hubert = torch.randn((1, n_frames, hubert_hidden_size))
+        mel2ph = torch.arange(end=n_frames).unsqueeze(0).long()
+        f0 = torch.randn((1, n_frames))
+        volume = torch.randn((1, n_frames))
+        spks = {}
+        for i in range(n_spk):
+            spks.update({i:1.0/float(self.n_spk)})
+        orgouttt = self.init_spkembed(hubert, f0.unsqueeze(-1), volume.unsqueeze(-1), spk_mix_dict=spks)
+
+    def forward(self, units, f0, volume, spk_id = None, spk_mix_dict = None, aug_shift = None,
+                gt_spec=None, infer=True, infer_speedup=10, method='dpm-solver', k_step=300, use_tqdm=True):
+        
+        '''
+        input: 
+            B x n_frames x n_unit
+        return: 
+            dict of B x n_frames x feat
+        '''
+
+        x = self.unit_embed(units) + self.f0_embed((1+ f0 / 700).log()) + self.volume_embed(volume)
+        if self.n_spk is not None and self.n_spk > 1:
+            if spk_mix_dict is not None:
+                for k, v in spk_mix_dict.items():
+                    spk_id_torch = torch.LongTensor(np.array([[k]])).to(units.device)
+                    x = x + v * self.spk_embed(spk_id_torch)
+            else:
+                if spk_id.shape[1] > 1:
+                    g = spk_id.reshape((spk_id.shape[0], spk_id.shape[1], 1, 1, 1))  # [N, S, B, 1, 1]
+                    g = g * self.speaker_map  # [N, S, B, 1, H]
+                    g = torch.sum(g, dim=1) # [N, 1, B, 1, H]
+                    g = g.transpose(0, -1).transpose(0, -2).squeeze(0) # [B, H, N]
+                    x = x + g
+                else:
+                    x = x + self.spk_embed(spk_id)
+        if self.aug_shift_embed is not None and aug_shift is not None:
+            x = x + self.aug_shift_embed(aug_shift / 5) 
+        x = self.decoder(x, gt_spec=gt_spec, infer=infer, infer_speedup=infer_speedup, method=method, k_step=k_step, use_tqdm=use_tqdm)
+    
+        return x
+

diffusion/vocoder.py

@@ -0,0 +1,94 @@
+import torch
+from vdecoder.nsf_hifigan.nvSTFT import STFT
+from vdecoder.nsf_hifigan.models import load_model,load_config
+from torchaudio.transforms import Resample
+
+    
+class Vocoder:
+    def __init__(self, vocoder_type, vocoder_ckpt, device = None):
+        if device is None:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.device = device
+        
+        if vocoder_type == 'nsf-hifigan':
+            self.vocoder = NsfHifiGAN(vocoder_ckpt, device = device)
+        elif vocoder_type == 'nsf-hifigan-log10':
+            self.vocoder = NsfHifiGANLog10(vocoder_ckpt, device = device)
+        else:
+            raise ValueError(f" [x] Unknown vocoder: {vocoder_type}")
+            
+        self.resample_kernel = {}
+        self.vocoder_sample_rate = self.vocoder.sample_rate()
+        self.vocoder_hop_size = self.vocoder.hop_size()
+        self.dimension = self.vocoder.dimension()
+        
+    def extract(self, audio, sample_rate, keyshift=0):
+                
+        # resample
+        if sample_rate == self.vocoder_sample_rate:
+            audio_res = audio
+        else:
+            key_str = str(sample_rate)
+            if key_str not in self.resample_kernel:
+                self.resample_kernel[key_str] = Resample(sample_rate, self.vocoder_sample_rate, lowpass_filter_width = 128).to(self.device)
+            audio_res = self.resample_kernel[key_str](audio)    
+        
+        # extract
+        mel = self.vocoder.extract(audio_res, keyshift=keyshift) # B, n_frames, bins
+        return mel
+   
+    def infer(self, mel, f0):
+        f0 = f0[:,:mel.size(1),0] # B, n_frames
+        audio = self.vocoder(mel, f0)
+        return audio
+        
+        
+class NsfHifiGAN(torch.nn.Module):
+    def __init__(self, model_path, device=None):
+        super().__init__()
+        if device is None:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.device = device
+        self.model_path = model_path
+        self.model = None
+        self.h = load_config(model_path)
+        self.stft = STFT(
+                self.h.sampling_rate, 
+                self.h.num_mels, 
+                self.h.n_fft, 
+                self.h.win_size, 
+                self.h.hop_size, 
+                self.h.fmin, 
+                self.h.fmax)
+    
+    def sample_rate(self):
+        return self.h.sampling_rate
+        
+    def hop_size(self):
+        return self.h.hop_size
+    
+    def dimension(self):
+        return self.h.num_mels
+        
+    def extract(self, audio, keyshift=0):       
+        mel = self.stft.get_mel(audio, keyshift=keyshift).transpose(1, 2) # B, n_frames, bins
+        return mel
+    
+    def forward(self, mel, f0):
+        if self.model is None:
+            print('| Load HifiGAN: ', self.model_path)
+            self.model, self.h = load_model(self.model_path, device=self.device)
+        with torch.no_grad():
+            c = mel.transpose(1, 2)
+            audio = self.model(c, f0)
+            return audio
+
+class NsfHifiGANLog10(NsfHifiGAN):    
+    def forward(self, mel, f0):
+        if self.model is None:
+            print('| Load HifiGAN: ', self.model_path)
+            self.model, self.h = load_model(self.model_path, device=self.device)
+        with torch.no_grad():
+            c = 0.434294 * mel.transpose(1, 2)
+            audio = self.model(c, f0)
+            return audio

diffusion/wavenet.py

@@ -0,0 +1,108 @@
+import math
+from math import sqrt
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import Mish
+
+
+class Conv1d(torch.nn.Conv1d):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        nn.init.kaiming_normal_(self.weight)
+
+
+class SinusoidalPosEmb(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        device = x.device
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
+        emb = x[:, None] * emb[None, :]
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, encoder_hidden, residual_channels, dilation):
+        super().__init__()
+        self.residual_channels = residual_channels
+        self.dilated_conv = nn.Conv1d(
+            residual_channels,
+            2 * residual_channels,
+            kernel_size=3,
+            padding=dilation,
+            dilation=dilation
+        )
+        self.diffusion_projection = nn.Linear(residual_channels, residual_channels)
+        self.conditioner_projection = nn.Conv1d(encoder_hidden, 2 * residual_channels, 1)
+        self.output_projection = nn.Conv1d(residual_channels, 2 * residual_channels, 1)
+
+    def forward(self, x, conditioner, diffusion_step):
+        diffusion_step = self.diffusion_projection(diffusion_step).unsqueeze(-1)
+        conditioner = self.conditioner_projection(conditioner)
+        y = x + diffusion_step
+
+        y = self.dilated_conv(y) + conditioner
+
+        # Using torch.split instead of torch.chunk to avoid using onnx::Slice
+        gate, filter = torch.split(y, [self.residual_channels, self.residual_channels], dim=1)
+        y = torch.sigmoid(gate) * torch.tanh(filter)
+
+        y = self.output_projection(y)
+
+        # Using torch.split instead of torch.chunk to avoid using onnx::Slice
+        residual, skip = torch.split(y, [self.residual_channels, self.residual_channels], dim=1)
+        return (x + residual) / math.sqrt(2.0), skip
+
+
+class WaveNet(nn.Module):
+    def __init__(self, in_dims=128, n_layers=20, n_chans=384, n_hidden=256):
+        super().__init__()
+        self.input_projection = Conv1d(in_dims, n_chans, 1)
+        self.diffusion_embedding = SinusoidalPosEmb(n_chans)
+        self.mlp = nn.Sequential(
+            nn.Linear(n_chans, n_chans * 4),
+            Mish(),
+            nn.Linear(n_chans * 4, n_chans)
+        )
+        self.residual_layers = nn.ModuleList([
+            ResidualBlock(
+                encoder_hidden=n_hidden,
+                residual_channels=n_chans,
+                dilation=1
+            )
+            for i in range(n_layers)
+        ])
+        self.skip_projection = Conv1d(n_chans, n_chans, 1)
+        self.output_projection = Conv1d(n_chans, in_dims, 1)
+        nn.init.zeros_(self.output_projection.weight)
+
+    def forward(self, spec, diffusion_step, cond):
+        """
+        :param spec: [B, 1, M, T]
+        :param diffusion_step: [B, 1]
+        :param cond: [B, M, T]
+        :return:
+        """
+        x = spec.squeeze(1)
+        x = self.input_projection(x)  # [B, residual_channel, T]
+
+        x = F.relu(x)
+        diffusion_step = self.diffusion_embedding(diffusion_step)
+        diffusion_step = self.mlp(diffusion_step)
+        skip = []
+        for layer in self.residual_layers:
+            x, skip_connection = layer(x, cond, diffusion_step)
+            skip.append(skip_connection)
+
+        x = torch.sum(torch.stack(skip), dim=0) / sqrt(len(self.residual_layers))
+        x = self.skip_projection(x)
+        x = F.relu(x)
+        x = self.output_projection(x)  # [B, mel_bins, T]
+        return x[:, None, :, :]