From 13242fc3bfdf23ef8706cdff5c57dffa689735b1 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sun, 22 Sep 2024 00:35:51 +0300
Subject: [PATCH] Add SetShakkerLabsUnionControlNetType -node

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

diff --git a/__init__.py b/__init__.py
index 50576e9..08e968b 100644
--- a/__init__.py
+++ b/__init__.py
@@ -103,6 +103,7 @@ NODE_CONFIG = {
     "EmptyLatentImagePresets": {"class": EmptyLatentImagePresets, "name": "Empty Latent Image Presets"},
     "EmptyLatentImageCustomPresets": {"class": EmptyLatentImageCustomPresets, "name": "Empty Latent Image Custom Presets"},
     "ModelPassThrough": {"class": ModelPassThrough, "name": "ModelPass"},
+    "SetShakkerLabsUnionControlNetType": {"class": SetShakkerLabsUnionControlNetType, "name": "Set Shakker Labs Union ControlNet Type"},
     #audioscheduler stuff
     "NormalizedAmplitudeToMask": {"class": NormalizedAmplitudeToMask},
     "NormalizedAmplitudeToFloatList": {"class": NormalizedAmplitudeToFloatList},
diff --git a/nodes/nodes.py b/nodes/nodes.py
index e2eaf94..638f73b 100644
--- a/nodes/nodes.py
+++ b/nodes/nodes.py
@@ -2023,4 +2023,36 @@ class CustomControlNetWeightsFluxFromList:
 
         weights = ControlWeights.controlnet(weights_input=list_of_floats, uncond_multiplier=uncond_multiplier, extras=cn_extras)
         print(weights.weights_input)
-        return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights)))
\ No newline at end of file
+        return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights)))
+    
+SHAKKERLABS_UNION_CONTROLNET_TYPES = {
+    "canny": 0,
+    "tile": 1,
+    "depth": 2,
+    "blur": 3,
+    "pose": 4,
+    "gray": 5,
+    "low quality": 6,
+}
+
+class SetShakkerLabsUnionControlNetType:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"control_net": ("CONTROL_NET", ),
+                             "type": (["auto"] + list(SHAKKERLABS_UNION_CONTROLNET_TYPES.keys()),)
+                             }}
+
+    CATEGORY = "conditioning/controlnet"
+    RETURN_TYPES = ("CONTROL_NET",)
+
+    FUNCTION = "set_controlnet_type"
+
+    def set_controlnet_type(self, control_net, type):
+        control_net = control_net.copy()
+        type_number = SHAKKERLABS_UNION_CONTROLNET_TYPES.get(type, -1)
+        if type_number >= 0:
+            control_net.set_extra_arg("control_type", [type_number])
+        else:
+            control_net.set_extra_arg("control_type", [])
+
+        return (control_net,)
\ No newline at end of file