Kikirilkov commited on
Commit
722dd3b
β€’
1 Parent(s): e24f073

Delete TTS/vocoder/models/fatchord_version.py

Browse files
TTS/vocoder/models/fatchord_version.py DELETED
@@ -1,413 +0,0 @@
1
- import torch
2
- import torch.nn as nn
3
- import torch.nn.functional as F
4
- from vocoder.distribution import sample_from_discretized_mix_logistic
5
- from vocoder.display import *
6
- from vocoder.audio import *
7
-
8
-
9
- class ResBlock(nn.Module):
10
- def __init__(self, dims):
11
- super().__init__()
12
- self.conv1 = nn.Conv1d(dims, dims, kernel_size=1, bias=False)
13
- self.conv2 = nn.Conv1d(dims, dims, kernel_size=1, bias=False)
14
- self.batch_norm1 = nn.BatchNorm1d(dims)
15
- self.batch_norm2 = nn.BatchNorm1d(dims)
16
-
17
- def forward(self, x):
18
- residual = x
19
- x = self.conv1(x)
20
- x = self.batch_norm1(x)
21
- x = F.relu(x)
22
- x = self.conv2(x)
23
- x = self.batch_norm2(x)
24
- return x + residual
25
-
26
-
27
- class MelResNet(nn.Module):
28
- def __init__(self, res_blocks, in_dims, compute_dims, res_out_dims, pad):
29
- super().__init__()
30
- k_size = pad * 2 + 1
31
- self.conv_in = nn.Conv1d(in_dims, compute_dims, kernel_size=k_size, bias=False)
32
- self.batch_norm = nn.BatchNorm1d(compute_dims)
33
- self.layers = nn.ModuleList()
34
- for i in range(res_blocks):
35
- self.layers.append(ResBlock(compute_dims))
36
- self.conv_out = nn.Conv1d(compute_dims, res_out_dims, kernel_size=1)
37
-
38
- def forward(self, x):
39
- x = self.conv_in(x)
40
- x = self.batch_norm(x)
41
- x = F.relu(x)
42
- for f in self.layers: x = f(x)
43
- x = self.conv_out(x)
44
- return x
45
-
46
-
47
- class Stretch2d(nn.Module):
48
- def __init__(self, x_scale, y_scale):
49
- super().__init__()
50
- self.x_scale = x_scale
51
- self.y_scale = y_scale
52
-
53
- def forward(self, x):
54
- b, c, h, w = x.size()
55
- x = x.unsqueeze(-1).unsqueeze(3)
56
- x = x.repeat(1, 1, 1, self.y_scale, 1, self.x_scale)
57
- return x.view(b, c, h * self.y_scale, w * self.x_scale)
58
-
59
-
60
- class UpsampleNetwork(nn.Module):
61
- def __init__(self, feat_dims, upsample_scales, compute_dims,
62
- res_blocks, res_out_dims, pad):
63
- super().__init__()
64
- total_scale = np.cumproduct(upsample_scales)[-1]
65
- self.indent = pad * total_scale
66
- self.resnet = MelResNet(res_blocks, feat_dims, compute_dims, res_out_dims, pad)
67
- self.resnet_stretch = Stretch2d(total_scale, 1)
68
- self.up_layers = nn.ModuleList()
69
- for scale in upsample_scales:
70
- k_size = (1, scale * 2 + 1)
71
- padding = (0, scale)
72
- stretch = Stretch2d(scale, 1)
73
- conv = nn.Conv2d(1, 1, kernel_size=k_size, padding=padding, bias=False)
74
- conv.weight.data.fill_(1. / k_size[1])
75
- self.up_layers.append(stretch)
76
- self.up_layers.append(conv)
77
-
78
- def forward(self, m):
79
- aux = self.resnet(m).unsqueeze(1)
80
- aux = self.resnet_stretch(aux)
81
- aux = aux.squeeze(1)
82
- m = m.unsqueeze(1)
83
- for f in self.up_layers: m = f(m)
84
- m = m.squeeze(1)[:, :, self.indent:-self.indent]
85
- return m.transpose(1, 2), aux.transpose(1, 2)
86
-
87
-
88
- class WaveRNN(nn.Module):
89
- def __init__(self, rnn_dims, fc_dims, bits, pad, upsample_factors,
90
- feat_dims, compute_dims, res_out_dims, res_blocks,
91
- hop_length, sample_rate, mode='RAW'):
92
- super().__init__()
93
- self.mode = mode
94
- self.pad = pad
95
- if self.mode == 'RAW' :
96
- self.n_classes = 2 ** bits
97
- elif self.mode == 'MOL' :
98
- self.n_classes = 30
99
- else :
100
- RuntimeError("Unknown model mode value - ", self.mode)
101
-
102
- self.rnn_dims = rnn_dims
103
- self.aux_dims = res_out_dims // 4
104
- self.hop_length = hop_length
105
- self.sample_rate = sample_rate
106
-
107
- self.upsample = UpsampleNetwork(feat_dims, upsample_factors, compute_dims, res_blocks, res_out_dims, pad)
108
- self.I = nn.Linear(feat_dims + self.aux_dims + 1, rnn_dims)
109
- self.rnn1 = nn.GRU(rnn_dims, rnn_dims, batch_first=True)
110
- self.rnn2 = nn.GRU(rnn_dims + self.aux_dims, rnn_dims, batch_first=True)
111
- self.fc1 = nn.Linear(rnn_dims + self.aux_dims, fc_dims)
112
- self.fc2 = nn.Linear(fc_dims + self.aux_dims, fc_dims)
113
- self.fc3 = nn.Linear(fc_dims, self.n_classes)
114
-
115
- self.step = nn.Parameter(torch.zeros(1).long(), requires_grad=False)
116
- self.num_params()
117
-
118
- def forward(self, x, mels):
119
- self.step += 1
120
- bsize = x.size(0)
121
- h1 = torch.zeros(1, bsize, self.rnn_dims).cuda()
122
- h2 = torch.zeros(1, bsize, self.rnn_dims).cuda()
123
- mels, aux = self.upsample(mels)
124
-
125
- aux_idx = [self.aux_dims * i for i in range(5)]
126
- a1 = aux[:, :, aux_idx[0]:aux_idx[1]]
127
- a2 = aux[:, :, aux_idx[1]:aux_idx[2]]
128
- a3 = aux[:, :, aux_idx[2]:aux_idx[3]]
129
- a4 = aux[:, :, aux_idx[3]:aux_idx[4]]
130
-
131
- x = torch.cat([x.unsqueeze(-1), mels, a1], dim=2)
132
- x = self.I(x)
133
- res = x
134
- x, _ = self.rnn1(x, h1)
135
-
136
- x = x + res
137
- res = x
138
- x = torch.cat([x, a2], dim=2)
139
- x, _ = self.rnn2(x, h2)
140
-
141
- x = x + res
142
- x = torch.cat([x, a3], dim=2)
143
- x = F.relu(self.fc1(x))
144
-
145
- x = torch.cat([x, a4], dim=2)
146
- x = F.relu(self.fc2(x))
147
- return self.fc3(x)
148
-
149
- def generate(self, mels, batched, target, overlap, mu_law, progress_callback=None):
150
- mu_law = mu_law if self.mode == 'RAW' else False
151
- progress_callback = progress_callback or self.gen_display
152
-
153
- self.eval()
154
- output = []
155
- start = time.time()
156
- rnn1 = self.get_gru_cell(self.rnn1)
157
- rnn2 = self.get_gru_cell(self.rnn2)
158
-
159
- with torch.no_grad():
160
- mels = mels.cuda()
161
- wave_len = (mels.size(-1) - 1) * self.hop_length
162
- mels = self.pad_tensor(mels.transpose(1, 2), pad=self.pad, side='both')
163
- mels, aux = self.upsample(mels.transpose(1, 2))
164
-
165
- if batched:
166
- mels = self.fold_with_overlap(mels, target, overlap)
167
- aux = self.fold_with_overlap(aux, target, overlap)
168
-
169
- b_size, seq_len, _ = mels.size()
170
-
171
- h1 = torch.zeros(b_size, self.rnn_dims).cuda()
172
- h2 = torch.zeros(b_size, self.rnn_dims).cuda()
173
- x = torch.zeros(b_size, 1).cuda()
174
-
175
- d = self.aux_dims
176
- aux_split = [aux[:, :, d * i:d * (i + 1)] for i in range(4)]
177
-
178
- for i in range(seq_len):
179
-
180
- m_t = mels[:, i, :]
181
-
182
- a1_t, a2_t, a3_t, a4_t = (a[:, i, :] for a in aux_split)
183
-
184
- x = torch.cat([x, m_t, a1_t], dim=1)
185
- x = self.I(x)
186
- h1 = rnn1(x, h1)
187
-
188
- x = x + h1
189
- inp = torch.cat([x, a2_t], dim=1)
190
- h2 = rnn2(inp, h2)
191
-
192
- x = x + h2
193
- x = torch.cat([x, a3_t], dim=1)
194
- x = F.relu(self.fc1(x))
195
-
196
- x = torch.cat([x, a4_t], dim=1)
197
- x = F.relu(self.fc2(x))
198
-
199
- logits = self.fc3(x)
200
-
201
- if self.mode == 'MOL':
202
- sample = sample_from_discretized_mix_logistic(logits.unsqueeze(0).transpose(1, 2))
203
- output.append(sample.view(-1))
204
- # x = torch.FloatTensor([[sample]]).cuda()
205
- x = sample.transpose(0, 1).cuda()
206
-
207
- elif self.mode == 'RAW' :
208
- posterior = F.softmax(logits, dim=1)
209
- distrib = torch.distributions.Categorical(posterior)
210
-
211
- sample = 2 * distrib.sample().float() / (self.n_classes - 1.) - 1.
212
- output.append(sample)
213
- x = sample.unsqueeze(-1)
214
- else:
215
- raise RuntimeError("Unknown model mode value - ", self.mode)
216
-
217
- if i % 100 == 0:
218
- gen_rate = (i + 1) / (time.time() - start) * b_size / 1000
219
- progress_callback(i, seq_len, b_size, gen_rate)
220
-
221
- output = torch.stack(output).transpose(0, 1)
222
- output = output.cpu().numpy()
223
- output = output.astype(np.float64)
224
-
225
- if batched:
226
- output = self.xfade_and_unfold(output, target, overlap)
227
- else:
228
- output = output[0]
229
-
230
- if mu_law:
231
- output = decode_mu_law(output, self.n_classes, False)
232
- if hp.apply_preemphasis:
233
- output = de_emphasis(output)
234
-
235
- # Fade-out at the end to avoid signal cutting out suddenly
236
- fade_out = np.linspace(1, 0, 20 * self.hop_length)
237
- output = output[:wave_len]
238
- output[-20 * self.hop_length:] *= fade_out
239
-
240
- self.train()
241
-
242
- return output
243
-
244
-
245
- def gen_display(self, i, seq_len, b_size, gen_rate):
246
- pbar = progbar(i, seq_len)
247
- msg = f'| {pbar} {i*b_size}/{seq_len*b_size} | Batch Size: {b_size} | Gen Rate: {gen_rate:.1f}kHz | '
248
- stream(msg)
249
-
250
- def get_gru_cell(self, gru):
251
- gru_cell = nn.GRUCell(gru.input_size, gru.hidden_size)
252
- gru_cell.weight_hh.data = gru.weight_hh_l0.data
253
- gru_cell.weight_ih.data = gru.weight_ih_l0.data
254
- gru_cell.bias_hh.data = gru.bias_hh_l0.data
255
- gru_cell.bias_ih.data = gru.bias_ih_l0.data
256
- return gru_cell
257
-
258
- def pad_tensor(self, x, pad, side='both'):
259
- # NB - this is just a quick method i need right now
260
- # i.e., it won't generalise to other shapes/dims
261
- b, t, c = x.size()
262
- total = t + 2 * pad if side == 'both' else t + pad
263
- padded = torch.zeros(b, total, c).cuda()
264
- if side == 'before' or side == 'both':
265
- padded[:, pad:pad + t, :] = x
266
- elif side == 'after':
267
- padded[:, :t, :] = x
268
- return padded
269
-
270
- def fold_with_overlap(self, x, target, overlap):
271
-
272
- ''' Fold the tensor with overlap for quick batched inference.
273
- Overlap will be used for crossfading in xfade_and_unfold()
274
-
275
- Args:
276
- x (tensor) : Upsampled conditioning features.
277
- shape=(1, timesteps, features)
278
- target (int) : Target timesteps for each index of batch
279
- overlap (int) : Timesteps for both xfade and rnn warmup
280
-
281
- Return:
282
- (tensor) : shape=(num_folds, target + 2 * overlap, features)
283
-
284
- Details:
285
- x = [[h1, h2, ... hn]]
286
-
287
- Where each h is a vector of conditioning features
288
-
289
- Eg: target=2, overlap=1 with x.size(1)=10
290
-
291
- folded = [[h1, h2, h3, h4],
292
- [h4, h5, h6, h7],
293
- [h7, h8, h9, h10]]
294
- '''
295
-
296
- _, total_len, features = x.size()
297
-
298
- # Calculate variables needed
299
- num_folds = (total_len - overlap) // (target + overlap)
300
- extended_len = num_folds * (overlap + target) + overlap
301
- remaining = total_len - extended_len
302
-
303
- # Pad if some time steps poking out
304
- if remaining != 0:
305
- num_folds += 1
306
- padding = target + 2 * overlap - remaining
307
- x = self.pad_tensor(x, padding, side='after')
308
-
309
- folded = torch.zeros(num_folds, target + 2 * overlap, features).cuda()
310
-
311
- # Get the values for the folded tensor
312
- for i in range(num_folds):
313
- start = i * (target + overlap)
314
- end = start + target + 2 * overlap
315
- folded[i] = x[:, start:end, :]
316
-
317
- return folded
318
-
319
- def xfade_and_unfold(self, y, target, overlap):
320
-
321
- ''' Applies a crossfade and unfolds into a 1d array.
322
-
323
- Args:
324
- y (ndarry) : Batched sequences of audio samples
325
- shape=(num_folds, target + 2 * overlap)
326
- dtype=np.float64
327
- overlap (int) : Timesteps for both xfade and rnn warmup
328
-
329
- Return:
330
- (ndarry) : audio samples in a 1d array
331
- shape=(total_len)
332
- dtype=np.float64
333
-
334
- Details:
335
- y = [[seq1],
336
- [seq2],
337
- [seq3]]
338
-
339
- Apply a gain envelope at both ends of the sequences
340
-
341
- y = [[seq1_in, seq1_target, seq1_out],
342
- [seq2_in, seq2_target, seq2_out],
343
- [seq3_in, seq3_target, seq3_out]]
344
-
345
- Stagger and add up the groups of samples:
346
-
347
- [seq1_in, seq1_target, (seq1_out + seq2_in), seq2_target, ...]
348
-
349
- '''
350
-
351
- num_folds, length = y.shape
352
- target = length - 2 * overlap
353
- total_len = num_folds * (target + overlap) + overlap
354
-
355
- # Need some silence for the rnn warmup
356
- silence_len = overlap // 2
357
- fade_len = overlap - silence_len
358
- silence = np.zeros((silence_len), dtype=np.float64)
359
-
360
- # Equal power crossfade
361
- t = np.linspace(-1, 1, fade_len, dtype=np.float64)
362
- fade_in = np.sqrt(0.5 * (1 + t))
363
- fade_out = np.sqrt(0.5 * (1 - t))
364
-
365
- # Concat the silence to the fades
366
- fade_in = np.concatenate([silence, fade_in])
367
- fade_out = np.concatenate([fade_out, silence])
368
-
369
- # Apply the gain to the overlap samples
370
- y[:, :overlap] *= fade_in
371
- y[:, -overlap:] *= fade_out
372
-
373
- unfolded = np.zeros((total_len), dtype=np.float64)
374
-
375
- # Loop to add up all the samples
376
- for i in range(num_folds):
377
- start = i * (target + overlap)
378
- end = start + target + 2 * overlap
379
- unfolded[start:end] += y[i]
380
-
381
- return unfolded
382
-
383
- def get_step(self) :
384
- return self.step.data.item()
385
-
386
- def checkpoint(self, model_dir, optimizer) :
387
- k_steps = self.get_step() // 1000
388
- self.save(model_dir.joinpath("checkpoint_%dk_steps.pt" % k_steps), optimizer)
389
-
390
- def log(self, path, msg) :
391
- with open(path, 'a') as f:
392
- print(msg, file=f)
393
-
394
- def load(self, path, optimizer) :
395
- checkpoint = torch.load(path)
396
- if "optimizer_state" in checkpoint:
397
- self.load_state_dict(checkpoint["model_state"])
398
- optimizer.load_state_dict(checkpoint["optimizer_state"])
399
- else:
400
- # Backwards compatibility
401
- self.load_state_dict(checkpoint)
402
-
403
- def save(self, path, optimizer) :
404
- torch.save({
405
- "model_state": self.state_dict(),
406
- "optimizer_state": optimizer.state_dict(),
407
- }, path)
408
-
409
- def num_params(self, print_out=True):
410
- parameters = filter(lambda p: p.requires_grad, self.parameters())
411
- parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
412
- if print_out :
413
- print('Trainable Parameters: %.3fM' % parameters)