Spaces:
Build error
Build error
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| def init_weights_func(m): | |
| classname = m.__class__.__name__ | |
| if classname.find("Conv1d") != -1: | |
| torch.nn.init.xavier_uniform_(m.weight) | |
| class LambdaLayer(nn.Module): | |
| def __init__(self, lambd): | |
| super(LambdaLayer, self).__init__() | |
| self.lambd = lambd | |
| def forward(self, x): | |
| return self.lambd(x) | |
| class LayerNorm(torch.nn.LayerNorm): | |
| """Layer normalization module. | |
| :param int nout: output dim size | |
| :param int dim: dimension to be normalized | |
| """ | |
| def __init__(self, nout, dim=-1, eps=1e-5): | |
| """Construct an LayerNorm object.""" | |
| super(LayerNorm, self).__init__(nout, eps=eps) | |
| self.dim = dim | |
| def forward(self, x): | |
| """Apply layer normalization. | |
| :param torch.Tensor x: input tensor | |
| :return: layer normalized tensor | |
| :rtype torch.Tensor | |
| """ | |
| if self.dim == -1: | |
| return super(LayerNorm, self).forward(x) | |
| return super(LayerNorm, self).forward(x.transpose(1, -1)).transpose(1, -1) | |
| class ResidualBlock(nn.Module): | |
| """Implements conv->PReLU->norm n-times""" | |
| def __init__(self, channels, kernel_size, dilation, n=2, norm_type='bn', dropout=0.0, | |
| c_multiple=2, ln_eps=1e-12, bias=False): | |
| super(ResidualBlock, self).__init__() | |
| if norm_type == 'bn': | |
| norm_builder = lambda: nn.BatchNorm1d(channels) | |
| elif norm_type == 'in': | |
| norm_builder = lambda: nn.InstanceNorm1d(channels, affine=True) | |
| elif norm_type == 'gn': | |
| norm_builder = lambda: nn.GroupNorm(8, channels) | |
| elif norm_type == 'ln': | |
| norm_builder = lambda: LayerNorm(channels, dim=1, eps=ln_eps) | |
| else: | |
| norm_builder = lambda: nn.Identity() | |
| self.blocks = [ | |
| nn.Sequential( | |
| norm_builder(), | |
| nn.Conv1d(channels, c_multiple * channels, kernel_size, dilation=dilation, | |
| padding=(dilation * (kernel_size - 1)) // 2, bias=bias), | |
| LambdaLayer(lambda x: x * kernel_size ** -0.5), | |
| nn.GELU(), | |
| nn.Conv1d(c_multiple * channels, channels, 1, dilation=dilation, bias=bias), | |
| ) | |
| for _ in range(n) | |
| ] | |
| self.blocks = nn.ModuleList(self.blocks) | |
| self.dropout = dropout | |
| def forward(self, x): | |
| nonpadding = (x.abs().sum(1) > 0).float()[:, None, :] | |
| for b in self.blocks: | |
| x_ = b(x) | |
| if self.dropout > 0 and self.training: | |
| x_ = F.dropout(x_, self.dropout, training=self.training) | |
| x = x + x_ | |
| x = x * nonpadding | |
| return x | |
| class ConvBlocks(nn.Module): | |
| """Decodes the expanded phoneme encoding into spectrograms""" | |
| def __init__(self, channels, out_dims, dilations, kernel_size, | |
| norm_type='ln', layers_in_block=2, c_multiple=2, | |
| dropout=0.0, ln_eps=1e-5, init_weights=True, is_BTC=True, bias=False): | |
| super(ConvBlocks, self).__init__() | |
| self.is_BTC = is_BTC | |
| self.res_blocks = nn.Sequential( | |
| *[ResidualBlock(channels, kernel_size, d, | |
| n=layers_in_block, norm_type=norm_type, c_multiple=c_multiple, | |
| dropout=dropout, ln_eps=ln_eps, bias=bias) | |
| for d in dilations], | |
| ) | |
| if norm_type == 'bn': | |
| norm = nn.BatchNorm1d(channels) | |
| elif norm_type == 'in': | |
| norm = nn.InstanceNorm1d(channels, affine=True) | |
| elif norm_type == 'gn': | |
| norm = nn.GroupNorm(8, channels) | |
| elif norm_type == 'ln': | |
| norm = LayerNorm(channels, dim=1, eps=ln_eps) | |
| self.last_norm = norm | |
| self.post_net1 = nn.Conv1d(channels, out_dims, kernel_size=3, padding=1, bias=bias) | |
| if init_weights: | |
| self.apply(init_weights_func) | |
| def forward(self, x): | |
| """ | |
| :param x: [B, T, H] | |
| :return: [B, T, H] | |
| """ | |
| if self.is_BTC: | |
| x = x.transpose(1, 2) # [B, C, T] | |
| nonpadding = (x.abs().sum(1) > 0).float()[:, None, :] | |
| x = self.res_blocks(x) * nonpadding | |
| x = self.last_norm(x) * nonpadding | |
| x = self.post_net1(x) * nonpadding | |
| if self.is_BTC: | |
| x = x.transpose(1, 2) | |
| return x | |
| class SeqLevelConvolutionalModel(nn.Module): | |
| def __init__(self, out_dim=64, dropout=0.5, audio_feat_type='ppg', backbone_type='unet', norm_type='bn'): | |
| nn.Module.__init__(self) | |
| self.audio_feat_type = audio_feat_type | |
| if audio_feat_type == 'ppg': | |
| self.audio_encoder = nn.Sequential(*[ | |
| nn.Conv1d(29, 48, 3, 1, 1, bias=False), | |
| nn.BatchNorm1d(48) if norm_type=='bn' else LayerNorm(48, dim=1), | |
| nn.GELU(), | |
| nn.Conv1d(48, 48, 3, 1, 1, bias=False) | |
| ]) | |
| self.energy_encoder = nn.Sequential(*[ | |
| nn.Conv1d(1, 16, 3, 1, 1, bias=False), | |
| nn.BatchNorm1d(16) if norm_type=='bn' else LayerNorm(16, dim=1), | |
| nn.GELU(), | |
| nn.Conv1d(16, 16, 3, 1, 1, bias=False) | |
| ]) | |
| elif audio_feat_type == 'mel': | |
| self.mel_encoder = nn.Sequential(*[ | |
| nn.Conv1d(80, 64, 3, 1, 1, bias=False), | |
| nn.BatchNorm1d(64) if norm_type=='bn' else LayerNorm(64, dim=1), | |
| nn.GELU(), | |
| nn.Conv1d(64, 64, 3, 1, 1, bias=False) | |
| ]) | |
| else: | |
| raise NotImplementedError("now only ppg or mel are supported!") | |
| self.style_encoder = nn.Sequential(*[ | |
| nn.Linear(135, 256), | |
| nn.GELU(), | |
| nn.Linear(256, 256) | |
| ]) | |
| if backbone_type == 'resnet': | |
| self.backbone = ResNetBackbone() | |
| elif backbone_type == 'unet': | |
| self.backbone = UNetBackbone() | |
| elif backbone_type == 'resblocks': | |
| self.backbone = ResBlocksBackbone() | |
| else: | |
| raise NotImplementedError("Now only resnet and unet are supported!") | |
| self.out_layer = nn.Sequential( | |
| nn.BatchNorm1d(512) if norm_type=='bn' else LayerNorm(512, dim=1), | |
| nn.Conv1d(512, 64, 3, 1, 1, bias=False), | |
| nn.PReLU(), | |
| nn.Conv1d(64, out_dim, 3, 1, 1, bias=False) | |
| ) | |
| self.feat_dropout = nn.Dropout(p=dropout) | |
| def device(self): | |
| return self.backbone.parameters().__next__().device | |
| def forward(self, batch, ret, log_dict=None): | |
| style, x_mask = batch['style'].to(self.device), batch['x_mask'].to(self.device) | |
| style_feat = self.style_encoder(style) # [B,C=135] => [B,C=128] | |
| if self.audio_feat_type == 'ppg': | |
| audio, energy = batch['audio'].to(self.device), batch['energy'].to(self.device) | |
| audio_feat = self.audio_encoder(audio.transpose(1,2)).transpose(1,2) * x_mask.unsqueeze(2) # [B,T,C=29] => [B,T,C=48] | |
| energy_feat = self.energy_encoder(energy.transpose(1,2)).transpose(1,2) * x_mask.unsqueeze(2) # [B,T,C=1] => [B,T,C=16] | |
| feat = torch.cat([audio_feat, energy_feat], dim=2) # [B,T,C=48+16] | |
| elif self.audio_feat_type == 'mel': | |
| mel = batch['mel'].to(self.device) | |
| feat = self.mel_encoder(mel.transpose(1,2)).transpose(1,2) * x_mask.unsqueeze(2) # [B,T,C=64] | |
| feat, x_mask = self.backbone(x=feat, sty=style_feat, x_mask=x_mask) | |
| out = self.out_layer(feat.transpose(1,2)).transpose(1,2) * x_mask.unsqueeze(2) # [B,T//2,C=256] => [B,T//2,C=64] | |
| ret['pred'] = out | |
| ret['mask'] = x_mask | |
| return out | |
| class ResBlocksBackbone(nn.Module): | |
| def __init__(self, in_dim=64, out_dim=512, p_dropout=0.5, norm_type='bn'): | |
| super(ResBlocksBackbone,self).__init__() | |
| self.resblocks_0 = ConvBlocks(channels=in_dim, out_dims=64, dilations=[1]*3, kernel_size=3, norm_type=norm_type, is_BTC=False) | |
| self.resblocks_1 = ConvBlocks(channels=64, out_dims=128, dilations=[1]*4, kernel_size=3, norm_type=norm_type, is_BTC=False) | |
| self.resblocks_2 = ConvBlocks(channels=128, out_dims=256, dilations=[1]*14, kernel_size=3, norm_type=norm_type, is_BTC=False) | |
| self.resblocks_3 = ConvBlocks(channels=512, out_dims=512, dilations=[1]*3, kernel_size=3, norm_type=norm_type, is_BTC=False) | |
| self.resblocks_4 = ConvBlocks(channels=512, out_dims=out_dim, dilations=[1]*3, kernel_size=3, norm_type=norm_type, is_BTC=False) | |
| self.downsampler = LambdaLayer(lambda x: F.interpolate(x, scale_factor=0.5, mode='linear')) | |
| self.upsampler = LambdaLayer(lambda x: F.interpolate(x, scale_factor=4, mode='linear')) | |
| self.dropout = nn.Dropout(p=p_dropout) | |
| def forward(self, x, sty, x_mask=1.): | |
| """ | |
| x: [B, T, C] | |
| sty: [B, C=256] | |
| x_mask: [B, T] | |
| ret: [B, T/2, C] | |
| """ | |
| x = x.transpose(1, 2) # [B, C, T] | |
| x_mask = x_mask[:, None, :] # [B, 1, T] | |
| x = self.resblocks_0(x) * x_mask # [B, C, T] | |
| x_mask = self.downsampler(x_mask) # [B, 1, T/2] | |
| x = self.downsampler(x) * x_mask # [B, C, T/2] | |
| x = self.resblocks_1(x) * x_mask # [B, C, T/2] | |
| x = self.resblocks_2(x) * x_mask # [B, C, T/2] | |
| x = self.dropout(x.transpose(1,2)).transpose(1,2) | |
| sty = sty[:, :, None].repeat([1,1,x_mask.shape[2]]) # [B,C=256,T/2] | |
| x = torch.cat([x, sty], dim=1) # [B, C=256+256, T/2] | |
| x = self.resblocks_3(x) * x_mask # [B, C, T/2] | |
| x = self.resblocks_4(x) * x_mask # [B, C, T/2] | |
| x = x.transpose(1,2) | |
| x_mask = x_mask.squeeze(1) | |
| return x, x_mask | |
| class ResNetBackbone(nn.Module): | |
| def __init__(self, in_dim=64, out_dim=512, p_dropout=0.5, norm_type='bn'): | |
| super(ResNetBackbone,self).__init__() | |
| self.resblocks_0 = ConvBlocks(channels=in_dim, out_dims=64, dilations=[1]*3, kernel_size=3, norm_type=norm_type, is_BTC=False) | |
| self.resblocks_1 = ConvBlocks(channels=64, out_dims=128, dilations=[1]*4, kernel_size=3, norm_type=norm_type, is_BTC=False) | |
| self.resblocks_2 = ConvBlocks(channels=128, out_dims=256, dilations=[1]*14, kernel_size=3, norm_type=norm_type, is_BTC=False) | |
| self.resblocks_3 = ConvBlocks(channels=512, out_dims=512, dilations=[1]*3, kernel_size=3, norm_type=norm_type, is_BTC=False) | |
| self.resblocks_4 = ConvBlocks(channels=512, out_dims=out_dim, dilations=[1]*3, kernel_size=3, norm_type=norm_type, is_BTC=False) | |
| self.downsampler = LambdaLayer(lambda x: F.interpolate(x, scale_factor=0.5, mode='linear')) | |
| self.upsampler = LambdaLayer(lambda x: F.interpolate(x, scale_factor=4, mode='linear')) | |
| self.dropout = nn.Dropout(p=p_dropout) | |
| def forward(self, x, sty, x_mask=1.): | |
| """ | |
| x: [B, T, C] | |
| sty: [B, C=256] | |
| x_mask: [B, T] | |
| ret: [B, T/2, C] | |
| """ | |
| x = x.transpose(1, 2) # [B, C, T] | |
| x_mask = x_mask[:, None, :] # [B, 1, T] | |
| x = self.resblocks_0(x) * x_mask # [B, C, T] | |
| x_mask = self.downsampler(x_mask) # [B, 1, T/2] | |
| x = self.downsampler(x) * x_mask # [B, C, T/2] | |
| x = self.resblocks_1(x) * x_mask # [B, C, T/2] | |
| x_mask = self.downsampler(x_mask) # [B, 1, T/4] | |
| x = self.downsampler(x) * x_mask # [B, C, T/4] | |
| x = self.resblocks_2(x) * x_mask # [B, C, T/4] | |
| x_mask = self.downsampler(x_mask) # [B, 1, T/8] | |
| x = self.downsampler(x) * x_mask # [B, C, T/8] | |
| x = self.dropout(x.transpose(1,2)).transpose(1,2) | |
| sty = sty[:, :, None].repeat([1,1,x_mask.shape[2]]) # [B,C=256,T/8] | |
| x = torch.cat([x, sty], dim=1) # [B, C=256+256, T/8] | |
| x = self.resblocks_3(x) * x_mask # [B, C, T/8] | |
| x_mask = self.upsampler(x_mask) # [B, 1, T/2] | |
| x = self.upsampler(x) * x_mask # [B, C, T/2] | |
| x = self.resblocks_4(x) * x_mask # [B, C, T/2] | |
| x = x.transpose(1,2) | |
| x_mask = x_mask.squeeze(1) | |
| return x, x_mask | |
| class UNetBackbone(nn.Module): | |
| def __init__(self, in_dim=64, out_dim=512, p_dropout=0.5, norm_type='bn'): | |
| super(UNetBackbone, self).__init__() | |
| self.resblocks_0 = ConvBlocks(channels=in_dim, out_dims=64, dilations=[1]*3, kernel_size=3, norm_type=norm_type, is_BTC=False) | |
| self.resblocks_1 = ConvBlocks(channels=64, out_dims=128, dilations=[1]*4, kernel_size=3, norm_type=norm_type, is_BTC=False) | |
| self.resblocks_2 = ConvBlocks(channels=128, out_dims=256, dilations=[1]*8, kernel_size=3, norm_type=norm_type, is_BTC=False) | |
| self.resblocks_3 = ConvBlocks(channels=512, out_dims=512, dilations=[1]*3, kernel_size=3, norm_type=norm_type, is_BTC=False) | |
| self.resblocks_4 = ConvBlocks(channels=768, out_dims=512, dilations=[1]*3, kernel_size=3, norm_type=norm_type, is_BTC=False) # [768 = c3(512) + c2(256)] | |
| self.resblocks_5 = ConvBlocks(channels=640, out_dims=out_dim, dilations=[1]*3, kernel_size=3, norm_type=norm_type, is_BTC=False) # [640 = c4(512) + c1(128)] | |
| self.downsampler = nn.Upsample(scale_factor=0.5, mode='linear') | |
| self.upsampler = nn.Upsample(scale_factor=2, mode='linear') | |
| self.dropout = nn.Dropout(p=p_dropout) | |
| def forward(self, x, sty, x_mask=1.): | |
| """ | |
| x: [B, T, C] | |
| sty: [B, C=256] | |
| x_mask: [B, T] | |
| ret: [B, T/2, C] | |
| """ | |
| x = x.transpose(1, 2) # [B, C, T] | |
| x_mask = x_mask[:, None, :] # [B, 1, T] | |
| x0 = self.resblocks_0(x) * x_mask # [B, C, T] | |
| x_mask = self.downsampler(x_mask) # [B, 1, T/2] | |
| x = self.downsampler(x0) * x_mask # [B, C, T/2] | |
| x1 = self.resblocks_1(x) * x_mask # [B, C, T/2] | |
| x_mask = self.downsampler(x_mask) # [B, 1, T/4] | |
| x = self.downsampler(x1) * x_mask # [B, C, T/4] | |
| x2 = self.resblocks_2(x) * x_mask # [B, C, T/4] | |
| x_mask = self.downsampler(x_mask) # [B, 1, T/8] | |
| x = self.downsampler(x2) * x_mask # [B, C, T/8] | |
| x = self.dropout(x.transpose(1,2)).transpose(1,2) | |
| sty = sty[:, :, None].repeat([1,1,x_mask.shape[2]]) # [B,C=256,T/8] | |
| x = torch.cat([x, sty], dim=1) # [B, C=256+256, T/8] | |
| x3 = self.resblocks_3(x) * x_mask # [B, C, T/8] | |
| x_mask = self.upsampler(x_mask) # [B, 1, T/4] | |
| x = self.upsampler(x3) * x_mask # [B, C, T/4] | |
| x = torch.cat([x, self.dropout(x2.transpose(1,2)).transpose(1,2)], dim=1) # | |
| x4 = self.resblocks_4(x) * x_mask # [B, C, T/4] | |
| x_mask = self.upsampler(x_mask) # [B, 1, T/2] | |
| x = self.upsampler(x4) * x_mask # [B, C, T/2] | |
| x = torch.cat([x, self.dropout(x1.transpose(1,2)).transpose(1,2)], dim=1) | |
| x5 = self.resblocks_5(x) * x_mask # [B, C, T/2] | |
| x = x5.transpose(1,2) | |
| x_mask = x_mask.squeeze(1) | |
| return x, x_mask | |
| if __name__ == '__main__': | |
| pass | |