From 602a59066e37b506dd518f168ef18d3aa55bb8c6 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Mon, 15 Apr 2024 12:37:30 +1000 Subject: [PATCH] fix(nodes): handle invert in alpha_mask_to_tensor --- invokeai/app/invocations/mask.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/invokeai/app/invocations/mask.py b/invokeai/app/invocations/mask.py index fcc5b04caf..4441ebdebb 100644 --- a/invokeai/app/invocations/mask.py +++ b/invokeai/app/invocations/mask.py @@ -48,12 +48,15 @@ class AlphaMaskToTensorInvocation(BaseInvocation): """Convert a mask image to a tensor. Opaque regions are 1 and transparent regions are 0.""" image: ImageField = InputField(description="The mask image to convert.") - invert: bool = InputField(default=False, description="Invert the mask (1s become 0s and 0s become 1s).") + invert: bool = InputField(default=False, description="Whether to invert the mask.") def invoke(self, context: InvocationContext) -> MaskOutput: image = context.images.get_pil(self.image.image_name) mask = torch.zeros((1, image.height, image.width), dtype=torch.bool) - mask[0] = torch.tensor(np.array(image)[:, :, 3] > 0, dtype=torch.bool) + if self.invert: + mask[0] = torch.tensor(np.array(image)[:, :, 3] == 0, dtype=torch.bool) + else: + mask[0] = torch.tensor(np.array(image)[:, :, 3] > 0, dtype=torch.bool) return MaskOutput( mask=TensorField(tensor_name=context.tensors.save(mask)), height=image.height, width=image.width