Add a MaskField primitive, and add a mask to the ConditioningField primitive type.

This commit is contained in:
Ryan Dick 2024-02-13 14:24:46 -05:00
parent 132aadca15
commit 46d83a3026
2 changed files with 52 additions and 1 deletions

View File

@ -0,0 +1,41 @@
import numpy as np
import torch
from PIL.Image import Image
from invokeai.app.invocations.baseinvocation import BaseInvocation, InputField, InvocationContext, invocation
from invokeai.app.invocations.primitives import ConditioningField, ConditioningOutput, ImageField
@invocation(
"add_conditioning_mask",
title="Add Conditioning Mask",
tags=["conditioning"],
category="conditioning",
version="1.0.0",
)
class AddConditioningMaskInvocation(BaseInvocation):
"""Add a mask to an existing conditioning tensor."""
conditioning: ConditioningField = InputField(description="The conditioning tensor to add a mask to.")
image: ImageField = InputField(
description="A mask image to add to the conditioning tensor. Only the first channel of the image is used. "
"Pixels <128 are excluded from the mask, pixels >=128 are included in the mask."
)
@staticmethod
def convert_image_to_mask(image: Image) -> torch.Tensor:
"""Convert a PIL image to a uint8 mask tensor."""
np_image = np.array(image)
torch_image = torch.from_numpy(np_image[0, :, :])
mask = torch_image >= 128
return mask.to(dtype=torch.uint8)
def invoke(self, context: InvocationContext) -> ConditioningOutput:
image = context.services.images.get_pil_image(self.image.image_name)
mask = self.convert_image_to_mask(image)
mask_name = f"{context.graph_execution_state_id}__{self.id}_conditioning_mask"
context.services.latents.save(mask_name, mask)
self.conditioning.mask_name = mask_name
return ConditioningOutput(conditioning=self.conditioning)

View File

@ -203,6 +203,12 @@ class DenoiseMaskField(BaseModel):
gradient: bool = Field(default=False, description="Used for gradient inpainting")
class MaskField(BaseModel):
"""A mask primitive field."""
mask_name: str = Field(description="The name of a spatial mask. dtype: bool, shape: (1, h, w).")
class LatentsField(BaseModel):
"""A latents tensor primitive field"""
@ -226,7 +232,11 @@ class ConditioningField(BaseModel):
"""A conditioning tensor primitive value"""
conditioning_name: str = Field(description="The name of conditioning tensor")
# endregion
mask: Optional[MaskField] = Field(
default=None,
description="The bool mask associated with this conditioning tensor. Excluded regions should be set to False, "
"included regions should be set to True.",
)
class MetadataField(RootModel[dict[str, Any]]):