878 lines
39 KiB
Python
878 lines
39 KiB
Python
from copy import deepcopy
|
|
from typing import Callable, Union
|
|
import torch
|
|
from torch import Tensor
|
|
import torch.nn.functional
|
|
from einops import rearrange
|
|
import numpy as np
|
|
import math
|
|
|
|
import comfy.ops
|
|
import comfy.utils
|
|
|
|
from comfy.controlnet import ControlBase
|
|
from comfy.model_patcher import ModelPatcher
|
|
from comfy.sd import VAE
|
|
|
|
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
|
|
comfy.utils.load_torch_file = orig_load_torch_file
|
|
return controlnet_data
|
|
return load_torch_file_with_dict
|
|
|
|
|
|
def get_properly_arranged_t2i_weights(initial_weights: list[float]):
|
|
new_weights = []
|
|
new_weights.extend([initial_weights[0]]*3)
|
|
new_weights.extend([initial_weights[1]]*3)
|
|
new_weights.extend([initial_weights[2]]*3)
|
|
new_weights.extend([initial_weights[3]]*3)
|
|
return new_weights
|
|
|
|
|
|
class ControlWeightType:
|
|
DEFAULT = "default"
|
|
UNIVERSAL = "universal"
|
|
T2IADAPTER = "t2iadapter"
|
|
CONTROLNET = "controlnet"
|
|
CONTROLNETPLUSPLUS = "controlnet++"
|
|
CONTROLLORA = "controllora"
|
|
CONTROLLLLITE = "controllllite"
|
|
SVD_CONTROLNET = "svd_controlnet"
|
|
SPARSECTRL = "sparsectrl"
|
|
|
|
|
|
class ControlWeights:
|
|
def __init__(self, weight_type: str, base_multiplier: float=1.0,
|
|
weights_input: list[float]=None, weights_middle: list[float]=None, weights_output: list[float]=None,
|
|
weight_func: Callable=None, weight_mask: Tensor=None,
|
|
uncond_multiplier=1.0, uncond_mask: Tensor=None, extras: dict[str]={},):
|
|
self.weight_type = weight_type
|
|
self.base_multiplier = base_multiplier
|
|
self.weights_input = weights_input
|
|
self.weights_middle = weights_middle
|
|
self.weights_output = weights_output
|
|
self.weight_func = weight_func
|
|
self.weight_mask = weight_mask
|
|
self.uncond_multiplier = float(uncond_multiplier)
|
|
self.has_uncond_multiplier = not math.isclose(self.uncond_multiplier, 1.0)
|
|
self.uncond_mask = uncond_mask if uncond_mask is not None else 1.0
|
|
self.has_uncond_mask = uncond_mask is not None
|
|
self.extras = extras
|
|
|
|
def get(self, idx: int, control: dict[str, list[Tensor]], key: str, default=1.0) -> Union[float, Tensor]:
|
|
# if weight_func present, use it
|
|
if self.weight_func is not None:
|
|
return self.weight_func(idx=idx, control=control, key=key)
|
|
# if weights is not none, return index
|
|
relevant_weights = None
|
|
if key == "middle":
|
|
relevant_weights = self.weights_middle
|
|
elif key == "input":
|
|
relevant_weights = self.weights_input
|
|
if relevant_weights is not None:
|
|
relevant_weights = list(reversed(relevant_weights))
|
|
else:
|
|
relevant_weights = self.weights_output
|
|
if relevant_weights is None:
|
|
return default
|
|
elif idx >= len(relevant_weights):
|
|
return default
|
|
return relevant_weights[idx]
|
|
|
|
def copy_with_new_weights(self, new_weights_input: list[float]=None, new_weights_middle: list[float]=None, new_weights_output: list[float]=None,
|
|
new_weight_func: Callable=None):
|
|
return ControlWeights(weight_type=self.weight_type, base_multiplier=self.base_multiplier,
|
|
weights_input=new_weights_input, weights_middle=new_weights_middle, weights_output=new_weights_output,
|
|
weight_func=new_weight_func, weight_mask=self.weight_mask,
|
|
uncond_multiplier=self.uncond_multiplier, extras=self.extras)
|
|
|
|
@classmethod
|
|
def default(cls, extras: dict[str]={}):
|
|
return cls(ControlWeightType.DEFAULT, extras=extras)
|
|
|
|
@classmethod
|
|
def universal(cls, base_multiplier: float, uncond_multiplier: float=1.0, extras: dict[str]={}):
|
|
return cls(ControlWeightType.UNIVERSAL, base_multiplier=base_multiplier, uncond_multiplier=uncond_multiplier, extras=extras)
|
|
|
|
@classmethod
|
|
def universal_mask(cls, weight_mask: Tensor, uncond_multiplier: float=1.0, extras: dict[str]={}):
|
|
return cls(ControlWeightType.UNIVERSAL, weight_mask=weight_mask, uncond_multiplier=uncond_multiplier, extras=extras)
|
|
|
|
@classmethod
|
|
def t2iadapter(cls, weights_input: list[float]=None, uncond_multiplier: float=1.0, extras: dict[str]={}):
|
|
return cls(ControlWeightType.T2IADAPTER, weights_input=weights_input, uncond_multiplier=uncond_multiplier, extras=extras)
|
|
|
|
@classmethod
|
|
def controlnet(cls, weights_output: list[float]=None, weights_middle: list[float]=None, weights_input: list[float]=None, uncond_multiplier: float=1.0, extras: dict[str]={}):
|
|
return cls(ControlWeightType.CONTROLNET, weights_output=weights_output, weights_middle=weights_middle, weights_input=weights_input, uncond_multiplier=uncond_multiplier, extras=extras)
|
|
|
|
@classmethod
|
|
def controllora(cls, weights_output: list[float]=None, weights_middle: list[float]=None, weights_input: list[float]=None, uncond_multiplier: float=1.0, extras: dict[str]={}):
|
|
return cls(ControlWeightType.CONTROLLORA, weights_output=weights_output, weights_middle=weights_middle, weights_input=weights_input, uncond_multiplier=uncond_multiplier, extras=extras)
|
|
|
|
@classmethod
|
|
def controllllite(cls, weights_output: list[float]=None, weights_middle: list[float]=None, weights_input: list[float]=None, uncond_multiplier: float=1.0, extras: dict[str]={}):
|
|
return cls(ControlWeightType.CONTROLLLLITE, weights_output=weights_output, weights_middle=weights_middle, weights_input=weights_input, uncond_multiplier=uncond_multiplier, extras=extras)
|
|
|
|
|
|
class StrengthInterpolation:
|
|
LINEAR = "linear"
|
|
EASE_IN = "ease-in"
|
|
EASE_OUT = "ease-out"
|
|
EASE_IN_OUT = "ease-in-out"
|
|
NONE = "none"
|
|
|
|
_LIST = [LINEAR, EASE_IN, EASE_OUT, EASE_IN_OUT]
|
|
_LIST_WITH_NONE = [LINEAR, EASE_IN, EASE_OUT, EASE_IN_OUT, NONE]
|
|
|
|
@classmethod
|
|
def get_weights(cls, num_from: float, num_to: float, length: int, method: str, reverse=False):
|
|
diff = num_to - num_from
|
|
if method == cls.LINEAR:
|
|
weights = torch.linspace(num_from, num_to, length)
|
|
elif method == cls.EASE_IN:
|
|
index = torch.linspace(0, 1, length)
|
|
weights = diff * np.power(index, 2) + num_from
|
|
elif method == cls.EASE_OUT:
|
|
index = torch.linspace(0, 1, length)
|
|
weights = diff * (1 - np.power(1 - index, 2)) + num_from
|
|
elif method == cls.EASE_IN_OUT:
|
|
index = torch.linspace(0, 1, length)
|
|
weights = diff * ((1 - np.cos(index * np.pi)) / 2) + num_from
|
|
else:
|
|
raise ValueError(f"Unrecognized interpolation method '{method}'.")
|
|
if reverse:
|
|
weights = weights.flip(dims=(0,))
|
|
return weights
|
|
|
|
|
|
class LatentKeyframe:
|
|
def __init__(self, batch_index: int, strength: float) -> None:
|
|
self.batch_index = batch_index
|
|
self.strength = strength
|
|
|
|
|
|
# always maintain sorted state (by batch_index of LatentKeyframe)
|
|
class LatentKeyframeGroup:
|
|
def __init__(self) -> None:
|
|
self.keyframes: list[LatentKeyframe] = []
|
|
|
|
def add(self, keyframe: LatentKeyframe) -> None:
|
|
added = False
|
|
# replace existing keyframe if same batch_index
|
|
for i in range(len(self.keyframes)):
|
|
if self.keyframes[i].batch_index == keyframe.batch_index:
|
|
self.keyframes[i] = keyframe
|
|
added = True
|
|
break
|
|
if not added:
|
|
self.keyframes.append(keyframe)
|
|
self.keyframes.sort(key=lambda k: k.batch_index)
|
|
|
|
def get_index(self, index: int) -> Union[LatentKeyframe, None]:
|
|
try:
|
|
return self.keyframes[index]
|
|
except IndexError:
|
|
return None
|
|
|
|
def __getitem__(self, index) -> LatentKeyframe:
|
|
return self.keyframes[index]
|
|
|
|
def is_empty(self) -> bool:
|
|
return len(self.keyframes) == 0
|
|
|
|
def clone(self) -> 'LatentKeyframeGroup':
|
|
cloned = LatentKeyframeGroup()
|
|
for tk in self.keyframes:
|
|
cloned.add(tk)
|
|
return cloned
|
|
|
|
|
|
class TimestepKeyframe:
|
|
def __init__(self,
|
|
start_percent: float = 0.0,
|
|
strength: float = 1.0,
|
|
control_weights: ControlWeights = None,
|
|
latent_keyframes: LatentKeyframeGroup = None,
|
|
null_latent_kf_strength: float = 0.0,
|
|
inherit_missing: bool = True,
|
|
guarantee_steps: int = 1,
|
|
mask_hint_orig: Tensor = None) -> None:
|
|
self.start_percent = float(start_percent)
|
|
self.start_t = 999999999.9
|
|
self.strength = strength
|
|
self.control_weights = control_weights
|
|
self.latent_keyframes = latent_keyframes
|
|
self.null_latent_kf_strength = null_latent_kf_strength
|
|
self.inherit_missing = inherit_missing
|
|
self.guarantee_steps = guarantee_steps
|
|
self.mask_hint_orig = mask_hint_orig
|
|
|
|
def has_control_weights(self):
|
|
return self.control_weights is not None
|
|
|
|
def has_latent_keyframes(self):
|
|
return self.latent_keyframes is not None
|
|
|
|
def has_mask_hint(self):
|
|
return self.mask_hint_orig is not None
|
|
|
|
|
|
@staticmethod
|
|
def default() -> 'TimestepKeyframe':
|
|
return TimestepKeyframe(start_percent=0.0, guarantee_steps=0)
|
|
|
|
|
|
# always maintain sorted state (by start_percent of TimestepKeyFrame)
|
|
class TimestepKeyframeGroup:
|
|
def __init__(self) -> None:
|
|
self.keyframes: list[TimestepKeyframe] = []
|
|
self.keyframes.append(TimestepKeyframe.default())
|
|
|
|
def add(self, keyframe: TimestepKeyframe) -> None:
|
|
# add to end of list, then sort
|
|
self.keyframes.append(keyframe)
|
|
self.keyframes = get_sorted_list_via_attr(self.keyframes, attr="start_percent")
|
|
|
|
def get_index(self, index: int) -> Union[TimestepKeyframe, None]:
|
|
try:
|
|
return self.keyframes[index]
|
|
except IndexError:
|
|
return None
|
|
|
|
def has_index(self, index: int) -> int:
|
|
return index >=0 and index < len(self.keyframes)
|
|
|
|
def __getitem__(self, index) -> TimestepKeyframe:
|
|
return self.keyframes[index]
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.keyframes)
|
|
|
|
def is_empty(self) -> bool:
|
|
return len(self.keyframes) == 0
|
|
|
|
def clone(self) -> 'TimestepKeyframeGroup':
|
|
cloned = TimestepKeyframeGroup()
|
|
# already sorted, so don't use add function to make cloning quicker
|
|
for tk in self.keyframes:
|
|
cloned.keyframes.append(tk)
|
|
return cloned
|
|
|
|
@classmethod
|
|
def default(cls, keyframe: TimestepKeyframe) -> 'TimestepKeyframeGroup':
|
|
group = cls()
|
|
group.keyframes[0] = keyframe
|
|
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):
|
|
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):
|
|
def forward_comfy_cast_weights(self, input):
|
|
weight, bias = comfy.ops.cast_bias_weight(self, input)
|
|
return torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps)
|
|
|
|
def forward(self, input):
|
|
if self.comfy_cast_weights:
|
|
return self.forward_comfy_cast_weights(input)
|
|
else:
|
|
return torch.nn.functional.group_norm(input, self.num_groups, self.weight, self.bias, self.eps)
|
|
|
|
class manual_cast_clean_groupnorm(comfy.ops.manual_cast):
|
|
class GroupNorm(disable_weight_init_clean_groupnorm.GroupNorm):
|
|
comfy_cast_weights = True
|
|
|
|
|
|
# adapted from comfy/sample.py
|
|
def prepare_mask_batch(mask: Tensor, shape: Tensor, multiplier: int=1, match_dim1=False, match_shape=False, flux_shape=None):
|
|
mask = mask.clone()
|
|
if flux_shape is not None:
|
|
multiplier = multiplier * 0.5
|
|
mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(round(flux_shape[-2]*multiplier), round(flux_shape[-1]*multiplier)), mode="bilinear")
|
|
mask = rearrange(mask, "b c h w -> b (h w) c")
|
|
else:
|
|
mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(round(shape[-2]*multiplier), round(shape[-1]*multiplier)), mode="bilinear")
|
|
if match_dim1:
|
|
if match_shape and len(shape) < 4:
|
|
raise Exception(f"match_dim1 cannot be True if shape is under 4 dims; was {len(shape)}.")
|
|
mask = torch.cat([mask] * shape[1], dim=1)
|
|
if match_shape and len(shape) == 3 and len(mask.shape) != 3:
|
|
mask = mask.squeeze(1)
|
|
return mask
|
|
|
|
|
|
# applies min-max normalization, from:
|
|
# https://stackoverflow.com/questions/68791508/min-max-normalization-of-a-tensor-in-pytorch
|
|
def normalize_min_max(x: Tensor, new_min = 0.0, new_max = 1.0):
|
|
x_min, x_max = x.min(), x.max()
|
|
return (((x - x_min)/(x_max - x_min)) * (new_max - new_min)) + new_min
|
|
|
|
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 extend_to_batch_size(tensor: Tensor, batch_size: int):
|
|
if tensor.shape[0] > batch_size:
|
|
return tensor[:batch_size]
|
|
elif tensor.shape[0] < batch_size:
|
|
remainder = batch_size-tensor.shape[0]
|
|
return torch.cat([tensor] + [tensor[-1:]]*remainder, dim=0)
|
|
return tensor
|
|
|
|
def broadcast_image_to_extend(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 = extend_to_batch_size(tensor=tensor, batch_size=per_batch)
|
|
|
|
current_batch_size = tensor.shape[0]
|
|
if current_batch_size == target_batch_size:
|
|
return tensor
|
|
else:
|
|
return torch.cat([tensor] * batched_number, dim=0)
|
|
|
|
|
|
# from https://stackoverflow.com/a/24621200
|
|
def deepcopy_with_sharing(obj, shared_attribute_names, memo=None):
|
|
'''
|
|
Deepcopy an object, except for a given list of attributes, which should
|
|
be shared between the original object and its copy.
|
|
|
|
obj is some object
|
|
shared_attribute_names: A list of strings identifying the attributes that
|
|
should be shared between the original and its copy.
|
|
memo is the dictionary passed into __deepcopy__. Ignore this argument if
|
|
not calling from within __deepcopy__.
|
|
'''
|
|
assert isinstance(shared_attribute_names, (list, tuple))
|
|
|
|
shared_attributes = {k: getattr(obj, k) for k in shared_attribute_names}
|
|
|
|
if hasattr(obj, '__deepcopy__'):
|
|
# Do hack to prevent infinite recursion in call to deepcopy
|
|
deepcopy_method = obj.__deepcopy__
|
|
obj.__deepcopy__ = None
|
|
|
|
for attr in shared_attribute_names:
|
|
del obj.__dict__[attr]
|
|
|
|
clone = deepcopy(obj)
|
|
|
|
for attr, val in shared_attributes.items():
|
|
setattr(obj, attr, val)
|
|
setattr(clone, attr, val)
|
|
|
|
if hasattr(obj, '__deepcopy__'):
|
|
# Undo hack
|
|
obj.__deepcopy__ = deepcopy_method
|
|
del clone.__deepcopy__
|
|
|
|
return clone
|
|
|
|
|
|
def get_sorted_list_via_attr(objects: list, attr: str) -> list:
|
|
if not objects:
|
|
return objects
|
|
elif len(objects) <= 1:
|
|
return [x for x in objects]
|
|
# now that we know we have to sort, do it following these rules:
|
|
# a) if objects have same value of attribute, maintain their relative order
|
|
# b) perform sorting of the groups of objects with same attributes
|
|
unique_attrs = {}
|
|
for o in objects:
|
|
val_attr = getattr(o, attr)
|
|
attr_list: list = unique_attrs.get(val_attr, list())
|
|
attr_list.append(o)
|
|
if val_attr not in unique_attrs:
|
|
unique_attrs[val_attr] = attr_list
|
|
# now that we have the unique attr values grouped together in relative order, sort them by key
|
|
sorted_attrs = dict(sorted(unique_attrs.items()))
|
|
# now flatten out the dict into a list to return
|
|
sorted_list = []
|
|
for object_list in sorted_attrs.values():
|
|
sorted_list.extend(object_list)
|
|
return sorted_list
|
|
|
|
|
|
# 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
|
|
|
|
|
|
class WeightTypeException(TypeError):
|
|
"Raised when weight not compatible with AdvancedControlBase object"
|
|
pass
|
|
|
|
|
|
class AdvancedControlBase:
|
|
def __init__(self, base: ControlBase, timestep_keyframes: TimestepKeyframeGroup, weights_default: ControlWeights, require_model=False, require_vae=False, allow_condhint_latents=False):
|
|
self.base = base
|
|
self.compatible_weights = [ControlWeightType.UNIVERSAL, ControlWeightType.DEFAULT]
|
|
self.add_compatible_weight(weights_default.weight_type)
|
|
# mask for which parts of controlnet output to keep
|
|
self.mask_cond_hint_original = None
|
|
self.mask_cond_hint = None
|
|
self.tk_mask_cond_hint_original = None
|
|
self.tk_mask_cond_hint = None
|
|
self.weight_mask_cond_hint = None
|
|
# actual index values
|
|
self.sub_idxs = None
|
|
self.full_latent_length = 0
|
|
self.context_length = 0
|
|
# timesteps
|
|
self.t: float = None
|
|
self.prev_t: float = None
|
|
self.batched_number: int = None
|
|
self.batch_size: int = 0
|
|
self.cond_or_uncond: list[int] = None
|
|
# weights + override
|
|
self.weights: ControlWeights = None
|
|
self.weights_default: ControlWeights = weights_default
|
|
self.weights_override: ControlWeights = None
|
|
# latent keyframe + override
|
|
self.latent_keyframes: LatentKeyframeGroup = None
|
|
self.latent_keyframe_override: LatentKeyframeGroup = None
|
|
# initialize timestep_keyframes
|
|
self.set_timestep_keyframes(timestep_keyframes)
|
|
# override some functions
|
|
self.get_control = self.get_control_inject
|
|
self.control_merge = self.control_merge_inject
|
|
self.pre_run = self.pre_run_inject
|
|
self.cleanup = self.cleanup_inject
|
|
self.set_previous_controlnet = self.set_previous_controlnet_inject
|
|
self.set_cond_hint = self.set_cond_hint_inject
|
|
# vae to store
|
|
self.adv_vae = None
|
|
# require model/vae to be passed into Apply Advanced ControlNet 🛂🅐🅒🅝 node
|
|
self.require_model = require_model
|
|
self.require_vae = require_vae
|
|
self.allow_condhint_latents = allow_condhint_latents
|
|
# disarm - when set to False, used to force usage of Apply Advanced ControlNet 🛂🅐🅒🅝 node (which will set it to True)
|
|
self.disarmed = not require_model
|
|
|
|
def patch_model(self, model: ModelPatcher):
|
|
pass
|
|
|
|
def add_compatible_weight(self, control_weight_type: str):
|
|
self.compatible_weights.append(control_weight_type)
|
|
|
|
def verify_all_weights(self, throw_error=True):
|
|
# first, check if override exists - if so, only need to check the override
|
|
if self.weights_override is not None:
|
|
if self.weights_override.weight_type not in self.compatible_weights:
|
|
msg = f"Weight override is type {self.weights_override.weight_type}, but loaded {type(self).__name__}" + \
|
|
f"only supports {self.compatible_weights} weights."
|
|
raise WeightTypeException(msg)
|
|
# otherwise, check all timestep keyframe weights
|
|
else:
|
|
for tk in self.timestep_keyframes.keyframes:
|
|
if tk.has_control_weights() and tk.control_weights.weight_type not in self.compatible_weights:
|
|
msg = f"Weight on Timestep Keyframe with start_percent={tk.start_percent} is type " + \
|
|
f"{tk.control_weights.weight_type}, but loaded {type(self).__name__} only supports {self.compatible_weights} weights."
|
|
raise WeightTypeException(msg)
|
|
|
|
def set_timestep_keyframes(self, timestep_keyframes: TimestepKeyframeGroup):
|
|
self.timestep_keyframes = timestep_keyframes if timestep_keyframes else TimestepKeyframeGroup()
|
|
# prepare first timestep_keyframe related stuff
|
|
self._current_timestep_keyframe = None
|
|
self._current_timestep_index = -1
|
|
self._current_used_steps = 0
|
|
self.weights = None
|
|
self.latent_keyframes = None
|
|
|
|
def prepare_current_timestep(self, t: Tensor, batched_number: int=1):
|
|
self.t = float(t[0])
|
|
# 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
|
|
# if met guaranteed steps (or no current keyframe), look for next keyframe in case need to switch
|
|
if self._current_timestep_keyframe is None or self._current_used_steps >= self._current_timestep_keyframe.guarantee_steps:
|
|
# if has next index, loop through and see if need to switch
|
|
if self.timestep_keyframes.has_index(self._current_timestep_index+1):
|
|
for i in range(self._current_timestep_index+1, len(self.timestep_keyframes)):
|
|
eval_tk = self.timestep_keyframes[i]
|
|
# check if start percent is less or equal to curr_t
|
|
if eval_tk.start_t >= curr_t:
|
|
self._current_timestep_index = i
|
|
self._current_timestep_keyframe = eval_tk
|
|
self._current_used_steps = 0
|
|
# keep track of control weights, latent keyframes, and masks,
|
|
# accounting for inherit_missing
|
|
if self._current_timestep_keyframe.has_control_weights():
|
|
self.weights = self._current_timestep_keyframe.control_weights
|
|
elif not self._current_timestep_keyframe.inherit_missing:
|
|
self.weights = self.weights_default
|
|
if self._current_timestep_keyframe.has_latent_keyframes():
|
|
self.latent_keyframes = self._current_timestep_keyframe.latent_keyframes
|
|
elif not self._current_timestep_keyframe.inherit_missing:
|
|
self.latent_keyframes = None
|
|
if self._current_timestep_keyframe.has_mask_hint():
|
|
self.tk_mask_cond_hint_original = self._current_timestep_keyframe.mask_hint_orig
|
|
elif not self._current_timestep_keyframe.inherit_missing:
|
|
del self.tk_mask_cond_hint_original
|
|
self.tk_mask_cond_hint_original = None
|
|
# if guarantee_steps greater than zero, stop searching for other keyframes
|
|
if self._current_timestep_keyframe.guarantee_steps > 0:
|
|
break
|
|
# 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
|
|
if prev_index != self._current_timestep_index:
|
|
if self.weights_override is not None:
|
|
self.weights = self.weights_override
|
|
if self.latent_keyframe_override is not None:
|
|
self.latent_keyframes = self.latent_keyframe_override
|
|
|
|
# make sure weights and latent_keyframes are in a workable state
|
|
# Note: each AdvancedControlBase should create their own get_universal_weights class
|
|
self.prepare_weights()
|
|
|
|
def prepare_weights(self):
|
|
if self.weights is None:
|
|
self.weights = self.weights_default
|
|
elif self.weights.weight_type == ControlWeightType.UNIVERSAL:
|
|
# if universal and weight_mask present, no need to convert
|
|
if self.weights.weight_mask is not None:
|
|
return
|
|
self.weights = self.get_universal_weights()
|
|
|
|
def get_universal_weights(self) -> ControlWeights:
|
|
return self.weights
|
|
|
|
def set_cond_hint_mask(self, mask_hint):
|
|
self.mask_cond_hint_original = mask_hint
|
|
return self
|
|
|
|
def set_cond_hint_inject(self, *args, **kwargs):
|
|
to_return = self.base.set_cond_hint(*args, **kwargs)
|
|
# if vae required, look in args and kwargs for it
|
|
if self.require_vae:
|
|
# check args first, as that's the default way vae param is used in ComfyUI
|
|
for arg in args:
|
|
if isinstance(arg, VAE):
|
|
self.adv_vae = arg
|
|
break
|
|
# if not in args, check kwargs now
|
|
if self.adv_vae is None:
|
|
if 'vae' in kwargs:
|
|
self.adv_vae = kwargs['vae']
|
|
return to_return
|
|
|
|
def pre_run_inject(self, model, percent_to_timestep_function):
|
|
self.base.pre_run(model, percent_to_timestep_function)
|
|
self.pre_run_advanced(model, percent_to_timestep_function)
|
|
|
|
def pre_run_advanced(self, model, percent_to_timestep_function):
|
|
# for each timestep keyframe, calculate the start_t
|
|
for tk in self.timestep_keyframes.keyframes:
|
|
tk.start_t = percent_to_timestep_function(tk.start_percent)
|
|
# clear variables
|
|
self.cleanup_advanced()
|
|
|
|
def set_previous_controlnet_inject(self, *args, **kwargs):
|
|
to_return = self.base.set_previous_controlnet(*args, **kwargs)
|
|
if not self.disarmed:
|
|
raise Exception(f"Type '{type(self).__name__}' must be used with Apply Advanced ControlNet 🛂🅐🅒🅝 node (with model_optional passed in); otherwise, it will not work.")
|
|
return to_return
|
|
|
|
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, transformer_options: dict):
|
|
self.batched_number = batched_number
|
|
self.batch_size = len(t)
|
|
self.cond_or_uncond = transformer_options.get("cond_or_uncond", None)
|
|
# 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
|
|
if self.strength == 0.0 or self._current_timestep_keyframe.strength == 0.0:
|
|
return self.default_control_actions(x_noisy, t, cond, batched_number)
|
|
# otherwise, perform normal function
|
|
return self.get_control_advanced(x_noisy, t, cond, batched_number, transformer_options)
|
|
|
|
def get_control_advanced(self, x_noisy, t, cond, batched_number, transformer_options):
|
|
return self.default_control_actions(x_noisy, t, cond, batched_number, transformer_options)
|
|
|
|
def default_control_actions(self, x_noisy, t, cond, batched_number, transformer_options):
|
|
control_prev = None
|
|
if self.previous_controlnet is not None:
|
|
control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number, transformer_options)
|
|
return control_prev
|
|
|
|
def calc_weight(self, idx: int, x: Tensor, control: dict[str, list[Tensor]], key: str) -> Union[float, Tensor]:
|
|
if self.weights.weight_mask is not None:
|
|
# prepare weight mask
|
|
self.prepare_weight_mask_cond_hint(x, self.batched_number)
|
|
# adjust mask for current layer and return
|
|
return torch.pow(self.weight_mask_cond_hint, self.get_calc_pow(idx=idx, control=control, key=key))
|
|
return self.weights.get(idx=idx, control=control, key=key)
|
|
|
|
def get_calc_pow(self, idx: int, control: dict[str, list[Tensor]], key: str) -> int:
|
|
if key == "middle":
|
|
return 0
|
|
else:
|
|
c_len = len(control[key])
|
|
real_idx = c_len-idx
|
|
if key == "input":
|
|
real_idx = c_len - real_idx + 1
|
|
return real_idx
|
|
|
|
def calc_latent_keyframe_mults(self, x: Tensor, batched_number: int) -> Tensor:
|
|
# apply strengths, and get batch indeces to null out
|
|
# AKA latents that should not be influenced by ControlNet
|
|
final_mults = [1.0] * x.shape[0]
|
|
if self.latent_keyframes:
|
|
latent_count = x.shape[0] // batched_number
|
|
indeces_to_null = set(range(latent_count))
|
|
mapped_indeces = None
|
|
# if expecting subdivision, will need to translate between subset and actual idx values
|
|
if self.sub_idxs:
|
|
mapped_indeces = {}
|
|
for i, actual in enumerate(self.sub_idxs):
|
|
mapped_indeces[actual] = i
|
|
for keyframe in self.latent_keyframes:
|
|
real_index = keyframe.batch_index
|
|
# if negative, count from end
|
|
if real_index < 0:
|
|
real_index += latent_count if self.sub_idxs is None else self.full_latent_length
|
|
|
|
# if not mapping indeces, what you see is what you get
|
|
if mapped_indeces is None:
|
|
if real_index in indeces_to_null:
|
|
indeces_to_null.remove(real_index)
|
|
# otherwise, see if batch_index is even included in this set of latents
|
|
else:
|
|
real_index = mapped_indeces.get(real_index, None)
|
|
if real_index is None:
|
|
continue
|
|
indeces_to_null.remove(real_index)
|
|
|
|
# if real_index is outside the bounds of latents, don't apply
|
|
if real_index >= latent_count or real_index < 0:
|
|
continue
|
|
|
|
# apply strength for each batched cond/uncond
|
|
for b in range(batched_number):
|
|
final_mults[(latent_count*b)+real_index] = keyframe.strength
|
|
# null them out by multiplying by null_latent_kf_strength
|
|
for batch_index in indeces_to_null:
|
|
# apply null for each batched cond/uncond
|
|
for b in range(batched_number):
|
|
final_mults[(latent_count*b)+batch_index] = self._current_timestep_keyframe.null_latent_kf_strength
|
|
# convert final_mults into tensor and match expected dimension count
|
|
final_tensor = torch.tensor(final_mults, dtype=x.dtype, device=x.device)
|
|
while len(final_tensor.shape) < len(x.shape):
|
|
final_tensor = final_tensor.unsqueeze(-1)
|
|
return final_tensor
|
|
|
|
def apply_advanced_strengths_and_masks(self, x: Tensor, batched_number: int, flux_shape: tuple=None):
|
|
# handle weight's uncond_multiplier, if applicable
|
|
if self.weights.has_uncond_multiplier:
|
|
actual_length = x.size(0) // batched_number
|
|
for idx, cond_type in enumerate(self.cond_or_uncond):
|
|
# if uncond, set to weight's uncond_multiplier
|
|
if cond_type == 1:
|
|
x[actual_length*idx:actual_length*(idx+1)] *= self.weights.uncond_multiplier
|
|
if self.weights.has_uncond_mask:
|
|
pass
|
|
|
|
if self.latent_keyframes is not None:
|
|
x[:] = x[:] * self.calc_latent_keyframe_mults(x=x, batched_number=batched_number)
|
|
# apply masks, resizing mask to required dims
|
|
if self.mask_cond_hint is not None:
|
|
masks = prepare_mask_batch(self.mask_cond_hint, x.shape, match_shape=True, flux_shape=flux_shape)
|
|
x[:] = x[:] * masks
|
|
if self.tk_mask_cond_hint is not None:
|
|
masks = prepare_mask_batch(self.tk_mask_cond_hint, x.shape, match_shape=True, flux_shape=flux_shape)
|
|
x[:] = x[:] * masks
|
|
# apply timestep keyframe strengths
|
|
if self._current_timestep_keyframe.strength != 1.0:
|
|
x[:] *= self._current_timestep_keyframe.strength
|
|
|
|
def control_merge_inject(self: 'AdvancedControlBase', control: dict[str, list[Tensor]], control_prev: dict, output_dtype):
|
|
out = {'input':[], 'middle':[], 'output': []}
|
|
|
|
for key in control:
|
|
control_output = control[key]
|
|
applied_to = set()
|
|
for i in range(len(control_output)):
|
|
x = control_output[i]
|
|
if x is not None:
|
|
if self.global_average_pooling:
|
|
x = torch.mean(x, dim=(2, 3), keepdim=True).repeat(1, 1, x.shape[2], x.shape[3])
|
|
|
|
if x not in applied_to: #memory saving strategy, allow shared tensors and only apply strength to shared tensors once
|
|
applied_to.add(x)
|
|
self.apply_advanced_strengths_and_masks(x, self.batched_number)
|
|
x *= self.strength * self.calc_weight(i, x, control, key)
|
|
|
|
if output_dtype is not None and x.dtype != output_dtype:
|
|
x = x.to(output_dtype)
|
|
|
|
out[key].append(x)
|
|
|
|
if control_prev is not None:
|
|
for x in ['input', 'middle', 'output']:
|
|
o = out[x]
|
|
for i in range(len(control_prev[x])):
|
|
prev_val = control_prev[x][i]
|
|
if i >= len(o):
|
|
o.append(prev_val)
|
|
elif prev_val is not None:
|
|
if o[i] is None:
|
|
o[i] = prev_val
|
|
else:
|
|
if o[i].shape[0] < prev_val.shape[0]:
|
|
o[i] = prev_val + o[i]
|
|
else:
|
|
o[i] = prev_val + o[i] # TODO from base ComfyUI: change back to inplace add if shared tensors stop being an issue
|
|
return out
|
|
|
|
def prepare_mask_cond_hint(self, x_noisy: Tensor, t, cond, batched_number, dtype=None, direct_attn=False):
|
|
self._prepare_mask("mask_cond_hint", self.mask_cond_hint_original, x_noisy, t, cond, batched_number, dtype, direct_attn=direct_attn)
|
|
self.prepare_tk_mask_cond_hint(x_noisy, t, cond, batched_number, dtype, direct_attn=direct_attn)
|
|
|
|
def prepare_tk_mask_cond_hint(self, x_noisy: Tensor, t, cond, batched_number, dtype=None, direct_attn=False):
|
|
return self._prepare_mask("tk_mask_cond_hint", self._current_timestep_keyframe.mask_hint_orig, x_noisy, t, cond, batched_number, dtype, direct_attn=direct_attn)
|
|
|
|
def prepare_weight_mask_cond_hint(self, x_noisy: Tensor, batched_number, dtype=None):
|
|
return self._prepare_mask("weight_mask_cond_hint", self.weights.weight_mask, x_noisy, t=None, cond=None, batched_number=batched_number, dtype=dtype, direct_attn=True)
|
|
|
|
def _prepare_mask(self, attr_name, orig_mask: Tensor, x_noisy: Tensor, t, cond, batched_number, dtype=None, direct_attn=False):
|
|
# make mask appropriate dimensions, if present
|
|
if orig_mask is not None:
|
|
out_mask = getattr(self, attr_name)
|
|
multiplier = 1 if direct_attn else 8
|
|
if self.sub_idxs is not None or out_mask is None or x_noisy.shape[2] * multiplier != out_mask.shape[1] or x_noisy.shape[3] * multiplier != out_mask.shape[2]:
|
|
self._reset_attr(attr_name)
|
|
del out_mask
|
|
# TODO: perform upscale on only the sub_idxs masks at a time instead of all to conserve RAM
|
|
# resize mask and match batch count
|
|
out_mask = prepare_mask_batch(orig_mask, x_noisy.shape, multiplier=multiplier, match_shape=True)
|
|
actual_latent_length = x_noisy.shape[0] // batched_number
|
|
out_mask = extend_to_batch_size(out_mask, actual_latent_length if self.sub_idxs is None else self.full_latent_length)
|
|
if self.sub_idxs is not None:
|
|
out_mask = out_mask[self.sub_idxs]
|
|
# make cond_hint_mask length match x_noise
|
|
if x_noisy.shape[0] != out_mask.shape[0]:
|
|
out_mask = broadcast_image_to_extend(out_mask, x_noisy.shape[0], batched_number)
|
|
# default dtype to be same as x_noisy
|
|
if dtype is None:
|
|
dtype = x_noisy.dtype
|
|
setattr(self, attr_name, out_mask.to(dtype=dtype).to(x_noisy.device))
|
|
del out_mask
|
|
|
|
def _reset_attr(self, attr_name, new_value=None):
|
|
if hasattr(self, attr_name):
|
|
delattr(self, attr_name)
|
|
setattr(self, attr_name, new_value)
|
|
|
|
def cleanup_inject(self):
|
|
self.base.cleanup()
|
|
self.cleanup_advanced()
|
|
|
|
def cleanup_advanced(self):
|
|
self.sub_idxs = None
|
|
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
|
|
self.latent_keyframes = None
|
|
# timestep stuff
|
|
self._current_timestep_keyframe = None
|
|
self._current_timestep_index = -1
|
|
self._current_used_steps = 0
|
|
# clear mask hints
|
|
if self.mask_cond_hint is not None:
|
|
del self.mask_cond_hint
|
|
self.mask_cond_hint = None
|
|
if self.tk_mask_cond_hint_original is not None:
|
|
del self.tk_mask_cond_hint_original
|
|
self.tk_mask_cond_hint_original = None
|
|
if self.tk_mask_cond_hint is not None:
|
|
del self.tk_mask_cond_hint
|
|
self.tk_mask_cond_hint = None
|
|
if self.weight_mask_cond_hint is not None:
|
|
del self.weight_mask_cond_hint
|
|
self.weight_mask_cond_hint = None
|
|
|
|
def copy_to_advanced(self, copied: 'AdvancedControlBase'):
|
|
copied.mask_cond_hint_original = self.mask_cond_hint_original
|
|
copied.weights_override = self.weights_override
|
|
copied.latent_keyframe_override = self.latent_keyframe_override
|
|
copied.adv_vae = self.adv_vae
|
|
copied.require_vae = self.require_vae
|
|
copied.allow_condhint_latents = self.allow_condhint_latents
|
|
copied.disarmed = self.disarmed
|