diff --git a/comfy/hooks.py b/comfy/hooks.py index 962cb81b7..96b37282c 100644 --- a/comfy/hooks.py +++ b/comfy/hooks.py @@ -1,4 +1,5 @@ -from typing import TYPE_CHECKING, List, Dict, Tuple, Callable +from __future__ import annotations +from typing import TYPE_CHECKING, Callable import enum import math import torch @@ -78,8 +79,8 @@ class Hook: class WeightHook(Hook): def __init__(self, strength_model=1.0, strength_clip=1.0): super().__init__(hook_type=EnumHookType.Weight) - self.weights: Dict = None - self.weights_clip: Dict = None + self.weights: dict = None + self.weights_clip: dict = None self.is_diff = False self.need_weight_init = True self._strength_model = strength_model @@ -132,7 +133,7 @@ class WeightHook(Hook): class PatchHook(Hook): def __init__(self): super().__init__(hook_type=EnumHookType.Patch) - self.patches: Dict = None + self.patches: dict = None def clone(self, subtype: Callable=None): if subtype is None: @@ -147,7 +148,7 @@ class PatchHook(Hook): class ObjectPatchHook(Hook): def __init__(self): super().__init__(hook_type=EnumHookType.ObjectPatch) - self.object_patches: Dict = None + self.object_patches: dict = None def clone(self, subtype: Callable=None): if subtype is None: @@ -160,7 +161,7 @@ class ObjectPatchHook(Hook): pass class AddModelsHook(Hook): - def __init__(self, key: str=None, models: List['ModelPatcher']=None): + def __init__(self, key: str=None, models: list['ModelPatcher']=None): super().__init__(hook_type=EnumHookType.AddModels) self.key = key self.models = models @@ -196,7 +197,7 @@ class AddCallbackHook(Hook): pass class SetInjectionsHook(Hook): - def __init__(self, key: str=None, injections: List['PatcherInjection']=None): + def __init__(self, key: str=None, injections: list['PatcherInjection']=None): super().__init__(hook_type=EnumHookType.SetInjections) self.key = key self.injections = injections @@ -231,7 +232,7 @@ class AddWrapperHook(Hook): class HookGroup: def __init__(self): - self.hooks: List[Hook] = [] + self.hooks: list[Hook] = [] def add(self, hook: Hook): if hook not in self.hooks: @@ -258,15 +259,15 @@ class HookGroup: hook.hook_keyframe = hook_kf def get_dict_repr(self): - d: Dict[EnumHookType, Dict[Hook, None]] = {} + d: dict[EnumHookType, dict[Hook, None]] = {} for hook in self.hooks: with_type = d.setdefault(hook.hook_type, {}) with_type[hook] = None return d @staticmethod - def combine_all_hooks(hooks_list: List['HookGroup'], require_count=0) -> 'HookGroup': - actual: List[HookGroup] = [] + def combine_all_hooks(hooks_list: list['HookGroup'], require_count=0) -> 'HookGroup': + actual: list[HookGroup] = [] for group in hooks_list: if group is not None: actual.append(group) @@ -303,7 +304,7 @@ class HookKeyframe: class HookKeyframeGroup: def __init__(self): - self.keyframes: List[HookKeyframe] = [] + self.keyframes: list[HookKeyframe] = [] self._current_keyframe: HookKeyframe = None self._current_used_steps = 0 self._current_index = 0 @@ -411,7 +412,7 @@ class InterpolationMethod: weights = weights.flip(dims=(0,)) return weights -def get_sorted_list_via_attr(objects: List, attr: str) -> List: +def get_sorted_list_via_attr(objects: list, attr: str) -> list: if not objects: return objects elif len(objects) <= 1: @@ -422,7 +423,7 @@ def get_sorted_list_via_attr(objects: List, attr: str) -> List: unique_attrs = {} for o in objects: val_attr = getattr(o, attr) - attr_list: List = unique_attrs.get(val_attr, list()) + 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 @@ -434,7 +435,7 @@ def get_sorted_list_via_attr(objects: List, attr: str) -> List: sorted_list.extend(object_list) return sorted_list -def create_hook_lora(lora: Dict[str, torch.Tensor], strength_model: float, strength_clip: float): +def create_hook_lora(lora: dict[str, torch.Tensor], strength_model: float, strength_clip: float): hook_group = HookGroup() hook = WeightHook(strength_model=strength_model, strength_clip=strength_clip) hook_group.add(hook) @@ -463,7 +464,7 @@ def create_hook_model_as_lora(weights_model, weights_clip, strength_model: float def get_patch_weights_from_model(model: 'ModelPatcher', discard_model_sampling=False): if model is None: return None - patches_model: Dict[str, torch.Tensor] = model.model.state_dict() + patches_model: dict[str, torch.Tensor] = model.model.state_dict() if discard_model_sampling: # do not include ANY model_sampling components of the model that should act as a patch for key in list(patches_model.keys()): @@ -479,7 +480,7 @@ def create_hook_model_as_lora_precalc(model: 'ModelPatcher', clip: 'CLIP', hook_group.add(hook) if model is not None and model_loaded is not None: expected_model_keys = set(model_loaded.model.state_dict().keys()) - patches_model: Dict[str, torch.Tensor] = model_loaded.model.state_dict() + patches_model: dict[str, torch.Tensor] = model_loaded.model.state_dict() # do not include ANY model_sampling components of the model that should act as a patch for key in list(patches_model.keys()): if key.startswith("model_sampling"): @@ -492,7 +493,7 @@ def create_hook_model_as_lora_precalc(model: 'ModelPatcher', clip: 'CLIP', if clip is not None and clip_loaded is not None: expected_clip_keys = clip_loaded.patcher.model.state_dict().copy() - patches_clip: Dict[str, torch.Tensor] = clip_loaded.cond_stage_model.state_dict() + patches_clip: dict[str, torch.Tensor] = clip_loaded.cond_stage_model.state_dict() weights_clip, k1 = clip.patcher.get_weight_diffs(patches_clip) else: weights_clip = {} @@ -514,7 +515,7 @@ def create_hook_model_as_lora_precalc(model: 'ModelPatcher', clip: 'CLIP', hook.need_weight_init = False return hook_group -def load_hook_lora_for_models(model: 'ModelPatcher', clip: 'CLIP', lora: Dict[str, torch.Tensor], +def load_hook_lora_for_models(model: 'ModelPatcher', clip: 'CLIP', lora: dict[str, torch.Tensor], strength_model: float, strength_clip: float): key_map = {} if model is not None: @@ -525,7 +526,7 @@ def load_hook_lora_for_models(model: 'ModelPatcher', clip: 'CLIP', lora: Dict[st hook_group = HookGroup() hook = WeightHook() hook_group.add(hook) - loaded: Dict[str] = comfy.lora.load_lora(lora, key_map) + loaded: dict[str] = comfy.lora.load_lora(lora, key_map) if model is not None: new_modelpatcher = model.clone() k = new_modelpatcher.add_hook_patches(hook=hook, patches=loaded, strength_patch=strength_model) @@ -555,7 +556,7 @@ def load_hook_model_as_lora_for_models(model: 'ModelPatcher', clip: 'CLIP', if model is not None and model_loaded is not None: new_modelpatcher = model.clone() expected_model_keys = set(model_loaded.model.state_dict().keys()) - patches_model: Dict[str, torch.Tensor] = model_loaded.model.state_dict() + patches_model: dict[str, torch.Tensor] = model_loaded.model.state_dict() # do not include ANY model_sampling components of the model that should act as a patch for key in list(patches_model.keys()): if key.startswith("model_sampling"): @@ -570,7 +571,7 @@ def load_hook_model_as_lora_for_models(model: 'ModelPatcher', clip: 'CLIP', new_clip = clip.clone() comfy.model_management.unload_model_clones(new_clip.patcher) expected_clip_keys = clip_loaded.patcher.model.state_dict().copy() - patches_clip: Dict[str, torch.Tensor] = clip_loaded.cond_stage_model.state_dict() + patches_clip: dict[str, torch.Tensor] = clip_loaded.cond_stage_model.state_dict() k1 = new_clip.patcher.add_hook_patches(hook=hook, patches=patches_clip, strength_patch=strength_clip, is_diff=True) else: k1 = () @@ -594,7 +595,7 @@ def set_hooks_for_conditioning(cond, hooks: HookGroup): return cond return conditioning_set_values(cond, {'hooks': hooks}) -def set_timesteps_for_conditioning(cond, timestep_range: Tuple[float,float]): +def set_timesteps_for_conditioning(cond, timestep_range: tuple[float,float]): if timestep_range is None: return cond return conditioning_set_values(cond, {"start_percent": timestep_range[0], @@ -612,14 +613,14 @@ def set_mask_for_conditioning(cond, mask: torch.Tensor, set_cond_area: str, stre 'set_area_to_bounds': set_area_to_bounds, 'mask_strength': strength}) -def combine_conditioning(conds: List): +def combine_conditioning(conds: list): combined_conds = [] for cond in conds: combined_conds.extend(cond) return combined_conds -def set_mask_conds(conds: List, strength: float, set_cond_area: str, - opt_mask: torch.Tensor=None, opt_hooks: HookGroup=None, opt_timestep_range: Tuple[float,float]=None): +def set_mask_conds(conds: list, strength: float, set_cond_area: str, + opt_mask: torch.Tensor=None, opt_hooks: HookGroup=None, opt_timestep_range: tuple[float,float]=None): masked_conds = [] for c in conds: # first, apply lora_hook to conditioning, if provided @@ -632,8 +633,8 @@ def set_mask_conds(conds: List, strength: float, set_cond_area: str, masked_conds.append(c) return masked_conds -def set_mask_and_combine_conds(conds: List, new_conds: List, strength: float=1.0, set_cond_area: str="default", - opt_mask: torch.Tensor=None, opt_hooks: HookGroup=None, opt_timestep_range: Tuple[float,float]=None): +def set_mask_and_combine_conds(conds: list, new_conds: list, strength: float=1.0, set_cond_area: str="default", + opt_mask: torch.Tensor=None, opt_hooks: HookGroup=None, opt_timestep_range: tuple[float,float]=None): combined_conds = [] for c, masked_c in zip(conds, new_conds): # first, apply lora_hook to new conditioning, if provided @@ -647,7 +648,7 @@ def set_mask_and_combine_conds(conds: List, new_conds: List, strength: float=1.0 return combined_conds def set_default_and_combine_conds(conds: list, new_conds: list, - opt_hooks: HookGroup=None, opt_timestep_range: Tuple[float,float]=None): + opt_hooks: HookGroup=None, opt_timestep_range: tuple[float,float]=None): combined_conds = [] for c, new_c in zip(conds, new_conds): # first, apply lora_hook to new conditioning, if provided diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index af5f41e2f..62ec72dd5 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -16,7 +16,8 @@ along with this program. If not, see . """ -from typing import Dict, List, Tuple, Optional, Callable +from __future__ import annotations +from typing import Optional, Callable import torch import copy import inspect @@ -78,6 +79,17 @@ def set_model_options_pre_cfg_function(model_options, pre_cfg_function, disable_ model_options["disable_cfg1_optimization"] = True return model_options +def create_model_options_clone(orig_model_options: dict): + def copy_nested_dicts(input_dict: dict): + new_dict = input_dict.copy() + for key, value in input_dict.items(): + if isinstance(value, dict): + new_dict[key] = copy_nested_dicts(value) + elif isinstance(value, list): + new_dict[key] = value.copy() + return new_dict + return copy_nested_dicts(orig_model_options) + def create_hook_patches_clone(orig_hook_patches): new_hook_patches = {} for hook_ref in orig_hook_patches: @@ -137,7 +149,7 @@ class WrappersMP: } class WrapperExecutor: - def __init__(self, original: Callable, wrappers: List[Callable], idx: int): + def __init__(self, original: Callable, wrappers: list[Callable], idx: int): self.original = original self.wrappers = wrappers.copy() self.idx = idx @@ -161,11 +173,11 @@ class WrapperExecutor: return WrapperExecutor(self.original, self.wrappers, new_idx) @classmethod - def new_executor(cls, original: Callable, wrappers: List[Callable]): + def new_executor(cls, original: Callable, wrappers: list[Callable]): return cls(original, wrappers, idx=0) class WrapperClassExecutor: - def __init__(self, original: Callable, wrappers: List[Callable], idx: int): + def __init__(self, original: Callable, wrappers: list[Callable], idx: int): self.original = original self.wrappers = wrappers.copy() self.idx = idx @@ -189,7 +201,7 @@ class WrapperClassExecutor: return WrapperClassExecutor(self.original, self.wrappers, new_idx) @classmethod - def new_executor(cls, original: Callable, wrappers: List[Callable]): + def new_executor(cls, original: Callable, wrappers: list[Callable]): return cls(original, wrappers, idx=0) class AutoPatcherEjector: @@ -261,19 +273,19 @@ class ModelPatcher: self.weight_inplace_update = weight_inplace_update self.patches_uuid = uuid.uuid4() - self.attachments: Dict[str] = {} - self.additional_models: Dict[str, List[ModelPatcher]] = {} - self.callbacks: Dict[str, Dict[str, List[Callable]]] = CallbacksMP.init_callbacks() - self.wrappers: Dict[str, Dict[str, List[Callable]]] = WrappersMP.init_wrappers() + self.attachments: dict[str] = {} + self.additional_models: dict[str, list[ModelPatcher]] = {} + self.callbacks: dict[str, dict[str, list[Callable]]] = CallbacksMP.init_callbacks() + self.wrappers: dict[str, dict[str, list[Callable]]] = WrappersMP.init_wrappers() self.is_injected = False self.skip_injection = False - self.injections: Dict[str, List[PatcherInjection]] = {} + self.injections: dict[str, list[PatcherInjection]] = {} - self.hook_patches: Dict[comfy.hooks._HookRef] = {} - self.hook_patches_backup: Dict[comfy.hooks._HookRef] = {} - self.hook_backup: Dict[str, Tuple[torch.Tensor, torch.device]] = {} - self.cached_hook_patches: Dict[comfy.hooks.HookGroup, Dict[str, torch.Tensor]] = {} + self.hook_patches: dict[comfy.hooks._HookRef] = {} + self.hook_patches_backup: dict[comfy.hooks._HookRef] = {} + self.hook_backup: dict[str, tuple[torch.Tensor, torch.device]] = {} + self.cached_hook_patches: dict[comfy.hooks.HookGroup, dict[str, torch.Tensor]] = {} self.current_hooks: Optional[comfy.hooks.HookGroup] = None self.forced_hooks: Optional[comfy.hooks.HookGroup] = None # NOTE: only used for CLIP # TODO: hook_mode should be entirely removed; behavior should be determined by remaining VRAM/memory @@ -850,14 +862,14 @@ class ModelPatcher: def get_attachment(self, key: str): return self.attachments.get(key, None) - def set_injections(self, key: str, injections: List[PatcherInjection]): + def set_injections(self, key: str, injections: list[PatcherInjection]): self.injections[key] = injections def remove_injections(self, key: str): if key in self.injections: self.injections.pop(key) - def set_additional_models(self, key: str, models: List['ModelPatcher']): + def set_additional_models(self, key: str, models: list['ModelPatcher']): self.additional_models[key] = models def remove_additional_models(self, key: str): @@ -927,9 +939,9 @@ class ModelPatcher: if cached_group.contains(hook): self.cached_hook_patches.pop(cached_group) - def register_all_hook_patches(self, hooks_dict: Dict[comfy.hooks.EnumHookType, Dict[comfy.hooks.Hook, None]], target: comfy.hooks.EnumWeightTarget): + def register_all_hook_patches(self, hooks_dict: dict[comfy.hooks.EnumHookType, dict[comfy.hooks.Hook, None]], target: comfy.hooks.EnumWeightTarget): self.restore_hook_patches() - weight_hooks_to_register: List[comfy.hooks.WeightHook] = [] + weight_hooks_to_register: list[comfy.hooks.WeightHook] = [] for hook in hooks_dict.get(comfy.hooks.EnumHookType.Weight, {}): if hook.hook_ref not in self.hook_patches: weight_hooks_to_register.append(hook) @@ -945,7 +957,7 @@ class ModelPatcher: # NOTE: this mirrors behavior of add_patches func if is_diff: comfy.model_management.unload_model_clones(self) - current_hook_patches: Dict[str,List] = self.hook_patches.get(hook.hook_ref, {}) + current_hook_patches: dict[str,list] = self.hook_patches.get(hook.hook_ref, {}) p = set() model_sd = self.model.state_dict() for k in patches: @@ -961,7 +973,7 @@ class ModelPatcher: if key in model_sd: p.add(k) - current_patches: List[Tuple] = current_hook_patches.get(key, []) + current_patches: list[tuple] = current_hook_patches.get(key, []) if is_diff: # take difference between desired weight and existing weight to get diff # TODO: try to implement diff via strength_path/strength_model diff @@ -982,7 +994,7 @@ class ModelPatcher: def get_weight_diffs(self, patches): with self.use_ejected(): comfy.model_management.unload_model_clones(self) - weights: Dict[str, Tuple] = {} + weights: dict[str, tuple] = {} p = set() model_sd = self.model.state_dict() for k in patches: @@ -1001,9 +1013,9 @@ class ModelPatcher: combined_patches = {} if hooks is not None: for hook in hooks.hooks: - hook_patches: Dict = self.hook_patches.get(hook.hook_ref, {}) + hook_patches: dict = self.hook_patches.get(hook.hook_ref, {}) for key in hook_patches.keys(): - current_patches: List[Tuple] = combined_patches.get(key, []) + current_patches: list[tuple] = combined_patches.get(key, []) if math.isclose(hook.strength, 1.0): current_patches.extend(hook_patches[key]) else: @@ -1050,7 +1062,7 @@ class ModelPatcher: memory_counter=memory_counter) self.current_hooks = hooks - def patch_cached_hook_weights(self, cached_weights: Dict, key: str, memory_counter: MemoryCounter): + def patch_cached_hook_weights(self, cached_weights: dict, key: str, memory_counter: MemoryCounter): if key not in self.hook_backup: weight: torch.Tensor = comfy.utils.get_attr(self.model, key) target_device = self.offload_device diff --git a/comfy/sampler_helpers.py b/comfy/sampler_helpers.py index 36325cfa2..0e911b3ff 100644 --- a/comfy/sampler_helpers.py +++ b/comfy/sampler_helpers.py @@ -1,8 +1,9 @@ +from __future__ import annotations import torch import comfy.model_management import comfy.conds import comfy.hooks -from typing import TYPE_CHECKING, Dict, List +from typing import TYPE_CHECKING if TYPE_CHECKING: from comfy.model_patcher import ModelPatcher from comfy.model_base import BaseModel @@ -26,7 +27,7 @@ def get_models_from_cond(cond, model_type): models += [c[model_type]] return models -def get_hooks_from_cond(cond, hooks_dict: Dict[comfy.hooks.EnumHookType, Dict[comfy.hooks.Hook, None]]): +def get_hooks_from_cond(cond, hooks_dict: dict[comfy.hooks.EnumHookType, dict[comfy.hooks.Hook, None]]): for c in cond: if 'hooks' in c: for hook in c['hooks'].hooks: @@ -49,10 +50,10 @@ def convert_cond(cond): def get_additional_models(conds, dtype): """loads additional models in conditioning""" - cnets: List[ControlBase] = [] + cnets: list[ControlBase] = [] gligen = [] add_models = [] - hooks: Dict[comfy.hooks.EnumHookType, Dict[comfy.hooks.Hook, None]] = {} + hooks: dict[comfy.hooks.EnumHookType, dict[comfy.hooks.Hook, None]] = {} for k in conds: cnets += get_models_from_cond(conds[k], "control") diff --git a/comfy/samplers.py b/comfy/samplers.py index 819e5c13a..e5ebda600 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -1,6 +1,6 @@ from .k_diffusion import sampling as k_diffusion_sampling from .extra_samplers import uni_pc -from typing import TYPE_CHECKING, Dict, List, Tuple +from typing import TYPE_CHECKING if TYPE_CHECKING: from comfy.model_patcher import ModelPatcher from comfy.model_base import BaseModel @@ -144,7 +144,7 @@ def cond_cat(c_list): return out -def finalize_default_conds(hooked_to_run: Dict[comfy.hooks.HookGroup,List[Tuple[Tuple,int]]], default_conds: List[List[Dict]], x_in, timestep): +def finalize_default_conds(hooked_to_run: dict[comfy.hooks.HookGroup,list[tuple[tuple,int]]], default_conds: list[list[dict]], x_in, timestep): # need to figure out remaining unmasked area for conds default_mults = [] for _ in default_conds: @@ -182,19 +182,19 @@ def finalize_default_conds(hooked_to_run: Dict[comfy.hooks.HookGroup,List[Tuple[ hooked_to_run.setdefault(hook, list()) hooked_to_run[hook] += [(p, i)] -def calc_cond_batch(model: 'BaseModel', conds: List[List[Dict]], x_in: torch.Tensor, timestep, model_options): +def calc_cond_batch(model: 'BaseModel', conds: list[list[dict]], x_in: torch.Tensor, timestep, model_options): executor = comfy.model_patcher.WrapperExecutor.new_executor( outer_calc_cond_batch, model.current_patcher.get_all_wrappers(comfy.model_patcher.WrappersMP.CALC_COND_BATCH) ) return executor._execute(model, conds, x_in, timestep, model_options) -def outer_calc_cond_batch(model: 'BaseModel', conds: List[List[Dict]], x_in: torch.Tensor, timestep, model_options): +def outer_calc_cond_batch(model: 'BaseModel', conds: list[list[dict]], x_in: torch.Tensor, timestep, model_options): out_conds = [] out_counts = [] # separate conds by matching hooks # TODO: implement default_conds support - hooked_to_run: Dict[comfy.hooks.HookGroup,List[Tuple[Tuple,int]]] = {} + hooked_to_run: dict[comfy.hooks.HookGroup,list[tuple[tuple,int]]] = {} default_conds = [] has_default_conds = False diff --git a/comfy_extras/nodes_hooks.py b/comfy_extras/nodes_hooks.py index b8b43c829..4ef639f34 100644 --- a/comfy_extras/nodes_hooks.py +++ b/comfy_extras/nodes_hooks.py @@ -1,4 +1,5 @@ -from typing import TYPE_CHECKING, Dict, List, Tuple, Union +from __future__ import annotations +from typing import TYPE_CHECKING, Union import torch from collections.abc import Iterable @@ -40,7 +41,7 @@ class PairConditioningSetProperties: def set_properties(self, positive_NEW, negative_NEW, strength: float, set_cond_area: str, - opt_mask: torch.Tensor=None, opt_hooks: comfy.hooks.HookGroup=None, opt_timesteps: Tuple=None): + opt_mask: torch.Tensor=None, opt_hooks: comfy.hooks.HookGroup=None, opt_timesteps: tuple=None): final_positive, final_negative = comfy.hooks.set_mask_conds(conds=[positive_NEW, negative_NEW], strength=strength, set_cond_area=set_cond_area, opt_mask=opt_mask, opt_hooks=opt_hooks, opt_timestep_range=opt_timesteps) @@ -70,7 +71,7 @@ class ConditioningSetProperties: def set_properties(self, cond_NEW, strength: float, set_cond_area: str, - opt_mask: torch.Tensor=None, opt_hooks: comfy.hooks.HookGroup=None, opt_timesteps: Tuple=None): + opt_mask: torch.Tensor=None, opt_hooks: comfy.hooks.HookGroup=None, opt_timesteps: tuple=None): (final_cond,) = comfy.hooks.set_mask_conds(conds=[cond_NEW], strength=strength, set_cond_area=set_cond_area, opt_mask=opt_mask, opt_hooks=opt_hooks, opt_timestep_range=opt_timesteps) @@ -596,7 +597,7 @@ class CreateHookKeyframesFromFloats: CATEGORY = "advanced/hooks/scheduling" FUNCTION = "create_hook_keyframes" - def create_hook_keyframes(self, floats_strength: Union[float, List[float]], + def create_hook_keyframes(self, floats_strength: Union[float, list[float]], start_percent: float, end_percent: float, prev_hook_kf: comfy.hooks.HookKeyframeGroup=None, print_keyframes=False): if prev_hook_kf is None: