FrankZxShen commited on
Commit
ce252ec
1 Parent(s): f77b083
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. G_228800.pth +3 -0
  2. app.py +378 -0
  3. cluster/__init__.py +29 -0
  4. cluster/kmeans.py +201 -0
  5. cluster/train_cluster.py +84 -0
  6. config.json +256 -0
  7. diffusion/__init__.py +0 -0
  8. diffusion/data_loaders.py +284 -0
  9. diffusion/diffusion.py +317 -0
  10. diffusion/diffusion_onnx.py +612 -0
  11. diffusion/dpm_solver_pytorch.py +1201 -0
  12. diffusion/how to export onnx.md +4 -0
  13. diffusion/infer_gt_mel.py +74 -0
  14. diffusion/logger/__init__.py +0 -0
  15. diffusion/logger/saver.py +150 -0
  16. diffusion/logger/utils.py +126 -0
  17. diffusion/onnx_export.py +226 -0
  18. diffusion/solver.py +195 -0
  19. diffusion/unit2mel.py +100 -0
  20. diffusion/vocoder.py +94 -0
  21. diffusion/wavenet.py +108 -0
  22. inference/__init__.py +0 -0
  23. inference/chunks_temp.json +1 -0
  24. inference/infer_tool.py +407 -0
  25. inference/infer_tool_grad.py +171 -0
  26. inference/slicer.py +142 -0
  27. models.py +420 -0
  28. modules/F0Predictor/CrepeF0Predictor.py +31 -0
  29. modules/F0Predictor/DioF0Predictor.py +85 -0
  30. modules/F0Predictor/F0Predictor.py +16 -0
  31. modules/F0Predictor/HarvestF0Predictor.py +81 -0
  32. modules/F0Predictor/PMF0Predictor.py +83 -0
  33. modules/F0Predictor/__init__.py +0 -0
  34. modules/F0Predictor/crepe.py +340 -0
  35. modules/__init__.py +0 -0
  36. modules/attentions.py +349 -0
  37. modules/commons.py +188 -0
  38. modules/enhancer.py +105 -0
  39. modules/losses.py +61 -0
  40. modules/mel_processing.py +112 -0
  41. modules/modules.py +342 -0
  42. requirements.txt +26 -0
  43. utils.py +446 -0
  44. vdecoder/__init__.py +0 -0
  45. vdecoder/hifigan/env.py +15 -0
  46. vdecoder/hifigan/models.py +503 -0
  47. vdecoder/hifigan/nvSTFT.py +111 -0
  48. vdecoder/hifigan/utils.py +68 -0
  49. vdecoder/nsf_hifigan/env.py +15 -0
  50. vdecoder/nsf_hifigan/models.py +439 -0
G_228800.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2635ff85250884517c9989ec5c46a2176c103ba22ade4a682e42759815b4c465
3
+ size 629383387
app.py ADDED
@@ -0,0 +1,378 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # -*- coding: utf-8 -*-
2
+ import traceback
3
+ import torch
4
+ from scipy.io import wavfile
5
+ import edge_tts
6
+ import subprocess
7
+ import gradio as gr
8
+ import gradio.processing_utils as gr_pu
9
+ import io
10
+ import os
11
+ import logging
12
+ import time
13
+ from pathlib import Path
14
+ import re
15
+ import json
16
+ import argparse
17
+
18
+ import librosa
19
+ import matplotlib.pyplot as plt
20
+ import numpy as np
21
+ import soundfile
22
+
23
+ from inference import infer_tool
24
+ from inference import slicer
25
+ from inference.infer_tool import Svc
26
+
27
+ logging.getLogger('numba').setLevel(logging.WARNING)
28
+ chunks_dict = infer_tool.read_temp("inference/chunks_temp.json")
29
+
30
+
31
+ logging.getLogger('numba').setLevel(logging.WARNING)
32
+ logging.getLogger('markdown_it').setLevel(logging.WARNING)
33
+ logging.getLogger('urllib3').setLevel(logging.WARNING)
34
+ logging.getLogger('matplotlib').setLevel(logging.WARNING)
35
+ logging.getLogger('multipart').setLevel(logging.WARNING)
36
+
37
+ model = None
38
+ spk = None
39
+ debug = False
40
+
41
+
42
+ class HParams():
43
+ def __init__(self, **kwargs):
44
+ for k, v in kwargs.items():
45
+ if type(v) == dict:
46
+ v = HParams(**v)
47
+ self[k] = v
48
+
49
+ def keys(self):
50
+ return self.__dict__.keys()
51
+
52
+ def items(self):
53
+ return self.__dict__.items()
54
+
55
+ def values(self):
56
+ return self.__dict__.values()
57
+
58
+ def __len__(self):
59
+ return len(self.__dict__)
60
+
61
+ def __getitem__(self, key):
62
+ return getattr(self, key)
63
+
64
+ def __setitem__(self, key, value):
65
+ return setattr(self, key, value)
66
+
67
+ def __contains__(self, key):
68
+ return key in self.__dict__
69
+
70
+ def __repr__(self):
71
+ return self.__dict__.__repr__()
72
+
73
+
74
+ def get_hparams_from_file(config_path):
75
+ with open(config_path, "r", encoding="utf-8") as f:
76
+ data = f.read()
77
+ config = json.loads(data)
78
+
79
+ hparams = HParams(**config)
80
+ return hparams
81
+
82
+
83
+ 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):
84
+ try:
85
+ if input_audio is None:
86
+ raise gr.Error("你需要上传音频")
87
+ if model is None:
88
+ raise gr.Error("你需要指定模型")
89
+ sampling_rate, audio = input_audio
90
+ # print(audio.shape,sampling_rate)
91
+ audio = (audio / np.iinfo(audio.dtype).max).astype(np.float32)
92
+ if len(audio.shape) > 1:
93
+ audio = librosa.to_mono(audio.transpose(1, 0))
94
+ temp_path = "temp.wav"
95
+ soundfile.write(temp_path, audio, sampling_rate, format="wav")
96
+ _audio = model.slice_inference(temp_path, sid, vc_transform, slice_db, cluster_ratio, auto_f0, noise_scale,
97
+ pad_seconds, cl_num, lg_num, lgr_num, f0_predictor, enhancer_adaptive_key, cr_threshold)
98
+ model.clear_empty()
99
+ os.remove(temp_path)
100
+ # 构建保存文件的路径,并保存到results文件夹内
101
+ try:
102
+ timestamp = str(int(time.time()))
103
+ filename = sid + "_" + timestamp + ".wav"
104
+ # output_file = os.path.join("./results", filename)
105
+ # soundfile.write(output_file, _audio, model.target_sample, format="wav")
106
+ soundfile.write('/tmp/'+filename, _audio,
107
+ model.target_sample, format="wav")
108
+ # return f"推理成功,音频文件保存为results/{filename}", (model.target_sample, _audio)
109
+ return f"推理成功,音频文件保存为{filename}", (model.target_sample, _audio)
110
+ except Exception as e:
111
+ if debug:
112
+ traceback.print_exc()
113
+ return f"文件保存失败,请手动保存", (model.target_sample, _audio)
114
+ except Exception as e:
115
+ if debug:
116
+ traceback.print_exc()
117
+ raise gr.Error(e)
118
+
119
+
120
+ def tts_func(_text, _rate, _voice):
121
+ # 使用edge-tts把文字转成音频
122
+ # voice = "zh-CN-XiaoyiNeural"#女性,较高音
123
+ # voice = "zh-CN-YunxiNeural"#男性
124
+ voice = "zh-CN-YunxiNeural" # 男性
125
+ if (_voice == "女"):
126
+ voice = "zh-CN-XiaoyiNeural"
127
+ output_file = "/tmp/"+_text[0:10]+".wav"
128
+ # communicate = edge_tts.Communicate(_text, voice)
129
+ # await communicate.save(output_file)
130
+ if _rate >= 0:
131
+ ratestr = "+{:.0%}".format(_rate)
132
+ elif _rate < 0:
133
+ ratestr = "{:.0%}".format(_rate) # 减号自带
134
+
135
+ p = subprocess.Popen("edge-tts " +
136
+ " --text "+_text +
137
+ " --write-media "+output_file +
138
+ " --voice "+voice +
139
+ " --rate="+ratestr, shell=True,
140
+ stdout=subprocess.PIPE,
141
+ stdin=subprocess.PIPE)
142
+ p.wait()
143
+ return output_file
144
+
145
+
146
+ def text_clear(text):
147
+ return re.sub(r"[\n\,\(\) ]", "", text)
148
+
149
+
150
+ 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):
151
+ # 使用edge-tts把文字转成音频
152
+ text2tts = text_clear(text2tts)
153
+ output_file = tts_func(text2tts, tts_rate, tts_voice)
154
+
155
+ # 调整采样率
156
+ sr2 = 44100
157
+ wav, sr = librosa.load(output_file)
158
+ wav2 = librosa.resample(wav, orig_sr=sr, target_sr=sr2)
159
+ save_path2 = text2tts[0:10]+"_44k"+".wav"
160
+ wavfile.write(save_path2, sr2,
161
+ (wav2 * np.iinfo(np.int16).max).astype(np.int16)
162
+ )
163
+
164
+ # 读取音频
165
+ sample_rate, data = gr_pu.audio_from_file(save_path2)
166
+ vc_input = (sample_rate, data)
167
+
168
+ a, b = vc_fn(sid, vc_input, vc_transform, auto_f0, cluster_ratio, slice_db, noise_scale,
169
+ pad_seconds, cl_num, lg_num, lgr_num, f0_predictor, enhancer_adaptive_key, cr_threshold)
170
+ os.remove(output_file)
171
+ os.remove(save_path2)
172
+ return a, b
173
+
174
+
175
+ models_info = [
176
+ {
177
+ "description": """
178
+ 这个模型包含公主连结的161名角色。\n\n
179
+ Space采用CPU推理,速度极慢,建议下载模型本地GPU推理。\n\n
180
+ """,
181
+ "model_path": "./G_228800.pth",
182
+ "config_path": "./config.json",
183
+ }
184
+ ]
185
+
186
+ model_inferall = []
187
+ if __name__ == "__main__":
188
+ parser = argparse.ArgumentParser()
189
+ parser.add_argument("--share", action="store_true",
190
+ default=False, help="share gradio app")
191
+ # 一定要设置的部分
192
+ parser.add_argument('-cl', '--clip', type=float,
193
+ default=0, help='音频强制切片,默认0为自动切片,单位为秒/s')
194
+ parser.add_argument('-n', '--clean_names', type=str, nargs='+',
195
+ default=["君の知らない物語-src.wav"], help='wav文件名列表,放在raw文件夹下')
196
+ parser.add_argument('-t', '--trans', type=int, nargs='+',
197
+ default=[0], help='音高调整,支持正负(半音)')
198
+ parser.add_argument('-s', '--spk_list', type=str,
199
+ nargs='+', default=['nen'], help='合成目标说话人名称')
200
+
201
+ # 可选项部分
202
+ parser.add_argument('-a', '--auto_predict_f0', action='store_true',
203
+ default=False, help='语音转换自动预测音高,转换歌声时不要打开这个会严重跑调')
204
+ parser.add_argument('-cm', '--cluster_model_path', type=str,
205
+ default="logs/44k/kmeans_10000.pt", help='聚类模型路径,如果没有训练聚类则随便填')
206
+ parser.add_argument('-cr', '--cluster_infer_ratio', type=float,
207
+ default=0, help='聚类方案占比,范围0-1,若没有训练聚类模型则默认0即可')
208
+ parser.add_argument('-lg', '--linear_gradient', type=float, default=0,
209
+ help='两段音频切片的交叉淡入长度,如果强制切片后出现人声不连贯可调整该数值,如果连贯建议采用默认值0,单位为秒')
210
+ parser.add_argument('-f0p', '--f0_predictor', type=str, default="pm",
211
+ help='选择F0预测器,可选择crepe,pm,dio,harvest,默认为pm(注意:crepe为原F0使用均值滤波器)')
212
+ parser.add_argument('-eh', '--enhance', action='store_true', default=False,
213
+ help='是否使用NSF_HIFIGAN增强器,该选项对部分训练集少的模型有一定的音质增强效果,但是对训练好的模型有反面效果,默认关闭')
214
+ parser.add_argument('-shd', '--shallow_diffusion', action='store_true',
215
+ default=False, help='是否使用浅层扩散,使用后可解决一部分电音问题,默认关闭,该选项打开时,NSF_HIFIGAN增强器将会被禁止')
216
+
217
+ # 浅扩散设置
218
+ parser.add_argument('-dm', '--diffusion_model_path', type=str,
219
+ default="logs/44k/diffusion/model_0.pt", help='扩散模型路径')
220
+ parser.add_argument('-dc', '--diffusion_config_path', type=str,
221
+ default="logs/44k/diffusion/config.yaml", help='扩散模型配置文件路径')
222
+ parser.add_argument('-ks', '--k_step', type=int,
223
+ default=100, help='扩散步数,越大越接近扩散模型的结果,默认100')
224
+ parser.add_argument('-od', '--only_diffusion', action='store_true',
225
+ default=False, help='纯扩散模式,该模式不会加载sovits模型,以扩散模型推理')
226
+
227
+ # 不用动的部分
228
+ parser.add_argument('-sd', '--slice_db', type=int,
229
+ default=-40, help='默认-40,嘈杂的音频可以-30,干声保留呼吸可以-50')
230
+ parser.add_argument('-d', '--device', type=str,
231
+ default=None, help='推理设备,None则为自动选择cpu和gpu')
232
+ parser.add_argument('-ns', '--noice_scale', type=float,
233
+ default=0.4, help='噪音级别,会影响咬字和音质,较为玄学')
234
+ parser.add_argument('-p', '--pad_seconds', type=float, default=0.5,
235
+ help='推理音频pad秒数,由于未��原因开头结尾会有异响,pad一小段静音段后就不会出现')
236
+ parser.add_argument('-wf', '--wav_format', type=str,
237
+ default='flac', help='音频输出格式')
238
+ parser.add_argument('-lgr', '--linear_gradient_retain', type=float,
239
+ default=0.75, help='自动音频切片后,需要舍弃每段切片的头尾。该参数设置交叉长度保留的比例,范围0-1,左开右闭')
240
+ parser.add_argument('-eak', '--enhancer_adaptive_key',
241
+ type=int, default=0, help='使增强器适应更高的音域(单位为半音数)|默认为0')
242
+ parser.add_argument('-ft', '--f0_filter_threshold', type=float, default=0.05,
243
+ help='F0过滤阈值,只有使用crepe时有效. 数值范围从0-1. 降低该值可减少跑调概率,但会增加哑音')
244
+ args = parser.parse_args()
245
+ categories = ["Princess Connect! Re:Dive"]
246
+ others = {
247
+ "None": "https://huggingface.co/spaces/FrankZxShen/vits-fast-finetuning-pcr",
248
+ }
249
+ for info in models_info:
250
+ config_path = info['config_path']
251
+ model_path = info['model_path']
252
+ description = info['description']
253
+ clean_names = args.clean_names
254
+ trans = args.trans
255
+ spk_list = list(get_hparams_from_file(config_path).spk.keys())
256
+ slice_db = args.slice_db
257
+ wav_format = args.wav_format
258
+ auto_predict_f0 = args.auto_predict_f0
259
+ cluster_infer_ratio = args.cluster_infer_ratio
260
+ noice_scale = args.noice_scale
261
+ pad_seconds = args.pad_seconds
262
+ clip = args.clip
263
+ lg = args.linear_gradient
264
+ lgr = args.linear_gradient_retain
265
+ f0p = args.f0_predictor
266
+ enhance = args.enhance
267
+ enhancer_adaptive_key = args.enhancer_adaptive_key
268
+ cr_threshold = args.f0_filter_threshold
269
+ diffusion_model_path = args.diffusion_model_path
270
+ diffusion_config_path = args.diffusion_config_path
271
+ k_step = args.k_step
272
+ only_diffusion = args.only_diffusion
273
+ shallow_diffusion = args.shallow_diffusion
274
+
275
+ model = Svc(model_path, config_path, args.device, args.cluster_model_path, enhance,
276
+ diffusion_model_path, diffusion_config_path, shallow_diffusion, only_diffusion)
277
+
278
+ model_inferall.append((description, spk_list, model))
279
+
280
+ app = gr.Blocks()
281
+ with app:
282
+ gr.Markdown(
283
+ "# <center> so-vits-svc-models-pcr\n"
284
+ "# <center> 注意!!!!!Space采用CPU推理,速度极慢,建议下载模型使用本地GPU推理。\n"
285
+ "## <center> Please do not generate content that could infringe upon the rights or cause harm to individuals or organizations.\n"
286
+ "## <center> 请不要生成会对个人以及组织造成侵害的内容\n\n"
287
+ "[![image](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1pn1xnFfdLK63gVXDwV4zCXfVeo8c-I-0?usp=sharing)\n\n"
288
+ "[![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"
289
+ "[![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)"
290
+ )
291
+ gr.Markdown("# Princess Connect! Re:Dive\n\n"
292
+ )
293
+ with gr.Tabs():
294
+ for category in categories:
295
+ with gr.TabItem(category):
296
+ for i, (description, speakers, model) in enumerate(
297
+ model_inferall):
298
+ gr.Markdown(description)
299
+ with gr.Row():
300
+ with gr.Column():
301
+ # textbox = gr.TextArea(label="Text",
302
+ # placeholder="Type your sentence here ",
303
+ # value="新たなキャラを解放できるようになったようですね。", elem_id=f"tts-input")
304
+
305
+ gr.Markdown(value="""
306
+ <font size=2> 推理设置</font>
307
+ """)
308
+ sid = gr.Dropdown(
309
+ choices=speakers, value=speakers[0], label='角色选择')
310
+ auto_f0 = gr.Checkbox(
311
+ label="自动f0预测,配合聚类模型f0预测效果更好,会导致变调功能失效(仅限转换语音,歌声勾选此项会究极跑调)", value=False)
312
+ f0_predictor = gr.Dropdown(label="选择F0预测器,可选择crepe,pm,dio,harvest,默认为pm(注意:crepe为原F0使用均值滤波器)", choices=[
313
+ "pm", "dio", "harvest", "crepe"], value="pm")
314
+ vc_transform = gr.Number(
315
+ label="变调(整数,可以正负,半音数量,升高八度就是12)", value=0)
316
+ cluster_ratio = gr.Number(
317
+ label="聚类模型混合比例,0-1之间,0即不启用聚类。使用聚类模型能提升音色相似度,但会导致咬字下降(如果使用建议0.5左右)", value=0)
318
+ slice_db = gr.Number(label="切片阈值", value=-40)
319
+ noise_scale = gr.Number(
320
+ label="noise_scale 建议不要动,会影响音质,玄学参数", value=0.4)
321
+ with gr.Column():
322
+ pad_seconds = gr.Number(
323
+ label="推理音频pad秒数,由于未知原因开头结尾会有异响,pad一小段静音段后就不会出现", value=0.5)
324
+ cl_num = gr.Number(
325
+ label="音频自动切片,0为不切片,单位为秒(s)", value=0)
326
+ lg_num = gr.Number(
327
+ label="两端音频切片的交叉淡入长度,如果自动切片后出现人声不连贯可调整该数值,如果连贯建议采用默认值0,注意,该设置会影响推理速度,单位为秒/s", value=0)
328
+ lgr_num = gr.Number(
329
+ label="自动音频切片后,需要舍弃每段切片的头尾。该参数设置交叉长度保留的比例,范围0-1,左开右闭", value=0.75)
330
+ enhancer_adaptive_key = gr.Number(
331
+ label="使增强器适应更高的音域(单位为半音数)|默认为0", value=0)
332
+ cr_threshold = gr.Number(
333
+ label="F0过滤阈值,只有启动crepe时有效. 数值范围从0-1. 降低该值可减少跑调概率,但会增加哑音", value=0.05)
334
+ with gr.Tabs():
335
+ with gr.TabItem("音频转音频"):
336
+ vc_input3 = gr.Audio(label="选择音频")
337
+ vc_submit = gr.Button(
338
+ "音频转换", variant="primary")
339
+ with gr.TabItem("文字转音频"):
340
+ text2tts = gr.Textbox(
341
+ label="在此输入要转译的文字。注意,使用该功能建议打开F0预测,不然会很怪")
342
+ tts_rate = gr.Number(label="tts语速", value=0)
343
+ tts_voice = gr.Radio(label="性别", choices=[
344
+ "男", "女"], value="男")
345
+ vc_submit2 = gr.Button(
346
+ "文字转换", variant="primary")
347
+ with gr.Row():
348
+ with gr.Column():
349
+ vc_output1 = gr.Textbox(label="Output Message")
350
+ with gr.Column():
351
+ vc_output2 = gr.Audio(
352
+ label="Output Audio", interactive=False)
353
+
354
+ vc_submit.click(vc_fn, [sid, vc_input3, vc_transform, auto_f0, cluster_ratio, slice_db, noise_scale, pad_seconds,
355
+ cl_num, lg_num, lgr_num, f0_predictor, enhancer_adaptive_key, cr_threshold], [vc_output1, vc_output2])
356
+ vc_submit2.click(vc_fn2, [sid, vc_input3, vc_transform, auto_f0, cluster_ratio, slice_db, noise_scale, pad_seconds, cl_num,
357
+ lg_num, lgr_num, text2tts, tts_rate, tts_voice, f0_predictor, enhancer_adaptive_key, cr_threshold], [vc_output1, vc_output2])
358
+ # gr.Examples(
359
+ # examples=example,
360
+ # inputs=[textbox, char_dropdown, language_dropdown,
361
+ # duration_slider, symbol_input],
362
+ # outputs=[text_output, audio_output],
363
+ # fn=tts_fn
364
+ # )
365
+ for category, link in others.items():
366
+ with gr.TabItem(category):
367
+ gr.Markdown(
368
+ f'''
369
+ <center>
370
+ <h2>Click to Go</h2>
371
+ <a href="{link}">
372
+ <img src="https://huggingface.co/datasets/huggingface/badges/raw/main/open-in-hf-spaces-xl-dark.svg"
373
+ </a>
374
+ </center>
375
+ '''
376
+ )
377
+
378
+ app.queue(concurrency_count=3).launch(show_api=False, share=args.share)
cluster/__init__.py ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import torch
3
+ from sklearn.cluster import KMeans
4
+
5
+ def get_cluster_model(ckpt_path):
6
+ checkpoint = torch.load(ckpt_path)
7
+ kmeans_dict = {}
8
+ for spk, ckpt in checkpoint.items():
9
+ km = KMeans(ckpt["n_features_in_"])
10
+ km.__dict__["n_features_in_"] = ckpt["n_features_in_"]
11
+ km.__dict__["_n_threads"] = ckpt["_n_threads"]
12
+ km.__dict__["cluster_centers_"] = ckpt["cluster_centers_"]
13
+ kmeans_dict[spk] = km
14
+ return kmeans_dict
15
+
16
+ def get_cluster_result(model, x, speaker):
17
+ """
18
+ x: np.array [t, 256]
19
+ return cluster class result
20
+ """
21
+ return model[speaker].predict(x)
22
+
23
+ def get_cluster_center_result(model, x,speaker):
24
+ """x: np.array [t, 256]"""
25
+ predict = model[speaker].predict(x)
26
+ return model[speaker].cluster_centers_[predict]
27
+
28
+ def get_center(model, x,speaker):
29
+ return model[speaker].cluster_centers_[x]
cluster/kmeans.py ADDED
@@ -0,0 +1,201 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math,pdb
2
+ import torch,pynvml
3
+ from torch.nn.functional import normalize
4
+ from time import time
5
+ import numpy as np
6
+ # device=torch.device("cuda:0")
7
+ def _kpp(data: torch.Tensor, k: int, sample_size: int = -1):
8
+ """ Picks k points in the data based on the kmeans++ method.
9
+
10
+ Parameters
11
+ ----------
12
+ data : torch.Tensor
13
+ Expect a rank 1 or 2 array. Rank 1 is assumed to describe 1-D
14
+ data, rank 2 multidimensional data, in which case one
15
+ row is one observation.
16
+ k : int
17
+ Number of samples to generate.
18
+ sample_size : int
19
+ sample data to avoid memory overflow during calculation
20
+
21
+ Returns
22
+ -------
23
+ init : ndarray
24
+ A 'k' by 'N' containing the initial centroids.
25
+
26
+ References
27
+ ----------
28
+ .. [1] D. Arthur and S. Vassilvitskii, "k-means++: the advantages of
29
+ careful seeding", Proceedings of the Eighteenth Annual ACM-SIAM Symposium
30
+ on Discrete Algorithms, 2007.
31
+ .. [2] scipy/cluster/vq.py: _kpp
32
+ """
33
+ batch_size=data.shape[0]
34
+ if batch_size>sample_size:
35
+ data = data[torch.randint(0, batch_size,[sample_size], device=data.device)]
36
+ dims = data.shape[1] if len(data.shape) > 1 else 1
37
+ init = torch.zeros((k, dims)).to(data.device)
38
+ r = torch.distributions.uniform.Uniform(0, 1)
39
+ for i in range(k):
40
+ if i == 0:
41
+ init[i, :] = data[torch.randint(data.shape[0], [1])]
42
+ else:
43
+ D2 = torch.cdist(init[:i, :][None, :], data[None, :], p=2)[0].amin(dim=0)
44
+ probs = D2 / torch.sum(D2)
45
+ cumprobs = torch.cumsum(probs, dim=0)
46
+ init[i, :] = data[torch.searchsorted(cumprobs, r.sample([1]).to(data.device))]
47
+ return init
48
+ class KMeansGPU:
49
+ '''
50
+ Kmeans clustering algorithm implemented with PyTorch
51
+
52
+ Parameters:
53
+ n_clusters: int,
54
+ Number of clusters
55
+
56
+ max_iter: int, default: 100
57
+ Maximum number of iterations
58
+
59
+ tol: float, default: 0.0001
60
+ Tolerance
61
+
62
+ verbose: int, default: 0
63
+ Verbosity
64
+
65
+ mode: {'euclidean', 'cosine'}, default: 'euclidean'
66
+ Type of distance measure
67
+
68
+ init_method: {'random', 'point', '++'}
69
+ Type of initialization
70
+
71
+ minibatch: {None, int}, default: None
72
+ Batch size of MinibatchKmeans algorithm
73
+ if None perform full KMeans algorithm
74
+
75
+ Attributes:
76
+ centroids: torch.Tensor, shape: [n_clusters, n_features]
77
+ cluster centroids
78
+ '''
79
+ def __init__(self, n_clusters, max_iter=200, tol=1e-4, verbose=0, mode="euclidean",device=torch.device("cuda:0")):
80
+ self.n_clusters = n_clusters
81
+ self.max_iter = max_iter
82
+ self.tol = tol
83
+ self.verbose = verbose
84
+ self.mode = mode
85
+ self.device=device
86
+ pynvml.nvmlInit()
87
+ gpu_handle = pynvml.nvmlDeviceGetHandleByIndex(device.index)
88
+ info = pynvml.nvmlDeviceGetMemoryInfo(gpu_handle)
89
+ self.minibatch=int(33e6/self.n_clusters*info.free/ 1024 / 1024 / 1024)
90
+ print("free_mem/GB:",info.free/ 1024 / 1024 / 1024,"minibatch:",self.minibatch)
91
+
92
+ @staticmethod
93
+ def cos_sim(a, b):
94
+ """
95
+ Compute cosine similarity of 2 sets of vectors
96
+
97
+ Parameters:
98
+ a: torch.Tensor, shape: [m, n_features]
99
+
100
+ b: torch.Tensor, shape: [n, n_features]
101
+ """
102
+ return normalize(a, dim=-1) @ normalize(b, dim=-1).transpose(-2, -1)
103
+
104
+ @staticmethod
105
+ def euc_sim(a, b):
106
+ """
107
+ Compute euclidean similarity of 2 sets of vectors
108
+ Parameters:
109
+ a: torch.Tensor, shape: [m, n_features]
110
+ b: torch.Tensor, shape: [n, n_features]
111
+ """
112
+ return 2 * a @ b.transpose(-2, -1) -(a**2).sum(dim=1)[..., :, None] - (b**2).sum(dim=1)[..., None, :]
113
+
114
+ def max_sim(self, a, b):
115
+ """
116
+ Compute maximum similarity (or minimum distance) of each vector
117
+ in a with all of the vectors in b
118
+ Parameters:
119
+ a: torch.Tensor, shape: [m, n_features]
120
+ b: torch.Tensor, shape: [n, n_features]
121
+ """
122
+ if self.mode == 'cosine':
123
+ sim_func = self.cos_sim
124
+ elif self.mode == 'euclidean':
125
+ sim_func = self.euc_sim
126
+ sim = sim_func(a, b)
127
+ max_sim_v, max_sim_i = sim.max(dim=-1)
128
+ return max_sim_v, max_sim_i
129
+
130
+ def fit_predict(self, X):
131
+ """
132
+ Combination of fit() and predict() methods.
133
+ This is faster than calling fit() and predict() seperately.
134
+ Parameters:
135
+ X: torch.Tensor, shape: [n_samples, n_features]
136
+ centroids: {torch.Tensor, None}, default: None
137
+ if given, centroids will be initialized with given tensor
138
+ if None, centroids will be randomly chosen from X
139
+ Return:
140
+ labels: torch.Tensor, shape: [n_samples]
141
+
142
+ mini_=33kk/k*remain
143
+ mini=min(mini_,fea_shape)
144
+ offset=log2(k/1000)*1.5
145
+ kpp_all=min(mini_*10/offset,fea_shape)
146
+ kpp_sample=min(mini_/12/offset,fea_shape)
147
+ """
148
+ assert isinstance(X, torch.Tensor), "input must be torch.Tensor"
149
+ assert X.dtype in [torch.half, torch.float, torch.double], "input must be floating point"
150
+ assert X.ndim == 2, "input must be a 2d tensor with shape: [n_samples, n_features] "
151
+ # print("verbose:%s"%self.verbose)
152
+
153
+ offset = np.power(1.5,np.log(self.n_clusters / 1000))/np.log(2)
154
+ with torch.no_grad():
155
+ batch_size= X.shape[0]
156
+ # print(self.minibatch, int(self.minibatch * 10 / offset), batch_size)
157
+ start_time = time()
158
+ if (self.minibatch*10//offset< batch_size):
159
+ x = X[torch.randint(0, batch_size,[int(self.minibatch*10/offset)])].to(self.device)
160
+ else:
161
+ x = X.to(self.device)
162
+ # print(x.device)
163
+ self.centroids = _kpp(x, self.n_clusters, min(int(self.minibatch/12/offset),batch_size))
164
+ del x
165
+ torch.cuda.empty_cache()
166
+ # self.centroids = self.centroids.to(self.device)
167
+ num_points_in_clusters = torch.ones(self.n_clusters, device=self.device, dtype=X.dtype)#全1
168
+ closest = None#[3098036]#int64
169
+ if(self.minibatch>=batch_size//2 and self.minibatch<batch_size):
170
+ X = X[torch.randint(0, batch_size,[self.minibatch])].to(self.device)
171
+ elif(self.minibatch>=batch_size):
172
+ X=X.to(self.device)
173
+ for i in range(self.max_iter):
174
+ iter_time = time()
175
+ if self.minibatch<batch_size//2:#可用minibatch数太小,每次都得从内存倒腾到显存
176
+ x = X[torch.randint(0, batch_size, [self.minibatch])].to(self.device)
177
+ else:#否则直接全部缓存
178
+ x = X
179
+
180
+ closest = self.max_sim(a=x, b=self.centroids)[1].to(torch.int16)#[3098036]#int64#0~999
181
+ matched_clusters, counts = closest.unique(return_counts=True)#int64#1k
182
+ expanded_closest = closest[None].expand(self.n_clusters, -1)#[1000, 3098036]#int16#0~999
183
+ mask = (expanded_closest==torch.arange(self.n_clusters, device=self.device)[:, None]).to(X.dtype)#==后者是int64*1000
184
+ c_grad = mask @ x / mask.sum(-1)[..., :, None]
185
+ c_grad[c_grad!=c_grad] = 0 # remove NaNs
186
+ error = (c_grad - self.centroids).pow(2).sum()
187
+ if self.minibatch is not None:
188
+ lr = 1/num_points_in_clusters[:,None] * 0.9 + 0.1
189
+ else:
190
+ lr = 1
191
+ matched_clusters=matched_clusters.long()
192
+ num_points_in_clusters[matched_clusters] += counts#IndexError: tensors used as indices must be long, byte or bool tensors
193
+ self.centroids = self.centroids * (1-lr) + c_grad * lr
194
+ if self.verbose >= 2:
195
+ print('iter:', i, 'error:', error.item(), 'time spent:', round(time()-iter_time, 4))
196
+ if error <= self.tol:
197
+ break
198
+
199
+ if self.verbose >= 1:
200
+ print(f'used {i+1} iterations ({round(time()-start_time, 4)}s) to cluster {batch_size} items into {self.n_clusters} clusters')
201
+ return closest
cluster/train_cluster.py ADDED
@@ -0,0 +1,84 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import time,pdb
2
+ import tqdm
3
+ from time import time as ttime
4
+ import os
5
+ from pathlib import Path
6
+ import logging
7
+ import argparse
8
+ from kmeans import KMeansGPU
9
+ import torch
10
+ import numpy as np
11
+ from sklearn.cluster import KMeans,MiniBatchKMeans
12
+
13
+ logging.basicConfig(level=logging.INFO)
14
+ logger = logging.getLogger(__name__)
15
+ from time import time as ttime
16
+ import pynvml,torch
17
+
18
+ def train_cluster(in_dir, n_clusters, use_minibatch=True, verbose=False,use_gpu=False):#gpu_minibatch真拉,虽然库支持但是也不考虑
19
+ logger.info(f"Loading features from {in_dir}")
20
+ features = []
21
+ nums = 0
22
+ for path in tqdm.tqdm(in_dir.glob("*.soft.pt")):
23
+ # for name in os.listdir(in_dir):
24
+ # path="%s/%s"%(in_dir,name)
25
+ features.append(torch.load(path,map_location="cpu").squeeze(0).numpy().T)
26
+ # print(features[-1].shape)
27
+ features = np.concatenate(features, axis=0)
28
+ print(nums, features.nbytes/ 1024**2, "MB , shape:",features.shape, features.dtype)
29
+ features = features.astype(np.float32)
30
+ logger.info(f"Clustering features of shape: {features.shape}")
31
+ t = time.time()
32
+ if(use_gpu==False):
33
+ if use_minibatch:
34
+ kmeans = MiniBatchKMeans(n_clusters=n_clusters,verbose=verbose, batch_size=4096, max_iter=80).fit(features)
35
+ else:
36
+ kmeans = KMeans(n_clusters=n_clusters,verbose=verbose).fit(features)
37
+ else:
38
+ kmeans = KMeansGPU(n_clusters=n_clusters, mode='euclidean', verbose=2 if verbose else 0,max_iter=500,tol=1e-2)#
39
+ features=torch.from_numpy(features)#.to(device)
40
+ labels = kmeans.fit_predict(features)#
41
+
42
+ print(time.time()-t, "s")
43
+
44
+ x = {
45
+ "n_features_in_": kmeans.n_features_in_ if use_gpu==False else features.shape[1],
46
+ "_n_threads": kmeans._n_threads if use_gpu==False else 4,
47
+ "cluster_centers_": kmeans.cluster_centers_ if use_gpu==False else kmeans.centroids.cpu().numpy(),
48
+ }
49
+ print("end")
50
+
51
+ return x
52
+
53
+ if __name__ == "__main__":
54
+ parser = argparse.ArgumentParser()
55
+ parser.add_argument('--dataset', type=Path, default="./dataset/44k",
56
+ help='path of training data directory')
57
+ parser.add_argument('--output', type=Path, default="logs/44k",
58
+ help='path of model output directory')
59
+ parser.add_argument('--gpu',action='store_true', default=False ,
60
+ help='to use GPU')
61
+
62
+
63
+ args = parser.parse_args()
64
+
65
+ checkpoint_dir = args.output
66
+ dataset = args.dataset
67
+ use_gpu = args.gpu
68
+ n_clusters = 10000
69
+
70
+ ckpt = {}
71
+ for spk in os.listdir(dataset):
72
+ if os.path.isdir(dataset/spk):
73
+ print(f"train kmeans for {spk}...")
74
+ in_dir = dataset/spk
75
+ x = train_cluster(in_dir, n_clusters,use_minibatch=False,verbose=False,use_gpu=use_gpu)
76
+ ckpt[spk] = x
77
+
78
+ checkpoint_path = checkpoint_dir / f"kmeans_{n_clusters}.pt"
79
+ checkpoint_path.parent.mkdir(exist_ok=True, parents=True)
80
+ torch.save(
81
+ ckpt,
82
+ checkpoint_path,
83
+ )
84
+
config.json ADDED
@@ -0,0 +1,256 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "train": {
3
+ "log_interval": 200,
4
+ "eval_interval": 800,
5
+ "seed": 1234,
6
+ "epochs": 10000,
7
+ "learning_rate": 0.0001,
8
+ "betas": [
9
+ 0.8,
10
+ 0.99
11
+ ],
12
+ "eps": 1e-09,
13
+ "batch_size": 6,
14
+ "fp16_run": false,
15
+ "lr_decay": 0.999875,
16
+ "segment_size": 10240,
17
+ "init_lr_ratio": 1,
18
+ "warmup_epochs": 0,
19
+ "c_mel": 45,
20
+ "c_kl": 1.0,
21
+ "use_sr": true,
22
+ "max_speclen": 512,
23
+ "port": "8976",
24
+ "keep_ckpts": 3,
25
+ "all_in_mem": false
26
+ },
27
+ "data": {
28
+ "training_files": "filelists/train.txt",
29
+ "validation_files": "filelists/val.txt",
30
+ "max_wav_value": 32768.0,
31
+ "sampling_rate": 44100,
32
+ "filter_length": 2048,
33
+ "hop_length": 512,
34
+ "win_length": 2048,
35
+ "n_mel_channels": 80,
36
+ "mel_fmin": 0.0,
37
+ "mel_fmax": 22050
38
+ },
39
+ "model": {
40
+ "inter_channels": 192,
41
+ "hidden_channels": 192,
42
+ "filter_channels": 768,
43
+ "n_heads": 2,
44
+ "n_layers": 6,
45
+ "kernel_size": 3,
46
+ "p_dropout": 0.1,
47
+ "resblock": "1",
48
+ "resblock_kernel_sizes": [
49
+ 3,
50
+ 7,
51
+ 11
52
+ ],
53
+ "resblock_dilation_sizes": [
54
+ [
55
+ 1,
56
+ 3,
57
+ 5
58
+ ],
59
+ [
60
+ 1,
61
+ 3,
62
+ 5
63
+ ],
64
+ [
65
+ 1,
66
+ 3,
67
+ 5
68
+ ]
69
+ ],
70
+ "upsample_rates": [
71
+ 8,
72
+ 8,
73
+ 2,
74
+ 2,
75
+ 2
76
+ ],
77
+ "upsample_initial_channel": 512,
78
+ "upsample_kernel_sizes": [
79
+ 16,
80
+ 16,
81
+ 4,
82
+ 4,
83
+ 4
84
+ ],
85
+ "n_layers_q": 3,
86
+ "use_spectral_norm": false,
87
+ "gin_channels": 768,
88
+ "ssl_dim": 768,
89
+ "n_speakers": 161,
90
+ "speech_encoder": "vec768l12",
91
+ "speaker_embedding": false
92
+ },
93
+ "spk": {
94
+ "\u83c8\u6bd4\u8389\u65af\u5854\uff08Overload\uff09": 0,
95
+ "\u94c3\u5948\uff08\u590f\u65e5\uff09": 1,
96
+ "\u674f\u5948": 2,
97
+ "\u96ea\u83f2": 3,
98
+ "\u53ef\u53ef\u841d": 4,
99
+ "\u83c8\u6bd4\u8389\u65af\u5854": 5,
100
+ "\u7eeb\u97f3\uff08\u5723\u8bde\u8282\uff09": 6,
101
+ "\u78a7\uff08\u63d2\u73ed\u751f\uff09": 7,
102
+ "\u5609\u591c": 8,
103
+ "\u955c\u534e\uff08\u4e07\u5723\u8282\uff09": 9,
104
+ "\u73e0\u5e0c": 10,
105
+ "\u671b\uff08\u590f\u65e5\uff09": 11,
106
+ "\u94c3\u8393": 12,
107
+ "\u771f\u7434\uff08\u590f\u65e5\uff09": 13,
108
+ "\u9759\u6d41\uff08\u60c5\u4eba\u8282\uff09": 14,
109
+ "\u54b2\u604b": 15,
110
+ "\u8389\u739b": 16,
111
+ "\u9999\u6f84\uff08\u590f\u65e5\uff09": 17,
112
+ "\u5343\u6b4c": 18,
113
+ "\u5fcd": 19,
114
+ "\u4f9d\u91cc": 20,
115
+ "\u4f69\u53ef\u8389\u59c6\uff08\u65b0\u5e74\uff09": 21,
116
+ "\u80e1\u6843": 22,
117
+ "\u667a\uff08\u9b54\u6cd5\u5c11\u5973\uff09": 23,
118
+ "\u4f18\u8863": 24,
119
+ "\u601c\uff08\u4e07\u5723\u8282\uff09": 25,
120
+ "\u681e": 26,
121
+ "\u9999\u7ec7\uff08\u590f\u65e5\uff09": 27,
122
+ "\u7948\u68a8\uff08\u65f6\u95f4\u65c5\u884c\uff09": 28,
123
+ "\u54b2\u604b\uff08\u590f\u65e5\uff09": 29,
124
+ "\u94c3\u5948": 30,
125
+ "\u771f\u6b65\uff08\u590f\u65e5\uff09": 31,
126
+ "\u4e9a\u91cc\u838e": 32,
127
+ "\u955c\u534e": 33,
128
+ "\u672a\u594f\u5e0c": 34,
129
+ "\u4f0a\u8389\u4e9a\uff08\u5723\u8bde\u8282\uff09": 35,
130
+ "\u7f8e\u91cc": 36,
131
+ "\u4f3c\u4f3c\u82b1": 37,
132
+ "\u514b\u8389\u4e1d\u63d0\u5a1c\uff08\u5723\u8bde\u8282\uff09": 38,
133
+ "\u7f8e\u51ac\uff08\u590f\u65e5\uff09": 39,
134
+ "\u8389\u739b\uff08\u7070\u59d1\u5a18\uff09": 40,
135
+ "\u94c3\u8393\uff08\u590f\u65e5\uff09": 41,
136
+ "\u53e4\u857e\u96c5": 42,
137
+ "\u7f8e\u7f8e\uff08\u4e07\u5723\u8282\uff09": 43,
138
+ "\u5343\u6b4c\uff08\u5723\u8bde\u8282\uff09": 44,
139
+ "\u78a7": 45,
140
+ "\u96ea": 46,
141
+ "\u60e0\u7406\u5b50\uff08\u60c5\u4eba\u8282\uff09": 47,
142
+ "\u4f0a\u7eea": 48,
143
+ "\u62c9\u59c6": 49,
144
+ "\u4e03\u4e03\u9999\uff08\u590f\u65e5\uff09": 50,
145
+ "\u94c3\u8393\uff08\u65b0\u5e74\uff09": 51,
146
+ "\u94c3\uff08\u6e38\u9a91\u5175\uff09": 52,
147
+ "\u73af\u5948": 53,
148
+ "\u60e0\u7406\u5b50": 54,
149
+ "\u53ef\u53ef\u841d\uff08\u793c\u670d\uff09": 55,
150
+ "\u4f18\u8863\uff08\u65b0\u5e74\uff09": 56,
151
+ "\u671b\uff08\u5723\u8bde\u8282\uff09": 57,
152
+ "\u53ef\u53ef\u841d\uff08\u516c\u4e3b\uff09": 58,
153
+ "\u7eaf\uff08\u590f\u65e5\uff09": 59,
154
+ "\u73af\u5948\uff08\u632f\u8896\uff09": 60,
155
+ "\u54b2\u604b\uff08\u5723\u8bde\u8282\uff09": 61,
156
+ "\u9732\u5a1c": 62,
157
+ "\u9999\u6f84\uff08\u9b54\u6cd5\u5c11\u5973\uff09": 63,
158
+ "\u59ec\u5854": 64,
159
+ "\u77db\u4f9d\u672a\uff08\u65b0\u5e74\uff09": 65,
160
+ "\u681e\uff08\u9b54\u6cd5\u5c11\u5973\uff09": 66,
161
+ "\u771f\u6b65\uff08\u7070\u59d1\u5a18\uff09": 67,
162
+ "\u601c\uff08\u65b0\u5e74\uff09": 68,
163
+ "\u5343\u6b4c\uff08\u590f\u65e5\uff09": 69,
164
+ "\u7948\u68a8": 70,
165
+ "\u7531\u52a0\u8389": 71,
166
+ "\u4f69\u53ef\u8389\u59c6\uff08\u590f\u65e5\uff09": 72,
167
+ "\u6b65\u7f8e\uff08\u4ed9\u5883\uff09": 73,
168
+ "\u53ef\u53ef\u841d\uff08\u590f\u65e5\uff09": 74,
169
+ "\u7eeb\u97f3": 75,
170
+ "\u7531\u52a0\u8389\uff08\u5723\u8bde\u8282\uff09": 76,
171
+ "\u771f\u9633": 77,
172
+ "\u96f7\u59c6": 78,
173
+ "\u831c\u91cc": 79,
174
+ "\u94c3": 80,
175
+ "\u51db\uff08\u5076\u50cf\u5927\u5e08\uff09": 81,
176
+ "\u51ef\u9732": 82,
177
+ "\u514b\u7f57\u4f9d\uff08\u5723\u5b66\u796d\uff09": 83,
178
+ "\u79cb\u4e43": 84,
179
+ "\u514b\u8389\u4e1d\u63d0\u5a1c": 85,
180
+ "\u4f69\u53ef\u8389\u59c6": 86,
181
+ "\u9732": 87,
182
+ "\u83ab\u59ae\u5361": 88,
183
+ "\u7483\u4e43": 89,
184
+ "\u9759\u6d41\uff08\u590f\u65e5\uff09": 90,
185
+ "\u7eba\u5e0c\uff08\u4e07\u5723\u8282\uff09": 91,
186
+ "\u9999\u6f84": 92,
187
+ "\u7f8e\u54b2": 93,
188
+ "\u7f8e\u91cc\uff08\u590f\u65e5\uff09": 94,
189
+ "\u667a": 95,
190
+ "\u521d\u97f3": 96,
191
+ "\u7483\u4e43\uff08\u4ed9\u5883\uff09": 97,
192
+ "\u536f\u6708\uff08\u5076\u50cf\u5927\u5e08\uff09": 98,
193
+ "\u7a7a\u82b1": 99,
194
+ "\u73e0\u5e0c\uff08\u590f\u65e5\uff09": 100,
195
+ "\u7eaf": 101,
196
+ "\u7f8e\u7f8e": 102,
197
+ "\u5fcd\uff08\u4e07\u5723\u8282\uff09": 103,
198
+ "\u59ae\u4fac\uff08\u5927\u6c5f\u6237\uff09": 104,
199
+ "\u5343\u7231\u7460": 105,
200
+ "\u53ef\u53ef\u841d\uff08\u65b0\u5e74\uff09": 106,
201
+ "\u601c\uff08\u516c\u4e3b\uff09": 107,
202
+ "\u672a\u594f\u5e0c\u3001\u7f8e\u7f8e\u3001\u955c\u534e": 108,
203
+ "\u4f0a\u8389\u4e9a": 109,
204
+ "\u671b": 110,
205
+ "\u4f18\u8863\uff08\u516c\u4e3b\uff09": 111,
206
+ "\u51ef\u9732\uff08\u516c\u4e3b\uff09": 112,
207
+ "\u514b\u7f57\u4f9d": 113,
208
+ "\u79cb\u4e43\uff08\u5723\u8bde\u8282\uff09": 114,
209
+ "\u4f18\u8863\uff08\u793c\u670d\uff09": 115,
210
+ "\u674f\u5948\uff08\u590f\u65e5\uff09": 116,
211
+ "\u9759\u6d41": 117,
212
+ "\u672a\u594f\u5e0c\uff08\u4e07\u5723\u8282\uff09": 118,
213
+ "\u831c\u91cc\uff08\u5929\u4f7f\uff09": 119,
214
+ "\u5bab\u5b50": 120,
215
+ "\u65e5\u548c\u8389\uff08\u65b0\u5e74\uff09": 121,
216
+ "\u4f18\u59ae": 122,
217
+ "\u5b89": 123,
218
+ "\u6d41\u590f\uff08\u590f\u65e5\uff09": 124,
219
+ "\u7f8e\u51ac\uff08\u5de5\u4f5c\u670d\uff09": 125,
220
+ "\u7f8e\u51ac": 126,
221
+ "\u521d\u97f3\uff08\u590f\u65e5\uff09": 127,
222
+ "\u77db\u4f9d\u672a": 128,
223
+ "\u51ef\u9732\uff08\u590f\u65e5\uff09": 129,
224
+ "\u83ab\u59ae\u5361\uff08\u9b54\u6cd5\u5c11\u5973\uff09": 130,
225
+ "\u9999\u7ec7": 131,
226
+ "\u5bab\u5b50\uff08\u4e07\u5723\u8282\uff09": 132,
227
+ "\u7f8e\u54b2\uff08\u4e07\u5723\u8282\uff09": 133,
228
+ "\u59ae\u4fac": 134,
229
+ "\u601c": 135,
230
+ "\u771f\u9633\uff08\u6e38\u9a91\u5175\uff09": 136,
231
+ "\u78a7\uff08\u5de5\u4f5c\u670d\uff09": 137,
232
+ "\u65e5\u548c\u8389": 138,
233
+ "\u4e03\u4e03\u9999": 139,
234
+ "\u771f\u7434": 140,
235
+ "\u6b65\u7f8e": 141,
236
+ "\u6df1\u6708": 142,
237
+ "\u65e5\u548c\u8389\uff08\u516c\u4e3b\uff09": 143,
238
+ "\u4f0a\u7eea\uff08\u590f\u65e5\uff09": 144,
239
+ "\u8309\u8389": 145,
240
+ "\u4f69\u53ef\u8389\u59c6\uff08\u516c\u4e3b\uff09": 146,
241
+ "\u7a7a\u82b1\uff08\u5927\u6c5f\u6237\uff09": 147,
242
+ "\u60e0\u7406\u5b50\uff08\u590f\u65e5\uff09": 148,
243
+ "\u51ef\u9732\uff08\u65b0\u5e74\uff09": 149,
244
+ "\u7231\u871c\u8389\u96c5": 150,
245
+ "\u4f9d\u91cc\uff08\u5929\u4f7f\uff09": 151,
246
+ "\u8309\u8389\uff08\u4e07\u5723\u8282\uff09": 152,
247
+ "\u771f\u7434\uff08\u7070\u59d1\u5a18\uff09": 153,
248
+ "\u7eba\u5e0c": 154,
249
+ "\u771f\u6b65": 155,
250
+ "\u6d41\u590f": 156,
251
+ "\u4f3c\u4f3c\u82b1\uff08\u65b0\u5e74\uff09": 157,
252
+ "\u672a\u592e\uff08\u5076\u50cf\u5927\u5e08\uff09": 158,
253
+ "\u80e1\u6843\uff08\u5723\u8bde\u8282\uff09": 159,
254
+ "\u5343\u7231\u7460\uff08\u5723\u5b66\u796d\uff09": 160
255
+ }
256
+ }
diffusion/__init__.py ADDED
File without changes
diffusion/data_loaders.py ADDED
@@ -0,0 +1,284 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import random
3
+ import re
4
+ import numpy as np
5
+ import librosa
6
+ import torch
7
+ import random
8
+ from utils import repeat_expand_2d
9
+ from tqdm import tqdm
10
+ from torch.utils.data import Dataset
11
+
12
+ def traverse_dir(
13
+ root_dir,
14
+ extensions,
15
+ amount=None,
16
+ str_include=None,
17
+ str_exclude=None,
18
+ is_pure=False,
19
+ is_sort=False,
20
+ is_ext=True):
21
+
22
+ file_list = []
23
+ cnt = 0
24
+ for root, _, files in os.walk(root_dir):
25
+ for file in files:
26
+ if any([file.endswith(f".{ext}") for ext in extensions]):
27
+ # path
28
+ mix_path = os.path.join(root, file)
29
+ pure_path = mix_path[len(root_dir)+1:] if is_pure else mix_path
30
+
31
+ # amount
32
+ if (amount is not None) and (cnt == amount):
33
+ if is_sort:
34
+ file_list.sort()
35
+ return file_list
36
+
37
+ # check string
38
+ if (str_include is not None) and (str_include not in pure_path):
39
+ continue
40
+ if (str_exclude is not None) and (str_exclude in pure_path):
41
+ continue
42
+
43
+ if not is_ext:
44
+ ext = pure_path.split('.')[-1]
45
+ pure_path = pure_path[:-(len(ext)+1)]
46
+ file_list.append(pure_path)
47
+ cnt += 1
48
+ if is_sort:
49
+ file_list.sort()
50
+ return file_list
51
+
52
+
53
+ def get_data_loaders(args, whole_audio=False):
54
+ data_train = AudioDataset(
55
+ filelists = args.data.training_files,
56
+ waveform_sec=args.data.duration,
57
+ hop_size=args.data.block_size,
58
+ sample_rate=args.data.sampling_rate,
59
+ load_all_data=args.train.cache_all_data,
60
+ whole_audio=whole_audio,
61
+ extensions=args.data.extensions,
62
+ n_spk=args.model.n_spk,
63
+ spk=args.spk,
64
+ device=args.train.cache_device,
65
+ fp16=args.train.cache_fp16,
66
+ use_aug=True)
67
+ loader_train = torch.utils.data.DataLoader(
68
+ data_train ,
69
+ batch_size=args.train.batch_size if not whole_audio else 1,
70
+ shuffle=True,
71
+ num_workers=args.train.num_workers if args.train.cache_device=='cpu' else 0,
72
+ persistent_workers=(args.train.num_workers > 0) if args.train.cache_device=='cpu' else False,
73
+ pin_memory=True if args.train.cache_device=='cpu' else False
74
+ )
75
+ data_valid = AudioDataset(
76
+ filelists = args.data.validation_files,
77
+ waveform_sec=args.data.duration,
78
+ hop_size=args.data.block_size,
79
+ sample_rate=args.data.sampling_rate,
80
+ load_all_data=args.train.cache_all_data,
81
+ whole_audio=True,
82
+ spk=args.spk,
83
+ extensions=args.data.extensions,
84
+ n_spk=args.model.n_spk)
85
+ loader_valid = torch.utils.data.DataLoader(
86
+ data_valid,
87
+ batch_size=1,
88
+ shuffle=False,
89
+ num_workers=0,
90
+ pin_memory=True
91
+ )
92
+ return loader_train, loader_valid
93
+
94
+
95
+ class AudioDataset(Dataset):
96
+ def __init__(
97
+ self,
98
+ filelists,
99
+ waveform_sec,
100
+ hop_size,
101
+ sample_rate,
102
+ spk,
103
+ load_all_data=True,
104
+ whole_audio=False,
105
+ extensions=['wav'],
106
+ n_spk=1,
107
+ device='cpu',
108
+ fp16=False,
109
+ use_aug=False,
110
+ ):
111
+ super().__init__()
112
+
113
+ self.waveform_sec = waveform_sec
114
+ self.sample_rate = sample_rate
115
+ self.hop_size = hop_size
116
+ self.filelists = filelists
117
+ self.whole_audio = whole_audio
118
+ self.use_aug = use_aug
119
+ self.data_buffer={}
120
+ self.pitch_aug_dict = {}
121
+ # np.load(os.path.join(self.path_root, 'pitch_aug_dict.npy'), allow_pickle=True).item()
122
+ if load_all_data:
123
+ print('Load all the data filelists:', filelists)
124
+ else:
125
+ print('Load the f0, volume data filelists:', filelists)
126
+ with open(filelists,"r") as f:
127
+ self.paths = f.read().splitlines()
128
+ for name_ext in tqdm(self.paths, total=len(self.paths)):
129
+ name = os.path.splitext(name_ext)[0]
130
+ path_audio = name_ext
131
+ duration = librosa.get_duration(filename = path_audio, sr = self.sample_rate)
132
+
133
+ path_f0 = name_ext + ".f0.npy"
134
+ f0,_ = np.load(path_f0,allow_pickle=True)
135
+ f0 = torch.from_numpy(np.array(f0,dtype=float)).float().unsqueeze(-1).to(device)
136
+
137
+ path_volume = name_ext + ".vol.npy"
138
+ volume = np.load(path_volume)
139
+ volume = torch.from_numpy(volume).float().unsqueeze(-1).to(device)
140
+
141
+ path_augvol = name_ext + ".aug_vol.npy"
142
+ aug_vol = np.load(path_augvol)
143
+ aug_vol = torch.from_numpy(aug_vol).float().unsqueeze(-1).to(device)
144
+
145
+ if n_spk is not None and n_spk > 1:
146
+ spk_name = name_ext.split("/")[-2]
147
+ spk_id = spk[spk_name] if spk_name in spk else 0
148
+ if spk_id < 0 or spk_id >= n_spk:
149
+ raise ValueError(' [x] Muiti-speaker traing error : spk_id must be a positive integer from 0 to n_spk-1 ')
150
+ else:
151
+ spk_id = 0
152
+ spk_id = torch.LongTensor(np.array([spk_id])).to(device)
153
+
154
+ if load_all_data:
155
+ '''
156
+ audio, sr = librosa.load(path_audio, sr=self.sample_rate)
157
+ if len(audio.shape) > 1:
158
+ audio = librosa.to_mono(audio)
159
+ audio = torch.from_numpy(audio).to(device)
160
+ '''
161
+ path_mel = name_ext + ".mel.npy"
162
+ mel = np.load(path_mel)
163
+ mel = torch.from_numpy(mel).to(device)
164
+
165
+ path_augmel = name_ext + ".aug_mel.npy"
166
+ aug_mel,keyshift = np.load(path_augmel, allow_pickle=True)
167
+ aug_mel = np.array(aug_mel,dtype=float)
168
+ aug_mel = torch.from_numpy(aug_mel).to(device)
169
+ self.pitch_aug_dict[name_ext] = keyshift
170
+
171
+ path_units = name_ext + ".soft.pt"
172
+ units = torch.load(path_units).to(device)
173
+ units = units[0]
174
+ units = repeat_expand_2d(units,f0.size(0)).transpose(0,1)
175
+
176
+ if fp16:
177
+ mel = mel.half()
178
+ aug_mel = aug_mel.half()
179
+ units = units.half()
180
+
181
+ self.data_buffer[name_ext] = {
182
+ 'duration': duration,
183
+ 'mel': mel,
184
+ 'aug_mel': aug_mel,
185
+ 'units': units,
186
+ 'f0': f0,
187
+ 'volume': volume,
188
+ 'aug_vol': aug_vol,
189
+ 'spk_id': spk_id
190
+ }
191
+ else:
192
+ path_augmel = name_ext + ".aug_mel.npy"
193
+ aug_mel,keyshift = np.load(path_augmel, allow_pickle=True)
194
+ self.pitch_aug_dict[name_ext] = keyshift
195
+ self.data_buffer[name_ext] = {
196
+ 'duration': duration,
197
+ 'f0': f0,
198
+ 'volume': volume,
199
+ 'aug_vol': aug_vol,
200
+ 'spk_id': spk_id
201
+ }
202
+
203
+
204
+ def __getitem__(self, file_idx):
205
+ name_ext = self.paths[file_idx]
206
+ data_buffer = self.data_buffer[name_ext]
207
+ # check duration. if too short, then skip
208
+ if data_buffer['duration'] < (self.waveform_sec + 0.1):
209
+ return self.__getitem__( (file_idx + 1) % len(self.paths))
210
+
211
+ # get item
212
+ return self.get_data(name_ext, data_buffer)
213
+
214
+ def get_data(self, name_ext, data_buffer):
215
+ name = os.path.splitext(name_ext)[0]
216
+ frame_resolution = self.hop_size / self.sample_rate
217
+ duration = data_buffer['duration']
218
+ waveform_sec = duration if self.whole_audio else self.waveform_sec
219
+
220
+ # load audio
221
+ idx_from = 0 if self.whole_audio else random.uniform(0, duration - waveform_sec - 0.1)
222
+ start_frame = int(idx_from / frame_resolution)
223
+ units_frame_len = int(waveform_sec / frame_resolution)
224
+ aug_flag = random.choice([True, False]) and self.use_aug
225
+ '''
226
+ audio = data_buffer.get('audio')
227
+ if audio is None:
228
+ path_audio = os.path.join(self.path_root, 'audio', name) + '.wav'
229
+ audio, sr = librosa.load(
230
+ path_audio,
231
+ sr = self.sample_rate,
232
+ offset = start_frame * frame_resolution,
233
+ duration = waveform_sec)
234
+ if len(audio.shape) > 1:
235
+ audio = librosa.to_mono(audio)
236
+ # clip audio into N seconds
237
+ audio = audio[ : audio.shape[-1] // self.hop_size * self.hop_size]
238
+ audio = torch.from_numpy(audio).float()
239
+ else:
240
+ audio = audio[start_frame * self.hop_size : (start_frame + units_frame_len) * self.hop_size]
241
+ '''
242
+ # load mel
243
+ mel_key = 'aug_mel' if aug_flag else 'mel'
244
+ mel = data_buffer.get(mel_key)
245
+ if mel is None:
246
+ mel = name_ext + ".mel.npy"
247
+ mel = np.load(mel)
248
+ mel = mel[start_frame : start_frame + units_frame_len]
249
+ mel = torch.from_numpy(mel).float()
250
+ else:
251
+ mel = mel[start_frame : start_frame + units_frame_len]
252
+
253
+ # load f0
254
+ f0 = data_buffer.get('f0')
255
+ aug_shift = 0
256
+ if aug_flag:
257
+ aug_shift = self.pitch_aug_dict[name_ext]
258
+ f0_frames = 2 ** (aug_shift / 12) * f0[start_frame : start_frame + units_frame_len]
259
+
260
+ # load units
261
+ units = data_buffer.get('units')
262
+ if units is None:
263
+ path_units = name_ext + ".soft.pt"
264
+ units = torch.load(path_units)
265
+ units = units[0]
266
+ units = repeat_expand_2d(units,f0.size(0)).transpose(0,1)
267
+
268
+ units = units[start_frame : start_frame + units_frame_len]
269
+
270
+ # load volume
271
+ vol_key = 'aug_vol' if aug_flag else 'volume'
272
+ volume = data_buffer.get(vol_key)
273
+ volume_frames = volume[start_frame : start_frame + units_frame_len]
274
+
275
+ # load spk_id
276
+ spk_id = data_buffer.get('spk_id')
277
+
278
+ # load shift
279
+ aug_shift = torch.from_numpy(np.array([[aug_shift]])).float()
280
+
281
+ 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)
282
+
283
+ def __len__(self):
284
+ return len(self.paths)
diffusion/diffusion.py ADDED
@@ -0,0 +1,317 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from collections import deque
2
+ from functools import partial
3
+ from inspect import isfunction
4
+ import torch.nn.functional as F
5
+ import librosa.sequence
6
+ import numpy as np
7
+ import torch
8
+ from torch import nn
9
+ from tqdm import tqdm
10
+
11
+
12
+ def exists(x):
13
+ return x is not None
14
+
15
+
16
+ def default(val, d):
17
+ if exists(val):
18
+ return val
19
+ return d() if isfunction(d) else d
20
+
21
+
22
+ def extract(a, t, x_shape):
23
+ b, *_ = t.shape
24
+ out = a.gather(-1, t)
25
+ return out.reshape(b, *((1,) * (len(x_shape) - 1)))
26
+
27
+
28
+ def noise_like(shape, device, repeat=False):
29
+ repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
30
+ noise = lambda: torch.randn(shape, device=device)
31
+ return repeat_noise() if repeat else noise()
32
+
33
+
34
+ def linear_beta_schedule(timesteps, max_beta=0.02):
35
+ """
36
+ linear schedule
37
+ """
38
+ betas = np.linspace(1e-4, max_beta, timesteps)
39
+ return betas
40
+
41
+
42
+ def cosine_beta_schedule(timesteps, s=0.008):
43
+ """
44
+ cosine schedule
45
+ as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
46
+ """
47
+ steps = timesteps + 1
48
+ x = np.linspace(0, steps, steps)
49
+ alphas_cumprod = np.cos(((x / steps) + s) / (1 + s) * np.pi * 0.5) ** 2
50
+ alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
51
+ betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
52
+ return np.clip(betas, a_min=0, a_max=0.999)
53
+
54
+
55
+ beta_schedule = {
56
+ "cosine": cosine_beta_schedule,
57
+ "linear": linear_beta_schedule,
58
+ }
59
+
60
+
61
+ class GaussianDiffusion(nn.Module):
62
+ def __init__(self,
63
+ denoise_fn,
64
+ out_dims=128,
65
+ timesteps=1000,
66
+ k_step=1000,
67
+ max_beta=0.02,
68
+ spec_min=-12,
69
+ spec_max=2):
70
+ super().__init__()
71
+ self.denoise_fn = denoise_fn
72
+ self.out_dims = out_dims
73
+ betas = beta_schedule['linear'](timesteps, max_beta=max_beta)
74
+
75
+ alphas = 1. - betas
76
+ alphas_cumprod = np.cumprod(alphas, axis=0)
77
+ alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
78
+
79
+ timesteps, = betas.shape
80
+ self.num_timesteps = int(timesteps)
81
+ self.k_step = k_step
82
+
83
+ self.noise_list = deque(maxlen=4)
84
+
85
+ to_torch = partial(torch.tensor, dtype=torch.float32)
86
+
87
+ self.register_buffer('betas', to_torch(betas))
88
+ self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
89
+ self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
90
+
91
+ # calculations for diffusion q(x_t | x_{t-1}) and others
92
+ self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
93
+ self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
94
+ self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
95
+ self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
96
+ self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))
97
+
98
+ # calculations for posterior q(x_{t-1} | x_t, x_0)
99
+ posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
100
+ # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
101
+ self.register_buffer('posterior_variance', to_torch(posterior_variance))
102
+ # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
103
+ self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20))))
104
+ self.register_buffer('posterior_mean_coef1', to_torch(
105
+ betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)))
106
+ self.register_buffer('posterior_mean_coef2', to_torch(
107
+ (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod)))
108
+
109
+ self.register_buffer('spec_min', torch.FloatTensor([spec_min])[None, None, :out_dims])
110
+ self.register_buffer('spec_max', torch.FloatTensor([spec_max])[None, None, :out_dims])
111
+
112
+ def q_mean_variance(self, x_start, t):
113
+ mean = extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
114
+ variance = extract(1. - self.alphas_cumprod, t, x_start.shape)
115
+ log_variance = extract(self.log_one_minus_alphas_cumprod, t, x_start.shape)
116
+ return mean, variance, log_variance
117
+
118
+ def predict_start_from_noise(self, x_t, t, noise):
119
+ return (
120
+ extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
121
+ extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
122
+ )
123
+
124
+ def q_posterior(self, x_start, x_t, t):
125
+ posterior_mean = (
126
+ extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +
127
+ extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
128
+ )
129
+ posterior_variance = extract(self.posterior_variance, t, x_t.shape)
130
+ posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)
131
+ return posterior_mean, posterior_variance, posterior_log_variance_clipped
132
+
133
+ def p_mean_variance(self, x, t, cond):
134
+ noise_pred = self.denoise_fn(x, t, cond=cond)
135
+ x_recon = self.predict_start_from_noise(x, t=t, noise=noise_pred)
136
+
137
+ x_recon.clamp_(-1., 1.)
138
+
139
+ model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
140
+ return model_mean, posterior_variance, posterior_log_variance
141
+
142
+ @torch.no_grad()
143
+ def p_sample(self, x, t, cond, clip_denoised=True, repeat_noise=False):
144
+ b, *_, device = *x.shape, x.device
145
+ model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, cond=cond)
146
+ noise = noise_like(x.shape, device, repeat_noise)
147
+ # no noise when t == 0
148
+ nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
149
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
150
+
151
+ @torch.no_grad()
152
+ def p_sample_plms(self, x, t, interval, cond, clip_denoised=True, repeat_noise=False):
153
+ """
154
+ Use the PLMS method from
155
+ [Pseudo Numerical Methods for Diffusion Models on Manifolds](https://arxiv.org/abs/2202.09778).
156
+ """
157
+
158
+ def get_x_pred(x, noise_t, t):
159
+ a_t = extract(self.alphas_cumprod, t, x.shape)
160
+ a_prev = extract(self.alphas_cumprod, torch.max(t - interval, torch.zeros_like(t)), x.shape)
161
+ a_t_sq, a_prev_sq = a_t.sqrt(), a_prev.sqrt()
162
+
163
+ x_delta = (a_prev - a_t) * ((1 / (a_t_sq * (a_t_sq + a_prev_sq))) * x - 1 / (
164
+ a_t_sq * (((1 - a_prev) * a_t).sqrt() + ((1 - a_t) * a_prev).sqrt())) * noise_t)
165
+ x_pred = x + x_delta
166
+
167
+ return x_pred
168
+
169
+ noise_list = self.noise_list
170
+ noise_pred = self.denoise_fn(x, t, cond=cond)
171
+
172
+ if len(noise_list) == 0:
173
+ x_pred = get_x_pred(x, noise_pred, t)
174
+ noise_pred_prev = self.denoise_fn(x_pred, max(t - interval, 0), cond=cond)
175
+ noise_pred_prime = (noise_pred + noise_pred_prev) / 2
176
+ elif len(noise_list) == 1:
177
+ noise_pred_prime = (3 * noise_pred - noise_list[-1]) / 2
178
+ elif len(noise_list) == 2:
179
+ noise_pred_prime = (23 * noise_pred - 16 * noise_list[-1] + 5 * noise_list[-2]) / 12
180
+ else:
181
+ noise_pred_prime = (55 * noise_pred - 59 * noise_list[-1] + 37 * noise_list[-2] - 9 * noise_list[-3]) / 24
182
+
183
+ x_prev = get_x_pred(x, noise_pred_prime, t)
184
+ noise_list.append(noise_pred)
185
+
186
+ return x_prev
187
+
188
+ def q_sample(self, x_start, t, noise=None):
189
+ noise = default(noise, lambda: torch.randn_like(x_start))
190
+ return (
191
+ extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
192
+ extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
193
+ )
194
+
195
+ def p_losses(self, x_start, t, cond, noise=None, loss_type='l2'):
196
+ noise = default(noise, lambda: torch.randn_like(x_start))
197
+
198
+ x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
199
+ x_recon = self.denoise_fn(x_noisy, t, cond)
200
+
201
+ if loss_type == 'l1':
202
+ loss = (noise - x_recon).abs().mean()
203
+ elif loss_type == 'l2':
204
+ loss = F.mse_loss(noise, x_recon)
205
+ else:
206
+ raise NotImplementedError()
207
+
208
+ return loss
209
+
210
+ def forward(self,
211
+ condition,
212
+ gt_spec=None,
213
+ infer=True,
214
+ infer_speedup=10,
215
+ method='dpm-solver',
216
+ k_step=300,
217
+ use_tqdm=True):
218
+ """
219
+ conditioning diffusion, use fastspeech2 encoder output as the condition
220
+ """
221
+ cond = condition.transpose(1, 2)
222
+ b, device = condition.shape[0], condition.device
223
+
224
+ if not infer:
225
+ spec = self.norm_spec(gt_spec)
226
+ t = torch.randint(0, self.k_step, (b,), device=device).long()
227
+ norm_spec = spec.transpose(1, 2)[:, None, :, :] # [B, 1, M, T]
228
+ return self.p_losses(norm_spec, t, cond=cond)
229
+ else:
230
+ shape = (cond.shape[0], 1, self.out_dims, cond.shape[2])
231
+
232
+ if gt_spec is None:
233
+ t = self.k_step
234
+ x = torch.randn(shape, device=device)
235
+ else:
236
+ t = k_step
237
+ norm_spec = self.norm_spec(gt_spec)
238
+ norm_spec = norm_spec.transpose(1, 2)[:, None, :, :]
239
+ x = self.q_sample(x_start=norm_spec, t=torch.tensor([t - 1], device=device).long())
240
+
241
+ if method is not None and infer_speedup > 1:
242
+ if method == 'dpm-solver':
243
+ from .dpm_solver_pytorch import NoiseScheduleVP, model_wrapper, DPM_Solver
244
+ # 1. Define the noise schedule.
245
+ noise_schedule = NoiseScheduleVP(schedule='discrete', betas=self.betas[:t])
246
+
247
+ # 2. Convert your discrete-time `model` to the continuous-time
248
+ # noise prediction model. Here is an example for a diffusion model
249
+ # `model` with the noise prediction type ("noise") .
250
+ def my_wrapper(fn):
251
+ def wrapped(x, t, **kwargs):
252
+ ret = fn(x, t, **kwargs)
253
+ if use_tqdm:
254
+ self.bar.update(1)
255
+ return ret
256
+
257
+ return wrapped
258
+
259
+ model_fn = model_wrapper(
260
+ my_wrapper(self.denoise_fn),
261
+ noise_schedule,
262
+ model_type="noise", # or "x_start" or "v" or "score"
263
+ model_kwargs={"cond": cond}
264
+ )
265
+
266
+ # 3. Define dpm-solver and sample by singlestep DPM-Solver.
267
+ # (We recommend singlestep DPM-Solver for unconditional sampling)
268
+ # You can adjust the `steps` to balance the computation
269
+ # costs and the sample quality.
270
+ dpm_solver = DPM_Solver(model_fn, noise_schedule)
271
+
272
+ steps = t // infer_speedup
273
+ if use_tqdm:
274
+ self.bar = tqdm(desc="sample time step", total=steps)
275
+ x = dpm_solver.sample(
276
+ x,
277
+ steps=steps,
278
+ order=3,
279
+ skip_type="time_uniform",
280
+ method="singlestep",
281
+ )
282
+ if use_tqdm:
283
+ self.bar.close()
284
+ elif method == 'pndm':
285
+ self.noise_list = deque(maxlen=4)
286
+ if use_tqdm:
287
+ for i in tqdm(
288
+ reversed(range(0, t, infer_speedup)), desc='sample time step',
289
+ total=t // infer_speedup,
290
+ ):
291
+ x = self.p_sample_plms(
292
+ x, torch.full((b,), i, device=device, dtype=torch.long),
293
+ infer_speedup, cond=cond
294
+ )
295
+ else:
296
+ for i in reversed(range(0, t, infer_speedup)):
297
+ x = self.p_sample_plms(
298
+ x, torch.full((b,), i, device=device, dtype=torch.long),
299
+ infer_speedup, cond=cond
300
+ )
301
+ else:
302
+ raise NotImplementedError(method)
303
+ else:
304
+ if use_tqdm:
305
+ for i in tqdm(reversed(range(0, t)), desc='sample time step', total=t):
306
+ x = self.p_sample(x, torch.full((b,), i, device=device, dtype=torch.long), cond)
307
+ else:
308
+ for i in reversed(range(0, t)):
309
+ x = self.p_sample(x, torch.full((b,), i, device=device, dtype=torch.long), cond)
310
+ x = x.squeeze(1).transpose(1, 2) # [B, T, M]
311
+ return self.denorm_spec(x)
312
+
313
+ def norm_spec(self, x):
314
+ return (x - self.spec_min) / (self.spec_max - self.spec_min) * 2 - 1
315
+
316
+ def denorm_spec(self, x):
317
+ return (x + 1) / 2 * (self.spec_max - self.spec_min) + self.spec_min
diffusion/diffusion_onnx.py ADDED
@@ -0,0 +1,612 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from collections import deque
2
+ from functools import partial
3
+ from inspect import isfunction
4
+ import torch.nn.functional as F
5
+ import librosa.sequence
6
+ import numpy as np
7
+ from torch.nn import Conv1d
8
+ from torch.nn import Mish
9
+ import torch
10
+ from torch import nn
11
+ from tqdm import tqdm
12
+ import math
13
+
14
+
15
+ def exists(x):
16
+ return x is not None
17
+
18
+
19
+ def default(val, d):
20
+ if exists(val):
21
+ return val
22
+ return d() if isfunction(d) else d
23
+
24
+
25
+ def extract(a, t):
26
+ return a[t].reshape((1, 1, 1, 1))
27
+
28
+
29
+ def noise_like(shape, device, repeat=False):
30
+ repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
31
+ noise = lambda: torch.randn(shape, device=device)
32
+ return repeat_noise() if repeat else noise()
33
+
34
+
35
+ def linear_beta_schedule(timesteps, max_beta=0.02):
36
+ """
37
+ linear schedule
38
+ """
39
+ betas = np.linspace(1e-4, max_beta, timesteps)
40
+ return betas
41
+
42
+
43
+ def cosine_beta_schedule(timesteps, s=0.008):
44
+ """
45
+ cosine schedule
46
+ as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
47
+ """
48
+ steps = timesteps + 1
49
+ x = np.linspace(0, steps, steps)
50
+ alphas_cumprod = np.cos(((x / steps) + s) / (1 + s) * np.pi * 0.5) ** 2
51
+ alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
52
+ betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
53
+ return np.clip(betas, a_min=0, a_max=0.999)
54
+
55
+
56
+ beta_schedule = {
57
+ "cosine": cosine_beta_schedule,
58
+ "linear": linear_beta_schedule,
59
+ }
60
+
61
+
62
+ def extract_1(a, t):
63
+ return a[t].reshape((1, 1, 1, 1))
64
+
65
+
66
+ def predict_stage0(noise_pred, noise_pred_prev):
67
+ return (noise_pred + noise_pred_prev) / 2
68
+
69
+
70
+ def predict_stage1(noise_pred, noise_list):
71
+ return (noise_pred * 3
72
+ - noise_list[-1]) / 2
73
+
74
+
75
+ def predict_stage2(noise_pred, noise_list):
76
+ return (noise_pred * 23
77
+ - noise_list[-1] * 16
78
+ + noise_list[-2] * 5) / 12
79
+
80
+
81
+ def predict_stage3(noise_pred, noise_list):
82
+ return (noise_pred * 55
83
+ - noise_list[-1] * 59
84
+ + noise_list[-2] * 37
85
+ - noise_list[-3] * 9) / 24
86
+
87
+
88
+ class SinusoidalPosEmb(nn.Module):
89
+ def __init__(self, dim):
90
+ super().__init__()
91
+ self.dim = dim
92
+ self.half_dim = dim // 2
93
+ self.emb = 9.21034037 / (self.half_dim - 1)
94
+ self.emb = torch.exp(torch.arange(self.half_dim) * torch.tensor(-self.emb)).unsqueeze(0)
95
+ self.emb = self.emb.cpu()
96
+
97
+ def forward(self, x):
98
+ emb = self.emb * x
99
+ emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
100
+ return emb
101
+
102
+
103
+ class ResidualBlock(nn.Module):
104
+ def __init__(self, encoder_hidden, residual_channels, dilation):
105
+ super().__init__()
106
+ self.residual_channels = residual_channels
107
+ self.dilated_conv = Conv1d(residual_channels, 2 * residual_channels, 3, padding=dilation, dilation=dilation)
108
+ self.diffusion_projection = nn.Linear(residual_channels, residual_channels)
109
+ self.conditioner_projection = Conv1d(encoder_hidden, 2 * residual_channels, 1)
110
+ self.output_projection = Conv1d(residual_channels, 2 * residual_channels, 1)
111
+
112
+ def forward(self, x, conditioner, diffusion_step):
113
+ diffusion_step = self.diffusion_projection(diffusion_step).unsqueeze(-1)
114
+ conditioner = self.conditioner_projection(conditioner)
115
+ y = x + diffusion_step
116
+ y = self.dilated_conv(y) + conditioner
117
+
118
+ gate, filter_1 = torch.split(y, [self.residual_channels, self.residual_channels], dim=1)
119
+
120
+ y = torch.sigmoid(gate) * torch.tanh(filter_1)
121
+ y = self.output_projection(y)
122
+
123
+ residual, skip = torch.split(y, [self.residual_channels, self.residual_channels], dim=1)
124
+
125
+ return (x + residual) / 1.41421356, skip
126
+
127
+
128
+ class DiffNet(nn.Module):
129
+ def __init__(self, in_dims, n_layers, n_chans, n_hidden):
130
+ super().__init__()
131
+ self.encoder_hidden = n_hidden
132
+ self.residual_layers = n_layers
133
+ self.residual_channels = n_chans
134
+ self.input_projection = Conv1d(in_dims, self.residual_channels, 1)
135
+ self.diffusion_embedding = SinusoidalPosEmb(self.residual_channels)
136
+ dim = self.residual_channels
137
+ self.mlp = nn.Sequential(
138
+ nn.Linear(dim, dim * 4),
139
+ Mish(),
140
+ nn.Linear(dim * 4, dim)
141
+ )
142
+ self.residual_layers = nn.ModuleList([
143
+ ResidualBlock(self.encoder_hidden, self.residual_channels, 1)
144
+ for i in range(self.residual_layers)
145
+ ])
146
+ self.skip_projection = Conv1d(self.residual_channels, self.residual_channels, 1)
147
+ self.output_projection = Conv1d(self.residual_channels, in_dims, 1)
148
+ nn.init.zeros_(self.output_projection.weight)
149
+
150
+ def forward(self, spec, diffusion_step, cond):
151
+ x = spec.squeeze(0)
152
+ x = self.input_projection(x) # x [B, residual_channel, T]
153
+ x = F.relu(x)
154
+ # skip = torch.randn_like(x)
155
+ diffusion_step = diffusion_step.float()
156
+ diffusion_step = self.diffusion_embedding(diffusion_step)
157
+ diffusion_step = self.mlp(diffusion_step)
158
+
159
+ x, skip = self.residual_layers[0](x, cond, diffusion_step)
160
+ # noinspection PyTypeChecker
161
+ for layer in self.residual_layers[1:]:
162
+ x, skip_connection = layer.forward(x, cond, diffusion_step)
163
+ skip = skip + skip_connection
164
+ x = skip / math.sqrt(len(self.residual_layers))
165
+ x = self.skip_projection(x)
166
+ x = F.relu(x)
167
+ x = self.output_projection(x) # [B, 80, T]
168
+ return x.unsqueeze(1)
169
+
170
+
171
+ class AfterDiffusion(nn.Module):
172
+ def __init__(self, spec_max, spec_min, v_type='a'):
173
+ super().__init__()
174
+ self.spec_max = spec_max
175
+ self.spec_min = spec_min
176
+ self.type = v_type
177
+
178
+ def forward(self, x):
179
+ x = x.squeeze(1).permute(0, 2, 1)
180
+ mel_out = (x + 1) / 2 * (self.spec_max - self.spec_min) + self.spec_min
181
+ if self.type == 'nsf-hifigan-log10':
182
+ mel_out = mel_out * 0.434294
183
+ return mel_out.transpose(2, 1)
184
+
185
+
186
+ class Pred(nn.Module):
187
+ def __init__(self, alphas_cumprod):
188
+ super().__init__()
189
+ self.alphas_cumprod = alphas_cumprod
190
+
191
+ def forward(self, x_1, noise_t, t_1, t_prev):
192
+ a_t = extract(self.alphas_cumprod, t_1).cpu()
193
+ a_prev = extract(self.alphas_cumprod, t_prev).cpu()
194
+ a_t_sq, a_prev_sq = a_t.sqrt().cpu(), a_prev.sqrt().cpu()
195
+ x_delta = (a_prev - a_t) * ((1 / (a_t_sq * (a_t_sq + a_prev_sq))) * x_1 - 1 / (
196
+ a_t_sq * (((1 - a_prev) * a_t).sqrt() + ((1 - a_t) * a_prev).sqrt())) * noise_t)
197
+ x_pred = x_1 + x_delta.cpu()
198
+
199
+ return x_pred
200
+
201
+
202
+ class GaussianDiffusion(nn.Module):
203
+ def __init__(self,
204
+ out_dims=128,
205
+ n_layers=20,
206
+ n_chans=384,
207
+ n_hidden=256,
208
+ timesteps=1000,
209
+ k_step=1000,
210
+ max_beta=0.02,
211
+ spec_min=-12,
212
+ spec_max=2):
213
+ super().__init__()
214
+ self.denoise_fn = DiffNet(out_dims, n_layers, n_chans, n_hidden)
215
+ self.out_dims = out_dims
216
+ self.mel_bins = out_dims
217
+ self.n_hidden = n_hidden
218
+ betas = beta_schedule['linear'](timesteps, max_beta=max_beta)
219
+
220
+ alphas = 1. - betas
221
+ alphas_cumprod = np.cumprod(alphas, axis=0)
222
+ alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
223
+ timesteps, = betas.shape
224
+ self.num_timesteps = int(timesteps)
225
+ self.k_step = k_step
226
+
227
+ self.noise_list = deque(maxlen=4)
228
+
229
+ to_torch = partial(torch.tensor, dtype=torch.float32)
230
+
231
+ self.register_buffer('betas', to_torch(betas))
232
+ self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
233
+ self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
234
+
235
+ # calculations for diffusion q(x_t | x_{t-1}) and others
236
+ self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
237
+ self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
238
+ self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
239
+ self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
240
+ self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))
241
+
242
+ # calculations for posterior q(x_{t-1} | x_t, x_0)
243
+ posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
244
+ # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
245
+ self.register_buffer('posterior_variance', to_torch(posterior_variance))
246
+ # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
247
+ self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20))))
248
+ self.register_buffer('posterior_mean_coef1', to_torch(
249
+ betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)))
250
+ self.register_buffer('posterior_mean_coef2', to_torch(
251
+ (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod)))
252
+
253
+ self.register_buffer('spec_min', torch.FloatTensor([spec_min])[None, None, :out_dims])
254
+ self.register_buffer('spec_max', torch.FloatTensor([spec_max])[None, None, :out_dims])
255
+ self.ad = AfterDiffusion(self.spec_max, self.spec_min)
256
+ self.xp = Pred(self.alphas_cumprod)
257
+
258
+ def q_mean_variance(self, x_start, t):
259
+ mean = extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
260
+ variance = extract(1. - self.alphas_cumprod, t, x_start.shape)
261
+ log_variance = extract(self.log_one_minus_alphas_cumprod, t, x_start.shape)
262
+ return mean, variance, log_variance
263
+
264
+ def predict_start_from_noise(self, x_t, t, noise):
265
+ return (
266
+ extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
267
+ extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
268
+ )
269
+
270
+ def q_posterior(self, x_start, x_t, t):
271
+ posterior_mean = (
272
+ extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +
273
+ extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
274
+ )
275
+ posterior_variance = extract(self.posterior_variance, t, x_t.shape)
276
+ posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)
277
+ return posterior_mean, posterior_variance, posterior_log_variance_clipped
278
+
279
+ def p_mean_variance(self, x, t, cond):
280
+ noise_pred = self.denoise_fn(x, t, cond=cond)
281
+ x_recon = self.predict_start_from_noise(x, t=t, noise=noise_pred)
282
+
283
+ x_recon.clamp_(-1., 1.)
284
+
285
+ model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
286
+ return model_mean, posterior_variance, posterior_log_variance
287
+
288
+ @torch.no_grad()
289
+ def p_sample(self, x, t, cond, clip_denoised=True, repeat_noise=False):
290
+ b, *_, device = *x.shape, x.device
291
+ model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, cond=cond)
292
+ noise = noise_like(x.shape, device, repeat_noise)
293
+ # no noise when t == 0
294
+ nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
295
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
296
+
297
+ @torch.no_grad()
298
+ def p_sample_plms(self, x, t, interval, cond, clip_denoised=True, repeat_noise=False):
299
+ """
300
+ Use the PLMS method from
301
+ [Pseudo Numerical Methods for Diffusion Models on Manifolds](https://arxiv.org/abs/2202.09778).
302
+ """
303
+
304
+ def get_x_pred(x, noise_t, t):
305
+ a_t = extract(self.alphas_cumprod, t)
306
+ a_prev = extract(self.alphas_cumprod, torch.max(t - interval, torch.zeros_like(t)))
307
+ a_t_sq, a_prev_sq = a_t.sqrt(), a_prev.sqrt()
308
+
309
+ x_delta = (a_prev - a_t) * ((1 / (a_t_sq * (a_t_sq + a_prev_sq))) * x - 1 / (
310
+ a_t_sq * (((1 - a_prev) * a_t).sqrt() + ((1 - a_t) * a_prev).sqrt())) * noise_t)
311
+ x_pred = x + x_delta
312
+
313
+ return x_pred
314
+
315
+ noise_list = self.noise_list
316
+ noise_pred = self.denoise_fn(x, t, cond=cond)
317
+
318
+ if len(noise_list) == 0:
319
+ x_pred = get_x_pred(x, noise_pred, t)
320
+ noise_pred_prev = self.denoise_fn(x_pred, max(t - interval, 0), cond=cond)
321
+ noise_pred_prime = (noise_pred + noise_pred_prev) / 2
322
+ elif len(noise_list) == 1:
323
+ noise_pred_prime = (3 * noise_pred - noise_list[-1]) / 2
324
+ elif len(noise_list) == 2:
325
+ noise_pred_prime = (23 * noise_pred - 16 * noise_list[-1] + 5 * noise_list[-2]) / 12
326
+ else:
327
+ noise_pred_prime = (55 * noise_pred - 59 * noise_list[-1] + 37 * noise_list[-2] - 9 * noise_list[-3]) / 24
328
+
329
+ x_prev = get_x_pred(x, noise_pred_prime, t)
330
+ noise_list.append(noise_pred)
331
+
332
+ return x_prev
333
+
334
+ def q_sample(self, x_start, t, noise=None):
335
+ noise = default(noise, lambda: torch.randn_like(x_start))
336
+ return (
337
+ extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
338
+ extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
339
+ )
340
+
341
+ def p_losses(self, x_start, t, cond, noise=None, loss_type='l2'):
342
+ noise = default(noise, lambda: torch.randn_like(x_start))
343
+
344
+ x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
345
+ x_recon = self.denoise_fn(x_noisy, t, cond)
346
+
347
+ if loss_type == 'l1':
348
+ loss = (noise - x_recon).abs().mean()
349
+ elif loss_type == 'l2':
350
+ loss = F.mse_loss(noise, x_recon)
351
+ else:
352
+ raise NotImplementedError()
353
+
354
+ return loss
355
+
356
+ def org_forward(self,
357
+ condition,
358
+ init_noise=None,
359
+ gt_spec=None,
360
+ infer=True,
361
+ infer_speedup=100,
362
+ method='pndm',
363
+ k_step=1000,
364
+ use_tqdm=True):
365
+ """
366
+ conditioning diffusion, use fastspeech2 encoder output as the condition
367
+ """
368
+ cond = condition
369
+ b, device = condition.shape[0], condition.device
370
+ if not infer:
371
+ spec = self.norm_spec(gt_spec)
372
+ t = torch.randint(0, self.k_step, (b,), device=device).long()
373
+ norm_spec = spec.transpose(1, 2)[:, None, :, :] # [B, 1, M, T]
374
+ return self.p_losses(norm_spec, t, cond=cond)
375
+ else:
376
+ shape = (cond.shape[0], 1, self.out_dims, cond.shape[2])
377
+
378
+ if gt_spec is None:
379
+ t = self.k_step
380
+ if init_noise is None:
381
+ x = torch.randn(shape, device=device)
382
+ else:
383
+ x = init_noise
384
+ else:
385
+ t = k_step
386
+ norm_spec = self.norm_spec(gt_spec)
387
+ norm_spec = norm_spec.transpose(1, 2)[:, None, :, :]
388
+ x = self.q_sample(x_start=norm_spec, t=torch.tensor([t - 1], device=device).long())
389
+
390
+ if method is not None and infer_speedup > 1:
391
+ if method == 'dpm-solver':
392
+ from .dpm_solver_pytorch import NoiseScheduleVP, model_wrapper, DPM_Solver
393
+ # 1. Define the noise schedule.
394
+ noise_schedule = NoiseScheduleVP(schedule='discrete', betas=self.betas[:t])
395
+
396
+ # 2. Convert your discrete-time `model` to the continuous-time
397
+ # noise prediction model. Here is an example for a diffusion model
398
+ # `model` with the noise prediction type ("noise") .
399
+ def my_wrapper(fn):
400
+ def wrapped(x, t, **kwargs):
401
+ ret = fn(x, t, **kwargs)
402
+ if use_tqdm:
403
+ self.bar.update(1)
404
+ return ret
405
+
406
+ return wrapped
407
+
408
+ model_fn = model_wrapper(
409
+ my_wrapper(self.denoise_fn),
410
+ noise_schedule,
411
+ model_type="noise", # or "x_start" or "v" or "score"
412
+ model_kwargs={"cond": cond}
413
+ )
414
+
415
+ # 3. Define dpm-solver and sample by singlestep DPM-Solver.
416
+ # (We recommend singlestep DPM-Solver for unconditional sampling)
417
+ # You can adjust the `steps` to balance the computation
418
+ # costs and the sample quality.
419
+ dpm_solver = DPM_Solver(model_fn, noise_schedule)
420
+
421
+ steps = t // infer_speedup
422
+ if use_tqdm:
423
+ self.bar = tqdm(desc="sample time step", total=steps)
424
+ x = dpm_solver.sample(
425
+ x,
426
+ steps=steps,
427
+ order=3,
428
+ skip_type="time_uniform",
429
+ method="singlestep",
430
+ )
431
+ if use_tqdm:
432
+ self.bar.close()
433
+ elif method == 'pndm':
434
+ self.noise_list = deque(maxlen=4)
435
+ if use_tqdm:
436
+ for i in tqdm(
437
+ reversed(range(0, t, infer_speedup)), desc='sample time step',
438
+ total=t // infer_speedup,
439
+ ):
440
+ x = self.p_sample_plms(
441
+ x, torch.full((b,), i, device=device, dtype=torch.long),
442
+ infer_speedup, cond=cond
443
+ )
444
+ else:
445
+ for i in reversed(range(0, t, infer_speedup)):
446
+ x = self.p_sample_plms(
447
+ x, torch.full((b,), i, device=device, dtype=torch.long),
448
+ infer_speedup, cond=cond
449
+ )
450
+ else:
451
+ raise NotImplementedError(method)
452
+ else:
453
+ if use_tqdm:
454
+ for i in tqdm(reversed(range(0, t)), desc='sample time step', total=t):
455
+ x = self.p_sample(x, torch.full((b,), i, device=device, dtype=torch.long), cond)
456
+ else:
457
+ for i in reversed(range(0, t)):
458
+ x = self.p_sample(x, torch.full((b,), i, device=device, dtype=torch.long), cond)
459
+ x = x.squeeze(1).transpose(1, 2) # [B, T, M]
460
+ return self.denorm_spec(x).transpose(2, 1)
461
+
462
+ def norm_spec(self, x):
463
+ return (x - self.spec_min) / (self.spec_max - self.spec_min) * 2 - 1
464
+
465
+ def denorm_spec(self, x):
466
+ return (x + 1) / 2 * (self.spec_max - self.spec_min) + self.spec_min
467
+
468
+ def get_x_pred(self, x_1, noise_t, t_1, t_prev):
469
+ a_t = extract(self.alphas_cumprod, t_1)
470
+ a_prev = extract(self.alphas_cumprod, t_prev)
471
+ a_t_sq, a_prev_sq = a_t.sqrt(), a_prev.sqrt()
472
+ x_delta = (a_prev - a_t) * ((1 / (a_t_sq * (a_t_sq + a_prev_sq))) * x_1 - 1 / (
473
+ a_t_sq * (((1 - a_prev) * a_t).sqrt() + ((1 - a_t) * a_prev).sqrt())) * noise_t)
474
+ x_pred = x_1 + x_delta
475
+ return x_pred
476
+
477
+ def OnnxExport(self, project_name=None, init_noise=None, hidden_channels=256, export_denoise=True, export_pred=True, export_after=True):
478
+ cond = torch.randn([1, self.n_hidden, 10]).cpu()
479
+ if init_noise is None:
480
+ x = torch.randn((1, 1, self.mel_bins, cond.shape[2]), dtype=torch.float32).cpu()
481
+ else:
482
+ x = init_noise
483
+ pndms = 100
484
+
485
+ org_y_x = self.org_forward(cond, init_noise=x)
486
+
487
+ device = cond.device
488
+ n_frames = cond.shape[2]
489
+ step_range = torch.arange(0, self.k_step, pndms, dtype=torch.long, device=device).flip(0)
490
+ plms_noise_stage = torch.tensor(0, dtype=torch.long, device=device)
491
+ noise_list = torch.zeros((0, 1, 1, self.mel_bins, n_frames), device=device)
492
+
493
+ ot = step_range[0]
494
+ ot_1 = torch.full((1,), ot, device=device, dtype=torch.long)
495
+ if export_denoise:
496
+ torch.onnx.export(
497
+ self.denoise_fn,
498
+ (x.cpu(), ot_1.cpu(), cond.cpu()),
499
+ f"{project_name}_denoise.onnx",
500
+ input_names=["noise", "time", "condition"],
501
+ output_names=["noise_pred"],
502
+ dynamic_axes={
503
+ "noise": [3],
504
+ "condition": [2]
505
+ },
506
+ opset_version=16
507
+ )
508
+
509
+ for t in step_range:
510
+ t_1 = torch.full((1,), t, device=device, dtype=torch.long)
511
+ noise_pred = self.denoise_fn(x, t_1, cond)
512
+ t_prev = t_1 - pndms
513
+ t_prev = t_prev * (t_prev > 0)
514
+ if plms_noise_stage == 0:
515
+ if export_pred:
516
+ torch.onnx.export(
517
+ self.xp,
518
+ (x.cpu(), noise_pred.cpu(), t_1.cpu(), t_prev.cpu()),
519
+ f"{project_name}_pred.onnx",
520
+ input_names=["noise", "noise_pred", "time", "time_prev"],
521
+ output_names=["noise_pred_o"],
522
+ dynamic_axes={
523
+ "noise": [3],
524
+ "noise_pred": [3]
525
+ },
526
+ opset_version=16
527
+ )
528
+
529
+ x_pred = self.get_x_pred(x, noise_pred, t_1, t_prev)
530
+ noise_pred_prev = self.denoise_fn(x_pred, t_prev, cond=cond)
531
+ noise_pred_prime = predict_stage0(noise_pred, noise_pred_prev)
532
+
533
+ elif plms_noise_stage == 1:
534
+ noise_pred_prime = predict_stage1(noise_pred, noise_list)
535
+
536
+ elif plms_noise_stage == 2:
537
+ noise_pred_prime = predict_stage2(noise_pred, noise_list)
538
+
539
+ else:
540
+ noise_pred_prime = predict_stage3(noise_pred, noise_list)
541
+
542
+ noise_pred = noise_pred.unsqueeze(0)
543
+
544
+ if plms_noise_stage < 3:
545
+ noise_list = torch.cat((noise_list, noise_pred), dim=0)
546
+ plms_noise_stage = plms_noise_stage + 1
547
+
548
+ else:
549
+ noise_list = torch.cat((noise_list[-2:], noise_pred), dim=0)
550
+
551
+ x = self.get_x_pred(x, noise_pred_prime, t_1, t_prev)
552
+ if export_after:
553
+ torch.onnx.export(
554
+ self.ad,
555
+ x.cpu(),
556
+ f"{project_name}_after.onnx",
557
+ input_names=["x"],
558
+ output_names=["mel_out"],
559
+ dynamic_axes={
560
+ "x": [3]
561
+ },
562
+ opset_version=16
563
+ )
564
+ x = self.ad(x)
565
+
566
+ print((x == org_y_x).all())
567
+ return x
568
+
569
+ def forward(self, condition=None, init_noise=None, pndms=None, k_step=None):
570
+ cond = condition
571
+ x = init_noise
572
+
573
+ device = cond.device
574
+ n_frames = cond.shape[2]
575
+ step_range = torch.arange(0, k_step.item(), pndms.item(), dtype=torch.long, device=device).flip(0)
576
+ plms_noise_stage = torch.tensor(0, dtype=torch.long, device=device)
577
+ noise_list = torch.zeros((0, 1, 1, self.mel_bins, n_frames), device=device)
578
+
579
+ ot = step_range[0]
580
+ ot_1 = torch.full((1,), ot, device=device, dtype=torch.long)
581
+
582
+ for t in step_range:
583
+ t_1 = torch.full((1,), t, device=device, dtype=torch.long)
584
+ noise_pred = self.denoise_fn(x, t_1, cond)
585
+ t_prev = t_1 - pndms
586
+ t_prev = t_prev * (t_prev > 0)
587
+ if plms_noise_stage == 0:
588
+ x_pred = self.get_x_pred(x, noise_pred, t_1, t_prev)
589
+ noise_pred_prev = self.denoise_fn(x_pred, t_prev, cond=cond)
590
+ noise_pred_prime = predict_stage0(noise_pred, noise_pred_prev)
591
+
592
+ elif plms_noise_stage == 1:
593
+ noise_pred_prime = predict_stage1(noise_pred, noise_list)
594
+
595
+ elif plms_noise_stage == 2:
596
+ noise_pred_prime = predict_stage2(noise_pred, noise_list)
597
+
598
+ else:
599
+ noise_pred_prime = predict_stage3(noise_pred, noise_list)
600
+
601
+ noise_pred = noise_pred.unsqueeze(0)
602
+
603
+ if plms_noise_stage < 3:
604
+ noise_list = torch.cat((noise_list, noise_pred), dim=0)
605
+ plms_noise_stage = plms_noise_stage + 1
606
+
607
+ else:
608
+ noise_list = torch.cat((noise_list[-2:], noise_pred), dim=0)
609
+
610
+ x = self.get_x_pred(x, noise_pred_prime, t_1, t_prev)
611
+ x = self.ad(x)
612
+ return x
diffusion/dpm_solver_pytorch.py ADDED
@@ -0,0 +1,1201 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+
3
+ import torch
4
+
5
+
6
+ class NoiseScheduleVP:
7
+ def __init__(
8
+ self,
9
+ schedule='discrete',
10
+ betas=None,
11
+ alphas_cumprod=None,
12
+ continuous_beta_0=0.1,
13
+ continuous_beta_1=20.,
14
+ ):
15
+ """Create a wrapper class for the forward SDE (VP type).
16
+
17
+ ***
18
+ Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t.
19
+ We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for high-resolution images.
20
+ ***
21
+
22
+ The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ).
23
+ We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper).
24
+ Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have:
25
+
26
+ log_alpha_t = self.marginal_log_mean_coeff(t)
27
+ sigma_t = self.marginal_std(t)
28
+ lambda_t = self.marginal_lambda(t)
29
+
30
+ Moreover, as lambda(t) is an invertible function, we also support its inverse function:
31
+
32
+ t = self.inverse_lambda(lambda_t)
33
+
34
+ ===============================================================
35
+
36
+ We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and continuous-time DPMs (trained on t in [t_0, T]).
37
+
38
+ 1. For discrete-time DPMs:
39
+
40
+ For discrete-time DPMs trained on n = 0, 1, ..., N-1, we convert the discrete steps to continuous time steps by:
41
+ t_i = (i + 1) / N
42
+ e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1.
43
+ We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3.
44
+
45
+ Args:
46
+ betas: A `torch.Tensor`. The beta array for the discrete-time DPM. (See the original DDPM paper for details)
47
+ alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. (See the original DDPM paper for details)
48
+
49
+ Note that we always have alphas_cumprod = cumprod(betas). Therefore, we only need to set one of `betas` and `alphas_cumprod`.
50
+
51
+ **Important**: Please pay special attention for the args for `alphas_cumprod`:
52
+ The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. Specifically, DDPMs assume that
53
+ q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ).
54
+ Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. In fact, we have
55
+ alpha_{t_n} = \sqrt{\hat{alpha_n}},
56
+ and
57
+ log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}).
58
+
59
+
60
+ 2. For continuous-time DPMs:
61
+
62
+ We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). The hyperparameters for the noise
63
+ schedule are the default settings in DDPM and improved-DDPM:
64
+
65
+ Args:
66
+ beta_min: A `float` number. The smallest beta for the linear schedule.
67
+ beta_max: A `float` number. The largest beta for the linear schedule.
68
+ cosine_s: A `float` number. The hyperparameter in the cosine schedule.
69
+ cosine_beta_max: A `float` number. The hyperparameter in the cosine schedule.
70
+ T: A `float` number. The ending time of the forward process.
71
+
72
+ ===============================================================
73
+
74
+ Args:
75
+ schedule: A `str`. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs,
76
+ 'linear' or 'cosine' for continuous-time DPMs.
77
+ Returns:
78
+ A wrapper object of the forward SDE (VP type).
79
+
80
+ ===============================================================
81
+
82
+ Example:
83
+
84
+ # For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1):
85
+ >>> ns = NoiseScheduleVP('discrete', betas=betas)
86
+
87
+ # For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array for n = 0, 1, ..., N - 1):
88
+ >>> ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod)
89
+
90
+ # For continuous-time DPMs (VPSDE), linear schedule:
91
+ >>> ns = NoiseScheduleVP('linear', continuous_beta_0=0.1, continuous_beta_1=20.)
92
+
93
+ """
94
+
95
+ if schedule not in ['discrete', 'linear', 'cosine']:
96
+ raise ValueError(
97
+ "Unsupported noise schedule {}. The schedule needs to be 'discrete' or 'linear' or 'cosine'".format(
98
+ schedule))
99
+
100
+ self.schedule = schedule
101
+ if schedule == 'discrete':
102
+ if betas is not None:
103
+ log_alphas = 0.5 * torch.log(1 - betas).cumsum(dim=0)
104
+ else:
105
+ assert alphas_cumprod is not None
106
+ log_alphas = 0.5 * torch.log(alphas_cumprod)
107
+ self.total_N = len(log_alphas)
108
+ self.T = 1.
109
+ self.t_array = torch.linspace(0., 1., self.total_N + 1)[1:].reshape((1, -1))
110
+ self.log_alpha_array = log_alphas.reshape((1, -1,))
111
+ else:
112
+ self.total_N = 1000
113
+ self.beta_0 = continuous_beta_0
114
+ self.beta_1 = continuous_beta_1
115
+ self.cosine_s = 0.008
116
+ self.cosine_beta_max = 999.
117
+ self.cosine_t_max = math.atan(self.cosine_beta_max * (1. + self.cosine_s) / math.pi) * 2. * (
118
+ 1. + self.cosine_s) / math.pi - self.cosine_s
119
+ self.cosine_log_alpha_0 = math.log(math.cos(self.cosine_s / (1. + self.cosine_s) * math.pi / 2.))
120
+ self.schedule = schedule
121
+ if schedule == 'cosine':
122
+ # For the cosine schedule, T = 1 will have numerical issues. So we manually set the ending time T.
123
+ # Note that T = 0.9946 may be not the optimal setting. However, we find it works well.
124
+ self.T = 0.9946
125
+ else:
126
+ self.T = 1.
127
+
128
+ def marginal_log_mean_coeff(self, t):
129
+ """
130
+ Compute log(alpha_t) of a given continuous-time label t in [0, T].
131
+ """
132
+ if self.schedule == 'discrete':
133
+ return interpolate_fn(t.reshape((-1, 1)), self.t_array.to(t.device),
134
+ self.log_alpha_array.to(t.device)).reshape((-1))
135
+ elif self.schedule == 'linear':
136
+ return -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
137
+ elif self.schedule == 'cosine':
138
+ log_alpha_fn = lambda s: torch.log(torch.cos((s + self.cosine_s) / (1. + self.cosine_s) * math.pi / 2.))
139
+ log_alpha_t = log_alpha_fn(t) - self.cosine_log_alpha_0
140
+ return log_alpha_t
141
+
142
+ def marginal_alpha(self, t):
143
+ """
144
+ Compute alpha_t of a given continuous-time label t in [0, T].
145
+ """
146
+ return torch.exp(self.marginal_log_mean_coeff(t))
147
+
148
+ def marginal_std(self, t):
149
+ """
150
+ Compute sigma_t of a given continuous-time label t in [0, T].
151
+ """
152
+ return torch.sqrt(1. - torch.exp(2. * self.marginal_log_mean_coeff(t)))
153
+
154
+ def marginal_lambda(self, t):
155
+ """
156
+ Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T].
157
+ """
158
+ log_mean_coeff = self.marginal_log_mean_coeff(t)
159
+ log_std = 0.5 * torch.log(1. - torch.exp(2. * log_mean_coeff))
160
+ return log_mean_coeff - log_std
161
+
162
+ def inverse_lambda(self, lamb):
163
+ """
164
+ Compute the continuous-time label t in [0, T] of a given half-logSNR lambda_t.
165
+ """
166
+ if self.schedule == 'linear':
167
+ tmp = 2. * (self.beta_1 - self.beta_0) * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
168
+ Delta = self.beta_0 ** 2 + tmp
169
+ return tmp / (torch.sqrt(Delta) + self.beta_0) / (self.beta_1 - self.beta_0)
170
+ elif self.schedule == 'discrete':
171
+ log_alpha = -0.5 * torch.logaddexp(torch.zeros((1,)).to(lamb.device), -2. * lamb)
172
+ t = interpolate_fn(log_alpha.reshape((-1, 1)), torch.flip(self.log_alpha_array.to(lamb.device), [1]),
173
+ torch.flip(self.t_array.to(lamb.device), [1]))
174
+ return t.reshape((-1,))
175
+ else:
176
+ log_alpha = -0.5 * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
177
+ t_fn = lambda log_alpha_t: torch.arccos(torch.exp(log_alpha_t + self.cosine_log_alpha_0)) * 2. * (
178
+ 1. + self.cosine_s) / math.pi - self.cosine_s
179
+ t = t_fn(log_alpha)
180
+ return t
181
+
182
+
183
+ def model_wrapper(
184
+ model,
185
+ noise_schedule,
186
+ model_type="noise",
187
+ model_kwargs={},
188
+ guidance_type="uncond",
189
+ condition=None,
190
+ unconditional_condition=None,
191
+ guidance_scale=1.,
192
+ classifier_fn=None,
193
+ classifier_kwargs={},
194
+ ):
195
+ """Create a wrapper function for the noise prediction model.
196
+
197
+ DPM-Solver needs to solve the continuous-time diffusion ODEs. For DPMs trained on discrete-time labels, we need to
198
+ firstly wrap the model function to a noise prediction model that accepts the continuous time as the input.
199
+
200
+ We support four types of the diffusion model by setting `model_type`:
201
+
202
+ 1. "noise": noise prediction model. (Trained by predicting noise).
203
+
204
+ 2. "x_start": data prediction model. (Trained by predicting the data x_0 at time 0).
205
+
206
+ 3. "v": velocity prediction model. (Trained by predicting the velocity).
207
+ The "v" prediction is derivation detailed in Appendix D of [1], and is used in Imagen-Video [2].
208
+
209
+ [1] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models."
210
+ arXiv preprint arXiv:2202.00512 (2022).
211
+ [2] Ho, Jonathan, et al. "Imagen Video: High Definition Video Generation with Diffusion Models."
212
+ arXiv preprint arXiv:2210.02303 (2022).
213
+
214
+ 4. "score": marginal score function. (Trained by denoising score matching).
215
+ Note that the score function and the noise prediction model follows a simple relationship:
216
+ ```
217
+ noise(x_t, t) = -sigma_t * score(x_t, t)
218
+ ```
219
+
220
+ We support three types of guided sampling by DPMs by setting `guidance_type`:
221
+ 1. "uncond": unconditional sampling by DPMs.
222
+ The input `model` has the following format:
223
+ ``
224
+ model(x, t_input, **model_kwargs) -> noise | x_start | v | score
225
+ ``
226
+
227
+ 2. "classifier": classifier guidance sampling [3] by DPMs and another classifier.
228
+ The input `model` has the following format:
229
+ ``
230
+ model(x, t_input, **model_kwargs) -> noise | x_start | v | score
231
+ ``
232
+
233
+ The input `classifier_fn` has the following format:
234
+ ``
235
+ classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond)
236
+ ``
237
+
238
+ [3] P. Dhariwal and A. Q. Nichol, "Diffusion models beat GANs on image synthesis,"
239
+ in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 8780-8794.
240
+
241
+ 3. "classifier-free": classifier-free guidance sampling by conditional DPMs.
242
+ The input `model` has the following format:
243
+ ``
244
+ model(x, t_input, cond, **model_kwargs) -> noise | x_start | v | score
245
+ ``
246
+ And if cond == `unconditional_condition`, the model output is the unconditional DPM output.
247
+
248
+ [4] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance."
249
+ arXiv preprint arXiv:2207.12598 (2022).
250
+
251
+
252
+ The `t_input` is the time label of the model, which may be discrete-time labels (i.e. 0 to 999)
253
+ or continuous-time labels (i.e. epsilon to T).
254
+
255
+ We wrap the model function to accept only `x` and `t_continuous` as inputs, and outputs the predicted noise:
256
+ ``
257
+ def model_fn(x, t_continuous) -> noise:
258
+ t_input = get_model_input_time(t_continuous)
259
+ return noise_pred(model, x, t_input, **model_kwargs)
260
+ ``
261
+ where `t_continuous` is the continuous time labels (i.e. epsilon to T). And we use `model_fn` for DPM-Solver.
262
+
263
+ ===============================================================
264
+
265
+ Args:
266
+ model: A diffusion model with the corresponding format described above.
267
+ noise_schedule: A noise schedule object, such as NoiseScheduleVP.
268
+ model_type: A `str`. The parameterization type of the diffusion model.
269
+ "noise" or "x_start" or "v" or "score".
270
+ model_kwargs: A `dict`. A dict for the other inputs of the model function.
271
+ guidance_type: A `str`. The type of the guidance for sampling.
272
+ "uncond" or "classifier" or "classifier-free".
273
+ condition: A pytorch tensor. The condition for the guided sampling.
274
+ Only used for "classifier" or "classifier-free" guidance type.
275
+ unconditional_condition: A pytorch tensor. The condition for the unconditional sampling.
276
+ Only used for "classifier-free" guidance type.
277
+ guidance_scale: A `float`. The scale for the guided sampling.
278
+ classifier_fn: A classifier function. Only used for the classifier guidance.
279
+ classifier_kwargs: A `dict`. A dict for the other inputs of the classifier function.
280
+ Returns:
281
+ A noise prediction model that accepts the noised data and the continuous time as the inputs.
282
+ """
283
+
284
+ def get_model_input_time(t_continuous):
285
+ """
286
+ Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time.
287
+ For discrete-time DPMs, we convert `t_continuous` in [1 / N, 1] to `t_input` in [0, 1000 * (N - 1) / N].
288
+ For continuous-time DPMs, we just use `t_continuous`.
289
+ """
290
+ if noise_schedule.schedule == 'discrete':
291
+ return (t_continuous - 1. / noise_schedule.total_N) * noise_schedule.total_N
292
+ else:
293
+ return t_continuous
294
+
295
+ def noise_pred_fn(x, t_continuous, cond=None):
296
+ if t_continuous.reshape((-1,)).shape[0] == 1:
297
+ t_continuous = t_continuous.expand((x.shape[0]))
298
+ t_input = get_model_input_time(t_continuous)
299
+ if cond is None:
300
+ output = model(x, t_input, **model_kwargs)
301
+ else:
302
+ output = model(x, t_input, cond, **model_kwargs)
303
+ if model_type == "noise":
304
+ return output
305
+ elif model_type == "x_start":
306
+ alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
307
+ dims = x.dim()
308
+ return (x - expand_dims(alpha_t, dims) * output) / expand_dims(sigma_t, dims)
309
+ elif model_type == "v":
310
+ alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
311
+ dims = x.dim()
312
+ return expand_dims(alpha_t, dims) * output + expand_dims(sigma_t, dims) * x
313
+ elif model_type == "score":
314
+ sigma_t = noise_schedule.marginal_std(t_continuous)
315
+ dims = x.dim()
316
+ return -expand_dims(sigma_t, dims) * output
317
+
318
+ def cond_grad_fn(x, t_input):
319
+ """
320
+ Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t).
321
+ """
322
+ with torch.enable_grad():
323
+ x_in = x.detach().requires_grad_(True)
324
+ log_prob = classifier_fn(x_in, t_input, condition, **classifier_kwargs)
325
+ return torch.autograd.grad(log_prob.sum(), x_in)[0]
326
+
327
+ def model_fn(x, t_continuous):
328
+ """
329
+ The noise predicition model function that is used for DPM-Solver.
330
+ """
331
+ if t_continuous.reshape((-1,)).shape[0] == 1:
332
+ t_continuous = t_continuous.expand((x.shape[0]))
333
+ if guidance_type == "uncond":
334
+ return noise_pred_fn(x, t_continuous)
335
+ elif guidance_type == "classifier":
336
+ assert classifier_fn is not None
337
+ t_input = get_model_input_time(t_continuous)
338
+ cond_grad = cond_grad_fn(x, t_input)
339
+ sigma_t = noise_schedule.marginal_std(t_continuous)
340
+ noise = noise_pred_fn(x, t_continuous)
341
+ return noise - guidance_scale * expand_dims(sigma_t, dims=cond_grad.dim()) * cond_grad
342
+ elif guidance_type == "classifier-free":
343
+ if guidance_scale == 1. or unconditional_condition is None:
344
+ return noise_pred_fn(x, t_continuous, cond=condition)
345
+ else:
346
+ x_in = torch.cat([x] * 2)
347
+ t_in = torch.cat([t_continuous] * 2)
348
+ c_in = torch.cat([unconditional_condition, condition])
349
+ noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
350
+ return noise_uncond + guidance_scale * (noise - noise_uncond)
351
+
352
+ assert model_type in ["noise", "x_start", "v"]
353
+ assert guidance_type in ["uncond", "classifier", "classifier-free"]
354
+ return model_fn
355
+
356
+
357
+ class DPM_Solver:
358
+ def __init__(self, model_fn, noise_schedule, predict_x0=False, thresholding=False, max_val=1.):
359
+ """Construct a DPM-Solver.
360
+
361
+ We support both the noise prediction model ("predicting epsilon") and the data prediction model ("predicting x0").
362
+ If `predict_x0` is False, we use the solver for the noise prediction model (DPM-Solver).
363
+ If `predict_x0` is True, we use the solver for the data prediction model (DPM-Solver++).
364
+ In such case, we further support the "dynamic thresholding" in [1] when `thresholding` is True.
365
+ The "dynamic thresholding" can greatly improve the sample quality for pixel-space DPMs with large guidance scales.
366
+
367
+ Args:
368
+ model_fn: A noise prediction model function which accepts the continuous-time input (t in [epsilon, T]):
369
+ ``
370
+ def model_fn(x, t_continuous):
371
+ return noise
372
+ ``
373
+ noise_schedule: A noise schedule object, such as NoiseScheduleVP.
374
+ predict_x0: A `bool`. If true, use the data prediction model; else, use the noise prediction model.
375
+ thresholding: A `bool`. Valid when `predict_x0` is True. Whether to use the "dynamic thresholding" in [1].
376
+ max_val: A `float`. Valid when both `predict_x0` and `thresholding` are True. The max value for thresholding.
377
+
378
+ [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.
379
+ """
380
+ self.model = model_fn
381
+ self.noise_schedule = noise_schedule
382
+ self.predict_x0 = predict_x0
383
+ self.thresholding = thresholding
384
+ self.max_val = max_val
385
+
386
+ def noise_prediction_fn(self, x, t):
387
+ """
388
+ Return the noise prediction model.
389
+ """
390
+ return self.model(x, t)
391
+
392
+ def data_prediction_fn(self, x, t):
393
+ """
394
+ Return the data prediction model (with thresholding).
395
+ """
396
+ noise = self.noise_prediction_fn(x, t)
397
+ dims = x.dim()
398
+ alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t)
399
+ x0 = (x - expand_dims(sigma_t, dims) * noise) / expand_dims(alpha_t, dims)
400
+ if self.thresholding:
401
+ p = 0.995 # A hyperparameter in the paper of "Imagen" [1].
402
+ s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
403
+ s = expand_dims(torch.maximum(s, self.max_val * torch.ones_like(s).to(s.device)), dims)
404
+ x0 = torch.clamp(x0, -s, s) / s
405
+ return x0
406
+
407
+ def model_fn(self, x, t):
408
+ """
409
+ Convert the model to the noise prediction model or the data prediction model.
410
+ """
411
+ if self.predict_x0:
412
+ return self.data_prediction_fn(x, t)
413
+ else:
414
+ return self.noise_prediction_fn(x, t)
415
+
416
+ def get_time_steps(self, skip_type, t_T, t_0, N, device):
417
+ """Compute the intermediate time steps for sampling.
418
+
419
+ Args:
420
+ skip_type: A `str`. The type for the spacing of the time steps. We support three types:
421
+ - 'logSNR': uniform logSNR for the time steps.
422
+ - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.)
423
+ - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)
424
+ t_T: A `float`. The starting time of the sampling (default is T).
425
+ t_0: A `float`. The ending time of the sampling (default is epsilon).
426
+ N: A `int`. The total number of the spacing of the time steps.
427
+ device: A torch device.
428
+ Returns:
429
+ A pytorch tensor of the time steps, with the shape (N + 1,).
430
+ """
431
+ if skip_type == 'logSNR':
432
+ lambda_T = self.noise_schedule.marginal_lambda(torch.tensor(t_T).to(device))
433
+ lambda_0 = self.noise_schedule.marginal_lambda(torch.tensor(t_0).to(device))
434
+ logSNR_steps = torch.linspace(lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1).to(device)
435
+ return self.noise_schedule.inverse_lambda(logSNR_steps)
436
+ elif skip_type == 'time_uniform':
437
+ return torch.linspace(t_T, t_0, N + 1).to(device)
438
+ elif skip_type == 'time_quadratic':
439
+ t_order = 2
440
+ t = torch.linspace(t_T ** (1. / t_order), t_0 ** (1. / t_order), N + 1).pow(t_order).to(device)
441
+ return t
442
+ else:
443
+ raise ValueError(
444
+ "Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(skip_type))
445
+
446
+ def get_orders_and_timesteps_for_singlestep_solver(self, steps, order, skip_type, t_T, t_0, device):
447
+ """
448
+ Get the order of each step for sampling by the singlestep DPM-Solver.
449
+
450
+ We combine both DPM-Solver-1,2,3 to use all the function evaluations, which is named as "DPM-Solver-fast".
451
+ Given a fixed number of function evaluations by `steps`, the sampling procedure by DPM-Solver-fast is:
452
+ - If order == 1:
453
+ We take `steps` of DPM-Solver-1 (i.e. DDIM).
454
+ - If order == 2:
455
+ - Denote K = (steps // 2). We take K or (K + 1) intermediate time steps for sampling.
456
+ - If steps % 2 == 0, we use K steps of DPM-Solver-2.
457
+ - If steps % 2 == 1, we use K steps of DPM-Solver-2 and 1 step of DPM-Solver-1.
458
+ - If order == 3:
459
+ - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling.
460
+ - 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.
461
+ - If steps % 3 == 1, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-1.
462
+ - If steps % 3 == 2, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-2.
463
+
464
+ ============================================
465
+ Args:
466
+ order: A `int`. The max order for the solver (2 or 3).
467
+ steps: A `int`. The total number of function evaluations (NFE).
468
+ skip_type: A `str`. The type for the spacing of the time steps. We support three types:
469
+ - 'logSNR': uniform logSNR for the time steps.
470
+ - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.)
471
+ - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)
472
+ t_T: A `float`. The starting time of the sampling (default is T).
473
+ t_0: A `float`. The ending time of the sampling (default is epsilon).
474
+ device: A torch device.
475
+ Returns:
476
+ orders: A list of the solver order of each step.
477
+ """
478
+ if order == 3:
479
+ K = steps // 3 + 1
480
+ if steps % 3 == 0:
481
+ orders = [3, ] * (K - 2) + [2, 1]
482
+ elif steps % 3 == 1:
483
+ orders = [3, ] * (K - 1) + [1]
484
+ else:
485
+ orders = [3, ] * (K - 1) + [2]
486
+ elif order == 2:
487
+ if steps % 2 == 0:
488
+ K = steps // 2
489
+ orders = [2, ] * K
490
+ else:
491
+ K = steps // 2 + 1
492
+ orders = [2, ] * (K - 1) + [1]
493
+ elif order == 1:
494
+ K = 1
495
+ orders = [1, ] * steps
496
+ else:
497
+ raise ValueError("'order' must be '1' or '2' or '3'.")
498
+ if skip_type == 'logSNR':
499
+ # To reproduce the results in DPM-Solver paper
500
+ timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, K, device)
501
+ else:
502
+ timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, steps, device)[
503
+ torch.cumsum(torch.tensor([0, ] + orders), dim=0).to(device)]
504
+ return timesteps_outer, orders
505
+
506
+ def denoise_fn(self, x, s):
507
+ """
508
+ Denoise at the final step, which is equivalent to solve the ODE from lambda_s to infty by first-order discretization.
509
+ """
510
+ return self.data_prediction_fn(x, s)
511
+
512
+ def dpm_solver_first_update(self, x, s, t, model_s=None, return_intermediate=False):
513
+ """
514
+ DPM-Solver-1 (equivalent to DDIM) from time `s` to time `t`.
515
+
516
+ Args:
517
+ x: A pytorch tensor. The initial value at time `s`.
518
+ s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
519
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
520
+ model_s: A pytorch tensor. The model function evaluated at time `s`.
521
+ If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
522
+ return_intermediate: A `bool`. If true, also return the model value at time `s`.
523
+ Returns:
524
+ x_t: A pytorch tensor. The approximated solution at time `t`.
525
+ """
526
+ ns = self.noise_schedule
527
+ dims = x.dim()
528
+ lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
529
+ h = lambda_t - lambda_s
530
+ log_alpha_s, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(t)
531
+ sigma_s, sigma_t = ns.marginal_std(s), ns.marginal_std(t)
532
+ alpha_t = torch.exp(log_alpha_t)
533
+
534
+ if self.predict_x0:
535
+ phi_1 = torch.expm1(-h)
536
+ if model_s is None:
537
+ model_s = self.model_fn(x, s)
538
+ x_t = (
539
+ expand_dims(sigma_t / sigma_s, dims) * x
540
+ - expand_dims(alpha_t * phi_1, dims) * model_s
541
+ )
542
+ if return_intermediate:
543
+ return x_t, {'model_s': model_s}
544
+ else:
545
+ return x_t
546
+ else:
547
+ phi_1 = torch.expm1(h)
548
+ if model_s is None:
549
+ model_s = self.model_fn(x, s)
550
+ x_t = (
551
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
552
+ - expand_dims(sigma_t * phi_1, dims) * model_s
553
+ )
554
+ if return_intermediate:
555
+ return x_t, {'model_s': model_s}
556
+ else:
557
+ return x_t
558
+
559
+ def singlestep_dpm_solver_second_update(self, x, s, t, r1=0.5, model_s=None, return_intermediate=False,
560
+ solver_type='dpm_solver'):
561
+ """
562
+ Singlestep solver DPM-Solver-2 from time `s` to time `t`.
563
+
564
+ Args:
565
+ x: A pytorch tensor. The initial value at time `s`.
566
+ s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
567
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
568
+ r1: A `float`. The hyperparameter of the second-order solver.
569
+ model_s: A pytorch tensor. The model function evaluated at time `s`.
570
+ If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
571
+ return_intermediate: A `bool`. If true, also return the model value at time `s` and `s1` (the intermediate time).
572
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
573
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
574
+ Returns:
575
+ x_t: A pytorch tensor. The approximated solution at time `t`.
576
+ """
577
+ if solver_type not in ['dpm_solver', 'taylor']:
578
+ raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
579
+ if r1 is None:
580
+ r1 = 0.5
581
+ ns = self.noise_schedule
582
+ dims = x.dim()
583
+ lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
584
+ h = lambda_t - lambda_s
585
+ lambda_s1 = lambda_s + r1 * h
586
+ s1 = ns.inverse_lambda(lambda_s1)
587
+ log_alpha_s, log_alpha_s1, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(
588
+ s1), ns.marginal_log_mean_coeff(t)
589
+ sigma_s, sigma_s1, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(t)
590
+ alpha_s1, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_t)
591
+
592
+ if self.predict_x0:
593
+ phi_11 = torch.expm1(-r1 * h)
594
+ phi_1 = torch.expm1(-h)
595
+
596
+ if model_s is None:
597
+ model_s = self.model_fn(x, s)
598
+ x_s1 = (
599
+ expand_dims(sigma_s1 / sigma_s, dims) * x
600
+ - expand_dims(alpha_s1 * phi_11, dims) * model_s
601
+ )
602
+ model_s1 = self.model_fn(x_s1, s1)
603
+ if solver_type == 'dpm_solver':
604
+ x_t = (
605
+ expand_dims(sigma_t / sigma_s, dims) * x
606
+ - expand_dims(alpha_t * phi_1, dims) * model_s
607
+ - (0.5 / r1) * expand_dims(alpha_t * phi_1, dims) * (model_s1 - model_s)
608
+ )
609
+ elif solver_type == 'taylor':
610
+ x_t = (
611
+ expand_dims(sigma_t / sigma_s, dims) * x
612
+ - expand_dims(alpha_t * phi_1, dims) * model_s
613
+ + (1. / r1) * expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * (
614
+ model_s1 - model_s)
615
+ )
616
+ else:
617
+ phi_11 = torch.expm1(r1 * h)
618
+ phi_1 = torch.expm1(h)
619
+
620
+ if model_s is None:
621
+ model_s = self.model_fn(x, s)
622
+ x_s1 = (
623
+ expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x
624
+ - expand_dims(sigma_s1 * phi_11, dims) * model_s
625
+ )
626
+ model_s1 = self.model_fn(x_s1, s1)
627
+ if solver_type == 'dpm_solver':
628
+ x_t = (
629
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
630
+ - expand_dims(sigma_t * phi_1, dims) * model_s
631
+ - (0.5 / r1) * expand_dims(sigma_t * phi_1, dims) * (model_s1 - model_s)
632
+ )
633
+ elif solver_type == 'taylor':
634
+ x_t = (
635
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
636
+ - expand_dims(sigma_t * phi_1, dims) * model_s
637
+ - (1. / r1) * expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * (model_s1 - model_s)
638
+ )
639
+ if return_intermediate:
640
+ return x_t, {'model_s': model_s, 'model_s1': model_s1}
641
+ else:
642
+ return x_t
643
+
644
+ def singlestep_dpm_solver_third_update(self, x, s, t, r1=1. / 3., r2=2. / 3., model_s=None, model_s1=None,
645
+ return_intermediate=False, solver_type='dpm_solver'):
646
+ """
647
+ Singlestep solver DPM-Solver-3 from time `s` to time `t`.
648
+
649
+ Args:
650
+ x: A pytorch tensor. The initial value at time `s`.
651
+ s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
652
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
653
+ r1: A `float`. The hyperparameter of the third-order solver.
654
+ r2: A `float`. The hyperparameter of the third-order solver.
655
+ model_s: A pytorch tensor. The model function evaluated at time `s`.
656
+ If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
657
+ model_s1: A pytorch tensor. The model function evaluated at time `s1` (the intermediate time given by `r1`).
658
+ If `model_s1` is None, we evaluate the model at `s1`; otherwise we directly use it.
659
+ return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times).
660
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
661
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
662
+ Returns:
663
+ x_t: A pytorch tensor. The approximated solution at time `t`.
664
+ """
665
+ if solver_type not in ['dpm_solver', 'taylor']:
666
+ raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
667
+ if r1 is None:
668
+ r1 = 1. / 3.
669
+ if r2 is None:
670
+ r2 = 2. / 3.
671
+ ns = self.noise_schedule
672
+ dims = x.dim()
673
+ lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
674
+ h = lambda_t - lambda_s
675
+ lambda_s1 = lambda_s + r1 * h
676
+ lambda_s2 = lambda_s + r2 * h
677
+ s1 = ns.inverse_lambda(lambda_s1)
678
+ s2 = ns.inverse_lambda(lambda_s2)
679
+ log_alpha_s, log_alpha_s1, log_alpha_s2, log_alpha_t = ns.marginal_log_mean_coeff(
680
+ s), ns.marginal_log_mean_coeff(s1), ns.marginal_log_mean_coeff(s2), ns.marginal_log_mean_coeff(t)
681
+ sigma_s, sigma_s1, sigma_s2, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(
682
+ s2), ns.marginal_std(t)
683
+ alpha_s1, alpha_s2, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_s2), torch.exp(log_alpha_t)
684
+
685
+ if self.predict_x0:
686
+ phi_11 = torch.expm1(-r1 * h)
687
+ phi_12 = torch.expm1(-r2 * h)
688
+ phi_1 = torch.expm1(-h)
689
+ phi_22 = torch.expm1(-r2 * h) / (r2 * h) + 1.
690
+ phi_2 = phi_1 / h + 1.
691
+ phi_3 = phi_2 / h - 0.5
692
+
693
+ if model_s is None:
694
+ model_s = self.model_fn(x, s)
695
+ if model_s1 is None:
696
+ x_s1 = (
697
+ expand_dims(sigma_s1 / sigma_s, dims) * x
698
+ - expand_dims(alpha_s1 * phi_11, dims) * model_s
699
+ )
700
+ model_s1 = self.model_fn(x_s1, s1)
701
+ x_s2 = (
702
+ expand_dims(sigma_s2 / sigma_s, dims) * x
703
+ - expand_dims(alpha_s2 * phi_12, dims) * model_s
704
+ + r2 / r1 * expand_dims(alpha_s2 * phi_22, dims) * (model_s1 - model_s)
705
+ )
706
+ model_s2 = self.model_fn(x_s2, s2)
707
+ if solver_type == 'dpm_solver':
708
+ x_t = (
709
+ expand_dims(sigma_t / sigma_s, dims) * x
710
+ - expand_dims(alpha_t * phi_1, dims) * model_s
711
+ + (1. / r2) * expand_dims(alpha_t * phi_2, dims) * (model_s2 - model_s)
712
+ )
713
+ elif solver_type == 'taylor':
714
+ D1_0 = (1. / r1) * (model_s1 - model_s)
715
+ D1_1 = (1. / r2) * (model_s2 - model_s)
716
+ D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)
717
+ D2 = 2. * (D1_1 - D1_0) / (r2 - r1)
718
+ x_t = (
719
+ expand_dims(sigma_t / sigma_s, dims) * x
720
+ - expand_dims(alpha_t * phi_1, dims) * model_s
721
+ + expand_dims(alpha_t * phi_2, dims) * D1
722
+ - expand_dims(alpha_t * phi_3, dims) * D2
723
+ )
724
+ else:
725
+ phi_11 = torch.expm1(r1 * h)
726
+ phi_12 = torch.expm1(r2 * h)
727
+ phi_1 = torch.expm1(h)
728
+ phi_22 = torch.expm1(r2 * h) / (r2 * h) - 1.
729
+ phi_2 = phi_1 / h - 1.
730
+ phi_3 = phi_2 / h - 0.5
731
+
732
+ if model_s is None:
733
+ model_s = self.model_fn(x, s)
734
+ if model_s1 is None:
735
+ x_s1 = (
736
+ expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x
737
+ - expand_dims(sigma_s1 * phi_11, dims) * model_s
738
+ )
739
+ model_s1 = self.model_fn(x_s1, s1)
740
+ x_s2 = (
741
+ expand_dims(torch.exp(log_alpha_s2 - log_alpha_s), dims) * x
742
+ - expand_dims(sigma_s2 * phi_12, dims) * model_s
743
+ - r2 / r1 * expand_dims(sigma_s2 * phi_22, dims) * (model_s1 - model_s)
744
+ )
745
+ model_s2 = self.model_fn(x_s2, s2)
746
+ if solver_type == 'dpm_solver':
747
+ x_t = (
748
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
749
+ - expand_dims(sigma_t * phi_1, dims) * model_s
750
+ - (1. / r2) * expand_dims(sigma_t * phi_2, dims) * (model_s2 - model_s)
751
+ )
752
+ elif solver_type == 'taylor':
753
+ D1_0 = (1. / r1) * (model_s1 - model_s)
754
+ D1_1 = (1. / r2) * (model_s2 - model_s)
755
+ D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)
756
+ D2 = 2. * (D1_1 - D1_0) / (r2 - r1)
757
+ x_t = (
758
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
759
+ - expand_dims(sigma_t * phi_1, dims) * model_s
760
+ - expand_dims(sigma_t * phi_2, dims) * D1
761
+ - expand_dims(sigma_t * phi_3, dims) * D2
762
+ )
763
+
764
+ if return_intermediate:
765
+ return x_t, {'model_s': model_s, 'model_s1': model_s1, 'model_s2': model_s2}
766
+ else:
767
+ return x_t
768
+
769
+ def multistep_dpm_solver_second_update(self, x, model_prev_list, t_prev_list, t, solver_type="dpm_solver"):
770
+ """
771
+ Multistep solver DPM-Solver-2 from time `t_prev_list[-1]` to time `t`.
772
+
773
+ Args:
774
+ x: A pytorch tensor. The initial value at time `s`.
775
+ model_prev_list: A list of pytorch tensor. The previous computed model values.
776
+ t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)
777
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
778
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
779
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
780
+ Returns:
781
+ x_t: A pytorch tensor. The approximated solution at time `t`.
782
+ """
783
+ if solver_type not in ['dpm_solver', 'taylor']:
784
+ raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
785
+ ns = self.noise_schedule
786
+ dims = x.dim()
787
+ model_prev_1, model_prev_0 = model_prev_list
788
+ t_prev_1, t_prev_0 = t_prev_list
789
+ lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_1), ns.marginal_lambda(
790
+ t_prev_0), ns.marginal_lambda(t)
791
+ log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
792
+ sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
793
+ alpha_t = torch.exp(log_alpha_t)
794
+
795
+ h_0 = lambda_prev_0 - lambda_prev_1
796
+ h = lambda_t - lambda_prev_0
797
+ r0 = h_0 / h
798
+ D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1)
799
+ if self.predict_x0:
800
+ if solver_type == 'dpm_solver':
801
+ x_t = (
802
+ expand_dims(sigma_t / sigma_prev_0, dims) * x
803
+ - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
804
+ - 0.5 * expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * D1_0
805
+ )
806
+ elif solver_type == 'taylor':
807
+ x_t = (
808
+ expand_dims(sigma_t / sigma_prev_0, dims) * x
809
+ - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
810
+ + expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1_0
811
+ )
812
+ else:
813
+ if solver_type == 'dpm_solver':
814
+ x_t = (
815
+ expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
816
+ - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
817
+ - 0.5 * expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * D1_0
818
+ )
819
+ elif solver_type == 'taylor':
820
+ x_t = (
821
+ expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
822
+ - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
823
+ - expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1_0
824
+ )
825
+ return x_t
826
+
827
+ def multistep_dpm_solver_third_update(self, x, model_prev_list, t_prev_list, t, solver_type='dpm_solver'):
828
+ """
829
+ Multistep solver DPM-Solver-3 from time `t_prev_list[-1]` to time `t`.
830
+
831
+ Args:
832
+ x: A pytorch tensor. The initial value at time `s`.
833
+ model_prev_list: A list of pytorch tensor. The previous computed model values.
834
+ t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)
835
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
836
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
837
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
838
+ Returns:
839
+ x_t: A pytorch tensor. The approximated solution at time `t`.
840
+ """
841
+ ns = self.noise_schedule
842
+ dims = x.dim()
843
+ model_prev_2, model_prev_1, model_prev_0 = model_prev_list
844
+ t_prev_2, t_prev_1, t_prev_0 = t_prev_list
845
+ lambda_prev_2, lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_2), ns.marginal_lambda(
846
+ t_prev_1), ns.marginal_lambda(t_prev_0), ns.marginal_lambda(t)
847
+ log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
848
+ sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
849
+ alpha_t = torch.exp(log_alpha_t)
850
+
851
+ h_1 = lambda_prev_1 - lambda_prev_2
852
+ h_0 = lambda_prev_0 - lambda_prev_1
853
+ h = lambda_t - lambda_prev_0
854
+ r0, r1 = h_0 / h, h_1 / h
855
+ D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1)
856
+ D1_1 = expand_dims(1. / r1, dims) * (model_prev_1 - model_prev_2)
857
+ D1 = D1_0 + expand_dims(r0 / (r0 + r1), dims) * (D1_0 - D1_1)
858
+ D2 = expand_dims(1. / (r0 + r1), dims) * (D1_0 - D1_1)
859
+ if self.predict_x0:
860
+ x_t = (
861
+ expand_dims(sigma_t / sigma_prev_0, dims) * x
862
+ - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
863
+ + expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1
864
+ - expand_dims(alpha_t * ((torch.exp(-h) - 1. + h) / h ** 2 - 0.5), dims) * D2
865
+ )
866
+ else:
867
+ x_t = (
868
+ expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
869
+ - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
870
+ - expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1
871
+ - expand_dims(sigma_t * ((torch.exp(h) - 1. - h) / h ** 2 - 0.5), dims) * D2
872
+ )
873
+ return x_t
874
+
875
+ def singlestep_dpm_solver_update(self, x, s, t, order, return_intermediate=False, solver_type='dpm_solver', r1=None,
876
+ r2=None):
877
+ """
878
+ Singlestep DPM-Solver with the order `order` from time `s` to time `t`.
879
+
880
+ Args:
881
+ x: A pytorch tensor. The initial value at time `s`.
882
+ s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
883
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
884
+ order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3.
885
+ return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times).
886
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
887
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
888
+ r1: A `float`. The hyperparameter of the second-order or third-order solver.
889
+ r2: A `float`. The hyperparameter of the third-order solver.
890
+ Returns:
891
+ x_t: A pytorch tensor. The approximated solution at time `t`.
892
+ """
893
+ if order == 1:
894
+ return self.dpm_solver_first_update(x, s, t, return_intermediate=return_intermediate)
895
+ elif order == 2:
896
+ return self.singlestep_dpm_solver_second_update(x, s, t, return_intermediate=return_intermediate,
897
+ solver_type=solver_type, r1=r1)
898
+ elif order == 3:
899
+ return self.singlestep_dpm_solver_third_update(x, s, t, return_intermediate=return_intermediate,
900
+ solver_type=solver_type, r1=r1, r2=r2)
901
+ else:
902
+ raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order))
903
+
904
+ def multistep_dpm_solver_update(self, x, model_prev_list, t_prev_list, t, order, solver_type='dpm_solver'):
905
+ """
906
+ Multistep DPM-Solver with the order `order` from time `t_prev_list[-1]` to time `t`.
907
+
908
+ Args:
909
+ x: A pytorch tensor. The initial value at time `s`.
910
+ model_prev_list: A list of pytorch tensor. The previous computed model values.
911
+ t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)
912
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
913
+ order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3.
914
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
915
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
916
+ Returns:
917
+ x_t: A pytorch tensor. The approximated solution at time `t`.
918
+ """
919
+ if order == 1:
920
+ return self.dpm_solver_first_update(x, t_prev_list[-1], t, model_s=model_prev_list[-1])
921
+ elif order == 2:
922
+ return self.multistep_dpm_solver_second_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type)
923
+ elif order == 3:
924
+ return self.multistep_dpm_solver_third_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type)
925
+ else:
926
+ raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order))
927
+
928
+ 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,
929
+ solver_type='dpm_solver'):
930
+ """
931
+ The adaptive step size solver based on singlestep DPM-Solver.
932
+
933
+ Args:
934
+ x: A pytorch tensor. The initial value at time `t_T`.
935
+ order: A `int`. The (higher) order of the solver. We only support order == 2 or 3.
936
+ t_T: A `float`. The starting time of the sampling (default is T).
937
+ t_0: A `float`. The ending time of the sampling (default is epsilon).
938
+ h_init: A `float`. The initial step size (for logSNR).
939
+ atol: A `float`. The absolute tolerance of the solver. For image data, the default setting is 0.0078, followed [1].
940
+ rtol: A `float`. The relative tolerance of the solver. The default setting is 0.05.
941
+ theta: A `float`. The safety hyperparameter for adapting the step size. The default setting is 0.9, followed [1].
942
+ t_err: A `float`. The tolerance for the time. We solve the diffusion ODE until the absolute error between the
943
+ current time and `t_0` is less than `t_err`. The default setting is 1e-5.
944
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
945
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
946
+ Returns:
947
+ x_0: A pytorch tensor. The approximated solution at time `t_0`.
948
+
949
+ [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.
950
+ """
951
+ ns = self.noise_schedule
952
+ s = t_T * torch.ones((x.shape[0],)).to(x)
953
+ lambda_s = ns.marginal_lambda(s)
954
+ lambda_0 = ns.marginal_lambda(t_0 * torch.ones_like(s).to(x))
955
+ h = h_init * torch.ones_like(s).to(x)
956
+ x_prev = x
957
+ nfe = 0
958
+ if order == 2:
959
+ r1 = 0.5
960
+ lower_update = lambda x, s, t: self.dpm_solver_first_update(x, s, t, return_intermediate=True)
961
+ higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1,
962
+ solver_type=solver_type,
963
+ **kwargs)
964
+ elif order == 3:
965
+ r1, r2 = 1. / 3., 2. / 3.
966
+ lower_update = lambda x, s, t: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1,
967
+ return_intermediate=True,
968
+ solver_type=solver_type)
969
+ higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_third_update(x, s, t, r1=r1, r2=r2,
970
+ solver_type=solver_type,
971
+ **kwargs)
972
+ else:
973
+ raise ValueError("For adaptive step size solver, order must be 2 or 3, got {}".format(order))
974
+ while torch.abs((s - t_0)).mean() > t_err:
975
+ t = ns.inverse_lambda(lambda_s + h)
976
+ x_lower, lower_noise_kwargs = lower_update(x, s, t)
977
+ x_higher = higher_update(x, s, t, **lower_noise_kwargs)
978
+ delta = torch.max(torch.ones_like(x).to(x) * atol, rtol * torch.max(torch.abs(x_lower), torch.abs(x_prev)))
979
+ norm_fn = lambda v: torch.sqrt(torch.square(v.reshape((v.shape[0], -1))).mean(dim=-1, keepdim=True))
980
+ E = norm_fn((x_higher - x_lower) / delta).max()
981
+ if torch.all(E <= 1.):
982
+ x = x_higher
983
+ s = t
984
+ x_prev = x_lower
985
+ lambda_s = ns.marginal_lambda(s)
986
+ h = torch.min(theta * h * torch.float_power(E, -1. / order).float(), lambda_0 - lambda_s)
987
+ nfe += order
988
+ print('adaptive solver nfe', nfe)
989
+ return x
990
+
991
+ def sample(self, x, steps=20, t_start=None, t_end=None, order=3, skip_type='time_uniform',
992
+ method='singlestep', denoise=False, solver_type='dpm_solver', atol=0.0078,
993
+ rtol=0.05,
994
+ ):
995
+ """
996
+ Compute the sample at time `t_end` by DPM-Solver, given the initial `x` at time `t_start`.
997
+
998
+ =====================================================
999
+
1000
+ We support the following algorithms for both noise prediction model and data prediction model:
1001
+ - 'singlestep':
1002
+ Singlestep DPM-Solver (i.e. "DPM-Solver-fast" in the paper), which combines different orders of singlestep DPM-Solver.
1003
+ We combine all the singlestep solvers with order <= `order` to use up all the function evaluations (steps).
1004
+ The total number of function evaluations (NFE) == `steps`.
1005
+ Given a fixed NFE == `steps`, the sampling procedure is:
1006
+ - If `order` == 1:
1007
+ - Denote K = steps. We use K steps of DPM-Solver-1 (i.e. DDIM).
1008
+ - If `order` == 2:
1009
+ - Denote K = (steps // 2) + (steps % 2). We take K intermediate time steps for sampling.
1010
+ - If steps % 2 == 0, we use K steps of singlestep DPM-Solver-2.
1011
+ - If steps % 2 == 1, we use (K - 1) steps of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1.
1012
+ - If `order` == 3:
1013
+ - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling.
1014
+ - 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.
1015
+ - If steps % 3 == 1, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of DPM-Solver-1.
1016
+ - If steps % 3 == 2, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of singlestep DPM-Solver-2.
1017
+ - 'multistep':
1018
+ Multistep DPM-Solver with the order of `order`. The total number of function evaluations (NFE) == `steps`.
1019
+ We initialize the first `order` values by lower order multistep solvers.
1020
+ Given a fixed NFE == `steps`, the sampling procedure is:
1021
+ Denote K = steps.
1022
+ - If `order` == 1:
1023
+ - We use K steps of DPM-Solver-1 (i.e. DDIM).
1024
+ - If `order` == 2:
1025
+ - We firstly use 1 step of DPM-Solver-1, then use (K - 1) step of multistep DPM-Solver-2.
1026
+ - If `order` == 3:
1027
+ - 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.
1028
+ - 'singlestep_fixed':
1029
+ Fixed order singlestep DPM-Solver (i.e. DPM-Solver-1 or singlestep DPM-Solver-2 or singlestep DPM-Solver-3).
1030
+ We use singlestep DPM-Solver-`order` for `order`=1 or 2 or 3, with total [`steps` // `order`] * `order` NFE.
1031
+ - 'adaptive':
1032
+ Adaptive step size DPM-Solver (i.e. "DPM-Solver-12" and "DPM-Solver-23" in the paper).
1033
+ We ignore `steps` and use adaptive step size DPM-Solver with a higher order of `order`.
1034
+ You can adjust the absolute tolerance `atol` and the relative tolerance `rtol` to balance the computatation costs
1035
+ (NFE) and the sample quality.
1036
+ - If `order` == 2, we use DPM-Solver-12 which combines DPM-Solver-1 and singlestep DPM-Solver-2.
1037
+ - If `order` == 3, we use DPM-Solver-23 which combines singlestep DPM-Solver-2 and singlestep DPM-Solver-3.
1038
+
1039
+ =====================================================
1040
+
1041
+ Some advices for choosing the algorithm:
1042
+ - For **unconditional sampling** or **guided sampling with small guidance scale** by DPMs:
1043
+ Use singlestep DPM-Solver ("DPM-Solver-fast" in the paper) with `order = 3`.
1044
+ e.g.
1045
+ >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=False)
1046
+ >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=3,
1047
+ skip_type='time_uniform', method='singlestep')
1048
+ - For **guided sampling with large guidance scale** by DPMs:
1049
+ Use multistep DPM-Solver with `predict_x0 = True` and `order = 2`.
1050
+ e.g.
1051
+ >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=True)
1052
+ >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=2,
1053
+ skip_type='time_uniform', method='multistep')
1054
+
1055
+ We support three types of `skip_type`:
1056
+ - 'logSNR': uniform logSNR for the time steps. **Recommended for low-resolutional images**
1057
+ - 'time_uniform': uniform time for the time steps. **Recommended for high-resolutional images**.
1058
+ - 'time_quadratic': quadratic time for the time steps.
1059
+
1060
+ =====================================================
1061
+ Args:
1062
+ x: A pytorch tensor. The initial value at time `t_start`
1063
+ e.g. if `t_start` == T, then `x` is a sample from the standard normal distribution.
1064
+ steps: A `int`. The total number of function evaluations (NFE).
1065
+ t_start: A `float`. The starting time of the sampling.
1066
+ If `T` is None, we use self.noise_schedule.T (default is 1.0).
1067
+ t_end: A `float`. The ending time of the sampling.
1068
+ If `t_end` is None, we use 1. / self.noise_schedule.total_N.
1069
+ e.g. if total_N == 1000, we have `t_end` == 1e-3.
1070
+ For discrete-time DPMs:
1071
+ - We recommend `t_end` == 1. / self.noise_schedule.total_N.
1072
+ For continuous-time DPMs:
1073
+ - We recommend `t_end` == 1e-3 when `steps` <= 15; and `t_end` == 1e-4 when `steps` > 15.
1074
+ order: A `int`. The order of DPM-Solver.
1075
+ skip_type: A `str`. The type for the spacing of the time steps. 'time_uniform' or 'logSNR' or 'time_quadratic'.
1076
+ method: A `str`. The method for sampling. 'singlestep' or 'multistep' or 'singlestep_fixed' or 'adaptive'.
1077
+ denoise: A `bool`. Whether to denoise at the final step. Default is False.
1078
+ If `denoise` is True, the total NFE is (`steps` + 1).
1079
+ solver_type: A `str`. The taylor expansion type for the solver. `dpm_solver` or `taylor`. We recommend `dpm_solver`.
1080
+ atol: A `float`. The absolute tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.
1081
+ rtol: A `float`. The relative tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.
1082
+ Returns:
1083
+ x_end: A pytorch tensor. The approximated solution at time `t_end`.
1084
+
1085
+ """
1086
+ t_0 = 1. / self.noise_schedule.total_N if t_end is None else t_end
1087
+ t_T = self.noise_schedule.T if t_start is None else t_start
1088
+ device = x.device
1089
+ if method == 'adaptive':
1090
+ with torch.no_grad():
1091
+ x = self.dpm_solver_adaptive(x, order=order, t_T=t_T, t_0=t_0, atol=atol, rtol=rtol,
1092
+ solver_type=solver_type)
1093
+ elif method == 'multistep':
1094
+ assert steps >= order
1095
+ timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device)
1096
+ assert timesteps.shape[0] - 1 == steps
1097
+ with torch.no_grad():
1098
+ vec_t = timesteps[0].expand((x.shape[0]))
1099
+ model_prev_list = [self.model_fn(x, vec_t)]
1100
+ t_prev_list = [vec_t]
1101
+ # Init the first `order` values by lower order multistep DPM-Solver.
1102
+ for init_order in range(1, order):
1103
+ vec_t = timesteps[init_order].expand(x.shape[0])
1104
+ x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, vec_t, init_order,
1105
+ solver_type=solver_type)
1106
+ model_prev_list.append(self.model_fn(x, vec_t))
1107
+ t_prev_list.append(vec_t)
1108
+ # Compute the remaining values by `order`-th order multistep DPM-Solver.
1109
+ for step in range(order, steps + 1):
1110
+ vec_t = timesteps[step].expand(x.shape[0])
1111
+ x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, vec_t, order,
1112
+ solver_type=solver_type)
1113
+ for i in range(order - 1):
1114
+ t_prev_list[i] = t_prev_list[i + 1]
1115
+ model_prev_list[i] = model_prev_list[i + 1]
1116
+ t_prev_list[-1] = vec_t
1117
+ # We do not need to evaluate the final model value.
1118
+ if step < steps:
1119
+ model_prev_list[-1] = self.model_fn(x, vec_t)
1120
+ elif method in ['singlestep', 'singlestep_fixed']:
1121
+ if method == 'singlestep':
1122
+ timesteps_outer, orders = self.get_orders_and_timesteps_for_singlestep_solver(steps=steps, order=order,
1123
+ skip_type=skip_type,
1124
+ t_T=t_T, t_0=t_0,
1125
+ device=device)
1126
+ elif method == 'singlestep_fixed':
1127
+ K = steps // order
1128
+ orders = [order, ] * K
1129
+ timesteps_outer = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=K, device=device)
1130
+ for i, order in enumerate(orders):
1131
+ t_T_inner, t_0_inner = timesteps_outer[i], timesteps_outer[i + 1]
1132
+ timesteps_inner = self.get_time_steps(skip_type=skip_type, t_T=t_T_inner.item(), t_0=t_0_inner.item(),
1133
+ N=order, device=device)
1134
+ lambda_inner = self.noise_schedule.marginal_lambda(timesteps_inner)
1135
+ vec_s, vec_t = t_T_inner.repeat(x.shape[0]), t_0_inner.repeat(x.shape[0])
1136
+ h = lambda_inner[-1] - lambda_inner[0]
1137
+ r1 = None if order <= 1 else (lambda_inner[1] - lambda_inner[0]) / h
1138
+ r2 = None if order <= 2 else (lambda_inner[2] - lambda_inner[0]) / h
1139
+ x = self.singlestep_dpm_solver_update(x, vec_s, vec_t, order, solver_type=solver_type, r1=r1, r2=r2)
1140
+ if denoise:
1141
+ x = self.denoise_fn(x, torch.ones((x.shape[0],)).to(device) * t_0)
1142
+ return x
1143
+
1144
+
1145
+ #############################################################
1146
+ # other utility functions
1147
+ #############################################################
1148
+
1149
+ def interpolate_fn(x, xp, yp):
1150
+ """
1151
+ A piecewise linear function y = f(x), using xp and yp as keypoints.
1152
+ We implement f(x) in a differentiable way (i.e. applicable for autograd).
1153
+ 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.)
1154
+
1155
+ Args:
1156
+ 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).
1157
+ xp: PyTorch tensor with shape [C, K], where K is the number of keypoints.
1158
+ yp: PyTorch tensor with shape [C, K].
1159
+ Returns:
1160
+ The function values f(x), with shape [N, C].
1161
+ """
1162
+ N, K = x.shape[0], xp.shape[1]
1163
+ all_x = torch.cat([x.unsqueeze(2), xp.unsqueeze(0).repeat((N, 1, 1))], dim=2)
1164
+ sorted_all_x, x_indices = torch.sort(all_x, dim=2)
1165
+ x_idx = torch.argmin(x_indices, dim=2)
1166
+ cand_start_idx = x_idx - 1
1167
+ start_idx = torch.where(
1168
+ torch.eq(x_idx, 0),
1169
+ torch.tensor(1, device=x.device),
1170
+ torch.where(
1171
+ torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
1172
+ ),
1173
+ )
1174
+ end_idx = torch.where(torch.eq(start_idx, cand_start_idx), start_idx + 2, start_idx + 1)
1175
+ start_x = torch.gather(sorted_all_x, dim=2, index=start_idx.unsqueeze(2)).squeeze(2)
1176
+ end_x = torch.gather(sorted_all_x, dim=2, index=end_idx.unsqueeze(2)).squeeze(2)
1177
+ start_idx2 = torch.where(
1178
+ torch.eq(x_idx, 0),
1179
+ torch.tensor(0, device=x.device),
1180
+ torch.where(
1181
+ torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
1182
+ ),
1183
+ )
1184
+ y_positions_expanded = yp.unsqueeze(0).expand(N, -1, -1)
1185
+ start_y = torch.gather(y_positions_expanded, dim=2, index=start_idx2.unsqueeze(2)).squeeze(2)
1186
+ end_y = torch.gather(y_positions_expanded, dim=2, index=(start_idx2 + 1).unsqueeze(2)).squeeze(2)
1187
+ cand = start_y + (x - start_x) * (end_y - start_y) / (end_x - start_x)
1188
+ return cand
1189
+
1190
+
1191
+ def expand_dims(v, dims):
1192
+ """
1193
+ Expand the tensor `v` to the dim `dims`.
1194
+
1195
+ Args:
1196
+ `v`: a PyTorch tensor with shape [N].
1197
+ `dim`: a `int`.
1198
+ Returns:
1199
+ a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`.
1200
+ """
1201
+ return v[(...,) + (None,) * (dims - 1)]
diffusion/how to export onnx.md ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ - Open [onnx_export](onnx_export.py)
2
+ - project_name = "dddsp" change "project_name" to your project name
3
+ - model_path = f'{project_name}/model_500000.pt' change "model_path" to your model path
4
+ - Run
diffusion/infer_gt_mel.py ADDED
@@ -0,0 +1,74 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import torch
3
+ import torch.nn.functional as F
4
+ from diffusion.unit2mel import load_model_vocoder
5
+
6
+
7
+ class DiffGtMel:
8
+ def __init__(self, project_path=None, device=None):
9
+ self.project_path = project_path
10
+ if device is not None:
11
+ self.device = device
12
+ else:
13
+ self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
14
+ self.model = None
15
+ self.vocoder = None
16
+ self.args = None
17
+
18
+ def flush_model(self, project_path, ddsp_config=None):
19
+ if (self.model is None) or (project_path != self.project_path):
20
+ model, vocoder, args = load_model_vocoder(project_path, device=self.device)
21
+ if self.check_args(ddsp_config, args):
22
+ self.model = model
23
+ self.vocoder = vocoder
24
+ self.args = args
25
+
26
+ def check_args(self, args1, args2):
27
+ if args1.data.block_size != args2.data.block_size:
28
+ raise ValueError("DDSP与DIFF模型的block_size不一致")
29
+ if args1.data.sampling_rate != args2.data.sampling_rate:
30
+ raise ValueError("DDSP与DIFF模型的sampling_rate不一致")
31
+ if args1.data.encoder != args2.data.encoder:
32
+ raise ValueError("DDSP与DIFF模型的encoder不一致")
33
+ return True
34
+
35
+ def __call__(self, audio, f0, hubert, volume, acc=1, spk_id=1, k_step=0, method='pndm',
36
+ spk_mix_dict=None, start_frame=0):
37
+ input_mel = self.vocoder.extract(audio, self.args.data.sampling_rate)
38
+ out_mel = self.model(
39
+ hubert,
40
+ f0,
41
+ volume,
42
+ spk_id=spk_id,
43
+ spk_mix_dict=spk_mix_dict,
44
+ gt_spec=input_mel,
45
+ infer=True,
46
+ infer_speedup=acc,
47
+ method=method,
48
+ k_step=k_step,
49
+ use_tqdm=False)
50
+ if start_frame > 0:
51
+ out_mel = out_mel[:, start_frame:, :]
52
+ f0 = f0[:, start_frame:, :]
53
+ output = self.vocoder.infer(out_mel, f0)
54
+ if start_frame > 0:
55
+ output = F.pad(output, (start_frame * self.vocoder.vocoder_hop_size, 0))
56
+ return output
57
+
58
+ def infer(self, audio, f0, hubert, volume, acc=1, spk_id=1, k_step=0, method='pndm', silence_front=0,
59
+ use_silence=False, spk_mix_dict=None):
60
+ start_frame = int(silence_front * self.vocoder.vocoder_sample_rate / self.vocoder.vocoder_hop_size)
61
+ if use_silence:
62
+ audio = audio[:, start_frame * self.vocoder.vocoder_hop_size:]
63
+ f0 = f0[:, start_frame:, :]
64
+ hubert = hubert[:, start_frame:, :]
65
+ volume = volume[:, start_frame:, :]
66
+ _start_frame = 0
67
+ else:
68
+ _start_frame = start_frame
69
+ audio = self.__call__(audio, f0, hubert, volume, acc=acc, spk_id=spk_id, k_step=k_step,
70
+ method=method, spk_mix_dict=spk_mix_dict, start_frame=_start_frame)
71
+ if use_silence:
72
+ if start_frame > 0:
73
+ audio = F.pad(audio, (start_frame * self.vocoder.vocoder_hop_size, 0))
74
+ return audio
diffusion/logger/__init__.py ADDED
File without changes
diffusion/logger/saver.py ADDED
@@ -0,0 +1,150 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ '''
2
+ author: wayn391@mastertones
3
+ '''
4
+
5
+ import os
6
+ import json
7
+ import time
8
+ import yaml
9
+ import datetime
10
+ import torch
11
+ import matplotlib.pyplot as plt
12
+ from . import utils
13
+ from torch.utils.tensorboard import SummaryWriter
14
+
15
+ class Saver(object):
16
+ def __init__(
17
+ self,
18
+ args,
19
+ initial_global_step=-1):
20
+
21
+ self.expdir = args.env.expdir
22
+ self.sample_rate = args.data.sampling_rate
23
+
24
+ # cold start
25
+ self.global_step = initial_global_step
26
+ self.init_time = time.time()
27
+ self.last_time = time.time()
28
+
29
+ # makedirs
30
+ os.makedirs(self.expdir, exist_ok=True)
31
+
32
+ # path
33
+ self.path_log_info = os.path.join(self.expdir, 'log_info.txt')
34
+
35
+ # ckpt
36
+ os.makedirs(self.expdir, exist_ok=True)
37
+
38
+ # writer
39
+ self.writer = SummaryWriter(os.path.join(self.expdir, 'logs'))
40
+
41
+ # save config
42
+ path_config = os.path.join(self.expdir, 'config.yaml')
43
+ with open(path_config, "w") as out_config:
44
+ yaml.dump(dict(args), out_config)
45
+
46
+
47
+ def log_info(self, msg):
48
+ '''log method'''
49
+ if isinstance(msg, dict):
50
+ msg_list = []
51
+ for k, v in msg.items():
52
+ tmp_str = ''
53
+ if isinstance(v, int):
54
+ tmp_str = '{}: {:,}'.format(k, v)
55
+ else:
56
+ tmp_str = '{}: {}'.format(k, v)
57
+
58
+ msg_list.append(tmp_str)
59
+ msg_str = '\n'.join(msg_list)
60
+ else:
61
+ msg_str = msg
62
+
63
+ # dsplay
64
+ print(msg_str)
65
+
66
+ # save
67
+ with open(self.path_log_info, 'a') as fp:
68
+ fp.write(msg_str+'\n')
69
+
70
+ def log_value(self, dict):
71
+ for k, v in dict.items():
72
+ self.writer.add_scalar(k, v, self.global_step)
73
+
74
+ def log_spec(self, name, spec, spec_out, vmin=-14, vmax=3.5):
75
+ spec_cat = torch.cat([(spec_out - spec).abs() + vmin, spec, spec_out], -1)
76
+ spec = spec_cat[0]
77
+ if isinstance(spec, torch.Tensor):
78
+ spec = spec.cpu().numpy()
79
+ fig = plt.figure(figsize=(12, 9))
80
+ plt.pcolor(spec.T, vmin=vmin, vmax=vmax)
81
+ plt.tight_layout()
82
+ self.writer.add_figure(name, fig, self.global_step)
83
+
84
+ def log_audio(self, dict):
85
+ for k, v in dict.items():
86
+ self.writer.add_audio(k, v, global_step=self.global_step, sample_rate=self.sample_rate)
87
+
88
+ def get_interval_time(self, update=True):
89
+ cur_time = time.time()
90
+ time_interval = cur_time - self.last_time
91
+ if update:
92
+ self.last_time = cur_time
93
+ return time_interval
94
+
95
+ def get_total_time(self, to_str=True):
96
+ total_time = time.time() - self.init_time
97
+ if to_str:
98
+ total_time = str(datetime.timedelta(
99
+ seconds=total_time))[:-5]
100
+ return total_time
101
+
102
+ def save_model(
103
+ self,
104
+ model,
105
+ optimizer,
106
+ name='model',
107
+ postfix='',
108
+ to_json=False):
109
+ # path
110
+ if postfix:
111
+ postfix = '_' + postfix
112
+ path_pt = os.path.join(
113
+ self.expdir , name+postfix+'.pt')
114
+
115
+ # check
116
+ print(' [*] model checkpoint saved: {}'.format(path_pt))
117
+
118
+ # save
119
+ if optimizer is not None:
120
+ torch.save({
121
+ 'global_step': self.global_step,
122
+ 'model': model.state_dict(),
123
+ 'optimizer': optimizer.state_dict()}, path_pt)
124
+ else:
125
+ torch.save({
126
+ 'global_step': self.global_step,
127
+ 'model': model.state_dict()}, path_pt)
128
+
129
+ # to json
130
+ if to_json:
131
+ path_json = os.path.join(
132
+ self.expdir , name+'.json')
133
+ utils.to_json(path_params, path_json)
134
+
135
+ def delete_model(self, name='model', postfix=''):
136
+ # path
137
+ if postfix:
138
+ postfix = '_' + postfix
139
+ path_pt = os.path.join(
140
+ self.expdir , name+postfix+'.pt')
141
+
142
+ # delete
143
+ if os.path.exists(path_pt):
144
+ os.remove(path_pt)
145
+ print(' [*] model checkpoint deleted: {}'.format(path_pt))
146
+
147
+ def global_step_increment(self):
148
+ self.global_step += 1
149
+
150
+
diffusion/logger/utils.py ADDED
@@ -0,0 +1,126 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import yaml
3
+ import json
4
+ import pickle
5
+ import torch
6
+
7
+ def traverse_dir(
8
+ root_dir,
9
+ extensions,
10
+ amount=None,
11
+ str_include=None,
12
+ str_exclude=None,
13
+ is_pure=False,
14
+ is_sort=False,
15
+ is_ext=True):
16
+
17
+ file_list = []
18
+ cnt = 0
19
+ for root, _, files in os.walk(root_dir):
20
+ for file in files:
21
+ if any([file.endswith(f".{ext}") for ext in extensions]):
22
+ # path
23
+ mix_path = os.path.join(root, file)
24
+ pure_path = mix_path[len(root_dir)+1:] if is_pure else mix_path
25
+
26
+ # amount
27
+ if (amount is not None) and (cnt == amount):
28
+ if is_sort:
29
+ file_list.sort()
30
+ return file_list
31
+
32
+ # check string
33
+ if (str_include is not None) and (str_include not in pure_path):
34
+ continue
35
+ if (str_exclude is not None) and (str_exclude in pure_path):
36
+ continue
37
+
38
+ if not is_ext:
39
+ ext = pure_path.split('.')[-1]
40
+ pure_path = pure_path[:-(len(ext)+1)]
41
+ file_list.append(pure_path)
42
+ cnt += 1
43
+ if is_sort:
44
+ file_list.sort()
45
+ return file_list
46
+
47
+
48
+
49
+ class DotDict(dict):
50
+ def __getattr__(*args):
51
+ val = dict.get(*args)
52
+ return DotDict(val) if type(val) is dict else val
53
+
54
+ __setattr__ = dict.__setitem__
55
+ __delattr__ = dict.__delitem__
56
+
57
+
58
+ def get_network_paras_amount(model_dict):
59
+ info = dict()
60
+ for model_name, model in model_dict.items():
61
+ # all_params = sum(p.numel() for p in model.parameters())
62
+ trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
63
+
64
+ info[model_name] = trainable_params
65
+ return info
66
+
67
+
68
+ def load_config(path_config):
69
+ with open(path_config, "r") as config:
70
+ args = yaml.safe_load(config)
71
+ args = DotDict(args)
72
+ # print(args)
73
+ return args
74
+
75
+ def save_config(path_config,config):
76
+ config = dict(config)
77
+ with open(path_config, "w") as f:
78
+ yaml.dump(config, f)
79
+
80
+ def to_json(path_params, path_json):
81
+ params = torch.load(path_params, map_location=torch.device('cpu'))
82
+ raw_state_dict = {}
83
+ for k, v in params.items():
84
+ val = v.flatten().numpy().tolist()
85
+ raw_state_dict[k] = val
86
+
87
+ with open(path_json, 'w') as outfile:
88
+ json.dump(raw_state_dict, outfile,indent= "\t")
89
+
90
+
91
+ def convert_tensor_to_numpy(tensor, is_squeeze=True):
92
+ if is_squeeze:
93
+ tensor = tensor.squeeze()
94
+ if tensor.requires_grad:
95
+ tensor = tensor.detach()
96
+ if tensor.is_cuda:
97
+ tensor = tensor.cpu()
98
+ return tensor.numpy()
99
+
100
+
101
+ def load_model(
102
+ expdir,
103
+ model,
104
+ optimizer,
105
+ name='model',
106
+ postfix='',
107
+ device='cpu'):
108
+ if postfix == '':
109
+ postfix = '_' + postfix
110
+ path = os.path.join(expdir, name+postfix)
111
+ path_pt = traverse_dir(expdir, ['pt'], is_ext=False)
112
+ global_step = 0
113
+ if len(path_pt) > 0:
114
+ steps = [s[len(path):] for s in path_pt]
115
+ maxstep = max([int(s) if s.isdigit() else 0 for s in steps])
116
+ if maxstep >= 0:
117
+ path_pt = path+str(maxstep)+'.pt'
118
+ else:
119
+ path_pt = path+'best.pt'
120
+ print(' [*] restoring model from', path_pt)
121
+ ckpt = torch.load(path_pt, map_location=torch.device(device))
122
+ global_step = ckpt['global_step']
123
+ model.load_state_dict(ckpt['model'], strict=False)
124
+ if ckpt.get('optimizer') != None:
125
+ optimizer.load_state_dict(ckpt['optimizer'])
126
+ return global_step, model, optimizer
diffusion/onnx_export.py ADDED
@@ -0,0 +1,226 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from diffusion_onnx import GaussianDiffusion
2
+ import os
3
+ import yaml
4
+ import torch
5
+ import torch.nn as nn
6
+ import numpy as np
7
+ from wavenet import WaveNet
8
+ import torch.nn.functional as F
9
+ import diffusion
10
+
11
+ class DotDict(dict):
12
+ def __getattr__(*args):
13
+ val = dict.get(*args)
14
+ return DotDict(val) if type(val) is dict else val
15
+
16
+ __setattr__ = dict.__setitem__
17
+ __delattr__ = dict.__delitem__
18
+
19
+
20
+ def load_model_vocoder(
21
+ model_path,
22
+ device='cpu'):
23
+ config_file = os.path.join(os.path.split(model_path)[0], 'config.yaml')
24
+ with open(config_file, "r") as config:
25
+ args = yaml.safe_load(config)
26
+ args = DotDict(args)
27
+
28
+ # load model
29
+ model = Unit2Mel(
30
+ args.data.encoder_out_channels,
31
+ args.model.n_spk,
32
+ args.model.use_pitch_aug,
33
+ 128,
34
+ args.model.n_layers,
35
+ args.model.n_chans,
36
+ args.model.n_hidden)
37
+
38
+ print(' [Loading] ' + model_path)
39
+ ckpt = torch.load(model_path, map_location=torch.device(device))
40
+ model.to(device)
41
+ model.load_state_dict(ckpt['model'])
42
+ model.eval()
43
+ return model, args
44
+
45
+
46
+ class Unit2Mel(nn.Module):
47
+ def __init__(
48
+ self,
49
+ input_channel,
50
+ n_spk,
51
+ use_pitch_aug=False,
52
+ out_dims=128,
53
+ n_layers=20,
54
+ n_chans=384,
55
+ n_hidden=256):
56
+ super().__init__()
57
+ self.unit_embed = nn.Linear(input_channel, n_hidden)
58
+ self.f0_embed = nn.Linear(1, n_hidden)
59
+ self.volume_embed = nn.Linear(1, n_hidden)
60
+ if use_pitch_aug:
61
+ self.aug_shift_embed = nn.Linear(1, n_hidden, bias=False)
62
+ else:
63
+ self.aug_shift_embed = None
64
+ self.n_spk = n_spk
65
+ if n_spk is not None and n_spk > 1:
66
+ self.spk_embed = nn.Embedding(n_spk, n_hidden)
67
+
68
+ # diffusion
69
+ self.decoder = GaussianDiffusion(out_dims, n_layers, n_chans, n_hidden)
70
+ self.hidden_size = n_hidden
71
+ self.speaker_map = torch.zeros((self.n_spk,1,1,n_hidden))
72
+
73
+
74
+
75
+ def forward(self, units, mel2ph, f0, volume, g = None):
76
+
77
+ '''
78
+ input:
79
+ B x n_frames x n_unit
80
+ return:
81
+ dict of B x n_frames x feat
82
+ '''
83
+
84
+ decoder_inp = F.pad(units, [0, 0, 1, 0])
85
+ mel2ph_ = mel2ph.unsqueeze(2).repeat([1, 1, units.shape[-1]])
86
+ units = torch.gather(decoder_inp, 1, mel2ph_) # [B, T, H]
87
+
88
+ x = self.unit_embed(units) + self.f0_embed((1 + f0.unsqueeze(-1) / 700).log()) + self.volume_embed(volume.unsqueeze(-1))
89
+
90
+ if self.n_spk is not None and self.n_spk > 1: # [N, S] * [S, B, 1, H]
91
+ g = g.reshape((g.shape[0], g.shape[1], 1, 1, 1)) # [N, S, B, 1, 1]
92
+ g = g * self.speaker_map # [N, S, B, 1, H]
93
+ g = torch.sum(g, dim=1) # [N, 1, B, 1, H]
94
+ g = g.transpose(0, -1).transpose(0, -2).squeeze(0) # [B, H, N]
95
+ x = x.transpose(1, 2) + g
96
+ return x
97
+ else:
98
+ return x.transpose(1, 2)
99
+
100
+
101
+ def init_spkembed(self, units, f0, volume, spk_id = None, spk_mix_dict = None, aug_shift = None,
102
+ gt_spec=None, infer=True, infer_speedup=10, method='dpm-solver', k_step=300, use_tqdm=True):
103
+
104
+ '''
105
+ input:
106
+ B x n_frames x n_unit
107
+ return:
108
+ dict of B x n_frames x feat
109
+ '''
110
+ x = self.unit_embed(units) + self.f0_embed((1+ f0 / 700).log()) + self.volume_embed(volume)
111
+ if self.n_spk is not None and self.n_spk > 1:
112
+ if spk_mix_dict is not None:
113
+ spk_embed_mix = torch.zeros((1,1,self.hidden_size))
114
+ for k, v in spk_mix_dict.items():
115
+ spk_id_torch = torch.LongTensor(np.array([[k]])).to(units.device)
116
+ spk_embeddd = self.spk_embed(spk_id_torch)
117
+ self.speaker_map[k] = spk_embeddd
118
+ spk_embed_mix = spk_embed_mix + v * spk_embeddd
119
+ x = x + spk_embed_mix
120
+ else:
121
+ x = x + self.spk_embed(spk_id - 1)
122
+ self.speaker_map = self.speaker_map.unsqueeze(0)
123
+ self.speaker_map = self.speaker_map.detach()
124
+ return x.transpose(1, 2)
125
+
126
+ def OnnxExport(self, project_name=None, init_noise=None, export_encoder=True, export_denoise=True, export_pred=True, export_after=True):
127
+ hubert_hidden_size = 768
128
+ n_frames = 100
129
+ hubert = torch.randn((1, n_frames, hubert_hidden_size))
130
+ mel2ph = torch.arange(end=n_frames).unsqueeze(0).long()
131
+ f0 = torch.randn((1, n_frames))
132
+ volume = torch.randn((1, n_frames))
133
+ spk_mix = []
134
+ spks = {}
135
+ if self.n_spk is not None and self.n_spk > 1:
136
+ for i in range(self.n_spk):
137
+ spk_mix.append(1.0/float(self.n_spk))
138
+ spks.update({i:1.0/float(self.n_spk)})
139
+ spk_mix = torch.tensor(spk_mix)
140
+ spk_mix = spk_mix.repeat(n_frames, 1)
141
+ orgouttt = self.init_spkembed(hubert, f0.unsqueeze(-1), volume.unsqueeze(-1), spk_mix_dict=spks)
142
+ outtt = self.forward(hubert, mel2ph, f0, volume, spk_mix)
143
+ if export_encoder:
144
+ torch.onnx.export(
145
+ self,
146
+ (hubert, mel2ph, f0, volume, spk_mix),
147
+ f"{project_name}_encoder.onnx",
148
+ input_names=["hubert", "mel2ph", "f0", "volume", "spk_mix"],
149
+ output_names=["mel_pred"],
150
+ dynamic_axes={
151
+ "hubert": [1],
152
+ "f0": [1],
153
+ "volume": [1],
154
+ "mel2ph": [1],
155
+ "spk_mix": [0],
156
+ },
157
+ opset_version=16
158
+ )
159
+
160
+ self.decoder.OnnxExport(project_name, init_noise=init_noise, export_denoise=export_denoise, export_pred=export_pred, export_after=export_after)
161
+
162
+ def ExportOnnx(self, project_name=None):
163
+ hubert_hidden_size = 768
164
+ n_frames = 100
165
+ hubert = torch.randn((1, n_frames, hubert_hidden_size))
166
+ mel2ph = torch.arange(end=n_frames).unsqueeze(0).long()
167
+ f0 = torch.randn((1, n_frames))
168
+ volume = torch.randn((1, n_frames))
169
+ spk_mix = []
170
+ spks = {}
171
+ if self.n_spk is not None and self.n_spk > 1:
172
+ for i in range(self.n_spk):
173
+ spk_mix.append(1.0/float(self.n_spk))
174
+ spks.update({i:1.0/float(self.n_spk)})
175
+ spk_mix = torch.tensor(spk_mix)
176
+ orgouttt = self.orgforward(hubert, f0.unsqueeze(-1), volume.unsqueeze(-1), spk_mix_dict=spks)
177
+ outtt = self.forward(hubert, mel2ph, f0, volume, spk_mix)
178
+
179
+ torch.onnx.export(
180
+ self,
181
+ (hubert, mel2ph, f0, volume, spk_mix),
182
+ f"{project_name}_encoder.onnx",
183
+ input_names=["hubert", "mel2ph", "f0", "volume", "spk_mix"],
184
+ output_names=["mel_pred"],
185
+ dynamic_axes={
186
+ "hubert": [1],
187
+ "f0": [1],
188
+ "volume": [1],
189
+ "mel2ph": [1]
190
+ },
191
+ opset_version=16
192
+ )
193
+
194
+ condition = torch.randn(1,self.decoder.n_hidden,n_frames)
195
+ noise = torch.randn((1, 1, self.decoder.mel_bins, condition.shape[2]), dtype=torch.float32)
196
+ pndm_speedup = torch.LongTensor([100])
197
+ K_steps = torch.LongTensor([1000])
198
+ self.decoder = torch.jit.script(self.decoder)
199
+ self.decoder(condition, noise, pndm_speedup, K_steps)
200
+
201
+ torch.onnx.export(
202
+ self.decoder,
203
+ (condition, noise, pndm_speedup, K_steps),
204
+ f"{project_name}_diffusion.onnx",
205
+ input_names=["condition", "noise", "pndm_speedup", "K_steps"],
206
+ output_names=["mel"],
207
+ dynamic_axes={
208
+ "condition": [2],
209
+ "noise": [3],
210
+ },
211
+ opset_version=16
212
+ )
213
+
214
+
215
+ if __name__ == "__main__":
216
+ project_name = "dddsp"
217
+ model_path = f'{project_name}/model_500000.pt'
218
+
219
+ model, _ = load_model_vocoder(model_path)
220
+
221
+ # 分开Diffusion导出(需要使用MoeSS/MoeVoiceStudio或者自己编写Pndm/Dpm采样)
222
+ model.OnnxExport(project_name, export_encoder=True, export_denoise=True, export_pred=True, export_after=True)
223
+
224
+ # 合并Diffusion导出(Encoder和Diffusion分开,直接将Encoder的结果和初始噪声输入Diffusion即可)
225
+ # model.ExportOnnx(project_name)
226
+
diffusion/solver.py ADDED
@@ -0,0 +1,195 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import time
3
+ import numpy as np
4
+ import torch
5
+ import librosa
6
+ from diffusion.logger.saver import Saver
7
+ from diffusion.logger import utils
8
+ from torch import autocast
9
+ from torch.cuda.amp import GradScaler
10
+
11
+ def test(args, model, vocoder, loader_test, saver):
12
+ print(' [*] testing...')
13
+ model.eval()
14
+
15
+ # losses
16
+ test_loss = 0.
17
+
18
+ # intialization
19
+ num_batches = len(loader_test)
20
+ rtf_all = []
21
+
22
+ # run
23
+ with torch.no_grad():
24
+ for bidx, data in enumerate(loader_test):
25
+ fn = data['name'][0].split("/")[-1]
26
+ speaker = data['name'][0].split("/")[-2]
27
+ print('--------')
28
+ print('{}/{} - {}'.format(bidx, num_batches, fn))
29
+
30
+ # unpack data
31
+ for k in data.keys():
32
+ if not k.startswith('name'):
33
+ data[k] = data[k].to(args.device)
34
+ print('>>', data['name'][0])
35
+
36
+ # forward
37
+ st_time = time.time()
38
+ mel = model(
39
+ data['units'],
40
+ data['f0'],
41
+ data['volume'],
42
+ data['spk_id'],
43
+ gt_spec=None,
44
+ infer=True,
45
+ infer_speedup=args.infer.speedup,
46
+ method=args.infer.method)
47
+ signal = vocoder.infer(mel, data['f0'])
48
+ ed_time = time.time()
49
+
50
+ # RTF
51
+ run_time = ed_time - st_time
52
+ song_time = signal.shape[-1] / args.data.sampling_rate
53
+ rtf = run_time / song_time
54
+ print('RTF: {} | {} / {}'.format(rtf, run_time, song_time))
55
+ rtf_all.append(rtf)
56
+
57
+ # loss
58
+ for i in range(args.train.batch_size):
59
+ loss = model(
60
+ data['units'],
61
+ data['f0'],
62
+ data['volume'],
63
+ data['spk_id'],
64
+ gt_spec=data['mel'],
65
+ infer=False)
66
+ test_loss += loss.item()
67
+
68
+ # log mel
69
+ saver.log_spec(f"{speaker}_{fn}.wav", data['mel'], mel)
70
+
71
+ # log audi
72
+ path_audio = data['name_ext'][0]
73
+ audio, sr = librosa.load(path_audio, sr=args.data.sampling_rate)
74
+ if len(audio.shape) > 1:
75
+ audio = librosa.to_mono(audio)
76
+ audio = torch.from_numpy(audio).unsqueeze(0).to(signal)
77
+ saver.log_audio({f"{speaker}_{fn}_gt.wav": audio,f"{speaker}_{fn}_pred.wav": signal})
78
+ # report
79
+ test_loss /= args.train.batch_size
80
+ test_loss /= num_batches
81
+
82
+ # check
83
+ print(' [test_loss] test_loss:', test_loss)
84
+ print(' Real Time Factor', np.mean(rtf_all))
85
+ return test_loss
86
+
87
+
88
+ def train(args, initial_global_step, model, optimizer, scheduler, vocoder, loader_train, loader_test):
89
+ # saver
90
+ saver = Saver(args, initial_global_step=initial_global_step)
91
+
92
+ # model size
93
+ params_count = utils.get_network_paras_amount({'model': model})
94
+ saver.log_info('--- model size ---')
95
+ saver.log_info(params_count)
96
+
97
+ # run
98
+ num_batches = len(loader_train)
99
+ model.train()
100
+ saver.log_info('======= start training =======')
101
+ scaler = GradScaler()
102
+ if args.train.amp_dtype == 'fp32':
103
+ dtype = torch.float32
104
+ elif args.train.amp_dtype == 'fp16':
105
+ dtype = torch.float16
106
+ elif args.train.amp_dtype == 'bf16':
107
+ dtype = torch.bfloat16
108
+ else:
109
+ raise ValueError(' [x] Unknown amp_dtype: ' + args.train.amp_dtype)
110
+ saver.log_info("epoch|batch_idx/num_batches|output_dir|batch/s|lr|time|step")
111
+ for epoch in range(args.train.epochs):
112
+ for batch_idx, data in enumerate(loader_train):
113
+ saver.global_step_increment()
114
+ optimizer.zero_grad()
115
+
116
+ # unpack data
117
+ for k in data.keys():
118
+ if not k.startswith('name'):
119
+ data[k] = data[k].to(args.device)
120
+
121
+ # forward
122
+ if dtype == torch.float32:
123
+ loss = model(data['units'].float(), data['f0'], data['volume'], data['spk_id'],
124
+ aug_shift = data['aug_shift'], gt_spec=data['mel'].float(), infer=False)
125
+ else:
126
+ with autocast(device_type=args.device, dtype=dtype):
127
+ loss = model(data['units'], data['f0'], data['volume'], data['spk_id'],
128
+ aug_shift = data['aug_shift'], gt_spec=data['mel'], infer=False)
129
+
130
+ # handle nan loss
131
+ if torch.isnan(loss):
132
+ raise ValueError(' [x] nan loss ')
133
+ else:
134
+ # backpropagate
135
+ if dtype == torch.float32:
136
+ loss.backward()
137
+ optimizer.step()
138
+ else:
139
+ scaler.scale(loss).backward()
140
+ scaler.step(optimizer)
141
+ scaler.update()
142
+ scheduler.step()
143
+
144
+ # log loss
145
+ if saver.global_step % args.train.interval_log == 0:
146
+ current_lr = optimizer.param_groups[0]['lr']
147
+ saver.log_info(
148
+ 'epoch: {} | {:3d}/{:3d} | {} | batch/s: {:.2f} | lr: {:.6} | loss: {:.3f} | time: {} | step: {}'.format(
149
+ epoch,
150
+ batch_idx,
151
+ num_batches,
152
+ args.env.expdir,
153
+ args.train.interval_log/saver.get_interval_time(),
154
+ current_lr,
155
+ loss.item(),
156
+ saver.get_total_time(),
157
+ saver.global_step
158
+ )
159
+ )
160
+
161
+ saver.log_value({
162
+ 'train/loss': loss.item()
163
+ })
164
+
165
+ saver.log_value({
166
+ 'train/lr': current_lr
167
+ })
168
+
169
+ # validation
170
+ if saver.global_step % args.train.interval_val == 0:
171
+ optimizer_save = optimizer if args.train.save_opt else None
172
+
173
+ # save latest
174
+ saver.save_model(model, optimizer_save, postfix=f'{saver.global_step}')
175
+ last_val_step = saver.global_step - args.train.interval_val
176
+ if last_val_step % args.train.interval_force_save != 0:
177
+ saver.delete_model(postfix=f'{last_val_step}')
178
+
179
+ # run testing set
180
+ test_loss = test(args, model, vocoder, loader_test, saver)
181
+
182
+ # log loss
183
+ saver.log_info(
184
+ ' --- <validation> --- \nloss: {:.3f}. '.format(
185
+ test_loss,
186
+ )
187
+ )
188
+
189
+ saver.log_value({
190
+ 'validation/loss': test_loss
191
+ })
192
+
193
+ model.train()
194
+
195
+
diffusion/unit2mel.py ADDED
@@ -0,0 +1,100 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import yaml
3
+ import torch
4
+ import torch.nn as nn
5
+ import numpy as np
6
+ from .diffusion import GaussianDiffusion
7
+ from .wavenet import WaveNet
8
+ from .vocoder import Vocoder
9
+
10
+ class DotDict(dict):
11
+ def __getattr__(*args):
12
+ val = dict.get(*args)
13
+ return DotDict(val) if type(val) is dict else val
14
+
15
+ __setattr__ = dict.__setitem__
16
+ __delattr__ = dict.__delitem__
17
+
18
+
19
+ def load_model_vocoder(
20
+ model_path,
21
+ device='cpu',
22
+ config_path = None
23
+ ):
24
+ if config_path is None: config_file = os.path.join(os.path.split(model_path)[0], 'config.yaml')
25
+ else: config_file = config_path
26
+
27
+ with open(config_file, "r") as config:
28
+ args = yaml.safe_load(config)
29
+ args = DotDict(args)
30
+
31
+ # load vocoder
32
+ vocoder = Vocoder(args.vocoder.type, args.vocoder.ckpt, device=device)
33
+
34
+ # load model
35
+ model = Unit2Mel(
36
+ args.data.encoder_out_channels,
37
+ args.model.n_spk,
38
+ args.model.use_pitch_aug,
39
+ vocoder.dimension,
40
+ args.model.n_layers,
41
+ args.model.n_chans,
42
+ args.model.n_hidden)
43
+
44
+ print(' [Loading] ' + model_path)
45
+ ckpt = torch.load(model_path, map_location=torch.device(device))
46
+ model.to(device)
47
+ model.load_state_dict(ckpt['model'])
48
+ model.eval()
49
+ return model, vocoder, args
50
+
51
+
52
+ class Unit2Mel(nn.Module):
53
+ def __init__(
54
+ self,
55
+ input_channel,
56
+ n_spk,
57
+ use_pitch_aug=False,
58
+ out_dims=128,
59
+ n_layers=20,
60
+ n_chans=384,
61
+ n_hidden=256):
62
+ super().__init__()
63
+ self.unit_embed = nn.Linear(input_channel, n_hidden)
64
+ self.f0_embed = nn.Linear(1, n_hidden)
65
+ self.volume_embed = nn.Linear(1, n_hidden)
66
+ if use_pitch_aug:
67
+ self.aug_shift_embed = nn.Linear(1, n_hidden, bias=False)
68
+ else:
69
+ self.aug_shift_embed = None
70
+ self.n_spk = n_spk
71
+ if n_spk is not None and n_spk > 1:
72
+ self.spk_embed = nn.Embedding(n_spk, n_hidden)
73
+
74
+ # diffusion
75
+ self.decoder = GaussianDiffusion(WaveNet(out_dims, n_layers, n_chans, n_hidden), out_dims=out_dims)
76
+
77
+ def forward(self, units, f0, volume, spk_id = None, spk_mix_dict = None, aug_shift = None,
78
+ gt_spec=None, infer=True, infer_speedup=10, method='dpm-solver', k_step=300, use_tqdm=True):
79
+
80
+ '''
81
+ input:
82
+ B x n_frames x n_unit
83
+ return:
84
+ dict of B x n_frames x feat
85
+ '''
86
+
87
+ x = self.unit_embed(units) + self.f0_embed((1+ f0 / 700).log()) + self.volume_embed(volume)
88
+ if self.n_spk is not None and self.n_spk > 1:
89
+ if spk_mix_dict is not None:
90
+ for k, v in spk_mix_dict.items():
91
+ spk_id_torch = torch.LongTensor(np.array([[k]])).to(units.device)
92
+ x = x + v * self.spk_embed(spk_id_torch)
93
+ else:
94
+ x = x + self.spk_embed(spk_id)
95
+ if self.aug_shift_embed is not None and aug_shift is not None:
96
+ x = x + self.aug_shift_embed(aug_shift / 5)
97
+ x = self.decoder(x, gt_spec=gt_spec, infer=infer, infer_speedup=infer_speedup, method=method, k_step=k_step, use_tqdm=use_tqdm)
98
+
99
+ return x
100
+
diffusion/vocoder.py ADDED
@@ -0,0 +1,94 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from vdecoder.nsf_hifigan.nvSTFT import STFT
3
+ from vdecoder.nsf_hifigan.models import load_model,load_config
4
+ from torchaudio.transforms import Resample
5
+
6
+
7
+ class Vocoder:
8
+ def __init__(self, vocoder_type, vocoder_ckpt, device = None):
9
+ if device is None:
10
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
11
+ self.device = device
12
+
13
+ if vocoder_type == 'nsf-hifigan':
14
+ self.vocoder = NsfHifiGAN(vocoder_ckpt, device = device)
15
+ elif vocoder_type == 'nsf-hifigan-log10':
16
+ self.vocoder = NsfHifiGANLog10(vocoder_ckpt, device = device)
17
+ else:
18
+ raise ValueError(f" [x] Unknown vocoder: {vocoder_type}")
19
+
20
+ self.resample_kernel = {}
21
+ self.vocoder_sample_rate = self.vocoder.sample_rate()
22
+ self.vocoder_hop_size = self.vocoder.hop_size()
23
+ self.dimension = self.vocoder.dimension()
24
+
25
+ def extract(self, audio, sample_rate, keyshift=0):
26
+
27
+ # resample
28
+ if sample_rate == self.vocoder_sample_rate:
29
+ audio_res = audio
30
+ else:
31
+ key_str = str(sample_rate)
32
+ if key_str not in self.resample_kernel:
33
+ self.resample_kernel[key_str] = Resample(sample_rate, self.vocoder_sample_rate, lowpass_filter_width = 128).to(self.device)
34
+ audio_res = self.resample_kernel[key_str](audio)
35
+
36
+ # extract
37
+ mel = self.vocoder.extract(audio_res, keyshift=keyshift) # B, n_frames, bins
38
+ return mel
39
+
40
+ def infer(self, mel, f0):
41
+ f0 = f0[:,:mel.size(1),0] # B, n_frames
42
+ audio = self.vocoder(mel, f0)
43
+ return audio
44
+
45
+
46
+ class NsfHifiGAN(torch.nn.Module):
47
+ def __init__(self, model_path, device=None):
48
+ super().__init__()
49
+ if device is None:
50
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
51
+ self.device = device
52
+ self.model_path = model_path
53
+ self.model = None
54
+ self.h = load_config(model_path)
55
+ self.stft = STFT(
56
+ self.h.sampling_rate,
57
+ self.h.num_mels,
58
+ self.h.n_fft,
59
+ self.h.win_size,
60
+ self.h.hop_size,
61
+ self.h.fmin,
62
+ self.h.fmax)
63
+
64
+ def sample_rate(self):
65
+ return self.h.sampling_rate
66
+
67
+ def hop_size(self):
68
+ return self.h.hop_size
69
+
70
+ def dimension(self):
71
+ return self.h.num_mels
72
+
73
+ def extract(self, audio, keyshift=0):
74
+ mel = self.stft.get_mel(audio, keyshift=keyshift).transpose(1, 2) # B, n_frames, bins
75
+ return mel
76
+
77
+ def forward(self, mel, f0):
78
+ if self.model is None:
79
+ print('| Load HifiGAN: ', self.model_path)
80
+ self.model, self.h = load_model(self.model_path, device=self.device)
81
+ with torch.no_grad():
82
+ c = mel.transpose(1, 2)
83
+ audio = self.model(c, f0)
84
+ return audio
85
+
86
+ class NsfHifiGANLog10(NsfHifiGAN):
87
+ def forward(self, mel, f0):
88
+ if self.model is None:
89
+ print('| Load HifiGAN: ', self.model_path)
90
+ self.model, self.h = load_model(self.model_path, device=self.device)
91
+ with torch.no_grad():
92
+ c = 0.434294 * mel.transpose(1, 2)
93
+ audio = self.model(c, f0)
94
+ return audio
diffusion/wavenet.py ADDED
@@ -0,0 +1,108 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ from math import sqrt
3
+
4
+ import torch
5
+ import torch.nn as nn
6
+ import torch.nn.functional as F
7
+ from torch.nn import Mish
8
+
9
+
10
+ class Conv1d(torch.nn.Conv1d):
11
+ def __init__(self, *args, **kwargs):
12
+ super().__init__(*args, **kwargs)
13
+ nn.init.kaiming_normal_(self.weight)
14
+
15
+
16
+ class SinusoidalPosEmb(nn.Module):
17
+ def __init__(self, dim):
18
+ super().__init__()
19
+ self.dim = dim
20
+
21
+ def forward(self, x):
22
+ device = x.device
23
+ half_dim = self.dim // 2
24
+ emb = math.log(10000) / (half_dim - 1)
25
+ emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
26
+ emb = x[:, None] * emb[None, :]
27
+ emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
28
+ return emb
29
+
30
+
31
+ class ResidualBlock(nn.Module):
32
+ def __init__(self, encoder_hidden, residual_channels, dilation):
33
+ super().__init__()
34
+ self.residual_channels = residual_channels
35
+ self.dilated_conv = nn.Conv1d(
36
+ residual_channels,
37
+ 2 * residual_channels,
38
+ kernel_size=3,
39
+ padding=dilation,
40
+ dilation=dilation
41
+ )
42
+ self.diffusion_projection = nn.Linear(residual_channels, residual_channels)
43
+ self.conditioner_projection = nn.Conv1d(encoder_hidden, 2 * residual_channels, 1)
44
+ self.output_projection = nn.Conv1d(residual_channels, 2 * residual_channels, 1)
45
+
46
+ def forward(self, x, conditioner, diffusion_step):
47
+ diffusion_step = self.diffusion_projection(diffusion_step).unsqueeze(-1)
48
+ conditioner = self.conditioner_projection(conditioner)
49
+ y = x + diffusion_step
50
+
51
+ y = self.dilated_conv(y) + conditioner
52
+
53
+ # Using torch.split instead of torch.chunk to avoid using onnx::Slice
54
+ gate, filter = torch.split(y, [self.residual_channels, self.residual_channels], dim=1)
55
+ y = torch.sigmoid(gate) * torch.tanh(filter)
56
+
57
+ y = self.output_projection(y)
58
+
59
+ # Using torch.split instead of torch.chunk to avoid using onnx::Slice
60
+ residual, skip = torch.split(y, [self.residual_channels, self.residual_channels], dim=1)
61
+ return (x + residual) / math.sqrt(2.0), skip
62
+
63
+
64
+ class WaveNet(nn.Module):
65
+ def __init__(self, in_dims=128, n_layers=20, n_chans=384, n_hidden=256):
66
+ super().__init__()
67
+ self.input_projection = Conv1d(in_dims, n_chans, 1)
68
+ self.diffusion_embedding = SinusoidalPosEmb(n_chans)
69
+ self.mlp = nn.Sequential(
70
+ nn.Linear(n_chans, n_chans * 4),
71
+ Mish(),
72
+ nn.Linear(n_chans * 4, n_chans)
73
+ )
74
+ self.residual_layers = nn.ModuleList([
75
+ ResidualBlock(
76
+ encoder_hidden=n_hidden,
77
+ residual_channels=n_chans,
78
+ dilation=1
79
+ )
80
+ for i in range(n_layers)
81
+ ])
82
+ self.skip_projection = Conv1d(n_chans, n_chans, 1)
83
+ self.output_projection = Conv1d(n_chans, in_dims, 1)
84
+ nn.init.zeros_(self.output_projection.weight)
85
+
86
+ def forward(self, spec, diffusion_step, cond):
87
+ """
88
+ :param spec: [B, 1, M, T]
89
+ :param diffusion_step: [B, 1]
90
+ :param cond: [B, M, T]
91
+ :return:
92
+ """
93
+ x = spec.squeeze(1)
94
+ x = self.input_projection(x) # [B, residual_channel, T]
95
+
96
+ x = F.relu(x)
97
+ diffusion_step = self.diffusion_embedding(diffusion_step)
98
+ diffusion_step = self.mlp(diffusion_step)
99
+ skip = []
100
+ for layer in self.residual_layers:
101
+ x, skip_connection = layer(x, cond, diffusion_step)
102
+ skip.append(skip_connection)
103
+
104
+ x = torch.sum(torch.stack(skip), dim=0) / sqrt(len(self.residual_layers))
105
+ x = self.skip_projection(x)
106
+ x = F.relu(x)
107
+ x = self.output_projection(x) # [B, mel_bins, T]
108
+ return x[:, None, :, :]
inference/__init__.py ADDED
File without changes
inference/chunks_temp.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"info": "temp_dict"}
inference/infer_tool.py ADDED
@@ -0,0 +1,407 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import hashlib
2
+ import io
3
+ import json
4
+ import logging
5
+ import os
6
+ import time
7
+ from pathlib import Path
8
+ from inference import slicer
9
+ import gc
10
+
11
+ import librosa
12
+ import numpy as np
13
+ # import onnxruntime
14
+ import soundfile
15
+ import torch
16
+ import torchaudio
17
+
18
+ import cluster
19
+ import utils
20
+ from models import SynthesizerTrn
21
+
22
+ from diffusion.unit2mel import load_model_vocoder
23
+ import yaml
24
+
25
+ logging.getLogger('matplotlib').setLevel(logging.WARNING)
26
+
27
+
28
+ def read_temp(file_name):
29
+ if not os.path.exists(file_name):
30
+ with open(file_name, "w") as f:
31
+ f.write(json.dumps({"info": "temp_dict"}))
32
+ return {}
33
+ else:
34
+ try:
35
+ with open(file_name, "r") as f:
36
+ data = f.read()
37
+ data_dict = json.loads(data)
38
+ if os.path.getsize(file_name) > 50 * 1024 * 1024:
39
+ f_name = file_name.replace("\\", "/").split("/")[-1]
40
+ print(f"clean {f_name}")
41
+ for wav_hash in list(data_dict.keys()):
42
+ if int(time.time()) - int(data_dict[wav_hash]["time"]) > 14 * 24 * 3600:
43
+ del data_dict[wav_hash]
44
+ except Exception as e:
45
+ print(e)
46
+ print(f"{file_name} error,auto rebuild file")
47
+ data_dict = {"info": "temp_dict"}
48
+ return data_dict
49
+
50
+
51
+ def write_temp(file_name, data):
52
+ with open(file_name, "w") as f:
53
+ f.write(json.dumps(data))
54
+
55
+
56
+ def timeit(func):
57
+ def run(*args, **kwargs):
58
+ t = time.time()
59
+ res = func(*args, **kwargs)
60
+ print('executing \'%s\' costed %.3fs' % (func.__name__, time.time() - t))
61
+ return res
62
+
63
+ return run
64
+
65
+
66
+ def format_wav(audio_path):
67
+ if Path(audio_path).suffix == '.wav':
68
+ return
69
+ raw_audio, raw_sample_rate = librosa.load(audio_path, mono=True, sr=None)
70
+ soundfile.write(Path(audio_path).with_suffix(".wav"), raw_audio, raw_sample_rate)
71
+
72
+
73
+ def get_end_file(dir_path, end):
74
+ file_lists = []
75
+ for root, dirs, files in os.walk(dir_path):
76
+ files = [f for f in files if f[0] != '.']
77
+ dirs[:] = [d for d in dirs if d[0] != '.']
78
+ for f_file in files:
79
+ if f_file.endswith(end):
80
+ file_lists.append(os.path.join(root, f_file).replace("\\", "/"))
81
+ return file_lists
82
+
83
+
84
+ def get_md5(content):
85
+ return hashlib.new("md5", content).hexdigest()
86
+
87
+ def fill_a_to_b(a, b):
88
+ if len(a) < len(b):
89
+ for _ in range(0, len(b) - len(a)):
90
+ a.append(a[0])
91
+
92
+ def mkdir(paths: list):
93
+ for path in paths:
94
+ if not os.path.exists(path):
95
+ os.mkdir(path)
96
+
97
+ def pad_array(arr, target_length):
98
+ current_length = arr.shape[0]
99
+ if current_length >= target_length:
100
+ return arr
101
+ else:
102
+ pad_width = target_length - current_length
103
+ pad_left = pad_width // 2
104
+ pad_right = pad_width - pad_left
105
+ padded_arr = np.pad(arr, (pad_left, pad_right), 'constant', constant_values=(0, 0))
106
+ return padded_arr
107
+
108
+ def split_list_by_n(list_collection, n, pre=0):
109
+ for i in range(0, len(list_collection), n):
110
+ yield list_collection[i-pre if i-pre>=0 else i: i + n]
111
+
112
+
113
+ class F0FilterException(Exception):
114
+ pass
115
+
116
+ class Svc(object):
117
+ def __init__(self, net_g_path, config_path,
118
+ device=None,
119
+ cluster_model_path="logs/44k/kmeans_10000.pt",
120
+ nsf_hifigan_enhance = False,
121
+ diffusion_model_path="logs/44k/diffusion/model_0.pt",
122
+ diffusion_config_path="configs/diffusion.yaml",
123
+ shallow_diffusion = False,
124
+ only_diffusion = False,
125
+ ):
126
+ self.net_g_path = net_g_path
127
+ self.only_diffusion = only_diffusion
128
+ self.shallow_diffusion = shallow_diffusion
129
+ if device is None:
130
+ # self.dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
131
+ self.dev = torch.device("cpu")
132
+ else:
133
+ self.dev = torch.device(device)
134
+ self.net_g_ms = None
135
+ if not self.only_diffusion:
136
+ self.hps_ms = utils.get_hparams_from_file(config_path)
137
+ self.target_sample = self.hps_ms.data.sampling_rate
138
+ self.hop_size = self.hps_ms.data.hop_length
139
+ self.spk2id = self.hps_ms.spk
140
+ try:
141
+ self.speech_encoder = self.hps_ms.model.speech_encoder
142
+ except Exception as e:
143
+ self.speech_encoder = 'vec768l12'
144
+
145
+ self.nsf_hifigan_enhance = nsf_hifigan_enhance
146
+ if self.shallow_diffusion or self.only_diffusion:
147
+ if os.path.exists(diffusion_model_path) and os.path.exists(diffusion_model_path):
148
+ self.diffusion_model,self.vocoder,self.diffusion_args = load_model_vocoder(diffusion_model_path,self.dev,config_path=diffusion_config_path)
149
+ if self.only_diffusion:
150
+ self.target_sample = self.diffusion_args.data.sampling_rate
151
+ self.hop_size = self.diffusion_args.data.block_size
152
+ self.spk2id = self.diffusion_args.spk
153
+ self.speech_encoder = self.diffusion_args.data.encoder
154
+ else:
155
+ print("No diffusion model or config found. Shallow diffusion mode will False")
156
+ self.shallow_diffusion = self.only_diffusion = False
157
+
158
+ # load hubert and model
159
+ if not self.only_diffusion:
160
+ self.load_model()
161
+ self.hubert_model = utils.get_speech_encoder(self.speech_encoder,device=self.dev)
162
+ self.volume_extractor = utils.Volume_Extractor(self.hop_size)
163
+ else:
164
+ self.hubert_model = utils.get_speech_encoder(self.diffusion_args.data.encoder,device=self.dev)
165
+ self.volume_extractor = utils.Volume_Extractor(self.diffusion_args.data.block_size)
166
+
167
+ if os.path.exists(cluster_model_path):
168
+ self.cluster_model = cluster.get_cluster_model(cluster_model_path)
169
+ if self.shallow_diffusion : self.nsf_hifigan_enhance = False
170
+ if self.nsf_hifigan_enhance:
171
+ from modules.enhancer import Enhancer
172
+ self.enhancer = Enhancer('nsf-hifigan', 'pretrain/nsf_hifigan/model',device=self.dev)
173
+
174
+ def load_model(self):
175
+ # get model configuration
176
+ self.net_g_ms = SynthesizerTrn(
177
+ self.hps_ms.data.filter_length // 2 + 1,
178
+ self.hps_ms.train.segment_size // self.hps_ms.data.hop_length,
179
+ **self.hps_ms.model)
180
+ _ = utils.load_checkpoint(self.net_g_path, self.net_g_ms, None)
181
+ if "half" in self.net_g_path and torch.cuda.is_available():
182
+ _ = self.net_g_ms.half().eval().to(self.dev)
183
+ else:
184
+ _ = self.net_g_ms.eval().to(self.dev)
185
+
186
+
187
+
188
+ def get_unit_f0(self, wav, tran, cluster_infer_ratio, speaker, f0_filter ,f0_predictor,cr_threshold=0.05):
189
+
190
+ 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)
191
+
192
+ f0, uv = f0_predictor_object.compute_f0_uv(wav)
193
+ if f0_filter and sum(f0) == 0:
194
+ raise F0FilterException("No voice detected")
195
+ f0 = torch.FloatTensor(f0).to(self.dev)
196
+ uv = torch.FloatTensor(uv).to(self.dev)
197
+
198
+ f0 = f0 * 2 ** (tran / 12)
199
+ f0 = f0.unsqueeze(0)
200
+ uv = uv.unsqueeze(0)
201
+
202
+ wav16k = librosa.resample(wav, orig_sr=self.target_sample, target_sr=16000)
203
+ wav16k = torch.from_numpy(wav16k).to(self.dev)
204
+ c = self.hubert_model.encoder(wav16k)
205
+ c = utils.repeat_expand_2d(c.squeeze(0), f0.shape[1])
206
+
207
+ if cluster_infer_ratio !=0:
208
+ cluster_c = cluster.get_cluster_center_result(self.cluster_model, c.cpu().numpy().T, speaker).T
209
+ cluster_c = torch.FloatTensor(cluster_c).to(self.dev)
210
+ c = cluster_infer_ratio * cluster_c + (1 - cluster_infer_ratio) * c
211
+
212
+ c = c.unsqueeze(0)
213
+ return c, f0, uv
214
+
215
+ def infer(self, speaker, tran, raw_path,
216
+ cluster_infer_ratio=0,
217
+ auto_predict_f0=False,
218
+ noice_scale=0.4,
219
+ f0_filter=False,
220
+ f0_predictor='pm',
221
+ enhancer_adaptive_key = 0,
222
+ cr_threshold = 0.05,
223
+ k_step = 100
224
+ ):
225
+
226
+ speaker_id = self.spk2id.get(speaker)
227
+ if not speaker_id and type(speaker) is int:
228
+ if len(self.spk2id.__dict__) >= speaker:
229
+ speaker_id = speaker
230
+ sid = torch.LongTensor([int(speaker_id)]).to(self.dev).unsqueeze(0)
231
+ wav, sr = librosa.load(raw_path, sr=self.target_sample)
232
+ c, f0, uv = self.get_unit_f0(wav, tran, cluster_infer_ratio, speaker, f0_filter,f0_predictor,cr_threshold=cr_threshold)
233
+ if "half" in self.net_g_path and torch.cuda.is_available():
234
+ c = c.half()
235
+ with torch.no_grad():
236
+ start = time.time()
237
+ if not self.only_diffusion:
238
+ audio,f0 = self.net_g_ms.infer(c, f0=f0, g=sid, uv=uv, predict_f0=auto_predict_f0, noice_scale=noice_scale)
239
+ audio = audio[0,0].data.float()
240
+ if self.shallow_diffusion:
241
+ audio_mel = self.vocoder.extract(audio[None,:],self.target_sample)
242
+ else:
243
+ audio = torch.FloatTensor(wav).to(self.dev)
244
+ audio_mel = None
245
+ if self.only_diffusion or self.shallow_diffusion:
246
+ vol = self.volume_extractor.extract(audio[None,:])[None,:,None].to(self.dev)
247
+ f0 = f0[:,:,None]
248
+ c = c.transpose(-1,-2)
249
+ audio_mel = self.diffusion_model(
250
+ c,
251
+ f0,
252
+ vol,
253
+ spk_id = sid,
254
+ spk_mix_dict = None,
255
+ gt_spec=audio_mel,
256
+ infer=True,
257
+ infer_speedup=self.diffusion_args.infer.speedup,
258
+ method=self.diffusion_args.infer.method,
259
+ k_step=k_step)
260
+ audio = self.vocoder.infer(audio_mel, f0).squeeze()
261
+ if self.nsf_hifigan_enhance:
262
+ audio, _ = self.enhancer.enhance(
263
+ audio[None,:],
264
+ self.target_sample,
265
+ f0[:,:,None],
266
+ self.hps_ms.data.hop_length,
267
+ adaptive_key = enhancer_adaptive_key)
268
+ use_time = time.time() - start
269
+ print("vits use time:{}".format(use_time))
270
+ return audio, audio.shape[-1]
271
+
272
+ def clear_empty(self):
273
+ # clean up vram
274
+ torch.cuda.empty_cache()
275
+
276
+ def unload_model(self):
277
+ # unload model
278
+ self.net_g_ms = self.net_g_ms.to("cpu")
279
+ del self.net_g_ms
280
+ if hasattr(self,"enhancer"):
281
+ self.enhancer.enhancer = self.enhancer.enhancer.to("cpu")
282
+ del self.enhancer.enhancer
283
+ del self.enhancer
284
+ gc.collect()
285
+
286
+ def slice_inference(self,
287
+ raw_audio_path,
288
+ spk,
289
+ tran,
290
+ slice_db,
291
+ cluster_infer_ratio,
292
+ auto_predict_f0,
293
+ noice_scale,
294
+ pad_seconds=0.5,
295
+ clip_seconds=0,
296
+ lg_num=0,
297
+ lgr_num =0.75,
298
+ f0_predictor='pm',
299
+ enhancer_adaptive_key = 0,
300
+ cr_threshold = 0.05,
301
+ k_step = 100
302
+ ):
303
+ wav_path = Path(raw_audio_path).with_suffix('.wav')
304
+ chunks = slicer.cut(wav_path, db_thresh=slice_db)
305
+ audio_data, audio_sr = slicer.chunks2audio(wav_path, chunks)
306
+ per_size = int(clip_seconds*audio_sr)
307
+ lg_size = int(lg_num*audio_sr)
308
+ lg_size_r = int(lg_size*lgr_num)
309
+ lg_size_c_l = (lg_size-lg_size_r)//2
310
+ lg_size_c_r = lg_size-lg_size_r-lg_size_c_l
311
+ lg = np.linspace(0,1,lg_size_r) if lg_size!=0 else 0
312
+
313
+ audio = []
314
+ for (slice_tag, data) in audio_data:
315
+ print(f'#=====segment start, {round(len(data) / audio_sr, 3)}s======')
316
+ # padd
317
+ length = int(np.ceil(len(data) / audio_sr * self.target_sample))
318
+ if slice_tag:
319
+ print('jump empty segment')
320
+ _audio = np.zeros(length)
321
+ audio.extend(list(pad_array(_audio, length)))
322
+ continue
323
+ if per_size != 0:
324
+ datas = split_list_by_n(data, per_size,lg_size)
325
+ else:
326
+ datas = [data]
327
+ for k,dat in enumerate(datas):
328
+ per_length = int(np.ceil(len(dat) / audio_sr * self.target_sample)) if clip_seconds!=0 else length
329
+ if clip_seconds!=0: print(f'###=====segment clip start, {round(len(dat) / audio_sr, 3)}s======')
330
+ # padd
331
+ pad_len = int(audio_sr * pad_seconds)
332
+ dat = np.concatenate([np.zeros([pad_len]), dat, np.zeros([pad_len])])
333
+ raw_path = io.BytesIO()
334
+ soundfile.write(raw_path, dat, audio_sr, format="wav")
335
+ raw_path.seek(0)
336
+ out_audio, out_sr = self.infer(spk, tran, raw_path,
337
+ cluster_infer_ratio=cluster_infer_ratio,
338
+ auto_predict_f0=auto_predict_f0,
339
+ noice_scale=noice_scale,
340
+ f0_predictor = f0_predictor,
341
+ enhancer_adaptive_key = enhancer_adaptive_key,
342
+ cr_threshold = cr_threshold,
343
+ k_step = k_step
344
+ )
345
+ _audio = out_audio.cpu().numpy()
346
+ pad_len = int(self.target_sample * pad_seconds)
347
+ _audio = _audio[pad_len:-pad_len]
348
+ _audio = pad_array(_audio, per_length)
349
+ if lg_size!=0 and k!=0:
350
+ lg1 = audio[-(lg_size_r+lg_size_c_r):-lg_size_c_r] if lgr_num != 1 else audio[-lg_size:]
351
+ lg2 = _audio[lg_size_c_l:lg_size_c_l+lg_size_r] if lgr_num != 1 else _audio[0:lg_size]
352
+ lg_pre = lg1*(1-lg)+lg2*lg
353
+ audio = audio[0:-(lg_size_r+lg_size_c_r)] if lgr_num != 1 else audio[0:-lg_size]
354
+ audio.extend(lg_pre)
355
+ _audio = _audio[lg_size_c_l+lg_size_r:] if lgr_num != 1 else _audio[lg_size:]
356
+ audio.extend(list(_audio))
357
+ return np.array(audio)
358
+
359
+ class RealTimeVC:
360
+ def __init__(self):
361
+ self.last_chunk = None
362
+ self.last_o = None
363
+ self.chunk_len = 16000 # chunk length
364
+ self.pre_len = 3840 # cross fade length, multiples of 640
365
+
366
+ # Input and output are 1-dimensional numpy waveform arrays
367
+
368
+ def process(self, svc_model, speaker_id, f_pitch_change, input_wav_path,
369
+ cluster_infer_ratio=0,
370
+ auto_predict_f0=False,
371
+ noice_scale=0.4,
372
+ f0_filter=False):
373
+
374
+ import maad
375
+ audio, sr = torchaudio.load(input_wav_path)
376
+ audio = audio.cpu().numpy()[0]
377
+ temp_wav = io.BytesIO()
378
+ if self.last_chunk is None:
379
+ input_wav_path.seek(0)
380
+
381
+ audio, sr = svc_model.infer(speaker_id, f_pitch_change, input_wav_path,
382
+ cluster_infer_ratio=cluster_infer_ratio,
383
+ auto_predict_f0=auto_predict_f0,
384
+ noice_scale=noice_scale,
385
+ f0_filter=f0_filter)
386
+
387
+ audio = audio.cpu().numpy()
388
+ self.last_chunk = audio[-self.pre_len:]
389
+ self.last_o = audio
390
+ return audio[-self.chunk_len:]
391
+ else:
392
+ audio = np.concatenate([self.last_chunk, audio])
393
+ soundfile.write(temp_wav, audio, sr, format="wav")
394
+ temp_wav.seek(0)
395
+
396
+ audio, sr = svc_model.infer(speaker_id, f_pitch_change, temp_wav,
397
+ cluster_infer_ratio=cluster_infer_ratio,
398
+ auto_predict_f0=auto_predict_f0,
399
+ noice_scale=noice_scale,
400
+ f0_filter=f0_filter)
401
+
402
+ audio = audio.cpu().numpy()
403
+ ret = maad.util.crossfade(self.last_o, audio, self.pre_len)
404
+ self.last_chunk = audio[-self.pre_len:]
405
+ self.last_o = audio
406
+ return ret[self.chunk_len:2 * self.chunk_len]
407
+
inference/infer_tool_grad.py ADDED
@@ -0,0 +1,171 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import hashlib
2
+ import json
3
+ import logging
4
+ import os
5
+ import time
6
+ from pathlib import Path
7
+ import io
8
+ import librosa
9
+ import maad
10
+ import numpy as np
11
+ from inference import slicer
12
+ import parselmouth
13
+ import soundfile
14
+ import torch
15
+ import torchaudio
16
+
17
+ # from hubert import hubert_model
18
+ import utils
19
+ from models import SynthesizerTrn
20
+ logging.getLogger('numba').setLevel(logging.WARNING)
21
+ logging.getLogger('matplotlib').setLevel(logging.WARNING)
22
+
23
+
24
+ def resize2d_f0(x, target_len):
25
+ source = np.array(x)
26
+ source[source < 0.001] = np.nan
27
+ target = np.interp(np.arange(0, len(source) * target_len, len(source)) / target_len, np.arange(0, len(source)),
28
+ source)
29
+ res = np.nan_to_num(target)
30
+ return res
31
+
32
+
33
+ def get_f0(x, p_len, f0_up_key=0):
34
+
35
+ time_step = 160 / 16000 * 1000
36
+ f0_min = 50
37
+ f0_max = 1100
38
+ f0_mel_min = 1127 * np.log(1 + f0_min / 700)
39
+ f0_mel_max = 1127 * np.log(1 + f0_max / 700)
40
+
41
+ f0 = parselmouth.Sound(x, 16000).to_pitch_ac(
42
+ time_step=time_step / 1000, voicing_threshold=0.6,
43
+ pitch_floor=f0_min, pitch_ceiling=f0_max).selected_array['frequency']
44
+
45
+ pad_size = (p_len - len(f0) + 1) // 2
46
+ if(pad_size > 0 or p_len - len(f0) - pad_size > 0):
47
+ f0 = np.pad(
48
+ f0, [[pad_size, p_len - len(f0) - pad_size]], mode='constant')
49
+
50
+ f0 *= pow(2, f0_up_key / 12)
51
+ f0_mel = 1127 * np.log(1 + f0 / 700)
52
+ f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - f0_mel_min) * \
53
+ 254 / (f0_mel_max - f0_mel_min) + 1
54
+ f0_mel[f0_mel <= 1] = 1
55
+ f0_mel[f0_mel > 255] = 255
56
+ f0_coarse = np.rint(f0_mel).astype(np.int)
57
+ return f0_coarse, f0
58
+
59
+
60
+ def clean_pitch(input_pitch):
61
+ num_nan = np.sum(input_pitch == 1)
62
+ if num_nan / len(input_pitch) > 0.9:
63
+ input_pitch[input_pitch != 1] = 1
64
+ return input_pitch
65
+
66
+
67
+ def plt_pitch(input_pitch):
68
+ input_pitch = input_pitch.astype(float)
69
+ input_pitch[input_pitch == 1] = np.nan
70
+ return input_pitch
71
+
72
+
73
+ def f0_to_pitch(ff):
74
+ f0_pitch = 69 + 12 * np.log2(ff / 440)
75
+ return f0_pitch
76
+
77
+
78
+ def fill_a_to_b(a, b):
79
+ if len(a) < len(b):
80
+ for _ in range(0, len(b) - len(a)):
81
+ a.append(a[0])
82
+
83
+
84
+ def mkdir(paths: list):
85
+ for path in paths:
86
+ if not os.path.exists(path):
87
+ os.mkdir(path)
88
+
89
+
90
+ class VitsSvc(object):
91
+ def __init__(self):
92
+ self.device = torch.device(
93
+ "cuda" if torch.cuda.is_available() else "cpu")
94
+ self.SVCVITS = None
95
+ self.hps = None
96
+ self.speakers = None
97
+ self.hubert_soft = utils.get_hubert_model()
98
+
99
+ def set_device(self, device):
100
+ self.device = torch.device(device)
101
+ self.hubert_soft.to(self.device)
102
+ if self.SVCVITS != None:
103
+ self.SVCVITS.to(self.device)
104
+
105
+ def loadCheckpoint(self, path):
106
+ self.hps = utils.get_hparams_from_file(
107
+ f"checkpoints/{path}/config.json")
108
+ self.SVCVITS = SynthesizerTrn(
109
+ self.hps.data.filter_length // 2 + 1,
110
+ self.hps.train.segment_size // self.hps.data.hop_length,
111
+ **self.hps.model)
112
+ _ = utils.load_checkpoint(
113
+ f"checkpoints/{path}/model.pth", self.SVCVITS, None)
114
+ _ = self.SVCVITS.eval().to(self.device)
115
+ self.speakers = self.hps.spk
116
+
117
+ def get_units(self, source, sr):
118
+ source = source.unsqueeze(0).to(self.device)
119
+ with torch.inference_mode():
120
+ units = self.hubert_soft.units(source)
121
+ return units
122
+
123
+ def get_unit_pitch(self, in_path, tran):
124
+ source, sr = torchaudio.load(in_path)
125
+ source = torchaudio.functional.resample(source, sr, 16000)
126
+ if len(source.shape) == 2 and source.shape[1] >= 2:
127
+ source = torch.mean(source, dim=0).unsqueeze(0)
128
+ soft = self.get_units(source, sr).squeeze(0).cpu().numpy()
129
+ f0_coarse, f0 = get_f0(source.cpu().numpy()[0], soft.shape[0]*2, tran)
130
+ return soft, f0
131
+
132
+ def infer(self, speaker_id, tran, raw_path):
133
+ speaker_id = self.speakers[speaker_id]
134
+ sid = torch.LongTensor([int(speaker_id)]).to(self.device).unsqueeze(0)
135
+ soft, pitch = self.get_unit_pitch(raw_path, tran)
136
+ f0 = torch.FloatTensor(clean_pitch(pitch)).unsqueeze(0).to(self.device)
137
+ stn_tst = torch.FloatTensor(soft)
138
+ with torch.no_grad():
139
+ x_tst = stn_tst.unsqueeze(0).to(self.device)
140
+ x_tst = torch.repeat_interleave(
141
+ x_tst, repeats=2, dim=1).transpose(1, 2)
142
+ audio, _ = self.SVCVITS.infer(x_tst, f0=f0, g=sid)[
143
+ 0, 0].data.float()
144
+ return audio, audio.shape[-1]
145
+
146
+ def inference(self, srcaudio, chara, tran, slice_db):
147
+ sampling_rate, audio = srcaudio
148
+ audio = (audio / np.iinfo(audio.dtype).max).astype(np.float32)
149
+ if len(audio.shape) > 1:
150
+ audio = librosa.to_mono(audio.transpose(1, 0))
151
+ if sampling_rate != 16000:
152
+ audio = librosa.resample(
153
+ audio, orig_sr=sampling_rate, target_sr=16000)
154
+ soundfile.write("tmpwav.wav", audio, 16000, format="wav")
155
+ chunks = slicer.cut("tmpwav.wav", db_thresh=slice_db)
156
+ audio_data, audio_sr = slicer.chunks2audio("tmpwav.wav", chunks)
157
+ audio = []
158
+ for (slice_tag, data) in audio_data:
159
+ length = int(np.ceil(len(data) / audio_sr *
160
+ self.hps.data.sampling_rate))
161
+ raw_path = io.BytesIO()
162
+ soundfile.write(raw_path, data, audio_sr, format="wav")
163
+ raw_path.seek(0)
164
+ if slice_tag:
165
+ _audio = np.zeros(length)
166
+ else:
167
+ out_audio, out_sr = self.infer(chara, tran, raw_path)
168
+ _audio = out_audio.cpu().numpy()
169
+ audio.extend(list(_audio))
170
+ audio = (np.array(audio) * 32768.0).astype('int16')
171
+ return (self.hps.data.sampling_rate, audio)
inference/slicer.py ADDED
@@ -0,0 +1,142 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import librosa
2
+ import torch
3
+ import torchaudio
4
+
5
+
6
+ class Slicer:
7
+ def __init__(self,
8
+ sr: int,
9
+ threshold: float = -40.,
10
+ min_length: int = 5000,
11
+ min_interval: int = 300,
12
+ hop_size: int = 20,
13
+ max_sil_kept: int = 5000):
14
+ if not min_length >= min_interval >= hop_size:
15
+ raise ValueError('The following condition must be satisfied: min_length >= min_interval >= hop_size')
16
+ if not max_sil_kept >= hop_size:
17
+ raise ValueError('The following condition must be satisfied: max_sil_kept >= hop_size')
18
+ min_interval = sr * min_interval / 1000
19
+ self.threshold = 10 ** (threshold / 20.)
20
+ self.hop_size = round(sr * hop_size / 1000)
21
+ self.win_size = min(round(min_interval), 4 * self.hop_size)
22
+ self.min_length = round(sr * min_length / 1000 / self.hop_size)
23
+ self.min_interval = round(min_interval / self.hop_size)
24
+ self.max_sil_kept = round(sr * max_sil_kept / 1000 / self.hop_size)
25
+
26
+ def _apply_slice(self, waveform, begin, end):
27
+ if len(waveform.shape) > 1:
28
+ return waveform[:, begin * self.hop_size: min(waveform.shape[1], end * self.hop_size)]
29
+ else:
30
+ return waveform[begin * self.hop_size: min(waveform.shape[0], end * self.hop_size)]
31
+
32
+ # @timeit
33
+ def slice(self, waveform):
34
+ if len(waveform.shape) > 1:
35
+ samples = librosa.to_mono(waveform)
36
+ else:
37
+ samples = waveform
38
+ if samples.shape[0] <= self.min_length:
39
+ return {"0": {"slice": False, "split_time": f"0,{len(waveform)}"}}
40
+ rms_list = librosa.feature.rms(y=samples, frame_length=self.win_size, hop_length=self.hop_size).squeeze(0)
41
+ sil_tags = []
42
+ silence_start = None
43
+ clip_start = 0
44
+ for i, rms in enumerate(rms_list):
45
+ # Keep looping while frame is silent.
46
+ if rms < self.threshold:
47
+ # Record start of silent frames.
48
+ if silence_start is None:
49
+ silence_start = i
50
+ continue
51
+ # Keep looping while frame is not silent and silence start has not been recorded.
52
+ if silence_start is None:
53
+ continue
54
+ # Clear recorded silence start if interval is not enough or clip is too short
55
+ is_leading_silence = silence_start == 0 and i > self.max_sil_kept
56
+ need_slice_middle = i - silence_start >= self.min_interval and i - clip_start >= self.min_length
57
+ if not is_leading_silence and not need_slice_middle:
58
+ silence_start = None
59
+ continue
60
+ # Need slicing. Record the range of silent frames to be removed.
61
+ if i - silence_start <= self.max_sil_kept:
62
+ pos = rms_list[silence_start: i + 1].argmin() + silence_start
63
+ if silence_start == 0:
64
+ sil_tags.append((0, pos))
65
+ else:
66
+ sil_tags.append((pos, pos))
67
+ clip_start = pos
68
+ elif i - silence_start <= self.max_sil_kept * 2:
69
+ pos = rms_list[i - self.max_sil_kept: silence_start + self.max_sil_kept + 1].argmin()
70
+ pos += i - self.max_sil_kept
71
+ pos_l = rms_list[silence_start: silence_start + self.max_sil_kept + 1].argmin() + silence_start
72
+ pos_r = rms_list[i - self.max_sil_kept: i + 1].argmin() + i - self.max_sil_kept
73
+ if silence_start == 0:
74
+ sil_tags.append((0, pos_r))
75
+ clip_start = pos_r
76
+ else:
77
+ sil_tags.append((min(pos_l, pos), max(pos_r, pos)))
78
+ clip_start = max(pos_r, pos)
79
+ else:
80
+ pos_l = rms_list[silence_start: silence_start + self.max_sil_kept + 1].argmin() + silence_start
81
+ pos_r = rms_list[i - self.max_sil_kept: i + 1].argmin() + i - self.max_sil_kept
82
+ if silence_start == 0:
83
+ sil_tags.append((0, pos_r))
84
+ else:
85
+ sil_tags.append((pos_l, pos_r))
86
+ clip_start = pos_r
87
+ silence_start = None
88
+ # Deal with trailing silence.
89
+ total_frames = rms_list.shape[0]
90
+ if silence_start is not None and total_frames - silence_start >= self.min_interval:
91
+ silence_end = min(total_frames, silence_start + self.max_sil_kept)
92
+ pos = rms_list[silence_start: silence_end + 1].argmin() + silence_start
93
+ sil_tags.append((pos, total_frames + 1))
94
+ # Apply and return slices.
95
+ if len(sil_tags) == 0:
96
+ return {"0": {"slice": False, "split_time": f"0,{len(waveform)}"}}
97
+ else:
98
+ chunks = []
99
+ # 第一段静音并非从头开始,补上有声片段
100
+ if sil_tags[0][0]:
101
+ chunks.append(
102
+ {"slice": False, "split_time": f"0,{min(waveform.shape[0], sil_tags[0][0] * self.hop_size)}"})
103
+ for i in range(0, len(sil_tags)):
104
+ # 标识有声片段(跳过第一段)
105
+ if i:
106
+ chunks.append({"slice": False,
107
+ "split_time": f"{sil_tags[i - 1][1] * self.hop_size},{min(waveform.shape[0], sil_tags[i][0] * self.hop_size)}"})
108
+ # 标识所有静音片段
109
+ chunks.append({"slice": True,
110
+ "split_time": f"{sil_tags[i][0] * self.hop_size},{min(waveform.shape[0], sil_tags[i][1] * self.hop_size)}"})
111
+ # 最后一段静音并非结尾,补上结尾片段
112
+ if sil_tags[-1][1] * self.hop_size < len(waveform):
113
+ chunks.append({"slice": False, "split_time": f"{sil_tags[-1][1] * self.hop_size},{len(waveform)}"})
114
+ chunk_dict = {}
115
+ for i in range(len(chunks)):
116
+ chunk_dict[str(i)] = chunks[i]
117
+ return chunk_dict
118
+
119
+
120
+ def cut(audio_path, db_thresh=-30, min_len=5000):
121
+ audio, sr = librosa.load(audio_path, sr=None)
122
+ slicer = Slicer(
123
+ sr=sr,
124
+ threshold=db_thresh,
125
+ min_length=min_len
126
+ )
127
+ chunks = slicer.slice(audio)
128
+ return chunks
129
+
130
+
131
+ def chunks2audio(audio_path, chunks):
132
+ chunks = dict(chunks)
133
+ audio, sr = torchaudio.load(audio_path)
134
+ if len(audio.shape) == 2 and audio.shape[1] >= 2:
135
+ audio = torch.mean(audio, dim=0).unsqueeze(0)
136
+ audio = audio.cpu().numpy()[0]
137
+ result = []
138
+ for k, v in chunks.items():
139
+ tag = v["split_time"].split(",")
140
+ if tag[0] != tag[1]:
141
+ result.append((v["slice"], audio[int(tag[0]):int(tag[1])]))
142
+ return result, sr
models.py ADDED
@@ -0,0 +1,420 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import copy
2
+ import math
3
+ import torch
4
+ from torch import nn
5
+ from torch.nn import functional as F
6
+
7
+ import modules.attentions as attentions
8
+ import modules.commons as commons
9
+ import modules.modules as modules
10
+
11
+ from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
12
+ from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
13
+
14
+ import utils
15
+ from modules.commons import init_weights, get_padding
16
+ from vdecoder.hifigan.models import Generator
17
+ from utils import f0_to_coarse
18
+
19
+ class ResidualCouplingBlock(nn.Module):
20
+ def __init__(self,
21
+ channels,
22
+ hidden_channels,
23
+ kernel_size,
24
+ dilation_rate,
25
+ n_layers,
26
+ n_flows=4,
27
+ gin_channels=0):
28
+ super().__init__()
29
+ self.channels = channels
30
+ self.hidden_channels = hidden_channels
31
+ self.kernel_size = kernel_size
32
+ self.dilation_rate = dilation_rate
33
+ self.n_layers = n_layers
34
+ self.n_flows = n_flows
35
+ self.gin_channels = gin_channels
36
+
37
+ self.flows = nn.ModuleList()
38
+ for i in range(n_flows):
39
+ self.flows.append(modules.ResidualCouplingLayer(channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels, mean_only=True))
40
+ self.flows.append(modules.Flip())
41
+
42
+ def forward(self, x, x_mask, g=None, reverse=False):
43
+ if not reverse:
44
+ for flow in self.flows:
45
+ x, _ = flow(x, x_mask, g=g, reverse=reverse)
46
+ else:
47
+ for flow in reversed(self.flows):
48
+ x = flow(x, x_mask, g=g, reverse=reverse)
49
+ return x
50
+
51
+
52
+ class Encoder(nn.Module):
53
+ def __init__(self,
54
+ in_channels,
55
+ out_channels,
56
+ hidden_channels,
57
+ kernel_size,
58
+ dilation_rate,
59
+ n_layers,
60
+ gin_channels=0):
61
+ super().__init__()
62
+ self.in_channels = in_channels
63
+ self.out_channels = out_channels
64
+ self.hidden_channels = hidden_channels
65
+ self.kernel_size = kernel_size
66
+ self.dilation_rate = dilation_rate
67
+ self.n_layers = n_layers
68
+ self.gin_channels = gin_channels
69
+
70
+ self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
71
+ self.enc = modules.WN(hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels)
72
+ self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
73
+
74
+ def forward(self, x, x_lengths, g=None):
75
+ # print(x.shape,x_lengths.shape)
76
+ x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)
77
+ x = self.pre(x) * x_mask
78
+ x = self.enc(x, x_mask, g=g)
79
+ stats = self.proj(x) * x_mask
80
+ m, logs = torch.split(stats, self.out_channels, dim=1)
81
+ z = (m + torch.randn_like(m) * torch.exp(logs)) * x_mask
82
+ return z, m, logs, x_mask
83
+
84
+
85
+ class TextEncoder(nn.Module):
86
+ def __init__(self,
87
+ out_channels,
88
+ hidden_channels,
89
+ kernel_size,
90
+ n_layers,
91
+ gin_channels=0,
92
+ filter_channels=None,
93
+ n_heads=None,
94
+ p_dropout=None):
95
+ super().__init__()
96
+ self.out_channels = out_channels
97
+ self.hidden_channels = hidden_channels
98
+ self.kernel_size = kernel_size
99
+ self.n_layers = n_layers
100
+ self.gin_channels = gin_channels
101
+ self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
102
+ self.f0_emb = nn.Embedding(256, hidden_channels)
103
+
104
+ self.enc_ = attentions.Encoder(
105
+ hidden_channels,
106
+ filter_channels,
107
+ n_heads,
108
+ n_layers,
109
+ kernel_size,
110
+ p_dropout)
111
+
112
+ def forward(self, x, x_mask, f0=None, noice_scale=1):
113
+ x = x + self.f0_emb(f0).transpose(1,2)
114
+ x = self.enc_(x * x_mask, x_mask)
115
+ stats = self.proj(x) * x_mask
116
+ m, logs = torch.split(stats, self.out_channels, dim=1)
117
+ z = (m + torch.randn_like(m) * torch.exp(logs) * noice_scale) * x_mask
118
+
119
+ return z, m, logs, x_mask
120
+
121
+
122
+
123
+ class DiscriminatorP(torch.nn.Module):
124
+ def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
125
+ super(DiscriminatorP, self).__init__()
126
+ self.period = period
127
+ self.use_spectral_norm = use_spectral_norm
128
+ norm_f = weight_norm if use_spectral_norm == False else spectral_norm
129
+ self.convs = nn.ModuleList([
130
+ norm_f(Conv2d(1, 32, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
131
+ norm_f(Conv2d(32, 128, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
132
+ norm_f(Conv2d(128, 512, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
133
+ norm_f(Conv2d(512, 1024, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
134
+ norm_f(Conv2d(1024, 1024, (kernel_size, 1), 1, padding=(get_padding(kernel_size, 1), 0))),
135
+ ])
136
+ self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
137
+
138
+ def forward(self, x):
139
+ fmap = []
140
+
141
+ # 1d to 2d
142
+ b, c, t = x.shape
143
+ if t % self.period != 0: # pad first
144
+ n_pad = self.period - (t % self.period)
145
+ x = F.pad(x, (0, n_pad), "reflect")
146
+ t = t + n_pad
147
+ x = x.view(b, c, t // self.period, self.period)
148
+
149
+ for l in self.convs:
150
+ x = l(x)
151
+ x = F.leaky_relu(x, modules.LRELU_SLOPE)
152
+ fmap.append(x)
153
+ x = self.conv_post(x)
154
+ fmap.append(x)
155
+ x = torch.flatten(x, 1, -1)
156
+
157
+ return x, fmap
158
+
159
+
160
+ class DiscriminatorS(torch.nn.Module):
161
+ def __init__(self, use_spectral_norm=False):
162
+ super(DiscriminatorS, self).__init__()
163
+ norm_f = weight_norm if use_spectral_norm == False else spectral_norm
164
+ self.convs = nn.ModuleList([
165
+ norm_f(Conv1d(1, 16, 15, 1, padding=7)),
166
+ norm_f(Conv1d(16, 64, 41, 4, groups=4, padding=20)),
167
+ norm_f(Conv1d(64, 256, 41, 4, groups=16, padding=20)),
168
+ norm_f(Conv1d(256, 1024, 41, 4, groups=64, padding=20)),
169
+ norm_f(Conv1d(1024, 1024, 41, 4, groups=256, padding=20)),
170
+ norm_f(Conv1d(1024, 1024, 5, 1, padding=2)),
171
+ ])
172
+ self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))
173
+
174
+ def forward(self, x):
175
+ fmap = []
176
+
177
+ for l in self.convs:
178
+ x = l(x)
179
+ x = F.leaky_relu(x, modules.LRELU_SLOPE)
180
+ fmap.append(x)
181
+ x = self.conv_post(x)
182
+ fmap.append(x)
183
+ x = torch.flatten(x, 1, -1)
184
+
185
+ return x, fmap
186
+
187
+
188
+ class MultiPeriodDiscriminator(torch.nn.Module):
189
+ def __init__(self, use_spectral_norm=False):
190
+ super(MultiPeriodDiscriminator, self).__init__()
191
+ periods = [2,3,5,7,11]
192
+
193
+ discs = [DiscriminatorS(use_spectral_norm=use_spectral_norm)]
194
+ discs = discs + [DiscriminatorP(i, use_spectral_norm=use_spectral_norm) for i in periods]
195
+ self.discriminators = nn.ModuleList(discs)
196
+
197
+ def forward(self, y, y_hat):
198
+ y_d_rs = []
199
+ y_d_gs = []
200
+ fmap_rs = []
201
+ fmap_gs = []
202
+ for i, d in enumerate(self.discriminators):
203
+ y_d_r, fmap_r = d(y)
204
+ y_d_g, fmap_g = d(y_hat)
205
+ y_d_rs.append(y_d_r)
206
+ y_d_gs.append(y_d_g)
207
+ fmap_rs.append(fmap_r)
208
+ fmap_gs.append(fmap_g)
209
+
210
+ return y_d_rs, y_d_gs, fmap_rs, fmap_gs
211
+
212
+
213
+ class SpeakerEncoder(torch.nn.Module):
214
+ def __init__(self, mel_n_channels=80, model_num_layers=3, model_hidden_size=256, model_embedding_size=256):
215
+ super(SpeakerEncoder, self).__init__()
216
+ self.lstm = nn.LSTM(mel_n_channels, model_hidden_size, model_num_layers, batch_first=True)
217
+ self.linear = nn.Linear(model_hidden_size, model_embedding_size)
218
+ self.relu = nn.ReLU()
219
+
220
+ def forward(self, mels):
221
+ self.lstm.flatten_parameters()
222
+ _, (hidden, _) = self.lstm(mels)
223
+ embeds_raw = self.relu(self.linear(hidden[-1]))
224
+ return embeds_raw / torch.norm(embeds_raw, dim=1, keepdim=True)
225
+
226
+ def compute_partial_slices(self, total_frames, partial_frames, partial_hop):
227
+ mel_slices = []
228
+ for i in range(0, total_frames-partial_frames, partial_hop):
229
+ mel_range = torch.arange(i, i+partial_frames)
230
+ mel_slices.append(mel_range)
231
+
232
+ return mel_slices
233
+
234
+ def embed_utterance(self, mel, partial_frames=128, partial_hop=64):
235
+ mel_len = mel.size(1)
236
+ last_mel = mel[:,-partial_frames:]
237
+
238
+ if mel_len > partial_frames:
239
+ mel_slices = self.compute_partial_slices(mel_len, partial_frames, partial_hop)
240
+ mels = list(mel[:,s] for s in mel_slices)
241
+ mels.append(last_mel)
242
+ mels = torch.stack(tuple(mels), 0).squeeze(1)
243
+
244
+ with torch.no_grad():
245
+ partial_embeds = self(mels)
246
+ embed = torch.mean(partial_embeds, axis=0).unsqueeze(0)
247
+ #embed = embed / torch.linalg.norm(embed, 2)
248
+ else:
249
+ with torch.no_grad():
250
+ embed = self(last_mel)
251
+
252
+ return embed
253
+
254
+ class F0Decoder(nn.Module):
255
+ def __init__(self,
256
+ out_channels,
257
+ hidden_channels,
258
+ filter_channels,
259
+ n_heads,
260
+ n_layers,
261
+ kernel_size,
262
+ p_dropout,
263
+ spk_channels=0):
264
+ super().__init__()
265
+ self.out_channels = out_channels
266
+ self.hidden_channels = hidden_channels
267
+ self.filter_channels = filter_channels
268
+ self.n_heads = n_heads
269
+ self.n_layers = n_layers
270
+ self.kernel_size = kernel_size
271
+ self.p_dropout = p_dropout
272
+ self.spk_channels = spk_channels
273
+
274
+ self.prenet = nn.Conv1d(hidden_channels, hidden_channels, 3, padding=1)
275
+ self.decoder = attentions.FFT(
276
+ hidden_channels,
277
+ filter_channels,
278
+ n_heads,
279
+ n_layers,
280
+ kernel_size,
281
+ p_dropout)
282
+ self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
283
+ self.f0_prenet = nn.Conv1d(1, hidden_channels , 3, padding=1)
284
+ self.cond = nn.Conv1d(spk_channels, hidden_channels, 1)
285
+
286
+ def forward(self, x, norm_f0, x_mask, spk_emb=None):
287
+ x = torch.detach(x)
288
+ if (spk_emb is not None):
289
+ x = x + self.cond(spk_emb)
290
+ x += self.f0_prenet(norm_f0)
291
+ x = self.prenet(x) * x_mask
292
+ x = self.decoder(x * x_mask, x_mask)
293
+ x = self.proj(x) * x_mask
294
+ return x
295
+
296
+
297
+ class SynthesizerTrn(nn.Module):
298
+ """
299
+ Synthesizer for Training
300
+ """
301
+
302
+ def __init__(self,
303
+ spec_channels,
304
+ segment_size,
305
+ inter_channels,
306
+ hidden_channels,
307
+ filter_channels,
308
+ n_heads,
309
+ n_layers,
310
+ kernel_size,
311
+ p_dropout,
312
+ resblock,
313
+ resblock_kernel_sizes,
314
+ resblock_dilation_sizes,
315
+ upsample_rates,
316
+ upsample_initial_channel,
317
+ upsample_kernel_sizes,
318
+ gin_channels,
319
+ ssl_dim,
320
+ n_speakers,
321
+ sampling_rate=44100,
322
+ **kwargs):
323
+
324
+ super().__init__()
325
+ self.spec_channels = spec_channels
326
+ self.inter_channels = inter_channels
327
+ self.hidden_channels = hidden_channels
328
+ self.filter_channels = filter_channels
329
+ self.n_heads = n_heads
330
+ self.n_layers = n_layers
331
+ self.kernel_size = kernel_size
332
+ self.p_dropout = p_dropout
333
+ self.resblock = resblock
334
+ self.resblock_kernel_sizes = resblock_kernel_sizes
335
+ self.resblock_dilation_sizes = resblock_dilation_sizes
336
+ self.upsample_rates = upsample_rates
337
+ self.upsample_initial_channel = upsample_initial_channel
338
+ self.upsample_kernel_sizes = upsample_kernel_sizes
339
+ self.segment_size = segment_size
340
+ self.gin_channels = gin_channels
341
+ self.ssl_dim = ssl_dim
342
+ self.emb_g = nn.Embedding(n_speakers, gin_channels)
343
+
344
+ self.pre = nn.Conv1d(ssl_dim, hidden_channels, kernel_size=5, padding=2)
345
+
346
+ self.enc_p = TextEncoder(
347
+ inter_channels,
348
+ hidden_channels,
349
+ filter_channels=filter_channels,
350
+ n_heads=n_heads,
351
+ n_layers=n_layers,
352
+ kernel_size=kernel_size,
353
+ p_dropout=p_dropout
354
+ )
355
+ hps = {
356
+ "sampling_rate": sampling_rate,
357
+ "inter_channels": inter_channels,
358
+ "resblock": resblock,
359
+ "resblock_kernel_sizes": resblock_kernel_sizes,
360
+ "resblock_dilation_sizes": resblock_dilation_sizes,
361
+ "upsample_rates": upsample_rates,
362
+ "upsample_initial_channel": upsample_initial_channel,
363
+ "upsample_kernel_sizes": upsample_kernel_sizes,
364
+ "gin_channels": gin_channels,
365
+ }
366
+ self.dec = Generator(h=hps)
367
+ self.enc_q = Encoder(spec_channels, inter_channels, hidden_channels, 5, 1, 16, gin_channels=gin_channels)
368
+ self.flow = ResidualCouplingBlock(inter_channels, hidden_channels, 5, 1, 4, gin_channels=gin_channels)
369
+ self.f0_decoder = F0Decoder(
370
+ 1,
371
+ hidden_channels,
372
+ filter_channels,
373
+ n_heads,
374
+ n_layers,
375
+ kernel_size,
376
+ p_dropout,
377
+ spk_channels=gin_channels
378
+ )
379
+ self.emb_uv = nn.Embedding(2, hidden_channels)
380
+
381
+ def forward(self, c, f0, uv, spec, g=None, c_lengths=None, spec_lengths=None):
382
+ g = self.emb_g(g).transpose(1,2)
383
+ # ssl prenet
384
+ x_mask = torch.unsqueeze(commons.sequence_mask(c_lengths, c.size(2)), 1).to(c.dtype)
385
+ x = self.pre(c) * x_mask + self.emb_uv(uv.long()).transpose(1,2)
386
+
387
+ # f0 predict
388
+ lf0 = 2595. * torch.log10(1. + f0.unsqueeze(1) / 700.) / 500
389
+ norm_lf0 = utils.normalize_f0(lf0, x_mask, uv)
390
+ pred_lf0 = self.f0_decoder(x, norm_lf0, x_mask, spk_emb=g)
391
+
392
+ # encoder
393
+ z_ptemp, m_p, logs_p, _ = self.enc_p(x, x_mask, f0=f0_to_coarse(f0))
394
+ z, m_q, logs_q, spec_mask = self.enc_q(spec, spec_lengths, g=g)
395
+
396
+ # flow
397
+ z_p = self.flow(z, spec_mask, g=g)
398
+ z_slice, pitch_slice, ids_slice = commons.rand_slice_segments_with_pitch(z, f0, spec_lengths, self.segment_size)
399
+
400
+ # nsf decoder
401
+ o = self.dec(z_slice, g=g, f0=pitch_slice)
402
+
403
+ return o, ids_slice, spec_mask, (z, z_p, m_p, logs_p, m_q, logs_q), pred_lf0, norm_lf0, lf0
404
+
405
+ def infer(self, c, f0, uv, g=None, noice_scale=0.35, predict_f0=False):
406
+ c_lengths = (torch.ones(c.size(0)) * c.size(-1)).to(c.device)
407
+ g = self.emb_g(g).transpose(1,2)
408
+ x_mask = torch.unsqueeze(commons.sequence_mask(c_lengths, c.size(2)), 1).to(c.dtype)
409
+ x = self.pre(c) * x_mask + self.emb_uv(uv.long()).transpose(1,2)
410
+
411
+ if predict_f0:
412
+ lf0 = 2595. * torch.log10(1. + f0.unsqueeze(1) / 700.) / 500
413
+ norm_lf0 = utils.normalize_f0(lf0, x_mask, uv, random_scale=False)
414
+ pred_lf0 = self.f0_decoder(x, norm_lf0, x_mask, spk_emb=g)
415
+ f0 = (700 * (torch.pow(10, pred_lf0 * 500 / 2595) - 1)).squeeze(1)
416
+
417
+ z_p, m_p, logs_p, c_mask = self.enc_p(x, x_mask, f0=f0_to_coarse(f0), noice_scale=noice_scale)
418
+ z = self.flow(z_p, c_mask, g=g, reverse=True)
419
+ o = self.dec(z * c_mask, g=g, f0=f0)
420
+ return o,f0
modules/F0Predictor/CrepeF0Predictor.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from modules.F0Predictor.F0Predictor import F0Predictor
2
+ from modules.F0Predictor.crepe import CrepePitchExtractor
3
+ import torch
4
+
5
+ class CrepeF0Predictor(F0Predictor):
6
+ def __init__(self,hop_length=512,f0_min=50,f0_max=1100,device=None,sampling_rate=44100,threshold=0.05,model="full"):
7
+ self.F0Creper = CrepePitchExtractor(hop_length=hop_length,f0_min=f0_min,f0_max=f0_max,device=device,threshold=threshold,model=model)
8
+ self.hop_length = hop_length
9
+ self.f0_min = f0_min
10
+ self.f0_max = f0_max
11
+ self.device = device
12
+ self.threshold = threshold
13
+ self.sampling_rate = sampling_rate
14
+
15
+ def compute_f0(self,wav,p_len=None):
16
+ x = torch.FloatTensor(wav).to(self.device)
17
+ if p_len is None:
18
+ p_len = x.shape[0]//self.hop_length
19
+ else:
20
+ assert abs(p_len-x.shape[0]//self.hop_length) < 4, "pad length error"
21
+ f0,uv = self.F0Creper(x[None,:].float(),self.sampling_rate,pad_to=p_len)
22
+ return f0
23
+
24
+ def compute_f0_uv(self,wav,p_len=None):
25
+ x = torch.FloatTensor(wav).to(self.device)
26
+ if p_len is None:
27
+ p_len = x.shape[0]//self.hop_length
28
+ else:
29
+ assert abs(p_len-x.shape[0]//self.hop_length) < 4, "pad length error"
30
+ f0,uv = self.F0Creper(x[None,:].float(),self.sampling_rate,pad_to=p_len)
31
+ return f0,uv
modules/F0Predictor/DioF0Predictor.py ADDED
@@ -0,0 +1,85 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from modules.F0Predictor.F0Predictor import F0Predictor
2
+ import pyworld
3
+ import numpy as np
4
+
5
+ class DioF0Predictor(F0Predictor):
6
+ def __init__(self,hop_length=512,f0_min=50,f0_max=1100,sampling_rate=44100):
7
+ self.hop_length = hop_length
8
+ self.f0_min = f0_min
9
+ self.f0_max = f0_max
10
+ self.sampling_rate = sampling_rate
11
+
12
+ def interpolate_f0(self,f0):
13
+ '''
14
+ 对F0进行插值处理
15
+ '''
16
+
17
+ data = np.reshape(f0, (f0.size, 1))
18
+
19
+ vuv_vector = np.zeros((data.size, 1), dtype=np.float32)
20
+ vuv_vector[data > 0.0] = 1.0
21
+ vuv_vector[data <= 0.0] = 0.0
22
+
23
+ ip_data = data
24
+
25
+ frame_number = data.size
26
+ last_value = 0.0
27
+ for i in range(frame_number):
28
+ if data[i] <= 0.0:
29
+ j = i + 1
30
+ for j in range(i + 1, frame_number):
31
+ if data[j] > 0.0:
32
+ break
33
+ if j < frame_number - 1:
34
+ if last_value > 0.0:
35
+ step = (data[j] - data[i - 1]) / float(j - i)
36
+ for k in range(i, j):
37
+ ip_data[k] = data[i - 1] + step * (k - i + 1)
38
+ else:
39
+ for k in range(i, j):
40
+ ip_data[k] = data[j]
41
+ else:
42
+ for k in range(i, frame_number):
43
+ ip_data[k] = last_value
44
+ else:
45
+ ip_data[i] = data[i] #这里可能存在一个没有必要的拷贝
46
+ last_value = data[i]
47
+
48
+ return ip_data[:,0], vuv_vector[:,0]
49
+
50
+ def resize_f0(self,x, target_len):
51
+ source = np.array(x)
52
+ source[source<0.001] = np.nan
53
+ target = np.interp(np.arange(0, len(source)*target_len, len(source))/ target_len, np.arange(0, len(source)), source)
54
+ res = np.nan_to_num(target)
55
+ return res
56
+
57
+ def compute_f0(self,wav,p_len=None):
58
+ if p_len is None:
59
+ p_len = wav.shape[0]//self.hop_length
60
+ f0, t = pyworld.dio(
61
+ wav.astype(np.double),
62
+ fs=self.sampling_rate,
63
+ f0_floor=self.f0_min,
64
+ f0_ceil=self.f0_max,
65
+ frame_period=1000 * self.hop_length / self.sampling_rate,
66
+ )
67
+ f0 = pyworld.stonemask(wav.astype(np.double), f0, t, self.sampling_rate)
68
+ for index, pitch in enumerate(f0):
69
+ f0[index] = round(pitch, 1)
70
+ return self.interpolate_f0(self.resize_f0(f0, p_len))[0]
71
+
72
+ def compute_f0_uv(self,wav,p_len=None):
73
+ if p_len is None:
74
+ p_len = wav.shape[0]//self.hop_length
75
+ f0, t = pyworld.dio(
76
+ wav.astype(np.double),
77
+ fs=self.sampling_rate,
78
+ f0_floor=self.f0_min,
79
+ f0_ceil=self.f0_max,
80
+ frame_period=1000 * self.hop_length / self.sampling_rate,
81
+ )
82
+ f0 = pyworld.stonemask(wav.astype(np.double), f0, t, self.sampling_rate)
83
+ for index, pitch in enumerate(f0):
84
+ f0[index] = round(pitch, 1)
85
+ return self.interpolate_f0(self.resize_f0(f0, p_len))
modules/F0Predictor/F0Predictor.py ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ class F0Predictor(object):
2
+ def compute_f0(self,wav,p_len):
3
+ '''
4
+ input: wav:[signal_length]
5
+ p_len:int
6
+ output: f0:[signal_length//hop_length]
7
+ '''
8
+ pass
9
+
10
+ def compute_f0_uv(self,wav,p_len):
11
+ '''
12
+ input: wav:[signal_length]
13
+ p_len:int
14
+ output: f0:[signal_length//hop_length],uv:[signal_length//hop_length]
15
+ '''
16
+ pass
modules/F0Predictor/HarvestF0Predictor.py ADDED
@@ -0,0 +1,81 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from modules.F0Predictor.F0Predictor import F0Predictor
2
+ import pyworld
3
+ import numpy as np
4
+
5
+ class HarvestF0Predictor(F0Predictor):
6
+ def __init__(self,hop_length=512,f0_min=50,f0_max=1100,sampling_rate=44100):
7
+ self.hop_length = hop_length
8
+ self.f0_min = f0_min
9
+ self.f0_max = f0_max
10
+ self.sampling_rate = sampling_rate
11
+
12
+ def interpolate_f0(self,f0):
13
+ '''
14
+ 对F0进行插值处理
15
+ '''
16
+
17
+ data = np.reshape(f0, (f0.size, 1))
18
+
19
+ vuv_vector = np.zeros((data.size, 1), dtype=np.float32)
20
+ vuv_vector[data > 0.0] = 1.0
21
+ vuv_vector[data <= 0.0] = 0.0
22
+
23
+ ip_data = data
24
+
25
+ frame_number = data.size
26
+ last_value = 0.0
27
+ for i in range(frame_number):
28
+ if data[i] <= 0.0:
29
+ j = i + 1
30
+ for j in range(i + 1, frame_number):
31
+ if data[j] > 0.0:
32
+ break
33
+ if j < frame_number - 1:
34
+ if last_value > 0.0:
35
+ step = (data[j] - data[i - 1]) / float(j - i)
36
+ for k in range(i, j):
37
+ ip_data[k] = data[i - 1] + step * (k - i + 1)
38
+ else:
39
+ for k in range(i, j):
40
+ ip_data[k] = data[j]
41
+ else:
42
+ for k in range(i, frame_number):
43
+ ip_data[k] = last_value
44
+ else:
45
+ ip_data[i] = data[i] #这里可能存在一个没有必要的拷贝
46
+ last_value = data[i]
47
+
48
+ return ip_data[:,0], vuv_vector[:,0]
49
+
50
+ def resize_f0(self,x, target_len):
51
+ source = np.array(x)
52
+ source[source<0.001] = np.nan
53
+ target = np.interp(np.arange(0, len(source)*target_len, len(source))/ target_len, np.arange(0, len(source)), source)
54
+ res = np.nan_to_num(target)
55
+ return res
56
+
57
+ def compute_f0(self,wav,p_len=None):
58
+ if p_len is None:
59
+ p_len = wav.shape[0]//self.hop_length
60
+ f0, t = pyworld.harvest(
61
+ wav.astype(np.double),
62
+ fs=self.hop_length,
63
+ f0_ceil=self.f0_max,
64
+ f0_floor=self.f0_min,
65
+ frame_period=1000 * self.hop_length / self.sampling_rate,
66
+ )
67
+ f0 = pyworld.stonemask(wav.astype(np.double), f0, t, self.fs)
68
+ return self.interpolate_f0(self.resize_f0(f0, p_len))[0]
69
+
70
+ def compute_f0_uv(self,wav,p_len=None):
71
+ if p_len is None:
72
+ p_len = wav.shape[0]//self.hop_length
73
+ f0, t = pyworld.harvest(
74
+ wav.astype(np.double),
75
+ fs=self.sampling_rate,
76
+ f0_floor=self.f0_min,
77
+ f0_ceil=self.f0_max,
78
+ frame_period=1000 * self.hop_length / self.sampling_rate,
79
+ )
80
+ f0 = pyworld.stonemask(wav.astype(np.double), f0, t, self.sampling_rate)
81
+ return self.interpolate_f0(self.resize_f0(f0, p_len))
modules/F0Predictor/PMF0Predictor.py ADDED
@@ -0,0 +1,83 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from modules.F0Predictor.F0Predictor import F0Predictor
2
+ import parselmouth
3
+ import numpy as np
4
+
5
+ class PMF0Predictor(F0Predictor):
6
+ def __init__(self,hop_length=512,f0_min=50,f0_max=1100,sampling_rate=44100):
7
+ self.hop_length = hop_length
8
+ self.f0_min = f0_min
9
+ self.f0_max = f0_max
10
+ self.sampling_rate = sampling_rate
11
+
12
+
13
+ def interpolate_f0(self,f0):
14
+ '''
15
+ 对F0进行插值处理
16
+ '''
17
+
18
+ data = np.reshape(f0, (f0.size, 1))
19
+
20
+ vuv_vector = np.zeros((data.size, 1), dtype=np.float32)
21
+ vuv_vector[data > 0.0] = 1.0
22
+ vuv_vector[data <= 0.0] = 0.0
23
+
24
+ ip_data = data
25
+
26
+ frame_number = data.size
27
+ last_value = 0.0
28
+ for i in range(frame_number):
29
+ if data[i] <= 0.0:
30
+ j = i + 1
31
+ for j in range(i + 1, frame_number):
32
+ if data[j] > 0.0:
33
+ break
34
+ if j < frame_number - 1:
35
+ if last_value > 0.0:
36
+ step = (data[j] - data[i - 1]) / float(j - i)
37
+ for k in range(i, j):
38
+ ip_data[k] = data[i - 1] + step * (k - i + 1)
39
+ else:
40
+ for k in range(i, j):
41
+ ip_data[k] = data[j]
42
+ else:
43
+ for k in range(i, frame_number):
44
+ ip_data[k] = last_value
45
+ else:
46
+ ip_data[i] = data[i] #这里可能存在一个没有必要的拷贝
47
+ last_value = data[i]
48
+
49
+ return ip_data[:,0], vuv_vector[:,0]
50
+
51
+ def compute_f0(self,wav,p_len=None):
52
+ x = wav
53
+ if p_len is None:
54
+ p_len = x.shape[0]//self.hop_length
55
+ else:
56
+ assert abs(p_len-x.shape[0]//self.hop_length) < 4, "pad length error"
57
+ time_step = self.hop_length / self.sampling_rate * 1000
58
+ f0 = parselmouth.Sound(x, self.sampling_rate).to_pitch_ac(
59
+ time_step=time_step / 1000, voicing_threshold=0.6,
60
+ pitch_floor=self.f0_min, pitch_ceiling=self.f0_max).selected_array['frequency']
61
+
62
+ pad_size=(p_len - len(f0) + 1) // 2
63
+ if(pad_size>0 or p_len - len(f0) - pad_size>0):
64
+ f0 = np.pad(f0,[[pad_size,p_len - len(f0) - pad_size]], mode='constant')
65
+ f0,uv = self.interpolate_f0(f0)
66
+ return f0
67
+
68
+ def compute_f0_uv(self,wav,p_len=None):
69
+ x = wav
70
+ if p_len is None:
71
+ p_len = x.shape[0]//self.hop_length
72
+ else:
73
+ assert abs(p_len-x.shape[0]//self.hop_length) < 4, "pad length error"
74
+ time_step = self.hop_length / self.sampling_rate * 1000
75
+ f0 = parselmouth.Sound(x, self.sampling_rate).to_pitch_ac(
76
+ time_step=time_step / 1000, voicing_threshold=0.6,
77
+ pitch_floor=self.f0_min, pitch_ceiling=self.f0_max).selected_array['frequency']
78
+
79
+ pad_size=(p_len - len(f0) + 1) // 2
80
+ if(pad_size>0 or p_len - len(f0) - pad_size>0):
81
+ f0 = np.pad(f0,[[pad_size,p_len - len(f0) - pad_size]], mode='constant')
82
+ f0,uv = self.interpolate_f0(f0)
83
+ return f0,uv
modules/F0Predictor/__init__.py ADDED
File without changes
modules/F0Predictor/crepe.py ADDED
@@ -0,0 +1,340 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Optional,Union
2
+ try:
3
+ from typing import Literal
4
+ except Exception as e:
5
+ from typing_extensions import Literal
6
+ import numpy as np
7
+ import torch
8
+ import torchcrepe
9
+ from torch import nn
10
+ from torch.nn import functional as F
11
+ import scipy
12
+
13
+ #from:https://github.com/fishaudio/fish-diffusion
14
+
15
+ def repeat_expand(
16
+ content: Union[torch.Tensor, np.ndarray], target_len: int, mode: str = "nearest"
17
+ ):
18
+ """Repeat content to target length.
19
+ This is a wrapper of torch.nn.functional.interpolate.
20
+
21
+ Args:
22
+ content (torch.Tensor): tensor
23
+ target_len (int): target length
24
+ mode (str, optional): interpolation mode. Defaults to "nearest".
25
+
26
+ Returns:
27
+ torch.Tensor: tensor
28
+ """
29
+
30
+ ndim = content.ndim
31
+
32
+ if content.ndim == 1:
33
+ content = content[None, None]
34
+ elif content.ndim == 2:
35
+ content = content[None]
36
+
37
+ assert content.ndim == 3
38
+
39
+ is_np = isinstance(content, np.ndarray)
40
+ if is_np:
41
+ content = torch.from_numpy(content)
42
+
43
+ results = torch.nn.functional.interpolate(content, size=target_len, mode=mode)
44
+
45
+ if is_np:
46
+ results = results.numpy()
47
+
48
+ if ndim == 1:
49
+ return results[0, 0]
50
+ elif ndim == 2:
51
+ return results[0]
52
+
53
+
54
+ class BasePitchExtractor:
55
+ def __init__(
56
+ self,
57
+ hop_length: int = 512,
58
+ f0_min: float = 50.0,
59
+ f0_max: float = 1100.0,
60
+ keep_zeros: bool = True,
61
+ ):
62
+ """Base pitch extractor.
63
+
64
+ Args:
65
+ hop_length (int, optional): Hop length. Defaults to 512.
66
+ f0_min (float, optional): Minimum f0. Defaults to 50.0.
67
+ f0_max (float, optional): Maximum f0. Defaults to 1100.0.
68
+ keep_zeros (bool, optional): Whether keep zeros in pitch. Defaults to True.
69
+ """
70
+
71
+ self.hop_length = hop_length
72
+ self.f0_min = f0_min
73
+ self.f0_max = f0_max
74
+ self.keep_zeros = keep_zeros
75
+
76
+ def __call__(self, x, sampling_rate=44100, pad_to=None):
77
+ raise NotImplementedError("BasePitchExtractor is not callable.")
78
+
79
+ def post_process(self, x, sampling_rate, f0, pad_to):
80
+ if isinstance(f0, np.ndarray):
81
+ f0 = torch.from_numpy(f0).float().to(x.device)
82
+
83
+ if pad_to is None:
84
+ return f0
85
+
86
+ f0 = repeat_expand(f0, pad_to)
87
+
88
+ if self.keep_zeros:
89
+ return f0
90
+
91
+ vuv_vector = torch.zeros_like(f0)
92
+ vuv_vector[f0 > 0.0] = 1.0
93
+ vuv_vector[f0 <= 0.0] = 0.0
94
+
95
+ # 去掉0频率, 并线性插值
96
+ nzindex = torch.nonzero(f0).squeeze()
97
+ f0 = torch.index_select(f0, dim=0, index=nzindex).cpu().numpy()
98
+ time_org = self.hop_length / sampling_rate * nzindex.cpu().numpy()
99
+ time_frame = np.arange(pad_to) * self.hop_length / sampling_rate
100
+
101
+ if f0.shape[0] <= 0:
102
+ return torch.zeros(pad_to, dtype=torch.float, device=x.device),torch.zeros(pad_to, dtype=torch.float, device=x.device)
103
+
104
+ if f0.shape[0] == 1:
105
+ return torch.ones(pad_to, dtype=torch.float, device=x.device) * f0[0],torch.ones(pad_to, dtype=torch.float, device=x.device)
106
+
107
+ # 大概可以用 torch 重写?
108
+ f0 = np.interp(time_frame, time_org, f0, left=f0[0], right=f0[-1])
109
+ vuv_vector = vuv_vector.cpu().numpy()
110
+ vuv_vector = np.ceil(scipy.ndimage.zoom(vuv_vector,pad_to/len(vuv_vector),order = 0))
111
+
112
+ return f0,vuv_vector
113
+
114
+
115
+ class MaskedAvgPool1d(nn.Module):
116
+ def __init__(
117
+ self, kernel_size: int, stride: Optional[int] = None, padding: Optional[int] = 0
118
+ ):
119
+ """An implementation of mean pooling that supports masked values.
120
+
121
+ Args:
122
+ kernel_size (int): The size of the median pooling window.
123
+ stride (int, optional): The stride of the median pooling window. Defaults to None.
124
+ padding (int, optional): The padding of the median pooling window. Defaults to 0.
125
+ """
126
+
127
+ super(MaskedAvgPool1d, self).__init__()
128
+ self.kernel_size = kernel_size
129
+ self.stride = stride or kernel_size
130
+ self.padding = padding
131
+
132
+ def forward(self, x, mask=None):
133
+ ndim = x.dim()
134
+ if ndim == 2:
135
+ x = x.unsqueeze(1)
136
+
137
+ assert (
138
+ x.dim() == 3
139
+ ), "Input tensor must have 2 or 3 dimensions (batch_size, channels, width)"
140
+
141
+ # Apply the mask by setting masked elements to zero, or make NaNs zero
142
+ if mask is None:
143
+ mask = ~torch.isnan(x)
144
+
145
+ # Ensure mask has the same shape as the input tensor
146
+ assert x.shape == mask.shape, "Input tensor and mask must have the same shape"
147
+
148
+ masked_x = torch.where(mask, x, torch.zeros_like(x))
149
+ # Create a ones kernel with the same number of channels as the input tensor
150
+ ones_kernel = torch.ones(x.size(1), 1, self.kernel_size, device=x.device)
151
+
152
+ # Perform sum pooling
153
+ sum_pooled = nn.functional.conv1d(
154
+ masked_x,
155
+ ones_kernel,
156
+ stride=self.stride,
157
+ padding=self.padding,
158
+ groups=x.size(1),
159
+ )
160
+
161
+ # Count the non-masked (valid) elements in each pooling window
162
+ valid_count = nn.functional.conv1d(
163
+ mask.float(),
164
+ ones_kernel,
165
+ stride=self.stride,
166
+ padding=self.padding,
167
+ groups=x.size(1),
168
+ )
169
+ valid_count = valid_count.clamp(min=1) # Avoid division by zero
170
+
171
+ # Perform masked average pooling
172
+ avg_pooled = sum_pooled / valid_count
173
+
174
+ # Fill zero values with NaNs
175
+ avg_pooled[avg_pooled == 0] = float("nan")
176
+
177
+ if ndim == 2:
178
+ return avg_pooled.squeeze(1)
179
+
180
+ return avg_pooled
181
+
182
+
183
+ class MaskedMedianPool1d(nn.Module):
184
+ def __init__(
185
+ self, kernel_size: int, stride: Optional[int] = None, padding: Optional[int] = 0
186
+ ):
187
+ """An implementation of median pooling that supports masked values.
188
+
189
+ This implementation is inspired by the median pooling implementation in
190
+ https://gist.github.com/rwightman/f2d3849281624be7c0f11c85c87c1598
191
+
192
+ Args:
193
+ kernel_size (int): The size of the median pooling window.
194
+ stride (int, optional): The stride of the median pooling window. Defaults to None.
195
+ padding (int, optional): The padding of the median pooling window. Defaults to 0.
196
+ """
197
+
198
+ super(MaskedMedianPool1d, self).__init__()
199
+ self.kernel_size = kernel_size
200
+ self.stride = stride or kernel_size
201
+ self.padding = padding
202
+
203
+ def forward(self, x, mask=None):
204
+ ndim = x.dim()
205
+ if ndim == 2:
206
+ x = x.unsqueeze(1)
207
+
208
+ assert (
209
+ x.dim() == 3
210
+ ), "Input tensor must have 2 or 3 dimensions (batch_size, channels, width)"
211
+
212
+ if mask is None:
213
+ mask = ~torch.isnan(x)
214
+
215
+ assert x.shape == mask.shape, "Input tensor and mask must have the same shape"
216
+
217
+ masked_x = torch.where(mask, x, torch.zeros_like(x))
218
+
219
+ x = F.pad(masked_x, (self.padding, self.padding), mode="reflect")
220
+ mask = F.pad(
221
+ mask.float(), (self.padding, self.padding), mode="constant", value=0
222
+ )
223
+
224
+ x = x.unfold(2, self.kernel_size, self.stride)
225
+ mask = mask.unfold(2, self.kernel_size, self.stride)
226
+
227
+ x = x.contiguous().view(x.size()[:3] + (-1,))
228
+ mask = mask.contiguous().view(mask.size()[:3] + (-1,)).to(x.device)
229
+
230
+ # Combine the mask with the input tensor
231
+ #x_masked = torch.where(mask.bool(), x, torch.fill_(torch.zeros_like(x),float("inf")))
232
+ x_masked = torch.where(mask.bool(), x, torch.FloatTensor([float("inf")]).to(x.device))
233
+
234
+ # Sort the masked tensor along the last dimension
235
+ x_sorted, _ = torch.sort(x_masked, dim=-1)
236
+
237
+ # Compute the count of non-masked (valid) values
238
+ valid_count = mask.sum(dim=-1)
239
+
240
+ # Calculate the index of the median value for each pooling window
241
+ median_idx = (torch.div((valid_count - 1), 2, rounding_mode='trunc')).clamp(min=0)
242
+
243
+ # Gather the median values using the calculated indices
244
+ median_pooled = x_sorted.gather(-1, median_idx.unsqueeze(-1).long()).squeeze(-1)
245
+
246
+ # Fill infinite values with NaNs
247
+ median_pooled[torch.isinf(median_pooled)] = float("nan")
248
+
249
+ if ndim == 2:
250
+ return median_pooled.squeeze(1)
251
+
252
+ return median_pooled
253
+
254
+
255
+ class CrepePitchExtractor(BasePitchExtractor):
256
+ def __init__(
257
+ self,
258
+ hop_length: int = 512,
259
+ f0_min: float = 50.0,
260
+ f0_max: float = 1100.0,
261
+ threshold: float = 0.05,
262
+ keep_zeros: bool = False,
263
+ device = None,
264
+ model: Literal["full", "tiny"] = "full",
265
+ use_fast_filters: bool = True,
266
+ decoder="viterbi"
267
+ ):
268
+ super().__init__(hop_length, f0_min, f0_max, keep_zeros)
269
+ if decoder == "viterbi":
270
+ self.decoder = torchcrepe.decode.viterbi
271
+ elif decoder == "argmax":
272
+ self.decoder = torchcrepe.decode.argmax
273
+ elif decoder == "weighted_argmax":
274
+ self.decoder = torchcrepe.decode.weighted_argmax
275
+ else:
276
+ raise "Unknown decoder"
277
+ self.threshold = threshold
278
+ self.model = model
279
+ self.use_fast_filters = use_fast_filters
280
+ self.hop_length = hop_length
281
+ if device is None:
282
+ self.dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
283
+ else:
284
+ self.dev = torch.device(device)
285
+ if self.use_fast_filters:
286
+ self.median_filter = MaskedMedianPool1d(3, 1, 1).to(device)
287
+ self.mean_filter = MaskedAvgPool1d(3, 1, 1).to(device)
288
+
289
+ def __call__(self, x, sampling_rate=44100, pad_to=None):
290
+ """Extract pitch using crepe.
291
+
292
+
293
+ Args:
294
+ x (torch.Tensor): Audio signal, shape (1, T).
295
+ sampling_rate (int, optional): Sampling rate. Defaults to 44100.
296
+ pad_to (int, optional): Pad to length. Defaults to None.
297
+
298
+ Returns:
299
+ torch.Tensor: Pitch, shape (T // hop_length,).
300
+ """
301
+
302
+ assert x.ndim == 2, f"Expected 2D tensor, got {x.ndim}D tensor."
303
+ assert x.shape[0] == 1, f"Expected 1 channel, got {x.shape[0]} channels."
304
+
305
+ x = x.to(self.dev)
306
+ f0, pd = torchcrepe.predict(
307
+ x,
308
+ sampling_rate,
309
+ self.hop_length,
310
+ self.f0_min,
311
+ self.f0_max,
312
+ pad=True,
313
+ model=self.model,
314
+ batch_size=1024,
315
+ device=x.device,
316
+ return_periodicity=True,
317
+ decoder=self.decoder
318
+ )
319
+
320
+ # Filter, remove silence, set uv threshold, refer to the original warehouse readme
321
+ if self.use_fast_filters:
322
+ pd = self.median_filter(pd)
323
+ else:
324
+ pd = torchcrepe.filter.median(pd, 3)
325
+
326
+ pd = torchcrepe.threshold.Silence(-60.0)(pd, x, sampling_rate, 512)
327
+ f0 = torchcrepe.threshold.At(self.threshold)(f0, pd)
328
+
329
+ if self.use_fast_filters:
330
+ f0 = self.mean_filter(f0)
331
+ else:
332
+ f0 = torchcrepe.filter.mean(f0, 3)
333
+
334
+ f0 = torch.where(torch.isnan(f0), torch.full_like(f0, 0), f0)[0]
335
+
336
+ if torch.all(f0 == 0):
337
+ rtn = f0.cpu().numpy() if pad_to==None else np.zeros(pad_to)
338
+ return rtn,rtn
339
+
340
+ return self.post_process(x, sampling_rate, f0, pad_to)
modules/__init__.py ADDED
File without changes
modules/attentions.py ADDED
@@ -0,0 +1,349 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import copy
2
+ import math
3
+ import numpy as np
4
+ import torch
5
+ from torch import nn
6
+ from torch.nn import functional as F
7
+
8
+ import modules.commons as commons
9
+ import modules.modules as modules
10
+ from modules.modules import LayerNorm
11
+
12
+
13
+ class FFT(nn.Module):
14
+ def __init__(self, hidden_channels, filter_channels, n_heads, n_layers=1, kernel_size=1, p_dropout=0.,
15
+ proximal_bias=False, proximal_init=True, **kwargs):
16
+ super().__init__()
17
+ self.hidden_channels = hidden_channels
18
+ self.filter_channels = filter_channels
19
+ self.n_heads = n_heads
20
+ self.n_layers = n_layers
21
+ self.kernel_size = kernel_size
22
+ self.p_dropout = p_dropout
23
+ self.proximal_bias = proximal_bias
24
+ self.proximal_init = proximal_init
25
+
26
+ self.drop = nn.Dropout(p_dropout)
27
+ self.self_attn_layers = nn.ModuleList()
28
+ self.norm_layers_0 = nn.ModuleList()
29
+ self.ffn_layers = nn.ModuleList()
30
+ self.norm_layers_1 = nn.ModuleList()
31
+ for i in range(self.n_layers):
32
+ self.self_attn_layers.append(
33
+ MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout, proximal_bias=proximal_bias,
34
+ proximal_init=proximal_init))
35
+ self.norm_layers_0.append(LayerNorm(hidden_channels))
36
+ self.ffn_layers.append(
37
+ FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout, causal=True))
38
+ self.norm_layers_1.append(LayerNorm(hidden_channels))
39
+
40
+ def forward(self, x, x_mask):
41
+ """
42
+ x: decoder input
43
+ h: encoder output
44
+ """
45
+ self_attn_mask = commons.subsequent_mask(x_mask.size(2)).to(device=x.device, dtype=x.dtype)
46
+ x = x * x_mask
47
+ for i in range(self.n_layers):
48
+ y = self.self_attn_layers[i](x, x, self_attn_mask)
49
+ y = self.drop(y)
50
+ x = self.norm_layers_0[i](x + y)
51
+
52
+ y = self.ffn_layers[i](x, x_mask)
53
+ y = self.drop(y)
54
+ x = self.norm_layers_1[i](x + y)
55
+ x = x * x_mask
56
+ return x
57
+
58
+
59
+ class Encoder(nn.Module):
60
+ def __init__(self, hidden_channels, filter_channels, n_heads, n_layers, kernel_size=1, p_dropout=0., window_size=4, **kwargs):
61
+ super().__init__()
62
+ self.hidden_channels = hidden_channels
63
+ self.filter_channels = filter_channels
64
+ self.n_heads = n_heads
65
+ self.n_layers = n_layers
66
+ self.kernel_size = kernel_size
67
+ self.p_dropout = p_dropout
68
+ self.window_size = window_size
69
+
70
+ self.drop = nn.Dropout(p_dropout)
71
+ self.attn_layers = nn.ModuleList()
72
+ self.norm_layers_1 = nn.ModuleList()
73
+ self.ffn_layers = nn.ModuleList()
74
+ self.norm_layers_2 = nn.ModuleList()
75
+ for i in range(self.n_layers):
76
+ self.attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout, window_size=window_size))
77
+ self.norm_layers_1.append(LayerNorm(hidden_channels))
78
+ self.ffn_layers.append(FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout))
79
+ self.norm_layers_2.append(LayerNorm(hidden_channels))
80
+
81
+ def forward(self, x, x_mask):
82
+ attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
83
+ x = x * x_mask
84
+ for i in range(self.n_layers):
85
+ y = self.attn_layers[i](x, x, attn_mask)
86
+ y = self.drop(y)
87
+ x = self.norm_layers_1[i](x + y)
88
+
89
+ y = self.ffn_layers[i](x, x_mask)
90
+ y = self.drop(y)
91
+ x = self.norm_layers_2[i](x + y)
92
+ x = x * x_mask
93
+ return x
94
+
95
+
96
+ class Decoder(nn.Module):
97
+ def __init__(self, hidden_channels, filter_channels, n_heads, n_layers, kernel_size=1, p_dropout=0., proximal_bias=False, proximal_init=True, **kwargs):
98
+ super().__init__()
99
+ self.hidden_channels = hidden_channels
100
+ self.filter_channels = filter_channels
101
+ self.n_heads = n_heads
102
+ self.n_layers = n_layers
103
+ self.kernel_size = kernel_size
104
+ self.p_dropout = p_dropout
105
+ self.proximal_bias = proximal_bias
106
+ self.proximal_init = proximal_init
107
+
108
+ self.drop = nn.Dropout(p_dropout)
109
+ self.self_attn_layers = nn.ModuleList()
110
+ self.norm_layers_0 = nn.ModuleList()
111
+ self.encdec_attn_layers = nn.ModuleList()
112
+ self.norm_layers_1 = nn.ModuleList()
113
+ self.ffn_layers = nn.ModuleList()
114
+ self.norm_layers_2 = nn.ModuleList()
115
+ for i in range(self.n_layers):
116
+ self.self_attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout, proximal_bias=proximal_bias, proximal_init=proximal_init))
117
+ self.norm_layers_0.append(LayerNorm(hidden_channels))
118
+ self.encdec_attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout))
119
+ self.norm_layers_1.append(LayerNorm(hidden_channels))
120
+ self.ffn_layers.append(FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout, causal=True))
121
+ self.norm_layers_2.append(LayerNorm(hidden_channels))
122
+
123
+ def forward(self, x, x_mask, h, h_mask):
124
+ """
125
+ x: decoder input
126
+ h: encoder output
127
+ """
128
+ self_attn_mask = commons.subsequent_mask(x_mask.size(2)).to(device=x.device, dtype=x.dtype)
129
+ encdec_attn_mask = h_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
130
+ x = x * x_mask
131
+ for i in range(self.n_layers):
132
+ y = self.self_attn_layers[i](x, x, self_attn_mask)
133
+ y = self.drop(y)
134
+ x = self.norm_layers_0[i](x + y)
135
+
136
+ y = self.encdec_attn_layers[i](x, h, encdec_attn_mask)
137
+ y = self.drop(y)
138
+ x = self.norm_layers_1[i](x + y)
139
+
140
+ y = self.ffn_layers[i](x, x_mask)
141
+ y = self.drop(y)
142
+ x = self.norm_layers_2[i](x + y)
143
+ x = x * x_mask
144
+ return x
145
+
146
+
147
+ class MultiHeadAttention(nn.Module):
148
+ 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):
149
+ super().__init__()
150
+ assert channels % n_heads == 0
151
+
152
+ self.channels = channels
153
+ self.out_channels = out_channels
154
+ self.n_heads = n_heads
155
+ self.p_dropout = p_dropout
156
+ self.window_size = window_size
157
+ self.heads_share = heads_share
158
+ self.block_length = block_length
159
+ self.proximal_bias = proximal_bias
160
+ self.proximal_init = proximal_init
161
+ self.attn = None
162
+
163
+ self.k_channels = channels // n_heads
164
+ self.conv_q = nn.Conv1d(channels, channels, 1)
165
+ self.conv_k = nn.Conv1d(channels, channels, 1)
166
+ self.conv_v = nn.Conv1d(channels, channels, 1)
167
+ self.conv_o = nn.Conv1d(channels, out_channels, 1)
168
+ self.drop = nn.Dropout(p_dropout)
169
+
170
+ if window_size is not None:
171
+ n_heads_rel = 1 if heads_share else n_heads
172
+ rel_stddev = self.k_channels**-0.5
173
+ self.emb_rel_k = nn.Parameter(torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) * rel_stddev)
174
+ self.emb_rel_v = nn.Parameter(torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) * rel_stddev)
175
+
176
+ nn.init.xavier_uniform_(self.conv_q.weight)
177
+ nn.init.xavier_uniform_(self.conv_k.weight)
178
+ nn.init.xavier_uniform_(self.conv_v.weight)
179
+ if proximal_init:
180
+ with torch.no_grad():
181
+ self.conv_k.weight.copy_(self.conv_q.weight)
182
+ self.conv_k.bias.copy_(self.conv_q.bias)
183
+
184
+ def forward(self, x, c, attn_mask=None):
185
+ q = self.conv_q(x)
186
+ k = self.conv_k(c)
187
+ v = self.conv_v(c)
188
+
189
+ x, self.attn = self.attention(q, k, v, mask=attn_mask)
190
+
191
+ x = self.conv_o(x)
192
+ return x
193
+
194
+ def attention(self, query, key, value, mask=None):
195
+ # reshape [b, d, t] -> [b, n_h, t, d_k]
196
+ b, d, t_s, t_t = (*key.size(), query.size(2))
197
+ query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
198
+ key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
199
+ value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
200
+
201
+ scores = torch.matmul(query / math.sqrt(self.k_channels), key.transpose(-2, -1))
202
+ if self.window_size is not None:
203
+ assert t_s == t_t, "Relative attention is only available for self-attention."
204
+ key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)
205
+ rel_logits = self._matmul_with_relative_keys(query /math.sqrt(self.k_channels), key_relative_embeddings)
206
+ scores_local = self._relative_position_to_absolute_position(rel_logits)
207
+ scores = scores + scores_local
208
+ if self.proximal_bias:
209
+ assert t_s == t_t, "Proximal bias is only available for self-attention."
210
+ scores = scores + self._attention_bias_proximal(t_s).to(device=scores.device, dtype=scores.dtype)
211
+ if mask is not None:
212
+ scores = scores.masked_fill(mask == 0, -1e4)
213
+ if self.block_length is not None:
214
+ assert t_s == t_t, "Local attention is only available for self-attention."
215
+ block_mask = torch.ones_like(scores).triu(-self.block_length).tril(self.block_length)
216
+ scores = scores.masked_fill(block_mask == 0, -1e4)
217
+ p_attn = F.softmax(scores, dim=-1) # [b, n_h, t_t, t_s]
218
+ p_attn = self.drop(p_attn)
219
+ output = torch.matmul(p_attn, value)
220
+ if self.window_size is not None:
221
+ relative_weights = self._absolute_position_to_relative_position(p_attn)
222
+ value_relative_embeddings = self._get_relative_embeddings(self.emb_rel_v, t_s)
223
+ output = output + self._matmul_with_relative_values(relative_weights, value_relative_embeddings)
224
+ output = output.transpose(2, 3).contiguous().view(b, d, t_t) # [b, n_h, t_t, d_k] -> [b, d, t_t]
225
+ return output, p_attn
226
+
227
+ def _matmul_with_relative_values(self, x, y):
228
+ """
229
+ x: [b, h, l, m]
230
+ y: [h or 1, m, d]
231
+ ret: [b, h, l, d]
232
+ """
233
+ ret = torch.matmul(x, y.unsqueeze(0))
234
+ return ret
235
+
236
+ def _matmul_with_relative_keys(self, x, y):
237
+ """
238
+ x: [b, h, l, d]
239
+ y: [h or 1, m, d]
240
+ ret: [b, h, l, m]
241
+ """
242
+ ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))
243
+ return ret
244
+
245
+ def _get_relative_embeddings(self, relative_embeddings, length):
246
+ max_relative_position = 2 * self.window_size + 1
247
+ # Pad first before slice to avoid using cond ops.
248
+ pad_length = max(length - (self.window_size + 1), 0)
249
+ slice_start_position = max((self.window_size + 1) - length, 0)
250
+ slice_end_position = slice_start_position + 2 * length - 1
251
+ if pad_length > 0:
252
+ padded_relative_embeddings = F.pad(
253
+ relative_embeddings,
254
+ commons.convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]))
255
+ else:
256
+ padded_relative_embeddings = relative_embeddings
257
+ used_relative_embeddings = padded_relative_embeddings[:,slice_start_position:slice_end_position]
258
+ return used_relative_embeddings
259
+
260
+ def _relative_position_to_absolute_position(self, x):
261
+ """
262
+ x: [b, h, l, 2*l-1]
263
+ ret: [b, h, l, l]
264
+ """
265
+ batch, heads, length, _ = x.size()
266
+ # Concat columns of pad to shift from relative to absolute indexing.
267
+ x = F.pad(x, commons.convert_pad_shape([[0,0],[0,0],[0,0],[0,1]]))
268
+
269
+ # Concat extra elements so to add up to shape (len+1, 2*len-1).
270
+ x_flat = x.view([batch, heads, length * 2 * length])
271
+ x_flat = F.pad(x_flat, commons.convert_pad_shape([[0,0],[0,0],[0,length-1]]))
272
+
273
+ # Reshape and slice out the padded elements.
274
+ x_final = x_flat.view([batch, heads, length+1, 2*length-1])[:, :, :length, length-1:]
275
+ return x_final
276
+
277
+ def _absolute_position_to_relative_position(self, x):
278
+ """
279
+ x: [b, h, l, l]
280
+ ret: [b, h, l, 2*l-1]
281
+ """
282
+ batch, heads, length, _ = x.size()
283
+ # padd along column
284
+ x = F.pad(x, commons.convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, length-1]]))
285
+ x_flat = x.view([batch, heads, length**2 + length*(length -1)])
286
+ # add 0's in the beginning that will skew the elements after reshape
287
+ x_flat = F.pad(x_flat, commons.convert_pad_shape([[0, 0], [0, 0], [length, 0]]))
288
+ x_final = x_flat.view([batch, heads, length, 2*length])[:,:,:,1:]
289
+ return x_final
290
+
291
+ def _attention_bias_proximal(self, length):
292
+ """Bias for self-attention to encourage attention to close positions.
293
+ Args:
294
+ length: an integer scalar.
295
+ Returns:
296
+ a Tensor with shape [1, 1, length, length]
297
+ """
298
+ r = torch.arange(length, dtype=torch.float32)
299
+ diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)
300
+ return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0)
301
+
302
+
303
+ class FFN(nn.Module):
304
+ def __init__(self, in_channels, out_channels, filter_channels, kernel_size, p_dropout=0., activation=None, causal=False):
305
+ super().__init__()
306
+ self.in_channels = in_channels
307
+ self.out_channels = out_channels
308
+ self.filter_channels = filter_channels
309
+ self.kernel_size = kernel_size
310
+ self.p_dropout = p_dropout
311
+ self.activation = activation
312
+ self.causal = causal
313
+
314
+ if causal:
315
+ self.padding = self._causal_padding
316
+ else:
317
+ self.padding = self._same_padding
318
+
319
+ self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size)
320
+ self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size)
321
+ self.drop = nn.Dropout(p_dropout)
322
+
323
+ def forward(self, x, x_mask):
324
+ x = self.conv_1(self.padding(x * x_mask))
325
+ if self.activation == "gelu":
326
+ x = x * torch.sigmoid(1.702 * x)
327
+ else:
328
+ x = torch.relu(x)
329
+ x = self.drop(x)
330
+ x = self.conv_2(self.padding(x * x_mask))
331
+ return x * x_mask
332
+
333
+ def _causal_padding(self, x):
334
+ if self.kernel_size == 1:
335
+ return x
336
+ pad_l = self.kernel_size - 1
337
+ pad_r = 0
338
+ padding = [[0, 0], [0, 0], [pad_l, pad_r]]
339
+ x = F.pad(x, commons.convert_pad_shape(padding))
340
+ return x
341
+
342
+ def _same_padding(self, x):
343
+ if self.kernel_size == 1:
344
+ return x
345
+ pad_l = (self.kernel_size - 1) // 2
346
+ pad_r = self.kernel_size // 2
347
+ padding = [[0, 0], [0, 0], [pad_l, pad_r]]
348
+ x = F.pad(x, commons.convert_pad_shape(padding))
349
+ return x
modules/commons.py ADDED
@@ -0,0 +1,188 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ import numpy as np
3
+ import torch
4
+ from torch import nn
5
+ from torch.nn import functional as F
6
+
7
+ def slice_pitch_segments(x, ids_str, segment_size=4):
8
+ ret = torch.zeros_like(x[:, :segment_size])
9
+ for i in range(x.size(0)):
10
+ idx_str = ids_str[i]
11
+ idx_end = idx_str + segment_size
12
+ ret[i] = x[i, idx_str:idx_end]
13
+ return ret
14
+
15
+ def rand_slice_segments_with_pitch(x, pitch, x_lengths=None, segment_size=4):
16
+ b, d, t = x.size()
17
+ if x_lengths is None:
18
+ x_lengths = t
19
+ ids_str_max = x_lengths - segment_size + 1
20
+ ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
21
+ ret = slice_segments(x, ids_str, segment_size)
22
+ ret_pitch = slice_pitch_segments(pitch, ids_str, segment_size)
23
+ return ret, ret_pitch, ids_str
24
+
25
+ def init_weights(m, mean=0.0, std=0.01):
26
+ classname = m.__class__.__name__
27
+ if classname.find("Conv") != -1:
28
+ m.weight.data.normal_(mean, std)
29
+
30
+
31
+ def get_padding(kernel_size, dilation=1):
32
+ return int((kernel_size*dilation - dilation)/2)
33
+
34
+
35
+ def convert_pad_shape(pad_shape):
36
+ l = pad_shape[::-1]
37
+ pad_shape = [item for sublist in l for item in sublist]
38
+ return pad_shape
39
+
40
+
41
+ def intersperse(lst, item):
42
+ result = [item] * (len(lst) * 2 + 1)
43
+ result[1::2] = lst
44
+ return result
45
+
46
+
47
+ def kl_divergence(m_p, logs_p, m_q, logs_q):
48
+ """KL(P||Q)"""
49
+ kl = (logs_q - logs_p) - 0.5
50
+ kl += 0.5 * (torch.exp(2. * logs_p) + ((m_p - m_q)**2)) * torch.exp(-2. * logs_q)
51
+ return kl
52
+
53
+
54
+ def rand_gumbel(shape):
55
+ """Sample from the Gumbel distribution, protect from overflows."""
56
+ uniform_samples = torch.rand(shape) * 0.99998 + 0.00001
57
+ return -torch.log(-torch.log(uniform_samples))
58
+
59
+
60
+ def rand_gumbel_like(x):
61
+ g = rand_gumbel(x.size()).to(dtype=x.dtype, device=x.device)
62
+ return g
63
+
64
+
65
+ def slice_segments(x, ids_str, segment_size=4):
66
+ ret = torch.zeros_like(x[:, :, :segment_size])
67
+ for i in range(x.size(0)):
68
+ idx_str = ids_str[i]
69
+ idx_end = idx_str + segment_size
70
+ ret[i] = x[i, :, idx_str:idx_end]
71
+ return ret
72
+
73
+
74
+ def rand_slice_segments(x, x_lengths=None, segment_size=4):
75
+ b, d, t = x.size()
76
+ if x_lengths is None:
77
+ x_lengths = t
78
+ ids_str_max = x_lengths - segment_size + 1
79
+ ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
80
+ ret = slice_segments(x, ids_str, segment_size)
81
+ return ret, ids_str
82
+
83
+
84
+ def rand_spec_segments(x, x_lengths=None, segment_size=4):
85
+ b, d, t = x.size()
86
+ if x_lengths is None:
87
+ x_lengths = t
88
+ ids_str_max = x_lengths - segment_size
89
+ ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
90
+ ret = slice_segments(x, ids_str, segment_size)
91
+ return ret, ids_str
92
+
93
+
94
+ def get_timing_signal_1d(
95
+ length, channels, min_timescale=1.0, max_timescale=1.0e4):
96
+ position = torch.arange(length, dtype=torch.float)
97
+ num_timescales = channels // 2
98
+ log_timescale_increment = (
99
+ math.log(float(max_timescale) / float(min_timescale)) /
100
+ (num_timescales - 1))
101
+ inv_timescales = min_timescale * torch.exp(
102
+ torch.arange(num_timescales, dtype=torch.float) * -log_timescale_increment)
103
+ scaled_time = position.unsqueeze(0) * inv_timescales.unsqueeze(1)
104
+ signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], 0)
105
+ signal = F.pad(signal, [0, 0, 0, channels % 2])
106
+ signal = signal.view(1, channels, length)
107
+ return signal
108
+
109
+
110
+ def add_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4):
111
+ b, channels, length = x.size()
112
+ signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
113
+ return x + signal.to(dtype=x.dtype, device=x.device)
114
+
115
+
116
+ def cat_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4, axis=1):
117
+ b, channels, length = x.size()
118
+ signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
119
+ return torch.cat([x, signal.to(dtype=x.dtype, device=x.device)], axis)
120
+
121
+
122
+ def subsequent_mask(length):
123
+ mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0)
124
+ return mask
125
+
126
+
127
+ @torch.jit.script
128
+ def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
129
+ n_channels_int = n_channels[0]
130
+ in_act = input_a + input_b
131
+ t_act = torch.tanh(in_act[:, :n_channels_int, :])
132
+ s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
133
+ acts = t_act * s_act
134
+ return acts
135
+
136
+
137
+ def convert_pad_shape(pad_shape):
138
+ l = pad_shape[::-1]
139
+ pad_shape = [item for sublist in l for item in sublist]
140
+ return pad_shape
141
+
142
+
143
+ def shift_1d(x):
144
+ x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [1, 0]]))[:, :, :-1]
145
+ return x
146
+
147
+
148
+ def sequence_mask(length, max_length=None):
149
+ if max_length is None:
150
+ max_length = length.max()
151
+ x = torch.arange(max_length, dtype=length.dtype, device=length.device)
152
+ return x.unsqueeze(0) < length.unsqueeze(1)
153
+
154
+
155
+ def generate_path(duration, mask):
156
+ """
157
+ duration: [b, 1, t_x]
158
+ mask: [b, 1, t_y, t_x]
159
+ """
160
+ device = duration.device
161
+
162
+ b, _, t_y, t_x = mask.shape
163
+ cum_duration = torch.cumsum(duration, -1)
164
+
165
+ cum_duration_flat = cum_duration.view(b * t_x)
166
+ path = sequence_mask(cum_duration_flat, t_y).to(mask.dtype)
167
+ path = path.view(b, t_x, t_y)
168
+ path = path - F.pad(path, convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
169
+ path = path.unsqueeze(1).transpose(2,3) * mask
170
+ return path
171
+
172
+
173
+ def clip_grad_value_(parameters, clip_value, norm_type=2):
174
+ if isinstance(parameters, torch.Tensor):
175
+ parameters = [parameters]
176
+ parameters = list(filter(lambda p: p.grad is not None, parameters))
177
+ norm_type = float(norm_type)
178
+ if clip_value is not None:
179
+ clip_value = float(clip_value)
180
+
181
+ total_norm = 0
182
+ for p in parameters:
183
+ param_norm = p.grad.data.norm(norm_type)
184
+ total_norm += param_norm.item() ** norm_type
185
+ if clip_value is not None:
186
+ p.grad.data.clamp_(min=-clip_value, max=clip_value)
187
+ total_norm = total_norm ** (1. / norm_type)
188
+ return total_norm
modules/enhancer.py ADDED
@@ -0,0 +1,105 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import torch
3
+ import torch.nn.functional as F
4
+ from vdecoder.nsf_hifigan.nvSTFT import STFT
5
+ from vdecoder.nsf_hifigan.models import load_model
6
+ from torchaudio.transforms import Resample
7
+
8
+ class Enhancer:
9
+ def __init__(self, enhancer_type, enhancer_ckpt, device=None):
10
+ if device is None:
11
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
12
+ self.device = device
13
+
14
+ if enhancer_type == 'nsf-hifigan':
15
+ self.enhancer = NsfHifiGAN(enhancer_ckpt, device=self.device)
16
+ else:
17
+ raise ValueError(f" [x] Unknown enhancer: {enhancer_type}")
18
+
19
+ self.resample_kernel = {}
20
+ self.enhancer_sample_rate = self.enhancer.sample_rate()
21
+ self.enhancer_hop_size = self.enhancer.hop_size()
22
+
23
+ def enhance(self,
24
+ audio, # 1, T
25
+ sample_rate,
26
+ f0, # 1, n_frames, 1
27
+ hop_size,
28
+ adaptive_key = 0,
29
+ silence_front = 0
30
+ ):
31
+ # enhancer start time
32
+ start_frame = int(silence_front * sample_rate / hop_size)
33
+ real_silence_front = start_frame * hop_size / sample_rate
34
+ audio = audio[:, int(np.round(real_silence_front * sample_rate)) : ]
35
+ f0 = f0[: , start_frame :, :]
36
+
37
+ # adaptive parameters
38
+ adaptive_factor = 2 ** ( -adaptive_key / 12)
39
+ adaptive_sample_rate = 100 * int(np.round(self.enhancer_sample_rate / adaptive_factor / 100))
40
+ real_factor = self.enhancer_sample_rate / adaptive_sample_rate
41
+
42
+ # resample the ddsp output
43
+ if sample_rate == adaptive_sample_rate:
44
+ audio_res = audio
45
+ else:
46
+ key_str = str(sample_rate) + str(adaptive_sample_rate)
47
+ if key_str not in self.resample_kernel:
48
+ self.resample_kernel[key_str] = Resample(sample_rate, adaptive_sample_rate, lowpass_filter_width = 128).to(self.device)
49
+ audio_res = self.resample_kernel[key_str](audio)
50
+
51
+ n_frames = int(audio_res.size(-1) // self.enhancer_hop_size + 1)
52
+
53
+ # resample f0
54
+ f0_np = f0.squeeze(0).squeeze(-1).cpu().numpy()
55
+ f0_np *= real_factor
56
+ time_org = (hop_size / sample_rate) * np.arange(len(f0_np)) / real_factor
57
+ time_frame = (self.enhancer_hop_size / self.enhancer_sample_rate) * np.arange(n_frames)
58
+ f0_res = np.interp(time_frame, time_org, f0_np, left=f0_np[0], right=f0_np[-1])
59
+ f0_res = torch.from_numpy(f0_res).unsqueeze(0).float().to(self.device) # 1, n_frames
60
+
61
+ # enhance
62
+ enhanced_audio, enhancer_sample_rate = self.enhancer(audio_res, f0_res)
63
+
64
+ # resample the enhanced output
65
+ if adaptive_factor != 0:
66
+ key_str = str(adaptive_sample_rate) + str(enhancer_sample_rate)
67
+ if key_str not in self.resample_kernel:
68
+ self.resample_kernel[key_str] = Resample(adaptive_sample_rate, enhancer_sample_rate, lowpass_filter_width = 128).to(self.device)
69
+ enhanced_audio = self.resample_kernel[key_str](enhanced_audio)
70
+
71
+ # pad the silence frames
72
+ if start_frame > 0:
73
+ enhanced_audio = F.pad(enhanced_audio, (int(np.round(enhancer_sample_rate * real_silence_front)), 0))
74
+
75
+ return enhanced_audio, enhancer_sample_rate
76
+
77
+
78
+ class NsfHifiGAN(torch.nn.Module):
79
+ def __init__(self, model_path, device=None):
80
+ super().__init__()
81
+ if device is None:
82
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
83
+ self.device = device
84
+ print('| Load HifiGAN: ', model_path)
85
+ self.model, self.h = load_model(model_path, device=self.device)
86
+
87
+ def sample_rate(self):
88
+ return self.h.sampling_rate
89
+
90
+ def hop_size(self):
91
+ return self.h.hop_size
92
+
93
+ def forward(self, audio, f0):
94
+ stft = STFT(
95
+ self.h.sampling_rate,
96
+ self.h.num_mels,
97
+ self.h.n_fft,
98
+ self.h.win_size,
99
+ self.h.hop_size,
100
+ self.h.fmin,
101
+ self.h.fmax)
102
+ with torch.no_grad():
103
+ mel = stft.get_mel(audio)
104
+ enhanced_audio = self.model(mel, f0[:,:mel.size(-1)]).view(-1)
105
+ return enhanced_audio, self.h.sampling_rate
modules/losses.py ADDED
@@ -0,0 +1,61 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from torch.nn import functional as F
3
+
4
+ import modules.commons as commons
5
+
6
+
7
+ def feature_loss(fmap_r, fmap_g):
8
+ loss = 0
9
+ for dr, dg in zip(fmap_r, fmap_g):
10
+ for rl, gl in zip(dr, dg):
11
+ rl = rl.float().detach()
12
+ gl = gl.float()
13
+ loss += torch.mean(torch.abs(rl - gl))
14
+
15
+ return loss * 2
16
+
17
+
18
+ def discriminator_loss(disc_real_outputs, disc_generated_outputs):
19
+ loss = 0
20
+ r_losses = []
21
+ g_losses = []
22
+ for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
23
+ dr = dr.float()
24
+ dg = dg.float()
25
+ r_loss = torch.mean((1-dr)**2)
26
+ g_loss = torch.mean(dg**2)
27
+ loss += (r_loss + g_loss)
28
+ r_losses.append(r_loss.item())
29
+ g_losses.append(g_loss.item())
30
+
31
+ return loss, r_losses, g_losses
32
+
33
+
34
+ def generator_loss(disc_outputs):
35
+ loss = 0
36
+ gen_losses = []
37
+ for dg in disc_outputs:
38
+ dg = dg.float()
39
+ l = torch.mean((1-dg)**2)
40
+ gen_losses.append(l)
41
+ loss += l
42
+
43
+ return loss, gen_losses
44
+
45
+
46
+ def kl_loss(z_p, logs_q, m_p, logs_p, z_mask):
47
+ """
48
+ z_p, logs_q: [b, h, t_t]
49
+ m_p, logs_p: [b, h, t_t]
50
+ """
51
+ z_p = z_p.float()
52
+ logs_q = logs_q.float()
53
+ m_p = m_p.float()
54
+ logs_p = logs_p.float()
55
+ z_mask = z_mask.float()
56
+ #print(logs_p)
57
+ kl = logs_p - logs_q - 0.5
58
+ kl += 0.5 * ((z_p - m_p)**2) * torch.exp(-2. * logs_p)
59
+ kl = torch.sum(kl * z_mask)
60
+ l = kl / torch.sum(z_mask)
61
+ return l
modules/mel_processing.py ADDED
@@ -0,0 +1,112 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ import os
3
+ import random
4
+ import torch
5
+ from torch import nn
6
+ import torch.nn.functional as F
7
+ import torch.utils.data
8
+ import numpy as np
9
+ import librosa
10
+ import librosa.util as librosa_util
11
+ from librosa.util import normalize, pad_center, tiny
12
+ from scipy.signal import get_window
13
+ from scipy.io.wavfile import read
14
+ from librosa.filters import mel as librosa_mel_fn
15
+
16
+ MAX_WAV_VALUE = 32768.0
17
+
18
+
19
+ def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
20
+ """
21
+ PARAMS
22
+ ------
23
+ C: compression factor
24
+ """
25
+ return torch.log(torch.clamp(x, min=clip_val) * C)
26
+
27
+
28
+ def dynamic_range_decompression_torch(x, C=1):
29
+ """
30
+ PARAMS
31
+ ------
32
+ C: compression factor used to compress
33
+ """
34
+ return torch.exp(x) / C
35
+
36
+
37
+ def spectral_normalize_torch(magnitudes):
38
+ output = dynamic_range_compression_torch(magnitudes)
39
+ return output
40
+
41
+
42
+ def spectral_de_normalize_torch(magnitudes):
43
+ output = dynamic_range_decompression_torch(magnitudes)
44
+ return output
45
+
46
+
47
+ mel_basis = {}
48
+ hann_window = {}
49
+
50
+
51
+ def spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center=False):
52
+ if torch.min(y) < -1.:
53
+ print('min value is ', torch.min(y))
54
+ if torch.max(y) > 1.:
55
+ print('max value is ', torch.max(y))
56
+
57
+ global hann_window
58
+ dtype_device = str(y.dtype) + '_' + str(y.device)
59
+ wnsize_dtype_device = str(win_size) + '_' + dtype_device
60
+ if wnsize_dtype_device not in hann_window:
61
+ hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(dtype=y.dtype, device=y.device)
62
+
63
+ y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect')
64
+ y = y.squeeze(1)
65
+
66
+ spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[wnsize_dtype_device],
67
+ center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=False)
68
+
69
+ spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
70
+ return spec
71
+
72
+
73
+ def spec_to_mel_torch(spec, n_fft, num_mels, sampling_rate, fmin, fmax):
74
+ global mel_basis
75
+ dtype_device = str(spec.dtype) + '_' + str(spec.device)
76
+ fmax_dtype_device = str(fmax) + '_' + dtype_device
77
+ if fmax_dtype_device not in mel_basis:
78
+ mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax)
79
+ mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(dtype=spec.dtype, device=spec.device)
80
+ spec = torch.matmul(mel_basis[fmax_dtype_device], spec)
81
+ spec = spectral_normalize_torch(spec)
82
+ return spec
83
+
84
+
85
+ def mel_spectrogram_torch(y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False):
86
+ if torch.min(y) < -1.:
87
+ print('min value is ', torch.min(y))
88
+ if torch.max(y) > 1.:
89
+ print('max value is ', torch.max(y))
90
+
91
+ global mel_basis, hann_window
92
+ dtype_device = str(y.dtype) + '_' + str(y.device)
93
+ fmax_dtype_device = str(fmax) + '_' + dtype_device
94
+ wnsize_dtype_device = str(win_size) + '_' + dtype_device
95
+ if fmax_dtype_device not in mel_basis:
96
+ mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax)
97
+ mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(dtype=y.dtype, device=y.device)
98
+ if wnsize_dtype_device not in hann_window:
99
+ hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(dtype=y.dtype, device=y.device)
100
+
101
+ y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect')
102
+ y = y.squeeze(1)
103
+
104
+ spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[wnsize_dtype_device],
105
+ center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=False)
106
+
107
+ spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
108
+
109
+ spec = torch.matmul(mel_basis[fmax_dtype_device], spec)
110
+ spec = spectral_normalize_torch(spec)
111
+
112
+ return spec
modules/modules.py ADDED
@@ -0,0 +1,342 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import copy
2
+ import math
3
+ import numpy as np
4
+ import scipy
5
+ import torch
6
+ from torch import nn
7
+ from torch.nn import functional as F
8
+
9
+ from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
10
+ from torch.nn.utils import weight_norm, remove_weight_norm
11
+
12
+ import modules.commons as commons
13
+ from modules.commons import init_weights, get_padding
14
+
15
+
16
+ LRELU_SLOPE = 0.1
17
+
18
+
19
+ class LayerNorm(nn.Module):
20
+ def __init__(self, channels, eps=1e-5):
21
+ super().__init__()
22
+ self.channels = channels
23
+ self.eps = eps
24
+
25
+ self.gamma = nn.Parameter(torch.ones(channels))
26
+ self.beta = nn.Parameter(torch.zeros(channels))
27
+
28
+ def forward(self, x):
29
+ x = x.transpose(1, -1)
30
+ x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
31
+ return x.transpose(1, -1)
32
+
33
+
34
+ class ConvReluNorm(nn.Module):
35
+ def __init__(self, in_channels, hidden_channels, out_channels, kernel_size, n_layers, p_dropout):
36
+ super().__init__()
37
+ self.in_channels = in_channels
38
+ self.hidden_channels = hidden_channels
39
+ self.out_channels = out_channels
40
+ self.kernel_size = kernel_size
41
+ self.n_layers = n_layers
42
+ self.p_dropout = p_dropout
43
+ assert n_layers > 1, "Number of layers should be larger than 0."
44
+
45
+ self.conv_layers = nn.ModuleList()
46
+ self.norm_layers = nn.ModuleList()
47
+ self.conv_layers.append(nn.Conv1d(in_channels, hidden_channels, kernel_size, padding=kernel_size//2))
48
+ self.norm_layers.append(LayerNorm(hidden_channels))
49
+ self.relu_drop = nn.Sequential(
50
+ nn.ReLU(),
51
+ nn.Dropout(p_dropout))
52
+ for _ in range(n_layers-1):
53
+ self.conv_layers.append(nn.Conv1d(hidden_channels, hidden_channels, kernel_size, padding=kernel_size//2))
54
+ self.norm_layers.append(LayerNorm(hidden_channels))
55
+ self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
56
+ self.proj.weight.data.zero_()
57
+ self.proj.bias.data.zero_()
58
+
59
+ def forward(self, x, x_mask):
60
+ x_org = x
61
+ for i in range(self.n_layers):
62
+ x = self.conv_layers[i](x * x_mask)
63
+ x = self.norm_layers[i](x)
64
+ x = self.relu_drop(x)
65
+ x = x_org + self.proj(x)
66
+ return x * x_mask
67
+
68
+
69
+ class DDSConv(nn.Module):
70
+ """
71
+ Dialted and Depth-Separable Convolution
72
+ """
73
+ def __init__(self, channels, kernel_size, n_layers, p_dropout=0.):
74
+ super().__init__()
75
+ self.channels = channels
76
+ self.kernel_size = kernel_size
77
+ self.n_layers = n_layers
78
+ self.p_dropout = p_dropout
79
+
80
+ self.drop = nn.Dropout(p_dropout)
81
+ self.convs_sep = nn.ModuleList()
82
+ self.convs_1x1 = nn.ModuleList()
83
+ self.norms_1 = nn.ModuleList()
84
+ self.norms_2 = nn.ModuleList()
85
+ for i in range(n_layers):
86
+ dilation = kernel_size ** i
87
+ padding = (kernel_size * dilation - dilation) // 2
88
+ self.convs_sep.append(nn.Conv1d(channels, channels, kernel_size,
89
+ groups=channels, dilation=dilation, padding=padding
90
+ ))
91
+ self.convs_1x1.append(nn.Conv1d(channels, channels, 1))
92
+ self.norms_1.append(LayerNorm(channels))
93
+ self.norms_2.append(LayerNorm(channels))
94
+
95
+ def forward(self, x, x_mask, g=None):
96
+ if g is not None:
97
+ x = x + g
98
+ for i in range(self.n_layers):
99
+ y = self.convs_sep[i](x * x_mask)
100
+ y = self.norms_1[i](y)
101
+ y = F.gelu(y)
102
+ y = self.convs_1x1[i](y)
103
+ y = self.norms_2[i](y)
104
+ y = F.gelu(y)
105
+ y = self.drop(y)
106
+ x = x + y
107
+ return x * x_mask
108
+
109
+
110
+ class WN(torch.nn.Module):
111
+ def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0):
112
+ super(WN, self).__init__()
113
+ assert(kernel_size % 2 == 1)
114
+ self.hidden_channels =hidden_channels
115
+ self.kernel_size = kernel_size,
116
+ self.dilation_rate = dilation_rate
117
+ self.n_layers = n_layers
118
+ self.gin_channels = gin_channels
119
+ self.p_dropout = p_dropout
120
+
121
+ self.in_layers = torch.nn.ModuleList()
122
+ self.res_skip_layers = torch.nn.ModuleList()
123
+ self.drop = nn.Dropout(p_dropout)
124
+
125
+ if gin_channels != 0:
126
+ cond_layer = torch.nn.Conv1d(gin_channels, 2*hidden_channels*n_layers, 1)
127
+ self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')
128
+
129
+ for i in range(n_layers):
130
+ dilation = dilation_rate ** i
131
+ padding = int((kernel_size * dilation - dilation) / 2)
132
+ in_layer = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, kernel_size,
133
+ dilation=dilation, padding=padding)
134
+ in_layer = torch.nn.utils.weight_norm(in_layer, name='weight')
135
+ self.in_layers.append(in_layer)
136
+
137
+ # last one is not necessary
138
+ if i < n_layers - 1:
139
+ res_skip_channels = 2 * hidden_channels
140
+ else:
141
+ res_skip_channels = hidden_channels
142
+
143
+ res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
144
+ res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight')
145
+ self.res_skip_layers.append(res_skip_layer)
146
+
147
+ def forward(self, x, x_mask, g=None, **kwargs):
148
+ output = torch.zeros_like(x)
149
+ n_channels_tensor = torch.IntTensor([self.hidden_channels])
150
+
151
+ if g is not None:
152
+ g = self.cond_layer(g)
153
+
154
+ for i in range(self.n_layers):
155
+ x_in = self.in_layers[i](x)
156
+ if g is not None:
157
+ cond_offset = i * 2 * self.hidden_channels
158
+ g_l = g[:,cond_offset:cond_offset+2*self.hidden_channels,:]
159
+ else:
160
+ g_l = torch.zeros_like(x_in)
161
+
162
+ acts = commons.fused_add_tanh_sigmoid_multiply(
163
+ x_in,
164
+ g_l,
165
+ n_channels_tensor)
166
+ acts = self.drop(acts)
167
+
168
+ res_skip_acts = self.res_skip_layers[i](acts)
169
+ if i < self.n_layers - 1:
170
+ res_acts = res_skip_acts[:,:self.hidden_channels,:]
171
+ x = (x + res_acts) * x_mask
172
+ output = output + res_skip_acts[:,self.hidden_channels:,:]
173
+ else:
174
+ output = output + res_skip_acts
175
+ return output * x_mask
176
+
177
+ def remove_weight_norm(self):
178
+ if self.gin_channels != 0:
179
+ torch.nn.utils.remove_weight_norm(self.cond_layer)
180
+ for l in self.in_layers:
181
+ torch.nn.utils.remove_weight_norm(l)
182
+ for l in self.res_skip_layers:
183
+ torch.nn.utils.remove_weight_norm(l)
184
+
185
+
186
+ class ResBlock1(torch.nn.Module):
187
+ def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
188
+ super(ResBlock1, self).__init__()
189
+ self.convs1 = nn.ModuleList([
190
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
191
+ padding=get_padding(kernel_size, dilation[0]))),
192
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
193
+ padding=get_padding(kernel_size, dilation[1]))),
194
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
195
+ padding=get_padding(kernel_size, dilation[2])))
196
+ ])
197
+ self.convs1.apply(init_weights)
198
+
199
+ self.convs2 = nn.ModuleList([
200
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
201
+ padding=get_padding(kernel_size, 1))),
202
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
203
+ padding=get_padding(kernel_size, 1))),
204
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
205
+ padding=get_padding(kernel_size, 1)))
206
+ ])
207
+ self.convs2.apply(init_weights)
208
+
209
+ def forward(self, x, x_mask=None):
210
+ for c1, c2 in zip(self.convs1, self.convs2):
211
+ xt = F.leaky_relu(x, LRELU_SLOPE)
212
+ if x_mask is not None:
213
+ xt = xt * x_mask
214
+ xt = c1(xt)
215
+ xt = F.leaky_relu(xt, LRELU_SLOPE)
216
+ if x_mask is not None:
217
+ xt = xt * x_mask
218
+ xt = c2(xt)
219
+ x = xt + x
220
+ if x_mask is not None:
221
+ x = x * x_mask
222
+ return x
223
+
224
+ def remove_weight_norm(self):
225
+ for l in self.convs1:
226
+ remove_weight_norm(l)
227
+ for l in self.convs2:
228
+ remove_weight_norm(l)
229
+
230
+
231
+ class ResBlock2(torch.nn.Module):
232
+ def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
233
+ super(ResBlock2, self).__init__()
234
+ self.convs = nn.ModuleList([
235
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
236
+ padding=get_padding(kernel_size, dilation[0]))),
237
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
238
+ padding=get_padding(kernel_size, dilation[1])))
239
+ ])
240
+ self.convs.apply(init_weights)
241
+
242
+ def forward(self, x, x_mask=None):
243
+ for c in self.convs:
244
+ xt = F.leaky_relu(x, LRELU_SLOPE)
245
+ if x_mask is not None:
246
+ xt = xt * x_mask
247
+ xt = c(xt)
248
+ x = xt + x
249
+ if x_mask is not None:
250
+ x = x * x_mask
251
+ return x
252
+
253
+ def remove_weight_norm(self):
254
+ for l in self.convs:
255
+ remove_weight_norm(l)
256
+
257
+
258
+ class Log(nn.Module):
259
+ def forward(self, x, x_mask, reverse=False, **kwargs):
260
+ if not reverse:
261
+ y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask
262
+ logdet = torch.sum(-y, [1, 2])
263
+ return y, logdet
264
+ else:
265
+ x = torch.exp(x) * x_mask
266
+ return x
267
+
268
+
269
+ class Flip(nn.Module):
270
+ def forward(self, x, *args, reverse=False, **kwargs):
271
+ x = torch.flip(x, [1])
272
+ if not reverse:
273
+ logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)
274
+ return x, logdet
275
+ else:
276
+ return x
277
+
278
+
279
+ class ElementwiseAffine(nn.Module):
280
+ def __init__(self, channels):
281
+ super().__init__()
282
+ self.channels = channels
283
+ self.m = nn.Parameter(torch.zeros(channels,1))
284
+ self.logs = nn.Parameter(torch.zeros(channels,1))
285
+
286
+ def forward(self, x, x_mask, reverse=False, **kwargs):
287
+ if not reverse:
288
+ y = self.m + torch.exp(self.logs) * x
289
+ y = y * x_mask
290
+ logdet = torch.sum(self.logs * x_mask, [1,2])
291
+ return y, logdet
292
+ else:
293
+ x = (x - self.m) * torch.exp(-self.logs) * x_mask
294
+ return x
295
+
296
+
297
+ class ResidualCouplingLayer(nn.Module):
298
+ def __init__(self,
299
+ channels,
300
+ hidden_channels,
301
+ kernel_size,
302
+ dilation_rate,
303
+ n_layers,
304
+ p_dropout=0,
305
+ gin_channels=0,
306
+ mean_only=False):
307
+ assert channels % 2 == 0, "channels should be divisible by 2"
308
+ super().__init__()
309
+ self.channels = channels
310
+ self.hidden_channels = hidden_channels
311
+ self.kernel_size = kernel_size
312
+ self.dilation_rate = dilation_rate
313
+ self.n_layers = n_layers
314
+ self.half_channels = channels // 2
315
+ self.mean_only = mean_only
316
+
317
+ self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
318
+ self.enc = WN(hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout=p_dropout, gin_channels=gin_channels)
319
+ self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
320
+ self.post.weight.data.zero_()
321
+ self.post.bias.data.zero_()
322
+
323
+ def forward(self, x, x_mask, g=None, reverse=False):
324
+ x0, x1 = torch.split(x, [self.half_channels]*2, 1)
325
+ h = self.pre(x0) * x_mask
326
+ h = self.enc(h, x_mask, g=g)
327
+ stats = self.post(h) * x_mask
328
+ if not self.mean_only:
329
+ m, logs = torch.split(stats, [self.half_channels]*2, 1)
330
+ else:
331
+ m = stats
332
+ logs = torch.zeros_like(m)
333
+
334
+ if not reverse:
335
+ x1 = m + x1 * torch.exp(logs) * x_mask
336
+ x = torch.cat([x0, x1], 1)
337
+ logdet = torch.sum(logs, [1,2])
338
+ return x, logdet
339
+ else:
340
+ x1 = (x1 - m) * torch.exp(-logs) * x_mask
341
+ x = torch.cat([x0, x1], 1)
342
+ return x
requirements.txt ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ffmpeg-python
2
+ Flask
3
+ Flask_Cors
4
+ gradio>=3.7.0
5
+ numpy==1.23.0
6
+ pyworld==0.2.5
7
+ scipy==1.10.0
8
+ SoundFile==0.12.1
9
+ torch==1.13.1
10
+ torchaudio==0.13.1
11
+ torchcrepe
12
+ tqdm
13
+ scikit-maad
14
+ praat-parselmouth
15
+ onnx
16
+ onnxsim
17
+ onnxoptimizer
18
+ fairseq==0.12.2
19
+ librosa==0.9.1
20
+ tensorboard
21
+ tensorboardX
22
+ transformers
23
+ edge_tts
24
+ pyyaml
25
+ pynvml
26
+ ffmpeg
utils.py ADDED
@@ -0,0 +1,446 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import glob
3
+ import re
4
+ import sys
5
+ import argparse
6
+ import logging
7
+ import json
8
+ import subprocess
9
+ import warnings
10
+ import random
11
+ import functools
12
+
13
+ import librosa
14
+ import numpy as np
15
+ from scipy.io.wavfile import read
16
+ import torch
17
+ from torch.nn import functional as F
18
+ from modules.commons import sequence_mask
19
+
20
+ MATPLOTLIB_FLAG = False
21
+
22
+ logging.basicConfig(stream=sys.stdout, level=logging.WARN)
23
+ logger = logging
24
+
25
+ f0_bin = 256
26
+ f0_max = 1100.0
27
+ f0_min = 50.0
28
+ f0_mel_min = 1127 * np.log(1 + f0_min / 700)
29
+ f0_mel_max = 1127 * np.log(1 + f0_max / 700)
30
+
31
+ def normalize_f0(f0, x_mask, uv, random_scale=True):
32
+ # calculate means based on x_mask
33
+ uv_sum = torch.sum(uv, dim=1, keepdim=True)
34
+ uv_sum[uv_sum == 0] = 9999
35
+ means = torch.sum(f0[:, 0, :] * uv, dim=1, keepdim=True) / uv_sum
36
+
37
+ if random_scale:
38
+ factor = torch.Tensor(f0.shape[0], 1).uniform_(0.8, 1.2).to(f0.device)
39
+ else:
40
+ factor = torch.ones(f0.shape[0], 1).to(f0.device)
41
+ # normalize f0 based on means and factor
42
+ f0_norm = (f0 - means.unsqueeze(-1)) * factor.unsqueeze(-1)
43
+ if torch.isnan(f0_norm).any():
44
+ exit(0)
45
+ return f0_norm * x_mask
46
+
47
+ def plot_data_to_numpy(x, y):
48
+ global MATPLOTLIB_FLAG
49
+ if not MATPLOTLIB_FLAG:
50
+ import matplotlib
51
+ matplotlib.use("Agg")
52
+ MATPLOTLIB_FLAG = True
53
+ mpl_logger = logging.getLogger('matplotlib')
54
+ mpl_logger.setLevel(logging.WARNING)
55
+ import matplotlib.pylab as plt
56
+ import numpy as np
57
+
58
+ fig, ax = plt.subplots(figsize=(10, 2))
59
+ plt.plot(x)
60
+ plt.plot(y)
61
+ plt.tight_layout()
62
+
63
+ fig.canvas.draw()
64
+ data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
65
+ data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
66
+ plt.close()
67
+ return data
68
+
69
+
70
+ def f0_to_coarse(f0):
71
+ is_torch = isinstance(f0, torch.Tensor)
72
+ f0_mel = 1127 * (1 + f0 / 700).log() if is_torch else 1127 * np.log(1 + f0 / 700)
73
+ f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - f0_mel_min) * (f0_bin - 2) / (f0_mel_max - f0_mel_min) + 1
74
+
75
+ f0_mel[f0_mel <= 1] = 1
76
+ f0_mel[f0_mel > f0_bin - 1] = f0_bin - 1
77
+ f0_coarse = (f0_mel + 0.5).int() if is_torch else np.rint(f0_mel).astype(np.int)
78
+ assert f0_coarse.max() <= 255 and f0_coarse.min() >= 1, (f0_coarse.max(), f0_coarse.min())
79
+ return f0_coarse
80
+
81
+ def get_content(cmodel, y):
82
+ with torch.no_grad():
83
+ c = cmodel.extract_features(y.squeeze(1))[0]
84
+ c = c.transpose(1, 2)
85
+ return c
86
+
87
+ def get_f0_predictor(f0_predictor,hop_length,sampling_rate,**kargs):
88
+ if f0_predictor == "pm":
89
+ from modules.F0Predictor.PMF0Predictor import PMF0Predictor
90
+ f0_predictor_object = PMF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate)
91
+ elif f0_predictor == "crepe":
92
+ from modules.F0Predictor.CrepeF0Predictor import CrepeF0Predictor
93
+ f0_predictor_object = CrepeF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate,device=kargs["device"],threshold=kargs["threshold"])
94
+ elif f0_predictor == "harvest":
95
+ from modules.F0Predictor.HarvestF0Predictor import HarvestF0Predictor
96
+ f0_predictor_object = HarvestF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate)
97
+ elif f0_predictor == "dio":
98
+ from modules.F0Predictor.DioF0Predictor import DioF0Predictor
99
+ f0_predictor_object = DioF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate)
100
+ else:
101
+ raise Exception("Unknown f0 predictor")
102
+ return f0_predictor_object
103
+
104
+ def get_speech_encoder(speech_encoder,device=None,**kargs):
105
+ if speech_encoder == "vec768l12":
106
+ from vencoder.ContentVec768L12 import ContentVec768L12
107
+ speech_encoder_object = ContentVec768L12(device = device)
108
+ elif speech_encoder == "vec256l9":
109
+ from vencoder.ContentVec256L9 import ContentVec256L9
110
+ speech_encoder_object = ContentVec256L9(device = device)
111
+ elif speech_encoder == "vec256l9-onnx":
112
+ from vencoder.ContentVec256L9_Onnx import ContentVec256L9_Onnx
113
+ speech_encoder_object = ContentVec256L9(device = device)
114
+ elif speech_encoder == "vec256l12-onnx":
115
+ from vencoder.ContentVec256L12_Onnx import ContentVec256L12_Onnx
116
+ speech_encoder_object = ContentVec256L9(device = device)
117
+ elif speech_encoder == "vec768l9-onnx":
118
+ from vencoder.ContentVec768L9_Onnx import ContentVec768L9_Onnx
119
+ speech_encoder_object = ContentVec256L9(device = device)
120
+ elif speech_encoder == "vec768l12-onnx":
121
+ from vencoder.ContentVec768L12_Onnx import ContentVec768L12_Onnx
122
+ speech_encoder_object = ContentVec256L9(device = device)
123
+ elif speech_encoder == "hubertsoft-onnx":
124
+ from vencoder.HubertSoft_Onnx import HubertSoft_Onnx
125
+ speech_encoder_object = HubertSoft(device = device)
126
+ elif speech_encoder == "hubertsoft":
127
+ from vencoder.HubertSoft import HubertSoft
128
+ speech_encoder_object = HubertSoft(device = device)
129
+ elif speech_encoder == "whisper-ppg":
130
+ from vencoder.WhisperPPG import WhisperPPG
131
+ speech_encoder_object = WhisperPPG(device = device)
132
+ else:
133
+ raise Exception("Unknown speech encoder")
134
+ return speech_encoder_object
135
+
136
+ def load_checkpoint(checkpoint_path, model, optimizer=None, skip_optimizer=False):
137
+ assert os.path.isfile(checkpoint_path)
138
+ checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')
139
+ iteration = checkpoint_dict['iteration']
140
+ learning_rate = checkpoint_dict['learning_rate']
141
+ if optimizer is not None and not skip_optimizer and checkpoint_dict['optimizer'] is not None:
142
+ optimizer.load_state_dict(checkpoint_dict['optimizer'])
143
+ saved_state_dict = checkpoint_dict['model']
144
+ if hasattr(model, 'module'):
145
+ state_dict = model.module.state_dict()
146
+ else:
147
+ state_dict = model.state_dict()
148
+ new_state_dict = {}
149
+ for k, v in state_dict.items():
150
+ try:
151
+ # assert "dec" in k or "disc" in k
152
+ # print("load", k)
153
+ new_state_dict[k] = saved_state_dict[k]
154
+ assert saved_state_dict[k].shape == v.shape, (saved_state_dict[k].shape, v.shape)
155
+ except:
156
+ print("error, %s is not in the checkpoint" % k)
157
+ logger.info("%s is not in the checkpoint" % k)
158
+ new_state_dict[k] = v
159
+ if hasattr(model, 'module'):
160
+ model.module.load_state_dict(new_state_dict)
161
+ else:
162
+ model.load_state_dict(new_state_dict)
163
+ print("load ")
164
+ logger.info("Loaded checkpoint '{}' (iteration {})".format(
165
+ checkpoint_path, iteration))
166
+ return model, optimizer, learning_rate, iteration
167
+
168
+
169
+ def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path):
170
+ logger.info("Saving model and optimizer state at iteration {} to {}".format(
171
+ iteration, checkpoint_path))
172
+ if hasattr(model, 'module'):
173
+ state_dict = model.module.state_dict()
174
+ else:
175
+ state_dict = model.state_dict()
176
+ torch.save({'model': state_dict,
177
+ 'iteration': iteration,
178
+ 'optimizer': optimizer.state_dict(),
179
+ 'learning_rate': learning_rate}, checkpoint_path)
180
+
181
+ def clean_checkpoints(path_to_models='logs/44k/', n_ckpts_to_keep=2, sort_by_time=True):
182
+ """Freeing up space by deleting saved ckpts
183
+
184
+ Arguments:
185
+ path_to_models -- Path to the model directory
186
+ n_ckpts_to_keep -- Number of ckpts to keep, excluding G_0.pth and D_0.pth
187
+ sort_by_time -- True -> chronologically delete ckpts
188
+ False -> lexicographically delete ckpts
189
+ """
190
+ ckpts_files = [f for f in os.listdir(path_to_models) if os.path.isfile(os.path.join(path_to_models, f))]
191
+ name_key = (lambda _f: int(re.compile('._(\d+)\.pth').match(_f).group(1)))
192
+ time_key = (lambda _f: os.path.getmtime(os.path.join(path_to_models, _f)))
193
+ sort_key = time_key if sort_by_time else name_key
194
+ x_sorted = lambda _x: sorted([f for f in ckpts_files if f.startswith(_x) and not f.endswith('_0.pth')], key=sort_key)
195
+ to_del = [os.path.join(path_to_models, fn) for fn in
196
+ (x_sorted('G')[:-n_ckpts_to_keep] + x_sorted('D')[:-n_ckpts_to_keep])]
197
+ del_info = lambda fn: logger.info(f".. Free up space by deleting ckpt {fn}")
198
+ del_routine = lambda x: [os.remove(x), del_info(x)]
199
+ rs = [del_routine(fn) for fn in to_del]
200
+
201
+ def summarize(writer, global_step, scalars={}, histograms={}, images={}, audios={}, audio_sampling_rate=22050):
202
+ for k, v in scalars.items():
203
+ writer.add_scalar(k, v, global_step)
204
+ for k, v in histograms.items():
205
+ writer.add_histogram(k, v, global_step)
206
+ for k, v in images.items():
207
+ writer.add_image(k, v, global_step, dataformats='HWC')
208
+ for k, v in audios.items():
209
+ writer.add_audio(k, v, global_step, audio_sampling_rate)
210
+
211
+
212
+ def latest_checkpoint_path(dir_path, regex="G_*.pth"):
213
+ f_list = glob.glob(os.path.join(dir_path, regex))
214
+ f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f))))
215
+ x = f_list[-1]
216
+ print(x)
217
+ return x
218
+
219
+
220
+ def plot_spectrogram_to_numpy(spectrogram):
221
+ global MATPLOTLIB_FLAG
222
+ if not MATPLOTLIB_FLAG:
223
+ import matplotlib
224
+ matplotlib.use("Agg")
225
+ MATPLOTLIB_FLAG = True
226
+ mpl_logger = logging.getLogger('matplotlib')
227
+ mpl_logger.setLevel(logging.WARNING)
228
+ import matplotlib.pylab as plt
229
+ import numpy as np
230
+
231
+ fig, ax = plt.subplots(figsize=(10,2))
232
+ im = ax.imshow(spectrogram, aspect="auto", origin="lower",
233
+ interpolation='none')
234
+ plt.colorbar(im, ax=ax)
235
+ plt.xlabel("Frames")
236
+ plt.ylabel("Channels")
237
+ plt.tight_layout()
238
+
239
+ fig.canvas.draw()
240
+ data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
241
+ data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
242
+ plt.close()
243
+ return data
244
+
245
+
246
+ def plot_alignment_to_numpy(alignment, info=None):
247
+ global MATPLOTLIB_FLAG
248
+ if not MATPLOTLIB_FLAG:
249
+ import matplotlib
250
+ matplotlib.use("Agg")
251
+ MATPLOTLIB_FLAG = True
252
+ mpl_logger = logging.getLogger('matplotlib')
253
+ mpl_logger.setLevel(logging.WARNING)
254
+ import matplotlib.pylab as plt
255
+ import numpy as np
256
+
257
+ fig, ax = plt.subplots(figsize=(6, 4))
258
+ im = ax.imshow(alignment.transpose(), aspect='auto', origin='lower',
259
+ interpolation='none')
260
+ fig.colorbar(im, ax=ax)
261
+ xlabel = 'Decoder timestep'
262
+ if info is not None:
263
+ xlabel += '\n\n' + info
264
+ plt.xlabel(xlabel)
265
+ plt.ylabel('Encoder timestep')
266
+ plt.tight_layout()
267
+
268
+ fig.canvas.draw()
269
+ data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
270
+ data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
271
+ plt.close()
272
+ return data
273
+
274
+
275
+ def load_wav_to_torch(full_path):
276
+ sampling_rate, data = read(full_path)
277
+ return torch.FloatTensor(data.astype(np.float32)), sampling_rate
278
+
279
+
280
+ def load_filepaths_and_text(filename, split="|"):
281
+ with open(filename, encoding='utf-8') as f:
282
+ filepaths_and_text = [line.strip().split(split) for line in f]
283
+ return filepaths_and_text
284
+
285
+
286
+ def get_hparams(init=True):
287
+ parser = argparse.ArgumentParser()
288
+ parser.add_argument('-c', '--config', type=str, default="./configs/config.json",
289
+ help='JSON file for configuration')
290
+ parser.add_argument('-m', '--model', type=str, required=True,
291
+ help='Model name')
292
+
293
+ args = parser.parse_args()
294
+ model_dir = os.path.join("./logs", args.model)
295
+
296
+ if not os.path.exists(model_dir):
297
+ os.makedirs(model_dir)
298
+
299
+ config_path = args.config
300
+ config_save_path = os.path.join(model_dir, "config.json")
301
+ if init:
302
+ with open(config_path, "r") as f:
303
+ data = f.read()
304
+ with open(config_save_path, "w") as f:
305
+ f.write(data)
306
+ else:
307
+ with open(config_save_path, "r") as f:
308
+ data = f.read()
309
+ config = json.loads(data)
310
+
311
+ hparams = HParams(**config)
312
+ hparams.model_dir = model_dir
313
+ return hparams
314
+
315
+
316
+ def get_hparams_from_dir(model_dir):
317
+ config_save_path = os.path.join(model_dir, "config.json")
318
+ with open(config_save_path, "r") as f:
319
+ data = f.read()
320
+ config = json.loads(data)
321
+
322
+ hparams =HParams(**config)
323
+ hparams.model_dir = model_dir
324
+ return hparams
325
+
326
+
327
+ def get_hparams_from_file(config_path):
328
+ with open(config_path, "r") as f:
329
+ data = f.read()
330
+ config = json.loads(data)
331
+ hparams =HParams(**config)
332
+ return hparams
333
+
334
+
335
+ def check_git_hash(model_dir):
336
+ source_dir = os.path.dirname(os.path.realpath(__file__))
337
+ if not os.path.exists(os.path.join(source_dir, ".git")):
338
+ logger.warn("{} is not a git repository, therefore hash value comparison will be ignored.".format(
339
+ source_dir
340
+ ))
341
+ return
342
+
343
+ cur_hash = subprocess.getoutput("git rev-parse HEAD")
344
+
345
+ path = os.path.join(model_dir, "githash")
346
+ if os.path.exists(path):
347
+ saved_hash = open(path).read()
348
+ if saved_hash != cur_hash:
349
+ logger.warn("git hash values are different. {}(saved) != {}(current)".format(
350
+ saved_hash[:8], cur_hash[:8]))
351
+ else:
352
+ open(path, "w").write(cur_hash)
353
+
354
+
355
+ def get_logger(model_dir, filename="train.log"):
356
+ global logger
357
+ logger = logging.getLogger(os.path.basename(model_dir))
358
+ logger.setLevel(logging.DEBUG)
359
+
360
+ formatter = logging.Formatter("%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s")
361
+ if not os.path.exists(model_dir):
362
+ os.makedirs(model_dir)
363
+ h = logging.FileHandler(os.path.join(model_dir, filename))
364
+ h.setLevel(logging.DEBUG)
365
+ h.setFormatter(formatter)
366
+ logger.addHandler(h)
367
+ return logger
368
+
369
+
370
+ def repeat_expand_2d(content, target_len):
371
+ # content : [h, t]
372
+
373
+ src_len = content.shape[-1]
374
+ target = torch.zeros([content.shape[0], target_len], dtype=torch.float).to(content.device)
375
+ temp = torch.arange(src_len+1) * target_len / src_len
376
+ current_pos = 0
377
+ for i in range(target_len):
378
+ if i < temp[current_pos+1]:
379
+ target[:, i] = content[:, current_pos]
380
+ else:
381
+ current_pos += 1
382
+ target[:, i] = content[:, current_pos]
383
+
384
+ return target
385
+
386
+
387
+ def mix_model(model_paths,mix_rate,mode):
388
+ mix_rate = torch.FloatTensor(mix_rate)/100
389
+ model_tem = torch.load(model_paths[0])
390
+ models = [torch.load(path)["model"] for path in model_paths]
391
+ if mode == 0:
392
+ mix_rate = F.softmax(mix_rate,dim=0)
393
+ for k in model_tem["model"].keys():
394
+ model_tem["model"][k] = torch.zeros_like(model_tem["model"][k])
395
+ for i,model in enumerate(models):
396
+ model_tem["model"][k] += model[k]*mix_rate[i]
397
+ torch.save(model_tem,os.path.join(os.path.curdir,"output.pth"))
398
+ return os.path.join(os.path.curdir,"output.pth")
399
+
400
+ class HParams():
401
+ def __init__(self, **kwargs):
402
+ for k, v in kwargs.items():
403
+ if type(v) == dict:
404
+ v = HParams(**v)
405
+ self[k] = v
406
+
407
+ def keys(self):
408
+ return self.__dict__.keys()
409
+
410
+ def items(self):
411
+ return self.__dict__.items()
412
+
413
+ def values(self):
414
+ return self.__dict__.values()
415
+
416
+ def __len__(self):
417
+ return len(self.__dict__)
418
+
419
+ def __getitem__(self, key):
420
+ return getattr(self, key)
421
+
422
+ def __setitem__(self, key, value):
423
+ return setattr(self, key, value)
424
+
425
+ def __contains__(self, key):
426
+ return key in self.__dict__
427
+
428
+ def __repr__(self):
429
+ return self.__dict__.__repr__()
430
+
431
+ def get(self,index):
432
+ return self.__dict__.get(index)
433
+
434
+ class Volume_Extractor:
435
+ def __init__(self, hop_size = 512):
436
+ self.hop_size = hop_size
437
+
438
+ def extract(self, audio): # audio: 2d tensor array
439
+ if not isinstance(audio,torch.Tensor):
440
+ audio = torch.Tensor(audio)
441
+ n_frames = int(audio.size(-1) // self.hop_size)
442
+ audio2 = audio ** 2
443
+ audio2 = torch.nn.functional.pad(audio2, (int(self.hop_size // 2), int((self.hop_size + 1) // 2)), mode = 'reflect')
444
+ volume = torch.FloatTensor([torch.mean(audio2[:,int(n * self.hop_size) : int((n + 1) * self.hop_size)]) for n in range(n_frames)])
445
+ volume = torch.sqrt(volume)
446
+ return volume
vdecoder/__init__.py ADDED
File without changes
vdecoder/hifigan/env.py ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import shutil
3
+
4
+
5
+ class AttrDict(dict):
6
+ def __init__(self, *args, **kwargs):
7
+ super(AttrDict, self).__init__(*args, **kwargs)
8
+ self.__dict__ = self
9
+
10
+
11
+ def build_env(config, config_name, path):
12
+ t_path = os.path.join(path, config_name)
13
+ if config != t_path:
14
+ os.makedirs(path, exist_ok=True)
15
+ shutil.copyfile(config, os.path.join(path, config_name))
vdecoder/hifigan/models.py ADDED
@@ -0,0 +1,503 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import json
3
+ from .env import AttrDict
4
+ import numpy as np
5
+ import torch
6
+ import torch.nn.functional as F
7
+ import torch.nn as nn
8
+ from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
9
+ from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
10
+ from .utils import init_weights, get_padding
11
+
12
+ LRELU_SLOPE = 0.1
13
+
14
+
15
+ def load_model(model_path, device='cuda'):
16
+ config_file = os.path.join(os.path.split(model_path)[0], 'config.json')
17
+ with open(config_file) as f:
18
+ data = f.read()
19
+
20
+ global h
21
+ json_config = json.loads(data)
22
+ h = AttrDict(json_config)
23
+
24
+ generator = Generator(h).to(device)
25
+
26
+ cp_dict = torch.load(model_path)
27
+ generator.load_state_dict(cp_dict['generator'])
28
+ generator.eval()
29
+ generator.remove_weight_norm()
30
+ del cp_dict
31
+ return generator, h
32
+
33
+
34
+ class ResBlock1(torch.nn.Module):
35
+ def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5)):
36
+ super(ResBlock1, self).__init__()
37
+ self.h = h
38
+ self.convs1 = nn.ModuleList([
39
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
40
+ padding=get_padding(kernel_size, dilation[0]))),
41
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
42
+ padding=get_padding(kernel_size, dilation[1]))),
43
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
44
+ padding=get_padding(kernel_size, dilation[2])))
45
+ ])
46
+ self.convs1.apply(init_weights)
47
+
48
+ self.convs2 = nn.ModuleList([
49
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
50
+ padding=get_padding(kernel_size, 1))),
51
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
52
+ padding=get_padding(kernel_size, 1))),
53
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
54
+ padding=get_padding(kernel_size, 1)))
55
+ ])
56
+ self.convs2.apply(init_weights)
57
+
58
+ def forward(self, x):
59
+ for c1, c2 in zip(self.convs1, self.convs2):
60
+ xt = F.leaky_relu(x, LRELU_SLOPE)
61
+ xt = c1(xt)
62
+ xt = F.leaky_relu(xt, LRELU_SLOPE)
63
+ xt = c2(xt)
64
+ x = xt + x
65
+ return x
66
+
67
+ def remove_weight_norm(self):
68
+ for l in self.convs1:
69
+ remove_weight_norm(l)
70
+ for l in self.convs2:
71
+ remove_weight_norm(l)
72
+
73
+
74
+ class ResBlock2(torch.nn.Module):
75
+ def __init__(self, h, channels, kernel_size=3, dilation=(1, 3)):
76
+ super(ResBlock2, self).__init__()
77
+ self.h = h
78
+ self.convs = nn.ModuleList([
79
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
80
+ padding=get_padding(kernel_size, dilation[0]))),
81
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
82
+ padding=get_padding(kernel_size, dilation[1])))
83
+ ])
84
+ self.convs.apply(init_weights)
85
+
86
+ def forward(self, x):
87
+ for c in self.convs:
88
+ xt = F.leaky_relu(x, LRELU_SLOPE)
89
+ xt = c(xt)
90
+ x = xt + x
91
+ return x
92
+
93
+ def remove_weight_norm(self):
94
+ for l in self.convs:
95
+ remove_weight_norm(l)
96
+
97
+
98
+ def padDiff(x):
99
+ return F.pad(F.pad(x, (0,0,-1,1), 'constant', 0) - x, (0,0,0,-1), 'constant', 0)
100
+
101
+ class SineGen(torch.nn.Module):
102
+ """ Definition of sine generator
103
+ SineGen(samp_rate, harmonic_num = 0,
104
+ sine_amp = 0.1, noise_std = 0.003,
105
+ voiced_threshold = 0,
106
+ flag_for_pulse=False)
107
+ samp_rate: sampling rate in Hz
108
+ harmonic_num: number of harmonic overtones (default 0)
109
+ sine_amp: amplitude of sine-wavefrom (default 0.1)
110
+ noise_std: std of Gaussian noise (default 0.003)
111
+ voiced_thoreshold: F0 threshold for U/V classification (default 0)
112
+ flag_for_pulse: this SinGen is used inside PulseGen (default False)
113
+ Note: when flag_for_pulse is True, the first time step of a voiced
114
+ segment is always sin(np.pi) or cos(0)
115
+ """
116
+
117
+ def __init__(self, samp_rate, harmonic_num=0,
118
+ sine_amp=0.1, noise_std=0.003,
119
+ voiced_threshold=0,
120
+ flag_for_pulse=False):
121
+ super(SineGen, self).__init__()
122
+ self.sine_amp = sine_amp
123
+ self.noise_std = noise_std
124
+ self.harmonic_num = harmonic_num
125
+ self.dim = self.harmonic_num + 1
126
+ self.sampling_rate = samp_rate
127
+ self.voiced_threshold = voiced_threshold
128
+ self.flag_for_pulse = flag_for_pulse
129
+
130
+ def _f02uv(self, f0):
131
+ # generate uv signal
132
+ uv = (f0 > self.voiced_threshold).type(torch.float32)
133
+ return uv
134
+
135
+ def _f02sine(self, f0_values):
136
+ """ f0_values: (batchsize, length, dim)
137
+ where dim indicates fundamental tone and overtones
138
+ """
139
+ # convert to F0 in rad. The interger part n can be ignored
140
+ # because 2 * np.pi * n doesn't affect phase
141
+ rad_values = (f0_values / self.sampling_rate) % 1
142
+
143
+ # initial phase noise (no noise for fundamental component)
144
+ rand_ini = torch.rand(f0_values.shape[0], f0_values.shape[2], \
145
+ device=f0_values.device)
146
+ rand_ini[:, 0] = 0
147
+ rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini
148
+
149
+ # instantanouse phase sine[t] = sin(2*pi \sum_i=1 ^{t} rad)
150
+ if not self.flag_for_pulse:
151
+ # for normal case
152
+
153
+ # To prevent torch.cumsum numerical overflow,
154
+ # it is necessary to add -1 whenever \sum_k=1^n rad_value_k > 1.
155
+ # Buffer tmp_over_one_idx indicates the time step to add -1.
156
+ # This will not change F0 of sine because (x-1) * 2*pi = x * 2*pi
157
+ tmp_over_one = torch.cumsum(rad_values, 1) % 1
158
+ tmp_over_one_idx = (padDiff(tmp_over_one)) < 0
159
+ cumsum_shift = torch.zeros_like(rad_values)
160
+ cumsum_shift[:, 1:, :] = tmp_over_one_idx * -1.0
161
+
162
+ sines = torch.sin(torch.cumsum(rad_values + cumsum_shift, dim=1)
163
+ * 2 * np.pi)
164
+ else:
165
+ # If necessary, make sure that the first time step of every
166
+ # voiced segments is sin(pi) or cos(0)
167
+ # This is used for pulse-train generation
168
+
169
+ # identify the last time step in unvoiced segments
170
+ uv = self._f02uv(f0_values)
171
+ uv_1 = torch.roll(uv, shifts=-1, dims=1)
172
+ uv_1[:, -1, :] = 1
173
+ u_loc = (uv < 1) * (uv_1 > 0)
174
+
175
+ # get the instantanouse phase
176
+ tmp_cumsum = torch.cumsum(rad_values, dim=1)
177
+ # different batch needs to be processed differently
178
+ for idx in range(f0_values.shape[0]):
179
+ temp_sum = tmp_cumsum[idx, u_loc[idx, :, 0], :]
180
+ temp_sum[1:, :] = temp_sum[1:, :] - temp_sum[0:-1, :]
181
+ # stores the accumulation of i.phase within
182
+ # each voiced segments
183
+ tmp_cumsum[idx, :, :] = 0
184
+ tmp_cumsum[idx, u_loc[idx, :, 0], :] = temp_sum
185
+
186
+ # rad_values - tmp_cumsum: remove the accumulation of i.phase
187
+ # within the previous voiced segment.
188
+ i_phase = torch.cumsum(rad_values - tmp_cumsum, dim=1)
189
+
190
+ # get the sines
191
+ sines = torch.cos(i_phase * 2 * np.pi)
192
+ return sines
193
+
194
+ def forward(self, f0):
195
+ """ sine_tensor, uv = forward(f0)
196
+ input F0: tensor(batchsize=1, length, dim=1)
197
+ f0 for unvoiced steps should be 0
198
+ output sine_tensor: tensor(batchsize=1, length, dim)
199
+ output uv: tensor(batchsize=1, length, 1)
200
+ """
201
+ with torch.no_grad():
202
+ f0_buf = torch.zeros(f0.shape[0], f0.shape[1], self.dim,
203
+ device=f0.device)
204
+ # fundamental component
205
+ fn = torch.multiply(f0, torch.FloatTensor([[range(1, self.harmonic_num + 2)]]).to(f0.device))
206
+
207
+ # generate sine waveforms
208
+ sine_waves = self._f02sine(fn) * self.sine_amp
209
+
210
+ # generate uv signal
211
+ # uv = torch.ones(f0.shape)
212
+ # uv = uv * (f0 > self.voiced_threshold)
213
+ uv = self._f02uv(f0)
214
+
215
+ # noise: for unvoiced should be similar to sine_amp
216
+ # std = self.sine_amp/3 -> max value ~ self.sine_amp
217
+ # . for voiced regions is self.noise_std
218
+ noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
219
+ noise = noise_amp * torch.randn_like(sine_waves)
220
+
221
+ # first: set the unvoiced part to 0 by uv
222
+ # then: additive noise
223
+ sine_waves = sine_waves * uv + noise
224
+ return sine_waves, uv, noise
225
+
226
+
227
+ class SourceModuleHnNSF(torch.nn.Module):
228
+ """ SourceModule for hn-nsf
229
+ SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1,
230
+ add_noise_std=0.003, voiced_threshod=0)
231
+ sampling_rate: sampling_rate in Hz
232
+ harmonic_num: number of harmonic above F0 (default: 0)
233
+ sine_amp: amplitude of sine source signal (default: 0.1)
234
+ add_noise_std: std of additive Gaussian noise (default: 0.003)
235
+ note that amplitude of noise in unvoiced is decided
236
+ by sine_amp
237
+ voiced_threshold: threhold to set U/V given F0 (default: 0)
238
+ Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
239
+ F0_sampled (batchsize, length, 1)
240
+ Sine_source (batchsize, length, 1)
241
+ noise_source (batchsize, length 1)
242
+ uv (batchsize, length, 1)
243
+ """
244
+
245
+ def __init__(self, sampling_rate, harmonic_num=0, sine_amp=0.1,
246
+ add_noise_std=0.003, voiced_threshod=0):
247
+ super(SourceModuleHnNSF, self).__init__()
248
+
249
+ self.sine_amp = sine_amp
250
+ self.noise_std = add_noise_std
251
+
252
+ # to produce sine waveforms
253
+ self.l_sin_gen = SineGen(sampling_rate, harmonic_num,
254
+ sine_amp, add_noise_std, voiced_threshod)
255
+
256
+ # to merge source harmonics into a single excitation
257
+ self.l_linear = torch.nn.Linear(harmonic_num + 1, 1)
258
+ self.l_tanh = torch.nn.Tanh()
259
+
260
+ def forward(self, x):
261
+ """
262
+ Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
263
+ F0_sampled (batchsize, length, 1)
264
+ Sine_source (batchsize, length, 1)
265
+ noise_source (batchsize, length 1)
266
+ """
267
+ # source for harmonic branch
268
+ sine_wavs, uv, _ = self.l_sin_gen(x)
269
+ sine_merge = self.l_tanh(self.l_linear(sine_wavs))
270
+
271
+ # source for noise branch, in the same shape as uv
272
+ noise = torch.randn_like(uv) * self.sine_amp / 3
273
+ return sine_merge, noise, uv
274
+
275
+
276
+ class Generator(torch.nn.Module):
277
+ def __init__(self, h):
278
+ super(Generator, self).__init__()
279
+ self.h = h
280
+
281
+ self.num_kernels = len(h["resblock_kernel_sizes"])
282
+ self.num_upsamples = len(h["upsample_rates"])
283
+ self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(h["upsample_rates"]))
284
+ self.m_source = SourceModuleHnNSF(
285
+ sampling_rate=h["sampling_rate"],
286
+ harmonic_num=8)
287
+ self.noise_convs = nn.ModuleList()
288
+ self.conv_pre = weight_norm(Conv1d(h["inter_channels"], h["upsample_initial_channel"], 7, 1, padding=3))
289
+ resblock = ResBlock1 if h["resblock"] == '1' else ResBlock2
290
+ self.ups = nn.ModuleList()
291
+ for i, (u, k) in enumerate(zip(h["upsample_rates"], h["upsample_kernel_sizes"])):
292
+ c_cur = h["upsample_initial_channel"] // (2 ** (i + 1))
293
+ self.ups.append(weight_norm(
294
+ ConvTranspose1d(h["upsample_initial_channel"] // (2 ** i), h["upsample_initial_channel"] // (2 ** (i + 1)),
295
+ k, u, padding=(k - u) // 2)))
296
+ if i + 1 < len(h["upsample_rates"]): #
297
+ stride_f0 = np.prod(h["upsample_rates"][i + 1:])
298
+ self.noise_convs.append(Conv1d(
299
+ 1, c_cur, kernel_size=stride_f0 * 2, stride=stride_f0, padding=stride_f0 // 2))
300
+ else:
301
+ self.noise_convs.append(Conv1d(1, c_cur, kernel_size=1))
302
+ self.resblocks = nn.ModuleList()
303
+ for i in range(len(self.ups)):
304
+ ch = h["upsample_initial_channel"] // (2 ** (i + 1))
305
+ for j, (k, d) in enumerate(zip(h["resblock_kernel_sizes"], h["resblock_dilation_sizes"])):
306
+ self.resblocks.append(resblock(h, ch, k, d))
307
+
308
+ self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3))
309
+ self.ups.apply(init_weights)
310
+ self.conv_post.apply(init_weights)
311
+ self.cond = nn.Conv1d(h['gin_channels'], h['upsample_initial_channel'], 1)
312
+
313
+ def forward(self, x, f0, g=None):
314
+ # print(1,x.shape,f0.shape,f0[:, None].shape)
315
+ f0 = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t
316
+ # print(2,f0.shape)
317
+ har_source, noi_source, uv = self.m_source(f0)
318
+ har_source = har_source.transpose(1, 2)
319
+ x = self.conv_pre(x)
320
+ x = x + self.cond(g)
321
+ # print(124,x.shape,har_source.shape)
322
+ for i in range(self.num_upsamples):
323
+ x = F.leaky_relu(x, LRELU_SLOPE)
324
+ # print(3,x.shape)
325
+ x = self.ups[i](x)
326
+ x_source = self.noise_convs[i](har_source)
327
+ # print(4,x_source.shape,har_source.shape,x.shape)
328
+ x = x + x_source
329
+ xs = None
330
+ for j in range(self.num_kernels):
331
+ if xs is None:
332
+ xs = self.resblocks[i * self.num_kernels + j](x)
333
+ else:
334
+ xs += self.resblocks[i * self.num_kernels + j](x)
335
+ x = xs / self.num_kernels
336
+ x = F.leaky_relu(x)
337
+ x = self.conv_post(x)
338
+ x = torch.tanh(x)
339
+
340
+ return x
341
+
342
+ def remove_weight_norm(self):
343
+ print('Removing weight norm...')
344
+ for l in self.ups:
345
+ remove_weight_norm(l)
346
+ for l in self.resblocks:
347
+ l.remove_weight_norm()
348
+ remove_weight_norm(self.conv_pre)
349
+ remove_weight_norm(self.conv_post)
350
+
351
+
352
+ class DiscriminatorP(torch.nn.Module):
353
+ def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
354
+ super(DiscriminatorP, self).__init__()
355
+ self.period = period
356
+ norm_f = weight_norm if use_spectral_norm == False else spectral_norm
357
+ self.convs = nn.ModuleList([
358
+ norm_f(Conv2d(1, 32, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
359
+ norm_f(Conv2d(32, 128, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
360
+ norm_f(Conv2d(128, 512, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
361
+ norm_f(Conv2d(512, 1024, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
362
+ norm_f(Conv2d(1024, 1024, (kernel_size, 1), 1, padding=(2, 0))),
363
+ ])
364
+ self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
365
+
366
+ def forward(self, x):
367
+ fmap = []
368
+
369
+ # 1d to 2d
370
+ b, c, t = x.shape
371
+ if t % self.period != 0: # pad first
372
+ n_pad = self.period - (t % self.period)
373
+ x = F.pad(x, (0, n_pad), "reflect")
374
+ t = t + n_pad
375
+ x = x.view(b, c, t // self.period, self.period)
376
+
377
+ for l in self.convs:
378
+ x = l(x)
379
+ x = F.leaky_relu(x, LRELU_SLOPE)
380
+ fmap.append(x)
381
+ x = self.conv_post(x)
382
+ fmap.append(x)
383
+ x = torch.flatten(x, 1, -1)
384
+
385
+ return x, fmap
386
+
387
+
388
+ class MultiPeriodDiscriminator(torch.nn.Module):
389
+ def __init__(self, periods=None):
390
+ super(MultiPeriodDiscriminator, self).__init__()
391
+ self.periods = periods if periods is not None else [2, 3, 5, 7, 11]
392
+ self.discriminators = nn.ModuleList()
393
+ for period in self.periods:
394
+ self.discriminators.append(DiscriminatorP(period))
395
+
396
+ def forward(self, y, y_hat):
397
+ y_d_rs = []
398
+ y_d_gs = []
399
+ fmap_rs = []
400
+ fmap_gs = []
401
+ for i, d in enumerate(self.discriminators):
402
+ y_d_r, fmap_r = d(y)
403
+ y_d_g, fmap_g = d(y_hat)
404
+ y_d_rs.append(y_d_r)
405
+ fmap_rs.append(fmap_r)
406
+ y_d_gs.append(y_d_g)
407
+ fmap_gs.append(fmap_g)
408
+
409
+ return y_d_rs, y_d_gs, fmap_rs, fmap_gs
410
+
411
+
412
+ class DiscriminatorS(torch.nn.Module):
413
+ def __init__(self, use_spectral_norm=False):
414
+ super(DiscriminatorS, self).__init__()
415
+ norm_f = weight_norm if use_spectral_norm == False else spectral_norm
416
+ self.convs = nn.ModuleList([
417
+ norm_f(Conv1d(1, 128, 15, 1, padding=7)),
418
+ norm_f(Conv1d(128, 128, 41, 2, groups=4, padding=20)),
419
+ norm_f(Conv1d(128, 256, 41, 2, groups=16, padding=20)),
420
+ norm_f(Conv1d(256, 512, 41, 4, groups=16, padding=20)),
421
+ norm_f(Conv1d(512, 1024, 41, 4, groups=16, padding=20)),
422
+ norm_f(Conv1d(1024, 1024, 41, 1, groups=16, padding=20)),
423
+ norm_f(Conv1d(1024, 1024, 5, 1, padding=2)),
424
+ ])
425
+ self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))
426
+
427
+ def forward(self, x):
428
+ fmap = []
429
+ for l in self.convs:
430
+ x = l(x)
431
+ x = F.leaky_relu(x, LRELU_SLOPE)
432
+ fmap.append(x)
433
+ x = self.conv_post(x)
434
+ fmap.append(x)
435
+ x = torch.flatten(x, 1, -1)
436
+
437
+ return x, fmap
438
+
439
+
440
+ class MultiScaleDiscriminator(torch.nn.Module):
441
+ def __init__(self):
442
+ super(MultiScaleDiscriminator, self).__init__()
443
+ self.discriminators = nn.ModuleList([
444
+ DiscriminatorS(use_spectral_norm=True),
445
+ DiscriminatorS(),
446
+ DiscriminatorS(),
447
+ ])
448
+ self.meanpools = nn.ModuleList([
449
+ AvgPool1d(4, 2, padding=2),
450
+ AvgPool1d(4, 2, padding=2)
451
+ ])
452
+
453
+ def forward(self, y, y_hat):
454
+ y_d_rs = []
455
+ y_d_gs = []
456
+ fmap_rs = []
457
+ fmap_gs = []
458
+ for i, d in enumerate(self.discriminators):
459
+ if i != 0:
460
+ y = self.meanpools[i - 1](y)
461
+ y_hat = self.meanpools[i - 1](y_hat)
462
+ y_d_r, fmap_r = d(y)
463
+ y_d_g, fmap_g = d(y_hat)
464
+ y_d_rs.append(y_d_r)
465
+ fmap_rs.append(fmap_r)
466
+ y_d_gs.append(y_d_g)
467
+ fmap_gs.append(fmap_g)
468
+
469
+ return y_d_rs, y_d_gs, fmap_rs, fmap_gs
470
+
471
+
472
+ def feature_loss(fmap_r, fmap_g):
473
+ loss = 0
474
+ for dr, dg in zip(fmap_r, fmap_g):
475
+ for rl, gl in zip(dr, dg):
476
+ loss += torch.mean(torch.abs(rl - gl))
477
+
478
+ return loss * 2
479
+
480
+
481
+ def discriminator_loss(disc_real_outputs, disc_generated_outputs):
482
+ loss = 0
483
+ r_losses = []
484
+ g_losses = []
485
+ for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
486
+ r_loss = torch.mean((1 - dr) ** 2)
487
+ g_loss = torch.mean(dg ** 2)
488
+ loss += (r_loss + g_loss)
489
+ r_losses.append(r_loss.item())
490
+ g_losses.append(g_loss.item())
491
+
492
+ return loss, r_losses, g_losses
493
+
494
+
495
+ def generator_loss(disc_outputs):
496
+ loss = 0
497
+ gen_losses = []
498
+ for dg in disc_outputs:
499
+ l = torch.mean((1 - dg) ** 2)
500
+ gen_losses.append(l)
501
+ loss += l
502
+
503
+ return loss, gen_losses
vdecoder/hifigan/nvSTFT.py ADDED
@@ -0,0 +1,111 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ import os
3
+ os.environ["LRU_CACHE_CAPACITY"] = "3"
4
+ import random
5
+ import torch
6
+ import torch.utils.data
7
+ import numpy as np
8
+ import librosa
9
+ from librosa.util import normalize
10
+ from librosa.filters import mel as librosa_mel_fn
11
+ from scipy.io.wavfile import read
12
+ import soundfile as sf
13
+
14
+ def load_wav_to_torch(full_path, target_sr=None, return_empty_on_exception=False):
15
+ sampling_rate = None
16
+ try:
17
+ data, sampling_rate = sf.read(full_path, always_2d=True)# than soundfile.
18
+ except Exception as ex:
19
+ print(f"'{full_path}' failed to load.\nException:")
20
+ print(ex)
21
+ if return_empty_on_exception:
22
+ return [], sampling_rate or target_sr or 32000
23
+ else:
24
+ raise Exception(ex)
25
+
26
+ if len(data.shape) > 1:
27
+ data = data[:, 0]
28
+ assert len(data) > 2# check duration of audio file is > 2 samples (because otherwise the slice operation was on the wrong dimension)
29
+
30
+ if np.issubdtype(data.dtype, np.integer): # if audio data is type int
31
+ max_mag = -np.iinfo(data.dtype).min # maximum magnitude = min possible value of intXX
32
+ else: # if audio data is type fp32
33
+ max_mag = max(np.amax(data), -np.amin(data))
34
+ 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
35
+
36
+ data = torch.FloatTensor(data.astype(np.float32))/max_mag
37
+
38
+ 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
39
+ return [], sampling_rate or target_sr or 32000
40
+ if target_sr is not None and sampling_rate != target_sr:
41
+ data = torch.from_numpy(librosa.core.resample(data.numpy(), orig_sr=sampling_rate, target_sr=target_sr))
42
+ sampling_rate = target_sr
43
+
44
+ return data, sampling_rate
45
+
46
+ def dynamic_range_compression(x, C=1, clip_val=1e-5):
47
+ return np.log(np.clip(x, a_min=clip_val, a_max=None) * C)
48
+
49
+ def dynamic_range_decompression(x, C=1):
50
+ return np.exp(x) / C
51
+
52
+ def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
53
+ return torch.log(torch.clamp(x, min=clip_val) * C)
54
+
55
+ def dynamic_range_decompression_torch(x, C=1):
56
+ return torch.exp(x) / C
57
+
58
+ class STFT():
59
+ 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):
60
+ self.target_sr = sr
61
+
62
+ self.n_mels = n_mels
63
+ self.n_fft = n_fft
64
+ self.win_size = win_size
65
+ self.hop_length = hop_length
66
+ self.fmin = fmin
67
+ self.fmax = fmax
68
+ self.clip_val = clip_val
69
+ self.mel_basis = {}
70
+ self.hann_window = {}
71
+
72
+ def get_mel(self, y, center=False):
73
+ sampling_rate = self.target_sr
74
+ n_mels = self.n_mels
75
+ n_fft = self.n_fft
76
+ win_size = self.win_size
77
+ hop_length = self.hop_length
78
+ fmin = self.fmin
79
+ fmax = self.fmax
80
+ clip_val = self.clip_val
81
+
82
+ if torch.min(y) < -1.:
83
+ print('min value is ', torch.min(y))
84
+ if torch.max(y) > 1.:
85
+ print('max value is ', torch.max(y))
86
+
87
+ if fmax not in self.mel_basis:
88
+ mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax)
89
+ self.mel_basis[str(fmax)+'_'+str(y.device)] = torch.from_numpy(mel).float().to(y.device)
90
+ self.hann_window[str(y.device)] = torch.hann_window(self.win_size).to(y.device)
91
+
92
+ y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_length)/2), int((n_fft-hop_length)/2)), mode='reflect')
93
+ y = y.squeeze(1)
94
+
95
+ spec = torch.stft(y, n_fft, hop_length=hop_length, win_length=win_size, window=self.hann_window[str(y.device)],
96
+ center=center, pad_mode='reflect', normalized=False, onesided=True)
97
+ # print(111,spec)
98
+ spec = torch.sqrt(spec.pow(2).sum(-1)+(1e-9))
99
+ # print(222,spec)
100
+ spec = torch.matmul(self.mel_basis[str(fmax)+'_'+str(y.device)], spec)
101
+ # print(333,spec)
102
+ spec = dynamic_range_compression_torch(spec, clip_val=clip_val)
103
+ # print(444,spec)
104
+ return spec
105
+
106
+ def __call__(self, audiopath):
107
+ audio, sr = load_wav_to_torch(audiopath, target_sr=self.target_sr)
108
+ spect = self.get_mel(audio.unsqueeze(0)).squeeze(0)
109
+ return spect
110
+
111
+ stft = STFT()
vdecoder/hifigan/utils.py ADDED
@@ -0,0 +1,68 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import glob
2
+ import os
3
+ import matplotlib
4
+ import torch
5
+ from torch.nn.utils import weight_norm
6
+ # matplotlib.use("Agg")
7
+ import matplotlib.pylab as plt
8
+
9
+
10
+ def plot_spectrogram(spectrogram):
11
+ fig, ax = plt.subplots(figsize=(10, 2))
12
+ im = ax.imshow(spectrogram, aspect="auto", origin="lower",
13
+ interpolation='none')
14
+ plt.colorbar(im, ax=ax)
15
+
16
+ fig.canvas.draw()
17
+ plt.close()
18
+
19
+ return fig
20
+
21
+
22
+ def init_weights(m, mean=0.0, std=0.01):
23
+ classname = m.__class__.__name__
24
+ if classname.find("Conv") != -1:
25
+ m.weight.data.normal_(mean, std)
26
+
27
+
28
+ def apply_weight_norm(m):
29
+ classname = m.__class__.__name__
30
+ if classname.find("Conv") != -1:
31
+ weight_norm(m)
32
+
33
+
34
+ def get_padding(kernel_size, dilation=1):
35
+ return int((kernel_size*dilation - dilation)/2)
36
+
37
+
38
+ def load_checkpoint(filepath, device):
39
+ assert os.path.isfile(filepath)
40
+ print("Loading '{}'".format(filepath))
41
+ checkpoint_dict = torch.load(filepath, map_location=device)
42
+ print("Complete.")
43
+ return checkpoint_dict
44
+
45
+
46
+ def save_checkpoint(filepath, obj):
47
+ print("Saving checkpoint to {}".format(filepath))
48
+ torch.save(obj, filepath)
49
+ print("Complete.")
50
+
51
+
52
+ def del_old_checkpoints(cp_dir, prefix, n_models=2):
53
+ pattern = os.path.join(cp_dir, prefix + '????????')
54
+ cp_list = glob.glob(pattern) # get checkpoint paths
55
+ cp_list = sorted(cp_list)# sort by iter
56
+ if len(cp_list) > n_models: # if more than n_models models are found
57
+ for cp in cp_list[:-n_models]:# delete the oldest models other than lastest n_models
58
+ open(cp, 'w').close()# empty file contents
59
+ os.unlink(cp)# delete file (move to trash when using Colab)
60
+
61
+
62
+ def scan_checkpoint(cp_dir, prefix):
63
+ pattern = os.path.join(cp_dir, prefix + '????????')
64
+ cp_list = glob.glob(pattern)
65
+ if len(cp_list) == 0:
66
+ return None
67
+ return sorted(cp_list)[-1]
68
+
vdecoder/nsf_hifigan/env.py ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import shutil
3
+
4
+
5
+ class AttrDict(dict):
6
+ def __init__(self, *args, **kwargs):
7
+ super(AttrDict, self).__init__(*args, **kwargs)
8
+ self.__dict__ = self
9
+
10
+
11
+ def build_env(config, config_name, path):
12
+ t_path = os.path.join(path, config_name)
13
+ if config != t_path:
14
+ os.makedirs(path, exist_ok=True)
15
+ shutil.copyfile(config, os.path.join(path, config_name))
vdecoder/nsf_hifigan/models.py ADDED
@@ -0,0 +1,439 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import json
3
+ from .env import AttrDict
4
+ import numpy as np
5
+ import torch
6
+ import torch.nn.functional as F
7
+ import torch.nn as nn
8
+ from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
9
+ from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
10
+ from .utils import init_weights, get_padding
11
+
12
+ LRELU_SLOPE = 0.1
13
+
14
+
15
+ def load_model(model_path, device='cuda'):
16
+ h = load_config(model_path)
17
+
18
+ generator = Generator(h).to(device)
19
+
20
+ cp_dict = torch.load(model_path, map_location=device)
21
+ generator.load_state_dict(cp_dict['generator'])
22
+ generator.eval()
23
+ generator.remove_weight_norm()
24
+ del cp_dict
25
+ return generator, h
26
+
27
+ def load_config(model_path):
28
+ config_file = os.path.join(os.path.split(model_path)[0], 'config.json')
29
+ with open(config_file) as f:
30
+ data = f.read()
31
+
32
+ json_config = json.loads(data)
33
+ h = AttrDict(json_config)
34
+ return h
35
+
36
+
37
+ class ResBlock1(torch.nn.Module):
38
+ def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5)):
39
+ super(ResBlock1, self).__init__()
40
+ self.h = h
41
+ self.convs1 = nn.ModuleList([
42
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
43
+ padding=get_padding(kernel_size, dilation[0]))),
44
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
45
+ padding=get_padding(kernel_size, dilation[1]))),
46
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
47
+ padding=get_padding(kernel_size, dilation[2])))
48
+ ])
49
+ self.convs1.apply(init_weights)
50
+
51
+ self.convs2 = nn.ModuleList([
52
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
53
+ padding=get_padding(kernel_size, 1))),
54
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
55
+ padding=get_padding(kernel_size, 1))),
56
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
57
+ padding=get_padding(kernel_size, 1)))
58
+ ])
59
+ self.convs2.apply(init_weights)
60
+
61
+ def forward(self, x):
62
+ for c1, c2 in zip(self.convs1, self.convs2):
63
+ xt = F.leaky_relu(x, LRELU_SLOPE)
64
+ xt = c1(xt)
65
+ xt = F.leaky_relu(xt, LRELU_SLOPE)
66
+ xt = c2(xt)
67
+ x = xt + x
68
+ return x
69
+
70
+ def remove_weight_norm(self):
71
+ for l in self.convs1:
72
+ remove_weight_norm(l)
73
+ for l in self.convs2:
74
+ remove_weight_norm(l)
75
+
76
+
77
+ class ResBlock2(torch.nn.Module):
78
+ def __init__(self, h, channels, kernel_size=3, dilation=(1, 3)):
79
+ super(ResBlock2, self).__init__()
80
+ self.h = h
81
+ self.convs = nn.ModuleList([
82
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
83
+ padding=get_padding(kernel_size, dilation[0]))),
84
+ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
85
+ padding=get_padding(kernel_size, dilation[1])))
86
+ ])
87
+ self.convs.apply(init_weights)
88
+
89
+ def forward(self, x):
90
+ for c in self.convs:
91
+ xt = F.leaky_relu(x, LRELU_SLOPE)
92
+ xt = c(xt)
93
+ x = xt + x
94
+ return x
95
+
96
+ def remove_weight_norm(self):
97
+ for l in self.convs:
98
+ remove_weight_norm(l)
99
+
100
+
101
+ class SineGen(torch.nn.Module):
102
+ """ Definition of sine generator
103
+ SineGen(samp_rate, harmonic_num = 0,
104
+ sine_amp = 0.1, noise_std = 0.003,
105
+ voiced_threshold = 0,
106
+ flag_for_pulse=False)
107
+ samp_rate: sampling rate in Hz
108
+ harmonic_num: number of harmonic overtones (default 0)
109
+ sine_amp: amplitude of sine-wavefrom (default 0.1)
110
+ noise_std: std of Gaussian noise (default 0.003)
111
+ voiced_thoreshold: F0 threshold for U/V classification (default 0)
112
+ flag_for_pulse: this SinGen is used inside PulseGen (default False)
113
+ Note: when flag_for_pulse is True, the first time step of a voiced
114
+ segment is always sin(np.pi) or cos(0)
115
+ """
116
+
117
+ def __init__(self, samp_rate, harmonic_num=0,
118
+ sine_amp=0.1, noise_std=0.003,
119
+ voiced_threshold=0):
120
+ super(SineGen, self).__init__()
121
+ self.sine_amp = sine_amp
122
+ self.noise_std = noise_std
123
+ self.harmonic_num = harmonic_num
124
+ self.dim = self.harmonic_num + 1
125
+ self.sampling_rate = samp_rate
126
+ self.voiced_threshold = voiced_threshold
127
+
128
+ def _f02uv(self, f0):
129
+ # generate uv signal
130
+ uv = torch.ones_like(f0)
131
+ uv = uv * (f0 > self.voiced_threshold)
132
+ return uv
133
+
134
+ @torch.no_grad()
135
+ def forward(self, f0, upp):
136
+ """ sine_tensor, uv = forward(f0)
137
+ input F0: tensor(batchsize=1, length, dim=1)
138
+ f0 for unvoiced steps should be 0
139
+ output sine_tensor: tensor(batchsize=1, length, dim)
140
+ output uv: tensor(batchsize=1, length, 1)
141
+ """
142
+ f0 = f0.unsqueeze(-1)
143
+ fn = torch.multiply(f0, torch.arange(1, self.dim + 1, device=f0.device).reshape((1, 1, -1)))
144
+ rad_values = (fn / self.sampling_rate) % 1 ###%1意味着n_har的乘积无法后处理优化
145
+ rand_ini = torch.rand(fn.shape[0], fn.shape[2], device=fn.device)
146
+ rand_ini[:, 0] = 0
147
+ rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini
148
+ is_half = rad_values.dtype is not torch.float32
149
+ tmp_over_one = torch.cumsum(rad_values.double(), 1) # % 1 #####%1意味着后面的cumsum无法再优化
150
+ if is_half:
151
+ tmp_over_one = tmp_over_one.half()
152
+ else:
153
+ tmp_over_one = tmp_over_one.float()
154
+ tmp_over_one *= upp
155
+ tmp_over_one = F.interpolate(
156
+ tmp_over_one.transpose(2, 1), scale_factor=upp,
157
+ mode='linear', align_corners=True
158
+ ).transpose(2, 1)
159
+ rad_values = F.interpolate(rad_values.transpose(2, 1), scale_factor=upp, mode='nearest').transpose(2, 1)
160
+ tmp_over_one %= 1
161
+ tmp_over_one_idx = (tmp_over_one[:, 1:, :] - tmp_over_one[:, :-1, :]) < 0
162
+ cumsum_shift = torch.zeros_like(rad_values)
163
+ cumsum_shift[:, 1:, :] = tmp_over_one_idx * -1.0
164
+ rad_values = rad_values.double()
165
+ cumsum_shift = cumsum_shift.double()
166
+ sine_waves = torch.sin(torch.cumsum(rad_values + cumsum_shift, dim=1) * 2 * np.pi)
167
+ if is_half:
168
+ sine_waves = sine_waves.half()
169
+ else:
170
+ sine_waves = sine_waves.float()
171
+ sine_waves = sine_waves * self.sine_amp
172
+ uv = self._f02uv(f0)
173
+ uv = F.interpolate(uv.transpose(2, 1), scale_factor=upp, mode='nearest').transpose(2, 1)
174
+ noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
175
+ noise = noise_amp * torch.randn_like(sine_waves)
176
+ sine_waves = sine_waves * uv + noise
177
+ return sine_waves, uv, noise
178
+
179
+
180
+ class SourceModuleHnNSF(torch.nn.Module):
181
+ """ SourceModule for hn-nsf
182
+ SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1,
183
+ add_noise_std=0.003, voiced_threshod=0)
184
+ sampling_rate: sampling_rate in Hz
185
+ harmonic_num: number of harmonic above F0 (default: 0)
186
+ sine_amp: amplitude of sine source signal (default: 0.1)
187
+ add_noise_std: std of additive Gaussian noise (default: 0.003)
188
+ note that amplitude of noise in unvoiced is decided
189
+ by sine_amp
190
+ voiced_threshold: threhold to set U/V given F0 (default: 0)
191
+ Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
192
+ F0_sampled (batchsize, length, 1)
193
+ Sine_source (batchsize, length, 1)
194
+ noise_source (batchsize, length 1)
195
+ uv (batchsize, length, 1)
196
+ """
197
+
198
+ def __init__(self, sampling_rate, harmonic_num=0, sine_amp=0.1,
199
+ add_noise_std=0.003, voiced_threshod=0):
200
+ super(SourceModuleHnNSF, self).__init__()
201
+
202
+ self.sine_amp = sine_amp
203
+ self.noise_std = add_noise_std
204
+
205
+ # to produce sine waveforms
206
+ self.l_sin_gen = SineGen(sampling_rate, harmonic_num,
207
+ sine_amp, add_noise_std, voiced_threshod)
208
+
209
+ # to merge source harmonics into a single excitation
210
+ self.l_linear = torch.nn.Linear(harmonic_num + 1, 1)
211
+ self.l_tanh = torch.nn.Tanh()
212
+
213
+ def forward(self, x, upp):
214
+ sine_wavs, uv, _ = self.l_sin_gen(x, upp)
215
+ sine_merge = self.l_tanh(self.l_linear(sine_wavs))
216
+ return sine_merge
217
+
218
+
219
+ class Generator(torch.nn.Module):
220
+ def __init__(self, h):
221
+ super(Generator, self).__init__()
222
+ self.h = h
223
+ self.num_kernels = len(h.resblock_kernel_sizes)
224
+ self.num_upsamples = len(h.upsample_rates)
225
+ self.m_source = SourceModuleHnNSF(
226
+ sampling_rate=h.sampling_rate,
227
+ harmonic_num=8
228
+ )
229
+ self.noise_convs = nn.ModuleList()
230
+ self.conv_pre = weight_norm(Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3))
231
+ resblock = ResBlock1 if h.resblock == '1' else ResBlock2
232
+
233
+ self.ups = nn.ModuleList()
234
+ for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
235
+ c_cur = h.upsample_initial_channel // (2 ** (i + 1))
236
+ self.ups.append(weight_norm(
237
+ ConvTranspose1d(h.upsample_initial_channel // (2 ** i), h.upsample_initial_channel // (2 ** (i + 1)),
238
+ k, u, padding=(k - u) // 2)))
239
+ if i + 1 < len(h.upsample_rates): #
240
+ stride_f0 = int(np.prod(h.upsample_rates[i + 1:]))
241
+ self.noise_convs.append(Conv1d(
242
+ 1, c_cur, kernel_size=stride_f0 * 2, stride=stride_f0, padding=stride_f0 // 2))
243
+ else:
244
+ self.noise_convs.append(Conv1d(1, c_cur, kernel_size=1))
245
+ self.resblocks = nn.ModuleList()
246
+ ch = h.upsample_initial_channel
247
+ for i in range(len(self.ups)):
248
+ ch //= 2
249
+ for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)):
250
+ self.resblocks.append(resblock(h, ch, k, d))
251
+
252
+ self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3))
253
+ self.ups.apply(init_weights)
254
+ self.conv_post.apply(init_weights)
255
+ self.upp = int(np.prod(h.upsample_rates))
256
+
257
+ def forward(self, x, f0):
258
+ har_source = self.m_source(f0, self.upp).transpose(1, 2)
259
+ x = self.conv_pre(x)
260
+ for i in range(self.num_upsamples):
261
+ x = F.leaky_relu(x, LRELU_SLOPE)
262
+ x = self.ups[i](x)
263
+ x_source = self.noise_convs[i](har_source)
264
+ x = x + x_source
265
+ xs = None
266
+ for j in range(self.num_kernels):
267
+ if xs is None:
268
+ xs = self.resblocks[i * self.num_kernels + j](x)
269
+ else:
270
+ xs += self.resblocks[i * self.num_kernels + j](x)
271
+ x = xs / self.num_kernels
272
+ x = F.leaky_relu(x)
273
+ x = self.conv_post(x)
274
+ x = torch.tanh(x)
275
+
276
+ return x
277
+
278
+ def remove_weight_norm(self):
279
+ print('Removing weight norm...')
280
+ for l in self.ups:
281
+ remove_weight_norm(l)
282
+ for l in self.resblocks:
283
+ l.remove_weight_norm()
284
+ remove_weight_norm(self.conv_pre)
285
+ remove_weight_norm(self.conv_post)
286
+
287
+
288
+ class DiscriminatorP(torch.nn.Module):
289
+ def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
290
+ super(DiscriminatorP, self).__init__()
291
+ self.period = period
292
+ norm_f = weight_norm if use_spectral_norm == False else spectral_norm
293
+ self.convs = nn.ModuleList([
294
+ norm_f(Conv2d(1, 32, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
295
+ norm_f(Conv2d(32, 128, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
296
+ norm_f(Conv2d(128, 512, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
297
+ norm_f(Conv2d(512, 1024, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
298
+ norm_f(Conv2d(1024, 1024, (kernel_size, 1), 1, padding=(2, 0))),
299
+ ])
300
+ self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
301
+
302
+ def forward(self, x):
303
+ fmap = []
304
+
305
+ # 1d to 2d
306
+ b, c, t = x.shape
307
+ if t % self.period != 0: # pad first
308
+ n_pad = self.period - (t % self.period)
309
+ x = F.pad(x, (0, n_pad), "reflect")
310
+ t = t + n_pad
311
+ x = x.view(b, c, t // self.period, self.period)
312
+
313
+ for l in self.convs:
314
+ x = l(x)
315
+ x = F.leaky_relu(x, LRELU_SLOPE)
316
+ fmap.append(x)
317
+ x = self.conv_post(x)
318
+ fmap.append(x)
319
+ x = torch.flatten(x, 1, -1)
320
+
321
+ return x, fmap
322
+
323
+
324
+ class MultiPeriodDiscriminator(torch.nn.Module):
325
+ def __init__(self, periods=None):
326
+ super(MultiPeriodDiscriminator, self).__init__()
327
+ self.periods = periods if periods is not None else [2, 3, 5, 7, 11]
328
+ self.discriminators = nn.ModuleList()
329
+ for period in self.periods:
330
+ self.discriminators.append(DiscriminatorP(period))
331
+
332
+ def forward(self, y, y_hat):
333
+ y_d_rs = []
334
+ y_d_gs = []
335
+ fmap_rs = []
336
+ fmap_gs = []
337
+ for i, d in enumerate(self.discriminators):
338
+ y_d_r, fmap_r = d(y)
339
+ y_d_g, fmap_g = d(y_hat)
340
+ y_d_rs.append(y_d_r)
341
+ fmap_rs.append(fmap_r)
342
+ y_d_gs.append(y_d_g)
343
+ fmap_gs.append(fmap_g)
344
+
345
+ return y_d_rs, y_d_gs, fmap_rs, fmap_gs
346
+
347
+
348
+ class DiscriminatorS(torch.nn.Module):
349
+ def __init__(self, use_spectral_norm=False):
350
+ super(DiscriminatorS, self).__init__()
351
+ norm_f = weight_norm if use_spectral_norm == False else spectral_norm
352
+ self.convs = nn.ModuleList([
353
+ norm_f(Conv1d(1, 128, 15, 1, padding=7)),
354
+ norm_f(Conv1d(128, 128, 41, 2, groups=4, padding=20)),
355
+ norm_f(Conv1d(128, 256, 41, 2, groups=16, padding=20)),
356
+ norm_f(Conv1d(256, 512, 41, 4, groups=16, padding=20)),
357
+ norm_f(Conv1d(512, 1024, 41, 4, groups=16, padding=20)),
358
+ norm_f(Conv1d(1024, 1024, 41, 1, groups=16, padding=20)),
359
+ norm_f(Conv1d(1024, 1024, 5, 1, padding=2)),
360
+ ])
361
+ self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))
362
+
363
+ def forward(self, x):
364
+ fmap = []
365
+ for l in self.convs:
366
+ x = l(x)
367
+ x = F.leaky_relu(x, LRELU_SLOPE)
368
+ fmap.append(x)
369
+ x = self.conv_post(x)
370
+ fmap.append(x)
371
+ x = torch.flatten(x, 1, -1)
372
+
373
+ return x, fmap
374
+
375
+
376
+ class MultiScaleDiscriminator(torch.nn.Module):
377
+ def __init__(self):
378
+ super(MultiScaleDiscriminator, self).__init__()
379
+ self.discriminators = nn.ModuleList([
380
+ DiscriminatorS(use_spectral_norm=True),
381
+ DiscriminatorS(),
382
+ DiscriminatorS(),
383
+ ])
384
+ self.meanpools = nn.ModuleList([
385
+ AvgPool1d(4, 2, padding=2),
386
+ AvgPool1d(4, 2, padding=2)
387
+ ])
388
+
389
+ def forward(self, y, y_hat):
390
+ y_d_rs = []
391
+ y_d_gs = []
392
+ fmap_rs = []
393
+ fmap_gs = []
394
+ for i, d in enumerate(self.discriminators):
395
+ if i != 0:
396
+ y = self.meanpools[i - 1](y)
397
+ y_hat = self.meanpools[i - 1](y_hat)
398
+ y_d_r, fmap_r = d(y)
399
+ y_d_g, fmap_g = d(y_hat)
400
+ y_d_rs.append(y_d_r)
401
+ fmap_rs.append(fmap_r)
402
+ y_d_gs.append(y_d_g)
403
+ fmap_gs.append(fmap_g)
404
+
405
+ return y_d_rs, y_d_gs, fmap_rs, fmap_gs
406
+
407
+
408
+ def feature_loss(fmap_r, fmap_g):
409
+ loss = 0
410
+ for dr, dg in zip(fmap_r, fmap_g):
411
+ for rl, gl in zip(dr, dg):
412
+ loss += torch.mean(torch.abs(rl - gl))
413
+
414
+ return loss * 2
415
+
416
+
417
+ def discriminator_loss(disc_real_outputs, disc_generated_outputs):
418
+ loss = 0
419
+ r_losses = []
420
+ g_losses = []
421
+ for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
422
+ r_loss = torch.mean((1 - dr) ** 2)
423
+ g_loss = torch.mean(dg ** 2)
424
+ loss += (r_loss + g_loss)
425
+ r_losses.append(r_loss.item())
426
+ g_losses.append(g_loss.item())
427
+
428
+ return loss, r_losses, g_losses
429
+
430
+
431
+ def generator_loss(disc_outputs):
432
+ loss = 0
433
+ gen_losses = []
434
+ for dg in disc_outputs:
435
+ l = torch.mean((1 - dg) ** 2)
436
+ gen_losses.append(l)
437
+ loss += l
438
+
439
+ return loss, gen_losses