369 lines
20 KiB
Python
369 lines
20 KiB
Python
from typing import Callable, Union
|
|
|
|
import comfy.model_patcher
|
|
import comfy.sample
|
|
import comfy.samplers
|
|
from comfy.model_patcher import ModelPatcher
|
|
from comfy.controlnet import ControlBase
|
|
from comfy.ldm.modules.attention import BasicTransformerBlock
|
|
|
|
|
|
from .control import convert_all_to_advanced, restore_all_controlnet_conns
|
|
from .control_reference import (ReferenceAdvanced, ReferenceInjections,
|
|
RefBasicTransformerBlock, RefTimestepEmbedSequential,
|
|
InjectionBasicTransformerBlockHolder, InjectionTimestepEmbedSequentialHolder,
|
|
_forward_inject_BasicTransformerBlock, factory_forward_inject_UNetModel,
|
|
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(conds) -> tuple[bool, list[dict]]:
|
|
# convert to advanced, with report if anything was actually modified
|
|
modified, new_conds = convert_all_to_advanced(conds)
|
|
return modified, new_conds
|
|
|
|
|
|
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: 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 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))
|
|
return cn_dict
|
|
|
|
def acn_sampler_sample_wrapper(executor, *args, **kwargs):
|
|
controlnets_modified = False
|
|
guider: comfy.samplers.CFGGuider = args[0]
|
|
model = guider.model_patcher
|
|
extra_args: dict = args[2]
|
|
orig_conds = guider.conds
|
|
orig_model_options = extra_args["model_options"]
|
|
try:
|
|
# if context options present, perform some special actions that may be required
|
|
context_refs = []
|
|
if has_sliding_context_windows(guider.model_patcher):
|
|
extra_args["model_options"] = comfy.model_patcher.create_model_options_clone(orig_model_options)
|
|
# convert all CNs to Advanced if needed
|
|
controlnets_modified, conds = support_sliding_context_windows(orig_conds.values())
|
|
if controlnets_modified:
|
|
guider.conds = conds
|
|
# enable ContextRef, if requested
|
|
existing_contextref_obj = get_contextref_obj(guider.model_patcher)
|
|
if existing_contextref_obj is not None:
|
|
context_refs = handle_context_ref_setup(existing_contextref_obj, extra_args["model_options"]["transformer_options"], guider.conds.values())
|
|
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
|
|
for outer_cond in guider.conds.values():
|
|
for cond in outer_cond:
|
|
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 executor(*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 executor(*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
|
|
extra_args["model_options"] = orig_model_options
|
|
# restore guider.conds
|
|
guider.conds = orig_conds
|
|
# restore controlnets in conds, if needed
|
|
if controlnets_modified:
|
|
restore_all_controlnet_conns(orig_conds.values())
|
|
del model
|
|
del guider
|
|
|
|
|
|
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]
|
|
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, conds = support_sliding_context_windows([positive, negative])
|
|
positive, negative = conds
|
|
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])
|
|
|
|
return acn_sample
|