From 4151fbfa8ab705d0a7e632c7d7c23ea5c436a60b Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Mon, 22 Jul 2024 11:27:32 -0400 Subject: [PATCH 1/3] 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 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..fd6e56e8f --- /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 ef7cfc6ab..566d51e33 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 From b2c995f623fd7169db87dfd6d356e535d0bb99ce Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 22 Jul 2024 11:30:38 -0400 Subject: [PATCH 2/3] "auto" type is only relevant to the SetUnionControlNetType node. --- comfy/cldm/control_types.py | 1 - comfy_extras/nodes_controlnet.py | 4 ++-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/comfy/cldm/control_types.py b/comfy/cldm/control_types.py index fd6e56e8f..4128631a3 100644 --- a/comfy/cldm/control_types.py +++ b/comfy/cldm/control_types.py @@ -1,5 +1,4 @@ UNION_CONTROLNET_TYPES = { - "auto": -1, "openpose": 0, "depth": 1, "hed/pidi/scribble/ted": 2, diff --git a/comfy_extras/nodes_controlnet.py b/comfy_extras/nodes_controlnet.py index 566d51e33..0773c8a5b 100644 --- a/comfy_extras/nodes_controlnet.py +++ b/comfy_extras/nodes_controlnet.py @@ -4,7 +4,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" @@ -14,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: From 14764aa2e2e2b282c4a4dffbfab4c01d3e46e8a7 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 22 Jul 2024 12:21:45 -0400 Subject: [PATCH 3/3] Rename LLAMATokenizer to SPieceTokenizer. --- comfy/text_encoders/aura_t5.py | 4 ++-- .../text_encoders/{llama_tokenizer.py => spiece_tokenizer.py} | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) rename comfy/text_encoders/{llama_tokenizer.py => spiece_tokenizer.py} (90%) 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