# Code ported and modified from the diffusers ControlNetPlus repo by Qi Xin:
# https://github.com/xinsir6/ControlNetPlus/blob/main/models/controlnet_union.py
from typing import Union
import os
import torch
import torch as th
import torch.nn as nn
from torch import Tensor
from collections import OrderedDict
from comfy.ldm.modules.diffusionmodules.util import (zero_module, timestep_embedding)
from comfy.cldm.cldm import ControlNet as ControlNetCLDM
import comfy.cldm.cldm
from comfy.controlnet import ControlNet
#from comfy.t2i_adapter.adapter import ResidualAttentionBlock
from comfy.ldm.modules.attention import optimized_attention
import comfy.ops
import comfy.model_management
import comfy.model_detection
import comfy.utils
from .utils import (AdvancedControlBase, ControlWeights, ControlWeightType, TimestepKeyframeGroup, AbstractPreprocWrapper,
extend_to_batch_size, broadcast_image_to_extend)
from .logger import logger
class PlusPlusType:
OPENPOSE = "openpose"
DEPTH = "depth"
THICKLINE = "hed/pidi/scribble/ted"
THINLINE = "canny/lineart/mlsd"
NORMAL = "normal"
SEGMENT = "segment"
TILE = "tile"
REPAINT = "inpaint/outpaint"
NONE = "none"
_LIST_WITH_NONE = [OPENPOSE, DEPTH, THICKLINE, THINLINE, NORMAL, SEGMENT, TILE, REPAINT, NONE]
_LIST = [OPENPOSE, DEPTH, THICKLINE, THINLINE, NORMAL, SEGMENT, TILE, REPAINT]
_DICT = {OPENPOSE: 0, DEPTH: 1, THICKLINE: 2, THINLINE: 3, NORMAL: 4, SEGMENT: 5, TILE: 6, REPAINT: 7, NONE: -1}
@classmethod
def to_idx(cls, control_type: str):
try:
return cls._DICT[control_type]
except KeyError:
raise Exception(f"Unknown control type '{control_type}'.")
class PlusPlusInput:
def __init__(self, image: Tensor, control_type: str, strength: float):
self.image = image
self.control_type = control_type
self.strength = strength
def clone(self):
return PlusPlusInput(self.image, self.control_type, self.strength)
class PlusPlusInputGroup:
def __init__(self):
self.controls: dict[str, PlusPlusInput] = {}
def add(self, pp_input: PlusPlusInput):
if pp_input.control_type in self.controls:
raise Exception(f"Control type '{pp_input.control_type}' is already present; ControlNet++ does not allow more than 1 of each type.")
self.controls[pp_input.control_type] = pp_input
def clone(self) -> 'PlusPlusInputGroup':
cloned = PlusPlusInputGroup()
for key, value in self.controls.items():
cloned.controls[key] = value.clone()
return cloned
class PlusPlusImageWrapper(AbstractPreprocWrapper):
error_msg = error_msg = "Invalid use of ControlNet++ Image Wrapper. The output of ControlNet++ Image Wrapper is NOT a usual image, but an object holding the images and extra info - you must connect the output directly to an Apply Advanced ControlNet node. It cannot be used for anything else that accepts IMAGE input."
def __init__(self, condhint: PlusPlusInputGroup):
super().__init__(condhint)
# just an IDE type hint
self.condhint: PlusPlusInputGroup
def movedim(self, source: int, destination: int):
condhint = self.condhint.clone()
for pp_input in condhint.controls.values():
pp_input.image = pp_input.image.movedim(source, destination)
return PlusPlusImageWrapper(condhint)
# parts taken from comfy/cldm/cldm.py
class OptimizedAttention(nn.Module):
def __init__(self, c, nhead, dropout=0.0, dtype=None, device=None, operations=None):
super().__init__()
self.heads = nhead
self.c = c
self.in_proj = operations.Linear(c, c * 3, bias=True, dtype=dtype, device=device)
self.out_proj = operations.Linear(c, c, bias=True, dtype=dtype, device=device)
def forward(self, x):
x = self.in_proj(x)
q, k, v = x.split(self.c, dim=2)
out = optimized_attention(q, k, v, self.heads)
return self.out_proj(out)
class QuickGELU(nn.Module):
def forward(self, x: torch.Tensor):
return x * torch.sigmoid(1.702 * x)
class ResBlockUnionControlnet(nn.Module):
def __init__(self, dim, nhead, dtype=None, device=None, operations=None):
super().__init__()
self.attn = OptimizedAttention(dim, nhead, dtype=dtype, device=device, operations=operations)
self.ln_1 = operations.LayerNorm(dim, dtype=dtype, device=device)
self.mlp = nn.Sequential(
OrderedDict([("c_fc", operations.Linear(dim, dim * 4, dtype=dtype, device=device)), ("gelu", QuickGELU()),
("c_proj", operations.Linear(dim * 4, dim, dtype=dtype, device=device))]))
self.ln_2 = operations.LayerNorm(dim, dtype=dtype, device=device)
def attention(self, x: torch.Tensor):
return self.attn(x)
def forward(self, x: torch.Tensor):
x = x + self.attention(self.ln_1(x))
x = x + self.mlp(self.ln_2(x))
return x
class ControlAddEmbeddingAdv(nn.Module):
def __init__(self, in_dim, out_dim, num_control_type, dtype=None, device=None, operations: comfy.ops.disable_weight_init=None):
super().__init__()
self.num_control_type = num_control_type
self.in_dim = in_dim
self.linear_1 = operations.Linear(in_dim * num_control_type, out_dim, dtype=dtype, device=device)
self.linear_2 = operations.Linear(out_dim, out_dim, dtype=dtype, device=device)
def forward(self, control_type, dtype, device):
if control_type is None:
control_type = torch.zeros((self.num_control_type,), device=device)
c_type = timestep_embedding(control_type.flatten(), self.in_dim, repeat_only=False).to(dtype).reshape((-1, self.num_control_type * self.in_dim))
return self.linear_2(torch.nn.functional.silu(self.linear_1(c_type)))
class ControlNetPlusPlus(ControlNetCLDM):
def __init__(self, *args,**kwargs):
super().__init__(*args, **kwargs)
operations: comfy.ops.disable_weight_init = kwargs.get("operations", comfy.ops.disable_weight_init)
device = kwargs.get("device", None)
time_embed_dim = self.model_channels * 4
control_add_embed_dim = 256
self.control_add_embedding = ControlAddEmbeddingAdv(control_add_embed_dim, time_embed_dim, self.num_control_type, dtype=self.dtype, device=device, operations=operations)
def union_controlnet_merge(self, hint: list[Tensor], control_type, emb, context):
# Equivalent to: https://github.com/xinsir6/ControlNetPlus/tree/main
indexes = torch.nonzero(control_type[0])
inputs = []
condition_list = []
for idx in range(indexes.shape[0]):
controlnet_cond = self.input_hint_block(hint[indexes[idx][0]], emb, context)
feat_seq = torch.mean(controlnet_cond, dim=(2, 3))
if idx < indexes.shape[0]:
feat_seq += self.task_embedding[indexes[idx][0]].to(dtype=feat_seq.dtype, device=feat_seq.device)
inputs.append(feat_seq.unsqueeze(1))
condition_list.append(controlnet_cond)
x = torch.cat(inputs, dim=1)
x = self.transformer_layes(x)
controlnet_cond_fuser = None
for idx in range(indexes.shape[0]):
alpha = self.spatial_ch_projs(x[:, idx])
alpha = alpha.unsqueeze(-1).unsqueeze(-1)
o = condition_list[idx] + alpha
if controlnet_cond_fuser is None:
controlnet_cond_fuser = o
else:
controlnet_cond_fuser += o
return controlnet_cond_fuser
def forward(self, x: Tensor, hint: list[Tensor], timesteps, context, y: Tensor=None, **kwargs):
t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
emb = self.time_embed(t_emb)
guided_hint = None
if self.control_add_embedding is not None:
control_type = kwargs.get("control_type", None)
emb += self.control_add_embedding(control_type, emb.dtype, emb.device)
if control_type is not None:
guided_hint = self.union_controlnet_merge(hint, control_type, emb, context)
if guided_hint is None:
guided_hint = self.input_hint_block(hint[0], emb, context)
out_output = []
out_middle = []
hs = []
if self.num_classes is not None:
assert y.shape[0] == x.shape[0]
emb = emb + self.label_emb(y)
h = x
for module, zero_conv in zip(self.input_blocks, self.zero_convs):
if guided_hint is not None:
h = module(h, emb, context)
h += guided_hint
guided_hint = None
else:
h = module(h, emb, context)
out_output.append(zero_conv(h, emb, context))
h = self.middle_block(h, emb, context)
out_middle.append(self.middle_block_out(h, emb, context))
return {"middle": out_middle, "output": out_output}
class ControlNetPlusPlusAdvanced(ControlNet, AdvancedControlBase):
def __init__(self, control_model: ControlNetPlusPlus, timestep_keyframes: TimestepKeyframeGroup, global_average_pooling=False, load_device=None, manual_cast_dtype=None):
super().__init__(control_model=control_model, global_average_pooling=global_average_pooling, load_device=load_device, manual_cast_dtype=manual_cast_dtype)
AdvancedControlBase.__init__(self, super(), timestep_keyframes=timestep_keyframes, weights_default=ControlWeights.controlnet())
self.add_compatible_weight(ControlWeightType.CONTROLNETPLUSPLUS)
# for IDE type hint purposes
self.control_model: ControlNetPlusPlus
self.cond_hint_original: Union[PlusPlusImageWrapper, PlusPlusInputGroup]
self.cond_hint: list[Union[Tensor, None]]
self.cond_hint_shape: Tensor = None
self.cond_hint_types: Tensor = None
# in case it is using the single loader
self.single_control_type: str = None
def get_universal_weights(self) -> ControlWeights:
def cn_weights_func(idx: int, control: dict[str, list[Tensor]], key: str):
if key == "middle":
return 1.0
c_len = len(control[key])
raw_weights = [(self.weights.base_multiplier ** float((c_len) - i)) for i in range(c_len+1)]
raw_weights = raw_weights[:-1]
if key == "input":
raw_weights.reverse()
return raw_weights[idx]
return self.weights.copy_with_new_weights(new_weight_func=cn_weights_func)
def verify_control_type(self, model_name: str, pp_group: PlusPlusInputGroup=None):
if pp_group is not None:
for pp_input in pp_group.controls.values():
if PlusPlusType.to_idx(pp_input.control_type) >= self.control_model.num_control_type:
raise Exception(f"ControlNet++ model '{model_name}' does not support control_type '{pp_input.control_type}'.")
if self.single_control_type is not None:
if PlusPlusType.to_idx(self.single_control_type) >= self.control_model.num_control_type:
raise Exception(f"ControlNet++ model '{model_name}' does not support control_type '{self.single_control_type}'.")
def set_cond_hint_inject(self, *args, **kwargs):
to_return = super().set_cond_hint_inject(*args, **kwargs)
# if not single_control_type, expect PlusPlusImageWrapper
if self.single_control_type is None:
# check that cond_hint is wrapped, and unwrap it
if type(self.cond_hint_original) != PlusPlusImageWrapper:
raise Exception("ControlNet++ (Multi) expects image input from the Load ControlNet++ Model node, NOT from anything else. Images are provided to that node via ControlNet++ Input nodes.")
self.cond_hint_original = self.cond_hint_original.condhint.clone()
# otherwise, expect single image input (AKA, usual controlnet input)
else:
# check that cond_hint is not a PlusPlusImageWrapper
if type(self.cond_hint_original) == PlusPlusImageWrapper:
raise Exception("ControlNet++ (Single) expects usual image input, NOT the image input from a Load ControlNet++ Model (Multi) node.")
pp_group = PlusPlusInputGroup()
pp_input = PlusPlusInput(self.cond_hint_original, self.single_control_type, 1.0)
pp_group.add(pp_input)
self.cond_hint_original = pp_group
return to_return
def get_control_advanced(self, x_noisy: Tensor, t, cond, batched_number):
control_prev = None
if self.previous_controlnet is not None:
control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number)
if self.timestep_range is not None:
if t[0] > self.timestep_range[0] or t[0] < self.timestep_range[1]:
if control_prev is not None:
return control_prev
else:
return None
dtype = self.control_model.dtype
if self.manual_cast_dtype is not None:
dtype = self.manual_cast_dtype
output_dtype = x_noisy.dtype
# make all cond_hints appropriate dimensions
# TODO: change this to not require cond_hint upscaling every step when self.sub_idxs is present
if self.sub_idxs is not None or self.cond_hint is None or x_noisy.shape[2] * self.compression_ratio != self.cond_hint_shape[2] or x_noisy.shape[3] * self.compression_ratio != self.cond_hint_shape[3]:
if self.cond_hint is not None:
del self.cond_hint
self.cond_hint = [None] * self.control_model.num_control_type
self.cond_hint_types = torch.tensor([0.0] * self.control_model.num_control_type)
self.cond_hint_shape = None
compression_ratio = self.compression_ratio
# unlike normal controlnet, need to handle each input image tensor (for each type)
for pp_type, pp_input in self.cond_hint_original.controls.items():
pp_idx = PlusPlusType.to_idx(pp_type)
# if negative, means no type should be selected (single only)
if pp_idx < 0:
pp_idx = 0
else:
self.cond_hint_types[pp_idx] = pp_input.strength
# if self.cond_hint_original lengths greater or equal to latent count, subdivide
if self.sub_idxs is not None:
actual_cond_hint_orig = pp_input.image
if pp_input.image.size(0) < self.full_latent_length:
actual_cond_hint_orig = extend_to_batch_size(tensor=actual_cond_hint_orig, batch_size=self.full_latent_length)
self.cond_hint[pp_idx] = comfy.utils.common_upscale(actual_cond_hint_orig[self.sub_idxs], x_noisy.shape[3] * compression_ratio, x_noisy.shape[2] * compression_ratio, 'nearest-exact', "center")
else:
self.cond_hint[pp_idx] = comfy.utils.common_upscale(pp_input.image, x_noisy.shape[3] * compression_ratio, x_noisy.shape[2] * compression_ratio, 'nearest-exact', "center")
self.cond_hint[pp_idx] = self.cond_hint[pp_idx].to(device=x_noisy.device, dtype=dtype)
self.cond_hint_shape = self.cond_hint[pp_idx].shape
# prepare cond_hint_controls to match batchsize
if self.cond_hint_types.count_nonzero() == 0:
self.cond_hint_types = None
else:
self.cond_hint_types = self.cond_hint_types.unsqueeze(0).to(device=x_noisy.device, dtype=dtype).repeat(x_noisy.shape[0], 1)
for i in range(len(self.cond_hint)):
if self.cond_hint[i] is not None:
if x_noisy.shape[0] != self.cond_hint[i].shape[0]:
self.cond_hint[i] = broadcast_image_to_extend(self.cond_hint[i], x_noisy.shape[0], batched_number)
if self.cond_hint_types is not None and x_noisy.shape[0] != self.cond_hint_types.shape[0]:
self.cond_hint_types = broadcast_image_to_extend(self.cond_hint_types, x_noisy.shape[0], batched_number, False)
# prepare mask_cond_hint
self.prepare_mask_cond_hint(x_noisy=x_noisy, t=t, cond=cond, batched_number=batched_number, dtype=dtype)
context = cond.get('crossattn_controlnet', cond['c_crossattn'])
y = cond.get('y', None)
if y is not None:
y = y.to(dtype)
timestep = self.model_sampling_current.timestep(t)
x_noisy = self.model_sampling_current.calculate_input(t, x_noisy)
control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.float(), context=context.to(dtype), y=y, control_type=self.cond_hint_types)
return self.control_merge(control, control_prev, output_dtype)
def copy(self):
c = ControlNetPlusPlusAdvanced(self.control_model, self.timestep_keyframes, global_average_pooling=self.global_average_pooling, load_device=self.load_device, manual_cast_dtype=self.manual_cast_dtype)
self.copy_to(c)
self.copy_to_advanced(c)
c.single_control_type = self.single_control_type
return c
def load_controlnetplusplus(ckpt_path: str, timestep_keyframe: TimestepKeyframeGroup=None, model=None):
controlnet_data = comfy.utils.load_torch_file(ckpt_path, safe_load=True)
# check that actually is ControlNet++ model
if "task_embedding" not in controlnet_data:
raise Exception(f"'{ckpt_path}' is not a valid ControlNet++ model.")
controlnet_config = None
supported_inference_dtypes = None
if "controlnet_cond_embedding.conv_in.weight" in controlnet_data: #diffusers format
controlnet_config = comfy.model_detection.unet_config_from_diffusers_unet(controlnet_data)
diffusers_keys = comfy.utils.unet_to_diffusers(controlnet_config)
diffusers_keys["controlnet_mid_block.weight"] = "middle_block_out.0.weight"
diffusers_keys["controlnet_mid_block.bias"] = "middle_block_out.0.bias"
count = 0
loop = True
while loop:
suffix = [".weight", ".bias"]
for s in suffix:
k_in = "controlnet_down_blocks.{}{}".format(count, s)
k_out = "zero_convs.{}.0{}".format(count, s)
if k_in not in controlnet_data:
loop = False
break
diffusers_keys[k_in] = k_out
count += 1
count = 0
loop = True
while loop:
suffix = [".weight", ".bias"]
for s in suffix:
if count == 0:
k_in = "controlnet_cond_embedding.conv_in{}".format(s)
else:
k_in = "controlnet_cond_embedding.blocks.{}{}".format(count - 1, s)
k_out = "input_hint_block.{}{}".format(count * 2, s)
if k_in not in controlnet_data:
k_in = "controlnet_cond_embedding.conv_out{}".format(s)
loop = False
diffusers_keys[k_in] = k_out
count += 1
new_sd = {}
for k in diffusers_keys:
if k in controlnet_data:
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_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)
leftover_keys = controlnet_data.keys()
if len(leftover_keys) > 0:
logger.warning("leftover ControlNet++ keys: {}".format(leftover_keys))
controlnet_data = new_sd
elif "controlnet_blocks.0.weight" in controlnet_data: #SD3 diffusers format
raise Exception("Unexpected SD3 diffusers format for ControlNet++ model. Something is very wrong.")
pth_key = 'control_model.zero_convs.0.0.weight'
pth = False
key = 'zero_convs.0.0.weight'
if pth_key in controlnet_data:
pth = True
key = pth_key
prefix = "control_model."
elif key in controlnet_data:
prefix = ""
else:
raise Exception("Unexpected T2IAdapter format for ControlNet++ model. Something is very wrong.")
if controlnet_config is None:
model_config = comfy.model_detection.model_config_from_unet(controlnet_data, prefix, True)
supported_inference_dtypes = model_config.supported_inference_dtypes
controlnet_config = model_config.unet_config
load_device = comfy.model_management.get_torch_device()
if supported_inference_dtypes is None:
unet_dtype = comfy.model_management.unet_dtype()
else:
unet_dtype = comfy.model_management.unet_dtype(supported_dtypes=supported_inference_dtypes)
manual_cast_dtype = comfy.model_management.unet_manual_cast(unet_dtype, load_device)
if manual_cast_dtype is not None:
controlnet_config["operations"] = comfy.ops.manual_cast
controlnet_config["dtype"] = unet_dtype
controlnet_config.pop("out_channels")
controlnet_config["hint_channels"] = controlnet_data["{}input_hint_block.0.weight".format(prefix)].shape[1]
control_model = ControlNetPlusPlus(**controlnet_config)
if pth:
if 'difference' in controlnet_data:
if model is not None:
comfy.model_management.load_models_gpu([model])
model_sd = model.model_state_dict()
for x in controlnet_data:
c_m = "control_model."
if x.startswith(c_m):
sd_key = "diffusion_model.{}".format(x[len(c_m):])
if sd_key in model_sd:
cd = controlnet_data[x]
cd += model_sd[sd_key].type(cd.dtype).to(cd.device)
else:
logger.warning("WARNING: Loaded a diff controlnet without a model. It will very likely not work.")
class WeightsLoader(torch.nn.Module):
pass
w = WeightsLoader()
w.control_model = control_model
missing, unexpected = w.load_state_dict(controlnet_data, strict=False)
else:
missing, unexpected = control_model.load_state_dict(controlnet_data, strict=False)
if len(missing) > 0:
logger.warning("missing ControlNet++ keys: {}".format(missing))
if len(unexpected) > 0:
logger.debug("unexpected ControlNet++ keys: {}".format(unexpected))
global_average_pooling = False
filename = os.path.splitext(ckpt_path)[0]
if filename.endswith("_shuffle") or filename.endswith("_shuffle_fp16"): #TODO: smarter way of enabling global_average_pooling
global_average_pooling = True
control = ControlNetPlusPlusAdvanced(control_model, timestep_keyframes=timestep_keyframe, global_average_pooling=global_average_pooling, load_device=load_device, manual_cast_dtype=manual_cast_dtype)
return control