diff --git a/__init__.py b/__init__.py index 5c37528..11ec40d 100644 --- a/__init__.py +++ b/__init__.py @@ -314,6 +314,7 @@ class MZ_LLavaImageInterrogator: "optional": { "clip": ("CLIP", ), "llama_cpp_options": ("LLamaCPPOptions", ), + "customize_instruct": ("CustomizeInstruct", ), }} RETURN_TYPES = ("STRING", "CONDITIONING",) RETURN_NAMES = ("text", "conditioning",) @@ -360,6 +361,7 @@ class MZ_BaseLLavaImageInterrogator: "optional": { "clip": ("CLIP", ), "llama_cpp_options": ("LLamaCPPOptions", ), + "customize_instruct": ("CustomizeInstruct", ), }} RETURN_TYPES = ("STRING", "CONDITIONING",) RETURN_NAMES = ("text", "conditioning",) diff --git a/mz_llava.py b/mz_llava.py index 06aa031..4642db1 100644 --- a/mz_llava.py +++ b/mz_llava.py @@ -135,19 +135,50 @@ def base_image_interrogator(args_dict): + + customize_instruct = args_dict.get("customize_instruct", None) + if customize_instruct is None: + response = mz_llama_cpp.llava_cpp_simple_interrogator( + model_file=model_file, + mmproj_file=mmproj_file, + image=image, + options=options, + ) + else: + system_prompt = customize_instruct.get("system", "") + question = customize_instruct.get("instruct", "{text}") + system_prompt = system_prompt.replace("{text}", "") + question = question.replace("{text}", "") - response = mz_llama_cpp.llava_cpp_simple_interrogator( - model_file=model_file, - mmproj_file=mmproj_file, - image=image, - options=options, - ) + full_response = mz_llama_cpp.llava_cpp_simple_interrogator( + model_file=model_file, + mmproj_file=mmproj_file, + image=image, + system=system_prompt, + question=question, + options=options, + ) + + + start_str = customize_instruct.get("start_str", "") + if start_str != "" and full_response.find(start_str) != -1: + full_response_list = full_response.split(start_str) + # 删除第一个元素 + full_response_list.pop(0) + full_response = start_str.join(full_response_list) + end_str = customize_instruct.get("end_str", "") + if end_str != "" and full_response.find(end_str) != -1: + full_response_list = full_response.split(end_str) + # 删除最后一个元素 + full_response_list.pop() + full_response = end_str.join(full_response_list) + sd_format = args_dict.get("sd_format", "v1") - if sd_format == "v1": + if sd_format == "v1" and customize_instruct is None: mz_prompt_utils.Utils.print_log(f"response: {response}")