Add SetUnionControlNetType to set the type of the union controlnet model.

2025-08-02 23:14:49 +08:00 · 2024-07-16 17:01:40 -04:00
parent 821f93872e
commit 8270c62530
3 changed files with 45 additions and 1 deletions
--- a/comfy_extras/nodes_controlnet.py
+++ b/comfy_extras/nodes_controlnet.py
@@ -0,0 +1,37 @@
+
+UNION_CONTROLNET_TYPES = {"auto": -1,
+                          "openpose": 0,
+                          "depth": 1,
+                          "hed/pidi/scribble/ted": 2,
+                          "canny/lineart/anime_lineart/mlsd": 3,
+                          "normal": 4,
+                          "segment": 5,
+                          "tile": 6,
+                          "repaint": 7,
+                        }
+
+class SetUnionControlNetType:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"control_net": ("CONTROL_NET", ),
+                             "type": (list(UNION_CONTROLNET_TYPES.keys()),)
+                             }}
+
+    CATEGORY = "conditioning"
+    RETURN_TYPES = ("CONTROL_NET",)
+
+    FUNCTION = "set_controlnet_type"
+
+    def set_controlnet_type(self, control_net, type):
+        control_net = control_net.copy()
+        type_number = UNION_CONTROLNET_TYPES[type]
+        if type_number >= 0:
+            control_net.set_extra_arg("control_type", [type_number])
+        else:
+            control_net.set_extra_arg("control_type", [])
+
+        return (control_net,)
+
+NODE_CLASS_MAPPINGS = {
+    "SetUnionControlNetType": SetUnionControlNetType,
+}