nets/layers.py

import os
import sys

sys.path.append(os.getcwd())

import torch
import torch.nn as nn
import numpy as np


# TODO: be aware of the actual netork structures

def get_log(x):
    log = 0
    while x > 1:
        if x % 2 == 0:
            x = x // 2
            log += 1
        else:
            raise ValueError('x is not a power of 2')

    return log


class ConvNormRelu(nn.Module):
    '''
    (B,C_in,H,W) -> (B, C_out, H, W) 
    there exist some kernel size that makes the result is not H/s
    #TODO: there might some problems with residual
    '''

    def __init__(self,
                 in_channels,
                 out_channels,
                 type='1d',
                 leaky=False,
                 downsample=False,
                 kernel_size=None,
                 stride=None,
                 padding=None,
                 p=0,
                 groups=1,
                 residual=False,
                 norm='bn'):
        '''
        conv-bn-relu
        '''
        super(ConvNormRelu, self).__init__()
        self.residual = residual
        self.norm_type = norm
        # kernel_size = k
        # stride = s

        if kernel_size is None and stride is None:
            if not downsample:
                kernel_size = 3
                stride = 1
            else:
                kernel_size = 4
                stride = 2

        if padding is None:
            if isinstance(kernel_size, int) and isinstance(stride, tuple):
                padding = tuple(int((kernel_size - st) / 2) for st in stride)
            elif isinstance(kernel_size, tuple) and isinstance(stride, int):
                padding = tuple(int((ks - stride) / 2) for ks in kernel_size)
            elif isinstance(kernel_size, tuple) and isinstance(stride, tuple):
                padding = tuple(int((ks - st) / 2) for ks, st in zip(kernel_size, stride))
            else:
                padding = int((kernel_size - stride) / 2)

        if self.residual:
            if downsample:
                if type == '1d':
                    self.residual_layer = nn.Sequential(
                        nn.Conv1d(
                            in_channels=in_channels,
                            out_channels=out_channels,
                            kernel_size=kernel_size,
                            stride=stride,
                            padding=padding
                        )
                    )
                elif type == '2d':
                    self.residual_layer = nn.Sequential(
                        nn.Conv2d(
                            in_channels=in_channels,
                            out_channels=out_channels,
                            kernel_size=kernel_size,
                            stride=stride,
                            padding=padding
                        )
                    )
            else:
                if in_channels == out_channels:
                    self.residual_layer = nn.Identity()
                else:
                    if type == '1d':
                        self.residual_layer = nn.Sequential(
                            nn.Conv1d(
                                in_channels=in_channels,
                                out_channels=out_channels,
                                kernel_size=kernel_size,
                                stride=stride,
                                padding=padding
                            )
                        )
                    elif type == '2d':
                        self.residual_layer = nn.Sequential(
                            nn.Conv2d(
                                in_channels=in_channels,
                                out_channels=out_channels,
                                kernel_size=kernel_size,
                                stride=stride,
                                padding=padding
                            )
                        )

        in_channels = in_channels * groups
        out_channels = out_channels * groups
        if type == '1d':
            self.conv = nn.Conv1d(in_channels=in_channels, out_channels=out_channels,
                                  kernel_size=kernel_size, stride=stride, padding=padding,
                                  groups=groups)
            self.norm = nn.BatchNorm1d(out_channels)
            self.dropout = nn.Dropout(p=p)
        elif type == '2d':
            self.conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
                                  kernel_size=kernel_size, stride=stride, padding=padding,
                                  groups=groups)
            self.norm = nn.BatchNorm2d(out_channels)
            self.dropout = nn.Dropout2d(p=p)
        if norm == 'gn':
            self.norm = nn.GroupNorm(2, out_channels)
        elif norm == 'ln':
            self.norm = nn.LayerNorm(out_channels)
        if leaky:
            self.relu = nn.LeakyReLU(negative_slope=0.2)
        else:
            self.relu = nn.ReLU()

    def forward(self, x, **kwargs):
        if self.norm_type == 'ln':
            out = self.dropout(self.conv(x))
            out = self.norm(out.transpose(1,2)).transpose(1,2)
        else:
            out = self.norm(self.dropout(self.conv(x)))
        if self.residual:
            residual = self.residual_layer(x)
            out += residual
        return self.relu(out)


