diff --git a/liveportrait/live_portrait_pipeline.py b/liveportrait/live_portrait_pipeline.py index 303dff7..dc4bf42 100644 --- a/liveportrait/live_portrait_pipeline.py +++ b/liveportrait/live_portrait_pipeline.py @@ -67,6 +67,7 @@ class LivePortraitPipeline(object): I_p_lst = [] I_p_paste_lst = [] driving_lmk_lst = [] + out_mask_list = [] R_d_0, x_d_0_info = None, None total_frames = driving_images_np.shape[0] @@ -282,6 +283,7 @@ class LivePortraitPipeline(object): I_p_i_to_ori_blend = I_p_i_to_ori I_p_paste_lst.append(I_p_i_to_ori_blend) + out_mask_list.append(mask_ori) pbar.update(1) - return I_p_lst, I_p_paste_lst + return I_p_lst, I_p_paste_lst, out_mask_list diff --git a/nodes.py b/nodes.py index 2d4a5be..c89df12 100644 --- a/nodes.py +++ b/nodes.py @@ -280,10 +280,12 @@ class LivePortraitProcess: RETURN_TYPES = ( "IMAGE", "IMAGE", + "MASK", ) RETURN_NAMES = ( "cropped_images", "full_images", + "mask", ) FUNCTION = "process" CATEGORY = "LivePortrait" @@ -330,7 +332,7 @@ class LivePortraitProcess: cropped_out_list = [] full_out_list = [] - cropped_out_list, full_out_list = pipeline.execute( + cropped_out_list, full_out_list, out_mask_list = pipeline.execute( source_np, driving_images_np, crop_info['crop_info'], mismatch_method ) cropped_tensors_out = ( @@ -341,8 +343,16 @@ class LivePortraitProcess: torch.stack([torch.from_numpy(np_array) for np_array in full_out_list]) / 255 ) + mask_tensors_out = ( + torch.stack([torch.from_numpy(np_array) for np_array in out_mask_list]) + )[:, :, :, 0] + print("mask out tensor: ", mask_tensors_out.shape) - return (cropped_tensors_out.cpu().float(), full_tensors_out.cpu().float()) + return ( + cropped_tensors_out.cpu().float(), + full_tensors_out.cpu().float(), + mask_tensors_out.cpu().float() + ) class LivePortraitCropper: