From 6dbb4ba989d344a31a7d74eed983475169fa6b27 Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Tue, 3 Sep 2024 05:28:49 +0200
Subject: [PATCH] Temp fix for disable smart memory

Model management seems to expect `unpatch_model` to put model on specified device after unpatching. Needs different way to avoid patch reload.
---
 nodes.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/nodes.py b/nodes.py
index e129988..061e2b9 100644
--- a/nodes.py
+++ b/nodes.py
@@ -176,7 +176,9 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
                 patches = getattr(p, "patches", [])
                 if len(patches) > 0:
                     p.patches = []
-        return super().unpatch_model(device_to=None, unpatch_weights=unpatch_weights)
+        # TODO: Find another way to not unload after patches
+        device_to = device_to if comfy.model_management.DISABLE_SMART_MEMORY else None
+        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