nets/spg/vqvae_1d.py

import os
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from .wav2vec import Wav2Vec2Model
from .vqvae_modules import VectorQuantizerEMA, ConvNormRelu, Res_CNR_Stack



class AudioEncoder(nn.Module):
    def __init__(self, in_dim, num_hiddens, num_residual_layers, num_residual_hiddens):
        super(AudioEncoder, self).__init__()
        self._num_hiddens = num_hiddens
        self._num_residual_layers = num_residual_layers
        self._num_residual_hiddens = num_residual_hiddens

        self.project = ConvNormRelu(in_dim, self._num_hiddens // 4, leaky=True)

        self._enc_1 = Res_CNR_Stack(self._num_hiddens // 4, self._num_residual_layers, leaky=True)
        self._down_1 = ConvNormRelu(self._num_hiddens // 4, self._num_hiddens // 2, leaky=True, residual=True,
                                    sample='down')
        self._enc_2 = Res_CNR_Stack(self._num_hiddens // 2, self._num_residual_layers, leaky=True)
        self._down_2 = ConvNormRelu(self._num_hiddens // 2, self._num_hiddens, leaky=True, residual=True, sample='down')
        self._enc_3 = Res_CNR_Stack(self._num_hiddens, self._num_residual_layers, leaky=True)

    def forward(self, x, frame_num=0):
        h = self.project(x)
        h = self._enc_1(h)
        h = self._down_1(h)
        h = self._enc_2(h)
        h = self._down_2(h)
        h = self._enc_3(h)
        return h


class Wav2VecEncoder(nn.Module):
    def __init__(self, num_hiddens, num_residual_layers):
        super(Wav2VecEncoder, self).__init__()
        self._num_hiddens = num_hiddens
        self._num_residual_layers = num_residual_layers

        self.audio_encoder = Wav2Vec2Model.from_pretrained(
            "facebook/wav2vec2-base-960h")  # "vitouphy/wav2vec2-xls-r-300m-phoneme""facebook/wav2vec2-base-960h"
        self.audio_encoder.feature_extractor._freeze_parameters()

        self.project = ConvNormRelu(768, self._num_hiddens, leaky=True)

        self._enc_1 = Res_CNR_Stack(self._num_hiddens, self._num_residual_layers, leaky=True)
        self._down_1 = ConvNormRelu(self._num_hiddens, self._num_hiddens, leaky=True, residual=True, sample='down')
        self._enc_2 = Res_CNR_Stack(self._num_hiddens, self._num_residual_layers, leaky=True)
        self._down_2 = ConvNormRelu(self._num_hiddens, self._num_hiddens, leaky=True, residual=True, sample='down')
        self._enc_3 = Res_CNR_Stack(self._num_hiddens, self._num_residual_layers, leaky=True)

    def forward(self, x, frame_num):
        h = self.audio_encoder(x.squeeze(), frame_num=frame_num).last_hidden_state.transpose(1, 2)
        h = self.project(h)
        h = self._enc_1(h)
        h = self._down_1(h)
        h = self._enc_2(h)
        h = self._down_2(h)
        h = self._enc_3(h)
        return h


class Encoder(nn.Module):
    def __init__(self, in_dim, embedding_dim, num_hiddens, num_residual_layers, num_residual_hiddens):
        super(Encoder, self).__init__()
        self._num_hiddens = num_hiddens
        self._num_residual_layers = num_residual_layers
        self._num_residual_hiddens = num_residual_hiddens

        self.project = ConvNormRelu(in_dim, self._num_hiddens // 4, leaky=True)

        self._enc_1 = Res_CNR_Stack(self._num_hiddens // 4, self._num_residual_layers, leaky=True)
        self._down_1 = ConvNormRelu(self._num_hiddens // 4, self._num_hiddens // 2, leaky=True, residual=True,
                                    sample='down')
        self._enc_2 = Res_CNR_Stack(self._num_hiddens // 2, self._num_residual_layers, leaky=True)
        self._down_2 = ConvNormRelu(self._num_hiddens // 2, self._num_hiddens, leaky=True, residual=True, sample='down')
        self._enc_3 = Res_CNR_Stack(self._num_hiddens, self._num_residual_layers, leaky=True)

        self.pre_vq_conv = nn.Conv1d(self._num_hiddens, embedding_dim, 1, 1)

    def forward(self, x):
        h = self.project(x)
        h = self._enc_1(h)
        h = self._down_1(h)
        h = self._enc_2(h)
        h = self._down_2(h)
        h = self._enc_3(h)
        h = self.pre_vq_conv(h)
        return h


class Frame_Enc(nn.Module):
    def __init__(self, in_dim, num_hiddens):
        super(Frame_Enc, self).__init__()
        self.in_dim = in_dim
        self.num_hiddens = num_hiddens

        # self.enc = transformer_Enc(in_dim, num_hiddens, 2, 8, 256, 256, 256, 256, 0, dropout=0.1, n_position=4)
        self.proj = nn.Conv1d(in_dim, num_hiddens, 1, 1)
        self.enc = Res_CNR_Stack(num_hiddens, 2, leaky=True)
        self.proj_1 = nn.Conv1d(256*4, num_hiddens, 1, 1)
        self.proj_2 = nn.Conv1d(256*4, num_hiddens*2, 1, 1)

    def forward(self, x):
        # x = self.enc(x, None)[0].reshape(x.shape[0], -1, 1)
        x = self.enc(self.proj(x)).reshape(x.shape[0], -1, 1)
        second_last = self.proj_2(x)
        last = self.proj_1(x)
        return second_last, last



class Decoder(nn.Module):
    def __init__(self, out_dim, embedding_dim, num_hiddens, num_residual_layers, num_residual_hiddens, ae=False):
        super(Decoder, self).__init__()
        self._num_hiddens = num_hiddens
        self._num_residual_layers = num_residual_layers
        self._num_residual_hiddens = num_residual_hiddens

        self.aft_vq_conv = nn.Conv1d(embedding_dim, self._num_hiddens, 1, 1)

        self._dec_1 = Res_CNR_Stack(self._num_hiddens, self._num_residual_layers, leaky=True)
        self._up_2 = ConvNormRelu(self._num_hiddens, self._num_hiddens // 2, leaky=True, residual=True, sample='up')
        self._dec_2 = Res_CNR_Stack(self._num_hiddens // 2, self._num_residual_layers, leaky=True)
        self._up_3 = ConvNormRelu(self._num_hiddens // 2, self._num_hiddens // 4, leaky=True, residual=True,
                                  sample='up')
        self._dec_3 = Res_CNR_Stack(self._num_hiddens // 4, self._num_residual_layers, leaky=True)

        if ae:
            self.frame_enc = Frame_Enc(out_dim, self._num_hiddens // 4)
            self.gru_sl = nn.GRU(self._num_hiddens // 2, self._num_hiddens // 2, 1, batch_first=True)
            self.gru_l = nn.GRU(self._num_hiddens // 4, self._num_hiddens // 4, 1, batch_first=True)

        self.project = nn.Conv1d(self._num_hiddens // 4, out_dim, 1, 1)

    def forward(self, h, last_frame=None):

        h = self.aft_vq_conv(h)
        h = self._dec_1(h)
        h = self._up_2(h)
        h = self._dec_2(h)
        h = self._up_3(h)
        h = self._dec_3(h)

        recon = self.project(h)
        return recon, None


class Pre_VQ(nn.Module):
    def __init__(self, num_hiddens, embedding_dim, num_chunks):
        super(Pre_VQ, self).__init__()
        self.conv = nn.Conv1d(num_hiddens, num_hiddens, 1, 1, 0, groups=num_chunks)
        self.bn = nn.GroupNorm(num_chunks, num_hiddens)
        self.relu = nn.ReLU()
        self.proj = nn.Conv1d(num_hiddens, embedding_dim, 1, 1, 0, groups=num_chunks)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        x = self.proj(x)
        return x


class VQVAE(nn.Module):
    """VQ-VAE"""

    def __init__(self, in_dim, embedding_dim, num_embeddings,
                 num_hiddens, num_residual_layers, num_residual_hiddens,
                 commitment_cost=0.25, decay=0.99, share=False):
        super().__init__()
        self.in_dim = in_dim
        self.embedding_dim = embedding_dim
        self.num_embeddings = num_embeddings
        self.share_code_vq = share

        self.encoder = Encoder(in_dim, embedding_dim, num_hiddens, num_residual_layers, num_residual_hiddens)
        self.vq_layer = VectorQuantizerEMA(embedding_dim, num_embeddings, commitment_cost, decay)
        self.decoder = Decoder(in_dim, embedding_dim, num_hiddens, num_residual_layers, num_residual_hiddens)

    def forward(self, gt_poses, id=None, pre_state=None):
        z = self.encoder(gt_poses.transpose(1, 2))
        if not self.training:
            e, _ = self.vq_layer(z)
            x_recon, cur_state = self.decoder(e, pre_state.transpose(1, 2) if pre_state is not None else None)
            return e, x_recon

        e, e_q_loss = self.vq_layer(z)
        gt_recon, cur_state = self.decoder(e, pre_state.transpose(1, 2) if pre_state is not None else None)

        return e_q_loss, gt_recon.transpose(1, 2)

    def encode(self, gt_poses, id=None):
        z = self.encoder(gt_poses.transpose(1, 2))
        e, latents = self.vq_layer(z)
        return e, latents

    def decode(self, b, w, e=None, latents=None, pre_state=None):
        if e is not None:
            x = self.decoder(e, pre_state.transpose(1, 2) if pre_state is not None else None)
        else:
            e = self.vq_layer.quantize(latents)
            e = e.view(b, w, -1).permute(0, 2, 1).contiguous()
            x = self.decoder(e, pre_state.transpose(1, 2) if pre_state is not None else None)
        return x


class AE(nn.Module):
    """VQ-VAE"""

    def __init__(self, in_dim, embedding_dim, num_embeddings,
                 num_hiddens, num_residual_layers, num_residual_hiddens):
        super().__init__()
        self.in_dim = in_dim
        self.embedding_dim = embedding_dim
        self.num_embeddings = num_embeddings

        self.encoder = Encoder(in_dim, embedding_dim, num_hiddens, num_residual_layers, num_residual_hiddens)
        self.decoder = Decoder(in_dim, embedding_dim, num_hiddens, num_residual_layers, num_residual_hiddens, True)

    def forward(self, gt_poses, id=None, pre_state=None):
        z = self.encoder(gt_poses.transpose(1, 2))
        if not self.training:
            x_recon, cur_state = self.decoder(z, pre_state.transpose(1, 2) if pre_state is not None else None)
            return z, x_recon
        gt_recon, cur_state = self.decoder(z, pre_state.transpose(1, 2) if pre_state is not None else None)

        return gt_recon.transpose(1, 2)

    def encode(self, gt_poses, id=None):
        z = self.encoder(gt_poses.transpose(1, 2))
        return z