mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-19 07:54:56 +00:00
Brandon Thomas
2d2b25d5a6
MIT license Commit: https://github.com/OpenTalker/SadTalker commit ae0b26359d44c9379450a8e02a38b45eeec15549 (HEAD -> main, origin/main, origin/HEAD) Author: shadowcun <shadowcun@tencent.com> Date: Mon Jul 3 13:48:35 2023 +0800 fixed bug: hacking facexlib of FAN
42 lines
1.2 KiB
Python
42 lines
1.2 KiB
Python
from tqdm import tqdm
|
|
import torch
|
|
from torch import nn
|
|
|
|
|
|
class Audio2Exp(nn.Module):
|
|
def __init__(self, netG, cfg, device, prepare_training_loss=False):
|
|
super(Audio2Exp, self).__init__()
|
|
self.cfg = cfg
|
|
self.device = device
|
|
self.netG = netG.to(device)
|
|
|
|
def test(self, batch):
|
|
|
|
mel_input = batch['indiv_mels'] # bs T 1 80 16
|
|
bs = mel_input.shape[0]
|
|
T = mel_input.shape[1]
|
|
|
|
exp_coeff_pred = []
|
|
|
|
for i in tqdm(range(0, T, 10),'audio2exp:'): # every 10 frames
|
|
|
|
current_mel_input = mel_input[:,i:i+10]
|
|
|
|
#ref = batch['ref'][:, :, :64].repeat((1,current_mel_input.shape[1],1)) #bs T 64
|
|
ref = batch['ref'][:, :, :64][:, i:i+10]
|
|
ratio = batch['ratio_gt'][:, i:i+10] #bs T
|
|
|
|
audiox = current_mel_input.view(-1, 1, 80, 16) # bs*T 1 80 16
|
|
|
|
curr_exp_coeff_pred = self.netG(audiox, ref, ratio) # bs T 64
|
|
|
|
exp_coeff_pred += [curr_exp_coeff_pred]
|
|
|
|
# BS x T x 64
|
|
results_dict = {
|
|
'exp_coeff_pred': torch.cat(exp_coeff_pred, axis=1)
|
|
}
|
|
return results_dict
|
|
|
|
|