Merge PR #145 from Kosinkadink/develop - ContextRef support for ADE

Added ContextRef support for ADE + ReferenceCN Refactor

Author: Kosinkadink

adv_control/control.py

@@ -16,13 +16,10 @@
 from .control_svd import svd_unet_config_from_diffusers_unet, SVDControlNet, svd_unet_to_diffusers
 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)
+                    broadcast_image_to_extend, extend_to_batch_size, ORIG_PREVIOUS_CONTROLNET, CONTROL_INIT_BY_ACN)
 from .logger import logger
 
 
-ORIG_PREVIOUS_CONTROLNET = "_orig_previous_controlnet"
-
-
 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):
         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)
@@ -574,7 +571,13 @@ def restore_all_controlnet_conns(conds: list[list[dict[str]]]):
         if main_cond is not None:
             for cond in main_cond:
                 if "control" in cond[1]:
-                    _restore_all_controlnet_conns(cond[1]["control"])
+                    # if ACN is the one to have initialized it, delete it
+                    # TODO: maybe check if someone else did a similar hack, and carefully pluck out our stuff?
+                    if CONTROL_INIT_BY_ACN in cond[1]:
+                        cond[1].pop("control")
+                        cond[1].pop(CONTROL_INIT_BY_ACN)
+                    else:
+                        _restore_all_controlnet_conns(cond[1]["control"])
 
 
 def _restore_all_controlnet_conns(input_object: ControlBase):

adv_control/control_reference.py

@@ -231,10 +231,10 @@ def apply_controlnet(self, positive, negative, control_net, image, strength, sta
 NODE_DISPLAY_NAME_MAPPINGS = {
     # Keyframes
     "TimestepKeyframe": "Timestep Keyframe 🛂🅐🅒🅝",
-    "ACN_TimestepKeyframeInterpolation": "Timestep Keyframe Interpolation 🛂🅐🅒🅝",
+    "ACN_TimestepKeyframeInterpolation": "Timestep Keyframe Interp. 🛂🅐🅒🅝",
     "ACN_TimestepKeyframeFromStrengthList": "Timestep Keyframe From List 🛂🅐🅒🅝",
     "LatentKeyframe": "Latent Keyframe 🛂🅐🅒🅝",
-    "LatentKeyframeTiming": "Latent Keyframe Interpolation 🛂🅐🅒🅝",
+    "LatentKeyframeTiming": "Latent Keyframe Interp. 🛂🅐🅒🅝",
     "LatentKeyframeBatchedGroup": "Latent Keyframe From List 🛂🅐🅒🅝",
     "LatentKeyframeGroup": "Latent Keyframe Group 🛂🅐🅒🅝",
     # Conditioning

adv_control/nodes.py

@@ -82,7 +82,7 @@ def INPUT_TYPES(s):
                 "inherit_missing": ("BOOLEAN", {"default": True},),
                 "mask_optional": ("MASK", ),
                 "print_keyframes": ("BOOLEAN", {"default": False}),
-                "autosize": ("ACNAUTOSIZE", {"padding": 70}),
+                "autosize": ("ACNAUTOSIZE", {"padding": 50}),
             }
         }
 
@@ -355,7 +355,7 @@ def INPUT_TYPES(s):
             "optional": {
                 "prev_latent_kf": ("LATENT_KEYFRAME", ),
                 "print_keyframes": ("BOOLEAN", {"default": False}),
-                "autosize": ("ACNAUTOSIZE", {"padding": 90}),
+                "autosize": ("ACNAUTOSIZE", {"padding": 50}),
             }
         }
 

adv_control/nodes_keyframes.py

@@ -11,30 +11,44 @@
                                 RefBasicTransformerBlock, RefTimestepEmbedSequential,
                                 InjectionBasicTransformerBlockHolder, InjectionTimestepEmbedSequentialHolder,
                                 _forward_inject_BasicTransformerBlock, factory_forward_inject_UNetModel,
