Merge branch 'comfyanonymous:master' into master

comfy/cldm/cldm.py

@@ -13,6 +13,7 @@ from ..ldm.modules.diffusionmodules.util import (
 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 @@ class ControlNet(nn.Module):
         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:

comfy/cldm/control_types.py

@@ -0,0 +1,10 @@
+UNION_CONTROLNET_TYPES = {
+    "openpose": 0,
+    "depth": 1,
+    "hed/pidi/scribble/ted": 2,
+    "canny/lineart/anime_lineart/mlsd": 3,
+    "normal": 4,
+    "segment": 5,
+    "tile": 6,
+    "repaint": 7,
+}

comfy/text_encoders/aura_t5.py

@@ -1,5 +1,5 @@
 from comfy import sd1_clip
-from .llama_tokenizer import LLAMATokenizer
+from .spiece_tokenizer import SPieceTokenizer
 import comfy.text_encoders.t5
 import os
 
@@ -11,7 +11,7 @@ class PT5XlModel(sd1_clip.SDClipModel):
 class PT5XlTokenizer(sd1_clip.SDTokenizer):
     def __init__(self, embedding_directory=None):
         tokenizer_path = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "t5_pile_tokenizer"), "tokenizer.model")
-        super().__init__(tokenizer_path, pad_with_end=False, embedding_size=2048, embedding_key='pile_t5xl', tokenizer_class=LLAMATokenizer, has_start_token=False, pad_to_max_length=False, max_length=99999999, min_length=256, pad_token=1)
+        super().__init__(tokenizer_path, pad_with_end=False, embedding_size=2048, embedding_key='pile_t5xl', tokenizer_class=SPieceTokenizer, has_start_token=False, pad_to_max_length=False, max_length=99999999, min_length=256, pad_token=1)
 
 class AuraT5Tokenizer(sd1_clip.SD1Tokenizer):
     def __init__(self, embedding_directory=None):

comfy/text_encoders/spiece_tokenizer.py

@@ -1,9 +1,9 @@
 import os
 
-class LLAMATokenizer:
+class SPieceTokenizer:
     @staticmethod
     def from_pretrained(path):
-        return LLAMATokenizer(path)
+        return SPieceTokenizer(path)
 
     def __init__(self, tokenizer_path):
         import sentencepiece

comfy_extras/nodes_controlnet.py

@@ -1,20 +1,10 @@
-
+from comfy.cldm.control_types import UNION_CONTROLNET_TYPES
-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()),)
+                             "type": (["auto"] + list(UNION_CONTROLNET_TYPES.keys()),)
                              }}
 
     CATEGORY = "conditioning/controlnet"
@@ -24,7 +14,7 @@ class SetUnionControlNetType:
 
     def set_controlnet_type(self, control_net, type):
         control_net = control_net.copy()
-        type_number = UNION_CONTROLNET_TYPES[type]
+        type_number = UNION_CONTROLNET_TYPES.get(type, -1)
         if type_number >= 0:
             control_net.set_extra_arg("control_type", [type_number])
         else: