Update difpoint/inference.py

difpoint/inference.py

@@ -7,6 +7,9 @@
 import glob
 
 import os
+os.environ['HYDRA_FULL_ERROR']='1'
+os.environ['CUDA_VISIBLE_DEVICES'] = '2'
+
 import os
 import time
 import shutil
@@ -54,7 +57,6 @@ import datetime
 import platform
 from omegaconf import OmegaConf
 from difpoint.src.pipelines.faster_live_portrait_pipeline import FasterLivePortraitPipeline
-import spaces
 
 FFMPEG = "ffmpeg"
 
@@ -160,19 +162,19 @@ class Inferencer(object):
         self.device = 'cuda'
         from difpoint.model import get_model
         self.point_diffusion = get_model()
-        ckpt = torch.load('./downloaded_repo/ckpts/KDTalker.pth', weights_only=False)
+        ckpt = torch.load('/home/yinuo/Gradio-UI_copy/difpoint/outputs/2024.08.26_dim_70_frame_64_vox1_selected_d6.5_c8.5/2024-08-26--16-52-34/checkpoint-500000.pth')
 
         self.point_diffusion.load_state_dict(ckpt['model'])
         print('model', self.point_diffusion.children())
         self.point_diffusion.eval()
         self.point_diffusion.to(self.device)
 
-        lm_croper_checkpoint = './downloaded_repo/ckpts/shape_predictor_68_face_landmarks.dat'
+        lm_croper_checkpoint = os.path.join('difpoint/dataset_process/ckpts/', 'shape_predictor_68_face_landmarks.dat')
         self.croper = Croper(lm_croper_checkpoint)
 
         self.norm_info = dict(np.load(r'difpoint/datasets/norm_info_d6.5_c8.5_vox1_train.npz'))
 
-        wav2lip_checkpoint = './downloaded_repo/ckpts/wav2lip.pth'
+        wav2lip_checkpoint = 'difpoint/dataset_process/ckpts/wav2lip.pth'
         self.wav2lip_model = AudioEncoder(wav2lip_checkpoint, 'cuda')
         self.wav2lip_model.cuda()
         self.wav2lip_model.eval()
@@ -270,7 +272,6 @@ class Inferencer(object):
         return combined_lip_ratio_tensor
 
     # 2024.06.26
-    @spaces.GPU
     @torch.no_grad()
     def generate_with_audio_img(self, upload_audio_path, tts_audio_path, audio_type, image_path, smoothed_pitch, smoothed_yaw, smoothed_roll, smoothed_t,  save_path='results'):
         print(audio_type)
@@ -305,12 +306,7 @@ class Inferencer(object):
 
         flag_lip_zero = self.inf_cfg.infer_params.flag_normalize_lip 
     
-        if flag_lip_zero:
-            # let lip-open scalar to be 0 at first
-            c_d_lip_before_animation = [0.]
-            
-            lip_delta_before_animation = self.live_portrait_pipeline.model_dict['stitching_lip_retarget'].predict(
-                                    concat_feat(x_s, combined_lip_ratio_tensor_before_animation))
+       
 
         ######## process driving info ########
         kp_info = {}
@@ -442,30 +438,16 @@ class Inferencer(object):
 
             # Algorithm 1:
             if not self.inf_cfg.infer_params.flag_stitching and not self.inf_cfg.infer_params.flag_eye_retargeting and not self.inf_cfg.infer_params.flag_lip_retargeting:
-                # without stitching or retargeting
-                if flag_lip_zero:
-                    x_d_i_new += lip_delta_before_animation.reshape(-1, x_s.shape[1], 3)
-                else:
-                    pass
+                pass
             elif self.inf_cfg.infer_params.flag_stitching and not self.inf_cfg.infer_params.flag_eye_retargeting and not self.inf_cfg.infer_params.flag_lip_retargeting:
                 # with stitching and without retargeting
-                if flag_lip_zero:
-                    x_d_i_new = self.live_portrait_pipeline.stitching(x_s, x_d_i_new) + lip_delta_before_animation.reshape(
-                            -1, x_s.shape[1], 3)
-                else:
-                    x_d_i_new = self.live_portrait_pipeline.stitching(x_s, x_d_i_new)
+                x_d_i_new = self.live_portrait_pipeline.stitching(x_s, x_d_i_new)
             else:
                 eyes_delta, lip_delta = None, None
                 if self.inf_cfg.infer_params.flag_eye_retargeting:
-                    c_d_eyes_i = template_dct['c_d_eyes_lst'][i]
-                    combined_eye_ratio_tensor = self.calc_combined_eye_ratio(c_d_eyes_i, c_s_eye)
-                    # ∆_eyes,i = R_eyes(x_s; c_s,eyes, c_d,eyes,i)
-                    eyes_delta = self.live_portrait_pipeline.retarget_eye(x_s, combined_eye_ratio_tensor)
+                    pass
                 if self.inf_cfg.infer_params.flag_lip_retargeting:
-                    c_d_lip_i = template_dct['c_d_lip_lst'][i]
-                    combined_lip_ratio_tensor = self.calc_combined_lip_ratio(c_d_lip_i, c_s_lip)
-                    # ∆_lip,i = R_lip(x_s; c_s,lip, c_d,lip,i)
-                    lip_delta = self.live_portrait_pipeline.retarget_lip(x_s, combined_lip_ratio_tensor)
+                    pass
 
                 if self.inf_cfg.infer_params.flag_relative_motion:  # use x_s
                     x_d_i_new = x_s + \