266 lines
9.7 KiB
Python
266 lines
9.7 KiB
Python
from abc import ABC, abstractmethod
|
|
from typing import Union
|
|
from torch import Tensor
|
|
import math
|
|
|
|
from .utils_motion import normalize_min_max
|
|
from .logger import logger
|
|
|
|
|
|
class AnimateDiffSettings:
|
|
def __init__(self,
|
|
adjust_pe: 'AdjustGroup'=None,
|
|
adjust_weight: 'AdjustGroup'=None,
|
|
attn_scale: float=1.0,
|
|
mask_attn_scale: Tensor=None,
|
|
mask_attn_scale_min: float=1.0,
|
|
mask_attn_scale_max: float=1.0,
|
|
):
|
|
# PE-interpolation settings
|
|
self.adjust_pe = adjust_pe if adjust_pe is not None else AdjustGroup()
|
|
# Weight settings
|
|
self.adjust_weight = adjust_weight if adjust_weight is not None else AdjustGroup()
|
|
# attention scale settings - DEPRECATED (part of scale_multival now)
|
|
self.attn_scale = attn_scale
|
|
# attention scale mask settings - DEPRECATED (part of scale_multival now)
|
|
self.mask_attn_scale = mask_attn_scale.clone() if mask_attn_scale is not None else mask_attn_scale
|
|
self.mask_attn_scale_min = mask_attn_scale_min
|
|
self.mask_attn_scale_max = mask_attn_scale_max
|
|
self._prepare_mask_attn_scale()
|
|
|
|
def _prepare_mask_attn_scale(self):
|
|
if self.mask_attn_scale is not None:
|
|
self.mask_attn_scale = normalize_min_max(self.mask_attn_scale, self.mask_attn_scale_min, self.mask_attn_scale_max)
|
|
|
|
def has_mask_attn_scale(self) -> bool:
|
|
return self.mask_attn_scale is not None
|
|
|
|
def has_anything_to_apply(self) -> bool:
|
|
return self.adjust_pe.has_anything_to_apply() \
|
|
or self.adjust_weight.has_anything_to_apply()
|
|
|
|
|
|
class AdjustAbstract(ABC):
|
|
def __init__(self, print_adjustment=False):
|
|
self.print_adjustment = print_adjustment
|
|
|
|
@abstractmethod
|
|
def has_anything_to_apply(self):
|
|
return False
|
|
|
|
|
|
class AdjustPE(AdjustAbstract):
|
|
def __init__(self,
|
|
cap_initial_pe_length: int=0, interpolate_pe_to_length: int=0,
|
|
initial_pe_idx_offset: int=0, final_pe_idx_offset: int=0,
|
|
motion_pe_stretch: int=0, print_adjustment=False):
|
|
super().__init__(print_adjustment=print_adjustment)
|
|
# PE-interpolation settings
|
|
self.cap_initial_pe_length = cap_initial_pe_length
|
|
self.interpolate_pe_to_length = interpolate_pe_to_length
|
|
self.initial_pe_idx_offset = initial_pe_idx_offset
|
|
self.final_pe_idx_offset = final_pe_idx_offset
|
|
self.motion_pe_stretch = motion_pe_stretch
|
|
|
|
def has_cap_initial_pe_length(self) -> bool:
|
|
return self.cap_initial_pe_length > 0
|
|
|
|
def has_interpolate_pe_to_length(self) -> bool:
|
|
return self.interpolate_pe_to_length > 0
|
|
|
|
def has_initial_pe_idx_offset(self) -> bool:
|
|
return self.initial_pe_idx_offset > 0
|
|
|
|
def has_final_pe_idx_offset(self) -> bool:
|
|
return self.final_pe_idx_offset > 0
|
|
|
|
def has_motion_pe_stretch(self) -> bool:
|
|
return self.motion_pe_stretch > 0
|
|
|
|
def has_anything_to_apply(self) -> bool:
|
|
return self.has_cap_initial_pe_length() \
|
|
or self.has_interpolate_pe_to_length() \
|
|
or self.has_initial_pe_idx_offset() \
|
|
or self.has_final_pe_idx_offset() \
|
|
or self.has_motion_pe_stretch()
|
|
|
|
|
|
class AdjustWeight(AdjustAbstract):
|
|
# possible operations
|
|
OP_ANY = "_____ANY"
|
|
OP_ADD = "_ADD"
|
|
OP_MULT = "_MULT"
|
|
OPS = [OP_ADD, OP_MULT]
|
|
# possible attributes
|
|
ATTR_ALL = "all"
|
|
ATTR_PE = "pe"
|
|
ATTR_ATTN = "attn"
|
|
ATTR_ATTN_Q = "attn_q"
|
|
ATTR_ATTN_K = "attn_k"
|
|
ATTR_ATTN_V = "attn_v"
|
|
ATTR_ATTN_OUT_WEIGHT = "attn_out_weight"
|
|
ATTR_ATTN_OUT_BIAS = "attn_out_bias"
|
|
ATTR_OTHER = "other"
|
|
ATTRS = [ATTR_ALL, ATTR_PE, ATTR_ATTN, ATTR_ATTN_Q, ATTR_ATTN_K, ATTR_ATTN_V, ATTR_ATTN_OUT_WEIGHT, ATTR_ATTN_OUT_BIAS, ATTR_OTHER]
|
|
|
|
def __init__(self,
|
|
all_ADD=0.0, all_MULT=1.0,
|
|
pe_ADD=0.0, pe_MULT=1.0,
|
|
attn_ADD=0.0, attn_MULT=1.0,
|
|
attn_q_ADD=0.0, attn_q_MULT=1.0,
|
|
attn_k_ADD=0.0, attn_k_MULT=1.0,
|
|
attn_v_ADD=0.0, attn_v_MULT=1.0,
|
|
attn_out_weight_ADD=0.0, attn_out_weight_MULT=1.0,
|
|
attn_out_bias_ADD=0.0, attn_out_bias_MULT=1.0,
|
|
other_ADD=0.0, other_MULT=1.0,
|
|
print_adjustment=False):
|
|
# all
|
|
self.all_ADD = all_ADD
|
|
self.all_MULT = all_MULT
|
|
# pe
|
|
self.pe_ADD = pe_ADD
|
|
self.pe_MULT = pe_MULT
|
|
# attn
|
|
self.attn_ADD = attn_ADD
|
|
self.attn_MULT = attn_MULT
|
|
# attn_q
|
|
self.attn_q_ADD = attn_q_ADD
|
|
self.attn_q_MULT = attn_q_MULT
|
|
# attn_k
|
|
self.attn_k_ADD = attn_k_ADD
|
|
self.attn_k_MULT = attn_k_MULT
|
|
# attn_v
|
|
self.attn_v_ADD = attn_v_ADD
|
|
self.attn_v_MULT = attn_v_MULT
|
|
# attn_out_weight
|
|
self.attn_out_weight_ADD = attn_out_weight_ADD
|
|
self.attn_out_weight_MULT = attn_out_weight_MULT
|
|
# attn_out_bias
|
|
self.attn_out_bias_ADD = attn_out_bias_ADD
|
|
self.attn_out_bias_MULT = attn_out_bias_MULT
|
|
# other
|
|
self.other_ADD = other_ADD
|
|
self.other_MULT = other_MULT
|
|
# additional vars
|
|
self.print_adjustment = print_adjustment
|
|
# temp var
|
|
self.already_printed: dict[str, bool] = {}
|
|
self.mark_attrs_as_unprinted()
|
|
|
|
def mark_attrs_as_unprinted(self):
|
|
for attr in self.ATTRS:
|
|
for op in self.OPS:
|
|
self.already_printed[attr+op] = False
|
|
|
|
def mask_as_printed(self, attr: str, op: str):
|
|
self.already_printed[attr+op] = True
|
|
|
|
def is_already_printed(self, attr: str, op: str):
|
|
return self.already_printed.get(attr+op, False)
|
|
|
|
def _get_val(self, op: str, attr: str) -> float:
|
|
try:
|
|
return getattr(self, attr+op)
|
|
except AttributeError:
|
|
raise Exception(f"Parameter '{attr+op}' could not be found in AdjustWeight class.")
|
|
|
|
def _has_OP(self, op: str, attr: str):
|
|
value = self._get_val(op=op, attr=attr)
|
|
if op == self.OP_ADD:
|
|
return not math.isclose(value, 0.0)
|
|
elif op == self.OP_MULT:
|
|
return not math.isclose(value, 1.0)
|
|
else:
|
|
raise Exception(f"Operation '{op}' not recognized in AdjustWeight.")
|
|
|
|
def _has_apply(self, op: str, attr: str):
|
|
# determine if attr with specific operation is to be applied
|
|
if op == self.OP_ANY:
|
|
any = False
|
|
for one_op in self.OPS:
|
|
any = any or self._has_OP(op=one_op, attr=attr)
|
|
return any
|
|
return self._has_OP(op=op, attr=attr)
|
|
|
|
def has_all(self, op: str) -> bool:
|
|
return self._has_apply(op, self.ATTR_ALL)
|
|
|
|
def has_pe(self, op: str) -> bool:
|
|
return self._has_apply(op, self.ATTR_PE)
|
|
|
|
def has_attn(self, op: str) -> bool:
|
|
return self._has_apply(op, self.ATTR_ATTN)
|
|
|
|
def has_attn_q(self, op: str) -> bool:
|
|
return self._has_apply(op, self.ATTR_ATTN_Q)
|
|
|
|
def has_attn_k(self, op: str) -> bool:
|
|
return self._has_apply(op, self.ATTR_ATTN_K)
|
|
|
|
def has_attn_v(self, op: str) -> bool:
|
|
return self._has_apply(op, self.ATTR_ATTN_V)
|
|
|
|
def has_attn_out_weight(self, op: str) -> bool:
|
|
return self._has_apply(op, self.ATTR_ATTN_OUT_WEIGHT)
|
|
|
|
def has_attn_out_bias(self, op: str) -> bool:
|
|
return self._has_apply(op, self.ATTR_ATTN_OUT_BIAS)
|
|
|
|
def has_other(self, op: str) -> bool:
|
|
return self._has_apply(op, self.ATTR_OTHER)
|
|
|
|
def has_anything_to_apply(self):
|
|
return self.has_all(self.OP_ANY) \
|
|
or self.has_pe(self.OP_ANY) \
|
|
or self.has_attn(self.OP_ANY) \
|
|
or self.has_attn_q(self.OP_ANY) \
|
|
or self.has_attn_k(self.OP_ANY) \
|
|
or self.has_attn_v(self.OP_ANY) \
|
|
or self.has_attn_out_weight(self.OP_ANY) \
|
|
or self.has_attn_out_bias(self.OP_ANY) \
|
|
or self.has_other(self.OP_ANY)
|
|
|
|
def _perform_op(self, model_dict: dict[str, Tensor], key: str, op: str, attr: str):
|
|
val = self._get_val(op=op, attr=attr)
|
|
specific_str = f"'{attr}' weights" if attr == self.ATTR_ALL else f"every '{attr}' weight"
|
|
if op == self.OP_ADD:
|
|
model_dict[key] += val
|
|
if self.print_adjustment and not self.is_already_printed(attr=attr, op=op):
|
|
logger.info(f"[Adjust Weight]: Adding to {specific_str} value {val}")
|
|
self.mask_as_printed(attr=attr, op=op)
|
|
elif op == self.OP_MULT:
|
|
model_dict[key] *= val
|
|
if self.print_adjustment and not self.is_already_printed(attr=attr, op=op):
|
|
logger.info(f"[Adjust Weight]: Multiplying {specific_str} by {val}")
|
|
self.mask_as_printed(attr=attr, op=op)
|
|
else:
|
|
raise Exception(f"Operation '{op}' not recognized in AdjustWeight.")
|
|
|
|
def perform_applicable_ops(self, attr: str, model_dict: dict[str, Tensor], key: str):
|
|
for op in self.OPS:
|
|
if self._has_apply(op=op, attr=attr):
|
|
self._perform_op(model_dict=model_dict, key=key, op=op, attr=attr)
|
|
|
|
|
|
ADJUST_TYPES = Union[AdjustPE, AdjustWeight]
|
|
class AdjustGroup:
|
|
def __init__(self, initial: ADJUST_TYPES=None):
|
|
self.adjusts: list[ADJUST_TYPES] = []
|
|
if initial is not None:
|
|
self.add(initial)
|
|
|
|
def add(self, adjust: ADJUST_TYPES):
|
|
self.adjusts.append(adjust)
|
|
|
|
def has_anything_to_apply(self):
|
|
for adjust in self.adjusts:
|
|
if adjust.has_anything_to_apply():
|
|
return True
|
|
return False
|
|
|
|
def clone(self):
|
|
new_group = AdjustGroup()
|
|
for adjust in self.adjusts:
|
|
new_group.add(adjust=adjust)
|
|
return new_group
|