From 1807cd64e4de87dbd09198f8782832eab099596d Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Sat, 24 Aug 2024 00:43:21 -0500 Subject: [PATCH] Updated ControlNetAdvanced to work with Flux controlnet, beginning of rework of weights --- adv_control/control.py | 24 +++++++++++++++--------- adv_control/nodes_weight.py | 12 ++++++------ adv_control/utils.py | 3 +++ 3 files changed, 24 insertions(+), 15 deletions(-) diff --git a/adv_control/control.py b/adv_control/control.py index 4eb0d18..1d81e39 100644 --- a/adv_control/control.py +++ b/adv_control/control.py @@ -8,7 +8,7 @@ import comfy.utils import comfy.model_management import comfy.model_detection import comfy.controlnet as comfy_cn -from comfy.controlnet import ControlBase, ControlNet, ControlLora, T2IAdapter +from comfy.controlnet import ControlBase, ControlNet, ControlLora, T2IAdapter, StrengthType from comfy.model_patcher import ModelPatcher from .control_sparsectrl import SparseModelPatcher, SparseControlNet, SparseCtrlMotionWrapper, SparseSettings, SparseConst @@ -21,7 +21,7 @@ from .logger import logger class ControlNetAdvanced(ControlNet, AdvancedControlBase): - def __init__(self, control_model, timestep_keyframes: TimestepKeyframeGroup, global_average_pooling=False, compression_ratio=8, latent_format=None, device=None, load_device=None, manual_cast_dtype=None): + def __init__(self, control_model, timestep_keyframes: TimestepKeyframeGroup, global_average_pooling=False, compression_ratio=8, latent_format=None, device=None, load_device=None, manual_cast_dtype=None, extra_conds=["y"], strength_type=StrengthType.CONSTANT): super().__init__(control_model=control_model, global_average_pooling=global_average_pooling, compression_ratio=compression_ratio, latent_format=latent_format, device=device, load_device=load_device, manual_cast_dtype=manual_cast_dtype) AdvancedControlBase.__init__(self, super(), timestep_keyframes=timestep_keyframes, weights_default=ControlWeights.controlnet()) @@ -64,9 +64,9 @@ class ControlNetAdvanced(ControlNet, AdvancedControlBase): actual_cond_hint_orig = self.cond_hint_original if self.cond_hint_original.size(0) < self.full_latent_length: actual_cond_hint_orig = extend_to_batch_size(tensor=actual_cond_hint_orig, batch_size=self.full_latent_length) - self.cond_hint = comfy.utils.common_upscale(actual_cond_hint_orig[self.sub_idxs], x_noisy.shape[3] * compression_ratio, x_noisy.shape[2] * compression_ratio, 'nearest-exact', "center") + self.cond_hint = comfy.utils.common_upscale(actual_cond_hint_orig[self.sub_idxs], x_noisy.shape[3] * compression_ratio, x_noisy.shape[2] * compression_ratio, self.upscale_algorithm, "center") else: - self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * compression_ratio, x_noisy.shape[2] * compression_ratio, 'nearest-exact', "center") + self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * compression_ratio, x_noisy.shape[2] * compression_ratio, self.upscale_algorithm, "center") if self.vae is not None: loaded_models = comfy.model_management.loaded_models(only_currently_used=True) self.cond_hint = self.vae.encode(self.cond_hint.movedim(1, -1)) @@ -81,17 +81,22 @@ class ControlNetAdvanced(ControlNet, AdvancedControlBase): self.prepare_mask_cond_hint(x_noisy=x_noisy, t=t, cond=cond, batched_number=batched_number, dtype=dtype) context = cond.get('crossattn_controlnet', cond['c_crossattn']) - y = cond.get('y', None) - if y is not None: - y = y.to(dtype) + extra = self.extra_args.copy() + for c in self.extra_conds: + temp = cond.get(c, None) + if temp is not None: + extra[c] = temp.to(dtype) + timestep = self.model_sampling_current.timestep(t) x_noisy = self.model_sampling_current.calculate_input(t, x_noisy) - control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.float(), context=context.to(dtype), y=y) + control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.to(dtype), context=context.to(dtype), **extra) return self.control_merge(control, control_prev, output_dtype) def copy(self): c = ControlNetAdvanced(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 @@ -99,7 +104,8 @@ class ControlNetAdvanced(ControlNet, AdvancedControlBase): @staticmethod def from_vanilla(v: ControlNet, timestep_keyframe: TimestepKeyframeGroup=None) -> 'ControlNetAdvanced': to_return = ControlNetAdvanced(control_model=v.control_model, timestep_keyframes=timestep_keyframe, - global_average_pooling=v.global_average_pooling, compression_ratio=v.compression_ratio, latent_format=v.latent_format, device=v.device, load_device=v.load_device, manual_cast_dtype=v.manual_cast_dtype) + global_average_pooling=v.global_average_pooling, compression_ratio=v.compression_ratio, latent_format=v.latent_format, device=v.device, load_device=v.load_device, + manual_cast_dtype=v.manual_cast_dtype) v.copy_to(to_return) return to_return diff --git a/adv_control/nodes_weight.py b/adv_control/nodes_weight.py index 5b9bd04..263c529 100644 --- a/adv_control/nodes_weight.py +++ b/adv_control/nodes_weight.py @@ -42,7 +42,7 @@ class ScaledSoftMaskedUniversalWeights: "optional": { "uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ), "cn_extras": ("CN_WEIGHTS_EXTRAS",), - "autosize": ("ACNAUTOSIZE", {"padding": 75}), + "autosize": ("ACNAUTOSIZE", {"padding": 0}), } } @@ -77,7 +77,7 @@ class ScaledSoftUniversalWeights: "optional": { "uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ), "cn_extras": ("CN_WEIGHTS_EXTRAS",), - "autosize": ("ACNAUTOSIZE", {"padding": 20}), + "autosize": ("ACNAUTOSIZE", {"padding": 0}), } } @@ -115,7 +115,7 @@ class SoftControlNetWeights: "optional": { "uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ), "cn_extras": ("CN_WEIGHTS_EXTRAS",), - "autosize": ("ACNAUTOSIZE", {"padding": 50}), + "autosize": ("ACNAUTOSIZE", {"padding": 0}), } } @@ -157,7 +157,7 @@ class CustomControlNetWeights: "optional": { "uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ), "cn_extras": ("CN_WEIGHTS_EXTRAS",), - "autosize": ("ACNAUTOSIZE", {"padding": 65}), + "autosize": ("ACNAUTOSIZE", {"padding": 0}), } } @@ -190,7 +190,7 @@ class SoftT2IAdapterWeights: "optional": { "uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ), "cn_extras": ("CN_WEIGHTS_EXTRAS",), - "autosize": ("ACNAUTOSIZE", {"padding": 55}), + "autosize": ("ACNAUTOSIZE", {"padding": 0}), } } @@ -223,7 +223,7 @@ class CustomT2IAdapterWeights: "optional": { "uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ), "cn_extras": ("CN_WEIGHTS_EXTRAS",), - "autosize": ("ACNAUTOSIZE", {"padding": 65}), + "autosize": ("ACNAUTOSIZE", {"padding": 0}), } } diff --git a/adv_control/utils.py b/adv_control/utils.py index b71e371..df7a67e 100644 --- a/adv_control/utils.py +++ b/adv_control/utils.py @@ -175,6 +175,9 @@ class ControlWeights: def get(self, idx: int, control: dict[str, list[Tensor]], key: str, default=1.0) -> Union[float, Tensor]: # 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"])