forked from speechbrain/speechbrain
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathvocoders.py
More file actions
390 lines (336 loc) · 12.6 KB
/
Copy pathvocoders.py
File metadata and controls
390 lines (336 loc) · 12.6 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
""" Specifies the inference interfaces for Text-To-Speech (TTS) modules.
Authors:
* Aku Rouhe 2021
* Peter Plantinga 2021
* Loren Lugosch 2020
* Mirco Ravanelli 2020
* Titouan Parcollet 2021
* Abdel Heba 2021
* Andreas Nautsch 2022, 2023
* Pooneh Mousavi 2023
* Sylvain de Langen 2023
* Adel Moumen 2023
* Pradnya Kandarkar 2023
"""
import torch
from speechbrain.dataio.dataio import length_to_mask
from speechbrain.inference.interfaces import Pretrained
from speechbrain.utils.logger import get_logger
logger = get_logger(__name__)
class HIFIGAN(Pretrained):
"""
A ready-to-use wrapper for HiFiGAN (mel_spec -> waveform).
Arguments
---------
*args : tuple
**kwargs : dict
Arguments are forwarded to ``Pretrained`` parent class.
Example
-------
>>> tmpdir_vocoder = getfixture('tmpdir') / "vocoder"
>>> hifi_gan = HIFIGAN.from_hparams(source="speechbrain/tts-hifigan-ljspeech", savedir=tmpdir_vocoder)
>>> mel_specs = torch.rand(2, 80,298)
>>> waveforms = hifi_gan.decode_batch(mel_specs)
>>> # You can use the vocoder coupled with a TTS system
>>> # Initialize TTS (tacotron2)
>>> tmpdir_tts = getfixture('tmpdir') / "tts"
>>> from speechbrain.inference.TTS import Tacotron2
>>> tacotron2 = Tacotron2.from_hparams(source="speechbrain/tts-tacotron2-ljspeech", savedir=tmpdir_tts)
>>> # Running the TTS
>>> mel_output, mel_length, alignment = tacotron2.encode_text("Mary had a little lamb")
>>> # Running Vocoder (spectrogram-to-waveform)
>>> waveforms = hifi_gan.decode_batch(mel_output)
"""
HPARAMS_NEEDED = ["generator"]
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.infer = self.hparams.generator.inference
self.first_call = True
def decode_batch(self, spectrogram, mel_lens=None, hop_len=None):
"""Computes waveforms from a batch of mel-spectrograms
Arguments
---------
spectrogram: torch.Tensor
Batch of mel-spectrograms [batch, mels, time]
mel_lens: torch.tensor
A list of lengths of mel-spectrograms for the batch
Can be obtained from the output of Tacotron/FastSpeech
hop_len: int
hop length used for mel-spectrogram extraction
should be the same value as in the .yaml file
Returns
-------
waveforms: torch.Tensor
Batch of mel-waveforms [batch, 1, time]
"""
# Prepare for inference by removing the weight norm
if self.first_call:
self.hparams.generator.remove_weight_norm()
self.first_call = False
with torch.no_grad():
waveform = self.infer(spectrogram.to(self.device))
# Mask the noise caused by padding during batch inference
if mel_lens is not None and hop_len is not None:
waveform = self.mask_noise(waveform, mel_lens, hop_len)
return waveform
def mask_noise(self, waveform, mel_lens, hop_len):
"""Mask the noise caused by padding during batch inference
Arguments
---------
waveform: torch.tensor
Batch of generated waveforms [batch, 1, time]
mel_lens: torch.tensor
A list of lengths of mel-spectrograms for the batch
Can be obtained from the output of Tacotron/FastSpeech
hop_len: int
hop length used for mel-spectrogram extraction
same value as in the .yaml file
Returns
-------
waveform: torch.tensor
Batch of waveforms without padded noise [batch, 1, time]
"""
waveform = waveform.squeeze(1)
# the correct audio length should be hop_len * mel_len
mask = length_to_mask(
mel_lens * hop_len, waveform.shape[1], device=waveform.device
).bool()
waveform.masked_fill_(~mask, 0.0)
return waveform.unsqueeze(1)
def decode_spectrogram(self, spectrogram):
"""Computes waveforms from a single mel-spectrogram
Arguments
---------
spectrogram: torch.Tensor
mel-spectrogram [mels, time]
Returns
-------
waveform: torch.Tensor
waveform [1, time]
audio can be saved by:
>>> import torchaudio
>>> waveform = torch.rand(1, 666666)
>>> sample_rate = 22050
>>> torchaudio.save(str(getfixture('tmpdir') / "test.wav"), waveform, sample_rate)
"""
if self.first_call:
self.hparams.generator.remove_weight_norm()
self.first_call = False
with torch.no_grad():
waveform = self.infer(spectrogram.unsqueeze(0).to(self.device))
return waveform.squeeze(0)
def forward(self, spectrogram):
"Decodes the input spectrograms"
return self.decode_batch(spectrogram)
class DiffWaveVocoder(Pretrained):
"""
A ready-to-use inference wrapper for DiffWave as vocoder.
The wrapper allows to perform generative tasks:
locally-conditional generation: mel_spec -> waveform
Arguments
---------
*args : tuple
**kwargs : dict
Arguments are forwarded to ``Pretrained`` parent class.
"""
HPARAMS_NEEDED = ["diffusion"]
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
if hasattr(self.hparams, "diffwave"):
self.infer = self.hparams.diffusion.inference
else:
raise NotImplementedError
def decode_batch(
self,
mel,
hop_len,
mel_lens=None,
fast_sampling=False,
fast_sampling_noise_schedule=None,
):
"""Generate waveforms from spectrograms
Arguments
---------
mel: torch.tensor
spectrogram [batch, mels, time]
hop_len: int
Hop length during mel-spectrogram extraction
Should be the same value as in the .yaml file
Used to determine the output wave length
Also used to mask the noise for vocoding task
mel_lens: torch.tensor
Used to mask the noise caused by padding
A list of lengths of mel-spectrograms for the batch
Can be obtained from the output of Tacotron/FastSpeech
fast_sampling: bool
whether to do fast sampling
fast_sampling_noise_schedule: list
the noise schedules used for fast sampling
Returns
-------
waveforms: torch.tensor
Batch of mel-waveforms [batch, 1, time]
"""
with torch.no_grad():
waveform = self.infer(
unconditional=False,
scale=hop_len,
condition=mel.to(self.device),
fast_sampling=fast_sampling,
fast_sampling_noise_schedule=fast_sampling_noise_schedule,
)
# Mask the noise caused by padding during batch inference
if mel_lens is not None and hop_len is not None:
waveform = self.mask_noise(waveform, mel_lens, hop_len)
return waveform
def mask_noise(self, waveform, mel_lens, hop_len):
"""Mask the noise caused by padding during batch inference
Arguments
---------
waveform: torch.tensor
Batch of generated waveforms [batch, 1, time]
mel_lens: torch.tensor
A list of lengths of mel-spectrograms for the batch
Can be obtained from the output of Tacotron/FastSpeech
hop_len: int
hop length used for mel-spectrogram extraction
same value as in the .yaml file
Returns
-------
waveform: torch.tensor
Batch of waveforms without padded noise [batch, 1, time]
"""
waveform = waveform.squeeze(1)
# the correct audio length should be hop_len * mel_len
mask = length_to_mask(
mel_lens * hop_len, waveform.shape[1], device=waveform.device
).bool()
waveform.masked_fill_(~mask, 0.0)
return waveform.unsqueeze(1)
def decode_spectrogram(
self,
spectrogram,
hop_len,
fast_sampling=False,
fast_sampling_noise_schedule=None,
):
"""Computes waveforms from a single mel-spectrogram
Arguments
---------
spectrogram: torch.tensor
mel-spectrogram [mels, time]
hop_len: int
hop length used for mel-spectrogram extraction
same value as in the .yaml file
fast_sampling: bool
whether to do fast sampling
fast_sampling_noise_schedule: list
the noise schedules used for fast sampling
Returns
-------
waveform: torch.tensor
waveform [1, time]
audio can be saved by:
>>> import torchaudio
>>> waveform = torch.rand(1, 666666)
>>> sample_rate = 22050
>>> torchaudio.save(str(getfixture('tmpdir') / "test.wav"), waveform, sample_rate)
"""
with torch.no_grad():
waveform = self.infer(
unconditional=False,
scale=hop_len,
condition=spectrogram.unsqueeze(0).to(self.device),
fast_sampling=fast_sampling,
fast_sampling_noise_schedule=fast_sampling_noise_schedule,
)
return waveform.squeeze(0)
def forward(self, spectrogram):
"""Decodes the input spectrograms"""
return self.decode_batch(spectrogram)
class UnitHIFIGAN(Pretrained):
"""
A ready-to-use wrapper for Unit HiFiGAN (discrete units -> waveform).
Arguments
---------
*args : tuple
See `Pretrained`
**kwargs : dict
See `Pretrained`
Example
-------
>>> tmpdir_vocoder = getfixture('tmpdir') / "vocoder"
>>> hifi_gan = UnitHIFIGAN.from_hparams(source="speechbrain/hifigan-hubert-l1-3-7-12-18-23-k1000-LibriTTS", savedir=tmpdir_vocoder)
>>> codes = torch.randint(0, 99, (100, 1))
>>> waveform = hifi_gan.decode_unit(codes)
"""
HPARAMS_NEEDED = ["generator"]
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.infer = self.hparams.generator.inference
self.first_call = True
# Temporary fix for mapping indices from the range [0, k] to [1, k+1]
self.tokenize = True
def decode_batch(self, units, spk=None):
"""Computes waveforms from a batch of discrete units
Arguments
---------
units: torch.tensor
Batch of discrete units [batch, codes]
spk: torch.tensor
Batch of speaker embeddings [batch, spk_dim]
Returns
-------
waveforms: torch.tensor
Batch of mel-waveforms [batch, 1, time]
"""
# Remove weight norm for inference if it's the first call
if self.first_call:
self.hparams.generator.remove_weight_norm()
self.first_call = False
# Ensure that the units sequence has a length of at least 3
if units.size(1) < 3:
raise ValueError(
"The 'units' argument should have a length of at least 3 because of padding size."
)
# Increment units if tokenization is enabled
if self.tokenize:
units += 1
if spk is not None:
spk = spk.to(self.device)
with torch.no_grad():
waveform = self.infer(units.to(self.device), spk=spk)
return waveform
def decode_unit(self, units, spk=None):
"""Computes waveforms from a single sequence of discrete units
Arguments
---------
units: torch.tensor
codes: [time]
spk: torch.tensor
spk: [spk_dim]
Returns
-------
waveform: torch.tensor
waveform [1, time]
"""
# Remove weight norm for inference if it's the first call
if self.first_call:
self.hparams.generator.remove_weight_norm()
self.first_call = False
# Ensure that the units sequence has a length of at least 4
if units.size(0) < 4:
raise ValueError(
"The 'units' argument should have a length of at least 4 because of padding size."
)
# Increment units if tokenization is enabled
if self.tokenize:
units = units + 1
if spk is not None:
spk = spk.unsqueeze(0).to(self.device)
with torch.no_grad():
waveform = self.infer(units.unsqueeze(0).to(self.device), spk=spk)
return waveform.squeeze(0)
def forward(self, units, spk=None):
"Decodes the input units"
return self.decode_batch(units, spk=spk)