output masks

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

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: