diff --git a/__init__.py b/__init__.py index e70bf90..d30f823 100644 --- a/__init__.py +++ b/__init__.py @@ -1,3 +1,3 @@ -from .control.nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +from .adv_control.nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/control/control.py b/adv_control/control.py similarity index 100% rename from control/control.py rename to adv_control/control.py diff --git a/control/control_lllite.py b/adv_control/control_lllite.py similarity index 100% rename from control/control_lllite.py rename to adv_control/control_lllite.py diff --git a/control/control_sparsectrl.py b/adv_control/control_sparsectrl.py similarity index 100% rename from control/control_sparsectrl.py rename to adv_control/control_sparsectrl.py diff --git a/control/logger.py b/adv_control/logger.py similarity index 100% rename from control/logger.py rename to adv_control/logger.py diff --git a/control/nodes.py b/adv_control/nodes.py similarity index 97% rename from control/nodes.py rename to adv_control/nodes.py index bdd5e7e..9c6ed91 100644 --- a/control/nodes.py +++ b/adv_control/nodes.py @@ -10,6 +10,7 @@ from .nodes_weight import (DefaultWeights, ScaledSoftMaskedUniversalWeights, Sca SoftT2IAdapterWeights, CustomT2IAdapterWeights) from .nodes_latent_keyframe import LatentKeyframeGroupNode, LatentKeyframeInterpolationNode, LatentKeyframeBatchedGroupNode, LatentKeyframeNode from .nodes_sparsectrl import SparseCtrlMergedLoaderAdvanced, SparseCtrlLoaderAdvanced, SparseIndexMethodNode, SparseSpreadMethodNode, RgbSparseCtrlPreprocessor +from .nodes_loosecontrol import ControlNetLoaderWithLoraAdvanced from .nodes_deprecated import LoadImagesFromDirectory from .logger import logger @@ -221,6 +222,8 @@ NODE_CLASS_MAPPINGS = { "ACN_SparseCtrlMergedLoaderAdvanced": SparseCtrlMergedLoaderAdvanced, "ACN_SparseCtrlIndexMethodNode": SparseIndexMethodNode, "ACN_SparseCtrlSpreadMethodNode": SparseSpreadMethodNode, + # LOOSEControl + #"ACN_ControlNetLoaderWithLoraAdvanced": ControlNetLoaderWithLoraAdvanced, # Deprecated "LoadImagesFromDirectory": LoadImagesFromDirectory, } @@ -251,6 +254,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ACN_SparseCtrlMergedLoaderAdvanced": "Load Merged SparseCtrl Model 🛂🅐🅒🅝", "ACN_SparseCtrlIndexMethodNode": "SparseCtrl Index Method 🛂🅐🅒🅝", "ACN_SparseCtrlSpreadMethodNode": "SparseCtrl Spread Method 🛂🅐🅒🅝", + # LOOSEControl + #"ACN_ControlNetLoaderWithLoraAdvanced": "Load Adv. ControlNet Model w/ LoRA 🛂🅐🅒🅝", # Deprecated "LoadImagesFromDirectory": "Load Images [DEPRECATED] 🛂🅐🅒🅝", } diff --git a/control/nodes_deprecated.py b/adv_control/nodes_deprecated.py similarity index 100% rename from control/nodes_deprecated.py rename to adv_control/nodes_deprecated.py diff --git a/control/nodes_latent_keyframe.py b/adv_control/nodes_latent_keyframe.py similarity index 100% rename from control/nodes_latent_keyframe.py rename to adv_control/nodes_latent_keyframe.py diff --git a/adv_control/nodes_loosecontrol.py b/adv_control/nodes_loosecontrol.py new file mode 100644 index 0000000..5c1d43e --- /dev/null +++ b/adv_control/nodes_loosecontrol.py @@ -0,0 +1,67 @@ +import folder_paths +import comfy.utils +import comfy.model_detection +import comfy.model_management +import comfy.lora +from comfy.model_patcher import ModelPatcher + +from .utils import TimestepKeyframeGroup +from .control import ControlNetAdvanced, load_controlnet + + + + +def convert_cn_lora_from_diffusers(cn_model: ModelPatcher, lora_path: str): + lora_data = comfy.utils.load_torch_file(lora_path, safe_load=True) + unet_dtype = comfy.model_management.unet_dtype() + for key, value in lora_data.items(): + lora_data[key] = value.to(unet_dtype) + diffusers_keys = comfy.utils.unet_to_diffusers(cn_model.model.state_dict()) + + #lora_data = comfy.model_detection.unet_config_from_diffusers_unet(lora_data, dtype=unet_dtype) + + + + #key_map = comfy.lora.model_lora_keys_unet(cn_model.model, key_map) + lora_data = comfy.lora.load_lora(lora_data, to_load=diffusers_keys) + + # TODO: detect if diffusers for sure? not sure if needed at this time, since cn loras are + # only used currently for LOOSEControl, and those are all in diffusers format + #unet_dtype = comfy.model_management.unet_dtype() + #lora_data = comfy.model_detection.unet_config_from_diffusers_unet(lora_data, unet_dtype) + return lora_data + + +class ControlNetLoaderWithLoraAdvanced: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "control_net_name": (folder_paths.get_filename_list("controlnet"), ), + "cn_lora_name": (folder_paths.get_filename_list("controlnet"), ), + "cn_lora_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}), + }, + "optional": { + "timestep_keyframe": ("TIMESTEP_KEYFRAME", ), + } + } + + RETURN_TYPES = ("CONTROL_NET", ) + FUNCTION = "load_controlnet" + + CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/LOOSEControl" + + def load_controlnet(self, control_net_name, cn_lora_name, cn_lora_strength: float, + timestep_keyframe: TimestepKeyframeGroup=None + ): + controlnet_path = folder_paths.get_full_path("controlnet", control_net_name) + controlnet: ControlNetAdvanced = load_controlnet(controlnet_path, timestep_keyframe) + if not isinstance(controlnet, ControlNetAdvanced): + raise ValueError("Type {} is not compatible with CN LoRA features at this time.") + # now, try to load CN LoRA + lora_path = folder_paths.get_full_path("controlnet", cn_lora_name) + lora_data = convert_cn_lora_from_diffusers(cn_model=controlnet.control_model_wrapped, lora_path=lora_path) + # apply patches to wrapped control_model + controlnet.control_model_wrapped.add_patches(lora_data, strength_patch=cn_lora_strength) + # all done + return (controlnet,) diff --git a/control/nodes_reference.py b/adv_control/nodes_reference.py similarity index 100% rename from control/nodes_reference.py rename to adv_control/nodes_reference.py diff --git a/control/nodes_sparsectrl.py b/adv_control/nodes_sparsectrl.py similarity index 100% rename from control/nodes_sparsectrl.py rename to adv_control/nodes_sparsectrl.py diff --git a/control/nodes_weight.py b/adv_control/nodes_weight.py similarity index 100% rename from control/nodes_weight.py rename to adv_control/nodes_weight.py diff --git a/control/utils.py b/adv_control/utils.py similarity index 100% rename from control/utils.py rename to adv_control/utils.py