Spaces:
Running
on
Zero
Running
on
Zero
| # 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 | |
| from .utils.retargeting_utils import calc_eye_close_ratio, calc_lip_close_ratio | |
| 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 | |
| # for single image retargeting | |
| self.start_prepare = False | |
| self.f_s_user = None | |
| self.x_c_s_info_user = None | |
| self.x_s_user = None | |
| self.source_lmk_user = None | |
| self.mask_ori = None | |
| self.img_rgb = None | |
| self.crop_M_c2o = None | |
| def execute_video( | |
| self, | |
| input_image_path, | |
| input_video_path, | |
| flag_relative_input, | |
| flag_do_crop_input, | |
| flag_remap_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, | |
| } | |
| # 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_image(self, input_eye_ratio: float, input_lip_ratio: float): | |
| """ for single image retargeting | |
| """ | |
| if input_eye_ratio is None or input_eye_ratio is None: | |
| raise gr.Error("Invalid ratio input 💥!", duration=5) | |
| elif self.f_s_user is None: | |
| if self.start_prepare: | |
| raise gr.Error( | |
| "The source portrait is under processing 💥! Please wait for a second.", | |
| duration=5 | |
| ) | |
| else: | |
| raise gr.Error( | |
| "The source portrait hasn't been prepared yet 💥! Please scroll to the top of the page to upload.", | |
| duration=5 | |
| ) | |
| else: | |
| x_s_user = self.x_s_user.to("cuda") | |
| f_s_user = self.f_s_user.to("cuda") | |
| # ∆_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]], self.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]], self.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, self.crop_M_c2o, self.img_rgb, self.mask_ori) | |
| # gr.Info("Run successfully!", duration=2) | |
| return out, out_to_ori_blend | |
| def prepare_retargeting(self, input_image_path, flag_do_crop = True): | |
| """ for single image retargeting | |
| """ | |
| if input_image_path is not None: | |
| # gr.Info("Upload successfully!", duration=2) | |
| self.start_prepare = True | |
| inference_cfg = self.live_portrait_wrapper.cfg | |
| ######## process source portrait ######## | |
| img_rgb = load_img_online(input_image_path, mode='rgb', max_dim=1280, n=16) | |
| log(f"Load source image from {input_image_path}.") | |
| crop_info = self.cropper.crop_single_image(img_rgb) | |
| 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']) | |
| ############################################ | |
| # record global info for next time use | |
| self.f_s_user = self.live_portrait_wrapper.extract_feature_3d(I_s) | |
| self.x_s_user = self.live_portrait_wrapper.transform_keypoint(x_s_info) | |
| self.x_s_info_user = x_s_info | |
| self.source_lmk_user = crop_info['lmk_crop'] | |
| self.img_rgb = img_rgb | |
| self.crop_M_c2o = crop_info['M_c2o'] | |
| self.mask_ori = prepare_paste_back(inference_cfg.mask_crop, crop_info['M_c2o'], dsize=(img_rgb.shape[1], img_rgb.shape[0])) | |
| # update slider | |
| eye_close_ratio = calc_eye_close_ratio(self.source_lmk_user[None]) | |
| eye_close_ratio = float(eye_close_ratio.squeeze(0).mean()) | |
| lip_close_ratio = calc_lip_close_ratio(self.source_lmk_user[None]) | |
| lip_close_ratio = float(lip_close_ratio.squeeze(0).mean()) | |
| # for vis | |
| self.I_s_vis = self.live_portrait_wrapper.parse_output(I_s)[0] | |
| return eye_close_ratio, lip_close_ratio, self.I_s_vis | |
| else: | |
| # when press the clear button, go here | |
| return 0.8, 0.8, self.I_s_vis | |