-                                REF_CONTROL_LIST_ALL)
+                                handle_context_ref_setup,
+                                REF_CONTROL_LIST_ALL, CONTEXTREF_CLEAN_FUNC)
 from .control_lllite import (ControlLLLiteAdvanced)
 from .utils import torch_dfs
 
 
 def support_sliding_context_windows(model, positive, negative) -> tuple[bool, dict, dict]:
-    if not hasattr(model, "motion_injection_params"):
-        return False, positive, negative
-    motion_injection_params = getattr(model, "motion_injection_params")
-    context_options = getattr(motion_injection_params, "context_options")
-    if context_options.context_length is None:
-        return False, positive, negative
     # convert to advanced, with report if anything was actually modified
     modified, new_conds = convert_all_to_advanced([positive, negative])
     positive, negative = new_conds
     return modified, positive, negative
 
 
+def has_sliding_context_windows(model):
+    motion_injection_params = getattr(model, "motion_injection_params", None)
+    if motion_injection_params is None:
+        return False
+    context_options = getattr(motion_injection_params, "context_options")
+    return context_options.context_length is not None
+
+
+def get_contextref_obj(model):
+    motion_injection_params = getattr(model, "motion_injection_params", None)
+    if motion_injection_params is None:
+        return None
+    context_options = getattr(motion_injection_params, "context_options")
+    extras = getattr(context_options, "extras", None)
+    if extras is None:
+        return None
+    return getattr(extras, "context_ref", None)
+
+
 def acn_sample_factory(orig_comfy_sample: Callable, is_custom=False) -> Callable:
     def get_refcn(control: ControlBase, order: int=-1):
         ref_set: set[ReferenceAdvanced] = set()
         if control is None:
             return ref_set
-        if type(control) == ReferenceAdvanced:
+        if type(control) == ReferenceAdvanced and not control.is_context_ref:
             control.order = order
             order -= 1
             ref_set.add(control)
@@ -59,13 +73,23 @@ def acn_sample(model: ModelPatcher, *args, **kwargs):
             # check if positive or negative conds contain ref cn
             positive = args[-3]
             negative = args[-2]
-            # if context options present, convert all CNs to Advanced if needed
-            controlnets_modified, positive, negative = support_sliding_context_windows(model, positive, negative)
-            if controlnets_modified:
-                args = list(args)
-                args[-3] = positive
-                args[-2] = negative
-                args = tuple(args)
+            # if context options present, perform some special actions that may be required
+            context_refs = []
+            if has_sliding_context_windows(model):
+                model.model_options = model.model_options.copy()
+                model.model_options["transformer_options"] = model.model_options["transformer_options"].copy()
+                # convert all CNs to Advanced if needed
+                controlnets_modified, positive, negative = support_sliding_context_windows(model, positive, negative)
+                if controlnets_modified:
+                    args = list(args)
+                    args[-3] = positive
+                    args[-2] = negative
+                    args = tuple(args)
+                # enable ContextRef, if requested
+                existing_contextref_obj = get_contextref_obj(model)
+                if existing_contextref_obj is not None:
+                    context_refs = handle_context_ref_setup(existing_contextref_obj, model.model_options["transformer_options"], positive, negative)
+                    controlnets_modified = True
             # look for Advanced ControlNets that will require intervention to work
             ref_set = set()
             lllite_dict: dict[ControlLLLiteAdvanced, None] = {} # dicts preserve insertion order since py3.7
@@ -88,7 +112,7 @@ def acn_sample(model: ModelPatcher, *args, **kwargs):
                 for lll in lllite_list:
                     lll.live_model_patches(model.model_options)
             # if no ref cn found, do original function immediately
-            if len(ref_set) == 0:
+            if len(ref_set) == 0 and len(context_refs) == 0:
                 return orig_comfy_sample(model, *args, **kwargs)
             # otherwise, injection time
             try:
@@ -158,6 +182,7 @@ def acn_sample(model: ModelPatcher, *args, **kwargs):
                 new_model_options["transformer_options"] = model.model_options["transformer_options"].copy()
                 ref_list: list[ReferenceAdvanced] = list(ref_set)
                 new_model_options["transformer_options"][REF_CONTROL_LIST_ALL] = sorted(ref_list, key=lambda x: x.order)
+                new_model_options["transformer_options"][CONTEXTREF_CLEAN_FUNC] = reference_injections.clean_contextref_module_mem
                 model.model_options = new_model_options
                 # continue with original function
                 return orig_comfy_sample(model, *args, **kwargs)
@@ -167,14 +192,14 @@ def acn_sample(model: ModelPatcher, *args, **kwargs):
                 attn_modules: list[RefBasicTransformerBlock] = reference_injections.attn_modules
                 for module in attn_modules:
                     module.injection_holder.restore(module)
-                    module.injection_holder.clean()
+                    module.injection_holder.clean_all()
                     del module.injection_holder
                 del attn_modules
                 # restore gn modules
                 gn_modules: list[RefTimestepEmbedSequential] = reference_injections.gn_modules
                 for module in gn_modules:
                     module.injection_holder.restore(module)
-                    module.injection_holder.clean()
+                    module.injection_holder.clean_all()
                     del module.injection_holder
                 del gn_modules
                 # restore diffusion_model forward function

adv_control/sampling.py

@@ -21,6 +21,10 @@
 BIGMIN = -(2**53-1)
 BIGMAX = (2**53-1)
 
+ORIG_PREVIOUS_CONTROLNET = "_orig_previous_controlnet"
+CONTROL_INIT_BY_ACN = "_control_init_by_ACN"
+
+
 def load_torch_file_with_dict_factory(controlnet_data: dict[str, Tensor], orig_load_torch_file: Callable):
     def load_torch_file_with_dict(*args, **kwargs):
         # immediately restore load_torch_file to original version
@@ -35,7 +39,7 @@ 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 (cond_or_uncond[0] in [0, 1]):
+        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)
@@ -569,7 +573,8 @@ def __init__(self, base: ControlBase, timestep_keyframes: TimestepKeyframeGroup,
         self.full_latent_length = 0
         self.context_length = 0
         # timesteps
-        self.t: Tensor = None
+        self.t: float = None
+        self.prev_t: float = None
         self.batched_number: Union[int, IntWithCondOrUncond] = None
         self.batch_size: int = 0
         # weights + override
@@ -627,10 +632,11 @@ def set_timestep_keyframes(self, timestep_keyframes: TimestepKeyframeGroup):
         self.weights = None
         self.latent_keyframes = None
 
-    def prepare_current_timestep(self, t: Tensor, batched_number: int):
+    def prepare_current_timestep(self, t: Tensor, batched_number: int=1):
         self.t = float(t[0])
-        self.batched_number = batched_number
-        self.batch_size = len(t)
+        # check if t has changed (otherwise do nothing, as step already accounted for)
+        if self.t == self.prev_t:
+            return
         # get current step percent
         curr_t: float = self.t
         prev_index = self._current_timestep_index
@@ -666,7 +672,8 @@ def prepare_current_timestep(self, t: Tensor, batched_number: int):
                     # if eval_tk is outside of percent range, stop looking further
                     else:
                         break
-        
+        # update prev_t
+        self.prev_t = self.t
         # update steps current keyframe is used
         self._current_used_steps += 1
         # if index changed, apply overrides
@@ -740,6 +747,8 @@ def should_run(self):
         return True
 
     def get_control_inject(self, x_noisy, t, cond, batched_number):
+        self.batched_number = batched_number
+        self.batch_size = len(t)
         # 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
@@ -932,6 +941,7 @@ def cleanup_advanced(self):
         self.full_latent_length = 0
         self.context_length = 0
         self.t = None
+        self.prev_t = None
         self.batched_number = None
         self.batch_size = 0
         self.weights = None

adv_control/utils.py

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-advanced-controlnet"
 description = "Nodes for scheduling ControlNet strength across timesteps and batched latents, as well as applying custom weights and attention masks."
-version = "1.1.5"
+version = "1.2.0"
 license = { file = "LICENSE" }
 dependencies = []