Spaces:
Build error
Build error
| 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 | |