File size: 6,055 Bytes
8c6b65f | 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 | import torch
import torch.nn as nn
from torch.nn.utils.parametrizations import weight_norm
from typing import List, Tuple, Any
class ResBlock1D(nn.Module):
"""Residual block with dilated 1D convolutions."""
def __init__(
self,
channels: int,
kernel_size: int = 3,
dilations: List[int] = [1, 3, 5]
):
super().__init__()
self.convs1 = nn.ModuleList()
self.convs2 = nn.ModuleList()
for d in dilations:
# Padding is adjusted to keep sequence length unchanged: (kernel_size - 1) * dilation // 2
padding1 = (kernel_size - 1) * d // 2
padding2 = (kernel_size - 1) // 2
self.convs1.append(
weight_norm(nn.Conv1d(
channels, channels, kernel_size,
stride=1, padding=padding1, dilation=d
))
)
self.convs2.append(
weight_norm(nn.Conv1d(
channels, channels, kernel_size,
stride=1, padding=padding2, dilation=1
))
)
self.activation = nn.LeakyReLU(0.1)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x: Input tensor of shape (batch, channels, time_steps).
Returns:
Output tensor of shape (batch, channels, time_steps).
"""
for c1, c2 in zip(self.convs1, self.convs2):
residual = x
x = self.activation(x)
x = c1(x)
x = self.activation(x)
x = c2(x)
x = x + residual
return x
class WaveformEncoder(nn.Module):
"""Waveform encoder that downsamples a time-domain signal into a latent space."""
def __init__(self, config_enc: Any):
super().__init__()
in_channels = config_enc.in_channels
base_channels = config_enc.channels
strides = config_enc.strides
kernel_sizes = config_enc.kernel_sizes
dilations = config_enc.dilations
# Initial convolution
self.conv_in = weight_norm(nn.Conv1d(in_channels, base_channels, kernel_size=7, stride=1, padding=3))
# Progressive downsampling blocks
self.down_blocks = nn.ModuleList()
curr_channels = base_channels
for stride, k_size in zip(strides, kernel_sizes):
next_channels = curr_channels * 2
self.down_blocks.append(
nn.Sequential(
nn.LeakyReLU(0.1),
weight_norm(nn.Conv1d(
curr_channels, next_channels, kernel_size=k_size,
stride=stride, padding=(k_size - 1) // 2
)),
ResBlock1D(next_channels, kernel_size=3, dilations=dilations)
)
)
curr_channels = next_channels
self.out_channels = curr_channels
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x: Input waveform tensor of shape (batch, 1, samples).
Returns:
Encoded features of shape (batch, out_channels, frames).
"""
x = self.conv_in(x)
for block in self.down_blocks:
x = block(x)
return x
class WaveformDecoder(nn.Module):
"""Waveform decoder that upsamples latent representations back into a time-domain waveform."""
def __init__(self, config_dec: Any, in_channels: int):
super().__init__()
upsample_rates = config_dec.upsample_rates
upsample_kernel_sizes = config_dec.upsample_kernel_sizes
base_channels = config_dec.channels
resblock_kernel_sizes = config_dec.resblock_kernel_sizes
resblock_dilations = config_dec.resblock_dilations
self.conv_in = weight_norm(nn.Conv1d(in_channels, base_channels * (2 ** len(upsample_rates)), kernel_size=3, stride=1, padding=1))
# Progressive upsampling blocks
self.up_blocks = nn.ModuleList()
curr_channels = base_channels * (2 ** len(upsample_rates))
for rate, k_size in zip(upsample_rates, upsample_kernel_sizes):
next_channels = curr_channels // 2
# Upsample block consists of a ConvTranspose1d followed by parallel ResBlocks
up_conv = weight_norm(nn.ConvTranspose1d(
curr_channels, next_channels, kernel_size=k_size,
stride=rate, padding=(k_size - rate) // 2
))
res_blocks = nn.ModuleList()
for r_k_size, r_dilations in zip(resblock_kernel_sizes, resblock_dilations):
res_blocks.append(ResBlock1D(next_channels, kernel_size=r_k_size, dilations=r_dilations))
self.up_blocks.append(nn.ModuleDict({
"upsample": up_conv,
"resblocks": res_blocks
}))
curr_channels = next_channels
self.activation = nn.LeakyReLU(0.1)
self.conv_out = weight_norm(nn.Conv1d(curr_channels, 1, kernel_size=7, stride=1, padding=3))
self.tanh = nn.Tanh()
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x: Latent features of shape (batch, in_channels, frames).
Returns:
Reconstructed waveform of shape (batch, 1, samples).
"""
x = self.conv_in(x)
for block in self.up_blocks:
x = self.activation(x)
x = block["upsample"](x)
# Sum predictions from multiple parallel residual blocks
res_out = torch.zeros_like(x)
for resblock in block["resblocks"]:
res_out = res_out + resblock(x)
x = res_out / len(block["resblocks"])
x = self.activation(x)
x = self.conv_out(x)
x = self.tanh(x)
return x
|