import scipy from scipy import linalg from torch.nn import functional as F import torch from torch import nn import numpy as np from modules.audio2motion.transformer_models import FFTBlocks import modules.audio2motion.utils as utils from modules.audio2motion.flow_base import Glow, WN, ResidualCouplingBlock import torch.distributions as dist from modules.audio2motion.cnn_models import LambdaLayer, LayerNorm from vector_quantize_pytorch import VectorQuantize class FVAEEncoder(nn.Module): def __init__(self, in_channels, hidden_channels, latent_channels, kernel_size, n_layers, gin_channels=0, p_dropout=0, strides=[4]): super().__init__() self.strides = strides self.hidden_size = hidden_channels self.pre_net = nn.Sequential(*[ nn.Conv1d(in_channels, hidden_channels, kernel_size=s * 2, stride=s, padding=s // 2) if i == 0 else nn.Conv1d(hidden_channels, hidden_channels, kernel_size=s * 2, stride=s, padding=s // 2) for i, s in enumerate(strides) ]) self.wn = WN(hidden_channels, kernel_size, 1, n_layers, gin_channels, p_dropout) self.out_proj = nn.Conv1d(hidden_channels, latent_channels * 2, 1) self.latent_channels = latent_channels def forward(self, x, x_mask, g): x = self.pre_net(x) x_mask = x_mask[:, :, ::np.prod(self.strides)][:, :, :x.shape[-1]] x = x * x_mask x = self.wn(x, x_mask, g) * x_mask x = self.out_proj(x) m, logs = torch.split(x, self.latent_channels, dim=1) z = (m + torch.randn_like(m) * torch.exp(logs)) return z, m, logs, x_mask class FVAEDecoder(nn.Module): def __init__(self, latent_channels, hidden_channels, out_channels, kernel_size, n_layers, gin_channels=0, p_dropout=0, strides=[4]): super().__init__() self.strides = strides self.hidden_size = hidden_channels self.pre_net = nn.Sequential(*[ nn.ConvTranspose1d(latent_channels, hidden_channels, kernel_size=s, stride=s) if i == 0 else nn.ConvTranspose1d(hidden_channels, hidden_channels, kernel_size=s, stride=s) for i, s in enumerate(strides) ]) self.wn = WN(hidden_channels, kernel_size, 1, n_layers, gin_channels, p_dropout) self.out_proj = nn.Conv1d(hidden_channels, out_channels, 1) def forward(self, x, x_mask, g): x = self.pre_net(x) x = x * x_mask x = self.wn(x, x_mask, g) * x_mask x = self.out_proj(x) return x class VQVAE(nn.Module): def __init__(self, in_out_channels=64, hidden_channels=256, latent_size=16, kernel_size=3, enc_n_layers=5, dec_n_layers=5, gin_channels=80, strides=[4,], sqz_prior=False): super().__init__() self.in_out_channels = in_out_channels self.strides = strides self.hidden_size = hidden_channels self.latent_size = latent_size self.g_pre_net = nn.Sequential(*[ nn.Conv1d(gin_channels, gin_channels, kernel_size=s * 2, stride=s, padding=s // 2) for i, s in enumerate(strides) ]) self.encoder = FVAEEncoder(in_out_channels, hidden_channels, hidden_channels, kernel_size, enc_n_layers, gin_channels, strides=strides) # if use_prior_glow: # self.prior_flow = ResidualCouplingBlock( # latent_size, glow_hidden, glow_kernel_size, 1, glow_n_blocks, 4, gin_channels=gin_channels) self.vq = VectorQuantize(dim=hidden_channels, codebook_size=256, codebook_dim=16) self.decoder = FVAEDecoder(hidden_channels, hidden_channels, in_out_channels, kernel_size, dec_n_layers, gin_channels, strides=strides) self.prior_dist = dist.Normal(0, 1) self.sqz_prior = sqz_prior def forward(self, x=None, x_mask=None, g=None, infer=False, **kwargs): """ :param x: [B, T, C_in_out] :param x_mask: [B, T] :param g: [B, T, C_g] :return: """ x_mask = x_mask[:, None, :] # [B, 1, T] g = g.transpose(1,2) # [B, C_g, T] g_for_sqz = g g_sqz = self.g_pre_net(g_for_sqz) if not infer: x = x.transpose(1,2) # [B, C, T] z_q, m_q, logs_q, x_mask_sqz = self.encoder(x, x_mask, g_sqz) if self.sqz_prior: z_q = F.interpolate(z_q, scale_factor=1/8) z_p, idx, commit_loss = self.vq(z_q.transpose(1,2)) if self.sqz_prior: z_p = F.interpolate(z_p.transpose(1,2),scale_factor=8).transpose(1,2) x_recon = self.decoder(z_p.transpose(1,2), x_mask, g) return x_recon.transpose(1,2), commit_loss, z_p.transpose(1,2), m_q.transpose(1,2), logs_q.transpose(1,2) else: bs, t = g_sqz.shape[0], g_sqz.shape[2] if self.sqz_prior: t = t // 8 latent_shape = [int(bs * t)] latent_idx = torch.randint(0,256,latent_shape).to(self.vq.codebook.device) # latent_idx = torch.ones_like(latent_idx, dtype=torch.long) # z_p = torch.gather(self.vq.codebook, 0, latent_idx)# self.vq.codebook[latent_idx] z_p = self.vq.codebook[latent_idx] z_p = z_p.reshape([bs, t, -1]) z_p = self.vq.project_out(z_p) if self.sqz_prior: z_p = F.interpolate(z_p.transpose(1,2),scale_factor=8).transpose(1,2) x_recon = self.decoder(z_p.transpose(1,2), 1, g) return x_recon.transpose(1,2), z_p.transpose(1,2) class VQVAEModel(nn.Module): def __init__(self, in_out_dim=71, sqz_prior=False, enc_no_cond=False): super().__init__() self.mel_encoder = nn.Sequential(*[ nn.Conv1d(80, 64, 3, 1, 1, bias=False), nn.BatchNorm1d(64), nn.GELU(), nn.Conv1d(64, 64, 3, 1, 1, bias=False) ]) self.in_dim, self.out_dim = in_out_dim, in_out_dim self.sqz_prior = sqz_prior self.enc_no_cond = enc_no_cond self.vae = VQVAE(in_out_channels=in_out_dim, hidden_channels=256, latent_size=16, kernel_size=5, enc_n_layers=8, dec_n_layers=4, gin_channels=64, strides=[4,], sqz_prior=sqz_prior) self.downsampler = LambdaLayer(lambda x: F.interpolate(x.transpose(1,2), scale_factor=0.5, mode='nearest').transpose(1,2)) @property def device(self): return self.vae.parameters().__next__().device def forward(self, batch, ret, log_dict=None, train=True): infer = not train mask = batch['y_mask'].to(self.device) mel = batch['mel'].to(self.device) mel = self.downsampler(mel) mel_feat = self.mel_encoder(mel.transpose(1,2)).transpose(1,2) if not infer: exp = batch['exp'].to(self.device) pose = batch['pose'].to(self.device) if self.in_dim == 71: x = torch.cat([exp, pose], dim=-1) # [B, T, C=64 + 7] elif self.in_dim == 64: x = exp elif self.in_dim == 7: x = pose if self.enc_no_cond: x_recon, loss_commit, z_p, m_q, logs_q = self.vae(x=x, x_mask=mask, g=torch.zeros_like(mel_feat), infer=False) else: x_recon, loss_commit, z_p, m_q, logs_q = self.vae(x=x, x_mask=mask, g=mel_feat, infer=False) loss_commit = loss_commit.reshape([]) ret['pred'] = x_recon ret['mask'] = mask ret['loss_commit'] = loss_commit return x_recon, loss_commit, m_q, logs_q else: x_recon, z_p = self.vae(x=None, x_mask=mask, g=mel_feat, infer=True) return x_recon # def __get_feat(self, exp, pose): # diff_exp = exp[:-1, :] - exp[1:, :] # exp_std = (np.std(exp, axis = 0) - self.exp_std_mean) / self.exp_std_std # diff_exp_std = (np.std(diff_exp, axis = 0) - self.exp_diff_std_mean) / self.exp_diff_std_std # diff_pose = pose[:-1, :] - pose[1:, :] # diff_pose_std = (np.std(diff_pose, axis = 0) - self.pose_diff_std_mean) / self.pose_diff_std_std # return np.concatenate((exp_std, diff_exp_std, diff_pose_std)) def num_params(self, model, print_out=True, model_name="model"): parameters = filter(lambda p: p.requires_grad, model.parameters()) parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000 if print_out: print(f'| {model_name} Trainable Parameters: %.3fM' % parameters) return parameters