From e800d52acf3db94953283f7eac08f34af49b4dc0 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Thu, 27 Jun 2024 05:03:10 -0500 Subject: [PATCH] Made all controlnets works with upcoming ComfyUI changes, started work on making latent preprocs obsolete, still need to make control weights apply the same way as before --- adv_control/control.py | 53 ++++++++++++++------------ adv_control/control_reference.py | 16 ++++---- adv_control/control_sparsectrl.py | 45 +++++----------------- adv_control/control_svd.py | 9 +++-- adv_control/nodes.py | 20 ++++++++-- adv_control/utils.py | 63 ++++++++++++++++++------------- 6 files changed, 104 insertions(+), 102 deletions(-) diff --git a/adv_control/control.py b/adv_control/control.py index d640d5d..66cd0e7 100644 --- a/adv_control/control.py +++ b/adv_control/control.py @@ -11,10 +11,10 @@ import comfy.controlnet as comfy_cn from comfy.controlnet import ControlBase, ControlNet, ControlLora, T2IAdapter from comfy.model_patcher import ModelPatcher -from .control_sparsectrl import SparseModelPatcher, SparseControlNet, SparseCtrlMotionWrapper, SparseMethod, SparseSettings, SparseSpreadMethod, PreprocSparseRGBWrapper, SparseConst +from .control_sparsectrl import SparseModelPatcher, SparseControlNet, SparseCtrlMotionWrapper, SparseSettings, SparseConst from .control_lllite import LLLiteModule, LLLitePatch from .control_svd import svd_unet_config_from_diffusers_unet, SVDControlNet, svd_unet_to_diffusers -from .utils import (AdvancedControlBase, TimestepKeyframeGroup, LatentKeyframeGroup, ControlWeightType, ControlWeights, WeightTypeException, +from .utils import (AdvancedControlBase, TimestepKeyframeGroup, LatentKeyframeGroup, AbstractPreprocWrapper, ControlWeightType, ControlWeights, WeightTypeException, manual_cast_clean_groupnorm, disable_weight_init_clean_groupnorm, prepare_mask_batch, get_properly_arranged_t2i_weights, load_torch_file_with_dict_factory, broadcast_image_to_extend, extend_to_batch_size) from .logger import logger @@ -81,7 +81,7 @@ class ControlNetAdvanced(ControlNet, AdvancedControlBase): 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) - return self.control_merge(None, control, control_prev, output_dtype) + 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) @@ -100,22 +100,17 @@ class T2IAdapterAdvanced(T2IAdapter, AdvancedControlBase): super().__init__(t2i_model=t2i_model, channels_in=channels_in, compression_ratio=compression_ratio, upscale_algorithm=upscale_algorithm, device=device) AdvancedControlBase.__init__(self, super(), timestep_keyframes=timestep_keyframes, weights_default=ControlWeights.t2iadapter()) - def control_merge_inject(self, control_input, control_output, control_prev, output_dtype): + def control_merge_inject(self, control: dict[str, list[Tensor]], control_prev, output_dtype): # if has uncond multiplier, need to make sure control shapes are the same batch size as expected if self.weights.has_uncond_multiplier or self.weights.has_uncond_mask: - if control_input is not None: - for i in range(len(control_input)): - x = control_input[i] + for key in control: + control_current = control[key] + for i in range(len(control_current)): + x = control_current[i] if x is not None: if x.size(0) < self.batch_size: - control_input[i] = x.repeat(self.batched_number, 1, 1, 1)[:self.batch_size] - if control_output is not None: - for i in range(len(control_output)): - x = control_output[i] - if x is not None: - if x.size(0) < self.batch_size: - control_output[i] = x.repeat(self.batched_number, 1, 1, 1)[:self.batch_size] - return AdvancedControlBase.control_merge_inject(self, control_input, control_output, control_prev, output_dtype) + control_current[i] = x.repeat(self.batched_number, 1, 1, 1)[:self.batch_size] + 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)] @@ -199,8 +194,8 @@ class SVDControlNetAdvanced(ControlNetAdvanced): def __init__(self, control_model: SVDControlNet, timestep_keyframes: TimestepKeyframeGroup, global_average_pooling=False, device=None, load_device=None, manual_cast_dtype=None): super().__init__(control_model=control_model, timestep_keyframes=timestep_keyframes, global_average_pooling=global_average_pooling, device=device, load_device=load_device, manual_cast_dtype=manual_cast_dtype) - def set_cond_hint(self, *args, **kwargs): - to_return = super().set_cond_hint(*args, **kwargs) + def set_cond_hint_inject(self, *args, **kwargs): + to_return = super().set_cond_hint_inject(*args, **kwargs) # cond hint for SVD-ControlNet needs to be scaled between (-1, 1) instead of (0, 1) self.cond_hint_original = self.cond_hint_original * 2.0 - 1.0 return to_return @@ -254,7 +249,7 @@ class SVDControlNetAdvanced(ControlNetAdvanced): x_noisy = torch.cat([x_noisy] + [cond['c_concat']], dim=1) control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.float(), context=context.to(dtype), y=y, cond=cond) - return self.control_merge(None, control, control_prev, output_dtype) + return self.control_merge(control, control_prev, output_dtype) def copy(self): c = SVDControlNetAdvanced(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) @@ -269,6 +264,9 @@ class SparseCtrlAdvanced(ControlNetAdvanced): self.control_model_wrapped = SparseModelPatcher(self.control_model, load_device=load_device, offload_device=comfy.model_management.unet_offload_device()) self.add_compatible_weight(ControlWeightType.SPARSECTRL) self.control_model: SparseControlNet = self.control_model # does nothing except help with IDE hints + if self.control_model.use_simplified_conditioning_embedding: + self.require_vae = True + self.allow_condhint_latents = True self.sparse_settings = sparse_settings if sparse_settings is not None else SparseSettings.default() self.model_latent_format = None # latent format for active SD model, NOT controlnet self.preprocessed = False @@ -366,7 +364,7 @@ class SparseCtrlAdvanced(ControlNetAdvanced): 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) - return self.control_merge(None, control, control_prev, output_dtype) + return self.control_merge(control, control_prev, output_dtype) def apply_advanced_strengths_and_masks(self, x: Tensor, batched_number: int): # apply mults to indexes with and without a direct condhint @@ -376,7 +374,7 @@ class SparseCtrlAdvanced(ControlNetAdvanced): def pre_run_advanced(self, model, percent_to_timestep_function): super().pre_run_advanced(model, percent_to_timestep_function) - if type(self.cond_hint_original) == PreprocSparseRGBWrapper: + if isinstance(self.cond_hint_original, AbstractPreprocWrapper): if not self.control_model.use_simplified_conditioning_embedding: raise ValueError("Any model besides RGB SparseCtrl should NOT have its images go through the RGB SparseCtrl preprocessor.") self.cond_hint_original = self.cond_hint_original.condhint @@ -415,8 +413,8 @@ class ControlLLLiteAdvanced(ControlBase, AdvancedControlBase): model.set_model_attn1_patch(self.patch_attn1) model.set_model_attn2_patch(self.patch_attn2) - def set_cond_hint(self, *args, **kwargs): - to_return = super().set_cond_hint(*args, **kwargs) + 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) self.cond_hint_original = self.cond_hint_original * 2.0 - 1.0 return to_return @@ -543,8 +541,6 @@ def load_controlnet(ckpt_path, timestep_keyframe: TimestepKeyframeGroup=None, mo control = load_sparsectrl(ckpt_path, controlnet_data=controlnet_data, timestep_keyframe=timestep_keyframe, model=model) elif controlnet_type == ControlWeightType.SVD_CONTROLNET: control = load_svdcontrolnet(ckpt_path, controlnet_data=controlnet_data, timestep_keyframe=timestep_keyframe) - #raise Exception(f"SVD-ControlNet is not supported yet!") - #control = comfy_cn.load_controlnet(ckpt_path, model=model) # otherwise, load vanilla ControlNet else: try: @@ -563,7 +559,10 @@ def convert_to_advanced(control, timestep_keyframe: TimestepKeyframeGroup=None): return control # if exactly ControlNet returned, transform it into ControlNetAdvanced if type(control) == ControlNet: - return ControlNetAdvanced.from_vanilla(v=control, timestep_keyframe=timestep_keyframe) + control = ControlNetAdvanced.from_vanilla(v=control, timestep_keyframe=timestep_keyframe) + if is_sd3_advanced_controlnet(control): + control.require_vae = True + return control # if exactly ControlLora returned, transform it into ControlLoraAdvanced elif type(control) == ControlLora: return ControlLoraAdvanced.from_vanilla(v=control, timestep_keyframe=timestep_keyframe) @@ -578,6 +577,10 @@ def is_advanced_controlnet(input_object): return hasattr(input_object, "sub_idxs") +def is_sd3_advanced_controlnet(input_object: ControlNetAdvanced): + return type(input_object) == ControlNetAdvanced and input_object.latent_format is not None + + def load_sparsectrl(ckpt_path: str, controlnet_data: dict[str, Tensor]=None, timestep_keyframe: TimestepKeyframeGroup=None, sparse_settings=SparseSettings.default(), model=None) -> SparseCtrlAdvanced: if controlnet_data is None: controlnet_data = comfy.utils.load_torch_file(ckpt_path, safe_load=True) diff --git a/adv_control/control_reference.py b/adv_control/control_reference.py index 24f1485..3561a16 100644 --- a/adv_control/control_reference.py +++ b/adv_control/control_reference.py @@ -218,7 +218,7 @@ class ReferenceOptions: class ReferencePreprocWrapper(AbstractPreprocWrapper): - error_msg = error_msg = "Invalid use of Reference Preprocess output. The output of RGB SparseCtrl preprocessor is NOT a usual image, but a latent pretending to be an image - you must connect the output directly to an Apply Advanced ControlNet node. It cannot be used for anything else that accepts IMAGE input." + error_msg = error_msg = "Invalid use of Reference Preprocess output. The output of Reference preprocessor is NOT a usual image, but a latent pretending to be an image - you must connect the output directly to an Apply Advanced ControlNet node. It cannot be used for anything else that accepts IMAGE input." def __init__(self, condhint: Tensor): super().__init__(condhint) @@ -228,10 +228,10 @@ class ReferenceAdvanced(ControlBase, AdvancedControlBase): def __init__(self, ref_opts: ReferenceOptions, timestep_keyframes: TimestepKeyframeGroup, device=None): super().__init__(device) - AdvancedControlBase.__init__(self, super(), timestep_keyframes=timestep_keyframes, weights_default=ControlWeights.controllllite()) + AdvancedControlBase.__init__(self, super(), timestep_keyframes=timestep_keyframes, weights_default=ControlWeights.controllllite(), allow_condhint_latents=True, require_vae=True) self.ref_opts = ref_opts self.order = 0 - self.latent_format = None + self.model_latent_format = None self.model_sampling_current = None self.should_apply_attn_effective_strength = False self.should_apply_adain_effective_strength = False @@ -288,9 +288,9 @@ class ReferenceAdvanced(ControlBase, AdvancedControlBase): def pre_run_advanced(self, model, percent_to_timestep_function): AdvancedControlBase.pre_run_advanced(self, model, percent_to_timestep_function) - if type(self.cond_hint_original) == ReferencePreprocWrapper: + if isinstance(self.cond_hint_original, AbstractPreprocWrapper): self.cond_hint_original = self.cond_hint_original.condhint - self.latent_format = model.latent_format # LatentFormat object, used to process_in latent cond_hint + self.model_latent_format = model.latent_format # LatentFormat object, used to process_in latent cond_hint self.model_sampling_current = model.model_sampling # SDXL is more sensitive to style_fidelity according to sd-webui-controlnet comments if type(model).__name__ == "SDXL": @@ -328,7 +328,7 @@ class ReferenceAdvanced(ControlBase, AdvancedControlBase): if x_noisy.shape[0] != self.cond_hint.shape[0]: self.cond_hint = broadcast_image_to_extend(self.cond_hint, x_noisy.shape[0], batched_number, except_one=False) # noise cond_hint based on sigma (current step) - self.cond_hint = self.latent_format.process_in(self.cond_hint) + self.cond_hint = self.model_latent_format.process_in(self.cond_hint) self.cond_hint = ref_noise_latents(self.cond_hint, sigma=t, noise=None) timestep = self.model_sampling_current.timestep(t) self.should_apply_attn_effective_strength = not (math.isclose(self.strength, 1.0) and math.isclose(self._current_timestep_keyframe.strength, 1.0) and math.isclose(self.ref_opts.attn_strength, 1.0)) @@ -343,8 +343,8 @@ class ReferenceAdvanced(ControlBase, AdvancedControlBase): def cleanup_advanced(self): super().cleanup_advanced() - del self.latent_format - self.latent_format = None + del self.model_latent_format + self.model_latent_format = None del self.model_sampling_current self.model_sampling_current = None self.should_apply_attn_effective_strength = False diff --git a/adv_control/control_sparsectrl.py b/adv_control/control_sparsectrl.py index 9532d47..c6e3e98 100644 --- a/adv_control/control_sparsectrl.py +++ b/adv_control/control_sparsectrl.py @@ -23,13 +23,13 @@ from comfy.cldm.cldm import ControlNet as ControlNetCLDM from comfy.ldm.modules.attention import SpatialTransformer from comfy.ldm.modules.attention import attention_basic, attention_pytorch, attention_split, attention_sub_quad, default from comfy.ldm.modules.attention import FeedForward, SpatialTransformer -from comfy.ldm.modules.diffusionmodules.openaimodel import TimestepEmbedSequential, ResBlock, Downsample +from comfy.ldm.modules.diffusionmodules.openaimodel import TimestepEmbedSequential from comfy.model_patcher import ModelPatcher import comfy.ops import comfy.model_management from .logger import logger -from .utils import (BIGMAX, TimestepKeyframeGroup, disable_weight_init_clean_groupnorm, +from .utils import (BIGMAX, AbstractPreprocWrapper, disable_weight_init_clean_groupnorm, prepare_mask_batch, broadcast_image_to_extend, extend_to_batch_size) @@ -85,7 +85,8 @@ class SparseControlNet(ControlNetCLDM): x = torch.zeros_like(x) guided_hint = self.input_hint_block(hint, emb, context) - outs = [] + out_output = [] + out_middle = [] hs = [] if self.num_classes is not None: @@ -100,12 +101,12 @@ class SparseControlNet(ControlNetCLDM): guided_hint = None else: h = module(h, emb, context) - outs.append(zero_conv(h, emb, context)) + out_output.append(zero_conv(h, emb, context)) h = self.middle_block(h, emb, context) - outs.append(self.middle_block_out(h, emb, context)) + out_middle.append(self.middle_block_out(h, emb, context)) - return outs + return {"middle": out_middle, "output": out_output} class SparseModelPatcher(ModelPatcher): @@ -154,36 +155,10 @@ class SparseModelPatcher(ModelPatcher): self.object_patches_backup = n.object_patches_backup -class PreprocSparseRGBWrapper: - error_msg = "Invalid use of RGB SparseCtrl output. The output of RGB SparseCtrl preprocessor is NOT a usual image, but a latent pretending to be an image - you must connect the output directly to an Apply ControlNet node (advanced or otherwise). It cannot be used for anything else that accepts IMAGE input." +class PreprocSparseRGBWrapper(AbstractPreprocWrapper): + error_msg = error_msg = "Invalid use of RGB SparseCtrl output. The output of RGB SparseCtrl preprocessor is NOT a usual image, but a latent pretending to be an image - you must connect the output directly to an Apply ControlNet node (advanced or otherwise). It cannot be used for anything else that accepts IMAGE input." def __init__(self, condhint: Tensor): - self.condhint = condhint - - def movedim(self, *args, **kwargs): - return self - - def __getattr__(self, *args, **kwargs): - raise AttributeError(self.error_msg) - - def __setattr__(self, name, value): - if name != "condhint": - raise AttributeError(self.error_msg) - super().__setattr__(name, value) - - def __iter__(self, *args, **kwargs): - raise AttributeError(self.error_msg) - - def __next__(self, *args, **kwargs): - raise AttributeError(self.error_msg) - - def __len__(self, *args, **kwargs): - raise AttributeError(self.error_msg) - - def __getitem__(self, *args, **kwargs): - raise AttributeError(self.error_msg) - - def __setitem__(self, *args, **kwargs): - raise AttributeError(self.error_msg) + super().__init__(condhint) class SparseContextAware: diff --git a/adv_control/control_svd.py b/adv_control/control_svd.py index 9458f6a..7c094fb 100644 --- a/adv_control/control_svd.py +++ b/adv_control/control_svd.py @@ -311,7 +311,8 @@ class SVDControlNet(nn.Module): guided_hint = self.input_hint_block(hint, emb, context, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator) - outs = [] + out_output = [] + out_middle = [] hs = [] if self.num_classes is not None: @@ -326,12 +327,12 @@ class SVDControlNet(nn.Module): guided_hint = None else: h = module(h, emb, context, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator) - outs.append(zero_conv(h, emb, context, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)) + out_output.append(zero_conv(h, emb, context, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)) h = self.middle_block(h, emb, context, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator) - outs.append(self.middle_block_out(h, emb, context, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)) + out_middle.append(self.middle_block_out(h, emb, context, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)) - return outs + return {"middle": out_middle, "output": out_output} TEMPORAL_TRANSFORMER_BLOCKS = { diff --git a/adv_control/nodes.py b/adv_control/nodes.py index 14659c7..66e7368 100644 --- a/adv_control/nodes.py +++ b/adv_control/nodes.py @@ -4,8 +4,8 @@ from torch import Tensor import folder_paths from comfy.model_patcher import ModelPatcher -from .control import load_controlnet, convert_to_advanced, is_advanced_controlnet -from .utils import ControlWeights, LatentKeyframeGroup, TimestepKeyframeGroup, BIGMAX +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_keyframes import (LatentKeyframeGroupNode, LatentKeyframeInterpolationNode, LatentKeyframeBatchedGroupNode, LatentKeyframeNode, @@ -89,6 +89,7 @@ class AdvancedControlNetApply: "latent_kf_override": ("LATENT_KEYFRAME", ), "weights_override": ("CONTROL_NET_WEIGHTS", ), "model_optional": ("MODEL",), + "vae_optional": ("VAE",), } } @@ -99,7 +100,7 @@ class AdvancedControlNetApply: CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝" def apply_controlnet(self, positive, negative, control_net, image, strength, start_percent, end_percent, - mask_optional: Tensor=None, model_optional: ModelPatcher=None, + mask_optional: Tensor=None, model_optional: ModelPatcher=None, vae_optional=None, timestep_kf: TimestepKeyframeGroup=None, latent_kf_override: LatentKeyframeGroup=None, weights_override: ControlWeights=None): if strength == 0: @@ -121,7 +122,7 @@ class AdvancedControlNetApply: c_net = cnets[prev_cnet] else: # copy, convert to advanced if needed, and set cond - c_net = convert_to_advanced(control_net.copy()).set_cond_hint(control_hint, strength, (start_percent, end_percent)) + c_net = convert_to_advanced(control_net.copy()).set_cond_hint(control_hint, strength, (start_percent, end_percent), vae_optional) if is_advanced_controlnet(c_net): # disarm node check c_net.disarm() @@ -130,6 +131,17 @@ class AdvancedControlNetApply: if not model_optional: raise Exception(f"Type '{type(c_net).__name__}' requires model_optional input, but got None.") c_net.patch_model(model=model_optional) + # if vae required, verify vae is passed in + if c_net.require_vae: + # if controlnet can accept preprocced condhint latents and is the case, ignore vae requirement + if c_net.allow_condhint_latents and isinstance(control_hint, AbstractPreprocWrapper): + pass + elif not vae_optional: + # make sure SD3 ControlNet will get a special message instead of generic type mention + if is_sd3_advanced_controlnet: + raise Exception(f"SD3 ControlNet requires vae_optional input, but got None.") + else: + raise Exception(f"Type '{type(c_net).__name__}' requires vae_optional input, but got None.") # apply optional parameters and overrides, if provided if timestep_kf is not None: c_net.set_timestep_keyframes(timestep_kf) diff --git a/adv_control/utils.py b/adv_control/utils.py index a92998e..ec64584 100644 --- a/adv_control/utils.py +++ b/adv_control/utils.py @@ -14,6 +14,7 @@ import comfy.model_base from comfy.controlnet import ControlBase from comfy.model_patcher import ModelPatcher +from comfy.sd import VAE from .logger import logger @@ -531,7 +532,7 @@ class WeightTypeException(TypeError): class AdvancedControlBase: - def __init__(self, base: ControlBase, timestep_keyframes: TimestepKeyframeGroup, weights_default: ControlWeights, require_model=False): + def __init__(self, base: ControlBase, timestep_keyframes: TimestepKeyframeGroup, weights_default: ControlWeights, require_model=False, require_vae=False, allow_condhint_latents=False): self.base = base self.compatible_weights = [ControlWeightType.UNIVERSAL, ControlWeightType.DEFAULT] self.add_compatible_weight(weights_default.weight_type) @@ -564,8 +565,13 @@ class AdvancedControlBase: self.pre_run = self.pre_run_inject self.cleanup = self.cleanup_inject self.set_previous_controlnet = self.set_previous_controlnet_inject - # require model to be passed into Apply Advanced ControlNet 🛂🅐🅒🅝 node + self.set_cond_hint = self.set_cond_hint_inject + # vae to store + self.adv_vae = None + # require model/vae to be passed into Apply Advanced ControlNet 🛂🅐🅒🅝 node self.require_model = require_model + self.require_vae = require_vae + self.allow_condhint_latents = allow_condhint_latents # disarm - when set to False, used to force usage of Apply Advanced ControlNet 🛂🅐🅒🅝 node (which will set it to True) self.disarmed = not require_model @@ -668,6 +674,21 @@ class AdvancedControlBase: self.mask_cond_hint_original = mask_hint return self + def set_cond_hint_inject(self, *args, **kwargs): + to_return = self.base.set_cond_hint(*args, **kwargs) + # if vae required, look in args and kwargs for it + if self.require_vae is not None: + # check args first, as that's the default way vae param is used in ComfyUI + for arg in args: + if isinstance(arg, VAE): + self.adv_vae = arg + break + # if not in args, check kwargs now + if self.adv_vae is None: + if 'vae' in kwargs: + self.adv_vae = kwargs['vae'] + return to_return + def pre_run_inject(self, model, percent_to_timestep_function): self.base.pre_run(model, percent_to_timestep_function) self.pre_run_advanced(model, percent_to_timestep_function) @@ -798,41 +819,28 @@ class AdvancedControlBase: if self._current_timestep_keyframe.strength != 1.0: x[:] *= self._current_timestep_keyframe.strength - def control_merge_inject(self: 'AdvancedControlBase', control_input, control_output, control_prev, output_dtype): + def control_merge_inject(self: 'AdvancedControlBase', control: dict[str, list[Tensor]], control_prev: dict, output_dtype): out = {'input':[], 'middle':[], 'output': []} - if control_input is not None: - for i in range(len(control_input)): - key = 'input' - x = control_input[i] - if x is not None: - self.apply_advanced_strengths_and_masks(x, self.batched_number) - - x *= self.strength * self.calc_weight(i, x, len(control_input)) - if x.dtype != output_dtype: - x = x.to(output_dtype) - out[key].insert(0, x) - - if control_output is not None: + for key in control: + control_output = control[key] + applied_to = set() for i in range(len(control_output)): - if i == (len(control_output) - 1): - key = 'middle' - index = 0 - else: - key = 'output' - index = i x = control_output[i] if x is not None: - self.apply_advanced_strengths_and_masks(x, self.batched_number) - if self.global_average_pooling: x = torch.mean(x, dim=(2, 3), keepdim=True).repeat(1, 1, x.shape[2], x.shape[3]) - x *= self.strength * self.calc_weight(i, x, len(control_output)) + if x not in applied_to: #memory saving strategy, allow shared tensors and only apply strength to shared tensors once + applied_to.add(x) + self.apply_advanced_strengths_and_masks(x, self.batched_number) + x *= self.strength * self.calc_weight(i, x, len(control_output)) + if x.dtype != output_dtype: x = x.to(output_dtype) out[key].append(x) + if control_prev is not None: for x in ['input', 'middle', 'output']: o = out[x] @@ -847,7 +855,7 @@ class AdvancedControlBase: if o[i].shape[0] < prev_val.shape[0]: o[i] = prev_val + o[i] else: - o[i] += prev_val + o[i] = prev_val + o[i] # TODO from base ComfyUI: change back to inplace add if shared tensors stop being an issue return out def prepare_mask_cond_hint(self, x_noisy: Tensor, t, cond, batched_number, dtype=None, direct_attn=False): @@ -924,4 +932,7 @@ class AdvancedControlBase: copied.mask_cond_hint_original = self.mask_cond_hint_original copied.weights_override = self.weights_override copied.latent_keyframe_override = self.latent_keyframe_override + copied.adv_vae = self.adv_vae + copied.require_vae = self.require_vae + copied.allow_condhint_latents = self.allow_condhint_latents copied.disarmed = self.disarmed