src/datasets/preprocess/extract_features/motion_processer.py

"""
Motion feature extractor 
"""
import os
import os.path as osp
import sys
import pickle
from omegaconf import OmegaConf

import torch

from PIL import Image
import numpy as np
import cv2
import imageio
import pickle
import time
from decord import VideoReader # must after import torch

from rich.progress import track




sys.path.append(osp.dirname(osp.dirname(osp.dirname(osp.dirname(osp.dirname(osp.realpath(__file__)))))))
from src.datasets.preprocess.extract_features.face_segmentation import build_face_parser, get_face_mask, vis_parsing_maps
from src.thirdparty.liveportrait.src.utils.helper import load_model, concat_feat
from src.thirdparty.liveportrait.src.utils.io import load_image_rgb, resize_to_limit, load_video
from src.thirdparty.liveportrait.src.utils.video import get_fps, images2video, add_audio_to_video
from src.thirdparty.liveportrait.src.utils.camera import headpose_pred_to_degree, get_rotation_matrix

from src.thirdparty.liveportrait.src.utils.cropper import Cropper
from src.thirdparty.liveportrait.src.utils.crop import prepare_paste_back, paste_back, paste_back_with_face_mask
from src.thirdparty.liveportrait.src.utils.retargeting_utils import calc_eye_close_ratio, calc_lip_close_ratio
from src.thirdparty.liveportrait.src.utils.helper import mkdir, basename, dct2device, is_image, calc_motion_multiplier
from src.utils.filter import smooth as ksmooth
from src.utils.filter import smooth_

from skimage.metrics import peak_signal_noise_ratio
import warnings


def psnr(imgs1, imgs2):
    psnrs = []
    for img1, img2 in zip(imgs1, imgs2):
        psnr = peak_signal_noise_ratio(img1, img2, data_range=255)
        psnrs.append(psnr)
    return psnrs


def suffix(filename):
    """a.jpg -> jpg"""
    pos = filename.rfind(".")
    if pos == -1:
        return ""
    return filename[pos + 1:]

def dump(wfp, obj):
    wd = osp.split(wfp)[0]
    if wd != "" and not osp.exists(wd):
        mkdir(wd)

    _suffix = suffix(wfp)
    if _suffix == "npy":
        np.save(wfp, obj)
    elif _suffix == "pkl":
        pickle.dump(obj, open(wfp, "wb"))
    else:
        raise Exception("Unknown type: {}".format(_suffix))

def load(fp):
    suffix_ = suffix(fp)

    if suffix_ == "npy":
        return np.load(fp)
    elif suffix_ == "pkl":
        return pickle.load(open(fp, "rb"))
    else:
        raise Exception(f"Unknown type: {suffix}")


def remove_suffix(filepath):
    """a/b/c.jpg -> a/b/c"""
    return osp.join(osp.dirname(filepath), basename(filepath))
    

