From 4a8dbb5d7a27c15193bc2c41ef5936355f222464 Mon Sep 17 00:00:00 2001
From: wailovet <wailovet@163.com>
Date: Tue, 23 Apr 2024 01:21:00 +0800
Subject: [PATCH] 'update'

---
 mz_llama3.py    |  2 ++
 mz_llama_cpp.py | 19 ++++++++++++++++++-
 mz_llava.py     |  4 +++-
 3 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/mz_llama3.py b/mz_llama3.py
index 74bc143..1f3dab5 100644
--- a/mz_llama3.py
+++ b/mz_llama3.py
@@ -91,6 +91,8 @@ def query_beautify_prompt_text(model_name, n_gpu_layers, text, style_presets, do
             max_tokens=2048,
         ) 
         
+        mz_llama_cpp.freed_gpu_memory(model_file=model_file)
+        
         
         response = json.loads(response_json)
         
diff --git a/mz_llama_cpp.py b/mz_llama_cpp.py
index b5349f3..11bfce7 100644
--- a/mz_llama_cpp.py
+++ b/mz_llama_cpp.py
@@ -68,16 +68,33 @@ def check_llama_cpp_requirements():
 
 
 
+def freed_gpu_memory(model_file):
+    check_llama_cpp_requirements() 
+    import llama_cpp
+    
+    model_and_opt = mz_utils.Utils.cache_get(f"llama_cpp_model_and_opt_{model_file}")
+
+    if model_and_opt is None:
+        return 0
+    
+    model = model_and_opt.get("model")
+
+    llama_cpp.llama_free(model.ctx())
+
+    mz_utils.Utils.cache_set(f"llama_cpp_model_and_opt_{model_file}", None)
+
 def llama_cpp_messages(model_file, n_gpu_layers, chat_handler=None, messages=[], options={}):
     print(f"Find local model file: {model_file}")
+    init_opts = ["n_ctx", "logits_all", "chat_format",]
 
     check_llama_cpp_requirements()
 
     from llama_cpp import Llama
 
+
+
     model_and_opt = mz_utils.Utils.cache_get(f"llama_cpp_model_and_opt_{model_file}")
     
-    init_opts = ["n_ctx"]
     # compared
     is_opts_changed = False
 
diff --git a/mz_llava.py b/mz_llava.py
index 6e0df7c..64071af 100644
--- a/mz_llava.py
+++ b/mz_llava.py
@@ -96,6 +96,8 @@ def image_interrogator(model_name, mmproj_name, n_gpu_layers, image, resolution,
         n_gpu_layers=n_gpu_layers,
         image=image,
     )
- 
+    
+
+    mz_llama_cpp.freed_gpu_memory(model_file=model_file)
     return response