class UNet1D(nn.Module):
    def __init__(self,
                 input_channels,
                 output_channels,
                 max_depth=5,
                 kernel_size=None,
                 stride=None,
                 p=0,
                 groups=1):
        super(UNet1D, self).__init__()
        self.pre_downsampling_conv = nn.ModuleList([])
        self.conv1 = nn.ModuleList([])
        self.conv2 = nn.ModuleList([])
        self.upconv = nn.Upsample(scale_factor=2, mode='nearest')
        self.max_depth = max_depth
        self.groups = groups

        self.pre_downsampling_conv.append(ConvNormRelu(input_channels, output_channels,
                                                       type='1d', leaky=True, downsample=False,
                                                       kernel_size=kernel_size, stride=stride, p=p, groups=groups))
        self.pre_downsampling_conv.append(ConvNormRelu(output_channels, output_channels,
                                                       type='1d', leaky=True, downsample=False,
                                                       kernel_size=kernel_size, stride=stride, p=p, groups=groups))

        for i in range(self.max_depth):
            self.conv1.append(ConvNormRelu(output_channels, output_channels,
                                           type='1d', leaky=True, downsample=True,
                                           kernel_size=kernel_size, stride=stride, p=p, groups=groups))

        for i in range(self.max_depth):
            self.conv2.append(ConvNormRelu(output_channels, output_channels,
                                           type='1d', leaky=True, downsample=False,
                                           kernel_size=kernel_size, stride=stride, p=p, groups=groups))

    def forward(self, x):

        input_size = x.shape[-1]

        assert get_log(
            input_size) >= self.max_depth, 'num_frames must be a power of 2 and its power must be greater than max_depth'

        x = nn.Sequential(*self.pre_downsampling_conv)(x)

        residuals = []
        residuals.append(x)
        for i, conv1 in enumerate(self.conv1):
            x = conv1(x)
            if i < self.max_depth - 1:
                residuals.append(x)

        for i, conv2 in enumerate(self.conv2):
            x = self.upconv(x) + residuals[self.max_depth - i - 1]
            x = conv2(x)

        return x


class UNet2D(nn.Module):
    def __init__(self):
        super(UNet2D, self).__init__()
        raise NotImplementedError('2D Unet is wierd')


class AudioPoseEncoder1D(nn.Module):
    '''
    (B, C, T) -> (B, C*2, T) -> ... -> (B, C_out, T)
    '''

    def __init__(self,
                 C_in,
                 C_out,
                 kernel_size=None,
                 stride=None,
                 min_layer_nums=None
                 ):
        super(AudioPoseEncoder1D, self).__init__()
        self.C_in = C_in
        self.C_out = C_out

        conv_layers = nn.ModuleList([])
        cur_C = C_in
        num_layers = 0
        while cur_C < self.C_out:
            conv_layers.append(ConvNormRelu(
                in_channels=cur_C,
                out_channels=cur_C * 2,
                kernel_size=kernel_size,
                stride=stride
            ))
            cur_C *= 2
            num_layers += 1

        if (cur_C != C_out) or (min_layer_nums is not None and num_layers < min_layer_nums):
            while (cur_C != C_out) or num_layers < min_layer_nums:
                conv_layers.append(ConvNormRelu(
                    in_channels=cur_C,
                    out_channels=C_out,
                    kernel_size=kernel_size,
                    stride=stride
                ))
                num_layers += 1
                cur_C = C_out

        self.conv_layers = nn.Sequential(*conv_layers)

    def forward(self, x):
        '''
        x: (B, C, T)
        '''
        x = self.conv_layers(x)
        return x


class AudioPoseEncoder2D(nn.Module):
    '''
    (B, C, T) -> (B, 1, C, T) -> ... -> (B, C_out, T)
    '''

    def __init__(self):
        raise NotImplementedError


