mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
Add a MaskField primitive, and add a mask to the ConditioningField primitive type.
This commit is contained in:
parent
132aadca15
commit
46d83a3026
41
invokeai/app/invocations/conditioning.py
Normal file
41
invokeai/app/invocations/conditioning.py
Normal 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)
|
@ -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]]):
|
||||
|
Loading…
Reference in New Issue
Block a user