From ca5d63cb335ec266f3146d3e0db25d343c5e07f7 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Fri, 19 Jul 2024 19:38:37 -0500 Subject: [PATCH] Fixed issues with ControlLLLite refactor --- adv_control/control_lllite.py | 36 ++++++++++++++++++++++++++++++++--- adv_control/sampling.py | 20 ++++++++++--------- 2 files changed, 44 insertions(+), 12 deletions(-) diff --git a/adv_control/control_lllite.py b/adv_control/control_lllite.py index 359c2a9..9fb2cbb 100644 --- a/adv_control/control_lllite.py +++ b/adv_control/control_lllite.py @@ -8,6 +8,8 @@ import torch import os import comfy.utils +import comfy.ops +import comfy.model_management from comfy.model_patcher import ModelPatcher from comfy.controlnet import ControlBase @@ -19,6 +21,11 @@ from .utils import (AdvancedControlBase, TimestepKeyframeGroup, ControlWeights, # based on set_model_patch code in comfy/model_patcher.py def set_model_patch(model_options, patch, name): to = model_options["transformer_options"] + # check if patch was already added + if "patches" in to: + current_patches = to["patches"].get(name, []) + if patch in current_patches: + return if "patches" not in to: to["patches"] = {} to["patches"][name] = to["patches"].get(name, []) + [patch] @@ -270,13 +277,24 @@ class LLLiteModule(torch.nn.Module): return cx * mask * control.strength * control._current_timestep_keyframe.strength +class ControlLLLiteModules(torch.nn.Module): + def __init__(self, patch_attn1: LLLitePatch, patch_attn2: LLLitePatch): + super().__init__() + self.patch_attn1_modules = torch.nn.Sequential(*list(patch_attn1.modules.values())) + self.patch_attn2_modules = torch.nn.Sequential(*list(patch_attn2.modules.values())) + + class ControlLLLiteAdvanced(ControlBase, AdvancedControlBase): # This ControlNet is more of an attention patch than a traditional controlnet - def __init__(self, patch_attn1: LLLitePatch, patch_attn2: LLLitePatch, timestep_keyframes: TimestepKeyframeGroup, device=None): + def __init__(self, patch_attn1: LLLitePatch, patch_attn2: LLLitePatch, timestep_keyframes: TimestepKeyframeGroup, device, ops: comfy.ops.disable_weight_init): super().__init__(device) AdvancedControlBase.__init__(self, super(), timestep_keyframes=timestep_keyframes, weights_default=ControlWeights.controllllite()) + self.device = device + self.ops = ops self.patch_attn1 = patch_attn1.clone_with_control(self) self.patch_attn2 = patch_attn2.clone_with_control(self) + self.control_model = ControlLLLiteModules(self.patch_attn1, self.patch_attn2) + self.control_model_wrapped = ModelPatcher(self.control_model, load_device=device, offload_device=comfy.model_management.unet_offload_device()) self.latent_dims_div2 = None self.latent_dims_div4 = None @@ -358,6 +376,12 @@ class ControlLLLiteAdvanced(ControlBase, AdvancedControlBase): # return normal controlnet stuff return control_prev + def get_models(self): + to_return: list = super().get_models() + to_return.append(self.control_model_wrapped) + return to_return + + def cleanup_advanced(self): super().cleanup_advanced() self.patch_attn1.cleanup() @@ -366,7 +390,7 @@ class ControlLLLiteAdvanced(ControlBase, AdvancedControlBase): self.latent_dims_div4 = None def copy(self): - c = ControlLLLiteAdvanced(self.patch_attn1, self.patch_attn2, self.timestep_keyframes) + c = ControlLLLiteAdvanced(self.patch_attn1, self.patch_attn2, self.timestep_keyframes, self.device, self.ops) self.copy_to(c) self.copy_to_advanced(c) return c @@ -424,9 +448,15 @@ def load_controllllite(ckpt_path: str, controlnet_data: dict[str, Tensor]=None, if len(modules) == 1: module.is_first = True + unet_dtype = comfy.model_management.unet_dtype() + load_device = comfy.model_management.get_torch_device() + manual_cast_dtype = comfy.model_management.unet_manual_cast(unet_dtype, load_device) + ops = comfy.ops.disable_weight_init + if manual_cast_dtype is not None: + ops = comfy.ops.manual_cast #logger.info(f"loaded {ckpt_path} successfully, {len(modules)} modules") patch_attn1 = LLLitePatch(modules=modules, patch_type=LLLitePatch.ATTN1) patch_attn2 = LLLitePatch(modules=modules, patch_type=LLLitePatch.ATTN2) - control = ControlLLLiteAdvanced(patch_attn1=patch_attn1, patch_attn2=patch_attn2, timestep_keyframes=timestep_keyframe) + control = ControlLLLiteAdvanced(patch_attn1=patch_attn1, patch_attn2=patch_attn2, timestep_keyframes=timestep_keyframe, device=load_device, ops=ops) return control diff --git a/adv_control/sampling.py b/adv_control/sampling.py index 940d36d..81f0407 100644 --- a/adv_control/sampling.py +++ b/adv_control/sampling.py @@ -42,13 +42,13 @@ def acn_sample_factory(orig_comfy_sample: Callable, is_custom=False) -> Callable return ref_set def get_lllitecn(control: ControlBase): - cn_list: list[ControlLLLiteAdvanced] = [] + cn_dict: dict[ControlLLLiteAdvanced,None] = {} if control is None: - return cn_list + return cn_dict if type(control) == ControlLLLiteAdvanced: - cn_list.append(control) - cn_list.extend(get_lllitecn(control.previous_controlnet)) - return cn_list + cn_dict[control] = None + cn_dict.update(get_lllitecn(control.previous_controlnet)) + return cn_dict def acn_sample(model: ModelPatcher, *args, **kwargs): controlnets_modified = False @@ -68,20 +68,22 @@ def acn_sample_factory(orig_comfy_sample: Callable, is_custom=False) -> Callable args = tuple(args) # look for Advanced ControlNets that will require intervention to work ref_set = set() - lllite_list: list[ControlLLLiteAdvanced] = [] + lllite_dict: dict[ControlLLLiteAdvanced, None] = {} # dicts preserve insertion order since py3.7 if positive is not None: for cond in positive: if "control" in cond[1]: ref_set.update(get_refcn(cond[1]["control"])) - lllite_list.extend(get_lllitecn(cond[1]["control"])) + lllite_dict.update(get_lllitecn(cond[1]["control"])) if negative is not None: for cond in negative: if "control" in cond[1]: ref_set.update(get_refcn(cond[1]["control"])) - lllite_list.extend(get_lllitecn(cond[1]["control"])) + lllite_dict.update(get_lllitecn(cond[1]["control"])) # if lllite found, apply patches to a cloned model_options, and continue - if len(lllite_list) > 0: + if len(lllite_dict) > 0: + lllite_list = list(lllite_dict.keys()) model.model_options = model.model_options.copy() + model.model_options["transformer_options"] = model.model_options["transformer_options"].copy() lllite_list.reverse() # reverse so that patches will be applied in expected order for lll in lllite_list: lll.live_model_patches(model.model_options)