diff --git a/comfy/cldm/cldm.py b/comfy/cldm/cldm.py index 1d7294bd6..9ec64a227 100644 --- a/comfy/cldm/cldm.py +++ b/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: diff --git a/comfy/cldm/control_types.py b/comfy/cldm/control_types.py new file mode 100644 index 000000000..4128631a3 --- /dev/null +++ b/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, +} diff --git a/comfy/text_encoders/aura_t5.py b/comfy/text_encoders/aura_t5.py index 6b9e4fe53..409867af3 100644 --- a/comfy/text_encoders/aura_t5.py +++ b/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): diff --git a/comfy/text_encoders/llama_tokenizer.py b/comfy/text_encoders/spiece_tokenizer.py similarity index 90% rename from comfy/text_encoders/llama_tokenizer.py rename to comfy/text_encoders/spiece_tokenizer.py index a6db1da62..d611d5bb7 100644 --- a/comfy/text_encoders/llama_tokenizer.py +++ b/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 diff --git a/comfy_extras/nodes_controlnet.py b/comfy_extras/nodes_controlnet.py index ef7cfc6ab..0773c8a5b 100644 --- a/comfy_extras/nodes_controlnet.py +++ b/comfy_extras/nodes_controlnet.py @@ -1,20 +1,10 @@ - -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 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: