from numpy import isin import torch import torch.nn as nn from modules.audio2motion.transformer_base import * DEFAULT_MAX_SOURCE_POSITIONS = 2000 DEFAULT_MAX_TARGET_POSITIONS = 2000 class TransformerEncoderLayer(nn.Module): def __init__(self, hidden_size, dropout, kernel_size=None, num_heads=2, norm='ln'): super().__init__() self.hidden_size = hidden_size self.dropout = dropout self.num_heads = num_heads self.op = EncSALayer( hidden_size, num_heads, dropout=dropout, attention_dropout=0.0, relu_dropout=dropout, kernel_size=kernel_size if kernel_size is not None else 9, padding='SAME', norm=norm, act='gelu' ) def forward(self, x, **kwargs): return self.op(x, **kwargs) ###################### # fastspeech modules ###################### 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 FFTBlocks(nn.Module): def __init__(self, hidden_size, num_layers, ffn_kernel_size=9, dropout=None, num_heads=2, use_pos_embed=True, use_last_norm=True, norm='ln', use_pos_embed_alpha=True): super().__init__() self.num_layers = num_layers embed_dim = self.hidden_size = hidden_size self.dropout = dropout if dropout is not None else 0.1 self.use_pos_embed = use_pos_embed self.use_last_norm = use_last_norm if use_pos_embed: self.max_source_positions = DEFAULT_MAX_TARGET_POSITIONS self.padding_idx = 0 self.pos_embed_alpha = nn.Parameter(torch.Tensor([1])) if use_pos_embed_alpha else 1 self.embed_positions = SinusoidalPositionalEmbedding( embed_dim, self.padding_idx, init_size=DEFAULT_MAX_TARGET_POSITIONS, ) self.layers = nn.ModuleList([]) self.layers.extend([ TransformerEncoderLayer(self.hidden_size, self.dropout, kernel_size=ffn_kernel_size, num_heads=num_heads, norm=norm) for _ in range(self.num_layers) ]) if self.use_last_norm: if norm == 'ln': self.layer_norm = nn.LayerNorm(embed_dim) elif norm == 'bn': self.layer_norm = BatchNorm1dTBC(embed_dim) elif norm == 'gn': self.layer_norm = GroupNorm1DTBC(8, embed_dim) else: self.layer_norm = None def forward(self, x, padding_mask=None, attn_mask=None, return_hiddens=False): """ :param x: [B, T, C] :param padding_mask: [B, T] :return: [B, T, C] or [L, B, T, C] """ padding_mask = x.abs().sum(-1).eq(0).data if padding_mask is None else padding_mask nonpadding_mask_TB = 1 - padding_mask.transpose(0, 1).float()[:, :, None] # [T, B, 1] if self.use_pos_embed: positions = self.pos_embed_alpha * self.embed_positions(x[..., 0]) x = x + positions x = F.dropout(x, p=self.dropout, training=self.training) # B x T x C -> T x B x C x = x.transpose(0, 1) * nonpadding_mask_TB hiddens = [] for layer in self.layers: x = layer(x, encoder_padding_mask=padding_mask, attn_mask=attn_mask) * nonpadding_mask_TB hiddens.append(x) if self.use_last_norm: x = self.layer_norm(x) * nonpadding_mask_TB if return_hiddens: x = torch.stack(hiddens, 0) # [L, T, B, C] x = x.transpose(1, 2) # [L, B, T, C] else: x = x.transpose(0, 1) # [B, T, C] return x class SequentialSA(nn.Module): def __init__(self,layers): super(SequentialSA,self).__init__() self.layers = nn.ModuleList(layers) def forward(self,x,x_mask): """ x: [batch, T, H] x_mask: [batch, T] """ pad_mask = 1. - x_mask for layer in self.layers: if isinstance(layer, EncSALayer): x = x.permute(1,0,2) x = layer(x,pad_mask) x = x.permute(1,0,2) elif isinstance(layer, nn.Linear): x = layer(x) * x_mask.unsqueeze(2) elif isinstance(layer, nn.AvgPool1d): x = x.permute(0,2,1) x = layer(x) x = x.permute(0,2,1) elif isinstance(layer, nn.PReLU): bs, t, hid = x.shape x = x.reshape([bs*t,hid]) x = layer(x) x = x.reshape([bs, t, hid]) else: # Relu x = layer(x) return x class TransformerStyleFusionModel(nn.Module): def __init__(self, num_heads=4, dropout = 0.1, out_dim = 64): super(TransformerStyleFusionModel, self).__init__() self.audio_layer = SequentialSA([ nn.Linear(29, 48), nn.ReLU(48), nn.Linear(48, 128), ]) self.energy_layer = SequentialSA([ nn.Linear(1, 16), nn.ReLU(16), nn.Linear(16, 64), ]) self.backbone1 = FFTBlocks(hidden_size=192,num_layers=3) self.sty_encoder = nn.Sequential(*[ nn.Linear(135, 64), nn.ReLU(), nn.Linear(64, 128) ]) self.backbone2 = FFTBlocks(hidden_size=320,num_layers=3) self.out_layer = SequentialSA([ nn.AvgPool1d(kernel_size=2,stride=2,padding=0), #[b,hid,t_audio]=>[b,hid,t_audio//2] nn.Linear(320,out_dim), nn.PReLU(out_dim), nn.Linear(out_dim,out_dim), ]) self.dropout = nn.Dropout(p = dropout) def forward(self, audio, energy, style, x_mask, y_mask): pad_mask = 1. - x_mask audio_feat = self.audio_layer(audio, x_mask) energy_feat = self.energy_layer(energy, x_mask) feat = torch.cat((audio_feat, energy_feat), dim=-1) # [batch, T, H=48+16] feat = self.backbone1(feat, pad_mask) feat = self.dropout(feat) sty_feat = self.sty_encoder(style) # [batch,135]=>[batch, H=64] sty_feat = sty_feat.unsqueeze(1).repeat(1, feat.shape[1], 1) # [batch, T, H=64] feat = torch.cat([feat, sty_feat], dim=-1) # [batch, T, H=64+64] feat = self.backbone2(feat, pad_mask) # [batch, T, H=128] out = self.out_layer(feat, y_mask) # [batch, T//2, H=out_dim] return out if __name__ == '__main__': model = TransformerStyleFusionModel() audio = torch.rand(4,200,29) # [B,T,H] energy = torch.rand(4,200,1) # [B,T,H] style = torch.ones(4,135) # [B,T] x_mask = torch.ones(4,200) # [B,T] x_mask[3,10:] = 0 ret = model(audio,energy,style, x_mask) print(" ")