Add error message on union controlnet (#4081)

This commit is contained in:
Chenlei Hu 2024-07-22 11:27:32 -04:00 committed by GitHub
parent 6045ed31f8
commit 4151fbfa8a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 25 additions and 11 deletions

View File

@ -13,6 +13,7 @@ from ..ldm.modules.diffusionmodules.util import (
from ..ldm.modules.attention import SpatialTransformer from ..ldm.modules.attention import SpatialTransformer
from ..ldm.modules.diffusionmodules.openaimodel import UNetModel, TimestepEmbedSequential, ResBlock, Downsample from ..ldm.modules.diffusionmodules.openaimodel import UNetModel, TimestepEmbedSequential, ResBlock, Downsample
from ..ldm.util import exists from ..ldm.util import exists
from .control_types import UNION_CONTROLNET_TYPES
from collections import OrderedDict from collections import OrderedDict
import comfy.ops import comfy.ops
from comfy.ldm.modules.attention import optimized_attention 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 if self.control_add_embedding is not None: #Union Controlnet
control_type = kwargs.get("control_type", []) 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) emb += self.control_add_embedding(control_type, emb.dtype, emb.device)
if len(control_type) > 0: if len(control_type) > 0:
if len(hint.shape) < 5: if len(hint.shape) < 5:

View File

@ -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,
}

View File

@ -1,14 +1,4 @@
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: class SetUnionControlNetType:
@classmethod @classmethod