From 92fad03ee507021040eeb9110c4149711c516e10 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Sun, 21 Jul 2024 20:20:46 +0300 Subject: [PATCH] restructure a bit for more caching --- examples/liveportrait_video_example_02.json | 289 +++++++++++--------- liveportrait/live_portrait_pipeline.py | 29 +- nodes.py | 32 ++- 3 files changed, 193 insertions(+), 157 deletions(-) diff --git a/examples/liveportrait_video_example_02.json b/examples/liveportrait_video_example_02.json index 8c77c7d..4684ca6 100644 --- a/examples/liveportrait_video_example_02.json +++ b/examples/liveportrait_video_example_02.json @@ -1,6 +1,6 @@ { - "last_node_id": 188, - "last_link_id": 443, + "last_node_id": 189, + "last_link_id": 447, "nodes": [ { "id": 81, @@ -252,8 +252,8 @@ "name": "image", "type": "IMAGE", "links": [ - 296, - 408 + 408, + 445 ], "shape": 3, "slot_index": 0 @@ -281,94 +281,6 @@ "Node name for S&R": "GetImageSizeAndCount" } }, - { - "id": 1, - "type": "DownloadAndLoadLivePortraitModels", - "pos": [ - 546, - -518 - ], - "size": { - "0": 302.43463134765625, - "1": 58 - }, - "flags": {}, - "order": 0, - "mode": 0, - "outputs": [ - { - "name": "live_portrait_pipe", - "type": "LIVEPORTRAITPIPE", - "links": [ - 407 - ], - "shape": 3, - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "DownloadAndLoadLivePortraitModels" - }, - "widgets_values": [ - "fp16" - ] - }, - { - "id": 131, - "type": "LivePortraitCropper", - "pos": [ - -3, - -755 - ], - "size": { - "0": 330, - "1": 246 - }, - "flags": {}, - "order": 16, - "mode": 0, - "inputs": [ - { - "name": "cropper", - "type": "LPCROPPER", - "link": 295 - }, - { - "name": "source_image", - "type": "IMAGE", - "link": 296 - } - ], - "outputs": [ - { - "name": "cropped_image", - "type": "IMAGE", - "links": null, - "shape": 3 - }, - { - "name": "crop_info", - "type": "CROPINFO", - "links": [ - 413 - ], - "shape": 3, - "slot_index": 1 - } - ], - "properties": { - "Node name for S&R": "LivePortraitCropper" - }, - "widgets_values": [ - 512, - 2.3, - 0, - -0.125, - 0, - "large-small", - true - ] - }, { "id": 129, "type": "LivePortraitLoadCropper", @@ -381,15 +293,15 @@ "1": 82 }, "flags": {}, - "order": 1, + "order": 0, "mode": 0, "outputs": [ { "name": "cropper", "type": "LPCROPPER", "links": [ - 295, - 403 + 403, + 444 ], "shape": 3, "slot_index": 0 @@ -440,7 +352,7 @@ "1": 270 }, "flags": {}, - "order": 2, + "order": 1, "mode": 0, "outputs": [ { @@ -588,7 +500,7 @@ { "name": "crop_info", "type": "CROPINFO", - "link": 413, + "link": 447, "slot_index": 1 }, { @@ -661,7 +573,7 @@ "1": 82 }, "flags": {}, - "order": 3, + "order": 2, "mode": 0, "outputs": [ { @@ -694,8 +606,8 @@ -371 ], "size": [ - 1103.397512192384, - 697.5853490318566 + 1103.3974609375, + 697.5853313577586 ], "flags": {}, "order": 20, @@ -747,7 +659,7 @@ "hidden": false, "paused": false, "params": { - "filename": "LivePortrait_00009.mp4", + "filename": "LivePortrait_00001.mp4", "subfolder": "", "type": "temp", "format": "video/h264-mp4", @@ -988,6 +900,11 @@ "name": "source_image", "type": "IMAGE", "link": 437 + }, + { + "name": "source_image", + "type": "IMAGE", + "link": null } ], "outputs": [ @@ -1076,7 +993,7 @@ "1": 58 }, "flags": {}, - "order": 4, + "order": 3, "mode": 0, "properties": { "text": "" @@ -1094,12 +1011,12 @@ 39, 137 ], - "size": [ - 229.16442871093705, - 95.92726440429658 - ], + "size": { + "0": 229.1644287109375, + "1": 95.92726135253906 + }, "flags": {}, - "order": 5, + "order": 4, "mode": 0, "properties": { "text": "" @@ -1109,6 +1026,102 @@ ], "color": "#432", "bgcolor": "#653" + }, + { + "id": 1, + "type": "DownloadAndLoadLivePortraitModels", + "pos": [ + -1244, + -870 + ], + "size": { + "0": 302.43463134765625, + "1": 58 + }, + "flags": {}, + "order": 5, + "mode": 0, + "outputs": [ + { + "name": "live_portrait_pipe", + "type": "LIVEPORTRAITPIPE", + "links": [ + 407, + 446 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "DownloadAndLoadLivePortraitModels" + }, + "widgets_values": [ + "fp16" + ] + }, + { + "id": 189, + "type": "LivePortraitCropper", + "pos": [ + -3, + -755 + ], + "size": { + "0": 330, + "1": 242 + }, + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "pipeline", + "type": "LIVEPORTRAITPIPE", + "link": 446, + "slot_index": 0 + }, + { + "name": "cropper", + "type": "LPCROPPER", + "link": 444 + }, + { + "name": "source_image", + "type": "IMAGE", + "link": 445 + } + ], + "outputs": [ + { + "name": "cropped_image", + "type": "IMAGE", + "links": null, + "shape": 3, + "slot_index": 0 + }, + { + "name": "crop_info", + "type": "CROPINFO", + "links": [ + 447 + ], + "shape": 3, + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "LivePortraitCropper" + }, + "widgets_values": [ + 512, + 2.3, + 0, + -0.125, + 0, + "large-small", + true + ] } ], "links": [ @@ -1128,22 +1141,6 @@ 0, "MASK" ], - [ - 295, - 129, - 0, - 131, - 0, - "LPCROPPER" - ], - [ - 296, - 78, - 0, - 131, - 1, - "IMAGE" - ], [ 302, 133, @@ -1200,14 +1197,6 @@ 4, "MASK" ], - [ - 413, - 131, - 1, - 176, - 1, - "CROPINFO" - ], [ 426, 174, @@ -1287,16 +1276,48 @@ 186, 0, "IMAGE" + ], + [ + 444, + 129, + 0, + 189, + 1, + "LPCROPPER" + ], + [ + 445, + 78, + 0, + 189, + 2, + "IMAGE" + ], + [ + 446, + 1, + 0, + 189, + 0, + "LIVEPORTRAITPIPE" + ], + [ + 447, + 189, + 1, + 176, + 1, + "CROPINFO" ] ], "groups": [], "config": {}, "extra": { "ds": { - "scale": 0.5644739300537777, + "scale": 0.6830134553650711, "offset": { - "0": 1874.4561767578125, - "1": 1039.090087890625 + "0": 1494.1269197578126, + "1": 923.177290890625 } } }, diff --git a/liveportrait/live_portrait_pipeline.py b/liveportrait/live_portrait_pipeline.py index 66b51ee..7195c84 100644 --- a/liveportrait/live_portrait_pipeline.py +++ b/liveportrait/live_portrait_pipeline.py @@ -10,7 +10,7 @@ import numpy as np from .config.inference_config import InferenceConfig import torch from .utils.camera import get_rotation_matrix -from .utils.crop import _transform_img, _transform_img_kornia +from .utils.crop import _transform_img_kornia from .live_portrait_wrapper import LivePortraitWrapper from .utils.retargeting_utils import calc_eye_close_ratio, calc_lip_close_ratio from .utils.filter import smooth @@ -68,23 +68,10 @@ class LivePortraitPipeline(object): else: total_frames = driving_images.shape[0] - - source_info = [] - source_rot_list = [] - f_s_list = [] - for i in tqdm(range(source_np.shape[0]), desc='Processing source images...', total=source_np.shape[0]): - #get source keypoints info - img_crop_256x256 = crop_info["crop_info_list"][i]["img_crop_256x256"] - I_s = self.live_portrait_wrapper.prepare_source(img_crop_256x256) - x_s_info = self.live_portrait_wrapper.get_kp_info(I_s) - f_s = self.live_portrait_wrapper.extract_feature_3d(I_s) - f_s_list.append(f_s) - source_info.append(x_s_info) - - R_s = get_rotation_matrix( - x_s_info["pitch"], x_s_info["yaw"], x_s_info["roll"] - ) - source_rot_list.append(R_s) + source_info = crop_info["source_info"] + source_rot_list = crop_info["source_rot_list"] + f_s_list = crop_info["f_s_list"] + x_s_list = crop_info["x_s_list"] driving_info = [] driving_exp_list = [] @@ -137,11 +124,11 @@ class LivePortraitPipeline(object): x_d_info = driving_info[i] x_s_info = source_info[safe_index] - x_c_s = x_s_info["kp"] - R_s = source_rot_list[safe_index] f_s = f_s_list[safe_index] - x_s = self.live_portrait_wrapper.transform_keypoint(x_s_info) + x_s = x_s_list[safe_index] + + x_c_s = x_s_info["kp"] #lip zero if inference_cfg.flag_lip_zero: diff --git a/nodes.py b/nodes.py index 59f51f4..bd24a64 100644 --- a/nodes.py +++ b/nodes.py @@ -21,6 +21,7 @@ from .liveportrait.modules.appearance_feature_extractor import ( from .liveportrait.modules.stitching_retargeting_network import ( StitchingRetargetingNetwork, ) +from .liveportrait.utils.camera import get_rotation_matrix import logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') @@ -402,6 +403,7 @@ class LivePortraitCropper: @classmethod def INPUT_TYPES(s): return {"required": { + "pipeline": ("LIVEPORTRAITPIPE",), "cropper": ("LPCROPPER",), "source_image": ("IMAGE",), "dsize": ("INT", {"default": 512, "min": 64, "max": 2048}), @@ -429,7 +431,7 @@ class LivePortraitCropper: FUNCTION = "process" CATEGORY = "LivePortrait" - def process(self, cropper, source_image, dsize, scale, vx_ratio, vy_ratio, face_index, face_index_order, rotate): + def process(self, pipeline, cropper, source_image, dsize, scale, vx_ratio, vy_ratio, face_index, face_index_order, rotate): source_image_np = (source_image * 255).byte().numpy() crop_info_list = [] @@ -446,6 +448,28 @@ class LivePortraitCropper: cropped_images_list.append(cropped_image) pbar.update(1) + + source_info = [] + source_rot_list = [] + f_s_list = [] + x_s_list = [] + + for i in tqdm(range(source_image_np.shape[0]), desc='Processing source images...', total=source_image_np.shape[0]): + #get source keypoints info + img_crop_256x256 = crop_info_list[i]["img_crop_256x256"] + I_s = pipeline.live_portrait_wrapper.prepare_source(img_crop_256x256) + x_s_info = pipeline.live_portrait_wrapper.get_kp_info(I_s) + f_s = pipeline.live_portrait_wrapper.extract_feature_3d(I_s) + f_s_list.append(f_s) + source_info.append(x_s_info) + + x_s = pipeline.live_portrait_wrapper.transform_keypoint(x_s_info) + x_s_list.append(x_s) + + R_s = get_rotation_matrix( + x_s_info["pitch"], x_s_info["yaw"], x_s_info["roll"] + ) + source_rot_list.append(R_s) cropped_tensors_out = ( torch.stack([torch.from_numpy(np_array) for np_array in cropped_images_list]) @@ -453,7 +477,11 @@ class LivePortraitCropper: ) crop_info_dict = { - 'crop_info_list': crop_info_list + 'crop_info_list': crop_info_list, + 'source_rot_list': source_rot_list, + 'f_s_list': f_s_list, + 'x_s_list': x_s_list, + 'source_info': source_info } return (cropped_tensors_out, crop_info_dict)