From e0f12c762ea136a6e48d02ce7cef69e3fa01d085 Mon Sep 17 00:00:00 2001 From: Ryan Dick Date: Thu, 29 Aug 2024 17:41:06 +0000 Subject: [PATCH] Update MaskTensorToImageInvocation to support input mask tensors with or without a channel dimension. --- invokeai/app/invocations/mask.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/invokeai/app/invocations/mask.py b/invokeai/app/invocations/mask.py index 64d1b48e38..f7cbc6d832 100644 --- a/invokeai/app/invocations/mask.py +++ b/invokeai/app/invocations/mask.py @@ -126,7 +126,7 @@ class ImageMaskToTensorInvocation(BaseInvocation, WithMetadata): title="Tensor Mask to Image", tags=["mask"], category="mask", - version="1.0.0", + version="1.1.0", ) class MaskTensorToImageInvocation(BaseInvocation, WithMetadata, WithBoard): """Convert a mask tensor to an image.""" @@ -135,6 +135,11 @@ class MaskTensorToImageInvocation(BaseInvocation, WithMetadata, WithBoard): def invoke(self, context: InvocationContext) -> ImageOutput: mask = context.tensors.load(self.mask.tensor_name) + + # Squeeze the channel dimension if it exists. + if mask.dim() == 3: + mask = mask.squeeze(0) + # Ensure that the mask is binary. if mask.dtype != torch.bool: mask = mask > 0.5