speechbrain.lobes.models.DiffWave module

Neural network modules for DIFFWAVE: A VERSATILE DIFFUSION MODEL FOR AUDIO SYNTHESIS

For more details: https://arxiv.org/pdf/2009.09761.pdf

Authors

Yingzhi WANG 2022

Summary

Classes:

`DiffWave`	DiffWave Model with dilated residual blocks
`DiffWaveDiffusion`	An enhanced diffusion implementation with DiffWave-specific inference
`DiffusionEmbedding`	Embeds the diffusion step into an input vector of DiffWave
`ResidualBlock`	Residual Block with dilated convolution
`SpectrogramUpsampler`	Upsampler for spectrograms with Transposed Conv Only the upsampling is done here, the layer-specific Conv can be found in residual block to map the mel bands into 2× residual channels

Functions:

diffwave_mel_spectogram

calculates MelSpectrogram for a raw audio signal and preprocesses it for diffwave training

Reference

speechbrain.lobes.models.DiffWave.diffwave_mel_spectogram(sample_rate, hop_length, win_length, n_fft, n_mels, f_min, f_max, power, normalized, norm, mel_scale, audio)[source]

calculates MelSpectrogram for a raw audio signal and preprocesses it for diffwave training

Parameters:

sample_rate (int) – Sample rate of audio signal.
hop_length (int) – Length of hop between STFT windows.
win_length (int) – Window size.
n_fft (int) – Size of FFT.
n_mels (int) – Number of mel filterbanks.
f_min (float) – Minimum frequency.
f_max (float) – Maximum frequency.
power (float) – Exponent for the magnitude spectrogram.
normalized (bool) – Whether to normalize by magnitude after stft.
norm (str or None) – If “slaney”, divide the triangular mel weights by the width of the mel band
mel_scale (str) – Scale to use: “htk” or “slaney”.
audio (torch.tensor) – input audio signal

Returns:

mel

Return type:

torch.Tensor

class speechbrain.lobes.models.DiffWave.DiffusionEmbedding(max_steps)[source]

Bases: Module

Embeds the diffusion step into an input vector of DiffWave

Parameters:: max_steps (int) – total diffusion steps

Example

>>> from speechbrain.lobes.models.DiffWave import DiffusionEmbedding
>>> diffusion_embedding = DiffusionEmbedding(max_steps=50)
>>> time_step = torch.randint(50, (1,))
>>> step_embedding = diffusion_embedding(time_step)
>>> step_embedding.shape
torch.Size([1, 512])

forward(diffusion_step)[source]

forward function of diffusion step embedding

Parameters:: diffusion_step (torch.Tensor) – which step of diffusion to execute
Returns:: diffusion step embedding
Return type:: tensor [bs, 512]

class speechbrain.lobes.models.DiffWave.SpectrogramUpsampler[source]

Bases: Module

Upsampler for spectrograms with Transposed Conv Only the upsampling is done here, the layer-specific Conv can be found in residual block to map the mel bands into 2× residual channels

Example

>>> from speechbrain.lobes.models.DiffWave import SpectrogramUpsampler
>>> spec_upsampler = SpectrogramUpsampler()
>>> mel_input = torch.rand(3, 80, 100)
>>> upsampled_mel = spec_upsampler(mel_input)
>>> upsampled_mel.shape
torch.Size([3, 80, 25600])

forward(x)[source]

Upsamples spectrograms 256 times to match the length of audios Hop length should be 256 when extracting mel spectrograms

Parameters:: x (torch.Tensor) – input mel spectrogram [bs, 80, mel_len]
Return type:: upsampled spectrogram [bs, 80, mel_len*256]

class speechbrain.lobes.models.DiffWave.ResidualBlock(n_mels, residual_channels, dilation, uncond=False)[source]

Bases: Module

Residual Block with dilated convolution

Parameters:

n_mels (int) – input mel channels of conv1x1 for conditional vocoding task
residual_channels (int) – channels of audio convolution
dilation (int) – dilation cycles of audio convolution
uncond (bool) – conditional/unconditional generation

Example

>>> from speechbrain.lobes.models.DiffWave import ResidualBlock
>>> res_block = ResidualBlock(n_mels=80, residual_channels=64, dilation=3)
>>> noisy_audio = torch.randn(1, 1, 22050)
>>> timestep_embedding = torch.rand(1, 512)
>>> upsampled_mel = torch.rand(1, 80, 22050)
>>> output = res_block(noisy_audio, timestep_embedding, upsampled_mel)
>>> output[0].shape
torch.Size([1, 64, 22050])

forward(x, diffusion_step, conditioner=None)[source]

forward function of Residual Block

Parameters:

x (torch.Tensor) – input sample [bs, 1, time]
diffusion_step (torch.Tensor) – the embedding of which step of diffusion to execute
conditioner (torch.Tensor) – the condition used for conditional generation

