328 lines
12 KiB
Python
328 lines
12 KiB
Python
import json
|
|
import os
|
|
from .. import mz_prompt_utils
|
|
from .. import mz_llama_cpp
|
|
from .. import mz_prompts
|
|
from .. import mz_llama_core_nodes
|
|
from .. import mz_prompt_webserver
|
|
|
|
import importlib
|
|
|
|
|
|
LLava_models = [
|
|
"llava-1.6-mistral-7b-gguf/llava-v1.6-mistral-7b.Q5_K_M.gguf",
|
|
"llava-v1.6-vicuna-13b-gguf/llava-v1.6-vicuna-13b.Q5_K_M.gguf",
|
|
"ggml_llava-v1.5-7b/ggml-model-q4_k.gguf",
|
|
"ggml_llava-v1.5-7b/ggml-model-q5_k.gguf",
|
|
"ggml_llava-v1.5-7b/ggml-model-f16.gguf",
|
|
"ggml_bakllava-1/ggml-model-q4_k.gguf",
|
|
"ggml_bakllava-1/ggml-model-q5_k.gguf",
|
|
"ggml_bakllava-1/ggml-model-f16.gguf",
|
|
]
|
|
|
|
LLava_mmproj_models = [
|
|
"llava-1.6-mistral-7b-gguf/mmproj-model-f16.gguf",
|
|
"llava-v1.6-vicuna-13b-gguf/mmproj-model-f16.gguf",
|
|
"ggml_llava-v1.5-7b/mmproj-model-f16.gguf",
|
|
"ggml_bakllava-1/mmproj-model-f16.gguf",
|
|
]
|
|
|
|
|
|
huggingface_models_map = {
|
|
"llava-v1.6-vicuna-13b-gguf": "cjpais",
|
|
"llava-1.6-mistral-7b-gguf": "cjpais",
|
|
"ggml_llava-v1.5-7b": "mys",
|
|
"ggml_llava-v1.5-13b": "mys",
|
|
"ggml_bakllava-1": "mys",
|
|
}
|
|
|
|
|
|
def get_exist_model(model_name):
|
|
modelscope_model_path = mz_prompt_utils.Utils.modelscope_download_model(
|
|
model_type="llava",
|
|
model_name=model_name,
|
|
only_get_path=True,
|
|
)
|
|
|
|
if modelscope_model_path is not None:
|
|
return modelscope_model_path
|
|
|
|
model_name = model_name.split("?")[0]
|
|
model_names = model_name.split("/")
|
|
|
|
author = huggingface_models_map.get(model_names[0], None)
|
|
if author is None:
|
|
return False
|
|
|
|
model_url = f"https://hf-mirror.com/{author}/{model_names[0]}/resolve/main/{model_names[1]}"
|
|
|
|
hf_model_path = mz_prompt_utils.Utils.hf_download_model(
|
|
model_url, only_get_path=True)
|
|
|
|
if hf_model_path is not None:
|
|
return hf_model_path
|
|
|
|
return None
|
|
|
|
|
|
def image_interrogator(args_dict):
|
|
args_dict = args_dict.copy()
|
|
|
|
captioner_config = args_dict.get("captioner_config", None)
|
|
if captioner_config is not None:
|
|
import PIL.Image as Image
|
|
directory = captioner_config.get("directory", None)
|
|
force_update = captioner_config.get("force_update", False)
|
|
caption_suffix = captioner_config.get("caption_suffix", "")
|
|
retry_keyword = captioner_config.get("retry_keyword", "")
|
|
retry_keywords = retry_keyword.split(",")
|
|
|
|
retry_keywords = [k.strip() for k in retry_keywords]
|
|
retry_keywords = [k for k in retry_keywords if k != ""]
|
|
|
|
pre_images = []
|
|
for root, dirs, files in os.walk(directory):
|
|
for file in files:
|
|
if file.endswith(".jpg") or file.endswith(".jpeg") or file.endswith(".png"):
|
|
image_path = os.path.join(root, file)
|
|
base_file_path = os.path.splitext(image_path)[0]
|
|
caption_file = os.path.join(
|
|
root, base_file_path + caption_suffix)
|
|
if os.path.exists(caption_file) and force_update is False:
|
|
continue
|
|
|
|
pre_images.append({
|
|
"image_path": image_path,
|
|
"caption_path": caption_file
|
|
})
|
|
|
|
result = []
|
|
|
|
pb = mz_prompt_utils.Utils.progress_bar(len(pre_images))
|
|
for i in range(len(pre_images)):
|
|
pre_image = pre_images[i]
|
|
image_path = pre_image["image_path"]
|
|
caption_file = pre_image["caption_path"]
|
|
|
|
onec_args_dict = args_dict.copy()
|
|
del onec_args_dict["captioner_config"]
|
|
|
|
pil_image = Image.open(image_path)
|
|
onec_args_dict["image"] = pil_image
|
|
|
|
if i < len(pre_images) - 1:
|
|
onec_args_dict["keep_device"] = True
|
|
|
|
pb.update(
|
|
i,
|
|
len(pre_images),
|
|
pil_image.copy(),
|
|
)
|
|
|
|
response = image_interrogator(onec_args_dict)
|
|
response = response.strip()
|
|
is_retry = response == ""
|
|
for k in retry_keywords:
|
|
if response.find(k) != -1:
|
|
print(f"存在需要重试的关键词 ; Retry keyword found: {k}")
|
|
is_retry = True
|
|
break
|
|
|
|
mz_prompt_utils.Utils.print_log(
|
|
"\n\nonec_args_dict: ", onec_args_dict)
|
|
if is_retry:
|
|
for retry_n in range(5):
|
|
print(f"Retry {retry_n+1}...")
|
|
onec_args_dict["seed"] = onec_args_dict["seed"] + 1
|
|
response = image_interrogator(onec_args_dict)
|
|
response = response.strip()
|
|
is_retry = response == ""
|
|
for k in retry_keywords:
|
|
if response.find(k) != -1:
|
|
print(f"存在需要重试的关键词 ; Retry keyword found: {k}")
|
|
is_retry = True
|
|
break
|
|
|
|
if is_retry is False:
|
|
break
|
|
if is_retry:
|
|
print(f"重试失败,图片被跳过 ; Retry failed")
|
|
response = ""
|
|
|
|
if response != "":
|
|
with open(caption_file, "w") as f:
|
|
f.write(response)
|
|
|
|
result.append(response)
|
|
|
|
# mz_prompt_webserver.show_toast_success(
|
|
# f"提示词保存成功(prompt saved successfully): {caption_file}",
|
|
# 1000,
|
|
# )
|
|
|
|
return result
|
|
|
|
model_name = args_dict.get("llama_cpp_model", "")
|
|
mmproj_name = args_dict.get("mmproj_model", "")
|
|
download_source = args_dict.get("download_source", None)
|
|
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":
|
|
if model_file is None:
|
|
model_file = mz_prompt_utils.Utils.modelscope_download_model(
|
|
model_type="llava",
|
|
model_name=model_name,
|
|
)
|
|
if mmproj_file is None:
|
|
mmproj_file = mz_prompt_utils.Utils.modelscope_download_model(
|
|
model_type="llava",
|
|
model_name=mmproj_name,
|
|
)
|
|
else:
|
|
model_name = model_name.split("?")[0]
|
|
model_names = model_name.split("/")
|
|
|
|
author = huggingface_models_map.get(model_names[0], None)
|
|
if author is None:
|
|
raise Exception(
|
|
f"Model {model_names[0]} is not supported for image_to_text.")
|
|
|
|
if download_source == "hf-mirror.com":
|
|
model_url = f"https://hf-mirror.com/{author}/{model_names[0]}/resolve/main/{model_names[1]}"
|
|
else:
|
|
model_url = f"https://huggingface.co/{author}/{model_names[0]}/resolve/main/{model_names[1]}"
|
|
|
|
if model_file is None:
|
|
model_file = mz_prompt_utils.Utils.hf_download_model(model_url)
|
|
|
|
mmproj_name = mmproj_name.split("?")[0]
|
|
mmproj_names = mmproj_name.split("/")
|
|
if download_source == "hf-mirror.com":
|
|
mmproj_url = f"https://hf-mirror.com/{author}/{mmproj_names[0]}/resolve/main/{mmproj_names[1]}"
|
|
else:
|
|
mmproj_url = f"https://huggingface.co/{author}/{mmproj_names[0]}/resolve/main/{mmproj_names[1]}"
|
|
|
|
if mmproj_file is None:
|
|
mmproj_file = mz_prompt_utils.Utils.hf_download_model(
|
|
mmproj_url)
|
|
|
|
args_dict["llama_cpp_model"] = model_file
|
|
args_dict["mmproj_model"] = mmproj_file
|
|
response = base_image_interrogator(args_dict=args_dict)
|
|
return response
|
|
|
|
|
|
def base_image_interrogator(args_dict):
|
|
model_file = args_dict.get("llama_cpp_model", "")
|
|
mmproj_file = args_dict.get("mmproj_model", "")
|
|
image = args_dict.get("image", None)
|
|
resolution = args_dict.get("resolution", 512)
|
|
keep_device = args_dict.get("keep_device", False)
|
|
seed = args_dict.get("seed", -1)
|
|
options = args_dict.get("llama_cpp_options", {})
|
|
options["seed"] = seed
|
|
|
|
mz_prompt_utils.Utils.print_log(
|
|
"base_image_interrogator options: ", options)
|
|
# input("Press Enter to continue...")
|
|
|
|
image = mz_prompt_utils.Utils.resize_max(image, resolution, resolution)
|
|
|
|
customize_instruct = args_dict.get("customize_instruct", None)
|
|
if customize_instruct is None:
|
|
|
|
system_text = mz_prompts.GPT4VImageCaptioner_System
|
|
question_text = mz_prompts.GPT4VImageCaptioner_Prompt
|
|
|
|
response = mz_llama_cpp.llava_cpp_simple_interrogator(
|
|
model_file=model_file,
|
|
mmproj_file=mmproj_file,
|
|
image=image,
|
|
options=options,
|
|
# system=system_text,
|
|
# question=question_text,
|
|
)
|
|
else:
|
|
|
|
system_prompt = customize_instruct.get("system", "")
|
|
question = customize_instruct.get("instruct", "%text%")
|
|
system_prompt = system_prompt.replace("%text%", "")
|
|
question = question.replace("%text%", "")
|
|
|
|
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)
|
|
|
|
response = full_response
|
|
|
|
sd_format = args_dict.get("sd_format", "v1")
|
|
|
|
if sd_format == "v1" and customize_instruct is None:
|
|
mz_prompt_utils.Utils.print_log(f"response v1: {response}")
|
|
schema = mz_llama_core_nodes.get_schema_obj(
|
|
keys_type={
|
|
"short_describes": mz_llama_core_nodes.get_schema_base_type("string"),
|
|
"subject_tags": mz_llama_core_nodes.get_schema_array("string"),
|
|
"action_tags": mz_llama_core_nodes.get_schema_array("string"),
|
|
"light_tags": mz_llama_core_nodes.get_schema_array("string"),
|
|
"scenes_tags": mz_llama_core_nodes.get_schema_array("string"),
|
|
"other_tags": mz_llama_core_nodes.get_schema_array("string"),
|
|
},
|
|
required=[
|
|
"short_describes",
|
|
"subject_tags",
|
|
"action_tags",
|
|
"light_tags",
|
|
"scenes_tags",
|
|
"other_tags",
|
|
]
|
|
)
|
|
|
|
response = mz_llama_cpp.llama_cpp_simple_interrogator_to_json(
|
|
model_file=model_file,
|
|
system=mz_prompts.Beautify_Prompt,
|
|
question=f"IDEA: {response}",
|
|
schema=schema,
|
|
options=options,
|
|
)
|
|
|
|
response_json = json.loads(response)
|
|
responses = []
|
|
for key, value in response_json.items():
|
|
if type(value) == list:
|
|
# 去除开头.和空格
|
|
value = [v.strip().lstrip(".") for v in value]
|
|
# 去除空字符串
|
|
value = [v for v in value if v != ""]
|
|
if len(value) > 0:
|
|
responses.append(f"({', '.join(value)})")
|
|
|
|
else:
|
|
if value != "":
|
|
responses.append(f"({value})")
|
|
response = ", ".join(responses)
|
|
|
|
if keep_device is False:
|
|
mz_llama_cpp.freed_gpu_memory(model_file=model_file)
|
|
return response
|