diff --git a/adv_control/control.py b/adv_control/control.py index 261d052..f7db9ae 100644 --- a/adv_control/control.py +++ b/adv_control/control.py @@ -12,6 +12,7 @@ from model_patcher import ModelPatcher from .control_sparsectrl import SparseControlNet, SparseCtrlMotionWrapper, SparseMethod, SparseSettings, SparseSpreadMethod, PreprocSparseRGBWrapper from .control_lllite import LLLiteModule, LLLitePatch +from .control_reference import MachineState, ReferenceOptions, ReferenceAttnPatch from .control_svd import svd_unet_config_from_diffusers_unet, SVDControlNet, svd_unet_to_diffusers from .utils import (AdvancedControlBase, TimestepKeyframeGroup, LatentKeyframeGroup, 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) @@ -348,60 +349,6 @@ class SparseCtrlAdvanced(ControlNetAdvanced): return c -class ReferenceAdvanced(ControlBase, AdvancedControlBase): - def __init__(self, timestep_keyframes: TimestepKeyframeGroup, device=None): - super().__init__(device) - AdvancedControlBase.__init__(self, super(), timestep_keyframes=timestep_keyframes, weights_default=ControlWeights.controllllite(), require_model=True) - # TODO: save attn patches here - - def patch_model(self, model: ModelPatcher): - # TODO: do model patching here - pass - - def pre_run_advanced(self, *args, **kwargs): - AdvancedControlBase.pre_run_advanced(self, *args, **kwargs) - # TODO: set control on patches - - def get_control_advanced(self, x_noisy: Tensor, t, cond, batched_number: int): - # normal ControlNet stuff - control_prev = None - if self.previous_controlnet is not None: - control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number) - - if self.timestep_range is not None: - if t[0] > self.timestep_range[0] or t[0] < self.timestep_range[1]: - return control_prev - - dtype = x_noisy.dtype - # prepare cond_hint - if self.sub_idxs is not None or self.cond_hint is None or x_noisy.shape[2] * 8 != self.cond_hint.shape[2] or x_noisy.shape[3] * 8 != 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] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device) - else: - self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device) - if x_noisy.shape[0] != self.cond_hint.shape[0]: - self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number) - # prepare mask - self.prepare_mask_cond_hint(x_noisy=x_noisy, t=t, cond=cond, batched_number=batched_number) - # done preparing; model patches will take care of everything now. - # return normal controlnet stuff - return control_prev - - def cleanup_advanced(self): - super().cleanup_advanced() - # TODO: cleanup patches here - - def copy(self): - c = ReferenceAdvanced(self.timestep_keyframes) - self.copy_to(c) - self.copy_to_advanced(c) - return c - - class ControlLLLiteAdvanced(ControlBase, AdvancedControlBase): # This ControlNet is more of an attention patch than a traditional controlnet def __init__(self, patch_attn1: LLLitePatch, patch_attn2: LLLitePatch, timestep_keyframes: TimestepKeyframeGroup, device=None): diff --git a/adv_control/control_reference.py b/adv_control/control_reference.py index e69de29..930d5b2 100644 --- a/adv_control/control_reference.py +++ b/adv_control/control_reference.py @@ -0,0 +1,545 @@ +from typing import Callable, Union + +import math +import torch +from torch import Tensor + +import comfy.model_patcher +import comfy.utils +from comfy.controlnet import ControlBase +from comfy.model_patcher import ModelPatcher +from comfy.ldm.modules.attention import BasicTransformerBlock + +from .logger import logger +from .utils import (AdvancedControlBase, ControlWeights, TimestepKeyframeGroup, AbstractPreprocWrapper, + deepcopy_with_sharing, prepare_mask_batch, broadcast_image_to_full, ddpm_noise_latents, simple_noise_latents) + + +REF_CONTROL_LIST = "ref_control_list" +REF_CONTROL_INFO = "ref_control_info" +REF_MACHINE_STATE = "ref_machine_state" + + +class MachineState: + WRITE = "write" + READ = "read" + STYLEALIGN = "stylealign" + + +class ReferenceType: + ATTN = "reference_attn" + ADAIN = "reference_adain" + ATTN_ADAIN = "reference_attn+adain" + STYLE_ALIGN = "StyleAlign" + + _LIST = [ATTN, ADAIN, ATTN_ADAIN] + + +class ReferenceOptions: + def __init__(self, reference_type: str, style_fidelity: float): + self.reference_type = reference_type + self.original_style_fidelity = style_fidelity + self.style_fidelity = style_fidelity + + def clone(self): + return ReferenceOptions(reference_type=self.reference_type, style_fidelity=self.original_style_fidelity) + + +class ReferencePreprocWrapper(AbstractPreprocWrapper): + error_msg = error_msg = "Invalid use of Reference Preprocess output. The output of RGB SparseCtrl preprocessor is NOT a usual image, but a latent pretending to be an image - you must connect the output directly to an Apply Advanced ControlNet node. It cannot be used for anything else that accepts IMAGE input." + def __init__(self, condhint: Tensor): + super().__init__(condhint) + + +class ReferenceAttnPatch: + def __init__(self, control: 'ReferenceAdvanced'=None): + self.control = control + + # def __call__(self, q: Tensor, k: Tensor, v: Tensor, extra_options: dict): + # # do nothing here - all ReferenceAttnPatch is trying to do is be a + # # ComfyUI-compliant way of tracking the corresponding ControlNet obj + # return q, k, v + + def __call__(self, x: Tensor, extra_options: dict): + # do nothing here - all ReferenceAttnPatch is trying to do is be a + # ComfyUI-compliant way of tracking the corresponding ControlNet obj + return x + + def set_control(self, control: 'ReferenceAdvanced') -> 'ReferenceAttnPatch': + self.control = control + return self + + def cleanup(self): + pass + + # make sure deepcopy does not copy control, and deepcopied patch should be assigned to control + def __deepcopy__(self, memo): + self.cleanup() + to_return: ReferenceAttnPatch = deepcopy_with_sharing(self, shared_attribute_names = ['control'], memo=memo) + #logger.warn(f"patch {id(self)} turned into {id(to_return)}") + try: + to_return.control.patch_attn1 = to_return + except Exception: + pass + return to_return + + +class ReferenceAdvanced(ControlBase, AdvancedControlBase): + def __init__(self, patch_attn1: ReferenceAttnPatch, ref_opts: ReferenceOptions, timestep_keyframes: TimestepKeyframeGroup, device=None): + super().__init__(device) + AdvancedControlBase.__init__(self, super(), timestep_keyframes=timestep_keyframes, weights_default=ControlWeights.controllllite(), require_model=True) + self.patch_attn1 = patch_attn1.set_control(self) + self.ref_opts = ref_opts + self.order = 0 + self.latent_format = None + + def get_effective_strength(self): + effective_strength = self.strength + if self.current_timestep_keyframe is not None: + effective_strength = effective_strength * self.current_timestep_keyframe.strength + return effective_strength + + def patch_model(self, model: ModelPatcher): + # need to patch model so that control can be found later from it + model.set_model_attn1_output_patch(self.patch_attn1) + #model.set_model_attn1_patch(self.patch_attn1) + # need to add model_options to make patch/unpatch injection know it has to run + if not REF_CONTROL_INFO in model.model_options: + model.model_options[REF_CONTROL_INFO] = 0 + self.order = model.model_options[REF_CONTROL_INFO] + model.model_options[REF_CONTROL_INFO] + + def pre_run_advanced(self, model, percent_to_timestep_function): + AdvancedControlBase.pre_run_advanced(self, model, percent_to_timestep_function) + if type(self.cond_hint_original) == ReferencePreprocWrapper: + self.cond_hint_original = self.cond_hint_original.condhint + self.latent_format = model.latent_format # LatentFormat object, used to process_in latent cond_hint + # SDXL is more sensitive to style_fidelity according to sd-webui-controlnet comments + if type(model).__name__ == "SDXL": + self.ref_opts.style_fidelity = self.ref_opts.style_fidelity ** 3.0 + else: + self.ref_opts.style_fidelity = self.ref_opts.style_fidelity + # set control on patches + self.patch_attn1.set_control(self) + + def get_control_advanced(self, x_noisy: Tensor, t, cond, batched_number: int): + # normal ControlNet stuff + control_prev = None + if self.previous_controlnet is not None: + control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number) + + if self.timestep_range is not None: + if t[0] > self.timestep_range[0] or t[0] < self.timestep_range[1]: + 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_full(self.cond_hint, x_noisy.shape[0], batched_number, except_one=False) + # noise cond_hint based on sigma (current step) + # TODO: how to handle noise? reproducibility is key... + # mess with the order here? + self.cond_hint = ddpm_noise_latents(self.cond_hint, sigma=t[0] / (self.latent_format.scale_factor), noise=None) + self.cond_hint = self.latent_format.process_in(self.cond_hint) + #self.cond_hint = simple_noise_latents(self.cond_hint, sigma=t[0], noise=None) + + # prepare mask + self.prepare_mask_cond_hint(x_noisy=x_noisy, t=t, cond=cond, batched_number=batched_number) + # done preparing; model patches will take care of everything now. + # return normal controlnet stuff + return control_prev + + def cleanup_advanced(self): + super().cleanup_advanced() + self.patch_attn1.cleanup() + if self.latent_format is not None: + del self.latent_format + self.latent_format = None + + def copy(self): + c = ReferenceAdvanced(self.patch_attn1, self.ref_opts, self.timestep_keyframes) + c.order = self.order + self.copy_to(c) + self.copy_to_advanced(c) + return c + + +class BankStylesBasicTransformerBlock: + def __init__(self): + self.bank = [] + self.style_cfgs = [] + + def get_avg_style_fidelity(self): + return sum(self.style_cfgs) / float(len(self.style_cfgs)) + + def clean(self): + del self.bank + self.bank = [] + del self.style_cfgs + self.style_cfgs = [] + + +class InjectionBasicTransformerBlockHolder: + def __init__(self, block: BasicTransformerBlock): + self.original_forward = block._forward + self.attn_weight = 1.0 + self.bank_styles: dict[int, BankStylesBasicTransformerBlock] = {} + + def restore(self, block: BasicTransformerBlock): + block._forward = self.original_forward + + def clean(self): + for bank_style in list(self.bank_styles.values()): + bank_style.clean() + self.bank_styles.clear() + + +# def factory_clone_injected_ModelPatcher(orig_clone_ModelPatcher: Callable): +# def clone_injected_ref(self, *args, **kwargs): +# cloned = orig_clone_ModelPatcher(*args, **kwargs) +# if InjectMP.is_injected(self): +# InjectMP.mark_injected(cloned) +# cloned.clone = factory_clone_injected_ModelPatcher(cloned.clone).__get__(cloned, type(cloned)) +# return cloned +# return clone_injected_ref + + +# inject ModelPatcher.clone so that necessary injection will happen when needed +# orig_modelpatcher_clone = comfy.model_patcher.ModelPatcher.clone +# def clone_injection_ref(self: ModelPatcher, *args, **kwargs): +# cloned = orig_modelpatcher_clone(self, *args, **kwargs) +# if InjectMP.is_injected(self): +# InjectMP.mark_injected(cloned) +# return cloned +# comfy.model_patcher.ModelPatcher.clone = clone_injection_ref + + +# inject ModelPatcher.patch_model to apply +orig_modelpatcher_patch_model = comfy.model_patcher.ModelPatcher.patch_model +def patch_model_injection_ref(self: ModelPatcher, *args, **kwargs): + if REF_CONTROL_INFO in self.model_options: + # storage for all Reference-related injections + reference_injections = ReferenceInjections() + # first, handle attn module injection + all_modules = torch_dfs(self.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) + injection_holder.attn_weight = float(i) / float(len(attn_modules)) + module._forward = _forward_inject_BasicTransformerBlock.__get__(module, type(module)) + module.injection_holder = injection_holder + reference_injections.attn_modules = attn_modules + # handle diffusion_model forward injection + reference_injections.diffusion_model_orig_forward = self.model.diffusion_model.forward + self.model.diffusion_model.forward = factory_forward_inject_UNetModel(reference_injections).__get__(self.model.diffusion_model, type(self.model.diffusion_model)) + InjectMP.set_injected(self, reference_injections) + to_return = orig_modelpatcher_patch_model(self, *args, **kwargs) + return to_return +comfy.model_patcher.ModelPatcher.patch_model = patch_model_injection_ref + + +orig_modelpatcher_unpatch_model = comfy.model_patcher.ModelPatcher.unpatch_model +def unpatch_model_injection_ref(self: ModelPatcher, *args, **kwargs): + if REF_CONTROL_INFO in self.model_options: + reference_injections: ReferenceInjections = InjectMP.get_injected(self) + # first, 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() + del module.injection_holder + del attn_modules + # restore diffusion_model forward function + self.model.diffusion_model.forward = reference_injections.diffusion_model_orig_forward.__get__(self.model.diffusion_model, type(self.model.diffusion_model)) + # cleanup + InjectMP.clean_injected(self) + reference_injections.cleanup() + to_return = orig_modelpatcher_unpatch_model(self, *args, **kwargs) + return to_return +comfy.model_patcher.ModelPatcher.unpatch_model = unpatch_model_injection_ref + + +class ReferenceInjections: + def __init__(self, attn_modules: list['RefBasicTransformerBlock']=None): + self.attn_modules = attn_modules if attn_modules else [] + self.diffusion_model_orig_forward: Callable = None + + def clean_module_mem(self): + for attn_module in self.attn_modules: + attn_module.injection_holder.clean() + + def cleanup(self): + del self.attn_modules + self.attn_modules = [] + self.diffusion_model_orig_forward = None + + +class InjectMP: + PARAM_INJECTED_REF = "___injected_ref" + + @staticmethod + def is_injected(model: ModelPatcher): + return getattr(model, InjectMP.PARAM_INJECTED_REF, False) + + def mark_injected(model: ModelPatcher): + setattr(model, InjectMP.PARAM_INJECTED_REF, True) + + def set_injected(model: ModelPatcher, value: ReferenceInjections): + setattr(model, InjectMP.PARAM_INJECTED_REF, value) + + def get_injected(model: ModelPatcher) -> ReferenceInjections: + return getattr(model, InjectMP.PARAM_INJECTED_REF) + + def clean_injected(model: ModelPatcher): + delattr(model, InjectMP.PARAM_INJECTED_REF) + InjectMP.mark_injected(model) + + +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) + # look for ReferenceAttnPatch objects to get ReferenceAdvanced objects + # and remove the patch from transformer_options so it won't be ran for no reason + patch_name = "attn1_output_patch" + #patch_name = "attn1_patch" + ref_patches: list[ReferenceAttnPatch] = [] + if "patches" in transformer_options: + if patch_name in transformer_options["patches"]: + patches: list = transformer_options["patches"][patch_name] + remove_idxs = [] + for i in range(len(patches)): + if isinstance(patches[i], ReferenceAttnPatch): + ref_patches.append(patches[i]) + remove_idxs.append(i) + # for i in reversed(remove_idxs): + # patches.pop(i) + # if len(transformer_options["patches"]["attn1_patch"]) == 0: + # transformer_options["patches"].pop("attn1_patch") + ref_controlnets: list[ReferenceAdvanced] = [x.control for x in ref_patches] + # discard any controlnets that should not run + ref_controlnets = [x for x in ref_controlnets if x.should_run()] + ref_controlnets = sorted(ref_controlnets, key=lambda x: x.order) + # 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: + # otherwise, need to handle ref controlnet stuff + for control in ref_controlnets: + transformer_options[REF_MACHINE_STATE] = MachineState.WRITE + transformer_options[REF_CONTROL_LIST] = [control] + # TODO: insert control's cond_hint + reference_injections.diffusion_model_orig_forward(control.cond_hint.to(dtype=x.dtype).to(device=x.device), *args, **kwargs) + transformer_options[REF_MACHINE_STATE] = MachineState.READ + transformer_options[REF_CONTROL_LIST] = ref_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() + + return forward_inject_UNetModel + + +# dummy class just to help IDE keep track of injected variables +class RefBasicTransformerBlock(BasicTransformerBlock): + injection_holder: InjectionBasicTransformerBlockHolder = None + +def _forward_inject_BasicTransformerBlock(self: RefBasicTransformerBlock, x: Tensor, context: Tensor=None, transformer_options: dict[str]={}): + extra_options = {} + block = transformer_options.get("block", None) + block_index = transformer_options.get("block_index", 0) + transformer_patches = {} + transformer_patches_replace = {} + + for k in transformer_options: + if k == "patches": + transformer_patches = transformer_options[k] + elif k == "patches_replace": + transformer_patches_replace = transformer_options[k] + else: + extra_options[k] = transformer_options[k] + + extra_options["n_heads"] = self.n_heads + extra_options["dim_head"] = self.d_head + + if self.ff_in: + x_skip = x + x = self.ff_in(self.norm_in(x)) + if self.is_res: + x += x_skip + + n = self.norm1(x) + if self.disable_self_attn: + context_attn1 = context + else: + context_attn1 = None + value_attn1 = None + + # Reference CN stuff + # WRITE mode will only have one ReferenceAdvanced, other modes will have all ReferenceAdvanced + ref_controlnets: list[ReferenceAdvanced] = transformer_options.get(REF_CONTROL_LIST, None) + ref_machine_state: str = transformer_options.get(REF_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].get_effective_strength() > self.injection_holder.attn_weight: + if not self.injection_holder.bank_styles.get(ref_controlnets[0].order, None): + self.injection_holder.bank_styles[ref_controlnets[0].order] = BankStylesBasicTransformerBlock() + bank_style = self.injection_holder.bank_styles[ref_controlnets[0].order] + bank_style.bank.append(n.detach().clone()) + bank_style.style_cfgs.append(ref_controlnets[0].ref_opts.style_fidelity) + # create uc_idx_mask + per_batch = x.shape[0] // len(transformer_options["cond_or_uncond"]) + indiv_conds = [] + for cond_type in transformer_options["cond_or_uncond"]: + indiv_conds.extend([cond_type] * per_batch) + uc_idx_mask = [i for i, x in enumerate(indiv_conds) if x == 0] + + if "attn1_patch" in transformer_patches: + patch = transformer_patches["attn1_patch"] + if context_attn1 is None: + context_attn1 = n + value_attn1 = context_attn1 + for p in patch: + n, context_attn1, value_attn1 = p(n, context_attn1, value_attn1, extra_options) + + if block is not None: + transformer_block = (block[0], block[1], block_index) + else: + transformer_block = None + attn1_replace_patch = transformer_patches_replace.get("attn1", {}) + block_attn1 = transformer_block + if block_attn1 not in attn1_replace_patch: + block_attn1 = block + + if block_attn1 in attn1_replace_patch: + if context_attn1 is None: + context_attn1 = n + value_attn1 = n + n = self.attn1.to_q(n) + # Reference CN READ - use attn1_replace_patch appropriately + if ref_machine_state == MachineState.READ: + bank_styles = self.injection_holder.bank_styles[ref_controlnets[0].order] + style_fidelity = bank_styles.get_avg_style_fidelity() + n_uc = self.attn1.to_out(attn1_replace_patch[block_attn1]( + n, + self.attn1.to_k(torch.cat([context_attn1] + bank_styles.bank, dim=1)), + self.attn1.to_v(torch.cat([value_attn1] + bank_styles.bank, dim=1)), + extra_options)) + n_c = n_uc.clone() + if len(uc_idx_mask) > 0 and not math.isclose(style_fidelity, 0.0): + n_c[uc_idx_mask] = self.attn1.to_out(attn1_replace_patch[block_attn1]( + n[uc_idx_mask], + self.attn1.to_k(context_attn1[uc_idx_mask]), + 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() + else: + context_attn1 = self.attn1.to_k(context_attn1) + value_attn1 = self.attn1.to_v(value_attn1) + n = attn1_replace_patch[block_attn1](n, context_attn1, value_attn1, extra_options) + n = self.attn1.to_out(n) + else: + # Reference CN READ - no attn1_replace_patch + if ref_machine_state == MachineState.READ: + if context_attn1 is None: + context_attn1 = n + bank_styles = self.injection_holder.bank_styles[ref_controlnets[0].order] + style_fidelity = bank_styles.get_avg_style_fidelity() + n_uc: Tensor = self.attn1( + n, + context=torch.cat([context_attn1] + bank_styles.bank, dim=1), + value=torch.cat([value_attn1] + bank_styles.bank, dim=1) if value_attn1 is not None else value_attn1) + n_c = n_uc.clone() + if len(uc_idx_mask) > 0 and not math.isclose(style_fidelity, 0.0): + n_c[uc_idx_mask] = self.attn1( + n[uc_idx_mask], + context=context_attn1[uc_idx_mask] if context_attn1 is not None else context_attn1, + 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() + else: + n = self.attn1(n, context=context_attn1, value=value_attn1) + + if "attn1_output_patch" in transformer_patches: + patch = transformer_patches["attn1_output_patch"] + for p in patch: + n = p(n, extra_options) + + x += n + if "middle_patch" in transformer_patches: + patch = transformer_patches["middle_patch"] + for p in patch: + x = p(x, extra_options) + + if self.attn2 is not None: + n = self.norm2(x) + if self.switch_temporal_ca_to_sa: + context_attn2 = n + else: + context_attn2 = context + value_attn2 = None + if "attn2_patch" in transformer_patches: + patch = transformer_patches["attn2_patch"] + value_attn2 = context_attn2 + for p in patch: + n, context_attn2, value_attn2 = p(n, context_attn2, value_attn2, extra_options) + + attn2_replace_patch = transformer_patches_replace.get("attn2", {}) + block_attn2 = transformer_block + if block_attn2 not in attn2_replace_patch: + block_attn2 = block + + if block_attn2 in attn2_replace_patch: + if value_attn2 is None: + value_attn2 = context_attn2 + n = self.attn2.to_q(n) + context_attn2 = self.attn2.to_k(context_attn2) + value_attn2 = self.attn2.to_v(value_attn2) + n = attn2_replace_patch[block_attn2](n, context_attn2, value_attn2, extra_options) + n = self.attn2.to_out(n) + else: + n = self.attn2(n, context=context_attn2, value=value_attn2) + + if "attn2_output_patch" in transformer_patches: + patch = transformer_patches["attn2_output_patch"] + for p in patch: + n = p(n, extra_options) + + x += n + if self.is_res: + x_skip = x + x = self.ff(self.norm3(x)) + if self.is_res: + x += x_skip + + return x + + +# DFS Search for Torch.nn.Module, Written by Lvmin +def torch_dfs(model: torch.nn.Module): + result = [model] + for child in model.children(): + result += torch_dfs(child) + return result diff --git a/adv_control/nodes.py b/adv_control/nodes.py index 1cea1af..ce4c0e7 100644 --- a/adv_control/nodes.py +++ b/adv_control/nodes.py @@ -11,6 +11,7 @@ from .nodes_weight import (DefaultWeights, ScaledSoftMaskedUniversalWeights, Sca SoftT2IAdapterWeights, CustomT2IAdapterWeights) from .nodes_latent_keyframe import LatentKeyframeGroupNode, LatentKeyframeInterpolationNode, LatentKeyframeBatchedGroupNode, LatentKeyframeNode from .nodes_sparsectrl import SparseCtrlMergedLoaderAdvanced, SparseCtrlLoaderAdvanced, SparseIndexMethodNode, SparseSpreadMethodNode, RgbSparseCtrlPreprocessor +from .nodes_reference import ReferenceControlNetNode, ReferencePreprocessorNode from .nodes_loosecontrol import ControlNetLoaderWithLoraAdvanced from .nodes_deprecated import LoadImagesFromDirectory from .logger import logger @@ -233,6 +234,9 @@ NODE_CLASS_MAPPINGS = { "ACN_SparseCtrlMergedLoaderAdvanced": SparseCtrlMergedLoaderAdvanced, "ACN_SparseCtrlIndexMethodNode": SparseIndexMethodNode, "ACN_SparseCtrlSpreadMethodNode": SparseSpreadMethodNode, + # Reference + "ACN_ReferencePreprocessor": ReferencePreprocessorNode, + "ACN_ReferenceControlNet": ReferenceControlNetNode, # LOOSEControl #"ACN_ControlNetLoaderWithLoraAdvanced": ControlNetLoaderWithLoraAdvanced, # Deprecated @@ -265,6 +269,9 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ACN_SparseCtrlMergedLoaderAdvanced": "Load Merged SparseCtrl Model 🛂🅐🅒🅝", "ACN_SparseCtrlIndexMethodNode": "SparseCtrl Index Method 🛂🅐🅒🅝", "ACN_SparseCtrlSpreadMethodNode": "SparseCtrl Spread Method 🛂🅐🅒🅝", + # Reference + "ACN_ReferencePreprocessor": "Reference Preproccessor 🛂🅐🅒🅝", + "ACN_ReferenceControlNet": "Reference ControlNet 🛂🅐🅒🅝", # LOOSEControl #"ACN_ControlNetLoaderWithLoraAdvanced": "Load Adv. ControlNet Model w/ LoRA 🛂🅐🅒🅝", # Deprecated diff --git a/adv_control/nodes_reference.py b/adv_control/nodes_reference.py index e69de29..0730647 100644 --- a/adv_control/nodes_reference.py +++ b/adv_control/nodes_reference.py @@ -0,0 +1,57 @@ +from torch import Tensor + +from nodes import VAEEncode +import comfy.utils + +from .control_reference import ReferenceAdvanced, ReferenceAttnPatch, ReferenceOptions, ReferenceType, ReferencePreprocWrapper + + +# node for ReferenceCN +class ReferenceControlNetNode: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "reference_type": (ReferenceType._LIST,), + "style_fidelity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}) + }, + } + + RETURN_TYPES = ("CONTROL_NET", ) + FUNCTION = "load_controlnet" + + CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/Reference" + + def load_controlnet(self, reference_type: str, style_fidelity: float): + ref_opts = ReferenceOptions(reference_type=reference_type, style_fidelity=style_fidelity) + ref_patch = ReferenceAttnPatch() + controlnet = ReferenceAdvanced(patch_attn1=ref_patch, ref_opts=ref_opts, timestep_keyframes=None) + return (controlnet,) + + +class ReferencePreprocessorNode: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE", ), + "vae": ("VAE", ), + "latent_size": ("LATENT", ), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("proc_IMAGE",) + FUNCTION = "preprocess_images" + + CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/Reference/preprocess" + + def preprocess_images(self, vae, image: Tensor, latent_size: Tensor): + # first, resize image to match latents + image = image.movedim(-1,1) + image = comfy.utils.common_upscale(image, latent_size["samples"].shape[3] * 8, latent_size["samples"].shape[2] * 8, 'nearest-exact', "center") + image = image.movedim(1,-1) + # then, vae encode + image = VAEEncode.vae_encode_crop_pixels(image) + encoded = vae.encode(image[:,:,:,:3]) + return (ReferencePreprocWrapper(condhint=encoded),) diff --git a/adv_control/utils.py b/adv_control/utils.py index 1d77707..85a35de 100644 --- a/adv_control/utils.py +++ b/adv_control/utils.py @@ -3,6 +3,7 @@ from typing import Callable, Union import torch from torch import Tensor import torch.nn.functional as F +import math import comfy.ops import comfy.utils @@ -232,6 +233,38 @@ class TimestepKeyframeGroup: return group +class AbstractPreprocWrapper: + error_msg = "Invalid use of [InsertHere] output. The output of [InsertHere] preprocessor is NOT a usual image, but a latent pretending to be an image - you must connect the output directly to an Apply ControlNet node (advanced or otherwise). It cannot be used for anything else that accepts IMAGE input." + def __init__(self, condhint: Tensor): + self.condhint = condhint + + def movedim(self, *args, **kwargs): + return self + + def __getattr__(self, *args, **kwargs): + raise AttributeError(self.error_msg) + + def __setattr__(self, name, value): + if name != "condhint": + raise AttributeError(self.error_msg) + super().__setattr__(name, value) + + def __iter__(self, *args, **kwargs): + raise AttributeError(self.error_msg) + + def __next__(self, *args, **kwargs): + raise AttributeError(self.error_msg) + + def __len__(self, *args, **kwargs): + raise AttributeError(self.error_msg) + + def __getitem__(self, *args, **kwargs): + raise AttributeError(self.error_msg) + + def __setitem__(self, *args, **kwargs): + raise AttributeError(self.error_msg) + + # depending on model, AnimateDiff may inject into GroupNorm, so make sure GroupNorm will be clean class disable_weight_init_clean_groupnorm(comfy.ops.disable_weight_init): class GroupNorm(comfy.ops.disable_weight_init.GroupNorm): @@ -264,6 +297,40 @@ def linear_conversion(x, x_min=0.0, x_max=1.0, new_min=0.0, new_max=1.0): return (((x - x_min)/(x_max - x_min)) * (new_max - new_min)) + new_min +def broadcast_image_to_full(tensor, target_batch_size, batched_number, except_one=True): + current_batch_size = tensor.shape[0] + #print(current_batch_size, target_batch_size) + if except_one and current_batch_size == 1: + return tensor + + per_batch = target_batch_size // batched_number + tensor = tensor[:per_batch] + + if per_batch > tensor.shape[0]: + tensor = torch.cat([tensor] * (per_batch // tensor.shape[0]) + [tensor[:(per_batch % tensor.shape[0])]], dim=0) + + current_batch_size = tensor.shape[0] + if current_batch_size == target_batch_size: + return tensor + else: + return torch.cat([tensor] * batched_number, dim=0) + + +def ddpm_noise_latents(latents: Tensor, sigma: float, noise: Tensor=None): + alpha_cumprod = 1 / ((sigma * sigma) + 1) + sqrt_alpha_prod = alpha_cumprod ** 0.5 + sqrt_one_minus_alpha_prod = (1 - alpha_cumprod) ** 0.5 + if noise is None: + noise = torch.rand_like(latents) + return latents * sqrt_alpha_prod + noise * sqrt_one_minus_alpha_prod + + +def simple_noise_latents(latents: Tensor, sigma: float, noise: Tensor=None): + if noise is None: + noise = torch.rand_like(latents) + return latents + noise * sigma + + # from https://stackoverflow.com/a/24621200 def deepcopy_with_sharing(obj, shared_attribute_names, memo=None): ''' @@ -458,6 +525,14 @@ class AdvancedControlBase: def disarm(self): self.disarmed = True + def should_run(self): + if math.isclose(self.strength, 0.0) or math.isclose(self.current_timestep_keyframe.strength, 0.0): + return False + if self.timestep_range is not None: + if self.t > self.timestep_range[0] or self.t < self.timestep_range[1]: + return False + return True + def get_control_inject(self, x_noisy, t, cond, batched_number): # prepare timestep and everything related self.prepare_current_timestep(t=t, batched_number=batched_number)