This commit is contained in:
wailovet
2024-04-24 13:29:40 +08:00
parent 87fdc835e9
commit 29439f3345
4 changed files with 70 additions and 23 deletions
+2 -1
View File
@@ -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)
+45
View File
@@ -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:
+1 -7
View File
@@ -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:
+22 -15
View File
@@ -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