init

G_228800.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2635ff85250884517c9989ec5c46a2176c103ba22ade4a682e42759815b4c465
+size 629383387

app.py

@@ -0,0 +1,378 @@
+# -*- coding: utf-8 -*-
+import traceback
+import torch
+from scipy.io import wavfile
+import edge_tts
+import subprocess
+import gradio as gr
+import gradio.processing_utils as gr_pu
+import io
+import os
+import logging
+import time
+from pathlib import Path
+import re
+import json
+import argparse
+
+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")
+
+
+logging.getLogger('numba').setLevel(logging.WARNING)
+logging.getLogger('markdown_it').setLevel(logging.WARNING)
+logging.getLogger('urllib3').setLevel(logging.WARNING)
+logging.getLogger('matplotlib').setLevel(logging.WARNING)
+logging.getLogger('multipart').setLevel(logging.WARNING)
+
+model = None
+spk = None
+debug = False
+
+
+class HParams():
+    def __init__(self, **kwargs):
+        for k, v in kwargs.items():
+            if type(v) == dict:
+                v = HParams(**v)
+            self[k] = v
+
+    def keys(self):
+        return self.__dict__.keys()
+
+    def items(self):
+        return self.__dict__.items()
+
+    def values(self):
+        return self.__dict__.values()
+
+    def __len__(self):
+        return len(self.__dict__)
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+    def __setitem__(self, key, value):
+        return setattr(self, key, value)
+
+    def __contains__(self, key):
+        return key in self.__dict__
+
+    def __repr__(self):
+        return self.__dict__.__repr__()
+
+
+def get_hparams_from_file(config_path):
+    with open(config_path, "r", encoding="utf-8") as f:
+        data = f.read()
+    config = json.loads(data)
+
+    hparams = HParams(**config)
+    return hparams
+
+
+def vc_fn(sid, input_audio, vc_transform, auto_f0, cluster_ratio, slice_db, noise_scale, pad_seconds, cl_num, lg_num, lgr_num, f0_predictor, enhancer_adaptive_key, cr_threshold):
+    try:
+        if input_audio is None:
+            raise gr.Error("你需要上传音频")
+        if model is None:
+            raise gr.Error("你需要指定模型")
+        sampling_rate, audio = input_audio
+        # print(audio.shape,sampling_rate)
+        audio = (audio / np.iinfo(audio.dtype).max).astype(np.float32)
+        if len(audio.shape) > 1:
+            audio = librosa.to_mono(audio.transpose(1, 0))
+        temp_path = "temp.wav"
+        soundfile.write(temp_path, audio, sampling_rate, format="wav")
+        _audio = model.slice_inference(temp_path, sid, vc_transform, slice_db, cluster_ratio, auto_f0, noise_scale,
+                                       pad_seconds, cl_num, lg_num, lgr_num, f0_predictor, enhancer_adaptive_key, cr_threshold)
+        model.clear_empty()
+        os.remove(temp_path)
+        # 构建保存文件的路径，并保存到results文件夹内
+        try:
+            timestamp = str(int(time.time()))
+            filename = sid + "_" + timestamp + ".wav"
+            # output_file = os.path.join("./results", filename)
+            # soundfile.write(output_file, _audio, model.target_sample, format="wav")
+            soundfile.write('/tmp/'+filename, _audio,
+                            model.target_sample, format="wav")
+            # return f"推理成功，音频文件保存为results/{filename}", (model.target_sample, _audio)
+            return f"推理成功，音频文件保存为{filename}", (model.target_sample, _audio)
+        except Exception as e:
+            if debug:
+                traceback.print_exc()
+            return f"文件保存失败，请手动保存", (model.target_sample, _audio)
+    except Exception as e:
+        if debug:
+            traceback.print_exc()
+        raise gr.Error(e)
+
+
+def tts_func(_text, _rate, _voice):
+    # 使用edge-tts把文字转成音频
+    # voice = "zh-CN-XiaoyiNeural"#女性，较高音
+    # voice = "zh-CN-YunxiNeural"#男性
+    voice = "zh-CN-YunxiNeural"  # 男性
+    if (_voice == "女"):
+        voice = "zh-CN-XiaoyiNeural"
+    output_file = "/tmp/"+_text[0:10]+".wav"
+    # communicate = edge_tts.Communicate(_text, voice)
+    # await communicate.save(output_file)
+    if _rate >= 0:
+        ratestr = "+{:.0%}".format(_rate)
+    elif _rate < 0:
+        ratestr = "{:.0%}".format(_rate)  # 减号自带
+
+    p = subprocess.Popen("edge-tts " +
+                         " --text "+_text +
+                         " --write-media "+output_file +
+                         " --voice "+voice +
+                         " --rate="+ratestr, shell=True,
+                         stdout=subprocess.PIPE,
+                         stdin=subprocess.PIPE)
+    p.wait()
+    return output_file
+
+
+def text_clear(text):
+    return re.sub(r"[\n\,\(\) ]", "", text)
+
+
+def vc_fn2(sid, input_audio, vc_transform, auto_f0, cluster_ratio, slice_db, noise_scale, pad_seconds, cl_num, lg_num, lgr_num, text2tts, tts_rate, tts_voice, f0_predictor, enhancer_adaptive_key, cr_threshold):
+    # 使用edge-tts把文字转成音频
+    text2tts = text_clear(text2tts)
+    output_file = tts_func(text2tts, tts_rate, tts_voice)
+
+    # 调整采样率
+    sr2 = 44100
+    wav, sr = librosa.load(output_file)
+    wav2 = librosa.resample(wav, orig_sr=sr, target_sr=sr2)
+    save_path2 = text2tts[0:10]+"_44k"+".wav"
+    wavfile.write(save_path2, sr2,
+                  (wav2 * np.iinfo(np.int16).max).astype(np.int16)
+                  )
+
+    # 读取音频
+    sample_rate, data = gr_pu.audio_from_file(save_path2)
+    vc_input = (sample_rate, data)
+
+    a, b = vc_fn(sid, vc_input, vc_transform, auto_f0, cluster_ratio, slice_db, noise_scale,
+                 pad_seconds, cl_num, lg_num, lgr_num, f0_predictor, enhancer_adaptive_key, cr_threshold)
+    os.remove(output_file)
+    os.remove(save_path2)
+    return a, b
+
+
+models_info = [
+    {
+        "description": """
+                       这个模型包含公主连结的161名角色。\n\n
+                       Space采用CPU推理，速度极慢，建议下载模型本地GPU推理。\n\n
+                       """,
+        "model_path": "./G_228800.pth",
+        "config_path": "./config.json",
+    }
+]
+
+model_inferall = []
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--share", action="store_true",
+                        default=False, help="share gradio app")
+    # 一定要设置的部分
+    parser.add_argument('-cl', '--clip', type=float,
+                        default=0, help='音频强制切片，默认0为自动切片，单位为秒/s')
+    parser.add_argument('-n', '--clean_names', type=str, nargs='+',
+                        default=["君の知らない物語-src.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=['nen'], 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="pm",
+                        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('-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('-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. 降低该值可减少跑调概率，但会增加哑音')
+    args = parser.parse_args()
+    categories = ["Princess Connect! Re:Dive"]
+    others = {
+        "None": "https://huggingface.co/spaces/FrankZxShen/vits-fast-finetuning-pcr",
+    }
+    for info in models_info:
+        config_path = info['config_path']
+        model_path = info['model_path']
+        description = info['description']
+        clean_names = args.clean_names
+        trans = args.trans
+        spk_list = list(get_hparams_from_file(config_path).spk.keys())
+        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
+
+        model = Svc(model_path, config_path, args.device, args.cluster_model_path, enhance,
+                    diffusion_model_path, diffusion_config_path, shallow_diffusion, only_diffusion)
+
+        model_inferall.append((description, spk_list, model))
+
+    app = gr.Blocks()
+    with app:
+        gr.Markdown(
+            "# <center> so-vits-svc-models-pcr\n"
+            "# <center> 注意！！！！！Space采用CPU推理，速度极慢，建议下载模型使用本地GPU推理。\n"
+            "## <center> Please do not generate content that could infringe upon the rights or cause harm to individuals or organizations.\n"
+            "## <center> 请不要生成会对个人以及组织造成侵害的内容\n\n"
+            "[![image](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1pn1xnFfdLK63gVXDwV4zCXfVeo8c-I-0?usp=sharing)\n\n"
+            "[![Duplicate this Space](https://huggingface.co/datasets/huggingface/badges/raw/main/duplicate-this-space-sm-dark.svg)](https://huggingface.co/spaces/FrankZxShen/vits-fast-finetuning-pcr?duplicate=true)\n\n"
+            "[![Finetune your own model](https://badgen.net/badge/icon/github?icon=github&label=Finetune%20your%20own%20model)](https://github.com/Plachtaa/VITS-fast-fine-tuning)"
+        )
+        gr.Markdown("# Princess Connect! Re:Dive\n\n"
+                    )
+        with gr.Tabs():
+            for category in categories:
+                with gr.TabItem(category):
+                    for i, (description, speakers, model) in enumerate(
+                            model_inferall):
+                        gr.Markdown(description)
+                        with gr.Row():
+                            with gr.Column():
+                                # textbox = gr.TextArea(label="Text",
+                                #                         placeholder="Type your sentence here ",
+                                #                         value="新たなキャラを解放できるようになったようですね。", elem_id=f"tts-input")
+
+                                gr.Markdown(value="""
+                                    <font size=2> 推理设置</font>
+                                    """)
+                                sid = gr.Dropdown(
+                                    choices=speakers, value=speakers[0], label='角色选择')
+                                auto_f0 = gr.Checkbox(
+                                    label="自动f0预测，配合聚类模型f0预测效果更好,会导致变调功能失效（仅限转换语音，歌声勾选此项会究极跑调）", value=False)
+                                f0_predictor = gr.Dropdown(label="选择F0预测器,可选择crepe,pm,dio,harvest,默认为pm(注意：crepe为原F0使用均值滤波器)", choices=[
+                                                           "pm", "dio", "harvest", "crepe"], value="pm")
+                                vc_transform = gr.Number(
+                                    label="变调（整数，可以正负，半音数量，升高八度就是12）", value=0)
+                                cluster_ratio = gr.Number(
+                                    label="聚类模型混合比例，0-1之间，0即不启用聚类。使用聚类模型能提升音色相似度，但会导致咬字下降（如果使用建议0.5左右）", value=0)
+                                slice_db = gr.Number(label="切片阈值", value=-40)
+                                noise_scale = gr.Number(
+                                    label="noise_scale 建议不要动，会影响音质，玄学参数", value=0.4)
+                            with gr.Column():
+                                pad_seconds = gr.Number(
+                                    label="推理音频pad秒数，由于未知原因开头结尾会有异响，pad一小段静音段后就不会出现", value=0.5)
+                                cl_num = gr.Number(
+                                    label="音频自动切片，0为不切片，单位为秒(s)", value=0)
+                                lg_num = gr.Number(
+                                    label="两端音频切片的交叉淡入长度，如果自动切片后出现人声不连贯可调整该数值，如果连贯建议采用默认值0，注意，该设置会影响推理速度，单位为秒/s", value=0)
+                                lgr_num = gr.Number(
+                                    label="自动音频切片后，需要舍弃每段切片的头尾。该参数设置交叉长度保留的比例，范围0-1,左开右闭", value=0.75)
+                                enhancer_adaptive_key = gr.Number(
+                                    label="使增强器适应更高的音域(单位为半音数)|默认为0", value=0)
+                                cr_threshold = gr.Number(
+                                    label="F0过滤阈值，只有启动crepe时有效. 数值范围从0-1. 降低该值可减少跑调概率，但会增加哑音", value=0.05)
+                        with gr.Tabs():
+                            with gr.TabItem("音频转音频"):
+                                vc_input3 = gr.Audio(label="选择音频")
+                                vc_submit = gr.Button(
+                                    "音频转换", variant="primary")
+                            with gr.TabItem("文字转音频"):
+                                text2tts = gr.Textbox(
+                                    label="在此输入要转译的文字。注意，使用该功能建议打开F0预测，不然会很怪")
+                                tts_rate = gr.Number(label="tts语速", value=0)
+                                tts_voice = gr.Radio(label="性别", choices=[
+                                                     "男", "女"], value="男")
+                                vc_submit2 = gr.Button(
+                                    "文字转换", variant="primary")
+                        with gr.Row():
+                            with gr.Column():
+                                vc_output1 = gr.Textbox(label="Output Message")
+                            with gr.Column():
+                                vc_output2 = gr.Audio(
+                                    label="Output Audio", interactive=False)
+
+                vc_submit.click(vc_fn, [sid, vc_input3, vc_transform, auto_f0, cluster_ratio, slice_db, noise_scale, pad_seconds,
+                                        cl_num, lg_num, lgr_num, f0_predictor, enhancer_adaptive_key, cr_threshold], [vc_output1, vc_output2])
+                vc_submit2.click(vc_fn2, [sid, vc_input3, vc_transform, auto_f0, cluster_ratio, slice_db, noise_scale, pad_seconds, cl_num,
+                                          lg_num, lgr_num, text2tts, tts_rate, tts_voice, f0_predictor, enhancer_adaptive_key, cr_threshold], [vc_output1, vc_output2])
+                # gr.Examples(
+                #     examples=example,
+                #     inputs=[textbox, char_dropdown, language_dropdown,
+                #             duration_slider, symbol_input],
+                #     outputs=[text_output, audio_output],
+                #     fn=tts_fn
+                # )
+            for category, link in others.items():
+                with gr.TabItem(category):
+                    gr.Markdown(
+                        f'''
+                        <center>
+                          <h2>Click to Go</h2>
+                          <a href="{link}">
+                            <img src="https://huggingface.co/datasets/huggingface/badges/raw/main/open-in-hf-spaces-xl-dark.svg"
+                          </a>
+                        </center>
+                        '''
+                    )
+
+    app.queue(concurrency_count=3).launch(show_api=False, share=args.share)

cluster/__init__.py

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

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

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

config.json

@@ -0,0 +1,256 @@
+{
+  "train": {
+    "log_interval": 200,
+    "eval_interval": 800,
+    "seed": 1234,
+    "epochs": 10000,
+    "learning_rate": 0.0001,
+    "betas": [
+      0.8,
+      0.99
+    ],
+    "eps": 1e-09,
+    "batch_size": 6,
+    "fp16_run": false,
+    "lr_decay": 0.999875,
+    "segment_size": 10240,
+    "init_lr_ratio": 1,
+    "warmup_epochs": 0,
+    "c_mel": 45,
+    "c_kl": 1.0,
+    "use_sr": true,
+    "max_speclen": 512,
+    "port": "8976",
+    "keep_ckpts": 3,
+    "all_in_mem": false
+  },
+  "data": {
+    "training_files": "filelists/train.txt",
+    "validation_files": "filelists/val.txt",
+    "max_wav_value": 32768.0,
+    "sampling_rate": 44100,
+    "filter_length": 2048,
+    "hop_length": 512,
+    "win_length": 2048,
+    "n_mel_channels": 80,
+    "mel_fmin": 0.0,
+    "mel_fmax": 22050
+  },
+  "model": {
+    "inter_channels": 192,
+    "hidden_channels": 192,
+    "filter_channels": 768,
+    "n_heads": 2,
+    "n_layers": 6,
+    "kernel_size": 3,
+    "p_dropout": 0.1,
+    "resblock": "1",
+    "resblock_kernel_sizes": [
+      3,
+      7,
+      11
+    ],
+    "resblock_dilation_sizes": [
+      [
+        1,
+        3,
+        5
+      ],
+      [
+        1,
+        3,
+        5
+      ],
+      [
+        1,
+        3,
+        5
+      ]
+    ],
+    "upsample_rates": [
+      8,
+      8,
+      2,
+      2,
+      2
+    ],
+    "upsample_initial_channel": 512,
+    "upsample_kernel_sizes": [
+      16,
+      16,
+      4,
+      4,
+      4
+    ],
+    "n_layers_q": 3,
+    "use_spectral_norm": false,
+    "gin_channels": 768,
+    "ssl_dim": 768,
+    "n_speakers": 161,
+    "speech_encoder": "vec768l12",
+    "speaker_embedding": false
+  },
+  "spk": {
+    "\u83c8\u6bd4\u8389\u65af\u5854\uff08Overload\uff09": 0,
+    "\u94c3\u5948\uff08\u590f\u65e5\uff09": 1,
+    "\u674f\u5948": 2,
+    "\u96ea\u83f2": 3,
+    "\u53ef\u53ef\u841d": 4,
+    "\u83c8\u6bd4\u8389\u65af\u5854": 5,
+    "\u7eeb\u97f3\uff08\u5723\u8bde\u8282\uff09": 6,
+    "\u78a7\uff08\u63d2\u73ed\u751f\uff09": 7,
+    "\u5609\u591c": 8,
+    "\u955c\u534e\uff08\u4e07\u5723\u8282\uff09": 9,
+    "\u73e0\u5e0c": 10,
+    "\u671b\uff08\u590f\u65e5\uff09": 11,
+    "\u94c3\u8393": 12,
+    "\u771f\u7434\uff08\u590f\u65e5\uff09": 13,
+    "\u9759\u6d41\uff08\u60c5\u4eba\u8282\uff09": 14,
+    "\u54b2\u604b": 15,
+    "\u8389\u739b": 16,
+    "\u9999\u6f84\uff08\u590f\u65e5\uff09": 17,
+    "\u5343\u6b4c": 18,
+    "\u5fcd": 19,
+    "\u4f9d\u91cc": 20,
+    "\u4f69\u53ef\u8389\u59c6\uff08\u65b0\u5e74\uff09": 21,
+    "\u80e1\u6843": 22,
+    "\u667a\uff08\u9b54\u6cd5\u5c11\u5973\uff09": 23,
+    "\u4f18\u8863": 24,
+    "\u601c\uff08\u4e07\u5723\u8282\uff09": 25,
+    "\u681e": 26,
+    "\u9999\u7ec7\uff08\u590f\u65e5\uff09": 27,
+    "\u7948\u68a8\uff08\u65f6\u95f4\u65c5\u884c\uff09": 28,
+    "\u54b2\u604b\uff08\u590f\u65e5\uff09": 29,
+    "\u94c3\u5948": 30,
+    "\u771f\u6b65\uff08\u590f\u65e5\uff09": 31,
+    "\u4e9a\u91cc\u838e": 32,
+    "\u955c\u534e": 33,
+    "\u672a\u594f\u5e0c": 34,
+    "\u4f0a\u8389\u4e9a\uff08\u5723\u8bde\u8282\uff09": 35,
+    "\u7f8e\u91cc": 36,
+    "\u4f3c\u4f3c\u82b1": 37,
+    "\u514b\u8389\u4e1d\u63d0\u5a1c\uff08\u5723\u8bde\u8282\uff09": 38,
+    "\u7f8e\u51ac\uff08\u590f\u65e5\uff09": 39,
+    "\u8389\u739b\uff08\u7070\u59d1\u5a18\uff09": 40,
+    "\u94c3\u8393\uff08\u590f\u65e5\uff09": 41,
+    "\u53e4\u857e\u96c5": 42,
+    "\u7f8e\u7f8e\uff08\u4e07\u5723\u8282\uff09": 43,
+    "\u5343\u6b4c\uff08\u5723\u8bde\u8282\uff09": 44,
+    "\u78a7": 45,
+    "\u96ea": 46,
+    "\u60e0\u7406\u5b50\uff08\u60c5\u4eba\u8282\uff09": 47,
+    "\u4f0a\u7eea": 48,
+    "\u62c9\u59c6": 49,
+    "\u4e03\u4e03\u9999\uff08\u590f\u65e5\uff09": 50,
+    "\u94c3\u8393\uff08\u65b0\u5e74\uff09": 51,
+    "\u94c3\uff08\u6e38\u9a91\u5175\uff09": 52,
+    "\u73af\u5948": 53,
+    "\u60e0\u7406\u5b50": 54,
+    "\u53ef\u53ef\u841d\uff08\u793c\u670d\uff09": 55,
+    "\u4f18\u8863\uff08\u65b0\u5e74\uff09": 56,
+    "\u671b\uff08\u5723\u8bde\u8282\uff09": 57,
+    "\u53ef\u53ef\u841d\uff08\u516c\u4e3b\uff09": 58,
+    "\u7eaf\uff08\u590f\u65e5\uff09": 59,
+    "\u73af\u5948\uff08\u632f\u8896\uff09": 60,
+    "\u54b2\u604b\uff08\u5723\u8bde\u8282\uff09": 61,
+    "\u9732\u5a1c": 62,
+    "\u9999\u6f84\uff08\u9b54\u6cd5\u5c11\u5973\uff09": 63,
+    "\u59ec\u5854": 64,
+    "\u77db\u4f9d\u672a\uff08\u65b0\u5e74\uff09": 65,
+    "\u681e\uff08\u9b54\u6cd5\u5c11\u5973\uff09": 66,
+    "\u771f\u6b65\uff08\u7070\u59d1\u5a18\uff09": 67,
+    "\u601c\uff08\u65b0\u5e74\uff09": 68,
+    "\u5343\u6b4c\uff08\u590f\u65e5\uff09": 69,
+    "\u7948\u68a8": 70,
+    "\u7531\u52a0\u8389": 71,
+    "\u4f69\u53ef\u8389\u59c6\uff08\u590f\u65e5\uff09": 72,
+    "\u6b65\u7f8e\uff08\u4ed9\u5883\uff09": 73,
+    "\u53ef\u53ef\u841d\uff08\u590f\u65e5\uff09": 74,
+    "\u7eeb\u97f3": 75,
+    "\u7531\u52a0\u8389\uff08\u5723\u8bde\u8282\uff09": 76,
+    "\u771f\u9633": 77,
+    "\u96f7\u59c6": 78,
+    "\u831c\u91cc": 79,
+    "\u94c3": 80,
+    "\u51db\uff08\u5076\u50cf\u5927\u5e08\uff09": 81,
+    "\u51ef\u9732": 82,
+    "\u514b\u7f57\u4f9d\uff08\u5723\u5b66\u796d\uff09": 83,
+    "\u79cb\u4e43": 84,
+    "\u514b\u8389\u4e1d\u63d0\u5a1c": 85,
+    "\u4f69\u53ef\u8389\u59c6": 86,
+    "\u9732": 87,
+    "\u83ab\u59ae\u5361": 88,
+    "\u7483\u4e43": 89,
+    "\u9759\u6d41\uff08\u590f\u65e5\uff09": 90,
+    "\u7eba\u5e0c\uff08\u4e07\u5723\u8282\uff09": 91,
+    "\u9999\u6f84": 92,
+    "\u7f8e\u54b2": 93,
+    "\u7f8e\u91cc\uff08\u590f\u65e5\uff09": 94,
+    "\u667a": 95,
+    "\u521d\u97f3": 96,
+    "\u7483\u4e43\uff08\u4ed9\u5883\uff09": 97,
+    "\u536f\u6708\uff08\u5076\u50cf\u5927\u5e08\uff09": 98,
+    "\u7a7a\u82b1": 99,
+    "\u73e0\u5e0c\uff08\u590f\u65e5\uff09": 100,
+    "\u7eaf": 101,
+    "\u7f8e\u7f8e": 102,
+    "\u5fcd\uff08\u4e07\u5723\u8282\uff09": 103,
+    "\u59ae\u4fac\uff08\u5927\u6c5f\u6237\uff09": 104,
+    "\u5343\u7231\u7460": 105,
+    "\u53ef\u53ef\u841d\uff08\u65b0\u5e74\uff09": 106,
+    "\u601c\uff08\u516c\u4e3b\uff09": 107,
+    "\u672a\u594f\u5e0c\u3001\u7f8e\u7f8e\u3001\u955c\u534e": 108,
+    "\u4f0a\u8389\u4e9a": 109,
+    "\u671b": 110,
+    "\u4f18\u8863\uff08\u516c\u4e3b\uff09": 111,
+    "\u51ef\u9732\uff08\u516c\u4e3b\uff09": 112,
+    "\u514b\u7f57\u4f9d": 113,
+    "\u79cb\u4e43\uff08\u5723\u8bde\u8282\uff09": 114,
+    "\u4f18\u8863\uff08\u793c\u670d\uff09": 115,
+    "\u674f\u5948\uff08\u590f\u65e5\uff09": 116,
+    "\u9759\u6d41": 117,
+    "\u672a\u594f\u5e0c\uff08\u4e07\u5723\u8282\uff09": 118,
+    "\u831c\u91cc\uff08\u5929\u4f7f\uff09": 119,
+    "\u5bab\u5b50": 120,
+    "\u65e5\u548c\u8389\uff08\u65b0\u5e74\uff09": 121,
+    "\u4f18\u59ae": 122,
+    "\u5b89": 123,
+    "\u6d41\u590f\uff08\u590f\u65e5\uff09": 124,
+    "\u7f8e\u51ac\uff08\u5de5\u4f5c\u670d\uff09": 125,
+    "\u7f8e\u51ac": 126,
+    "\u521d\u97f3\uff08\u590f\u65e5\uff09": 127,
+    "\u77db\u4f9d\u672a": 128,
+    "\u51ef\u9732\uff08\u590f\u65e5\uff09": 129,
+    "\u83ab\u59ae\u5361\uff08\u9b54\u6cd5\u5c11\u5973\uff09": 130,
+    "\u9999\u7ec7": 131,
+    "\u5bab\u5b50\uff08\u4e07\u5723\u8282\uff09": 132,
+    "\u7f8e\u54b2\uff08\u4e07\u5723\u8282\uff09": 133,
+    "\u59ae\u4fac": 134,
+    "\u601c": 135,
+    "\u771f\u9633\uff08\u6e38\u9a91\u5175\uff09": 136,
+    "\u78a7\uff08\u5de5\u4f5c\u670d\uff09": 137,
+    "\u65e5\u548c\u8389": 138,
+    "\u4e03\u4e03\u9999": 139,
+    "\u771f\u7434": 140,
+    "\u6b65\u7f8e": 141,
+    "\u6df1\u6708": 142,
+    "\u65e5\u548c\u8389\uff08\u516c\u4e3b\uff09": 143,
+    "\u4f0a\u7eea\uff08\u590f\u65e5\uff09": 144,
+    "\u8309\u8389": 145,
+    "\u4f69\u53ef\u8389\u59c6\uff08\u516c\u4e3b\uff09": 146,
+    "\u7a7a\u82b1\uff08\u5927\u6c5f\u6237\uff09": 147,
+    "\u60e0\u7406\u5b50\uff08\u590f\u65e5\uff09": 148,
+    "\u51ef\u9732\uff08\u65b0\u5e74\uff09": 149,
+    "\u7231\u871c\u8389\u96c5": 150,
+    "\u4f9d\u91cc\uff08\u5929\u4f7f\uff09": 151,
+    "\u8309\u8389\uff08\u4e07\u5723\u8282\uff09": 152,
+    "\u771f\u7434\uff08\u7070\u59d1\u5a18\uff09": 153,
+    "\u7eba\u5e0c": 154,
+    "\u771f\u6b65": 155,
+    "\u6d41\u590f": 156,
+    "\u4f3c\u4f3c\u82b1\uff08\u65b0\u5e74\uff09": 157,
+    "\u672a\u592e\uff08\u5076\u50cf\u5927\u5e08\uff09": 158,
+    "\u80e1\u6843\uff08\u5723\u8bde\u8282\uff09": 159,
+    "\u5343\u7231\u7460\uff08\u5723\u5b66\u796d\uff09": 160
+  }
+}

diffusion/data_loaders.py

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

diffusion/diffusion.py

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

diffusion/diffusion_onnx.py

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

diffusion/dpm_solver_pytorch.py

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

diffusion/how to export onnx.md

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

diffusion/infer_gt_mel.py

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

diffusion/logger/saver.py

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

diffusion/logger/utils.py

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

diffusion/onnx_export.py

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

diffusion/solver.py

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

diffusion/unit2mel.py

@@ -0,0 +1,100 @@
+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)
+            
+        # diffusion
+        self.decoder = GaussianDiffusion(WaveNet(out_dims, n_layers, n_chans, n_hidden), out_dims=out_dims)
+
+    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:
+                x = x + self.spk_embed(spk_id)
+        if self.aug_shift_embed is not None and aug_shift is not None:
+            x = x + self.aug_shift_embed(aug_shift / 5) 
+        x = self.decoder(x, gt_spec=gt_spec, infer=infer, infer_speedup=infer_speedup, method=method, k_step=k_step, use_tqdm=use_tqdm)
+    
+        return x
+

diffusion/vocoder.py

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

diffusion/wavenet.py

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

inference/chunks_temp.json

@@ -0,0 +1 @@
+{"info": "temp_dict"}

inference/infer_tool.py

@@ -0,0 +1,407 @@
+import hashlib
+import io
+import json
+import logging
+import os
+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 models import SynthesizerTrn
+
+from diffusion.unit2mel import load_model_vocoder
+import yaml
+
+logging.getLogger('matplotlib').setLevel(logging.WARNING)
+
+
+def read_temp(file_name):
+    if not os.path.exists(file_name):
+        with open(file_name, "w") as f:
+            f.write(json.dumps({"info": "temp_dict"}))
+        return {}
+    else:
+        try:
+            with open(file_name, "r") as f:
+                data = f.read()
+            data_dict = json.loads(data)
+            if os.path.getsize(file_name) > 50 * 1024 * 1024:
+                f_name = file_name.replace("\\", "/").split("/")[-1]
+                print(f"clean {f_name}")
+                for wav_hash in list(data_dict.keys()):
+                    if int(time.time()) - int(data_dict[wav_hash]["time"]) > 14 * 24 * 3600:
+                        del data_dict[wav_hash]
+        except Exception as e:
+            print(e)
+            print(f"{file_name} error,auto rebuild file")
+            data_dict = {"info": "temp_dict"}
+        return data_dict
+
+
+def write_temp(file_name, data):
+    with open(file_name, "w") as f:
+        f.write(json.dumps(data))
+
+
+def timeit(func):
+    def run(*args, **kwargs):
+        t = time.time()
+        res = func(*args, **kwargs)
+        print('executing \'%s\' costed %.3fs' % (func.__name__, time.time() - t))
+        return res
+
+    return run
+
+
+def format_wav(audio_path):
+    if Path(audio_path).suffix == '.wav':
+        return
+    raw_audio, raw_sample_rate = librosa.load(audio_path, mono=True, sr=None)
+    soundfile.write(Path(audio_path).with_suffix(".wav"), raw_audio, raw_sample_rate)
+
+
+def get_end_file(dir_path, end):
+    file_lists = []
+    for root, dirs, files in os.walk(dir_path):
+        files = [f for f in files if f[0] != '.']
+        dirs[:] = [d for d in dirs if d[0] != '.']
+        for f_file in files:
+            if f_file.endswith(end):
+                file_lists.append(os.path.join(root, f_file).replace("\\", "/"))
+    return file_lists
+
+
+def get_md5(content):
+    return hashlib.new("md5", content).hexdigest()
+
+def fill_a_to_b(a, b):
+    if len(a) < len(b):
+        for _ in range(0, len(b) - len(a)):
+            a.append(a[0])
+
+def mkdir(paths: list):
+    for path in paths:
+        if not os.path.exists(path):
+            os.mkdir(path)
+
+def pad_array(arr, target_length):
+    current_length = arr.shape[0]
+    if current_length >= target_length:
+        return arr
+    else:
+        pad_width = target_length - current_length
+        pad_left = pad_width // 2
+        pad_right = pad_width - pad_left
+        padded_arr = np.pad(arr, (pad_left, pad_right), 'constant', constant_values=(0, 0))
+        return padded_arr
+    
+def split_list_by_n(list_collection, n, pre=0):
+    for i in range(0, len(list_collection), n):
+        yield list_collection[i-pre if i-pre>=0 else i: i + n]
+
+
+class F0FilterException(Exception):
+    pass
+
+class Svc(object):
+    def __init__(self, net_g_path, config_path,
+                 device=None,
+                 cluster_model_path="logs/44k/kmeans_10000.pt",
+                 nsf_hifigan_enhance = False,
+                 diffusion_model_path="logs/44k/diffusion/model_0.pt",
+                 diffusion_config_path="configs/diffusion.yaml",
+                 shallow_diffusion = False,
+                 only_diffusion = False,
+                 ):
+        self.net_g_path = net_g_path
+        self.only_diffusion = only_diffusion
+        self.shallow_diffusion = shallow_diffusion
+        if device is None:
+            # self.dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            self.dev = torch.device("cpu")
+        else:
+            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.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.speech_encoder = self.hps_ms.model.speech_encoder
+            except Exception as e:
+                self.speech_encoder = '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
+            else:
+                print("No diffusion model or config found. Shallow diffusion mode will False")
+                self.shallow_diffusion = self.only_diffusion = False
+                
+        # load hubert and model
+        if not self.only_diffusion:
+            self.load_model()
+            self.hubert_model = utils.get_speech_encoder(self.speech_encoder,device=self.dev)
+            self.volume_extractor = utils.Volume_Extractor(self.hop_size)
+        else:
+            self.hubert_model = utils.get_speech_encoder(self.diffusion_args.data.encoder,device=self.dev)
+            self.volume_extractor = utils.Volume_Extractor(self.diffusion_args.data.block_size)
+            
+        if os.path.exists(cluster_model_path):
+            self.cluster_model = cluster.get_cluster_model(cluster_model_path)
+        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)
+            
+    def load_model(self):
+        # get model configuration
+        self.net_g_ms = SynthesizerTrn(
+            self.hps_ms.data.filter_length // 2 + 1,
+            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)
+        if "half" in self.net_g_path and torch.cuda.is_available():
+            _ = self.net_g_ms.half().eval().to(self.dev)
+        else:
+            _ = self.net_g_ms.eval().to(self.dev)
+
+
+
+    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)
+        
+        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)
+        uv = torch.FloatTensor(uv).to(self.dev)
+
+        f0 = f0 * 2 ** (tran / 12)
+        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)
+        c = self.hubert_model.encoder(wav16k)
+        c = utils.repeat_expand_2d(c.squeeze(0), f0.shape[1])
+
+        if cluster_infer_ratio !=0:
+            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 = c.unsqueeze(0)
+        return c, f0, uv
+
+    def infer(self, speaker, tran, raw_path,
+              cluster_infer_ratio=0,
+              auto_predict_f0=False,
+              noice_scale=0.4,
+              f0_filter=False,
+              f0_predictor='pm',
+              enhancer_adaptive_key = 0,
+              cr_threshold = 0.05,
+              k_step = 100
+              ):
+
+        speaker_id = self.spk2id.get(speaker)
+        if not speaker_id and type(speaker) is int:
+            if len(self.spk2id.__dict__) >= speaker:
+                speaker_id = speaker
+        sid = torch.LongTensor([int(speaker_id)]).to(self.dev).unsqueeze(0)
+        wav, sr = librosa.load(raw_path, sr=self.target_sample)
+        c, f0, uv = self.get_unit_f0(wav, tran, cluster_infer_ratio, speaker, f0_filter,f0_predictor,cr_threshold=cr_threshold)
+        if "half" in self.net_g_path and torch.cuda.is_available():
+            c = c.half()
+        with torch.no_grad():
+            start = time.time()
+            if not self.only_diffusion:
+                audio,f0 = self.net_g_ms.infer(c, f0=f0, g=sid, uv=uv, predict_f0=auto_predict_f0, noice_scale=noice_scale)
+                audio = audio[0,0].data.float()
+                if self.shallow_diffusion:
+                    audio_mel = self.vocoder.extract(audio[None,:],self.target_sample)
+            else:
+                audio = torch.FloatTensor(wav).to(self.dev)
+                audio_mel = None
+            if self.only_diffusion or self.shallow_diffusion:
+                vol = self.volume_extractor.extract(audio[None,:])[None,:,None].to(self.dev)
+                f0 = f0[:,:,None]
+                c = c.transpose(-1,-2)
+                audio_mel = self.diffusion_model(
+                c, 
+                f0, 
+                vol, 
+                spk_id = sid, 
+                spk_mix_dict = None,
+                gt_spec=audio_mel,
+                infer=True, 
+                infer_speedup=self.diffusion_args.infer.speedup, 
+                method=self.diffusion_args.infer.method,
+                k_step=k_step)
+                audio = self.vocoder.infer(audio_mel, f0).squeeze()
+            if self.nsf_hifigan_enhance:
+                audio, _ = self.enhancer.enhance(
+                                    audio[None,:], 
+                                    self.target_sample, 
+                                    f0[:,:,None], 
+                                    self.hps_ms.data.hop_length, 
+                                    adaptive_key = enhancer_adaptive_key)
+            use_time = time.time() - start
+            print("vits use time:{}".format(use_time))
+        return audio, audio.shape[-1]
+
+    def clear_empty(self):
+        # clean up vram
+        torch.cuda.empty_cache()
+
+    def unload_model(self):
+        # unload model
+        self.net_g_ms = self.net_g_ms.to("cpu")
+        del self.net_g_ms
+        if hasattr(self,"enhancer"): 
+            self.enhancer.enhancer = self.enhancer.enhancer.to("cpu")
+            del self.enhancer.enhancer
+            del self.enhancer
+        gc.collect()
+
+    def slice_inference(self,
+                        raw_audio_path,
+                        spk,
+                        tran,
+                        slice_db,
+                        cluster_infer_ratio,
+                        auto_predict_f0,
+                        noice_scale,
+                        pad_seconds=0.5,
+                        clip_seconds=0,
+                        lg_num=0,
+                        lgr_num =0.75,
+                        f0_predictor='pm',
+                        enhancer_adaptive_key = 0,
+                        cr_threshold = 0.05,
+                        k_step = 100
+                        ):
+        wav_path = Path(raw_audio_path).with_suffix('.wav')
+        chunks = slicer.cut(wav_path, db_thresh=slice_db)
+        audio_data, audio_sr = slicer.chunks2audio(wav_path, chunks)
+        per_size = int(clip_seconds*audio_sr)
+        lg_size = int(lg_num*audio_sr)
+        lg_size_r = int(lg_size*lgr_num)
+        lg_size_c_l = (lg_size-lg_size_r)//2
+        lg_size_c_r = lg_size-lg_size_r-lg_size_c_l
+        lg = np.linspace(0,1,lg_size_r) if lg_size!=0 else 0
+        
+        audio = []
+        for (slice_tag, data) in audio_data:
+            print(f'#=====segment start, {round(len(data) / audio_sr, 3)}s======')
+            # padd
+            length = int(np.ceil(len(data) / audio_sr * self.target_sample))
+            if slice_tag:
+                print('jump empty segment')
+                _audio = np.zeros(length)
+                audio.extend(list(pad_array(_audio, length)))
+                continue
+            if per_size != 0:
+                datas = split_list_by_n(data, per_size,lg_size)
+            else:
+                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======')
+                # padd
+                pad_len = int(audio_sr * pad_seconds)
+                dat = np.concatenate([np.zeros([pad_len]), dat, np.zeros([pad_len])])
+                raw_path = io.BytesIO()
+                soundfile.write(raw_path, dat, audio_sr, format="wav")
+                raw_path.seek(0)
+                out_audio, out_sr = self.infer(spk, tran, raw_path,
+                                                    cluster_infer_ratio=cluster_infer_ratio,
+                                                    auto_predict_f0=auto_predict_f0,
+                                                    noice_scale=noice_scale,
+                                                    f0_predictor = f0_predictor,
+                                                    enhancer_adaptive_key = enhancer_adaptive_key,
+                                                    cr_threshold = cr_threshold,
+                                                    k_step = k_step
+                                                    )
+                _audio = out_audio.cpu().numpy()
+                pad_len = int(self.target_sample * pad_seconds)
+                _audio = _audio[pad_len:-pad_len]
+                _audio = pad_array(_audio, per_length)
+                if lg_size!=0 and k!=0:
+                    lg1 = audio[-(lg_size_r+lg_size_c_r):-lg_size_c_r] if lgr_num != 1 else audio[-lg_size:]
+                    lg2 = _audio[lg_size_c_l:lg_size_c_l+lg_size_r]  if lgr_num != 1 else _audio[0:lg_size]
+                    lg_pre = lg1*(1-lg)+lg2*lg
+                    audio = audio[0:-(lg_size_r+lg_size_c_r)] if lgr_num != 1 else audio[0:-lg_size]
+                    audio.extend(lg_pre)
+                    _audio = _audio[lg_size_c_l+lg_size_r:] if lgr_num != 1 else _audio[lg_size:]
+                audio.extend(list(_audio))
+        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

@@ -0,0 +1,171 @@
+import hashlib
+import json
+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 models import SynthesizerTrn
+logging.getLogger('numba').setLevel(logging.WARNING)
+logging.getLogger('matplotlib').setLevel(logging.WARNING)
+
+
+def resize2d_f0(x, target_len):
+    source = np.array(x)
+    source[source < 0.001] = np.nan
+    target = np.interp(np.arange(0, len(source) * target_len, len(source)) / target_len, np.arange(0, len(source)),
+                       source)
+    res = np.nan_to_num(target)
+    return res
+
+
+def get_f0(x, p_len, f0_up_key=0):
+
+    time_step = 160 / 16000 * 1000
+    f0_min = 50
+    f0_max = 1100
+    f0_mel_min = 1127 * np.log(1 + f0_min / 700)
+    f0_mel_max = 1127 * np.log(1 + f0_max / 700)
+
+    f0 = parselmouth.Sound(x, 16000).to_pitch_ac(
+        time_step=time_step / 1000, voicing_threshold=0.6,
+        pitch_floor=f0_min, pitch_ceiling=f0_max).selected_array['frequency']
+
+    pad_size = (p_len - len(f0) + 1) // 2
+    if(pad_size > 0 or p_len - len(f0) - pad_size > 0):
+        f0 = np.pad(
+            f0, [[pad_size, p_len - len(f0) - pad_size]], mode='constant')
+
+    f0 *= pow(2, f0_up_key / 12)
+    f0_mel = 1127 * np.log(1 + f0 / 700)
+    f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - f0_mel_min) * \
+        254 / (f0_mel_max - f0_mel_min) + 1
+    f0_mel[f0_mel <= 1] = 1
+    f0_mel[f0_mel > 255] = 255
+    f0_coarse = np.rint(f0_mel).astype(np.int)
+    return f0_coarse, f0
+
+
+def clean_pitch(input_pitch):
+    num_nan = np.sum(input_pitch == 1)
+    if num_nan / len(input_pitch) > 0.9:
+        input_pitch[input_pitch != 1] = 1
+    return input_pitch
+
+
+def plt_pitch(input_pitch):
+    input_pitch = input_pitch.astype(float)
+    input_pitch[input_pitch == 1] = np.nan
+    return input_pitch
+
+
+def f0_to_pitch(ff):
+    f0_pitch = 69 + 12 * np.log2(ff / 440)
+    return f0_pitch
+
+
+def fill_a_to_b(a, b):
+    if len(a) < len(b):
+        for _ in range(0, len(b) - len(a)):
+            a.append(a[0])
+
+
+def mkdir(paths: list):
+    for path in paths:
+        if not os.path.exists(path):
+            os.mkdir(path)
+
+
+class VitsSvc(object):
+    def __init__(self):
+        self.device = torch.device(
+            "cuda" if torch.cuda.is_available() else "cpu")
+        self.SVCVITS = None
+        self.hps = None
+        self.speakers = None
+        self.hubert_soft = utils.get_hubert_model()
+
+    def set_device(self, device):
+        self.device = torch.device(device)
+        self.hubert_soft.to(self.device)
+        if self.SVCVITS != None:
+            self.SVCVITS.to(self.device)
+
+    def loadCheckpoint(self, path):
+        self.hps = utils.get_hparams_from_file(
+            f"checkpoints/{path}/config.json")
+        self.SVCVITS = SynthesizerTrn(
+            self.hps.data.filter_length // 2 + 1,
+            self.hps.train.segment_size // self.hps.data.hop_length,
+            **self.hps.model)
+        _ = utils.load_checkpoint(
+            f"checkpoints/{path}/model.pth", self.SVCVITS, None)
+        _ = self.SVCVITS.eval().to(self.device)
+        self.speakers = self.hps.spk
+
+    def get_units(self, source, sr):
+        source = source.unsqueeze(0).to(self.device)
+        with torch.inference_mode():
+            units = self.hubert_soft.units(source)
+            return units
+
+    def get_unit_pitch(self, in_path, tran):
+        source, sr = torchaudio.load(in_path)
+        source = torchaudio.functional.resample(source, sr, 16000)
+        if len(source.shape) == 2 and source.shape[1] >= 2:
+            source = torch.mean(source, dim=0).unsqueeze(0)
+        soft = self.get_units(source, sr).squeeze(0).cpu().numpy()
+        f0_coarse, f0 = get_f0(source.cpu().numpy()[0], soft.shape[0]*2, tran)
+        return soft, f0
+
+    def infer(self, speaker_id, tran, raw_path):
+        speaker_id = self.speakers[speaker_id]
+        sid = torch.LongTensor([int(speaker_id)]).to(self.device).unsqueeze(0)
+        soft, pitch = self.get_unit_pitch(raw_path, tran)
+        f0 = torch.FloatTensor(clean_pitch(pitch)).unsqueeze(0).to(self.device)
+        stn_tst = torch.FloatTensor(soft)
+        with torch.no_grad():
+            x_tst = stn_tst.unsqueeze(0).to(self.device)
+            x_tst = torch.repeat_interleave(
+                x_tst, repeats=2, dim=1).transpose(1, 2)
+            audio, _ = self.SVCVITS.infer(x_tst, f0=f0, g=sid)[
+                0, 0].data.float()
+        return audio, audio.shape[-1]
+
+    def inference(self, srcaudio, chara, tran, slice_db):
+        sampling_rate, audio = srcaudio
+        audio = (audio / np.iinfo(audio.dtype).max).astype(np.float32)
+        if len(audio.shape) > 1:
+            audio = librosa.to_mono(audio.transpose(1, 0))
+        if sampling_rate != 16000:
+            audio = librosa.resample(
+                audio, orig_sr=sampling_rate, target_sr=16000)
+        soundfile.write("tmpwav.wav", audio, 16000, format="wav")
+        chunks = slicer.cut("tmpwav.wav", db_thresh=slice_db)
+        audio_data, audio_sr = slicer.chunks2audio("tmpwav.wav", chunks)
+        audio = []
+        for (slice_tag, data) in audio_data:
+            length = int(np.ceil(len(data) / audio_sr *
+                                 self.hps.data.sampling_rate))
+            raw_path = io.BytesIO()
+            soundfile.write(raw_path, data, audio_sr, format="wav")
+            raw_path.seek(0)
+            if slice_tag:
+                _audio = np.zeros(length)
+            else:
+                out_audio, out_sr = self.infer(chara, tran, raw_path)
+                _audio = out_audio.cpu().numpy()
+            audio.extend(list(_audio))
+        audio = (np.array(audio) * 32768.0).astype('int16')
+        return (self.hps.data.sampling_rate, audio)

inference/slicer.py

@@ -0,0 +1,142 @@
+import librosa
+import torch
+import torchaudio
+
+
+class Slicer:
+    def __init__(self,
+                 sr: int,
+                 threshold: float = -40.,
+                 min_length: int = 5000,
+                 min_interval: int = 300,
+                 hop_size: int = 20,
+                 max_sil_kept: int = 5000):
+        if not min_length >= min_interval >= hop_size:
+            raise ValueError('The following condition must be satisfied: min_length >= min_interval >= hop_size')
+        if not max_sil_kept >= hop_size:
+            raise ValueError('The following condition must be satisfied: max_sil_kept >= hop_size')
+        min_interval = sr * min_interval / 1000
+        self.threshold = 10 ** (threshold / 20.)
+        self.hop_size = round(sr * hop_size / 1000)
+        self.win_size = min(round(min_interval), 4 * self.hop_size)
+        self.min_length = round(sr * min_length / 1000 / self.hop_size)
+        self.min_interval = round(min_interval / self.hop_size)
+        self.max_sil_kept = round(sr * max_sil_kept / 1000 / self.hop_size)
+
+    def _apply_slice(self, waveform, begin, end):
+        if len(waveform.shape) > 1:
+            return waveform[:, begin * self.hop_size: min(waveform.shape[1], end * self.hop_size)]
+        else:
+            return waveform[begin * self.hop_size: min(waveform.shape[0], end * self.hop_size)]
+
+    # @timeit
+    def slice(self, waveform):
+        if len(waveform.shape) > 1:
+            samples = librosa.to_mono(waveform)
+        else:
+            samples = waveform
+        if samples.shape[0] <= self.min_length:
+            return {"0": {"slice": False, "split_time": f"0,{len(waveform)}"}}
+        rms_list = librosa.feature.rms(y=samples, frame_length=self.win_size, hop_length=self.hop_size).squeeze(0)
+        sil_tags = []
+        silence_start = None
+        clip_start = 0
+        for i, rms in enumerate(rms_list):
+            # Keep looping while frame is silent.
+            if rms < self.threshold:
+                # Record start of silent frames.
+                if silence_start is None:
+                    silence_start = i
+                continue
+            # Keep looping while frame is not silent and silence start has not been recorded.
+            if silence_start is None:
+                continue
+            # Clear recorded silence start if interval is not enough or clip is too short
+            is_leading_silence = silence_start == 0 and i > self.max_sil_kept
+            need_slice_middle = i - silence_start >= self.min_interval and i - clip_start >= self.min_length
+            if not is_leading_silence and not need_slice_middle:
+                silence_start = None
+                continue
+            # Need slicing. Record the range of silent frames to be removed.
+            if i - silence_start <= self.max_sil_kept:
+                pos = rms_list[silence_start: i + 1].argmin() + silence_start
+                if silence_start == 0:
+                    sil_tags.append((0, pos))
+                else:
+                    sil_tags.append((pos, pos))
+                clip_start = pos
+            elif i - silence_start <= self.max_sil_kept * 2:
+                pos = rms_list[i - self.max_sil_kept: silence_start + self.max_sil_kept + 1].argmin()
+                pos += i - self.max_sil_kept
+                pos_l = rms_list[silence_start: silence_start + self.max_sil_kept + 1].argmin() + silence_start
+                pos_r = rms_list[i - self.max_sil_kept: i + 1].argmin() + i - self.max_sil_kept
+                if silence_start == 0:
+                    sil_tags.append((0, pos_r))
+                    clip_start = pos_r
+                else:
+                    sil_tags.append((min(pos_l, pos), max(pos_r, pos)))
+                    clip_start = max(pos_r, pos)
+            else:
+                pos_l = rms_list[silence_start: silence_start + self.max_sil_kept + 1].argmin() + silence_start
+                pos_r = rms_list[i - self.max_sil_kept: i + 1].argmin() + i - self.max_sil_kept
+                if silence_start == 0:
+                    sil_tags.append((0, pos_r))
+                else:
+                    sil_tags.append((pos_l, pos_r))
+                clip_start = pos_r
+            silence_start = None
+        # Deal with trailing silence.
+        total_frames = rms_list.shape[0]
+        if silence_start is not None and total_frames - silence_start >= self.min_interval:
+            silence_end = min(total_frames, silence_start + self.max_sil_kept)
+            pos = rms_list[silence_start: silence_end + 1].argmin() + silence_start
+            sil_tags.append((pos, total_frames + 1))
+        # Apply and return slices.
+        if len(sil_tags) == 0:
+            return {"0": {"slice": False, "split_time": f"0,{len(waveform)}"}}
+        else:
+            chunks = []
+            # 第一段静音并非从头开始，补上有声片段
+            if sil_tags[0][0]:
+                chunks.append(
+                    {"slice": False, "split_time": f"0,{min(waveform.shape[0], sil_tags[0][0] * self.hop_size)}"})
+            for i in range(0, len(sil_tags)):
+                # 标识有声片段（跳过第一段）
+                if i:
+                    chunks.append({"slice": False,
+                                   "split_time": f"{sil_tags[i - 1][1] * self.hop_size},{min(waveform.shape[0], sil_tags[i][0] * self.hop_size)}"})
+                # 标识所有静音片段
+                chunks.append({"slice": True,
+                               "split_time": f"{sil_tags[i][0] * self.hop_size},{min(waveform.shape[0], sil_tags[i][1] * self.hop_size)}"})
+            # 最后一段静音并非结尾，补上结尾片段
+            if sil_tags[-1][1] * self.hop_size < len(waveform):
+                chunks.append({"slice": False, "split_time": f"{sil_tags[-1][1] * self.hop_size},{len(waveform)}"})
+            chunk_dict = {}
+            for i in range(len(chunks)):
+                chunk_dict[str(i)] = chunks[i]
+            return chunk_dict
+
+
+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,
+        min_length=min_len
+    )
+    chunks = slicer.slice(audio)
+    return chunks
+
+
+def chunks2audio(audio_path, chunks):
+    chunks = dict(chunks)
+    audio, sr = torchaudio.load(audio_path)
+    if len(audio.shape) == 2 and audio.shape[1] >= 2:
+        audio = torch.mean(audio, dim=0).unsqueeze(0)
+    audio = audio.cpu().numpy()[0]
+    result = []
+    for k, v in chunks.items():
+        tag = v["split_time"].split(",")
+        if tag[0] != tag[1]:
+            result.append((v["slice"], audio[int(tag[0]):int(tag[1])]))
+    return result, sr

models.py

@@ -0,0 +1,420 @@
+import copy
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+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 vdecoder.hifigan.models import Generator
+from utils import f0_to_coarse
+
+class ResidualCouplingBlock(nn.Module):
+  def __init__(self,
+      channels,
+      hidden_channels,
+      kernel_size,
+      dilation_rate,
+      n_layers,
+      n_flows=4,
+      gin_channels=0):
+    super().__init__()
+    self.channels = channels
+    self.hidden_channels = hidden_channels
+    self.kernel_size = kernel_size
+    self.dilation_rate = dilation_rate
+    self.n_layers = n_layers
+    self.n_flows = n_flows
+    self.gin_channels = gin_channels
+
+    self.flows = nn.ModuleList()
+    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))
+      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 Encoder(nn.Module):
+  def __init__(self,
+      in_channels,
+      out_channels,
+      hidden_channels,
+      kernel_size,
+      dilation_rate,
+      n_layers,
+      gin_channels=0):
+    super().__init__()
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+    self.hidden_channels = hidden_channels
+    self.kernel_size = kernel_size
+    self.dilation_rate = dilation_rate
+    self.n_layers = n_layers
+    self.gin_channels = gin_channels
+
+    self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
+    self.enc = modules.WN(hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels)
+    self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+  def forward(self, x, x_lengths, g=None):
+    # print(x.shape,x_lengths.shape)
+    x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)
+    x = self.pre(x) * x_mask
+    x = self.enc(x, x_mask, g=g)
+    stats = self.proj(x) * x_mask
+    m, logs = torch.split(stats, self.out_channels, dim=1)
+    z = (m + torch.randn_like(m) * torch.exp(logs)) * x_mask
+    return z, m, logs, x_mask
+
+
+class TextEncoder(nn.Module):
+  def __init__(self,
+      out_channels,
+      hidden_channels,
+      kernel_size,
+      n_layers,
+      gin_channels=0,
+      filter_channels=None,
+      n_heads=None,
+      p_dropout=None):
+    super().__init__()
+    self.out_channels = out_channels
+    self.hidden_channels = hidden_channels
+    self.kernel_size = kernel_size
+    self.n_layers = n_layers
+    self.gin_channels = gin_channels
+    self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+    self.f0_emb = nn.Embedding(256, hidden_channels)
+
+    self.enc_ =  attentions.Encoder(
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size,
+        p_dropout)
+
+  def forward(self, x, x_mask, f0=None, noice_scale=1):
+    x = x + self.f0_emb(f0).transpose(1,2)
+    x = self.enc_(x * x_mask, x_mask)
+    stats = self.proj(x) * x_mask
+    m, logs = torch.split(stats, self.out_channels, dim=1)
+    z = (m + torch.randn_like(m) * torch.exp(logs) * noice_scale) * x_mask
+
+    return z, m, logs, x_mask
+
+
+
+class DiscriminatorP(torch.nn.Module):
+    def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
+        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
+        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))),
+            norm_f(Conv2d(128, 512, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
+            norm_f(Conv2d(512, 1024, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
+            norm_f(Conv2d(1024, 1024, (kernel_size, 1), 1, padding=(get_padding(kernel_size, 1), 0))),
+        ])
+        self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
+
+    def forward(self, x):
+        fmap = []
+
+        # 1d to 2d
+        b, c, t = x.shape
+        if t % self.period != 0: # pad first
+            n_pad = self.period - (t % self.period)
+            x = F.pad(x, (0, n_pad), "reflect")
+            t = t + n_pad
+        x = x.view(b, c, t // self.period, self.period)
+
+        for l in self.convs:
+            x = l(x)
+            x = F.leaky_relu(x, modules.LRELU_SLOPE)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+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
+        self.convs = nn.ModuleList([
+            norm_f(Conv1d(1, 16, 15, 1, padding=7)),
+            norm_f(Conv1d(16, 64, 41, 4, groups=4, padding=20)),
+            norm_f(Conv1d(64, 256, 41, 4, groups=16, padding=20)),
+            norm_f(Conv1d(256, 1024, 41, 4, groups=64, padding=20)),
+            norm_f(Conv1d(1024, 1024, 41, 4, groups=256, padding=20)),
+            norm_f(Conv1d(1024, 1024, 5, 1, padding=2)),
+        ])
+        self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))
+
+    def forward(self, x):
+        fmap = []
+
+        for l in self.convs:
+            x = l(x)
+            x = F.leaky_relu(x, modules.LRELU_SLOPE)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+class MultiPeriodDiscriminator(torch.nn.Module):
+    def __init__(self, use_spectral_norm=False):
+        super(MultiPeriodDiscriminator, self).__init__()
+        periods = [2,3,5,7,11]
+
+        discs = [DiscriminatorS(use_spectral_norm=use_spectral_norm)]
+        discs = discs + [DiscriminatorP(i, use_spectral_norm=use_spectral_norm) for i in periods]
+        self.discriminators = nn.ModuleList(discs)
+
+    def forward(self, y, y_hat):
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        for i, d in enumerate(self.discriminators):
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            y_d_rs.append(y_d_r)
+            y_d_gs.append(y_d_g)
+            fmap_rs.append(fmap_r)
+            fmap_gs.append(fmap_g)
+
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+class SpeakerEncoder(torch.nn.Module):
+    def __init__(self, mel_n_channels=80, model_num_layers=3, model_hidden_size=256, model_embedding_size=256):
+        super(SpeakerEncoder, self).__init__()
+        self.lstm = nn.LSTM(mel_n_channels, model_hidden_size, model_num_layers, batch_first=True)
+        self.linear = nn.Linear(model_hidden_size, model_embedding_size)
+        self.relu = nn.ReLU()
+
+    def forward(self, mels):
+        self.lstm.flatten_parameters()
+        _, (hidden, _) = self.lstm(mels)
+        embeds_raw = self.relu(self.linear(hidden[-1]))
+        return embeds_raw / torch.norm(embeds_raw, dim=1, keepdim=True)
+
+    def compute_partial_slices(self, total_frames, partial_frames, partial_hop):
+        mel_slices = []
+        for i in range(0, total_frames-partial_frames, partial_hop):
+            mel_range = torch.arange(i, i+partial_frames)
+            mel_slices.append(mel_range)
+
+        return mel_slices
+
+    def embed_utterance(self, mel, partial_frames=128, partial_hop=64):
+        mel_len = mel.size(1)
+        last_mel = mel[:,-partial_frames:]
+
+        if mel_len > partial_frames:
+            mel_slices = self.compute_partial_slices(mel_len, partial_frames, partial_hop)
+            mels = list(mel[:,s] for s in mel_slices)
+            mels.append(last_mel)
+            mels = torch.stack(tuple(mels), 0).squeeze(1)
+
+            with torch.no_grad():
+                partial_embeds = self(mels)
+            embed = torch.mean(partial_embeds, axis=0).unsqueeze(0)
+            #embed = embed / torch.linalg.norm(embed, 2)
+        else:
+            with torch.no_grad():
+                embed = self(last_mel)
+
+        return embed
+
+class F0Decoder(nn.Module):
+    def __init__(self,
+                 out_channels,
+                 hidden_channels,
+                 filter_channels,
+                 n_heads,
+                 n_layers,
+                 kernel_size,
+                 p_dropout,
+                 spk_channels=0):
+        super().__init__()
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.spk_channels = spk_channels
+
+        self.prenet = nn.Conv1d(hidden_channels, hidden_channels, 3, padding=1)
+        self.decoder = attentions.FFT(
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            p_dropout)
+        self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
+        self.f0_prenet = nn.Conv1d(1, hidden_channels , 3, padding=1)
+        self.cond = nn.Conv1d(spk_channels, hidden_channels, 1)
+
+    def forward(self, x, norm_f0, x_mask, spk_emb=None):
+        x = torch.detach(x)
+        if (spk_emb is not None):
+            x = x + self.cond(spk_emb)
+        x += self.f0_prenet(norm_f0)
+        x = self.prenet(x) * x_mask
+        x = self.decoder(x * x_mask, x_mask)
+        x = self.proj(x) * x_mask
+        return x
+
+
+class SynthesizerTrn(nn.Module):
+  """
+  Synthesizer for Training
+  """
+
+  def __init__(self,
+    spec_channels,
+    segment_size,
+    inter_channels,
+    hidden_channels,
+    filter_channels,
+    n_heads,
+    n_layers,
+    kernel_size,
+    p_dropout,
+    resblock,
+    resblock_kernel_sizes,
+    resblock_dilation_sizes,
+    upsample_rates,
+    upsample_initial_channel,
+    upsample_kernel_sizes,
+    gin_channels,
+    ssl_dim,
+    n_speakers,
+    sampling_rate=44100,
+    **kwargs):
+
+    super().__init__()
+    self.spec_channels = spec_channels
+    self.inter_channels = inter_channels
+    self.hidden_channels = hidden_channels
+    self.filter_channels = filter_channels
+    self.n_heads = n_heads
+    self.n_layers = n_layers
+    self.kernel_size = kernel_size
+    self.p_dropout = p_dropout
+    self.resblock = resblock
+    self.resblock_kernel_sizes = resblock_kernel_sizes
+    self.resblock_dilation_sizes = resblock_dilation_sizes
+    self.upsample_rates = upsample_rates
+    self.upsample_initial_channel = upsample_initial_channel
+    self.upsample_kernel_sizes = upsample_kernel_sizes
+    self.segment_size = segment_size
+    self.gin_channels = gin_channels
+    self.ssl_dim = ssl_dim
+    self.emb_g = nn.Embedding(n_speakers, gin_channels)
+
+    self.pre = nn.Conv1d(ssl_dim, hidden_channels, kernel_size=5, padding=2)
+
+    self.enc_p = TextEncoder(
+        inter_channels,
+        hidden_channels,
+        filter_channels=filter_channels,
+        n_heads=n_heads,
+        n_layers=n_layers,
+        kernel_size=kernel_size,
+        p_dropout=p_dropout
+    )
+    hps = {
+        "sampling_rate": sampling_rate,
+        "inter_channels": inter_channels,
+        "resblock": resblock,
+        "resblock_kernel_sizes": resblock_kernel_sizes,
+        "resblock_dilation_sizes": resblock_dilation_sizes,
+        "upsample_rates": upsample_rates,
+        "upsample_initial_channel": upsample_initial_channel,
+        "upsample_kernel_sizes": upsample_kernel_sizes,
+        "gin_channels": gin_channels,
+    }
+    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
+    )
+    self.emb_uv = nn.Embedding(2, hidden_channels)
+
+  def forward(self, c, f0, uv, spec, g=None, c_lengths=None, spec_lengths=None):
+    g = self.emb_g(g).transpose(1,2)
+    # 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)
+
+    # 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)
+
+    # 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) 
+
+    # flow
+    z_p = self.flow(z, spec_mask, g=g)
+    z_slice, pitch_slice, ids_slice = commons.rand_slice_segments_with_pitch(z, f0, spec_lengths, self.segment_size)
+
+    # nsf decoder
+    o = self.dec(z_slice, g=g, f0=pitch_slice)
+
+    return o, ids_slice, spec_mask, (z, z_p, m_p, logs_p, m_q, logs_q), pred_lf0, norm_lf0, lf0
+
+  def infer(self, c, f0, uv, g=None, noice_scale=0.35, predict_f0=False):
+    c_lengths = (torch.ones(c.size(0)) * c.size(-1)).to(c.device)
+    g = self.emb_g(g).transpose(1,2)
+    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)
+
+    if 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)
+        f0 = (700 * (torch.pow(10, pred_lf0 * 500 / 2595) - 1)).squeeze(1)
+
+    z_p, m_p, logs_p, c_mask = self.enc_p(x, x_mask, f0=f0_to_coarse(f0), noice_scale=noice_scale)
+    z = self.flow(z_p, c_mask, g=g, reverse=True)
+    o = self.dec(z * c_mask, g=g, f0=f0)
+    return o,f0

modules/F0Predictor/CrepeF0Predictor.py

@@ -0,0 +1,31 @@
+from modules.F0Predictor.F0Predictor import F0Predictor
+from modules.F0Predictor.crepe import CrepePitchExtractor
+import torch
+
+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)
+        self.hop_length = hop_length
+        self.f0_min = f0_min
+        self.f0_max = f0_max
+        self.device = device
+        self.threshold = threshold
+        self.sampling_rate = sampling_rate
+
+    def compute_f0(self,wav,p_len=None):
+        x = torch.FloatTensor(wav).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,uv = self.F0Creper(x[None,:].float(),self.sampling_rate,pad_to=p_len)
+        return f0
+    
+    def compute_f0_uv(self,wav,p_len=None):
+        x = torch.FloatTensor(wav).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,uv = self.F0Creper(x[None,:].float(),self.sampling_rate,pad_to=p_len)
+        return f0,uv

modules/F0Predictor/DioF0Predictor.py

@@ -0,0 +1,85 @@
+from modules.F0Predictor.F0Predictor import F0Predictor
+import pyworld
+import numpy as np
+
+class DioF0Predictor(F0Predictor):
+    def __init__(self,hop_length=512,f0_min=50,f0_max=1100,sampling_rate=44100):
+        self.hop_length = hop_length
+        self.f0_min = f0_min
+        self.f0_max = f0_max
+        self.sampling_rate = sampling_rate
+
+    def interpolate_f0(self,f0):
+        '''
+        对F0进行插值处理
+        '''
+    
+        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]
+
+    def resize_f0(self,x, target_len):
+        source = np.array(x)
+        source[source<0.001] = np.nan
+        target = np.interp(np.arange(0, len(source)*target_len, len(source))/ target_len, np.arange(0, len(source)), source)
+        res = np.nan_to_num(target)
+        return res
+        
+    def compute_f0(self,wav,p_len=None):
+        if p_len is None:
+            p_len = wav.shape[0]//self.hop_length
+        f0, t = pyworld.dio(
+            wav.astype(np.double),
+            fs=self.sampling_rate,
+            f0_floor=self.f0_min,
+            f0_ceil=self.f0_max,
+            frame_period=1000 * self.hop_length / self.sampling_rate,
+        )
+        f0 = pyworld.stonemask(wav.astype(np.double), f0, t, self.sampling_rate)
+        for index, pitch in enumerate(f0):
+            f0[index] = round(pitch, 1)
+        return self.interpolate_f0(self.resize_f0(f0, p_len))[0]
+
+    def compute_f0_uv(self,wav,p_len=None):
+        if p_len is None:
+            p_len = wav.shape[0]//self.hop_length
+        f0, t = pyworld.dio(
+            wav.astype(np.double),
+            fs=self.sampling_rate,
+            f0_floor=self.f0_min,
+            f0_ceil=self.f0_max,
+            frame_period=1000 * self.hop_length / self.sampling_rate,
+        )
+        f0 = pyworld.stonemask(wav.astype(np.double), f0, t, self.sampling_rate)
+        for index, pitch in enumerate(f0):
+            f0[index] = round(pitch, 1)
+        return self.interpolate_f0(self.resize_f0(f0, p_len))

modules/F0Predictor/F0Predictor.py

@@ -0,0 +1,16 @@
+class F0Predictor(object):
+    def compute_f0(self,wav,p_len):
+        '''
+        input: wav:[signal_length]
+               p_len:int
+        output: f0:[signal_length//hop_length]
+        '''
+        pass
+
+    def compute_f0_uv(self,wav,p_len):
+        '''
+        input: wav:[signal_length]
+               p_len:int
+        output: f0:[signal_length//hop_length],uv:[signal_length//hop_length]
+        '''
+        pass

modules/F0Predictor/HarvestF0Predictor.py

@@ -0,0 +1,81 @@
+from modules.F0Predictor.F0Predictor import F0Predictor
+import pyworld
+import numpy as np
+
+class HarvestF0Predictor(F0Predictor):
+    def __init__(self,hop_length=512,f0_min=50,f0_max=1100,sampling_rate=44100):
+        self.hop_length = hop_length
+        self.f0_min = f0_min
+        self.f0_max = f0_max
+        self.sampling_rate = sampling_rate
+
+    def interpolate_f0(self,f0):
+        '''
+        对F0进行插值处理
+        '''
+    
+        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]
+
+    def resize_f0(self,x, target_len):
+        source = np.array(x)
+        source[source<0.001] = np.nan
+        target = np.interp(np.arange(0, len(source)*target_len, len(source))/ target_len, np.arange(0, len(source)), source)
+        res = np.nan_to_num(target)
+        return res
+        
+    def compute_f0(self,wav,p_len=None):
+        if p_len is None:
+            p_len = wav.shape[0]//self.hop_length
+        f0, t = pyworld.harvest(
+                wav.astype(np.double),
+                fs=self.hop_length,
+                f0_ceil=self.f0_max,
+                f0_floor=self.f0_min,
+                frame_period=1000 * self.hop_length / self.sampling_rate,
+            )
+        f0 = pyworld.stonemask(wav.astype(np.double), f0, t, self.fs)
+        return self.interpolate_f0(self.resize_f0(f0, p_len))[0]
+
+    def compute_f0_uv(self,wav,p_len=None):
+        if p_len is None:
+            p_len = wav.shape[0]//self.hop_length
+        f0, t = pyworld.harvest(
+            wav.astype(np.double),
+            fs=self.sampling_rate,
+            f0_floor=self.f0_min,
+            f0_ceil=self.f0_max,
+            frame_period=1000 * self.hop_length / self.sampling_rate,
+        )
+        f0 = pyworld.stonemask(wav.astype(np.double), f0, t, self.sampling_rate)
+        return self.interpolate_f0(self.resize_f0(f0, p_len))

modules/F0Predictor/PMF0Predictor.py

@@ -0,0 +1,83 @@
+from modules.F0Predictor.F0Predictor import F0Predictor
+import parselmouth
+import numpy as np
+
+class PMF0Predictor(F0Predictor):
+    def __init__(self,hop_length=512,f0_min=50,f0_max=1100,sampling_rate=44100):
+        self.hop_length = hop_length
+        self.f0_min = f0_min
+        self.f0_max = f0_max
+        self.sampling_rate = sampling_rate
+
+    
+    def interpolate_f0(self,f0):
+        '''
+        对F0进行插值处理
+        '''
+    
+        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]
+
+    def compute_f0(self,wav,p_len=None):
+        x = wav
+        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"
+        time_step = self.hop_length / self.sampling_rate * 1000
+        f0 = parselmouth.Sound(x, self.sampling_rate).to_pitch_ac(
+            time_step=time_step / 1000, voicing_threshold=0.6,
+            pitch_floor=self.f0_min, pitch_ceiling=self.f0_max).selected_array['frequency']
+
+        pad_size=(p_len - len(f0) + 1) // 2
+        if(pad_size>0 or p_len - len(f0) - pad_size>0):
+            f0 = np.pad(f0,[[pad_size,p_len - len(f0) - pad_size]], mode='constant')
+        f0,uv = self.interpolate_f0(f0)
+        return f0
+
+    def compute_f0_uv(self,wav,p_len=None):
+        x = wav
+        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"
+        time_step = self.hop_length / self.sampling_rate * 1000
+        f0 = parselmouth.Sound(x, self.sampling_rate).to_pitch_ac(
+            time_step=time_step / 1000, voicing_threshold=0.6,
+            pitch_floor=self.f0_min, pitch_ceiling=self.f0_max).selected_array['frequency']
+
+        pad_size=(p_len - len(f0) + 1) // 2
+        if(pad_size>0 or p_len - len(f0) - pad_size>0):
+            f0 = np.pad(f0,[[pad_size,p_len - len(f0) - pad_size]], mode='constant')
+        f0,uv = self.interpolate_f0(f0)
+        return f0,uv

modules/F0Predictor/crepe.py

@@ -0,0 +1,340 @@
+from typing import Optional,Union
+try:
+    from typing import Literal
+except Exception as e:
+    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
+
+def repeat_expand(
+    content: Union[torch.Tensor, np.ndarray], target_len: int, mode: str = "nearest"
+):
+    """Repeat content to target length.
+    This is a wrapper of torch.nn.functional.interpolate.
+
+    Args:
+        content (torch.Tensor): tensor
+        target_len (int): target length
+        mode (str, optional): interpolation mode. Defaults to "nearest".
+
+    Returns:
+        torch.Tensor: tensor
+    """
+
+    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]
+
+
+class BasePitchExtractor:
+    def __init__(
+        self,
+        hop_length: int = 512,
+        f0_min: float = 50.0,
+        f0_max: float = 1100.0,
+        keep_zeros: bool = True,
+    ):
+        """Base pitch extractor.
+
+        Args:
+            hop_length (int, optional): Hop length. Defaults to 512.
+            f0_min (float, optional): Minimum f0. Defaults to 50.0.
+            f0_max (float, optional): Maximum f0. Defaults to 1100.0.
+            keep_zeros (bool, optional): Whether keep zeros in pitch. Defaults to True.
+        """
+
+        self.hop_length = hop_length
+        self.f0_min = f0_min
+        self.f0_max = f0_max
+        self.keep_zeros = keep_zeros
+
+    def __call__(self, x, sampling_rate=44100, pad_to=None):
+        raise NotImplementedError("BasePitchExtractor is not callable.")
+
+    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 = repeat_expand(f0, pad_to)
+
+        if self.keep_zeros:
+            return f0
+        
+        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
+
+        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)
+
+        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)
+    
+        # 大概可以用 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))
+        
+        return f0,vuv_vector
+
+
+class MaskedAvgPool1d(nn.Module):
+    def __init__(
+        self, kernel_size: int, stride: Optional[int] = None, padding: Optional[int] = 0
+    ):
+        """An implementation of mean pooling that supports masked values.
+
+        Args:
+            kernel_size (int): The size of the median pooling window.
+            stride (int, optional): The stride of the median pooling window. Defaults to None.
+            padding (int, optional): The padding of the median pooling window. Defaults to 0.
+        """
+
+        super(MaskedAvgPool1d, self).__init__()
+        self.kernel_size = kernel_size
+        self.stride = stride or kernel_size
+        self.padding = padding
+
+    def forward(self, x, mask=None):
+        ndim = x.dim()
+        if ndim == 2:
+            x = x.unsqueeze(1)
+
+        assert (
+            x.dim() == 3
+        ), "Input tensor must have 2 or 3 dimensions (batch_size, channels, width)"
+
+        # Apply the mask by setting masked elements to zero, or make NaNs zero
+        if mask is None:
+            mask = ~torch.isnan(x)
+
+        # Ensure mask has the same shape as the input tensor
+        assert x.shape == mask.shape, "Input tensor and mask must have the same shape"
+
+        masked_x = torch.where(mask, x, torch.zeros_like(x))
+        # Create a ones kernel with the same number of channels as the input tensor
+        ones_kernel = torch.ones(x.size(1), 1, self.kernel_size, device=x.device)
+
+        # Perform sum pooling
+        sum_pooled = nn.functional.conv1d(
+            masked_x,
+            ones_kernel,
+            stride=self.stride,
+            padding=self.padding,
+            groups=x.size(1),
+        )
+
+        # Count the non-masked (valid) elements in each pooling window
+        valid_count = nn.functional.conv1d(
+            mask.float(),
+            ones_kernel,
+            stride=self.stride,
+            padding=self.padding,
+            groups=x.size(1),
+        )
+        valid_count = valid_count.clamp(min=1)  # Avoid division by zero
+
+        # Perform masked average pooling
+        avg_pooled = sum_pooled / valid_count
+
+        # Fill zero values with NaNs
+        avg_pooled[avg_pooled == 0] = float("nan")
+
+        if ndim == 2:
+            return avg_pooled.squeeze(1)
+
+        return avg_pooled
+
+
+class MaskedMedianPool1d(nn.Module):
+    def __init__(
+        self, kernel_size: int, stride: Optional[int] = None, padding: Optional[int] = 0
+    ):
+        """An implementation of median pooling that supports masked values.
+
+        This implementation is inspired by the median pooling implementation in
+        https://gist.github.com/rwightman/f2d3849281624be7c0f11c85c87c1598
+
+        Args:
+            kernel_size (int): The size of the median pooling window.
+            stride (int, optional): The stride of the median pooling window. Defaults to None.
+            padding (int, optional): The padding of the median pooling window. Defaults to 0.
+        """
+
+        super(MaskedMedianPool1d, self).__init__()
+        self.kernel_size = kernel_size
+        self.stride = stride or kernel_size
+        self.padding = padding
+
+    def forward(self, x, mask=None):
+        ndim = x.dim()
+        if ndim == 2:
+            x = x.unsqueeze(1)
+
+        assert (
+            x.dim() == 3
+        ), "Input tensor must have 2 or 3 dimensions (batch_size, channels, width)"
+
+        if mask is None:
+            mask = ~torch.isnan(x)
+
+        assert x.shape == mask.shape, "Input tensor and mask must have the same shape"
+
+        masked_x = torch.where(mask, x, torch.zeros_like(x))
+
+        x = F.pad(masked_x, (self.padding, self.padding), mode="reflect")
+        mask = F.pad(
+            mask.float(), (self.padding, self.padding), mode="constant", value=0
+        )
+
+        x = x.unfold(2, self.kernel_size, self.stride)
+        mask = mask.unfold(2, self.kernel_size, self.stride)
+
+        x = x.contiguous().view(x.size()[:3] + (-1,))
+        mask = mask.contiguous().view(mask.size()[:3] + (-1,)).to(x.device)
+
+        # Combine the mask with the input tensor
+        #x_masked = torch.where(mask.bool(), x, torch.fill_(torch.zeros_like(x),float("inf")))
+        x_masked = torch.where(mask.bool(), x, torch.FloatTensor([float("inf")]).to(x.device))
+
+        # Sort the masked tensor along the last dimension
+        x_sorted, _ = torch.sort(x_masked, dim=-1)
+
+        # Compute the count of non-masked (valid) values
+        valid_count = mask.sum(dim=-1)
+
+        # Calculate the index of the median value for each pooling window
+        median_idx = (torch.div((valid_count - 1), 2, rounding_mode='trunc')).clamp(min=0)
+
+        # Gather the median values using the calculated indices
+        median_pooled = x_sorted.gather(-1, median_idx.unsqueeze(-1).long()).squeeze(-1)
+
+        # Fill infinite values with NaNs
+        median_pooled[torch.isinf(median_pooled)] = float("nan")
+        
+        if ndim == 2:
+            return median_pooled.squeeze(1)
+
+        return median_pooled
+
+
+class CrepePitchExtractor(BasePitchExtractor):
+    def __init__(
+        self,
+        hop_length: int = 512,
+        f0_min: float = 50.0,
+        f0_max: float = 1100.0,
+        threshold: float = 0.05,
+        keep_zeros: bool = False,
+        device = None,
+        model: Literal["full", "tiny"] = "full",
+        use_fast_filters: bool = True,
+        decoder="viterbi"
+    ):
+        super().__init__(hop_length, f0_min, f0_max, keep_zeros)
+        if decoder == "viterbi":
+            self.decoder = torchcrepe.decode.viterbi
+        elif decoder == "argmax":
+            self.decoder = torchcrepe.decode.argmax
+        elif decoder == "weighted_argmax":
+            self.decoder = torchcrepe.decode.weighted_argmax
+        else:
+            raise "Unknown decoder"
+        self.threshold = threshold
+        self.model = model
+        self.use_fast_filters = use_fast_filters
+        self.hop_length = hop_length
+        if device is None:
+            self.dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        else:
+            self.dev = torch.device(device)
+        if self.use_fast_filters:
+            self.median_filter = MaskedMedianPool1d(3, 1, 1).to(device)
+            self.mean_filter = MaskedAvgPool1d(3, 1, 1).to(device)
+
+    def __call__(self, x, sampling_rate=44100, pad_to=None):
+        """Extract pitch using crepe.
+
+
+        Args:
+            x (torch.Tensor): Audio signal, shape (1, T).
+            sampling_rate (int, optional): Sampling rate. Defaults to 44100.
+            pad_to (int, optional): Pad to length. Defaults to None.
+
+        Returns:
+            torch.Tensor: Pitch, shape (T // hop_length,).
+        """
+
+        assert x.ndim == 2, f"Expected 2D tensor, got {x.ndim}D tensor."
+        assert x.shape[0] == 1, f"Expected 1 channel, got {x.shape[0]} channels."
+
+        x = x.to(self.dev)
+        f0, pd = torchcrepe.predict(
+            x,
+            sampling_rate,
+            self.hop_length,
+            self.f0_min,
+            self.f0_max,
+            pad=True,
+            model=self.model,
+            batch_size=1024,
+            device=x.device,
+            return_periodicity=True,
+            decoder=self.decoder
+        )
+
+        # Filter, remove silence, set uv threshold, refer to the original warehouse readme
+        if self.use_fast_filters:
+            pd = self.median_filter(pd)
+        else:
+            pd = torchcrepe.filter.median(pd, 3)
+
+        pd = torchcrepe.threshold.Silence(-60.0)(pd, x, sampling_rate, 512)
+        f0 = torchcrepe.threshold.At(self.threshold)(f0, pd)
+        
+        if self.use_fast_filters:
+            f0 = self.mean_filter(f0)
+        else:
+            f0 = torchcrepe.filter.mean(f0, 3)
+
+        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)
+            return rtn,rtn
+        
+        return self.post_process(x, sampling_rate, f0, pad_to)

modules/attentions.py

@@ -0,0 +1,349 @@
+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.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):
+    super().__init__()
+    self.hidden_channels = hidden_channels
+    self.filter_channels = filter_channels
+    self.n_heads = n_heads
+    self.n_layers = n_layers
+    self.kernel_size = kernel_size
+    self.p_dropout = p_dropout
+    self.proximal_bias = proximal_bias
+    self.proximal_init = proximal_init
+
+    self.drop = nn.Dropout(p_dropout)
+    self.self_attn_layers = nn.ModuleList()
+    self.norm_layers_0 = nn.ModuleList()
+    self.ffn_layers = nn.ModuleList()
+    self.norm_layers_1 = nn.ModuleList()
+    for i in range(self.n_layers):
+      self.self_attn_layers.append(
+        MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout, proximal_bias=proximal_bias,
+                           proximal_init=proximal_init))
+      self.norm_layers_0.append(LayerNorm(hidden_channels))
+      self.ffn_layers.append(
+        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):
+    """
+    x: decoder input
+    h: encoder output
+    """
+    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):
+      y = self.self_attn_layers[i](x, x, self_attn_mask)
+      y = self.drop(y)
+      x = self.norm_layers_0[i](x + y)
+
+      y = self.ffn_layers[i](x, x_mask)
+      y = self.drop(y)
+      x = self.norm_layers_1[i](x + y)
+    x = x * x_mask
+    return x
+
+
+class Encoder(nn.Module):
+  def __init__(self, hidden_channels, filter_channels, n_heads, n_layers, kernel_size=1, p_dropout=0., window_size=4, **kwargs):
+    super().__init__()
+    self.hidden_channels = hidden_channels
+    self.filter_channels = filter_channels
+    self.n_heads = n_heads
+    self.n_layers = n_layers
+    self.kernel_size = kernel_size
+    self.p_dropout = p_dropout
+    self.window_size = window_size
+
+    self.drop = nn.Dropout(p_dropout)
+    self.attn_layers = nn.ModuleList()
+    self.norm_layers_1 = nn.ModuleList()
+    self.ffn_layers = nn.ModuleList()
+    self.norm_layers_2 = nn.ModuleList()
+    for i in range(self.n_layers):
+      self.attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout, window_size=window_size))
+      self.norm_layers_1.append(LayerNorm(hidden_channels))
+      self.ffn_layers.append(FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout))
+      self.norm_layers_2.append(LayerNorm(hidden_channels))
+
+  def forward(self, x, x_mask):
+    attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+    x = x * x_mask
+    for i in range(self.n_layers):
+      y = self.attn_layers[i](x, x, attn_mask)
+      y = self.drop(y)
+      x = self.norm_layers_1[i](x + y)
+
+      y = self.ffn_layers[i](x, x_mask)
+      y = self.drop(y)
+      x = self.norm_layers_2[i](x + y)
+    x = x * x_mask
+    return x
+
+
+class Decoder(nn.Module):
+  def __init__(self, hidden_channels, filter_channels, n_heads, n_layers, kernel_size=1, p_dropout=0., proximal_bias=False, proximal_init=True, **kwargs):
+    super().__init__()
+    self.hidden_channels = hidden_channels
+    self.filter_channels = filter_channels
+    self.n_heads = n_heads
+    self.n_layers = n_layers
+    self.kernel_size = kernel_size
+    self.p_dropout = p_dropout
+    self.proximal_bias = proximal_bias
+    self.proximal_init = proximal_init
+
+    self.drop = nn.Dropout(p_dropout)
+    self.self_attn_layers = nn.ModuleList()
+    self.norm_layers_0 = nn.ModuleList()
+    self.encdec_attn_layers = nn.ModuleList()
+    self.norm_layers_1 = nn.ModuleList()
+    self.ffn_layers = nn.ModuleList()
+    self.norm_layers_2 = nn.ModuleList()
+    for i in range(self.n_layers):
+      self.self_attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout, proximal_bias=proximal_bias, proximal_init=proximal_init))
+      self.norm_layers_0.append(LayerNorm(hidden_channels))
+      self.encdec_attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout))
+      self.norm_layers_1.append(LayerNorm(hidden_channels))
+      self.ffn_layers.append(FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout, causal=True))
+      self.norm_layers_2.append(LayerNorm(hidden_channels))
+
+  def forward(self, x, x_mask, h, h_mask):
+    """
+    x: decoder input
+    h: encoder output
+    """
+    self_attn_mask = commons.subsequent_mask(x_mask.size(2)).to(device=x.device, dtype=x.dtype)
+    encdec_attn_mask = h_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+    x = x * x_mask
+    for i in range(self.n_layers):
+      y = self.self_attn_layers[i](x, x, self_attn_mask)
+      y = self.drop(y)
+      x = self.norm_layers_0[i](x + y)
+
+      y = self.encdec_attn_layers[i](x, h, encdec_attn_mask)
+      y = self.drop(y)
+      x = self.norm_layers_1[i](x + y)
+      
+      y = self.ffn_layers[i](x, x_mask)
+      y = self.drop(y)
+      x = self.norm_layers_2[i](x + y)
+    x = x * x_mask
+    return x
+
+
+class MultiHeadAttention(nn.Module):
+  def __init__(self, channels, out_channels, n_heads, p_dropout=0., window_size=None, heads_share=True, block_length=None, proximal_bias=False, proximal_init=False):
+    super().__init__()
+    assert channels % n_heads == 0
+
+    self.channels = channels
+    self.out_channels = out_channels
+    self.n_heads = n_heads
+    self.p_dropout = p_dropout
+    self.window_size = window_size
+    self.heads_share = heads_share
+    self.block_length = block_length
+    self.proximal_bias = proximal_bias
+    self.proximal_init = proximal_init
+    self.attn = None
+
+    self.k_channels = channels // n_heads
+    self.conv_q = nn.Conv1d(channels, channels, 1)
+    self.conv_k = nn.Conv1d(channels, channels, 1)
+    self.conv_v = nn.Conv1d(channels, channels, 1)
+    self.conv_o = nn.Conv1d(channels, out_channels, 1)
+    self.drop = nn.Dropout(p_dropout)
+
+    if window_size is not None:
+      n_heads_rel = 1 if heads_share else n_heads
+      rel_stddev = self.k_channels**-0.5
+      self.emb_rel_k = nn.Parameter(torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) * rel_stddev)
+      self.emb_rel_v = nn.Parameter(torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) * rel_stddev)
+
+    nn.init.xavier_uniform_(self.conv_q.weight)
+    nn.init.xavier_uniform_(self.conv_k.weight)
+    nn.init.xavier_uniform_(self.conv_v.weight)
+    if proximal_init:
+      with torch.no_grad():
+        self.conv_k.weight.copy_(self.conv_q.weight)
+        self.conv_k.bias.copy_(self.conv_q.bias)
+      
+  def forward(self, x, c, attn_mask=None):
+    q = self.conv_q(x)
+    k = self.conv_k(c)
+    v = self.conv_v(c)
+    
+    x, self.attn = self.attention(q, k, v, mask=attn_mask)
+
+    x = self.conv_o(x)
+    return x
+
+  def attention(self, query, key, value, mask=None):
+    # reshape [b, d, t] -> [b, n_h, t, d_k]
+    b, d, t_s, t_t = (*key.size(), query.size(2))
+    query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
+    key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+    value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+
+    scores = torch.matmul(query / math.sqrt(self.k_channels), key.transpose(-2, -1))
+    if self.window_size is not None:
+      assert t_s == t_t, "Relative attention is only available for self-attention."
+      key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)
+      rel_logits = self._matmul_with_relative_keys(query /math.sqrt(self.k_channels), key_relative_embeddings)
+      scores_local = self._relative_position_to_absolute_position(rel_logits)
+      scores = scores + scores_local
+    if self.proximal_bias:
+      assert t_s == t_t, "Proximal bias is only available for self-attention."
+      scores = scores + self._attention_bias_proximal(t_s).to(device=scores.device, dtype=scores.dtype)
+    if mask is not None:
+      scores = scores.masked_fill(mask == 0, -1e4)
+      if self.block_length is not None:
+        assert t_s == t_t, "Local attention is only available for self-attention."
+        block_mask = torch.ones_like(scores).triu(-self.block_length).tril(self.block_length)
+        scores = scores.masked_fill(block_mask == 0, -1e4)
+    p_attn = F.softmax(scores, dim=-1) # [b, n_h, t_t, t_s]
+    p_attn = self.drop(p_attn)
+    output = torch.matmul(p_attn, value)
+    if self.window_size is not None:
+      relative_weights = self._absolute_position_to_relative_position(p_attn)
+      value_relative_embeddings = self._get_relative_embeddings(self.emb_rel_v, t_s)
+      output = output + self._matmul_with_relative_values(relative_weights, value_relative_embeddings)
+    output = output.transpose(2, 3).contiguous().view(b, d, t_t) # [b, n_h, t_t, d_k] -> [b, d, t_t]
+    return output, p_attn
+
+  def _matmul_with_relative_values(self, x, y):
+    """
+    x: [b, h, l, m]
+    y: [h or 1, m, d]
+    ret: [b, h, l, d]
+    """
+    ret = torch.matmul(x, y.unsqueeze(0))
+    return ret
+
+  def _matmul_with_relative_keys(self, x, y):
+    """
+    x: [b, h, l, d]
+    y: [h or 1, m, d]
+    ret: [b, h, l, m]
+    """
+    ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))
+    return ret
+
+  def _get_relative_embeddings(self, relative_embeddings, length):
+    max_relative_position = 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)
+    slice_end_position = slice_start_position + 2 * length - 1
+    if pad_length > 0:
+      padded_relative_embeddings = F.pad(
+          relative_embeddings,
+          commons.convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]))
+    else:
+      padded_relative_embeddings = relative_embeddings
+    used_relative_embeddings = padded_relative_embeddings[:,slice_start_position:slice_end_position]
+    return used_relative_embeddings
+
+  def _relative_position_to_absolute_position(self, x):
+    """
+    x: [b, h, l, 2*l-1]
+    ret: [b, h, l, l]
+    """
+    batch, heads, length, _ = x.size()
+    # Concat columns of pad to shift from relative to absolute indexing.
+    x = F.pad(x, commons.convert_pad_shape([[0,0],[0,0],[0,0],[0,1]]))
+
+    # Concat extra elements so to add up to shape (len+1, 2*len-1).
+    x_flat = x.view([batch, heads, length * 2 * length])
+    x_flat = F.pad(x_flat, commons.convert_pad_shape([[0,0],[0,0],[0,length-1]]))
+
+    # Reshape and slice out the padded elements.
+    x_final = x_flat.view([batch, heads, length+1, 2*length-1])[:, :, :length, length-1:]
+    return x_final
+
+  def _absolute_position_to_relative_position(self, x):
+    """
+    x: [b, h, l, l]
+    ret: [b, h, l, 2*l-1]
+    """
+    batch, heads, length, _ = x.size()
+    # padd along column
+    x = F.pad(x, commons.convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, length-1]]))
+    x_flat = x.view([batch, heads, length**2 + length*(length -1)])
+    # add 0's in the beginning that will skew the elements after reshape
+    x_flat = F.pad(x_flat, commons.convert_pad_shape([[0, 0], [0, 0], [length, 0]]))
+    x_final = x_flat.view([batch, heads, length, 2*length])[:,:,:,1:]
+    return x_final
+
+  def _attention_bias_proximal(self, length):
+    """Bias for self-attention to encourage attention to close positions.
+    Args:
+      length: an integer scalar.
+    Returns:
+      a Tensor with shape [1, 1, length, length]
+    """
+    r = torch.arange(length, dtype=torch.float32)
+    diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)
+    return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0)
+
+
+class FFN(nn.Module):
+  def __init__(self, in_channels, out_channels, filter_channels, kernel_size, p_dropout=0., activation=None, causal=False):
+    super().__init__()
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+    self.filter_channels = filter_channels
+    self.kernel_size = kernel_size
+    self.p_dropout = p_dropout
+    self.activation = activation
+    self.causal = causal
+
+    if causal:
+      self.padding = self._causal_padding
+    else:
+      self.padding = self._same_padding
+
+    self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size)
+    self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size)
+    self.drop = nn.Dropout(p_dropout)
+
+  def forward(self, x, x_mask):
+    x = self.conv_1(self.padding(x * x_mask))
+    if self.activation == "gelu":
+      x = x * torch.sigmoid(1.702 * x)
+    else:
+      x = torch.relu(x)
+    x = self.drop(x)
+    x = self.conv_2(self.padding(x * x_mask))
+    return x * x_mask
+  
+  def _causal_padding(self, x):
+    if self.kernel_size == 1:
+      return x
+    pad_l = self.kernel_size - 1
+    pad_r = 0
+    padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+    x = F.pad(x, commons.convert_pad_shape(padding))
+    return x
+
+  def _same_padding(self, x):
+    if self.kernel_size == 1:
+      return x
+    pad_l = (self.kernel_size - 1) // 2
+    pad_r = self.kernel_size // 2
+    padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+    x = F.pad(x, commons.convert_pad_shape(padding))
+    return x

modules/commons.py

@@ -0,0 +1,188 @@
+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)):
+    idx_str = ids_str[i]
+    idx_end = idx_str + segment_size
+    ret[i] = x[i, idx_str:idx_end]
+  return ret
+
+def rand_slice_segments_with_pitch(x, pitch, x_lengths=None, segment_size=4):
+  b, d, t = x.size()
+  if x_lengths is None:
+    x_lengths = t
+  ids_str_max = x_lengths - segment_size + 1
+  ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
+  ret = slice_segments(x, ids_str, segment_size)
+  ret_pitch = slice_pitch_segments(pitch, ids_str, segment_size)
+  return ret, ret_pitch, ids_str
+
+def init_weights(m, mean=0.0, std=0.01):
+  classname = m.__class__.__name__
+  if classname.find("Conv") != -1:
+    m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size, dilation=1):
+  return int((kernel_size*dilation - dilation)/2)
+
+
+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 intersperse(lst, item):
+  result = [item] * (len(lst) * 2 + 1)
+  result[1::2] = lst
+  return result
+
+
+def kl_divergence(m_p, logs_p, m_q, logs_q):
+  """KL(P||Q)"""
+  kl = (logs_q - logs_p) - 0.5
+  kl += 0.5 * (torch.exp(2. * logs_p) + ((m_p - m_q)**2)) * torch.exp(-2. * logs_q)
+  return kl
+
+
+def rand_gumbel(shape):
+  """Sample from the Gumbel distribution, protect from overflows."""
+  uniform_samples = torch.rand(shape) * 0.99998 + 0.00001
+  return -torch.log(-torch.log(uniform_samples))
+
+
+def rand_gumbel_like(x):
+  g = rand_gumbel(x.size()).to(dtype=x.dtype, device=x.device)
+  return g
+
+
+def slice_segments(x, ids_str, segment_size=4):
+  ret = torch.zeros_like(x[:, :, :segment_size])
+  for i in range(x.size(0)):
+    idx_str = ids_str[i]
+    idx_end = idx_str + segment_size
+    ret[i] = x[i, :, idx_str:idx_end]
+  return ret
+
+
+def rand_slice_segments(x, x_lengths=None, segment_size=4):
+  b, d, t = x.size()
+  if x_lengths is None:
+    x_lengths = t
+  ids_str_max = x_lengths - segment_size + 1
+  ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
+  ret = slice_segments(x, ids_str, segment_size)
+  return ret, ids_str
+
+
+def rand_spec_segments(x, x_lengths=None, segment_size=4):
+  b, d, t = x.size()
+  if x_lengths is None:
+    x_lengths = t
+  ids_str_max = x_lengths - segment_size
+  ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
+  ret = slice_segments(x, ids_str, segment_size)
+  return ret, ids_str
+
+
+def get_timing_signal_1d(
+    length, channels, min_timescale=1.0, max_timescale=1.0e4):
+  position = torch.arange(length, dtype=torch.float)
+  num_timescales = channels // 2
+  log_timescale_increment = (
+      math.log(float(max_timescale) / float(min_timescale)) /
+      (num_timescales - 1))
+  inv_timescales = min_timescale * torch.exp(
+      torch.arange(num_timescales, dtype=torch.float) * -log_timescale_increment)
+  scaled_time = position.unsqueeze(0) * inv_timescales.unsqueeze(1)
+  signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], 0)
+  signal = F.pad(signal, [0, 0, 0, channels % 2])
+  signal = signal.view(1, channels, length)
+  return signal
+
+
+def add_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4):
+  b, channels, length = x.size()
+  signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
+  return x + signal.to(dtype=x.dtype, device=x.device)
+
+
+def cat_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4, axis=1):
+  b, channels, length = x.size()
+  signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
+  return torch.cat([x, signal.to(dtype=x.dtype, device=x.device)], axis)
+
+
+def subsequent_mask(length):
+  mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0)
+  return mask
+
+
+@torch.jit.script
+def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
+  n_channels_int = n_channels[0]
+  in_act = input_a + input_b
+  t_act = torch.tanh(in_act[:, :n_channels_int, :])
+  s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+  acts = t_act * s_act
+  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
+
+
+def sequence_mask(length, max_length=None):
+  if max_length is None:
+    max_length = length.max()
+  x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+  return x.unsqueeze(0) < length.unsqueeze(1)
+
+
+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)
+  
+  cum_duration_flat = cum_duration.view(b * t_x)
+  path = sequence_mask(cum_duration_flat, t_y).to(mask.dtype)
+  path = path.view(b, t_x, t_y)
+  path = path - F.pad(path, convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
+  path = path.unsqueeze(1).transpose(2,3) * mask
+  return path
+
+
+def clip_grad_value_(parameters, clip_value, norm_type=2):
+  if isinstance(parameters, torch.Tensor):
+    parameters = [parameters]
+  parameters = list(filter(lambda p: p.grad is not None, parameters))
+  norm_type = float(norm_type)
+  if clip_value is not None:
+    clip_value = float(clip_value)
+
+  total_norm = 0
+  for p in parameters:
+    param_norm = p.grad.data.norm(norm_type)
+    total_norm += param_norm.item() ** norm_type
+    if clip_value is not None:
+      p.grad.data.clamp_(min=-clip_value, max=clip_value)
+  total_norm = total_norm ** (1. / norm_type)
+  return total_norm

modules/enhancer.py

@@ -0,0 +1,105 @@
+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
+
+class Enhancer:
+    def __init__(self, enhancer_type, enhancer_ckpt, device=None):
+        if device is None:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.device = device
+        
+        if enhancer_type == 'nsf-hifigan':
+            self.enhancer = NsfHifiGAN(enhancer_ckpt, device=self.device)
+        else:
+            raise ValueError(f" [x] Unknown enhancer: {enhancer_type}")
+        
+        self.resample_kernel = {}
+        self.enhancer_sample_rate = self.enhancer.sample_rate()
+        self.enhancer_hop_size = self.enhancer.hop_size()
+        
+    def enhance(self,
+                audio, # 1, T
+                sample_rate,
+                f0, # 1, n_frames, 1
+                hop_size,
+                adaptive_key = 0,
+                silence_front = 0
+                ):
+        # enhancer start time 
+        start_frame = int(silence_front * sample_rate / hop_size)
+        real_silence_front = start_frame * hop_size / sample_rate
+        audio = audio[:, int(np.round(real_silence_front * sample_rate)) : ]
+        f0 = f0[: , start_frame :, :]
+        
+        # adaptive parameters
+        adaptive_factor = 2 ** ( -adaptive_key / 12)
+        adaptive_sample_rate = 100 * int(np.round(self.enhancer_sample_rate / adaptive_factor / 100))
+        real_factor = self.enhancer_sample_rate / adaptive_sample_rate
+        
+        # resample the ddsp output
+        if sample_rate == adaptive_sample_rate:
+            audio_res = audio
+        else:
+            key_str = str(sample_rate) + str(adaptive_sample_rate)
+            if key_str not in self.resample_kernel:
+                self.resample_kernel[key_str] = Resample(sample_rate, adaptive_sample_rate, lowpass_filter_width = 128).to(self.device)
+            audio_res = self.resample_kernel[key_str](audio)
+        
+        n_frames = int(audio_res.size(-1) // self.enhancer_hop_size + 1)
+        
+        # resample f0
+        f0_np = f0.squeeze(0).squeeze(-1).cpu().numpy()
+        f0_np *= real_factor
+        time_org = (hop_size / sample_rate) * np.arange(len(f0_np)) / real_factor
+        time_frame = (self.enhancer_hop_size / self.enhancer_sample_rate) * np.arange(n_frames)
+        f0_res = np.interp(time_frame, time_org, f0_np, left=f0_np[0], right=f0_np[-1])
+        f0_res = torch.from_numpy(f0_res).unsqueeze(0).float().to(self.device) # 1, n_frames
+
+        # enhance
+        enhanced_audio, enhancer_sample_rate = self.enhancer(audio_res, f0_res)
+        
+        # resample the enhanced output
+        if adaptive_factor != 0:
+            key_str = str(adaptive_sample_rate) + str(enhancer_sample_rate)
+            if key_str not in self.resample_kernel:
+                self.resample_kernel[key_str] = Resample(adaptive_sample_rate, enhancer_sample_rate, lowpass_filter_width = 128).to(self.device)
+            enhanced_audio =  self.resample_kernel[key_str](enhanced_audio)
+        
+        # pad the silence frames
+        if start_frame > 0:
+            enhanced_audio = F.pad(enhanced_audio, (int(np.round(enhancer_sample_rate * real_silence_front)), 0))
+            
+        return enhanced_audio, enhancer_sample_rate
+        
+        
+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
+        print('| Load HifiGAN: ', model_path)
+        self.model, self.h = load_model(model_path, device=self.device)
+    
+    def sample_rate(self):
+        return self.h.sampling_rate
+        
+    def hop_size(self):
+        return self.h.hop_size
+        
+    def forward(self, audio, f0):
+        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)
+        with torch.no_grad():
+            mel = stft.get_mel(audio)
+            enhanced_audio = self.model(mel, f0[:,:mel.size(-1)]).view(-1)
+            return enhanced_audio, self.h.sampling_rate

modules/losses.py

@@ -0,0 +1,61 @@
+import torch 
+from torch.nn import functional as F
+
+import modules.commons as commons
+
+
+def feature_loss(fmap_r, fmap_g):
+  loss = 0
+  for dr, dg in zip(fmap_r, fmap_g):
+    for rl, gl in zip(dr, dg):
+      rl = rl.float().detach()
+      gl = gl.float()
+      loss += torch.mean(torch.abs(rl - gl))
+
+  return loss * 2 
+
+
+def discriminator_loss(disc_real_outputs, disc_generated_outputs):
+  loss = 0
+  r_losses = []
+  g_losses = []
+  for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
+    dr = dr.float()
+    dg = dg.float()
+    r_loss = torch.mean((1-dr)**2)
+    g_loss = torch.mean(dg**2)
+    loss += (r_loss + g_loss)
+    r_losses.append(r_loss.item())
+    g_losses.append(g_loss.item())
+
+  return loss, r_losses, g_losses
+
+
+def generator_loss(disc_outputs):
+  loss = 0
+  gen_losses = []
+  for dg in disc_outputs:
+    dg = dg.float()
+    l = torch.mean((1-dg)**2)
+    gen_losses.append(l)
+    loss += l
+
+  return loss, gen_losses
+
+
+def kl_loss(z_p, logs_q, m_p, logs_p, z_mask):
+  """
+  z_p, logs_q: [b, h, t_t]
+  m_p, logs_p: [b, h, t_t]
+  """
+  z_p = z_p.float()
+  logs_q = logs_q.float()
+  m_p = m_p.float()
+  logs_p = logs_p.float()
+  z_mask = z_mask.float()
+  #print(logs_p)
+  kl = logs_p - logs_q - 0.5
+  kl += 0.5 * ((z_p - m_p)**2) * torch.exp(-2. * logs_p)
+  kl = torch.sum(kl * z_mask)
+  l = kl / torch.sum(z_mask)
+  return l

modules/mel_processing.py

@@ -0,0 +1,112 @@
+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
+
+
+def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
+    """
+    PARAMS
+    ------
+    C: compression factor
+    """
+    return torch.log(torch.clamp(x, min=clip_val) * C)
+
+
+def dynamic_range_decompression_torch(x, C=1):
+    """
+    PARAMS
+    ------
+    C: compression factor used to compress
+    """
+    return torch.exp(x) / C
+
+
+def spectral_normalize_torch(magnitudes):
+    output = dynamic_range_compression_torch(magnitudes)
+    return output
+
+
+def spectral_de_normalize_torch(magnitudes):
+    output = dynamic_range_decompression_torch(magnitudes)
+    return output
+
+
+mel_basis = {}
+hann_window = {}
+
+
+def spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, 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 hann_window
+    dtype_device = str(y.dtype) + '_' + str(y.device)
+    wnsize_dtype_device = str(win_size) + '_' + dtype_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)
+    return spec
+
+
+def spec_to_mel_torch(spec, n_fft, num_mels, sampling_rate, fmin, fmax):
+    global mel_basis
+    dtype_device = str(spec.dtype) + '_' + str(spec.device)
+    fmax_dtype_device = str(fmax) + '_' + 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=spec.dtype, device=spec.device)
+    spec = torch.matmul(mel_basis[fmax_dtype_device], spec)
+    spec = spectral_normalize_torch(spec)
+    return spec
+
+
+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)
+
+    return spec

modules/modules.py

@@ -0,0 +1,342 @@
+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.commons as commons
+from modules.commons import init_weights, get_padding
+
+
+LRELU_SLOPE = 0.1
+
+
+class LayerNorm(nn.Module):
+  def __init__(self, channels, eps=1e-5):
+    super().__init__()
+    self.channels = channels
+    self.eps = eps
+
+    self.gamma = nn.Parameter(torch.ones(channels))
+    self.beta = nn.Parameter(torch.zeros(channels))
+
+  def forward(self, x):
+    x = x.transpose(1, -1)
+    x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+    return x.transpose(1, -1)
+
+ 
+class ConvReluNorm(nn.Module):
+  def __init__(self, in_channels, hidden_channels, out_channels, kernel_size, n_layers, p_dropout):
+    super().__init__()
+    self.in_channels = in_channels
+    self.hidden_channels = hidden_channels
+    self.out_channels = out_channels
+    self.kernel_size = kernel_size
+    self.n_layers = n_layers
+    self.p_dropout = p_dropout
+    assert n_layers > 1, "Number of layers should be larger than 0."
+
+    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.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.norm_layers.append(LayerNorm(hidden_channels))
+    self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
+    self.proj.weight.data.zero_()
+    self.proj.bias.data.zero_()
+
+  def forward(self, x, x_mask):
+    x_org = x
+    for i in range(self.n_layers):
+      x = self.conv_layers[i](x * x_mask)
+      x = self.norm_layers[i](x)
+      x = self.relu_drop(x)
+    x = x_org + self.proj(x)
+    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__()
+    assert(kernel_size % 2 == 1)
+    self.hidden_channels =hidden_channels
+    self.kernel_size = kernel_size,
+    self.dilation_rate = dilation_rate
+    self.n_layers = n_layers
+    self.gin_channels = gin_channels
+    self.p_dropout = p_dropout
+
+    self.in_layers = torch.nn.ModuleList()
+    self.res_skip_layers = torch.nn.ModuleList()
+    self.drop = nn.Dropout(p_dropout)
+
+    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')
+
+    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,
+                                 dilation=dilation, padding=padding)
+      in_layer = torch.nn.utils.weight_norm(in_layer, name='weight')
+      self.in_layers.append(in_layer)
+
+      # last one is not necessary
+      if i < n_layers - 1:
+        res_skip_channels = 2 * hidden_channels
+      else:
+        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')
+      self.res_skip_layers.append(res_skip_layer)
+
+  def forward(self, x, x_mask, g=None, **kwargs):
+    output = torch.zeros_like(x)
+    n_channels_tensor = torch.IntTensor([self.hidden_channels])
+
+    if g is not None:
+      g = self.cond_layer(g)
+
+    for i in range(self.n_layers):
+      x_in = self.in_layers[i](x)
+      if g is not None:
+        cond_offset = i * 2 * self.hidden_channels
+        g_l = g[:,cond_offset:cond_offset+2*self.hidden_channels,:]
+      else:
+        g_l = torch.zeros_like(x_in)
+
+      acts = commons.fused_add_tanh_sigmoid_multiply(
+          x_in,
+          g_l,
+          n_channels_tensor)
+      acts = self.drop(acts)
+
+      res_skip_acts = self.res_skip_layers[i](acts)
+      if i < self.n_layers - 1:
+        res_acts = res_skip_acts[:,:self.hidden_channels,:]
+        x = (x + res_acts) * x_mask
+        output = output + res_skip_acts[:,self.hidden_channels:,:]
+      else:
+        output = output + res_skip_acts
+    return output * x_mask
+
+  def remove_weight_norm(self):
+    if self.gin_channels != 0:
+      torch.nn.utils.remove_weight_norm(self.cond_layer)
+    for l in self.in_layers:
+      torch.nn.utils.remove_weight_norm(l)
+    for l in self.res_skip_layers:
+     torch.nn.utils.remove_weight_norm(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],
+                               padding=get_padding(kernel_size, dilation[0]))),
+            weight_norm(Conv1d(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],
+                               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,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1)))
+        ])
+        self.convs2.apply(init_weights)
+
+    def forward(self, x, x_mask=None):
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c1(xt)
+            xt = F.leaky_relu(xt, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c2(xt)
+            x = xt + x
+        if x_mask is not None:
+            x = x * x_mask
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(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],
+                               padding=get_padding(kernel_size, dilation[0]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
+                               padding=get_padding(kernel_size, dilation[1])))
+        ])
+        self.convs.apply(init_weights)
+
+    def forward(self, x, x_mask=None):
+        for c in self.convs:
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c(xt)
+            x = xt + x
+        if x_mask is not None:
+            x = x * x_mask
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs:
+            remove_weight_norm(l)
+
+
+class Log(nn.Module):
+  def forward(self, x, x_mask, reverse=False, **kwargs):
+    if not reverse:
+      y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask
+      logdet = torch.sum(-y, [1, 2])
+      return y, logdet
+    else:
+      x = torch.exp(x) * x_mask
+      return x
+    
+
+class Flip(nn.Module):
+  def forward(self, x, *args, reverse=False, **kwargs):
+    x = torch.flip(x, [1])
+    if not reverse:
+      logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)
+      return x, logdet
+    else:
+      return x
+
+
+class ElementwiseAffine(nn.Module):
+  def __init__(self, channels):
+    super().__init__()
+    self.channels = channels
+    self.m = nn.Parameter(torch.zeros(channels,1))
+    self.logs = nn.Parameter(torch.zeros(channels,1))
+
+  def forward(self, x, x_mask, reverse=False, **kwargs):
+    if not reverse:
+      y = self.m + torch.exp(self.logs) * x
+      y = y * x_mask
+      logdet = torch.sum(self.logs * x_mask, [1,2])
+      return y, logdet
+    else:
+      x = (x - self.m) * torch.exp(-self.logs) * x_mask
+      return x
+
+
+class ResidualCouplingLayer(nn.Module):
+  def __init__(self,
+      channels,
+      hidden_channels,
+      kernel_size,
+      dilation_rate,
+      n_layers,
+      p_dropout=0,
+      gin_channels=0,
+      mean_only=False):
+    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.dilation_rate = dilation_rate
+    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 = WN(hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout=p_dropout, gin_channels=gin_channels)
+    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

requirements.txt

@@ -0,0 +1,26 @@
+ffmpeg-python
+Flask
+Flask_Cors
+gradio>=3.7.0
+numpy==1.23.0
+pyworld==0.2.5
+scipy==1.10.0
+SoundFile==0.12.1
+torch==1.13.1
+torchaudio==0.13.1
+torchcrepe
+tqdm
+scikit-maad
+praat-parselmouth
+onnx
+onnxsim
+onnxoptimizer
+fairseq==0.12.2
+librosa==0.9.1
+tensorboard
+tensorboardX
+transformers
+edge_tts
+pyyaml
+pynvml
+ffmpeg

utils.py

