From 7685abd5925c38d6e21072d072fcb5fe0240cbfe Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Sat, 27 Jul 2024 18:54:41 -0500 Subject: [PATCH 01/14] Massive rework of ReferenceCN code to soon support ContextRef for AnimateDiff-Evolved --- adv_control/control_reference.py | 596 ++++++++++++++++++++++++++----- adv_control/sampling.py | 57 ++- 2 files changed, 538 insertions(+), 115 deletions(-) diff --git a/adv_control/control_reference.py b/adv_control/control_reference.py index 28f1105..c36d747 100644 --- a/adv_control/control_reference.py +++ b/adv_control/control_reference.py @@ -17,6 +17,11 @@ from .utils import (AdvancedControlBase, ControlWeights, TimestepKeyframeGroup, broadcast_image_to_extend) +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,6 +31,13 @@ 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_ATTN_MACHINE_STATE = "contextref_attn_machine_state" +CONTEXTREF_ADAIN_MACHINE_STATE = "contextref_adain_machine_state" + class MachineState: WRITE = "write" @@ -111,6 +123,8 @@ 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 def any_attn_strength_to_apply(self): return self.should_apply_attn_effective_strength or self.should_apply_effective_masks @@ -185,25 +199,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 +244,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 +255,18 @@ class ReferenceAdvanced(ControlBase, AdvancedControlBase): return self +def handle_context_ref_setup(transformer_options): + transformer_options[CONTEXTREF_ATTN_MACHINE_STATE] = MachineState.OFF + transformer_options[CONTEXTREF_ADAIN_MACHINE_STATE] = MachineState.OFF + opts = ReferenceOptions(ReferenceType.ATTN, attn_style_fidelity=1.0, attn_ref_weight=1.0, attn_strength=1.0, adain_style_fidelity=0.0, adain_ref_weight=0.0) + cref = ReferenceAdvanced(ref_opts=opts, timestep_keyframes=None) + cref.order = -1 + context_ref_list = [cref] + transformer_options[CONTEXTREF_CONTROL_LIST_ALL] = context_ref_list + transformer_options[CONTEXTREF_OPTIONS_CLASS] = ReferenceOptions + return context_ref_list + + 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 +291,109 @@ 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 + self.c_bank = [] + self.c_style_cfgs = [] + self.c_cn_idx = [] + + def get_bank(self, ignore_contextref=False): + if ignore_contextref: + return self.bank + return self.bank + self.c_bank + + def get_avg_style_fidelity(self, ignore_contextref=False): + if ignore_contextref: + return sum(self.style_cfgs) / float(len(self.style_cfgs)) + combined = self.style_cfgs + self.c_style_cfgs + 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, ignore_contxtref=False): + if ignore_contxtref: + return self.cn_idx + return self.cn_idx + self.c_cn_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 = [] + self.c_mean_bank = [] + self.c_style_cfgs = [] + self.c_cn_idx: list[int] = [] - def get_avg_var_bank(self): - return sum(self.var_bank) / float(len(self.var_bank)) + def get_var_bank(self, ignore_contextref=False): + if ignore_contextref: + return self.var_bank + return self.var_bank + self.c_var_bank - 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, ignore_contextref=False): + if ignore_contextref: + return self.mean_bank + return self.mean_bank + self.c_mean_bank - def clean(self): + def get_style_cfgs(self, ignore_contextref=False): + if ignore_contextref: + return self.style_cfgs + return self.style_cfgs + self.c_style_cfgs + + def get_cn_idx(self, ignore_contextref=False): + if ignore_contextref: + return self.cn_idx + return self.cn_idx + self.c_cn_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 +412,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 +435,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 +451,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 @@ -370,13 +496,262 @@ class ReferenceInjections: self.diffusion_model_orig_forward = None +def HACK_factory_forward_inject_UNetModel(reference_injections: ReferenceInjections): + def forward_inject_UNetModel(self, x: Tensor, *args, **kwargs): + # get control and transformer_options from kwargs + real_args = list(args) + real_kwargs = list(kwargs.keys()) + control = kwargs.get("control", None) + transformer_options = kwargs.get("transformer_options", {}) + # look for ReferenceAttnPatch objects to get ReferenceAdvanced objects + ref_controlnets: list[ReferenceAdvanced] = transformer_options[REF_CONTROL_LIST_ALL] + # discard any controlnets that should not run + ref_controlnets = [x for x in ref_controlnets if x.should_run()] + # if nothing related to reference controlnets, do nothing special + if len(ref_controlnets) == 0: + return reference_injections.diffusion_model_orig_forward(x, *args, **kwargs) + try: + # assign cond and uncond idxs + batched_number = len(transformer_options["cond_or_uncond"]) + per_batch = x.shape[0] // batched_number + 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] + # check which controlnets do which thing + attn_controlnets = [] + adain_controlnets = [] + for control in ref_controlnets: + if ReferenceType.is_attn(control.ref_opts.reference_type): + attn_controlnets.append(control) + if ReferenceType.is_adain(control.ref_opts.reference_type): + adain_controlnets.append(control) + if len(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) + + context_ref = ref_controlnets[0] + if transformer_options[CONTEXTREF_ATTN_MACHINE_STATE] in [MachineState.WRITE, MachineState.OFF]: + reference_injections.clean_module_mem() + transformer_options[REF_ATTN_MACHINE_STATE] = transformer_options[CONTEXTREF_ATTN_MACHINE_STATE] + transformer_options[REF_ADAIN_MACHINE_STATE] = transformer_options[CONTEXTREF_ADAIN_MACHINE_STATE] + transformer_options[REF_ATTN_CONTROL_LIST] = [context_ref] + transformer_options[REF_ADAIN_CONTROL_LIST] = [context_ref] + return reference_injections.diffusion_model_orig_forward(x, *args, **kwargs) + + # 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] + + 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) + finally: + # make sure banks are cleared no matter what happens - otherwise, RIP VRAM + #reference_injections.clean_module_mem() + if len(adain_controlnets) > 0: + openaimodel.forward_timestep_embed = orig_forward_timestep_embed + + + + if len(ref_controlnets) == 0: + return reference_injections.diffusion_model_orig_forward(x, *args, **kwargs) + try: + # assign cond and uncond idxs + batched_number = len(transformer_options["cond_or_uncond"]) + per_batch = x.shape[0] // batched_number + 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] + # check which controlnets do which thing + attn_controlnets = [] + adain_controlnets = [] + for control in ref_controlnets: + if ReferenceType.is_attn(control.ref_opts.reference_type): + attn_controlnets.append(control) + if ReferenceType.is_adain(control.ref_opts.reference_type): + adain_controlnets.append(control) + if len(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] + + 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) + finally: + # make sure banks are cleared no matter what happens - otherwise, RIP VRAM + reference_injections.clean_module_mem() + if len(adain_controlnets) > 0: + openaimodel.forward_timestep_embed = orig_forward_timestep_embed + + return forward_inject_UNetModel + + def factory_forward_inject_UNetModel(reference_injections: ReferenceInjections): def forward_inject_UNetModel(self, x: Tensor, *args, **kwargs): # get control and transformer_options from kwargs real_args = list(args) real_kwargs = list(kwargs.keys()) control = kwargs.get("control", None) - transformer_options = kwargs.get("transformer_options", None) + 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, []) + # discard any controlnets that should not 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 and len(context_controlnets) == 0: + return reference_injections.diffusion_model_orig_forward(x, *args, **kwargs) + try: + # assign cond and uncond idxs + batched_number = len(transformer_options["cond_or_uncond"]) + per_batch = x.shape[0] // batched_number + 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, 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 = [] + for control in ref_controlnets: + if ReferenceType.is_attn(control.ref_opts.reference_type): + attn_controlnets.append(control) + if ReferenceType.is_adain(control.ref_opts.reference_type): + adain_controlnets.append(control) + context_attn_controlnets = [] + context_adain_controlnets = [] + 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) + + # 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 = [] + + 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: + # TODO: clean contextref stuff if attn writing or off + if transformer_options[CONTEXTREF_ATTN_MACHINE_STATE] in [MachineState.WRITE, MachineState.OFF]: + reference_injections.clean_contextref_module_mem() + ### add ContextRef to appropriate lists + # attn + if transformer_options[CONTEXTREF_ATTN_MACHINE_STATE] == MachineState.WRITE: + write_attn_list.extend(context_attn_controlnets) + elif transformer_options[CONTEXTREF_ATTN_MACHINE_STATE] == MachineState.READ: + read_attn_list.extend(context_attn_controlnets) + # adain + if transformer_options[CONTEXTREF_ADAIN_MACHINE_STATE] == MachineState.WRITE: + write_attn_list.extend(context_adain_controlnets) + elif transformer_options[CONTEXTREF_ADAIN_MACHINE_STATE] == MachineState.READ: + read_attn_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 + + return reference_injections.diffusion_model_orig_forward(x, *args, **kwargs) + finally: + # make sure banks are cleared no matter what happens - otherwise, RIP VRAM + reference_injections.clean_ref_module_mem() + if len(adain_controlnets) > 0: + openaimodel.forward_timestep_embed = orig_forward_timestep_embed + + return forward_inject_UNetModel + + + +def ORIG_factory_forward_inject_UNetModel(reference_injections: ReferenceInjections): + def forward_inject_UNetModel(self, x: Tensor, *args, **kwargs): + # get control and transformer_options from kwargs + real_args = list(args) + real_kwargs = list(kwargs.keys()) + control = kwargs.get("control", None) + transformer_options = kwargs.get("transformer_options", {}) # look for ReferenceAttnPatch objects to get ReferenceAdvanced objects ref_controlnets: list[ReferenceAdvanced] = transformer_options[REF_CONTROL_LIST_ALL] # discard any controlnets that should not run @@ -476,16 +851,29 @@ 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, []) + #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) + 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, []) + ignore_contextref_read = False # if writing to bank, should NOT be read in the same execution + + # if any refs to WRITE, save n and style_fidelity + if len(ref_write_cns) > 0: + cached_n = None + 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() + if refcn.is_context_ref: # store separately for RefCN and ContextRef + self.injection_holder.bank_styles.c_bank.append(cached_n) + self.injection_holder.bank_styles.c_style_cfgs.append(ref_write_cns[0].ref_opts.attn_style_fidelity) + self.injection_holder.bank_styles.c_cn_idx.append(ref_write_cns[0].order) + ignore_contextref_read = True + else: + self.injection_holder.bank_styles.bank.append(cached_n) + self.injection_holder.bank_styles.style_cfgs.append(ref_write_cns[0].ref_opts.attn_style_fidelity) + self.injection_holder.bank_styles.cn_idx.append(ref_write_cns[0].order) + del cached_n if "attn1_patch" in transformer_patches: patch = transformer_patches["attn1_patch"] @@ -510,20 +898,20 @@ 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(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(ignore_contextref_read) + real_bank = bank_styles.get_bank(ignore_contextref_read).copy() cn_idx = 0 - for idx, order in enumerate(bank_styles.cn_idx): + for idx, order in enumerate(bank_styles.get_cn_idxs(ignore_contextref_read)): # 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 +926,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 +934,22 @@ 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(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(ignore_contextref_read) + real_bank = bank_styles.get_bank(ignore_contextref_read).copy() cn_idx = 0 - for idx, order in enumerate(bank_styles.cn_idx): + for idx, order in enumerate(bank_styles.get_cn_idxs(ignore_contextref_read)): # 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,7 +962,7 @@ 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) @@ -647,41 +1035,53 @@ 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: + 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, []) + ignore_contextref_read = False # if writing to bank, should NOT be read in the same execution + + # 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: 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: + if refcn.is_context_ref: + ts.injection_holder.bank_styles.c_var_bank.append(var) + ts.injection_holder.bank_styles.c_mean_bank.append(mean) + ts.injection_holder.bank_styles.c_style_cfgs.append(refcn.ref_opts.adain_style_fidelity) + ts.injection_holder.bank_styles.c_cn_idx.append(refcn.order) + ignore_contextref_read = True + else: + 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(refcn.ref_opts.adain_style_fidelity) + ts.injection_holder.bank_styles.cn_idx.append(refcn.order) + + # 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(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(ignore_contextref_read) + real_var_bank = bank_styles.get_var_bank(ignore_contextref_read) + real_mean_bank = bank_styles.get_mean_bank(ignore_contextref_read) + for idx, order in enumerate(bank_styles.get_cn_idx(ignore_contextref_read)): # 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 @@ -691,7 +1091,7 @@ def forward_timestep_embed_ref_inject_factory(orig_timestep_embed_inject_factory 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() if y is None: y = x diff --git a/adv_control/sampling.py b/adv_control/sampling.py index 81f0407..bec3c44 100644 --- a/adv_control/sampling.py +++ b/adv_control/sampling.py @@ -11,24 +11,39 @@ 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 has_contextref_enabled(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") + extras = getattr(context_options, "extras", None) + if extras is None: + return False + context_ref = getattr(extras, "context_ref", None) + return context_ref is not 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() @@ -59,13 +74,20 @@ 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() + if has_contextref_enabled(model): + context_refs = handle_context_ref_setup(model.model_options["transformer_options"]) + # 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) # 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 +110,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 +180,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 +190,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 From ee3490d4b027cd482329d405b93b6ca0e14547a1 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Sat, 27 Jul 2024 21:08:46 -0500 Subject: [PATCH 02/14] Initial working implementation of ContextRef --- adv_control/control.py | 13 ++++--- adv_control/control_reference.py | 64 ++++++++++++++++++++++++-------- adv_control/sampling.py | 8 ++-- adv_control/utils.py | 4 ++ 4 files changed, 65 insertions(+), 24 deletions(-) diff --git a/adv_control/control.py b/adv_control/control.py index 9cc4497..4eb0d18 100644 --- a/adv_control/control.py +++ b/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): diff --git a/adv_control/control_reference.py b/adv_control/control_reference.py index c36d747..c9cd151 100644 --- a/adv_control/control_reference.py +++ b/adv_control/control_reference.py @@ -14,7 +14,7 @@ from comfy.ldm.modules.diffusionmodules import openaimodel from .logger import logger from .utils import (AdvancedControlBase, ControlWeights, TimestepKeyframeGroup, AbstractPreprocWrapper, - broadcast_image_to_extend) + broadcast_image_to_extend, ORIG_PREVIOUS_CONTROLNET, CONTROL_INIT_BY_ACN) REF_READ_ATTN_CONTROL_LIST = "ref_read_attn_control_list" @@ -255,18 +255,50 @@ class ReferenceAdvanced(ControlBase, AdvancedControlBase): return self -def handle_context_ref_setup(transformer_options): - transformer_options[CONTEXTREF_ATTN_MACHINE_STATE] = MachineState.OFF - transformer_options[CONTEXTREF_ADAIN_MACHINE_STATE] = MachineState.OFF +def handle_context_ref_setup(transformer_options, positive, negative): + transformer_options[CONTEXTREF_MACHINE_STATE] = MachineState.OFF opts = ReferenceOptions(ReferenceType.ATTN, attn_style_fidelity=1.0, attn_ref_weight=1.0, attn_strength=1.0, adain_style_fidelity=0.0, adain_ref_weight=0.0) cref = ReferenceAdvanced(ref_opts=opts, timestep_keyframes=None) - cref.order = -1 + 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 _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) @@ -714,29 +746,29 @@ def factory_forward_inject_UNetModel(reference_injections: ReferenceInjections): # do contextref stuff, if needed if len(context_controlnets) > 0: - # TODO: clean contextref stuff if attn writing or off - if transformer_options[CONTEXTREF_ATTN_MACHINE_STATE] in [MachineState.WRITE, MachineState.OFF]: + # TODO: clean contextref stuff if contextref is writing or off + if transformer_options[CONTEXTREF_MACHINE_STATE] in [MachineState.WRITE, MachineState.OFF]: reference_injections.clean_contextref_module_mem() ### add ContextRef to appropriate lists # attn - if transformer_options[CONTEXTREF_ATTN_MACHINE_STATE] == MachineState.WRITE: - write_attn_list.extend(context_attn_controlnets) - elif transformer_options[CONTEXTREF_ATTN_MACHINE_STATE] == MachineState.READ: + if transformer_options[CONTEXTREF_MACHINE_STATE] == MachineState.READ: read_attn_list.extend(context_attn_controlnets) + elif transformer_options[CONTEXTREF_MACHINE_STATE] == MachineState.WRITE: + write_attn_list.extend(context_attn_controlnets) # adain - if transformer_options[CONTEXTREF_ADAIN_MACHINE_STATE] == MachineState.WRITE: - write_attn_list.extend(context_adain_controlnets) - elif transformer_options[CONTEXTREF_ADAIN_MACHINE_STATE] == MachineState.READ: - read_attn_list.extend(context_adain_controlnets) + if transformer_options[CONTEXTREF_MACHINE_STATE] == MachineState.READ: + read_adain_list.extend(context_adain_controlnets) + elif transformer_options[CONTEXTREF_MACHINE_STATE] == MachineState.WRITE: + 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 return reference_injections.diffusion_model_orig_forward(x, *args, **kwargs) finally: - # make sure banks are cleared no matter what happens - otherwise, RIP VRAM + # make sure ref banks are cleared no matter what happens - otherwise, RIP VRAM reference_injections.clean_ref_module_mem() if len(adain_controlnets) > 0: openaimodel.forward_timestep_embed = orig_forward_timestep_embed diff --git a/adv_control/sampling.py b/adv_control/sampling.py index bec3c44..aefc1cd 100644 --- a/adv_control/sampling.py +++ b/adv_control/sampling.py @@ -49,7 +49,7 @@ def acn_sample_factory(orig_comfy_sample: Callable, is_custom=False) -> Callable 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) @@ -79,8 +79,6 @@ def acn_sample_factory(orig_comfy_sample: Callable, is_custom=False) -> Callable if has_sliding_context_windows(model): model.model_options = model.model_options.copy() model.model_options["transformer_options"] = model.model_options["transformer_options"].copy() - if has_contextref_enabled(model): - context_refs = handle_context_ref_setup(model.model_options["transformer_options"]) # convert all CNs to Advanced if needed controlnets_modified, positive, negative = support_sliding_context_windows(model, positive, negative) if controlnets_modified: @@ -88,6 +86,10 @@ def acn_sample_factory(orig_comfy_sample: Callable, is_custom=False) -> Callable args[-3] = positive args[-2] = negative args = tuple(args) + # enable ContextRef, if requested + if has_contextref_enabled(model): + context_refs = handle_context_ref_setup(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 diff --git a/adv_control/utils.py b/adv_control/utils.py index f385133..02ed6e7 100644 --- a/adv_control/utils.py +++ b/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 From 413c95edcc7e1d3bcd884fe57f54535ff580ebb3 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Sun, 28 Jul 2024 19:11:23 -0500 Subject: [PATCH 03/14] Modified ControlRef to keep track of one cond at a time - ensured by AnimateDiff-Evolved --- adv_control/control_reference.py | 332 +++++++------------------------ 1 file changed, 76 insertions(+), 256 deletions(-) diff --git a/adv_control/control_reference.py b/adv_control/control_reference.py index c9cd151..fbfe750 100644 --- a/adv_control/control_reference.py +++ b/adv_control/control_reference.py @@ -35,8 +35,7 @@ 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_ATTN_MACHINE_STATE = "contextref_attn_machine_state" -CONTEXTREF_ADAIN_MACHINE_STATE = "contextref_adain_machine_state" +CONTEXTREF_TEMP_COND_IDX = "contextref_temp_cond_idx" class MachineState: @@ -125,6 +124,7 @@ class ReferenceAdvanced(ControlBase, AdvancedControlBase): self.latent_shape = None # ContextRef stuff self.is_context_ref = False + self.contextref_cond_idx = -1 def any_attn_strength_to_apply(self): return self.should_apply_attn_effective_strength or self.should_apply_effective_masks @@ -327,26 +327,26 @@ class BankStylesBasicTransformerBlock: self.bank = [] self.style_cfgs = [] self.cn_idx: list[int] = [] - # contextref - self.c_bank = [] - self.c_style_cfgs = [] - self.c_cn_idx = [] + # 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, ignore_contextref=False): + def get_bank(self, cref_idx, ignore_contextref): if ignore_contextref: return self.bank - return self.bank + self.c_bank + return self.bank + self.c_bank[cref_idx] - def get_avg_style_fidelity(self, ignore_contextref=False): + def get_avg_style_fidelity(self, cref_idx, ignore_contextref): if ignore_contextref: return sum(self.style_cfgs) / float(len(self.style_cfgs)) - combined = self.style_cfgs + self.c_style_cfgs + combined = self.style_cfgs + self.c_style_cfgs[cref_idx] return sum(combined) / float(len(combined)) - def get_cn_idxs(self, ignore_contxtref=False): + def get_cn_idxs(self, cref_idx, ignore_contxtref): if ignore_contxtref: return self.cn_idx - return self.cn_idx + self.c_cn_idx + return self.cn_idx + self.c_cn_idx[cref_idx] def clean_ref(self): del self.bank @@ -377,30 +377,30 @@ class BankStylesTimestepEmbedSequential: self.style_cfgs = [] self.cn_idx: list[int] = [] # cref - self.c_var_bank = [] - self.c_mean_bank = [] - self.c_style_cfgs = [] - self.c_cn_idx: list[int] = [] + 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_var_bank(self, ignore_contextref=False): + def get_var_bank(self, cref_idx, ignore_contextref): if ignore_contextref: return self.var_bank - return self.var_bank + self.c_var_bank + return self.var_bank + self.c_var_bank[cref_idx] - def get_mean_bank(self, ignore_contextref=False): + def get_mean_bank(self, cref_idx, ignore_contextref): if ignore_contextref: return self.mean_bank - return self.mean_bank + self.c_mean_bank + return self.mean_bank + self.c_mean_bank[cref_idx] - def get_style_cfgs(self, ignore_contextref=False): + def get_style_cfgs(self, cref_idx, ignore_contextref): if ignore_contextref: return self.style_cfgs - return self.style_cfgs + self.c_style_cfgs + return self.style_cfgs + self.c_style_cfgs[cref_idx] - def get_cn_idx(self, ignore_contextref=False): + def get_cn_idxs(self, cref_idx, ignore_contextref): if ignore_contextref: return self.cn_idx - return self.cn_idx + self.c_cn_idx + return self.cn_idx + self.c_cn_idx[cref_idx] def clean_ref(self): del self.mean_bank @@ -528,141 +528,6 @@ class ReferenceInjections: self.diffusion_model_orig_forward = None -def HACK_factory_forward_inject_UNetModel(reference_injections: ReferenceInjections): - def forward_inject_UNetModel(self, x: Tensor, *args, **kwargs): - # get control and transformer_options from kwargs - real_args = list(args) - real_kwargs = list(kwargs.keys()) - control = kwargs.get("control", None) - transformer_options = kwargs.get("transformer_options", {}) - # look for ReferenceAttnPatch objects to get ReferenceAdvanced objects - ref_controlnets: list[ReferenceAdvanced] = transformer_options[REF_CONTROL_LIST_ALL] - # discard any controlnets that should not run - ref_controlnets = [x for x in ref_controlnets if x.should_run()] - # if nothing related to reference controlnets, do nothing special - if len(ref_controlnets) == 0: - return reference_injections.diffusion_model_orig_forward(x, *args, **kwargs) - try: - # assign cond and uncond idxs - batched_number = len(transformer_options["cond_or_uncond"]) - per_batch = x.shape[0] // batched_number - 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] - # check which controlnets do which thing - attn_controlnets = [] - adain_controlnets = [] - for control in ref_controlnets: - if ReferenceType.is_attn(control.ref_opts.reference_type): - attn_controlnets.append(control) - if ReferenceType.is_adain(control.ref_opts.reference_type): - adain_controlnets.append(control) - if len(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) - - context_ref = ref_controlnets[0] - if transformer_options[CONTEXTREF_ATTN_MACHINE_STATE] in [MachineState.WRITE, MachineState.OFF]: - reference_injections.clean_module_mem() - transformer_options[REF_ATTN_MACHINE_STATE] = transformer_options[CONTEXTREF_ATTN_MACHINE_STATE] - transformer_options[REF_ADAIN_MACHINE_STATE] = transformer_options[CONTEXTREF_ADAIN_MACHINE_STATE] - transformer_options[REF_ATTN_CONTROL_LIST] = [context_ref] - transformer_options[REF_ADAIN_CONTROL_LIST] = [context_ref] - return reference_injections.diffusion_model_orig_forward(x, *args, **kwargs) - - # 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] - - 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) - finally: - # make sure banks are cleared no matter what happens - otherwise, RIP VRAM - #reference_injections.clean_module_mem() - if len(adain_controlnets) > 0: - openaimodel.forward_timestep_embed = orig_forward_timestep_embed - - - - if len(ref_controlnets) == 0: - return reference_injections.diffusion_model_orig_forward(x, *args, **kwargs) - try: - # assign cond and uncond idxs - batched_number = len(transformer_options["cond_or_uncond"]) - per_batch = x.shape[0] // batched_number - 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] - # check which controlnets do which thing - attn_controlnets = [] - adain_controlnets = [] - for control in ref_controlnets: - if ReferenceType.is_attn(control.ref_opts.reference_type): - attn_controlnets.append(control) - if ReferenceType.is_adain(control.ref_opts.reference_type): - adain_controlnets.append(control) - if len(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] - - 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) - finally: - # make sure banks are cleared no matter what happens - otherwise, RIP VRAM - reference_injections.clean_module_mem() - if len(adain_controlnets) > 0: - openaimodel.forward_timestep_embed = orig_forward_timestep_embed - - return forward_inject_UNetModel - - def factory_forward_inject_UNetModel(reference_injections: ReferenceInjections): def forward_inject_UNetModel(self, x: Tensor, *args, **kwargs): # get control and transformer_options from kwargs @@ -674,6 +539,10 @@ def factory_forward_inject_UNetModel(reference_injections: ReferenceInjections): # 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 = [z for z in ref_controlnets if z.should_run()] context_controlnets = [z for z in context_controlnets if z.should_run()] @@ -699,6 +568,13 @@ def factory_forward_inject_UNetModel(reference_injections: ReferenceInjections): adain_controlnets.append(control) 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) @@ -746,8 +622,8 @@ def factory_forward_inject_UNetModel(reference_injections: ReferenceInjections): # do contextref stuff, if needed if len(context_controlnets) > 0: - # TODO: clean contextref stuff if contextref is writing or off - if transformer_options[CONTEXTREF_MACHINE_STATE] in [MachineState.WRITE, MachineState.OFF]: + # clean contextref stuff if first WRITE + if context_controlnets[0].contextref_cond_idx == 0 and transformer_options[CONTEXTREF_MACHINE_STATE] == MachineState.WRITE: reference_injections.clean_contextref_module_mem() ### add ContextRef to appropriate lists # attn @@ -766,82 +642,19 @@ def factory_forward_inject_UNetModel(reference_injections: ReferenceInjections): transformer_options[REF_READ_ADAIN_CONTROL_LIST] = read_adain_list transformer_options[REF_WRITE_ADAIN_CONTROL_LIST] = write_adain_list # run diffusion for real - return reference_injections.diffusion_model_orig_forward(x, *args, **kwargs) + 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 ref banks are cleared no matter what happens - otherwise, RIP VRAM reference_injections.clean_ref_module_mem() if len(adain_controlnets) > 0: openaimodel.forward_timestep_embed = orig_forward_timestep_embed - return forward_inject_UNetModel - - - -def ORIG_factory_forward_inject_UNetModel(reference_injections: ReferenceInjections): - def forward_inject_UNetModel(self, x: Tensor, *args, **kwargs): - # get control and transformer_options from kwargs - real_args = list(args) - real_kwargs = list(kwargs.keys()) - control = kwargs.get("control", None) - transformer_options = kwargs.get("transformer_options", {}) - # look for ReferenceAttnPatch objects to get ReferenceAdvanced objects - ref_controlnets: list[ReferenceAdvanced] = transformer_options[REF_CONTROL_LIST_ALL] - # discard any controlnets that should not run - ref_controlnets = [x for x in ref_controlnets if x.should_run()] - # if nothing related to reference controlnets, do nothing special - if len(ref_controlnets) == 0: - return reference_injections.diffusion_model_orig_forward(x, *args, **kwargs) - try: - # assign cond and uncond idxs - batched_number = len(transformer_options["cond_or_uncond"]) - per_batch = x.shape[0] // batched_number - 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] - # check which controlnets do which thing - attn_controlnets = [] - adain_controlnets = [] - for control in ref_controlnets: - if ReferenceType.is_attn(control.ref_opts.reference_type): - attn_controlnets.append(control) - if ReferenceType.is_adain(control.ref_opts.reference_type): - adain_controlnets.append(control) - if len(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] - - 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) - finally: - # make sure banks are cleared no matter what happens - otherwise, RIP VRAM - reference_injections.clean_module_mem() - if len(adain_controlnets) > 0: - openaimodel.forward_timestep_embed = orig_forward_timestep_embed return forward_inject_UNetModel @@ -887,7 +700,8 @@ def _forward_inject_BasicTransformerBlock(self: RefBasicTransformerBlock, x: Ten # WRITE mode will only have one ReferenceAdvanced, 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, []) - ignore_contextref_read = False # if writing to bank, should NOT be read in the same execution + 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 # if any refs to WRITE, save n and style_fidelity if len(ref_write_cns) > 0: @@ -897,10 +711,11 @@ def _forward_inject_BasicTransformerBlock(self: RefBasicTransformerBlock, x: Ten if cached_n is None: cached_n = n.detach().clone() if refcn.is_context_ref: # store separately for RefCN and ContextRef - self.injection_holder.bank_styles.c_bank.append(cached_n) - self.injection_holder.bank_styles.c_style_cfgs.append(ref_write_cns[0].ref_opts.attn_style_fidelity) - self.injection_holder.bank_styles.c_cn_idx.append(ref_write_cns[0].order) - ignore_contextref_read = True + # add a whole list to match expected type when combining + self.injection_holder.bank_styles.c_bank.append([cached_n]) + self.injection_holder.bank_styles.c_style_cfgs.append([ref_write_cns[0].ref_opts.attn_style_fidelity]) + self.injection_holder.bank_styles.c_cn_idx.append([ref_write_cns[0].order]) + ignore_contextref_read = True else: self.injection_holder.bank_styles.bank.append(cached_n) self.injection_holder.bank_styles.style_cfgs.append(ref_write_cns[0].ref_opts.attn_style_fidelity) @@ -930,12 +745,13 @@ 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 len(ref_read_cns) > 0 and len(self.injection_holder.bank_styles.get_bank(ignore_contextref_read)) > 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(ignore_contextref_read) - real_bank = bank_styles.get_bank(ignore_contextref_read).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).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.get_cn_idxs(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_read_cns)): if ref_read_cns[i].order == order: @@ -966,14 +782,15 @@ def _forward_inject_BasicTransformerBlock(self: RefBasicTransformerBlock, x: Ten n = self.attn1.to_out(n) else: # Reference CN READ - no attn1_replace_patch - if len(ref_read_cns) > 0 and len(self.injection_holder.bank_styles.get_bank(ignore_contextref_read)) > 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(ignore_contextref_read) - real_bank = bank_styles.get_bank(ignore_contextref_read).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).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.get_cn_idxs(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_read_cns)): if ref_read_cns[i].order == order: @@ -1071,17 +888,19 @@ def forward_timestep_embed_ref_inject_factory(orig_timestep_embed_inject_factory # WRITE mode will only have one ReferenceAdvanced, other modes will have all ReferenceAdvanced 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, []) - ignore_contextref_read = False # if writing to bank, should NOT be read in the same execution + 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 # 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: var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0) if refcn.is_context_ref: - ts.injection_holder.bank_styles.c_var_bank.append(var) - ts.injection_holder.bank_styles.c_mean_bank.append(mean) - ts.injection_holder.bank_styles.c_style_cfgs.append(refcn.ref_opts.adain_style_fidelity) - ts.injection_holder.bank_styles.c_cn_idx.append(refcn.order) + # add a whole list to match expected type when combining + ts.injection_holder.bank_styles.c_var_bank.append([var]) + ts.injection_holder.bank_styles.c_mean_bank.append([mean]) + ts.injection_holder.bank_styles.c_style_cfgs.append([refcn.ref_opts.adain_style_fidelity]) + ts.injection_holder.bank_styles.c_cn_idx.append([refcn.order]) ignore_contextref_read = True else: ts.injection_holder.bank_styles.var_bank.append(var) @@ -1091,16 +910,17 @@ def forward_timestep_embed_ref_inject_factory(orig_timestep_embed_inject_factory # 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(ignore_contextref_read)) > 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 - real_style_cfgs = bank_styles.get_style_cfgs(ignore_contextref_read) - real_var_bank = bank_styles.get_var_bank(ignore_contextref_read) - real_mean_bank = bank_styles.get_mean_bank(ignore_contextref_read) - for idx, order in enumerate(bank_styles.get_cn_idx(ignore_contextref_read)): + 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_read_cns)): if ref_read_cns[i].order == order: @@ -1117,8 +937,8 @@ def forward_timestep_embed_ref_inject_factory(orig_timestep_embed_inject_factory 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] From f1916a4a0b14aa6a0261636aba5b4cb113a46e23 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Wed, 31 Jul 2024 06:32:21 -0500 Subject: [PATCH 04/14] Added ability to READ and WRITE ContextRef in one go, for the purposes of adding a sliding reference window --- adv_control/control_reference.py | 72 ++++++++++++++++++++++++-------- 1 file changed, 54 insertions(+), 18 deletions(-) diff --git a/adv_control/control_reference.py b/adv_control/control_reference.py index fbfe750..84f7973 100644 --- a/adv_control/control_reference.py +++ b/adv_control/control_reference.py @@ -41,9 +41,16 @@ CONTEXTREF_TEMP_COND_IDX = "contextref_temp_cond_idx" 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" @@ -348,6 +355,22 @@ class BankStylesBasicTransformerBlock: 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 del self.style_cfgs @@ -623,18 +646,18 @@ def factory_forward_inject_UNetModel(reference_injections: ReferenceInjections): # 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 transformer_options[CONTEXTREF_MACHINE_STATE] == MachineState.WRITE: - reference_injections.clean_contextref_module_mem() + # 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 transformer_options[CONTEXTREF_MACHINE_STATE] == MachineState.READ: + if is_read(transformer_options[CONTEXTREF_MACHINE_STATE]): read_attn_list.extend(context_attn_controlnets) - elif transformer_options[CONTEXTREF_MACHINE_STATE] == MachineState.WRITE: + elif is_write(transformer_options[CONTEXTREF_MACHINE_STATE]): write_attn_list.extend(context_attn_controlnets) # adain - if transformer_options[CONTEXTREF_MACHINE_STATE] == MachineState.READ: + if is_read(transformer_options[CONTEXTREF_MACHINE_STATE]): read_adain_list.extend(context_adain_controlnets) - elif transformer_options[CONTEXTREF_MACHINE_STATE] == MachineState.WRITE: + elif 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 @@ -697,30 +720,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 + # 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 if len(ref_write_cns) > 0: - cached_n = None 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() - if refcn.is_context_ref: # store separately for RefCN and ContextRef - # add a whole list to match expected type when combining - self.injection_holder.bank_styles.c_bank.append([cached_n]) - self.injection_holder.bank_styles.c_style_cfgs.append([ref_write_cns[0].ref_opts.attn_style_fidelity]) - self.injection_holder.bank_styles.c_cn_idx.append([ref_write_cns[0].order]) - ignore_contextref_read = True - else: + # 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(ref_write_cns[0].ref_opts.attn_style_fidelity) - self.injection_holder.bank_styles.cn_idx.append(ref_write_cns[0].order) - del 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"] @@ -815,6 +840,17 @@ def _forward_inject_BasicTransformerBlock(self: RefBasicTransformerBlock, x: Ten 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) + 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: From f25e35f1fe299f57e9c7980b36e4d96809543c43 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Wed, 31 Jul 2024 21:04:52 -0500 Subject: [PATCH 05/14] Fixed READ_WRITE state for ContextRef behaving like READ state instead --- adv_control/control_reference.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/adv_control/control_reference.py b/adv_control/control_reference.py index 84f7973..1fcb5d0 100644 --- a/adv_control/control_reference.py +++ b/adv_control/control_reference.py @@ -652,12 +652,12 @@ def factory_forward_inject_UNetModel(reference_injections: ReferenceInjections): # attn if is_read(transformer_options[CONTEXTREF_MACHINE_STATE]): read_attn_list.extend(context_attn_controlnets) - elif is_write(transformer_options[CONTEXTREF_MACHINE_STATE]): + 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) - elif is_write(transformer_options[CONTEXTREF_MACHINE_STATE]): + 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 From 2af647de56cdd5f52b167ef3f32130af40e6c3aa Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Fri, 2 Aug 2024 17:04:22 -0500 Subject: [PATCH 06/14] Added adain support for ContextRef, fixed ref_weight less than 1.0 index error --- adv_control/control_reference.py | 120 ++++++++++++++++++++++--------- adv_control/sampling.py | 14 ++-- 2 files changed, 92 insertions(+), 42 deletions(-) diff --git a/adv_control/control_reference.py b/adv_control/control_reference.py index 1fcb5d0..051a323 100644 --- a/adv_control/control_reference.py +++ b/adv_control/control_reference.py @@ -104,7 +104,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): @@ -262,9 +278,10 @@ class ReferenceAdvanced(ControlBase, AdvancedControlBase): return self -def handle_context_ref_setup(transformer_options, positive, negative): +def handle_context_ref_setup(contextref_obj, transformer_options: dict, positive, negative): transformer_options[CONTEXTREF_MACHINE_STATE] = MachineState.OFF - opts = ReferenceOptions(ReferenceType.ATTN, attn_style_fidelity=1.0, attn_ref_weight=1.0, attn_strength=1.0, adain_style_fidelity=0.0, adain_ref_weight=0.0) + cref_opt_dict = contextref_obj.params.create_dict() # ContextRefParams obj from ADE + opts = ReferenceOptions.create_from_kwargs(**cref_opt_dict) cref = ReferenceAdvanced(ref_opts=opts, timestep_keyframes=None) cref.order = 99 cref.is_context_ref = True @@ -340,18 +357,18 @@ class BankStylesBasicTransformerBlock: self.c_cn_idx: list[list[int]] = [] def get_bank(self, cref_idx, ignore_contextref): - if ignore_contextref: + if ignore_contextref or cref_idx >= len(self.c_bank): return self.bank return self.bank + self.c_bank[cref_idx] def get_avg_style_fidelity(self, cref_idx, ignore_contextref): - if 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_cn_idxs(self, cref_idx, ignore_contxtref): - if 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] @@ -406,25 +423,43 @@ class BankStylesTimestepEmbedSequential: self.c_cn_idx: list[list[int]] = [] def get_var_bank(self, cref_idx, ignore_contextref): - if 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_mean_bank(self, cref_idx, ignore_contextref): - if 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 get_style_cfgs(self, cref_idx, ignore_contextref): - if 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: + 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 del self.var_bank @@ -730,22 +765,21 @@ def _forward_inject_BasicTransformerBlock(self: RefBasicTransformerBlock, x: Ten 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 - if len(ref_write_cns) > 0: - 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 + 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"] @@ -927,22 +961,25 @@ def forward_timestep_embed_ref_inject_factory(orig_timestep_embed_inject_factory 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: - var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0) + 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: - # add a whole list to match expected type when combining - ts.injection_holder.bank_styles.c_var_bank.append([var]) - ts.injection_holder.bank_styles.c_mean_bank.append([mean]) - ts.injection_holder.bank_styles.c_style_cfgs.append([refcn.ref_opts.adain_style_fidelity]) - ts.injection_holder.bank_styles.c_cn_idx.append([refcn.order]) - ignore_contextref_read = True + 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(var) - ts.injection_holder.bank_styles.mean_bank.append(mean) + 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: @@ -981,6 +1018,19 @@ def forward_timestep_embed_ref_inject_factory(orig_timestep_embed_inject_factory y = style_fidelity * y_c + (1.0 - style_fidelity) * y_uc 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 return y.to(x.dtype) diff --git a/adv_control/sampling.py b/adv_control/sampling.py index aefc1cd..ff5ee6c 100644 --- a/adv_control/sampling.py +++ b/adv_control/sampling.py @@ -32,16 +32,15 @@ def has_sliding_context_windows(model): return context_options.context_length is not None -def has_contextref_enabled(model): +def get_contextref_obj(model): motion_injection_params = getattr(model, "motion_injection_params", None) if motion_injection_params is None: - return False + return None context_options = getattr(motion_injection_params, "context_options") extras = getattr(context_options, "extras", None) if extras is None: - return False - context_ref = getattr(extras, "context_ref", None) - return context_ref is not None + return None + return getattr(extras, "context_ref", None) def acn_sample_factory(orig_comfy_sample: Callable, is_custom=False) -> Callable: @@ -87,8 +86,9 @@ def acn_sample_factory(orig_comfy_sample: Callable, is_custom=False) -> Callable args[-2] = negative args = tuple(args) # enable ContextRef, if requested - if has_contextref_enabled(model): - context_refs = handle_context_ref_setup(model.model_options["transformer_options"], positive, negative) + 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() From c4aa684cb3f4dc421ec8ac1ff07619b65cd6ae4b Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Sun, 4 Aug 2024 16:23:45 -0500 Subject: [PATCH 07/14] Fixed TimestepKeyframes being evaluated each time get_control is run instead of just once per step (technically could change results of some workflows, but now works as intended) --- adv_control/utils.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/adv_control/utils.py b/adv_control/utils.py index 02ed6e7..30746fe 100644 --- a/adv_control/utils.py +++ b/adv_control/utils.py @@ -573,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 @@ -635,6 +636,9 @@ class AdvancedControlBase: 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 @@ -670,7 +674,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 @@ -936,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 From 17b70c421ab0d6e60cc0402748a65f9e7a9bb709 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Sun, 4 Aug 2024 21:08:43 -0500 Subject: [PATCH 08/14] Add additional type check in len wrapper hack to prevent issues with Tensors in list --- adv_control/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/adv_control/utils.py b/adv_control/utils.py index 30746fe..8ece2db 100644 --- a/adv_control/utils.py +++ b/adv_control/utils.py @@ -39,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) From 15308a39f3e5b6a66041a5963eede21d85c7c0bd Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Sun, 4 Aug 2024 22:35:46 -0500 Subject: [PATCH 09/14] Added offloading for stored cached_n value to avoid OOMs when a lot of conds are present + higher resolutions --- adv_control/control_reference.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/adv_control/control_reference.py b/adv_control/control_reference.py index 051a323..9d969f8 100644 --- a/adv_control/control_reference.py +++ b/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 @@ -356,10 +357,15 @@ class BankStylesBasicTransformerBlock: self.c_style_cfgs: list[list] = [] self.c_cn_idx: list[list[int]] = [] - def get_bank(self, cref_idx, ignore_contextref): + def get_bank(self, cref_idx, ignore_contextref, cdevice=None): if ignore_contextref or cref_idx >= len(self.c_bank): return self.bank - return self.bank + self.c_bank[cref_idx] + 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): @@ -807,7 +813,7 @@ def _forward_inject_BasicTransformerBlock(self: RefBasicTransformerBlock, x: Ten 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(cref_cond_idx, ignore_contextref_read) - real_bank = bank_styles.get_bank(cref_cond_idx, ignore_contextref_read).copy() + 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(real_cn_idxs): @@ -846,7 +852,7 @@ def _forward_inject_BasicTransformerBlock(self: RefBasicTransformerBlock, x: Ten context_attn1 = n bank_styles = self.injection_holder.bank_styles 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).copy() + 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(real_cn_idxs): @@ -880,7 +886,7 @@ def _forward_inject_BasicTransformerBlock(self: RefBasicTransformerBlock, x: Ten 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) + 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 From d4a0a113478f26b99a5ad05fb8f96320f7eb0116 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Tue, 6 Aug 2024 05:28:06 -0500 Subject: [PATCH 10/14] Added keyframe support for ContextRef, small changes to timestep keyframe code, initial ContextRef version check --- adv_control/control_reference.py | 40 +++++++++++++++++++++++++++++--- adv_control/utils.py | 6 ++--- 2 files changed, 40 insertions(+), 6 deletions(-) diff --git a/adv_control/control_reference.py b/adv_control/control_reference.py index 9d969f8..19a4122 100644 --- a/adv_control/control_reference.py +++ b/adv_control/control_reference.py @@ -14,7 +14,7 @@ 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, +from .utils import (AdvancedControlBase, ControlWeights, TimestepKeyframeGroup, TimestepKeyframe, AbstractPreprocWrapper, broadcast_image_to_extend, ORIG_PREVIOUS_CONTROLNET, CONTROL_INIT_BY_ACN) @@ -38,6 +38,8 @@ 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 + class MachineState: WRITE = "write" @@ -281,9 +283,20 @@ class ReferenceAdvanced(ControlBase, AdvancedControlBase): def handle_context_ref_setup(contextref_obj, transformer_options: dict, positive, negative): transformer_options[CONTEXTREF_MACHINE_STATE] = MachineState.OFF - cref_opt_dict = contextref_obj.params.create_dict() # ContextRefParams obj from ADE + # 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) - cref = ReferenceAdvanced(ref_opts=opts, timestep_keyframes=None) + # 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] @@ -293,6 +306,27 @@ def handle_context_ref_setup(contextref_obj, transformer_options: dict, positive 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"] + # create keyframe + tk = TimestepKeyframe(start_percent=start_percent, guarantee_steps=guarantee_steps, inherit_missing=inherit_missing, + strength=strength, mask_hint_orig=mask) + 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 diff --git a/adv_control/utils.py b/adv_control/utils.py index 8ece2db..b71e371 100644 --- a/adv_control/utils.py +++ b/adv_control/utils.py @@ -632,10 +632,8 @@ 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 @@ -749,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 From 6b77eab34878a746417bfc8787937d42aef4bc98 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Tue, 6 Aug 2024 20:06:21 -0500 Subject: [PATCH 11/14] Renamed Interpolation to Interp. in node names to let the nodes be visually smaller --- adv_control/nodes.py | 4 ++-- adv_control/nodes_keyframes.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/adv_control/nodes.py b/adv_control/nodes.py index b39fbfb..55c2db2 100644 --- a/adv_control/nodes.py +++ b/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 diff --git a/adv_control/nodes_keyframes.py b/adv_control/nodes_keyframes.py index 8c9d21f..9ff333c 100644 --- a/adv_control/nodes_keyframes.py +++ b/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}), } } From 800af8f69f5811d756563751a76d835d012d66f4 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Wed, 7 Aug 2024 02:40:36 -0500 Subject: [PATCH 12/14] Added support for ContextRef mode_replace and tune_replace, fixed orig_forware_timestep_embed not being properly restored for ContextRef adain, ReferenceCN options are schedulable in the backed (no nodes added yet) --- adv_control/control_reference.py | 52 ++++++++++++++++++++++++++------ 1 file changed, 42 insertions(+), 10 deletions(-) diff --git a/adv_control/control_reference.py b/adv_control/control_reference.py index 19a4122..b3a0f5b 100644 --- a/adv_control/control_reference.py +++ b/adv_control/control_reference.py @@ -41,6 +41,11 @@ CONTEXTREF_TEMP_COND_IDX = "contextref_temp_cond_idx" HIGHEST_VERSION_SUPPORT = 1 +class RefConst: + OPTS = "refcn_opts" + CREF_MODE = "contextref_mode" + + class MachineState: WRITE = "write" READ = "read" @@ -140,7 +145,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 @@ -152,6 +157,12 @@ class ReferenceAdvanced(ControlBase, AdvancedControlBase): self.is_context_ref = False self.contextref_cond_idx = -1 + @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 @@ -186,6 +197,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: @@ -206,13 +223,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 @@ -320,9 +343,18 @@ def _create_tks_from_dict_list(dlist: list[dict[str]]) -> TimestepKeyframeGroup: 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) + strength=strength, mask_hint_orig=mask, control_weights=weights) tks.add(tk) return tks @@ -750,7 +782,7 @@ def factory_forward_inject_UNetModel(reference_injections: ReferenceInjections): finally: # make sure ref banks are cleared no matter what happens - otherwise, RIP VRAM reference_injections.clean_ref_module_mem() - if len(adain_controlnets) > 0: + if len(adain_controlnets) > 0 or len(context_adain_controlnets) > 0: openaimodel.forward_timestep_embed = orig_forward_timestep_embed From 2303845a0264a24ed6f3a4a69ee66ad9d849ccae Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Thu, 8 Aug 2024 02:50:04 -0500 Subject: [PATCH 13/14] Added a value for future ContextRef version compatibility checks --- adv_control/control_reference.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/adv_control/control_reference.py b/adv_control/control_reference.py index b3a0f5b..3bea56e 100644 --- a/adv_control/control_reference.py +++ b/adv_control/control_reference.py @@ -39,6 +39,7 @@ CONTEXTREF_MACHINE_STATE = "contextref_machine_state" CONTEXTREF_TEMP_COND_IDX = "contextref_temp_cond_idx" HIGHEST_VERSION_SUPPORT = 1 +RETURNED_CONTEXTREF_VERSION = 1 class RefConst: @@ -156,6 +157,7 @@ class ReferenceAdvanced(ControlBase, AdvancedControlBase): # ContextRef stuff self.is_context_ref = False self.contextref_cond_idx = -1 + self.contextref_version = RETURNED_CONTEXTREF_VERSION @property def ref_opts(self): From 890e8cbce0d6eeefa7548b13c316b1af90cf8a79 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Thu, 8 Aug 2024 03:03:48 -0500 Subject: [PATCH 14/14] version bump --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index e5484f9..c9dfe25 100644 --- a/pyproject.toml +++ b/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 = []