diff --git a/adv_control/nodes_keyframes.py b/adv_control/nodes_keyframes.py index 9ff333c..5f92977 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": 50}), + "autosize": ("ACNAUTOSIZE", {"padding": 0}), } } diff --git a/adv_control/sampling.py b/adv_control/sampling.py index ff5ee6c..0434c5b 100644 --- a/adv_control/sampling.py +++ b/adv_control/sampling.py @@ -1,6 +1,7 @@ from typing import Callable, Union import comfy.sample +import comfy.samplers from comfy.model_patcher import ModelPatcher from comfy.controlnet import ControlBase from comfy.ldm.modules.attention import BasicTransformerBlock @@ -24,46 +25,206 @@ def support_sliding_context_windows(model, positive, negative) -> tuple[bool, di 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") +def has_sliding_context_windows(model: ModelPatcher): + params = model.get_attachment("ADE_params") + if params is None: + # backwards compatibility + params = getattr(model, "motion_injection_params", None) + if params is None: + return False + context_options = getattr(params, "context_options") return context_options.context_length is not None -def get_contextref_obj(model): - motion_injection_params = getattr(model, "motion_injection_params", None) - if motion_injection_params is None: - return None - context_options = getattr(motion_injection_params, "context_options") +def get_contextref_obj(model: ModelPatcher): + params = model.get_attachment("ADE_params") + if params is None: + # backwards compatibility + params = getattr(model, "motion_injection_params", None) + if params is None: + return None + context_options = getattr(params, "context_options") extras = getattr(context_options, "extras", None) if extras is None: return None return getattr(extras, "context_ref", None) -def acn_sample_factory(orig_comfy_sample: Callable, is_custom=False) -> Callable: - def get_refcn(control: ControlBase, order: int=-1): - ref_set: set[ReferenceAdvanced] = set() - if control is None: - return ref_set - if type(control) == ReferenceAdvanced and not control.is_context_ref: - control.order = order - order -= 1 - ref_set.add(control) - ref_set.update(get_refcn(control.previous_controlnet, order=order)) +def get_refcn(control: ControlBase, order: int=-1): + ref_set: set[ReferenceAdvanced] = set() + if control is None: return ref_set + if type(control) == ReferenceAdvanced and not control.is_context_ref: + control.order = order + order -= 1 + ref_set.add(control) + ref_set.update(get_refcn(control.previous_controlnet, order=order)) + return ref_set - def get_lllitecn(control: ControlBase): - cn_dict: dict[ControlLLLiteAdvanced,None] = {} - if control is None: - return cn_dict - if type(control) == ControlLLLiteAdvanced: - cn_dict[control] = None - cn_dict.update(get_lllitecn(control.previous_controlnet)) +def get_lllitecn(control: ControlBase): + cn_dict: dict[ControlLLLiteAdvanced,None] = {} + if control is None: return cn_dict + if type(control) == ControlLLLiteAdvanced: + cn_dict[control] = None + cn_dict.update(get_lllitecn(control.previous_controlnet)) + return cn_dict +def acn_sample_wrapper(executor, guider: comfy.samplers.CFGGuider, *args, **kwargs): + controlnets_modified = False + orig_conds = guider.conds + orig_positive = args[-3] + orig_negative = args[-2] + try: + orig_model_options = model.model_options + # check if positive or negative conds contain ref cn + positive = args[-3] + negative = args[-2] + # if context options present, perform some special actions that may be required + context_refs = [] + if has_sliding_context_windows(model): + model.model_options = model.model_options.copy() + model.model_options["transformer_options"] = model.model_options["transformer_options"].copy() + # convert all CNs to Advanced if needed + controlnets_modified, positive, negative = support_sliding_context_windows(model, positive, negative) + if controlnets_modified: + args = list(args) + args[-3] = positive + args[-2] = negative + args = tuple(args) + # enable ContextRef, if requested + existing_contextref_obj = get_contextref_obj(model) + if existing_contextref_obj is not None: + context_refs = handle_context_ref_setup(existing_contextref_obj, model.model_options["transformer_options"], positive, negative) + controlnets_modified = True + # look for Advanced ControlNets that will require intervention to work + ref_set = set() + lllite_dict: dict[ControlLLLiteAdvanced, None] = {} # dicts preserve insertion order since py3.7 + if positive is not None: + for cond in positive: + if "control" in cond[1]: + ref_set.update(get_refcn(cond[1]["control"])) + lllite_dict.update(get_lllitecn(cond[1]["control"])) + if negative is not None: + for cond in negative: + if "control" in cond[1]: + ref_set.update(get_refcn(cond[1]["control"])) + lllite_dict.update(get_lllitecn(cond[1]["control"])) + # if lllite found, apply patches to a cloned model_options, and continue + if len(lllite_dict) > 0: + lllite_list = list(lllite_dict.keys()) + model.model_options = model.model_options.copy() + model.model_options["transformer_options"] = model.model_options["transformer_options"].copy() + lllite_list.reverse() # reverse so that patches will be applied in expected order + 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 and len(context_refs) == 0: + return orig_comfy_sample(model, *args, **kwargs) + # otherwise, injection time + try: + # inject + # storage for all Reference-related injections + reference_injections = ReferenceInjections() + + # first, handle attn module injection + all_modules = torch_dfs(model.model) + attn_modules: list[RefBasicTransformerBlock] = [] + for module in all_modules: + if isinstance(module, BasicTransformerBlock): + attn_modules.append(module) + attn_modules = [module for module in all_modules if isinstance(module, BasicTransformerBlock)] + attn_modules = sorted(attn_modules, key=lambda x: -x.norm1.normalized_shape[0]) + for i, module in enumerate(attn_modules): + injection_holder = InjectionBasicTransformerBlockHolder(block=module, idx=i) + injection_holder.attn_weight = float(i) / float(len(attn_modules)) + if hasattr(module, "_forward"): # backward compatibility + module._forward = _forward_inject_BasicTransformerBlock.__get__(module, type(module)) + else: + module.forward = _forward_inject_BasicTransformerBlock.__get__(module, type(module)) + module.injection_holder = injection_holder + reference_injections.attn_modules.append(module) + # figure out which module is middle block + if hasattr(model.model.diffusion_model, "middle_block"): + mid_modules = torch_dfs(model.model.diffusion_model.middle_block) + mid_attn_modules: list[RefBasicTransformerBlock] = [module for module in mid_modules if isinstance(module, BasicTransformerBlock)] + for module in mid_attn_modules: + module.injection_holder.is_middle = True + + # next, handle gn module injection (TimestepEmbedSequential) + # TODO: figure out the logic behind these hardcoded indexes + if type(model.model).__name__ == "SDXL": + input_block_indices = [4, 5, 7, 8] + output_block_indices = [0, 1, 2, 3, 4, 5] + else: + input_block_indices = [4, 5, 7, 8, 10, 11] + output_block_indices = [0, 1, 2, 3, 4, 5, 6, 7] + if hasattr(model.model.diffusion_model, "middle_block"): + module = model.model.diffusion_model.middle_block + injection_holder = InjectionTimestepEmbedSequentialHolder(block=module, idx=0, is_middle=True) + injection_holder.gn_weight = 0.0 + module.injection_holder = injection_holder + reference_injections.gn_modules.append(module) + for w, i in enumerate(input_block_indices): + module = model.model.diffusion_model.input_blocks[i] + injection_holder = InjectionTimestepEmbedSequentialHolder(block=module, idx=i, is_input=True) + injection_holder.gn_weight = 1.0 - float(w) / float(len(input_block_indices)) + module.injection_holder = injection_holder + reference_injections.gn_modules.append(module) + for w, i in enumerate(output_block_indices): + module = model.model.diffusion_model.output_blocks[i] + injection_holder = InjectionTimestepEmbedSequentialHolder(block=module, idx=i, is_output=True) + injection_holder.gn_weight = float(w) / float(len(output_block_indices)) + module.injection_holder = injection_holder + reference_injections.gn_modules.append(module) + # hack gn_module forwards and update weights + for i, module in enumerate(reference_injections.gn_modules): + module.injection_holder.gn_weight *= 2 + + # handle diffusion_model forward injection + reference_injections.diffusion_model_orig_forward = model.model.diffusion_model.forward + model.model.diffusion_model.forward = factory_forward_inject_UNetModel(reference_injections).__get__(model.model.diffusion_model, type(model.model.diffusion_model)) + # store ordered ref cns in model's transformer options + new_model_options = model.model_options.copy() + 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) + finally: + # cleanup injections + # restore attn modules + attn_modules: list[RefBasicTransformerBlock] = reference_injections.attn_modules + for module in attn_modules: + module.injection_holder.restore(module) + 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_all() + del module.injection_holder + del gn_modules + # restore diffusion_model forward function + model.model.diffusion_model.forward = reference_injections.diffusion_model_orig_forward.__get__(model.model.diffusion_model, type(model.model.diffusion_model)) + # cleanup + reference_injections.cleanup() + finally: + # restore model_options + model.model_options = orig_model_options + # restore controlnets in conds, if needed + if controlnets_modified: + restore_all_controlnet_conns([orig_positive, orig_negative]) + + + + + +def acn_sample_factory(orig_comfy_sample: Callable, is_custom=False) -> Callable: def acn_sample(model: ModelPatcher, *args, **kwargs): controlnets_modified = False orig_positive = args[-3]