kevinwang676 commited on
Commit
0ce44cb
·
verified ·
1 Parent(s): c4e8e20

Delete usr/diff/shallow_diffusion_tts_gpu.py

Browse files
Files changed (1) hide show
  1. usr/diff/shallow_diffusion_tts_gpu.py +0 -321
usr/diff/shallow_diffusion_tts_gpu.py DELETED
@@ -1,321 +0,0 @@
1
- import math
2
- import random
3
- from collections import deque
4
- from functools import partial
5
- from inspect import isfunction
6
- from pathlib import Path
7
- import numpy as np
8
- import torch
9
- import torch.nn.functional as F
10
- from torch import nn
11
- from tqdm import tqdm
12
- from einops import rearrange
13
-
14
- from modules.fastspeech.fs2 import FastSpeech2
15
- from modules.diffsinger_midi.fs2 import FastSpeech2MIDI
16
- from utils.hparams import hparams
17
-
18
- import spaces
19
-
20
- def exists(x):
21
- return x is not None
22
-
23
-
24
- def default(val, d):
25
- if exists(val):
26
- return val
27
- return d() if isfunction(d) else d
28
-
29
-
30
- # gaussian diffusion trainer class
31
-
32
- def extract(a, t, x_shape):
33
- b, *_ = t.shape
34
- out = a.gather(-1, t)
35
- return out.reshape(b, *((1,) * (len(x_shape) - 1)))
36
-
37
-
38
- def noise_like(shape, device, repeat=False):
39
- repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
40
- noise = lambda: torch.randn(shape, device=device)
41
- return repeat_noise() if repeat else noise()
42
-
43
-
44
- def linear_beta_schedule(timesteps, max_beta=hparams.get('max_beta', 0.01)):
45
- """
46
- linear schedule
47
- """
48
- betas = np.linspace(1e-4, max_beta, timesteps)
49
- return betas
50
-
51
-
52
- def cosine_beta_schedule(timesteps, s=0.008):
53
- """
54
- cosine schedule
55
- as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
56
- """
57
- steps = timesteps + 1
58
- x = np.linspace(0, steps, steps)
59
- alphas_cumprod = np.cos(((x / steps) + s) / (1 + s) * np.pi * 0.5) ** 2
60
- alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
61
- betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
62
- return np.clip(betas, a_min=0, a_max=0.999)
63
-
64
-
65
- beta_schedule = {
66
- "cosine": cosine_beta_schedule,
67
- "linear": linear_beta_schedule,
68
- }
69
-
70
-
71
- class GaussianDiffusion(nn.Module):
72
- def __init__(self, phone_encoder, out_dims, denoise_fn,
73
- timesteps=1000, K_step=1000, loss_type=hparams.get('diff_loss_type', 'l1'), betas=None, spec_min=None, spec_max=None):
74
- super().__init__()
75
- self.denoise_fn = denoise_fn
76
- if hparams.get('use_midi') is not None and hparams['use_midi']:
77
- self.fs2 = FastSpeech2MIDI(phone_encoder, out_dims)
78
- else:
79
- self.fs2 = FastSpeech2(phone_encoder, out_dims)
80
- self.mel_bins = out_dims
81
-
82
- if exists(betas):
83
- betas = betas.detach().cpu().numpy() if isinstance(betas, torch.Tensor) else betas
84
- else:
85
- if 'schedule_type' in hparams.keys():
86
- betas = beta_schedule[hparams['schedule_type']](timesteps)
87
- else:
88
- betas = cosine_beta_schedule(timesteps)
89
-
90
- alphas = 1. - betas
91
- alphas_cumprod = np.cumprod(alphas, axis=0)
92
- alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
93
-
94
- timesteps, = betas.shape
95
- self.num_timesteps = int(timesteps)
96
- self.K_step = K_step
97
- self.loss_type = loss_type
98
-
99
- self.noise_list = deque(maxlen=4)
100
-
101
- to_torch = partial(torch.tensor, dtype=torch.float32)
102
-
103
- self.register_buffer('betas', to_torch(betas))
104
- self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
105
- self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
106
-
107
- # calculations for diffusion q(x_t | x_{t-1}) and others
108
- self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
109
- self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
110
- self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
111
- self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
112
- self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))
113
-
114
- # calculations for posterior q(x_{t-1} | x_t, x_0)
115
- posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
116
- # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
117
- self.register_buffer('posterior_variance', to_torch(posterior_variance))
118
- # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
119
- self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20))))
120
- self.register_buffer('posterior_mean_coef1', to_torch(
121
- betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)))
122
- self.register_buffer('posterior_mean_coef2', to_torch(
123
- (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod)))
124
-
125
- self.register_buffer('spec_min', torch.FloatTensor(spec_min)[None, None, :hparams['keep_bins']])
126
- self.register_buffer('spec_max', torch.FloatTensor(spec_max)[None, None, :hparams['keep_bins']])
127
-
128
- def q_mean_variance(self, x_start, t):
129
- mean = extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
130
- variance = extract(1. - self.alphas_cumprod, t, x_start.shape)
131
- log_variance = extract(self.log_one_minus_alphas_cumprod, t, x_start.shape)
132
- return mean, variance, log_variance
133
-
134
- def predict_start_from_noise(self, x_t, t, noise):
135
- return (
136
- extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
137
- extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
138
- )
139
-
140
- def q_posterior(self, x_start, x_t, t):
141
- posterior_mean = (
142
- extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +
143
- extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
144
- )
145
- posterior_variance = extract(self.posterior_variance, t, x_t.shape)
146
- posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)
147
- return posterior_mean, posterior_variance, posterior_log_variance_clipped
148
-
149
- def p_mean_variance(self, x, t, cond, clip_denoised: bool):
150
- noise_pred = self.denoise_fn(x, t, cond=cond)
151
- x_recon = self.predict_start_from_noise(x, t=t, noise=noise_pred)
152
-
153
- if clip_denoised:
154
- x_recon.clamp_(-1., 1.)
155
-
156
- model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
157
- return model_mean, posterior_variance, posterior_log_variance
158
-
159
- @torch.no_grad()
160
- def p_sample(self, x, t, cond, clip_denoised=True, repeat_noise=False):
161
- b, *_, device = *x.shape, x.device
162
- model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, cond=cond, clip_denoised=clip_denoised)
163
- noise = noise_like(x.shape, device, repeat_noise)
164
- # no noise when t == 0
165
- nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
166
- return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
167
-
168
- @torch.no_grad()
169
- def p_sample_plms(self, x, t, interval, cond, clip_denoised=True, repeat_noise=False):
170
- """
171
- Use the PLMS method from [Pseudo Numerical Methods for Diffusion Models on Manifolds](https://arxiv.org/abs/2202.09778).
172
- """
173
-
174
- def get_x_pred(x, noise_t, t):
175
- a_t = extract(self.alphas_cumprod, t, x.shape)
176
- a_prev = extract(self.alphas_cumprod, torch.max(t-interval, torch.zeros_like(t)), x.shape)
177
- a_t_sq, a_prev_sq = a_t.sqrt(), a_prev.sqrt()
178
-
179
- x_delta = (a_prev - a_t) * ((1 / (a_t_sq * (a_t_sq + a_prev_sq))) * x - 1 / (a_t_sq * (((1 - a_prev) * a_t).sqrt() + ((1 - a_t) * a_prev).sqrt())) * noise_t)
180
- x_pred = x + x_delta
181
-
182
- return x_pred
183
-
184
- noise_list = self.noise_list
185
- noise_pred = self.denoise_fn(x, t, cond=cond)
186
-
187
- if len(noise_list) == 0:
188
- x_pred = get_x_pred(x, noise_pred, t)
189
- noise_pred_prev = self.denoise_fn(x_pred, max(t-interval, 0), cond=cond)
190
- noise_pred_prime = (noise_pred + noise_pred_prev) / 2
191
- elif len(noise_list) == 1:
192
- noise_pred_prime = (3 * noise_pred - noise_list[-1]) / 2
193
- elif len(noise_list) == 2:
194
- noise_pred_prime = (23 * noise_pred - 16 * noise_list[-1] + 5 * noise_list[-2]) / 12
195
- elif len(noise_list) >= 3:
196
- noise_pred_prime = (55 * noise_pred - 59 * noise_list[-1] + 37 * noise_list[-2] - 9 * noise_list[-3]) / 24
197
-
198
- x_prev = get_x_pred(x, noise_pred_prime, t)
199
- noise_list.append(noise_pred)
200
-
201
- return x_prev
202
-
203
- def q_sample(self, x_start, t, noise=None):
204
- noise = default(noise, lambda: torch.randn_like(x_start))
205
- return (
206
- extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
207
- extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
208
- )
209
-
210
- def p_losses(self, x_start, t, cond, noise=None, nonpadding=None):
211
- noise = default(noise, lambda: torch.randn_like(x_start))
212
-
213
- x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
214
- x_recon = self.denoise_fn(x_noisy, t, cond)
215
-
216
- if self.loss_type == 'l1':
217
- if nonpadding is not None:
218
- loss = ((noise - x_recon).abs() * nonpadding.unsqueeze(1)).mean()
219
- else:
220
- # print('are you sure w/o nonpadding?')
221
- loss = (noise - x_recon).abs().mean()
222
-
223
- elif self.loss_type == 'l2':
224
- loss = F.mse_loss(noise, x_recon)
225
- else:
226
- raise NotImplementedError()
227
-
228
- return loss
229
-
230
- @spaces.GPU
231
- def forward(self, txt_tokens, mel2ph=None, spk_embed=None,
232
- ref_mels=None, f0=None, uv=None, energy=None, infer=False, **kwargs):
233
- b, *_, device = *txt_tokens.shape, txt_tokens.device
234
- ret = self.fs2(txt_tokens, mel2ph, spk_embed, ref_mels, f0, uv, energy,
235
- skip_decoder=(not infer), infer=infer, **kwargs)
236
- cond = ret['decoder_inp'].transpose(1, 2)
237
-
238
- if not infer:
239
- t = torch.randint(0, self.K_step, (b,), device=device).long()
240
- x = ref_mels
241
- x = self.norm_spec(x)
242
- x = x.transpose(1, 2)[:, None, :, :] # [B, 1, M, T]
243
- ret['diff_loss'] = self.p_losses(x, t, cond)
244
- # nonpadding = (mel2ph != 0).float()
245
- # ret['diff_loss'] = self.p_losses(x, t, cond, nonpadding=nonpadding)
246
- else:
247
- ret['fs2_mel'] = ret['mel_out']
248
- fs2_mels = ret['mel_out']
249
- t = self.K_step
250
- fs2_mels = self.norm_spec(fs2_mels)
251
- fs2_mels = fs2_mels.transpose(1, 2)[:, None, :, :]
252
-
253
- x = self.q_sample(x_start=fs2_mels, t=torch.tensor([t - 1], device=device).long())
254
- if hparams.get('gaussian_start') is not None and hparams['gaussian_start']:
255
- print('===> gaussion start.')
256
- shape = (cond.shape[0], 1, self.mel_bins, cond.shape[2])
257
- x = torch.randn(shape, device=device)
258
-
259
- if hparams.get('pndm_speedup'):
260
- self.noise_list = deque(maxlen=4)
261
- iteration_interval = hparams['pndm_speedup']
262
- for i in tqdm(reversed(range(0, t, iteration_interval)), desc='sample time step',
263
- total=t // iteration_interval):
264
- x = self.p_sample_plms(x, torch.full((b,), i, device=device, dtype=torch.long), iteration_interval,
265
- cond)
266
- else:
267
- for i in tqdm(reversed(range(0, t)), desc='sample time step', total=t):
268
- x = self.p_sample(x, torch.full((b,), i, device=device, dtype=torch.long), cond)
269
- x = x[:, 0].transpose(1, 2)
270
- if mel2ph is not None: # for singing
271
- ret['mel_out'] = self.denorm_spec(x) * ((mel2ph > 0).float()[:, :, None])
272
- else:
273
- ret['mel_out'] = self.denorm_spec(x)
274
- return ret
275
-
276
- def norm_spec(self, x):
277
- return (x - self.spec_min) / (self.spec_max - self.spec_min) * 2 - 1
278
-
279
- def denorm_spec(self, x):
280
- return (x + 1) / 2 * (self.spec_max - self.spec_min) + self.spec_min
281
-
282
- def cwt2f0_norm(self, cwt_spec, mean, std, mel2ph):
283
- return self.fs2.cwt2f0_norm(cwt_spec, mean, std, mel2ph)
284
-
285
- def out2mel(self, x):
286
- return x
287
-
288
-
289
- class OfflineGaussianDiffusion(GaussianDiffusion):
290
- def forward(self, txt_tokens, mel2ph=None, spk_embed=None,
291
- ref_mels=None, f0=None, uv=None, energy=None, infer=False, **kwargs):
292
- b, *_, device = *txt_tokens.shape, txt_tokens.device
293
-
294
- ret = self.fs2(txt_tokens, mel2ph, spk_embed, ref_mels, f0, uv, energy,
295
- skip_decoder=True, infer=True, **kwargs)
296
- cond = ret['decoder_inp'].transpose(1, 2)
297
- fs2_mels = ref_mels[1]
298
- ref_mels = ref_mels[0]
299
-
300
- if not infer:
301
- t = torch.randint(0, self.K_step, (b,), device=device).long()
302
- x = ref_mels
303
- x = self.norm_spec(x)
304
- x = x.transpose(1, 2)[:, None, :, :] # [B, 1, M, T]
305
- ret['diff_loss'] = self.p_losses(x, t, cond)
306
- else:
307
- t = self.K_step
308
- fs2_mels = self.norm_spec(fs2_mels)
309
- fs2_mels = fs2_mels.transpose(1, 2)[:, None, :, :]
310
-
311
- x = self.q_sample(x_start=fs2_mels, t=torch.tensor([t - 1], device=device).long())
312
-
313
- if hparams.get('gaussian_start') is not None and hparams['gaussian_start']:
314
- print('===> gaussion start.')
315
- shape = (cond.shape[0], 1, self.mel_bins, cond.shape[2])
316
- x = torch.randn(shape, device=device)
317
- for i in tqdm(reversed(range(0, t)), desc='sample time step', total=t):
318
- x = self.p_sample(x, torch.full((b,), i, device=device, dtype=torch.long), cond)
319
- x = x[:, 0].transpose(1, 2)
320
- ret['mel_out'] = self.denorm_spec(x)
321
- return ret