class AudioPoseEncoderRNN(nn.Module):
    '''
    (B, C, T)->(B, T, C)->(B, T, C_out)->(B, C_out, T)
    '''

    def __init__(self,
                 C_in,
                 hidden_size,
                 num_layers,
                 rnn_cell='gru',
                 bidirectional=False
                 ):
        super(AudioPoseEncoderRNN, self).__init__()
        if rnn_cell == 'gru':
            self.cell = nn.GRU(input_size=C_in, hidden_size=hidden_size, num_layers=num_layers, batch_first=True,
                               bidirectional=bidirectional)
        elif rnn_cell == 'lstm':
            self.cell = nn.LSTM(input_size=C_in, hidden_size=hidden_size, num_layers=num_layers, batch_first=True,
                                bidirectional=bidirectional)
        else:
            raise ValueError('invalid rnn cell:%s' % (rnn_cell))

    def forward(self, x, state=None):

        x = x.permute(0, 2, 1)
        x, state = self.cell(x, state)
        x = x.permute(0, 2, 1)

        return x


class AudioPoseEncoderGraph(nn.Module):
    '''
    (B, C, T)->(B, 2, V, T)->(B, 2, T, V)->(B, D, T, V)
    '''

    def __init__(self,
                 layers_config,  # 理应是(C_in, C_out, kernel_size)的list
                 A,  # adjacent matrix (num_parts, V, V)
                 residual,
                 local_bn=False,
                 share_weights=False
                 ) -> None:
        super().__init__()
        self.A = A
        self.num_joints = A.shape[1]
        self.num_parts = A.shape[0]
        self.C_in = layers_config[0][0]
        self.C_out = layers_config[-1][1]

        self.conv_layers = nn.ModuleList([
            GraphConvNormRelu(
                C_in=c_in,
                C_out=c_out,
                A=self.A,
                residual=residual,
                local_bn=local_bn,
                kernel_size=k,
                share_weights=share_weights
            ) for (c_in, c_out, k) in layers_config
        ])

        self.conv_layers = nn.Sequential(*self.conv_layers)

    def forward(self, x):
        '''
        x: (B, C, T), C should be num_joints*D
        output: (B, D, T, V)
        '''
        B, C, T = x.shape
        x = x.view(B, self.num_joints, self.C_in, T)  # (B, V, D, T)，D：每个joint的特征维度，注意这里V在前面
        x = x.permute(0, 2, 3, 1)  # (B, D, T, V)
        assert x.shape[1] == self.C_in

        x_conved = self.conv_layers(x)

        # x_conved = x_conved.permute(0, 3, 1, 2).contiguous().view(B, self.C_out*self.num_joints, T)#(B, V*C_out, T)

        return x_conved