Returns:

residual output [bs, residual_channels, time]
a skip of residual branch [bs, residual_channels, time]

class speechbrain.lobes.models.DiffWave.DiffWave(input_channels, residual_layers, residual_channels, dilation_cycle_length, total_steps, unconditional=False)[source]

Bases: Module

DiffWave Model with dilated residual blocks

Parameters:

input_channels (int) – input mel channels of conv1x1 for conditional vocoding task
residual_layers (int) – number of residual blocks
residual_channels (int) – channels of audio convolution
dilation_cycle_length (int) – dilation cycles of audio convolution
total_steps (int) – total steps of diffusion
unconditional (bool) – conditional/unconditional generation

Example

>>> from speechbrain.lobes.models.DiffWave import DiffWave
>>> diffwave = DiffWave(
...     input_channels=80,
...     residual_layers=30,
...     residual_channels=64,
...     dilation_cycle_length=10,
...     total_steps=50,
... )
>>> noisy_audio = torch.randn(1, 1, 25600)
>>> timestep = torch.randint(50, (1,))
>>> input_mel = torch.rand(1, 80, 100)
>>> predicted_noise = diffwave(noisy_audio, timestep, input_mel)
>>> predicted_noise.shape
torch.Size([1, 1, 25600])

forward(audio, diffusion_step, spectrogram=None, length=None)[source]

DiffWave forward function

Parameters:

audio (torch.Tensor) – input gaussian sample [bs, 1, time]
diffusion_step (torch.Tensor) – which timestep of diffusion to execute [bs, 1]
spectrogram (torch.Tensor) – spectrogram data [bs, 80, mel_len]
length (torch.Tensor) – sample lengths - not used - provided for compatibility only

Return type:

predicted noise [bs, 1, time]

diffusion_forward(x, timesteps, cond_emb=None, length=None, out_mask_value=None, latent_mask_value=None)[source]: Forward function suitable for wrapping by diffusion. For this model, out_mask_value/latent_mask_value are unused and discarded. See forward() for details.

class speechbrain.lobes.models.DiffWave.DiffWaveDiffusion(model, timesteps=None, noise=None, beta_start=None, beta_end=None, sample_min=None, sample_max=None, show_progress=False)[source]

Bases: DenoisingDiffusion

An enhanced diffusion implementation with DiffWave-specific inference

Parameters:

model (nn.Module) – the underlying model
timesteps (int) – the total number of timesteps
noise (str|nn.Module) – the type of noise being used “gaussian” will produce standard Gaussian noise
beta_start (float) – the value of the “beta” parameter at the beginning of the process (see DiffWave paper)
beta_end (float) – the value of the “beta” parameter at the end of the process
sample_min (float)
sample_max (float) – Used to clip the output.
show_progress (bool) – whether to show progress during inference

Example

>>> from speechbrain.lobes.models.DiffWave import DiffWave
>>> diffwave = DiffWave(
...     input_channels=80,
...     residual_layers=30,
...     residual_channels=64,
...     dilation_cycle_length=10,
...     total_steps=50,
... )
>>> from speechbrain.lobes.models.DiffWave import DiffWaveDiffusion
>>> from speechbrain.nnet.diffusion import GaussianNoise
>>> diffusion = DiffWaveDiffusion(
...     model=diffwave,
...     beta_start=0.0001,
...     beta_end=0.05,
...     timesteps=50,
...     noise=GaussianNoise,
... )
>>> input_mel = torch.rand(1, 80, 100)
>>> output = diffusion.inference(
...     unconditional=False,
...     scale=256,
...     condition=input_mel,
...     fast_sampling=True,
...     fast_sampling_noise_schedule=[0.0001, 0.001, 0.01, 0.05, 0.2, 0.5],
... )
>>> output.shape
torch.Size([1, 25600])

inference(unconditional, scale, condition=None, fast_sampling=False, fast_sampling_noise_schedule=None, device=None)[source]

Processes the inference for diffwave One inference function for all the locally/globally conditional generation and unconditional generation tasks

Parameters:

unconditional (bool) – do unconditional generation if True, else do conditional generation
scale (int) – scale to get the final output wave length for conditional generation, the output wave length is scale * condition.shape[-1] for example, if the condition is spectrogram (bs, n_mel, time), scale should be hop length for unconditional generation, scale should be the desired audio length
condition (torch.Tensor) – input spectrogram for vocoding or other conditions for other conditional generation, should be None for unconditional generation
fast_sampling (bool) – whether to do fast sampling
fast_sampling_noise_schedule (list) – the noise schedules used for fast sampling
device (str|torch.device) – inference device

Returns:

predicted_sample – the predicted audio (bs, 1, t)

Return type:

torch.Tensor