From 43427eff52db19265368d0c17d7cffb0ac11c0fc Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Mon, 2 Sep 2024 01:37:16 +0200 Subject: [PATCH 1/3] Better patch clear logic Reduces RAM overhead --- nodes.py | 28 +++++++++++++++++++--------- 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/nodes.py b/nodes.py index 9b7bede..58934b0 100644 --- a/nodes.py +++ b/nodes.py @@ -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) From c7b902bf9a7f8a020a370249d56f9b339c447c65 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Mon, 2 Sep 2024 02:00:10 +0200 Subject: [PATCH 2/3] Patch load/unload fix May break other things. For some reason T5 decided to run on the CPU with this change, but only like once and I can't for the life of me reproduce it. --- nodes.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/nodes.py b/nodes.py index 58934b0..062aafe 100644 --- a/nodes.py +++ b/nodes.py @@ -140,8 +140,7 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher): continue if len(patches) > 0: p.patches = [] - - return super().unpatch_model(device_to=device_to, unpatch_weights=unpatch_weights) + return super().unpatch_model(device_to=None, unpatch_weights=unpatch_weights) def load(self, *args, force_patch_weights=False, **kwargs): # always call `patch_weight_to_device` even for lowvram From ba8ce44b4c3e0ce8c5a07529983e6fb58b9d1f1e Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Mon, 2 Sep 2024 02:22:32 +0200 Subject: [PATCH 3/3] Rename move_patch function --- nodes.py | 5 ++--- ops.py | 8 ++++---- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/nodes.py b/nodes.py index 062aafe..094d8e9 100644 --- a/nodes.py +++ b/nodes.py @@ -10,7 +10,7 @@ import comfy.model_management import comfy.model_patcher import folder_paths -from .ops import GGMLTensor, GGMLOps, move_patch_to_cuda +from .ops import GGMLTensor, GGMLOps, move_patch_to_device from .dequant import dequantize_tensor # Add a custom keys for files ending in .gguf @@ -108,8 +108,7 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher): qtype = getattr(weight, "tensor_type", None) if qtype not in (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16): out_weight = weight.to(device_to) - if self.patch_on_device: - patches = move_patch_to_cuda(patches, self.load_device) + patches = move_patch_to_device(patches, self.load_device if self.patch_on_device else self.offload_device) out_weight.patches.append((calculate_weight, patches, key)) else: inplace_update = self.weight_inplace_update or inplace_update diff --git a/ops.py b/ops.py index bc17c3a..82a3e8e 100644 --- a/ops.py +++ b/ops.py @@ -91,7 +91,7 @@ class GGMLLayer(torch.nn.Module): patch_list = [] device = tensor.device for function, patches, key in getattr(tensor, "patches", []): - patch_list += move_patch_to_cuda(patches, device) + patch_list += move_patch_to_device(patches, device) # dequantize tensor while patches load weight = dequantize_tensor(tensor, dtype, self.dequant_dtype) @@ -149,12 +149,12 @@ class GGMLOps(comfy.ops.manual_cast): input, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse ).to(dtype=output_dtype) -def move_patch_to_cuda(item, device): +def move_patch_to_device(item, device): if isinstance(item, torch.Tensor): return item.to(device, non_blocking=True) elif isinstance(item, tuple): - return tuple(move_patch_to_cuda(x, device) for x in item) + return tuple(move_patch_to_device(x, device) for x in item) elif isinstance(item, list): - return [move_patch_to_cuda(x, device) for x in item] + return [move_patch_to_device(x, device) for x in item] else: return item