class SeqEncoder2D(nn.Module):
    '''
    seq_encoder, encoding a seq to a vector
    (B, C, T)->(B, 2, V, T)->(B, 2, T, V) -> (B, 32, )->...->(B, C_out)
    '''

    def __init__(self,
                 C_in,  # should be 2
                 T_in,
                 C_out,
                 num_joints,
                 min_layer_num=None,
                 residual=False
                 ):
        super(SeqEncoder2D, self).__init__()
        self.C_in = C_in
        self.C_out = C_out
        self.T_in = T_in
        self.num_joints = num_joints

        conv_layers = nn.ModuleList([])
        conv_layers.append(ConvNormRelu(
            in_channels=C_in,
            out_channels=32,
            type='2d',
            residual=residual
        ))

        cur_C = 32
        cur_H = T_in
        cur_W = num_joints
        num_layers = 1
        while (cur_C < C_out) or (cur_H > 1) or (cur_W > 1):
            ks = [3, 3]
            st = [1, 1]

            if cur_H > 1:
                if cur_H > 4:
                    ks[0] = 4
                    st[0] = 2
                else:
                    ks[0] = cur_H
                    st[0] = cur_H
            if cur_W > 1:
                if cur_W > 4:
                    ks[1] = 4
                    st[1] = 2
                else:
                    ks[1] = cur_W
                    st[1] = cur_W

            conv_layers.append(ConvNormRelu(
                in_channels=cur_C,
                out_channels=min(C_out, cur_C * 2),
                type='2d',
                kernel_size=tuple(ks),
                stride=tuple(st),
                residual=residual
            ))
            cur_C = min(cur_C * 2, C_out)
            if cur_H > 1:
                if cur_H > 4:
                    cur_H //= 2
                else:
                    cur_H = 1
            if cur_W > 1:
                if cur_W > 4:
                    cur_W //= 2
                else:
                    cur_W = 1
            num_layers += 1

        if min_layer_num is not None and (num_layers < min_layer_num):
            while num_layers < min_layer_num:
                conv_layers.append(ConvNormRelu(
                    in_channels=C_out,
                    out_channels=C_out,
                    type='2d',
                    kernel_size=1,
                    stride=1,
                    residual=residual
                ))
                num_layers += 1

        self.conv_layers = nn.Sequential(*conv_layers)
        self.num_layers = num_layers

    def forward(self, x):
        B, C, T = x.shape
        x = x.view(B, self.num_joints, self.C_in, T)  # (B, V, D, T) V in front
        x = x.permute(0, 2, 3, 1)  # (B, D, T, V)
        assert x.shape[1] == self.C_in and x.shape[-1] == self.num_joints

        x = self.conv_layers(x)
        return x.squeeze()


class SeqEncoder1D(nn.Module):
    '''
    (B, C, T)->(B, D)
    '''

    def __init__(self,
                 C_in,
                 C_out,
                 T_in,
                 min_layer_nums=None
                 ):
        super(SeqEncoder1D, self).__init__()
        conv_layers = nn.ModuleList([])
        cur_C = C_in
        cur_T = T_in
        self.num_layers = 0
        while (cur_C < C_out) or (cur_T > 1):
            ks = 3
            st = 1
            if cur_T > 1:
                if cur_T > 4:
                    ks = 4
                    st = 2
                else:
                    ks = cur_T
                    st = cur_T

            conv_layers.append(ConvNormRelu(
                in_channels=cur_C,
                out_channels=min(C_out, cur_C * 2),
                type='1d',
                kernel_size=ks,
                stride=st
            ))
            cur_C = min(cur_C * 2, C_out)
            if cur_T > 1:
                if cur_T > 4:
                    cur_T = cur_T // 2
                else:
                    cur_T = 1
            self.num_layers += 1

        if min_layer_nums is not None and (self.num_layers < min_layer_nums):
            while self.num_layers < min_layer_nums:
                conv_layers.append(ConvNormRelu(
                    in_channels=C_out,
                    out_channels=C_out,
                    type='1d',
                    kernel_size=1,
                    stride=1
                ))
                self.num_layers += 1
        self.conv_layers = nn.Sequential(*conv_layers)

    def forward(self, x):
        x = self.conv_layers(x)
        return x.squeeze()


class SeqEncoderRNN(nn.Module):
    '''
    (B, C, T) -> (B, T, C) -> (B, D)
    LSTM/GRU-FC
    '''

    def __init__(self,
                 hidden_size,
                 in_size,
                 num_rnn_layers,
                 rnn_cell='gru',
                 bidirectional=False
                 ):
        super(SeqEncoderRNN, self).__init__()
        self.hidden_size = hidden_size
        self.in_size = in_size
        self.num_rnn_layers = num_rnn_layers
        self.bidirectional = bidirectional

        if rnn_cell == 'gru':
            self.cell = nn.GRU(input_size=self.in_size, hidden_size=self.hidden_size, num_layers=self.num_rnn_layers,
                               batch_first=True, bidirectional=bidirectional)
        elif rnn_cell == 'lstm':
            self.cell = nn.LSTM(input_size=self.in_size, hidden_size=self.hidden_size, num_layers=self.num_rnn_layers,
                                batch_first=True, bidirectional=bidirectional)

    def forward(self, x, state=None):

        x = x.permute(0, 2, 1)
        B, T, C = x.shape
        x, _ = self.cell(x, state)
        if self.bidirectional:
            out = torch.cat([x[:, -1, :self.hidden_size], x[:, 0, self.hidden_size:]], dim=-1)
        else:
            out = x[:, -1, :]
        assert out.shape[0] == B
        return out


