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

Added ContextRef support for ADE + ReferenceCN Refactor

adv_control/control.py

@@ -16,13 +16,10 @@ from .control_lllite import LLLiteModule, LLLitePatch, load_controllllite
 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

@@ -4,6 +4,7 @@ import math
 import torch
 from torch import Tensor
 
+import comfy.model_management
 import comfy.sample
 import comfy.model_patcher
 import comfy.utils
@@ -13,10 +14,15 @@ from comfy.ldm.modules.attention import BasicTransformerBlock
 from comfy.ldm.modules.diffusionmodules import openaimodel
 
 from .logger import logger
-from .utils import (AdvancedControlBase, ControlWeights, TimestepKeyframeGroup, AbstractPreprocWrapper,
-                    broadcast_image_to_extend)
+from .utils import (AdvancedControlBase, ControlWeights, TimestepKeyframeGroup, TimestepKeyframe, AbstractPreprocWrapper,
+                    broadcast_image_to_extend, ORIG_PREVIOUS_CONTROLNET, CONTROL_INIT_BY_ACN)
 
 
+REF_READ_ATTN_CONTROL_LIST = "ref_read_attn_control_list"
+REF_WRITE_ATTN_CONTROL_LIST = "ref_write_attn_control_list"
+REF_READ_ADAIN_CONTROL_LIST = "ref_read_adain_control_list"
+REF_WRITE_ADAIN_CONTROL_LIST = "ref_write_adain_control_list"
+
 REF_ATTN_CONTROL_LIST = "ref_attn_control_list"
 REF_ADAIN_CONTROL_LIST = "ref_adain_control_list"
 REF_CONTROL_LIST_ALL = "ref_control_list_all"
@@ -26,13 +32,34 @@ REF_ADAIN_MACHINE_STATE = "ref_adain_machine_state"
 REF_COND_IDXS = "ref_cond_idxs"
 REF_UNCOND_IDXS = "ref_uncond_idxs"
 
+CONTEXTREF_OPTIONS_CLASS = "contextref_options_class"
+CONTEXTREF_CLEAN_FUNC = "contextref_clean_func"
+CONTEXTREF_CONTROL_LIST_ALL = "contextref_control_list_all"
+CONTEXTREF_MACHINE_STATE = "contextref_machine_state"
+CONTEXTREF_TEMP_COND_IDX = "contextref_temp_cond_idx"
+
+HIGHEST_VERSION_SUPPORT = 1
+RETURNED_CONTEXTREF_VERSION = 1
+
+
+class RefConst:
+    OPTS = "refcn_opts"
+    CREF_MODE = "contextref_mode"
+
 
 class MachineState:
     WRITE = "write"
     READ = "read"
+    READ_WRITE = "read_write"
     STYLEALIGN = "stylealign"
     OFF = "off"
 
+def is_read(state: str):
+    return state in [MachineState.READ, MachineState.READ_WRITE]
+
+def is_write(state: str):
+    return state in [MachineState.WRITE, MachineState.READ_WRITE]
+
 
 class ReferenceType:
     ATTN = "reference_attn"
@@ -86,7 +113,23 @@ class ReferenceOptions:
                                 attn_style_fidelity=style_fidelity, adain_style_fidelity=style_fidelity,
                                 attn_ref_weight=ref_weight, adain_ref_weight=ref_weight,
                                 ref_with_other_cns=ref_with_other_cns)
-
+    
+    @staticmethod
+    def create_from_kwargs(attn_style_fidelity=0.0, adain_style_fidelity=0.0,
+                         attn_ref_weight=0.0, adain_ref_weight=0.0,
+                         attn_strength=0.0, adain_strength=0.0, **kwargs):
+        has_attn = attn_strength > 0.0
+        has_adain = adain_strength > 0.0
+        if has_attn and has_adain:
+            reference_type = ReferenceType.ATTN_ADAIN
+        elif has_adain:
+            reference_type = ReferenceType.ADAIN
+        else:
+            reference_type = ReferenceType.ATTN
+        return ReferenceOptions(reference_type=reference_type,
+                                attn_style_fidelity=float(attn_style_fidelity), adain_style_fidelity=float(adain_style_fidelity),
+                                attn_ref_weight=float(attn_ref_weight), adain_ref_weight=float(adain_ref_weight),
+                                attn_strength=float(attn_strength), adain_strength=float(adain_strength))
 
 
 class ReferencePreprocWrapper(AbstractPreprocWrapper):
@@ -103,7 +146,7 @@ class ReferenceAdvanced(ControlBase, AdvancedControlBase):
         AdvancedControlBase.__init__(self, super(), timestep_keyframes=timestep_keyframes, weights_default=ControlWeights.controllllite(), allow_condhint_latents=True)
         # TODO: allow vae_optional to be used instead of preprocessor
         #require_vae=True
-        self.ref_opts = ref_opts
+        self._ref_opts = ref_opts
         self.order = 0
         self.model_latent_format = None
         self.model_sampling_current = None
@@ -111,6 +154,16 @@ class ReferenceAdvanced(ControlBase, AdvancedControlBase):
         self.should_apply_adain_effective_strength = False
         self.should_apply_effective_masks = False
         self.latent_shape = None
+        # ContextRef stuff
+        self.is_context_ref = False
+        self.contextref_cond_idx = -1
+        self.contextref_version = RETURNED_CONTEXTREF_VERSION
+
+    @property
+    def ref_opts(self):
+        if self._current_timestep_keyframe is not None and self._current_timestep_keyframe.has_control_weights():
+            return self._current_timestep_keyframe.control_weights.extras.get(RefConst.OPTS, self._ref_opts)
+        return self._ref_opts
 
     def any_attn_strength_to_apply(self):
         return self.should_apply_attn_effective_strength or self.should_apply_effective_masks
@@ -146,6 +199,12 @@ class ReferenceAdvanced(ControlBase, AdvancedControlBase):
         self.apply_advanced_strengths_and_masks(x=real_mask, batched_number=self.batched_number)
         return real_mask
 
+    def get_contextref_mode_replace(self):
+        # used by ADE to get mode_replace for current keyframe
+        if self._current_timestep_keyframe.has_control_weights():
+            return self._current_timestep_keyframe.control_weights.extras.get(RefConst.CREF_MODE, None)
+        return None
+
     def should_run(self):
         running = super().should_run()
         if not running:
@@ -166,13 +225,19 @@ class ReferenceAdvanced(ControlBase, AdvancedControlBase):
             self.cond_hint_original = self.cond_hint_original.condhint
         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":
