From 4412c67aa690d14fdb8acf673ac4f7ee1eb4c6b1 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sat, 4 May 2024 15:12:43 +0300
Subject: [PATCH] Add AppendInstanceDiffusionTracking -node

---
 __init__.py          |  1 +
 nodes/curve_nodes.py | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/__init__.py b/__init__.py
index 811716c..15732f9 100644
--- a/__init__.py
+++ b/__init__.py
@@ -114,6 +114,7 @@ NODE_CONFIG = {
     "GLIGENTextBoxApplyBatchCoords": {"class": GLIGENTextBoxApplyBatchCoords},
     "Intrinsic_lora_sampling": {"class": Intrinsic_lora_sampling, "name": "Intrinsic Lora Sampling"},
     "CreateInstanceDiffusionTracking": {"class": CreateInstanceDiffusionTracking},
+    "AppendInstanceDiffusionTracking": {"class": AppendInstanceDiffusionTracking, "name": "Append Diffusion Tracking"},
 }
 
 def generate_node_mappings(node_config):
diff --git a/nodes/curve_nodes.py b/nodes/curve_nodes.py
index 4eee557..338a873 100644
--- a/nodes/curve_nodes.py
+++ b/nodes/curve_nodes.py
@@ -740,7 +740,41 @@ for example:
 
         print(tracked)
         return (tracked, )
+
+class AppendInstanceDiffusionTracking:
     
+    RETURN_TYPES = ("TRACKING",)
+    RETURN_NAMES = ("tracking",)
+    FUNCTION = "append"
+    CATEGORY = "KJNodes/experimental"
+    DESCRIPTION = """
+Appends tracking data to be used with InstanceDiffusion:  
+https://github.com/logtd/ComfyUI-InstanceDiffusion  
+
+"""
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "tracking_1": ("TRACKING", {"forceInput": True}),
+                "tracking_2": ("TRACKING", {"forceInput": True}),
+        },
+    } 
+
+    def append(self, tracking_1, tracking_2):
+        tracking_copy = tracking_1.copy()
+        # Check for existing class names and class IDs, and raise an error if they exist
+        for class_name, class_data in tracking_2.items():
+            if class_name in tracking_copy:
+                for class_id in class_data.keys():
+                    if class_id in tracking_copy[class_name]:
+                        raise ValueError(f"Class ID {class_id} already exists for class name {class_name}. Cannot append tracking data.")
+            # If class name does not exist, add it
+            tracking_copy[class_name] = class_data
+
+        return (tracking_copy, )
+        
 class InterpolateCoords:
     
     RETURN_TYPES = ("STRING",)