From b681132da4cd05e9a40660863db0ea8ce52e992c Mon Sep 17 00:00:00 2001 From: Ryan Dick Date: Fri, 24 Jan 2025 22:04:46 +0000 Subject: [PATCH] Update InvertTensorMaskInvocation to handle mask tensors with dim 2 or 3. --- invokeai/app/invocations/mask.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/invokeai/app/invocations/mask.py b/invokeai/app/invocations/mask.py index e23216c8d5..8edcbb1f94 100644 --- a/invokeai/app/invocations/mask.py +++ b/invokeai/app/invocations/mask.py @@ -86,7 +86,7 @@ class AlphaMaskToTensorInvocation(BaseInvocation): title="Invert Tensor Mask", tags=["conditioning"], category="conditioning", - version="1.0.0", + version="1.1.0", classification=Classification.Beta, ) class InvertTensorMaskInvocation(BaseInvocation): @@ -96,6 +96,15 @@ class InvertTensorMaskInvocation(BaseInvocation): def invoke(self, context: InvocationContext) -> MaskOutput: mask = context.tensors.load(self.mask.tensor_name) + + # Verify dtype and shape. + assert mask.dtype == torch.bool + assert mask.dim() in [2, 3] + + # Unsqueeze the channel dimension if it is missing. The MaskOutput type expects a single channel. + if mask.dim() == 2: + mask = mask.unsqueeze(0) + inverted = ~mask return MaskOutput(