From 978efbd075e534bdd8bde4c7e8f70085246feca1 Mon Sep 17 00:00:00 2001
From: Rei Koshka <purrogurama@gmail.com>
Date: Wed, 22 Nov 2023 02:18:25 +0300
Subject: [PATCH] Add custom node `ComfyCouple`

---
 __init__.py     |  5 ++-
 comfy_couple.py | 97 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 100 insertions(+), 2 deletions(-)
 create mode 100644 comfy_couple.py

diff --git a/__init__.py b/__init__.py
index 2dee0d6..f0ac486 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,2 +1,3 @@
-from .attention_couple import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
-__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
\ No newline at end of file
+from .comfy_couple import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
diff --git a/comfy_couple.py b/comfy_couple.py
new file mode 100644
index 0000000..174df6d
--- /dev/null
+++ b/comfy_couple.py
@@ -0,0 +1,97 @@
+from nodes import MAX_RESOLUTION, ConditioningCombine, ConditioningSetMask
+from comfy_extras.nodes_mask import MaskComposite, SolidMask
+
+from .attention_couple import AttentionCouple
+
+class ComfyCouple:
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "positive_1": ("CONDITIONING",),
+                "positive_2": ("CONDITIONING",),
+                "negative": ("CONDITIONING",),
+                "orientation": (["horizontal", "vertical"],),
+                "center": ("FLOAT", {"default": 0.5, "min": 0, "max": 1.0, "step": 0.01}),
+                "width": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8}),
+                "height": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8}),
+            }
+        }
+
+    RETURN_TYPES = (
+        "MODEL",
+        "CONDITIONING",
+        "CONDITIONING",
+    )
+
+    FUNCTION = "process"
+    CATEGORY = "loaders"
+
+    def process(
+            self,
+            model,
+            positive_1,
+            positive_2,
+            negative,
+            orientation,
+            center,
+            width,
+            height,
+    ):
+        mask_rect_first_x = None
+        mask_rect_first_y = None
+        mask_rect_first_width = None
+        mask_rect_first_height = None
+
+        mask_rect_second_x = None
+        mask_rect_second_y = None
+        mask_rect_second_width = None
+        mask_rect_second_height = None
+
+        if orientation == "horizontal":
+            width_first = int(width * center)
+
+            mask_rect_first_x = width_first
+            mask_rect_first_y = 0
+            mask_rect_first_width = width - width_first
+            mask_rect_first_height = height
+            mask_rect_second_x = 0
+            mask_rect_second_y = 0
+            mask_rect_second_width = width_first
+            mask_rect_second_height = height
+        elif orientation == "vertical":
+            height_first = int(width * center)
+
+            mask_rect_first_x = 0
+            mask_rect_first_y = height_first
+            mask_rect_first_width = width
+            mask_rect_first_height = height - height_first
+            mask_rect_second_x = 0
+            mask_rect_second_y = 0
+            mask_rect_second_width = width
+            mask_rect_second_height = height_first
+
+        solid_mask_zero = SolidMask().solid(0.0, width, height)[0]
+
+        solid_mask_first = SolidMask().solid(1.0, mask_rect_first_width, mask_rect_first_height)[0]
+        solid_mask_second = SolidMask().solid(1.0, mask_rect_second_width, mask_rect_second_height)[0]
+
+        mask_composite_first = MaskComposite().combine(solid_mask_zero, solid_mask_first, mask_rect_first_x, mask_rect_first_y, "add")[0]
+        mask_composite_second = MaskComposite().combine(solid_mask_zero, solid_mask_second, mask_rect_second_x, mask_rect_second_y, "add")[0]
+
+        conditioning_mask_first = ConditioningSetMask().append(positive_1, mask_composite_second, "default", 1.0)[0]
+        conditioning_mask_second = ConditioningSetMask().append(positive_2, mask_composite_first, "default", 1.0)[0]
+
+        positive_combined = ConditioningCombine().combine(conditioning_mask_first, conditioning_mask_second)[0]
+
+        return AttentionCouple().attention_couple(model, positive_combined, negative, "Attention")
+
+NODE_CLASS_MAPPINGS = {
+    "Comfy Couple": ComfyCouple
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "Comfy Couple": "Comfy Couple",
+}