feat(nodes): add alpha mask to tensor invocation

This commit is contained in:
psychedelicious 2024-04-09 20:27:03 +10:00 committed by Kent Keirsey
parent c89a24d1ea
commit fc26f3e430

View File

@ -1,7 +1,8 @@
import numpy as np
import torch import torch
from invokeai.app.invocations.baseinvocation import BaseInvocation, InvocationContext, invocation from invokeai.app.invocations.baseinvocation import BaseInvocation, InvocationContext, invocation
from invokeai.app.invocations.fields import InputField, TensorField, WithMetadata from invokeai.app.invocations.fields import ImageField, InputField, TensorField, WithMetadata
from invokeai.app.invocations.primitives import MaskOutput from invokeai.app.invocations.primitives import MaskOutput
@ -34,3 +35,25 @@ class RectangleMaskInvocation(BaseInvocation, WithMetadata):
width=self.width, width=self.width,
height=self.height, height=self.height,
) )
@invocation(
"alpha_mask_to_tensor",
title="Alpha Mask to Tensor",
tags=["conditioning"],
category="conditioning",
version="1.0.0",
)
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.")
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)
return MaskOutput(
mask=TensorField(tensor_name=context.tensors.save(mask)), height=image.height, width=image.width
)