Update app.py

app.py

@@ -1,72 +1,387 @@
-import os, sys
-import tempfile
-import gradio as gr
-from modules.text2speech import text2speech 
-from modules.gfpgan_inference import gfpgan
-from modules.sadtalker_test import SadTalker  
-
-def get_driven_audio(audio):  
-    if os.path.isfile(audio):
-        return audio
-    else:
-        save_path = tempfile.NamedTemporaryFile(
-                delete=False,
-                suffix=("." + "wav"),
-            )
-        gen_audio = text2speech(audio, save_path.name)
-        return gen_audio, gen_audio 
-
-def get_source_image(image):   
-        return image
-
-def sadtalker_demo(result_dir):
-
-    sad_talker = SadTalker()
-    with gr.Blocks(analytics_enabled=False) as sadtalker_interface:
-        with gr.Row().style(equal_height=False):
-            with gr.Column(variant='panel'):
-                with gr.Tabs(elem_id="sadtalker_source_image"):
-                    with gr.TabItem('Upload image'):
-                        with gr.Row():
-                            source_image = gr.Image(label="Source image", source="upload", type="filepath").style(height=256,width=256)
- 
-                with gr.Tabs(elem_id="sadtalker_driven_audio"):
-                    with gr.TabItem('Upload audio'):
-                        with gr.Column(variant='panel'):
-                            driven_audio = gr.Audio(label="Input audio", source="upload", type="filepath")
-                            # submit_audio_1 = gr.Button('Submit', variant='primary')
-                        # submit_audio_1.click(fn=get_driven_audio, inputs=input_audio1, outputs=driven_audio)
-                
-                with gr.Tabs(elem_id="sadtalker_checkbox"):
-                    with gr.TabItem('Settings'):
-                        with gr.Column(variant='panel'):
-                            is_still_mode = gr.Checkbox(label="w/ Still Mode")
-                            enhancer = gr.Checkbox(label="w/ GFPGAN as Face enhancer")
-
-            with gr.Column(variant='panel'): 
-                gen_video = gr.Video(label="Generated video", format="mp4").style(height=256,width=256)
-                gen_text = gr.Textbox(visible=False)
-                submit = gr.Button('Generate', elem_id="sadtalker_generate", variant='primary')
-
-        submit.click(
-                    fn=sad_talker.test, 
-                    inputs=[source_image,
-                            driven_audio,
-                            is_still_mode,
-                            enhancer,
-                            gr.Textbox(value=result_dir, visible=False)], 
-                    outputs=[gen_video, gen_text]
-                    )
-
-    return sadtalker_interface
+import pickle
+import time
+import numpy as np
+import scipy, cv2, os, sys, argparse
+from tqdm import tqdm
+import torch
+import librosa
+from networks import define_G
+from pcavs.config.AudioConfig import AudioConfig
+
+sys.path.append('spectre')
+from config import cfg as spectre_cfg
+from src.spectre import SPECTRE
+
+from audio2mesh_helper import *
+from pcavs.models import create_model, networks
+
+torch.manual_seed(0)
+from scipy.signal import savgol_filter
+
+
+class SimpleWrapperV2(nn.Module):
+    def __init__(self, cfg, use_ref=True, exp_dim=53, noload=False) -> None:
+        super().__init__()
+       
+        self.audio_encoder = networks.define_A_sync(cfg)
+
+        self.mapping1 = nn.Linear(512+exp_dim, exp_dim)
+        nn.init.constant_(self.mapping1.weight, 0.)
+        nn.init.constant_(self.mapping1.bias, 0.)
+        self.use_ref = use_ref
+
+    def forward(self, x, ref, use_tanh=False):
+        x = self.audio_encoder.forward_feature(x).view(x.size(0), -1)
+        ref_reshape = ref.reshape(x.size(0), -1)  #20, -1
+        
+        y = self.mapping1(torch.cat([x, ref_reshape], dim=1)) 
+
+        if self.use_ref:
+            out = y.reshape(ref.shape[0], ref.shape[1], -1) + ref # resudial
+        else:
+            out = y.reshape(ref.shape[0], ref.shape[1], -1)
+
+        if use_tanh:
+            out[:, :50] = torch.tanh(out[:, :50]) * 3
+            
+        return out
+
+class Audio2Mesh(object):
+    def __init__(self, args) -> None:
+        self.args = args
+
+        spectre_cfg.model.use_tex = True
+        spectre_cfg.model.mask_type = args.mask_type
+        spectre_cfg.debug = self.args.debug
+        spectre_cfg.model.netA_sync = 'ressesync'
+        spectre_cfg.model.gpu_ids = [0]
+
+        self.spectre = SPECTRE(spectre_cfg)
+        self.spectre.eval()
+        self.face_tracker = None #FaceTrackerV2() # face landmark detection
+        self.mel_step_size = 16
+        self.fps = args.fps
+        self.Nw = args.tframes
+        self.device = self.args.device
+        self.image_size = self.args.image_size
+
+        ### only audio
+        args.netA_sync = 'ressesync'
+        args.gpu_ids = [0]
+        args.exp_dim = 53
+        args.use_tanh = False
+        args.K = 20 
+
+        self.audio2exp = 'pcavs'
+
+        # 
+        self.avmodel = SimpleWrapperV2(args, exp_dim=args.exp_dim).cuda() 
+        self.avmodel.load_state_dict(torch.load('../packages/pretrained/audio2expression_v2_model.tar')['opt'])
+
+        # 5, 160 = 25fps
+        self.audio = AudioConfig(frame_rate=args.fps, num_frames_per_clip=5, hop_size=160)
+        
+        with open(os.path.join(args.source_dir, 'deca_infos.pkl'), 'rb') as f: # ?
+            self.fitting_coeffs = pickle.load(f, encoding='bytes')
+
+        self.coeffs_dict  = { key: torch.Tensor(self.fitting_coeffs[key]).cuda().squeeze(1) for key in ['cam', 'pose', 'light', 'tex', 'shape', 'exp']}
+
+        #### find the close month
+        exp_tensors = torch.sum(self.coeffs_dict['exp'], dim=1)
+        ssss, sorted_indices = torch.sort(exp_tensors)
+        self.exp_id = sorted_indices[0].item()
+
+        if '.ts' in args.render_path:
+            self.render = torch.jit.load(args.render_path).cuda()
+            self.trt = True
+        else:
+            self.render = define_G(self.Nw*6, 3, args.ngf, args.netR).eval().cuda()
+            self.render.load_state_dict(torch.load(args.render_path))
+            self.trt = False
+
+        print('loaded cached images...')
+
+    @torch.no_grad()
+    def cg2real(self, rendedimages, start_frame=0):
+
+        ## load original image and the mask
+        self.source_images = np.concatenate(load_image_from_dir(os.path.join(self.args.source_dir, 'original_frame'),\
+             resize=self.image_size, limit=len(rendedimages)+start_frame))[start_frame:]
+        self.source_masks = np.concatenate(load_image_from_dir(os.path.join(self.args.source_dir, 'original_mask'),\
+             resize=self.image_size, limit=len(rendedimages)+start_frame))[start_frame:]
+
+        self.source_masks = torch.FloatTensor(np.transpose(self.source_masks,(0,3,1,2))/255.)
+        self.padded_real_tensor = torch.FloatTensor(np.transpose(self.source_images,(0,3,1,2))/255.)
+
+        ## padding the rended_imgs
+        paded_tensor = torch.cat([rendedimages[0:1]]* (self.Nw // 2) + [rendedimages] + [rendedimages[-1:]]* (self.Nw // 2)).contiguous()
+        paded_mask_tensor = torch.cat([self.source_masks[0:1]]* (self.Nw // 2) + [self.source_masks] + [self.source_masks[-1:]]* (self.Nw // 2)).contiguous()
+        paded_real_tensor = torch.cat([self.padded_real_tensor[0:1]]* (self.Nw // 2) + [self.padded_real_tensor] + [self.padded_real_tensor[-1:]]* (self.Nw // 2)).contiguous()
+
+        # paded_mask_tensor = maskErosion(paded_mask_tensor, offY=self.args.mask)
+        padded_input = ((paded_real_tensor-0.5)*2 ) # *(1-paded_mask_tensor)
+        padded_input = torch.nn.functional.interpolate(padded_input, (self.image_size, self.image_size), mode='bilinear', align_corners=False)
+        paded_tensor = torch.nn.functional.interpolate(paded_tensor, (self.image_size, self.image_size), mode='bilinear', align_corners=False)
+        paded_tensor = (paded_tensor-0.5)*2
+
+        result = []
+        for index in tqdm(range(0, len(rendedimages), self.args.renderbs), desc='CG2REAL:'):
+            list_A = []
+            list_R = []
+            list_M = []
+            for i in range(self.args.renderbs):
+                idx = index + i
+                if idx+self.Nw > len(padded_input):
+                    list_A.append(torch.zeros(self.Nw*3,self.image_size,self.image_size).unsqueeze(0))
+                    list_R.append(torch.zeros(self.Nw*3,self.image_size,self.image_size).unsqueeze(0))
+                    list_M.append(torch.zeros(self.Nw*3,self.image_size,self.image_size).unsqueeze(0))
+                else:
+                    list_A.append(padded_input[idx:idx+self.Nw].view(-1, self.image_size, self.image_size).unsqueeze(0))
+                    list_R.append(paded_tensor[idx:idx+self.Nw].view(-1, self.image_size, self.image_size).unsqueeze(0))
+                    list_M.append(paded_mask_tensor[idx:idx+self.Nw].view(-1, self.image_size, self.image_size).unsqueeze(0))
+
+            list_A = torch.cat(list_A)
+            list_R = torch.cat(list_R)
+            list_M = torch.cat(list_M)
+
+            idx = (self.Nw//2) * 3
+            mask = list_M[:, idx:idx+3]
+
+            # list_A = padded_input
+            mask = maskErosion(mask, offY=self.args.mask)
+            list_A = list_A * (1 - mask[:,0:1])
+            A = torch.cat([list_A, list_R], 1)
+
+            if self.trt:
+                B = self.render(A.half().cuda())
+            elif self.args.netR == 'unet_256':
+                # import pdb; pdb.set_trace()
+                idx = (self.Nw//2) * 3
+                mask = list_M[:, idx:idx+3].cuda()
+                mask = maskErosion(mask, offY=self.args.mask)
+                B0 = list_A[:, idx:idx+3].cuda()
+                B = self.render(A.cuda()) * mask[:,0:1] + (1 - mask[:,0:1]) * B0
+            elif  self.args.netR == 's2am':
+                # import pdb; pdb.set_trace()
+                idx = (self.Nw//2) * 3
+                mask = list_M[:, idx:idx+3].cuda()
+                mask = maskErosion(mask, offY=self.args.mask)
+                B0 = list_A[:, idx:idx+3].cuda()
+                B = self.render(A.cuda(), mask[:,0:1] ) * mask[:,0:1] + (1 - mask[:,0:1]) * B0
+            else:
+                B = self.render(A.cuda()) 
+
+            result.append((B.cpu() + 1) * 0.5) # -1,1 -> 0,1
+        
+        return torch.cat(result)[:len(rendedimages)]
+
+    @torch.no_grad()
+    def coeffs_to_img(self, vertices, coeffs, zero_pose=False, XK = 20):
+
+        xlen = vertices.shape[0]
+        all_shape_images = []
+        landmark2d = []
+
+        #### find the most larger pose 51 in the coeffs.
+        max_pose_51 = torch.max(self.coeffs_dict['pose'][..., 3:4].squeeze(-1))
+
+        for i in tqdm(range(0, xlen, XK)):
+            
+            if i + XK > xlen:
+                XK = xlen - i
+
+            codedictdecoder = {}
+            codedictdecoder['shape'] = torch.zeros_like(self.coeffs_dict['shape'][i:i+XK].cuda())
+            codedictdecoder['tex'] = self.coeffs_dict['tex'][i:i+XK].cuda()
+            codedictdecoder['exp'] =  torch.zeros_like(self.coeffs_dict['exp'][i:i+XK].cuda()) #  all_exps[i:i+XK, :50].cuda()  # # # vid_exps[i:i+1].cuda() i:i+XK
+            codedictdecoder['pose'] = self.coeffs_dict['pose'][i:i+XK]  # vid_poses[i:i+1].cuda()
+            codedictdecoder['cam'] =  self.coeffs_dict['cam'][i:i+XK].cuda() # vid_poses[i:i+1].cuda()
+            codedictdecoder['light'] = self.coeffs_dict['light'][i:i+XK].cuda() # vid_poses[i:i+1].cuda()
+            codedictdecoder['images'] = torch.zeros((XK,3,256,256)).cuda()
+
+            codedictdecoder['pose'][..., 3:4] = torch.clip(coeffs[i:i+XK, 50:51], 0, max_pose_51*0.9) # torch.zeros_like(self.coeffs_dict['pose'][i:i+XK, 3:])
+            codedictdecoder['pose'][..., 4:6] = 0 # coeffs[i:i+XK, 50:]*( - 0.25) # torch.zeros_like(self.coeffs_dict['pose'][i:i+XK, 3:])
+
+            sub_vertices = vertices[i:i+XK].cuda()
+
+            opdict = self.spectre.decode_verts(codedictdecoder, sub_vertices, rendering=True, vis_lmk=False, return_vis=False)
+
+            landmark2d.append(opdict['landmarks2d'].cpu())
+
+            all_shape_images.append(opdict['rendered_images'].cpu())
+
+        rendedimages = torch.cat(all_shape_images)
+
+        lmk2d = torch.cat(landmark2d)
+
+        return rendedimages, lmk2d
+
+    
+    @torch.no_grad()
+    def run_spectre_v3(self, wav=None, ds_features=None, L=20):
+
+        wav = audio_normalize(wav)
+        all_mel = self.audio.melspectrogram(wav).astype(np.float32).T
+        frames_from_audio = np.arange(2, len(all_mel) // self.audio.num_bins_per_frame - 2) # 2,[]mmmmmmmmmmmmmmmmmmmmmmmmmmmm
+        audio_inds = frame2audio_indexs(frames_from_audio, self.audio.num_frames_per_clip, self.audio.num_bins_per_frame)
+
+        vid_exps = self.coeffs_dict['exp'][self.exp_id:self.exp_id+1]
+        vid_poses = self.coeffs_dict['pose'][self.exp_id:self.exp_id+1]
+        
+        ref = torch.cat([vid_exps.view(1, 50), vid_poses[:, 3:].view(1, 3)], dim=-1)
+        ref = ref[...,:self.args.exp_dim]
+
+        K = 20
+        xlens = len(audio_inds) # len(self.coeffs_dict['exp'])
+
+        exps = []
+        for i in tqdm(range(0, xlens, K), desc='S2 DECODER:'+ str(xlens) + ' '):
+            
+            mels = []
+            for j in range(K):
+                if i + j < xlens:
+                    idx = i+j # //3 * 3 
+                    mel = load_spectrogram(all_mel, audio_inds[idx], self.audio.num_frames_per_clip * self.audio.num_bins_per_frame).cuda()
+                    mel = mel.view(-1, 1, 80, self.audio.num_frames_per_clip * self.audio.num_bins_per_frame)
+                    mels.append(mel)
+                else:
+                    break
+
+            mels = torch.cat(mels, dim=0)
+            new_exp = self.avmodel(mels, ref.repeat(mels.shape[0], 1, 1).cuda(), self.args.use_tanh) # exp 53
+            exps+= [new_exp.view(-1, 53)]
+            
+        all_exps = torch.cat(exps,axis=0)
+        
+        return all_exps
+
+    @torch.no_grad()
+    def test_model(self, wav_path):   
+        
+        sys.path.append('../FaceFormer')
+        from faceformer import Faceformer
+        from transformers import Wav2Vec2FeatureExtractor,Wav2Vec2Processor
+        from faceformer import PeriodicPositionalEncoding, init_biased_mask
+        
+        #build model
+        self.args.train_subjects = " ".join(["A"]*8) # suitable for pre-trained faceformer checkpoint
+        model = Faceformer(self.args)
+        model.load_state_dict(torch.load('/apdcephfs/private_shadowcun/Avatar2dFF/medias/videos/c8/mask5000_l2/6_model.pth'))  # ../packages/pretrained/28_ff_model.pth
+        model = model.to(torch.device(self.device))
+        model.eval()
+
+        # hacking for long audio generation
+        model.PPE = PeriodicPositionalEncoding(self.args.feature_dim, period = self.args.period, max_seq_len=6000).cuda()
+        model.biased_mask = init_biased_mask(n_head = 4, max_seq_len = 6000, period=self.args.period).cuda()
+
+        train_subjects_list = ["A"] * 8
+
+        one_hot_labels = np.eye(len(train_subjects_list))
+        one_hot = one_hot_labels[0]
+        one_hot = np.reshape(one_hot,(-1,one_hot.shape[0]))
+        one_hot = torch.FloatTensor(one_hot).to(device=self.device)
+
+        vertices_npy = np.load(self.args.source_dir + '/mesh_pose0.npy')
+        vertices_npy = np.array(vertices_npy).reshape(-1, 5023*3)
+
+        temp = vertices_npy[33] # 829
+
+        template = temp.reshape((-1))
+        template = np.reshape(template,(-1,template.shape[0]))
+        template = torch.FloatTensor(template).to(device=self.device)
+
+        speech_array, sampling_rate = librosa.load(os.path.join(wav_path), sr=16000)
+        processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
+        audio_feature = np.squeeze(processor(speech_array,sampling_rate=16000).input_values)
+        audio_feature = np.reshape(audio_feature,(-1,audio_feature.shape[0]))
+        audio_feature = torch.FloatTensor(audio_feature).to(device=self.device)
  
+        prediction = model.predict(audio_feature, template, one_hot, 1.0) # (1, seq_len, V*3)
+        
+        return prediction.squeeze()
+
+    @torch.no_grad()
+    def run(self, face, audio, start_frame=0):
+
+        wav, sr = librosa.load(audio, sr=16000) # 16*80 ? 20*80
+        wav_tensor = torch.FloatTensor(wav).unsqueeze(0) if len(wav.shape) == 1 else torch.FloatTensor(wav)
+        _, frames = parse_audio_length(wav_tensor.shape[1], 16000, self.args.fps)
+
+        #####  audio-guided, only use the jaw movement
+        all_exps = self.run_spectre_v3(wav)
+
+        # #### temp. interpolation
+        all_exps = torch.nn.functional.interpolate(all_exps.unsqueeze(0).permute([0,2,1]), size=frames, mode='linear')
+        all_exps = all_exps.permute([0,2,1]).squeeze(0)
+
+        # run faceformer for face mesh generation
+        predicted_vertices = self.test_model(audio)
+        predicted_vertices = predicted_vertices.view(-1, 5023*3)
+
+        #### temp. interpolation
+        predicted_vertices = torch.nn.functional.interpolate(predicted_vertices.unsqueeze(0).permute([0,2,1]), size=frames, mode='linear')
+        predicted_vertices = predicted_vertices.permute([0,2,1]).squeeze(0).view(-1, 5023, 3)
+
+        all_exps = torch.Tensor(savgol_filter(all_exps.cpu().numpy(), 5, 3, axis=0)).cpu() # smooth GT
+        
+        rendedimages, lm2d = self.coeffs_to_img(predicted_vertices, all_exps, zero_pose=True)
+        debug_video_gen(rendedimages, self.args.result_dir+"/debug_before_ff.mp4", wav_tensor, self.args.fps, sr)
+
+        # cg2real
+        debug_video_gen(self.cg2real(rendedimages, start_frame=start_frame), self.args.result_dir+"/debug_cg2real_raw.mp4", wav_tensor, self.args.fps, sr)
+        
+        exit()
+
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Stylization and Seamless Video Dubbing')
+    parser.add_argument('--face', default='examples', type=str, help='')
+    parser.add_argument('--audio', default='examples', type=str, help='')
+    parser.add_argument('--source_dir', default='examples', type=str,help='TODO')
+    parser.add_argument('--result_dir', default='examples', type=str,help='TODO')
+    parser.add_argument('--backend', default='wav2lip', type=str,help='wav2lip or pcavs')
+    parser.add_argument('--result_tag', default='result', type=str,help='TODO')
+    parser.add_argument('--netR', default='unet_256', type=str,help='TODO')
+    parser.add_argument('--render_path', default='', type=str,help='TODO')
+    parser.add_argument('--ngf', default=16, type=int,help='TODO')
+    parser.add_argument('--fps', default=20, type=int,help='TODO')
+    parser.add_argument('--mask', default=100, type=int,help='TODO')
+    parser.add_argument('--mask_type', default='v3', type=str,help='TODO')
+    parser.add_argument('--image_size', default=256, type=int,help='TODO')
+    parser.add_argument('--input_nc', default=21, type=int,help='TODO')
+    parser.add_argument('--output_nc', default=3, type=int,help='TODO')
+    parser.add_argument('--renderbs', default=16, type=int,help='TODO')
+    parser.add_argument('--tframes', default=1, type=int,help='TODO')
+    parser.add_argument('--debug', action='store_true')
+    parser.add_argument('--enhance', action='store_true')
+    parser.add_argument('--phone', action='store_true')
+
+    #### faceformer
+    parser.add_argument("--model_name", type=str, default="VOCA")
+    parser.add_argument("--dataset", type=str, default="vocaset", help='vocaset or BIWI')
+    parser.add_argument("--feature_dim", type=int, default=64, help='64 for vocaset; 128 for BIWI')
+    parser.add_argument("--period", type=int, default=30, help='period in PPE - 30 for vocaset; 25 for BIWI')
+    parser.add_argument("--vertice_dim", type=int, default=5023*3, help='number of vertices - 5023*3 for vocaset; 23370*3 for BIWI')
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--train_subjects", type=str, default="FaceTalk_170728_03272_TA ")
+    parser.add_argument("--test_subjects", type=str, default="FaceTalk_170809_00138_TA FaceTalk_170731_00024_TA")
+    parser.add_argument("--condition", type=str, default="FaceTalk_170904_00128_TA", help='select a conditioning subject from train_subjects')
+    parser.add_argument("--subject", type=str, default="FaceTalk_170731_00024_TA", help='select a subject from test_subjects or train_subjects')
+    parser.add_argument("--background_black", type=bool, default=True, help='whether to use black background')
+    parser.add_argument("--template_path", type=str, default="templates.pkl", help='path of the personalized templates')
+    parser.add_argument("--render_template_path", type=str, default="templates", help='path of the mesh in BIWI/FLAME topology')
 
-if __name__ == "__main__":
+    opt = parser.parse_args()
 
-    current_code_path = sys.argv[0]
-    current_root_dir = os.path.split(current_code_path)[0] 
-    sadtalker_result_dir = os.path.join(current_root_dir, 'results', 'sadtalker')
-    demo = sadtalker_demo(sadtalker_result_dir)
-    demo.launch()
+    opt.img_size = 96
+    opt.static = True
+    opt.device = torch.device("cuda")
 
+    a2m = Audio2Mesh(opt)
 
+    print('link start!')
+    t = time.time()
+    # 02780
+    a2m.run(opt.face, opt.audio, 0)
+    print(time.time() - t)