add

app.py

@@ -242,383 +242,6 @@ class BaseTrainer(object):
             original_shape_t[i, selected_indices] = filtered_t[i]
         return original_shape_t
     
-    def _load_data(self, dict_data):
-        tar_pose_raw = dict_data["pose"]
-        tar_pose = tar_pose_raw[:, :, :165].to(self.rank)
-        tar_contact = tar_pose_raw[:, :, 165:169].to(self.rank)
-        tar_trans = dict_data["trans"].to(self.rank)
-        tar_trans_v = dict_data["trans_v"].to(self.rank)
-        tar_exps = dict_data["facial"].to(self.rank)
-        in_audio = dict_data["audio"].to(self.rank) 
-        in_word = dict_data["word"].to(self.rank)
-        tar_beta = dict_data["beta"].to(self.rank)
-        tar_id = dict_data["id"].to(self.rank).long()
-        bs, n, j = tar_pose.shape[0], tar_pose.shape[1], self.joints
-
-        tar_pose_jaw = tar_pose[:, :, 66:69]
-        tar_pose_jaw = rc.axis_angle_to_matrix(tar_pose_jaw.reshape(bs, n, 1, 3))
-        tar_pose_jaw = rc.matrix_to_rotation_6d(tar_pose_jaw).reshape(bs, n, 1*6)
-        tar_pose_face = torch.cat([tar_pose_jaw, tar_exps], dim=2)
-
-        tar_pose_hands = tar_pose[:, :, 25*3:55*3]
-        tar_pose_hands = rc.axis_angle_to_matrix(tar_pose_hands.reshape(bs, n, 30, 3))
-        tar_pose_hands = rc.matrix_to_rotation_6d(tar_pose_hands).reshape(bs, n, 30*6)
-
-        tar_pose_upper = tar_pose[:, :, self.joint_mask_upper.astype(bool)]
-        tar_pose_upper = rc.axis_angle_to_matrix(tar_pose_upper.reshape(bs, n, 13, 3))
-        tar_pose_upper = rc.matrix_to_rotation_6d(tar_pose_upper).reshape(bs, n, 13*6)
-
-        tar_pose_leg = tar_pose[:, :, self.joint_mask_lower.astype(bool)]
-        tar_pose_leg = rc.axis_angle_to_matrix(tar_pose_leg.reshape(bs, n, 9, 3))
-        tar_pose_leg = rc.matrix_to_rotation_6d(tar_pose_leg).reshape(bs, n, 9*6)
-
-        tar_pose_lower = tar_pose_leg
-
-
-        tar4dis = torch.cat([tar_pose_jaw, tar_pose_upper, tar_pose_hands, tar_pose_leg], dim=2)
-
-      
-        if self.args.pose_norm:
-            tar_pose_upper = (tar_pose_upper - self.mean_upper) / self.std_upper
-            tar_pose_hands = (tar_pose_hands - self.mean_hands) / self.std_hands
-            tar_pose_lower = (tar_pose_lower - self.mean_lower) / self.std_lower
-        
-        if self.use_trans:
-            tar_trans_v = (tar_trans_v - self.trans_mean)/self.trans_std
-            tar_pose_lower = torch.cat([tar_pose_lower,tar_trans_v], dim=-1)
-      
-        latent_face_top = None#self.vq_model_face.map2latent(tar_pose_face) # bs*n/4
-        latent_upper_top = self.vq_model_upper.map2latent(tar_pose_upper)
-        latent_hands_top = self.vq_model_hands.map2latent(tar_pose_hands)
-        latent_lower_top = self.vq_model_lower.map2latent(tar_pose_lower)
-        
-        latent_in = torch.cat([latent_upper_top, latent_hands_top, latent_lower_top], dim=2)/self.args.vqvae_latent_scale
-        
-        
-        tar_pose_6d = rc.axis_angle_to_matrix(tar_pose.reshape(bs, n, 55, 3))
-        tar_pose_6d = rc.matrix_to_rotation_6d(tar_pose_6d).reshape(bs, n, 55*6)
-        latent_all = torch.cat([tar_pose_6d, tar_trans, tar_contact], dim=-1)
-        style_feature = None
-        if self.args.use_motionclip:
-            motionclip_feat = tar_pose_6d[...,:22*6]
-            batch = {}
-            bs,seq,feat = motionclip_feat.shape
-            batch['x']=motionclip_feat.permute(0,2,1).contiguous()
-            batch['y']=torch.zeros(bs).int().cuda()
-            batch['mask']=torch.ones([bs,seq]).bool().cuda()
-            style_feature = self.motionclip.encoder(batch)['mu'].detach().float()
-            
-        
-        
-        # print(tar_index_value_upper_top.shape, index_in.shape)
-        return {
-            "tar_pose_jaw": tar_pose_jaw,
-            "tar_pose_face": tar_pose_face,
-            "tar_pose_upper": tar_pose_upper,
-            "tar_pose_lower": tar_pose_lower,
-            "tar_pose_hands": tar_pose_hands,
-            'tar_pose_leg': tar_pose_leg,
-            "in_audio": in_audio,
-            "in_word": in_word,
-            "tar_trans": tar_trans,
-            "tar_exps": tar_exps,
-            "tar_beta": tar_beta,
-            "tar_pose": tar_pose,
-            "tar4dis": tar4dis,
-            "latent_face_top": latent_face_top,
-            "latent_upper_top": latent_upper_top,
-            "latent_hands_top": latent_hands_top,
-            "latent_lower_top": latent_lower_top,
-            "latent_in":  latent_in,
-            "tar_id": tar_id,
-            "latent_all": latent_all,
-            "tar_pose_6d": tar_pose_6d,
-            "tar_contact": tar_contact,
-            "style_feature":style_feature,
-        }
-    
-    def _g_test(self, loaded_data):
-        sample_fn = self.diffusion.p_sample_loop
-        if self.args.use_ddim:
-            sample_fn = self.diffusion.ddim_sample_loop
-        mode = 'test'
-        bs, n, j = loaded_data["tar_pose"].shape[0], loaded_data["tar_pose"].shape[1], self.joints 
-        tar_pose = loaded_data["tar_pose"]
-        tar_beta = loaded_data["tar_beta"]
-        tar_exps = loaded_data["tar_exps"]
-        tar_contact = loaded_data["tar_contact"]
-        tar_trans = loaded_data["tar_trans"]
-        in_word = loaded_data["in_word"]
-        in_audio = loaded_data["in_audio"]
-        in_x0 = loaded_data['latent_in']
-        in_seed = loaded_data['latent_in']
-        
-        remain = n%8
-        if remain != 0:
-            tar_pose = tar_pose[:, :-remain, :]
-            tar_beta = tar_beta[:, :-remain, :]
-            tar_trans = tar_trans[:, :-remain, :]
-            in_word = in_word[:, :-remain]
-            tar_exps = tar_exps[:, :-remain, :]
-            tar_contact = tar_contact[:, :-remain, :]
-            in_x0 = in_x0[:, :in_x0.shape[1]-(remain//self.args.vqvae_squeeze_scale), :]
-            in_seed = in_seed[:, :in_x0.shape[1]-(remain//self.args.vqvae_squeeze_scale), :]
-            n = n - remain
-
-        tar_pose_jaw = tar_pose[:, :, 66:69]
-        tar_pose_jaw = rc.axis_angle_to_matrix(tar_pose_jaw.reshape(bs, n, 1, 3))
-        tar_pose_jaw = rc.matrix_to_rotation_6d(tar_pose_jaw).reshape(bs, n, 1*6)
-        tar_pose_face = torch.cat([tar_pose_jaw, tar_exps], dim=2)
-
-        tar_pose_hands = tar_pose[:, :, 25*3:55*3]
-        tar_pose_hands = rc.axis_angle_to_matrix(tar_pose_hands.reshape(bs, n, 30, 3))
-        tar_pose_hands = rc.matrix_to_rotation_6d(tar_pose_hands).reshape(bs, n, 30*6)
-
-        tar_pose_upper = tar_pose[:, :, self.joint_mask_upper.astype(bool)]
-        tar_pose_upper = rc.axis_angle_to_matrix(tar_pose_upper.reshape(bs, n, 13, 3))
-        tar_pose_upper = rc.matrix_to_rotation_6d(tar_pose_upper).reshape(bs, n, 13*6)
-
-        tar_pose_leg = tar_pose[:, :, self.joint_mask_lower.astype(bool)]
-        tar_pose_leg = rc.axis_angle_to_matrix(tar_pose_leg.reshape(bs, n, 9, 3))
-        tar_pose_leg = rc.matrix_to_rotation_6d(tar_pose_leg).reshape(bs, n, 9*6)
-        tar_pose_lower = torch.cat([tar_pose_leg, tar_trans, tar_contact], dim=2)
-        
-        tar_pose_6d = rc.axis_angle_to_matrix(tar_pose.reshape(bs, n, 55, 3))
-        tar_pose_6d = rc.matrix_to_rotation_6d(tar_pose_6d).reshape(bs, n, 55*6)
-        latent_all = torch.cat([tar_pose_6d, tar_trans, tar_contact], dim=-1)
-        
-        rec_all_face = []
-        rec_all_upper = []
-        rec_all_lower = []
-        rec_all_hands = []
-        vqvae_squeeze_scale = self.args.vqvae_squeeze_scale
-        roundt = (n - self.args.pre_frames * vqvae_squeeze_scale) // (self.args.pose_length - self.args.pre_frames * vqvae_squeeze_scale)
-        remain = (n - self.args.pre_frames * vqvae_squeeze_scale) % (self.args.pose_length - self.args.pre_frames * vqvae_squeeze_scale)
-        round_l = self.args.pose_length - self.args.pre_frames * vqvae_squeeze_scale
-         
-
-        for i in range(0, roundt):
-            in_word_tmp = in_word[:, i*(round_l):(i+1)*(round_l)+self.args.pre_frames * vqvae_squeeze_scale]
-
-            in_audio_tmp = in_audio[:, i*(16000//30*round_l):(i+1)*(16000//30*round_l)+16000//30*self.args.pre_frames * vqvae_squeeze_scale]
-            in_id_tmp = loaded_data['tar_id'][:, i*(round_l):(i+1)*(round_l)+self.args.pre_frames]
-            in_seed_tmp = in_seed[:, i*(round_l)//vqvae_squeeze_scale:(i+1)*(round_l)//vqvae_squeeze_scale+self.args.pre_frames]
-            in_x0_tmp = in_x0[:, i*(round_l)//vqvae_squeeze_scale:(i+1)*(round_l)//vqvae_squeeze_scale+self.args.pre_frames]
-            mask_val = torch.ones(bs, self.args.pose_length, self.args.pose_dims+3+4).float().cuda()
-            mask_val[:, :self.args.pre_frames, :] = 0.0
-            if i == 0:
-                in_seed_tmp = in_seed_tmp[:, :self.args.pre_frames, :]
-            else:
-                in_seed_tmp = last_sample[:, -self.args.pre_frames:, :]
-
-            cond_ = {'y':{}}
-            cond_['y']['audio'] = in_audio_tmp
-            cond_['y']['word'] = in_word_tmp
-            cond_['y']['id'] = in_id_tmp
-            cond_['y']['seed'] =in_seed_tmp
-            cond_['y']['mask'] = (torch.zeros([self.args.batch_size, 1, 1, self.args.pose_length]) < 1).cuda()
-            
-
-        
-            cond_['y']['style_feature'] = torch.zeros([bs, 512]).cuda()
-
-            shape_ = (bs, 1536, 1, 32)
-            sample = sample_fn(
-                self.model,
-                shape_,
-                clip_denoised=False,
-                model_kwargs=cond_,
-                skip_timesteps=0,
-                init_image=None,
-                progress=True,
-                dump_steps=None,
-                noise=None,
-                const_noise=False,
-            )
-            sample = sample.squeeze().permute(1,0).unsqueeze(0)
-
-            last_sample = sample.clone()
-            
-            rec_latent_upper = sample[...,:512]
-            rec_latent_hands = sample[...,512:1024]
-            rec_latent_lower = sample[...,1024:1536]
-            
-           
-
-            if i == 0:
-                rec_all_upper.append(rec_latent_upper)
-                rec_all_hands.append(rec_latent_hands)
-                rec_all_lower.append(rec_latent_lower)
-            else:
-                rec_all_upper.append(rec_latent_upper[:, self.args.pre_frames:])
-                rec_all_hands.append(rec_latent_hands[:, self.args.pre_frames:])
-                rec_all_lower.append(rec_latent_lower[:, self.args.pre_frames:])
-
-        rec_all_upper = torch.cat(rec_all_upper, dim=1) * self.vqvae_latent_scale
-        rec_all_hands = torch.cat(rec_all_hands, dim=1) * self.vqvae_latent_scale
-        rec_all_lower = torch.cat(rec_all_lower, dim=1) * self.vqvae_latent_scale
-
-        rec_upper = self.vq_model_upper.latent2origin(rec_all_upper)[0]
-        rec_hands = self.vq_model_hands.latent2origin(rec_all_hands)[0]
-        rec_lower = self.vq_model_lower.latent2origin(rec_all_lower)[0]
-        
-        
-        if self.use_trans:
-            rec_trans_v = rec_lower[...,-3:]
-            rec_trans_v = rec_trans_v * self.trans_std + self.trans_mean
-            rec_trans = torch.zeros_like(rec_trans_v)
-            rec_trans = torch.cumsum(rec_trans_v, dim=-2)
-            rec_trans[...,1]=rec_trans_v[...,1]
-            rec_lower = rec_lower[...,:-3]
-        
-        if self.args.pose_norm:
-            rec_upper = rec_upper * self.std_upper + self.mean_upper
-            rec_hands = rec_hands * self.std_hands + self.mean_hands
-            rec_lower = rec_lower * self.std_lower + self.mean_lower
-
-
-
-
-        n = n - remain
-        tar_pose = tar_pose[:, :n, :]
-        tar_exps = tar_exps[:, :n, :]
-        tar_trans = tar_trans[:, :n, :]
-        tar_beta = tar_beta[:, :n, :]
-
-
-        rec_exps = tar_exps
-        #rec_pose_jaw = rec_face[:, :, :6]
-        rec_pose_legs = rec_lower[:, :, :54]
-        bs, n = rec_pose_legs.shape[0], rec_pose_legs.shape[1]
-        rec_pose_upper = rec_upper.reshape(bs, n, 13, 6)
-        rec_pose_upper = rc.rotation_6d_to_matrix(rec_pose_upper)#
-        rec_pose_upper = rc.matrix_to_axis_angle(rec_pose_upper).reshape(bs*n, 13*3)
-        rec_pose_upper_recover = self.inverse_selection_tensor(rec_pose_upper, self.joint_mask_upper, bs*n)
-        rec_pose_lower = rec_pose_legs.reshape(bs, n, 9, 6)
-        rec_pose_lower = rc.rotation_6d_to_matrix(rec_pose_lower)
-        rec_lower2global = rc.matrix_to_rotation_6d(rec_pose_lower.clone()).reshape(bs, n, 9*6)
-        rec_pose_lower = rc.matrix_to_axis_angle(rec_pose_lower).reshape(bs*n, 9*3)
-        rec_pose_lower_recover = self.inverse_selection_tensor(rec_pose_lower, self.joint_mask_lower, bs*n)
-        rec_pose_hands = rec_hands.reshape(bs, n, 30, 6)
-        rec_pose_hands = rc.rotation_6d_to_matrix(rec_pose_hands)
-        rec_pose_hands = rc.matrix_to_axis_angle(rec_pose_hands).reshape(bs*n, 30*3)
-        rec_pose_hands_recover = self.inverse_selection_tensor(rec_pose_hands, self.joint_mask_hands, bs*n)
-        rec_pose = rec_pose_upper_recover + rec_pose_lower_recover + rec_pose_hands_recover 
-        rec_pose[:, 66:69] = tar_pose.reshape(bs*n, 55*3)[:, 66:69]
-
-        rec_pose = rc.axis_angle_to_matrix(rec_pose.reshape(bs*n, j, 3))
-        rec_pose = rc.matrix_to_rotation_6d(rec_pose).reshape(bs, n, j*6)
-        tar_pose = rc.axis_angle_to_matrix(tar_pose.reshape(bs*n, j, 3))
-        tar_pose = rc.matrix_to_rotation_6d(tar_pose).reshape(bs, n, j*6)
-        
-        return {
-            'rec_pose': rec_pose,
-            'rec_trans': rec_trans,
-            'tar_pose': tar_pose,
-            'tar_exps': tar_exps,
-            'tar_beta': tar_beta,
-            'tar_trans': tar_trans,
-            'rec_exps': rec_exps,
-        }
-
-
-    def _create_cuda_model(self):
-        args = self.args
-        other_tools.load_checkpoints(self.model, args.test_ckpt, args.g_name)
-        args.num_quantizers = 6
-        args.shared_codebook =  False
-        args.quantize_dropout_prob = 0.2
-        args.mu = 0.99
-
-        args.nb_code = 512
-        args.code_dim = 512
-        args.code_dim = 512
-        args.down_t = 2
-        args.stride_t = 2
-        args.width = 512
-        args.depth = 3
-        args.dilation_growth_rate = 3
-        args.vq_act = "relu"
-        args.vq_norm = None
-                            
-        dim_pose = 78  
-        args.body_part = "upper"
-        self.vq_model_upper = RVQVAE(args,
-                            dim_pose,
-                            args.nb_code,
-                            args.code_dim,
-                            args.code_dim,
-                            args.down_t,
-                            args.stride_t,
-                            args.width,
-                            args.depth,
-                            args.dilation_growth_rate,
-                            args.vq_act,
-                            args.vq_norm)
-
-        dim_pose = 180
-        args.body_part = "hands"
-        self.vq_model_hands = RVQVAE(args,
-                            dim_pose,
-                            args.nb_code,
-                            args.code_dim,
-                            args.code_dim,
-                            args.down_t,
-                            args.stride_t,
-                            args.width,
-                            args.depth,
-                            args.dilation_growth_rate,
-                            args.vq_act,
-                            args.vq_norm)
-
-        dim_pose = 54
-        if args.use_trans:
-            dim_pose = 57
-            self.args.vqvae_lower_path = self.args.vqvae_lower_trans_path
-        args.body_part = "lower"
-        self.vq_model_lower = RVQVAE(args,
-                            dim_pose,
-                            args.nb_code,
-                            args.code_dim,
-                            args.code_dim,
-                            args.down_t,
-                            args.stride_t,
-                            args.width,
-                            args.depth,
-                            args.dilation_growth_rate,
-                            args.vq_act,
-                            args.vq_norm)
-                
-        self.vq_model_upper.load_state_dict(torch.load(self.args.vqvae_upper_path)['net'])
-        self.vq_model_hands.load_state_dict(torch.load(self.args.vqvae_hands_path)['net'])
-        self.vq_model_lower.load_state_dict(torch.load(self.args.vqvae_lower_path)['net'])
-        
-        self.vqvae_latent_scale = self.args.vqvae_latent_scale 
-
-        self.vq_model_upper.eval().to(self.rank)
-        self.vq_model_hands.eval().to(self.rank)
-        self.vq_model_lower.eval().to(self.rank)
-            
-        self.model = self.model.cuda()
-        self.model.eval()
-        
-        self.mean_upper = torch.from_numpy(self.mean_upper).cuda()
-        self.mean_hands = torch.from_numpy(self.mean_hands).cuda()
-        self.mean_lower = torch.from_numpy(self.mean_lower).cuda()
-        self.std_upper = torch.from_numpy(self.std_upper).cuda()
-        self.std_hands = torch.from_numpy(self.std_hands).cuda()
-        self.std_lower = torch.from_numpy(self.std_lower).cuda()
-        self.trans_mean = torch.from_numpy(self.trans_mean).cuda()
-        self.trans_std = torch.from_numpy(self.trans_std).cuda()
-        
-    @spaces.GPU(duration=149)
-    def _warp(self, batch_data):
-        self._create_cuda_model()
-        
-        
-        loaded_data = self._load_data(batch_data)    
-        net_out = self._g_test(loaded_data)
-        return net_out
         
     def test_demo(self, epoch):
         '''
@@ -644,7 +267,7 @@ class BaseTrainer(object):
             for its, batch_data in enumerate(self.test_loader):
                 # loaded_data = self._load_data(batch_data)    
                 # net_out = self._g_test(loaded_data)
-                net_out = self._warp(batch_data)
+                net_out = _warp(self.args,self.model, batch_data,self.joints,self.joint_mask_upper,self.joint_mask_hands,self.joint_mask_lower,self.use_trans,self.diffusion)
                 tar_pose = net_out['tar_pose']
                 rec_pose = net_out['rec_pose']
                 tar_exps = net_out['tar_exps']
@@ -708,7 +331,402 @@ class BaseTrainer(object):
         end_time = time.time() - start_time
         logger.info(f"total inference time: {int(end_time)} s for {int(total_length/self.args.pose_fps)} s motion")
         return result
-       
+
+
+@spaces.GPU(duration=149)
+def _warp(args,model, batch_data,joints,joint_mask_upper,joint_mask_hands,joint_mask_lower,use_trans,diffusion):
+    args,model,vq_model_upper,vq_model_hands,vq_model_lower,mean_upper,mean_hands,mean_lower,std_upper,std_hands,std_lower,trans_mean,trans_std,vqvae_latent_scale=_warp_create_cuda_model(args,model)
+    
+    
+    loaded_data = _warp_load_data(
+        batch_data,joints,joint_mask_upper,joint_mask_hands,joint_mask_lower,args,use_trans,mean_upper,mean_hands,mean_lower,std_upper,std_hands,std_lower,trans_mean,trans_std,vq_model_upper,vq_model_hands,vq_model_lower        
+    )    
+    net_out = _warp_g_test(loaded_data,diffusion,args,joints,joint_mask_upper,joint_mask_hands,joint_mask_lower,model,vqvae_latent_scale,vq_model_upper,vq_model_hands,vq_model_lower,use_trans,trans_std,trans_mean,std_upper,std_hands,std_lower,mean_upper,mean_hands,mean_lower)
+    return net_out
+    
+def _warp_inverse_selection_tensor(filtered_t, selection_array, n):
+    selection_array = torch.from_numpy(selection_array).cuda()
+    original_shape_t = torch.zeros((n, 165)).cuda()
+    selected_indices = torch.where(selection_array == 1)[0]
+    for i in range(n):
+        original_shape_t[i, selected_indices] = filtered_t[i]
+    return original_shape_t
+
+def _warp_g_test(loaded_data,diffusion,args,joints,joint_mask_upper,joint_mask_hands,joint_mask_lower,model,vqvae_latent_scale,vq_model_upper,vq_model_hands,vq_model_lower,use_trans,trans_std,trans_mean,std_upper,std_hands,std_lower,mean_upper,mean_hands,mean_lower):
+    sample_fn = diffusion.p_sample_loop
+    if args.use_ddim:
+        sample_fn = diffusion.ddim_sample_loop
+    mode = 'test'
+    bs, n, j = loaded_data["tar_pose"].shape[0], loaded_data["tar_pose"].shape[1], joints 
+    tar_pose = loaded_data["tar_pose"]
+    tar_beta = loaded_data["tar_beta"]
+    tar_exps = loaded_data["tar_exps"]
+    tar_contact = loaded_data["tar_contact"]
+    tar_trans = loaded_data["tar_trans"]
+    in_word = loaded_data["in_word"]
+    in_audio = loaded_data["in_audio"]
+    in_x0 = loaded_data['latent_in']
+    in_seed = loaded_data['latent_in']
+    
+    remain = n%8
+    if remain != 0:
+        tar_pose = tar_pose[:, :-remain, :]
+        tar_beta = tar_beta[:, :-remain, :]
+        tar_trans = tar_trans[:, :-remain, :]
+        in_word = in_word[:, :-remain]
+        tar_exps = tar_exps[:, :-remain, :]
+        tar_contact = tar_contact[:, :-remain, :]
+        in_x0 = in_x0[:, :in_x0.shape[1]-(remain//args.vqvae_squeeze_scale), :]
+        in_seed = in_seed[:, :in_x0.shape[1]-(remain//args.vqvae_squeeze_scale), :]
+        n = n - remain
+
+    tar_pose_jaw = tar_pose[:, :, 66:69]
+    tar_pose_jaw = rc.axis_angle_to_matrix(tar_pose_jaw.reshape(bs, n, 1, 3))
+    tar_pose_jaw = rc.matrix_to_rotation_6d(tar_pose_jaw).reshape(bs, n, 1*6)
+    tar_pose_face = torch.cat([tar_pose_jaw, tar_exps], dim=2)
+
+    tar_pose_hands = tar_pose[:, :, 25*3:55*3]
+    tar_pose_hands = rc.axis_angle_to_matrix(tar_pose_hands.reshape(bs, n, 30, 3))
+    tar_pose_hands = rc.matrix_to_rotation_6d(tar_pose_hands).reshape(bs, n, 30*6)
+
+    tar_pose_upper = tar_pose[:, :, joint_mask_upper.astype(bool)]
+    tar_pose_upper = rc.axis_angle_to_matrix(tar_pose_upper.reshape(bs, n, 13, 3))
+    tar_pose_upper = rc.matrix_to_rotation_6d(tar_pose_upper).reshape(bs, n, 13*6)
+
+    tar_pose_leg = tar_pose[:, :, joint_mask_lower.astype(bool)]
+    tar_pose_leg = rc.axis_angle_to_matrix(tar_pose_leg.reshape(bs, n, 9, 3))
+    tar_pose_leg = rc.matrix_to_rotation_6d(tar_pose_leg).reshape(bs, n, 9*6)
+    tar_pose_lower = torch.cat([tar_pose_leg, tar_trans, tar_contact], dim=2)
+    
+    tar_pose_6d = rc.axis_angle_to_matrix(tar_pose.reshape(bs, n, 55, 3))
+    tar_pose_6d = rc.matrix_to_rotation_6d(tar_pose_6d).reshape(bs, n, 55*6)
+    latent_all = torch.cat([tar_pose_6d, tar_trans, tar_contact], dim=-1)
+    
+    rec_all_face = []
+    rec_all_upper = []
+    rec_all_lower = []
+    rec_all_hands = []
+    vqvae_squeeze_scale = args.vqvae_squeeze_scale
+    roundt = (n - args.pre_frames * vqvae_squeeze_scale) // (args.pose_length - args.pre_frames * vqvae_squeeze_scale)
+    remain = (n - args.pre_frames * vqvae_squeeze_scale) % (args.pose_length - args.pre_frames * vqvae_squeeze_scale)
+    round_l = args.pose_length - args.pre_frames * vqvae_squeeze_scale
+        
+
+    for i in range(0, roundt):
+        in_word_tmp = in_word[:, i*(round_l):(i+1)*(round_l)+args.pre_frames * vqvae_squeeze_scale]
+
+        in_audio_tmp = in_audio[:, i*(16000//30*round_l):(i+1)*(16000//30*round_l)+16000//30*args.pre_frames * vqvae_squeeze_scale]
+        in_id_tmp = loaded_data['tar_id'][:, i*(round_l):(i+1)*(round_l)+args.pre_frames]
+        in_seed_tmp = in_seed[:, i*(round_l)//vqvae_squeeze_scale:(i+1)*(round_l)//vqvae_squeeze_scale+args.pre_frames]
+        in_x0_tmp = in_x0[:, i*(round_l)//vqvae_squeeze_scale:(i+1)*(round_l)//vqvae_squeeze_scale+args.pre_frames]
+        mask_val = torch.ones(bs, args.pose_length, args.pose_dims+3+4).float().cuda()
+        mask_val[:, :args.pre_frames, :] = 0.0
+        if i == 0:
+            in_seed_tmp = in_seed_tmp[:, :args.pre_frames, :]
+        else:
+            in_seed_tmp = last_sample[:, -args.pre_frames:, :]
+
+        cond_ = {'y':{}}
+        cond_['y']['audio'] = in_audio_tmp
+        cond_['y']['word'] = in_word_tmp
+        cond_['y']['id'] = in_id_tmp
+        cond_['y']['seed'] =in_seed_tmp
+        cond_['y']['mask'] = (torch.zeros([args.batch_size, 1, 1, args.pose_length]) < 1).cuda()
+        
+
+    
+        cond_['y']['style_feature'] = torch.zeros([bs, 512]).cuda()
+
+        shape_ = (bs, 1536, 1, 32)
+        sample = sample_fn(
+            model,
+            shape_,
+            clip_denoised=False,
+            model_kwargs=cond_,
+            skip_timesteps=0,
+            init_image=None,
+            progress=True,
+            dump_steps=None,
+            noise=None,
+            const_noise=False,
+        )
+        sample = sample.squeeze().permute(1,0).unsqueeze(0)
+
+        last_sample = sample.clone()
+        
+        rec_latent_upper = sample[...,:512]
+        rec_latent_hands = sample[...,512:1024]
+        rec_latent_lower = sample[...,1024:1536]
+        
+        
+
+        if i == 0:
+            rec_all_upper.append(rec_latent_upper)
+            rec_all_hands.append(rec_latent_hands)
+            rec_all_lower.append(rec_latent_lower)
+        else:
+            rec_all_upper.append(rec_latent_upper[:, args.pre_frames:])
+            rec_all_hands.append(rec_latent_hands[:, args.pre_frames:])
+            rec_all_lower.append(rec_latent_lower[:, args.pre_frames:])
+
+    rec_all_upper = torch.cat(rec_all_upper, dim=1) * vqvae_latent_scale
+    rec_all_hands = torch.cat(rec_all_hands, dim=1) * vqvae_latent_scale
+    rec_all_lower = torch.cat(rec_all_lower, dim=1) * vqvae_latent_scale
+
+    rec_upper = vq_model_upper.latent2origin(rec_all_upper)[0]
+    rec_hands = vq_model_hands.latent2origin(rec_all_hands)[0]
+    rec_lower = vq_model_lower.latent2origin(rec_all_lower)[0]
+    
+    
+    if use_trans:
+        rec_trans_v = rec_lower[...,-3:]
+        rec_trans_v = rec_trans_v * trans_std + trans_mean
+        rec_trans = torch.zeros_like(rec_trans_v)
+        rec_trans = torch.cumsum(rec_trans_v, dim=-2)
+        rec_trans[...,1]=rec_trans_v[...,1]
+        rec_lower = rec_lower[...,:-3]
+    
+    if args.pose_norm:
+        rec_upper = rec_upper * std_upper + mean_upper
+        rec_hands = rec_hands * std_hands + mean_hands
+        rec_lower = rec_lower * std_lower + mean_lower
+
+
+
+
+    n = n - remain
+    tar_pose = tar_pose[:, :n, :]
+    tar_exps = tar_exps[:, :n, :]
+    tar_trans = tar_trans[:, :n, :]
+    tar_beta = tar_beta[:, :n, :]
+
+
+    rec_exps = tar_exps
+    #rec_pose_jaw = rec_face[:, :, :6]
+    rec_pose_legs = rec_lower[:, :, :54]
+    bs, n = rec_pose_legs.shape[0], rec_pose_legs.shape[1]
+    rec_pose_upper = rec_upper.reshape(bs, n, 13, 6)
+    rec_pose_upper = rc.rotation_6d_to_matrix(rec_pose_upper)#
+    rec_pose_upper = rc.matrix_to_axis_angle(rec_pose_upper).reshape(bs*n, 13*3)
+    rec_pose_upper_recover = _warp_inverse_selection_tensor(rec_pose_upper, joint_mask_upper, bs*n)
+    rec_pose_lower = rec_pose_legs.reshape(bs, n, 9, 6)
+    rec_pose_lower = rc.rotation_6d_to_matrix(rec_pose_lower)
+    rec_lower2global = rc.matrix_to_rotation_6d(rec_pose_lower.clone()).reshape(bs, n, 9*6)
+    rec_pose_lower = rc.matrix_to_axis_angle(rec_pose_lower).reshape(bs*n, 9*3)
+    rec_pose_lower_recover = _warp_inverse_selection_tensor(rec_pose_lower, joint_mask_lower, bs*n)
+    rec_pose_hands = rec_hands.reshape(bs, n, 30, 6)
+    rec_pose_hands = rc.rotation_6d_to_matrix(rec_pose_hands)
+    rec_pose_hands = rc.matrix_to_axis_angle(rec_pose_hands).reshape(bs*n, 30*3)
+    rec_pose_hands_recover = _warp_inverse_selection_tensor(rec_pose_hands, joint_mask_hands, bs*n)
+    rec_pose = rec_pose_upper_recover + rec_pose_lower_recover + rec_pose_hands_recover 
+    rec_pose[:, 66:69] = tar_pose.reshape(bs*n, 55*3)[:, 66:69]
+
+    rec_pose = rc.axis_angle_to_matrix(rec_pose.reshape(bs*n, j, 3))
+    rec_pose = rc.matrix_to_rotation_6d(rec_pose).reshape(bs, n, j*6)
+    tar_pose = rc.axis_angle_to_matrix(tar_pose.reshape(bs*n, j, 3))
+    tar_pose = rc.matrix_to_rotation_6d(tar_pose).reshape(bs, n, j*6)
+    
+    return {
+        'rec_pose': rec_pose,
+        'rec_trans': rec_trans,
+        'tar_pose': tar_pose,
+        'tar_exps': tar_exps,
+        'tar_beta': tar_beta,
+        'tar_trans': tar_trans,
+        'rec_exps': rec_exps,
+    }
+
+
+
+def _warp_load_data(dict_data,joints,joint_mask_upper,joint_mask_hands,joint_mask_lower,args,use_trans,mean_upper,mean_hands,mean_lower,std_upper,std_hands,std_lower,trans_mean,trans_std,vq_model_upper,vq_model_hands,vq_model_lower):
+    tar_pose_raw = dict_data["pose"]
+    tar_pose = tar_pose_raw[:, :, :165].cuda()
+    tar_contact = tar_pose_raw[:, :, 165:169].cuda()
+    tar_trans = dict_data["trans"].cuda()
+    tar_trans_v = dict_data["trans_v"].cuda()
+    tar_exps = dict_data["facial"].cuda()
+    in_audio = dict_data["audio"].cuda() 
+    in_word = dict_data["word"].cuda()
+    tar_beta = dict_data["beta"].cuda()
+    tar_id = dict_data["id"].cuda().long()
+    bs, n, j = tar_pose.shape[0], tar_pose.shape[1], joints
+
+    tar_pose_jaw = tar_pose[:, :, 66:69]
+    tar_pose_jaw = rc.axis_angle_to_matrix(tar_pose_jaw.reshape(bs, n, 1, 3))
+    tar_pose_jaw = rc.matrix_to_rotation_6d(tar_pose_jaw).reshape(bs, n, 1*6)
+    tar_pose_face = torch.cat([tar_pose_jaw, tar_exps], dim=2)
+
+    tar_pose_hands = tar_pose[:, :, 25*3:55*3]
+    tar_pose_hands = rc.axis_angle_to_matrix(tar_pose_hands.reshape(bs, n, 30, 3))
+    tar_pose_hands = rc.matrix_to_rotation_6d(tar_pose_hands).reshape(bs, n, 30*6)
+
+    tar_pose_upper = tar_pose[:, :, joint_mask_upper.astype(bool)]
+    tar_pose_upper = rc.axis_angle_to_matrix(tar_pose_upper.reshape(bs, n, 13, 3))
+    tar_pose_upper = rc.matrix_to_rotation_6d(tar_pose_upper).reshape(bs, n, 13*6)
+
+    tar_pose_leg = tar_pose[:, :, joint_mask_lower.astype(bool)]
+    tar_pose_leg = rc.axis_angle_to_matrix(tar_pose_leg.reshape(bs, n, 9, 3))
+    tar_pose_leg = rc.matrix_to_rotation_6d(tar_pose_leg).reshape(bs, n, 9*6)
+
+    tar_pose_lower = tar_pose_leg
+
+
+    tar4dis = torch.cat([tar_pose_jaw, tar_pose_upper, tar_pose_hands, tar_pose_leg], dim=2)
+
+    
+    if args.pose_norm:
+        tar_pose_upper = (tar_pose_upper - mean_upper) / std_upper
+        tar_pose_hands = (tar_pose_hands - mean_hands) / std_hands
+        tar_pose_lower = (tar_pose_lower - mean_lower) / std_lower
+    
+    if use_trans:
+        tar_trans_v = (tar_trans_v - trans_mean)/trans_std
+        tar_pose_lower = torch.cat([tar_pose_lower,tar_trans_v], dim=-1)
+    
+    latent_face_top = None#self.vq_model_face.map2latent(tar_pose_face) # bs*n/4
+    latent_upper_top = vq_model_upper.map2latent(tar_pose_upper)
+    latent_hands_top = vq_model_hands.map2latent(tar_pose_hands)
+    latent_lower_top = vq_model_lower.map2latent(tar_pose_lower)
+    
+    latent_in = torch.cat([latent_upper_top, latent_hands_top, latent_lower_top], dim=2)/args.vqvae_latent_scale
+    
+    
+    tar_pose_6d = rc.axis_angle_to_matrix(tar_pose.reshape(bs, n, 55, 3))
+    tar_pose_6d = rc.matrix_to_rotation_6d(tar_pose_6d).reshape(bs, n, 55*6)
+    latent_all = torch.cat([tar_pose_6d, tar_trans, tar_contact], dim=-1)
+    style_feature = None
+    if args.use_motionclip:
+        motionclip_feat = tar_pose_6d[...,:22*6]
+        batch = {}
+        bs,seq,feat = motionclip_feat.shape
+        batch['x']=motionclip_feat.permute(0,2,1).contiguous()
+        batch['y']=torch.zeros(bs).int().cuda()
+        batch['mask']=torch.ones([bs,seq]).bool().cuda()
+        style_feature = motionclip.encoder(batch)['mu'].detach().float()
+        
+    
+    
+    # print(tar_index_value_upper_top.shape, index_in.shape)
+    return {
+        "tar_pose_jaw": tar_pose_jaw,
+        "tar_pose_face": tar_pose_face,
+        "tar_pose_upper": tar_pose_upper,
+        "tar_pose_lower": tar_pose_lower,
+        "tar_pose_hands": tar_pose_hands,
+        'tar_pose_leg': tar_pose_leg,
+        "in_audio": in_audio,
+        "in_word": in_word,
+        "tar_trans": tar_trans,
+        "tar_exps": tar_exps,
+        "tar_beta": tar_beta,
+        "tar_pose": tar_pose,
+        "tar4dis": tar4dis,
+        "latent_face_top": latent_face_top,
+        "latent_upper_top": latent_upper_top,
+        "latent_hands_top": latent_hands_top,
+        "latent_lower_top": latent_lower_top,
+        "latent_in":  latent_in,
+        "tar_id": tar_id,
+        "latent_all": latent_all,
+        "tar_pose_6d": tar_pose_6d,
+        "tar_contact": tar_contact,
+        "style_feature":style_feature,
+    }
+
+
+def _warp_create_cuda_model(args,model):
+    args = args
+    other_tools.load_checkpoints(model, args.test_ckpt, args.g_name)
+    args.num_quantizers = 6
+    args.shared_codebook =  False
+    args.quantize_dropout_prob = 0.2
+    args.mu = 0.99
+
+    args.nb_code = 512
+    args.code_dim = 512
+    args.code_dim = 512
+    args.down_t = 2
+    args.stride_t = 2
+    args.width = 512
+    args.depth = 3
+    args.dilation_growth_rate = 3
+    args.vq_act = "relu"
+    args.vq_norm = None
+                        
+    dim_pose = 78  
+    args.body_part = "upper"
+    vq_model_upper = RVQVAE(args,
+                        dim_pose,
+                        args.nb_code,
+                        args.code_dim,
+                        args.code_dim,
+                        args.down_t,
+                        args.stride_t,
+                        args.width,
+                        args.depth,
+                        args.dilation_growth_rate,
+                        args.vq_act,
+                        args.vq_norm)
+
+    dim_pose = 180
+    args.body_part = "hands"
+    vq_model_hands = RVQVAE(args,
+                        dim_pose,
+                        args.nb_code,
+                        args.code_dim,
+                        args.code_dim,
+                        args.down_t,
+                        args.stride_t,
+                        args.width,
+                        args.depth,
+                        args.dilation_growth_rate,
+                        args.vq_act,
+                        args.vq_norm)
+
+    dim_pose = 54
+    if args.use_trans:
+        dim_pose = 57
+        args.vqvae_lower_path = args.vqvae_lower_trans_path
+    args.body_part = "lower"
+    vq_model_lower = RVQVAE(args,
+                        dim_pose,
+                        args.nb_code,
+                        args.code_dim,
+                        args.code_dim,
+                        args.down_t,
+                        args.stride_t,
+                        args.width,
+                        args.depth,
+                        args.dilation_growth_rate,
+                        args.vq_act,
+                        args.vq_norm)
+            
+    vq_model_upper.load_state_dict(torch.load(args.vqvae_upper_path)['net'])
+    vq_model_hands.load_state_dict(torch.load(args.vqvae_hands_path)['net'])
+    vq_model_lower.load_state_dict(torch.load(args.vqvae_lower_path)['net'])
+    
+    vqvae_latent_scale = args.vqvae_latent_scale 
+
+    vq_model_upper.eval().cuda()
+    vq_model_hands.eval().cuda()
+    vq_model_lower.eval().cuda()
+        
+    model = model.cuda()
+    model.eval()
+    
+    mean_upper = torch.from_numpy(mean_upper).cuda()
+    mean_hands = torch.from_numpy(mean_hands).cuda()
+    mean_lower = torch.from_numpy(mean_lower).cuda()
+    std_upper = torch.from_numpy(std_upper).cuda()
+    std_hands = torch.from_numpy(std_hands).cuda()
+    std_lower = torch.from_numpy(std_lower).cuda()
+    trans_mean = torch.from_numpy(trans_mean).cuda()
+    trans_std = torch.from_numpy(trans_std).cuda()
+
+    return args,model,vq_model_upper,vq_model_hands,vq_model_lower,mean_upper,mean_hands,mean_lower,std_upper,std_hands,std_lower,trans_mean,trans_std,vqvae_latent_scale
+
+
+
 @logger.catch
 def syntalker(audio_path,sample_stratege):
     args = config.parse_args()