Added create_model_options_clone func, modified type annotations to use __future__ so that I can use the better type annotations

This commit is contained in:
kosinkadink1@gmail.com 2024-09-25 20:29:49 +09:00
parent fd2d572447
commit 09cbd69161
5 changed files with 81 additions and 66 deletions

View File

@ -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

View File

@ -16,7 +16,8 @@
along with this program. If not, see <https://www.gnu.org/licenses/>.
"""
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

View File

@ -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")

View File

@ -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

View File

@ -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: