mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
62 lines
2.2 KiB
Python
62 lines
2.2 KiB
Python
import torch
|
|
|
|
from invokeai.app.invocations.baseinvocation import (
|
|
BaseInvocation,
|
|
InputField,
|
|
InvocationContext,
|
|
WithMetadata,
|
|
invocation,
|
|
)
|
|
from invokeai.app.invocations.primitives import ConditioningField, ConditioningOutput, MaskField, MaskOutput
|
|
|
|
|
|
@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.")
|
|
mask: MaskField = InputField(description="A mask to add to the conditioning tensor.")
|
|
|
|
def invoke(self, context: InvocationContext) -> ConditioningOutput:
|
|
self.conditioning.mask = self.mask
|
|
return ConditioningOutput(conditioning=self.conditioning)
|
|
|
|
|
|
@invocation(
|
|
"rectangle_mask",
|
|
title="Create Rectangle Mask",
|
|
tags=["conditioning"],
|
|
category="conditioning",
|
|
version="1.0.0",
|
|
)
|
|
class RectangleMaskInvocation(BaseInvocation, WithMetadata):
|
|
"""Create a rectangular mask."""
|
|
|
|
height: int = InputField(description="The height of the entire mask.")
|
|
width: int = InputField(description="The width of the entire mask.")
|
|
y_top: int = InputField(description="The top y-coordinate of the rectangular masked region (inclusive).")
|
|
x_left: int = InputField(description="The left x-coordinate of the rectangular masked region (inclusive).")
|
|
rectangle_height: int = InputField(description="The height of the rectangular masked region.")
|
|
rectangle_width: int = InputField(description="The width of the rectangular masked region.")
|
|
|
|
def invoke(self, context: InvocationContext) -> MaskOutput:
|
|
mask = torch.zeros((1, self.height, self.width), dtype=torch.bool)
|
|
mask[
|
|
:, self.y_top : self.y_top + self.rectangle_height, self.x_left : self.x_left + self.rectangle_width
|
|
] = True
|
|
|
|
mask_name = f"{context.graph_execution_state_id}__{self.id}_mask"
|
|
context.services.latents.save(mask_name, mask)
|
|
|
|
return MaskOutput(
|
|
mask=MaskField(mask_name=mask_name),
|
|
width=self.width,
|
|
height=self.height,
|
|
)
|