From bd653a7245ec57177f56023aebd06caa240e4f74 Mon Sep 17 00:00:00 2001
From: Zuellni <123005779+Zuellni@users.noreply.github.com>
Date: Fri, 15 Sep 2023 22:03:31 +0200
Subject: [PATCH] Empty cache on model reload

---
 nodes.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/nodes.py b/nodes.py
index 6be6c46..544178c 100644
--- a/nodes.py
+++ b/nodes.py
@@ -2,6 +2,7 @@ from pathlib import Path
 
 import torch
 from comfy.utils import ProgressBar
+from comfy.model_management import soft_empty_cache
 from exllama.alt_generator import ExLlamaAltGenerator
 from exllama.model import ExLlama, ExLlamaCache, ExLlamaConfig
 from exllama.tokenizer import ExLlamaTokenizer
@@ -73,6 +74,7 @@ class Loader:
 
     def load(self, model_dir, max_seq_len):
         model_dir = Path(model_dir).expanduser()
+        soft_empty_cache()
 
         config = ExLlamaConfig(model_dir / "config.json")
         config.model_path = model_dir.glob("*.safetensors")