meng2003/music2dance

1import math2import torch3import torch.nn as nn4 5from models.flowplusplus import log_dist as logistic6from models.flowplusplus.nn import NN7from models.flowplusplus.transformer_nn import TransformerNN8 9class Coupling(nn.Module):10    """Mixture-of-Logistics Coupling layer in Flow++11 12    Args:13        in_channels (int): Number of channels in the input.14        mid_channels (int): Number of channels in the transformation network.15        num_blocks (int): Number of residual blocks in the transformation network.16        num_components (int): Number of components in the mixture.17        drop_prob (float): Dropout probability.18        use_attn (bool): Use attention in the NN blocks.19        aux_channels (int): Number of channels in optional auxiliary input.20    """21    def __init__(self, in_channels, cond_dim, out_channels, mid_channels, num_blocks, num_components, drop_prob, seq_length, output_length,22                 use_attn=True, use_logmix=True, use_transformer_nn=False, use_pos_emb=False, use_rel_pos_emb=False, num_heads=10, aux_channels=None, concat_dims=True):23        super(Coupling, self).__init__()24 25        if use_transformer_nn:26            if concat_dims:27                self.nn = TransformerNN(in_channels, out_channels, mid_channels, num_blocks, num_heads, num_components, drop_prob=drop_prob, use_pos_emb=use_pos_emb, use_rel_pos_emb=use_rel_pos_emb, input_length=seq_length, concat_dims=concat_dims, output_length=output_length)28            else:29                self.nn = TransformerNN(cond_dim, out_channels, mid_channels, num_blocks, num_heads, num_components, drop_prob=drop_prob, use_pos_emb=use_pos_emb, use_rel_pos_emb=use_rel_pos_emb, input_length=seq_length, concat_dims=concat_dims, output_length=output_length)30        else:31            self.nn = NN(in_channels, out_channels, mid_channels, num_blocks, num_components, drop_prob, use_attn, aux_channels)32 33        if not concat_dims:34            self.input_encoder = nn.Linear(in_channels,cond_dim)35        self.use_logmix = use_logmix36        self.offset = 2.037        self.sigmoid_offset = 1 - 1 / (1 + math.exp(-self.offset))38        self.cond_dim = cond_dim39        self.concat_dims = concat_dims40 41    def forward(self, x, cond, sldj=None, reverse=False, aux=None):42        x_change, x_id = x43 44        if self.concat_dims:45            x_id_cond = torch.cat((x_id, cond), dim=1)46        else:47            # import pdb;pdb.set_trace()48            x_id_enc = self.input_encoder(x_id.permute(0,2,3,1)).permute(0,3,1,2)49            #import pdb;pdb.set_trace()50            x_id_cond = torch.cat((x_id_enc, cond), dim=2)51        #import pdb;pdb.set_trace()52        a, b, pi, mu, s = self.nn(x_id_cond, aux)53        # import pdb;pdb.set_trace()54        scale = (torch.sigmoid(a+self.offset)+self.sigmoid_offset)55 56        if reverse:57            out = x_change / scale - b58            if self.use_logmix:59                out, scale_ldj = logistic.inverse(out, reverse=True)60                #out = out.clamp(1e-5, 1. - 1e-5)61                out = logistic.mixture_inv_cdf(out, pi, mu, s)62                logistic_ldj = logistic.mixture_log_pdf(out, pi, mu, s)63                sldj = sldj - (torch.log(scale) + scale_ldj + logistic_ldj).flatten(1).sum(-1)64            else:65                sldj = sldj - torch.log(scale).flatten(1).sum(-1)66        else:67            if self.use_logmix:68                out = logistic.mixture_log_cdf(x_change, pi, mu, s).exp()69                out, scale_ldj = logistic.inverse(out)70                logistic_ldj = logistic.mixture_log_pdf(x_change, pi, mu, s)71                sldj = sldj + (logistic_ldj + scale_ldj + torch.log(scale)).flatten(1).sum(-1)72            else:73                out = x_change74                sldj = sldj + torch.log(scale).flatten(1).sum(-1)75            76            out = (out + b) * scale77            78 79        x = (out, x_id)80 81        return x, sldj82