diff --git a/adv_control/control.py b/adv_control/control.py index 1d81e39..4225a47 100644 --- a/adv_control/control.py +++ b/adv_control/control.py @@ -26,8 +26,16 @@ class ControlNetAdvanced(ControlNet, AdvancedControlBase): AdvancedControlBase.__init__(self, super(), timestep_keyframes=timestep_keyframes, weights_default=ControlWeights.controlnet()) def get_universal_weights(self) -> ControlWeights: - raw_weights = [(self.weights.base_multiplier ** float(12 - i)) for i in range(13)] - return self.weights.copy_with_new_weights(raw_weights) + def cn_weights_func(idx: int, control: dict[str, list[Tensor]], key: str): + if key == "middle": + return 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] + if key == "input": + raw_weights.reverse() + 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): # perform special version of get_control that supports sliding context and masks @@ -127,18 +135,28 @@ class T2IAdapterAdvanced(T2IAdapter, AdvancedControlBase): return AdvancedControlBase.control_merge_inject(self, control, control_prev, output_dtype) def get_universal_weights(self) -> ControlWeights: - raw_weights = [(self.weights.base_multiplier ** float(7 - i)) for i in range(8)] - raw_weights = [raw_weights[-8], raw_weights[-3], raw_weights[-2], raw_weights[-1]] - raw_weights = get_properly_arranged_t2i_weights(raw_weights) - raw_weights.reverse() # need to reverse to match recent ComfyUI changes - return self.weights.copy_with_new_weights(raw_weights) + def t2i_weights_func(idx: int, control: dict[str, list[Tensor]], key: str): + if key == "middle": + return 1.0 + c_len = 8 #len(control[key]) + raw_weights = [(self.weights.base_multiplier ** float((c_len-1) - i)) for i in range(c_len)] + raw_weights = [raw_weights[-c_len], raw_weights[-3], raw_weights[-2], raw_weights[-1]] + raw_weights = get_properly_arranged_t2i_weights(raw_weights) + if key == "input": + raw_weights.reverse() + return raw_weights[idx] + return self.weights.copy_with_new_weights(new_weight_func=t2i_weights_func) def get_calc_pow(self, idx: int, control: dict[str, list[Tensor]], key: str) -> int: + if key == "middle": + return 0 # match how T2IAdapterAdvanced deals with universal weights - indeces = [7 - i for i in range(8)] - indeces = [indeces[-8], indeces[-3], indeces[-2], indeces[-1]] + c_len = 8 #len(control[key]) + indeces = [(c_len-1) - i for i in range(c_len)] + indeces = [indeces[-c_len], indeces[-3], indeces[-2], indeces[-1]] indeces = get_properly_arranged_t2i_weights(indeces) - indeces.reverse() # need to reverse to match recent ComfyUI changes + if key == "input": + indeces.reverse() # need to reverse to match recent ComfyUI changes return indeces[idx] def get_control_advanced(self, x_noisy, t, cond, batched_number): diff --git a/adv_control/control_plusplus.py b/adv_control/control_plusplus.py index a6f5ca9..cb5c80b 100644 --- a/adv_control/control_plusplus.py +++ b/adv_control/control_plusplus.py @@ -237,9 +237,16 @@ class ControlNetPlusPlusAdvanced(ControlNet, AdvancedControlBase): self.single_control_type: str = None def get_universal_weights(self) -> ControlWeights: - # TODO: match actual layer count of model - raw_weights = [(self.weights.base_multiplier ** float(12 - i)) for i in range(13)] - return self.weights.copy_with_new_weights(raw_weights) + def cn_weights_func(idx: int, control: dict[str, list[Tensor]], key: str): + if key == "middle": + return 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] + if key == "input": + raw_weights.reverse() + return raw_weights[idx] + return self.weights.copy_with_new_weights(new_weight_func=cn_weights_func) def verify_control_type(self, model_name: str, pp_group: PlusPlusInputGroup=None): if pp_group is not None: diff --git a/adv_control/nodes.py b/adv_control/nodes.py index 55c2db2..470149f 100644 --- a/adv_control/nodes.py +++ b/adv_control/nodes.py @@ -7,8 +7,8 @@ from comfy.model_patcher import ModelPatcher from .control import load_controlnet, convert_to_advanced, is_advanced_controlnet, is_sd3_advanced_controlnet from .utils import ControlWeights, LatentKeyframeGroup, TimestepKeyframeGroup, AbstractPreprocWrapper, BIGMAX -from .nodes_weight import (DefaultWeights, ScaledSoftMaskedUniversalWeights, ScaledSoftUniversalWeights, SoftControlNetWeights, CustomControlNetWeights, - SoftT2IAdapterWeights, CustomT2IAdapterWeights) +from .nodes_weight import (DefaultWeights, ScaledSoftMaskedUniversalWeights, ScaledSoftUniversalWeightsDeprecated, SoftControlNetWeightsDeprecated, CustomControlNetWeightsDeprecated, + SoftT2IAdapterWeightsDeprecated, CustomT2IAdapterWeightsDeprecated) from .nodes_keyframes import (LatentKeyframeGroupNode, LatentKeyframeInterpolationNode, LatentKeyframeBatchedGroupNode, LatentKeyframeNode, TimestepKeyframeNode, TimestepKeyframeInterpolationNode, TimestepKeyframeFromStrengthListNode) from .nodes_sparsectrl import SparseCtrlMergedLoaderAdvanced, SparseCtrlLoaderAdvanced, SparseIndexMethodNode, SparseSpreadMethodNode, RgbSparseCtrlPreprocessor, SparseWeightExtras @@ -200,12 +200,12 @@ NODE_CLASS_MAPPINGS = { "ControlNetLoaderAdvanced": ControlNetLoaderAdvanced, "DiffControlNetLoaderAdvanced": DiffControlNetLoaderAdvanced, # Weights - "ScaledSoftControlNetWeights": ScaledSoftUniversalWeights, + "ScaledSoftControlNetWeights": ScaledSoftUniversalWeightsDeprecated, "ScaledSoftMaskedUniversalWeights": ScaledSoftMaskedUniversalWeights, - "SoftControlNetWeights": SoftControlNetWeights, - "CustomControlNetWeights": CustomControlNetWeights, - "SoftT2IAdapterWeights": SoftT2IAdapterWeights, - "CustomT2IAdapterWeights": CustomT2IAdapterWeights, + "SoftControlNetWeights": SoftControlNetWeightsDeprecated, + "CustomControlNetWeights": CustomControlNetWeightsDeprecated, + "SoftT2IAdapterWeights": SoftT2IAdapterWeightsDeprecated, + "CustomT2IAdapterWeights": CustomT2IAdapterWeightsDeprecated, "ACN_DefaultUniversalWeights": DefaultWeights, # SparseCtrl "ACN_SparseCtrlRGBPreprocessor": RgbSparseCtrlPreprocessor, diff --git a/adv_control/nodes_weight.py b/adv_control/nodes_weight.py index 263c529..85f2af7 100644 --- a/adv_control/nodes_weight.py +++ b/adv_control/nodes_weight.py @@ -66,7 +66,7 @@ class ScaledSoftMaskedUniversalWeights: return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) -class ScaledSoftUniversalWeights: +class ScaledSoftUniversalWeightsDeprecated: @classmethod def INPUT_TYPES(s): return { @@ -88,11 +88,11 @@ class ScaledSoftUniversalWeights: CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/weights" def load_weights(self, base_multiplier, flip_weights, uncond_multiplier: float=1.0, cn_extras: dict[str]={}): - weights = ControlWeights.universal(base_multiplier=base_multiplier, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=cn_extras) - return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) + weights = ControlWeights.universal(base_multiplier=base_multiplier, uncond_multiplier=uncond_multiplier, extras=cn_extras) + return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) -class SoftControlNetWeights: +class SoftControlNetWeightsDeprecated: @classmethod def INPUT_TYPES(s): return { @@ -128,13 +128,14 @@ class SoftControlNetWeights: def load_weights(self, weight_00, weight_01, weight_02, weight_03, weight_04, weight_05, weight_06, weight_07, weight_08, weight_09, weight_10, weight_11, weight_12, flip_weights, uncond_multiplier: float=1.0, cn_extras: dict[str]={}): - weights = [weight_00, weight_01, weight_02, weight_03, weight_04, weight_05, weight_06, - weight_07, weight_08, weight_09, weight_10, weight_11, weight_12] - weights = ControlWeights.controlnet(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=cn_extras) + weights_output = [weight_00, weight_01, weight_02, weight_03, weight_04, weight_05, weight_06, + weight_07, weight_08, weight_09, weight_10, weight_11] + weights_middle = [weight_12] + weights = ControlWeights.controlnet(weights_output=weights_output, weights_middle=weights_middle, uncond_multiplier=uncond_multiplier, extras=cn_extras) return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) -class CustomControlNetWeights: +class CustomControlNetWeightsDeprecated: @classmethod def INPUT_TYPES(s): return { @@ -170,13 +171,14 @@ class CustomControlNetWeights: def load_weights(self, weight_00, weight_01, weight_02, weight_03, weight_04, weight_05, weight_06, weight_07, weight_08, weight_09, weight_10, weight_11, weight_12, flip_weights, uncond_multiplier: float=1.0, cn_extras: dict[str]={}): - weights = [weight_00, weight_01, weight_02, weight_03, weight_04, weight_05, weight_06, - weight_07, weight_08, weight_09, weight_10, weight_11, weight_12] - weights = ControlWeights.controlnet(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=cn_extras) + weights_output = [weight_00, weight_01, weight_02, weight_03, weight_04, weight_05, weight_06, + weight_07, weight_08, weight_09, weight_10, weight_11] + weights_middle = [weight_12] + weights = ControlWeights.controlnet(weights_output=weights_output, weights_middle=weights_middle, uncond_multiplier=uncond_multiplier, extras=cn_extras) return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) -class SoftT2IAdapterWeights: +class SoftT2IAdapterWeightsDeprecated: @classmethod def INPUT_TYPES(s): return { @@ -205,11 +207,11 @@ class SoftT2IAdapterWeights: weights = [weight_00, weight_01, weight_02, weight_03] weights = get_properly_arranged_t2i_weights(weights) weights.reverse() # to account for recent ComfyUI changes - weights = ControlWeights.t2iadapter(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=cn_extras) + weights = ControlWeights.t2iadapter(weights_input=weights, uncond_multiplier=uncond_multiplier, extras=cn_extras) return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) -class CustomT2IAdapterWeights: +class CustomT2IAdapterWeightsDeprecated: @classmethod def INPUT_TYPES(s): return { @@ -238,5 +240,5 @@ class CustomT2IAdapterWeights: weights = [weight_00, weight_01, weight_02, weight_03] weights = get_properly_arranged_t2i_weights(weights) weights.reverse() # to account for recent ComfyUI changes - weights = ControlWeights.t2iadapter(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=cn_extras) + weights = ControlWeights.t2iadapter(weights_input=weights, uncond_multiplier=uncond_multiplier, extras=cn_extras) return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) diff --git a/adv_control/utils.py b/adv_control/utils.py index df7a67e..e2d5a48 100644 --- a/adv_control/utils.py +++ b/adv_control/utils.py @@ -157,14 +157,16 @@ class ControlWeightType: class ControlWeights: - def __init__(self, weight_type: str, base_multiplier: float=1.0, flip_weights: bool=False, weights: list[float]=None, weight_mask: Tensor=None, + def __init__(self, weight_type: str, base_multiplier: float=1.0, + weights_input: list[float]=None, weights_middle: list[float]=None, weights_output: list[float]=None, + weight_func: Callable=None, weight_mask: Tensor=None, uncond_multiplier=1.0, uncond_mask: Tensor=None, extras: dict[str]={},): self.weight_type = weight_type self.base_multiplier = base_multiplier - self.flip_weights = flip_weights - self.weights = weights - if self.weights is not None and self.flip_weights: - self.weights.reverse() + self.weights_input = weights_input + self.weights_middle = weights_middle + self.weights_output = weights_output + self.weight_func = weight_func self.weight_mask = weight_mask self.uncond_multiplier = float(uncond_multiplier) self.has_uncond_multiplier = not math.isclose(self.uncond_multiplier, 1.0) @@ -173,66 +175,74 @@ class ControlWeights: self.extras = extras def get(self, idx: int, control: dict[str, list[Tensor]], key: str, default=1.0) -> Union[float, Tensor]: + # if weight_func present, use it + if self.weight_func is not None: + return self.weight_func(idx=idx, control=control, key=key) # if weights is not none, return index - if self.weights is not None: - # support live calculation - if callable(self.weights): - return self.weights(idx=idx, control=control, key=key) - # if middle weight, need to pretend index is actually after all the output weights (if applicable) - if key == "middle" and "output" in control: - idx += len(control["output"]) + relevant_weights = None + if key == "middle": + relevant_weights = self.weights_middle + elif key == "input": + relevant_weights = self.weights_input + else: + relevant_weights = self.weights_output + if relevant_weights is None: + return default + return relevant_weights[idx] - if key == "output": - # this implies weights list is not aligning with expectations - will need to adjust code - if idx >= len(self.weights)-1: - return default - else: - # this implies weights list is not aligning with expectations - will need to adjust code - if idx >= len(self.weights): - return default - return self.weights[idx] - return 1.0 - - def copy_with_new_weights(self, new_weights: list[float]): - return ControlWeights(weight_type=self.weight_type, base_multiplier=self.base_multiplier, flip_weights=self.flip_weights, - weights=new_weights, weight_mask=self.weight_mask, uncond_multiplier=self.uncond_multiplier, extras=self.extras) + def copy_with_new_weights(self, new_weights_input: list[float]=None, new_weights_middle: list[float]=None, new_weights_output: list[float]=None, + new_weight_func: Callable=None): + return ControlWeights(weight_type=self.weight_type, base_multiplier=self.base_multiplier, + weights_input=new_weights_input, weights_middle=new_weights_middle, weights_output=new_weights_output, + weight_func=new_weight_func, weight_mask=self.weight_mask, + uncond_multiplier=self.uncond_multiplier, extras=self.extras) @classmethod def default(cls, extras: dict[str]={}): return cls(ControlWeightType.DEFAULT, extras=extras) @classmethod - def universal(cls, base_multiplier: float, flip_weights: bool=False, uncond_multiplier: float=1.0, extras: dict[str]={}): - return cls(ControlWeightType.UNIVERSAL, base_multiplier=base_multiplier, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=extras) + def universal(cls, base_multiplier: float, uncond_multiplier: float=1.0, extras: dict[str]={}): + return cls(ControlWeightType.UNIVERSAL, base_multiplier=base_multiplier, uncond_multiplier=uncond_multiplier, extras=extras) @classmethod def universal_mask(cls, weight_mask: Tensor, uncond_multiplier: float=1.0, extras: dict[str]={}): return cls(ControlWeightType.UNIVERSAL, weight_mask=weight_mask, uncond_multiplier=uncond_multiplier, extras=extras) @classmethod - def t2iadapter(cls, weights: list[float]=None, flip_weights: bool=False, uncond_multiplier: float=1.0, extras: dict[str]={}): - if weights is None: - weights = [1.0]*12 - return cls(ControlWeightType.T2IADAPTER, weights=weights,flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=extras) + def t2iadapter(cls, weights_input: list[float]=None, uncond_multiplier: float=1.0, extras: dict[str]={}): + if weights_input is None: + weights_input = [1.0]*12 + return cls(ControlWeightType.T2IADAPTER, weights_input=weights_input, uncond_multiplier=uncond_multiplier, extras=extras) @classmethod - def controlnet(cls, weights: list[float]=None, flip_weights: bool=False, uncond_multiplier: float=1.0, extras: dict[str]={}): - if weights is None: - weights = [1.0]*13 - return cls(ControlWeightType.CONTROLNET, weights=weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=extras) + def controlnet(cls, weights_output: list[float]=None, weights_middle: list[float]=None, uncond_multiplier: float=1.0, extras: dict[str]={}): + if weights_output is None: + weights_output = [1.0]*12 + if weights_middle is None: + weights_middle = [1.0]*1 + return cls(ControlWeightType.CONTROLNET, weights_output=weights_output, weights_middle=weights_middle, uncond_multiplier=uncond_multiplier, extras=extras) @classmethod - def controllora(cls, weights: list[float]=None, flip_weights: bool=False, uncond_multiplier: float=1.0, extras: dict[str]={}): - if weights is None: - weights = [1.0]*10 - return cls(ControlWeightType.CONTROLLORA, weights=weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=extras) + def controllora(cls, weights_output: list[float]=None, weights_middle: list[float]=None, weights_input: list[float]=None, uncond_multiplier: float=1.0, extras: dict[str]={}): + if weights_output is None: + weights_output = [1.0]*10 + if weights_middle is None: + weights_middle = [1.0]*10 + if weights_input is None: + weights_input = [1.0]*10 + return cls(ControlWeightType.CONTROLLORA, weights_output=weights_output, weights_middle=weights_middle, weights_input=weights_input, uncond_multiplier=uncond_multiplier, extras=extras) @classmethod - def controllllite(cls, weights: list[float]=None, flip_weights: bool=False, uncond_multiplier: float=1.0, extras: dict[str]={}): - if weights is None: - # TODO: make this have a real value - weights = [1.0]*200 - return cls(ControlWeightType.CONTROLLLLITE, weights=weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=extras) + def controllllite(cls, weights_output: list[float]=None, weights_middle: list[float]=None, weights_input: list[float]=None, uncond_multiplier: float=1.0, extras: dict[str]={}): + # TODO: make the lenghts make sense + if weights_output is None: + weights_output = [1.0]*200 + if weights_middle is None: + weights_middle = [1.0]*200 + if weights_input is None: + weights_input = [1.0]*200 + return cls(ControlWeightType.CONTROLLLLITE, weights_output=weights_output, weights_middle=weights_middle, weights_input=weights_input, uncond_multiplier=uncond_multiplier, extras=extras) class StrengthInterpolation: @@ -778,11 +788,14 @@ class AdvancedControlBase: return self.weights.get(idx=idx, control=control, key=key) def get_calc_pow(self, idx: int, control: dict[str, list[Tensor]], key: str) -> int: - c_len = len(control[key])-1 - if key == "output": - if "middle" in control: - c_len += len(control["middle"]) - return c_len-idx + if key == "middle": + return 0 + else: + c_len = len(control[key]) + real_idx = c_len-idx + if key == "input": + real_idx = c_len - real_idx + 1 + return real_idx def calc_latent_keyframe_mults(self, x: Tensor, batched_number: int) -> Tensor: # apply strengths, and get batch indeces to null out @@ -916,7 +929,7 @@ class AdvancedControlBase: del out_mask # TODO: perform upscale on only the sub_idxs masks at a time instead of all to conserve RAM # resize mask and match batch count - out_mask = prepare_mask_batch(orig_mask, x_noisy.shape, multiplier=multiplier) + out_mask = prepare_mask_batch(orig_mask, x_noisy.shape, multiplier=multiplier, match_shape=True) actual_latent_length = x_noisy.shape[0] // batched_number out_mask = extend_to_batch_size(out_mask, actual_latent_length if self.sub_idxs is None else self.full_latent_length) if self.sub_idxs is not None: