Spaces:
Running
Running
| # Copyright 2022 Christian J. Steinmetz | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # TCN implementation adapted from: | |
| # https://github.com/csteinmetz1/micro-tcn/blob/main/microtcn/tcn.py | |
| import torch | |
| from argparse import ArgumentParser | |
| from deepafx_st.utils import center_crop, causal_crop | |
| class FiLM(torch.nn.Module): | |
| def __init__(self, num_features, cond_dim): | |
| super().__init__() | |
| self.num_features = num_features | |
| self.bn = torch.nn.BatchNorm1d(num_features, affine=False) | |
| self.adaptor = torch.nn.Linear(cond_dim, num_features * 2) | |
| def forward(self, x, cond): | |
| # project conditioning to 2 x num. conv channels | |
| cond = self.adaptor(cond) | |
| # split the projection into gain and bias | |
| g, b = torch.chunk(cond, 2, dim=-1) | |
| # add virtual channel dim if needed | |
| if g.ndim == 2: | |
| g = g.unsqueeze(1) | |
| b = b.unsqueeze(1) | |
| # reshape for application | |
| g = g.permute(0, 2, 1) | |
| b = b.permute(0, 2, 1) | |
| x = self.bn(x) # apply BatchNorm without affine | |
| x = (x * g) + b # then apply conditional affine | |
| return x | |
| class ConditionalTCNBlock(torch.nn.Module): | |
| def __init__( | |
| self, in_ch, out_ch, cond_dim, kernel_size=3, dilation=1, causal=False, **kwargs | |
| ): | |
| super().__init__() | |
| self.in_ch = in_ch | |
| self.out_ch = out_ch | |
| self.kernel_size = kernel_size | |
| self.dilation = dilation | |
| self.causal = causal | |
| self.conv1 = torch.nn.Conv1d( | |
| in_ch, | |
| out_ch, | |
| kernel_size=kernel_size, | |
| padding=0, | |
| dilation=dilation, | |
| bias=True, | |
| ) | |
| self.film = FiLM(out_ch, cond_dim) | |
| self.relu = torch.nn.PReLU(out_ch) | |
| self.res = torch.nn.Conv1d( | |
| in_ch, out_ch, kernel_size=1, groups=in_ch, bias=False | |
| ) | |
| def forward(self, x, p): | |
| x_in = x | |
| x = self.conv1(x) | |
| x = self.film(x, p) # apply FiLM conditioning | |
| x = self.relu(x) | |
| x_res = self.res(x_in) | |
| if self.causal: | |
| x = x + causal_crop(x_res, x.shape[-1]) | |
| else: | |
| x = x + center_crop(x_res, x.shape[-1]) | |
| return x | |
| class ConditionalTCN(torch.nn.Module): | |
| """Temporal convolutional network with conditioning module. | |
| Args: | |
| sample_rate (float): Audio sample rate. | |
| num_control_params (int, optional): Dimensionality of the conditioning signal. Default: 24 | |
| ninputs (int, optional): Number of input channels (mono = 1, stereo 2). Default: 1 | |
| noutputs (int, optional): Number of output channels (mono = 1, stereo 2). Default: 1 | |
| nblocks (int, optional): Number of total TCN blocks. Default: 10 | |
| kernel_size (int, optional: Width of the convolutional kernels. Default: 3 | |
| dialation_growth (int, optional): Compute the dilation factor at each block as dilation_growth ** (n % stack_size). Default: 1 | |
| channel_growth (int, optional): Compute the output channels at each black as in_ch * channel_growth. Default: 2 | |
| channel_width (int, optional): When channel_growth = 1 all blocks use convolutions with this many channels. Default: 64 | |
| stack_size (int, optional): Number of blocks that constitute a single stack of blocks. Default: 10 | |
| causal (bool, optional): Causal TCN configuration does not consider future input values. Default: False | |
| """ | |
| def __init__( | |
| self, | |
| sample_rate, | |
| num_control_params=24, | |
| ninputs=1, | |
| noutputs=1, | |
| nblocks=10, | |
| kernel_size=15, | |
| dilation_growth=2, | |
| channel_growth=1, | |
| channel_width=64, | |
| stack_size=10, | |
| causal=False, | |
| skip_connections=False, | |
| **kwargs, | |
| ): | |
| super().__init__() | |
| self.num_control_params = num_control_params | |
| self.ninputs = ninputs | |
| self.noutputs = noutputs | |
| self.nblocks = nblocks | |
| self.kernel_size = kernel_size | |
| self.dilation_growth = dilation_growth | |
| self.channel_growth = channel_growth | |
| self.channel_width = channel_width | |
| self.stack_size = stack_size | |
| self.causal = causal | |
| self.skip_connections = skip_connections | |
| self.sample_rate = sample_rate | |
| self.blocks = torch.nn.ModuleList() | |
| for n in range(nblocks): | |
| in_ch = out_ch if n > 0 else ninputs | |
| if self.channel_growth > 1: | |
| out_ch = in_ch * self.channel_growth | |
| else: | |
| out_ch = self.channel_width | |
| dilation = self.dilation_growth ** (n % self.stack_size) | |
| self.blocks.append( | |
| ConditionalTCNBlock( | |
| in_ch, | |
| out_ch, | |
| self.num_control_params, | |
| kernel_size=self.kernel_size, | |
| dilation=dilation, | |
| padding="same" if self.causal else "valid", | |
| causal=self.causal, | |
| ) | |
| ) | |
| self.output = torch.nn.Conv1d(out_ch, noutputs, kernel_size=1) | |
| self.receptive_field = self.compute_receptive_field() | |
| # print( | |
| # f"TCN receptive field: {self.receptive_field} samples", | |
| # f" or {(self.receptive_field/self.sample_rate)*1e3:0.3f} ms", | |
| # ) | |
| def forward(self, x, p, **kwargs): | |
| # causally pad input signal | |
| x = torch.nn.functional.pad(x, (self.receptive_field - 1, 0)) | |
| # iterate over blocks passing conditioning | |
| for idx, block in enumerate(self.blocks): | |
| x = block(x, p) | |
| if self.skip_connections: | |
| if idx == 0: | |
| skips = x | |
| else: | |
| skips = center_crop(skips, x[-1]) + x | |
| else: | |
| skips = 0 | |
| # final 1x1 convolution to collapse channels | |
| out = self.output(x + skips) | |
| return out | |
| def compute_receptive_field(self): | |
| """Compute the receptive field in samples.""" | |
| rf = self.kernel_size | |
| for n in range(1, self.nblocks): | |
| dilation = self.dilation_growth ** (n % self.stack_size) | |
| rf = rf + ((self.kernel_size - 1) * dilation) | |
| return rf | |