From f778bb2a70b5ecbbafa0ec1a23f207486a8a0abf Mon Sep 17 00:00:00 2001 From: Kosinkadink Date: Wed, 25 Sep 2024 20:50:20 +0900 Subject: [PATCH] Modified code to accept transformer_options to match ComfyUI's get_control change and removed need to a terribly ugly hack that was required before to get the cond_or_uncond; now can just be gotten from transformer_options --- adv_control/control.py | 20 ++--- adv_control/control_lllite.py | 4 +- adv_control/control_plusplus.py | 4 +- adv_control/control_reference.py | 4 +- adv_control/utils.py | 124 +++---------------------------- 5 files changed, 26 insertions(+), 130 deletions(-) diff --git a/adv_control/control.py b/adv_control/control.py index 715d725..9778975 100644 --- a/adv_control/control.py +++ b/adv_control/control.py @@ -39,14 +39,14 @@ class ControlNetAdvanced(ControlNet, AdvancedControlBase): return raw_weights[idx] return self.weights.copy_with_new_weights(new_weight_func=cn_weights_func) - def get_control_advanced(self, x_noisy, t, cond, batched_number): + def get_control_advanced(self, x_noisy, t, cond, batched_number, transformer_options): # perform special version of get_control that supports sliding context and masks - return self.sliding_get_control(x_noisy, t, cond, batched_number) + return self.sliding_get_control(x_noisy, t, cond, batched_number, transformer_options) - def sliding_get_control(self, x_noisy: Tensor, t, cond, batched_number): + def sliding_get_control(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]: @@ -173,7 +173,7 @@ class T2IAdapterAdvanced(T2IAdapter, AdvancedControlBase): indeces.reverse() # need to reverse to match recent ComfyUI changes return indeces[idx] - def get_control_advanced(self, x_noisy, t, cond, batched_number): + def get_control_advanced(self, x_noisy, t, cond, batched_number, transformer_options): try: # if sub indexes present, replace original hint with subsection if self.sub_idxs is not None: @@ -187,7 +187,7 @@ class T2IAdapterAdvanced(T2IAdapter, AdvancedControlBase): self.cond_hint_original = actual_cond_hint_orig[self.sub_idxs] # mask hints self.prepare_mask_cond_hint(x_noisy=x_noisy, t=t, cond=cond, batched_number=batched_number) - return super().get_control(x_noisy, t, cond, batched_number) + return super().get_control(x_noisy, t, cond, batched_number, transformer_options) finally: if self.sub_idxs is not None: # replace original cond hint @@ -252,10 +252,10 @@ class SVDControlNetAdvanced(ControlNetAdvanced): self.cond_hint_original = self.cond_hint_original * 2.0 - 1.0 return to_return - def get_control_advanced(self, x_noisy, t, cond, batched_number): + def get_control_advanced(self, x_noisy, 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]: @@ -324,11 +324,11 @@ class SparseCtrlAdvanced(ControlNetAdvanced): self.model_latent_format = None # latent format for active SD model, NOT controlnet self.preprocessed = False - 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): # 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]: diff --git a/adv_control/control_lllite.py b/adv_control/control_lllite.py index 4d9dcc1..e1a367e 100644 --- a/adv_control/control_lllite.py +++ b/adv_control/control_lllite.py @@ -319,11 +319,11 @@ class ControlLLLiteAdvanced(ControlBase, AdvancedControlBase): 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): # 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]: diff --git a/adv_control/control_plusplus.py b/adv_control/control_plusplus.py index cb5c80b..87d3bdd 100644 --- a/adv_control/control_plusplus.py +++ b/adv_control/control_plusplus.py @@ -276,10 +276,10 @@ class ControlNetPlusPlusAdvanced(ControlNet, AdvancedControlBase): 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]: diff --git a/adv_control/control_reference.py b/adv_control/control_reference.py index 3bea56e..3c54503 100644 --- a/adv_control/control_reference.py +++ b/adv_control/control_reference.py @@ -239,11 +239,11 @@ class ReferenceAdvanced(ControlBase, AdvancedControlBase): ropts.attn_style_fidelity = ropts.original_attn_style_fidelity ropts.adain_style_fidelity = ropts.original_adain_style_fidelity - 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): # 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]: diff --git a/adv_control/utils.py b/adv_control/utils.py index 9c7c1bb..91dfd79 100644 --- a/adv_control/utils.py +++ b/adv_control/utils.py @@ -9,9 +9,6 @@ import math import comfy.ops import comfy.utils -import comfy.sample -import comfy.samplers -import comfy.model_base from comfy.controlnet import ControlBase from comfy.model_patcher import ModelPatcher @@ -33,108 +30,6 @@ def load_torch_file_with_dict_factory(controlnet_data: dict[str, Tensor], orig_l return controlnet_data return load_torch_file_with_dict -# wrapping len function so that it will save the thing len is trying to get the length of; -# this will be assumed to be the cond_or_uncond variable; -# automatically restores len to original function after running -def wrapper_len_factory(orig_len: Callable) -> Callable: - def wrapper_len(*args, **kwargs): - cond_or_uncond = args[0] - real_length = orig_len(*args, **kwargs) - if real_length > 0 and type(cond_or_uncond) == list and isinstance(cond_or_uncond[0], int) and (cond_or_uncond[0] in [0, 1]): - try: - to_return = IntWithCondOrUncond(real_length) - setattr(to_return, "cond_or_uncond", cond_or_uncond) - return to_return - finally: - __builtins__["len"] = orig_len - else: - return real_length - return wrapper_len - -# wrapping cond_cat function so that it will wrap around len function to get cond_or_uncond variable value -# from comfy.samplers.calc_conds_batch -def wrapper_cond_cat_factory(orig_cond_cat: Callable): - def wrapper_cond_cat(*args, **kwargs): - __builtins__["len"] = wrapper_len_factory(__builtins__["len"]) - return orig_cond_cat(*args, **kwargs) - return wrapper_cond_cat -orig_cond_cat = comfy.samplers.cond_cat -comfy.samplers.cond_cat = wrapper_cond_cat_factory(orig_cond_cat) - - -# wrapping apply_model so that len function will be cleaned up fairly soon after being injected -def apply_model_uncond_cleanup_factory(orig_apply_model, orig_len): - def apply_model_uncond_cleanup_wrapper(self, *args, **kwargs): - __builtins__["len"] = orig_len - return orig_apply_model(self, *args, **kwargs) - return apply_model_uncond_cleanup_wrapper -global_orig_len = __builtins__["len"] -orig_apply_model = comfy.model_base.BaseModel.apply_model -comfy.model_base.BaseModel.apply_model = apply_model_uncond_cleanup_factory(orig_apply_model, global_orig_len) - - -def uncond_multiplier_check_cn_sample_factory(orig_comfy_sample: Callable, is_custom=False) -> Callable: - def contains_uncond_multiplier(control: Union[ControlBase, 'AdvancedControlBase']): - if control is None: - return False - if not isinstance(control, AdvancedControlBase): - return contains_uncond_multiplier(control.previous_controlnet) - # check if weights_override has an uncond_multiplier - if control.weights_override is not None and control.weights_override.has_uncond_multiplier: - return True - # check if any timestep_keyframes have an uncond_multiplier on their weights - if control.timestep_keyframes is not None: - for tk in control.timestep_keyframes.keyframes: - if tk.has_control_weights() and tk.control_weights.has_uncond_multiplier: - return True - return contains_uncond_multiplier(control.previous_controlnet) - - # check if positive or negative conds contain Adv. Cns that use multiply_negative on weights - def uncond_multiplier_check_cn_sample(model: ModelPatcher, *args, **kwargs): - positive = args[-3] - negative = args[-2] - has_uncond_multiplier = False - if positive is not None: - for cond in positive: - if "control" in cond[1]: - has_uncond_multiplier = contains_uncond_multiplier(cond[1]["control"]) - if has_uncond_multiplier: - break - if negative is not None and not has_uncond_multiplier: - for cond in negative: - if "control" in cond[1]: - has_uncond_multiplier = contains_uncond_multiplier(cond[1]["control"]) - if has_uncond_multiplier: - break - try: - # if uncond_multiplier found, continue to use wrapped version of function - if has_uncond_multiplier: - return orig_comfy_sample(model, *args, **kwargs) - # otherwise, use original version of function to prevent even the smallest of slowdowns (0.XX%) - try: - wrapped_cond_cat = comfy.samplers.cond_cat - comfy.samplers.cond_cat = orig_cond_cat - return orig_comfy_sample(model, *args, **kwargs) - finally: - comfy.samplers.cond_cat = wrapped_cond_cat - finally: - # make sure len function is unwrapped by the time sampling is done, just in case - __builtins__["len"] = global_orig_len - return uncond_multiplier_check_cn_sample -# inject sample functions -comfy.sample.sample = uncond_multiplier_check_cn_sample_factory(comfy.sample.sample) -comfy.sample.sample_custom = uncond_multiplier_check_cn_sample_factory(comfy.sample.sample_custom, is_custom=True) - - -class IntWithCondOrUncond(int): - def __new__(cls, *args, **kwargs): - return super(IntWithCondOrUncond, cls).__new__(cls, *args, **kwargs) - - def __init__(self, *args, **kwargs): - super().__init__() - self.cond_or_uncond = None - - def get_properly_arranged_t2i_weights(initial_weights: list[float]): new_weights = [] @@ -579,8 +474,9 @@ class AdvancedControlBase: # timesteps self.t: float = None self.prev_t: float = None - self.batched_number: Union[int, IntWithCondOrUncond] = None + self.batched_number: int = None self.batch_size: int = 0 + self.cond_or_uncond: list[int] = None # weights + override self.weights: ControlWeights = None self.weights_default: ControlWeights = weights_default @@ -750,24 +646,25 @@ class AdvancedControlBase: return False return True - def get_control_inject(self, x_noisy, t, cond, batched_number): + def get_control_inject(self, x_noisy, t, cond, batched_number, transformer_options: dict): self.batched_number = batched_number self.batch_size = len(t) + self.cond_or_uncond = transformer_options.get("cond_or_uncond", None) # prepare timestep and everything related self.prepare_current_timestep(t=t, batched_number=batched_number) # if should not perform any actions for the controlnet, exit without doing any work if self.strength == 0.0 or self._current_timestep_keyframe.strength == 0.0: return self.default_control_actions(x_noisy, t, cond, batched_number) # otherwise, perform normal function - return self.get_control_advanced(x_noisy, t, cond, batched_number) + return self.get_control_advanced(x_noisy, t, cond, batched_number, transformer_options) - def get_control_advanced(self, x_noisy, t, cond, batched_number): - return self.default_control_actions(x_noisy, t, cond, batched_number) + def get_control_advanced(self, x_noisy, t, cond, batched_number, transformer_options): + return self.default_control_actions(x_noisy, t, cond, batched_number, transformer_options) - def default_control_actions(self, x_noisy, t, cond, batched_number): + def default_control_actions(self, x_noisy, 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) return control_prev def calc_weight(self, idx: int, x: Tensor, control: dict[str, list[Tensor]], key: str) -> Union[float, Tensor]: @@ -839,9 +736,8 @@ class AdvancedControlBase: def apply_advanced_strengths_and_masks(self, x: Tensor, batched_number: int, flux_shape: tuple=None): # handle weight's uncond_multiplier, if applicable if self.weights.has_uncond_multiplier: - cond_or_uncond = self.batched_number.cond_or_uncond actual_length = x.size(0) // batched_number - for idx, cond_type in enumerate(cond_or_uncond): + for idx, cond_type in enumerate(self.cond_or_uncond): # if uncond, set to weight's uncond_multiplier if cond_type == 1: x[actual_length*idx:actual_length*(idx+1)] *= self.weights.uncond_multiplier