nodes/Zuellni-ComfyUI-ExLlama-Nodes
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Empty cache on model reload

Browse Source
This commit is contained in:
Zuellni
2023-09-15 22:03:31 +02:00
parent 251bff4510
commit bd653a7245
1 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes.py
+2
View File
@@ -2,6 +2,7 @@ from pathlib import Path
import torch import torch
from comfy.utils import ProgressBar from comfy.utils import ProgressBar
from comfy.model_management import soft_empty_cache
from exllama.alt_generator import ExLlamaAltGenerator from exllama.alt_generator import ExLlamaAltGenerator
from exllama.model import ExLlama, ExLlamaCache, ExLlamaConfig from exllama.model import ExLlama, ExLlamaCache, ExLlamaConfig
from exllama.tokenizer import ExLlamaTokenizer from exllama.tokenizer import ExLlamaTokenizer
@@ -73,6 +74,7 @@ class Loader:
def load(self, model_dir, max_seq_len): def load(self, model_dir, max_seq_len):
model_dir = Path(model_dir).expanduser() model_dir = Path(model_dir).expanduser()
soft_empty_cache()
config = ExLlamaConfig(model_dir / "config.json") config = ExLlamaConfig(model_dir / "config.json")
config.model_path = model_dir.glob("*.safetensors") config.model_path = model_dir.glob("*.safetensors")