@@ -0,0 +1,446 @@
+import os
+import glob
+import re
+import sys
+import argparse
+import logging
+import json
+import subprocess
+import warnings
+import random
+import functools
+
+import librosa
+import numpy as np
+from scipy.io.wavfile import read
+import torch
+from torch.nn import functional as F
+from modules.commons import sequence_mask
+
+MATPLOTLIB_FLAG = False
+
+logging.basicConfig(stream=sys.stdout, level=logging.WARN)
+logger = logging
+
+f0_bin = 256
+f0_max = 1100.0
+f0_min = 50.0
+f0_mel_min = 1127 * np.log(1 + f0_min / 700)
+f0_mel_max = 1127 * np.log(1 + f0_max / 700)
+
+def normalize_f0(f0, x_mask, uv, random_scale=True):
+    # calculate means based on x_mask
+    uv_sum = torch.sum(uv, dim=1, keepdim=True)
+    uv_sum[uv_sum == 0] = 9999
+    means = torch.sum(f0[:, 0, :] * uv, dim=1, keepdim=True) / uv_sum
+
+    if random_scale:
+        factor = torch.Tensor(f0.shape[0], 1).uniform_(0.8, 1.2).to(f0.device)
+    else:
+        factor = torch.ones(f0.shape[0], 1).to(f0.device)
+    # normalize f0 based on means and factor
+    f0_norm = (f0 - means.unsqueeze(-1)) * factor.unsqueeze(-1)
+    if torch.isnan(f0_norm).any():
+        exit(0)
+    return f0_norm * x_mask
+
+def plot_data_to_numpy(x, y):
+    global MATPLOTLIB_FLAG
+    if not MATPLOTLIB_FLAG:
+        import matplotlib
+        matplotlib.use("Agg")
+        MATPLOTLIB_FLAG = True
+        mpl_logger = logging.getLogger('matplotlib')
+        mpl_logger.setLevel(logging.WARNING)
+    import matplotlib.pylab as plt
+    import numpy as np
+
+    fig, ax = plt.subplots(figsize=(10, 2))
+    plt.plot(x)
+    plt.plot(y)
+    plt.tight_layout()
+
+    fig.canvas.draw()
+    data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
+    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    plt.close()
+    return data
+
+
+def f0_to_coarse(f0):
+  is_torch = isinstance(f0, torch.Tensor)
+  f0_mel = 1127 * (1 + f0 / 700).log() if is_torch else 1127 * np.log(1 + f0 / 700)
+  f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - f0_mel_min) * (f0_bin - 2) / (f0_mel_max - f0_mel_min) + 1
+
+  f0_mel[f0_mel <= 1] = 1
+  f0_mel[f0_mel > f0_bin - 1] = f0_bin - 1
+  f0_coarse = (f0_mel + 0.5).int() if is_torch else np.rint(f0_mel).astype(np.int)
+  assert f0_coarse.max() <= 255 and f0_coarse.min() >= 1, (f0_coarse.max(), f0_coarse.min())
+  return f0_coarse
+
+def get_content(cmodel, y):
+    with torch.no_grad():
+        c = cmodel.extract_features(y.squeeze(1))[0]
+    c = c.transpose(1, 2)
+    return c
+
+def get_f0_predictor(f0_predictor,hop_length,sampling_rate,**kargs):
+    if f0_predictor == "pm":
+        from modules.F0Predictor.PMF0Predictor import PMF0Predictor
+        f0_predictor_object = PMF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate)
+    elif f0_predictor == "crepe":
+        from modules.F0Predictor.CrepeF0Predictor import CrepeF0Predictor
+        f0_predictor_object = CrepeF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate,device=kargs["device"],threshold=kargs["threshold"])
+    elif f0_predictor == "harvest":
+        from modules.F0Predictor.HarvestF0Predictor import HarvestF0Predictor
+        f0_predictor_object = HarvestF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate)
+    elif f0_predictor == "dio":
+        from modules.F0Predictor.DioF0Predictor import DioF0Predictor
+        f0_predictor_object = DioF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate)
+    else:
+        raise Exception("Unknown f0 predictor")
+    return f0_predictor_object
+
+def get_speech_encoder(speech_encoder,device=None,**kargs):
+    if speech_encoder == "vec768l12":
+        from vencoder.ContentVec768L12 import ContentVec768L12
+        speech_encoder_object = ContentVec768L12(device = device)
+    elif speech_encoder == "vec256l9":
+        from vencoder.ContentVec256L9 import ContentVec256L9
+        speech_encoder_object = ContentVec256L9(device = device)
+    elif speech_encoder == "vec256l9-onnx":
+        from vencoder.ContentVec256L9_Onnx import ContentVec256L9_Onnx
+        speech_encoder_object = ContentVec256L9(device = device)
+    elif speech_encoder == "vec256l12-onnx":
+        from vencoder.ContentVec256L12_Onnx import ContentVec256L12_Onnx
+        speech_encoder_object = ContentVec256L9(device = device)
+    elif speech_encoder == "vec768l9-onnx":
+        from vencoder.ContentVec768L9_Onnx import ContentVec768L9_Onnx
+        speech_encoder_object = ContentVec256L9(device = device)
+    elif speech_encoder == "vec768l12-onnx":
+        from vencoder.ContentVec768L12_Onnx import ContentVec768L12_Onnx
+        speech_encoder_object = ContentVec256L9(device = device)
+    elif speech_encoder == "hubertsoft-onnx":
+        from vencoder.HubertSoft_Onnx import HubertSoft_Onnx
+        speech_encoder_object = HubertSoft(device = device)
+    elif speech_encoder == "hubertsoft":
+        from vencoder.HubertSoft import HubertSoft
+        speech_encoder_object = HubertSoft(device = device)
+    elif speech_encoder == "whisper-ppg":
+        from vencoder.WhisperPPG import WhisperPPG
+        speech_encoder_object = WhisperPPG(device = device)
+    else:
+        raise Exception("Unknown speech encoder")
+    return speech_encoder_object 
+
+def load_checkpoint(checkpoint_path, model, optimizer=None, skip_optimizer=False):
+    assert os.path.isfile(checkpoint_path)
+    checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')
+    iteration = checkpoint_dict['iteration']
+    learning_rate = checkpoint_dict['learning_rate']
+    if optimizer is not None and not skip_optimizer and checkpoint_dict['optimizer'] is not None:
+        optimizer.load_state_dict(checkpoint_dict['optimizer'])
+    saved_state_dict = checkpoint_dict['model']
+    if hasattr(model, 'module'):
+        state_dict = model.module.state_dict()
+    else:
+        state_dict = model.state_dict()
+    new_state_dict = {}
+    for k, v in state_dict.items():
+        try:
+            # assert "dec" in k or "disc" in k
+            # print("load", k)
+            new_state_dict[k] = saved_state_dict[k]
+            assert saved_state_dict[k].shape == v.shape, (saved_state_dict[k].shape, v.shape)
+        except:
+            print("error, %s is not in the checkpoint" % k)
+            logger.info("%s is not in the checkpoint" % k)
+            new_state_dict[k] = v
+    if hasattr(model, 'module'):
+        model.module.load_state_dict(new_state_dict)
+    else:
+        model.load_state_dict(new_state_dict)
+    print("load ")
+    logger.info("Loaded checkpoint '{}' (iteration {})".format(
+        checkpoint_path, iteration))
+    return model, optimizer, learning_rate, iteration
+
+
+def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path):
+  logger.info("Saving model and optimizer state at iteration {} to {}".format(
+    iteration, checkpoint_path))
+  if hasattr(model, 'module'):
+    state_dict = model.module.state_dict()
+  else:
+    state_dict = model.state_dict()
+  torch.save({'model': state_dict,
+              'iteration': iteration,
+              'optimizer': optimizer.state_dict(),
+              'learning_rate': learning_rate}, checkpoint_path)
+
+def clean_checkpoints(path_to_models='logs/44k/', n_ckpts_to_keep=2, sort_by_time=True):
+  """Freeing up space by deleting saved ckpts
+
+  Arguments:
+  path_to_models    --  Path to the model directory
+  n_ckpts_to_keep   --  Number of ckpts to keep, excluding G_0.pth and D_0.pth
+  sort_by_time      --  True -> chronologically delete ckpts
+                        False -> lexicographically delete ckpts
+  """
+  ckpts_files = [f for f in os.listdir(path_to_models) if os.path.isfile(os.path.join(path_to_models, f))]
+  name_key = (lambda _f: int(re.compile('._(\d+)\.pth').match(_f).group(1)))
+  time_key = (lambda _f: os.path.getmtime(os.path.join(path_to_models, _f)))
+  sort_key = time_key if sort_by_time else name_key
+  x_sorted = lambda _x: sorted([f for f in ckpts_files if f.startswith(_x) and not f.endswith('_0.pth')], key=sort_key)
+  to_del = [os.path.join(path_to_models, fn) for fn in
+            (x_sorted('G')[:-n_ckpts_to_keep] + x_sorted('D')[:-n_ckpts_to_keep])]
+  del_info = lambda fn: logger.info(f".. Free up space by deleting ckpt {fn}")
+  del_routine = lambda x: [os.remove(x), del_info(x)]
+  rs = [del_routine(fn) for fn in to_del]
+
+def summarize(writer, global_step, scalars={}, histograms={}, images={}, audios={}, audio_sampling_rate=22050):
+  for k, v in scalars.items():
+    writer.add_scalar(k, v, global_step)
+  for k, v in histograms.items():
+    writer.add_histogram(k, v, global_step)
+  for k, v in images.items():
+    writer.add_image(k, v, global_step, dataformats='HWC')
+  for k, v in audios.items():
+    writer.add_audio(k, v, global_step, audio_sampling_rate)
+
+
+def latest_checkpoint_path(dir_path, regex="G_*.pth"):
+  f_list = glob.glob(os.path.join(dir_path, regex))
+  f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f))))
+  x = f_list[-1]
+  print(x)
+  return x
+
+
+def plot_spectrogram_to_numpy(spectrogram):
+  global MATPLOTLIB_FLAG
+  if not MATPLOTLIB_FLAG:
+    import matplotlib
+    matplotlib.use("Agg")
+    MATPLOTLIB_FLAG = True
+    mpl_logger = logging.getLogger('matplotlib')
+    mpl_logger.setLevel(logging.WARNING)
+  import matplotlib.pylab as plt
+  import numpy as np
+
+  fig, ax = plt.subplots(figsize=(10,2))
+  im = ax.imshow(spectrogram, aspect="auto", origin="lower",
+                  interpolation='none')
+  plt.colorbar(im, ax=ax)
+  plt.xlabel("Frames")
+  plt.ylabel("Channels")
+  plt.tight_layout()
+
+  fig.canvas.draw()
+  data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
+  data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+  plt.close()
+  return data
+
+
+def plot_alignment_to_numpy(alignment, info=None):
+  global MATPLOTLIB_FLAG
+  if not MATPLOTLIB_FLAG:
+    import matplotlib
+    matplotlib.use("Agg")
+    MATPLOTLIB_FLAG = True
+    mpl_logger = logging.getLogger('matplotlib')
+    mpl_logger.setLevel(logging.WARNING)
+  import matplotlib.pylab as plt
+  import numpy as np
+
+  fig, ax = plt.subplots(figsize=(6, 4))
+  im = ax.imshow(alignment.transpose(), aspect='auto', origin='lower',
+                  interpolation='none')
+  fig.colorbar(im, ax=ax)
+  xlabel = 'Decoder timestep'
+  if info is not None:
+      xlabel += '\n\n' + info
+  plt.xlabel(xlabel)
+  plt.ylabel('Encoder timestep')
+  plt.tight_layout()
+
+  fig.canvas.draw()
+  data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
+  data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+  plt.close()
+  return data
+
+
+def load_wav_to_torch(full_path):
+  sampling_rate, data = read(full_path)
+  return torch.FloatTensor(data.astype(np.float32)), sampling_rate
+
+
+def load_filepaths_and_text(filename, split="|"):
+  with open(filename, encoding='utf-8') as f:
+    filepaths_and_text = [line.strip().split(split) for line in f]
+  return filepaths_and_text
+
+
+def get_hparams(init=True):
+  parser = argparse.ArgumentParser()
+  parser.add_argument('-c', '--config', type=str, default="./configs/config.json",
+                      help='JSON file for configuration')
+  parser.add_argument('-m', '--model', type=str, required=True,
+                      help='Model name')
+
+  args = parser.parse_args()
+  model_dir = os.path.join("./logs", args.model)
+
+  if not os.path.exists(model_dir):
+    os.makedirs(model_dir)
+
+  config_path = args.config
+  config_save_path = os.path.join(model_dir, "config.json")
+  if init:
+    with open(config_path, "r") as f:
+      data = f.read()
+    with open(config_save_path, "w") as f:
+      f.write(data)
+  else:
+    with open(config_save_path, "r") as f:
+      data = f.read()
+  config = json.loads(data)
+
+  hparams = HParams(**config)
+  hparams.model_dir = model_dir
+  return hparams
+
+
+def get_hparams_from_dir(model_dir):
+  config_save_path = os.path.join(model_dir, "config.json")
+  with open(config_save_path, "r") as f:
+    data = f.read()
+  config = json.loads(data)
+
+  hparams =HParams(**config)
+  hparams.model_dir = model_dir
+  return hparams
+
+
+def get_hparams_from_file(config_path):
+  with open(config_path, "r") as f:
+    data = f.read()
+  config = json.loads(data)
+  hparams =HParams(**config)
+  return hparams
+
+
+def check_git_hash(model_dir):
+  source_dir = os.path.dirname(os.path.realpath(__file__))
+  if not os.path.exists(os.path.join(source_dir, ".git")):
+    logger.warn("{} is not a git repository, therefore hash value comparison will be ignored.".format(
+      source_dir
+    ))
+    return
+
+  cur_hash = subprocess.getoutput("git rev-parse HEAD")
+
+  path = os.path.join(model_dir, "githash")
+  if os.path.exists(path):
+    saved_hash = open(path).read()
+    if saved_hash != cur_hash:
+      logger.warn("git hash values are different. {}(saved) != {}(current)".format(
+        saved_hash[:8], cur_hash[:8]))
+  else:
+    open(path, "w").write(cur_hash)
+
+
+def get_logger(model_dir, filename="train.log"):
+  global logger
+  logger = logging.getLogger(os.path.basename(model_dir))
+  logger.setLevel(logging.DEBUG)
+
+  formatter = logging.Formatter("%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s")
+  if not os.path.exists(model_dir):
+    os.makedirs(model_dir)
+  h = logging.FileHandler(os.path.join(model_dir, filename))
+  h.setLevel(logging.DEBUG)
+  h.setFormatter(formatter)
+  logger.addHandler(h)
+  return logger
+
+
+def repeat_expand_2d(content, target_len):
+    # content : [h, t]
+
+    src_len = content.shape[-1]
+    target = torch.zeros([content.shape[0], target_len], dtype=torch.float).to(content.device)
+    temp = torch.arange(src_len+1) * target_len / src_len
+    current_pos = 0
+    for i in range(target_len):
+        if i < temp[current_pos+1]:
+            target[:, i] = content[:, current_pos]
+        else:
+            current_pos += 1
+            target[:, i] = content[:, current_pos]
+
+    return target
+
+
+def mix_model(model_paths,mix_rate,mode):
+  mix_rate = torch.FloatTensor(mix_rate)/100
+  model_tem = torch.load(model_paths[0])
+  models = [torch.load(path)["model"] for path in model_paths]
+  if mode == 0:
+     mix_rate = F.softmax(mix_rate,dim=0)
+  for k in model_tem["model"].keys():
+     model_tem["model"][k] = torch.zeros_like(model_tem["model"][k])
+     for i,model in enumerate(models):
+        model_tem["model"][k] += model[k]*mix_rate[i]
+  torch.save(model_tem,os.path.join(os.path.curdir,"output.pth"))
+  return os.path.join(os.path.curdir,"output.pth")
+  
+class HParams():
+  def __init__(self, **kwargs):
+    for k, v in kwargs.items():
+      if type(v) == dict:
+        v = HParams(**v)
+      self[k] = v
+
+  def keys(self):
+    return self.__dict__.keys()
+
+  def items(self):
+    return self.__dict__.items()
+
+  def values(self):
+    return self.__dict__.values()
+
+  def __len__(self):
+    return len(self.__dict__)
+
+  def __getitem__(self, key):
+    return getattr(self, key)
+
+  def __setitem__(self, key, value):
+    return setattr(self, key, value)
+
+  def __contains__(self, key):
+    return key in self.__dict__
+
+  def __repr__(self):
+    return self.__dict__.__repr__()
+
+  def get(self,index):
+    return self.__dict__.get(index)
+
+class Volume_Extractor:
+    def __init__(self, hop_size = 512):
+        self.hop_size = hop_size
+        
+    def extract(self, audio): # audio: 2d tensor array
+        if not isinstance(audio,torch.Tensor):
+           audio = torch.Tensor(audio)
+        n_frames = int(audio.size(-1) // self.hop_size)
+        audio2 = audio ** 2
+        audio2 = torch.nn.functional.pad(audio2, (int(self.hop_size // 2), int((self.hop_size + 1) // 2)), mode = 'reflect')
+        volume = torch.FloatTensor([torch.mean(audio2[:,int(n * self.hop_size) : int((n + 1) * self.hop_size)]) for n in range(n_frames)])
+        volume = torch.sqrt(volume)
+        return volume

vdecoder/hifigan/env.py

@@ -0,0 +1,15 @@
+import os
+import shutil
+
+
+class AttrDict(dict):
+    def __init__(self, *args, **kwargs):
+        super(AttrDict, self).__init__(*args, **kwargs)
+        self.__dict__ = self
+
+
+def build_env(config, config_name, path):
+    t_path = os.path.join(path, config_name)
+    if config != t_path:
+        os.makedirs(path, exist_ok=True)
+        shutil.copyfile(config, os.path.join(path, config_name))

vdecoder/hifigan/models.py

@@ -0,0 +1,503 @@
+import os
+import json
+from .env import AttrDict
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
+from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
+from .utils import init_weights, get_padding
+
+LRELU_SLOPE = 0.1
+
+
+def load_model(model_path, device='cuda'):
+    config_file = os.path.join(os.path.split(model_path)[0], 'config.json')
+    with open(config_file) as f:
+        data = f.read()
+
+    global h
+    json_config = json.loads(data)
+    h = AttrDict(json_config)
+
+    generator = Generator(h).to(device)
+
+    cp_dict = torch.load(model_path)
+    generator.load_state_dict(cp_dict['generator'])
+    generator.eval()
+    generator.remove_weight_norm()
+    del cp_dict
+    return generator, h
+
+
+class ResBlock1(torch.nn.Module):
+    def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5)):
+        super(ResBlock1, self).__init__()
+        self.h = h
+        self.convs1 = nn.ModuleList([
+            weight_norm(Conv1d(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],
+                               padding=get_padding(kernel_size, dilation[1]))),
+            weight_norm(Conv1d(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,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1)))
+        ])
+        self.convs2.apply(init_weights)
+
+    def forward(self, x):
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            xt = c1(xt)
+            xt = F.leaky_relu(xt, LRELU_SLOPE)
+            xt = c2(xt)
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+
+
+class ResBlock2(torch.nn.Module):
+    def __init__(self, h, channels, kernel_size=3, dilation=(1, 3)):
+        super(ResBlock2, self).__init__()
+        self.h = h
+        self.convs = nn.ModuleList([
+            weight_norm(Conv1d(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],
+                               padding=get_padding(kernel_size, dilation[1])))
+        ])
+        self.convs.apply(init_weights)
+
+    def forward(self, x):
+        for c in self.convs:
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            xt = c(xt)
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs:
+            remove_weight_norm(l)
+
+
+def padDiff(x):
+    return F.pad(F.pad(x, (0,0,-1,1), 'constant', 0) - x, (0,0,0,-1), 'constant', 0)
+
+class SineGen(torch.nn.Module):
+    """ Definition of sine generator
+    SineGen(samp_rate, harmonic_num = 0,
+            sine_amp = 0.1, noise_std = 0.003,
+            voiced_threshold = 0,
+            flag_for_pulse=False)
+    samp_rate: sampling rate in Hz
+    harmonic_num: number of harmonic overtones (default 0)
+    sine_amp: amplitude of sine-wavefrom (default 0.1)
+    noise_std: std of Gaussian noise (default 0.003)
+    voiced_thoreshold: F0 threshold for U/V classification (default 0)
+    flag_for_pulse: this SinGen is used inside PulseGen (default False)
+    Note: when flag_for_pulse is True, the first time step of a voiced
+        segment is always sin(np.pi) or cos(0)
+    """
+
+    def __init__(self, samp_rate, harmonic_num=0,
+                 sine_amp=0.1, noise_std=0.003,
+                 voiced_threshold=0,
+                 flag_for_pulse=False):
+        super(SineGen, self).__init__()
+        self.sine_amp = sine_amp
+        self.noise_std = noise_std
+        self.harmonic_num = harmonic_num
+        self.dim = self.harmonic_num + 1
+        self.sampling_rate = samp_rate
+        self.voiced_threshold = voiced_threshold
+        self.flag_for_pulse = flag_for_pulse
+
+    def _f02uv(self, f0):
+        # generate uv signal
+        uv = (f0 > self.voiced_threshold).type(torch.float32)
+        return uv
+
+    def _f02sine(self, f0_values):
+        """ f0_values: (batchsize, length, dim)
+            where dim indicates fundamental tone and overtones
+        """
+        # convert to F0 in rad. The interger part n can be ignored
+        # because 2 * np.pi * n doesn't affect phase
+        rad_values = (f0_values / self.sampling_rate) % 1
+
+        # initial phase noise (no noise for fundamental component)
+        rand_ini = torch.rand(f0_values.shape[0], f0_values.shape[2], \
+                              device=f0_values.device)
+        rand_ini[:, 0] = 0
+        rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini
+
+        # instantanouse phase sine[t] = sin(2*pi \sum_i=1 ^{t} rad)
+        if not self.flag_for_pulse:
+            # for normal case
+
+            # To prevent torch.cumsum numerical overflow,
+            # it is necessary to add -1 whenever \sum_k=1^n rad_value_k > 1.
+            # Buffer tmp_over_one_idx indicates the time step to add -1.
+            # This will not change F0 of sine because (x-1) * 2*pi = x * 2*pi
+            tmp_over_one = torch.cumsum(rad_values, 1) % 1
+            tmp_over_one_idx = (padDiff(tmp_over_one)) < 0
+            cumsum_shift = torch.zeros_like(rad_values)
+            cumsum_shift[:, 1:, :] = tmp_over_one_idx * -1.0
+
+            sines = torch.sin(torch.cumsum(rad_values + cumsum_shift, dim=1)
+                              * 2 * np.pi)
+        else:
+            # If necessary, make sure that the first time step of every
+            # voiced segments is sin(pi) or cos(0)
+            # This is used for pulse-train generation
+
+            # identify the last time step in unvoiced segments
+            uv = self._f02uv(f0_values)
+            uv_1 = torch.roll(uv, shifts=-1, dims=1)
+            uv_1[:, -1, :] = 1
+            u_loc = (uv < 1) * (uv_1 > 0)
+
+            # get the instantanouse phase
+            tmp_cumsum = torch.cumsum(rad_values, dim=1)
+            # different batch needs to be processed differently
+            for idx in range(f0_values.shape[0]):
+                temp_sum = tmp_cumsum[idx, u_loc[idx, :, 0], :]
+                temp_sum[1:, :] = temp_sum[1:, :] - temp_sum[0:-1, :]
+                # stores the accumulation of i.phase within
+                # each voiced segments
+                tmp_cumsum[idx, :, :] = 0
+                tmp_cumsum[idx, u_loc[idx, :, 0], :] = temp_sum
+
+            # rad_values - tmp_cumsum: remove the accumulation of i.phase
+            # within the previous voiced segment.
+            i_phase = torch.cumsum(rad_values - tmp_cumsum, dim=1)
+
+            # get the sines
+            sines = torch.cos(i_phase * 2 * np.pi)
+        return sines
+
+    def forward(self, f0):
+        """ sine_tensor, uv = forward(f0)
+        input F0: tensor(batchsize=1, length, dim=1)
+                  f0 for unvoiced steps should be 0
+        output sine_tensor: tensor(batchsize=1, length, dim)
+        output uv: tensor(batchsize=1, length, 1)
+        """
+        with torch.no_grad():
+            f0_buf = torch.zeros(f0.shape[0], f0.shape[1], self.dim,
+                                 device=f0.device)
+            # fundamental component
+            fn = torch.multiply(f0, torch.FloatTensor([[range(1, self.harmonic_num + 2)]]).to(f0.device))
+
+            # generate sine waveforms
+            sine_waves = self._f02sine(fn) * self.sine_amp
+
+            # generate uv signal
+            # uv = torch.ones(f0.shape)
+            # uv = uv * (f0 > self.voiced_threshold)
+            uv = self._f02uv(f0)
+
+            # noise: for unvoiced should be similar to sine_amp
+            #        std = self.sine_amp/3 -> max value ~ self.sine_amp
+            # .       for voiced regions is self.noise_std
+            noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
+            noise = noise_amp * torch.randn_like(sine_waves)
+
+            # first: set the unvoiced part to 0 by uv
+            # then: additive noise
+            sine_waves = sine_waves * uv + noise
+        return sine_waves, uv, noise
+
+
+class SourceModuleHnNSF(torch.nn.Module):
+    """ SourceModule for hn-nsf
+    SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1,
+                 add_noise_std=0.003, voiced_threshod=0)
+    sampling_rate: sampling_rate in Hz
+    harmonic_num: number of harmonic above F0 (default: 0)
+    sine_amp: amplitude of sine source signal (default: 0.1)
+    add_noise_std: std of additive Gaussian noise (default: 0.003)
+        note that amplitude of noise in unvoiced is decided
+        by sine_amp
+    voiced_threshold: threhold to set U/V given F0 (default: 0)
+    Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
+    F0_sampled (batchsize, length, 1)
+    Sine_source (batchsize, length, 1)
+    noise_source (batchsize, length 1)
+    uv (batchsize, length, 1)
+    """
+
+    def __init__(self, sampling_rate, harmonic_num=0, sine_amp=0.1,
+                 add_noise_std=0.003, voiced_threshod=0):
+        super(SourceModuleHnNSF, self).__init__()
+
+        self.sine_amp = sine_amp
+        self.noise_std = add_noise_std
+
+        # to produce sine waveforms
+        self.l_sin_gen = SineGen(sampling_rate, harmonic_num,
+                                 sine_amp, add_noise_std, voiced_threshod)
+
+        # to merge source harmonics into a single excitation
+        self.l_linear = torch.nn.Linear(harmonic_num + 1, 1)
+        self.l_tanh = torch.nn.Tanh()
+
+    def forward(self, x):
+        """
+        Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
+        F0_sampled (batchsize, length, 1)
+        Sine_source (batchsize, length, 1)
+        noise_source (batchsize, length 1)
+        """
+        # source for harmonic branch
+        sine_wavs, uv, _ = self.l_sin_gen(x)
+        sine_merge = self.l_tanh(self.l_linear(sine_wavs))
+
+        # source for noise branch, in the same shape as uv
+        noise = torch.randn_like(uv) * self.sine_amp / 3
+        return sine_merge, noise, uv
+
+
+class Generator(torch.nn.Module):
+    def __init__(self, h):
+        super(Generator, self).__init__()
+        self.h = h
+
+        self.num_kernels = len(h["resblock_kernel_sizes"])
+        self.num_upsamples = len(h["upsample_rates"])
+        self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(h["upsample_rates"]))
+        self.m_source = SourceModuleHnNSF(
+            sampling_rate=h["sampling_rate"],
+            harmonic_num=8)
+        self.noise_convs = nn.ModuleList()
+        self.conv_pre = weight_norm(Conv1d(h["inter_channels"], h["upsample_initial_channel"], 7, 1, padding=3))
+        resblock = ResBlock1 if h["resblock"] == '1' else ResBlock2
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(h["upsample_rates"], h["upsample_kernel_sizes"])):
+            c_cur = h["upsample_initial_channel"] // (2 ** (i + 1))
+            self.ups.append(weight_norm(
+                ConvTranspose1d(h["upsample_initial_channel"] // (2 ** i), h["upsample_initial_channel"] // (2 ** (i + 1)),
+                                k, u, padding=(k - u) // 2)))
+            if i + 1 < len(h["upsample_rates"]):  #
+                stride_f0 = np.prod(h["upsample_rates"][i + 1:])
+                self.noise_convs.append(Conv1d(
+                    1, c_cur, kernel_size=stride_f0 * 2, stride=stride_f0, padding=stride_f0 // 2))
+            else:
+                self.noise_convs.append(Conv1d(1, c_cur, kernel_size=1))
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = h["upsample_initial_channel"] // (2 ** (i + 1))
+            for j, (k, d) in enumerate(zip(h["resblock_kernel_sizes"], h["resblock_dilation_sizes"])):
+                self.resblocks.append(resblock(h, ch, k, d))
+
+        self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3))
+        self.ups.apply(init_weights)
+        self.conv_post.apply(init_weights)
+        self.cond = nn.Conv1d(h['gin_channels'], h['upsample_initial_channel'], 1)
+
+    def forward(self, x, f0, g=None):
+        # print(1,x.shape,f0.shape,f0[:, None].shape)
+        f0 = self.f0_upsamp(f0[:, None]).transpose(1, 2)  # bs,n,t
+        # print(2,f0.shape)
+        har_source, noi_source, uv = self.m_source(f0)
+        har_source = har_source.transpose(1, 2)
+        x = self.conv_pre(x)
+        x = x + self.cond(g)
+        # print(124,x.shape,har_source.shape)
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            # print(3,x.shape)
+            x = self.ups[i](x)
+            x_source = self.noise_convs[i](har_source)
+            # print(4,x_source.shape,har_source.shape,x.shape)
+            x = x + x_source
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        x = torch.tanh(x)
+
+        return x
+
+    def remove_weight_norm(self):
+        print('Removing weight norm...')
+        for l in self.ups:
+            remove_weight_norm(l)
+        for l in self.resblocks:
+            l.remove_weight_norm()
+        remove_weight_norm(self.conv_pre)
+        remove_weight_norm(self.conv_post)
+
+
+class DiscriminatorP(torch.nn.Module):
+    def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
+        super(DiscriminatorP, self).__init__()
+        self.period = period
+        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        self.convs = nn.ModuleList([
+            norm_f(Conv2d(1, 32, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
+            norm_f(Conv2d(32, 128, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
+            norm_f(Conv2d(128, 512, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
+            norm_f(Conv2d(512, 1024, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
+            norm_f(Conv2d(1024, 1024, (kernel_size, 1), 1, padding=(2, 0))),
+        ])
+        self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
+
+    def forward(self, x):
+        fmap = []
+
+        # 1d to 2d
+        b, c, t = x.shape
+        if t % self.period != 0:  # pad first
+            n_pad = self.period - (t % self.period)
+            x = F.pad(x, (0, n_pad), "reflect")
+            t = t + n_pad
+        x = x.view(b, c, t // self.period, self.period)
+
+        for l in self.convs:
+            x = l(x)
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+class MultiPeriodDiscriminator(torch.nn.Module):
+    def __init__(self, periods=None):
+        super(MultiPeriodDiscriminator, self).__init__()
+        self.periods = periods if periods is not None else [2, 3, 5, 7, 11]
+        self.discriminators = nn.ModuleList()
+        for period in self.periods:
+            self.discriminators.append(DiscriminatorP(period))
+
+    def forward(self, y, y_hat):
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        for i, d in enumerate(self.discriminators):
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            y_d_rs.append(y_d_r)
+            fmap_rs.append(fmap_r)
+            y_d_gs.append(y_d_g)
+            fmap_gs.append(fmap_g)
+
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+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
+        self.convs = nn.ModuleList([
+            norm_f(Conv1d(1, 128, 15, 1, padding=7)),
+            norm_f(Conv1d(128, 128, 41, 2, groups=4, padding=20)),
+            norm_f(Conv1d(128, 256, 41, 2, groups=16, padding=20)),
+            norm_f(Conv1d(256, 512, 41, 4, groups=16, padding=20)),
+            norm_f(Conv1d(512, 1024, 41, 4, groups=16, padding=20)),
+            norm_f(Conv1d(1024, 1024, 41, 1, groups=16, padding=20)),
+            norm_f(Conv1d(1024, 1024, 5, 1, padding=2)),
+        ])
+        self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))
+
+    def forward(self, x):
+        fmap = []
+        for l in self.convs:
+            x = l(x)
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+class MultiScaleDiscriminator(torch.nn.Module):
+    def __init__(self):
+        super(MultiScaleDiscriminator, self).__init__()
+        self.discriminators = nn.ModuleList([
+            DiscriminatorS(use_spectral_norm=True),
+            DiscriminatorS(),
+            DiscriminatorS(),
+        ])
+        self.meanpools = nn.ModuleList([
+            AvgPool1d(4, 2, padding=2),
+            AvgPool1d(4, 2, padding=2)
+        ])
+
+    def forward(self, y, y_hat):
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        for i, d in enumerate(self.discriminators):
+            if i != 0:
+                y = self.meanpools[i - 1](y)
+                y_hat = self.meanpools[i - 1](y_hat)
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            y_d_rs.append(y_d_r)
+            fmap_rs.append(fmap_r)
+            y_d_gs.append(y_d_g)
+            fmap_gs.append(fmap_g)
+
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+def feature_loss(fmap_r, fmap_g):
+    loss = 0
+    for dr, dg in zip(fmap_r, fmap_g):
+        for rl, gl in zip(dr, dg):
+            loss += torch.mean(torch.abs(rl - gl))
+
+    return loss * 2
+
+
+def discriminator_loss(disc_real_outputs, disc_generated_outputs):
+    loss = 0
+    r_losses = []
+    g_losses = []
+    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
+        r_loss = torch.mean((1 - dr) ** 2)
+        g_loss = torch.mean(dg ** 2)
+        loss += (r_loss + g_loss)
+        r_losses.append(r_loss.item())
+        g_losses.append(g_loss.item())
+
+    return loss, r_losses, g_losses
+
+
+def generator_loss(disc_outputs):
+    loss = 0
+    gen_losses = []
+    for dg in disc_outputs:
+        l = torch.mean((1 - dg) ** 2)
+        gen_losses.append(l)
+        loss += l
+
+    return loss, gen_losses

vdecoder/hifigan/nvSTFT.py

@@ -0,0 +1,111 @@
+import math
+import os
+os.environ["LRU_CACHE_CAPACITY"] = "3"
+import random
+import torch
+import torch.utils.data
+import numpy as np
+import librosa
+from librosa.util import normalize
+from librosa.filters import mel as librosa_mel_fn
+from scipy.io.wavfile import read
+import soundfile as sf
+
+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 32000
+        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 32000
+    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, center=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
+        
+        if torch.min(y) < -1.:
+            print('min value is ', torch.min(y))
+        if torch.max(y) > 1.:
+            print('max value is ', torch.max(y))
+        
+        if fmax not in self.mel_basis:
+            mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax)
+            self.mel_basis[str(fmax)+'_'+str(y.device)] = torch.from_numpy(mel).float().to(y.device)
+            self.hann_window[str(y.device)] = torch.hann_window(self.win_size).to(y.device)
+        
+        y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_length)/2), int((n_fft-hop_length)/2)), mode='reflect')
+        y = y.squeeze(1)
+        
+        spec = torch.stft(y, n_fft, hop_length=hop_length, win_length=win_size, window=self.hann_window[str(y.device)],
+                          center=center, pad_mode='reflect', normalized=False, onesided=True)
+        # print(111,spec)
+        spec = torch.sqrt(spec.pow(2).sum(-1)+(1e-9))
+        # print(222,spec)
+        spec = torch.matmul(self.mel_basis[str(fmax)+'_'+str(y.device)], spec)
+        # print(333,spec)
+        spec = dynamic_range_compression_torch(spec, clip_val=clip_val)
+        # print(444,spec)
+        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()

vdecoder/hifigan/utils.py

@@ -0,0 +1,68 @@
+import glob
+import os
+import matplotlib
+import torch
+from torch.nn.utils import weight_norm
+# matplotlib.use("Agg")
+import matplotlib.pylab as plt
+
+
+def plot_spectrogram(spectrogram):
+    fig, ax = plt.subplots(figsize=(10, 2))
+    im = ax.imshow(spectrogram, aspect="auto", origin="lower",
+                   interpolation='none')
+    plt.colorbar(im, ax=ax)
+
+    fig.canvas.draw()
+    plt.close()
+
+    return fig
+
+
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def apply_weight_norm(m):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        weight_norm(m)
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size*dilation - dilation)/2)
+
+
+def load_checkpoint(filepath, device):
+    assert os.path.isfile(filepath)
+    print("Loading '{}'".format(filepath))
+    checkpoint_dict = torch.load(filepath, map_location=device)
+    print("Complete.")
+    return checkpoint_dict
+
+
+def save_checkpoint(filepath, obj):
+    print("Saving checkpoint to {}".format(filepath))
+    torch.save(obj, filepath)
+    print("Complete.")
+
+
+def del_old_checkpoints(cp_dir, prefix, n_models=2):
+    pattern = os.path.join(cp_dir, prefix + '????????')
+    cp_list = glob.glob(pattern) # get checkpoint paths
+    cp_list = sorted(cp_list)# sort by iter
+    if len(cp_list) > n_models: # if more than n_models models are found
+        for cp in cp_list[:-n_models]:# delete the oldest models other than lastest n_models
+            open(cp, 'w').close()# empty file contents
+            os.unlink(cp)# delete file (move to trash when using Colab)
+
+
+def scan_checkpoint(cp_dir, prefix):
+    pattern = os.path.join(cp_dir, prefix + '????????')
+    cp_list = glob.glob(pattern)
+    if len(cp_list) == 0:
+        return None
+    return sorted(cp_list)[-1]
+

vdecoder/nsf_hifigan/env.py

@@ -0,0 +1,15 @@
+import os
+import shutil
+
+
+class AttrDict(dict):
+    def __init__(self, *args, **kwargs):
+        super(AttrDict, self).__init__(*args, **kwargs)
+        self.__dict__ = self
+
+
+def build_env(config, config_name, path):
+    t_path = os.path.join(path, config_name)
+    if config != t_path:
+        os.makedirs(path, exist_ok=True)
+        shutil.copyfile(config, os.path.join(path, config_name))

vdecoder/nsf_hifigan/models.py

@@ -0,0 +1,439 @@
+import os
+import json
+from .env import AttrDict
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
+from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
+from .utils import init_weights, get_padding
+
+LRELU_SLOPE = 0.1
+
+
+def load_model(model_path, device='cuda'):
+    h = load_config(model_path)
+
+    generator = Generator(h).to(device)
+
+    cp_dict = torch.load(model_path, map_location=device)
+    generator.load_state_dict(cp_dict['generator'])
+    generator.eval()
+    generator.remove_weight_norm()
+    del cp_dict
+    return generator, h
+
+def load_config(model_path):
+    config_file = os.path.join(os.path.split(model_path)[0], 'config.json')
+    with open(config_file) as f:
+        data = f.read()
+
+    json_config = json.loads(data)
+    h = AttrDict(json_config)
+    return h
+
+
+class ResBlock1(torch.nn.Module):
+    def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5)):
+        super(ResBlock1, self).__init__()
+        self.h = h
+        self.convs1 = nn.ModuleList([
+            weight_norm(Conv1d(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],
+                               padding=get_padding(kernel_size, dilation[1]))),
+            weight_norm(Conv1d(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,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1)))
+        ])
+        self.convs2.apply(init_weights)
+
+    def forward(self, x):
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            xt = c1(xt)
+            xt = F.leaky_relu(xt, LRELU_SLOPE)
+            xt = c2(xt)
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+
+
+class ResBlock2(torch.nn.Module):
+    def __init__(self, h, channels, kernel_size=3, dilation=(1, 3)):
+        super(ResBlock2, self).__init__()
+        self.h = h
+        self.convs = nn.ModuleList([
+            weight_norm(Conv1d(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],
+                               padding=get_padding(kernel_size, dilation[1])))
+        ])
+        self.convs.apply(init_weights)
+
+    def forward(self, x):
+        for c in self.convs:
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            xt = c(xt)
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs:
+            remove_weight_norm(l)
+
+
+class SineGen(torch.nn.Module):
+    """ Definition of sine generator
+    SineGen(samp_rate, harmonic_num = 0,
+            sine_amp = 0.1, noise_std = 0.003,
+            voiced_threshold = 0,
+            flag_for_pulse=False)
+    samp_rate: sampling rate in Hz
+    harmonic_num: number of harmonic overtones (default 0)
+    sine_amp: amplitude of sine-wavefrom (default 0.1)
+    noise_std: std of Gaussian noise (default 0.003)
+    voiced_thoreshold: F0 threshold for U/V classification (default 0)
+    flag_for_pulse: this SinGen is used inside PulseGen (default False)
+    Note: when flag_for_pulse is True, the first time step of a voiced
+        segment is always sin(np.pi) or cos(0)
+    """
+
+    def __init__(self, samp_rate, harmonic_num=0,
+                 sine_amp=0.1, noise_std=0.003,
+                 voiced_threshold=0):
+        super(SineGen, self).__init__()
+        self.sine_amp = sine_amp
+        self.noise_std = noise_std
+        self.harmonic_num = harmonic_num
+        self.dim = self.harmonic_num + 1
+        self.sampling_rate = samp_rate
+        self.voiced_threshold = voiced_threshold
+
+    def _f02uv(self, f0):
+        # generate uv signal
+        uv = torch.ones_like(f0)
+        uv = uv * (f0 > self.voiced_threshold)
+        return uv
+
+    @torch.no_grad()
+    def forward(self, f0, upp):
+        """ sine_tensor, uv = forward(f0)
+        input F0: tensor(batchsize=1, length, dim=1)
+                  f0 for unvoiced steps should be 0
+        output sine_tensor: tensor(batchsize=1, length, dim)
+        output uv: tensor(batchsize=1, length, 1)
+        """
+        f0 = f0.unsqueeze(-1)
+        fn = torch.multiply(f0, torch.arange(1, self.dim + 1, device=f0.device).reshape((1, 1, -1)))
+        rad_values = (fn / self.sampling_rate) % 1  ###%1意味着n_har的乘积无法后处理优化
+        rand_ini = torch.rand(fn.shape[0], fn.shape[2], device=fn.device)
+        rand_ini[:, 0] = 0
+        rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini
+        is_half = rad_values.dtype is not torch.float32
+        tmp_over_one = torch.cumsum(rad_values.double(), 1)  # % 1  #####%1意味着后面的cumsum无法再优化
+        if is_half:
+            tmp_over_one = tmp_over_one.half()
+        else:
+            tmp_over_one = tmp_over_one.float()
+        tmp_over_one *= upp
+        tmp_over_one = F.interpolate(
+            tmp_over_one.transpose(2, 1), scale_factor=upp,
+            mode='linear', align_corners=True
+        ).transpose(2, 1)
+        rad_values = F.interpolate(rad_values.transpose(2, 1), scale_factor=upp, mode='nearest').transpose(2, 1)
+        tmp_over_one %= 1
+        tmp_over_one_idx = (tmp_over_one[:, 1:, :] - tmp_over_one[:, :-1, :]) < 0
+        cumsum_shift = torch.zeros_like(rad_values)
+        cumsum_shift[:, 1:, :] = tmp_over_one_idx * -1.0
+        rad_values = rad_values.double()
+        cumsum_shift = cumsum_shift.double()
+        sine_waves = torch.sin(torch.cumsum(rad_values + cumsum_shift, dim=1) * 2 * np.pi)
+        if is_half:
+            sine_waves = sine_waves.half()
+        else:
+            sine_waves = sine_waves.float()
+        sine_waves = sine_waves * self.sine_amp
+        uv = self._f02uv(f0)
+        uv = F.interpolate(uv.transpose(2, 1), scale_factor=upp, mode='nearest').transpose(2, 1)
+        noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
+        noise = noise_amp * torch.randn_like(sine_waves)
+        sine_waves = sine_waves * uv + noise
+        return sine_waves, uv, noise
+
+
+class SourceModuleHnNSF(torch.nn.Module):
+    """ SourceModule for hn-nsf
+    SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1,
+                 add_noise_std=0.003, voiced_threshod=0)
+    sampling_rate: sampling_rate in Hz
+    harmonic_num: number of harmonic above F0 (default: 0)
+    sine_amp: amplitude of sine source signal (default: 0.1)
+    add_noise_std: std of additive Gaussian noise (default: 0.003)
+        note that amplitude of noise in unvoiced is decided
+        by sine_amp
+    voiced_threshold: threhold to set U/V given F0 (default: 0)
+    Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
+    F0_sampled (batchsize, length, 1)
+    Sine_source (batchsize, length, 1)
+    noise_source (batchsize, length 1)
+    uv (batchsize, length, 1)
+    """
+
+    def __init__(self, sampling_rate, harmonic_num=0, sine_amp=0.1,
+                 add_noise_std=0.003, voiced_threshod=0):
+        super(SourceModuleHnNSF, self).__init__()
+
+        self.sine_amp = sine_amp
+        self.noise_std = add_noise_std
+
+        # to produce sine waveforms
+        self.l_sin_gen = SineGen(sampling_rate, harmonic_num,
+                                 sine_amp, add_noise_std, voiced_threshod)
+
+        # to merge source harmonics into a single excitation
+        self.l_linear = torch.nn.Linear(harmonic_num + 1, 1)
+        self.l_tanh = torch.nn.Tanh()
+
+    def forward(self, x, upp):
+        sine_wavs, uv, _ = self.l_sin_gen(x, upp)
+        sine_merge = self.l_tanh(self.l_linear(sine_wavs))
+        return sine_merge
+
+
+class Generator(torch.nn.Module):
+    def __init__(self, h):
+        super(Generator, self).__init__()
+        self.h = h
+        self.num_kernels = len(h.resblock_kernel_sizes)
+        self.num_upsamples = len(h.upsample_rates)
+        self.m_source = SourceModuleHnNSF(
+            sampling_rate=h.sampling_rate,
+            harmonic_num=8
+        )
+        self.noise_convs = nn.ModuleList()
+        self.conv_pre = weight_norm(Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3))
+        resblock = ResBlock1 if h.resblock == '1' else ResBlock2
+
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
+            c_cur = h.upsample_initial_channel // (2 ** (i + 1))
+            self.ups.append(weight_norm(
+                ConvTranspose1d(h.upsample_initial_channel // (2 ** i), h.upsample_initial_channel // (2 ** (i + 1)),
+                                k, u, padding=(k - u) // 2)))
+            if i + 1 < len(h.upsample_rates):  #
+                stride_f0 = int(np.prod(h.upsample_rates[i + 1:]))
+                self.noise_convs.append(Conv1d(
+                    1, c_cur, kernel_size=stride_f0 * 2, stride=stride_f0, padding=stride_f0 // 2))
+            else:
+                self.noise_convs.append(Conv1d(1, c_cur, kernel_size=1))
+        self.resblocks = nn.ModuleList()
+        ch = h.upsample_initial_channel
+        for i in range(len(self.ups)):
+            ch //= 2
+            for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)):
+                self.resblocks.append(resblock(h, ch, k, d))
+
+        self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3))
+        self.ups.apply(init_weights)
+        self.conv_post.apply(init_weights)
+        self.upp = int(np.prod(h.upsample_rates))
+
+    def forward(self, x, f0):
+        har_source = self.m_source(f0, self.upp).transpose(1, 2)
+        x = self.conv_pre(x)
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            x = self.ups[i](x)
+            x_source = self.noise_convs[i](har_source)
+            x = x + x_source
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        x = torch.tanh(x)
+
+        return x
+
+    def remove_weight_norm(self):
+        print('Removing weight norm...')
+        for l in self.ups:
+            remove_weight_norm(l)
+        for l in self.resblocks:
+            l.remove_weight_norm()
+        remove_weight_norm(self.conv_pre)
+        remove_weight_norm(self.conv_post)
+
+
+class DiscriminatorP(torch.nn.Module):
+    def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
+        super(DiscriminatorP, self).__init__()
+        self.period = period
+        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        self.convs = nn.ModuleList([
+            norm_f(Conv2d(1, 32, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
+            norm_f(Conv2d(32, 128, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
+            norm_f(Conv2d(128, 512, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
+            norm_f(Conv2d(512, 1024, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
+            norm_f(Conv2d(1024, 1024, (kernel_size, 1), 1, padding=(2, 0))),
+        ])
+        self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
+
+    def forward(self, x):
+        fmap = []
+
+        # 1d to 2d
+        b, c, t = x.shape
+        if t % self.period != 0:  # pad first
+            n_pad = self.period - (t % self.period)
+            x = F.pad(x, (0, n_pad), "reflect")
+            t = t + n_pad
+        x = x.view(b, c, t // self.period, self.period)
+
+        for l in self.convs:
+            x = l(x)
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+class MultiPeriodDiscriminator(torch.nn.Module):
+    def __init__(self, periods=None):
+        super(MultiPeriodDiscriminator, self).__init__()
+        self.periods = periods if periods is not None else [2, 3, 5, 7, 11]
+        self.discriminators = nn.ModuleList()
+        for period in self.periods:
+            self.discriminators.append(DiscriminatorP(period))
+
+    def forward(self, y, y_hat):
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        for i, d in enumerate(self.discriminators):
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            y_d_rs.append(y_d_r)
+            fmap_rs.append(fmap_r)
+            y_d_gs.append(y_d_g)
+            fmap_gs.append(fmap_g)
+
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+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
+        self.convs = nn.ModuleList([
+            norm_f(Conv1d(1, 128, 15, 1, padding=7)),
+            norm_f(Conv1d(128, 128, 41, 2, groups=4, padding=20)),
+            norm_f(Conv1d(128, 256, 41, 2, groups=16, padding=20)),
+            norm_f(Conv1d(256, 512, 41, 4, groups=16, padding=20)),
+            norm_f(Conv1d(512, 1024, 41, 4, groups=16, padding=20)),
+            norm_f(Conv1d(1024, 1024, 41, 1, groups=16, padding=20)),
+            norm_f(Conv1d(1024, 1024, 5, 1, padding=2)),
+        ])
+        self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))
+
+    def forward(self, x):
+        fmap = []
+        for l in self.convs:
+            x = l(x)
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+class MultiScaleDiscriminator(torch.nn.Module):
+    def __init__(self):
+        super(MultiScaleDiscriminator, self).__init__()
+        self.discriminators = nn.ModuleList([
+            DiscriminatorS(use_spectral_norm=True),
+            DiscriminatorS(),
+            DiscriminatorS(),
+        ])
+        self.meanpools = nn.ModuleList([
+            AvgPool1d(4, 2, padding=2),
+            AvgPool1d(4, 2, padding=2)
+        ])
+
+    def forward(self, y, y_hat):
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        for i, d in enumerate(self.discriminators):
+            if i != 0:
+                y = self.meanpools[i - 1](y)
+                y_hat = self.meanpools[i - 1](y_hat)
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            y_d_rs.append(y_d_r)
+            fmap_rs.append(fmap_r)
+            y_d_gs.append(y_d_g)
+            fmap_gs.append(fmap_g)
+
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+def feature_loss(fmap_r, fmap_g):
+    loss = 0
+    for dr, dg in zip(fmap_r, fmap_g):
+        for rl, gl in zip(dr, dg):
+            loss += torch.mean(torch.abs(rl - gl))
+
+    return loss * 2
+
+
+def discriminator_loss(disc_real_outputs, disc_generated_outputs):
+    loss = 0
+    r_losses = []
+    g_losses = []
+    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
+        r_loss = torch.mean((1 - dr) ** 2)
+        g_loss = torch.mean(dg ** 2)
+        loss += (r_loss + g_loss)
+        r_losses.append(r_loss.item())
+        g_losses.append(g_loss.item())
+
+    return loss, r_losses, g_losses
+
+
+def generator_loss(disc_outputs):
+    loss = 0
+    gen_losses = []
+    for dg in disc_outputs:
+        l = torch.mean((1 - dg) ** 2)
+        gen_losses.append(l)
+        loss += l
+
+    return loss, gen_losses