From 63bcc99ac592b02c9b5ec9ce88b013c8cf0a7ec7 Mon Sep 17 00:00:00 2001
From: Luis Quesada <elezeta@gmail.com>
Date: Thu, 25 Apr 2024 08:48:15 +0200
Subject: [PATCH] Add Prompt Combinator Merger node

---
 README.md            |  2 ++
 __init__.py          |  7 +++--
 prompt_combinator.py | 65 ++++++++++++++++++++++++++++++++++++++++++--
 3 files changed, 70 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 98bf739..9e22612 100644
--- a/README.md
+++ b/README.md
@@ -7,6 +7,8 @@ Check ComfyUI here: https://github.com/comfyanonymous/ComfyUI
 
 "🔢 Prompt Combinator" is a node that generates all possible combinations of prompts from several lists of strings.
 
+"🔢 Prompt Combinator Merger" is a node that enables merging the output of two different "🔢 Prompt Combinator" nodes.
+
 Download the following example workflow from [here](prompt-combinator_example_workflow.json) or drag and drop the screenshot into ComfyUI.
 
 ![Workflow](prompt-combinator_example_workflow.png)
diff --git a/__init__.py b/__init__.py
index f68e095..e48eee0 100755
--- a/__init__.py
+++ b/__init__.py
@@ -1,11 +1,14 @@
 from .prompt_combinator import PromptCombinator
+from .prompt_combinator import PromptCombinatorMerger
 
 NODE_CLASS_MAPPINGS = {
-    "PromptCombinator": PromptCombinator
+    "PromptCombinator": PromptCombinator,
+    "PromptCombinatorMerger": PromptCombinatorMerger
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "PromptCombinator": "🔢 Prompt Combinator"
+    "PromptCombinator": "🔢 Prompt Combinator",
+    "PromptCombinatorMerger": "🔢 Prompt Combinator Merger"
 }
 
 __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
diff --git a/prompt_combinator.py b/prompt_combinator.py
index c5ccf0b..fdc7737 100755
--- a/prompt_combinator.py
+++ b/prompt_combinator.py
@@ -7,11 +7,10 @@ class PromptCombinator:
     https://github.com/lquesada/ComfyUI-Prompt-Combinator
 
     Node that generates all possible combinations of prompts from several lists of strings.
-
     """
     def __init__(self):
         pass
-    
+
     @classmethod
     def INPUT_TYPES(cls):
         return {
@@ -82,3 +81,65 @@ class PromptCombinator:
         combinations, (ids1, ids2, ids3, ids4) = self.combine_descriptions_and_ids(input_list_1, input_list_2, input_list_3, input_list_4, id_separator=id_separator, comment_prefix=comment_prefix)
         return (combinations, ids1, ids2, ids3, ids4)
 
+
+class PromptCombinatorMerger:
+    """
+    ComfyUI-Prompt-Combinator
+    https://github.com/lquesada/ComfyUI-Prompt-Combinator
+
+    Node that merges the prompts and IDs out of two PromptCombinators.
+    """
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "combinations_1": ("STRING", {"default": '', "multiline": True, "forceInput": True}),
+                "combinations_2": ("STRING", {"default": '', "multiline": True, "forceInput": True}),
+            },
+            "optional": {
+                "input_id1_list_1": ("STRING", {"default": '', "multiline": True, "forceInput": True}),
+                "input_id2_list_1": ("STRING", {"default": '', "multiline": True, "forceInput": True}),
+                "input_id3_list_1": ("STRING", {"default": '', "multiline": True, "forceInput": True}),
+                "input_id4_list_1": ("STRING", {"default": '', "multiline": True, "forceInput": True}),
+                "input_id1_list_2": ("STRING", {"default": '', "multiline": True, "forceInput": True}),
+                "input_id2_list_2": ("STRING", {"default": '', "multiline": True, "forceInput": True}),
+                "input_id3_list_2": ("STRING", {"default": '', "multiline": True, "forceInput": True}),
+                "input_id4_list_2": ("STRING", {"default": '', "multiline": True, "forceInput": True}),
+            }
+        }
+    INPUT_IS_LIST = (True, True, True, True, True, True, True, True, True, True)
+
+    RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING", "STRING")
+    RETURN_NAMES = ("combinations", "id_1", "id_2", "id_3", "id_4")
+    OUTPUT_IS_LIST = (True, True, True, True, True)
+
+    FUNCTION = "execute"
+
+    CATEGORY = "batch"
+
+    def merge_lists(self, list1, list2):
+        """
+        Helper method to merge two lists, handling cases where one or both may be undefined or empty.
+        """
+        if not list1:
+            list1 = []
+        if not list2:
+            list2 = []
+        return list1 + list2
+
+    def execute(self, combinations_1, input_id1_list_1, input_id2_list_1, input_id3_list_1, input_id4_list_1, combinations_2, input_id1_list_2, input_id2_list_2, input_id3_list_2, input_id4_list_2):
+        # Combine the combinations and IDs lists from two sets of inputs
+        combined_combinations = self.merge_lists(combinations_1, combinations_2)
+        combined_ids1 = self.merge_lists(input_id1_list_1, input_id1_list_2)
+        combined_ids2 = self.merge_lists(input_id2_list_1, input_id2_list_2)
+        combined_ids3 = self.merge_lists(input_id3_list_1, input_id3_list_2)
+        combined_ids4 = self.merge_lists(input_id4_list_1, input_id4_list_2)
+
+        # Check if all outputs are empty, return None if so
+        if not any([combined_combinations, combined_ids1, combined_ids2, combined_ids3, combined_ids4]):
+            return None
+
+        return (combined_combinations, combined_ids1, combined_ids2, combined_ids3, combined_ids4)