RabbitRUI/ruispace

1import torch2from torch import nn3from src.audio2pose_models.cvae import CVAE4from src.audio2pose_models.discriminator import PoseSequenceDiscriminator5from src.audio2pose_models.audio_encoder import AudioEncoder6 7class Audio2Pose(nn.Module):8    def __init__(self, cfg, wav2lip_checkpoint, device='cuda'):9        super().__init__()10        self.cfg = cfg11        self.seq_len = cfg.MODEL.CVAE.SEQ_LEN12        self.latent_dim = cfg.MODEL.CVAE.LATENT_SIZE13        self.device = device14 15        self.audio_encoder = AudioEncoder(wav2lip_checkpoint, device)16        self.audio_encoder.eval()17        for param in self.audio_encoder.parameters():18            param.requires_grad = False19 20        self.netG = CVAE(cfg)21        self.netD_motion = PoseSequenceDiscriminator(cfg)22        23        24    def forward(self, x):25 26        batch = {}27        coeff_gt = x['gt'].cuda().squeeze(0)           #bs frame_len+1 7328        batch['pose_motion_gt'] = coeff_gt[:, 1:, -9:-3] - coeff_gt[:, :1, -9:-3] #bs frame_len 629        batch['ref'] = coeff_gt[:, 0, -9:-3]  #bs  630        batch['class'] = x['class'].squeeze(0).cuda() # bs31        indiv_mels= x['indiv_mels'].cuda().squeeze(0) # bs seq_len+1 80 1632 33        # forward34        audio_emb_list = []35        audio_emb = self.audio_encoder(indiv_mels[:, 1:, :, :].unsqueeze(2)) #bs seq_len 51236        batch['audio_emb'] = audio_emb37        batch = self.netG(batch)38 39        pose_motion_pred = batch['pose_motion_pred']           # bs frame_len 640        pose_gt = coeff_gt[:, 1:, -9:-3].clone()               # bs frame_len 641        pose_pred = coeff_gt[:, :1, -9:-3] + pose_motion_pred  # bs frame_len 642 43        batch['pose_pred'] = pose_pred44        batch['pose_gt'] = pose_gt45 46        return batch47 48    def test(self, x):49 50        batch = {}51        ref = x['ref']                            #bs 1 7052        batch['ref'] = x['ref'][:,0,-6:]  53        batch['class'] = x['class']  54        bs = ref.shape[0]55        56        indiv_mels= x['indiv_mels']               # bs T 1 80 1657        indiv_mels_use = indiv_mels[:, 1:]        # we regard the ref as the first frame58        num_frames = x['num_frames']59        num_frames = int(num_frames) - 160 61        #  62        div = num_frames//self.seq_len63        re = num_frames%self.seq_len64        audio_emb_list = []65        pose_motion_pred_list = [torch.zeros(batch['ref'].unsqueeze(1).shape, dtype=batch['ref'].dtype, 66                                                device=batch['ref'].device)]67 68        for i in range(div):69            z = torch.randn(bs, self.latent_dim).to(ref.device)70            batch['z'] = z71            audio_emb = self.audio_encoder(indiv_mels_use[:, i*self.seq_len:(i+1)*self.seq_len,:,:,:]) #bs seq_len 51272            batch['audio_emb'] = audio_emb73            batch = self.netG.test(batch)74            pose_motion_pred_list.append(batch['pose_motion_pred'])  #list of bs seq_len 675        76        if re != 0:77            z = torch.randn(bs, self.latent_dim).to(ref.device)78            batch['z'] = z79            audio_emb = self.audio_encoder(indiv_mels_use[:, -1*self.seq_len:,:,:,:]) #bs seq_len  51280            if audio_emb.shape[1] != self.seq_len:81                pad_dim = self.seq_len-audio_emb.shape[1]82                pad_audio_emb = audio_emb[:, :1].repeat(1, pad_dim, 1) 83                audio_emb = torch.cat([pad_audio_emb, audio_emb], 1) 84            batch['audio_emb'] = audio_emb85            batch = self.netG.test(batch)86            pose_motion_pred_list.append(batch['pose_motion_pred'][:,-1*re:,:])   87        88        pose_motion_pred = torch.cat(pose_motion_pred_list, dim = 1)89        batch['pose_motion_pred'] = pose_motion_pred90 91        pose_pred = ref[:, :1, -6:] + pose_motion_pred  # bs T 692 93        batch['pose_pred'] = pose_pred94        return batch95