同步官方仓库的更新

app.py

@@ -11,7 +11,9 @@ from scipy.io import wavfile
 import tempfile
 import edge_tts
 import utils
-import torch
+import matplotlib
+
+matplotlib.use('TkAgg')
 
 from inference.infer_tool import Svc
 
@@ -67,9 +69,7 @@ def create_fn(model, spk):
         input_text = re.sub(r"[\n\,\(\) ]", "", input_text)
         voice = tts_voice[gender]
         ratestr = "+{:.0%}".format(tts_rate) if tts_rate >= 0 else "{:.0%}".format(tts_rate)
-        communicate = edge_tts.Communicate(text=input_text,
-                                           voice=voice,
-                                           rate=ratestr)
+        communicate = edge_tts.Communicate(text=input_text, voice=voice, rate=ratestr)
         with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_file:
             temp_path = tmp_file.name
         await communicate.save(temp_path)
@@ -100,10 +100,12 @@ if __name__ == '__main__':
         models.append((name, cover, create_fn(model, name)))
     with gr.Blocks() as app:
         gr.Markdown(
-            "# <center> 圣安地列斯角色语音生成\n"
-            "## <center> 模型作者：B站[Cyber蝈蝈总](https://space.bilibili.com/37706580)\n"
-            "#### <center> 罪恶都市人物AI语音请移步[GTAVC](https://huggingface.co/spaces/GroveStreet/GTAVC_SOVITS) \n"
-            "<center> 使用此资源创作的作品请标明出处，CJ有两个模型，carl1更清晰，carl2音域广\n"
+            """
+            # <center> 圣安地列斯角色语音生成
+            ## <center> 模型作者：B站[Cyber蝈蝈总](https://space.bilibili.com/37706580)
+            #### <center> 罪恶都市人物AI语音请移步[GTAVC](https://huggingface.co/spaces/GroveStreet/GTAVC_SOVITS)
+            <center> 使用此资源创作的作品请标出处，CJ有两个模型，carl1更清晰，carl2音域广
+            """
         )
         with gr.Tabs():
             for (name, cover, (svc_fn, tts_fn)) in models:
@@ -112,8 +114,8 @@ if __name__ == '__main__':
                         with gr.Column():
                             with gr.Row():
                                 vc_transform = gr.Number(label="音高调整 (正负半音，12为1个八度)", value=0)
-                                f0_predictor = gr.Radio(label="f0预测器 (harvest适合讲话，crepe适合唱歌)",
-                                                        choices=['crepe', 'harvest', 'dio', 'pm'], value='crepe')
+                                f0_predictor = gr.Radio(label="f0预测器 (推荐rmvpe)",
+                                                        choices=['crepe', 'harvest', 'rmvpe'], value='rmvpe')
                             auto_f0 = gr.Checkbox(label="自动音高预测 (文本转语音或讲话可选,会导致唱歌跑调)",
                                                   value=False)
                             with gr.Tabs():
@@ -132,13 +134,9 @@ if __name__ == '__main__':
                                     tts_submit = gr.Button("生成", variant="primary")
 
                         with gr.Column():
-                            gr.Markdown(
-                                '<div align="center">'
-                                f'<img style="width:auto;height:300px;" src="file/{cover}">' if cover else ""
-                                                                                                           '</div>'
-                            )
+                            gr.Image(value=os.path.join(os.path.dirname(__file__), cover), height=300, width=300)
                             vc_output = gr.Audio(label="输出音频")
                     svc_submit.click(svc_fn, [svc_input, vc_transform, auto_f0, f0_predictor], vc_output)
                     tts_submit.click(tts_fn, [tts_input, gender, tts_rate, vc_transform, auto_f0, f0_predictor],
                                      vc_output)
-        app.queue(concurrency_count=1, api_open=args.api).launch(share=args.share)
+        app.queue(default_concurrency_limit=1, api_open=args.api).launch(share=args.share)

cluster/__init__.py

@@ -1,29 +0,0 @@
-import numpy as np
-import torch
-from sklearn.cluster import KMeans
-
-def get_cluster_model(ckpt_path):
-    checkpoint = torch.load(ckpt_path)
-    kmeans_dict = {}
-    for spk, ckpt in checkpoint.items():
-        km = KMeans(ckpt["n_features_in_"])
-        km.__dict__["n_features_in_"] = ckpt["n_features_in_"]
-        km.__dict__["_n_threads"] = ckpt["_n_threads"]
-        km.__dict__["cluster_centers_"] = ckpt["cluster_centers_"]
-        kmeans_dict[spk] = km
-    return kmeans_dict
-
-def get_cluster_result(model, x, speaker):
-    """
-        x: np.array [t, 256]
-        return cluster class result
-    """
-    return model[speaker].predict(x)
-
-def get_cluster_center_result(model, x,speaker):
-    """x: np.array [t, 256]"""
-    predict = model[speaker].predict(x)
-    return model[speaker].cluster_centers_[predict]
-
-def get_center(model, x,speaker):
-    return model[speaker].cluster_centers_[x]

cluster/kmeans.py

