Code cleanup.

main
comfyanonymous 6 months ago
parent b29b3b86c5
commit 20ace7c853

@ -644,40 +644,46 @@ class ModelPatcher:
def partially_unload(self, device_to, memory_to_free=0): def partially_unload(self, device_to, memory_to_free=0):
memory_freed = 0 memory_freed = 0
patch_counter = 0 patch_counter = 0
unload_list = []
for n, m in list(self.model.named_modules())[::-1]: for n, m in self.model.named_modules():
if memory_to_free < memory_freed:
break
shift_lowvram = False shift_lowvram = False
if hasattr(m, "comfy_cast_weights"): if hasattr(m, "comfy_cast_weights"):
module_mem = comfy.model_management.module_size(m) module_mem = comfy.model_management.module_size(m)
weight_key = "{}.weight".format(n) unload_list.append((module_mem, n, m))
bias_key = "{}.bias".format(n)
unload_list.sort()
for unload in unload_list:
if m.weight is not None and m.weight.device != device_to: if memory_to_free < memory_freed:
for key in [weight_key, bias_key]: break
bk = self.backup.get(key, None) module_mem = unload[0]
if bk is not None: n = unload[1]
if bk.inplace_update: m = unload[2]
comfy.utils.copy_to_param(self.model, key, bk.weight) weight_key = "{}.weight".format(n)
else: bias_key = "{}.bias".format(n)
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 if m.weight is not None and m.weight.device != device_to:
m.comfy_cast_weights = True for key in [weight_key, bias_key]:
memory_freed += module_mem bk = self.backup.get(key, None)
logging.debug("freed {}".format(n)) 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.model_lowvram = True
self.model.lowvram_patch_counter += patch_counter self.model.lowvram_patch_counter += patch_counter

Loading…
Cancel
Save