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

@@ -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)

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

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:

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 = {

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)

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