Add error message on union controlnet (#4081)

comfyanonymous · Jul 22, 2024 · 4151fbf · 4151fbf
1 parent 6045ed3
commit 4151fbf
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 11 deletions.
diff --git 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
@@ -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
@@ -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