diff --git a/liveportrait/live_portrait_pipeline.py b/liveportrait/live_portrait_pipeline.py index 3f544cd..29649ca 100644 --- a/liveportrait/live_portrait_pipeline.py +++ b/liveportrait/live_portrait_pipeline.py @@ -37,15 +37,16 @@ class LivePortraitPipeline(object): appearance_feature_extractor, motion_extractor, warping_module, spade_generator, stitching_retargeting_module, cfg=inference_cfg) - def execute(self, img_rgb, driving_images_np): + def execute(self, img_rgb, driving_images_np, crop_info): inference_cfg = self.live_portrait_wrapper.cfg # for convenience ######## process reference portrait ######## #img_rgb = load_image_rgb(args.source_image) - img_rgb = resize_to_limit(img_rgb, inference_cfg.ref_max_shape, inference_cfg.ref_shape_n) + #img_rgb = resize_to_limit(img_rgb, inference_cfg.ref_max_shape, inference_cfg.ref_shape_n) #log(f"Load source image from {args.source_image}") - crop_info = self.cropper.crop_single_image(img_rgb) + #crop_info = self.cropper.crop_single_image(img_rgb) source_lmk = crop_info['lmk_crop'] - _, img_crop_256x256 = crop_info['img_crop'], crop_info['img_crop_256x256'] + #_, img_crop_256x256 = crop_info['img_crop'], crop_info['img_crop_256x256'] + img_crop_256x256 = img_rgb if inference_cfg.flag_do_crop: I_s = self.live_portrait_wrapper.prepare_source(img_crop_256x256) else: @@ -59,7 +60,7 @@ class LivePortraitPipeline(object): if inference_cfg.flag_lip_zero: # let lip-open scalar to be 0 at first c_d_lip_before_animation = [0.] - combined_lip_ratio_tensor_before_animation = self.live_portrait_wrapper.calc_combined_lip_ratio(c_d_lip_before_animation, source_lmk) + combined_lip_ratio_tensor_before_animation = self.live_portrait_wrapper.calc_combined_lip_ratio(c_d_lip_before_animation, source_lmk, inference_cfg) if combined_lip_ratio_tensor_before_animation[0][0] < inference_cfg.lip_zero_threshold: inference_cfg.flag_lip_zero = False else: @@ -79,7 +80,7 @@ class LivePortraitPipeline(object): n_frames = I_d_lst.shape[0] if inference_cfg.flag_eye_retargeting or inference_cfg.flag_lip_retargeting: driving_lmk_lst = self.cropper.get_retargeting_lmk_info(driving_rgb_lst) - input_eye_ratio_lst, input_lip_ratio_lst = self.live_portrait_wrapper.calc_retargeting_ratio(source_lmk, driving_lmk_lst) + input_eye_ratio_lst, input_lip_ratio_lst = self.live_portrait_wrapper.calc_retargeting_ratio(source_lmk, driving_lmk_lst, inference_cfg) # elif is_template(args.driving_info): # log(f"Load from video templates {args.driving_info}") @@ -150,14 +151,14 @@ class LivePortraitPipeline(object): eyes_delta, lip_delta = None, None if inference_cfg.flag_eye_retargeting: c_d_eyes_i = input_eye_ratio_lst[i] - combined_eye_ratio_tensor = self.live_portrait_wrapper.calc_combined_eye_ratio(c_d_eyes_i, source_lmk) - combined_eye_ratio_tensor = combined_eye_ratio_tensor * inference_cfg.eyes_retargeting_multiplier + combined_eye_ratio_tensor = self.live_portrait_wrapper.calc_combined_eye_ratio(c_d_eyes_i, source_lmk, inference_cfg) + combined_eye_ratio_tensor = combined_eye_ratio_tensor # ∆_eyes,i = R_eyes(x_s; c_s,eyes, c_d,eyes,i) eyes_delta = self.live_portrait_wrapper.retarget_eye(x_s, combined_eye_ratio_tensor) if inference_cfg.flag_lip_retargeting: c_d_lip_i = input_lip_ratio_lst[i] - combined_lip_ratio_tensor = self.live_portrait_wrapper.calc_combined_lip_ratio(c_d_lip_i, source_lmk) - combined_lip_ratio_tensor = combined_lip_ratio_tensor * inference_cfg.lip_retargeting_multiplier + combined_lip_ratio_tensor = self.live_portrait_wrapper.calc_combined_lip_ratio(c_d_lip_i, source_lmk, inference_cfg) + combined_lip_ratio_tensor = combined_lip_ratio_tensor # ∆_lip,i = R_lip(x_s; c_s,lip, c_d,lip,i) lip_delta = self.live_portrait_wrapper.retarget_lip(x_s, combined_lip_ratio_tensor) diff --git a/liveportrait/live_portrait_wrapper.py b/liveportrait/live_portrait_wrapper.py index 1cbd474..448cf27 100644 --- a/liveportrait/live_portrait_wrapper.py +++ b/liveportrait/live_portrait_wrapper.py @@ -301,20 +301,20 @@ class LivePortraitWrapper(object): input_lip_ratio_lst.append(calc_lip_close_ratio(lmk[None])) return input_eye_ratio_lst, input_lip_ratio_lst - def calc_combined_eye_ratio(self, input_eye_ratio, source_lmk): + def calc_combined_eye_ratio(self, input_eye_ratio, source_lmk, inference_cfg): eye_close_ratio = calc_eye_close_ratio(source_lmk[None]) eye_close_ratio_tensor = torch.from_numpy(eye_close_ratio).float().to(self.device_id) - input_eye_ratio_tensor = torch.Tensor([input_eye_ratio[0][0]]).reshape(1, 1).to(self.device_id) + input_eye_ratio_tensor = torch.Tensor([input_eye_ratio[0][0]]).reshape(1, 1).to(self.device_id) * inference_cfg.eyes_retargeting_multiplier # [c_s,eyes, c_d,eyes,i] combined_eye_ratio_tensor = torch.cat([eye_close_ratio_tensor, input_eye_ratio_tensor], dim=1) return combined_eye_ratio_tensor - def calc_combined_lip_ratio(self, input_lip_ratio, source_lmk): + def calc_combined_lip_ratio(self, input_lip_ratio, source_lmk, inference_cfg): lip_close_ratio = calc_lip_close_ratio(source_lmk[None]) lip_close_ratio_tensor = torch.from_numpy(lip_close_ratio).float().to(self.device_id) # [c_s,lip, c_d,lip,i] input_lip_ratio_tensor = torch.Tensor([input_lip_ratio[0]]).to(self.device_id) if input_lip_ratio_tensor.shape != [1, 1]: input_lip_ratio_tensor = input_lip_ratio_tensor.reshape(1, 1) - combined_lip_ratio_tensor = torch.cat([lip_close_ratio_tensor, input_lip_ratio_tensor], dim=1) + combined_lip_ratio_tensor = torch.cat([lip_close_ratio_tensor, input_lip_ratio_tensor], dim=1) * inference_cfg.lip_retargeting_multiplier return combined_lip_ratio_tensor diff --git a/liveportrait/utils/cropper.py b/liveportrait/utils/cropper.py index 024d9b7..185c43b 100644 --- a/liveportrait/utils/cropper.py +++ b/liveportrait/utils/cropper.py @@ -82,6 +82,7 @@ class Cropper(object): src_face = src_face[0] pts = src_face.landmark_2d_106 + # crop the face ret_dct = crop_image( @@ -99,7 +100,17 @@ class Cropper(object): lmk = recon_ret['pts'] ret_dct['lmk_crop'] = lmk - return ret_dct + # Draw each landmark as a circle + width, height = 512, 512 + blank_image = np.zeros((height, width, 3), dtype=np.uint8) * 255 + for (x, y) in lmk: + # Ensure the coordinates are within the dimensions of the blank image + if 0 <= x < width and 0 <= y < height: + cv2.circle(blank_image, (int(x), int(y)), radius=2, color=(0, 0, 255)) + + keypoints_image = cv2.cvtColor(blank_image, cv2.COLOR_BGR2RGB) + + return ret_dct, keypoints_image def get_retargeting_lmk_info(self, driving_rgb_lst): # TODO: implement a tracking-based version diff --git a/nodes.py b/nodes.py index b9bf637..445b9d7 100644 --- a/nodes.py +++ b/nodes.py @@ -4,6 +4,8 @@ import yaml import folder_paths import comfy.model_management as mm import comfy.utils +import numpy as np +import cv2 script_directory = os.path.dirname(os.path.abspath(__file__)) @@ -223,12 +225,9 @@ class LivePortraitProcess: return {"required": { "pipeline": ("LIVEPORTRAITPIPE",), + "crop_info": ("CROPINFO", {"default": {}}), "source_image": ("IMAGE",), "driving_images": ("IMAGE",), - "dsize": ("INT", {"default": 512, "min": 64, "max": 2048}), - "scale": ("FLOAT", {"default": 2.3, "min": 1.0, "max": 4.0, "step": 0.01}), - "vx_ratio": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}), - "vy_ratio": ("FLOAT", {"default": -0.125, "min": -1.0, "max": 1.0, "step": 0.01}), "lip_zero": ("BOOLEAN", {"default": True}), "eye_retargeting": ("BOOLEAN", {"default": False}), "eyes_retargeting_multiplier": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 10.0, "step": 0.001}), @@ -237,17 +236,6 @@ class LivePortraitProcess: "stitching": ("BOOLEAN", {"default": True}), "relative": ("BOOLEAN", {"default": True}), }, - "optional": { - "onnx_device": ( - [ - 'CPU', - 'CUDA', - ], { - "default": 'CPU' - }), - } - - } RETURN_TYPES = ("IMAGE", "IMAGE",) @@ -255,20 +243,11 @@ class LivePortraitProcess: FUNCTION = "process" CATEGORY = "LivePortrait" - def process(self, source_image, driving_images, dsize, scale, vx_ratio, vy_ratio, pipeline, - lip_zero, eye_retargeting, lip_retargeting, stitching, relative, eyes_retargeting_multiplier, lip_retargeting_multiplier, onnx_device='CUDA'): + def process(self, source_image, driving_images, pipeline, + lip_zero, eye_retargeting, lip_retargeting, stitching, relative, eyes_retargeting_multiplier, lip_retargeting_multiplier, crop_info = {}): source_image_np = (source_image * 255).byte().numpy() driving_images_np = (driving_images * 255).byte().numpy() - - crop_cfg = CropConfig( - dsize = dsize, - scale = scale, - vx_ratio = vx_ratio, - vy_ratio = vy_ratio, - ) - cropper = Cropper(crop_cfg=crop_cfg, provider=onnx_device) - pipeline.cropper = cropper pipeline.live_portrait_wrapper.cfg.flag_eye_retargeting = eye_retargeting pipeline.live_portrait_wrapper.cfg.eyes_retargeting_multiplier = eyes_retargeting_multiplier pipeline.live_portrait_wrapper.cfg.flag_lip_retargeting = lip_retargeting @@ -280,7 +259,7 @@ class LivePortraitProcess: cropped_out_list = [] full_out_list = [] for img in source_image_np: - cropped_frames, full_frame = pipeline.execute(img, driving_images_np) + cropped_frames, full_frame = pipeline.execute(img, driving_images_np, crop_info) cropped_tensors = [torch.from_numpy(np_array) for np_array in cropped_frames] cropped_tensors_out = torch.stack(cropped_tensors) / 255 cropped_tensors_out = cropped_tensors_out.cpu().float() @@ -297,11 +276,119 @@ class LivePortraitProcess: return (cropped_tensors_out, full_tensors_out) +class LivePortraitCropper: + @classmethod + def INPUT_TYPES(s): + return {"required": { + + "source_image": ("IMAGE",), + "dsize": ("INT", {"default": 512, "min": 64, "max": 2048}), + "scale": ("FLOAT", {"default": 2.3, "min": 1.0, "max": 4.0, "step": 0.01}), + "vx_ratio": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}), + "vy_ratio": ("FLOAT", {"default": -0.125, "min": -1.0, "max": 1.0, "step": 0.01}), + }, + "optional": { + "onnx_device": ( + [ + 'CPU', + 'CUDA', + ], { + "default": 'CPU' + }), + } + } + + RETURN_TYPES = ("IMAGE", "CROPINFO", "IMAGE",) + RETURN_NAMES = ("cropped_image", "crop_info", "keypoints_image",) + FUNCTION = "process" + CATEGORY = "LivePortrait" + + def process(self, source_image, dsize, scale, vx_ratio, vy_ratio, onnx_device='CUDA'): + source_image_np = (source_image * 255).byte().numpy() + + crop_cfg = CropConfig( + dsize = dsize, + scale = scale, + vx_ratio = vx_ratio, + vy_ratio = vy_ratio, + ) + + cropper = Cropper(crop_cfg=crop_cfg, provider=onnx_device) + crop_info, keypoints_img = cropper.crop_single_image(source_image_np[0]) + + keypoints_image_tensor = torch.from_numpy(keypoints_img) / 255 + keypoints_image_tensor = keypoints_image_tensor.unsqueeze(0).cpu().float() + print(keypoints_image_tensor.shape) + + + cropped_image = crop_info['img_crop_256x256'] + cropped_tensors = torch.from_numpy(cropped_image) / 255 + cropped_tensors = cropped_tensors.unsqueeze(0).cpu().float() + + print(cropped_tensors.shape) + + return (cropped_tensors, crop_info, keypoints_image_tensor) + +class KeypointScaler: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "crop_info": ("CROPINFO", {"default": {}}), + "scale": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 10.0, "step": 0.001}), + "offset_x": ("INT", {"default": 0, "min": -1024, "max": 1024, "step": 1}), + "offset_y": ("INT", {"default": 0, "min": -1024, "max": 1024, "step": 1}), + + } + } + + RETURN_TYPES = ("CROPINFO", "IMAGE",) + RETURN_NAMES = ("crop_info", "keypoints_image",) + FUNCTION = "process" + CATEGORY = "LivePortrait" + + def process(self, crop_info, offset_x, offset_y, scale): + + keypoints = crop_info['lmk_crop'].copy() + + # Create an offset array + # Calculate the centroid of the keypoints + centroid = keypoints.mean(axis=0) + + # Translate keypoints to origin by subtracting the centroid + translated_keypoints = keypoints - centroid + + # Scale the translated keypoints + scaled_keypoints = translated_keypoints * scale + + # Translate scaled keypoints back to original position and then apply the offset + final_keypoints = scaled_keypoints + centroid + np.array([offset_x, offset_y]) + + crop_info['lmk_crop'] = final_keypoints + + # Draw each landmark as a circle + width, height = 512, 512 + blank_image = np.zeros((height, width, 3), dtype=np.uint8) * 255 + for (x, y) in final_keypoints: + # Ensure the coordinates are within the dimensions of the blank image + if 0 <= x < width and 0 <= y < height: + cv2.circle(blank_image, (int(x), int(y)), radius=2, color=(0, 0, 255)) + + keypoints_image = cv2.cvtColor(blank_image, cv2.COLOR_BGR2RGB) + keypoints_image_tensor = torch.from_numpy(keypoints_image) / 255 + keypoints_image_tensor = keypoints_image_tensor.unsqueeze(0).cpu().float() + print(keypoints_image_tensor.shape) + + return (crop_info, keypoints_image_tensor,) + NODE_CLASS_MAPPINGS = { "DownloadAndLoadLivePortraitModels": DownloadAndLoadLivePortraitModels, "LivePortraitProcess": LivePortraitProcess, + "LivePortraitCropper": LivePortraitCropper, + "KeypointScaler": KeypointScaler } NODE_DISPLAY_NAME_MAPPINGS = { "DownloadAndLoadLivePortraitModels": "(Down)Load LivePortraitModels", "LivePortraitProcess": "LivePortraitProcess", + "LivePortraitCropper": "LivePortraitCropper", + "KeypointScaler": "KeypointScaler" } \ No newline at end of file