From 29439f3345baf4afba3c7698b050b3734832ef24 Mon Sep 17 00:00:00 2001
From: wailovet <wailovet@163.com>
Date: Wed, 24 Apr 2024 13:29:40 +0800
Subject: [PATCH] 'update'

---
 README.md    |  3 ++-
 __init__.py  | 45 +++++++++++++++++++++++++++++++++++++++++++++
 mz_llama3.py |  8 +-------
 mz_llava.py  | 37 ++++++++++++++++++++++---------------
 4 files changed, 70 insertions(+), 23 deletions(-)

diff --git a/README.md b/README.md
index dd5c88a..8775d93 100644
--- a/README.md
+++ b/README.md
@@ -11,6 +11,7 @@ Use llama.cpp to assist in generating some nodes related to prompt words, includ
 ## Nodes
 + LLama3CLIPTextEncode 
 + LLavaImageInterrogator 
-
++ BaseLLamaCPPCLIPTextEncode (You can directly pass in the model path)
++ BaseLLavaImageInterrogator (You can directly pass in the model path)
 
 
diff --git a/__init__.py b/__init__.py
index 32e3401..04b4c54 100644
--- a/__init__.py
+++ b/__init__.py
@@ -231,6 +231,51 @@ class MZ_LLavaImageInterrogator:
 
 NODE_CLASS_MAPPINGS["MZ_LLavaImageInterrogator"] = MZ_LLavaImageInterrogator
 NODE_DISPLAY_NAME_MAPPINGS["MZ_LLavaImageInterrogator"] = f"{AUTHOR_NAME} - LLavaImageInterrogator"
+class MZ_BaseLLavaImageInterrogator:
+    @classmethod
+    def INPUT_TYPES(s):  
+        return {"required": {
+            "llama_cpp_model": ("STRING", {"default": ""}),
+            "mmproj_model": ("STRING", {"default": ""}), 
+            "image":  ("IMAGE",),
+            "resolution": ("INT", {"default": 512, "min": 128, "max": 2048}), 
+            "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+            "n_gpu_layers": ("INT", {"default": 40, "min": -1, "max": 0xffffffffffffffff}),
+        },
+        "optional": {
+            "clip": ("CLIP", ),
+            "llama_cpp_options": ("LLamaCPPOptions", ),
+        }}
+    RETURN_TYPES = ("STRING", "CONDITIONING",)
+    FUNCTION = "interrogate"
+    CATEGORY = CATEGORY_NAME
+    def interrogate(self, llama_cpp_model, mmproj_model, image, resolution, seed=0, clip=None, n_gpu_layers=40, llama_cpp_options=None):
+        importlib.reload(mz_llava)
+ 
+        image_pil = Utils.tensor2pil(image)
+
+
+        options = {}
+        if llama_cpp_options is not None:
+            options = llama_cpp_options.value
+
+        response = mz_llava.base_image_interrogator(
+            llama_cpp_model,
+            mmproj_model,
+            n_gpu_layers,
+            image_pil, 
+            resolution,  
+            options,
+        ) 
+        conditionings = None
+        if clip is not None:
+            tokens = clip.tokenize(response)
+            cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True,)
+            conditionings = [[cond, {"pooled_output": pooled}]]
+        return (response, conditionings)
+
+NODE_CLASS_MAPPINGS["MZ_BaseLLavaImageInterrogator"] = MZ_BaseLLavaImageInterrogator
+NODE_DISPLAY_NAME_MAPPINGS["MZ_BaseLLavaImageInterrogator"] = f"{AUTHOR_NAME} - BaseLLavaImageInterrogator"
 
  
 class MZ_LLamaCPPInterrogator:
diff --git a/mz_llama3.py b/mz_llama3.py
index 8b07bd0..9f8b504 100644
--- a/mz_llama3.py
+++ b/mz_llama3.py
@@ -33,13 +33,7 @@ def query_beautify_prompt_text(model_name, n_gpu_layers, text, style_presets, do
 
 
 
-    try: 
-            
-        model_file = None 
-        if options.get("customize_model_file", None) is not None:
-            model_file = options.get("customize_model_file")
-
-
+    try:  
         model_file = get_exist_model(model_name)
         
         if model_file is None:
diff --git a/mz_llava.py b/mz_llava.py
index c6af876..4e88dc7 100644
--- a/mz_llava.py
+++ b/mz_llava.py
@@ -49,21 +49,9 @@ def image_interrogator(model_name, mmproj_name, n_gpu_layers, image, resolution,
     if options is None:
         options = {}
     image = mz_prompt_utils.Utils.resize_max(image, resolution, resolution)
-
-    model_file = None
-    mmproj_file = None
-    if options.get("customize_model_file", None) is not None:
-        model_file = options.get("customize_model_file")
-
-    if options.get("customize_mmproj_file", None) is not None:
-        mmproj_file = options.get("customize_mmproj_file")
-
-
-
-    if model_file is None:
-        model_file = get_exist_model(model_name)
-    if mmproj_file is None:
-        mmproj_file = get_exist_model(mmproj_name) 
+ 
+    model_file = get_exist_model(model_name) 
+    mmproj_file = get_exist_model(mmproj_name) 
     
     if model_file is None or mmproj_file is None:
         if download_source == "modelscope":
@@ -117,3 +105,22 @@ def image_interrogator(model_name, mmproj_name, n_gpu_layers, image, resolution,
     mz_llama_cpp.freed_gpu_memory(model_file=model_file)
     return response
         
+
+def base_image_interrogator(model_file, mmproj_file, n_gpu_layers, image, resolution, options={}):  
+    if options is None:
+        options = {}
+    image = mz_prompt_utils.Utils.resize_max(image, resolution, resolution)
+ 
+  
+    response = mz_llama_cpp.llava_cpp_simple_interrogator(
+        model_file=model_file,
+        mmproj_file=mmproj_file,
+        n_gpu_layers=n_gpu_layers,
+        image=image,
+        options=options,
+    )
+    
+
+    mz_llama_cpp.freed_gpu_memory(model_file=model_file)
+    return response
+