@@ -1,201 +0,0 @@
-import math,pdb
-import torch,pynvml
-from torch.nn.functional import normalize
-from time import time
-import numpy as np
-# device=torch.device("cuda:0")
-def _kpp(data: torch.Tensor, k: int, sample_size: int = -1):
-    """ Picks k points in the data based on the kmeans++ method.
-
-    Parameters
-    ----------
-    data : torch.Tensor
-        Expect a rank 1 or 2 array. Rank 1 is assumed to describe 1-D
-        data, rank 2 multidimensional data, in which case one
-        row is one observation.
-    k : int
-        Number of samples to generate.
-    sample_size : int
-        sample data to avoid memory overflow during calculation
-
-    Returns
-    -------
-    init : ndarray
-        A 'k' by 'N' containing the initial centroids.
-
-    References
-    ----------
-    .. [1] D. Arthur and S. Vassilvitskii, "k-means++: the advantages of
-       careful seeding", Proceedings of the Eighteenth Annual ACM-SIAM Symposium
-       on Discrete Algorithms, 2007.
-    .. [2] scipy/cluster/vq.py: _kpp
-    """
-    batch_size=data.shape[0]
-    if batch_size>sample_size:
-        data = data[torch.randint(0, batch_size,[sample_size], device=data.device)]
-    dims = data.shape[1] if len(data.shape) > 1 else 1
-    init = torch.zeros((k, dims)).to(data.device)
-    r = torch.distributions.uniform.Uniform(0, 1)
-    for i in range(k):
-        if i == 0:
-            init[i, :] = data[torch.randint(data.shape[0], [1])]
-        else:
-            D2 = torch.cdist(init[:i, :][None, :], data[None, :], p=2)[0].amin(dim=0)
-            probs = D2 / torch.sum(D2)
-            cumprobs = torch.cumsum(probs, dim=0)
-            init[i, :] = data[torch.searchsorted(cumprobs, r.sample([1]).to(data.device))]
-    return init
-class KMeansGPU:
-  '''
-  Kmeans clustering algorithm implemented with PyTorch
-
-  Parameters:
-    n_clusters: int, 
-      Number of clusters
-
-    max_iter: int, default: 100
-      Maximum number of iterations
-
-    tol: float, default: 0.0001
-      Tolerance
-    
-    verbose: int, default: 0
-      Verbosity
-
-    mode: {'euclidean', 'cosine'}, default: 'euclidean'
-      Type of distance measure
-      
-    init_method: {'random', 'point', '++'}
-      Type of initialization
-
-    minibatch: {None, int}, default: None
-      Batch size of MinibatchKmeans algorithm
-      if None perform full KMeans algorithm
-      
-  Attributes:
-    centroids: torch.Tensor, shape: [n_clusters, n_features]
-      cluster centroids
-  '''
-  def __init__(self, n_clusters, max_iter=200, tol=1e-4, verbose=0, mode="euclidean",device=torch.device("cuda:0")):
-    self.n_clusters = n_clusters
-    self.max_iter = max_iter
-    self.tol = tol
-    self.verbose = verbose
-    self.mode = mode
-    self.device=device
-    pynvml.nvmlInit()
-    gpu_handle = pynvml.nvmlDeviceGetHandleByIndex(device.index)
-    info = pynvml.nvmlDeviceGetMemoryInfo(gpu_handle)
-    self.minibatch=int(33e6/self.n_clusters*info.free/ 1024 / 1024 / 1024)
-    print("free_mem/GB:",info.free/ 1024 / 1024 / 1024,"minibatch:",self.minibatch)
-    
-  @staticmethod
-  def cos_sim(a, b):
-    """
-      Compute cosine similarity of 2 sets of vectors
-
-      Parameters:
-      a: torch.Tensor, shape: [m, n_features]
-
-      b: torch.Tensor, shape: [n, n_features]
-    """
-    return normalize(a, dim=-1) @ normalize(b, dim=-1).transpose(-2, -1)
-
-  @staticmethod
-  def euc_sim(a, b):
-    """
-      Compute euclidean similarity of 2 sets of vectors
-      Parameters:
-      a: torch.Tensor, shape: [m, n_features]
-      b: torch.Tensor, shape: [n, n_features]
-    """
-    return 2 * a @ b.transpose(-2, -1) -(a**2).sum(dim=1)[..., :, None] - (b**2).sum(dim=1)[..., None, :]
-
-  def max_sim(self, a, b):
-    """
-      Compute maximum similarity (or minimum distance) of each vector
-      in a with all of the vectors in b
-      Parameters:
-      a: torch.Tensor, shape: [m, n_features]
-      b: torch.Tensor, shape: [n, n_features]
-    """
-    if self.mode == 'cosine':
-      sim_func = self.cos_sim
-    elif self.mode == 'euclidean':
-      sim_func = self.euc_sim
-    sim = sim_func(a, b)
-    max_sim_v, max_sim_i = sim.max(dim=-1)
-    return max_sim_v, max_sim_i
-
-  def fit_predict(self, X):
-    """
-      Combination of fit() and predict() methods.
-      This is faster than calling fit() and predict() seperately.
-      Parameters:
-      X: torch.Tensor, shape: [n_samples, n_features]
-      centroids: {torch.Tensor, None}, default: None
-        if given, centroids will be initialized with given tensor
-        if None, centroids will be randomly chosen from X
-      Return:
-      labels: torch.Tensor, shape: [n_samples]
-
-            mini_=33kk/k*remain
-            mini=min(mini_,fea_shape)
-            offset=log2(k/1000)*1.5
-            kpp_all=min(mini_*10/offset,fea_shape)
-            kpp_sample=min(mini_/12/offset,fea_shape)
-    """
-    assert isinstance(X, torch.Tensor), "input must be torch.Tensor"
-    assert X.dtype in [torch.half, torch.float, torch.double], "input must be floating point"
-    assert X.ndim == 2, "input must be a 2d tensor with shape: [n_samples, n_features] "
-    # print("verbose:%s"%self.verbose)
-
-    offset = np.power(1.5,np.log(self.n_clusters / 1000))/np.log(2)
-    with torch.no_grad():
-      batch_size= X.shape[0]
-      # print(self.minibatch, int(self.minibatch * 10 / offset), batch_size)
-      start_time = time()
-      if (self.minibatch*10//offset< batch_size):
-        x = X[torch.randint(0, batch_size,[int(self.minibatch*10/offset)])].to(self.device)
-      else:
-        x = X.to(self.device)
-      # print(x.device)
-      self.centroids = _kpp(x, self.n_clusters, min(int(self.minibatch/12/offset),batch_size))
-      del x
-      torch.cuda.empty_cache()
-      # self.centroids = self.centroids.to(self.device)
-      num_points_in_clusters = torch.ones(self.n_clusters, device=self.device, dtype=X.dtype)#全1
-      closest = None#[3098036]#int64
-      if(self.minibatch>=batch_size//2 and self.minibatch<batch_size):
-        X = X[torch.randint(0, batch_size,[self.minibatch])].to(self.device)
-      elif(self.minibatch>=batch_size):
-        X=X.to(self.device)
-      for i in range(self.max_iter):
-        iter_time = time()
-        if self.minibatch<batch_size//2:#可用minibatch数太小，每次都得从内存倒腾到显存
-          x = X[torch.randint(0, batch_size, [self.minibatch])].to(self.device)
-        else:#否则直接全部缓存
-          x = X
-
-        closest = self.max_sim(a=x, b=self.centroids)[1].to(torch.int16)#[3098036]#int64#0~999
-        matched_clusters, counts = closest.unique(return_counts=True)#int64#1k
-        expanded_closest = closest[None].expand(self.n_clusters, -1)#[1000, 3098036]#int16#0~999
-        mask = (expanded_closest==torch.arange(self.n_clusters, device=self.device)[:, None]).to(X.dtype)#==后者是int64*1000
-        c_grad = mask @ x / mask.sum(-1)[..., :, None]
-        c_grad[c_grad!=c_grad] = 0 # remove NaNs
-        error = (c_grad - self.centroids).pow(2).sum()
-        if self.minibatch is not None:
-          lr = 1/num_points_in_clusters[:,None] * 0.9 + 0.1
-        else:
-          lr = 1
-        matched_clusters=matched_clusters.long()
-        num_points_in_clusters[matched_clusters] += counts#IndexError: tensors used as indices must be long, byte or bool tensors
-        self.centroids = self.centroids * (1-lr) + c_grad * lr
-        if self.verbose >= 2:
-          print('iter:', i, 'error:', error.item(), 'time spent:', round(time()-iter_time, 4))
-        if error <= self.tol:
-          break
-
-      if self.verbose >= 1:
-        print(f'used {i+1} iterations ({round(time()-start_time, 4)}s) to cluster {batch_size} items into {self.n_clusters} clusters')
-    return closest

cluster/train_cluster.py

@@ -1,84 +0,0 @@
-import time,pdb
-import tqdm
-from time import time as ttime
-import os
-from pathlib import Path
-import logging
-import argparse
-from kmeans import KMeansGPU
-import torch
-import numpy as np
-from sklearn.cluster import KMeans,MiniBatchKMeans
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-from time import time as ttime
-import pynvml,torch
-
-def train_cluster(in_dir, n_clusters, use_minibatch=True, verbose=False,use_gpu=False):#gpu_minibatch真拉，虽然库支持但是也不考虑
-    logger.info(f"Loading features from {in_dir}")
-    features = []
-    nums = 0
-    for path in tqdm.tqdm(in_dir.glob("*.soft.pt")):
-    # for name in os.listdir(in_dir):
-    #     path="%s/%s"%(in_dir,name)
-        features.append(torch.load(path,map_location="cpu").squeeze(0).numpy().T)
-        # print(features[-1].shape)
-    features = np.concatenate(features, axis=0)
-    print(nums, features.nbytes/ 1024**2, "MB , shape:",features.shape, features.dtype)
-    features = features.astype(np.float32)
-    logger.info(f"Clustering features of shape: {features.shape}")
-    t = time.time()
-    if(use_gpu==False):
-        if use_minibatch:
-            kmeans = MiniBatchKMeans(n_clusters=n_clusters,verbose=verbose, batch_size=4096, max_iter=80).fit(features)
-        else:
-            kmeans = KMeans(n_clusters=n_clusters,verbose=verbose).fit(features)
-    else:
-            kmeans = KMeansGPU(n_clusters=n_clusters, mode='euclidean', verbose=2 if verbose else 0,max_iter=500,tol=1e-2)#
-            features=torch.from_numpy(features)#.to(device)
-            labels = kmeans.fit_predict(features)#
-
-    print(time.time()-t, "s")
-
-    x = {
-            "n_features_in_": kmeans.n_features_in_ if use_gpu==False else features.shape[1],
-            "_n_threads": kmeans._n_threads if use_gpu==False else 4,
-            "cluster_centers_": kmeans.cluster_centers_ if use_gpu==False else kmeans.centroids.cpu().numpy(),
-    }
-    print("end")
-
-    return x
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--dataset', type=Path, default="./dataset/44k",
-                        help='path of training data directory')
-    parser.add_argument('--output', type=Path, default="logs/44k",
-                        help='path of model output directory')
-    parser.add_argument('--gpu',action='store_true', default=False ,
-                        help='to use GPU')
-
-
-    args = parser.parse_args()
-
-    checkpoint_dir = args.output
-    dataset = args.dataset
-    use_gpu = args.gpu
-    n_clusters = 10000
-    
-    ckpt = {}
-    for spk in os.listdir(dataset):
-        if os.path.isdir(dataset/spk):
-            print(f"train kmeans for {spk}...")
-            in_dir = dataset/spk
-            x = train_cluster(in_dir, n_clusters,use_minibatch=False,verbose=False,use_gpu=use_gpu)
-            ckpt[spk] = x
-
-    checkpoint_path = checkpoint_dir / f"kmeans_{n_clusters}.pt"
-    checkpoint_path.parent.mkdir(exist_ok=True, parents=True)
-    torch.save(
-        ckpt,
-        checkpoint_path,
-    )
-    

diffusion/data_loaders.py

@@ -1,284 +0,0 @@
-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

@@ -1,317 +0,0 @@
-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

@@ -1,612 +0,0 @@
-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

@@ -1,1201 +0,0 @@
-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

@@ -1,4 +0,0 @@
-- 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

@@ -1,74 +0,0 @@
-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

@@ -1,150 +0,0 @@
-'''
-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

@@ -1,126 +0,0 @@
-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

@@ -1,226 +0,0 @@
-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

@@ -1,195 +0,0 @@
-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

@@ -1,147 +0,0 @@
-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

@@ -1,94 +0,0 @@
-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

@@ -1,108 +0,0 @@
-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, :, :]

inference/infer_tool.py

@@ -1,27 +1,24 @@
+import gc
 import hashlib
 import io
 import json
 import logging
 import os
+import pickle
 import time
 from pathlib import Path
-from inference import slicer
-import gc
 
 import librosa
 import numpy as np
+
 # import onnxruntime
 import soundfile
 import torch
 import torchaudio
 
-import cluster
 import utils
+from inference import slicer
 from models import SynthesizerTrn
-import pickle
-
-from diffusion.unit2mel import load_model_vocoder
-import yaml
 
 logging.getLogger('matplotlib').setLevel(logging.WARNING)
 
@@ -142,53 +139,26 @@ class Svc(object):
             self.dev = torch.device(device)
         self.net_g_ms = None
         if not self.only_diffusion:
-            self.hps_ms = utils.get_hparams_from_file(config_path)
+            self.hps_ms = utils.get_hparams_from_file(config_path, True)
             self.target_sample = self.hps_ms.data.sampling_rate
             self.hop_size = self.hps_ms.data.hop_length
             self.spk2id = self.hps_ms.spk
-            try:
-                self.vol_embedding = self.hps_ms.model.vol_embedding
-            except Exception as e:
-                self.vol_embedding = False
-            try:
-                self.speech_encoder = self.hps_ms.model.speech_encoder
-            except Exception as e:
-                self.speech_encoder = 'vec768l12'
+            self.unit_interpolate_mode = self.hps_ms.data.unit_interpolate_mode if self.hps_ms.data.unit_interpolate_mode is not None else 'left'
+            self.vol_embedding = self.hps_ms.model.vol_embedding if self.hps_ms.model.vol_embedding is not None else False
+            self.speech_encoder = self.hps_ms.model.speech_encoder if self.hps_ms.model.speech_encoder is not None else 'vec768l12'
 
         self.nsf_hifigan_enhance = nsf_hifigan_enhance
-        if self.shallow_diffusion or self.only_diffusion:
-            if os.path.exists(diffusion_model_path) and os.path.exists(diffusion_model_path):
-                self.diffusion_model, self.vocoder, self.diffusion_args = load_model_vocoder(diffusion_model_path,
-                                                                                             self.dev,
-                                                                                             config_path=diffusion_config_path)
-                if self.only_diffusion:
-                    self.target_sample = self.diffusion_args.data.sampling_rate
-                    self.hop_size = self.diffusion_args.data.block_size
-                    self.spk2id = self.diffusion_args.spk
-                    self.speech_encoder = self.diffusion_args.data.encoder
-                if spk_mix_enable:
-                    self.diffusion_model.init_spkmix(len(self.spk2id))
-            else:
-                print("No diffusion model or config found. Shallow diffusion mode will False")
-                self.shallow_diffusion = self.only_diffusion = False
 
         # load hubert and model
         self.load_model(spk_mix_enable)
-        # self.hubert_model = utils.get_speech_encoder(self.speech_encoder, device=self.dev) // ram optimize
+        # self.hubert_model = utils.get_speech_encoder(self.speech_encoder, device=self.dev)
         self.volume_extractor = utils.Volume_Extractor(self.hop_size)
 
-        if os.path.exists(cluster_model_path):
-            if self.feature_retrieval:
-                with open(cluster_model_path, "rb") as f:
-                    self.cluster_model = pickle.load(f)
-                self.big_npy = None
-                self.now_spk_id = -1
-            else:
-                self.cluster_model = cluster.get_cluster_model(cluster_model_path)
-        else:
-            self.feature_retrieval = False
 
-        if self.shallow_diffusion: self.nsf_hifigan_enhance = False
+        self.feature_retrieval = False
+
+        if self.shallow_diffusion:
+            self.nsf_hifigan_enhance = False
         if self.nsf_hifigan_enhance:
             from modules.enhancer import Enhancer
             self.enhancer = Enhancer('nsf-hifigan', 'pretrain/nsf_hifigan/model', device=self.dev)
@@ -200,6 +170,7 @@ class Svc(object):
             self.hps_ms.train.segment_size // self.hps_ms.data.hop_length,
             **self.hps_ms.model)
         _ = utils.load_checkpoint(self.net_g_path, self.net_g_ms, None)
+        self.dtype = list(self.net_g_ms.parameters())[0].dtype
         if "half" in self.net_g_path and torch.cuda.is_available():
             _ = self.net_g_ms.half().eval().to(self.dev)
         else:
@@ -209,11 +180,13 @@ class Svc(object):
 
     def get_unit_f0(self, wav, tran, cluster_infer_ratio, speaker, f0_filter, f0_predictor, cr_threshold=0.05):
 
-        f0_predictor_object = utils.get_f0_predictor(f0_predictor, hop_length=self.hop_size,
-                                                     sampling_rate=self.target_sample, device=self.dev,
-                                                     threshold=cr_threshold)
+        if not hasattr(self,
+                       "f0_predictor_object") or self.f0_predictor_object is None or f0_predictor != self.f0_predictor_object.name:
+            self.f0_predictor_object = utils.get_f0_predictor(f0_predictor, hop_length=self.hop_size,
+                                                              sampling_rate=self.target_sample, device=self.dev,
+                                                              threshold=cr_threshold)
+        f0, uv = self.f0_predictor_object.compute_f0_uv(wav)
 
-        f0, uv = f0_predictor_object.compute_f0_uv(wav)
         if f0_filter and sum(f0) == 0:
             raise F0FilterException("No voice detected")
         f0 = torch.FloatTensor(f0).to(self.dev)
@@ -223,36 +196,13 @@ class Svc(object):
         f0 = f0.unsqueeze(0)
         uv = uv.unsqueeze(0)
 
-        wav16k = librosa.resample(wav, orig_sr=self.target_sample, target_sr=16000)
-        wav16k = torch.from_numpy(wav16k).to(self.dev)
+        wav = torch.from_numpy(wav).to(self.dev)
+        if not hasattr(self, "audio16k_resample_transform"):
+            self.audio16k_resample_transform = torchaudio.transforms.Resample(self.target_sample, 16000).to(self.dev)
+        wav16k = self.audio16k_resample_transform(wav[None, :])[0]
+
         c = self.hubert_model.encoder(wav16k)
-        c = utils.repeat_expand_2d(c.squeeze(0), f0.shape[1])
-
-        if cluster_infer_ratio != 0:
-            if self.feature_retrieval:
-                speaker_id = self.spk2id.get(speaker)
-                if speaker_id is None:
-                    raise RuntimeError("The name you entered is not in the speaker list!")
-                if not speaker_id and type(speaker) is int:
-                    if len(self.spk2id.__dict__) >= speaker:
-                        speaker_id = speaker
-                feature_index = self.cluster_model[speaker_id]
-                feat_np = c.transpose(0, 1).cpu().numpy()
-                if self.big_npy is None or self.now_spk_id != speaker_id:
-                    self.big_npy = feature_index.reconstruct_n(0, feature_index.ntotal)
-                    self.now_spk_id = speaker_id
-                print("starting feature retrieval...")
-                score, ix = feature_index.search(feat_np, k=8)
-                weight = np.square(1 / score)
-                weight /= weight.sum(axis=1, keepdims=True)
-                npy = np.sum(self.big_npy[ix] * np.expand_dims(weight, axis=2), axis=1)
-                c = cluster_infer_ratio * npy + (1 - cluster_infer_ratio) * feat_np
-                c = torch.FloatTensor(c).to(self.dev).transpose(0, 1)
-                print("end feature retrieval...")
-            else:
-                cluster_c = cluster.get_cluster_center_result(self.cluster_model, c.cpu().numpy().T, speaker).T
-                cluster_c = torch.FloatTensor(cluster_c).to(self.dev)
-                c = cluster_infer_ratio * cluster_c + (1 - cluster_infer_ratio) * c
+        c = utils.repeat_expand_2d(c.squeeze(0), f0.shape[1], self.unit_interpolate_mode)
 
         c = c.unsqueeze(0)
         return c, f0, uv
@@ -271,7 +221,11 @@ class Svc(object):
               second_encoding=False,
               loudness_envelope_adjustment=1
               ):
-        wav, sr = librosa.load(raw_path, sr=self.target_sample)
+        torchaudio.set_audio_backend("soundfile")
+        wav, sr = torchaudio.load(raw_path)
+        if not hasattr(self, "audio_resample_transform") or self.audio16k_resample_transform.orig_freq != sr:
+            self.audio_resample_transform = torchaudio.transforms.Resample(sr, self.target_sample)
+        wav = self.audio_resample_transform(wav).numpy()[0]
         if spk_mix:
             c, f0, uv = self.get_unit_f0(wav, tran, 0, None, f0_filter, f0_predictor, cr_threshold=cr_threshold)
             n_frames = f0.size(1)
@@ -287,8 +241,9 @@ class Svc(object):
             c, f0, uv = self.get_unit_f0(wav, tran, cluster_infer_ratio, speaker, f0_filter, f0_predictor,
                                          cr_threshold=cr_threshold)
             n_frames = f0.size(1)
-        if "half" in self.net_g_path and torch.cuda.is_available():
-            c = c.half()
+        c = c.to(self.dtype)
+        f0 = f0.to(self.dtype)
+        uv = uv.to(self.dtype)
         with torch.no_grad():
             start = time.time()
             vol = None
@@ -302,17 +257,22 @@ class Svc(object):
             else:
                 audio = torch.FloatTensor(wav).to(self.dev)
                 audio_mel = None
+            if self.dtype != torch.float32:
+                c = c.to(torch.float32)
+                f0 = f0.to(torch.float32)
+                uv = uv.to(torch.float32)
             if self.only_diffusion or self.shallow_diffusion:
-                vol = self.volume_extractor.extract(audio[None, :])[None, :, None].to(self.dev) if vol == None else vol[
+                vol = self.volume_extractor.extract(audio[None, :])[None, :, None].to(self.dev) if vol is None else vol[
                                                                                                                     :,
                                                                                                                     :,
                                                                                                                     None]
                 if self.shallow_diffusion and second_encoding:
-                    audio16k = librosa.resample(audio.detach().cpu().numpy(), orig_sr=self.target_sample,
-                                                target_sr=16000)
-                    audio16k = torch.from_numpy(audio16k).to(self.dev)
+                    if not hasattr(self, "audio16k_resample_transform"):
+                        self.audio16k_resample_transform = torchaudio.transforms.Resample(self.target_sample, 16000).to(
+                            self.dev)
+                    audio16k = self.audio16k_resample_transform(audio[None, :])[0]
                     c = self.hubert_model.encoder(audio16k)
-                    c = utils.repeat_expand_2d(c.squeeze(0), f0.shape[1])
+                    c = utils.repeat_expand_2d(c.squeeze(0), f0.shape[1], self.unit_interpolate_mode)
                 f0 = f0[:, :, None]
                 c = c.transpose(-1, -2)
                 audio_mel = self.diffusion_model(
@@ -461,7 +421,8 @@ class Svc(object):
                 datas = [data]
             for k, dat in enumerate(datas):
                 per_length = int(np.ceil(len(dat) / audio_sr * self.target_sample)) if clip_seconds != 0 else length
-                if clip_seconds != 0: print(f'###=====segment clip start, {round(len(dat) / audio_sr, 3)}s======')
+                if clip_seconds != 0:
+                    print(f'###=====segment clip start, {round(len(dat) / audio_sr, 3)}s======')
                 # padd
                 pad_len = int(audio_sr * pad_seconds)
                 dat = np.concatenate([np.zeros([pad_len]), dat, np.zeros([pad_len])])
@@ -497,51 +458,3 @@ class Svc(object):
         return np.array(audio)
 
 
-class RealTimeVC:
-    def __init__(self):
-        self.last_chunk = None
-        self.last_o = None
-        self.chunk_len = 16000  # chunk length
-        self.pre_len = 3840  # cross fade length, multiples of 640
-
-    # Input and output are 1-dimensional numpy waveform arrays
-
-    def process(self, svc_model, speaker_id, f_pitch_change, input_wav_path,
-                cluster_infer_ratio=0,
-                auto_predict_f0=False,
-                noice_scale=0.4,
-                f0_filter=False):
-
-        import maad
-        audio, sr = torchaudio.load(input_wav_path)
-        audio = audio.cpu().numpy()[0]
-        temp_wav = io.BytesIO()
-        if self.last_chunk is None:
-            input_wav_path.seek(0)
-
-            audio, sr = svc_model.infer(speaker_id, f_pitch_change, input_wav_path,
-                                        cluster_infer_ratio=cluster_infer_ratio,
-                                        auto_predict_f0=auto_predict_f0,
-                                        noice_scale=noice_scale,
-                                        f0_filter=f0_filter)
-
-            audio = audio.cpu().numpy()
-            self.last_chunk = audio[-self.pre_len:]
-            self.last_o = audio
-            return audio[-self.chunk_len:]
-        else:
-            audio = np.concatenate([self.last_chunk, audio])
-            soundfile.write(temp_wav, audio, sr, format="wav")
-            temp_wav.seek(0)
-
-            audio, sr = svc_model.infer(speaker_id, f_pitch_change, temp_wav,
-                                        cluster_infer_ratio=cluster_infer_ratio,
-                                        auto_predict_f0=auto_predict_f0,
-                                        noice_scale=noice_scale,
-                                        f0_filter=f0_filter)
-
-            audio = audio.cpu().numpy()
-            ret = maad.util.crossfade(self.last_o, audio, self.pre_len)
-            self.last_chunk = audio[-self.pre_len:]
-            self.last_o = audio
-            return ret[self.chunk_len:2 * self.chunk_len]

inference/infer_tool_grad.py

@@ -1,22 +1,18 @@
-import hashlib
-import json
+import io
 import logging
 import os
-import time
-from pathlib import Path
-import io
+
 import librosa
-import maad
 import numpy as np
-from inference import slicer
 import parselmouth
 import soundfile
 import torch
 import torchaudio
 
-from hubert import hubert_model
 import utils
+from inference import slicer
 from models import SynthesizerTrn
+
 logging.getLogger('numba').setLevel(logging.WARNING)
 logging.getLogger('matplotlib').setLevel(logging.WARNING)
 
@@ -93,7 +89,7 @@ class VitsSvc(object):
     def set_device(self, device):
         self.device = torch.device(device)
         self.hubert_soft.to(self.device)
-        if self.SVCVITS != None:
+        if self.SVCVITS is not None:
             self.SVCVITS.to(self.device)
 
     def loadCheckpoint(self, path):

inference/slicer.py

@@ -117,8 +117,8 @@ class Slicer:
             return chunk_dict
 
 
-def cut(input_audio, db_thresh=-30, min_len=5000):
-    audio, sr = librosa.load(input_audio, sr=None)
+def cut(audio_path, db_thresh=-30, min_len=5000):
+    audio, sr = librosa.load(audio_path, sr=None)
     slicer = Slicer(
         sr=sr,
         threshold=db_thresh,

inference_main.py

@@ -1,181 +0,0 @@
-import io
-import logging
-import time
-from pathlib import Path
-from spkmix import spk_mix_map
-import librosa
-import matplotlib.pyplot as plt
-import numpy as np
-import soundfile
-from inference import infer_tool
-from inference import slicer
-from inference.infer_tool import Svc
-
-logging.getLogger('numba').setLevel(logging.WARNING)
-chunks_dict = infer_tool.read_temp("inference/chunks_temp.json")
-
-
-def main():
-    import argparse
-
-    parser = argparse.ArgumentParser(description='sovits4 inference')
-
-    # 一定要设置的部分
-    parser.add_argument('-m', '--model_path', type=str, default="logs/44k/", help='模型路径')
-    parser.add_argument('-c', '--config_path', type=str, default="configs/", help='配置文件路径')
-    parser.add_argument('-cl', '--clip', type=float, default=0, help='音频强制切片，默认0为自动切片，单位为秒/s')
-    parser.add_argument('-n', '--clean_names', type=str, nargs='+', default=["test.wav"],
-                        help='wav文件名列表，放在raw文件夹下')
-    parser.add_argument('-t', '--trans', type=int, nargs='+', default=[0], help='音高调整，支持正负（半音）')
-    parser.add_argument('-s', '--spk_list', type=str, nargs='+', default=['buyizi'], help='合成目标说话人名称')
-
-    # 可选项部分
-    parser.add_argument('-a', '--auto_predict_f0', action='store_true', default=False,
-                        help='语音转换自动预测音高，转换歌声时不要打开这个会严重跑调')
-    parser.add_argument('-cm', '--cluster_model_path', type=str, default="logs/44k/kmeans_10000.pt",
-                        help='聚类模型或特征检索索引路径，如果没有训练聚类或特征检索则随便填')
-    parser.add_argument('-cr', '--cluster_infer_ratio', type=float, default=0,
-                        help='聚类方案或特征检索占比，范围0-1，若没有训练聚类模型或特征检索则默认0即可')
-    parser.add_argument('-lg', '--linear_gradient', type=float, default=0,
-                        help='两段音频切片的交叉淡入长度，如果强制切片后出现人声不连贯可调整该数值，如果连贯建议采用默认值0，单位为秒')
-    parser.add_argument('-f0p', '--f0_predictor', type=str, default="harvest",
-                        help='选择F0预测器,可选择crepe,pm,dio,harvest,默认为pm(注意：crepe为原F0使用均值滤波器)')
-    parser.add_argument('-eh', '--enhance', action='store_true', default=False,
-                        help='是否使用NSF_HIFIGAN增强器,该选项对部分训练集少的模型有一定的音质增强效果，但是对训练好的模型有反面效果，默认关闭')
-    parser.add_argument('-shd', '--shallow_diffusion', action='store_true', default=False,
-                        help='是否使用浅层扩散，使用后可解决一部分电音问题，默认关闭，该选项打开时，NSF_HIFIGAN增强器将会被禁止')
-    parser.add_argument('-usm', '--use_spk_mix', action='store_true', default=False, help='是否使用角色融合')
-    parser.add_argument('-lea', '--loudness_envelope_adjustment', type=float, default=1,
-                        help='输入源响度包络替换输出响度包络融合比例，越靠近1越使用输出响度包络')
-    parser.add_argument('-fr', '--feature_retrieval', action='store_true', default=False,
-                        help='是否使用特征检索，如果使用聚类模型将被禁用，且cm与cr参数将会变成特征检索的索引路径与混合比例')
-
-    # 浅扩散设置
-    parser.add_argument('-dm', '--diffusion_model_path', type=str, default="logs/44k/diffusion/model_0.pt",
-                        help='扩散模型路径')
-    parser.add_argument('-dc', '--diffusion_config_path', type=str, default="logs/44k/diffusion/config.yaml",
-                        help='扩散模型配置文件路径')
-    parser.add_argument('-ks', '--k_step', type=int, default=100, help='扩散步数，越大越接近扩散模型的结果，默认100')
-    parser.add_argument('-se', '--second_encoding', action='store_true', default=False,
-                        help='二次编码，浅扩散前会对原始音频进行二次编码，玄学选项，有时候效果好，有时候效果差')
-    parser.add_argument('-od', '--only_diffusion', action='store_true', default=False,
-                        help='纯扩散模式，该模式不会加载sovits模型，以扩散模型推理')
-
-    # 不用动的部分
-    parser.add_argument('-sd', '--slice_db', type=int, default=-40,
-                        help='默认-40，嘈杂的音频可以-30，干声保留呼吸可以-50')
-    parser.add_argument('-d', '--device', type=str, default=None, help='推理设备，None则为自动选择cpu和gpu')
-    parser.add_argument('-ns', '--noice_scale', type=float, default=0.4, help='噪音级别，会影响咬字和音质，较为玄学')
-    parser.add_argument('-p', '--pad_seconds', type=float, default=0.5,
-                        help='推理音频pad秒数，由于未知原因开头结尾会有异响，pad一小段静音段后就不会出现')
-    parser.add_argument('-wf', '--wav_format', type=str, default='flac', help='音频输出格式')
-    parser.add_argument('-lgr', '--linear_gradient_retain', type=float, default=0.75,
-                        help='自动音频切片后，需要舍弃每段切片的头尾。该参数设置交叉长度保留的比例，范围0-1,左开右闭')
-    parser.add_argument('-eak', '--enhancer_adaptive_key', type=int, default=0,
-                        help='使增强器适应更高的音域(单位为半音数)|默认为0')
-    parser.add_argument('-ft', '--f0_filter_threshold', type=float, default=0.05,
-                        help='F0过滤阈值，只有使用crepe时有效. 数值范围从0-1. 降低该值可减少跑调概率，但会增加哑音')
-
-    def preprocess_args(args1):
-        spk1 = args1.spk_list[0]
-        args1.model_path += f"{spk1}.pth"
-        args1.config_path += f"config_{spk1}.json"
-        args1.clip = 30
-
-        if spk1 == 'tomori':
-            args1.feature_retrieval = True
-            args1.cluster_model_path = "logs/44k/tomori_index.pkl"
-            args1.cluster_infer_ratio = 0.5
-            args1.f0_predictor = 'crepe'
-
-        return args1
-
-    args = parser.parse_args()
-    args = preprocess_args(args)
-
-    clean_names = args.clean_names
-    trans = args.trans
-    spk_list = args.spk_list
-    slice_db = args.slice_db
-    wav_format = args.wav_format
-    auto_predict_f0 = args.auto_predict_f0
-    cluster_infer_ratio = args.cluster_infer_ratio
-    noice_scale = args.noice_scale
-    pad_seconds = args.pad_seconds
-    clip = args.clip
-    lg = args.linear_gradient
-    lgr = args.linear_gradient_retain
-    f0p = args.f0_predictor
-    enhance = args.enhance
-    enhancer_adaptive_key = args.enhancer_adaptive_key
-    cr_threshold = args.f0_filter_threshold
-    diffusion_model_path = args.diffusion_model_path
-    diffusion_config_path = args.diffusion_config_path
-    k_step = args.k_step
-    only_diffusion = args.only_diffusion
-    shallow_diffusion = args.shallow_diffusion
-    use_spk_mix = args.use_spk_mix
-    second_encoding = args.second_encoding
-    loudness_envelope_adjustment = args.loudness_envelope_adjustment
-
-    svc_model = Svc(args.model_path,
-                    args.config_path,
-                    args.device,
-                    args.cluster_model_path,
-                    enhance,
-                    diffusion_model_path,
-                    diffusion_config_path,
-                    shallow_diffusion,
-                    only_diffusion,
-                    use_spk_mix,
-                    args.feature_retrieval)
-
-    infer_tool.mkdir(["raw", "results"])
-
-    if len(spk_mix_map) <= 1:
-        use_spk_mix = False
-    if use_spk_mix:
-        spk_list = [spk_mix_map]
-
-    infer_tool.fill_a_to_b(trans, clean_names)
-    for clean_name, tran in zip(clean_names, trans):
-        raw_audio_path = f"raw/{clean_name}"
-        if "." not in raw_audio_path:
-            raw_audio_path += ".wav"
-        infer_tool.format_wav(raw_audio_path)
-        for spk in spk_list:
-            kwarg = {
-                "raw_audio_path": raw_audio_path,
-                "spk": spk,
-                "tran": tran,
-                "slice_db": slice_db,
-                "cluster_infer_ratio": cluster_infer_ratio,
-                "auto_predict_f0": auto_predict_f0,
-                "noice_scale": noice_scale,
-                "pad_seconds": pad_seconds,
-                "clip_seconds": clip,
-                "lg_num": lg,
-                "lgr_num": lgr,
-                "f0_predictor": f0p,
-                "enhancer_adaptive_key": enhancer_adaptive_key,
-                "cr_threshold": cr_threshold,
-                "k_step": k_step,
-                "use_spk_mix": use_spk_mix,
-                "second_encoding": second_encoding,
-                "loudness_envelope_adjustment": loudness_envelope_adjustment
-            }
-            audio = svc_model.slice_inference(**kwarg)
-            key = "auto" if auto_predict_f0 else f"{tran}key"
-            cluster_name = "" if cluster_infer_ratio == 0 else f"_{cluster_infer_ratio}"
-            isdiffusion = "sovits"
-            if shallow_diffusion: isdiffusion = "sovdiff"
-            if only_diffusion: isdiffusion = "diff"
-            if use_spk_mix:
-                spk = "spk_mix"
-            res_path = f'results/{clean_name}_{key}_{spk}{cluster_name}_{isdiffusion}.{wav_format}'
-            soundfile.write(res_path, audio, svc_model.target_sample, format=wav_format)
-            svc_model.clear_empty()
-
-
-if __name__ == '__main__':
-    main()

models.py

@@ -1,20 +1,17 @@
-import copy
-import math
 import torch
 from torch import nn
+from torch.nn import Conv1d, Conv2d
 from torch.nn import functional as F
+from torch.nn.utils import spectral_norm, weight_norm
 
 import modules.attentions as attentions
 import modules.commons as commons
 import modules.modules as modules
-
-from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
-from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
-
 import utils
-from modules.commons import init_weights, get_padding
+from modules.commons import get_padding
 from utils import f0_to_coarse
 
+
 class ResidualCouplingBlock(nn.Module):
     def __init__(self,
                  channels,
@@ -23,7 +20,9 @@ class ResidualCouplingBlock(nn.Module):
                  dilation_rate,
                  n_layers,
                  n_flows=4,
-                 gin_channels=0):
+                 gin_channels=0,
+                 share_parameter=False
+                 ):
         super().__init__()
         self.channels = channels
         self.hidden_channels = hidden_channels
@@ -34,10 +33,53 @@ class ResidualCouplingBlock(nn.Module):
         self.gin_channels = gin_channels
 
         self.flows = nn.ModuleList()
+
+        self.wn = modules.WN(hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout=0, gin_channels=gin_channels) if share_parameter else None
+
         for i in range(n_flows):
             self.flows.append(
                 modules.ResidualCouplingLayer(channels, hidden_channels, kernel_size, dilation_rate, n_layers,
-                                              gin_channels=gin_channels, mean_only=True))
+                                              gin_channels=gin_channels, mean_only=True, wn_sharing_parameter=self.wn))
+            self.flows.append(modules.Flip())
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        if not reverse:
+            for flow in self.flows:
+                x, _ = flow(x, x_mask, g=g, reverse=reverse)
+        else:
+            for flow in reversed(self.flows):
+                x = flow(x, x_mask, g=g, reverse=reverse)
+        return x
+
+class TransformerCouplingBlock(nn.Module):
+    def __init__(self,
+                 channels,
+                 hidden_channels,
+                 filter_channels,
+                 n_heads,
+                 n_layers,
+                 kernel_size,
+                 p_dropout,
+                 n_flows=4,
+                 gin_channels=0,
+                 share_parameter=False
+                 ):
+            
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.n_flows = n_flows
+        self.gin_channels = gin_channels
+
+        self.flows = nn.ModuleList()
+
+        self.wn = attentions.FFT(hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, isflow = True, gin_channels = self.gin_channels) if share_parameter else None
+
+        for i in range(n_flows):
+            self.flows.append(
+                modules.TransformerCouplingLayer(channels, hidden_channels, kernel_size, n_layers, n_heads, p_dropout, filter_channels, mean_only=True, wn_sharing_parameter=self.wn, gin_channels = self.gin_channels))
             self.flows.append(modules.Flip())
 
     def forward(self, x, x_mask, g=None, reverse=False):
@@ -125,7 +167,7 @@ class DiscriminatorP(torch.nn.Module):
         super(DiscriminatorP, self).__init__()
         self.period = period
         self.use_spectral_norm = use_spectral_norm
-        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        norm_f = weight_norm if use_spectral_norm is False else spectral_norm
         self.convs = nn.ModuleList([
             norm_f(Conv2d(1, 32, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
             norm_f(Conv2d(32, 128, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
@@ -160,7 +202,7 @@ class DiscriminatorP(torch.nn.Module):
 class DiscriminatorS(torch.nn.Module):
     def __init__(self, use_spectral_norm=False):
         super(DiscriminatorS, self).__init__()
-        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        norm_f = weight_norm if use_spectral_norm is False else spectral_norm
         self.convs = nn.ModuleList([
             norm_f(Conv1d(1, 16, 15, 1, padding=7)),
             norm_f(Conv1d(16, 64, 41, 4, groups=4, padding=20)),
@@ -321,6 +363,12 @@ class SynthesizerTrn(nn.Module):
                  sampling_rate=44100,
                  vol_embedding=False,
                  vocoder_name = "nsf-hifigan",
+                 use_depthwise_conv = False,
+                 use_automatic_f0_prediction = True,
+                 flow_share_parameter = False,
+                 n_flow_layer = 4,
+                 n_layers_trans_flow = 3,
+                 use_transformer_flow = False,
                  **kwargs):
 
         super().__init__()
@@ -343,6 +391,9 @@ class SynthesizerTrn(nn.Module):
         self.ssl_dim = ssl_dim
         self.vol_embedding = vol_embedding
         self.emb_g = nn.Embedding(n_speakers, gin_channels)
+        self.use_depthwise_conv = use_depthwise_conv
+        self.use_automatic_f0_prediction = use_automatic_f0_prediction
+        self.n_layers_trans_flow = n_layers_trans_flow
         if vol_embedding:
            self.emb_vol = nn.Linear(1, hidden_channels)
 
@@ -367,9 +418,11 @@ class SynthesizerTrn(nn.Module):
             "upsample_initial_channel": upsample_initial_channel,
             "upsample_kernel_sizes": upsample_kernel_sizes,
             "gin_channels": gin_channels,
+            "use_depthwise_conv":use_depthwise_conv
         }
         
-        
+        modules.set_Conv1dModel(self.use_depthwise_conv)
+
         if vocoder_name == "nsf-hifigan":
             from vdecoder.hifigan.models import Generator
             self.dec = Generator(h=hps)
@@ -382,17 +435,21 @@ class SynthesizerTrn(nn.Module):
             self.dec = Generator(h=hps)
 
         self.enc_q = Encoder(spec_channels, inter_channels, hidden_channels, 5, 1, 16, gin_channels=gin_channels)
-        self.flow = ResidualCouplingBlock(inter_channels, hidden_channels, 5, 1, 4, gin_channels=gin_channels)
-        self.f0_decoder = F0Decoder(
-            1,
-            hidden_channels,
-            filter_channels,
-            n_heads,
-            n_layers,
-            kernel_size,
-            p_dropout,
-            spk_channels=gin_channels
-        )
+        if use_transformer_flow:
+            self.flow = TransformerCouplingBlock(inter_channels, hidden_channels, filter_channels, n_heads, n_layers_trans_flow, 5, p_dropout, n_flow_layer,  gin_channels=gin_channels, share_parameter= flow_share_parameter)
+        else:
+            self.flow = ResidualCouplingBlock(inter_channels, hidden_channels, 5, 1, n_flow_layer, gin_channels=gin_channels, share_parameter= flow_share_parameter)
+        if self.use_automatic_f0_prediction:
+            self.f0_decoder = F0Decoder(
+                1,
+                hidden_channels,
+                filter_channels,
+                n_heads,
+                n_layers,
+                kernel_size,
+                p_dropout,
+                spk_channels=gin_channels
+            )
         self.emb_uv = nn.Embedding(2, hidden_channels)
         self.character_mix = False
 
@@ -407,17 +464,21 @@ class SynthesizerTrn(nn.Module):
         g = self.emb_g(g).transpose(1,2)
 
         # vol proj
-        vol = self.emb_vol(vol[:,:,None]).transpose(1,2) if vol!=None and self.vol_embedding else 0
+        vol = self.emb_vol(vol[:,:,None]).transpose(1,2) if vol is not None and self.vol_embedding else 0
 
         # ssl prenet
         x_mask = torch.unsqueeze(commons.sequence_mask(c_lengths, c.size(2)), 1).to(c.dtype)
         x = self.pre(c) * x_mask + self.emb_uv(uv.long()).transpose(1,2) + vol
-
+        
         # f0 predict
-        lf0 = 2595. * torch.log10(1. + f0.unsqueeze(1) / 700.) / 500
-        norm_lf0 = utils.normalize_f0(lf0, x_mask, uv)
-        pred_lf0 = self.f0_decoder(x, norm_lf0, x_mask, spk_emb=g)
-
+        if self.use_automatic_f0_prediction:
+            lf0 = 2595. * torch.log10(1. + f0.unsqueeze(1) / 700.) / 500
+            norm_lf0 = utils.normalize_f0(lf0, x_mask, uv)
+            pred_lf0 = self.f0_decoder(x, norm_lf0, x_mask, spk_emb=g)
+        else:
+            lf0 = 0
+            norm_lf0 = 0
+            pred_lf0 = 0
         # encoder
         z_ptemp, m_p, logs_p, _ = self.enc_p(x, x_mask, f0=f0_to_coarse(f0))
         z, m_q, logs_q, spec_mask = self.enc_q(spec, spec_lengths, g=g)
@@ -431,6 +492,7 @@ class SynthesizerTrn(nn.Module):
 
         return o, ids_slice, spec_mask, (z, z_p, m_p, logs_p, m_q, logs_q), pred_lf0, norm_lf0, lf0
 
+    @torch.no_grad()
     def infer(self, c, f0, uv, g=None, noice_scale=0.35, seed=52468, predict_f0=False, vol = None):
 
         if c.device == torch.device("cuda"):
@@ -452,11 +514,13 @@ class SynthesizerTrn(nn.Module):
         
         x_mask = torch.unsqueeze(commons.sequence_mask(c_lengths, c.size(2)), 1).to(c.dtype)
         # vol proj
-        vol = self.emb_vol(vol[:,:,None]).transpose(1,2) if vol!=None and self.vol_embedding else 0
-           
-        x = self.pre(c) * x_mask + self.emb_uv(uv.long()).transpose(1,2) + vol
         
-        if predict_f0:
+        vol = self.emb_vol(vol[:,:,None]).transpose(1,2) if vol is not None and self.vol_embedding else 0
+
+        x = self.pre(c) * x_mask + self.emb_uv(uv.long()).transpose(1, 2) + vol
+
+        
+        if self.use_automatic_f0_prediction and predict_f0:
             lf0 = 2595. * torch.log10(1. + f0.unsqueeze(1) / 700.) / 500
             norm_lf0 = utils.normalize_f0(lf0, x_mask, uv, random_scale=False)
             pred_lf0 = self.f0_decoder(x, norm_lf0, x_mask, spk_emb=g)

modules/DSConv.py

@@ -0,0 +1,76 @@
+import torch.nn as nn
+from torch.nn.utils import remove_weight_norm, weight_norm
+
+
+class Depthwise_Separable_Conv1D(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride = 1,
+        padding = 0,
+        dilation = 1,
+        bias = True,
+        padding_mode = 'zeros',  # TODO: refine this type
+        device=None,
+        dtype=None
+    ):
+      super().__init__()
+      self.depth_conv = nn.Conv1d(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size, groups=in_channels,stride = stride,padding=padding,dilation=dilation,bias=bias,padding_mode=padding_mode,device=device,dtype=dtype)
+      self.point_conv = nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias, device=device,dtype=dtype)
+    
+    def forward(self, input):
+      return self.point_conv(self.depth_conv(input))
+
+    def weight_norm(self):
+      self.depth_conv = weight_norm(self.depth_conv, name = 'weight')
+      self.point_conv = weight_norm(self.point_conv, name = 'weight')
+
+    def remove_weight_norm(self):
+      self.depth_conv = remove_weight_norm(self.depth_conv, name = 'weight')
+      self.point_conv = remove_weight_norm(self.point_conv, name = 'weight')
+
+class Depthwise_Separable_TransposeConv1D(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride = 1,
+        padding = 0, 
+        output_padding = 0,
+        bias = True,
+        dilation = 1,
+        padding_mode = 'zeros',  # TODO: refine this type
+        device=None,
+        dtype=None
+    ):
+      super().__init__()
+      self.depth_conv = nn.ConvTranspose1d(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size, groups=in_channels,stride = stride,output_padding=output_padding,padding=padding,dilation=dilation,bias=bias,padding_mode=padding_mode,device=device,dtype=dtype)
+      self.point_conv = nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias, device=device,dtype=dtype)
+    
+    def forward(self, input):
+      return self.point_conv(self.depth_conv(input))
+
+    def weight_norm(self):
+      self.depth_conv = weight_norm(self.depth_conv, name = 'weight')
+      self.point_conv = weight_norm(self.point_conv, name = 'weight')
+
+    def remove_weight_norm(self):
+      remove_weight_norm(self.depth_conv, name = 'weight')
+      remove_weight_norm(self.point_conv, name = 'weight')
+
+
+def weight_norm_modules(module, name = 'weight', dim = 0):
+    if isinstance(module,Depthwise_Separable_Conv1D) or isinstance(module,Depthwise_Separable_TransposeConv1D):
+      module.weight_norm()
+      return module
+    else:
+      return weight_norm(module,name,dim)
+
+def remove_weight_norm_modules(module, name = 'weight'):
+    if isinstance(module,Depthwise_Separable_Conv1D) or isinstance(module,Depthwise_Separable_TransposeConv1D):
+      module.remove_weight_norm()
+    else:
+      remove_weight_norm(module,name)

modules/F0Predictor/CrepeF0Predictor.py

@@ -1,7 +1,9 @@
-from modules.F0Predictor.F0Predictor import F0Predictor
-from modules.F0Predictor.crepe import CrepePitchExtractor
 import torch
 
+from modules.F0Predictor.crepe import CrepePitchExtractor
+from modules.F0Predictor.F0Predictor import F0Predictor
+
+
 class CrepeF0Predictor(F0Predictor):
     def __init__(self,hop_length=512,f0_min=50,f0_max=1100,device=None,sampling_rate=44100,threshold=0.05,model="full"):
         self.F0Creper = CrepePitchExtractor(hop_length=hop_length,f0_min=f0_min,f0_max=f0_max,device=device,threshold=threshold,model=model)
@@ -11,6 +13,7 @@ class CrepeF0Predictor(F0Predictor):
         self.device = device
         self.threshold = threshold
         self.sampling_rate = sampling_rate
+        self.name = "crepe"
 
     def compute_f0(self,wav,p_len=None):
         x = torch.FloatTensor(wav).to(self.device)

modules/F0Predictor/DioF0Predictor.py

@@ -1,6 +1,8 @@
-from modules.F0Predictor.F0Predictor import F0Predictor
-import pyworld
 import numpy as np
+import pyworld
+
+from modules.F0Predictor.F0Predictor import F0Predictor
+
 
 class DioF0Predictor(F0Predictor):
     def __init__(self,hop_length=512,f0_min=50,f0_max=1100,sampling_rate=44100):
@@ -8,44 +10,31 @@ class DioF0Predictor(F0Predictor):
         self.f0_min = f0_min
         self.f0_max = f0_max
         self.sampling_rate = sampling_rate
+        self.name = "dio"
 
     def interpolate_f0(self,f0):
         '''
         对F0进行插值处理
         '''
+        vuv_vector = np.zeros_like(f0, dtype=np.float32)
+        vuv_vector[f0 > 0.0] = 1.0
+        vuv_vector[f0 <= 0.0] = 0.0
     
-        data = np.reshape(f0, (f0.size, 1))
-    
-        vuv_vector = np.zeros((data.size, 1), dtype=np.float32)
-        vuv_vector[data > 0.0] = 1.0
-        vuv_vector[data <= 0.0] = 0.0
-    
-        ip_data = data
-    
-        frame_number = data.size
-        last_value = 0.0
-        for i in range(frame_number):
-            if data[i] <= 0.0:
-                j = i + 1
-                for j in range(i + 1, frame_number):
-                    if data[j] > 0.0:
-                        break
-                if j < frame_number - 1:
-                    if last_value > 0.0:
-                        step = (data[j] - data[i - 1]) / float(j - i)
-                        for k in range(i, j):
-                            ip_data[k] = data[i - 1] + step * (k - i + 1)
-                    else:
-                        for k in range(i, j):
-                            ip_data[k] = data[j]
-                else:
-                    for k in range(i, frame_number):
-                        ip_data[k] = last_value
-            else:
-                ip_data[i] = data[i] #这里可能存在一个没有必要的拷贝
-                last_value = data[i]
-    
-        return ip_data[:,0], vuv_vector[:,0]
+        nzindex = np.nonzero(f0)[0]
+        data = f0[nzindex]
+        nzindex = nzindex.astype(np.float32)
+        time_org = self.hop_length / self.sampling_rate * nzindex
+        time_frame = np.arange(f0.shape[0]) * self.hop_length / self.sampling_rate
+
+        if data.shape[0] <= 0:
+            return np.zeros(f0.shape[0], dtype=np.float32),vuv_vector
+
+        if data.shape[0] == 1:
+            return np.ones(f0.shape[0], dtype=np.float32) * f0[0],vuv_vector
+
+        f0 = np.interp(time_frame, time_org, data, left=data[0], right=data[-1])
+        
+        return f0,vuv_vector
 
     def resize_f0(self,x, target_len):
         source = np.array(x)

modules/F0Predictor/FCPEF0Predictor.py

@@ -0,0 +1,109 @@
+from typing import Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from modules.F0Predictor.F0Predictor import F0Predictor
+
+from .fcpe.model import FCPEInfer
+
+
+class FCPEF0Predictor(F0Predictor):
+    def __init__(self, hop_length=512, f0_min=50, f0_max=1100, dtype=torch.float32, device=None, sampling_rate=44100,
+                 threshold=0.05):
+        self.fcpe = FCPEInfer(model_path="pretrain/fcpe.pt", device=device, dtype=dtype)
+        self.hop_length = hop_length
+        self.f0_min = f0_min
+        self.f0_max = f0_max
+        if device is None:
+            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        else:
+            self.device = device
+        self.threshold = threshold
+        self.sampling_rate = sampling_rate
+        self.dtype = dtype
+        self.name = "fcpe"
+
+    def repeat_expand(
+            self, content: Union[torch.Tensor, np.ndarray], target_len: int, mode: str = "nearest"
+    ):
+        ndim = content.ndim
+
+        if content.ndim == 1:
+            content = content[None, None]
+        elif content.ndim == 2:
+            content = content[None]
+
+        assert content.ndim == 3
+
+        is_np = isinstance(content, np.ndarray)
+        if is_np:
+            content = torch.from_numpy(content)
+
+        results = torch.nn.functional.interpolate(content, size=target_len, mode=mode)
+
+        if is_np:
+            results = results.numpy()
+
+        if ndim == 1:
+            return results[0, 0]
+        elif ndim == 2:
+            return results[0]
+
+    def post_process(self, x, sampling_rate, f0, pad_to):
+        if isinstance(f0, np.ndarray):
+            f0 = torch.from_numpy(f0).float().to(x.device)
+
+        if pad_to is None:
+            return f0
+
+        f0 = self.repeat_expand(f0, pad_to)
+
+        vuv_vector = torch.zeros_like(f0)
+        vuv_vector[f0 > 0.0] = 1.0
+        vuv_vector[f0 <= 0.0] = 0.0
+
+        # 去掉0频率, 并线性插值
+        nzindex = torch.nonzero(f0).squeeze()
+        f0 = torch.index_select(f0, dim=0, index=nzindex).cpu().numpy()
+        time_org = self.hop_length / sampling_rate * nzindex.cpu().numpy()
+        time_frame = np.arange(pad_to) * self.hop_length / sampling_rate
+
+        vuv_vector = F.interpolate(vuv_vector[None, None, :], size=pad_to)[0][0]
+
+        if f0.shape[0] <= 0:
+            return torch.zeros(pad_to, dtype=torch.float, device=x.device).cpu().numpy(), vuv_vector.cpu().numpy()
+        if f0.shape[0] == 1:
+            return (torch.ones(pad_to, dtype=torch.float, device=x.device) * f0[
+                0]).cpu().numpy(), vuv_vector.cpu().numpy()
+
+        # 大概可以用 torch 重写?
+        f0 = np.interp(time_frame, time_org, f0, left=f0[0], right=f0[-1])
+        # vuv_vector = np.ceil(scipy.ndimage.zoom(vuv_vector,pad_to/len(vuv_vector),order = 0))
+
+        return f0, vuv_vector.cpu().numpy()
+
+    def compute_f0(self, wav, p_len=None):
+        x = torch.FloatTensor(wav).to(self.dtype).to(self.device)
+        if p_len is None:
+            p_len = x.shape[0] // self.hop_length
+        else:
+            assert abs(p_len - x.shape[0] // self.hop_length) < 4, "pad length error"
+        f0 = self.fcpe(x, sr=self.sampling_rate, threshold=self.threshold)[0,:,0]
+        if torch.all(f0 == 0):
+            rtn = f0.cpu().numpy() if p_len is None else np.zeros(p_len)
+            return rtn, rtn
+        return self.post_process(x, self.sampling_rate, f0, p_len)[0]
+
+    def compute_f0_uv(self, wav, p_len=None):
+        x = torch.FloatTensor(wav).to(self.dtype).to(self.device)
+        if p_len is None:
+            p_len = x.shape[0] // self.hop_length
+        else:
+            assert abs(p_len - x.shape[0] // self.hop_length) < 4, "pad length error"
+        f0 = self.fcpe(x, sr=self.sampling_rate, threshold=self.threshold)[0,:,0]
+        if torch.all(f0 == 0):
+            rtn = f0.cpu().numpy() if p_len is None else np.zeros(p_len)
+            return rtn, rtn
+        return self.post_process(x, self.sampling_rate, f0, p_len)

modules/F0Predictor/HarvestF0Predictor.py

@@ -1,6 +1,8 @@
-from modules.F0Predictor.F0Predictor import F0Predictor
-import pyworld
 import numpy as np
+import pyworld
+
+from modules.F0Predictor.F0Predictor import F0Predictor
+
 
 class HarvestF0Predictor(F0Predictor):
     def __init__(self,hop_length=512,f0_min=50,f0_max=1100,sampling_rate=44100):
@@ -8,45 +10,31 @@ class HarvestF0Predictor(F0Predictor):
         self.f0_min = f0_min
         self.f0_max = f0_max
         self.sampling_rate = sampling_rate
+        self.name = "harvest"
 
     def interpolate_f0(self,f0):
         '''
         对F0进行插值处理
         '''
+        vuv_vector = np.zeros_like(f0, dtype=np.float32)
+        vuv_vector[f0 > 0.0] = 1.0
+        vuv_vector[f0 <= 0.0] = 0.0
     
-        data = np.reshape(f0, (f0.size, 1))
-    
-        vuv_vector = np.zeros((data.size, 1), dtype=np.float32)
-        vuv_vector[data > 0.0] = 1.0
-        vuv_vector[data <= 0.0] = 0.0
-    
-        ip_data = data
-    
-        frame_number = data.size
-        last_value = 0.0
-        for i in range(frame_number):
-            if data[i] <= 0.0:
-                j = i + 1
-                for j in range(i + 1, frame_number):
-                    if data[j] > 0.0:
-                        break
-                if j < frame_number - 1:
-                    if last_value > 0.0:
-                        step = (data[j] - data[i - 1]) / float(j - i)
-                        for k in range(i, j):
-                            ip_data[k] = data[i - 1] + step * (k - i + 1)
-                    else:
-                        for k in range(i, j):
-                            ip_data[k] = data[j]
-                else:
-                    for k in range(i, frame_number):
-                        ip_data[k] = last_value
-            else:
-                ip_data[i] = data[i] #这里可能存在一个没有必要的拷贝
-                last_value = data[i]
-    
-        return ip_data[:,0], vuv_vector[:,0]
+        nzindex = np.nonzero(f0)[0]
+        data = f0[nzindex]
+        nzindex = nzindex.astype(np.float32)
+        time_org = self.hop_length / self.sampling_rate * nzindex
+        time_frame = np.arange(f0.shape[0]) * self.hop_length / self.sampling_rate
 
+        if data.shape[0] <= 0:
+            return np.zeros(f0.shape[0], dtype=np.float32),vuv_vector
+
+        if data.shape[0] == 1:
+            return np.ones(f0.shape[0], dtype=np.float32) * f0[0],vuv_vector
+
+        f0 = np.interp(time_frame, time_org, data, left=data[0], right=data[-1])
+        
+        return f0,vuv_vector
     def resize_f0(self,x, target_len):
         source = np.array(x)
         source[source<0.001] = np.nan

modules/F0Predictor/PMF0Predictor.py

@@ -1,6 +1,8 @@
-from modules.F0Predictor.F0Predictor import F0Predictor
-import parselmouth
 import numpy as np
+import parselmouth
+
+from modules.F0Predictor.F0Predictor import F0Predictor
+
 
 class PMF0Predictor(F0Predictor):
     def __init__(self,hop_length=512,f0_min=50,f0_max=1100,sampling_rate=44100):
@@ -8,45 +10,32 @@ class PMF0Predictor(F0Predictor):
         self.f0_min = f0_min
         self.f0_max = f0_max
         self.sampling_rate = sampling_rate
-
+        self.name = "pm"
     
     def interpolate_f0(self,f0):
         '''
         对F0进行插值处理
         '''
+        vuv_vector = np.zeros_like(f0, dtype=np.float32)
+        vuv_vector[f0 > 0.0] = 1.0
+        vuv_vector[f0 <= 0.0] = 0.0
     
-        data = np.reshape(f0, (f0.size, 1))
-    
-        vuv_vector = np.zeros((data.size, 1), dtype=np.float32)
-        vuv_vector[data > 0.0] = 1.0
-        vuv_vector[data <= 0.0] = 0.0
-    
-        ip_data = data
-    
-        frame_number = data.size
-        last_value = 0.0
-        for i in range(frame_number):
-            if data[i] <= 0.0:
-                j = i + 1
-                for j in range(i + 1, frame_number):
-                    if data[j] > 0.0:
-                        break
-                if j < frame_number - 1:
-                    if last_value > 0.0:
-                        step = (data[j] - data[i - 1]) / float(j - i)
-                        for k in range(i, j):
-                            ip_data[k] = data[i - 1] + step * (k - i + 1)
-                    else:
-                        for k in range(i, j):
-                            ip_data[k] = data[j]
-                else:
-                    for k in range(i, frame_number):
-                        ip_data[k] = last_value
-            else:
-                ip_data[i] = data[i] #这里可能存在一个没有必要的拷贝
-                last_value = data[i]
+        nzindex = np.nonzero(f0)[0]
+        data = f0[nzindex]
+        nzindex = nzindex.astype(np.float32)
+        time_org = self.hop_length / self.sampling_rate * nzindex
+        time_frame = np.arange(f0.shape[0]) * self.hop_length / self.sampling_rate
+
+        if data.shape[0] <= 0:
+            return np.zeros(f0.shape[0], dtype=np.float32),vuv_vector
+
+        if data.shape[0] == 1:
+            return np.ones(f0.shape[0], dtype=np.float32) * f0[0],vuv_vector
+
+        f0 = np.interp(time_frame, time_org, data, left=data[0], right=data[-1])
+        
+        return f0,vuv_vector
     
-        return ip_data[:,0], vuv_vector[:,0]
 
     def compute_f0(self,wav,p_len=None):
         x = wav

modules/F0Predictor/RMVPEF0Predictor.py

@@ -0,0 +1,107 @@
+from typing import Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from modules.F0Predictor.F0Predictor import F0Predictor
+
+from .rmvpe import RMVPE
+
+
+class RMVPEF0Predictor(F0Predictor):
+    def __init__(self,hop_length=512,f0_min=50,f0_max=1100, dtype=torch.float32, device=None,sampling_rate=44100,threshold=0.05):
+        self.rmvpe = RMVPE(model_path="pretrain/rmvpe.pt",dtype=dtype,device=device)
+        self.hop_length = hop_length
+        self.f0_min = f0_min
+        self.f0_max = f0_max
+        if device is None:
+            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        else:
+            self.device = device
+        self.threshold = threshold
+        self.sampling_rate = sampling_rate
+        self.dtype = dtype
+        self.name = "rmvpe"
+
+    def repeat_expand(
+        self, content: Union[torch.Tensor, np.ndarray], target_len: int, mode: str = "nearest"
+    ):
+        ndim = content.ndim
+
+        if content.ndim == 1:
+            content = content[None, None]
+        elif content.ndim == 2:
+            content = content[None]
+
+        assert content.ndim == 3
+
+        is_np = isinstance(content, np.ndarray)
+        if is_np:
+            content = torch.from_numpy(content)
+
+        results = torch.nn.functional.interpolate(content, size=target_len, mode=mode)
+
+        if is_np:
+            results = results.numpy()
+
+        if ndim == 1:
+            return results[0, 0]
+        elif ndim == 2:
+            return results[0]
+
+    def post_process(self, x, sampling_rate, f0, pad_to):
+        if isinstance(f0, np.ndarray):
+            f0 = torch.from_numpy(f0).float().to(x.device)
+
+        if pad_to is None:
+            return f0
+
+        f0 = self.repeat_expand(f0, pad_to)
+        
+        vuv_vector = torch.zeros_like(f0)
+        vuv_vector[f0 > 0.0] = 1.0
+        vuv_vector[f0 <= 0.0] = 0.0
+        
+        # 去掉0频率, 并线性插值
+        nzindex = torch.nonzero(f0).squeeze()
+        f0 = torch.index_select(f0, dim=0, index=nzindex).cpu().numpy()
+        time_org = self.hop_length / sampling_rate * nzindex.cpu().numpy()
+        time_frame = np.arange(pad_to) * self.hop_length / sampling_rate
+        
+        vuv_vector = F.interpolate(vuv_vector[None,None,:],size=pad_to)[0][0]
+
+        if f0.shape[0] <= 0:
+            return torch.zeros(pad_to, dtype=torch.float, device=x.device).cpu().numpy(),vuv_vector.cpu().numpy()
+        if f0.shape[0] == 1:
+            return (torch.ones(pad_to, dtype=torch.float, device=x.device) * f0[0]).cpu().numpy() ,vuv_vector.cpu().numpy()
+    
+        # 大概可以用 torch 重写?
+        f0 = np.interp(time_frame, time_org, f0, left=f0[0], right=f0[-1])
+        #vuv_vector = np.ceil(scipy.ndimage.zoom(vuv_vector,pad_to/len(vuv_vector),order = 0))
+        
+        return f0,vuv_vector.cpu().numpy()
+
+    def compute_f0(self,wav,p_len=None):
+        x = torch.FloatTensor(wav).to(self.dtype).to(self.device)
+        if p_len is None:
+            p_len = x.shape[0]//self.hop_length
+        else:
+            assert abs(p_len-x.shape[0]//self.hop_length) < 4, "pad length error"
+        f0 = self.rmvpe.infer_from_audio(x,self.sampling_rate,self.threshold)
+        if torch.all(f0 == 0):
+            rtn = f0.cpu().numpy() if p_len is None else np.zeros(p_len)
+            return rtn,rtn
+        return self.post_process(x,self.sampling_rate,f0,p_len)[0]
+    
+    def compute_f0_uv(self,wav,p_len=None):
+        x = torch.FloatTensor(wav).to(self.dtype).to(self.device)
+        if p_len is None:
+            p_len = x.shape[0]//self.hop_length
+        else:
+            assert abs(p_len-x.shape[0]//self.hop_length) < 4, "pad length error"
+        f0 = self.rmvpe.infer_from_audio(x,self.sampling_rate,self.threshold)
+        if torch.all(f0 == 0):
+            rtn = f0.cpu().numpy() if p_len is None else np.zeros(p_len)
+            return rtn,rtn
+        return self.post_process(x,self.sampling_rate,f0,p_len)

modules/F0Predictor/crepe.py

@@ -1,14 +1,14 @@
-from typing import Optional,Union
+from typing import Optional, Union
+
 try:
     from typing import Literal
-except Exception as e:
+except Exception:
     from typing_extensions import Literal
 import numpy as np
 import torch
 import torchcrepe
 from torch import nn
 from torch.nn import functional as F
-import scipy
 
 #from:https://github.com/fishaudio/fish-diffusion
 
@@ -97,19 +97,19 @@ class BasePitchExtractor:
         f0 = torch.index_select(f0, dim=0, index=nzindex).cpu().numpy()
         time_org = self.hop_length / sampling_rate * nzindex.cpu().numpy()
         time_frame = np.arange(pad_to) * self.hop_length / sampling_rate
+        
+        vuv_vector = F.interpolate(vuv_vector[None,None,:],size=pad_to)[0][0]
 
         if f0.shape[0] <= 0:
-            return torch.zeros(pad_to, dtype=torch.float, device=x.device),torch.zeros(pad_to, dtype=torch.float, device=x.device)
-
+            return torch.zeros(pad_to, dtype=torch.float, device=x.device),vuv_vector.cpu().numpy()
         if f0.shape[0] == 1:
-            return torch.ones(pad_to, dtype=torch.float, device=x.device) * f0[0],torch.ones(pad_to, dtype=torch.float, device=x.device)
+            return torch.ones(pad_to, dtype=torch.float, device=x.device) * f0[0],vuv_vector.cpu().numpy()
     
         # 大概可以用 torch 重写?
         f0 = np.interp(time_frame, time_org, f0, left=f0[0], right=f0[-1])
-        vuv_vector = vuv_vector.cpu().numpy()
-        vuv_vector = np.ceil(scipy.ndimage.zoom(vuv_vector,pad_to/len(vuv_vector),order = 0))
+        #vuv_vector = np.ceil(scipy.ndimage.zoom(vuv_vector,pad_to/len(vuv_vector),order = 0))
         
-        return f0,vuv_vector
+        return f0,vuv_vector.cpu().numpy()
 
 
 class MaskedAvgPool1d(nn.Module):
@@ -323,7 +323,7 @@ class CrepePitchExtractor(BasePitchExtractor):
         else:
             pd = torchcrepe.filter.median(pd, 3)
 
-        pd = torchcrepe.threshold.Silence(-60.0)(pd, x, sampling_rate, 512)
+        pd = torchcrepe.threshold.Silence(-60.0)(pd, x, sampling_rate, self.hop_length)
         f0 = torchcrepe.threshold.At(self.threshold)(f0, pd)
         
         if self.use_fast_filters:
@@ -334,7 +334,7 @@ class CrepePitchExtractor(BasePitchExtractor):
         f0 = torch.where(torch.isnan(f0), torch.full_like(f0, 0), f0)[0]
 
         if torch.all(f0 == 0):
-            rtn = f0.cpu().numpy() if pad_to==None else np.zeros(pad_to)
+            rtn = f0.cpu().numpy() if pad_to is None else np.zeros(pad_to)
             return rtn,rtn
         
         return self.post_process(x, sampling_rate, f0, pad_to)

modules/F0Predictor/fcpe/__init__.py

@@ -0,0 +1,3 @@
+from .model import FCPEInfer  # noqa: F401
+from .nvSTFT import STFT  # noqa: F401
+from .pcmer import PCmer  # noqa: F401

modules/F0Predictor/fcpe/model.py

@@ -0,0 +1,262 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils import weight_norm
+from torchaudio.transforms import Resample
+
+from .nvSTFT import STFT
+from .pcmer import PCmer
+
+
+def l2_regularization(model, l2_alpha):
+    l2_loss = []
+    for module in model.modules():
+        if type(module) is nn.Conv2d:
+            l2_loss.append((module.weight ** 2).sum() / 2.0)
+    return l2_alpha * sum(l2_loss)
+
+
+class FCPE(nn.Module):
+    def __init__(
+            self,
+            input_channel=128,
+            out_dims=360,
+            n_layers=12,
+            n_chans=512,
+            use_siren=False,
+            use_full=False,
+            loss_mse_scale=10,
+            loss_l2_regularization=False,
+            loss_l2_regularization_scale=1,
+            loss_grad1_mse=False,
+            loss_grad1_mse_scale=1,
+            f0_max=1975.5,
+            f0_min=32.70,
+            confidence=False,
+            threshold=0.05,
+            use_input_conv=True
+    ):
+        super().__init__()
+        if use_siren is True:
+            raise ValueError("Siren is not supported yet.")
+        if use_full is True:
+            raise ValueError("Full model is not supported yet.")
+
+        self.loss_mse_scale = loss_mse_scale if (loss_mse_scale is not None) else 10
+        self.loss_l2_regularization = loss_l2_regularization if (loss_l2_regularization is not None) else False
+        self.loss_l2_regularization_scale = loss_l2_regularization_scale if (loss_l2_regularization_scale
+                                                                             is not None) else 1
+        self.loss_grad1_mse = loss_grad1_mse if (loss_grad1_mse is not None) else False
+        self.loss_grad1_mse_scale = loss_grad1_mse_scale if (loss_grad1_mse_scale is not None) else 1
+        self.f0_max = f0_max if (f0_max is not None) else 1975.5
+        self.f0_min = f0_min if (f0_min is not None) else 32.70
+        self.confidence = confidence if (confidence is not None) else False
+        self.threshold = threshold if (threshold is not None) else 0.05
+        self.use_input_conv = use_input_conv if (use_input_conv is not None) else True
+
+        self.cent_table_b = torch.Tensor(
+            np.linspace(self.f0_to_cent(torch.Tensor([f0_min]))[0], self.f0_to_cent(torch.Tensor([f0_max]))[0],
+                        out_dims))
+        self.register_buffer("cent_table", self.cent_table_b)
+
+        # conv in stack
+        _leaky = nn.LeakyReLU()
+        self.stack = nn.Sequential(
+            nn.Conv1d(input_channel, n_chans, 3, 1, 1),
+            nn.GroupNorm(4, n_chans),
+            _leaky,
+            nn.Conv1d(n_chans, n_chans, 3, 1, 1))
+
+        # transformer
+        self.decoder = PCmer(
+            num_layers=n_layers,
+            num_heads=8,
+            dim_model=n_chans,
+            dim_keys=n_chans,
+            dim_values=n_chans,
+            residual_dropout=0.1,
+            attention_dropout=0.1)
+        self.norm = nn.LayerNorm(n_chans)
+
+        # out
+        self.n_out = out_dims
+        self.dense_out = weight_norm(
+            nn.Linear(n_chans, self.n_out))
+
+    def forward(self, mel, infer=True, gt_f0=None, return_hz_f0=False, cdecoder = "local_argmax"):
+        """
+        input:
+            B x n_frames x n_unit
+        return:
+            dict of B x n_frames x feat
+        """
+        if cdecoder == "argmax":
+            self.cdecoder = self.cents_decoder
+        elif cdecoder == "local_argmax":
+            self.cdecoder = self.cents_local_decoder
+        if self.use_input_conv:
+            x = self.stack(mel.transpose(1, 2)).transpose(1, 2)
+        else:
+            x = mel
+        x = self.decoder(x)
+        x = self.norm(x)
+        x = self.dense_out(x)  # [B,N,D]
+        x = torch.sigmoid(x)
+        if not infer:
+            gt_cent_f0 = self.f0_to_cent(gt_f0)  # mel f0  #[B,N,1]
+            gt_cent_f0 = self.gaussian_blurred_cent(gt_cent_f0)  # #[B,N,out_dim]
+            loss_all = self.loss_mse_scale * F.binary_cross_entropy(x, gt_cent_f0)  # bce loss
+            # l2 regularization
+            if self.loss_l2_regularization:
+                loss_all = loss_all + l2_regularization(model=self, l2_alpha=self.loss_l2_regularization_scale)
+            x = loss_all
+        if infer:
+            x = self.cdecoder(x)
+            x = self.cent_to_f0(x)
+            if not return_hz_f0:
+                x = (1 + x / 700).log()
+        return x
+
+    def cents_decoder(self, y, mask=True):
+        B, N, _ = y.size()
+        ci = self.cent_table[None, None, :].expand(B, N, -1)
+        rtn = torch.sum(ci * y, dim=-1, keepdim=True) / torch.sum(y, dim=-1, keepdim=True)  # cents: [B,N,1]
+        if mask:
+            confident = torch.max(y, dim=-1, keepdim=True)[0]
+            confident_mask = torch.ones_like(confident)
+            confident_mask[confident <= self.threshold] = float("-INF")
+            rtn = rtn * confident_mask
+        if self.confidence:
+            return rtn, confident
+        else:
+            return rtn
+        
+    def cents_local_decoder(self, y, mask=True):
+        B, N, _ = y.size()
+        ci = self.cent_table[None, None, :].expand(B, N, -1)
+        confident, max_index = torch.max(y, dim=-1, keepdim=True)
+        local_argmax_index = torch.arange(0,9).to(max_index.device) + (max_index - 4)
+        local_argmax_index[local_argmax_index<0] = 0
+        local_argmax_index[local_argmax_index>=self.n_out] = self.n_out - 1
+        ci_l = torch.gather(ci,-1,local_argmax_index)
+        y_l = torch.gather(y,-1,local_argmax_index)
+        rtn = torch.sum(ci_l * y_l, dim=-1, keepdim=True) / torch.sum(y_l, dim=-1, keepdim=True)  # cents: [B,N,1]
+        if mask:
+            confident_mask = torch.ones_like(confident)
+            confident_mask[confident <= self.threshold] = float("-INF")
+            rtn = rtn * confident_mask
+        if self.confidence:
+            return rtn, confident
+        else:
+            return rtn
+
+    def cent_to_f0(self, cent):
+        return 10. * 2 ** (cent / 1200.)
+
+    def f0_to_cent(self, f0):
+        return 1200. * torch.log2(f0 / 10.)
+
+    def gaussian_blurred_cent(self, cents):  # cents: [B,N,1]
+        mask = (cents > 0.1) & (cents < (1200. * np.log2(self.f0_max / 10.)))
+        B, N, _ = cents.size()
+        ci = self.cent_table[None, None, :].expand(B, N, -1)
+        return torch.exp(-torch.square(ci - cents) / 1250) * mask.float()
+
+
+class FCPEInfer:
+    def __init__(self, model_path, device=None, dtype=torch.float32):
+        if device is None:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.device = device
+        ckpt = torch.load(model_path, map_location=torch.device(self.device))
+        self.args = DotDict(ckpt["config"])
+        self.dtype = dtype
+        model = FCPE(
+            input_channel=self.args.model.input_channel,
+            out_dims=self.args.model.out_dims,
+            n_layers=self.args.model.n_layers,
+            n_chans=self.args.model.n_chans,
+            use_siren=self.args.model.use_siren,
+            use_full=self.args.model.use_full,
+            loss_mse_scale=self.args.loss.loss_mse_scale,
+            loss_l2_regularization=self.args.loss.loss_l2_regularization,
+            loss_l2_regularization_scale=self.args.loss.loss_l2_regularization_scale,
+            loss_grad1_mse=self.args.loss.loss_grad1_mse,
+            loss_grad1_mse_scale=self.args.loss.loss_grad1_mse_scale,
+            f0_max=self.args.model.f0_max,
+            f0_min=self.args.model.f0_min,
+            confidence=self.args.model.confidence,
+        )
+        model.to(self.device).to(self.dtype)
+        model.load_state_dict(ckpt['model'])
+        model.eval()
+        self.model = model
+        self.wav2mel = Wav2Mel(self.args, dtype=self.dtype, device=self.device)
+
+    @torch.no_grad()
+    def __call__(self, audio, sr, threshold=0.05):
+        self.model.threshold = threshold
+        audio = audio[None,:]
+        mel = self.wav2mel(audio=audio, sample_rate=sr).to(self.dtype)
+        f0 = self.model(mel=mel, infer=True, return_hz_f0=True)
+        return f0
+
+
+class Wav2Mel:
+
+    def __init__(self, args, device=None, dtype=torch.float32):
+        # self.args = args
+        self.sampling_rate = args.mel.sampling_rate
+        self.hop_size = args.mel.hop_size
+        if device is None:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.device = device
+        self.dtype = dtype
+        self.stft = STFT(
+            args.mel.sampling_rate,
+            args.mel.num_mels,
+            args.mel.n_fft,
+            args.mel.win_size,
+            args.mel.hop_size,
+            args.mel.fmin,
+            args.mel.fmax
+        )
+        self.resample_kernel = {}
+
+    def extract_nvstft(self, audio, keyshift=0, train=False):
+        mel = self.stft.get_mel(audio, keyshift=keyshift, train=train).transpose(1, 2)  # B, n_frames, bins
+        return mel
+
+    def extract_mel(self, audio, sample_rate, keyshift=0, train=False):
+        audio = audio.to(self.dtype).to(self.device)
+        # resample
+        if sample_rate == self.sampling_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.sampling_rate, lowpass_filter_width=128)
+            self.resample_kernel[key_str] = self.resample_kernel[key_str].to(self.dtype).to(self.device)
+            audio_res = self.resample_kernel[key_str](audio)
+
+        # extract
+        mel = self.extract_nvstft(audio_res, keyshift=keyshift, train=train)  # B, n_frames, bins
+        n_frames = int(audio.shape[1] // self.hop_size) + 1
+        if n_frames > int(mel.shape[1]):
+            mel = torch.cat((mel, mel[:, -1:, :]), 1)
+        if n_frames < int(mel.shape[1]):
+            mel = mel[:, :n_frames, :]
+        return mel
+
+    def __call__(self, audio, sample_rate, keyshift=0, train=False):
+        return self.extract_mel(audio, sample_rate, keyshift=keyshift, train=train)
+
+
+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__

modules/F0Predictor/fcpe/nvSTFT.py

@@ -0,0 +1,133 @@
+import os
+
+import librosa
+import numpy as np
+import soundfile as sf
+import torch
+import torch.nn.functional as F
+import torch.utils.data
+from librosa.filters import mel as librosa_mel_fn
+
+os.environ["LRU_CACHE_CAPACITY"] = "3"
+
+def load_wav_to_torch(full_path, target_sr=None, return_empty_on_exception=False):
+    sampling_rate = None
+    try:
+        data, sampling_rate = sf.read(full_path, always_2d=True)# than soundfile.
+    except Exception as ex:
+        print(f"'{full_path}' failed to load.\nException:")
+        print(ex)
+        if return_empty_on_exception:
+            return [], sampling_rate or target_sr or 48000
+        else:
+            raise Exception(ex)
+    
+    if len(data.shape) > 1:
+        data = data[:, 0]
+        assert len(data) > 2# check duration of audio file is > 2 samples (because otherwise the slice operation was on the wrong dimension)
+    
+    if np.issubdtype(data.dtype, np.integer): # if audio data is type int
+        max_mag = -np.iinfo(data.dtype).min # maximum magnitude = min possible value of intXX
+    else: # if audio data is type fp32
+        max_mag = max(np.amax(data), -np.amin(data))
+        max_mag = (2**31)+1 if max_mag > (2**15) else ((2**15)+1 if max_mag > 1.01 else 1.0) # data should be either 16-bit INT, 32-bit INT or [-1 to 1] float32
+    
+    data = torch.FloatTensor(data.astype(np.float32))/max_mag
+    
+    if (torch.isinf(data) | torch.isnan(data)).any() and return_empty_on_exception:# resample will crash with inf/NaN inputs. return_empty_on_exception will return empty arr instead of except
+        return [], sampling_rate or target_sr or 48000
+    if target_sr is not None and sampling_rate != target_sr:
+        data = torch.from_numpy(librosa.core.resample(data.numpy(), orig_sr=sampling_rate, target_sr=target_sr))
+        sampling_rate = target_sr
+    
+    return data, sampling_rate
+
+def dynamic_range_compression(x, C=1, clip_val=1e-5):
+    return np.log(np.clip(x, a_min=clip_val, a_max=None) * C)
+
+def dynamic_range_decompression(x, C=1):
+    return np.exp(x) / C
+
+def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
+    return torch.log(torch.clamp(x, min=clip_val) * C)
+
+def dynamic_range_decompression_torch(x, C=1):
+    return torch.exp(x) / C
+
+class STFT():
+    def __init__(self, sr=22050, n_mels=80, n_fft=1024, win_size=1024, hop_length=256, fmin=20, fmax=11025, clip_val=1e-5):
+        self.target_sr = sr
+        
+        self.n_mels     = n_mels
+        self.n_fft      = n_fft
+        self.win_size   = win_size
+        self.hop_length = hop_length
+        self.fmin     = fmin
+        self.fmax     = fmax
+        self.clip_val = clip_val
+        self.mel_basis = {}
+        self.hann_window = {}
+    
+    def get_mel(self, y, keyshift=0, speed=1, center=False, train=False):
+        sampling_rate = self.target_sr
+        n_mels     = self.n_mels
+        n_fft      = self.n_fft
+        win_size   = self.win_size
+        hop_length = self.hop_length
+        fmin       = self.fmin
+        fmax       = self.fmax
+        clip_val   = self.clip_val
+        
+        factor = 2 ** (keyshift / 12)       
+        n_fft_new = int(np.round(n_fft * factor))
+        win_size_new = int(np.round(win_size * factor))
+        hop_length_new = int(np.round(hop_length * speed))
+        if not train:
+            mel_basis = self.mel_basis
+            hann_window = self.hann_window
+        else:
+            mel_basis = {}
+            hann_window = {}
+        
+        if torch.min(y) < -1.:
+            print('min value is ', torch.min(y))
+        if torch.max(y) > 1.:
+            print('max value is ', torch.max(y))
+        
+        mel_basis_key = str(fmax)+'_'+str(y.device)
+        if mel_basis_key not in mel_basis:
+            mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax)
+            mel_basis[mel_basis_key] = torch.from_numpy(mel).float().to(y.device)
+        
+        keyshift_key = str(keyshift)+'_'+str(y.device)
+        if keyshift_key not in hann_window:
+            hann_window[keyshift_key] = torch.hann_window(win_size_new).to(y.device)
+        
+        pad_left = (win_size_new - hop_length_new) //2
+        pad_right = max((win_size_new- hop_length_new + 1) //2, win_size_new - y.size(-1) - pad_left)
+        if pad_right < y.size(-1):
+            mode = 'reflect'
+        else:
+            mode = 'constant'
+        y = torch.nn.functional.pad(y.unsqueeze(1), (pad_left, pad_right), mode = mode)
+        y = y.squeeze(1)
+        
+        spec = torch.stft(y, n_fft_new, hop_length=hop_length_new, win_length=win_size_new, window=hann_window[keyshift_key],
+                          center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=True)                          
+        spec = torch.sqrt(spec.real.pow(2) + spec.imag.pow(2) + (1e-9))
+        if keyshift != 0:
+            size = n_fft // 2 + 1
+            resize = spec.size(1)
+            if resize < size:
+                spec = F.pad(spec, (0, 0, 0, size-resize))
+            spec = spec[:, :size, :] * win_size / win_size_new   
+        spec = torch.matmul(mel_basis[mel_basis_key], spec)
+        spec = dynamic_range_compression_torch(spec, clip_val=clip_val)
+        return spec
+    
+    def __call__(self, audiopath):
+        audio, sr = load_wav_to_torch(audiopath, target_sr=self.target_sr)
+        spect = self.get_mel(audio.unsqueeze(0)).squeeze(0)
+        return spect
+
+stft = STFT()

modules/F0Predictor/fcpe/pcmer.py

@@ -0,0 +1,369 @@
+import math
+from functools import partial
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from local_attention import LocalAttention
+from torch import nn
+
+#import fast_transformers.causal_product.causal_product_cuda
+
+def softmax_kernel(data, *, projection_matrix, is_query, normalize_data=True, eps=1e-4, device = None):
+    b, h, *_ = data.shape
+    # (batch size, head, length, model_dim)
+
+    # normalize model dim
+    data_normalizer = (data.shape[-1] ** -0.25) if normalize_data else 1.
+
+    # what is ration?, projection_matrix.shape[0] --> 266
+    
+    ratio = (projection_matrix.shape[0] ** -0.5)
+
+    projection = repeat(projection_matrix, 'j d -> b h j d', b = b, h = h)
+    projection = projection.type_as(data)
+
+    #data_dash = w^T x
+    data_dash = torch.einsum('...id,...jd->...ij', (data_normalizer * data), projection)
+
+    
+    # diag_data = D**2 
+    diag_data = data ** 2
+    diag_data = torch.sum(diag_data, dim=-1)
+    diag_data = (diag_data / 2.0) * (data_normalizer ** 2)
+    diag_data = diag_data.unsqueeze(dim=-1)
+    
+    #print ()
+    if is_query:
+        data_dash = ratio * (
+            torch.exp(data_dash - diag_data -
+                    torch.max(data_dash, dim=-1, keepdim=True).values) + eps)
+    else:
+        data_dash = ratio * (
+            torch.exp(data_dash - diag_data + eps))#- torch.max(data_dash)) + eps)
+
+    return data_dash.type_as(data)
+
+def orthogonal_matrix_chunk(cols, qr_uniform_q = False, device = None):
+    unstructured_block = torch.randn((cols, cols), device = device)
+    q, r = torch.linalg.qr(unstructured_block.cpu(), mode='reduced')
+    q, r = map(lambda t: t.to(device), (q, r))
+
+    # proposed by @Parskatt
+    # to make sure Q is uniform https://arxiv.org/pdf/math-ph/0609050.pdf
+    if qr_uniform_q:
+        d = torch.diag(r, 0)
+        q *= d.sign()
+    return q.t()
+def exists(val):
+    return val is not None
+
+def empty(tensor):
+    return tensor.numel() == 0
+
+def default(val, d):
+    return val if exists(val) else d
+
+def cast_tuple(val):
+    return (val,) if not isinstance(val, tuple) else val
+
+class PCmer(nn.Module):
+    """The encoder that is used in the Transformer model."""
+    
+    def __init__(self, 
+                num_layers,
+                num_heads,
+                dim_model,
+                dim_keys,
+                dim_values,
+                residual_dropout,
+                attention_dropout):
+        super().__init__()
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.dim_model = dim_model
+        self.dim_values = dim_values
+        self.dim_keys = dim_keys
+        self.residual_dropout = residual_dropout
+        self.attention_dropout = attention_dropout
+
+        self._layers = nn.ModuleList([_EncoderLayer(self) for _ in range(num_layers)])
+        
+    #  METHODS  ########################################################################################################
+    
+    def forward(self, phone, mask=None):
+        
+        # apply all layers to the input
+        for (i, layer) in enumerate(self._layers):
+            phone = layer(phone, mask)
+        # provide the final sequence
+        return phone
+
+
+# ==================================================================================================================== #
+#  CLASS  _ E N C O D E R  L A Y E R                                                                                   #
+# ==================================================================================================================== #
+
+
+class _EncoderLayer(nn.Module):
+    """One layer of the encoder.
+    
+    Attributes:
+        attn: (:class:`mha.MultiHeadAttention`): The attention mechanism that is used to read the input sequence.
+        feed_forward (:class:`ffl.FeedForwardLayer`): The feed-forward layer on top of the attention mechanism.
+    """
+    
+    def __init__(self, parent: PCmer):
+        """Creates a new instance of ``_EncoderLayer``.
+        
+        Args:
+            parent (Encoder): The encoder that the layers is created for.
+        """
+        super().__init__()
+        
+        
+        self.conformer = ConformerConvModule(parent.dim_model)
+        self.norm = nn.LayerNorm(parent.dim_model)
+        self.dropout = nn.Dropout(parent.residual_dropout)
+        
+        # selfatt -> fastatt: performer!
+        self.attn = SelfAttention(dim = parent.dim_model,
+                                  heads = parent.num_heads,
+                                  causal = False)
+        
+    #  METHODS  ########################################################################################################
+
+    def forward(self, phone, mask=None):
+        
+        # compute attention sub-layer
+        phone = phone + (self.attn(self.norm(phone), mask=mask))
+        
+        phone = phone + (self.conformer(phone))
+        
+        return phone 
+
+def calc_same_padding(kernel_size):
+    pad = kernel_size // 2
+    return (pad, pad - (kernel_size + 1) % 2)
+
+# helper classes
+
+class Swish(nn.Module):
+    def forward(self, x):
+        return x * x.sigmoid()
+
+class Transpose(nn.Module):
+    def __init__(self, dims):
+        super().__init__()
+        assert len(dims) == 2, 'dims must be a tuple of two dimensions'
+        self.dims = dims
+
+    def forward(self, x):
+        return x.transpose(*self.dims)
+
+class GLU(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        out, gate = x.chunk(2, dim=self.dim)
+        return out * gate.sigmoid()
+
+class DepthWiseConv1d(nn.Module):
+    def __init__(self, chan_in, chan_out, kernel_size, padding):
+        super().__init__()
+        self.padding = padding
+        self.conv = nn.Conv1d(chan_in, chan_out, kernel_size, groups = chan_in)
+
+    def forward(self, x):
+        x = F.pad(x, self.padding)
+        return self.conv(x)
+
+class ConformerConvModule(nn.Module):
+    def __init__(
+        self,
+        dim,
+        causal = False,
+        expansion_factor = 2,
+        kernel_size = 31,
+        dropout = 0.):
+        super().__init__()
+
+        inner_dim = dim * expansion_factor
+        padding = calc_same_padding(kernel_size) if not causal else (kernel_size - 1, 0)
+
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            Transpose((1, 2)),
+            nn.Conv1d(dim, inner_dim * 2, 1),
+            GLU(dim=1),
+            DepthWiseConv1d(inner_dim, inner_dim, kernel_size = kernel_size, padding = padding),
+            #nn.BatchNorm1d(inner_dim) if not causal else nn.Identity(),
+            Swish(),
+            nn.Conv1d(inner_dim, dim, 1),
+            Transpose((1, 2)),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+def linear_attention(q, k, v):
+    if v is None:
+        #print (k.size(), q.size())
+        out = torch.einsum('...ed,...nd->...ne', k, q)
+        return out
+
+    else:
+        k_cumsum = k.sum(dim = -2) 
+        #k_cumsum = k.sum(dim = -2)
+        D_inv = 1. / (torch.einsum('...nd,...d->...n', q, k_cumsum.type_as(q)) + 1e-8)
+
+        context = torch.einsum('...nd,...ne->...de', k, v)
+        #print ("TRUEEE: ", context.size(), q.size(), D_inv.size())
+        out = torch.einsum('...de,...nd,...n->...ne', context, q, D_inv)
+        return out
+
+def gaussian_orthogonal_random_matrix(nb_rows, nb_columns, scaling = 0, qr_uniform_q = False, device = None):
+    nb_full_blocks = int(nb_rows / nb_columns)
+    #print (nb_full_blocks)
+    block_list = []
+
+    for _ in range(nb_full_blocks):
+        q = orthogonal_matrix_chunk(nb_columns, qr_uniform_q = qr_uniform_q, device = device)
+        block_list.append(q)
+    # block_list[n] is a orthogonal matrix ... (model_dim * model_dim)
+    #print (block_list[0].size(), torch.einsum('...nd,...nd->...n', block_list[0], torch.roll(block_list[0],1,1)))
+    #print (nb_rows, nb_full_blocks, nb_columns)
+    remaining_rows = nb_rows - nb_full_blocks * nb_columns
+    #print (remaining_rows)
+    if remaining_rows > 0:
+        q = orthogonal_matrix_chunk(nb_columns, qr_uniform_q = qr_uniform_q, device = device)
+        #print (q[:remaining_rows].size())
+        block_list.append(q[:remaining_rows])
+
+    final_matrix = torch.cat(block_list)
+    
+    if scaling == 0:
+        multiplier = torch.randn((nb_rows, nb_columns), device = device).norm(dim = 1)
+    elif scaling == 1:
+        multiplier = math.sqrt((float(nb_columns))) * torch.ones((nb_rows,), device = device)
+    else:
+        raise ValueError(f'Invalid scaling {scaling}')
+
+    return torch.diag(multiplier) @ final_matrix
+
+class FastAttention(nn.Module):
+    def __init__(self, dim_heads, nb_features = None, ortho_scaling = 0, causal = False, generalized_attention = False, kernel_fn = nn.ReLU(), qr_uniform_q = False, no_projection = False):
+        super().__init__()
+        nb_features = default(nb_features, int(dim_heads * math.log(dim_heads)))
+
+        self.dim_heads = dim_heads
+        self.nb_features = nb_features
+        self.ortho_scaling = ortho_scaling
+
+        self.create_projection = partial(gaussian_orthogonal_random_matrix, nb_rows = self.nb_features, nb_columns = dim_heads, scaling = ortho_scaling, qr_uniform_q = qr_uniform_q)
+        projection_matrix = self.create_projection()
+        self.register_buffer('projection_matrix', projection_matrix)
+
+        self.generalized_attention = generalized_attention
+        self.kernel_fn = kernel_fn
+
+        # if this is turned on, no projection will be used
+        # queries and keys will be softmax-ed as in the original efficient attention paper
+        self.no_projection = no_projection
+
+        self.causal = causal
+
+    @torch.no_grad()
+    def redraw_projection_matrix(self):
+        projections = self.create_projection()
+        self.projection_matrix.copy_(projections)
+        del projections
+
+    def forward(self, q, k, v):
+        device = q.device
+
+        if self.no_projection:
+            q = q.softmax(dim = -1)
+            k = torch.exp(k) if self.causal else k.softmax(dim = -2)
+        else:
+            create_kernel = partial(softmax_kernel, projection_matrix = self.projection_matrix, device = device)
+            
+            q = create_kernel(q, is_query = True)
+            k = create_kernel(k, is_query = False)
+
+        attn_fn = linear_attention if not self.causal else self.causal_linear_fn
+        if v is None:
+            out = attn_fn(q, k, None)
+            return out
+        else:
+            out = attn_fn(q, k, v)
+            return out
+class SelfAttention(nn.Module):
+    def __init__(self, dim, causal = False, heads = 8, dim_head = 64, local_heads = 0, local_window_size = 256, nb_features = None, feature_redraw_interval = 1000, generalized_attention = False, kernel_fn = nn.ReLU(), qr_uniform_q = False, dropout = 0., no_projection = False):
+        super().__init__()
+        assert dim % heads == 0, 'dimension must be divisible by number of heads'
+        dim_head = default(dim_head, dim // heads)
+        inner_dim = dim_head * heads
+        self.fast_attention = FastAttention(dim_head, nb_features, causal = causal, generalized_attention = generalized_attention, kernel_fn = kernel_fn, qr_uniform_q = qr_uniform_q, no_projection = no_projection)
+
+        self.heads = heads
+        self.global_heads = heads - local_heads
+        self.local_attn = LocalAttention(window_size = local_window_size, causal = causal, autopad = True, dropout = dropout, look_forward = int(not causal), rel_pos_emb_config = (dim_head, local_heads)) if local_heads > 0 else None
+
+        #print (heads, nb_features, dim_head)
+        #name_embedding = torch.zeros(110, heads, dim_head, dim_head)
+        #self.name_embedding = nn.Parameter(name_embedding, requires_grad=True)
+        
+
+        self.to_q = nn.Linear(dim, inner_dim)
+        self.to_k = nn.Linear(dim, inner_dim)
+        self.to_v = nn.Linear(dim, inner_dim)
+        self.to_out = nn.Linear(inner_dim, dim)
+        self.dropout = nn.Dropout(dropout)
+
+    @torch.no_grad()
+    def redraw_projection_matrix(self):
+        self.fast_attention.redraw_projection_matrix()
+        #torch.nn.init.zeros_(self.name_embedding)
+        #print (torch.sum(self.name_embedding))
+    def forward(self, x, context = None, mask = None, context_mask = None, name=None, inference=False, **kwargs):
+        _, _, _, h, gh = *x.shape, self.heads, self.global_heads
+        
+        cross_attend = exists(context)
+
+        context = default(context, x)
+        context_mask = default(context_mask, mask) if not cross_attend else context_mask
+        #print (torch.sum(self.name_embedding))
+        q, k, v = self.to_q(x), self.to_k(context), self.to_v(context)
+
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, k, v))
+        (q, lq), (k, lk), (v, lv) = map(lambda t: (t[:, :gh], t[:, gh:]), (q, k, v))
+
+        attn_outs = []
+        #print (name)
+        #print (self.name_embedding[name].size())
+        if not empty(q):
+            if exists(context_mask):
+                global_mask = context_mask[:, None, :, None]
+                v.masked_fill_(~global_mask, 0.)
+            if cross_attend:
+                pass
+                #print (torch.sum(self.name_embedding))
+                #out = self.fast_attention(q,self.name_embedding[name],None)
+                #print (torch.sum(self.name_embedding[...,-1:]))
+            else:
+                out = self.fast_attention(q, k, v)
+            attn_outs.append(out)
+
+        if not empty(lq):
+            assert not cross_attend, 'local attention is not compatible with cross attention'
+            out = self.local_attn(lq, lk, lv, input_mask = mask)
+            attn_outs.append(out)
+
+        out = torch.cat(attn_outs, dim = 1)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        out =  self.to_out(out)
+        return self.dropout(out)

modules/F0Predictor/rmvpe/__init__.py

@@ -0,0 +1,10 @@
+from .constants import *  # noqa: F403
+from .inference import RMVPE  # noqa: F401
+from .model import E2E, E2E0  # noqa: F401
+from .spec import MelSpectrogram  # noqa: F401
+from .utils import (  # noqa: F401
+    cycle,
+    summary,
+    to_local_average_cents,
+    to_viterbi_cents,
+)

modules/F0Predictor/rmvpe/constants.py

@@ -0,0 +1,9 @@
+SAMPLE_RATE = 16000
+
+N_CLASS = 360
+
+N_MELS = 128
+MEL_FMIN = 30
+MEL_FMAX = SAMPLE_RATE // 2
+WINDOW_LENGTH = 1024
+CONST = 1997.3794084376191

modules/F0Predictor/rmvpe/deepunet.py

@@ -0,0 +1,190 @@
+import torch
+import torch.nn as nn
+
+from .constants import N_MELS
+
+
+class ConvBlockRes(nn.Module):
+    def __init__(self, in_channels, out_channels, momentum=0.01):
+        super(ConvBlockRes, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels=in_channels,
+                      out_channels=out_channels,
+                      kernel_size=(3, 3),
+                      stride=(1, 1),
+                      padding=(1, 1),
+                      bias=False),
+            nn.BatchNorm2d(out_channels, momentum=momentum),
+            nn.ReLU(),
+
+            nn.Conv2d(in_channels=out_channels,
+                      out_channels=out_channels,
+                      kernel_size=(3, 3),
+                      stride=(1, 1),
+                      padding=(1, 1),
+                      bias=False),
+            nn.BatchNorm2d(out_channels, momentum=momentum),
+            nn.ReLU(),
+        )
+        if in_channels != out_channels:
+            self.shortcut = nn.Conv2d(in_channels, out_channels, (1, 1))
+            self.is_shortcut = True
+        else:
+            self.is_shortcut = False
+
+    def forward(self, x):
+        if self.is_shortcut:
+            return self.conv(x) + self.shortcut(x)
+        else:
+            return self.conv(x) + x
+
+
+class ResEncoderBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, n_blocks=1, momentum=0.01):
+        super(ResEncoderBlock, self).__init__()
+        self.n_blocks = n_blocks
+        self.conv = nn.ModuleList()
+        self.conv.append(ConvBlockRes(in_channels, out_channels, momentum))
+        for i in range(n_blocks - 1):
+            self.conv.append(ConvBlockRes(out_channels, out_channels, momentum))
+        self.kernel_size = kernel_size
+        if self.kernel_size is not None:
+            self.pool = nn.AvgPool2d(kernel_size=kernel_size)
+
+    def forward(self, x):
+        for i in range(self.n_blocks):
+            x = self.conv[i](x)
+        if self.kernel_size is not None:
+            return x, self.pool(x)
+        else:
+            return x
+
+
+class ResDecoderBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride, n_blocks=1, momentum=0.01):
+        super(ResDecoderBlock, self).__init__()
+        out_padding = (0, 1) if stride == (1, 2) else (1, 1)
+        self.n_blocks = n_blocks
+        self.conv1 = nn.Sequential(
+            nn.ConvTranspose2d(in_channels=in_channels,
+                               out_channels=out_channels,
+                               kernel_size=(3, 3),
+                               stride=stride,
+                               padding=(1, 1),
+                               output_padding=out_padding,
+                               bias=False),
+            nn.BatchNorm2d(out_channels, momentum=momentum),
+            nn.ReLU(),
+        )
+        self.conv2 = nn.ModuleList()
+        self.conv2.append(ConvBlockRes(out_channels * 2, out_channels, momentum))
+        for i in range(n_blocks-1):
+            self.conv2.append(ConvBlockRes(out_channels, out_channels, momentum))
+
+    def forward(self, x, concat_tensor):
+        x = self.conv1(x)
+        x = torch.cat((x, concat_tensor), dim=1)
+        for i in range(self.n_blocks):
+            x = self.conv2[i](x)
+        return x
+
+
+class Encoder(nn.Module):
+    def __init__(self, in_channels, in_size, n_encoders, kernel_size, n_blocks, out_channels=16, momentum=0.01):
+        super(Encoder, self).__init__()
+        self.n_encoders = n_encoders
+        self.bn = nn.BatchNorm2d(in_channels, momentum=momentum)
+        self.layers = nn.ModuleList()
+        self.latent_channels = []
+        for i in range(self.n_encoders):
+            self.layers.append(ResEncoderBlock(in_channels, out_channels, kernel_size, n_blocks, momentum=momentum))
+            self.latent_channels.append([out_channels, in_size])
+            in_channels = out_channels
+            out_channels *= 2
+            in_size //= 2
+        self.out_size = in_size
+        self.out_channel = out_channels
+
+    def forward(self, x):
+        concat_tensors = []
+        x = self.bn(x)
+        for i in range(self.n_encoders):
+            _, x = self.layers[i](x)
+            concat_tensors.append(_)
+        return x, concat_tensors
+
+
+class Intermediate(nn.Module):
+    def __init__(self, in_channels, out_channels, n_inters, n_blocks, momentum=0.01):
+        super(Intermediate, self).__init__()
+        self.n_inters = n_inters
+        self.layers = nn.ModuleList()
+        self.layers.append(ResEncoderBlock(in_channels, out_channels, None, n_blocks, momentum))
+        for i in range(self.n_inters-1):
+            self.layers.append(ResEncoderBlock(out_channels, out_channels, None, n_blocks, momentum))
+
+    def forward(self, x):
+        for i in range(self.n_inters):
+            x = self.layers[i](x)
+        return x
+
+
+class Decoder(nn.Module):
+    def __init__(self, in_channels, n_decoders, stride, n_blocks, momentum=0.01):
+        super(Decoder, self).__init__()
+        self.layers = nn.ModuleList()
+        self.n_decoders = n_decoders
+        for i in range(self.n_decoders):
+            out_channels = in_channels // 2
+            self.layers.append(ResDecoderBlock(in_channels, out_channels, stride, n_blocks, momentum))
+            in_channels = out_channels
+
+    def forward(self, x, concat_tensors):
+        for i in range(self.n_decoders):
+            x = self.layers[i](x, concat_tensors[-1-i])
+        return x
+
+
+class TimbreFilter(nn.Module):
+    def __init__(self, latent_rep_channels):
+        super(TimbreFilter, self).__init__()
+        self.layers = nn.ModuleList()
+        for latent_rep in latent_rep_channels:
+            self.layers.append(ConvBlockRes(latent_rep[0], latent_rep[0]))
+
+    def forward(self, x_tensors):
+        out_tensors = []
+        for i, layer in enumerate(self.layers):
+            out_tensors.append(layer(x_tensors[i]))
+        return out_tensors
+
+
+class DeepUnet(nn.Module):
+    def __init__(self, kernel_size, n_blocks, en_de_layers=5, inter_layers=4, in_channels=1, en_out_channels=16):
+        super(DeepUnet, self).__init__()
+        self.encoder = Encoder(in_channels, N_MELS, en_de_layers, kernel_size, n_blocks, en_out_channels)
+        self.intermediate = Intermediate(self.encoder.out_channel // 2, self.encoder.out_channel, inter_layers, n_blocks)
+        self.tf = TimbreFilter(self.encoder.latent_channels)
+        self.decoder = Decoder(self.encoder.out_channel, en_de_layers, kernel_size, n_blocks)
+
+    def forward(self, x):
+        x, concat_tensors = self.encoder(x)
+        x = self.intermediate(x)
+        concat_tensors = self.tf(concat_tensors)
+        x = self.decoder(x, concat_tensors)
+        return x
+
+      
+class DeepUnet0(nn.Module):
+    def __init__(self, kernel_size, n_blocks, en_de_layers=5, inter_layers=4, in_channels=1, en_out_channels=16):
+        super(DeepUnet0, self).__init__()
+        self.encoder = Encoder(in_channels, N_MELS, en_de_layers, kernel_size, n_blocks, en_out_channels)
+        self.intermediate = Intermediate(self.encoder.out_channel // 2, self.encoder.out_channel, inter_layers, n_blocks)
+        self.tf = TimbreFilter(self.encoder.latent_channels)
+        self.decoder = Decoder(self.encoder.out_channel, en_de_layers, kernel_size, n_blocks)
+
+    def forward(self, x):
+        x, concat_tensors = self.encoder(x)
+        x = self.intermediate(x)
+        x = self.decoder(x, concat_tensors)
+        return x

modules/F0Predictor/rmvpe/inference.py

@@ -0,0 +1,57 @@
+import torch
+import torch.nn.functional as F
+from torchaudio.transforms import Resample
+
+from .constants import *  # noqa: F403
+from .model import E2E0
+from .spec import MelSpectrogram
+from .utils import to_local_average_cents, to_viterbi_cents
+
+
+class RMVPE:
+    def __init__(self, model_path, device=None, dtype = torch.float32, hop_length=160):
+        self.resample_kernel = {}
+        if device is None:
+            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        else:
+            self.device = device
+        model = E2E0(4, 1, (2, 2))
+        ckpt = torch.load(model_path, map_location=torch.device(self.device))
+        model.load_state_dict(ckpt['model'])
+        model = model.to(dtype).to(self.device)
+        model.eval()
+        self.model = model
+        self.dtype = dtype
+        self.mel_extractor = MelSpectrogram(N_MELS, SAMPLE_RATE, WINDOW_LENGTH, hop_length, None, MEL_FMIN, MEL_FMAX)  # noqa: F405
+        self.resample_kernel = {}
+
+    def mel2hidden(self, mel):
+        with torch.no_grad():
+            n_frames = mel.shape[-1]
+            mel = F.pad(mel, (0, 32 * ((n_frames - 1) // 32 + 1) - n_frames), mode='constant')
+            hidden = self.model(mel)
+            return hidden[:, :n_frames]
+
+    def decode(self, hidden, thred=0.03, use_viterbi=False):
+        if use_viterbi:
+            cents_pred = to_viterbi_cents(hidden, thred=thred)
+        else:
+            cents_pred = to_local_average_cents(hidden, thred=thred)
+        f0 = torch.Tensor([10 * (2 ** (cent_pred / 1200)) if cent_pred else 0 for cent_pred in cents_pred]).to(self.device)
+        return f0
+
+    def infer_from_audio(self, audio, sample_rate=16000, thred=0.05, use_viterbi=False):
+        audio = audio.unsqueeze(0).to(self.dtype).to(self.device)
+        if sample_rate == 16000:
+            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, 16000, lowpass_filter_width=128)
+            self.resample_kernel[key_str] = self.resample_kernel[key_str].to(self.dtype).to(self.device)
+            audio_res = self.resample_kernel[key_str](audio)
+        mel_extractor = self.mel_extractor.to(self.device)
+        mel = mel_extractor(audio_res, center=True).to(self.dtype)
+        hidden = self.mel2hidden(mel)
+        f0 = self.decode(hidden.squeeze(0), thred=thred, use_viterbi=use_viterbi)
+        return f0

modules/F0Predictor/rmvpe/model.py

@@ -0,0 +1,67 @@
+from torch import nn
+
+from .constants import *  # noqa: F403
+from .deepunet import DeepUnet, DeepUnet0
+from .seq import BiGRU
+from .spec import MelSpectrogram
+
+
+class E2E(nn.Module):
+    def __init__(self, hop_length, n_blocks, n_gru, kernel_size, en_de_layers=5, inter_layers=4, in_channels=1,
+                 en_out_channels=16):
+        super(E2E, self).__init__()
+        self.mel = MelSpectrogram(N_MELS, SAMPLE_RATE, WINDOW_LENGTH, hop_length, None, MEL_FMIN, MEL_FMAX)  # noqa: F405
+        self.unet = DeepUnet(kernel_size, n_blocks, en_de_layers, inter_layers, in_channels, en_out_channels)
+        self.cnn = nn.Conv2d(en_out_channels, 3, (3, 3), padding=(1, 1))
+        if n_gru:
+            self.fc = nn.Sequential(
+                BiGRU(3 * N_MELS, 256, n_gru),   # noqa: F405
+                nn.Linear(512, N_CLASS),   # noqa: F405
+                nn.Dropout(0.25),
+                nn.Sigmoid()
+            )
+        else:
+            self.fc = nn.Sequential(
+                nn.Linear(3 * N_MELS, N_CLASS),  # noqa: F405
+                nn.Dropout(0.25),
+                nn.Sigmoid()
+            )
+
+    def forward(self, x):
+        mel = self.mel(x.reshape(-1, x.shape[-1])).transpose(-1, -2).unsqueeze(1)
+        x = self.cnn(self.unet(mel)).transpose(1, 2).flatten(-2)
+        # x = self.fc(x)
+        hidden_vec = 0
+        if len(self.fc) == 4:
+            for i in range(len(self.fc)):
+                x = self.fc[i](x)
+                if i == 0:
+                    hidden_vec = x
+        return hidden_vec, x
+
+
+class E2E0(nn.Module):
+    def __init__(self, n_blocks, n_gru, kernel_size, en_de_layers=5, inter_layers=4, in_channels=1,
+                 en_out_channels=16):
+        super(E2E0, self).__init__()
+        self.unet = DeepUnet0(kernel_size, n_blocks, en_de_layers, inter_layers, in_channels, en_out_channels)
+        self.cnn = nn.Conv2d(en_out_channels, 3, (3, 3), padding=(1, 1))
+        if n_gru:
+            self.fc = nn.Sequential(
+                BiGRU(3 * N_MELS, 256, n_gru),  # noqa: F405
+                nn.Linear(512, N_CLASS),  # noqa: F405
+                nn.Dropout(0.25),
+                nn.Sigmoid()
+            )
+        else:
+            self.fc = nn.Sequential(
+                nn.Linear(3 * N_MELS, N_CLASS),  # noqa: F405
+                nn.Dropout(0.25),
+                nn.Sigmoid()
+            )
+
+    def forward(self, mel):
+        mel = mel.transpose(-1, -2).unsqueeze(1)
+        x = self.cnn(self.unet(mel)).transpose(1, 2).flatten(-2)
+        x = self.fc(x)
+        return x

modules/F0Predictor/rmvpe/seq.py

@@ -0,0 +1,20 @@
+import torch.nn as nn
+
+
+class BiGRU(nn.Module):
+    def __init__(self, input_features, hidden_features, num_layers):
+        super(BiGRU, self).__init__()
+        self.gru = nn.GRU(input_features, hidden_features, num_layers=num_layers, batch_first=True, bidirectional=True)
+
+    def forward(self, x):
+        return self.gru(x)[0]
+
+
+class BiLSTM(nn.Module):
+    def __init__(self, input_features, hidden_features, num_layers):
+        super(BiLSTM, self).__init__()
+        self.lstm = nn.LSTM(input_features, hidden_features, num_layers=num_layers, batch_first=True, bidirectional=True)
+
+    def forward(self, x):
+        return self.lstm(x)[0]
+

modules/F0Predictor/rmvpe/spec.py

@@ -0,0 +1,67 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+from librosa.filters import mel
+
+
+class MelSpectrogram(torch.nn.Module):
+    def __init__(
+        self,
+        n_mel_channels,
+        sampling_rate,
+        win_length,
+        hop_length,
+        n_fft=None,
+        mel_fmin=0,
+        mel_fmax=None,
+        clamp = 1e-5
+    ):
+        super().__init__()
+        n_fft = win_length if n_fft is None else n_fft
+        self.hann_window = {}
+        mel_basis = mel(
+            sr=sampling_rate,
+            n_fft=n_fft, 
+            n_mels=n_mel_channels, 
+            fmin=mel_fmin, 
+            fmax=mel_fmax,
+            htk=True)
+        mel_basis = torch.from_numpy(mel_basis).float()
+        self.register_buffer("mel_basis", mel_basis)
+        self.n_fft = win_length if n_fft is None else n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.sampling_rate = sampling_rate
+        self.n_mel_channels = n_mel_channels
+        self.clamp = clamp
+
+    def forward(self, audio, keyshift=0, speed=1, center=True):
+        factor = 2 ** (keyshift / 12)       
+        n_fft_new = int(np.round(self.n_fft * factor))
+        win_length_new = int(np.round(self.win_length * factor))
+        hop_length_new = int(np.round(self.hop_length * speed))
+        
+        keyshift_key = str(keyshift)+'_'+str(audio.device)
+        if keyshift_key not in self.hann_window:
+            self.hann_window[keyshift_key] = torch.hann_window(win_length_new).to(audio.device)
+            
+        fft = torch.stft(
+            audio,
+            n_fft=n_fft_new,
+            hop_length=hop_length_new,
+            win_length=win_length_new,
+            window=self.hann_window[keyshift_key],
+            center=center,
+            return_complex=True)
+        magnitude = torch.sqrt(fft.real.pow(2) + fft.imag.pow(2))
+        
+        if keyshift != 0:
+            size = self.n_fft // 2 + 1
+            resize = magnitude.size(1)
+            if resize < size:
+                magnitude = F.pad(magnitude, (0, 0, 0, size-resize))
+            magnitude = magnitude[:, :size, :] * self.win_length / win_length_new
+            
+        mel_output = torch.matmul(self.mel_basis, magnitude)
+        log_mel_spec = torch.log(torch.clamp(mel_output, min=self.clamp))
+        return log_mel_spec

modules/F0Predictor/rmvpe/utils.py

@@ -0,0 +1,107 @@
+import sys
+from functools import reduce
+
+import librosa
+import numpy as np
+import torch
+from torch.nn.modules.module import _addindent
+
+from .constants import *  # noqa: F403
+
+
+def cycle(iterable):
+    while True:
+        for item in iterable:
+            yield item
+
+
+def summary(model, file=sys.stdout):
+    def repr(model):
+        # We treat the extra repr like the sub-module, one item per line
+        extra_lines = []
+        extra_repr = model.extra_repr()
+        # empty string will be split into list ['']
+        if extra_repr:
+            extra_lines = extra_repr.split('\n')
+        child_lines = []
+        total_params = 0
+        for key, module in model._modules.items():
+            mod_str, num_params = repr(module)
+            mod_str = _addindent(mod_str, 2)
+            child_lines.append('(' + key + '): ' + mod_str)
+            total_params += num_params
+        lines = extra_lines + child_lines
+
+        for name, p in model._parameters.items():
+            if hasattr(p, 'shape'):
+                total_params += reduce(lambda x, y: x * y, p.shape)
+
+        main_str = model._get_name() + '('
+        if lines:
+            # simple one-liner info, which most builtin Modules will use
+            if len(extra_lines) == 1 and not child_lines:
+                main_str += extra_lines[0]
+            else:
+                main_str += '\n  ' + '\n  '.join(lines) + '\n'
+
+        main_str += ')'
+        if file is sys.stdout:
+            main_str += ', \033[92m{:,}\033[0m params'.format(total_params)
+        else:
+            main_str += ', {:,} params'.format(total_params)
+        return main_str, total_params
+
+    string, count = repr(model)
+    if file is not None:
+        if isinstance(file, str):
+            file = open(file, 'w')
+        print(string, file=file)
+        file.flush()
+
+    return count
+
+    
+def to_local_average_cents(salience, center=None, thred=0.05):
+    """
+    find the weighted average cents near the argmax bin
+    """
+
+    if not hasattr(to_local_average_cents, 'cents_mapping'):
+        # the bin number-to-cents mapping
+        to_local_average_cents.cents_mapping = (
+                20 * torch.arange(N_CLASS) + CONST).to(salience.device)  # noqa: F405
+
+    if salience.ndim == 1:
+        if center is None:
+            center = int(torch.argmax(salience))
+        start = max(0, center - 4)
+        end = min(len(salience), center + 5)
+        salience = salience[start:end]
+        product_sum = torch.sum(
+            salience * to_local_average_cents.cents_mapping[start:end])
+        weight_sum = torch.sum(salience)
+        return product_sum / weight_sum if torch.max(salience) > thred else 0
+    if salience.ndim == 2:
+        return torch.Tensor([to_local_average_cents(salience[i, :], None, thred) for i in
+                         range(salience.shape[0])]).to(salience.device)
+
+    raise Exception("label should be either 1d or 2d ndarray")
+
+def to_viterbi_cents(salience, thred=0.05):
+    # Create viterbi transition matrix
+    if not hasattr(to_viterbi_cents, 'transition'):
+        xx, yy = torch.meshgrid(range(N_CLASS), range(N_CLASS))  # noqa: F405
+        transition = torch.maximum(30 - abs(xx - yy), 0)
+        transition = transition / transition.sum(axis=1, keepdims=True)
+        to_viterbi_cents.transition = transition
+
+    # Convert to probability
+    prob = salience.T
+    prob = prob / prob.sum(axis=0)    
+
+    # Perform viterbi decoding
+    path = librosa.sequence.viterbi(prob.detach().cpu().numpy(), to_viterbi_cents.transition).astype(np.int64)
+
+    return torch.Tensor([to_local_average_cents(salience[i, :], path[i], thred) for i in
+                     range(len(path))]).to(salience.device)
+                     

modules/attentions.py

@@ -1,18 +1,17 @@
-import copy
 import math
-import numpy as np
+
 import torch
 from torch import nn
 from torch.nn import functional as F
 
 import modules.commons as commons
-import modules.modules as modules
+from modules.DSConv import weight_norm_modules
 from modules.modules import LayerNorm
 
 
 class FFT(nn.Module):
   def __init__(self, hidden_channels, filter_channels, n_heads, n_layers=1, kernel_size=1, p_dropout=0.,
-               proximal_bias=False, proximal_init=True, **kwargs):
+               proximal_bias=False, proximal_init=True, isflow = False, **kwargs):
     super().__init__()
     self.hidden_channels = hidden_channels
     self.filter_channels = filter_channels
@@ -22,7 +21,11 @@ class FFT(nn.Module):
     self.p_dropout = p_dropout
     self.proximal_bias = proximal_bias
     self.proximal_init = proximal_init
-
+    if isflow:
+      cond_layer = torch.nn.Conv1d(kwargs["gin_channels"], 2*hidden_channels*n_layers, 1)
+      self.cond_pre = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, 1)
+      self.cond_layer = weight_norm_modules(cond_layer, name='weight')
+      self.gin_channels = kwargs["gin_channels"]
     self.drop = nn.Dropout(p_dropout)
     self.self_attn_layers = nn.ModuleList()
     self.norm_layers_0 = nn.ModuleList()
@@ -37,14 +40,25 @@ class FFT(nn.Module):
         FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout, causal=True))
       self.norm_layers_1.append(LayerNorm(hidden_channels))
 
-  def forward(self, x, x_mask):
+  def forward(self, x, x_mask, g = None):
     """
     x: decoder input
     h: encoder output
     """
+    if g is not None:
+      g = self.cond_layer(g)
+
     self_attn_mask = commons.subsequent_mask(x_mask.size(2)).to(device=x.device, dtype=x.dtype)
     x = x * x_mask
     for i in range(self.n_layers):
+      if g is not None:
+        x = self.cond_pre(x)
+        cond_offset = i * 2 * self.hidden_channels
+        g_l = g[:,cond_offset:cond_offset+2*self.hidden_channels,:]
+        x = commons.fused_add_tanh_sigmoid_multiply(
+          x,
+          g_l,
+          torch.IntTensor([self.hidden_channels]))
       y = self.self_attn_layers[i](x, x, self_attn_mask)
       y = self.drop(y)
       x = self.norm_layers_0[i](x + y)
@@ -243,7 +257,7 @@ class MultiHeadAttention(nn.Module):
     return ret
 
   def _get_relative_embeddings(self, relative_embeddings, length):
-    max_relative_position = 2 * self.window_size + 1
+    2 * self.window_size + 1
     # Pad first before slice to avoid using cond ops.
     pad_length = max(length - (self.window_size + 1), 0)
     slice_start_position = max((self.window_size + 1) - length, 0)

modules/commons.py

@@ -1,9 +1,9 @@
 import math
-import numpy as np
+
 import torch
-from torch import nn
 from torch.nn import functional as F
 
+
 def slice_pitch_segments(x, ids_str, segment_size=4):
   ret = torch.zeros_like(x[:, :segment_size])
   for i in range(x.size(0)):
@@ -24,10 +24,12 @@ def rand_slice_segments_with_pitch(x, pitch, x_lengths=None, segment_size=4):
 
 def init_weights(m, mean=0.0, std=0.01):
   classname = m.__class__.__name__
-  if classname.find("Conv") != -1:
+  if "Depthwise_Separable" in classname:
+    m.depth_conv.weight.data.normal_(mean, std)
+    m.point_conv.weight.data.normal_(mean, std) 
+  elif classname.find("Conv") != -1:
     m.weight.data.normal_(mean, std)
 
-
 def get_padding(kernel_size, dilation=1):
   return int((kernel_size*dilation - dilation)/2)
 
@@ -134,12 +136,6 @@ def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
   return acts
 
 
-def convert_pad_shape(pad_shape):
-  l = pad_shape[::-1]
-  pad_shape = [item for sublist in l for item in sublist]
-  return pad_shape
-
-
 def shift_1d(x):
   x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [1, 0]]))[:, :, :-1]
   return x
@@ -157,7 +153,6 @@ def generate_path(duration, mask):
   duration: [b, 1, t_x]
   mask: [b, 1, t_y, t_x]
   """
-  device = duration.device
   
   b, _, t_y, t_x = mask.shape
   cum_duration = torch.cumsum(duration, -1)

modules/enhancer.py

@@ -1,10 +1,12 @@
 import numpy as np
 import torch
 import torch.nn.functional as F
-from vdecoder.nsf_hifigan.nvSTFT import STFT
-from vdecoder.nsf_hifigan.models import load_model
 from torchaudio.transforms import Resample
 
+from vdecoder.nsf_hifigan.models import load_model
+from vdecoder.nsf_hifigan.nvSTFT import STFT
+
+
 class Enhancer:
     def __init__(self, enhancer_type, enhancer_ckpt, device=None):
         if device is None:

modules/losses.py

@@ -1,7 +1,4 @@
-import torch 
-from torch.nn import functional as F
-
-import modules.commons as commons
+import torch
 
 
 def feature_loss(fmap_r, fmap_g):

modules/mel_processing.py

@@ -1,16 +1,5 @@
-import math
-import os
-import random
 import torch
-from torch import nn
-import torch.nn.functional as F
 import torch.utils.data
-import numpy as np
-import librosa
-import librosa.util as librosa_util
-from librosa.util import normalize, pad_center, tiny
-from scipy.signal import get_window
-from scipy.io.wavfile import read
 from librosa.filters import mel as librosa_mel_fn
 
 MAX_WAV_VALUE = 32768.0
@@ -62,9 +51,14 @@ def spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center=False)
 
     y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect')
     y = y.squeeze(1)
+    
+    y_dtype = y.dtype
+    if y.dtype == torch.bfloat16:
+        y = y.to(torch.float32)
 
     spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[wnsize_dtype_device],
-                      center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=False)
+                      center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=True)
+    spec = torch.view_as_real(spec).to(y_dtype)
 
     spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
     return spec
@@ -83,30 +77,7 @@ def spec_to_mel_torch(spec, n_fft, num_mels, sampling_rate, fmin, fmax):
 
 
 def mel_spectrogram_torch(y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False):
-    if torch.min(y) < -1.:
-        print('min value is ', torch.min(y))
-    if torch.max(y) > 1.:
-        print('max value is ', torch.max(y))
-
-    global mel_basis, hann_window
-    dtype_device = str(y.dtype) + '_' + str(y.device)
-    fmax_dtype_device = str(fmax) + '_' + dtype_device
-    wnsize_dtype_device = str(win_size) + '_' + dtype_device
-    if fmax_dtype_device not in mel_basis:
-        mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax)
-        mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(dtype=y.dtype, device=y.device)
-    if wnsize_dtype_device not in hann_window:
-        hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(dtype=y.dtype, device=y.device)
-
-    y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect')
-    y = y.squeeze(1)
-
-    spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[wnsize_dtype_device],
-                      center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=False)
-
-    spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
-
-    spec = torch.matmul(mel_basis[fmax_dtype_device], spec)
-    spec = spectral_normalize_torch(spec)
-
+    spec = spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center)
+    spec = spec_to_mel_torch(spec, n_fft, num_mels, sampling_rate, fmin, fmax)
+    
     return spec

modules/modules.py

@@ -1,20 +1,24 @@
-import copy
-import math
-import numpy as np
-import scipy
 import torch
 from torch import nn
 from torch.nn import functional as F
 
-from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
-from torch.nn.utils import weight_norm, remove_weight_norm
-
+import modules.attentions as attentions
 import modules.commons as commons
-from modules.commons import init_weights, get_padding
-
+from modules.commons import get_padding, init_weights
+from modules.DSConv import (
+    Depthwise_Separable_Conv1D,
+    remove_weight_norm_modules,
+    weight_norm_modules,
+)
 
 LRELU_SLOPE = 0.1
 
+Conv1dModel = nn.Conv1d
+
+def set_Conv1dModel(use_depthwise_conv):
+    global Conv1dModel
+    Conv1dModel = Depthwise_Separable_Conv1D if use_depthwise_conv else nn.Conv1d
+
 
 class LayerNorm(nn.Module):
   def __init__(self, channels, eps=1e-5):
@@ -44,13 +48,13 @@ class ConvReluNorm(nn.Module):
 
     self.conv_layers = nn.ModuleList()
     self.norm_layers = nn.ModuleList()
-    self.conv_layers.append(nn.Conv1d(in_channels, hidden_channels, kernel_size, padding=kernel_size//2))
+    self.conv_layers.append(Conv1dModel(in_channels, hidden_channels, kernel_size, padding=kernel_size//2))
     self.norm_layers.append(LayerNorm(hidden_channels))
     self.relu_drop = nn.Sequential(
         nn.ReLU(),
         nn.Dropout(p_dropout))
     for _ in range(n_layers-1):
-      self.conv_layers.append(nn.Conv1d(hidden_channels, hidden_channels, kernel_size, padding=kernel_size//2))
+      self.conv_layers.append(Conv1dModel(hidden_channels, hidden_channels, kernel_size, padding=kernel_size//2))
       self.norm_layers.append(LayerNorm(hidden_channels))
     self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
     self.proj.weight.data.zero_()
@@ -66,47 +70,6 @@ class ConvReluNorm(nn.Module):
     return x * x_mask
 
 
-class DDSConv(nn.Module):
-  """
-  Dialted and Depth-Separable Convolution
-  """
-  def __init__(self, channels, kernel_size, n_layers, p_dropout=0.):
-    super().__init__()
-    self.channels = channels
-    self.kernel_size = kernel_size
-    self.n_layers = n_layers
-    self.p_dropout = p_dropout
-
-    self.drop = nn.Dropout(p_dropout)
-    self.convs_sep = nn.ModuleList()
-    self.convs_1x1 = nn.ModuleList()
-    self.norms_1 = nn.ModuleList()
-    self.norms_2 = nn.ModuleList()
-    for i in range(n_layers):
-      dilation = kernel_size ** i
-      padding = (kernel_size * dilation - dilation) // 2
-      self.convs_sep.append(nn.Conv1d(channels, channels, kernel_size, 
-          groups=channels, dilation=dilation, padding=padding
-      ))
-      self.convs_1x1.append(nn.Conv1d(channels, channels, 1))
-      self.norms_1.append(LayerNorm(channels))
-      self.norms_2.append(LayerNorm(channels))
-
-  def forward(self, x, x_mask, g=None):
-    if g is not None:
-      x = x + g
-    for i in range(self.n_layers):
-      y = self.convs_sep[i](x * x_mask)
-      y = self.norms_1[i](y)
-      y = F.gelu(y)
-      y = self.convs_1x1[i](y)
-      y = self.norms_2[i](y)
-      y = F.gelu(y)
-      y = self.drop(y)
-      x = x + y
-    return x * x_mask
-
-
 class WN(torch.nn.Module):
   def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0):
     super(WN, self).__init__()
@@ -124,14 +87,14 @@ class WN(torch.nn.Module):
 
     if gin_channels != 0:
       cond_layer = torch.nn.Conv1d(gin_channels, 2*hidden_channels*n_layers, 1)
-      self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')
+      self.cond_layer = weight_norm_modules(cond_layer, name='weight')
 
     for i in range(n_layers):
       dilation = dilation_rate ** i
       padding = int((kernel_size * dilation - dilation) / 2)
-      in_layer = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, kernel_size,
+      in_layer = Conv1dModel(hidden_channels, 2*hidden_channels, kernel_size,
                                  dilation=dilation, padding=padding)
-      in_layer = torch.nn.utils.weight_norm(in_layer, name='weight')
+      in_layer = weight_norm_modules(in_layer, name='weight')
       self.in_layers.append(in_layer)
 
       # last one is not necessary
@@ -141,7 +104,7 @@ class WN(torch.nn.Module):
         res_skip_channels = hidden_channels
 
       res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
-      res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight')
+      res_skip_layer = weight_norm_modules(res_skip_layer, name='weight')
       self.res_skip_layers.append(res_skip_layer)
 
   def forward(self, x, x_mask, g=None, **kwargs):
@@ -176,32 +139,32 @@ class WN(torch.nn.Module):
 
   def remove_weight_norm(self):
     if self.gin_channels != 0:
-      torch.nn.utils.remove_weight_norm(self.cond_layer)
+      remove_weight_norm_modules(self.cond_layer)
     for l in self.in_layers:
-      torch.nn.utils.remove_weight_norm(l)
+      remove_weight_norm_modules(l)
     for l in self.res_skip_layers:
-     torch.nn.utils.remove_weight_norm(l)
+      remove_weight_norm_modules(l)
 
 
 class ResBlock1(torch.nn.Module):
     def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
         super(ResBlock1, self).__init__()
         self.convs1 = nn.ModuleList([
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
+            weight_norm_modules(Conv1dModel(channels, channels, kernel_size, 1, dilation=dilation[0],
                                padding=get_padding(kernel_size, dilation[0]))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
+            weight_norm_modules(Conv1dModel(channels, channels, kernel_size, 1, dilation=dilation[1],
                                padding=get_padding(kernel_size, dilation[1]))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
+            weight_norm_modules(Conv1dModel(channels, channels, kernel_size, 1, dilation=dilation[2],
                                padding=get_padding(kernel_size, dilation[2])))
         ])
         self.convs1.apply(init_weights)
 
         self.convs2 = nn.ModuleList([
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+            weight_norm_modules(Conv1dModel(channels, channels, kernel_size, 1, dilation=1,
                                padding=get_padding(kernel_size, 1))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+            weight_norm_modules(Conv1dModel(channels, channels, kernel_size, 1, dilation=1,
                                padding=get_padding(kernel_size, 1))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+            weight_norm_modules(Conv1dModel(channels, channels, kernel_size, 1, dilation=1,
                                padding=get_padding(kernel_size, 1)))
         ])
         self.convs2.apply(init_weights)
@@ -223,18 +186,18 @@ class ResBlock1(torch.nn.Module):
 
     def remove_weight_norm(self):
         for l in self.convs1:
-            remove_weight_norm(l)
+            remove_weight_norm_modules(l)
         for l in self.convs2:
-            remove_weight_norm(l)
+            remove_weight_norm_modules(l)
 
 
 class ResBlock2(torch.nn.Module):
     def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
         super(ResBlock2, self).__init__()
         self.convs = nn.ModuleList([
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
+            weight_norm_modules(Conv1dModel(channels, channels, kernel_size, 1, dilation=dilation[0],
                                padding=get_padding(kernel_size, dilation[0]))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
+            weight_norm_modules(Conv1dModel(channels, channels, kernel_size, 1, dilation=dilation[1],
                                padding=get_padding(kernel_size, dilation[1])))
         ])
         self.convs.apply(init_weights)
@@ -252,7 +215,7 @@ class ResBlock2(torch.nn.Module):
 
     def remove_weight_norm(self):
         for l in self.convs:
-            remove_weight_norm(l)
+            remove_weight_norm_modules(l)
 
 
 class Log(nn.Module):
@@ -303,7 +266,9 @@ class ResidualCouplingLayer(nn.Module):
       n_layers,
       p_dropout=0,
       gin_channels=0,
-      mean_only=False):
+      mean_only=False,
+      wn_sharing_parameter=None
+      ):
     assert channels % 2 == 0, "channels should be divisible by 2"
     super().__init__()
     self.channels = channels
@@ -315,7 +280,56 @@ class ResidualCouplingLayer(nn.Module):
     self.mean_only = mean_only
 
     self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
-    self.enc = WN(hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout=p_dropout, gin_channels=gin_channels)
+    self.enc = WN(hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout=p_dropout, gin_channels=gin_channels) if wn_sharing_parameter is None else wn_sharing_parameter
+    self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
+    self.post.weight.data.zero_()
+    self.post.bias.data.zero_()
+
+  def forward(self, x, x_mask, g=None, reverse=False):
+    x0, x1 = torch.split(x, [self.half_channels]*2, 1)
+    h = self.pre(x0) * x_mask
+    h = self.enc(h, x_mask, g=g)
+    stats = self.post(h) * x_mask
+    if not self.mean_only:
+      m, logs = torch.split(stats, [self.half_channels]*2, 1)
+    else:
+      m = stats
+      logs = torch.zeros_like(m)
+
+    if not reverse:
+      x1 = m + x1 * torch.exp(logs) * x_mask
+      x = torch.cat([x0, x1], 1)
+      logdet = torch.sum(logs, [1,2])
+      return x, logdet
+    else:
+      x1 = (x1 - m) * torch.exp(-logs) * x_mask
+      x = torch.cat([x0, x1], 1)
+      return x
+
+class TransformerCouplingLayer(nn.Module):
+  def __init__(self,
+      channels,
+      hidden_channels,
+      kernel_size,
+      n_layers,
+      n_heads,
+      p_dropout=0,
+      filter_channels=0,
+      mean_only=False,
+      wn_sharing_parameter=None,
+      gin_channels = 0
+      ):
+    assert channels % 2 == 0, "channels should be divisible by 2"
+    super().__init__()
+    self.channels = channels
+    self.hidden_channels = hidden_channels
+    self.kernel_size = kernel_size
+    self.n_layers = n_layers
+    self.half_channels = channels // 2
+    self.mean_only = mean_only
+
+    self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
+    self.enc = attentions.FFT(hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, isflow = True, gin_channels = gin_channels) if wn_sharing_parameter is None else wn_sharing_parameter
     self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
     self.post.weight.data.zero_()
     self.post.bias.data.zero_()