Spaces:
Running
Running
File size: 3,207 Bytes
0d80816 |
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 |
# Copyright (c) 2023 Amphion.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import torch
import numpy as np
from torch import nn
from torch.nn import functional as F
class Stretch2d(nn.Module):
def __init__(self, x_scale, y_scale, mode="nearest"):
super(Stretch2d, self).__init__()
self.x_scale = x_scale
self.y_scale = y_scale
self.mode = mode
def forward(self, x):
return F.interpolate(
x, scale_factor=(self.y_scale, self.x_scale), mode=self.mode
)
def _get_activation(upsample_activation):
nonlinear = getattr(nn, upsample_activation)
return nonlinear
class UpsampleNetwork(nn.Module):
def __init__(
self,
upsample_scales,
upsample_activation="none",
upsample_activation_params={},
mode="nearest",
freq_axis_kernel_size=1,
cin_pad=0,
cin_channels=128,
):
super(UpsampleNetwork, self).__init__()
self.up_layers = nn.ModuleList()
total_scale = np.prod(upsample_scales)
self.indent = cin_pad * total_scale
for scale in upsample_scales:
freq_axis_padding = (freq_axis_kernel_size - 1) // 2
k_size = (freq_axis_kernel_size, scale * 2 + 1)
padding = (freq_axis_padding, scale)
stretch = Stretch2d(scale, 1, mode)
conv = nn.Conv2d(1, 1, kernel_size=k_size, padding=padding, bias=False)
conv.weight.data.fill_(1.0 / np.prod(k_size))
conv = nn.utils.weight_norm(conv)
self.up_layers.append(stretch)
self.up_layers.append(conv)
if upsample_activation != "none":
nonlinear = _get_activation(upsample_activation)
self.up_layers.append(nonlinear(**upsample_activation_params))
def forward(self, c):
"""
Args:
c : B x C x T
"""
# B x 1 x C x T
c = c.unsqueeze(1)
for f in self.up_layers:
c = f(c)
# B x C x T
c = c.squeeze(1)
if self.indent > 0:
c = c[:, :, self.indent : -self.indent]
return c
class ConvInUpsampleNetwork(nn.Module):
def __init__(
self,
upsample_scales,
upsample_activation="none",
upsample_activation_params={},
mode="nearest",
freq_axis_kernel_size=1,
cin_pad=0,
cin_channels=128,
):
super(ConvInUpsampleNetwork, self).__init__()
# To capture wide-context information in conditional features
# meaningless if cin_pad == 0
ks = 2 * cin_pad + 1
self.conv_in = nn.Conv1d(
cin_channels, cin_channels, kernel_size=ks, padding=cin_pad, bias=False
)
self.upsample = UpsampleNetwork(
upsample_scales,
upsample_activation,
upsample_activation_params,
mode,
freq_axis_kernel_size,
cin_pad=cin_pad,
cin_channels=cin_channels,
)
def forward(self, c):
c_up = self.upsample(self.conv_in(c))
return c_up
|