From edc0a2081503175f011649a19e142248a0829299 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Sun, 20 Aug 2023 00:41:35 -0500 Subject: [PATCH] updated control.py to account for most recent changes in base ComfyUI ControlNet code --- control.py | 25 ++++++++++++++----------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/control.py b/control.py index b4e2cff..a67687d 100644 --- a/control.py +++ b/control.py @@ -13,7 +13,7 @@ sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "co from comfy.cldm import cldm from comfy.t2i_adapter import adapter -from comfy.sd import ControlBase, broadcast_image_to +from comfy.sd import ControlBase, ModelPatcher, broadcast_image_to, ControlLora import comfy.utils as utils import comfy.model_management as model_management import comfy.model_detection as model_detection @@ -28,6 +28,7 @@ class ControlNetAdvanced(ControlBase): def __init__(self, control_model, weights: ControlNetWeightsType, global_average_pooling=False, device=None): super().__init__(device) self.control_model = control_model + self.control_model_wrapped = ModelPatcher(self.control_model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device()) self.weights = weights if weights else [1.0]*13 self.global_average_pooling = global_average_pooling @@ -58,11 +59,9 @@ class ControlNetAdvanced(ControlBase): precision_scope = contextlib.nullcontext with precision_scope(model_management.get_autocast_device(self.device)): - self.control_model = model_management.load_if_low_vram(self.control_model) context = torch.cat(cond['c_crossattn'], 1) y = cond.get('c_adm', None) control = self.control_model(x=x_noisy, hint=self.cond_hint, timesteps=t, context=context, y=y) - self.control_model = model_management.unload_if_low_vram(self.control_model) out = {'middle':[], 'output': []} autocast_enabled = torch.is_autocast_enabled() @@ -77,7 +76,7 @@ class ControlNetAdvanced(ControlBase): if self.global_average_pooling: x = torch.mean(x, dim=(2, 3), keepdim=True).repeat(1, 1, x.shape[2], x.shape[3]) - x *= self.strength * self.weights[i] # apply layer weight + x *= self.strength * self.weights[i] if x.dtype != output_dtype and not autocast_enabled: x = x.to(output_dtype) @@ -97,17 +96,19 @@ class ControlNetAdvanced(ControlBase): def get_models(self): out = super().get_models() - out.append(self.control_model) + out.append(self.control_model_wrapped) return out def load_controlnet(ckpt_path, control_net_weights: ControlNetWeightsType=None, t2i_adapter_weights: T2IAdapterWeightsType=None, model=None): controlnet_data = utils.load_torch_file(ckpt_path, safe_load=True) + if "lora_controlnet" in controlnet_data: + return ControlLora(controlnet_data) # TODO: apply weights to ControlLora controlnet_config = None if "controlnet_cond_embedding.conv_in.weight" in controlnet_data: #diffusers format use_fp16 = model_management.should_use_fp16() - controlnet_config = model_detection.model_config_from_diffusers_unet(controlnet_data, use_fp16).unet_config + controlnet_config = model_detection.unet_config_from_diffusers_unet(controlnet_data, use_fp16) diffusers_keys = 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" @@ -146,6 +147,9 @@ def load_controlnet(ckpt_path, control_net_weights: ControlNetWeightsType=None, if k in controlnet_data: new_sd[diffusers_keys[k]] = controlnet_data.pop(k) + leftover_keys = controlnet_data.keys() + if len(leftover_keys) > 0: + print("leftover keys:", leftover_keys) controlnet_data = new_sd pth_key = 'control_model.zero_convs.0.0.weight' @@ -158,7 +162,7 @@ def load_controlnet(ckpt_path, control_net_weights: ControlNetWeightsType=None, elif key in controlnet_data: prefix = "" else: - net = load_t2i_adapter(controlnet_data, t2i_adapter_weights) + net = load_t2i_adapter(controlnet_data) if net is None: print("error checkpoint does not contain controlnet or t2i adapter data", ckpt_path) return net @@ -167,14 +171,14 @@ def load_controlnet(ckpt_path, control_net_weights: ControlNetWeightsType=None, use_fp16 = model_management.should_use_fp16() controlnet_config = model_detection.model_config_from_unet(controlnet_data, prefix, use_fp16).unet_config controlnet_config.pop("out_channels") - controlnet_config["hint_channels"] = 3 + controlnet_config["hint_channels"] = controlnet_data["{}input_hint_block.0.weight".format(prefix)].shape[1] control_model = cldm.ControlNet(**controlnet_config) if pth: if 'difference' in controlnet_data: if model is not None: - m = model.patch_model() - model_sd = m.state_dict() + 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): @@ -182,7 +186,6 @@ def load_controlnet(ckpt_path, control_net_weights: ControlNetWeightsType=None, if sd_key in model_sd: cd = controlnet_data[x] cd += model_sd[sd_key].type(cd.dtype).to(cd.device) - model.unpatch_model() else: print("WARNING: Loaded a diff controlnet without a model. It will very likely not work.")