class MotionProcesser(object):
    def __init__(self, cfg_path, device_id=0) -> None:
        device = f"cuda:{device_id}"
        cfg = OmegaConf.load(cfg_path)
        print(f"Load cfg from {osp.realpath(cfg_path)} done.")
        print(f"=============================== Driven CFG ===============================")
        print(OmegaConf.to_yaml(cfg))
        print(f"=============================== ========== ===============================")
        models_config = OmegaConf.load(cfg.models_config)

        # 1. init appearance feature extractor
        self.appearance_feature_extractor = load_model(
            cfg.appearance_feature_extractor_path, 
            models_config, 
            device, 
            'appearance_feature_extractor'
        )
        print(f'1. Load appearance_feature_extractor from {osp.realpath(cfg.appearance_feature_extractor_path)} done.')

        # 2. # init motion extractor
        self.motion_extractor = load_model(
            cfg.motion_extractor_path, 
            models_config, 
            device, 
            'motion_extractor'
        )
        print(f'2. Load motion_extractor from {osp.realpath(cfg.motion_extractor_path)} done.')
        
        # 3. init S and R
        if cfg.stitching_retargeting_module_path is not None and osp.exists(cfg.stitching_retargeting_module_path):
            self.stitching_retargeting_module = load_model(
                cfg.stitching_retargeting_module_path, 
                models_config, 
                device, 
                'stitching_retargeting_module'
            )
            print(f'3. Load stitching_retargeting_module from {osp.realpath(cfg.stitching_retargeting_module_path)} done.')
        else:
            self.stitching_retargeting_module = None
        
        # 4. init motion warper
        self.warping_module = load_model(
            cfg.warping_module_path, 
            models_config, 
            device, 
            'warping_module'
        )
        print(f"4. Load warping_module from {osp.realpath(cfg.warping_module_path)} done.")

        # 5. init decoder
        self.spade_generator = load_model(
            cfg.spade_generator_path, 
            models_config, 
            device, 
            'spade_generator'
        )
        print(f"Load generator from {osp.realpath(cfg.spade_generator_path)} done.")

        # # Optimize for inference
        self.compile = cfg.flag_do_torch_compile
        if self.compile:
            torch._dynamo.config.suppress_errors = True  # Suppress errors and fall back to eager execution
            self.warping_module = torch.compile(self.warping_module, mode='max-autotune')
            self.spade_generator = torch.compile(self.spade_generator, mode='max-autotune')

        # 6. init cropper
        crop_cfg = OmegaConf.load(cfg.crop_cfg)
        self.cropper = Cropper(crop_cfg=crop_cfg, image_type="human_face", device_id=device_id)
    
        self.cfg = cfg
        self.models_config = models_config
        self.device = device
    

        # 7. load crop mask
        self.mask_crop = cv2.imread(cfg.mask_crop, cv2.IMREAD_COLOR)
        # 8. load lib array
        with open(cfg.lip_array, 'rb') as f:
            self.lip_array = pickle.load(f)

        # 9. load face parser
        self.face_parser, self.to_tensor = build_face_parser(weight_path=cfg.face_parser_weight_path, resnet_weight_path=cfg.resnet_weight_path, device_id=device_id)

    def inference_ctx(self):    
        ctx = torch.autocast(device_type=self.device[:4], dtype=torch.float16,
                                 enabled=self.cfg.flag_use_half_precision)
        return ctx

    @torch.no_grad()
    def extract_feature_3d(self, x: torch.Tensor) -> torch.Tensor:
        """ get the appearance feature of the image by F
        x: Bx3xHxW, normalized to 0~1
        """
        with self.inference_ctx():
            feature_3d = self.appearance_feature_extractor(x)

        return feature_3d.float()

    @torch.no_grad()
    def get_kp_info(self, x: torch.Tensor, **kwargs) -> dict:
        """ get the implicit keypoint information
        x: Bx3xHxW, normalized to 0~1
        flag_refine_info: whether to trandform the pose to degrees and the dimention of the reshape
        return: A dict contains keys: 'pitch', 'yaw', 'roll', 't', 'exp', 'scale', 'kp'
        """
        with self.inference_ctx():
            kp_info = self.motion_extractor(x)

            if self.cfg.flag_use_half_precision:
                # float the dict
                for k, v in kp_info.items():
                    if isinstance(v, torch.Tensor):
                        kp_info[k] = v.float()

        return kp_info

    @torch.no_grad()
    def refine_kp(self, kp_info):
        bs = kp_info['exp'].shape[0]
        kp_info['pitch'] = headpose_pred_to_degree(kp_info['pitch'])[:, None]  # Bx1
        kp_info['yaw'] = headpose_pred_to_degree(kp_info['yaw'])[:, None]  # Bx1
        kp_info['roll'] = headpose_pred_to_degree(kp_info['roll'])[:, None]  # Bx1
        kp_info['exp'] = kp_info['exp'].reshape(bs, -1, 3)  # BxNx3
        if 'kp' in kp_info.keys():
            kp_info['kp'] = kp_info['kp'].reshape(bs, -1, 3)  # BxNx3

        return kp_info

    @torch.no_grad()
    def transform_keypoint(self, kp_info: dict):
        """
        transform the implicit keypoints with the pose, shift, and expression deformation
        kp: BxNx3
        """
        kp = kp_info['kp']    # (bs, k, 3)
        pitch, yaw, roll = kp_info['pitch'], kp_info['yaw'], kp_info['roll']

        t, exp = kp_info['t'], kp_info['exp']
        scale = kp_info['scale']

        pitch = headpose_pred_to_degree(pitch)
        yaw = headpose_pred_to_degree(yaw)
        roll = headpose_pred_to_degree(roll)

        bs = kp.shape[0]
        if kp.ndim == 2:
            num_kp = kp.shape[1] // 3  # Bx(num_kpx3)
        else:
            num_kp = kp.shape[1]  # Bxnum_kpx3

        rot_mat = get_rotation_matrix(pitch, yaw, roll)    # (bs, 3, 3)

        # Eqn.2: s * (R * x_c,s + exp) + t
        kp_transformed = kp.view(bs, num_kp, 3) @ rot_mat + exp.view(bs, num_kp, 3)
        kp_transformed *= scale[..., None]  # (bs, k, 3) * (bs, 1, 1) = (bs, k, 3)
        kp_transformed[:, :, 0:2] += t[:, None, 0:2]  # remove z, only apply tx ty

        return kp_transformed

    @torch.no_grad()
    def stitching(self, kp_source: torch.Tensor, kp_driving: torch.Tensor) -> torch.Tensor:
        """ conduct the stitching
        kp_source: Bxnum_kpx3
        kp_driving: Bxnum_kpx3
        """

        if self.stitching_retargeting_module is not None:
            bs, num_kp = kp_source.shape[:2]
            kp_driving_new = kp_driving.clone()
            # stich
            feat_stiching = concat_feat(kp_source, kp_driving_new)
            delta = self.stitching_retargeting_module['stitching'](feat_stiching) # Bxnum_kpx3

            delta_exp = delta[..., :3*num_kp].reshape(bs, num_kp, 3)  # 1x20x3
            delta_tx_ty = delta[..., 3*num_kp:3*num_kp+2].reshape(bs, 1, 2)  # 1x1x2

            kp_driving_new += delta_exp
            kp_driving_new[..., :2] += delta_tx_ty

            return kp_driving_new

        return kp_driving

    @torch.no_grad()
    def warp_decode(self, feature_3d: torch.Tensor, kp_source: torch.Tensor, kp_driving: torch.Tensor) -> dict[str, torch.Tensor]:
        """ get the image after the warping of the implicit keypoints
        feature_3d: Bx32x16x64x64, feature volume
        kp_source: BxNx3
        kp_driving: BxNx3
        """
        # The line 18 in Algorithm 1: D(W(f_s; x_s, x′_d,i)）
        with self.inference_ctx():
            if self.compile:
                # Mark the beginning of a new CUDA Graph step
                torch.compiler.cudagraph_mark_step_begin()
            # get decoder input
            ret_dct = self.warping_module(feature_3d, kp_source=kp_source, kp_driving=kp_driving)

            # print(f"=============================================================================")
            # for out_key, out_value in ret_dct.items():
            #     if isinstance(out_value, str) or isinstance(out_value, int) or isinstance(out_value, float):
            #         print(f"{out_key}: {out_value}")
            #     elif isinstance(out_value, torch.Tensor):
            #         print(f"{out_key}: tensor shape {out_value.shape}, min: {torch.min(out_value)}, max: {torch.max(out_value)}, mean: {torch.mean(out_value)}, std: {torch.std(out_value)}")
            #     else:
            #         print(f"{out_key}: data type {type(out_value)}")
            # decode
            ret_dct['out'] = self.spade_generator(feature=ret_dct['out'])

            # float the dict
            if self.cfg.flag_use_half_precision:
                for k, v in ret_dct.items():
                    if isinstance(v, torch.Tensor):
                        ret_dct[k] = v.float()

        return ret_dct
    
    def parse_output(self, out: torch.Tensor) -> np.ndarray:
        """ construct the output as standard
        return: 1xHxWx3, uint8
        """
        out = np.transpose(out.cpu().numpy(), [0, 2, 3, 1])  # 1x3xHxW -> 1xHxWx3
        out = np.clip(out, 0, 1)  # clip to 0~1
        out = np.clip(out * 255, 0, 255).astype(np.uint8)  # 0~1 -> 0~255

        return out

    @torch.no_grad()
    def calc_combined_eye_ratio(self, c_d_eyes_i, source_lmk):
        c_s_eyes = calc_eye_close_ratio(source_lmk[None])
        c_s_eyes_tensor = torch.from_numpy(c_s_eyes).float().to(self.device)
        c_d_eyes_i_tensor = torch.Tensor([c_d_eyes_i[0][0]]).reshape(1, 1).to(self.device)
        # [c_s,eyes, c_d,eyes,i]
        combined_eye_ratio_tensor = torch.cat([c_s_eyes_tensor, c_d_eyes_i_tensor], dim=1)
        return combined_eye_ratio_tensor

    @torch.no_grad()
    def calc_combined_lip_ratio(self, c_d_lip_i, source_lmk):
        c_s_lip = calc_lip_close_ratio(source_lmk[None])
        c_s_lip_tensor = torch.from_numpy(c_s_lip).float().to(self.device)
        c_d_lip_i_tensor = torch.Tensor([c_d_lip_i[0]]).to(self.device).reshape(1, 1) # 1x1
        # [c_s,lip, c_d,lip,i]
        combined_lip_ratio_tensor = torch.cat([c_s_lip_tensor, c_d_lip_i_tensor], dim=1) # 1x2
        return combined_lip_ratio_tensor

    def calc_ratio(self, lmk_lst):
        input_eye_ratio_lst = []
        input_lip_ratio_lst = []
        for lmk in lmk_lst:
            # for eyes retargeting
            input_eye_ratio_lst.append(calc_eye_close_ratio(lmk[None]))
            # for lip retargeting
            input_lip_ratio_lst.append(calc_lip_close_ratio(lmk[None]))
        return input_eye_ratio_lst, input_lip_ratio_lst

    @torch.no_grad()
    def retarget_lip(self, kp_source: torch.Tensor, lip_close_ratio: torch.Tensor) -> torch.Tensor:
        """
        kp_source: BxNx3
        lip_close_ratio: Bx2
        Return: Bx(3*num_kp)
        """
        feat_lip = concat_feat(kp_source, lip_close_ratio)

        delta = self.stitching_retargeting_module['lip'](feat_lip)

        return delta.reshape(-1, kp_source.shape[1], 3)

    @torch.no_grad()
    def retarget_eye(self, kp_source: torch.Tensor, eye_close_ratio: torch.Tensor) -> torch.Tensor:
        """
        kp_source: BxNx3
        eye_close_ratio: Bx3
        Return: Bx(3*num_kp)
        """
        feat_eye = concat_feat(kp_source, eye_close_ratio)

        delta = self.stitching_retargeting_module['eye'](feat_eye)

        return delta.reshape(-1, kp_source.shape[1], 3)

    def crop_image(self, img, do_crop=False):
        ######## process source info ########
        if do_crop:
            crop_info = self.cropper.crop_source_image(img, self.cropper.crop_cfg)
            if crop_info is None:
                raise Exception("No face detected in the source image!")
            lmk = crop_info['lmk_crop']
            img_crop_256x256 = crop_info['img_crop_256x256']
        else:
            crop_info = None
            lmk = self.cropper.calc_lmk_from_cropped_image(img)
            img_crop_256x256 = cv2.resize(img, (256, 256))  # force to resize to 256x256
        return img_crop_256x256, lmk, crop_info

    def crop_source_video(self, img_lst, do_crop=False):
        if do_crop:
            ret_s = self.cropper.crop_source_video(img_lst, self.cropper.crop_cfg)
            print(f'Source video is cropped, {len(ret_s["frame_crop_lst"])} frames are processed.')
            img_crop_256x256_lst, lmk_crop_lst, M_c2o_lst = ret_s['frame_crop_lst'], ret_s['lmk_crop_lst'], ret_s['M_c2o_lst']
        else:
            M_c2o_lst = None
            lmk_crop_lst = self.cropper.calc_lmks_from_cropped_video(img_lst)
            img_crop_256x256_lst = [cv2.resize(_, (256, 256)) for _ in img_lst]  # force to resize to 256x256
        return img_crop_256x256_lst, lmk_crop_lst, M_c2o_lst
    
    def crop_driving_videos(self, img_lst, do_crop=False):
        if do_crop:
            ret_d = self.cropper.crop_driving_video(img_lst)
            print(f'Driving video is cropped, {len(ret_d["frame_crop_lst"])} frames are processed.')
            img_crop_lst, lmk_crop_lst = ret_d['frame_crop_lst'], ret_d['lmk_crop_lst']
            img_crop_256x256_lst = [cv2.resize(_, (256, 256)) for _ in img_lst]
        else:
            lmk_crop_lst = self.cropper.calc_lmks_from_cropped_video(img_lst)
            img_crop_256x256_lst = [cv2.resize(_, (256, 256)) for _ in img_lst]  # force to resize to 256x256
        return img_crop_256x256_lst, lmk_crop_lst

    def prepare_source(self, src_img):
        """ construct the input as standard
        img: HxWx3, uint8, 256x256
        """
        # processing source image to tensor
        h, w = src_img.shape[:2]
        if h != self.cfg.input_height or w != self.cfg.input_width:
            x = cv2.resize(src_img, (self.cfg.input_width, self.cfg.input_height))
        else:
            x = src_img.copy()
        
        if x.ndim == 3:
            x = x[np.newaxis].astype(np.float32) / 255.  # HxWx3 -> 1xHxWx3, normalized to 0~1
        elif x.ndim == 4:
            x = x.astype(np.float32) / 255.  # BxHxWx3, normalized to 0~1
        else:
            raise ValueError(f'img ndim should be 3 or 4: {x.ndim}')
        
        x = np.clip(x, 0, 1)  # clip to 0~1
        x = torch.from_numpy(x).permute(0, 3, 1, 2)  # 1xHxWx3 -> 1x3xHxW
        x = x.to(self.device)

        # extract features
        I_s = x
        f_s = self.extract_feature_3d(I_s)
        x_s_info = self.get_kp_info(I_s)
        
        return f_s, x_s_info
    
    def process_clips(self, clips):
        """ construct the input as standard
        clips: NxBxHxWx3, uint8
        """
        # resize to 256 x 256
        imgs = []
        for img in clips:
            h, w = img.shape[:2]
            if h != self.cfg.input_height or w != self.cfg.input_width:
                img = cv2.resize(img, (self.cfg.input_width, self.cfg.input_height))
            else:
                img = img.copy()
            imgs.append(img)

        # processing video frames to tensor
        if isinstance(imgs, list):
            _imgs = np.array(imgs)[..., np.newaxis]  # TxHxWx3x1
        elif isinstance(imgs, np.ndarray):
            _imgs = imgs
        else:
            raise ValueError(f'imgs type error: {type(imgs)}')

        y = _imgs.astype(np.float32) / 255.
        y = np.clip(y, 0, 1)  # clip to 0~1
        y = torch.from_numpy(y).permute(0, 4, 3, 1, 2)  # TxHxWx3x1 -> Tx1x3xHxW
        y = y.to(self.device)

        return y

    def prepare_driving_videos(self, vid_frames, feat_type="tensor"):
        """ get driving kp infos
        vid_frames: image list of HxWx3, uint8
        """
        # extract features
        total_len = len(vid_frames)
        kp_infos = {"pitch": [], "yaw": [], "roll": [], "t": [], "exp": [], "scale": [], "kp": []}
        for start_idx in range(0, total_len, self.cfg.batch_size):
            frames = vid_frames[start_idx: min(start_idx + self.cfg.batch_size, total_len)]
            frames = self.process_clips(frames).squeeze(1)
            kp_info = self.get_kp_info(frames)

            for k, v in kp_info.items():
                kp_infos[k].append(v)

        # combine the kp_infos
        for k, v in kp_infos.items():
            kp_infos[k] = torch.cat(v, dim=0)

        if feat_type == "np":
            for k, v in kp_infos.items():
                kp_infos[k] = v.cpu().numpy()

        return kp_infos

    def get_driving_template(self, kp_infos, smooth=False, dtype="pt_tensor"):
        kp_infos = self.refine_kp(kp_infos)
        motion_list = []
        n_frames = len(kp_infos["exp"])
        for idx in range(n_frames):
            exp = kp_infos["exp"][idx]
            scale = kp_infos["scale"][idx]
            t = kp_infos["t"][idx]
            pitch = kp_infos["pitch"][idx]
            yaw = kp_infos["yaw"][idx]
            roll = kp_infos["roll"][idx]
            
            R = get_rotation_matrix(pitch, yaw, roll)
            R = R.reshape(1, 3, 3)    
            
            exp = exp.reshape(1, 21, 3)
            scale = scale.reshape(1, 1)
            t = t.reshape(1, 3)
            pitch = pitch.reshape(1, 1)
            yaw = yaw.reshape(1, 1)
            roll = roll.reshape(1, 1)

            if dtype == "np":
                R = R.cpu().numpy().astype(np.float32)
                exp = exp.cpu().numpy().astype(np.float32)
                scale = scale.cpu().numpy().astype(np.float32)
                t = t.cpu().numpy().astype(np.float32)
                pitch = pitch.cpu().numpy().astype(np.float32)
                yaw = yaw.cpu().numpy().astype(np.float32)
                roll = roll.cpu().numpy().astype(np.float32)
            
            motion_list.append(
                {"exp": exp, "scale": scale, "R": R, "t": t, "pitch": pitch, "yaw": yaw, "roll": roll}
            )
        tgt_motion = {'n_frames': n_frames, 'output_fps': 25, 'motion': motion_list}

        if smooth:
            print("Smoothing motion sequence...")
            tgt_motion = smooth_(tgt_motion, method="ema")
        return tgt_motion

    @torch.no_grad()
    def update_delta_new_eyeball_direction(self, eyeball_direction_x, eyeball_direction_y, delta_new, **kwargs):
        if eyeball_direction_x > 0:
                delta_new[0, 11, 0] += eyeball_direction_x * 0.0007
                delta_new[0, 15, 0] += eyeball_direction_x * 0.001
        else:
            delta_new[0, 11, 0] += eyeball_direction_x * 0.001
            delta_new[0, 15, 0] += eyeball_direction_x * 0.0007

        delta_new[0, 11, 1] += eyeball_direction_y * -0.001
        delta_new[0, 15, 1] += eyeball_direction_y * -0.001
        blink = -eyeball_direction_y / 2.

        delta_new[0, 11, 1] += blink * -0.001
        delta_new[0, 13, 1] += blink * 0.0003
        delta_new[0, 15, 1] += blink * -0.001
        delta_new[0, 16, 1] += blink * 0.0003

        return delta_new

    def driven(self, f_s, x_s_info, s_lmk, c_s_eyes_lst, kp_infos, c_d_eyes_lst=None, c_d_lip_lst=None, smooth=False):
        # source kp info
        x_d_i_news=[]
        x_ss=[]
        f_ss=[]
        x_s_info = self.refine_kp(x_s_info)
        R_s = get_rotation_matrix(x_s_info['pitch'], x_s_info['yaw'], x_s_info['roll'])
        x_s = self.transform_keypoint(x_s_info)
        x_c_s = x_s_info["kp"]

        # driving kp infos
        driving_template_dct = self.get_driving_template(kp_infos, smooth)
        n_frames = driving_template_dct['n_frames']

        # driving params
        flag_normalize_lip = self.cfg.flag_normalize_lip
        flag_relative_motion = self.cfg.flag_relative_motion
        flag_source_video_eye_retargeting = self.cfg.flag_source_video_eye_retargeting
        lip_normalize_threshold = self.cfg.lip_normalize_threshold
        source_video_eye_retargeting_threshold = self.cfg.source_video_eye_retargeting_threshold
        animation_region = self.cfg.animation_region
        driving_option = self.cfg.driving_option
        flag_stitching = self.cfg.flag_stitching
        flag_eye_retargeting = self.cfg.flag_eye_retargeting
        flag_lip_retargeting = self.cfg.flag_lip_retargeting
        driving_multiplier = self.cfg.driving_multiplier
        lib_multiplier = self.cfg.lib_multiplier

        # let lip-open scalar to be 0 at first
        lip_delta_before_animation, eye_delta_before_animation = None, None
        if flag_normalize_lip and flag_relative_motion and s_lmk is not None:
            c_d_lip_before_animation = [0.]
            combined_lip_ratio_tensor_before_animation = self.calc_combined_lip_ratio(c_d_lip_before_animation, s_lmk)
            if combined_lip_ratio_tensor_before_animation[0][0] >= lip_normalize_threshold:
                lip_delta_before_animation = self.retarget_lip(x_s, combined_lip_ratio_tensor_before_animation)

        # let eye-open scalar to be the same as the first frame if the latter is eye-open state
        if flag_source_video_eye_retargeting and s_lmk is not None:
            combined_eye_ratio_tensor_frame_zero = c_s_eyes_lst[0]
            c_d_eye_before_animation_frame_zero = [[combined_eye_ratio_tensor_frame_zero[0][:2].mean()]]
            if c_d_eye_before_animation_frame_zero[0][0] < source_video_eye_retargeting_threshold:
                c_d_eye_before_animation_frame_zero = [[0.39]]
            combined_eye_ratio_tensor_before_animation = self.calc_combined_eye_ratio(c_d_eye_before_animation_frame_zero, s_lmk)
            eye_delta_before_animation = self.retarget_eye(x_s, combined_eye_ratio_tensor_before_animation)
        
        # animate 
        I_p_lst = []
        for i in range(n_frames):
            x_d_i_info = driving_template_dct['motion'][i]
            x_d_i_info = dct2device(x_d_i_info, self.device)
            # R
            R_d_i = x_d_i_info['R']
            if i == 0:  # cache the first frame
                R_d_0 = R_d_i
                x_d_0_info = x_d_i_info.copy()
            
            # enhance lip
            # if i > 0:
            #     for lip_idx in [6, 12, 14, 17, 19, 20]:
            #         x_d_i_info['exp'][:, lip_idx, :] = x_d_0_info['exp'][:, lip_idx, :] + (x_d_i_info['exp'][:, lip_idx, :] - x_d_0_info['exp'][:, lip_idx, :]) * lib_multiplier
            
            # normalize eye_ball, TODO
            x_d_i_info['exp'] = self.update_delta_new_eyeball_direction(0, -5, x_d_i_info['exp'])

            # debug
            #print(f"frame {i:03d}, src scale {x_s_info['scale']}, 0 scale {x_d_0_info['scale']}, i scale {x_d_i_info['scale']}")
            # delta
            delta_new = x_s_info['exp'].clone()
            if flag_relative_motion:
                # R
                if animation_region == "all" or animation_region == "pose":
                    R_new = (R_d_i @ R_d_0.permute(0, 2, 1)) @ R_s
                else:
                    R_new = R_s

                # exp
                if animation_region == "all" or animation_region == "exp":
                    delta_new = x_s_info['exp'] + (x_d_i_info['exp'] - x_d_0_info['exp'])
                elif animation_region == "lip":
                    for lip_idx in [6, 12, 14, 17, 19, 20]:
                        delta_new[:, lip_idx, :] = (x_s_info['exp'] + (x_d_i_info['exp'] - x_d_0_info['exp']))[:, lip_idx, :]
                elif animation_region == "eyes":
                    for eyes_idx in [11, 13, 15, 16, 18]:
                        delta_new[:, eyes_idx, :] = (x_s_info['exp'] + (x_d_i_info['exp'] - x_d_0_info['exp']))[:, eyes_idx, :]
                
                # scale
                if animation_region == "all":
                    scale_new = x_s_info['scale'] * (x_d_i_info['scale'] / x_d_0_info['scale'])
                else:
                    scale_new = x_s_info['scale']

                # translation
                if animation_region == "all" or animation_region == "pose":
                    t_new = x_s_info['t'] + (x_d_i_info['t'] - x_d_0_info['t'])
                else:
                    t_new = x_s_info['t']
            else:
                # R
                if animation_region == "all" or animation_region == "pose":
                    R_new = R_d_i
                else:
                    R_new = R_s

                # exp
                if animation_region == "all" or animation_region == "exp":
                    EYE_IDX=[1,2,6,11,12,13,14,15,16,17,18,19,20]
                    delta_new[:, EYE_IDX, :] = x_d_i_info['exp'][:, EYE_IDX, :]
                    # for idx in [1,2,6,11,12,13,14,15,16,17,18,19,20]:
                    #     delta_new[:, idx, :] = x_d_i_info['exp'][:, idx, :]
                    delta_new[:, 3:5, 1] = x_d_i_info['exp'][:, 3:5, 1]
                    delta_new[:, 5, 2] = x_d_i_info['exp'][:, 5, 2]
                    delta_new[:, 8, 2] = x_d_i_info['exp'][:, 8, 2]
                    delta_new[:, 9, 1:] = x_d_i_info['exp'][:, 9, 1:]
                elif animation_region == "lip":
                    for lip_idx in [6, 12, 14, 17, 19, 20]:
                        delta_new[:, lip_idx, :] = x_d_i_info['exp'][:, lip_idx, :]
                elif animation_region == "eyes":
                    for eyes_idx in [11, 13, 15, 16, 18]:
                        delta_new[:, eyes_idx, :] = x_d_i_info['exp'][:, eyes_idx, :]
                
                # scale
                scale_new = x_s_info['scale']

                # translation
                if animation_region == "all" or animation_region == "pose":
                    t_new = x_d_i_info['t']
                else:
                    t_new = x_s_info['t']

            t_new[..., 2].fill_(0)  # zero tz

            x_d_i_new = scale_new * (x_c_s @ R_new + delta_new) + t_new

            if flag_relative_motion and driving_option == "expression-friendly":
                if i == 0:
                    x_d_0_new = x_d_i_new
                    motion_multiplier = calc_motion_multiplier(x_s, x_d_0_new)
                x_d_diff = (x_d_i_new - x_d_0_new) * motion_multiplier
                x_d_i_new = x_d_diff + x_s
            
            # Algorithm 1 in Liveportrait:
            if not flag_stitching and not flag_eye_retargeting and not flag_lip_retargeting:
                # without stitching or retargeting
                if flag_normalize_lip and lip_delta_before_animation is not None:
                    x_d_i_new += lip_delta_before_animation
                if flag_source_video_eye_retargeting and eye_delta_before_animation is not None:
                    x_d_i_new += eye_delta_before_animation
                else:
                    pass
            elif flag_stitching and not flag_eye_retargeting and not flag_lip_retargeting:
                # with stitching and without retargeting
                if flag_normalize_lip and lip_delta_before_animation is not None:
                    x_d_i_new = self.stitching(x_s, x_d_i_new) + lip_delta_before_animation
                else:
                    x_d_i_new = self.stitching(x_s, x_d_i_new)
                if flag_source_video_eye_retargeting and eye_delta_before_animation is not None:
                    x_d_i_new += eye_delta_before_animation
            else:
                eyes_delta, lip_delta = None, None
                if flag_eye_retargeting and s_lmk is not None and c_d_eyes_lst is not None:
                    c_d_eyes_i = c_d_eyes_lst[i]
                    combined_eye_ratio_tensor = self.calc_combined_eye_ratio(c_d_eyes_i, s_lmk)
                    eyes_delta = self.retarget_eye(x_s, combined_eye_ratio_tensor)

                if flag_lip_retargeting and s_lmk is not None and c_d_lip_lst is not None:
                    c_d_lip_i = c_d_lip_lst[i]
                    combined_lip_ratio_tensor = self.calc_combined_lip_ratio(c_d_lip_i, s_lmk)
                    # ∆_lip,i = R_lip(x_s; c_s,lip, c_d,lip,i)
                    lip_delta = self.retarget_lip(x_s, combined_lip_ratio_tensor)

                if flag_relative_motion:  # use x_s
                    x_d_i_new = x_s + \
                        (eyes_delta if eyes_delta is not None else 0) + \
                        (lip_delta if lip_delta is not None else 0)
                else:  # use x_d,i
                    x_d_i_new = x_d_i_new + \
                        (eyes_delta if eyes_delta is not None else 0) + \
                        (lip_delta if lip_delta is not None else 0)

                if flag_stitching:
                    x_d_i_new = self.stitching(x_s, x_d_i_new)

            x_d_i_new = x_s + (x_d_i_new - x_s) * driving_multiplier
            x_d_i_news.append(x_d_i_new)
        f_s_s= f_s.expand(n_frames, *f_s.shape[1:]) 
        x_s_s = x_s.expand(n_frames, *x_s.shape[1:])  
        x_d_i_new = torch.cat(x_d_i_news, dim=0)        
        for start in range(0, n_frames, 100):
            end = min(start + 100,n_frames)
            with torch.no_grad(), torch.autocast('cuda'):
                out = self.warp_decode(f_s_s[start:end], x_s_s[start:end], x_d_i_new[start:end])        
                I_p_lst.append(out['out'])
        I_p=torch.cat(I_p_lst, dim=0) 
        I_p_i = self.parse_output(I_p)
        return I_p_i 

    def driven_debug(self, f_s, x_s_info, s_lmk, c_s_eyes_lst, driving_template_dct, c_d_eyes_lst=None, c_d_lip_lst=None):
        # source kp info
        x_s_info = self.refine_kp(x_s_info)
        x_c_s = x_s_info["kp"]
        R_s = get_rotation_matrix(x_s_info['pitch'], x_s_info['yaw'], x_s_info['roll'])
        x_s = self.transform_keypoint(x_s_info)
        
        n_frames = driving_template_dct['n_frames']

        # driving params
        flag_normalize_lip = self.cfg.flag_normalize_lip
        flag_relative_motion = self.cfg.flag_relative_motion
        flag_source_video_eye_retargeting = self.cfg.flag_source_video_eye_retargeting
        lip_normalize_threshold = self.cfg.lip_normalize_threshold
        source_video_eye_retargeting_threshold = self.cfg.source_video_eye_retargeting_threshold
        animation_region = self.cfg.animation_region
        driving_option = self.cfg.driving_option
        flag_stitching = self.cfg.flag_stitching
        flag_eye_retargeting = self.cfg.flag_eye_retargeting
        flag_lip_retargeting = self.cfg.flag_lip_retargeting
        driving_multiplier = self.cfg.driving_multiplier

        # let lip-open scalar to be 0 at first
        lip_delta_before_animation, eye_delta_before_animation = None, None
        if flag_normalize_lip and flag_relative_motion and s_lmk is not None:
            c_d_lip_before_animation = [0.]
            combined_lip_ratio_tensor_before_animation = self.calc_combined_lip_ratio(c_d_lip_before_animation, s_lmk)
            if combined_lip_ratio_tensor_before_animation[0][0] >= lip_normalize_threshold:
                lip_delta_before_animation = self.retarget_lip(x_s, combined_lip_ratio_tensor_before_animation)

        # let eye-open scalar to be the same as the first frame if the latter is eye-open state
        if flag_source_video_eye_retargeting and s_lmk is not None:
            combined_eye_ratio_tensor_frame_zero = c_s_eyes_lst[0]
            c_d_eye_before_animation_frame_zero = [[combined_eye_ratio_tensor_frame_zero[0][:2].mean()]]
            if c_d_eye_before_animation_frame_zero[0][0] < source_video_eye_retargeting_threshold:
                c_d_eye_before_animation_frame_zero = [[0.39]]
            combined_eye_ratio_tensor_before_animation = self.calc_combined_eye_ratio(c_d_eye_before_animation_frame_zero, s_lmk)
            eye_delta_before_animation = self.retarget_eye(x_s, combined_eye_ratio_tensor_before_animation)
        
        # animate 
        I_p_lst = []
        for i in range(n_frames):
            x_d_i_info = driving_template_dct['motion'][i]
            x_d_i_info = dct2device(x_d_i_info, self.device)
            # R
            R_d_i = x_d_i_info['R'] if 'R' in x_d_i_info.keys() else x_d_i_info['R_d']  # compatible with previous keys
            if i == 0:  # cache the first frame
                R_d_0 = R_d_i
                x_d_0_info = x_d_i_info.copy()
            
            # debug
            #print(f"frame {i:03d}, src scale {x_s_info['scale']}, 0 scale {x_d_0_info['scale']}, i scale {x_d_i_info['scale']}")
            # delta
            delta_new = x_s_info['exp'].clone()
            if flag_relative_motion:
                # R
                if animation_region == "all" or animation_region == "pose":
                    R_new = (R_d_i @ R_d_0.permute(0, 2, 1)) @ R_s
                else:
                    R_new = R_s

                # exp
                if animation_region == "all" or animation_region == "exp":
                    delta_new = x_s_info['exp'] + (x_d_i_info['exp'] - x_d_0_info['exp'])
                elif animation_region == "lip":
                    for lip_idx in [6, 12, 14, 17, 19, 20]:
                        delta_new[:, lip_idx, :] = (x_s_info['exp'] + (x_d_i_info['exp'] - x_d_0_info['exp']))[:, lip_idx, :]
                elif animation_region == "eyes":
                    for eyes_idx in [11, 13, 15, 16, 18]:
                        delta_new[:, eyes_idx, :] = (x_s_info['exp'] + (x_d_i_info['exp'] - x_d_0_info['exp']))[:, eyes_idx, :]
                
                # scale
                if animation_region == "all":
                    scale_new = x_s_info['scale'] * (x_d_i_info['scale'] / x_d_0_info['scale'])
                else:
                    scale_new = x_s_info['scale']

                # translation
                if animation_region == "all" or animation_region == "pose":
                    t_new = x_s_info['t'] + (x_d_i_info['t'] - x_d_0_info['t'])
                else:
                    t_new = x_s_info['t']
            else:
                # R
                if animation_region == "all" or animation_region == "pose":
                    R_new = R_d_i
                else:
                    R_new = R_s

                # exp
                if animation_region == "all" or animation_region == "exp":
                    for idx in [1,2,6,11,12,13,14,15,16,17,18,19,20]:
                        delta_new[:, idx, :] = x_d_i_info['exp'][:, idx, :]
                    delta_new[:, 3:5, 1] = x_d_i_info['exp'][:, 3:5, 1]
                    delta_new[:, 5, 2] = x_d_i_info['exp'][:, 5, 2]
                    delta_new[:, 8, 2] = x_d_i_info['exp'][:, 8, 2]
                    delta_new[:, 9, 1:] = x_d_i_info['exp'][:, 9, 1:]
                elif animation_region == "lip":
                    for lip_idx in [6, 12, 14, 17, 19, 20]:
                        delta_new[:, lip_idx, :] = x_d_i_info['exp'][:, lip_idx, :]
                elif animation_region == "eyes":
                    for eyes_idx in [11, 13, 15, 16, 18]:
                        delta_new[:, eyes_idx, :] = x_d_i_info['exp'][:, eyes_idx, :]
                
                # scale
                scale_new = x_s_info['scale']

                # translation
                if animation_region == "all" or animation_region == "pose":
                    t_new = x_d_i_info['t']
                else:
                    t_new = x_s_info['t']

            t_new[..., 2].fill_(0)  # zero tz

            x_d_i_new = scale_new * (x_c_s @ R_new + delta_new) + t_new

            if flag_relative_motion and driving_option == "expression-friendly":
                if i == 0:
                    x_d_0_new = x_d_i_new
                    motion_multiplier = calc_motion_multiplier(x_s, x_d_0_new)
                x_d_diff = (x_d_i_new - x_d_0_new) * motion_multiplier
                x_d_i_new = x_d_diff + x_s
            
            # Algorithm 1 in Liveportrait:
            if not flag_stitching and not flag_eye_retargeting and not flag_lip_retargeting:
                # without stitching or retargeting
                if flag_normalize_lip and lip_delta_before_animation is not None:
                    x_d_i_new += lip_delta_before_animation
                if flag_source_video_eye_retargeting and eye_delta_before_animation is not None:
                    x_d_i_new += eye_delta_before_animation
                else:
                    pass
            elif flag_stitching and not flag_eye_retargeting and not flag_lip_retargeting:
                # with stitching and without retargeting
                if flag_normalize_lip and lip_delta_before_animation is not None:
                    x_d_i_new = self.stitching(x_s, x_d_i_new) + lip_delta_before_animation
                else:
                    x_d_i_new = self.stitching(x_s, x_d_i_new)
                if flag_source_video_eye_retargeting and eye_delta_before_animation is not None:
                    x_d_i_new += eye_delta_before_animation
            else:
                eyes_delta, lip_delta = None, None
                if flag_eye_retargeting and s_lmk is not None and c_d_eyes_lst is not None:
                    c_d_eyes_i = c_d_eyes_lst[i]
                    combined_eye_ratio_tensor = self.calc_combined_eye_ratio(c_d_eyes_i, s_lmk)
                    eyes_delta = self.retarget_eye(x_s, combined_eye_ratio_tensor)

                if flag_lip_retargeting and s_lmk is not None and c_d_lip_lst is not None:
                    c_d_lip_i = c_d_lip_lst[i]
                    combined_lip_ratio_tensor = self.calc_combined_lip_ratio(c_d_lip_i, s_lmk)
                    # ∆_lip,i = R_lip(x_s; c_s,lip, c_d,lip,i)
                    lip_delta = self.retarget_lip(x_s, combined_lip_ratio_tensor)

                if flag_relative_motion:  # use x_s
                    x_d_i_new = x_s + \
                        (eyes_delta if eyes_delta is not None else 0) + \
                        (lip_delta if lip_delta is not None else 0)
                else:  # use x_d,i
                    x_d_i_new = x_d_i_new + \
                        (eyes_delta if eyes_delta is not None else 0) + \
                        (lip_delta if lip_delta is not None else 0)

                if flag_stitching:
                    x_d_i_new = self.stitching(x_s, x_d_i_new)

            x_d_i_new = x_s + (x_d_i_new - x_s) * driving_multiplier
            out = self.warp_decode(f_s, x_s, x_d_i_new)
            I_p_i = self.parse_output(out['out'])[0]
            I_p_lst.append(I_p_i)
        
        return I_p_lst 

    def read_image(self, image_path: str) -> list:
        img_rgb = load_image_rgb(image_path)
        img_rgb = resize_to_limit(img_rgb, self.cfg.source_max_dim, self.cfg.source_division)
        source_rgb_list = [img_rgb]
        print(f"Load image from {osp.realpath(image_path)} done.")
        return source_rgb_list

    def read_video(self, video_path: str, interval=None) -> list:
        vr = VideoReader(video_path)
        if interval is not None:
            video_frames = vr.get_batch(np.arange(0, len(vr), interval)).numpy()
        else:
            video_frames = [vr[0].numpy(), vr[len(vr) // 2].numpy(), vr[-1].numpy()]
        vr.seek(0)
        driving_rgb_list = []
        for video_frame in video_frames:
            # h, w = video_frame.shape[:2]
            # if h != self.cfg.output_height or w != self.cfg.output_width:
            #     video_frame = cv2.resize(video_frame, (self.cfg.output_height, self.cfg.output_width))
            driving_rgb_list.append(video_frame)

        return driving_rgb_list

    def prepare_videos(self, imgs) -> torch.Tensor:
        """ construct the input as standard
        imgs: NxBxHxWx3, uint8
        """
        if isinstance(imgs, list):
            _imgs = np.array(imgs)[..., np.newaxis]  # TxHxWx3x1
        elif isinstance(imgs, np.ndarray):
            _imgs = imgs
        else:
            raise ValueError(f'imgs type error: {type(imgs)}')

        y = _imgs.astype(np.float32) / 255.
        y = np.clip(y, 0, 1)  # clip to 0~1
        y = torch.from_numpy(y).permute(0, 4, 3, 1, 2)  # TxHxWx3x1 -> Tx1x3xHxW
        y = y.to(self.device)

        return y

    def make_motion_template(self, I_lst, c_eyes_lst, c_lip_lst, **kwargs):
        n_frames = I_lst.shape[0]
        template_dct = {
            'n_frames': n_frames,
            'output_fps': kwargs.get('output_fps', 25),
            'motion': [],
            'c_eyes_lst': [],
            'c_lip_lst': [],
        }

        for i in track(range(n_frames), description='Making motion templates...', total=n_frames):
            # collect s, R, δ and t for inference
            I_i = I_lst[i]
            x_i_info = self.refine_kp(self.get_kp_info(I_i))
            x_s = self.transform_keypoint(x_i_info)
            R_i = get_rotation_matrix(x_i_info['pitch'], x_i_info['yaw'], x_i_info['roll'])

            item_dct = {
                'scale': x_i_info['scale'].cpu().numpy().astype(np.float32),
                'R': R_i.cpu().numpy().astype(np.float32),
                'exp': x_i_info['exp'].cpu().numpy().astype(np.float32),
                't': x_i_info['t'].cpu().numpy().astype(np.float32),
                'kp': x_i_info['kp'].cpu().numpy().astype(np.float32),
                'x_s': x_s.cpu().numpy().astype(np.float32),
            }

            template_dct['motion'].append(item_dct)

            c_eyes = c_eyes_lst[i].astype(np.float32)
            template_dct['c_eyes_lst'].append(c_eyes)

            c_lip = c_lip_lst[i].astype(np.float32)
            template_dct['c_lip_lst'].append(c_lip)

        return template_dct

    def load_template(self, wfp_template):
        print(f"Load from template: {wfp_template}, NOT the video, so the cropping video and audio are both NULL.")
        driving_template_dct = load(wfp_template)
        c_d_eyes_lst = driving_template_dct['c_eyes_lst'] if 'c_eyes_lst' in driving_template_dct.keys() else driving_template_dct['c_d_eyes_lst'] # compatible with previous keys
        c_d_lip_lst = driving_template_dct['c_lip_lst'] if 'c_lip_lst' in driving_template_dct.keys() else driving_template_dct['c_d_lip_lst']
        driving_n_frames = driving_template_dct['n_frames']
        flag_is_driving_video = True if driving_n_frames > 1 else False
        n_frames = driving_n_frames

        # set output_fps
        output_fps = driving_template_dct.get('output_fps', 25)
        print(f'The FPS of template: {output_fps}')
        return driving_template_dct

    def reconstruction(self, src_img, dst_imgs, video_path="template"):
        # prepare source
        src_img_256x256, s_lmk, _ = self.crop_image(src_img, do_crop=False)
        #c_s_eyes_lst, c_s_lip_lst = self.calc_ratio([s_lmk])
        c_s_eyes_lst = None
        f_s, x_s_info = self.prepare_source(src_img_256x256)
        
        # prepare driving videos
        dst_imgs_256x256, d_lmk_lst = self.crop_driving_videos(dst_imgs, do_crop=False)
        c_d_eyes_lst, c_d_lip_lst = self.calc_ratio(d_lmk_lst)
        kp_infos = self.prepare_driving_videos(dst_imgs_256x256)


        recs = self.driven(f_s, x_s_info, s_lmk, c_s_eyes_lst, kp_infos, c_d_eyes_lst, c_d_lip_lst)
        return recs

    def save_results(self, results, save_path, audio_path=None):
        save_dir = osp.dirname(save_path)
        save_name = osp.basename(save_path)
        final_video = osp.join(save_dir, f'final_{save_name}')

        images2video(results, wfp=save_path, fps=self.cfg.output_fps)

        if audio_path is not None:
            add_audio_to_video(save_path, audio_path, final_video)
            os.remove(save_path)
    
    def rec_score(self, video_path: str, interval=None, save_path=None):
        video_frames = self.read_video(video_path, interval=interval)
        #print(f"len frames: {len(video_frames)}, shape: {video_frames[0].shape}")
        recs = self.reconstruction(video_frames[0], video_frames[1:], video_path)
        if save_path is not None:
            self.save_results(recs, save_path)
        #print(f"len rec: {len(recs)}, shape: {recs[0].shape}")
        psnrs = psnr(video_frames[1:], recs)
        psnrs_np = np.array(psnrs)
        psnr_mean, psnr_std = np.mean(psnrs_np), np.std(psnrs_np)
        rec_score = {"mean": psnr_mean, "std": psnr_std}
        return rec_score

    @torch.no_grad()
    def paste_back_by_face_mask(self, result, crop_info, src_img, crop_src_image, use_laplacian=False):
        """
        paste back the result to the original image with face mask
        """
        # detect src mask
        crop_src_tensor = self.to_tensor(crop_src_image).unsqueeze(0).to(self.device)
        src_msks = get_face_mask(self.face_parser, crop_src_tensor)
        result_tensor = self.to_tensor(result).unsqueeze(0).to(self.device)
        result_msks = get_face_mask(self.face_parser, result_tensor)
        # combine masks
        masks = []
        for src_msk, result_msk in zip(src_msks, result_msks):
            mask = np.clip(src_msk + result_msk, 0, 1)
            masks.append(mask)
        result = paste_back_with_face_mask(result, crop_info, src_img, masks[0], use_laplacian=use_laplacian)
        return result

    def driven_by_audio(self, src_img, kp_infos, save_path, audio_path=None, smooth=False):
        # prepare source
        # prepare source
        src_img_256x256, s_lmk, crop_info = self.crop_image(src_img, do_crop=True)
        #c_s_eyes_lst, c_s_lip_lst = self.calc_ratio([s_lmk])
        c_s_eyes_lst = None
        f_s, x_s_info = self.prepare_source(src_img_256x256)

        mask_ori_float = prepare_paste_back(self.mask_crop, crop_info['M_c2o'], dsize=(src_img.shape[1], src_img.shape[0]))
        
        # prepare driving videos
        results = self.driven(f_s, x_s_info, s_lmk, c_s_eyes_lst, kp_infos, smooth=smooth)
        frames=results.shape[0]
        results = [paste_back(results[i], crop_info['M_c2o'], src_img, mask_ori_float) for i in range(frames)]
        self.save_results(results, save_path, audio_path)
    def mix_kp_infos(self, emo_kp_infos, lip_kp_infos, smooth=False, dtype="pt_tensor"):
        driving_emo_template_dct = self.get_driving_template(emo_kp_infos, smooth=False, dtype=dtype)
        if lip_kp_infos is not None:
            driving_lip_template_dct = self.get_driving_template(lip_kp_infos, smooth=smooth, dtype=dtype)
            driving_template_dct = {**driving_emo_template_dct}
            n_frames = min(driving_emo_template_dct['n_frames'], driving_lip_template_dct['n_frames'])
            driving_template_dct['n_frames'] = n_frames
            for i in range(n_frames):
                emo_motion = driving_emo_template_dct['motion'][i]['exp']
                lib_motion = driving_lip_template_dct['motion'][i]['exp']
                for lip_idx in [6, 12, 14, 17, 19, 20]:
                    emo_motion[:, lip_idx, :] = lib_motion[:, lip_idx, :]
                driving_template_dct['motion'][i]['exp'] = emo_motion
        else:
            driving_template_dct = driving_emo_template_dct
        
        return driving_template_dct

    def driven_by_mix(self, src_img, driving_video_path, kp_infos, save_path, audio_path=None, smooth=False):
        # prepare source
        src_img_256x256, s_lmk, crop_info = self.crop_image(src_img, do_crop=True)
        c_s_eyes_lst, c_s_lip_lst = self.calc_ratio([s_lmk])
        f_s, x_s_info = self.prepare_source(src_img_256x256)
        mask_ori_float = prepare_paste_back(self.mask_crop, crop_info['M_c2o'], dsize=(src_img.shape[1], src_img.shape[0]))
        # prepare driving videos
        driving_imgs = self.read_video(driving_video_path, interval=1)
        dst_imgs_256x256, d_lmk_lst = self.crop_driving_videos(driving_imgs, do_crop=True)
        c_d_eyes_lst, c_d_lip_lst = self.calc_ratio(d_lmk_lst)
        emo_kp_infos = self.prepare_driving_videos(dst_imgs_256x256)
        # mix kp_infos
        driving_template_dct = self.mix_kp_infos(emo_kp_infos, kp_infos, smooth=smooth)
        # driven 
        results = self.driven_debug(f_s, x_s_info, s_lmk, c_s_eyes_lst, driving_template_dct, c_d_eyes_lst=c_d_eyes_lst, c_d_lip_lst=c_d_lip_lst)
        results = [paste_back(result, crop_info['M_c2o'], src_img, mask_ori_float) for result in results]
        print(results.shape)
        self.save_results(results, save_path, audio_path)
    
    def drive_video_by_mix(self, video_path, driving_video_path, kp_infos, save_path, audio_path):
        # prepare driving videos
        driving_imgs = self.read_video(driving_video_path, interval=1)
        dst_imgs_256x256, d_lmk_lst = self.crop_driving_videos(driving_imgs, do_crop=True)
        emo_kp_infos = self.prepare_driving_videos(dst_imgs_256x256)
        # mix kp_infos
        #driving_template_dct = self.get_driving_template(emo_kp_infos, smooth=True, dtype="np")
        driving_template_dct = self.mix_kp_infos(emo_kp_infos, kp_infos, smooth=True, dtype="np")
        # driven
        self.video_lip_retargeting(
            video_path, None, 
            save_path, audio_path, 
            driving_template_dct=driving_template_dct, retargeting_ragion="exp"
        )

    def load_source_video(self, video_info, n_frames=-1):
        reader = imageio.get_reader(video_info, "ffmpeg")

        ret = []
        for idx, frame_rgb in enumerate(reader):
            if n_frames > 0 and idx >= n_frames:
                break
            ret.append(frame_rgb)

        reader.close()
        
        return ret

    def video_lip_retargeting(self, video_path, kp_infos, save_path, audio_path, c_d_eyes_lst=None, c_d_lip_lst=None, smooth=False, driving_template_dct=None, retargeting_ragion="exp"):
        # 0. process source motion template
        source_rgb_lst = load_video(video_path)
        source_rgb_lst = [resize_to_limit(img, self.cfg.source_max_dim, self.cfg.source_division) for img in source_rgb_lst]
        img_crop_256x256_lst, source_lmk_crop_lst, source_M_c2o_lst = self.crop_source_video(source_rgb_lst, do_crop=True)
        c_s_eyes_lst, c_s_lip_lst = self.calc_ratio(source_lmk_crop_lst)
        I_s_lst = self.prepare_videos(img_crop_256x256_lst)
        source_template_dct = self.make_motion_template(I_s_lst, c_s_eyes_lst, c_s_lip_lst, output_fps=25)
        # 1. prepare driving template
        if driving_template_dct is None:
            driving_template_dct = self.get_driving_template(kp_infos, smooth=smooth, dtype="np")
        # 2. driving
        n_frames = min(source_template_dct['n_frames'], driving_template_dct['n_frames'])
        # driving params
        I_p_lst = []
        I_p_pstbk_lst = []
        R_d_0, x_d_0_info = None, None
        flag_normalize_lip = self.cfg.flag_normalize_lip
        flag_relative_motion = True #self.cfg.flag_relative_motion
        flag_source_video_eye_retargeting = self.cfg.flag_source_video_eye_retargeting
        lip_normalize_threshold = self.cfg.lip_normalize_threshold
        source_video_eye_retargeting_threshold = self.cfg.source_video_eye_retargeting_threshold
        animation_region = 'lip' #self.cfg.animation_region
        driving_option = self.cfg.driving_option
        flag_stitching = self.cfg.flag_stitching
        flag_eye_retargeting = self.cfg.flag_eye_retargeting
        flag_lip_retargeting = self.cfg.flag_lip_retargeting
        driving_multiplier = self.cfg.driving_multiplier
        driving_smooth_observation_variance = self.cfg.driving_smooth_observation_variance
        
        key_r = 'R' if 'R' in driving_template_dct['motion'][0].keys() else 'R_d'
        if flag_relative_motion:
            x_d_exp_lst = [source_template_dct['motion'][i]['exp'] + driving_template_dct['motion'][i]['exp'] - driving_template_dct['motion'][0]['exp'] for i in range(n_frames)]
            for i in range(n_frames):
                for idx in [6, 12, 14, 17, 19, 20]:
                    # lip motion use abs motion
                    x_d_exp_lst[i][:, idx, :] = driving_template_dct['motion'][i]['exp'][:, idx, :]
            x_d_exp_lst_smooth = ksmooth(x_d_exp_lst, source_template_dct['motion'][0]['exp'].shape, self.device, driving_smooth_observation_variance)
            
            if animation_region == "all" or animation_region == "pose" or "all" in animation_region:
                x_d_r_lst = [(np.dot(driving_template_dct['motion'][i][key_r], driving_template_dct['motion'][0][key_r].transpose(0, 2, 1))) @ source_template_dct['motion'][i]['R'] for i in range(n_frames)]
                x_d_r_lst_smooth = ksmooth(x_d_r_lst, source_template_dct['motion'][0]['R'].shape, self.device, driving_smooth_observation_variance)
        else:
            x_d_exp_lst = [driving_template_dct['motion'][i]['exp'] for i in range(n_frames)]
            x_d_exp_lst_smooth = ksmooth(x_d_exp_lst, source_template_dct['motion'][0]['exp'].shape, self.device, driving_smooth_observation_variance)

            if animation_region == "all" or animation_region == "pose" or "all" in animation_region:
                x_d_r_lst = [driving_template_dct['motion'][i][key_r] for i in range(n_frames)]
                x_d_r_lst_smooth = ksmooth(x_d_r_lst, source_template_dct['motion'][0]['R'].shape, self.device, driving_smooth_observation_variance)
        
        # driving all
        for i in track(range(n_frames), description='🚀Retargeting...', total=n_frames):
            x_s_info = source_template_dct['motion'][i]
            x_s_info = dct2device(x_s_info, self.device)

            source_lmk = source_lmk_crop_lst[i]
            img_crop_256x256 = img_crop_256x256_lst[i]
            I_s = I_s_lst[i]
            f_s = self.extract_feature_3d(I_s)

            x_c_s = x_s_info['kp']
            R_s = x_s_info['R']
            x_s =x_s_info['x_s']

            # let lip-open scalar to be 0 at first if the input is a video
            lip_delta_before_animation = None
            if flag_normalize_lip and flag_relative_motion and source_lmk is not None:
                c_d_lip_before_animation = [0.]
                combined_lip_ratio_tensor_before_animation = self.calc_combined_lip_ratio(c_d_lip_before_animation, source_lmk)
                if combined_lip_ratio_tensor_before_animation[0][0] >= lip_normalize_threshold:
                    lip_delta_before_animation = self.retarget_lip(x_s, combined_lip_ratio_tensor_before_animation)
                else:
                    lip_delta_before_animation = None

            # let eye-open scalar to be the same as the first frame if the latter is eye-open state
            eye_delta_before_animation = None
            if flag_source_video_eye_retargeting and source_lmk is not None:
                if i == 0:
                    combined_eye_ratio_tensor_frame_zero = c_s_eyes_lst[0]
                    c_d_eye_before_animation_frame_zero = [[combined_eye_ratio_tensor_frame_zero[0][:2].mean()]]
                    if c_d_eye_before_animation_frame_zero[0][0] < source_video_eye_retargeting_threshold:
                        c_d_eye_before_animation_frame_zero = [[0.39]]
                combined_eye_ratio_tensor_before_animation = self.calc_combined_eye_ratio(c_d_eye_before_animation_frame_zero, source_lmk)
                eye_delta_before_animation = self.retarget_eye(x_s, combined_eye_ratio_tensor_before_animation)
            
            if flag_stitching:  # prepare for paste back
                mask_ori_float = prepare_paste_back(self.mask_crop, source_M_c2o_lst[i], dsize=(source_rgb_lst[i].shape[1], source_rgb_lst[i].shape[0]))
            
            x_d_i_info = driving_template_dct['motion'][i]
            x_d_i_info = dct2device(x_d_i_info, self.device)
            R_d_i = x_d_i_info['R'] if 'R' in x_d_i_info.keys() else x_d_i_info['R_d']  # compatible with previous keys

            if i == 0:  # cache the first frame
                R_d_0 = R_d_i
                x_d_0_info = x_d_i_info.copy()
            
            delta_new = x_s_info['exp'].clone()
            if flag_relative_motion:
                if animation_region == "all" or animation_region == "pose" or "all" in animation_region:
                    R_new = x_d_r_lst_smooth[i]
                else:
                    R_new = R_s
                if animation_region == "all" or animation_region == "exp":
                    for idx in [1,2,6,11,12,13,14,15,16,17,18,19,20]:
                        delta_new[:, idx, :] = x_d_exp_lst_smooth[i][idx, :]
                    delta_new[:, 3:5, 1] = x_d_exp_lst_smooth[i][3:5, 1]
                    delta_new[:, 5, 2] = x_d_exp_lst_smooth[i][5, 2]
                    delta_new[:, 8, 2] = x_d_exp_lst_smooth[i][8, 2]
                    delta_new[:, 9, 1:] = x_d_exp_lst_smooth[i][9, 1:]
                elif animation_region == "all_wo_lip" or animation_region == "exp_wo_lip":
                    for idx in [1, 2, 11, 13, 15, 16, 18]:
                        delta_new[:, idx, :] = x_d_exp_lst_smooth[i][idx, :]
                    delta_new[:, 3:5, 1] = x_d_exp_lst_smooth[i][3:5, 1]
                    delta_new[:, 5, 2] = x_d_exp_lst_smooth[i][5, 2]
                    delta_new[:, 8, 2] = x_d_exp_lst_smooth[i][8, 2]
                    delta_new[:, 9, 1:] = x_d_exp_lst_smooth[i][9, 1:]
                elif animation_region == "lip":
                    for lip_idx in [6, 12, 14, 17, 19, 20]:
                        delta_new[:, lip_idx, :] = x_d_exp_lst_smooth[i][lip_idx, :]
                elif animation_region == "eyes":
                    for eyes_idx in [11, 13, 15, 16, 18]:
                        delta_new[:, eyes_idx, :] = x_d_exp_lst_smooth[i][eyes_idx, :]
                
                scale_new = x_s_info['scale']
                t_new = x_s_info['t']
            else:
                if animation_region == "all" or animation_region == "pose" or "all" in animation_region:
                    R_new = x_d_r_lst_smooth[i] 
                else:
                    R_new = R_s
                if animation_region == "all" or animation_region == "exp":
                    for idx in [1,2,6,11,12,13,14,15,16,17,18,19,20]:
                        delta_new[:, idx, :] = x_d_exp_lst_smooth[i][idx, :]
                    delta_new[:, 3:5, 1] = x_d_exp_lst_smooth[i][3:5, 1] 
                    delta_new[:, 5, 2] = x_d_exp_lst_smooth[i][5, 2] 
                    delta_new[:, 8, 2] = x_d_exp_lst_smooth[i][8, 2] 
                    delta_new[:, 9, 1:] = x_d_exp_lst_smooth[i][9, 1:]
                elif animation_region == "all_wo_lip" or animation_region == "exp_wo_lip":
                    for idx in [1, 2, 11, 13, 15, 16, 18]:
                        delta_new[:, idx, :] = x_d_exp_lst_smooth[i][idx, :]
                    delta_new[:, 3:5, 1] = x_d_exp_lst_smooth[i][3:5, 1]
                    delta_new[:, 5, 2] = x_d_exp_lst_smooth[i][5, 2]
                    delta_new[:, 8, 2] = x_d_exp_lst_smooth[i][8, 2]
                    delta_new[:, 9, 1:] = x_d_exp_lst_smooth[i][9, 1:]
                elif animation_region == "lip":
                    for lip_idx in [6, 12, 14, 17, 19, 20]:
                        delta_new[:, lip_idx, :] = x_d_exp_lst_smooth[i][lip_idx, :]
                elif animation_region == "eyes":
                    for eyes_idx in [11, 13, 15, 16, 18]:
                        delta_new[:, eyes_idx, :] = x_d_exp_lst_smooth[i][eyes_idx, :]
                scale_new = x_s_info['scale']
                if animation_region == "all" or animation_region == "pose" or "all" in animation_region:
                    t_new = x_d_i_info['t']
                else:
                    t_new = x_s_info['t']

            t_new[..., 2].fill_(0)  # zero tz
            x_d_i_new = scale_new * (x_c_s @ R_new + delta_new) + t_new

            # Algorithm 1:
            if not flag_stitching and not flag_eye_retargeting and not flag_lip_retargeting:
                # without stitching or retargeting
                if flag_normalize_lip and lip_delta_before_animation is not None:
                    x_d_i_new += lip_delta_before_animation
                if flag_source_video_eye_retargeting and eye_delta_before_animation is not None:
                    x_d_i_new += eye_delta_before_animation
                else:
                    pass
            elif flag_stitching and not flag_eye_retargeting and not flag_lip_retargeting:
                # with stitching and without retargeting
                if flag_normalize_lip and lip_delta_before_animation is not None:
                    x_d_i_new = self.stitching(x_s, x_d_i_new) + lip_delta_before_animation
                else:
                    x_d_i_new = self.stitching(x_s, x_d_i_new)
                if flag_source_video_eye_retargeting and eye_delta_before_animation is not None:
                    x_d_i_new += eye_delta_before_animation
            else:
                eyes_delta, lip_delta = None, None
                if flag_eye_retargeting and source_lmk is not None and c_d_eyes_lst is not None:
                    c_d_eyes_i = c_d_eyes_lst[i]
                    combined_eye_ratio_tensor = self.calc_combined_eye_ratio(c_d_eyes_i, source_lmk)
                    # ∆_eyes,i = R_eyes(x_s; c_s,eyes, c_d,eyes,i)
                    eyes_delta = self.retarget_eye(x_s, combined_eye_ratio_tensor)
                if flag_lip_retargeting and source_lmk is not None and c_d_lip_lst is not None:
                    c_d_lip_i = c_d_lip_lst[i]
                    combined_lip_ratio_tensor = self.calc_combined_lip_ratio(c_d_lip_i, source_lmk)
                    # ∆_lip,i = R_lip(x_s; c_s,lip, c_d,lip,i)
                    lip_delta = self.retarget_lip(x_s, combined_lip_ratio_tensor)

                if flag_relative_motion:  # use x_s
                    x_d_i_new = x_s + \
                        (eyes_delta if eyes_delta is not None else 0) + \
                        (lip_delta if lip_delta is not None else 0)
                else:  # use x_d,i
                    x_d_i_new = x_d_i_new + \
                        (eyes_delta if eyes_delta is not None else 0) + \
                        (lip_delta if lip_delta is not None else 0)

                if flag_stitching:
                    x_d_i_new = self.stitching(x_s, x_d_i_new)

            x_d_i_new = x_s + (x_d_i_new - x_s) * driving_multiplier
            out = self.warp_decode(f_s, x_s, x_d_i_new)
            I_p_i = self.parse_output(out['out'])[0]
            I_p_lst.append(I_p_i)

            if flag_stitching:
                # TODO: the paste back procedure is slow, considering optimize it using multi-threading or GPU
                #I_p_pstbk = self.paste_back_by_face_mask(I_p_i, source_M_c2o_lst[i], source_rgb_lst[i], img_crop_256x256, use_laplacian=True)
                I_p_pstbk = paste_back(I_p_i, source_M_c2o_lst[i], source_rgb_lst[i], mask_ori_float, use_laplacian=True)
                I_p_pstbk_lst.append(I_p_pstbk)
            
        if len(I_p_pstbk_lst) > 0:
            self.save_results(I_p_pstbk_lst, save_path, audio_path)
        else:
            self.save_results(I_p_lst, save_path, audio_path)

    @torch.no_grad()
    def video_reconstruction_test(self, video_tensor, xs, save_path):
        # video_tensor, (1, F, C, H, W), [-1, 1]
        # xs, (1, F, 63)
        result_lst = []
        #ori_videos = []
        video_tensor = video_tensor[0:1] * 0.5 + 0.5  # [-1, 1] -> [0, 1], 1xTx3xHxW
        video_tensor = torch.clip(video_tensor, 0, 1)
        video_tensor = video_tensor.permute(1, 0, 2, 3, 4) # 1xTx3xHxW -> Tx1x3xHxW
        video = video_tensor.to(self.device)
        xs = xs[0:1].permute(1, 0, 2)    # 1xTx63 -> Tx1x63
        xs = xs.reshape(-1, 1, 21, 3)
        xs = xs.to(self.device)

        x_s_0 = xs[0]
        I_s_0 = torch.nn.functional.interpolate(video[0], size=(256, 256), mode='bilinear')
        f_s_0 = self.extract_feature_3d(I_s_0)

        for i in range(video_tensor.shape[0]):
            #I_s = video[i]   # 1x3xHxW
            #ori_videos.append((I_s.squeeze(0).squeeze(0).permute(1, 2, 0).cpu().numpy()*255).astype(np.uint8))
            x_s = self.stitching(x_s_0, xs[i])
            out = self.warp_decode(f_s_0, x_s_0, x_s)
            I_p_i = self.parse_output(out['out'])[0]
            result_lst.append(I_p_i)

        #save_dir = osp.dirname(save_path)
        #ori_path = osp.join(save_dir, "ori.mp4")
        #save_path = osp.join(save_dir, "rec.mp4")
        self.save_results(result_lst, save_path, audio_path=None)
        #self.save_results(ori_videos, ori_path, audio_path=None)

    @torch.no_grad()
    def self_driven(self, image_tensor, xs, save_path, length):
        result_lst = []
        image_tensor = image_tensor[0:1] * 0.5 + 0.5    # [-1, 1] -> [0, 1], 1x3xHxW
        image_tensor = torch.clip(image_tensor, 0, 1)
        image = image_tensor.to(self.device)
        I_s_0 = torch.nn.functional.interpolate(image, size=(256, 256), mode='bilinear')

        xs = xs[0:1].permute(1, 0, 2)    # 1xTx63 -> Tx1x63
        xs = xs.reshape(-1, 1, 21, 3)
        xs = xs.to(self.device)

        x_s_0 = xs[0]
        f_s_0 = self.extract_feature_3d(I_s_0)

        for i in range(xs.shape[0]):
            x_d = self.stitching(x_s_0, xs[i])
            out = self.warp_decode(f_s_0, x_s_0, x_d)
            I_p_i = self.parse_output(out['out'])[0]
            result_lst.append(I_p_i)

        assert len(result_lst) == length, f"length of result_lst is {len(result_lst)}, but length is {length}"

        self.save_results(result_lst, save_path, audio_path=None)