Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from ..commons.layers import LayerNorm, Embedding | |
| from ..commons.conv import ConvBlocks, ResidualBlock, get_norm_builder, get_act_builder | |
| class UnetDown(nn.Module): | |
| def __init__(self, hidden_size, n_layers, kernel_size, down_rates, channel_multiples=None, dropout=0.0, | |
| is_BTC=True, constant_channels=False): | |
| super(UnetDown, self).__init__() | |
| assert n_layers == len(down_rates) # downs, down sample rate | |
| down_rates = [int(i) for i in down_rates] | |
| self.n_layers = n_layers | |
| self.hidden_size = hidden_size | |
| self.is_BTC = is_BTC | |
| channel_multiples = channel_multiples if channel_multiples is not None else down_rates | |
| self.layers = nn.ModuleList() | |
| self.downs = nn.ModuleList() | |
| in_channels = hidden_size | |
| for i in range(self.n_layers): | |
| out_channels = int(in_channels * channel_multiples[i]) if not constant_channels else in_channels | |
| self.layers.append(nn.Sequential( | |
| ResidualBlock(in_channels, kernel_size, dilation=1, n=1, norm_type='ln', dropout=dropout, | |
| c_multiple=1, ln_eps=1e-5, act_type='leakyrelu'), | |
| nn.Conv1d(in_channels, out_channels, kernel_size, padding=(kernel_size - 1) // 2), | |
| ResidualBlock(out_channels, kernel_size, dilation=1, n=1, norm_type='ln', | |
| dropout=dropout, c_multiple=1, ln_eps=1e-5, act_type='leakyrelu') | |
| )) | |
| self.downs.append(nn.Sequential( | |
| nn.AvgPool1d(down_rates[i]) | |
| )) | |
| in_channels = out_channels | |
| self.last_norm = get_norm_builder('ln', out_channels)() | |
| self.post_net = nn.Conv1d(out_channels, out_channels, kernel_size=kernel_size, | |
| padding=kernel_size // 2) | |
| def forward(self, x, **kwargs): | |
| # x [B, T, C] | |
| if self.is_BTC: | |
| x = x.transpose(1, 2) # [B, C, T] | |
| skip_xs = [] | |
| for i in range(self.n_layers): | |
| skip_x = self.layers[i](x) | |
| x = self.downs[i](skip_x) | |
| if self.is_BTC: | |
| skip_xs.append(skip_x.transpose(1, 2)) # [B, T, C] | |
| else: | |
| skip_xs.append(skip_x) | |
| x = self.post_net(self.last_norm(x)) | |
| if self.is_BTC: | |
| x = x.transpose(1, 2) | |
| return x, skip_xs | |
| class UnetMid(nn.Module): | |
| def __init__(self, hidden_size, kernel_size, n_layers=None, in_dims=None, out_dims=None, | |
| dropout=0.0, is_BTC=True, net=None): | |
| super(UnetMid, self).__init__() | |
| in_dims = in_dims if in_dims is not None else hidden_size | |
| out_dims = out_dims if out_dims is not None else hidden_size | |
| self.pre = nn.Conv1d(in_dims, hidden_size, kernel_size, padding=kernel_size // 2) | |
| self.post = nn.Conv1d(hidden_size, out_dims, kernel_size, padding=kernel_size // 2) | |
| self.is_BTC = is_BTC | |
| if net is not None: | |
| self.net = net | |
| else: | |
| self.net = ConvBlocks(hidden_size, out_dims=hidden_size, dilations=None, kernel_size=kernel_size, | |
| layers_in_block=2, c_multiple=2, dropout=dropout, num_layers=n_layers, | |
| post_net_kernel=3, act_type='leakyrelu', is_BTC=is_BTC) | |
| def forward(self, x, cond=None, **kwargs): | |
| # x [B, T, C] | |
| if self.is_BTC: | |
| x = self.pre(x.transpose(1, 2)).transpose(1, 2) | |
| else: | |
| x = self.pre(x) | |
| if cond is None: | |
| cond = 0 | |
| x = self.net(x + cond) | |
| if self.is_BTC: | |
| x = self.post(x.transpose(1, 2)).transpose(1, 2) | |
| else: | |
| x = self.post(x) | |
| return x | |
| class UnetUp(nn.Module): | |
| def __init__(self, hidden_size, n_layers, kernel_size, up_rates, channel_multiples=None, dropout=0.0, | |
| is_BTC=True, constant_channels=False, use_skip_layer=False, skip_scale=1.0): | |
| super(UnetUp, self).__init__() | |
| assert n_layers == len(up_rates) # this is reversed in up module, from the output to the interface with middle | |
| up_rates = [int(i) for i in up_rates] | |
| self.n_layers = n_layers | |
| self.hidden_size = hidden_size | |
| self.is_BTC = is_BTC | |
| self.skip_scale = skip_scale | |
| channel_multiples = channel_multiples if channel_multiples is not None else up_rates | |
| # in_channels = int(np.cumprod(channel_multiples)[-1] * hidden_size) if not constant_channels else hidden_size | |
| self.in_channels_lst = (np.cumprod([1] + channel_multiples) * hidden_size).astype(int) if not constant_channels \ | |
| else [hidden_size for _ in range(self.n_layers + 1)] | |
| in_channels = self.in_channels_lst[-1] | |
| self.ups = nn.ModuleList() | |
| self.skip_layers = nn.ModuleList() | |
| self.layers = nn.ModuleList() | |
| for i in range(self.n_layers-1, -1, -1): | |
| out_channels = self.in_channels_lst[i] if not constant_channels else in_channels | |
| self.ups.append(nn.Sequential( | |
| nn.ConvTranspose1d(in_channels, in_channels, kernel_size=kernel_size, stride=up_rates[i], | |
| padding=kernel_size//2, output_padding=up_rates[i]-1), | |
| get_norm_builder('ln', in_channels)(), | |
| get_act_builder('leakyrelu')() | |
| )) | |
| self.layers.append(nn.Sequential( | |
| # ResidualBlock(in_channels*2, kernel_size, dilation=1, n=1, norm_type='ln', dropout=dropout, | |
| # c_multiple=1, ln_eps=1e-5, act_type='leakyrelu'), | |
| nn.Conv1d(in_channels*2, out_channels, kernel_size, padding=(kernel_size - 1) // 2), | |
| ResidualBlock(out_channels, kernel_size, dilation=1, n=1, norm_type='ln', | |
| dropout=dropout, c_multiple=1, ln_eps=1e-5, act_type='leakyrelu') | |
| )) | |
| if use_skip_layer: | |
| self.skip_layers.append( | |
| ResidualBlock(in_channels, kernel_size, dilation=1, n=1, norm_type='ln', dropout=dropout, | |
| c_multiple=1, ln_eps=1e-5, act_type='leakyrelu') | |
| ) | |
| else: | |
| self.skip_layers.append(nn.Identity()) | |
| in_channels = out_channels | |
| self.out_channels = out_channels | |
| self.last_norm = get_norm_builder('ln', out_channels)() | |
| self.post_net = nn.Conv1d(out_channels, out_channels, kernel_size=kernel_size, | |
| padding=kernel_size // 2) | |
| def forward(self, x, skips, **kwargs): | |
| # x [B, T, C] | |
| if self.is_BTC: | |
| x = x.transpose(1, 2) # [B, C, T] | |
| for i in range(self.n_layers): | |
| x = self.ups[i](x) | |
| skip_x = skips[self.n_layers - i - 1] if not self.is_BTC \ | |
| else skips[self.n_layers - i - 1].transpose(1, 2) # [B, T, C] -> [B, C, T] | |
| skip_x = self.skip_layers[i](skip_x) * self.skip_scale | |
| x = torch.cat((x, skip_x), dim=1) # [B, C, T] | |
| x = self.layers[i](x) | |
| x = self.post_net(self.last_norm(x)) | |
| if self.is_BTC: | |
| x = x.transpose(1, 2) | |
| return x | |
| class Unet(nn.Module): | |
| def __init__(self, hidden_size, down_layers, up_layers, kernel_size, | |
| updown_rates, mid_layers=None, channel_multiples=None, dropout=0.0, | |
| is_BTC=True, constant_channels=False, mid_net=None, use_skip_layer=False, skip_scale=1.0): | |
| super(Unet, self).__init__() | |
| assert len(updown_rates) == down_layers == up_layers, f"{len(updown_rates)}, {down_layers}, {up_layers}" | |
| if channel_multiples is not None: | |
| assert len(channel_multiples) == len(updown_rates) | |
| else: | |
| channel_multiples = updown_rates | |
| self.down = UnetDown(hidden_size, down_layers, kernel_size, updown_rates, | |
| channel_multiples, dropout, is_BTC, constant_channels) | |
| down_out_dims = int(np.cumprod(channel_multiples)[-1] * hidden_size) if not constant_channels else hidden_size | |
| self.mid = UnetMid(hidden_size, kernel_size, mid_layers, | |
| in_dims=down_out_dims, out_dims=down_out_dims, dropout=dropout, is_BTC=is_BTC, net=mid_net) | |
| self.up = UnetUp(hidden_size, up_layers, kernel_size, updown_rates, | |
| channel_multiples, dropout, is_BTC, constant_channels, use_skip_layer, skip_scale) | |
| def forward(self, x, mid_cond=None, **kwargs): | |
| x, skips = self.down(x) | |
| x = self.mid(x, mid_cond) | |
| x = self.up(x, skips) | |
| return x | |