WIP: Regional Prompting

invokeai/app/invocations/compel.py

@@ -5,7 +5,15 @@
 from compel.prompt_parser import Blend, Conjunction, CrossAttentionControlSubstitute, FlattenedPrompt, Fragment
 from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
 
-from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField, UIComponent
+from invokeai.app.invocations.fields import (
+    ConditioningField,
+    FieldDescriptions,
+    Input,
+    InputField,
+    MaskField,
+    OutputField,
+    UIComponent,
+)
 from invokeai.app.invocations.primitives import ConditioningOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.util.ti_utils import generate_ti_list
@@ -36,7 +44,7 @@
     title="Prompt",
     tags=["prompt", "compel"],
     category="conditioning",
-    version="1.0.1",
+    version="1.1.0",
 )
 class CompelInvocation(BaseInvocation):
     """Parse prompt using compel package to conditioning."""
@@ -51,6 +59,12 @@ class CompelInvocation(BaseInvocation):
         description=FieldDescriptions.clip,
         input=Input.Connection,
     )
+    mask: Optional[MaskField] = InputField(
+        default=None, description="A mask defining the region that this conditioning prompt applies to."
+    )
+    positive_cross_attn_mask_score: float = InputField(default=0.0, description="")
+    positive_self_attn_mask_score: float = InputField(default=1.0, description="")
+    self_attn_adjustment_end_step_percent: float = InputField(default=0.0, description="")
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> ConditioningOutput:
@@ -118,7 +132,15 @@ def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
 
         conditioning_name = context.conditioning.save(conditioning_data)
 
-        return ConditioningOutput.build(conditioning_name)
+        return ConditioningOutput(
+            conditioning=ConditioningField(
+                conditioning_name=conditioning_name,
+                mask=self.mask,
+                positive_cross_attn_mask_score=self.positive_cross_attn_mask_score,
+                positive_self_attn_mask_score=self.positive_self_attn_mask_score,
+                self_attn_adjustment_end_step_percent=self.self_attn_adjustment_end_step_percent,
+            )
+        )
 
 
 class SDXLPromptInvocationBase:
@@ -232,7 +254,7 @@ def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
     title="SDXL Prompt",
     tags=["sdxl", "compel", "prompt"],
     category="conditioning",
-    version="1.0.1",
+    version="1.1.0",
 )
 class SDXLCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase):
     """Parse prompt using compel package to conditioning."""
@@ -256,6 +278,13 @@ class SDXLCompelPromptInvocation(BaseInvocation, SDXLPromptInvocationBase):
     clip: ClipField = InputField(description=FieldDescriptions.clip, input=Input.Connection, title="CLIP 1")
     clip2: ClipField = InputField(description=FieldDescriptions.clip, input=Input.Connection, title="CLIP 2")
 
+    mask: Optional[MaskField] = InputField(
+        default=None, description="A mask defining the region that this conditioning prompt applies to."
+    )
+    positive_cross_attn_mask_score: float = InputField(default=0.0, description="")
+    positive_self_attn_mask_score: float = InputField(default=1.0, description="")
+    self_attn_adjustment_end_step_percent: float = InputField(default=0.0, description="")
+
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> ConditioningOutput:
         c1, c1_pooled, ec1 = self.run_clip_compel(
@@ -317,7 +346,15 @@ def invoke(self, context: InvocationContext) -> ConditioningOutput:
 
         conditioning_name = context.conditioning.save(conditioning_data)
 
-        return ConditioningOutput.build(conditioning_name)
+        return ConditioningOutput(
+            conditioning=ConditioningField(
+                conditioning_name=conditioning_name,
+                mask=self.mask,
+                positive_cross_attn_mask_score=self.positive_cross_attn_mask_score,
+                positive_self_attn_mask_score=self.positive_self_attn_mask_score,
+                self_attn_adjustment_end_step_percent=self.self_attn_adjustment_end_step_percent,
+            )
+        )
 
 
 @invocation(
@@ -366,7 +403,7 @@ def invoke(self, context: InvocationContext) -> ConditioningOutput:
 
         conditioning_name = context.conditioning.save(conditioning_data)
 
-        return ConditioningOutput.build(conditioning_name)
+        return ConditioningOutput(conditioning=ConditioningField(conditioning_name=conditioning_name))
 
 
 @invocation_output("clip_skip_output")

invokeai/app/invocations/conditioning.py

@@ -0,0 +1,40 @@
+import torch
+
+from invokeai.app.invocations.baseinvocation import (
+    BaseInvocation,
+    InvocationContext,
+    invocation,
+)
+from invokeai.app.invocations.fields import InputField, WithMetadata
+from invokeai.app.invocations.primitives import MaskField, MaskOutput
+
+
+@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 = context.tensors.save(mask)
+        return MaskOutput(
+            mask=MaskField(mask_name=mask_name),
+            width=self.width,
+            height=self.height,
+        )

invokeai/app/invocations/fields.py

@@ -194,6 +194,12 @@ class BoardField(BaseModel):
     board_id: str = Field(description="The id of the board")
 
 
+class MaskField(BaseModel):
+    """A mask primitive field."""
+
+    mask_name: str = Field(description="The name of the mask.")
+
+
 class DenoiseMaskField(BaseModel):
     """An inpaint mask field"""
 
@@ -225,7 +231,18 @@ 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.",
+    )
+    positive_cross_attn_mask_score: float = Field(
+        default=0.0,
+        # TODO(ryand): Add more details to this description
+        description="The weight of this conditioning tensor's mask relative to overlapping masks.",
+    )
+    positive_self_attn_mask_score: float = Field(default=1.0, description="")
+    self_attn_adjustment_end_step_percent: float = Field(default=0.0, description="")
 
 
 class MetadataField(RootModel):