class SeqEncoderGraph(nn.Module):
    '''
    '''

    def __init__(self,
                 embedding_size,
                 layer_configs,
                 residual,
                 local_bn,
                 A,
                 T,
                 share_weights=False
                 ) -> None:
        super().__init__()

        self.C_in = layer_configs[0][0]
        self.C_out = embedding_size

        self.num_joints = A.shape[1]

        self.graph_encoder = AudioPoseEncoderGraph(
            layers_config=layer_configs,
            A=A,
            residual=residual,
            local_bn=local_bn,
            share_weights=share_weights
        )

        cur_C = layer_configs[-1][1]
        self.spatial_pool = ConvNormRelu(
            in_channels=cur_C,
            out_channels=cur_C,
            type='2d',
            kernel_size=(1, self.num_joints),
            stride=(1, 1),
            padding=(0, 0)
        )

        temporal_pool = nn.ModuleList([])
        cur_H = T
        num_layers = 0
        self.temporal_conv_info = []
        while cur_C < self.C_out or cur_H > 1:
            self.temporal_conv_info.append(cur_C)
            ks = [3, 1]
            st = [1, 1]

            if cur_H > 1:
                if cur_H > 4:
                    ks[0] = 4
                    st[0] = 2
                else:
                    ks[0] = cur_H
                    st[0] = cur_H

            temporal_pool.append(ConvNormRelu(
                in_channels=cur_C,
                out_channels=min(self.C_out, cur_C * 2),
                type='2d',
                kernel_size=tuple(ks),
                stride=tuple(st)
            ))
            cur_C = min(cur_C * 2, self.C_out)

            if cur_H > 1:
                if cur_H > 4:
                    cur_H //= 2
                else:
                    cur_H = 1

            num_layers += 1

        self.temporal_pool = nn.Sequential(*temporal_pool)
        print("graph seq encoder info: temporal pool:", self.temporal_conv_info)
        self.num_layers = num_layers
        # need fc?

    def forward(self, x):
        '''
        x: (B, C, T)
        '''
        B, C, T = x.shape
        x = self.graph_encoder(x)
        x = self.spatial_pool(x)
        x = self.temporal_pool(x)
        x = x.view(B, self.C_out)

        return x


class SeqDecoder2D(nn.Module):
    '''
    (B, D)->(B, D, 1, 1)->(B, C_out, C, T)->(B, C_out, T)
    '''

    def __init__(self):
        super(SeqDecoder2D, self).__init__()
        raise NotImplementedError


class SeqDecoder1D(nn.Module):
    '''
    (B, D)->(B, D, 1)->...->(B, C_out, T)
    '''

    def __init__(self,
                 D_in,
                 C_out,
                 T_out,
                 min_layer_num=None
                 ):
        super(SeqDecoder1D, self).__init__()
        self.T_out = T_out
        self.min_layer_num = min_layer_num

        cur_t = 1

        self.pre_conv = ConvNormRelu(
            in_channels=D_in,
            out_channels=C_out,
            type='1d'
        )
        self.num_layers = 1
        self.upconv = nn.Upsample(scale_factor=2, mode='nearest')
        self.conv_layers = nn.ModuleList([])
        cur_t *= 2
        while cur_t <= T_out:
            self.conv_layers.append(ConvNormRelu(
                in_channels=C_out,
                out_channels=C_out,
                type='1d'
            ))
            cur_t *= 2
            self.num_layers += 1

        post_conv = nn.ModuleList([ConvNormRelu(
            in_channels=C_out,
            out_channels=C_out,
            type='1d'
        )])
        self.num_layers += 1
        if min_layer_num is not None and self.num_layers < min_layer_num:
            while self.num_layers < min_layer_num:
                post_conv.append(ConvNormRelu(
                    in_channels=C_out,
                    out_channels=C_out,
                    type='1d'
                ))
                self.num_layers += 1
        self.post_conv = nn.Sequential(*post_conv)

    def forward(self, x):

        x = x.unsqueeze(-1)
        x = self.pre_conv(x)
        for conv in self.conv_layers:
            x = self.upconv(x)
            x = conv(x)

        x = torch.nn.functional.interpolate(x, size=self.T_out, mode='nearest')
        x = self.post_conv(x)
        return x


