Merge PR #198 from Kosinkadink/rework-modelpatcher

Rework ModelPatcher for upcoming ComfyUI update

Author: Kosinkadink

__init__.py

@@ -1,6 +1,11 @@
 from .adv_control.nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
 from .adv_control import documentation
+from .adv_control.dinklink import init_dinklink
+from .adv_control.sampling import prepare_dinklink_acn_wrapper
 
 WEB_DIRECTORY = "./web"
 __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', "WEB_DIRECTORY"]
 documentation.format_descriptions(NODE_CLASS_MAPPINGS)
+
+init_dinklink()
+prepare_dinklink_acn_wrapper()

adv_control/control.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)

adv_control/control_ctrlora.py

@@ -15,12 +15,12 @@
 
 from .logger import logger
 from .utils import (AdvancedControlBase, TimestepKeyframeGroup, ControlWeights, broadcast_image_to_extend, extend_to_batch_size,
-                    deepcopy_with_sharing, prepare_mask_batch)
+                    prepare_mask_batch)
 
 
 # based on set_model_patch code in comfy/model_patcher.py
-def set_model_patch(model_options, patch, name):
-    to = model_options["transformer_options"]
+def set_model_patch(transformer_options, patch, name):
+    to = transformer_options
     # check if patch was already added
     if "patches" in to:
         current_patches = to["patches"].get(name, [])
@@ -30,11 +30,11 @@ def set_model_patch(model_options, patch, name):
         to["patches"] = {}
     to["patches"][name] = to["patches"].get(name, []) + [patch]
 
-def set_model_attn1_patch(model_options, patch):
-    set_model_patch(model_options, patch, "attn1_patch")
+def set_model_attn1_patch(transformer_options, patch):
+    set_model_patch(transformer_options, patch, "attn1_patch")
 
-def set_model_attn2_patch(model_options, patch):
-    set_model_patch(model_options, patch, "attn2_patch")
+def set_model_attn2_patch(transformer_options, patch):
+    set_model_patch(transformer_options, patch, "attn2_patch")
 
 
 def extra_options_to_module_prefix(extra_options):
@@ -115,26 +115,8 @@ def clone_with_control(self, control: AdvancedControlBase):
         return LLLitePatch(self.modules, self.patch_type, control)
 
     def cleanup(self):
-        #total_cleaned = 0
         for module in self.modules.values():
             module.cleanup()
-        #    total_cleaned += 1
-        #logger.info(f"cleaned modules: {total_cleaned}, {id(self)}")
-        #logger.error(f"cleanup LLLitePatch: {id(self)}")
-
-    # make sure deepcopy does not copy control, and deepcopied LLLitePatch should be assigned to control
-    # def __deepcopy__(self, memo):
-    #     self.cleanup()
-    #     to_return: LLLitePatch = deepcopy_with_sharing(self, shared_attribute_names = ['control'], memo=memo)
-    #     #logger.warn(f"patch {id(self)} turned into {id(to_return)}")
-    #     try:
-    #         if self.patch_type == self.ATTN1:
-    #             to_return.control.patch_attn1 = to_return
-    #         elif self.patch_type == self.ATTN2:
-    #             to_return.control.patch_attn2 = to_return
-    #     except Exception:
-    #         pass
-    #     return to_return
 
 
 # TODO: use comfy.ops to support fp8 properly
@@ -298,14 +280,6 @@ def __init__(self, patch_attn1: LLLitePatch, patch_attn2: LLLitePatch, timestep_
         self.latent_dims_div2 = None
         self.latent_dims_div4 = None
 
-    def live_model_patches(self, model_options):
-        set_model_attn1_patch(model_options, self.patch_attn1.set_control(self))
-        set_model_attn2_patch(model_options, self.patch_attn2.set_control(self))
-
-    # def patch_model(self, model: ModelPatcher):
-    #     model.set_model_attn1_patch(self.patch_attn1)
-    #     model.set_model_attn2_patch(self.patch_attn2)
-
     def set_cond_hint_inject(self, *args, **kwargs):
         to_return = super().set_cond_hint_inject(*args, **kwargs)
         # cond hint for LLLite needs to be scaled between (-1, 1) instead of (0, 1)
@@ -319,11 +293,11 @@ def pre_run_advanced(self, *args, **kwargs):
         self.patch_attn2.set_control(self)
         #logger.warn(f"in pre_run_advanced: {id(self)}")
 
-    def get_control_advanced(self, x_noisy: Tensor, t, cond, batched_number: int):
+    def get_control_advanced(self, x_noisy: Tensor, t, cond, batched_number: int, transformer_options: dict):
         # normal ControlNet stuff
         control_prev = None
         if self.previous_controlnet is not None:
-            control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number)
+            control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number, transformer_options)
 
         if self.timestep_range is not None:
             if t[0] > self.timestep_range[0] or t[0] < self.timestep_range[1]:
@@ -372,7 +346,9 @@ def get_control_advanced(self, x_noisy: Tensor, t, cond, batched_number: int):
             self.latent_dims_div4 = (new_h, new_w)
         # prepare mask
         self.prepare_mask_cond_hint(x_noisy=x_noisy, t=t, cond=cond, batched_number=batched_number)
-        # done preparing; model patches will take care of everything now.
+        # done preparing; model patches will take care of everything now
+        set_model_attn1_patch(transformer_options, self.patch_attn1.set_control(self))
+        set_model_attn2_patch(transformer_options, self.patch_attn2.set_control(self))
         # return normal controlnet stuff
         return control_prev
 
@@ -393,17 +369,6 @@ def copy(self):
         self.copy_to(c)
         self.copy_to_advanced(c)
         return c
-    
-    # deepcopy needs to properly keep track of objects to work between model.clone calls!
-    # def __deepcopy__(self, *args, **kwargs):
-    #     self.cleanup_advanced()
-    #     return self
-
-    # def get_models(self):
-    #     # get_models is called once at the start of every KSampler run - use to reset already_patched status
-    #     out = super().get_models()
-    #     logger.error(f"in get_models! {id(self)}")
-    #     return out
 
 
 def load_controllllite(ckpt_path: str, controlnet_data: dict[str, Tensor]=None, timestep_keyframe: TimestepKeyframeGroup=None):

adv_control/control_lllite.py

@@ -22,7 +22,7 @@
 import comfy.model_detection
 import comfy.utils
 
-from .utils import (AdvancedControlBase, ControlWeights, ControlWeightType, TimestepKeyframeGroup, AbstractPreprocWrapper,
+from .utils import (AdvancedControlBase, ControlWeights, ControlWeightType, TimestepKeyframeGroup, AbstractPreprocWrapper, Extras,
                     extend_to_batch_size, broadcast_image_to_extend)
 from .logger import logger
 
@@ -239,7 +239,7 @@ def __init__(self, control_model: ControlNetPlusPlus, timestep_keyframes: Timest
     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
+                return 1.0 * self.weights.extras.get(Extras.MIDDLE_MULT, 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]
@@ -276,10 +276,10 @@ def set_cond_hint_inject(self, *args, **kwargs):
             self.cond_hint_original = pp_group
         return to_return
 
-    def get_control_advanced(self, x_noisy: Tensor, t, cond, batched_number):
+    def get_control_advanced(self, x_noisy: Tensor, t, cond, batched_number, transformer_options):
         control_prev = None
         if self.previous_controlnet is not None:
-            control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number)
+            control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number, transformer_options)
 
         if self.timestep_range is not None:
             if t[0] > self.timestep_range[0] or t[0] < self.timestep_range[1]: