ComfyUI/comfy/model_patcher.py

685 lines
29 KiB
Python
Raw Normal View History

2024-08-08 07:27:37 +00:00
"""
This file is part of ComfyUI.
Copyright (C) 2024 Comfy
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
"""
2023-08-28 18:49:18 +00:00
import torch
import copy
import inspect
import logging
import uuid
2024-08-08 07:27:37 +00:00
import collections
2023-08-28 18:49:18 +00:00
import comfy.utils
import comfy.model_management
from comfy.types import UnetWrapperFunction
2023-08-28 18:49:18 +00:00
def weight_decompose(dora_scale, weight, lora_diff, alpha, strength):
dora_scale = comfy.model_management.cast_to_device(dora_scale, weight.device, torch.float32)
lora_diff *= alpha
weight_calc = weight + lora_diff.type(weight.dtype)
2024-03-25 22:09:23 +00:00
weight_norm = (
weight_calc.transpose(0, 1)
.reshape(weight_calc.shape[1], -1)
2024-03-25 22:09:23 +00:00
.norm(dim=1, keepdim=True)
.reshape(weight_calc.shape[1], *[1] * (weight_calc.dim() - 1))
2024-03-25 22:09:23 +00:00
.transpose(0, 1)
)
weight_calc *= (dora_scale / weight_norm).type(weight.dtype)
if strength != 1.0:
weight_calc -= weight
weight += strength * (weight_calc)
else:
weight[:] = weight_calc
return weight
2024-03-25 22:09:23 +00:00
2024-04-15 03:34:25 +00:00
def set_model_options_patch_replace(model_options, patch, name, block_name, number, transformer_index=None):
to = model_options["transformer_options"].copy()
if "patches_replace" not in to:
to["patches_replace"] = {}
else:
to["patches_replace"] = to["patches_replace"].copy()
if name not in to["patches_replace"]:
to["patches_replace"][name] = {}
else:
to["patches_replace"][name] = to["patches_replace"][name].copy()
if transformer_index is not None:
block = (block_name, number, transformer_index)
else:
block = (block_name, number)
to["patches_replace"][name][block] = patch
model_options["transformer_options"] = to
return model_options
2024-03-25 22:09:23 +00:00
def set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=False):
model_options["sampler_post_cfg_function"] = model_options.get("sampler_post_cfg_function", []) + [post_cfg_function]
if disable_cfg1_optimization:
model_options["disable_cfg1_optimization"] = True
return model_options
def set_model_options_pre_cfg_function(model_options, pre_cfg_function, disable_cfg1_optimization=False):
model_options["sampler_pre_cfg_function"] = model_options.get("sampler_pre_cfg_function", []) + [pre_cfg_function]
if disable_cfg1_optimization:
model_options["disable_cfg1_optimization"] = True
return model_options
2024-08-08 07:27:37 +00:00
def wipe_lowvram_weight(m):
if hasattr(m, "prev_comfy_cast_weights"):
m.comfy_cast_weights = m.prev_comfy_cast_weights
del m.prev_comfy_cast_weights
m.weight_function = None
m.bias_function = None
class LowVramPatch:
def __init__(self, key, model_patcher):
self.key = key
self.model_patcher = model_patcher
def __call__(self, weight):
return self.model_patcher.calculate_weight(self.model_patcher.patches[self.key], weight, self.key)
2023-08-28 18:49:18 +00:00
class ModelPatcher:
def __init__(self, model, load_device, offload_device, size=0, weight_inplace_update=False):
2023-08-28 18:49:18 +00:00
self.size = size
self.model = model
if not hasattr(self.model, 'device'):
logging.debug("Model doesn't have a device attribute.")
self.model.device = offload_device
elif self.model.device is None:
self.model.device = offload_device
2023-08-28 18:49:18 +00:00
self.patches = {}
self.backup = {}
self.object_patches = {}
self.object_patches_backup = {}
2023-08-28 18:49:18 +00:00
self.model_options = {"transformer_options":{}}
self.model_size()
self.load_device = load_device
self.offload_device = offload_device
self.weight_inplace_update = weight_inplace_update
self.patches_uuid = uuid.uuid4()
2024-08-08 07:27:37 +00:00
if not hasattr(self.model, 'model_loaded_weight_memory'):
self.model.model_loaded_weight_memory = 0
if not hasattr(self.model, 'lowvram_patch_counter'):
self.model.lowvram_patch_counter = 0
if not hasattr(self.model, 'model_lowvram'):
self.model.model_lowvram = False
2023-08-28 18:49:18 +00:00
def model_size(self):
if self.size > 0:
return self.size
self.size = comfy.model_management.module_size(self.model)
return self.size
2023-08-28 18:49:18 +00:00
2024-08-08 07:27:37 +00:00
def loaded_size(self):
return self.model.model_loaded_weight_memory
def lowvram_patch_counter(self):
return self.model.lowvram_patch_counter
2023-08-28 18:49:18 +00:00
def clone(self):
n = ModelPatcher(self.model, self.load_device, self.offload_device, self.size, weight_inplace_update=self.weight_inplace_update)
2023-08-28 18:49:18 +00:00
n.patches = {}
for k in self.patches:
n.patches[k] = self.patches[k][:]
n.patches_uuid = self.patches_uuid
2023-08-28 18:49:18 +00:00
n.object_patches = self.object_patches.copy()
2023-08-28 18:49:18 +00:00
n.model_options = copy.deepcopy(self.model_options)
n.backup = self.backup
n.object_patches_backup = self.object_patches_backup
2023-08-28 18:49:18 +00:00
return n
def is_clone(self, other):
if hasattr(other, 'model') and self.model is other.model:
return True
return False
def clone_has_same_weights(self, clone):
if not self.is_clone(clone):
return False
if len(self.patches) == 0 and len(clone.patches) == 0:
return True
if self.patches_uuid == clone.patches_uuid:
if len(self.patches) != len(clone.patches):
logging.warning("WARNING: something went wrong, same patch uuid but different length of patches.")
else:
return True
def memory_required(self, input_shape):
return self.model.memory_required(input_shape=input_shape)
2023-12-18 22:03:32 +00:00
def set_model_sampler_cfg_function(self, sampler_cfg_function, disable_cfg1_optimization=False):
2023-08-28 18:49:18 +00:00
if len(inspect.signature(sampler_cfg_function).parameters) == 3:
self.model_options["sampler_cfg_function"] = lambda args: sampler_cfg_function(args["cond"], args["uncond"], args["cond_scale"]) #Old way
else:
self.model_options["sampler_cfg_function"] = sampler_cfg_function
2023-12-18 22:03:32 +00:00
if disable_cfg1_optimization:
self.model_options["disable_cfg1_optimization"] = True
2023-08-28 18:49:18 +00:00
2023-12-18 22:03:32 +00:00
def set_model_sampler_post_cfg_function(self, post_cfg_function, disable_cfg1_optimization=False):
self.model_options = set_model_options_post_cfg_function(self.model_options, post_cfg_function, disable_cfg1_optimization)
def set_model_sampler_pre_cfg_function(self, pre_cfg_function, disable_cfg1_optimization=False):
self.model_options = set_model_options_pre_cfg_function(self.model_options, pre_cfg_function, disable_cfg1_optimization)
def set_model_unet_function_wrapper(self, unet_wrapper_function: UnetWrapperFunction):
2023-08-28 18:49:18 +00:00
self.model_options["model_function_wrapper"] = unet_wrapper_function
2024-03-03 20:11:13 +00:00
def set_model_denoise_mask_function(self, denoise_mask_function):
self.model_options["denoise_mask_function"] = denoise_mask_function
2023-08-28 18:49:18 +00:00
def set_model_patch(self, patch, name):
to = self.model_options["transformer_options"]
if "patches" not in to:
to["patches"] = {}
to["patches"][name] = to["patches"].get(name, []) + [patch]
def set_model_patch_replace(self, patch, name, block_name, number, transformer_index=None):
2024-04-15 03:34:25 +00:00
self.model_options = set_model_options_patch_replace(self.model_options, patch, name, block_name, number, transformer_index=transformer_index)
2023-08-28 18:49:18 +00:00
def set_model_attn1_patch(self, patch):
self.set_model_patch(patch, "attn1_patch")
def set_model_attn2_patch(self, patch):
self.set_model_patch(patch, "attn2_patch")
def set_model_attn1_replace(self, patch, block_name, number, transformer_index=None):
self.set_model_patch_replace(patch, "attn1", block_name, number, transformer_index)
2023-08-28 18:49:18 +00:00
def set_model_attn2_replace(self, patch, block_name, number, transformer_index=None):
self.set_model_patch_replace(patch, "attn2", block_name, number, transformer_index)
2023-08-28 18:49:18 +00:00
def set_model_attn1_output_patch(self, patch):
self.set_model_patch(patch, "attn1_output_patch")
def set_model_attn2_output_patch(self, patch):
self.set_model_patch(patch, "attn2_output_patch")
def set_model_input_block_patch(self, patch):
self.set_model_patch(patch, "input_block_patch")
def set_model_input_block_patch_after_skip(self, patch):
self.set_model_patch(patch, "input_block_patch_after_skip")
def set_model_output_block_patch(self, patch):
self.set_model_patch(patch, "output_block_patch")
def add_object_patch(self, name, obj):
self.object_patches[name] = obj
def get_model_object(self, name):
if name in self.object_patches:
return self.object_patches[name]
else:
if name in self.object_patches_backup:
return self.object_patches_backup[name]
else:
return comfy.utils.get_attr(self.model, name)
2023-08-28 18:49:18 +00:00
def model_patches_to(self, device):
to = self.model_options["transformer_options"]
if "patches" in to:
patches = to["patches"]
for name in patches:
patch_list = patches[name]
for i in range(len(patch_list)):
if hasattr(patch_list[i], "to"):
patch_list[i] = patch_list[i].to(device)
if "patches_replace" in to:
patches = to["patches_replace"]
for name in patches:
patch_list = patches[name]
for k in patch_list:
if hasattr(patch_list[k], "to"):
patch_list[k] = patch_list[k].to(device)
if "model_function_wrapper" in self.model_options:
wrap_func = self.model_options["model_function_wrapper"]
if hasattr(wrap_func, "to"):
self.model_options["model_function_wrapper"] = wrap_func.to(device)
2023-08-28 18:49:18 +00:00
def model_dtype(self):
if hasattr(self.model, "get_dtype"):
return self.model.get_dtype()
def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):
p = set()
2024-06-02 23:21:53 +00:00
model_sd = self.model.state_dict()
2023-08-28 18:49:18 +00:00
for k in patches:
offset = None
function = None
if isinstance(k, str):
key = k
else:
offset = k[1]
key = k[0]
if len(k) > 2:
function = k[2]
if key in model_sd:
2023-08-28 18:49:18 +00:00
p.add(k)
current_patches = self.patches.get(key, [])
current_patches.append((strength_patch, patches[k], strength_model, offset, function))
self.patches[key] = current_patches
2023-08-28 18:49:18 +00:00
self.patches_uuid = uuid.uuid4()
2023-08-28 18:49:18 +00:00
return list(p)
def get_key_patches(self, filter_prefix=None):
comfy.model_management.unload_model_clones(self)
2023-08-28 18:49:18 +00:00
model_sd = self.model_state_dict()
p = {}
for k in model_sd:
if filter_prefix is not None:
if not k.startswith(filter_prefix):
continue
if k in self.patches:
p[k] = [model_sd[k]] + self.patches[k]
else:
p[k] = (model_sd[k],)
return p
def model_state_dict(self, filter_prefix=None):
sd = self.model.state_dict()
keys = list(sd.keys())
if filter_prefix is not None:
for k in keys:
if not k.startswith(filter_prefix):
sd.pop(k)
return sd
2024-08-08 07:27:37 +00:00
def patch_weight_to_device(self, key, device_to=None, inplace_update=False):
if key not in self.patches:
return
weight = comfy.utils.get_attr(self.model, key)
2024-08-08 07:27:37 +00:00
inplace_update = self.weight_inplace_update or inplace_update
if key not in self.backup:
2024-08-08 07:27:37 +00:00
self.backup[key] = collections.namedtuple('Dimension', ['weight', 'inplace_update'])(weight.to(device=self.offload_device, copy=inplace_update), inplace_update)
if device_to is not None:
temp_weight = comfy.model_management.cast_to_device(weight, device_to, torch.float32, copy=True)
else:
temp_weight = weight.to(torch.float32, copy=True)
out_weight = self.calculate_weight(self.patches[key], temp_weight, key).to(weight.dtype)
if inplace_update:
comfy.utils.copy_to_param(self.model, key, out_weight)
else:
comfy.utils.set_attr_param(self.model, key, out_weight)
def patch_model(self, device_to=None, patch_weights=True):
for k in self.object_patches:
old = comfy.utils.set_attr(self.model, k, self.object_patches[k])
if k not in self.object_patches_backup:
self.object_patches_backup[k] = old
if patch_weights:
model_sd = self.model_state_dict()
for key in self.patches:
if key not in model_sd:
logging.warning("could not patch. key doesn't exist in model: {}".format(key))
continue
2023-08-28 18:49:18 +00:00
self.patch_weight_to_device(key, device_to)
2023-08-28 18:49:18 +00:00
if device_to is not None:
self.model.to(device_to)
self.model.device = device_to
2024-08-08 07:27:37 +00:00
self.model.model_loaded_weight_memory = self.model_size()
2023-08-28 18:49:18 +00:00
return self.model
2024-08-12 15:38:06 +00:00
def lowvram_load(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False, full_load=False):
mem_counter = 0
patch_counter = 0
2024-08-08 07:27:37 +00:00
lowvram_counter = 0
for n, m in self.model.named_modules():
lowvram_weight = False
2024-08-12 15:38:06 +00:00
if not full_load and hasattr(m, "comfy_cast_weights"):
module_mem = comfy.model_management.module_size(m)
if mem_counter + module_mem >= lowvram_model_memory:
lowvram_weight = True
2024-08-08 07:27:37 +00:00
lowvram_counter += 1
if m.comfy_cast_weights:
continue
weight_key = "{}.weight".format(n)
bias_key = "{}.bias".format(n)
if lowvram_weight:
if weight_key in self.patches:
if force_patch_weights:
self.patch_weight_to_device(weight_key)
else:
m.weight_function = LowVramPatch(weight_key, self)
patch_counter += 1
if bias_key in self.patches:
if force_patch_weights:
self.patch_weight_to_device(bias_key)
else:
m.bias_function = LowVramPatch(bias_key, self)
patch_counter += 1
m.prev_comfy_cast_weights = m.comfy_cast_weights
m.comfy_cast_weights = True
else:
2024-08-08 07:27:37 +00:00
if hasattr(m, "comfy_cast_weights"):
if m.comfy_cast_weights:
wipe_lowvram_weight(m)
2024-08-09 16:16:25 +00:00
if hasattr(m, "weight"):
2024-08-08 07:27:37 +00:00
mem_counter += comfy.model_management.module_size(m)
2024-08-09 16:16:25 +00:00
param = list(m.parameters())
2024-08-09 18:02:38 +00:00
if len(param) > 0:
weight = param[0]
if weight.device == device_to:
continue
2024-08-09 07:58:28 +00:00
2024-08-12 15:38:06 +00:00
weight_to = None
if full_load:#TODO
weight_to = device_to
self.patch_weight_to_device(weight_key, device_to=weight_to) #TODO: speed this up without OOM
self.patch_weight_to_device(bias_key, device_to=weight_to)
m.to(device_to)
2024-06-14 22:18:53 +00:00
logging.debug("lowvram: loaded module regularly {} {}".format(n, m))
2024-08-08 07:27:37 +00:00
if lowvram_counter > 0:
logging.info("loaded partially {} {} {}".format(lowvram_model_memory / (1024 * 1024), mem_counter / (1024 * 1024), patch_counter))
2024-08-08 07:27:37 +00:00
self.model.model_lowvram = True
else:
2024-08-08 21:05:16 +00:00
logging.info("loaded completely {} {}".format(lowvram_model_memory / (1024 * 1024), mem_counter / (1024 * 1024)))
2024-08-08 07:27:37 +00:00
self.model.model_lowvram = False
self.model.lowvram_patch_counter += patch_counter
self.model.device = device_to
2024-08-08 07:27:37 +00:00
self.model.model_loaded_weight_memory = mem_counter
def patch_model_lowvram(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False):
self.patch_model(device_to, patch_weights=False)
self.lowvram_load(device_to, lowvram_model_memory=lowvram_model_memory, force_patch_weights=force_patch_weights)
return self.model
2023-08-28 18:49:18 +00:00
def calculate_weight(self, patches, weight, key):
for p in patches:
strength = p[0]
2023-08-28 18:49:18 +00:00
v = p[1]
strength_model = p[2]
offset = p[3]
function = p[4]
if function is None:
function = lambda a: a
old_weight = None
if offset is not None:
old_weight = weight
weight = weight.narrow(offset[0], offset[1], offset[2])
2023-08-28 18:49:18 +00:00
if strength_model != 1.0:
weight *= strength_model
if isinstance(v, list):
v = (self.calculate_weight(v[1:], v[0].clone(), key), )
if len(v) == 1:
2023-12-09 19:15:09 +00:00
patch_type = "diff"
elif len(v) == 2:
patch_type = v[0]
v = v[1]
if patch_type == "diff":
2023-08-28 18:49:18 +00:00
w1 = v[0]
if strength != 0.0:
2023-08-28 18:49:18 +00:00
if w1.shape != weight.shape:
logging.warning("WARNING SHAPE MISMATCH {} WEIGHT NOT MERGED {} != {}".format(key, w1.shape, weight.shape))
2023-08-28 18:49:18 +00:00
else:
weight += function(strength * comfy.model_management.cast_to_device(w1, weight.device, weight.dtype))
2023-12-09 19:15:09 +00:00
elif patch_type == "lora": #lora/locon
mat1 = comfy.model_management.cast_to_device(v[0], weight.device, torch.float32)
mat2 = comfy.model_management.cast_to_device(v[1], weight.device, torch.float32)
2024-03-25 22:09:23 +00:00
dora_scale = v[4]
2023-08-28 18:49:18 +00:00
if v[2] is not None:
alpha = v[2] / mat2.shape[0]
else:
alpha = 1.0
2023-08-28 18:49:18 +00:00
if v[3] is not None:
#locon mid weights, hopefully the math is fine because I didn't properly test it
mat3 = comfy.model_management.cast_to_device(v[3], weight.device, torch.float32)
2023-08-28 18:49:18 +00:00
final_shape = [mat2.shape[1], mat2.shape[0], mat3.shape[2], mat3.shape[3]]
mat2 = torch.mm(mat2.transpose(0, 1).flatten(start_dim=1), mat3.transpose(0, 1).flatten(start_dim=1)).reshape(final_shape).transpose(0, 1)
try:
lora_diff = torch.mm(mat1.flatten(start_dim=1), mat2.flatten(start_dim=1)).reshape(weight.shape)
2024-03-25 22:09:23 +00:00
if dora_scale is not None:
weight = function(weight_decompose(dora_scale, weight, lora_diff, alpha, strength))
else:
weight += function(((strength * alpha) * lora_diff).type(weight.dtype))
2023-08-28 18:49:18 +00:00
except Exception as e:
logging.error("ERROR {} {} {}".format(patch_type, key, e))
2023-12-09 19:15:09 +00:00
elif patch_type == "lokr":
2023-08-28 18:49:18 +00:00
w1 = v[0]
w2 = v[1]
w1_a = v[3]
w1_b = v[4]
w2_a = v[5]
w2_b = v[6]
t2 = v[7]
2024-03-25 22:09:23 +00:00
dora_scale = v[8]
2023-08-28 18:49:18 +00:00
dim = None
if w1 is None:
dim = w1_b.shape[0]
w1 = torch.mm(comfy.model_management.cast_to_device(w1_a, weight.device, torch.float32),
comfy.model_management.cast_to_device(w1_b, weight.device, torch.float32))
2023-08-28 18:49:18 +00:00
else:
w1 = comfy.model_management.cast_to_device(w1, weight.device, torch.float32)
2023-08-28 18:49:18 +00:00
if w2 is None:
dim = w2_b.shape[0]
if t2 is None:
w2 = torch.mm(comfy.model_management.cast_to_device(w2_a, weight.device, torch.float32),
comfy.model_management.cast_to_device(w2_b, weight.device, torch.float32))
2023-08-28 18:49:18 +00:00
else:
w2 = torch.einsum('i j k l, j r, i p -> p r k l',
comfy.model_management.cast_to_device(t2, weight.device, torch.float32),
comfy.model_management.cast_to_device(w2_b, weight.device, torch.float32),
comfy.model_management.cast_to_device(w2_a, weight.device, torch.float32))
2023-08-28 18:49:18 +00:00
else:
w2 = comfy.model_management.cast_to_device(w2, weight.device, torch.float32)
2023-08-28 18:49:18 +00:00
if len(w2.shape) == 4:
w1 = w1.unsqueeze(2).unsqueeze(2)
if v[2] is not None and dim is not None:
alpha = v[2] / dim
else:
alpha = 1.0
2023-08-28 18:49:18 +00:00
try:
lora_diff = torch.kron(w1, w2).reshape(weight.shape)
2024-03-25 22:09:23 +00:00
if dora_scale is not None:
weight = function(weight_decompose(dora_scale, weight, lora_diff, alpha, strength))
else:
weight += function(((strength * alpha) * lora_diff).type(weight.dtype))
2023-08-28 18:49:18 +00:00
except Exception as e:
logging.error("ERROR {} {} {}".format(patch_type, key, e))
2023-12-09 19:15:09 +00:00
elif patch_type == "loha":
2023-08-28 18:49:18 +00:00
w1a = v[0]
w1b = v[1]
if v[2] is not None:
alpha = v[2] / w1b.shape[0]
else:
alpha = 1.0
2023-08-28 18:49:18 +00:00
w2a = v[3]
w2b = v[4]
2024-03-25 22:09:23 +00:00
dora_scale = v[7]
2023-08-28 18:49:18 +00:00
if v[5] is not None: #cp decomposition
t1 = v[5]
t2 = v[6]
m1 = torch.einsum('i j k l, j r, i p -> p r k l',
comfy.model_management.cast_to_device(t1, weight.device, torch.float32),
comfy.model_management.cast_to_device(w1b, weight.device, torch.float32),
comfy.model_management.cast_to_device(w1a, weight.device, torch.float32))
m2 = torch.einsum('i j k l, j r, i p -> p r k l',
comfy.model_management.cast_to_device(t2, weight.device, torch.float32),
comfy.model_management.cast_to_device(w2b, weight.device, torch.float32),
comfy.model_management.cast_to_device(w2a, weight.device, torch.float32))
2023-08-28 18:49:18 +00:00
else:
m1 = torch.mm(comfy.model_management.cast_to_device(w1a, weight.device, torch.float32),
comfy.model_management.cast_to_device(w1b, weight.device, torch.float32))
m2 = torch.mm(comfy.model_management.cast_to_device(w2a, weight.device, torch.float32),
comfy.model_management.cast_to_device(w2b, weight.device, torch.float32))
2023-08-28 18:49:18 +00:00
try:
lora_diff = (m1 * m2).reshape(weight.shape)
2024-03-25 22:09:23 +00:00
if dora_scale is not None:
weight = function(weight_decompose(dora_scale, weight, lora_diff, alpha, strength))
else:
weight += function(((strength * alpha) * lora_diff).type(weight.dtype))
2023-08-28 18:49:18 +00:00
except Exception as e:
logging.error("ERROR {} {} {}".format(patch_type, key, e))
2023-12-09 23:15:26 +00:00
elif patch_type == "glora":
if v[4] is not None:
alpha = v[4] / v[0].shape[0]
else:
alpha = 1.0
2023-12-09 23:15:26 +00:00
2024-03-25 22:09:23 +00:00
dora_scale = v[5]
2023-12-09 23:15:26 +00:00
a1 = comfy.model_management.cast_to_device(v[0].flatten(start_dim=1), weight.device, torch.float32)
a2 = comfy.model_management.cast_to_device(v[1].flatten(start_dim=1), weight.device, torch.float32)
b1 = comfy.model_management.cast_to_device(v[2].flatten(start_dim=1), weight.device, torch.float32)
b2 = comfy.model_management.cast_to_device(v[3].flatten(start_dim=1), weight.device, torch.float32)
try:
lora_diff = (torch.mm(b2, b1) + torch.mm(torch.mm(weight.flatten(start_dim=1), a2), a1)).reshape(weight.shape)
2024-03-25 22:09:23 +00:00
if dora_scale is not None:
weight = function(weight_decompose(dora_scale, weight, lora_diff, alpha, strength))
else:
weight += function(((strength * alpha) * lora_diff).type(weight.dtype))
except Exception as e:
logging.error("ERROR {} {} {}".format(patch_type, key, e))
2023-12-09 19:15:09 +00:00
else:
logging.warning("patch type not recognized {} {}".format(patch_type, key))
2023-08-28 18:49:18 +00:00
if old_weight is not None:
weight = old_weight
2023-08-28 18:49:18 +00:00
return weight
def unpatch_model(self, device_to=None, unpatch_weights=True):
if unpatch_weights:
2024-08-08 07:27:37 +00:00
if self.model.model_lowvram:
for m in self.model.modules():
2024-08-08 07:27:37 +00:00
wipe_lowvram_weight(m)
2024-08-08 07:27:37 +00:00
self.model.model_lowvram = False
self.model.lowvram_patch_counter = 0
keys = list(self.backup.keys())
2023-08-28 18:49:18 +00:00
2024-08-08 07:27:37 +00:00
for k in keys:
bk = self.backup[k]
if bk.inplace_update:
comfy.utils.copy_to_param(self.model, k, bk.weight)
else:
comfy.utils.set_attr_param(self.model, k, bk.weight)
2023-08-28 18:49:18 +00:00
self.backup.clear()
2023-08-28 18:49:18 +00:00
if device_to is not None:
self.model.to(device_to)
self.model.device = device_to
2024-08-08 07:27:37 +00:00
self.model.model_loaded_weight_memory = 0
keys = list(self.object_patches_backup.keys())
for k in keys:
2024-03-02 16:43:27 +00:00
comfy.utils.set_attr(self.model, k, self.object_patches_backup[k])
self.object_patches_backup.clear()
2024-08-08 07:27:37 +00:00
def partially_unload(self, device_to, memory_to_free=0):
memory_freed = 0
patch_counter = 0
for n, m in list(self.model.named_modules())[::-1]:
if memory_to_free < memory_freed:
break
shift_lowvram = False
if hasattr(m, "comfy_cast_weights"):
module_mem = comfy.model_management.module_size(m)
weight_key = "{}.weight".format(n)
bias_key = "{}.bias".format(n)
if m.weight is not None and m.weight.device != device_to:
for key in [weight_key, bias_key]:
bk = self.backup.get(key, None)
if bk is not None:
if bk.inplace_update:
comfy.utils.copy_to_param(self.model, key, bk.weight)
else:
comfy.utils.set_attr_param(self.model, key, bk.weight)
self.backup.pop(key)
m.to(device_to)
if weight_key in self.patches:
m.weight_function = LowVramPatch(weight_key, self)
patch_counter += 1
if bias_key in self.patches:
m.bias_function = LowVramPatch(bias_key, self)
patch_counter += 1
m.prev_comfy_cast_weights = m.comfy_cast_weights
m.comfy_cast_weights = True
memory_freed += module_mem
logging.debug("freed {}".format(n))
self.model.model_lowvram = True
self.model.lowvram_patch_counter += patch_counter
self.model.model_loaded_weight_memory -= memory_freed
return memory_freed
def partially_load(self, device_to, extra_memory=0):
self.unpatch_model(unpatch_weights=False)
self.patch_model(patch_weights=False)
2024-08-12 15:38:06 +00:00
full_load = False
2024-08-08 07:27:37 +00:00
if self.model.model_lowvram == False:
return 0
if self.model.model_loaded_weight_memory + extra_memory > self.model_size():
2024-08-12 15:38:06 +00:00
full_load = True
2024-08-08 07:27:37 +00:00
current_used = self.model.model_loaded_weight_memory
2024-08-12 15:38:06 +00:00
self.lowvram_load(device_to, lowvram_model_memory=current_used + extra_memory, full_load=full_load)
2024-08-08 07:27:37 +00:00
return self.model.model_loaded_weight_memory - current_used
def current_loaded_device(self):
return self.model.device