From 483d173d23309f77d197951ad9f21632955fd13a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=CE=A6=CF=86?= <42910943+Brawlence@users.noreply.github.com>
Date: Tue, 21 Mar 2023 20:19:38 +0300
Subject: [PATCH] Code reuse + indication

Now shows the message in the console when unloading weights. Also reload_model() calls unload_model() first to free the memory so that multiple reloads won't overfill it.
---
 server.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/server.py b/server.py
index 1309c17..4c3497c 100644
--- a/server.py
+++ b/server.py
@@ -64,9 +64,7 @@ def load_model_wrapper(selected_model):
     return selected_model
 
 def reload_model():
-    if not shared.args.cpu:
-        gc.collect()
-        torch.cuda.empty_cache()
+    unload_model()
     shared.model, shared.tokenizer = load_model(shared.model_name)
 
 def unload_model():
@@ -74,6 +72,7 @@ def unload_model():
     if not shared.args.cpu:
         gc.collect()
         torch.cuda.empty_cache()
+        print("Model weights unloaded.")
 
 def load_lora_wrapper(selected_lora):
     shared.lora_name = selected_lora