From 4151fbfa8ab705d0a7e632c7d7c23ea5c436a60b Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Mon, 22 Jul 2024 11:27:32 -0400
Subject: [PATCH] Add error message on union controlnet (#4081)

---
 comfy/cldm/cldm.py               | 13 +++++++++++++
 comfy/cldm/control_types.py      | 11 +++++++++++
 comfy_extras/nodes_controlnet.py | 12 +-----------
 3 files changed, 25 insertions(+), 11 deletions(-)
 create mode 100644 comfy/cldm/control_types.py

diff --git a/comfy/cldm/cldm.py b/comfy/cldm/cldm.py
index 1d7294bd63d..9ec64a22751 100644
--- a/comfy/cldm/cldm.py
+++ b/comfy/cldm/cldm.py
@@ -13,6 +13,7 @@
 from ..ldm.modules.attention import SpatialTransformer
 from ..ldm.modules.diffusionmodules.openaimodel import UNetModel, TimestepEmbedSequential, ResBlock, Downsample
 from ..ldm.util import exists
+from .control_types import UNION_CONTROLNET_TYPES
 from collections import OrderedDict
 import comfy.ops
 from comfy.ldm.modules.attention import optimized_attention
@@ -390,6 +391,18 @@ def forward(self, x, hint, timesteps, context, y=None, **kwargs):
         if self.control_add_embedding is not None: #Union Controlnet
             control_type = kwargs.get("control_type", [])
 
+            if any([c >= self.num_control_type for c in control_type]):
+                max_type = max(control_type)
+                max_type_name = {
+                    v: k for k, v in UNION_CONTROLNET_TYPES.items()
+                }[max_type]
+                raise ValueError(
+                    f"Control type {max_type_name}({max_type}) is out of range for the number of control types" +
+                    f"({self.num_control_type}) supported.\n" +
+                    "Please consider using the ProMax ControlNet Union model.\n" +
+                    "https://huggingface.co/xinsir/controlnet-union-sdxl-1.0/tree/main"
+                )
+
             emb += self.control_add_embedding(control_type, emb.dtype, emb.device)
             if len(control_type) > 0:
                 if len(hint.shape) < 5:
diff --git a/comfy/cldm/control_types.py b/comfy/cldm/control_types.py
new file mode 100644
index 00000000000..fd6e56e8fa8
--- /dev/null
+++ b/comfy/cldm/control_types.py
@@ -0,0 +1,11 @@
+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,
+}
diff --git a/comfy_extras/nodes_controlnet.py b/comfy_extras/nodes_controlnet.py
index ef7cfc6ab49..566d51e33c6 100644
--- a/comfy_extras/nodes_controlnet.py
+++ b/comfy_extras/nodes_controlnet.py
@@ -1,14 +1,4 @@
-
-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,
-                        }
+from comfy.cldm.control_types import UNION_CONTROLNET_TYPES
 
 class SetUnionControlNetType:
     @classmethod