# coding: utf-8 """ Pipeline for gradio """ import gradio as gr from .config.argument_config import ArgumentConfig from .live_portrait_pipeline import LivePortraitPipeline from .utils.io import load_img_online from .utils.rprint import rlog as log from .utils.crop import prepare_paste_back, paste_back from .utils.camera import get_rotation_matrix def update_args(args, user_args): """update the args according to user inputs """ for k, v in user_args.items(): if hasattr(args, k): setattr(args, k, v) return args class GradioPipeline(LivePortraitPipeline): def __init__(self, inference_cfg, crop_cfg, args: ArgumentConfig): super().__init__(inference_cfg, crop_cfg) # self.live_portrait_wrapper = self.live_portrait_wrapper self.args = args def execute_video( self, input_image_path, input_video_path, flag_relative_input, flag_do_crop_input, flag_remap_input, flag_crop_driving_video_input ): """ for video driven potrait animation """ if input_image_path is not None and input_video_path is not None: args_user = { 'source_image': input_image_path, 'driving_info': input_video_path, 'flag_relative': flag_relative_input, 'flag_do_crop': flag_do_crop_input, 'flag_pasteback': flag_remap_input, 'flag_crop_driving_video': flag_crop_driving_video_input } # update config from user input self.args = update_args(self.args, args_user) self.live_portrait_wrapper.update_config(self.args.__dict__) self.cropper.update_config(self.args.__dict__) # video driven animation video_path, video_path_concat = self.execute(self.args) gr.Info("Run successfully!", duration=2) return video_path, video_path_concat, else: raise gr.Error("The input source portrait or driving video hasn't been prepared yet 💥!", duration=5) def execute_s_video( self, input_s_video_path, input_video_path, flag_relative_input, flag_do_crop_input, flag_remap_input, flag_crop_driving_video_input, progress=gr.Progress(track_tqdm=True) ): """ for video driven source to video animation """ if input_s_video_path is not None and input_video_path is not None: args_user = { 'source_driving_info': input_s_video_path, 'driving_info': input_video_path, 'flag_relative': flag_relative_input, 'flag_do_crop': flag_do_crop_input, 'flag_pasteback': flag_remap_input, 'flag_crop_driving_video': flag_crop_driving_video_input } # update config from user input self.args = update_args(self.args, args_user) self.live_portrait_wrapper.update_config(self.args.__dict__) self.cropper.update_config(self.args.__dict__) # video driven animation video_path, video_path_concat = self.execute_source_video(self.args) gr.Info("Run successfully!", duration=3) return video_path, video_path_concat, else: raise gr.Error("The input source video or driving video hasn't been prepared yet 💥!", duration=5) def execute_image(self, input_eye_ratio: float, input_lip_ratio: float, input_image, flag_do_crop=True): """ for single image retargeting """ # disposable feature f_s_user, x_s_user, source_lmk_user, crop_M_c2o, mask_ori, img_rgb = \ self.prepare_retargeting(input_image, flag_do_crop) if input_eye_ratio is None or input_lip_ratio is None: raise gr.Error("Invalid ratio input 💥!", duration=5) else: inference_cfg = self.live_portrait_wrapper.inference_cfg x_s_user = x_s_user.to(self.live_portrait_wrapper.device) f_s_user = f_s_user.to(self.live_portrait_wrapper.device) # ∆_eyes,i = R_eyes(x_s; c_s,eyes, c_d,eyes,i) combined_eye_ratio_tensor = self.live_portrait_wrapper.calc_combined_eye_ratio([[input_eye_ratio]], source_lmk_user) eyes_delta = self.live_portrait_wrapper.retarget_eye(x_s_user, combined_eye_ratio_tensor) # ∆_lip,i = R_lip(x_s; c_s,lip, c_d,lip,i) combined_lip_ratio_tensor = self.live_portrait_wrapper.calc_combined_lip_ratio([[input_lip_ratio]], source_lmk_user) lip_delta = self.live_portrait_wrapper.retarget_lip(x_s_user, combined_lip_ratio_tensor) num_kp = x_s_user.shape[1] # default: use x_s x_d_new = x_s_user + eyes_delta.reshape(-1, num_kp, 3) + lip_delta.reshape(-1, num_kp, 3) # D(W(f_s; x_s, x′_d)) out = self.live_portrait_wrapper.warp_decode(f_s_user, x_s_user, x_d_new) out = self.live_portrait_wrapper.parse_output(out['out'])[0] out_to_ori_blend = paste_back(out, crop_M_c2o, img_rgb, mask_ori) gr.Info("Run successfully!", duration=2) return out, out_to_ori_blend def prepare_retargeting(self, input_image, flag_do_crop=True): """ for single image retargeting """ if input_image is not None: # gr.Info("Upload successfully!", duration=2) inference_cfg = self.live_portrait_wrapper.inference_cfg ######## process source portrait ######## img_rgb = load_img_online(input_image, mode='rgb', max_dim=1280, n=16) log(f"Load source image from {input_image}.") crop_info = self.cropper.crop_source_image(img_rgb, self.cropper.crop_cfg) if flag_do_crop: I_s = self.live_portrait_wrapper.prepare_source(crop_info['img_crop_256x256']) else: I_s = self.live_portrait_wrapper.prepare_source(img_rgb) x_s_info = self.live_portrait_wrapper.get_kp_info(I_s) R_s = get_rotation_matrix(x_s_info['pitch'], x_s_info['yaw'], x_s_info['roll']) ############################################ f_s_user = self.live_portrait_wrapper.extract_feature_3d(I_s) x_s_user = self.live_portrait_wrapper.transform_keypoint(x_s_info) source_lmk_user = crop_info['lmk_crop'] crop_M_c2o = crop_info['M_c2o'] mask_ori = prepare_paste_back(inference_cfg.mask_crop, crop_info['M_c2o'], dsize=(img_rgb.shape[1], img_rgb.shape[0])) return f_s_user, x_s_user, source_lmk_user, crop_M_c2o, mask_ori, img_rgb else: # when press the clear button, go here raise gr.Error("The retargeting input hasn't been prepared yet 💥!", duration=5)