608 lines
25 KiB
Python
608 lines
25 KiB
Python
from __future__ import annotations
|
|
from typing import TYPE_CHECKING, Union
|
|
import math
|
|
import torch
|
|
from torch import Tensor
|
|
|
|
from comfy.model_base import BaseModel
|
|
|
|
from .dinklink import get_acn_dinklink_version
|
|
from .utils_model import BIGMAX_TENSOR, MachineState
|
|
from .utils_motion import (prepare_mask_batch, extend_to_batch_size, get_combined_multival, resize_multival,
|
|
get_sorted_list_via_attr)
|
|
if TYPE_CHECKING:
|
|
from .sampling import AnimateDiffGlobalState
|
|
|
|
|
|
CONTEXTREF_VERSION = 1
|
|
|
|
|
|
class ContextExtra:
|
|
def __init__(self, start_percent: float, end_percent: float):
|
|
# scheduling
|
|
self.start_percent = float(start_percent)
|
|
self.start_t = 999999999.9
|
|
self.end_percent = float(end_percent)
|
|
self.end_t = 0.0
|
|
self.curr_t = 999999999.9
|
|
|
|
def initialize_timesteps(self, model: BaseModel):
|
|
self.start_t = model.model_sampling.percent_to_sigma(self.start_percent)
|
|
self.end_t = model.model_sampling.percent_to_sigma(self.end_percent)
|
|
|
|
def prepare_current(self, t: Tensor, transformer_options: dict[str, Tensor]):
|
|
self.curr_t = t[0]
|
|
|
|
def should_run(self):
|
|
if self.curr_t > self.start_t or self.curr_t < self.end_t:
|
|
return False
|
|
return True
|
|
|
|
def cleanup(self):
|
|
pass
|
|
|
|
|
|
################################
|
|
# ContextRef
|
|
class ContextRefTune:
|
|
def __init__(self,
|
|
attn_style_fidelity=0.0, attn_ref_weight=0.0, attn_strength=0.0,
|
|
adain_style_fidelity=0.0, adain_ref_weight=0.0, adain_strength=0.0):
|
|
# attn1
|
|
self.attn_style_fidelity = float(attn_style_fidelity)
|
|
self.attn_ref_weight = float(attn_ref_weight)
|
|
self.attn_strength = float(attn_strength)
|
|
# adain
|
|
self.adain_style_fidelity = float(adain_style_fidelity)
|
|
self.adain_ref_weight = float(adain_ref_weight)
|
|
self.adain_strength = float(adain_strength)
|
|
|
|
def create_dict(self):
|
|
return {
|
|
"attn_style_fidelity": self.attn_style_fidelity,
|
|
"attn_ref_weight": self.attn_ref_weight,
|
|
"attn_strength": self.attn_strength,
|
|
"adain_style_fidelity": self.adain_style_fidelity,
|
|
"adain_ref_weight": self.adain_ref_weight,
|
|
"adain_strength": self.adain_strength,
|
|
}
|
|
|
|
|
|
class ContextRefMode:
|
|
FIRST = "first"
|
|
SLIDING = "sliding"
|
|
INDEXES = "indexes"
|
|
_LIST = [FIRST, SLIDING, INDEXES]
|
|
|
|
def __init__(self, mode: str, sliding_width=2, indexes: set[int]=set([0])):
|
|
self.mode = mode
|
|
self.sliding_width = sliding_width
|
|
self.indexes = indexes
|
|
self.single_trigger = True
|
|
|
|
@classmethod
|
|
def init_first(cls):
|
|
return ContextRefMode(cls.FIRST)
|
|
|
|
@classmethod
|
|
def init_sliding(cls, sliding_width: int):
|
|
return ContextRefMode(cls.SLIDING, sliding_width=sliding_width)
|
|
|
|
@classmethod
|
|
def init_indexes(cls, indexes: set[int]):
|
|
return ContextRefMode(cls.INDEXES, indexes=indexes)
|
|
|
|
|
|
class ContextRefKeyframe:
|
|
def __init__(self, mult=1.0, mult_multival: Union[float, Tensor]=None, tune_replace: ContextRefTune=None, mode_replace: ContextRefMode=None,
|
|
start_percent=0.0, guarantee_steps=1, inherit_missing=True):
|
|
self.mult = mult
|
|
self.orig_mult_multival = mult_multival
|
|
self.orig_tune_replace = tune_replace
|
|
self.orig_mode_replace = mode_replace
|
|
self.mult_multival = self.orig_mult_multival
|
|
self.tune_replace = self.orig_tune_replace
|
|
self.mode_replace = self.orig_mode_replace
|
|
# scheduling
|
|
self.start_percent = float(start_percent)
|
|
self.guarantee_steps = guarantee_steps
|
|
self.inherit_missing = inherit_missing
|
|
|
|
def clone(self):
|
|
c = ContextRefKeyframe(mult=self.mult, mult_multival=self.orig_mult_multival, tune_replace=self.orig_tune_replace, mode_replace=self.orig_mode_replace,
|
|
start_percent=self.start_percent, guarantee_steps=self.guarantee_steps, inherit_missing=self.inherit_missing)
|
|
return c
|
|
|
|
|
|
class ContextRefKeyframeGroup:
|
|
def __init__(self):
|
|
self.keyframes: list[ContextRefKeyframe] = []
|
|
self._current_keyframe: NaiveReuseKeyframe = None
|
|
self._current_used_steps: int = 0
|
|
self._current_index: int = 0
|
|
self._previous_t = -1
|
|
|
|
def reset(self):
|
|
self._current_keyframe = None
|
|
self._current_used_steps = 0
|
|
self._current_index = 0
|
|
self._set_first_as_current()
|
|
|
|
def add(self, keyframe: ContextRefKeyframe):
|
|
# add to end of list, then sort
|
|
self.keyframes.append(keyframe)
|
|
self.keyframes = get_sorted_list_via_attr(self.keyframes, "start_percent")
|
|
self._set_first_as_current()
|
|
self._prepare_all_keyframe_vals()
|
|
|
|
def _set_first_as_current(self):
|
|
if len(self.keyframes) > 0:
|
|
self._current_keyframe = self.keyframes[0]
|
|
else:
|
|
self._current_keyframe = None
|
|
|
|
def _prepare_all_keyframe_vals(self):
|
|
if self.is_empty():
|
|
return
|
|
multival = None
|
|
tune = None
|
|
mode = None
|
|
for kf in self.keyframes:
|
|
# if shouldn't inherit, clear cache
|
|
if not kf.inherit_missing:
|
|
multival = None
|
|
tune = None
|
|
mode = None
|
|
# assign cached values, if origs were None
|
|
# Mult #################
|
|
if kf.orig_mult_multival is None:
|
|
kf.mult_multival = multival
|
|
else:
|
|
kf.mult_multival = kf.orig_mult_multival
|
|
# Tune #################
|
|
if kf.orig_tune_replace is None:
|
|
kf.tune_replace = tune
|
|
else:
|
|
kf.tune_replace = kf.orig_tune_replace
|
|
# Mode #################
|
|
if kf.orig_mode_replace is None:
|
|
kf.mode_replace = mode
|
|
else:
|
|
kf.mode_replace = kf.orig_mode_replace
|
|
# save new caches, in case next keyframe inherits missing
|
|
if kf.mult_multival is not None:
|
|
multival = kf.mult_multival
|
|
if kf.tune_replace is not None:
|
|
tune = kf.tune_replace
|
|
if kf.mode_replace is not None:
|
|
mode = kf.mode_replace
|
|
|
|
def has_index(self, index: int) -> int:
|
|
return index >=0 and index < len(self.keyframes)
|
|
|
|
def is_empty(self) -> bool:
|
|
return len(self.keyframes) == 0
|
|
|
|
def clone(self):
|
|
cloned = ContextRefKeyframeGroup()
|
|
for keyframe in self.keyframes:
|
|
cloned.keyframes.append(keyframe.clone())
|
|
cloned._set_first_as_current()
|
|
cloned._prepare_all_keyframe_vals()
|
|
return cloned
|
|
|
|
def create_list_of_dicts(self):
|
|
# for each keyframe, create a dict representing values relevant to TimestepKeyframe creation in ACN
|
|
c = []
|
|
for kf in self.keyframes:
|
|
d = {}
|
|
# scheduling
|
|
d["start_percent"] = kf.start_percent
|
|
d["guarantee_steps"] = kf.guarantee_steps
|
|
d["inherit_missing"] = kf.inherit_missing
|
|
# values
|
|
if type(kf.mult_multival) == Tensor:
|
|
d["strength"] = kf.mult
|
|
d["mask"] = kf.mult_multival
|
|
else:
|
|
if kf.mult_multival is None:
|
|
d["strength"] = kf.mult
|
|
else:
|
|
d["strength"] = kf.mult * kf.mult_multival
|
|
d["mask"] = None
|
|
d["tune"] = kf.tune_replace
|
|
d["mode"] = kf.mode_replace
|
|
# add to list
|
|
c.append(d)
|
|
return c
|
|
|
|
|
|
class ContextRef(ContextExtra):
|
|
def __init__(self, start_percent: float, end_percent: float,
|
|
strength_multival: Union[float, Tensor], tune: ContextRefTune, mode: ContextRefMode,
|
|
keyframe: ContextRefKeyframeGroup=None):
|
|
super().__init__(start_percent=start_percent, end_percent=end_percent)
|
|
self.tune = tune
|
|
self.mode = mode
|
|
self.keyframe = keyframe if keyframe else ContextRefKeyframeGroup()
|
|
self.version = CONTEXTREF_VERSION
|
|
# stuff for ACN usage
|
|
self.strength = 1.0
|
|
self.mask = None
|
|
self._strength_multival = strength_multival
|
|
self.strength_multival = strength_multival
|
|
|
|
@property
|
|
def strength_multival(self):
|
|
return self.strength_multival
|
|
@strength_multival.setter
|
|
def strength_multival(self, value):
|
|
if value is None:
|
|
value = 1.0
|
|
if type(value) == Tensor:
|
|
self.strength = 1.0
|
|
self.mask = value
|
|
else:
|
|
self.strength = value
|
|
self.mask = None
|
|
self._strength_multival = value
|
|
|
|
def should_run(self):
|
|
return super().should_run()
|
|
|
|
class ContextRefHandler:
|
|
CONTEXTREF_CONTROL_LIST_ALL = "contextref_control_list_all"
|
|
CONTEXTREF_MACHINE_STATE = "contextref_machine_state"
|
|
CONTEXTREF_CLEAN_FUNC = "contextref_clean_func"
|
|
|
|
def __init__(self):
|
|
self.contextref_active = False
|
|
self.contextref_mode: ContextRefMode = None
|
|
self.contextref_idxs_set: set[int] = None
|
|
self.first_context = True
|
|
|
|
def initialize_step(self, timestep, model_options, ADGS: AnimateDiffGlobalState):
|
|
# check that ACN provided ContextRef as requested
|
|
temp_refcn_list = model_options["transformer_options"].get(self.CONTEXTREF_CONTROL_LIST_ALL, None)
|
|
if temp_refcn_list is None:
|
|
raise Exception("Advanced-ControlNet nodes are either missing or too outdated to support ContextRef. Update/install ComfyUI-Advanced-ControlNet to use ContextRef.")
|
|
if len(temp_refcn_list) == 0:
|
|
raise Exception("Unexpected ContextRef issue; Advanced-ControlNet did not provide any ContextRef objs for AnimateDiff-Evolved.")
|
|
del temp_refcn_list
|
|
# check if ContextRef ReferenceAdvanced ACN objs should_run
|
|
actually_should_run = True
|
|
for refcn in model_options["transformer_options"][self.CONTEXTREF_CONTROL_LIST_ALL]:
|
|
acn_dl_version = get_acn_dinklink_version()
|
|
if acn_dl_version > 10000:
|
|
refcn.prepare_current_timestep(timestep, model_options["transformer_options"])
|
|
else:
|
|
refcn.prepare_current_timestep(timestep)
|
|
if not refcn.should_run():
|
|
actually_should_run = False
|
|
if actually_should_run:
|
|
self.contextref_active = True
|
|
for refcn in model_options["transformer_options"][self.CONTEXTREF_CONTROL_LIST_ALL]:
|
|
# get mode_override if present, mode otherwise
|
|
self.contextref_mode = refcn.get_contextref_mode_replace() or ADGS.params.context_options.extras.context_ref.mode
|
|
self.contextref_idxs_set = self.contextref_mode.indexes.copy()
|
|
|
|
def prepare_referencecn(self, ctx_idxs: list[int], window_idx: int, model_options):
|
|
if self.contextref_active:
|
|
# set cond counter to 0 (each cond encountered will increment it by 1)
|
|
for refcn in model_options["transformer_options"][self.CONTEXTREF_CONTROL_LIST_ALL]:
|
|
refcn.contextref_cond_idx = 0
|
|
if self.first_context:
|
|
model_options["transformer_options"][self.CONTEXTREF_MACHINE_STATE] = MachineState.WRITE
|
|
else:
|
|
model_options["transformer_options"][self.CONTEXTREF_MACHINE_STATE] = MachineState.READ
|
|
if self.contextref_mode.mode == ContextRefMode.SLIDING: # if sliding, check if time to READ and WRITE
|
|
if window_idx % (self.contextref_mode.sliding_width-1) == 0:
|
|
model_options["transformer_options"][self.CONTEXTREF_MACHINE_STATE] = MachineState.READ_WRITE
|
|
# override with indexes mode, if set
|
|
if self.contextref_mode.mode == ContextRefMode.INDEXES:
|
|
contains_idx = False
|
|
for i in ctx_idxs:
|
|
if i in self.contextref_idxs_set:
|
|
contains_idx = True
|
|
# single trigger decides if each index should only trigger READ_WRITE once per step
|
|
if not self.contextref_mode.single_trigger:
|
|
break
|
|
self.contextref_idxs_set.remove(i)
|
|
if contains_idx:
|
|
model_options["transformer_options"][self.CONTEXTREF_MACHINE_STATE] = MachineState.READ_WRITE
|
|
if self.first_context:
|
|
model_options["transformer_options"][self.CONTEXTREF_MACHINE_STATE] = MachineState.WRITE
|
|
else:
|
|
model_options["transformer_options"][self.CONTEXTREF_MACHINE_STATE] = MachineState.READ
|
|
else:
|
|
model_options["transformer_options"][self.CONTEXTREF_MACHINE_STATE] = MachineState.OFF
|
|
#logger.info(f"window: {window_idx} - {model_options['transformer_options'][CONTEXTREF_MACHINE_STATE]}")
|
|
|
|
|
|
def finalize_step(self):
|
|
'Toggle first_context off, if needed.'
|
|
if self.first_context:
|
|
self.first_context = False
|
|
|
|
def cleanup(self, model_options):
|
|
'Clean contextref stuff with provided ACN function, if applicable.'
|
|
if self.contextref_active:
|
|
model_options["transformer_options"][self.CONTEXTREF_CLEAN_FUNC]()
|
|
#--------------------------------
|
|
|
|
|
|
################################
|
|
# NaiveReuse
|
|
class NaiveReuseKeyframe:
|
|
def __init__(self, mult=1.0, mult_multival: Union[float, Tensor]=None, start_percent=0.0, guarantee_steps=1, inherit_missing=True):
|
|
self.mult = mult
|
|
self.orig_mult_multival = mult_multival
|
|
self.mult_multival = mult_multival
|
|
# scheduling
|
|
self.start_percent = float(start_percent)
|
|
self.start_t = 999999999.9
|
|
self.guarantee_steps = guarantee_steps
|
|
self.inherit_missing = inherit_missing
|
|
|
|
def get_effective_guarantee_steps(self, max_sigma: torch.Tensor):
|
|
'''If keyframe starts before current sampling range (max_sigma), treat as 0.'''
|
|
if self.start_t > max_sigma:
|
|
return 0
|
|
return self.guarantee_steps
|
|
|
|
def clone(self):
|
|
c = NaiveReuseKeyframe(mult=self.mult, mult_multival=self.mult_multival,
|
|
start_percent=self.start_percent, guarantee_steps=self.guarantee_steps)
|
|
c.start_t = self.start_t
|
|
return c
|
|
|
|
|
|
class NaiveReuseKeyframeGroup:
|
|
def __init__(self):
|
|
self.keyframes: list[NaiveReuseKeyframe] = []
|
|
self._current_keyframe: NaiveReuseKeyframe = None
|
|
self._current_used_steps: int = 0
|
|
self._current_index: int = 0
|
|
self._previous_t = -1
|
|
|
|
def reset(self):
|
|
self._current_keyframe = None
|
|
self._current_used_steps = 0
|
|
self._current_index = 0
|
|
self._set_first_as_current()
|
|
|
|
def add(self, keyframe: NaiveReuseKeyframe):
|
|
# add to end of list, then sort
|
|
self.keyframes.append(keyframe)
|
|
self.keyframes = get_sorted_list_via_attr(self.keyframes, "start_percent")
|
|
self._set_first_as_current()
|
|
self._prepare_all_keyframe_vals()
|
|
|
|
def _set_first_as_current(self):
|
|
if len(self.keyframes) > 0:
|
|
self._current_keyframe = self.keyframes[0]
|
|
else:
|
|
self._current_keyframe = None
|
|
|
|
def _prepare_all_keyframe_vals(self):
|
|
if self.is_empty():
|
|
return
|
|
multival = None
|
|
for kf in self.keyframes:
|
|
# if shouldn't inherit, clear cache
|
|
if not kf.inherit_missing:
|
|
multival = None
|
|
# assign cached values, if origs were None
|
|
# Mult #################
|
|
if kf.orig_mult_multival is None:
|
|
kf.mult_multival = multival
|
|
else:
|
|
kf.mult_multival = kf.orig_mult_multival
|
|
# save new caches, in case next keyframe inherits missing
|
|
if kf.mult_multival is not None:
|
|
multival = kf.mult_multival
|
|
|
|
def has_index(self, index: int) -> int:
|
|
return index >=0 and index < len(self.keyframes)
|
|
|
|
def is_empty(self) -> bool:
|
|
return len(self.keyframes) == 0
|
|
|
|
def clone(self):
|
|
cloned = NaiveReuseKeyframeGroup()
|
|
for keyframe in self.keyframes:
|
|
cloned.keyframes.append(keyframe)
|
|
cloned._set_first_as_current()
|
|
cloned._prepare_all_keyframe_vals()
|
|
return cloned
|
|
|
|
def initialize_timesteps(self, model: BaseModel):
|
|
for keyframe in self.keyframes:
|
|
keyframe.start_t = model.model_sampling.percent_to_sigma(keyframe.start_percent)
|
|
|
|
def prepare_current_keyframe(self, t: Tensor, transformer_options: dict[str, Tensor]):
|
|
if self.is_empty():
|
|
return
|
|
curr_t: float = t[0]
|
|
# if curr_t same as before, do nothing as step already accounted for
|
|
if curr_t == self._previous_t:
|
|
return
|
|
prev_index = self._current_index
|
|
max_sigma = torch.max(transformer_options.get("sample_sigmas", BIGMAX_TENSOR))
|
|
# if met guaranteed steps, look for next keyframe in case need to switch
|
|
if self._current_used_steps >= self._current_keyframe.get_effective_guarantee_steps(max_sigma):
|
|
# if has next index, loop through and see if need t oswitch
|
|
if self.has_index(self._current_index+1):
|
|
for i in range(self._current_index+1, len(self.keyframes)):
|
|
eval_c = self.keyframes[i]
|
|
# check if start_t is greater or equal to curr_t
|
|
# NOTE: t is in terms of sigmas, not percent, so bigger number = earlier step in sampling
|
|
if eval_c.start_t >= curr_t:
|
|
self._current_index = i
|
|
self._current_keyframe = eval_c
|
|
self._current_used_steps = 0
|
|
# if guarantee_steps greater than zero, stop searching for other keyframes
|
|
if self._current_keyframe.get_effective_guarantee_steps(max_sigma) > 0:
|
|
break
|
|
# if eval_c is outside the percent range, stop looking further
|
|
else: break
|
|
# update steps current context is used
|
|
self._current_used_steps += 1
|
|
# update previous_t
|
|
self._previous_t = curr_t
|
|
|
|
# properties shadow those of NaiveReuseKeyframe
|
|
@property
|
|
def mult(self):
|
|
if self._current_keyframe != None:
|
|
return self._current_keyframe.mult
|
|
return 1.0
|
|
|
|
@property
|
|
def mult_multival(self):
|
|
if self._current_keyframe != None:
|
|
return self._current_keyframe.mult_multival
|
|
return None
|
|
|
|
|
|
class NaiveReuse(ContextExtra):
|
|
def __init__(self, start_percent: float, end_percent: float, weighted_mean: float, multival_opt: Union[float, Tensor]=None, naivereuse_kf: NaiveReuseKeyframeGroup=None):
|
|
super().__init__(start_percent=start_percent, end_percent=end_percent)
|
|
self.weighted_mean = weighted_mean
|
|
self.orig_multival = multival_opt
|
|
self.mask: Tensor = None
|
|
self.keyframe = naivereuse_kf if naivereuse_kf else NaiveReuseKeyframeGroup()
|
|
self._prev_keyframe = None
|
|
|
|
def cleanup(self):
|
|
super().cleanup()
|
|
del self.mask
|
|
self.mask = None
|
|
self._prev_keyframe = None
|
|
self.keyframe.reset()
|
|
|
|
def initialize_timesteps(self, model: BaseModel):
|
|
super().initialize_timesteps(model)
|
|
self.keyframe.initialize_timesteps(model)
|
|
|
|
def prepare_current(self, t: Tensor, transformer_options: dict[str, Tensor]):
|
|
super().prepare_current(t, transformer_options)
|
|
self.keyframe.prepare_current_keyframe(t, transformer_options)
|
|
|
|
def get_effective_weighted_mean(self, x: Tensor, idxs: list[int]):
|
|
if self.orig_multival is None and self.keyframe.mult_multival is None:
|
|
return self.weighted_mean * self.keyframe.mult
|
|
# check if keyframe changed
|
|
keyframe_changed = False
|
|
if self.keyframe._current_keyframe != self._prev_keyframe:
|
|
keyframe_changed = True
|
|
self._prev_keyframe = self.keyframe._current_keyframe
|
|
|
|
if type(self.orig_multival) != Tensor and type(self.keyframe.mult_multival) != Tensor:
|
|
return self.weighted_mean * self.keyframe.mult * get_combined_multival(self.orig_multival, self.keyframe.mult_multival)
|
|
|
|
if self.mask is None or keyframe_changed or self.mask.shape[0] != x.shape[0] or self.mask.shape[-1] != x.shape[-1] or self.mask.shape[-2] != x.shape[-2]:
|
|
del self.mask
|
|
real_mult_multival = resize_multival(self.keyframe.mult_multival, batch_size=x.shape[0], height=x.shape[-1], width=x.shape[-2])
|
|
self.mask = resize_multival(self.orig_multival, batch_size=x.shape[0], height=x.shape[-1], width=x.shape[-2])
|
|
self.mask = get_combined_multival(self.mask, real_mult_multival)
|
|
return self.weighted_mean * self.keyframe.mult * self.mask[idxs].to(dtype=x.dtype, device=x.device)
|
|
|
|
def should_run(self):
|
|
to_return = super().should_run()
|
|
# if keyframe has 0.0 val, should not run
|
|
if self.keyframe.mult_multival is not None and type(self.keyframe.mult_multival) != Tensor and math.isclose(self.keyframe.mult_multival, 0.0):
|
|
return False
|
|
# if weighted_mean is 0.0, then reuse will take no effect anyway
|
|
return to_return and self.weighted_mean > 0.0 and self.keyframe.mult > 0.0
|
|
|
|
class NaiveReuseHandler:
|
|
def __init__(self):
|
|
self.naivereuse_active = False
|
|
self.cached_naive_conds = None
|
|
self.cached_naive_ctx_idxs = None
|
|
|
|
def initialize_step(self, x_in: Tensor, conds):
|
|
self.cached_naive_conds = [torch.zeros_like(x_in) for _ in conds]
|
|
#cached_naive_counts = [torch.zeros((x_in.shape[0], 1, 1, 1), device=x_in.device) for _ in conds]
|
|
self.naivereuse_active = True
|
|
|
|
def cache_first_context_results(self, window_idx, ctx_idxs: list[int], sub_conds: list, conds_final: list[Tensor], counts_final: list[Tensor]):
|
|
if self.naivereuse_active and window_idx == 0:
|
|
self.cached_naive_ctx_idxs = ctx_idxs
|
|
for i in range(len(sub_conds)):
|
|
self.cached_naive_conds[i][ctx_idxs] = conds_final[i][ctx_idxs] / counts_final[i][ctx_idxs]
|
|
self.naivereuse_active = False
|
|
|
|
def apply_cached(self, x_in: Tensor, conds_final: list[Tensor], counts_final: list[Tensor], ADGS: AnimateDiffGlobalState):
|
|
if self.cached_naive_conds is not None:
|
|
start_idx = self.cached_naive_ctx_idxs[0]
|
|
for z in range(0, ADGS.params.full_length, len(self.cached_naive_ctx_idxs)):
|
|
for i in range(len(self.cached_naive_conds)):
|
|
# get the 'true' idxs of this window
|
|
new_ctx_idxs = [(zz+start_idx) % ADGS.params.full_length for zz in list(range(z, z+len(self.cached_naive_ctx_idxs))) if zz < ADGS.params.full_length]
|
|
# make sure when getting cached_naive idxs, they are adjusted for actual length leftover length
|
|
adjusted_naive_ctx_idxs = self.cached_naive_ctx_idxs[:len(new_ctx_idxs)]
|
|
weighted_mean = ADGS.params.context_options.extras.naive_reuse.get_effective_weighted_mean(x_in, new_ctx_idxs)
|
|
conds_final[i][new_ctx_idxs] = (weighted_mean * (self.cached_naive_conds[i][adjusted_naive_ctx_idxs]*counts_final[i][new_ctx_idxs])) + ((1.-weighted_mean) * conds_final[i][new_ctx_idxs])
|
|
self.cleanup()
|
|
|
|
def cleanup(self):
|
|
del self.cached_naive_conds
|
|
self.cached_naive_conds = None
|
|
#--------------------------------
|
|
|
|
|
|
################################
|
|
# DenoiseReuse
|
|
|
|
|
|
class ContextExtrasGroup:
|
|
def __init__(self):
|
|
self.context_ref: ContextRef = None
|
|
self.naive_reuse: NaiveReuse = None
|
|
|
|
def get_extras_list(self) -> list[ContextExtra]:
|
|
extras_list = []
|
|
if self.context_ref is not None:
|
|
extras_list.append(self.context_ref)
|
|
if self.naive_reuse is not None:
|
|
extras_list.append(self.naive_reuse)
|
|
return extras_list
|
|
|
|
def initialize_timesteps(self, model: BaseModel):
|
|
for extra in self.get_extras_list():
|
|
extra.initialize_timesteps(model)
|
|
|
|
def prepare_current(self, t: Tensor, transformer_options):
|
|
for extra in self.get_extras_list():
|
|
extra.prepare_current(t, transformer_options)
|
|
|
|
def should_run_context_ref(self):
|
|
if not self.context_ref:
|
|
return False
|
|
return self.context_ref.should_run()
|
|
|
|
def should_run_naive_reuse(self):
|
|
if not self.naive_reuse:
|
|
return False
|
|
return self.naive_reuse.should_run()
|
|
|
|
def add(self, extra: ContextExtra):
|
|
if type(extra) == ContextRef:
|
|
self.context_ref = extra
|
|
elif type(extra) == NaiveReuse:
|
|
self.naive_reuse = extra
|
|
else:
|
|
raise Exception(f"Unrecognized ContextExtras type: {type(extra)}")
|
|
|
|
def cleanup(self):
|
|
for extra in self.get_extras_list():
|
|
extra.cleanup()
|
|
|
|
def clone(self):
|
|
cloned = ContextExtrasGroup()
|
|
cloned.context_ref = self.context_ref
|
|
cloned.naive_reuse = self.naive_reuse
|
|
return cloned
|