| | import math |
| |
|
| | import torch |
| | import torch.nn as nn |
| | import torch.nn.functional as F |
| |
|
| | from math import sqrt |
| |
|
| | from .diffusion import Mish |
| | from utils.hparams import hparams |
| |
|
| | Linear = nn.Linear |
| | ConvTranspose2d = nn.ConvTranspose2d |
| |
|
| |
|
| | class AttrDict(dict): |
| | def __init__(self, *args, **kwargs): |
| | super(AttrDict, self).__init__(*args, **kwargs) |
| | self.__dict__ = self |
| |
|
| | def override(self, attrs): |
| | if isinstance(attrs, dict): |
| | self.__dict__.update(**attrs) |
| | elif isinstance(attrs, (list, tuple, set)): |
| | for attr in attrs: |
| | self.override(attr) |
| | elif attrs is not None: |
| | raise NotImplementedError |
| | return self |
| |
|
| |
|
| | class SinusoidalPosEmb(nn.Module): |
| | def __init__(self, dim): |
| | super().__init__() |
| | self.dim = dim |
| |
|
| | def forward(self, x): |
| | device = x.device |
| | half_dim = self.dim // 2 |
| | emb = math.log(10000) / (half_dim - 1) |
| | emb = torch.exp(torch.arange(half_dim, device=device) * -emb) |
| | emb = x[:, None] * emb[None, :] |
| | emb = torch.cat((emb.sin(), emb.cos()), dim=-1) |
| | return emb |
| |
|
| |
|
| | def Conv1d(*args, **kwargs): |
| | layer = nn.Conv1d(*args, **kwargs) |
| | nn.init.kaiming_normal_(layer.weight) |
| | return layer |
| |
|
| |
|
| | @torch.jit.script |
| | def silu(x): |
| | return x * torch.sigmoid(x) |
| |
|
| |
|
| | class ResidualBlock(nn.Module): |
| | def __init__(self, encoder_hidden, residual_channels, dilation): |
| | super().__init__() |
| | self.dilated_conv = Conv1d(residual_channels, 2 * residual_channels, 3, padding=dilation, dilation=dilation) |
| | self.diffusion_projection = Linear(residual_channels, residual_channels) |
| | self.conditioner_projection = Conv1d(encoder_hidden, 2 * residual_channels, 1) |
| | self.output_projection = Conv1d(residual_channels, 2 * residual_channels, 1) |
| |
|
| | def forward(self, x, conditioner, diffusion_step): |
| | diffusion_step = self.diffusion_projection(diffusion_step).unsqueeze(-1) |
| | conditioner = self.conditioner_projection(conditioner) |
| | y = x + diffusion_step |
| |
|
| | y = self.dilated_conv(y) + conditioner |
| |
|
| | gate, filter = torch.chunk(y, 2, dim=1) |
| | y = torch.sigmoid(gate) * torch.tanh(filter) |
| |
|
| | y = self.output_projection(y) |
| | residual, skip = torch.chunk(y, 2, dim=1) |
| | return (x + residual) / sqrt(2.0), skip |
| |
|
| |
|
| | class DiffNet(nn.Module): |
| | def __init__(self, in_dims=80): |
| | super().__init__() |
| | self.params = params = AttrDict( |
| | |
| | encoder_hidden=hparams['hidden_size'], |
| | residual_layers=hparams['residual_layers'], |
| | residual_channels=hparams['residual_channels'], |
| | dilation_cycle_length=hparams['dilation_cycle_length'], |
| | ) |
| | self.input_projection = Conv1d(in_dims, params.residual_channels, 1) |
| | self.diffusion_embedding = SinusoidalPosEmb(params.residual_channels) |
| | dim = params.residual_channels |
| | self.mlp = nn.Sequential( |
| | nn.Linear(dim, dim * 4), |
| | Mish(), |
| | nn.Linear(dim * 4, dim) |
| | ) |
| | self.residual_layers = nn.ModuleList([ |
| | ResidualBlock(params.encoder_hidden, params.residual_channels, 2 ** (i % params.dilation_cycle_length)) |
| | for i in range(params.residual_layers) |
| | ]) |
| | self.skip_projection = Conv1d(params.residual_channels, params.residual_channels, 1) |
| | self.output_projection = Conv1d(params.residual_channels, in_dims, 1) |
| | nn.init.zeros_(self.output_projection.weight) |
| |
|
| | def forward(self, spec, diffusion_step, cond): |
| | """ |
| | |
| | :param spec: [B, 1, M, T] |
| | :param diffusion_step: [B, 1] |
| | :param cond: [B, M, T] |
| | :return: |
| | """ |
| | x = spec[:, 0] |
| | x = self.input_projection(x) |
| |
|
| | x = F.relu(x) |
| | diffusion_step = self.diffusion_embedding(diffusion_step) |
| | diffusion_step = self.mlp(diffusion_step) |
| | skip = [] |
| | for layer_id, layer in enumerate(self.residual_layers): |
| | x, skip_connection = layer(x, cond, diffusion_step) |
| | skip.append(skip_connection) |
| |
|
| | x = torch.sum(torch.stack(skip), dim=0) / sqrt(len(self.residual_layers)) |
| | x = self.skip_projection(x) |
| | x = F.relu(x) |
| | x = self.output_projection(x) |
| | return x[:, None, :, :] |
| |
|
