Nathan9/xcodec_mini_infer
0
1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the license found in the5# LICENSE file in the root directory of this source tree.6 7"""Encodec SEANet-based encoder and decoder implementation."""8 9import typing as tp10 11import numpy as np12import torch.nn as nn13 14from . import (15 SConv1d,16 SConvTranspose1d,17 SLSTM18)19 20 21class SEANetResnetBlock(nn.Module):22 """Residual block from SEANet model.23 Args:24 dim (int): Dimension of the input/output25 kernel_sizes (list): List of kernel sizes for the convolutions.26 dilations (list): List of dilations for the convolutions.27 activation (str): Activation function.28 activation_params (dict): Parameters to provide to the activation function29 norm (str): Normalization method.30 norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.31 causal (bool): Whether to use fully causal convolution.32 pad_mode (str): Padding mode for the convolutions.33 compress (int): Reduced dimensionality in residual branches (from Demucs v3)34 true_skip (bool): Whether to use true skip connection or a simple convolution as the skip connection.35 """36 def __init__(self, dim: int, kernel_sizes: tp.List[int] = [3, 1], dilations: tp.List[int] = [1, 1],37 activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},38 norm: str = 'weight_norm', norm_params: tp.Dict[str, tp.Any] = {}, causal: bool = False,39 pad_mode: str = 'reflect', compress: int = 2, true_skip: bool = True):40 super().__init__()41 assert len(kernel_sizes) == len(dilations), 'Number of kernel sizes should match number of dilations'42 act = getattr(nn, activation)43 hidden = dim // compress44 block = []45 for i, (kernel_size, dilation) in enumerate(zip(kernel_sizes, dilations)):46 in_chs = dim if i == 0 else hidden47 out_chs = dim if i == len(kernel_sizes) - 1 else hidden48 block += [49 act(**activation_params),50 SConv1d(in_chs, out_chs, kernel_size=kernel_size, dilation=dilation,51 norm=norm, norm_kwargs=norm_params,52 causal=causal, pad_mode=pad_mode),53 ]54 self.block = nn.Sequential(*block)55 self.shortcut: nn.Module56 if true_skip:57 self.shortcut = nn.Identity()58 else:59 self.shortcut = SConv1d(dim, dim, kernel_size=1, norm=norm, norm_kwargs=norm_params,60 causal=causal, pad_mode=pad_mode)61 62 def forward(self, x):63 return self.shortcut(x) + self.block(x)64 65 66class SEANetEncoder(nn.Module):67 """SEANet encoder.68 Args:69 channels (int): Audio channels.70 dimension (int): Intermediate representation dimension.71 n_filters (int): Base width for the model.72 n_residual_layers (int): nb of residual layers.73 ratios (Sequence[int]): kernel size and stride ratios. The encoder uses downsampling ratios instead of74 upsampling ratios, hence it will use the ratios in the reverse order to the ones specified here75 that must match the decoder order76 activation (str): Activation function.77 activation_params (dict): Parameters to provide to the activation function78 norm (str): Normalization method.79 norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.80 kernel_size (int): Kernel size for the initial convolution.81 last_kernel_size (int): Kernel size for the initial convolution.82 residual_kernel_size (int): Kernel size for the residual layers.83 dilation_base (int): How much to increase the dilation with each layer.84 causal (bool): Whether to use fully causal convolution.85 pad_mode (str): Padding mode for the convolutions.86 true_skip (bool): Whether to use true skip connection or a simple87 (streamable) convolution as the skip connection in the residual network blocks.88 compress (int): Reduced dimensionality in residual branches (from Demucs v3).89 lstm (int): Number of LSTM layers at the end of the encoder.90 """91 def __init__(self, channels: int = 1, dimension: int = 128, n_filters: int = 32, n_residual_layers: int = 1,92 ratios: tp.List[int] = [8, 5, 4, 2], activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},93 norm: str = 'weight_norm', norm_params: tp.Dict[str, tp.Any] = {}, kernel_size: int = 7,94 last_kernel_size: int = 7, residual_kernel_size: int = 3, dilation_base: int = 2, causal: bool = False,95 pad_mode: str = 'reflect', true_skip: bool = False, compress: int = 2, lstm: int = 2):96 super().__init__()97 self.channels = channels98 self.dimension = dimension99 self.n_filters = n_filters100 self.ratios = list(reversed(ratios))101 del ratios102 self.n_residual_layers = n_residual_layers103 self.hop_length = np.prod(self.ratios) # 计算乘积104 105 act = getattr(nn, activation)106 mult = 1107 model: tp.List[nn.Module] = [108 SConv1d(channels, mult * n_filters, kernel_size, norm=norm, norm_kwargs=norm_params,109 causal=causal, pad_mode=pad_mode)110 ]111 # Downsample to raw audio scale112 for i, ratio in enumerate(self.ratios):113 # Add residual layers114 for j in range(n_residual_layers):115 model += [116 SEANetResnetBlock(mult * n_filters, kernel_sizes=[residual_kernel_size, 1],117 dilations=[dilation_base ** j, 1],118 norm=norm, norm_params=norm_params,119 activation=activation, activation_params=activation_params,120 causal=causal, pad_mode=pad_mode, compress=compress, true_skip=true_skip)]121 122 # Add downsampling layers123 model += [124 act(**activation_params),125 SConv1d(mult * n_filters, mult * n_filters * 2,126 kernel_size=ratio * 2, stride=ratio,127 norm=norm, norm_kwargs=norm_params,128 causal=causal, pad_mode=pad_mode),129 ]130 mult *= 2131 132 if lstm:133 model += [SLSTM(mult * n_filters, num_layers=lstm)]134 135 model += [136 act(**activation_params),137 SConv1d(mult * n_filters, dimension, last_kernel_size, norm=norm, norm_kwargs=norm_params,138 causal=causal, pad_mode=pad_mode)139 ]140 141 self.model = nn.Sequential(*model)142 143 def forward(self, x):144 return self.model(x)145 146 147class SEANetDecoder(nn.Module):148 """SEANet decoder.149 Args:150 channels (int): Audio channels.151 dimension (int): Intermediate representation dimension.152 n_filters (int): Base width for the model.153 n_residual_layers (int): nb of residual layers.154 ratios (Sequence[int]): kernel size and stride ratios155 activation (str): Activation function.156 activation_params (dict): Parameters to provide to the activation function157 final_activation (str): Final activation function after all convolutions.158 final_activation_params (dict): Parameters to provide to the activation function159 norm (str): Normalization method.160 norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.161 kernel_size (int): Kernel size for the initial convolution.162 last_kernel_size (int): Kernel size for the initial convolution.163 residual_kernel_size (int): Kernel size for the residual layers.164 dilation_base (int): How much to increase the dilation with each layer.165 causal (bool): Whether to use fully causal convolution.166 pad_mode (str): Padding mode for the convolutions.167 true_skip (bool): Whether to use true skip connection or a simple168 (streamable) convolution as the skip connection in the residual network blocks.169 compress (int): Reduced dimensionality in residual branches (from Demucs v3).170 lstm (int): Number of LSTM layers at the end of the encoder.171 trim_right_ratio (float): Ratio for trimming at the right of the transposed convolution under the causal setup.172 If equal to 1.0, it means that all the trimming is done at the right.173 """174 def __init__(self, channels: int = 1, dimension: int = 128, n_filters: int = 32, n_residual_layers: int = 1,175 ratios: tp.List[int] = [8, 5, 4, 2], activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},176 final_activation: tp.Optional[str] = None, final_activation_params: tp.Optional[dict] = None,177 norm: str = 'weight_norm', norm_params: tp.Dict[str, tp.Any] = {}, kernel_size: int = 7,178 last_kernel_size: int = 7, residual_kernel_size: int = 3, dilation_base: int = 2, causal: bool = False,179 pad_mode: str = 'reflect', true_skip: bool = False, compress: int = 2, lstm: int = 2,180 trim_right_ratio: float = 1.0):181 super().__init__()182 self.dimension = dimension183 self.channels = channels184 self.n_filters = n_filters185 self.ratios = ratios186 del ratios187 self.n_residual_layers = n_residual_layers188 self.hop_length = np.prod(self.ratios)189 190 act = getattr(nn, activation)191 mult = int(2 ** len(self.ratios))192 model: tp.List[nn.Module] = [193 SConv1d(dimension, mult * n_filters, kernel_size, norm=norm, norm_kwargs=norm_params,194 causal=causal, pad_mode=pad_mode)195 ]196 197 if lstm:198 model += [SLSTM(mult * n_filters, num_layers=lstm)]199 200 # Upsample to raw audio scale201 for i, ratio in enumerate(self.ratios):202 # Add upsampling layers203 model += [204 act(**activation_params),205 SConvTranspose1d(mult * n_filters, mult * n_filters // 2,206 kernel_size=ratio * 2, stride=ratio,207 norm=norm, norm_kwargs=norm_params,208 causal=causal, trim_right_ratio=trim_right_ratio),209 ]210 # Add residual layers211 for j in range(n_residual_layers):212 model += [213 SEANetResnetBlock(mult * n_filters // 2, kernel_sizes=[residual_kernel_size, 1],214 dilations=[dilation_base ** j, 1],215 activation=activation, activation_params=activation_params,216 norm=norm, norm_params=norm_params, causal=causal,217 pad_mode=pad_mode, compress=compress, true_skip=true_skip)]218 219 mult //= 2220 221 # Add final layers222 model += [223 act(**activation_params),224 SConv1d(n_filters, channels, last_kernel_size, norm=norm, norm_kwargs=norm_params,225 causal=causal, pad_mode=pad_mode)226 ]227 # Add optional final activation to decoder (eg. tanh)228 if final_activation is not None:229 final_act = getattr(nn, final_activation)230 final_activation_params = final_activation_params or {}231 model += [232 final_act(**final_activation_params)233 ]234 self.model = nn.Sequential(*model)235 236 def forward(self, z):237 y = self.model(z)238 return y239 240 241def test():242 import torch243 encoder = SEANetEncoder()244 decoder = SEANetDecoder()245 x = torch.randn(1, 1, 24000)246 z = encoder(x)247 print('z ', z.shape)248 assert 1==2249 assert list(z.shape) == [1, 128, 75], z.shape250 y = decoder(z)251 assert y.shape == x.shape, (x.shape, y.shape)252 253 254if __name__ == '__main__':255 test()256