From 7914c47d5afb1c5ffab49d665f9a1ac86a458821 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sun, 14 Jul 2024 10:07:36 -0400 Subject: [PATCH] Quick fix for the promax controlnet. --- comfy/cldm/cldm.py | 6 +++--- comfy/controlnet.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/comfy/cldm/cldm.py b/comfy/cldm/cldm.py index d4d32b87ec9..4a58c823f5f 100644 --- a/comfy/cldm/cldm.py +++ b/comfy/cldm/cldm.py @@ -92,7 +92,7 @@ def __init__( transformer_depth_middle=None, transformer_depth_output=None, attn_precision=None, - union_controlnet=False, + union_controlnet_num_control_type=None, device=None, operations=comfy.ops.disable_weight_init, **kwargs, @@ -320,8 +320,8 @@ def __init__( self.middle_block_out = self.make_zero_conv(ch, operations=operations, dtype=self.dtype, device=device) self._feature_size += ch - if union_controlnet: - self.num_control_type = 6 + if union_controlnet_num_control_type is not None: + self.num_control_type = union_controlnet_num_control_type num_trans_channel = 320 num_trans_head = 8 num_trans_layer = 1 diff --git a/comfy/controlnet.py b/comfy/controlnet.py index 84286f1fac7..b8e27c71f8c 100644 --- a/comfy/controlnet.py +++ b/comfy/controlnet.py @@ -414,7 +414,7 @@ def load_controlnet(ckpt_path, model=None): new_sd[diffusers_keys[k]] = controlnet_data.pop(k) if "control_add_embedding.linear_1.bias" in controlnet_data: #Union Controlnet - controlnet_config["union_controlnet"] = True + controlnet_config["union_controlnet_num_control_type"] = controlnet_data["task_embedding"].shape[0] for k in list(controlnet_data.keys()): new_k = k.replace('.attn.in_proj_', '.attn.in_proj.') new_sd[new_k] = controlnet_data.pop(k)