fred-dev/comfy_ui_ali
0
1from __future__ import annotations2from typing import TYPE_CHECKING, Callable3import enum4import math5import torch6import numpy as np7import itertools8import logging9 10if TYPE_CHECKING:11 from comfy.model_patcher import ModelPatcher, PatcherInjection12 from comfy.model_base import BaseModel13 from comfy.sd import CLIP14import comfy.lora15import comfy.model_management16import comfy.patcher_extension17from node_helpers import conditioning_set_values18 19# #######################################################################################################20# Hooks explanation21# -------------------22# The purpose of hooks is to allow conds to influence sampling without the need for ComfyUI core code to23# make explicit special cases like it does for ControlNet and GLIGEN.24#25# This is necessary for nodes/features that are intended for use with masked or scheduled conds, or those26# that should run special code when a 'marked' cond is used in sampling.27# #######################################################################################################28 29class EnumHookMode(enum.Enum):30 '''31 Priority of hook memory optimization vs. speed, mostly related to WeightHooks.32 33 MinVram: No caching will occur for any operations related to hooks.34 MaxSpeed: Excess VRAM (and RAM, once VRAM is sufficiently depleted) will be used to cache hook weights when switching hook groups.35 '''36 MinVram = "minvram"37 MaxSpeed = "maxspeed"38 39class EnumHookType(enum.Enum):40 '''41 Hook types, each of which has different expected behavior.42 '''43 Weight = "weight"44 ObjectPatch = "object_patch"45 AdditionalModels = "add_models"46 TransformerOptions = "transformer_options"47 Injections = "add_injections"48 49class EnumWeightTarget(enum.Enum):50 Model = "model"51 Clip = "clip"52 53class EnumHookScope(enum.Enum):54 '''55 Determines if hook should be limited in its influence over sampling.56 57 AllConditioning: hook will affect all conds used in sampling.58 HookedOnly: hook will only affect the conds it was attached to.59 '''60 AllConditioning = "all_conditioning"61 HookedOnly = "hooked_only"62 63 64class _HookRef:65 pass66 67 68def default_should_register(hook: Hook, model: ModelPatcher, model_options: dict, target_dict: dict[str], registered: HookGroup):69 '''Example for how custom_should_register function can look like.'''70 return True71 72 73def create_target_dict(target: EnumWeightTarget=None, **kwargs) -> dict[str]:74 '''Creates base dictionary for use with Hooks' target param.'''75 d = {}76 if target is not None:77 d['target'] = target78 d.update(kwargs)79 return d80 81 82class Hook:83 def __init__(self, hook_type: EnumHookType=None, hook_ref: _HookRef=None, hook_id: str=None,84 hook_keyframe: HookKeyframeGroup=None, hook_scope=EnumHookScope.AllConditioning):85 self.hook_type = hook_type86 '''Enum identifying the general class of this hook.'''87 self.hook_ref = hook_ref if hook_ref else _HookRef()88 '''Reference shared between hook clones that have the same value. Should NOT be modified.'''89 self.hook_id = hook_id90 '''Optional string ID to identify hook; useful if need to consolidate duplicates at registration time.'''91 self.hook_keyframe = hook_keyframe if hook_keyframe else HookKeyframeGroup()92 '''Keyframe storage that can be referenced to get strength for current sampling step.'''93 self.hook_scope = hook_scope94 '''Scope of where this hook should apply in terms of the conds used in sampling run.'''95 self.custom_should_register = default_should_register96 '''Can be overriden with a compatible function to decide if this hook should be registered without the need to override .should_register'''97 98 @property99 def strength(self):100 return self.hook_keyframe.strength101 102 def initialize_timesteps(self, model: BaseModel):103 self.reset()104 self.hook_keyframe.initialize_timesteps(model)105 106 def reset(self):107 self.hook_keyframe.reset()108 109 def clone(self):110 c: Hook = self.__class__()111 c.hook_type = self.hook_type112 c.hook_ref = self.hook_ref113 c.hook_id = self.hook_id114 c.hook_keyframe = self.hook_keyframe115 c.hook_scope = self.hook_scope116 c.custom_should_register = self.custom_should_register117 return c118 119 def should_register(self, model: ModelPatcher, model_options: dict, target_dict: dict[str], registered: HookGroup):120 return self.custom_should_register(self, model, model_options, target_dict, registered)121 122 def add_hook_patches(self, model: ModelPatcher, model_options: dict, target_dict: dict[str], registered: HookGroup):123 raise NotImplementedError("add_hook_patches should be defined for Hook subclasses")124 125 def __eq__(self, other: Hook):126 return self.__class__ == other.__class__ and self.hook_ref == other.hook_ref127 128 def __hash__(self):129 return hash(self.hook_ref)130 131class WeightHook(Hook):132 '''133 Hook responsible for tracking weights to be applied to some model/clip.134 135 Note, value of hook_scope is ignored and is treated as HookedOnly.136 '''137 def __init__(self, strength_model=1.0, strength_clip=1.0):138 super().__init__(hook_type=EnumHookType.Weight, hook_scope=EnumHookScope.HookedOnly)139 self.weights: dict = None140 self.weights_clip: dict = None141 self.need_weight_init = True142 self._strength_model = strength_model143 self._strength_clip = strength_clip144 self.hook_scope = EnumHookScope.HookedOnly # this value does not matter for WeightHooks, just for docs145 146 @property147 def strength_model(self):148 return self._strength_model * self.strength149 150 @property151 def strength_clip(self):152 return self._strength_clip * self.strength153 154 def add_hook_patches(self, model: ModelPatcher, model_options: dict, target_dict: dict[str], registered: HookGroup):155 if not self.should_register(model, model_options, target_dict, registered):156 return False157 weights = None158 159 target = target_dict.get('target', None)160 if target == EnumWeightTarget.Clip:161 strength = self._strength_clip162 else:163 strength = self._strength_model164 165 if self.need_weight_init:166 key_map = {}167 if target == EnumWeightTarget.Clip:168 key_map = comfy.lora.model_lora_keys_clip(model.model, key_map)169 else:170 key_map = comfy.lora.model_lora_keys_unet(model.model, key_map)171 weights = comfy.lora.load_lora(self.weights, key_map, log_missing=False)172 else:173 if target == EnumWeightTarget.Clip:174 weights = self.weights_clip175 else:176 weights = self.weights177 model.add_hook_patches(hook=self, patches=weights, strength_patch=strength)178 registered.add(self)179 return True180 # TODO: add logs about any keys that were not applied181 182 def clone(self):183 c: WeightHook = super().clone()184 c.weights = self.weights185 c.weights_clip = self.weights_clip186 c.need_weight_init = self.need_weight_init187 c._strength_model = self._strength_model188 c._strength_clip = self._strength_clip189 return c190 191class ObjectPatchHook(Hook):192 def __init__(self, object_patches: dict[str]=None,193 hook_scope=EnumHookScope.AllConditioning):194 super().__init__(hook_type=EnumHookType.ObjectPatch)195 self.object_patches = object_patches196 self.hook_scope = hook_scope197 198 def clone(self):199 c: ObjectPatchHook = super().clone()200 c.object_patches = self.object_patches201 return c202 203 def add_hook_patches(self, model: ModelPatcher, model_options: dict, target_dict: dict[str], registered: HookGroup):204 raise NotImplementedError("ObjectPatchHook is not supported yet in ComfyUI.")205 206class AdditionalModelsHook(Hook):207 '''208 Hook responsible for telling model management any additional models that should be loaded.209 210 Note, value of hook_scope is ignored and is treated as AllConditioning.211 '''212 def __init__(self, models: list[ModelPatcher]=None, key: str=None):213 super().__init__(hook_type=EnumHookType.AdditionalModels)214 self.models = models215 self.key = key216 217 def clone(self):218 c: AdditionalModelsHook = super().clone()219 c.models = self.models.copy() if self.models else self.models220 c.key = self.key221 return c222 223 def add_hook_patches(self, model: ModelPatcher, model_options: dict, target_dict: dict[str], registered: HookGroup):224 if not self.should_register(model, model_options, target_dict, registered):225 return False226 registered.add(self)227 return True228 229class TransformerOptionsHook(Hook):230 '''231 Hook responsible for adding wrappers, callbacks, patches, or anything else related to transformer_options.232 '''233 def __init__(self, transformers_dict: dict[str, dict[str, dict[str, list[Callable]]]]=None,234 hook_scope=EnumHookScope.AllConditioning):235 super().__init__(hook_type=EnumHookType.TransformerOptions)236 self.transformers_dict = transformers_dict237 self.hook_scope = hook_scope238 self._skip_adding = False239 '''Internal value used to avoid double load of transformer_options when hook_scope is AllConditioning.'''240 241 def clone(self):242 c: TransformerOptionsHook = super().clone()243 c.transformers_dict = self.transformers_dict244 c._skip_adding = self._skip_adding245 return c246 247 def add_hook_patches(self, model: ModelPatcher, model_options: dict, target_dict: dict[str], registered: HookGroup):248 if not self.should_register(model, model_options, target_dict, registered):249 return False250 # NOTE: to_load_options will be used to manually load patches/wrappers/callbacks from hooks251 self._skip_adding = False252 if self.hook_scope == EnumHookScope.AllConditioning:253 add_model_options = {"transformer_options": self.transformers_dict,254 "to_load_options": self.transformers_dict}255 # skip_adding if included in AllConditioning to avoid double loading256 self._skip_adding = True257 else:258 add_model_options = {"to_load_options": self.transformers_dict}259 registered.add(self)260 comfy.patcher_extension.merge_nested_dicts(model_options, add_model_options, copy_dict1=False)261 return True262 263 def on_apply_hooks(self, model: ModelPatcher, transformer_options: dict[str]):264 if not self._skip_adding:265 comfy.patcher_extension.merge_nested_dicts(transformer_options, self.transformers_dict, copy_dict1=False)266 267WrapperHook = TransformerOptionsHook268'''Only here for backwards compatibility, WrapperHook is identical to TransformerOptionsHook.'''269 270class InjectionsHook(Hook):271 def __init__(self, key: str=None, injections: list[PatcherInjection]=None,272 hook_scope=EnumHookScope.AllConditioning):273 super().__init__(hook_type=EnumHookType.Injections)274 self.key = key275 self.injections = injections276 self.hook_scope = hook_scope277 278 def clone(self):279 c: InjectionsHook = super().clone()280 c.key = self.key281 c.injections = self.injections.copy() if self.injections else self.injections282 return c283 284 def add_hook_patches(self, model: ModelPatcher, model_options: dict, target_dict: dict[str], registered: HookGroup):285 raise NotImplementedError("InjectionsHook is not supported yet in ComfyUI.")286 287class HookGroup:288 '''289 Stores groups of hooks, and allows them to be queried by type.290 291 To prevent breaking their functionality, never modify the underlying self.hooks or self._hook_dict vars directly;292 always use the provided functions on HookGroup.293 '''294 def __init__(self):295 self.hooks: list[Hook] = []296 self._hook_dict: dict[EnumHookType, list[Hook]] = {}297 298 def __len__(self):299 return len(self.hooks)300 301 def add(self, hook: Hook):302 if hook not in self.hooks:303 self.hooks.append(hook)304 self._hook_dict.setdefault(hook.hook_type, []).append(hook)305 306 def remove(self, hook: Hook):307 if hook in self.hooks:308 self.hooks.remove(hook)309 self._hook_dict[hook.hook_type].remove(hook)310 311 def get_type(self, hook_type: EnumHookType):312 return self._hook_dict.get(hook_type, [])313 314 def contains(self, hook: Hook):315 return hook in self.hooks316 317 def is_subset_of(self, other: HookGroup):318 self_hooks = set(self.hooks)319 other_hooks = set(other.hooks)320 return self_hooks.issubset(other_hooks)321 322 def new_with_common_hooks(self, other: HookGroup):323 c = HookGroup()324 for hook in self.hooks:325 if other.contains(hook):326 c.add(hook.clone())327 return c328 329 def clone(self):330 c = HookGroup()331 for hook in self.hooks:332 c.add(hook.clone())333 return c334 335 def clone_and_combine(self, other: HookGroup):336 c = self.clone()337 if other is not None:338 for hook in other.hooks:339 c.add(hook.clone())340 return c341 342 def set_keyframes_on_hooks(self, hook_kf: HookKeyframeGroup):343 if hook_kf is None:344 hook_kf = HookKeyframeGroup()345 else:346 hook_kf = hook_kf.clone()347 for hook in self.hooks:348 hook.hook_keyframe = hook_kf349 350 def get_hooks_for_clip_schedule(self):351 scheduled_hooks: dict[WeightHook, list[tuple[tuple[float,float], HookKeyframe]]] = {}352 # only care about WeightHooks, for now353 for hook in self.get_type(EnumHookType.Weight):354 hook: WeightHook355 hook_schedule = []356 # if no hook keyframes, assign default value357 if len(hook.hook_keyframe.keyframes) == 0:358 hook_schedule.append(((0.0, 1.0), None))359 scheduled_hooks[hook] = hook_schedule360 continue361 # find ranges of values362 prev_keyframe = hook.hook_keyframe.keyframes[0]363 for keyframe in hook.hook_keyframe.keyframes:364 if keyframe.start_percent > prev_keyframe.start_percent and not math.isclose(keyframe.strength, prev_keyframe.strength):365 hook_schedule.append(((prev_keyframe.start_percent, keyframe.start_percent), prev_keyframe))366 prev_keyframe = keyframe367 elif keyframe.start_percent == prev_keyframe.start_percent:368 prev_keyframe = keyframe369 # create final range, assuming last start_percent was not 1.0370 if not math.isclose(prev_keyframe.start_percent, 1.0):371 hook_schedule.append(((prev_keyframe.start_percent, 1.0), prev_keyframe))372 scheduled_hooks[hook] = hook_schedule373 # hooks should not have their schedules in a list of tuples374 all_ranges: list[tuple[float, float]] = []375 for range_kfs in scheduled_hooks.values():376 for t_range, keyframe in range_kfs:377 all_ranges.append(t_range)378 # turn list of ranges into boundaries379 boundaries_set = set(itertools.chain.from_iterable(all_ranges))380 boundaries_set.add(0.0)381 boundaries = sorted(boundaries_set)382 real_ranges = [(boundaries[i], boundaries[i + 1]) for i in range(len(boundaries) - 1)]383 # with real ranges defined, give appropriate hooks w/ keyframes for each range384 scheduled_keyframes: list[tuple[tuple[float,float], list[tuple[WeightHook, HookKeyframe]]]] = []385 for t_range in real_ranges:386 hooks_schedule = []387 for hook, val in scheduled_hooks.items():388 keyframe = None389 # check if is a keyframe that works for the current t_range390 for stored_range, stored_kf in val:391 # if stored start is less than current end, then fits - give it assigned keyframe392 if stored_range[0] < t_range[1] and stored_range[1] > t_range[0]:393 keyframe = stored_kf394 break395 hooks_schedule.append((hook, keyframe))396 scheduled_keyframes.append((t_range, hooks_schedule))397 return scheduled_keyframes398 399 def reset(self):400 for hook in self.hooks:401 hook.reset()402 403 @staticmethod404 def combine_all_hooks(hooks_list: list[HookGroup], require_count=0) -> HookGroup:405 actual: list[HookGroup] = []406 for group in hooks_list:407 if group is not None:408 actual.append(group)409 if len(actual) < require_count:410 raise Exception(f"Need at least {require_count} hooks to combine, but only had {len(actual)}.")411 # if no hooks, then return None412 if len(actual) == 0:413 return None414 # if only 1 hook, just return itself without cloning415 elif len(actual) == 1:416 return actual[0]417 final_hook: HookGroup = None418 for hook in actual:419 if final_hook is None:420 final_hook = hook.clone()421 else:422 final_hook = final_hook.clone_and_combine(hook)423 return final_hook424 425 426class HookKeyframe:427 def __init__(self, strength: float, start_percent=0.0, guarantee_steps=1):428 self.strength = strength429 # scheduling430 self.start_percent = float(start_percent)431 self.start_t = 999999999.9432 self.guarantee_steps = guarantee_steps433 434 def get_effective_guarantee_steps(self, max_sigma: torch.Tensor):435 '''If keyframe starts before current sampling range (max_sigma), treat as 0.'''436 if self.start_t > max_sigma:437 return 0438 return self.guarantee_steps439 440 def clone(self):441 c = HookKeyframe(strength=self.strength,442 start_percent=self.start_percent, guarantee_steps=self.guarantee_steps)443 c.start_t = self.start_t444 return c445 446class HookKeyframeGroup:447 def __init__(self):448 self.keyframes: list[HookKeyframe] = []449 self._current_keyframe: HookKeyframe = None450 self._current_used_steps = 0451 self._current_index = 0452 self._current_strength = None453 self._curr_t = -1.454 455 # properties shadow those of HookWeightsKeyframe456 @property457 def strength(self):458 if self._current_keyframe is not None:459 return self._current_keyframe.strength460 return 1.0461 462 def reset(self):463 self._current_keyframe = None464 self._current_used_steps = 0465 self._current_index = 0466 self._current_strength = None467 self.curr_t = -1.468 self._set_first_as_current()469 470 def add(self, keyframe: HookKeyframe):471 # add to end of list, then sort472 self.keyframes.append(keyframe)473 self.keyframes = get_sorted_list_via_attr(self.keyframes, "start_percent")474 self._set_first_as_current()475 476 def _set_first_as_current(self):477 if len(self.keyframes) > 0:478 self._current_keyframe = self.keyframes[0]479 else:480 self._current_keyframe = None481 482 def has_guarantee_steps(self):483 for kf in self.keyframes:484 if kf.guarantee_steps > 0:485 return True486 return False487 488 def has_index(self, index: int):489 return index >= 0 and index < len(self.keyframes)490 491 def is_empty(self):492 return len(self.keyframes) == 0493 494 def clone(self):495 c = HookKeyframeGroup()496 for keyframe in self.keyframes:497 c.keyframes.append(keyframe.clone())498 c._set_first_as_current()499 return c500 501 def initialize_timesteps(self, model: BaseModel):502 for keyframe in self.keyframes:503 keyframe.start_t = model.model_sampling.percent_to_sigma(keyframe.start_percent)504 505 def prepare_current_keyframe(self, curr_t: float, transformer_options: dict[str, torch.Tensor]) -> bool:506 if self.is_empty():507 return False508 if curr_t == self._curr_t:509 return False510 max_sigma = torch.max(transformer_options["sample_sigmas"])511 prev_index = self._current_index512 prev_strength = self._current_strength513 # if met guaranteed steps, look for next keyframe in case need to switch514 if self._current_used_steps >= self._current_keyframe.get_effective_guarantee_steps(max_sigma):515 # if has next index, loop through and see if need to switch516 if self.has_index(self._current_index+1):517 for i in range(self._current_index+1, len(self.keyframes)):518 eval_c = self.keyframes[i]519 # check if start_t is greater or equal to curr_t520 # NOTE: t is in terms of sigmas, not percent, so bigger number = earlier step in sampling521 if eval_c.start_t >= curr_t:522 self._current_index = i523 self._current_strength = eval_c.strength524 self._current_keyframe = eval_c525 self._current_used_steps = 0526 # if guarantee_steps greater than zero, stop searching for other keyframes527 if self._current_keyframe.get_effective_guarantee_steps(max_sigma) > 0:528 break529 # if eval_c is outside the percent range, stop looking further530 else: break531 # update steps current context is used532 self._current_used_steps += 1533 # update current timestep this was performed on534 self._curr_t = curr_t535 # return True if keyframe changed, False if no change536 return prev_index != self._current_index and prev_strength != self._current_strength537 538 539class InterpolationMethod:540 LINEAR = "linear"541 EASE_IN = "ease_in"542 EASE_OUT = "ease_out"543 EASE_IN_OUT = "ease_in_out"544 545 _LIST = [LINEAR, EASE_IN, EASE_OUT, EASE_IN_OUT]546 547 @classmethod548 def get_weights(cls, num_from: float, num_to: float, length: int, method: str, reverse=False):549 diff = num_to - num_from550 if method == cls.LINEAR:551 weights = torch.linspace(num_from, num_to, length)552 elif method == cls.EASE_IN:553 index = torch.linspace(0, 1, length)554 weights = diff * np.power(index, 2) + num_from555 elif method == cls.EASE_OUT:556 index = torch.linspace(0, 1, length)557 weights = diff * (1 - np.power(1 - index, 2)) + num_from558 elif method == cls.EASE_IN_OUT:559 index = torch.linspace(0, 1, length)560 weights = diff * ((1 - np.cos(index * np.pi)) / 2) + num_from561 else:562 raise ValueError(f"Unrecognized interpolation method '{method}'.")563 if reverse:564 weights = weights.flip(dims=(0,))565 return weights566 567def get_sorted_list_via_attr(objects: list, attr: str) -> list:568 if not objects:569 return objects570 elif len(objects) <= 1:571 return [x for x in objects]572 # now that we know we have to sort, do it following these rules:573 # a) if objects have same value of attribute, maintain their relative order574 # b) perform sorting of the groups of objects with same attributes575 unique_attrs = {}576 for o in objects:577 val_attr = getattr(o, attr)578 attr_list: list = unique_attrs.get(val_attr, list())579 attr_list.append(o)580 if val_attr not in unique_attrs:581 unique_attrs[val_attr] = attr_list582 # now that we have the unique attr values grouped together in relative order, sort them by key583 sorted_attrs = dict(sorted(unique_attrs.items()))584 # now flatten out the dict into a list to return585 sorted_list = []586 for object_list in sorted_attrs.values():587 sorted_list.extend(object_list)588 return sorted_list589 590def create_transformer_options_from_hooks(model: ModelPatcher, hooks: HookGroup, transformer_options: dict[str]=None):591 # if no hooks or is not a ModelPatcher for sampling, return empty dict592 if hooks is None or model.is_clip:593 return {}594 if transformer_options is None:595 transformer_options = {}596 for hook in hooks.get_type(EnumHookType.TransformerOptions):597 hook: TransformerOptionsHook598 hook.on_apply_hooks(model, transformer_options)599 return transformer_options600 601def create_hook_lora(lora: dict[str, torch.Tensor], strength_model: float, strength_clip: float):602 hook_group = HookGroup()603 hook = WeightHook(strength_model=strength_model, strength_clip=strength_clip)604 hook_group.add(hook)605 hook.weights = lora606 return hook_group607 608def create_hook_model_as_lora(weights_model, weights_clip, strength_model: float, strength_clip: float):609 hook_group = HookGroup()610 hook = WeightHook(strength_model=strength_model, strength_clip=strength_clip)611 hook_group.add(hook)612 patches_model = None613 patches_clip = None614 if weights_model is not None:615 patches_model = {}616 for key in weights_model:617 patches_model[key] = ("model_as_lora", (weights_model[key],))618 if weights_clip is not None:619 patches_clip = {}620 for key in weights_clip:621 patches_clip[key] = ("model_as_lora", (weights_clip[key],))622 hook.weights = patches_model623 hook.weights_clip = patches_clip624 hook.need_weight_init = False625 return hook_group626 627def get_patch_weights_from_model(model: ModelPatcher, discard_model_sampling=True):628 if model is None:629 return None630 patches_model: dict[str, torch.Tensor] = model.model.state_dict()631 if discard_model_sampling:632 # do not include ANY model_sampling components of the model that should act as a patch633 for key in list(patches_model.keys()):634 if key.startswith("model_sampling"):635 patches_model.pop(key, None)636 return patches_model637 638# NOTE: this function shows how to register weight hooks directly on the ModelPatchers639def load_hook_lora_for_models(model: ModelPatcher, clip: CLIP, lora: dict[str, torch.Tensor],640 strength_model: float, strength_clip: float):641 key_map = {}642 if model is not None:643 key_map = comfy.lora.model_lora_keys_unet(model.model, key_map)644 if clip is not None:645 key_map = comfy.lora.model_lora_keys_clip(clip.cond_stage_model, key_map)646 647 hook_group = HookGroup()648 hook = WeightHook()649 hook_group.add(hook)650 loaded: dict[str] = comfy.lora.load_lora(lora, key_map)651 if model is not None:652 new_modelpatcher = model.clone()653 k = new_modelpatcher.add_hook_patches(hook=hook, patches=loaded, strength_patch=strength_model)654 else:655 k = ()656 new_modelpatcher = None657 658 if clip is not None:659 new_clip = clip.clone()660 k1 = new_clip.patcher.add_hook_patches(hook=hook, patches=loaded, strength_patch=strength_clip)661 else:662 k1 = ()663 new_clip = None664 k = set(k)665 k1 = set(k1)666 for x in loaded:667 if (x not in k) and (x not in k1):668 logging.warning(f"NOT LOADED {x}")669 return (new_modelpatcher, new_clip, hook_group)670 671def _combine_hooks_from_values(c_dict: dict[str, HookGroup], values: dict[str, HookGroup], cache: dict[tuple[HookGroup, HookGroup], HookGroup]):672 hooks_key = 'hooks'673 # if hooks only exist in one dict, do what's needed so that it ends up in c_dict674 if hooks_key not in values:675 return676 if hooks_key not in c_dict:677 hooks_value = values.get(hooks_key, None)678 if hooks_value is not None:679 c_dict[hooks_key] = hooks_value680 return681 # otherwise, need to combine with minimum duplication via cache682 hooks_tuple = (c_dict[hooks_key], values[hooks_key])683 cached_hooks = cache.get(hooks_tuple, None)684 if cached_hooks is None:685 new_hooks = hooks_tuple[0].clone_and_combine(hooks_tuple[1])686 cache[hooks_tuple] = new_hooks687 c_dict[hooks_key] = new_hooks688 else:689 c_dict[hooks_key] = cache[hooks_tuple]690 691def conditioning_set_values_with_hooks(conditioning, values={}, append_hooks=True,692 cache: dict[tuple[HookGroup, HookGroup], HookGroup]=None):693 c = []694 if cache is None:695 cache = {}696 for t in conditioning:697 n = [t[0], t[1].copy()]698 for k in values:699 if append_hooks and k == 'hooks':700 _combine_hooks_from_values(n[1], values, cache)701 else:702 n[1][k] = values[k]703 c.append(n)704 705 return c706 707def set_hooks_for_conditioning(cond, hooks: HookGroup, append_hooks=True, cache: dict[tuple[HookGroup, HookGroup], HookGroup]=None):708 if hooks is None:709 return cond710 return conditioning_set_values_with_hooks(cond, {'hooks': hooks}, append_hooks=append_hooks, cache=cache)711 712def set_timesteps_for_conditioning(cond, timestep_range: tuple[float,float]):713 if timestep_range is None:714 return cond715 return conditioning_set_values(cond, {"start_percent": timestep_range[0],716 "end_percent": timestep_range[1]})717 718def set_mask_for_conditioning(cond, mask: torch.Tensor, set_cond_area: str, strength: float):719 if mask is None:720 return cond721 set_area_to_bounds = False722 if set_cond_area != 'default':723 set_area_to_bounds = True724 if len(mask.shape) < 3:725 mask = mask.unsqueeze(0)726 return conditioning_set_values(cond, {'mask': mask,727 'set_area_to_bounds': set_area_to_bounds,728 'mask_strength': strength})729 730def combine_conditioning(conds: list):731 combined_conds = []732 for cond in conds:733 combined_conds.extend(cond)734 return combined_conds735 736def combine_with_new_conds(conds: list, new_conds: list):737 combined_conds = []738 for c, new_c in zip(conds, new_conds):739 combined_conds.append(combine_conditioning([c, new_c]))740 return combined_conds741 742def set_conds_props(conds: list, strength: float, set_cond_area: str,743 mask: torch.Tensor=None, hooks: HookGroup=None, timesteps_range: tuple[float,float]=None, append_hooks=True):744 final_conds = []745 cache = {}746 for c in conds:747 # first, apply lora_hook to conditioning, if provided748 c = set_hooks_for_conditioning(c, hooks, append_hooks=append_hooks, cache=cache)749 # next, apply mask to conditioning750 c = set_mask_for_conditioning(cond=c, mask=mask, strength=strength, set_cond_area=set_cond_area)751 # apply timesteps, if present752 c = set_timesteps_for_conditioning(cond=c, timestep_range=timesteps_range)753 # finally, apply mask to conditioning and store754 final_conds.append(c)755 return final_conds756 757def set_conds_props_and_combine(conds: list, new_conds: list, strength: float=1.0, set_cond_area: str="default",758 mask: torch.Tensor=None, hooks: HookGroup=None, timesteps_range: tuple[float,float]=None, append_hooks=True):759 combined_conds = []760 cache = {}761 for c, masked_c in zip(conds, new_conds):762 # first, apply lora_hook to new conditioning, if provided763 masked_c = set_hooks_for_conditioning(masked_c, hooks, append_hooks=append_hooks, cache=cache)764 # next, apply mask to new conditioning, if provided765 masked_c = set_mask_for_conditioning(cond=masked_c, mask=mask, set_cond_area=set_cond_area, strength=strength)766 # apply timesteps, if present767 masked_c = set_timesteps_for_conditioning(cond=masked_c, timestep_range=timesteps_range)768 # finally, combine with existing conditioning and store769 combined_conds.append(combine_conditioning([c, masked_c]))770 return combined_conds771 772def set_default_conds_and_combine(conds: list, new_conds: list,773 hooks: HookGroup=None, timesteps_range: tuple[float,float]=None, append_hooks=True):774 combined_conds = []775 cache = {}776 for c, new_c in zip(conds, new_conds):777 # first, apply lora_hook to new conditioning, if provided778 new_c = set_hooks_for_conditioning(new_c, hooks, append_hooks=append_hooks, cache=cache)779 # next, add default_cond key to cond so that during sampling, it can be identified780 new_c = conditioning_set_values(new_c, {'default': True})781 # apply timesteps, if present782 new_c = set_timesteps_for_conditioning(cond=new_c, timestep_range=timesteps_range)783 # finally, combine with existing conditioning and store784 combined_conds.append(combine_conditioning([c, new_c]))785 return combined_conds786