From 0e04480507930209d44f54030acc131708bb029b Mon Sep 17 00:00:00 2001 From: Mario Klingemann Date: Thu, 25 Jul 2024 13:39:44 +0200 Subject: [PATCH] Update nodes.py Some custom nodes that generate masks do not return them in (batch_size,h,w) format but only (h,w) - this change handles this gracefully. --- nodes.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/nodes.py b/nodes.py index 83cdfd8..d3e5f2d 100644 --- a/nodes.py +++ b/nodes.py @@ -395,7 +395,10 @@ class LivePortraitComposite: cropped_image = cropped_image.permute(0, 3, 1, 2) if mask is not None: - crop_mask = mask.unsqueeze(-1).expand(-1, -1, -1, 3) + if len(mask.size())==2: + crop_mask = mask.unsqueeze(0).unsqueeze(-1).expand(-1, -1, -1, 3) + else: + crop_mask = mask.unsqueeze(-1).expand(-1, -1, -1, 3) else: log.info("Using default mask template") crop_mask = cv2.imread(os.path.join(script_directory, "liveportrait", "utils", "resources", "mask_template.png"), cv2.IMREAD_COLOR) @@ -829,4 +832,4 @@ NODE_DISPLAY_NAME_MAPPINGS = { "LivePortraitLoadMediaPipeCropper": "LivePortrait Load MediaPipeCropper", "LivePortraitLoadFaceAlignmentCropper": "LivePortrait Load FaceAlignmentCropper", "LivePortraitComposite": "LivePortrait Composite", - } \ No newline at end of file + }