class SeqDecoderRNN(nn.Module):
    '''
    (B, D)->(B, C_out, T)
    '''

    def __init__(self,
                 hidden_size,
                 C_out,
                 T_out,
                 num_layers,
                 rnn_cell='gru'
                 ):
        super(SeqDecoderRNN, self).__init__()
        self.num_steps = T_out
        if rnn_cell == 'gru':
            self.cell = nn.GRU(input_size=C_out, hidden_size=hidden_size, num_layers=num_layers, batch_first=True,
                               bidirectional=False)
        elif rnn_cell == 'lstm':
            self.cell = nn.LSTM(input_size=C_out, hidden_size=hidden_size, num_layers=num_layers, batch_first=True,
                                bidirectional=False)
        else:
            raise ValueError('invalid rnn cell:%s' % (rnn_cell))

        self.fc = nn.Linear(hidden_size, C_out)

    def forward(self, hidden, frame_0):
        frame_0 = frame_0.permute(0, 2, 1)
        dec_input = frame_0
        outputs = []
        for i in range(self.num_steps):
            frame_out, hidden = self.cell(dec_input, hidden)
            frame_out = self.fc(frame_out)
            dec_input = frame_out
            outputs.append(frame_out)
        output = torch.cat(outputs, dim=1)
        return output.permute(0, 2, 1)


class SeqTranslator2D(nn.Module):
    '''
    (B, C, T)->(B, 1, C, T)-> ... -> (B, 1, C_out, T_out)
    '''

    def __init__(self,
                 C_in=64,
                 C_out=108,
                 T_in=75,
                 T_out=25,
                 residual=True
                 ):
        super(SeqTranslator2D, self).__init__()
        print("Warning: hard coded")
        self.C_in = C_in
        self.C_out = C_out
        self.T_in = T_in
        self.T_out = T_out
        self.residual = residual

        self.conv_layers = nn.Sequential(
            ConvNormRelu(1, 32, '2d', kernel_size=5, stride=1),
            ConvNormRelu(32, 32, '2d', kernel_size=5, stride=1, residual=self.residual),
            ConvNormRelu(32, 32, '2d', kernel_size=5, stride=1, residual=self.residual),

            ConvNormRelu(32, 64, '2d', kernel_size=5, stride=(4, 3)),
            ConvNormRelu(64, 64, '2d', kernel_size=5, stride=1, residual=self.residual),
            ConvNormRelu(64, 64, '2d', kernel_size=5, stride=1, residual=self.residual),

            ConvNormRelu(64, 128, '2d', kernel_size=5, stride=(4, 1)),
            ConvNormRelu(128, 108, '2d', kernel_size=3, stride=(4, 1)),
            ConvNormRelu(108, 108, '2d', kernel_size=(1, 3), stride=1, residual=self.residual),

            ConvNormRelu(108, 108, '2d', kernel_size=(1, 3), stride=1, residual=self.residual),
            ConvNormRelu(108, 108, '2d', kernel_size=(1, 3), stride=1),
        )

    def forward(self, x):
        assert len(x.shape) == 3 and x.shape[1] == self.C_in and x.shape[2] == self.T_in
        x = x.view(x.shape[0], 1, x.shape[1], x.shape[2])
        x = self.conv_layers(x)
        x = x.squeeze(2)
        return x


class SeqTranslator1D(nn.Module):
    '''
    (B, C, T)->(B, C_out, T)
    '''

    def __init__(self,
                 C_in,
                 C_out,
                 kernel_size=None,
                 stride=None,
                 min_layers_num=None,
                 residual=True,
                 norm='bn'
                 ):
        super(SeqTranslator1D, self).__init__()

        conv_layers = nn.ModuleList([])
        conv_layers.append(ConvNormRelu(
            in_channels=C_in,
            out_channels=C_out,
            type='1d',
            kernel_size=kernel_size,
            stride=stride,
            residual=residual,
            norm=norm
        ))
        self.num_layers = 1
        if min_layers_num is not None and self.num_layers < min_layers_num:
            while self.num_layers < min_layers_num:
                conv_layers.append(ConvNormRelu(
                    in_channels=C_out,
                    out_channels=C_out,
                    type='1d',
                    kernel_size=kernel_size,
                    stride=stride,
                    residual=residual,
                    norm=norm
                ))
                self.num_layers += 1
        self.conv_layers = nn.Sequential(*conv_layers)

    def forward(self, x):
        return self.conv_layers(x)


class SeqTranslatorRNN(nn.Module):
    '''
    (B, C, T)->(B, C_out, T)
    LSTM-FC
    '''

    def __init__(self,
                 C_in,
                 C_out,
                 hidden_size,
                 num_layers,
                 rnn_cell='gru'
                 ):
        super(SeqTranslatorRNN, self).__init__()

        if rnn_cell == 'gru':
            self.enc_cell = nn.GRU(input_size=C_in, hidden_size=hidden_size, num_layers=num_layers, batch_first=True,
                                   bidirectional=False)
            self.dec_cell = nn.GRU(input_size=C_out, hidden_size=hidden_size, num_layers=num_layers, batch_first=True,
                                   bidirectional=False)
        elif rnn_cell == 'lstm':
            self.enc_cell = nn.LSTM(input_size=C_in, hidden_size=hidden_size, num_layers=num_layers, batch_first=True,
                                    bidirectional=False)
            self.dec_cell = nn.LSTM(input_size=C_out, hidden_size=hidden_size, num_layers=num_layers, batch_first=True,
                                    bidirectional=False)
        else:
            raise ValueError('invalid rnn cell:%s' % (rnn_cell))

        self.fc = nn.Linear(hidden_size, C_out)

    def forward(self, x, frame_0):

        num_steps = x.shape[-1]
        x = x.permute(0, 2, 1)
        frame_0 = frame_0.permute(0, 2, 1)
        _, hidden = self.enc_cell(x, None)

        outputs = []
        for i in range(num_steps):
            inputs = frame_0
            output_frame, hidden = self.dec_cell(inputs, hidden)
            output_frame = self.fc(output_frame)
            frame_0 = output_frame
            outputs.append(output_frame)
        outputs = torch.cat(outputs, dim=1)
        return outputs.permute(0, 2, 1)


class ResBlock(nn.Module):
    def __init__(self,
                 input_dim,
                 fc_dim,
                 afn,
                 nfn
                 ):
        '''
        afn: activation fn
        nfn: normalization fn
        '''
        super(ResBlock, self).__init__()

        self.input_dim = input_dim
        self.fc_dim = fc_dim
        self.afn = afn
        self.nfn = nfn

        if self.afn != 'relu':
            raise ValueError('Wrong')

        if self.nfn == 'layer_norm':
            raise ValueError('wrong')

        self.layers = nn.Sequential(
            nn.Linear(self.input_dim, self.fc_dim // 2),
            nn.ReLU(),
            nn.Linear(self.fc_dim // 2, self.fc_dim // 2),
            nn.ReLU(),
            nn.Linear(self.fc_dim // 2, self.fc_dim),
            nn.ReLU()
        )

        self.shortcut_layer = nn.Sequential(
            nn.Linear(self.input_dim, self.fc_dim),
            nn.ReLU(),
        )

    def forward(self, inputs):
        return self.layers(inputs) + self.shortcut_layer(inputs)


class AudioEncoder(nn.Module):
    def __init__(self, channels, padding=3, kernel_size=8, conv_stride=2, conv_pool=None, augmentation=False):
        super(AudioEncoder, self).__init__()
        self.in_channels = channels[0]
        self.augmentation = augmentation

        model = []
        acti = nn.LeakyReLU(0.2)

        nr_layer = len(channels) - 1

        for i in range(nr_layer):
            if conv_pool is None:
                model.append(nn.ReflectionPad1d(padding))
                model.append(nn.Conv1d(channels[i], channels[i + 1], kernel_size=kernel_size, stride=conv_stride))
                model.append(acti)
            else:
                model.append(nn.ReflectionPad1d(padding))
                model.append(nn.Conv1d(channels[i], channels[i + 1], kernel_size=kernel_size, stride=conv_stride))
                model.append(acti)
                model.append(conv_pool(kernel_size=2, stride=2))

        if self.augmentation:
            model.append(
                nn.Conv1d(channels[-1], channels[-1], kernel_size=kernel_size, stride=conv_stride)
            )
            model.append(acti)

        self.model = nn.Sequential(*model)

    def forward(self, x):

        x = x[:, :self.in_channels, :]
        x = self.model(x)
        return x


class AudioDecoder(nn.Module):
    def __init__(self, channels, kernel_size=7, ups=25):
        super(AudioDecoder, self).__init__()

        model = []
        pad = (kernel_size - 1) // 2
        acti = nn.LeakyReLU(0.2)

        for i in range(len(channels) - 2):
            model.append(nn.Upsample(scale_factor=2, mode='nearest'))
            model.append(nn.ReflectionPad1d(pad))
            model.append(nn.Conv1d(channels[i], channels[i + 1],
                                   kernel_size=kernel_size, stride=1))
            if i == 0 or i == 1:
                model.append(nn.Dropout(p=0.2))
            if not i == len(channels) - 2:
                model.append(acti)

        model.append(nn.Upsample(size=ups, mode='nearest'))
        model.append(nn.ReflectionPad1d(pad))
        model.append(nn.Conv1d(channels[-2], channels[-1],
                               kernel_size=kernel_size, stride=1))

        self.model = nn.Sequential(*model)

    def forward(self, x):
        return self.model(x)


class Audio2Pose(nn.Module):
    def __init__(self, pose_dim, embed_size, augmentation, ups=25):
        super(Audio2Pose, self).__init__()
        self.pose_dim = pose_dim
        self.embed_size = embed_size
        self.augmentation = augmentation

        self.aud_enc = AudioEncoder(channels=[13, 64, 128, 256], padding=2, kernel_size=7, conv_stride=1,
                                    conv_pool=nn.AvgPool1d, augmentation=self.augmentation)
        if self.augmentation:
            self.aud_dec = AudioDecoder(channels=[512, 256, 128, pose_dim])
        else:
            self.aud_dec = AudioDecoder(channels=[256, 256, 128, pose_dim], ups=ups)

        if self.augmentation:
            self.pose_enc = nn.Sequential(
                nn.Linear(self.embed_size // 2, 256),
                nn.LayerNorm(256)
            )

    def forward(self, audio_feat, dec_input=None):

        B = audio_feat.shape[0]

        aud_embed = self.aud_enc.forward(audio_feat)

        if self.augmentation:
            dec_input = dec_input.squeeze(0)
            dec_embed = self.pose_enc(dec_input)
            dec_embed = dec_embed.unsqueeze(2)
            dec_embed = dec_embed.expand(dec_embed.shape[0], dec_embed.shape[1], aud_embed.shape[-1])
            aud_embed = torch.cat([aud_embed, dec_embed], dim=1)

        out = self.aud_dec.forward(aud_embed)
        return out


if __name__ == '__main__':
    import numpy as np
    import os
    import sys

    test_model = SeqEncoder2D(
        C_in=2,
        T_in=25,
        C_out=512,
        num_joints=54,
    )
    print(test_model.num_layers)

    input = torch.randn((64, 108, 25))
    output = test_model(input)
    print(output.shape)