Spaces:
Sleeping
Sleeping
File size: 4,079 Bytes
c42fe7e |
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 |
import math
from math import sqrt
import torch
import torch.nn as nn
import torch.nn.functional as F
from modules.commons.common_layers import SinusoidalPosEmb
from utils.hparams import hparams
class Conv1d(torch.nn.Conv1d):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
nn.init.kaiming_normal_(self.weight)
class ResidualBlock(nn.Module):
def __init__(self, encoder_hidden, residual_channels, dilation):
super().__init__()
self.residual_channels = residual_channels
self.dilated_conv = nn.Conv1d(
residual_channels,
2 * residual_channels,
kernel_size=3,
padding=dilation,
dilation=dilation
)
self.diffusion_projection = nn.Linear(residual_channels, residual_channels)
self.conditioner_projection = nn.Conv1d(encoder_hidden, 2 * residual_channels, 1)
self.output_projection = nn.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
# Using torch.split instead of torch.chunk to avoid using onnx::Slice
gate, filter = torch.split(y, [self.residual_channels, self.residual_channels], dim=1)
y = torch.sigmoid(gate) * torch.tanh(filter)
y = self.output_projection(y)
# Using torch.split instead of torch.chunk to avoid using onnx::Slice
residual, skip = torch.split(y, [self.residual_channels, self.residual_channels], dim=1)
return (x + residual) / math.sqrt(2.0), skip
class WaveNet(nn.Module):
def __init__(self, in_dims, n_feats, *, num_layers=20, num_channels=256, dilation_cycle_length=4):
super().__init__()
self.in_dims = in_dims
self.n_feats = n_feats
self.input_projection = Conv1d(in_dims * n_feats, num_channels, 1)
self.diffusion_embedding = SinusoidalPosEmb(num_channels)
self.mlp = nn.Sequential(
nn.Linear(num_channels, num_channels * 4),
nn.Mish(),
nn.Linear(num_channels * 4, num_channels)
)
self.residual_layers = nn.ModuleList([
ResidualBlock(
encoder_hidden=hparams['hidden_size'],
residual_channels=num_channels,
dilation=2 ** (i % dilation_cycle_length)
)
for i in range(num_layers)
])
self.skip_projection = Conv1d(num_channels, num_channels, 1)
self.output_projection = Conv1d(num_channels, in_dims * n_feats, 1)
nn.init.zeros_(self.output_projection.weight)
def forward(self, spec, diffusion_step, cond):
"""
:param spec: [B, F, M, T]
:param diffusion_step: [B, 1]
:param cond: [B, H, T]
:return:
"""
if self.n_feats == 1:
x = spec.squeeze(1) # [B, M, T]
else:
x = spec.flatten(start_dim=1, end_dim=2) # [B, F x M, T]
x = self.input_projection(x) # [B, C, T]
x = F.relu(x)
diffusion_step = self.diffusion_embedding(diffusion_step)
diffusion_step = self.mlp(diffusion_step)
skip = []
for layer in 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) # [B, M, T]
if self.n_feats == 1:
x = x[:, None, :, :]
else:
# This is the temporary solution since PyTorch 1.13
# does not support exporting aten::unflatten to ONNX
# x = x.unflatten(dim=1, sizes=(self.n_feats, self.in_dims))
x = x.reshape(-1, self.n_feats, self.in_dims, x.shape[2])
return x
|