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.
This commit is contained in:
Mario Klingemann
2024-07-25 13:39:44 +02:00
committed by GitHub
parent 2b23af0508
commit 0e04480507
+5 -2
View File
@@ -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",
}
}