Better patch clear logic

Reduces RAM overhead
This commit is contained in:
City
2024-09-02 01:37:16 +02:00
parent 6618801f1c
commit 43427eff52
+19 -9
View File
@@ -98,9 +98,6 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
if key not in self.patches:
return
weight = comfy.utils.get_attr(self.model, key)
inplace_update = self.weight_inplace_update or inplace_update
if key not in self.backup:
self.backup[key] = collections.namedtuple('Dimension', ['weight', 'inplace_update'])(weight.to(device=self.offload_device, copy=inplace_update), inplace_update)
try:
from comfy.lora import calculate_weight
@@ -110,16 +107,17 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
patches = self.patches[key]
qtype = getattr(weight, "tensor_type", None)
if qtype not in (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16):
if device_to is not None:
out_weight = weight.to(device_to, copy=True)
else:
out_weight = weight.to(weight.device, copy=True) # make sure patches are copied
out_weight = weight.to(device_to)
if self.patch_on_device:
patches = move_patch_to_cuda(patches, self.load_device)
out_weight.patches.append((calculate_weight, patches, key))
else:
inplace_update = self.weight_inplace_update or inplace_update
if key not in self.backup:
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:
@@ -133,6 +131,18 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
else:
comfy.utils.set_attr_param(self.model, key, out_weight)
def unpatch_model(self, device_to=None, unpatch_weights=True):
if unpatch_weights:
for p in self.model.parameters():
qtype = getattr(p, "tensor_type", None)
patches = getattr(p, "patches", [])
if qtype in (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16):
continue
if len(patches) > 0:
p.patches = []
return super().unpatch_model(device_to=device_to, unpatch_weights=unpatch_weights)
def load(self, *args, force_patch_weights=False, **kwargs):
# always call `patch_weight_to_device` even for lowvram
return super().load(*args, force_patch_weights=True, **kwargs)