-            self.ref_opts.attn_style_fidelity = self.ref_opts.original_attn_style_fidelity ** 3.0
-            self.ref_opts.adain_style_fidelity = self.ref_opts.original_adain_style_fidelity ** 3.0
-        else:
-            self.ref_opts.attn_style_fidelity = self.ref_opts.original_attn_style_fidelity
-            self.ref_opts.adain_style_fidelity = self.ref_opts.original_adain_style_fidelity
+        # SDXL is more sensitive to style_fidelity according to sd-webui-controlnet comments;
+        # prepare all ref_opts accordingly
+        all_ref_opts = [self._ref_opts]
+        for kf in self.timestep_keyframes.keyframes:
+            if kf.has_control_weights() and RefConst.OPTS in kf.control_weights.extras:
+                all_ref_opts.append(kf.control_weights.extras[RefConst.OPTS])
+        for ropts in all_ref_opts:
+            if type(model).__name__ == "SDXL":
+                ropts.attn_style_fidelity = ropts.original_attn_style_fidelity ** 3.0
+                ropts.adain_style_fidelity = ropts.original_adain_style_fidelity ** 3.0
+            else:
+                ropts.attn_style_fidelity = ropts.original_attn_style_fidelity
+                ropts.adain_style_fidelity = ropts.original_adain_style_fidelity
 
     def get_control_advanced(self, x_noisy: Tensor, t, cond, batched_number: int):
         # normal ControlNet stuff
@@ -185,25 +250,27 @@ class ReferenceAdvanced(ControlBase, AdvancedControlBase):
                 return control_prev
 
         dtype = x_noisy.dtype
-        # prepare cond_hint - it is a latent, NOT an image
-        #if self.sub_idxs is not None or self.cond_hint is None or x_noisy.shape[2] != self.cond_hint.shape[2] or x_noisy.shape[3] != self.cond_hint.shape[3]:
-        if self.cond_hint is not None:
-            del self.cond_hint
-        self.cond_hint = None
-        # if self.cond_hint_original length greater or equal to real latent count, subdivide it before scaling
-        if self.sub_idxs is not None and self.cond_hint_original.size(0) >= self.full_latent_length:
-            self.cond_hint = comfy.utils.common_upscale(
-                self.cond_hint_original[self.sub_idxs],
-                x_noisy.shape[3], x_noisy.shape[2], 'nearest-exact', "center").to(dtype).to(self.device)
-        else:
-            self.cond_hint = comfy.utils.common_upscale(
-                self.cond_hint_original,
-                x_noisy.shape[3], x_noisy.shape[2], 'nearest-exact', "center").to(dtype).to(self.device)
-        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.model_latent_format.process_in(self.cond_hint)
-        self.cond_hint = ref_noise_latents(self.cond_hint, sigma=t, noise=None)
+        # cond_hint_original only matters for RefCN, NOT ContextRef
+        if self.cond_hint_original is not None:
+            # prepare cond_hint - it is a latent, NOT an image
+            #if self.sub_idxs is not None or self.cond_hint is None or x_noisy.shape[2] != self.cond_hint.shape[2] or x_noisy.shape[3] != self.cond_hint.shape[3]:
+            if self.cond_hint is not None:
+                del self.cond_hint
+            self.cond_hint = None
+            # if self.cond_hint_original length greater or equal to real latent count, subdivide it before scaling
+            if self.sub_idxs is not None and self.cond_hint_original.size(0) >= self.full_latent_length:
+                self.cond_hint = comfy.utils.common_upscale(
+                    self.cond_hint_original[self.sub_idxs],
+                    x_noisy.shape[3], x_noisy.shape[2], 'nearest-exact', "center").to(dtype).to(self.device)
+            else:
+                self.cond_hint = comfy.utils.common_upscale(
+                    self.cond_hint_original,
+                    x_noisy.shape[3], x_noisy.shape[2], 'nearest-exact', "center").to(dtype).to(self.device)
+            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.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))
         self.should_apply_adain_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.adain_strength, 1.0))
@@ -228,6 +295,7 @@ class ReferenceAdvanced(ControlBase, AdvancedControlBase):
     def copy(self):
         c = ReferenceAdvanced(self.ref_opts, self.timestep_keyframes)
         c.order = self.order
+        c.is_context_ref = self.is_context_ref
         self.copy_to(c)
         self.copy_to_advanced(c)
         return c
@@ -238,6 +306,92 @@ class ReferenceAdvanced(ControlBase, AdvancedControlBase):
         return self
 
 
+def handle_context_ref_setup(contextref_obj, transformer_options: dict, positive, negative):
+    transformer_options[CONTEXTREF_MACHINE_STATE] = MachineState.OFF
+    # verify version is compatible
+    if contextref_obj.version > HIGHEST_VERSION_SUPPORT:
+        raise Exception(f"AnimateDiff-Evolved's ContextRef v{contextref_obj.version} is not supported in currently-installed Advanced-ControlNet (only supports ContextRef up to v{HIGHEST_VERSION_SUPPORT}); " +
+                        f"update your Advanced-ControlNet nodes for ContextRef to work.")
+    # init ReferenceOptions
+    cref_opt_dict = contextref_obj.tune.create_dict() # ContextRefTune obj from ADE
+    opts = ReferenceOptions.create_from_kwargs(**cref_opt_dict)
+    # init TimestepKeyframes
+    cref_tks_list = contextref_obj.keyframe.create_list_of_dicts() # ContextRefKeyframeGroup obj from ADE
+    timestep_keyframes = _create_tks_from_dict_list(cref_tks_list)
+    # create ReferenceAdvanced
+    cref = ReferenceAdvanced(ref_opts=opts, timestep_keyframes=timestep_keyframes)
+    cref.strength = contextref_obj.strength # ContextRef obj from ADE
+    cref.set_cond_hint_mask(contextref_obj.mask)
+    cref.order = 99
+    cref.is_context_ref = True
+    context_ref_list = [cref]
+    transformer_options[CONTEXTREF_CONTROL_LIST_ALL] = context_ref_list
+    transformer_options[CONTEXTREF_OPTIONS_CLASS] = ReferenceOptions
+    _add_context_ref_to_conds([positive, negative], cref)
+    return context_ref_list
+
+
+def _create_tks_from_dict_list(dlist: list[dict[str]]) -> TimestepKeyframeGroup:
+    tks = TimestepKeyframeGroup()
+    if dlist is None or len(dlist) == 0:
+        return tks
+    for d in dlist:
+        # scheduling
+        start_percent = d["start_percent"]
+        guarantee_steps = d["guarantee_steps"]
+        inherit_missing = d["inherit_missing"]
+        # values
+        strength = d["strength"]
+        mask = d["mask"]
+        tune = d["tune"]
+        mode = d["mode"]
+        weights = None
+        extras = {}
+        if tune is not None:
+            cref_opt_dict = tune.create_dict() # ContextRefTune obj from ADE
+            opts = ReferenceOptions.create_from_kwargs(**cref_opt_dict)
+            extras[RefConst.OPTS] = opts
+        if mode is not None:
+            extras[RefConst.CREF_MODE] = mode
+        weights = ControlWeights.default(extras=extras)
+        # create keyframe
+        tk = TimestepKeyframe(start_percent=start_percent, guarantee_steps=guarantee_steps, inherit_missing=inherit_missing,
+                              strength=strength, mask_hint_orig=mask, control_weights=weights)
+        tks.add(tk)
+    return tks
+
+
+def _add_context_ref_to_conds(conds: list[list[dict[str]]], context_ref: ReferenceAdvanced):
+    def _add_context_ref_to_existing_control(control: ControlBase, context_ref: ReferenceAdvanced):
+        curr_cn = control
+        while curr_cn is not None:
+            if type(curr_cn) == ReferenceAdvanced and curr_cn.is_context_ref:
+                break
+            if curr_cn.previous_controlnet is not None:
+                curr_cn = curr_cn.previous_controlnet
+                continue
+            orig_previous_controlnet = curr_cn.previous_controlnet
+            # NOTE: code is already in place to restore any ORIG_PREVIOUS_CONTROLNET props
+            setattr(curr_cn, ORIG_PREVIOUS_CONTROLNET, orig_previous_controlnet)
+            curr_cn.previous_controlnet = context_ref
+            curr_cn = orig_previous_controlnet
+
+    def _add_context_ref(actual_cond: dict[str], context_ref: ReferenceAdvanced):
+        # if controls already present on cond, add it to the last previous_controlnet
+        if "control" in actual_cond:
+            return _add_context_ref_to_existing_control(actual_cond["control"], context_ref)
+        # otherwise, need to add it to begin with, and should mark that it should be cleaned after
+        actual_cond["control"] = context_ref
+        actual_cond[CONTROL_INIT_BY_ACN] = True
+    
+    # either add context_ref to end of existing cnet chain, or init 'control' key on actual cond
+    for cond in conds:
+        if cond is not None:
+            for sub_cond in cond:
+                actual_cond = sub_cond[1]
+                _add_context_ref(actual_cond, context_ref)
+
+
 def ref_noise_latents(latents: Tensor, sigma: Tensor, noise: Tensor=None):
     sigma = sigma.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
     alpha_cumprod = 1 / ((sigma * sigma) + 1)
@@ -262,48 +416,148 @@ def simple_noise_latents(latents: Tensor, sigma: float, noise: Tensor=None):
 
 class BankStylesBasicTransformerBlock:
     def __init__(self):
+        # ref
         self.bank = []
         self.style_cfgs = []
         self.cn_idx: list[int] = []
+        # contextref - list of lists as each cond/uncond stored separately
+        self.c_bank: list[list] = []
+        self.c_style_cfgs: list[list] = []
+        self.c_cn_idx: list[list[int]] = []
+
+    def get_bank(self, cref_idx, ignore_contextref, cdevice=None):
+        if ignore_contextref or cref_idx >= len(self.c_bank):
+            return self.bank
+        real_c_bank_list = self.c_bank[cref_idx]
+        if cdevice != None:
+            real_c_bank_list = real_c_bank_list.copy()
+            for i in range(len(real_c_bank_list)):
+                real_c_bank_list[i] = real_c_bank_list[i].to(cdevice)
+        return self.bank + real_c_bank_list
+
+    def get_avg_style_fidelity(self, cref_idx, ignore_contextref):
+        if ignore_contextref or cref_idx >= len(self.c_style_cfgs):
+            return sum(self.style_cfgs) / float(len(self.style_cfgs))
+        combined = self.style_cfgs + self.c_style_cfgs[cref_idx]
+        return sum(combined) / float(len(combined))
     
-    def get_avg_style_fidelity(self):
-        return sum(self.style_cfgs) / float(len(self.style_cfgs))
-    
-    def clean(self):
+    def get_cn_idxs(self, cref_idx, ignore_contxtref):
+        if ignore_contxtref or cref_idx >= len(self.c_cn_idx):
+            return self.cn_idx
+        return self.cn_idx + self.c_cn_idx[cref_idx]
+
+    def init_cref_for_idx(self, cref_idx: int):
+        # makes sure cref lists can accommodate cref_idx 
+        if cref_idx < 0:
+            return
+        while cref_idx >= len(self.c_bank):
+            self.c_bank.append([])
+            self.c_style_cfgs.append([])
+            self.c_cn_idx.append([])
+
+    def clear_cref_for_idx(self, cref_idx: int):
+        if cref_idx < 0 or cref_idx >= len(self.c_bank):
+            return
+        self.c_bank[cref_idx] = []
+        self.c_style_cfgs[cref_idx] = []
+        self.c_cn_idx[cref_idx] = []
+
+    def clean_ref(self):
         del self.bank
-        self.bank = []
         del self.style_cfgs
-        self.style_cfgs = []
         del self.cn_idx
+        self.bank = []
+        self.style_cfgs = []
         self.cn_idx = []
 
+    def clean_contextref(self):
+        del self.c_bank
+        del self.c_style_cfgs
+        del self.c_cn_idx
+        self.c_bank = []
+        self.c_style_cfgs = []
+        self.c_cn_idx = []
+
+    def clean_all(self):
+        self.clean_ref()
+        self.clean_contextref()
+
 
 class BankStylesTimestepEmbedSequential:
     def __init__(self):
+        # ref
         self.var_bank = []
         self.mean_bank = []
         self.style_cfgs = []
         self.cn_idx: list[int] = []
+        # cref
+        self.c_var_bank: list[list] = []
+        self.c_mean_bank: list[list] = []
+        self.c_style_cfgs: list[list] = []
+        self.c_cn_idx: list[list[int]] = []
 
-    def get_avg_var_bank(self):
-        return sum(self.var_bank) / float(len(self.var_bank))
+    def get_var_bank(self, cref_idx, ignore_contextref):
+        if ignore_contextref or cref_idx >= len(self.c_var_bank):
+            return self.var_bank
+        return self.var_bank + self.c_var_bank[cref_idx]
 
-    def get_avg_mean_bank(self):
-        return sum(self.mean_bank) / float(len(self.mean_bank))
-    
-    def get_avg_style_fidelity(self):
-        return sum(self.style_cfgs) / float(len(self.style_cfgs))
+    def get_mean_bank(self, cref_idx, ignore_contextref):
+        if ignore_contextref or cref_idx >= len(self.c_mean_bank):
+            return self.mean_bank
+        return self.mean_bank + self.c_mean_bank[cref_idx]
 
-    def clean(self):
+    def get_style_cfgs(self, cref_idx, ignore_contextref):
+        if ignore_contextref or cref_idx >= len(self.c_style_cfgs):
+            return self.style_cfgs
+        return self.style_cfgs + self.c_style_cfgs[cref_idx]
+
+    def get_cn_idxs(self, cref_idx, ignore_contextref):
+        if ignore_contextref or cref_idx >= len(self.c_cn_idx):
+            return self.cn_idx
+        return self.cn_idx + self.c_cn_idx[cref_idx]
+
+    def init_cref_for_idx(self, cref_idx: int):
+        # makes sure cref lists can accommodate cref_idx 
+        if cref_idx < 0:
+            return
+        while cref_idx >= len(self.c_var_bank):
+            self.c_var_bank.append([])
+            self.c_mean_bank.append([])
+            self.c_style_cfgs.append([])
+            self.c_cn_idx.append([])
+
+    def clear_cref_for_idx(self, cref_idx: int):
+        if cref_idx < 0 or cref_idx >= len(self.c_var_bank):
+            return
+        self.c_var_bank[cref_idx] = []
+        self.c_mean_bank[cref_idx] = []
+        self.c_style_cfgs[cref_idx] = []
+        self.c_cn_idx[cref_idx] = []
+
+    def clean_ref(self):
         del self.mean_bank
-        self.mean_bank = []
         del self.var_bank
-        self.var_bank = []
         del self.style_cfgs
-        self.style_cfgs = []
         del self.cn_idx
+        self.mean_bank = []
+        self.var_bank = []
+        self.style_cfgs = []
         self.cn_idx = []
 
+    def clean_contextref(self):
+        del self.c_var_bank
+        del self.c_mean_bank
+        del self.c_style_cfgs
+        del self.c_cn_idx
+        self.c_var_bank = []
+        self.c_mean_bank = []
+        self.c_style_cfgs = []
+        self.c_cn_idx = []
+
+    def clean_all(self):
+        self.clean_ref()
+        self.clean_contextref()
+
 
 class InjectionBasicTransformerBlockHolder:
     def __init__(self, block: BasicTransformerBlock, idx=None):
@@ -322,8 +576,14 @@ class InjectionBasicTransformerBlockHolder:
         else:
             block.forward = self.original_forward
 
-    def clean(self):
-        self.bank_styles.clean()
+    def clean_ref(self):
+        self.bank_styles.clean_ref()
+    
+    def clean_contextref(self):
+        self.bank_styles.clean_contextref()
+
+    def clean_all(self):
+        self.bank_styles.clean_all()
 
 
 class InjectionTimestepEmbedSequentialHolder:
@@ -339,8 +599,14 @@ class InjectionTimestepEmbedSequentialHolder:
     def restore(self, block: openaimodel.TimestepEmbedSequential):
         block.forward = self.original_forward
     
-    def clean(self):
-        self.bank_styles.clean()
+    def clean_ref(self):
+        self.bank_styles.clean_ref()
+    
+    def clean_contextref(self):
+        self.bank_styles.clean_contextref()
+
+    def clean_all(self):
+        self.bank_styles.clean_all()
 
 
 class ReferenceInjections:
@@ -349,20 +615,44 @@ class ReferenceInjections:
         self.gn_modules = gn_modules if gn_modules else []
         self.diffusion_model_orig_forward: Callable = None
     
-    def clean_module_mem(self):
+    def clean_ref_module_mem(self):
         for attn_module in self.attn_modules:
             try:
-                attn_module.injection_holder.clean()
+                attn_module.injection_holder.clean_ref()
             except Exception:
                 pass
         for gn_module in self.gn_modules:
             try:
-                gn_module.injection_holder.clean()
+                gn_module.injection_holder.clean_ref()
+            except Exception:
+                pass
+
+    def clean_contextref_module_mem(self):
+        for attn_module in self.attn_modules:
+            try:
+                attn_module.injection_holder.clean_contextref()
+            except Exception:
+                pass
+        for gn_module in self.gn_modules:
+            try:
+                gn_module.injection_holder.clean_contextref()
+            except Exception:
+                pass
+
+    def clean_all_module_mem(self):
+        for attn_module in self.attn_modules:
+            try:
+                attn_module.injection_holder.clean_all()
+            except Exception:
+                pass
+        for gn_module in self.gn_modules:
+            try:
+                gn_module.injection_holder.clean_all()
             except Exception:
                 pass
 
     def cleanup(self):
-        self.clean_module_mem()
+        self.clean_all_module_mem()
         del self.attn_modules
         self.attn_modules = []
         del self.gn_modules
@@ -376,13 +666,20 @@ def factory_forward_inject_UNetModel(reference_injections: ReferenceInjections):
         real_args = list(args)
         real_kwargs = list(kwargs.keys())
         control = kwargs.get("control", None)
-        transformer_options = kwargs.get("transformer_options", None)
-        # look for ReferenceAttnPatch objects to get ReferenceAdvanced objects
-        ref_controlnets: list[ReferenceAdvanced] = transformer_options[REF_CONTROL_LIST_ALL]
+        transformer_options: dict[str] = kwargs.get("transformer_options", {})
+        # NOTE: adds support for both ReferenceCN and ContextRef, so need to track them separately
+        # get ReferenceAdvanced objects
+        ref_controlnets: list[ReferenceAdvanced] = transformer_options.get(REF_CONTROL_LIST_ALL, [])
+        context_controlnets: list[ReferenceAdvanced] = transformer_options.get(CONTEXTREF_CONTROL_LIST_ALL, [])
+        # clean contextref stuff if OFF
+        if len(context_controlnets) > 0 and transformer_options[CONTEXTREF_MACHINE_STATE] == MachineState.OFF:
+            reference_injections.clean_contextref_module_mem()
+            context_controlnets = []
         # discard any controlnets that should not run
-        ref_controlnets = [x for x in ref_controlnets if x.should_run()]
+        ref_controlnets = [z for z in ref_controlnets if z.should_run()]
+        context_controlnets = [z for z in context_controlnets if z.should_run()]
         # if nothing related to reference controlnets, do nothing special
