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] 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)