103 lines
2.9 KiB
Python
103 lines
2.9 KiB
Python
# 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 math
|
|
|
|
import torch.nn as nn
|
|
|
|
from modules.general.utils import Conv1d, zero_module
|
|
from .residual_block import ResidualBlock
|
|
|
|
|
|
class BiDilConv(nn.Module):
|
|
r"""Dilated CNN architecture with residual connections, default diffusion decoder.
|
|
|
|
Args:
|
|
input_channel: The number of input channels.
|
|
base_channel: The number of base channels.
|
|
n_res_block: The number of residual blocks.
|
|
conv_kernel_size: The kernel size of convolutional layers.
|
|
dilation_cycle_length: The cycle length of dilation.
|
|
conditioner_size: The size of conditioner.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
input_channel,
|
|
base_channel,
|
|
n_res_block,
|
|
conv_kernel_size,
|
|
dilation_cycle_length,
|
|
conditioner_size,
|
|
output_channel: int = -1,
|
|
):
|
|
super().__init__()
|
|
|
|
self.input_channel = input_channel
|
|
self.base_channel = base_channel
|
|
self.n_res_block = n_res_block
|
|
self.conv_kernel_size = conv_kernel_size
|
|
self.dilation_cycle_length = dilation_cycle_length
|
|
self.conditioner_size = conditioner_size
|
|
self.output_channel = output_channel if output_channel > 0 else input_channel
|
|
|
|
self.input = nn.Sequential(
|
|
Conv1d(
|
|
input_channel,
|
|
base_channel,
|
|
1,
|
|
),
|
|
nn.ReLU(),
|
|
)
|
|
|
|
self.residual_blocks = nn.ModuleList(
|
|
[
|
|
ResidualBlock(
|
|
channels=base_channel,
|
|
kernel_size=conv_kernel_size,
|
|
dilation=2 ** (i % dilation_cycle_length),
|
|
d_context=conditioner_size,
|
|
)
|
|
for i in range(n_res_block)
|
|
]
|
|
)
|
|
|
|
self.out_proj = nn.Sequential(
|
|
Conv1d(
|
|
base_channel,
|
|
base_channel,
|
|
1,
|
|
),
|
|
nn.ReLU(),
|
|
zero_module(
|
|
Conv1d(
|
|
base_channel,
|
|
self.output_channel,
|
|
1,
|
|
),
|
|
),
|
|
)
|
|
|
|
def forward(self, x, y, context=None):
|
|
"""
|
|
Args:
|
|
x: Noisy mel-spectrogram [B x ``n_mel`` x L]
|
|
y: FILM embeddings with the shape of (B, ``base_channel``)
|
|
context: Context with the shape of [B x ``d_context`` x L], default to None.
|
|
"""
|
|
|
|
h = self.input(x)
|
|
|
|
skip = None
|
|
for i in range(self.n_res_block):
|
|
h, skip_connection = self.residual_blocks[i](h, y, context)
|
|
skip = skip_connection if skip is None else skip_connection + skip
|
|
|
|
out = skip / math.sqrt(self.n_res_block)
|
|
|
|
out = self.out_proj(out)
|
|
|
|
return out
|