fred-dev/comfy_ui_ali
0
1"""2 This file is part of ComfyUI.3 Copyright (C) 2024 Comfy4 5 This program is free software: you can redistribute it and/or modify6 it under the terms of the GNU General Public License as published by7 the Free Software Foundation, either version 3 of the License, or8 (at your option) any later version.9 10 This program is distributed in the hope that it will be useful,11 but WITHOUT ANY WARRANTY; without even the implied warranty of12 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the13 GNU General Public License for more details.14 15 You should have received a copy of the GNU General Public License16 along with this program. If not, see <https://www.gnu.org/licenses/>.17"""18 19from __future__ import annotations20from typing import Optional, Callable21import torch22import copy23import inspect24import logging25import uuid26import collections27import math28 29import comfy.utils30import comfy.float31import comfy.model_management32import comfy.lora33import comfy.hooks34import comfy.patcher_extension35from comfy.patcher_extension import CallbacksMP, WrappersMP, PatcherInjection36from comfy.comfy_types import UnetWrapperFunction37 38def string_to_seed(data):39 crc = 0xFFFFFFFF40 for byte in data:41 if isinstance(byte, str):42 byte = ord(byte)43 crc ^= byte44 for _ in range(8):45 if crc & 1:46 crc = (crc >> 1) ^ 0xEDB8832047 else:48 crc >>= 149 return crc ^ 0xFFFFFFFF50 51def set_model_options_patch_replace(model_options, patch, name, block_name, number, transformer_index=None):52 to = model_options["transformer_options"].copy()53 54 if "patches_replace" not in to:55 to["patches_replace"] = {}56 else:57 to["patches_replace"] = to["patches_replace"].copy()58 59 if name not in to["patches_replace"]:60 to["patches_replace"][name] = {}61 else:62 to["patches_replace"][name] = to["patches_replace"][name].copy()63 64 if transformer_index is not None:65 block = (block_name, number, transformer_index)66 else:67 block = (block_name, number)68 to["patches_replace"][name][block] = patch69 model_options["transformer_options"] = to70 return model_options71 72def set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=False):73 model_options["sampler_post_cfg_function"] = model_options.get("sampler_post_cfg_function", []) + [post_cfg_function]74 if disable_cfg1_optimization:75 model_options["disable_cfg1_optimization"] = True76 return model_options77 78def set_model_options_pre_cfg_function(model_options, pre_cfg_function, disable_cfg1_optimization=False):79 model_options["sampler_pre_cfg_function"] = model_options.get("sampler_pre_cfg_function", []) + [pre_cfg_function]80 if disable_cfg1_optimization:81 model_options["disable_cfg1_optimization"] = True82 return model_options83 84def create_model_options_clone(orig_model_options: dict):85 return comfy.patcher_extension.copy_nested_dicts(orig_model_options)86 87def create_hook_patches_clone(orig_hook_patches):88 new_hook_patches = {}89 for hook_ref in orig_hook_patches:90 new_hook_patches[hook_ref] = {}91 for k in orig_hook_patches[hook_ref]:92 new_hook_patches[hook_ref][k] = orig_hook_patches[hook_ref][k][:]93 return new_hook_patches94 95def wipe_lowvram_weight(m):96 if hasattr(m, "prev_comfy_cast_weights"):97 m.comfy_cast_weights = m.prev_comfy_cast_weights98 del m.prev_comfy_cast_weights99 100 if hasattr(m, "weight_function"):101 m.weight_function = []102 103 if hasattr(m, "bias_function"):104 m.bias_function = []105 106def move_weight_functions(m, device):107 if device is None:108 return 0109 110 memory = 0111 if hasattr(m, "weight_function"):112 for f in m.weight_function:113 if hasattr(f, "move_to"):114 memory += f.move_to(device=device)115 116 if hasattr(m, "bias_function"):117 for f in m.bias_function:118 if hasattr(f, "move_to"):119 memory += f.move_to(device=device)120 return memory121 122class LowVramPatch:123 def __init__(self, key, patches):124 self.key = key125 self.patches = patches126 def __call__(self, weight):127 intermediate_dtype = weight.dtype128 if intermediate_dtype not in [torch.float32, torch.float16, torch.bfloat16]: #intermediate_dtype has to be one that is supported in math ops129 intermediate_dtype = torch.float32130 return comfy.float.stochastic_rounding(comfy.lora.calculate_weight(self.patches[self.key], weight.to(intermediate_dtype), self.key, intermediate_dtype=intermediate_dtype), weight.dtype, seed=string_to_seed(self.key))131 132 return comfy.lora.calculate_weight(self.patches[self.key], weight, self.key, intermediate_dtype=intermediate_dtype)133 134def get_key_weight(model, key):135 set_func = None136 convert_func = None137 op_keys = key.rsplit('.', 1)138 if len(op_keys) < 2:139 weight = comfy.utils.get_attr(model, key)140 else:141 op = comfy.utils.get_attr(model, op_keys[0])142 try:143 set_func = getattr(op, "set_{}".format(op_keys[1]))144 except AttributeError:145 pass146 147 try:148 convert_func = getattr(op, "convert_{}".format(op_keys[1]))149 except AttributeError:150 pass151 152 weight = getattr(op, op_keys[1])153 if convert_func is not None:154 weight = comfy.utils.get_attr(model, key)155 156 return weight, set_func, convert_func157 158class AutoPatcherEjector:159 def __init__(self, model: 'ModelPatcher', skip_and_inject_on_exit_only=False):160 self.model = model161 self.was_injected = False162 self.prev_skip_injection = False163 self.skip_and_inject_on_exit_only = skip_and_inject_on_exit_only164 165 def __enter__(self):166 self.was_injected = False167 self.prev_skip_injection = self.model.skip_injection168 if self.skip_and_inject_on_exit_only:169 self.model.skip_injection = True170 if self.model.is_injected:171 self.model.eject_model()172 self.was_injected = True173 174 def __exit__(self, *args):175 if self.skip_and_inject_on_exit_only:176 self.model.skip_injection = self.prev_skip_injection177 self.model.inject_model()178 if self.was_injected and not self.model.skip_injection:179 self.model.inject_model()180 self.model.skip_injection = self.prev_skip_injection181 182class MemoryCounter:183 def __init__(self, initial: int, minimum=0):184 self.value = initial185 self.minimum = minimum186 # TODO: add a safe limit besides 0187 188 def use(self, weight: torch.Tensor):189 weight_size = weight.nelement() * weight.element_size()190 if self.is_useable(weight_size):191 self.decrement(weight_size)192 return True193 return False194 195 def is_useable(self, used: int):196 return self.value - used > self.minimum197 198 def decrement(self, used: int):199 self.value -= used200 201class ModelPatcher:202 def __init__(self, model, load_device, offload_device, size=0, weight_inplace_update=False):203 self.size = size204 self.model = model205 if not hasattr(self.model, 'device'):206 logging.debug("Model doesn't have a device attribute.")207 self.model.device = offload_device208 elif self.model.device is None:209 self.model.device = offload_device210 211 self.patches = {}212 self.backup = {}213 self.object_patches = {}214 self.object_patches_backup = {}215 self.weight_wrapper_patches = {}216 self.model_options = {"transformer_options":{}}217 self.model_size()218 self.load_device = load_device219 self.offload_device = offload_device220 self.weight_inplace_update = weight_inplace_update221 self.force_cast_weights = False222 self.patches_uuid = uuid.uuid4()223 self.parent = None224 225 self.attachments: dict[str] = {}226 self.additional_models: dict[str, list[ModelPatcher]] = {}227 self.callbacks: dict[str, dict[str, list[Callable]]] = CallbacksMP.init_callbacks()228 self.wrappers: dict[str, dict[str, list[Callable]]] = WrappersMP.init_wrappers()229 230 self.is_injected = False231 self.skip_injection = False232 self.injections: dict[str, list[PatcherInjection]] = {}233 234 self.hook_patches: dict[comfy.hooks._HookRef] = {}235 self.hook_patches_backup: dict[comfy.hooks._HookRef] = None236 self.hook_backup: dict[str, tuple[torch.Tensor, torch.device]] = {}237 self.cached_hook_patches: dict[comfy.hooks.HookGroup, dict[str, torch.Tensor]] = {}238 self.current_hooks: Optional[comfy.hooks.HookGroup] = None239 self.forced_hooks: Optional[comfy.hooks.HookGroup] = None # NOTE: only used for CLIP at this time240 self.is_clip = False241 self.hook_mode = comfy.hooks.EnumHookMode.MaxSpeed242 243 if not hasattr(self.model, 'model_loaded_weight_memory'):244 self.model.model_loaded_weight_memory = 0245 246 if not hasattr(self.model, 'lowvram_patch_counter'):247 self.model.lowvram_patch_counter = 0248 249 if not hasattr(self.model, 'model_lowvram'):250 self.model.model_lowvram = False251 252 if not hasattr(self.model, 'current_weight_patches_uuid'):253 self.model.current_weight_patches_uuid = None254 255 def model_size(self):256 if self.size > 0:257 return self.size258 self.size = comfy.model_management.module_size(self.model)259 return self.size260 261 def loaded_size(self):262 return self.model.model_loaded_weight_memory263 264 def lowvram_patch_counter(self):265 return self.model.lowvram_patch_counter266 267 def clone(self):268 n = self.__class__(self.model, self.load_device, self.offload_device, self.size, weight_inplace_update=self.weight_inplace_update)269 n.patches = {}270 for k in self.patches:271 n.patches[k] = self.patches[k][:]272 n.patches_uuid = self.patches_uuid273 274 n.object_patches = self.object_patches.copy()275 n.weight_wrapper_patches = self.weight_wrapper_patches.copy()276 n.model_options = copy.deepcopy(self.model_options)277 n.backup = self.backup278 n.object_patches_backup = self.object_patches_backup279 n.parent = self280 281 n.force_cast_weights = self.force_cast_weights282 283 # attachments284 n.attachments = {}285 for k in self.attachments:286 if hasattr(self.attachments[k], "on_model_patcher_clone"):287 n.attachments[k] = self.attachments[k].on_model_patcher_clone()288 else:289 n.attachments[k] = self.attachments[k]290 # additional models291 for k, c in self.additional_models.items():292 n.additional_models[k] = [x.clone() for x in c]293 # callbacks294 for k, c in self.callbacks.items():295 n.callbacks[k] = {}296 for k1, c1 in c.items():297 n.callbacks[k][k1] = c1.copy()298 # sample wrappers299 for k, w in self.wrappers.items():300 n.wrappers[k] = {}301 for k1, w1 in w.items():302 n.wrappers[k][k1] = w1.copy()303 # injection304 n.is_injected = self.is_injected305 n.skip_injection = self.skip_injection306 for k, i in self.injections.items():307 n.injections[k] = i.copy()308 # hooks309 n.hook_patches = create_hook_patches_clone(self.hook_patches)310 n.hook_patches_backup = create_hook_patches_clone(self.hook_patches_backup) if self.hook_patches_backup else self.hook_patches_backup311 for group in self.cached_hook_patches:312 n.cached_hook_patches[group] = {}313 for k in self.cached_hook_patches[group]:314 n.cached_hook_patches[group][k] = self.cached_hook_patches[group][k]315 n.hook_backup = self.hook_backup316 n.current_hooks = self.current_hooks.clone() if self.current_hooks else self.current_hooks317 n.forced_hooks = self.forced_hooks.clone() if self.forced_hooks else self.forced_hooks318 n.is_clip = self.is_clip319 n.hook_mode = self.hook_mode320 321 for callback in self.get_all_callbacks(CallbacksMP.ON_CLONE):322 callback(self, n)323 return n324 325 def is_clone(self, other):326 if hasattr(other, 'model') and self.model is other.model:327 return True328 return False329 330 def clone_has_same_weights(self, clone: 'ModelPatcher'):331 if not self.is_clone(clone):332 return False333 334 if self.current_hooks != clone.current_hooks:335 return False336 if self.forced_hooks != clone.forced_hooks:337 return False338 if self.hook_patches.keys() != clone.hook_patches.keys():339 return False340 if self.attachments.keys() != clone.attachments.keys():341 return False342 if self.additional_models.keys() != clone.additional_models.keys():343 return False344 for key in self.callbacks:345 if len(self.callbacks[key]) != len(clone.callbacks[key]):346 return False347 for key in self.wrappers:348 if len(self.wrappers[key]) != len(clone.wrappers[key]):349 return False350 if self.injections.keys() != clone.injections.keys():351 return False352 353 if len(self.patches) == 0 and len(clone.patches) == 0:354 return True355 356 if self.patches_uuid == clone.patches_uuid:357 if len(self.patches) != len(clone.patches):358 logging.warning("WARNING: something went wrong, same patch uuid but different length of patches.")359 else:360 return True361 362 def memory_required(self, input_shape):363 return self.model.memory_required(input_shape=input_shape)364 365 def set_model_sampler_cfg_function(self, sampler_cfg_function, disable_cfg1_optimization=False):366 if len(inspect.signature(sampler_cfg_function).parameters) == 3:367 self.model_options["sampler_cfg_function"] = lambda args: sampler_cfg_function(args["cond"], args["uncond"], args["cond_scale"]) #Old way368 else:369 self.model_options["sampler_cfg_function"] = sampler_cfg_function370 if disable_cfg1_optimization:371 self.model_options["disable_cfg1_optimization"] = True372 373 def set_model_sampler_post_cfg_function(self, post_cfg_function, disable_cfg1_optimization=False):374 self.model_options = set_model_options_post_cfg_function(self.model_options, post_cfg_function, disable_cfg1_optimization)375 376 def set_model_sampler_pre_cfg_function(self, pre_cfg_function, disable_cfg1_optimization=False):377 self.model_options = set_model_options_pre_cfg_function(self.model_options, pre_cfg_function, disable_cfg1_optimization)378 379 def set_model_unet_function_wrapper(self, unet_wrapper_function: UnetWrapperFunction):380 self.model_options["model_function_wrapper"] = unet_wrapper_function381 382 def set_model_denoise_mask_function(self, denoise_mask_function):383 self.model_options["denoise_mask_function"] = denoise_mask_function384 385 def set_model_patch(self, patch, name):386 to = self.model_options["transformer_options"]387 if "patches" not in to:388 to["patches"] = {}389 to["patches"][name] = to["patches"].get(name, []) + [patch]390 391 def set_model_patch_replace(self, patch, name, block_name, number, transformer_index=None):392 self.model_options = set_model_options_patch_replace(self.model_options, patch, name, block_name, number, transformer_index=transformer_index)393 394 def set_model_attn1_patch(self, patch):395 self.set_model_patch(patch, "attn1_patch")396 397 def set_model_attn2_patch(self, patch):398 self.set_model_patch(patch, "attn2_patch")399 400 def set_model_attn1_replace(self, patch, block_name, number, transformer_index=None):401 self.set_model_patch_replace(patch, "attn1", block_name, number, transformer_index)402 403 def set_model_attn2_replace(self, patch, block_name, number, transformer_index=None):404 self.set_model_patch_replace(patch, "attn2", block_name, number, transformer_index)405 406 def set_model_attn1_output_patch(self, patch):407 self.set_model_patch(patch, "attn1_output_patch")408 409 def set_model_attn2_output_patch(self, patch):410 self.set_model_patch(patch, "attn2_output_patch")411 412 def set_model_input_block_patch(self, patch):413 self.set_model_patch(patch, "input_block_patch")414 415 def set_model_input_block_patch_after_skip(self, patch):416 self.set_model_patch(patch, "input_block_patch_after_skip")417 418 def set_model_output_block_patch(self, patch):419 self.set_model_patch(patch, "output_block_patch")420 421 def set_model_emb_patch(self, patch):422 self.set_model_patch(patch, "emb_patch")423 424 def set_model_forward_timestep_embed_patch(self, patch):425 self.set_model_patch(patch, "forward_timestep_embed_patch")426 427 def add_object_patch(self, name, obj):428 self.object_patches[name] = obj429 430 def set_model_compute_dtype(self, dtype):431 self.add_object_patch("manual_cast_dtype", dtype)432 if dtype is not None:433 self.force_cast_weights = True434 self.patches_uuid = uuid.uuid4() #TODO: optimize by preventing a full model reload for this435 436 def add_weight_wrapper(self, name, function):437 self.weight_wrapper_patches[name] = self.weight_wrapper_patches.get(name, []) + [function]438 self.patches_uuid = uuid.uuid4()439 440 def get_model_object(self, name: str) -> torch.nn.Module:441 """Retrieves a nested attribute from an object using dot notation considering442 object patches.443 444 Args:445 name (str): The attribute path using dot notation (e.g. "model.layer.weight")446 447 Returns:448 The value of the requested attribute449 450 Example:451 patcher = ModelPatcher()452 weight = patcher.get_model_object("layer1.conv.weight")453 """454 if name in self.object_patches:455 return self.object_patches[name]456 else:457 if name in self.object_patches_backup:458 return self.object_patches_backup[name]459 else:460 return comfy.utils.get_attr(self.model, name)461 462 def model_patches_to(self, device):463 to = self.model_options["transformer_options"]464 if "patches" in to:465 patches = to["patches"]466 for name in patches:467 patch_list = patches[name]468 for i in range(len(patch_list)):469 if hasattr(patch_list[i], "to"):470 patch_list[i] = patch_list[i].to(device)471 if "patches_replace" in to:472 patches = to["patches_replace"]473 for name in patches:474 patch_list = patches[name]475 for k in patch_list:476 if hasattr(patch_list[k], "to"):477 patch_list[k] = patch_list[k].to(device)478 if "model_function_wrapper" in self.model_options:479 wrap_func = self.model_options["model_function_wrapper"]480 if hasattr(wrap_func, "to"):481 self.model_options["model_function_wrapper"] = wrap_func.to(device)482 483 def model_dtype(self):484 if hasattr(self.model, "get_dtype"):485 return self.model.get_dtype()486 487 def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):488 with self.use_ejected():489 p = set()490 model_sd = self.model.state_dict()491 for k in patches:492 offset = None493 function = None494 if isinstance(k, str):495 key = k496 else:497 offset = k[1]498 key = k[0]499 if len(k) > 2:500 function = k[2]501 502 if key in model_sd:503 p.add(k)504 current_patches = self.patches.get(key, [])505 current_patches.append((strength_patch, patches[k], strength_model, offset, function))506 self.patches[key] = current_patches507 508 self.patches_uuid = uuid.uuid4()509 return list(p)510 511 def get_key_patches(self, filter_prefix=None):512 model_sd = self.model_state_dict()513 p = {}514 for k in model_sd:515 if filter_prefix is not None:516 if not k.startswith(filter_prefix):517 continue518 bk = self.backup.get(k, None)519 hbk = self.hook_backup.get(k, None)520 weight, set_func, convert_func = get_key_weight(self.model, k)521 if bk is not None:522 weight = bk.weight523 if hbk is not None:524 weight = hbk[0]525 if convert_func is None:526 convert_func = lambda a, **kwargs: a527 528 if k in self.patches:529 p[k] = [(weight, convert_func)] + self.patches[k]530 else:531 p[k] = [(weight, convert_func)]532 return p533 534 def model_state_dict(self, filter_prefix=None):535 with self.use_ejected():536 sd = self.model.state_dict()537 keys = list(sd.keys())538 if filter_prefix is not None:539 for k in keys:540 if not k.startswith(filter_prefix):541 sd.pop(k)542 return sd543 544 def patch_weight_to_device(self, key, device_to=None, inplace_update=False):545 if key not in self.patches:546 return547 548 weight, set_func, convert_func = get_key_weight(self.model, key)549 inplace_update = self.weight_inplace_update or inplace_update550 551 if key not in self.backup:552 self.backup[key] = collections.namedtuple('Dimension', ['weight', 'inplace_update'])(weight.to(device=self.offload_device, copy=inplace_update), inplace_update)553 554 if device_to is not None:555 temp_weight = comfy.model_management.cast_to_device(weight, device_to, torch.float32, copy=True)556 else:557 temp_weight = weight.to(torch.float32, copy=True)558 if convert_func is not None:559 temp_weight = convert_func(temp_weight, inplace=True)560 561 out_weight = comfy.lora.calculate_weight(self.patches[key], temp_weight, key)562 if set_func is None:563 out_weight = comfy.float.stochastic_rounding(out_weight, weight.dtype, seed=string_to_seed(key))564 if inplace_update:565 comfy.utils.copy_to_param(self.model, key, out_weight)566 else:567 comfy.utils.set_attr_param(self.model, key, out_weight)568 else:569 set_func(out_weight, inplace_update=inplace_update, seed=string_to_seed(key))570 571 def _load_list(self):572 loading = []573 for n, m in self.model.named_modules():574 params = []575 skip = False576 for name, param in m.named_parameters(recurse=False):577 params.append(name)578 for name, param in m.named_parameters(recurse=True):579 if name not in params:580 skip = True # skip random weights in non leaf modules581 break582 if not skip and (hasattr(m, "comfy_cast_weights") or len(params) > 0):583 loading.append((comfy.model_management.module_size(m), n, m, params))584 return loading585 586 def load(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False, full_load=False):587 with self.use_ejected():588 self.unpatch_hooks()589 mem_counter = 0590 patch_counter = 0591 lowvram_counter = 0592 loading = self._load_list()593 594 load_completely = []595 loading.sort(reverse=True)596 for x in loading:597 n = x[1]598 m = x[2]599 params = x[3]600 module_mem = x[0]601 602 lowvram_weight = False603 604 weight_key = "{}.weight".format(n)605 bias_key = "{}.bias".format(n)606 607 if not full_load and hasattr(m, "comfy_cast_weights"):608 if mem_counter + module_mem >= lowvram_model_memory:609 lowvram_weight = True610 lowvram_counter += 1611 if hasattr(m, "prev_comfy_cast_weights"): #Already lowvramed612 continue613 614 cast_weight = self.force_cast_weights615 if lowvram_weight:616 if hasattr(m, "comfy_cast_weights"):617 m.weight_function = []618 m.bias_function = []619 620 if weight_key in self.patches:621 if force_patch_weights:622 self.patch_weight_to_device(weight_key)623 else:624 m.weight_function = [LowVramPatch(weight_key, self.patches)]625 patch_counter += 1626 if bias_key in self.patches:627 if force_patch_weights:628 self.patch_weight_to_device(bias_key)629 else:630 m.bias_function = [LowVramPatch(bias_key, self.patches)]631 patch_counter += 1632 633 cast_weight = True634 else:635 if hasattr(m, "comfy_cast_weights"):636 wipe_lowvram_weight(m)637 638 if full_load or mem_counter + module_mem < lowvram_model_memory:639 mem_counter += module_mem640 load_completely.append((module_mem, n, m, params))641 642 if cast_weight and hasattr(m, "comfy_cast_weights"):643 m.prev_comfy_cast_weights = m.comfy_cast_weights644 m.comfy_cast_weights = True645 646 if weight_key in self.weight_wrapper_patches:647 m.weight_function.extend(self.weight_wrapper_patches[weight_key])648 649 if bias_key in self.weight_wrapper_patches:650 m.bias_function.extend(self.weight_wrapper_patches[bias_key])651 652 mem_counter += move_weight_functions(m, device_to)653 654 load_completely.sort(reverse=True)655 for x in load_completely:656 n = x[1]657 m = x[2]658 params = x[3]659 if hasattr(m, "comfy_patched_weights"):660 if m.comfy_patched_weights == True:661 continue662 663 for param in params:664 self.patch_weight_to_device("{}.{}".format(n, param), device_to=device_to)665 666 logging.debug("lowvram: loaded module regularly {} {}".format(n, m))667 m.comfy_patched_weights = True668 669 for x in load_completely:670 x[2].to(device_to)671 672 if lowvram_counter > 0:673 logging.info("loaded partially {} {} {}".format(lowvram_model_memory / (1024 * 1024), mem_counter / (1024 * 1024), patch_counter))674 self.model.model_lowvram = True675 else:676 logging.info("loaded completely {} {} {}".format(lowvram_model_memory / (1024 * 1024), mem_counter / (1024 * 1024), full_load))677 self.model.model_lowvram = False678 if full_load:679 self.model.to(device_to)680 mem_counter = self.model_size()681 682 self.model.lowvram_patch_counter += patch_counter683 self.model.device = device_to684 self.model.model_loaded_weight_memory = mem_counter685 self.model.current_weight_patches_uuid = self.patches_uuid686 687 for callback in self.get_all_callbacks(CallbacksMP.ON_LOAD):688 callback(self, device_to, lowvram_model_memory, force_patch_weights, full_load)689 690 self.apply_hooks(self.forced_hooks, force_apply=True)691 692 def patch_model(self, device_to=None, lowvram_model_memory=0, load_weights=True, force_patch_weights=False):693 with self.use_ejected():694 for k in self.object_patches:695 old = comfy.utils.set_attr(self.model, k, self.object_patches[k])696 if k not in self.object_patches_backup:697 self.object_patches_backup[k] = old698 699 if lowvram_model_memory == 0:700 full_load = True701 else:702 full_load = False703 704 if load_weights:705 self.load(device_to, lowvram_model_memory=lowvram_model_memory, force_patch_weights=force_patch_weights, full_load=full_load)706 self.inject_model()707 return self.model708 709 def unpatch_model(self, device_to=None, unpatch_weights=True):710 self.eject_model()711 if unpatch_weights:712 self.unpatch_hooks()713 if self.model.model_lowvram:714 for m in self.model.modules():715 move_weight_functions(m, device_to)716 wipe_lowvram_weight(m)717 718 self.model.model_lowvram = False719 self.model.lowvram_patch_counter = 0720 721 keys = list(self.backup.keys())722 723 for k in keys:724 bk = self.backup[k]725 if bk.inplace_update:726 comfy.utils.copy_to_param(self.model, k, bk.weight)727 else:728 comfy.utils.set_attr_param(self.model, k, bk.weight)729 730 self.model.current_weight_patches_uuid = None731 self.backup.clear()732 733 if device_to is not None:734 self.model.to(device_to)735 self.model.device = device_to736 self.model.model_loaded_weight_memory = 0737 738 for m in self.model.modules():739 if hasattr(m, "comfy_patched_weights"):740 del m.comfy_patched_weights741 742 keys = list(self.object_patches_backup.keys())743 for k in keys:744 comfy.utils.set_attr(self.model, k, self.object_patches_backup[k])745 746 self.object_patches_backup.clear()747 748 def partially_unload(self, device_to, memory_to_free=0):749 with self.use_ejected():750 memory_freed = 0751 patch_counter = 0752 unload_list = self._load_list()753 unload_list.sort()754 for unload in unload_list:755 if memory_to_free < memory_freed:756 break757 module_mem = unload[0]758 n = unload[1]759 m = unload[2]760 params = unload[3]761 762 lowvram_possible = hasattr(m, "comfy_cast_weights")763 if hasattr(m, "comfy_patched_weights") and m.comfy_patched_weights == True:764 move_weight = True765 for param in params:766 key = "{}.{}".format(n, param)767 bk = self.backup.get(key, None)768 if bk is not None:769 if not lowvram_possible:770 move_weight = False771 break772 773 if bk.inplace_update:774 comfy.utils.copy_to_param(self.model, key, bk.weight)775 else:776 comfy.utils.set_attr_param(self.model, key, bk.weight)777 self.backup.pop(key)778 779 weight_key = "{}.weight".format(n)780 bias_key = "{}.bias".format(n)781 if move_weight:782 cast_weight = self.force_cast_weights783 m.to(device_to)784 module_mem += move_weight_functions(m, device_to)785 if lowvram_possible:786 if weight_key in self.patches:787 m.weight_function.append(LowVramPatch(weight_key, self.patches))788 patch_counter += 1789 if bias_key in self.patches:790 m.bias_function.append(LowVramPatch(bias_key, self.patches))791 patch_counter += 1792 cast_weight = True793 794 if cast_weight:795 m.prev_comfy_cast_weights = m.comfy_cast_weights796 m.comfy_cast_weights = True797 m.comfy_patched_weights = False798 memory_freed += module_mem799 logging.debug("freed {}".format(n))800 801 self.model.model_lowvram = True802 self.model.lowvram_patch_counter += patch_counter803 self.model.model_loaded_weight_memory -= memory_freed804 return memory_freed805 806 def partially_load(self, device_to, extra_memory=0, force_patch_weights=False):807 with self.use_ejected(skip_and_inject_on_exit_only=True):808 unpatch_weights = self.model.current_weight_patches_uuid is not None and (self.model.current_weight_patches_uuid != self.patches_uuid or force_patch_weights)809 # TODO: force_patch_weights should not unload + reload full model810 used = self.model.model_loaded_weight_memory811 self.unpatch_model(self.offload_device, unpatch_weights=unpatch_weights)812 if unpatch_weights:813 extra_memory += (used - self.model.model_loaded_weight_memory)814 815 self.patch_model(load_weights=False)816 full_load = False817 if self.model.model_lowvram == False and self.model.model_loaded_weight_memory > 0:818 self.apply_hooks(self.forced_hooks, force_apply=True)819 return 0820 if self.model.model_loaded_weight_memory + extra_memory > self.model_size():821 full_load = True822 current_used = self.model.model_loaded_weight_memory823 try:824 self.load(device_to, lowvram_model_memory=current_used + extra_memory, force_patch_weights=force_patch_weights, full_load=full_load)825 except Exception as e:826 self.detach()827 raise e828 829 return self.model.model_loaded_weight_memory - current_used830 831 def detach(self, unpatch_all=True):832 self.eject_model()833 self.model_patches_to(self.offload_device)834 if unpatch_all:835 self.unpatch_model(self.offload_device, unpatch_weights=unpatch_all)836 for callback in self.get_all_callbacks(CallbacksMP.ON_DETACH):837 callback(self, unpatch_all)838 return self.model839 840 def current_loaded_device(self):841 return self.model.device842 843 def calculate_weight(self, patches, weight, key, intermediate_dtype=torch.float32):844 logging.warning("The ModelPatcher.calculate_weight function is deprecated, please use: comfy.lora.calculate_weight instead")845 return comfy.lora.calculate_weight(patches, weight, key, intermediate_dtype=intermediate_dtype)846 847 def cleanup(self):848 self.clean_hooks()849 if hasattr(self.model, "current_patcher"):850 self.model.current_patcher = None851 for callback in self.get_all_callbacks(CallbacksMP.ON_CLEANUP):852 callback(self)853 854 def add_callback(self, call_type: str, callback: Callable):855 self.add_callback_with_key(call_type, None, callback)856 857 def add_callback_with_key(self, call_type: str, key: str, callback: Callable):858 c = self.callbacks.setdefault(call_type, {}).setdefault(key, [])859 c.append(callback)860 861 def remove_callbacks_with_key(self, call_type: str, key: str):862 c = self.callbacks.get(call_type, {})863 if key in c:864 c.pop(key)865 866 def get_callbacks(self, call_type: str, key: str):867 return self.callbacks.get(call_type, {}).get(key, [])868 869 def get_all_callbacks(self, call_type: str):870 c_list = []871 for c in self.callbacks.get(call_type, {}).values():872 c_list.extend(c)873 return c_list874 875 def add_wrapper(self, wrapper_type: str, wrapper: Callable):876 self.add_wrapper_with_key(wrapper_type, None, wrapper)877 878 def add_wrapper_with_key(self, wrapper_type: str, key: str, wrapper: Callable):879 w = self.wrappers.setdefault(wrapper_type, {}).setdefault(key, [])880 w.append(wrapper)881 882 def remove_wrappers_with_key(self, wrapper_type: str, key: str):883 w = self.wrappers.get(wrapper_type, {})884 if key in w:885 w.pop(key)886 887 def get_wrappers(self, wrapper_type: str, key: str):888 return self.wrappers.get(wrapper_type, {}).get(key, [])889 890 def get_all_wrappers(self, wrapper_type: str):891 w_list = []892 for w in self.wrappers.get(wrapper_type, {}).values():893 w_list.extend(w)894 return w_list895 896 def set_attachments(self, key: str, attachment):897 self.attachments[key] = attachment898 899 def remove_attachments(self, key: str):900 if key in self.attachments:901 self.attachments.pop(key)902 903 def get_attachment(self, key: str):904 return self.attachments.get(key, None)905 906 def set_injections(self, key: str, injections: list[PatcherInjection]):907 self.injections[key] = injections908 909 def remove_injections(self, key: str):910 if key in self.injections:911 self.injections.pop(key)912 913 def get_injections(self, key: str):914 return self.injections.get(key, None)915 916 def set_additional_models(self, key: str, models: list['ModelPatcher']):917 self.additional_models[key] = models918 919 def remove_additional_models(self, key: str):920 if key in self.additional_models:921 self.additional_models.pop(key)922 923 def get_additional_models_with_key(self, key: str):924 return self.additional_models.get(key, [])925 926 def get_additional_models(self):927 all_models = []928 for models in self.additional_models.values():929 all_models.extend(models)930 return all_models931 932 def get_nested_additional_models(self):933 def _evaluate_sub_additional_models(prev_models: list[ModelPatcher], cache_set: set[ModelPatcher]):934 '''Make sure circular references do not cause infinite recursion.'''935 next_models = []936 for model in prev_models:937 candidates = model.get_additional_models()938 for c in candidates:939 if c not in cache_set:940 next_models.append(c)941 cache_set.add(c)942 if len(next_models) == 0:943 return prev_models944 return prev_models + _evaluate_sub_additional_models(next_models, cache_set)945 946 all_models = self.get_additional_models()947 models_set = set(all_models)948 real_all_models = _evaluate_sub_additional_models(prev_models=all_models, cache_set=models_set)949 return real_all_models950 951 def use_ejected(self, skip_and_inject_on_exit_only=False):952 return AutoPatcherEjector(self, skip_and_inject_on_exit_only=skip_and_inject_on_exit_only)953 954 def inject_model(self):955 if self.is_injected or self.skip_injection:956 return957 for injections in self.injections.values():958 for inj in injections:959 inj.inject(self)960 self.is_injected = True961 if self.is_injected:962 for callback in self.get_all_callbacks(CallbacksMP.ON_INJECT_MODEL):963 callback(self)964 965 def eject_model(self):966 if not self.is_injected:967 return968 for injections in self.injections.values():969 for inj in injections:970 inj.eject(self)971 self.is_injected = False972 for callback in self.get_all_callbacks(CallbacksMP.ON_EJECT_MODEL):973 callback(self)974 975 def pre_run(self):976 if hasattr(self.model, "current_patcher"):977 self.model.current_patcher = self978 for callback in self.get_all_callbacks(CallbacksMP.ON_PRE_RUN):979 callback(self)980 981 def prepare_state(self, timestep):982 for callback in self.get_all_callbacks(CallbacksMP.ON_PREPARE_STATE):983 callback(self, timestep)984 985 def restore_hook_patches(self):986 if self.hook_patches_backup is not None:987 self.hook_patches = self.hook_patches_backup988 self.hook_patches_backup = None989 990 def set_hook_mode(self, hook_mode: comfy.hooks.EnumHookMode):991 self.hook_mode = hook_mode992 993 def prepare_hook_patches_current_keyframe(self, t: torch.Tensor, hook_group: comfy.hooks.HookGroup, model_options: dict[str]):994 curr_t = t[0]995 reset_current_hooks = False996 transformer_options = model_options.get("transformer_options", {})997 for hook in hook_group.hooks:998 changed = hook.hook_keyframe.prepare_current_keyframe(curr_t=curr_t, transformer_options=transformer_options)999 # if keyframe changed, remove any cached HookGroups that contain hook with the same hook_ref;1000 # this will cause the weights to be recalculated when sampling1001 if changed:1002 # reset current_hooks if contains hook that changed1003 if self.current_hooks is not None:1004 for current_hook in self.current_hooks.hooks:1005 if current_hook == hook:1006 reset_current_hooks = True1007 break1008 for cached_group in list(self.cached_hook_patches.keys()):1009 if cached_group.contains(hook):1010 self.cached_hook_patches.pop(cached_group)1011 if reset_current_hooks:1012 self.patch_hooks(None)1013 1014 def register_all_hook_patches(self, hooks: comfy.hooks.HookGroup, target_dict: dict[str], model_options: dict=None,1015 registered: comfy.hooks.HookGroup = None):1016 self.restore_hook_patches()1017 if registered is None:1018 registered = comfy.hooks.HookGroup()1019 # handle WeightHooks1020 weight_hooks_to_register: list[comfy.hooks.WeightHook] = []1021 for hook in hooks.get_type(comfy.hooks.EnumHookType.Weight):1022 if hook.hook_ref not in self.hook_patches:1023 weight_hooks_to_register.append(hook)1024 else:1025 registered.add(hook)1026 if len(weight_hooks_to_register) > 0:1027 # clone hook_patches to become backup so that any non-dynamic hooks will return to their original state1028 self.hook_patches_backup = create_hook_patches_clone(self.hook_patches)1029 for hook in weight_hooks_to_register:1030 hook.add_hook_patches(self, model_options, target_dict, registered)1031 for callback in self.get_all_callbacks(CallbacksMP.ON_REGISTER_ALL_HOOK_PATCHES):1032 callback(self, hooks, target_dict, model_options, registered)1033 return registered1034 1035 def add_hook_patches(self, hook: comfy.hooks.WeightHook, patches, strength_patch=1.0, strength_model=1.0):1036 with self.use_ejected():1037 # NOTE: this mirrors behavior of add_patches func1038 current_hook_patches: dict[str,list] = self.hook_patches.get(hook.hook_ref, {})1039 p = set()1040 model_sd = self.model.state_dict()1041 for k in patches:1042 offset = None1043 function = None1044 if isinstance(k, str):1045 key = k1046 else:1047 offset = k[1]1048 key = k[0]1049 if len(k) > 2:1050 function = k[2]1051 1052 if key in model_sd:1053 p.add(k)1054 current_patches: list[tuple] = current_hook_patches.get(key, [])1055 current_patches.append((strength_patch, patches[k], strength_model, offset, function))1056 current_hook_patches[key] = current_patches1057 self.hook_patches[hook.hook_ref] = current_hook_patches1058 # since should care about these patches too to determine if same model, reroll patches_uuid1059 self.patches_uuid = uuid.uuid4()1060 return list(p)1061 1062 def get_combined_hook_patches(self, hooks: comfy.hooks.HookGroup):1063 # combined_patches will contain weights of all relevant hooks, per key1064 combined_patches = {}1065 if hooks is not None:1066 for hook in hooks.hooks:1067 hook_patches: dict = self.hook_patches.get(hook.hook_ref, {})1068 for key in hook_patches.keys():1069 current_patches: list[tuple] = combined_patches.get(key, [])1070 if math.isclose(hook.strength, 1.0):1071 current_patches.extend(hook_patches[key])1072 else:1073 # patches are stored as tuples: (strength_patch, (tuple_with_weights,), strength_model)1074 for patch in hook_patches[key]:1075 new_patch = list(patch)1076 new_patch[0] *= hook.strength1077 current_patches.append(tuple(new_patch))1078 combined_patches[key] = current_patches1079 return combined_patches1080 1081 def apply_hooks(self, hooks: comfy.hooks.HookGroup, transformer_options: dict=None, force_apply=False):1082 # TODO: return transformer_options dict with any additions from hooks1083 if self.current_hooks == hooks and (not force_apply or (not self.is_clip and hooks is None)):1084 return comfy.hooks.create_transformer_options_from_hooks(self, hooks, transformer_options)1085 self.patch_hooks(hooks=hooks)1086 for callback in self.get_all_callbacks(CallbacksMP.ON_APPLY_HOOKS):1087 callback(self, hooks)1088 return comfy.hooks.create_transformer_options_from_hooks(self, hooks, transformer_options)1089 1090 def patch_hooks(self, hooks: comfy.hooks.HookGroup):1091 with self.use_ejected():1092 if hooks is not None:1093 model_sd_keys = list(self.model_state_dict().keys())1094 memory_counter = None1095 if self.hook_mode == comfy.hooks.EnumHookMode.MaxSpeed:1096 # TODO: minimum_counter should have a minimum that conforms to loaded model requirements1097 memory_counter = MemoryCounter(initial=comfy.model_management.get_free_memory(self.load_device),1098 minimum=comfy.model_management.minimum_inference_memory()*2)1099 # if have cached weights for hooks, use it1100 cached_weights = self.cached_hook_patches.get(hooks, None)1101 if cached_weights is not None:1102 model_sd_keys_set = set(model_sd_keys)1103 for key in cached_weights:1104 if key not in model_sd_keys:1105 logging.warning(f"Cached hook could not patch. Key does not exist in model: {key}")1106 continue1107 self.patch_cached_hook_weights(cached_weights=cached_weights, key=key, memory_counter=memory_counter)1108 model_sd_keys_set.remove(key)1109 self.unpatch_hooks(model_sd_keys_set)1110 else:1111 self.unpatch_hooks()1112 relevant_patches = self.get_combined_hook_patches(hooks=hooks)1113 original_weights = None1114 if len(relevant_patches) > 0:1115 original_weights = self.get_key_patches()1116 for key in relevant_patches:1117 if key not in model_sd_keys:1118 logging.warning(f"Cached hook would not patch. Key does not exist in model: {key}")1119 continue1120 self.patch_hook_weight_to_device(hooks=hooks, combined_patches=relevant_patches, key=key, original_weights=original_weights,1121 memory_counter=memory_counter)1122 else:1123 self.unpatch_hooks()1124 self.current_hooks = hooks1125 1126 def patch_cached_hook_weights(self, cached_weights: dict, key: str, memory_counter: MemoryCounter):1127 if key not in self.hook_backup:1128 weight: torch.Tensor = comfy.utils.get_attr(self.model, key)1129 target_device = self.offload_device1130 if self.hook_mode == comfy.hooks.EnumHookMode.MaxSpeed:1131 used = memory_counter.use(weight)1132 if used:1133 target_device = weight.device1134 self.hook_backup[key] = (weight.to(device=target_device, copy=True), weight.device)1135 comfy.utils.copy_to_param(self.model, key, cached_weights[key][0].to(device=cached_weights[key][1]))1136 1137 def clear_cached_hook_weights(self):1138 self.cached_hook_patches.clear()1139 self.patch_hooks(None)1140 1141 def patch_hook_weight_to_device(self, hooks: comfy.hooks.HookGroup, combined_patches: dict, key: str, original_weights: dict, memory_counter: MemoryCounter):1142 if key not in combined_patches:1143 return1144 1145 weight, set_func, convert_func = get_key_weight(self.model, key)1146 weight: torch.Tensor1147 if key not in self.hook_backup:1148 target_device = self.offload_device1149 if self.hook_mode == comfy.hooks.EnumHookMode.MaxSpeed:1150 used = memory_counter.use(weight)1151 if used:1152 target_device = weight.device1153 self.hook_backup[key] = (weight.to(device=target_device, copy=True), weight.device)1154 # TODO: properly handle LowVramPatch, if it ends up an issue1155 temp_weight = comfy.model_management.cast_to_device(weight, weight.device, torch.float32, copy=True)1156 if convert_func is not None:1157 temp_weight = convert_func(temp_weight, inplace=True)1158 1159 out_weight = comfy.lora.calculate_weight(combined_patches[key],1160 temp_weight,1161 key, original_weights=original_weights)1162 del original_weights[key]1163 if set_func is None:1164 out_weight = comfy.float.stochastic_rounding(out_weight, weight.dtype, seed=string_to_seed(key))1165 comfy.utils.copy_to_param(self.model, key, out_weight)1166 else:1167 set_func(out_weight, inplace_update=True, seed=string_to_seed(key))1168 if self.hook_mode == comfy.hooks.EnumHookMode.MaxSpeed:1169 # TODO: disable caching if not enough system RAM to do so1170 target_device = self.offload_device1171 used = memory_counter.use(weight)1172 if used:1173 target_device = weight.device1174 self.cached_hook_patches.setdefault(hooks, {})1175 self.cached_hook_patches[hooks][key] = (out_weight.to(device=target_device, copy=False), weight.device)1176 del temp_weight1177 del out_weight1178 del weight1179 1180 def unpatch_hooks(self, whitelist_keys_set: set[str]=None) -> None:1181 with self.use_ejected():1182 if len(self.hook_backup) == 0:1183 self.current_hooks = None1184 return1185 keys = list(self.hook_backup.keys())1186 if whitelist_keys_set:1187 for k in keys:1188 if k in whitelist_keys_set:1189 comfy.utils.copy_to_param(self.model, k, self.hook_backup[k][0].to(device=self.hook_backup[k][1]))1190 self.hook_backup.pop(k)1191 else:1192 for k in keys:1193 comfy.utils.copy_to_param(self.model, k, self.hook_backup[k][0].to(device=self.hook_backup[k][1]))1194 1195 self.hook_backup.clear()1196 self.current_hooks = None1197 1198 def clean_hooks(self):1199 self.unpatch_hooks()1200 self.clear_cached_hook_weights()