From 72f0e6dbac936be2cd5e87a2ba8270f14f600723 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Mon, 25 Nov 2024 20:33:53 -0600 Subject: [PATCH] Add initial CtrLoRA support, fix issue with vae not being set, add a way to disable multiplying compression_ratio by vae.downscale_ratio when not needed --- adv_control/control.py | 8 +- adv_control/control_ctrlora.py | 228 +++++++++++++++++++++++++++++++++ adv_control/nodes.py | 5 + adv_control/nodes_ctrlora.py | 25 ++++ adv_control/nodes_main.py | 2 +- adv_control/utils.py | 5 + 6 files changed, 271 insertions(+), 2 deletions(-) create mode 100644 adv_control/control_ctrlora.py create mode 100644 adv_control/nodes_ctrlora.py diff --git a/adv_control/control.py b/adv_control/control.py index ee83072..e31f0e5 100644 --- a/adv_control/control.py +++ b/adv_control/control.py @@ -66,7 +66,7 @@ class ControlNetAdvanced(ControlNet, AdvancedControlBase): del self.cond_hint self.cond_hint = None compression_ratio = self.compression_ratio - if self.vae is not None: + if self.vae is not None and self.mult_by_ratio_when_vae: compression_ratio *= self.vae.downscale_ratio # if self.cond_hint_original length greater or equal to real latent count, subdivide it before scaling if self.sub_idxs is not None: @@ -483,6 +483,10 @@ def load_controlnet(ckpt_path, timestep_keyframe: TimestepKeyframeGroup=None, mo # ControlNet++ check elif "task_embedding" in key: pass + # CtrLoRA check + elif "lora_layer" in key: + controlnet_type = ControlWeightType.CTRLORA + break if has_controlnet_key and has_motion_modules_key: controlnet_type = ControlWeightType.SPARSECTRL @@ -496,6 +500,8 @@ def load_controlnet(ckpt_path, timestep_keyframe: TimestepKeyframeGroup=None, mo control = load_sparsectrl(ckpt_path, controlnet_data=controlnet_data, timestep_keyframe=timestep_keyframe, model=model) elif controlnet_type == ControlWeightType.SVD_CONTROLNET: control = load_svdcontrolnet(ckpt_path, controlnet_data=controlnet_data, timestep_keyframe=timestep_keyframe) + elif controlnet_type == ControlWeightType.CTRLORA: + raise Exception("This is a CtrLoRA; use the Load CtrLoRA Model node.") # otherwise, load vanilla ControlNet else: try: diff --git a/adv_control/control_ctrlora.py b/adv_control/control_ctrlora.py new file mode 100644 index 0000000..03e192f --- /dev/null +++ b/adv_control/control_ctrlora.py @@ -0,0 +1,228 @@ +import torch +from torch import Tensor + +from comfy.cldm.cldm import ControlNet as ControlNetCLDM +import comfy.model_detection +import comfy.model_management +import comfy.ops +import comfy.utils + +from comfy.ldm.modules.diffusionmodules.util import ( + zero_module, + timestep_embedding, +) + +from .control import ControlNetAdvanced +from .utils import TimestepKeyframeGroup +from .logger import logger + + +class ControlNetCtrLoRA(ControlNetCLDM): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # delete input hint block + del self.input_hint_block + + def forward(self, x: Tensor, hint: Tensor, timesteps, context, y=None, **kwargs): + t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype) + emb = self.time_embed(t_emb) + + out_output = [] + out_middle = [] + + if self.num_classes is not None: + assert y.shape[0] == x.shape[0] + emb = emb + self.label_emb(y) + + h = hint.to(dtype=x.dtype) + for module, zero_conv in zip(self.input_blocks, self.zero_convs): + 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 CtrLoRAAdvanced(ControlNetAdvanced): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.require_vae = True + self.mult_by_ratio_when_vae = False + + def pre_run_advanced(self, model, percent_to_timestep_function): + super().pre_run_advanced(model, percent_to_timestep_function) + self.latent_format = model.latent_format # LatentFormat object, used to process_in latent cond hint + + def cleanup_advanced(self): + super().cleanup_advanced() + if self.latent_format is not None: + del self.latent_format + self.latent_format = None + + def copy(self): + c = CtrLoRAAdvanced(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) + c.control_model = self.control_model + c.control_model_wrapped = self.control_model_wrapped + self.copy_to(c) + self.copy_to_advanced(c) + return c + + +def load_ctrlora(base_path: str, lora_path: str, + base_data: dict[str, Tensor]=None, lora_data: dict[str, Tensor]=None, + timestep_keyframe: TimestepKeyframeGroup=None, model=None, model_options={}): + if base_data is None: + base_data = comfy.utils.load_torch_file(base_path, safe_load=True) + controlnet_data = base_data + + # first, check that base_data contains keys with lora_layer + contains_lora_layers = False + for key in base_data: + if "lora_layer" in key: + contains_lora_layers = True + if not contains_lora_layers: + raise Exception(f"File '{base_path}' is not a valid CtrLoRA base model; does not contain any lora_layer keys.") + + controlnet_config = None + supported_inference_dtypes = None + + 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("") + net = load_t2i_adapter(controlnet_data, model_options=model_options) + if net is None: + logging.error("error could not detect control model type.") + return net + + if controlnet_config is None: + model_config = comfy.model_detection.model_config_from_unet(controlnet_data, prefix, True) + supported_inference_dtypes = list(model_config.supported_inference_dtypes) + controlnet_config = model_config.unet_config + + unet_dtype = model_options.get("dtype", None) + if unet_dtype is None: + weight_dtype = comfy.utils.weight_dtype(controlnet_data) + + if supported_inference_dtypes is None: + supported_inference_dtypes = [comfy.model_management.unet_dtype()] + + if weight_dtype is not None: + supported_inference_dtypes.append(weight_dtype) + + unet_dtype = comfy.model_management.unet_dtype(model_params=-1, supported_dtypes=supported_inference_dtypes) + + load_device = comfy.model_management.get_torch_device() + + manual_cast_dtype = comfy.model_management.unet_manual_cast(unet_dtype, load_device) + operations = model_options.get("custom_operations", None) + if operations is None: + operations = comfy.ops.pick_operations(unet_dtype, manual_cast_dtype) + + controlnet_config["operations"] = operations + controlnet_config["dtype"] = unet_dtype + controlnet_config["device"] = comfy.model_management.unet_offload_device() + 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 = ControlNetCtrLoRA(**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 = model_options.get("global_average_pooling", False) + control = CtrLoRAAdvanced(control_model, timestep_keyframe, global_average_pooling=global_average_pooling, + load_device=load_device, manual_cast_dtype=manual_cast_dtype) + # load lora data onto the controlnet + if lora_path is not None: + load_lora_data(control, lora_path) + + return control + + +def load_lora_data(control: CtrLoRAAdvanced, lora_path: str, loaded_data: dict[str, Tensor]=None, lora_strength=1.0): + if loaded_data is None: + loaded_data = comfy.utils.load_torch_file(lora_path, safe_load=True) + # check that lora_data contains keys with lora_layer + contains_lora_layers = False + for key in loaded_data: + if "lora_layer" in key: + contains_lora_layers = True + if not contains_lora_layers: + raise Exception(f"File '{lora_path}' is not a valid CtrLoRA lora model; does not contain any lora_layer keys.") + + # now that we know we have a ctrlora file, separate keys into 'set' and 'lora' keys + data_set: dict[str, Tensor] = {} + data_lora: dict[str, Tensor] = {} + + for key in list(loaded_data.keys()): + if 'lora_layer' in key: + data_lora[key] = loaded_data.pop(key) + else: + data_set[key] = loaded_data.pop(key) + # no keys should be left over + if len(loaded_data) > 0: + logger.warning("Not all keys from CtrlLoRA lora model's loaded data were parsed!") + + # turn set/lora data into corresponding patches; + patches = {} + # set will replace the values + for key, value in data_set.items(): + # prase model key from key; + # remove "control_model." + model_key = key.replace("control_model.", "") + patches[model_key] = ("set", (value,)) + # lora will do mm of up and down tensors + for down_key in data_lora: + # only process lora down keys; we will process both up+down at the same time + if ".up." in key: + continue + # get up version of down key + up_key = down_key.replace(".down.", ".up.") + # get key that will match up with model key; + # remove "lora_layer.down." and "control_model." + model_key = down_key.replace("lora_layer.down.", "").replace("control_model.", "") + + weight_down = data_lora[down_key] + weight_up = data_lora[up_key] + # currently, ComfyUI expects 6 elements in 'lora' type, but for future-proofing add a bunch more with None + patches[model_key] = ("lora", (weight_up, weight_down, None, None, None, None, + None, None, None, None, None, None, None, None)) + + # now that patches are made, add them to model + control.control_model_wrapped.add_patches(patches, strength_patch=lora_strength) diff --git a/adv_control/nodes.py b/adv_control/nodes.py index 32593a5..2b9d384 100644 --- a/adv_control/nodes.py +++ b/adv_control/nodes.py @@ -10,6 +10,7 @@ from .nodes_keyframes import (LatentKeyframeGroupNode, LatentKeyframeInterpolati from .nodes_sparsectrl import SparseCtrlMergedLoaderAdvanced, SparseCtrlLoaderAdvanced, SparseIndexMethodNode, SparseSpreadMethodNode, RgbSparseCtrlPreprocessor, SparseWeightExtras from .nodes_reference import ReferenceControlNetNode, ReferenceControlFinetune, ReferencePreprocessorNode from .nodes_plusplus import PlusPlusLoaderAdvanced, PlusPlusLoaderSingle, PlusPlusInputNode +from .nodes_ctrlora import CtrLoRALoader from .nodes_loosecontrol import ControlNetLoaderWithLoraAdvanced from .nodes_deprecated import (LoadImagesFromDirectory, ScaledSoftUniversalWeightsDeprecated, SoftControlNetWeightsDeprecated, CustomControlNetWeightsDeprecated, @@ -55,6 +56,8 @@ NODE_CLASS_MAPPINGS = { "ACN_ControlNet++LoaderSingle": PlusPlusLoaderSingle, "ACN_ControlNet++LoaderAdvanced": PlusPlusLoaderAdvanced, "ACN_ControlNet++InputNode": PlusPlusInputNode, + # CtrLoRA + "ACN_CtrLoRALoader": CtrLoRALoader, # Reference "ACN_ReferencePreprocessor": ReferencePreprocessorNode, "ACN_ReferenceControlNet": ReferenceControlNetNode, @@ -109,6 +112,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ACN_ControlNet++LoaderSingle": "Load ControlNet++ Model (Single) 🛂🅐🅒🅝", "ACN_ControlNet++LoaderAdvanced": "Load ControlNet++ Model (Multi) 🛂🅐🅒🅝", "ACN_ControlNet++InputNode": "ControlNet++ Input 🛂🅐🅒🅝", + # CtrLoRA + "ACN_CtrLoRALoader": "Load CtrLoRA Model 🛂🅐🅒🅝", # Reference "ACN_ReferencePreprocessor": "Reference Preproccessor 🛂🅐🅒🅝", "ACN_ReferenceControlNet": "Reference ControlNet 🛂🅐🅒🅝", diff --git a/adv_control/nodes_ctrlora.py b/adv_control/nodes_ctrlora.py new file mode 100644 index 0000000..261dde4 --- /dev/null +++ b/adv_control/nodes_ctrlora.py @@ -0,0 +1,25 @@ +import folder_paths + +from .control_ctrlora import load_ctrlora + + +class CtrLoRALoader: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "base": (folder_paths.get_filename_list("controlnet"), ), + "lora": (folder_paths.get_filename_list("controlnet"), ), + } + } + + RETURN_TYPES = ("CONTROL_NET",) + FUNCTION = "load_controlnet_plusplus" + + CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/ControlNet++" + + def load_controlnet_plusplus(self, base: str, lora: str): + base_path = folder_paths.get_full_path("controlnet", base) + lora_path = folder_paths.get_full_path("controlnet", lora) + controlnet = load_ctrlora(base_path, lora_path) + return (controlnet,) diff --git a/adv_control/nodes_main.py b/adv_control/nodes_main.py index 650d7b7..acb9ef2 100644 --- a/adv_control/nodes_main.py +++ b/adv_control/nodes_main.py @@ -139,7 +139,7 @@ class AdvancedControlNetApply: pass elif not vae_optional: # make sure SD3 ControlNet will get a special message instead of generic type mention - if is_sd3_advanced_controlnet: + if is_sd3_advanced_controlnet(c_net): raise Exception(f"SD3 ControlNet requires vae_optional input, but got None.") else: raise Exception(f"Type '{type(c_net).__name__}' requires vae_optional input, but got None.") diff --git a/adv_control/utils.py b/adv_control/utils.py index acc26d6..869cafb 100644 --- a/adv_control/utils.py +++ b/adv_control/utils.py @@ -57,6 +57,7 @@ class ControlWeightType: CONTROLLLLITE = "controllllite" SVD_CONTROLNET = "svd_controlnet" SPARSECTRL = "sparsectrl" + CTRLORA = "ctrlora" class ControlWeights: @@ -502,6 +503,7 @@ class AdvancedControlBase: self.set_cond_hint = self.set_cond_hint_inject # vae to store self.adv_vae = None + self.mult_by_ratio_when_vae = True # require model/vae to be passed into Apply Advanced ControlNet 🛂🅐🅒🅝 node self.require_vae = require_vae self.allow_condhint_latents = allow_condhint_latents @@ -615,11 +617,13 @@ class AdvancedControlBase: for arg in args: if isinstance(arg, VAE): self.adv_vae = arg + self.vae = arg break # if not in args, check kwargs now if self.adv_vae is None: if 'vae' in kwargs: self.adv_vae = kwargs['vae'] + self.vae = kwargs['vae'] return to_return def pre_run_inject(self, model, percent_to_timestep_function): @@ -876,6 +880,7 @@ class AdvancedControlBase: copied.weights_override = self.weights_override copied.latent_keyframe_override = self.latent_keyframe_override copied.adv_vae = self.adv_vae + copied.mult_by_ratio_when_vae = self.mult_by_ratio_when_vae copied.require_vae = self.require_vae copied.allow_condhint_latents = self.allow_condhint_latents copied.postpone_condhint_latents_check = self.postpone_condhint_latents_check