-        if len(ref_controlnets) == 0:
+        if len(ref_controlnets) == 0 and len(context_controlnets) == 0:
             return reference_injections.diffusion_model_orig_forward(x, *args, **kwargs)
         try:
             # assign cond and uncond idxs
@@ -391,8 +688,8 @@ def factory_forward_inject_UNetModel(reference_injections: ReferenceInjections):
             indiv_conds = []
             for cond_type in transformer_options["cond_or_uncond"]:
                 indiv_conds.extend([cond_type] * per_batch)
-            transformer_options[REF_UNCOND_IDXS] = [i for i, x in enumerate(indiv_conds) if x == 1]
-            transformer_options[REF_COND_IDXS] = [i for i, x in enumerate(indiv_conds) if x == 0]
+            transformer_options[REF_UNCOND_IDXS] = [i for i, z in enumerate(indiv_conds) if z == 1]
+            transformer_options[REF_COND_IDXS] = [i for i, z in enumerate(indiv_conds) if z == 0]
             # check which controlnets do which thing
             attn_controlnets = []
             adain_controlnets = []
@@ -401,41 +698,96 @@ def factory_forward_inject_UNetModel(reference_injections: ReferenceInjections):
                     attn_controlnets.append(control)
                 if ReferenceType.is_adain(control.ref_opts.reference_type):
                     adain_controlnets.append(control)
-            if len(adain_controlnets) > 0:
+            context_attn_controlnets = []
+            context_adain_controlnets = []
+            # for ease of access, store current contextref_cond_idx value
+            if len(context_controlnets) == 0:
+                transformer_options[CONTEXTREF_TEMP_COND_IDX] = -1
+            else:
+                transformer_options[CONTEXTREF_TEMP_COND_IDX] = context_controlnets[0].contextref_cond_idx
+                # logger.info(f"{transformer_options[CONTEXTREF_MACHINE_STATE]}: {transformer_options[CONTEXTREF_TEMP_COND_IDX]}")
+            
+            for control in context_controlnets:
+                if ReferenceType.is_attn(control.ref_opts.reference_type):
+                    context_attn_controlnets.append(control)
+                if ReferenceType.is_adain(control.ref_opts.reference_type):
+                    context_adain_controlnets.append(control)
+            if len(adain_controlnets) > 0 or len(context_adain_controlnets) > 0:
                 # ComfyUI uses forward_timestep_embed with the TimestepEmbedSequential passed into it
                 orig_forward_timestep_embed = openaimodel.forward_timestep_embed
                 openaimodel.forward_timestep_embed = forward_timestep_embed_ref_inject_factory(orig_forward_timestep_embed)
-            # handle running diffusion with ref cond hints
-            for control in ref_controlnets:
-                if ReferenceType.is_attn(control.ref_opts.reference_type):
-                    transformer_options[REF_ATTN_MACHINE_STATE] = MachineState.WRITE
-                else:
-                    transformer_options[REF_ATTN_MACHINE_STATE] = MachineState.OFF
-                if ReferenceType.is_adain(control.ref_opts.reference_type):
-                    transformer_options[REF_ADAIN_MACHINE_STATE] = MachineState.WRITE
-                else:
-                    transformer_options[REF_ADAIN_MACHINE_STATE] = MachineState.OFF
-                transformer_options[REF_ATTN_CONTROL_LIST] = [control]
-                transformer_options[REF_ADAIN_CONTROL_LIST] = [control]
+            
+            # if RefCN to be used, handle running diffusion with ref cond hints
+            if len(ref_controlnets) > 0:
+                for control in ref_controlnets:
+                    read_attn_list = []
+                    write_attn_list = []
+                    read_adain_list = []
+                    write_adain_list = []
 
-                orig_kwargs = kwargs
-                if not control.ref_opts.ref_with_other_cns:
-                    kwargs = kwargs.copy()
-                    kwargs["control"] = None
-                reference_injections.diffusion_model_orig_forward(control.cond_hint.to(dtype=x.dtype).to(device=x.device), *args, **kwargs)
-                kwargs = orig_kwargs
-            # run diffusion for real now
-            transformer_options[REF_ATTN_MACHINE_STATE] = MachineState.READ
-            transformer_options[REF_ADAIN_MACHINE_STATE] = MachineState.READ
-            transformer_options[REF_ATTN_CONTROL_LIST] = attn_controlnets
-            transformer_options[REF_ADAIN_CONTROL_LIST] = adain_controlnets
-            return reference_injections.diffusion_model_orig_forward(x, *args, **kwargs)
+                    if ReferenceType.is_attn(control.ref_opts.reference_type):
+                        write_attn_list.append(control)
+                    if ReferenceType.is_adain(control.ref_opts.reference_type):
+                        write_adain_list.append(control)
+                    # apply lists
+                    transformer_options[REF_READ_ATTN_CONTROL_LIST] = read_attn_list
+                    transformer_options[REF_WRITE_ATTN_CONTROL_LIST] = write_attn_list
+                    transformer_options[REF_READ_ADAIN_CONTROL_LIST] = read_adain_list
+                    transformer_options[REF_WRITE_ADAIN_CONTROL_LIST] = write_adain_list
+
+                    orig_kwargs = kwargs
+                    # disable other controlnets for this run, if specified
+                    if not control.ref_opts.ref_with_other_cns:
+                        kwargs = kwargs.copy()
+                        kwargs["control"] = None
+                    reference_injections.diffusion_model_orig_forward(control.cond_hint.to(dtype=x.dtype).to(device=x.device), *args, **kwargs)
+                    kwargs = orig_kwargs
+            # prepare running diffusion for real now
+            read_attn_list = []
+            write_attn_list = []
+            read_adain_list = []
+            write_adain_list = []
+
+            # add RefCNs to read lists
+            read_attn_list.extend(attn_controlnets)
+            read_adain_list.extend(adain_controlnets)
+            
+            # do contextref stuff, if needed
+            if len(context_controlnets) > 0:
+                # clean contextref stuff if first WRITE
+                # if context_controlnets[0].contextref_cond_idx == 0 and is_write(transformer_options[CONTEXTREF_MACHINE_STATE]):
+                #     reference_injections.clean_contextref_module_mem()
+                ### add ContextRef to appropriate lists
+                # attn
+                if is_read(transformer_options[CONTEXTREF_MACHINE_STATE]):
+                    read_attn_list.extend(context_attn_controlnets)
+                if is_write(transformer_options[CONTEXTREF_MACHINE_STATE]):
+                    write_attn_list.extend(context_attn_controlnets)
+                # adain
+                if is_read(transformer_options[CONTEXTREF_MACHINE_STATE]):
+                    read_adain_list.extend(context_adain_controlnets)
+                if is_write(transformer_options[CONTEXTREF_MACHINE_STATE]):
+                    write_adain_list.extend(context_adain_controlnets)
+            # apply lists, containing both RefCN and ContextRef
+            transformer_options[REF_READ_ATTN_CONTROL_LIST] = read_attn_list
+            transformer_options[REF_WRITE_ATTN_CONTROL_LIST] = write_attn_list
+            transformer_options[REF_READ_ADAIN_CONTROL_LIST] = read_adain_list
+            transformer_options[REF_WRITE_ADAIN_CONTROL_LIST] = write_adain_list
+            # run diffusion for real
+            try:
+                return reference_injections.diffusion_model_orig_forward(x, *args, **kwargs)
+            finally:
+                # increment current cond idx
+                if len(context_controlnets) > 0:
+                    for cn in context_controlnets:
+                        cn.contextref_cond_idx += 1
         finally:
-            # make sure banks are cleared no matter what happens - otherwise, RIP VRAM
-            reference_injections.clean_module_mem()
-            if len(adain_controlnets) > 0:
+            # make sure ref banks are cleared no matter what happens - otherwise, RIP VRAM
+            reference_injections.clean_ref_module_mem()
+            if len(adain_controlnets) > 0 or len(context_adain_controlnets) > 0:
                 openaimodel.forward_timestep_embed = orig_forward_timestep_embed
 
+
     return forward_inject_UNetModel
 
 
@@ -476,16 +828,32 @@ def _forward_inject_BasicTransformerBlock(self: RefBasicTransformerBlock, x: Ten
 
     # Reference CN stuff
     uc_idx_mask = transformer_options.get(REF_UNCOND_IDXS, [])
-    c_idx_mask = transformer_options.get(REF_COND_IDXS, [])
-    # WRITE mode will only have one ReferenceAdvanced, other modes will have all ReferenceAdvanced
-    ref_controlnets: list[ReferenceAdvanced] = transformer_options.get(REF_ATTN_CONTROL_LIST, None)
-    ref_machine_state: str = transformer_options.get(REF_ATTN_MACHINE_STATE, None)
-    # if in WRITE mode, save n and style_fidelity
-    if ref_controlnets and ref_machine_state == MachineState.WRITE:
-        if ref_controlnets[0].ref_opts.attn_ref_weight > self.injection_holder.attn_weight:
-            self.injection_holder.bank_styles.bank.append(n.detach().clone())
-            self.injection_holder.bank_styles.style_cfgs.append(ref_controlnets[0].ref_opts.attn_style_fidelity)
-            self.injection_holder.bank_styles.cn_idx.append(ref_controlnets[0].order)
+    #c_idx_mask = transformer_options.get(REF_COND_IDXS, [])
+    # WRITE mode may have only 1 ReferenceAdvanced for RefCN at a time, other modes will have all ReferenceAdvanced
+    ref_write_cns: list[ReferenceAdvanced] = transformer_options.get(REF_WRITE_ATTN_CONTROL_LIST, [])
+    ref_read_cns: list[ReferenceAdvanced] = transformer_options.get(REF_READ_ATTN_CONTROL_LIST, [])
+    cref_cond_idx: int = transformer_options.get(CONTEXTREF_TEMP_COND_IDX, -1)
+    ignore_contextref_read = cref_cond_idx < 0 # if writing to bank, should NOT be read in the same execution
+
+    cached_n = None
+    cref_write_cns: list[ReferenceAdvanced] = []
+    # check if any WRITE cns are applicable; Reference CN WRITEs immediately, ContextREF WRITEs after READ completed
+    # if any refs to WRITE, save n and style_fidelity
+    for refcn in ref_write_cns:
+        if refcn.ref_opts.attn_ref_weight > self.injection_holder.attn_weight:
+            if cached_n is None:
+                cached_n = n.detach().clone()
+            # for ContextRef, make sure relevant lists are long enough to cond_idx
+            # store RefCN and ContextRef stuff separately
+            if refcn.is_context_ref:
+                cref_write_cns.append(refcn)
+                self.injection_holder.bank_styles.init_cref_for_idx(cref_cond_idx)
+            else: # Reference CN WRITE
+                self.injection_holder.bank_styles.bank.append(cached_n)
+                self.injection_holder.bank_styles.style_cfgs.append(refcn.ref_opts.attn_style_fidelity)
+                self.injection_holder.bank_styles.cn_idx.append(refcn.order)
+    if len(cref_write_cns) == 0:
+        del cached_n
 
     if "attn1_patch" in transformer_patches:
         patch = transformer_patches["attn1_patch"]
@@ -510,20 +878,21 @@ def _forward_inject_BasicTransformerBlock(self: RefBasicTransformerBlock, x: Ten
             value_attn1 = n
         n = self.attn1.to_q(n)
         # Reference CN READ - use attn1_replace_patch appropriately
-        if ref_machine_state == MachineState.READ and len(self.injection_holder.bank_styles.bank) > 0:
+        if len(ref_read_cns) > 0 and len(self.injection_holder.bank_styles.get_bank(cref_cond_idx, ignore_contextref_read)) > 0:
             bank_styles = self.injection_holder.bank_styles
-            style_fidelity = bank_styles.get_avg_style_fidelity()
-            real_bank = bank_styles.bank.copy()
+            style_fidelity = bank_styles.get_avg_style_fidelity(cref_cond_idx, ignore_contextref_read)
+            real_bank = bank_styles.get_bank(cref_cond_idx, ignore_contextref_read, cdevice=n.device).copy()
+            real_cn_idxs = bank_styles.get_cn_idxs(cref_cond_idx, ignore_contextref_read)
             cn_idx = 0
-            for idx, order in enumerate(bank_styles.cn_idx):
+            for idx, order in enumerate(real_cn_idxs):
                 # make sure matching ref cn is selected
-                for i in range(cn_idx, len(ref_controlnets)):
-                    if ref_controlnets[i].order == order:
+                for i in range(cn_idx, len(ref_read_cns)):
+                    if ref_read_cns[i].order == order:
                         cn_idx = i
                         break
-                assert order == ref_controlnets[cn_idx].order
-                if ref_controlnets[cn_idx].any_attn_strength_to_apply():
-                    effective_strength = ref_controlnets[cn_idx].get_effective_attn_mask_or_float(x=n, channels=n.shape[2], is_mid=self.injection_holder.is_middle)
+                assert order == ref_read_cns[cn_idx].order
+                if ref_read_cns[cn_idx].any_attn_strength_to_apply():
+                    effective_strength = ref_read_cns[cn_idx].get_effective_attn_mask_or_float(x=n, channels=n.shape[2], is_mid=self.injection_holder.is_middle)
                     real_bank[idx] = real_bank[idx] * effective_strength + context_attn1 * (1-effective_strength)
             n_uc = self.attn1.to_out(attn1_replace_patch[block_attn1](
                 n,
@@ -538,7 +907,7 @@ def _forward_inject_BasicTransformerBlock(self: RefBasicTransformerBlock, x: Ten
                     self.attn1.to_v(value_attn1[uc_idx_mask]),
                     extra_options))
             n = style_fidelity * n_c + (1.0-style_fidelity) * n_uc
-            bank_styles.clean()
+            bank_styles.clean_ref()
         else:
             context_attn1 = self.attn1.to_k(context_attn1)
             value_attn1 = self.attn1.to_v(value_attn1)
@@ -546,22 +915,23 @@ def _forward_inject_BasicTransformerBlock(self: RefBasicTransformerBlock, x: Ten
             n = self.attn1.to_out(n)
     else:
         # Reference CN READ - no attn1_replace_patch
-        if ref_machine_state == MachineState.READ and len(self.injection_holder.bank_styles.bank) > 0:
+        if len(ref_read_cns) > 0 and len(self.injection_holder.bank_styles.get_bank(cref_cond_idx, ignore_contextref_read)) > 0:
             if context_attn1 is None:
                 context_attn1 = n
             bank_styles = self.injection_holder.bank_styles
-            style_fidelity = bank_styles.get_avg_style_fidelity()
-            real_bank = bank_styles.bank.copy()
+            style_fidelity = bank_styles.get_avg_style_fidelity(cref_cond_idx, ignore_contextref_read)
+            real_bank = bank_styles.get_bank(cref_cond_idx, ignore_contextref_read, cdevice=n.device).copy()
+            real_cn_idxs = bank_styles.get_cn_idxs(cref_cond_idx, ignore_contextref_read)
             cn_idx = 0
-            for idx, order in enumerate(bank_styles.cn_idx):
+            for idx, order in enumerate(real_cn_idxs):
                 # make sure matching ref cn is selected
-                for i in range(cn_idx, len(ref_controlnets)):
-                    if ref_controlnets[i].order == order:
+                for i in range(cn_idx, len(ref_read_cns)):
+                    if ref_read_cns[i].order == order:
                         cn_idx = i
                         break
-                assert order == ref_controlnets[cn_idx].order
-                if ref_controlnets[cn_idx].any_attn_strength_to_apply():
-                    effective_strength = ref_controlnets[cn_idx].get_effective_attn_mask_or_float(x=n, channels=n.shape[2], is_mid=self.injection_holder.is_middle)
+                assert order == ref_read_cns[cn_idx].order
+                if ref_read_cns[cn_idx].any_attn_strength_to_apply():
+                    effective_strength = ref_read_cns[cn_idx].get_effective_attn_mask_or_float(x=n, channels=n.shape[2], is_mid=self.injection_holder.is_middle)
                     real_bank[idx] = real_bank[idx] * effective_strength + context_attn1 * (1-effective_strength)
             n_uc: Tensor = self.attn1(
                 n,
@@ -574,10 +944,21 @@ def _forward_inject_BasicTransformerBlock(self: RefBasicTransformerBlock, x: Ten
                     context=context_attn1[uc_idx_mask],
                     value=value_attn1[uc_idx_mask] if value_attn1 is not None else value_attn1)
             n = style_fidelity * n_c + (1.0-style_fidelity) * n_uc
-            bank_styles.clean()
+            bank_styles.clean_ref()
         else:
             n = self.attn1(n, context=context_attn1, value=value_attn1)
 
+    # ContextRef CN WRITE
+    if len(cref_write_cns) > 0:
+        # clear so that ContextRef CNs can properly 'replace' previous value at cond_idx
+        self.injection_holder.bank_styles.clear_cref_for_idx(cref_cond_idx)
+        for refcn in cref_write_cns:
+            # add a whole list to match expected type when combining
+            self.injection_holder.bank_styles.c_bank[cref_cond_idx].append(cached_n.to(comfy.model_management.unet_offload_device()))
+            self.injection_holder.bank_styles.c_style_cfgs[cref_cond_idx].append(refcn.ref_opts.attn_style_fidelity)
+            self.injection_holder.bank_styles.c_cn_idx[cref_cond_idx].append(refcn.order)
+        del cached_n
+
     if "attn1_output_patch" in transformer_patches:
         patch = transformer_patches["attn1_output_patch"]
         for p in patch:
@@ -647,51 +1028,82 @@ def forward_timestep_embed_ref_inject_factory(orig_timestep_embed_inject_factory
         transformer_options: dict[str] = args[4]
         # Reference CN stuff
         uc_idx_mask = transformer_options.get(REF_UNCOND_IDXS, [])
-        c_idx_mask = transformer_options.get(REF_COND_IDXS, [])
+        #c_idx_mask = transformer_options.get(REF_COND_IDXS, [])
         # WRITE mode will only have one ReferenceAdvanced, other modes will have all ReferenceAdvanced
-        ref_controlnets: list[ReferenceAdvanced] = transformer_options.get(REF_ADAIN_CONTROL_LIST, None)
-        ref_machine_state: str = transformer_options.get(REF_ADAIN_MACHINE_STATE, None)
-        
-        # if in WRITE mode, save var, mean, and style_cfg
-        if ref_machine_state == MachineState.WRITE:
-            if ref_controlnets[0].ref_opts.adain_ref_weight > ts.injection_holder.gn_weight:
-                var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)
-                ts.injection_holder.bank_styles.var_bank.append(var)
-                ts.injection_holder.bank_styles.mean_bank.append(mean)
-                ts.injection_holder.bank_styles.style_cfgs.append(ref_controlnets[0].ref_opts.adain_style_fidelity)
-                ts.injection_holder.bank_styles.cn_idx.append(ref_controlnets[0].order)
-        # if in READ mode, do math with saved var, mean, and style_cfg
-        if ref_machine_state == MachineState.READ:
-            if len(ts.injection_holder.bank_styles.var_bank) > 0:
+        ref_write_cns: list[ReferenceAdvanced] = transformer_options.get(REF_WRITE_ADAIN_CONTROL_LIST, [])
+        ref_read_cns: list[ReferenceAdvanced] = transformer_options.get(REF_READ_ADAIN_CONTROL_LIST, [])
+        cref_cond_idx: int = transformer_options.get(CONTEXTREF_TEMP_COND_IDX, -1)
+        ignore_contextref_read = cref_cond_idx < 0 # if writing to bank, should NOT be read in the same execution
+
+        cached_var = None
+        cached_mean = None
+        cref_write_cns: list[ReferenceAdvanced] = []
+        # if any refs to WRITE, save var, mean, and style_cfg
+        for refcn in ref_write_cns:
+            if refcn.ref_opts.adain_ref_weight > ts.injection_holder.gn_weight:
+                if cached_var is None:
+                    cached_var, cached_mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)
+                if refcn.is_context_ref:
+                    cref_write_cns.append(refcn)
+                    ts.injection_holder.bank_styles.init_cref_for_idx(cref_cond_idx)
+                else:
+                    ts.injection_holder.bank_styles.var_bank.append(cached_var)
+                    ts.injection_holder.bank_styles.mean_bank.append(cached_mean)
+                    ts.injection_holder.bank_styles.style_cfgs.append(refcn.ref_opts.adain_style_fidelity)
+                    ts.injection_holder.bank_styles.cn_idx.append(refcn.order)
+        if len(cref_write_cns) == 0:
+            del cached_var
+            del cached_mean
+
+        # if any refs to READ, do math with saved var, mean, and style_cfg
+        if len(ref_read_cns) > 0:
+            if len(ts.injection_holder.bank_styles.get_var_bank(cref_cond_idx, ignore_contextref_read)) > 0:
                 bank_styles = ts.injection_holder.bank_styles
                 var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)
                 std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5
                 y_uc = torch.zeros_like(x)
                 cn_idx = 0
-                for idx, order in enumerate(bank_styles.cn_idx):
+                real_style_cfgs = bank_styles.get_style_cfgs(cref_cond_idx, ignore_contextref_read)
+                real_var_bank = bank_styles.get_var_bank(cref_cond_idx, ignore_contextref_read)
+                real_mean_bank = bank_styles.get_mean_bank(cref_cond_idx, ignore_contextref_read)
+                real_cn_idxs = bank_styles.get_cn_idxs(cref_cond_idx, ignore_contextref_read)
+                for idx, order in enumerate(real_cn_idxs):
                     # make sure matching ref cn is selected
-                    for i in range(cn_idx, len(ref_controlnets)):
-                        if ref_controlnets[i].order == order:
+                    for i in range(cn_idx, len(ref_read_cns)):
+                        if ref_read_cns[i].order == order:
                             cn_idx = i
                             break
-                    assert order == ref_controlnets[cn_idx].order
-                    style_fidelity = bank_styles.style_cfgs[idx]
-                    var_acc = bank_styles.var_bank[idx]
-                    mean_acc = bank_styles.mean_bank[idx]
+                    assert order == ref_read_cns[cn_idx].order
+                    style_fidelity = real_style_cfgs[idx]
+                    var_acc = real_var_bank[idx]
+                    mean_acc = real_mean_bank[idx]
                     std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5
                     sub_y_uc = (((x - mean) / std) * std_acc) + mean_acc
-                    if ref_controlnets[cn_idx].any_adain_strength_to_apply():
-                        effective_strength = ref_controlnets[cn_idx].get_effective_adain_mask_or_float(x=x)
+                    if ref_read_cns[cn_idx].any_adain_strength_to_apply():
+                        effective_strength = ref_read_cns[cn_idx].get_effective_adain_mask_or_float(x=x)
                         sub_y_uc = sub_y_uc * effective_strength + x * (1-effective_strength)
                     y_uc += sub_y_uc
                 # get average, if more than one
-                if len(bank_styles.cn_idx) > 1:
-                    y_uc /= len(bank_styles.cn_idx)
+                if len(real_cn_idxs) > 1:
+                    y_uc /= len(real_cn_idxs)
                 y_c = y_uc.clone()
                 if len(uc_idx_mask) > 0 and not math.isclose(style_fidelity, 0.0):
                     y_c[uc_idx_mask] = x.to(y_c.dtype)[uc_idx_mask]
                 y = style_fidelity * y_c + (1.0 - style_fidelity) * y_uc
-            ts.injection_holder.bank_styles.clean()
+            ts.injection_holder.bank_styles.clean_ref()
+
+        # ContextRef CN WRITE
+        if len(cref_write_cns) > 0:
+            # clear so that ContextRef CNs can properly 'replace' previous value at cond_idx
+            ts.injection_holder.bank_styles.clear_cref_for_idx(cref_cond_idx)
+            for refcn in cref_write_cns:
+                # add a whole list to match expected type when combining
+                ts.injection_holder.bank_styles.c_var_bank[cref_cond_idx].append(cached_var)
+                ts.injection_holder.bank_styles.c_mean_bank[cref_cond_idx].append(cached_mean)
+                ts.injection_holder.bank_styles.c_style_cfgs[cref_cond_idx].append(refcn.ref_opts.adain_style_fidelity)
+                ts.injection_holder.bank_styles.c_cn_idx[cref_cond_idx].append(refcn.order)
+            del cached_var
+            del cached_mean
 
         if y is None:
             y = x

adv_control/nodes.py

@@ -231,10 +231,10 @@ NODE_CLASS_MAPPINGS = {
 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_keyframes.py

@@ -82,7 +82,7 @@ class TimestepKeyframeInterpolationNode:
                 "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 @@ class LatentKeyframeInterpolationNode:
             "optional": {
                 "prev_latent_kf": ("LATENT_KEYFRAME", ),
                 "print_keyframes": ("BOOLEAN", {"default": False}),
-                "autosize": ("ACNAUTOSIZE", {"padding": 90}),
+                "autosize": ("ACNAUTOSIZE", {"padding": 50}),
             }
         }
 

adv_control/sampling.py

@@ -11,30 +11,44 @@ from .control_reference import (ReferenceAdvanced, ReferenceInjections,
                                 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_factory(orig_comfy_sample: Callable, is_custom=False) -> Callable
             # 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_factory(orig_comfy_sample: Callable, is_custom=False) -> Callable
                 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_factory(orig_comfy_sample: Callable, is_custom=False) -> Callable
                 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_factory(orig_comfy_sample: Callable, is_custom=False) -> Callable
                 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/utils.py

@@ -21,6 +21,10 @@ from .logger import logger
 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 @@ class AdvancedControlBase:
         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 @@ class AdvancedControlBase:
         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 @@ class AdvancedControlBase:
                     # 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 @@ class AdvancedControlBase:
         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 @@ class AdvancedControlBase:
         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

pyproject.toml

@@ -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 = []