import importlib import json import os from . import mz_prompt_utils from . import mz_llama_cpp from . import mz_prompts def get_schema_base_type(t): return { "type": t, } def get_schema_obj(keys_type={}, required=[]): item = {} for key, value in keys_type.items(): if type(value) == str: value = get_schema_base_type(value) item[key] = value return { "type": "object", "properties": item, "required": required } def get_schema_array(item_type="string"): if type(item_type) == str: item_type = get_schema_base_type(item_type) return { "type": "array", "items": item_type, } high_quality_prompt = "((high quality:1.4), (best quality:1.4), (masterpiece:1.4), (8K resolution), (2k wallpaper))" style_presets_prompt = { "none": "", "high_quality": high_quality_prompt, "photography": f"{high_quality_prompt}, (RAW photo, best quality), (realistic, photo-realistic:1.2), (bokeh, cinematic shot, dynamic composition, incredibly detailed, sharpen, details, intricate detail, professional lighting, film lighting, 35mm, anamorphic, lightroom, cinematography, bokeh, lens flare, film grain, HDR10, 8K)", "illustration": f"{high_quality_prompt}, ((detailed matte painting, intricate detail, splash screen, complementary colors), (detailed),(intricate details),illustration,an extremely delicate and beautiful,ultra-detailed,highres,extremely detailed)", } def get_style_presets(): return [ "none", "high_quality", "photography", "illustration", ] def llama_cpp_node_encode(args_dict): importlib.reload(mz_prompts) importlib.reload(mz_llama_cpp) # importlib.reload(mz_prompt_utils) model_config = args_dict.get("llama_cpp_model", {}) mz_prompt_utils.Utils.print_log(f"model_config: {model_config}") chat_format = model_config.get("chat_format", None) select_model_type = model_config.get("type", "ManualSelect") if select_model_type == "ManualSelect": model_file = model_config.get("model_path", "auto") if model_file == "auto": model_file = mz_prompt_utils.Utils.get_auto_model_fullpath( "Meta-Llama-3-8B-Instruct.Q4_K_M") if "llama-3" in mz_llama_cpp.get_llama_cpp_chat_handlers(): chat_format = "llama-3" elif select_model_type == "DownloaderSelect": model_name = model_config.get("model_name", "") model_file = mz_prompt_utils.Utils.get_auto_model_fullpath( model_name) else: raise Exception("Unknown select_model_type") mz_prompt_utils.Utils.print_log(f"model_file: {model_file}") text = args_dict.get("text", "") style_presets = args_dict.get("style_presets", "") options = args_dict.get("llama_cpp_options", {}) keep_device = args_dict.get("keep_device", False) seed = args_dict.get("seed", -1) options["seed"] = seed options["chat_format"] = chat_format customize_instruct = args_dict.get("customize_instruct", None) mz_prompt_utils.Utils.print_log( f"customize_instruct: {customize_instruct}") try: schema = None if customize_instruct is None: schema = get_schema_obj( keys_type={ "description": get_schema_base_type("string"), "long_prompt": get_schema_base_type("string"), "main_color_word": get_schema_base_type("string"), "camera_angle_word": get_schema_base_type("string"), "style_words": get_schema_array("string"), "subject_words": get_schema_array("string"), "light_words": get_schema_array("string"), "environment_words": get_schema_array("string"), }, required=[ "description", "long_prompt", "main_color_word", "camera_angle_word", "style_words", "subject_words", "light_words", "environment_words", ] ) question = f"IDEA: {style_presets},{text}" if style_presets == "none": question = f"IDEA: {text}" system_prompt = mz_prompts.Beautify_Prompt + mz_prompts.Long_prompt + "\n" else: system_prompt = customize_instruct.get("system", "") question = customize_instruct.get("instruct", "%text%") system_prompt = system_prompt.replace("%text%", text) question = question.replace("%text%", text) mz_prompt_utils.Utils.print_log(f"system_prompt: {system_prompt}") mz_prompt_utils.Utils.print_log(f"question: {question}") if schema is not None: response_text = mz_llama_cpp.llama_cpp_simple_interrogator_to_json( model_file=model_file, system=system_prompt, question=question, schema=schema, options=options, ) try: response_json = json.loads(response_text) except Exception as e: from . import half_json print("json.loads failed, try fix response_text: ", response_text) json_fixer = half_json.JSONFixer() fix_resp = json_fixer.fix(response_text) if fix_resp.success: print("fix success, use fixed response_text: ", fix_resp.line) response_json = json.loads(fix_resp.line) else: raise e mz_prompt_utils.Utils.print_log( f"response_json: {json.dumps(response_json, indent=2)}") 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) else: response = mz_llama_cpp.llama_cpp_simple_interrogator( model_file=model_file, system=system_prompt, question=question, options=options, ) start_str = customize_instruct.get("start_str", "") if start_str != "" and response.find(start_str) != -1: full_response_list = response.split(start_str) # 删除第一个元素 full_response_list.pop(0) response = start_str.join(full_response_list) end_str = customize_instruct.get("end_str", "") if end_str != "" and response.find(end_str) != -1: full_response_list = response.split(end_str) # 删除最后一个元素 full_response_list.pop() response = end_str.join(full_response_list) if keep_device is False: mz_llama_cpp.freed_gpu_memory(model_file=model_file) # 去除换行 while response.find("\n") != -1: response = response.replace("\n", " ") # 句号换成逗号 while response.find(".") != -1: response = response.replace(".", ",") # 去除多余逗号 while response.find(",,") != -1: response = response.replace(",,", ",") while response.find(", ,") != -1: response = response.replace(", ,", ",") response = mz_prompt_utils.Utils.prompt_zh_to_en(response) style_presets_prompt_text = style_presets_prompt.get(style_presets, "") if style_presets_prompt_text != "": response = f"{style_presets_prompt_text}, {response}" except Exception as e: mz_llama_cpp.freed_gpu_memory(model_file=model_file) raise e conditionings = None clip = args_dict.get("clip", None) if clip is not None: conditionings = mz_prompt_utils.Utils.a1111_clip_text_encode( clip, response, ) return {"ui": {"string": [mz_prompt_utils.Utils.to_debug_prompt(response),]}, "result": (response, conditionings)} def image_interrogator_captioner(args_dict): import PIL.Image as Image captioner_config = args_dict.get("captioner_config", {}) 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"] = mz_prompt_utils.Utils.pil2tensor(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_node_encode(onec_args_dict) response = response.get("result", ())[0] 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_node_encode(onec_args_dict) response = response.get("result", ())[0] 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: prompt_fixed_beginning = captioner_config.get( "prompt_fixed_beginning", "") f.write(prompt_fixed_beginning + response) result.append(response) # mz_prompt_webserver.show_toast_success( # f"提示词保存成功(prompt saved successfully): {caption_file}", # 1000, # ) return result def image_interrogator_node_encode(args_dict): importlib.reload(mz_prompts) captioner_config = args_dict.get("captioner_config", None) if captioner_config is not None: image_interrogator_captioner(args_dict) # raise Exception( # "图片批量反推任务已完成 ; Image batch reverse push task completed") return {"ui": {"string": ["图片批量反推任务已完成 ; Image batch reverse push task completed",]}, "result": ("", None)} model_config = args_dict.get("image_interrogator_model", {}) chat_format = model_config.get("chat_format", None) llama_cpp_model = model_config.get("llama_cpp_model", "auto") mmproj_model = model_config.get("mmproj_model", "auto") select_model_type = model_config.get("type", "ManualSelect") if select_model_type == "ManualSelect": llama_cpp_model = model_config.get("model_path", "auto") if llama_cpp_model == "auto": llama_cpp_model = mz_prompt_utils.Utils.get_auto_model_fullpath( "ggml_llava1_5-7b-q4_k_m") else: llama_cpp_model = os.path.join( mz_prompt_utils.Utils.get_gguf_models_path(), llama_cpp_model) if mmproj_model.endswith("auto"): llama_cpp_model_sha256 = mz_prompt_utils.Utils.file_sha256( llama_cpp_model) mmproj_model_name = mz_prompt_utils.Utils.get_model_zoo( tags_filter=llama_cpp_model_sha256) if len(mmproj_model_name) == 0: mmproj_model_name = None else: mmproj_model_name = mmproj_model_name[0].get("model", None) if mmproj_model_name is None: mz_prompt_utils.Utils.print_log( "llama_cpp_model_sha256: ", llama_cpp_model_sha256) raise Exception( "未能自动找到对应的mmproj文件 ; Failed to automatically find the corresponding mmproj file.") else: pass mmproj_model = mz_prompt_utils.Utils.get_auto_model_fullpath( mmproj_model_name) else: # mmproj_model = os.path.join( # mz_prompt_utils.Utils.get_gguf_models_path(), mmproj_model) pass elif select_model_type == "DownloaderSelect": model_name = model_config.get("model_name") llama_cpp_model = mz_prompt_utils.Utils.get_auto_model_fullpath( model_name) mmproj_model = model_config.get("mmproj_model_name", "auto") mmproj_model_name = mmproj_model if mmproj_model == "auto": llama_cpp_model_sha256 = mz_prompt_utils.Utils.file_sha256( llama_cpp_model) mz_prompt_utils.Utils.print_log( "llama_cpp_model_sha256: ", llama_cpp_model_sha256) mmproj_model_name = mz_prompt_utils.Utils.get_model_zoo( tags_filter=llama_cpp_model_sha256) if len(mmproj_model_name) == 0: mmproj_model_name = None else: mmproj_model_name = mmproj_model_name[0].get("model", None) if mmproj_model_name is None: raise Exception( "未能自动找到对应的mmproj文件 ; Failed to automatically find the corresponding mmproj file") mmproj_model = mz_prompt_utils.Utils.get_auto_model_fullpath( mmproj_model_name) else: raise Exception("Unknown select_model_type") image = args_dict.get("image", None) image = mz_prompt_utils.Utils.tensor2pil(image) 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 options["chat_format"] = chat_format image = mz_prompt_utils.Utils.resize_max(image, resolution, resolution) customize_instruct = args_dict.get("customize_instruct", None) if customize_instruct is None: # system_prompt = mz_prompts.GPT4VImageCaptioner_System # question = mz_prompts.GPT4VImageCaptioner_Prompt # system_prompt = mz_prompts.M_ImageCaptioner2_System # question = mz_prompts.M_ImageCaptioner2_Prompt system_prompt = "You are an assistant who perfectly describes images." question = "Describe this image in detail please." else: system_prompt = customize_instruct.get("system", "") question = customize_instruct.get("instruct", "") mz_prompt_utils.Utils.print_log(f"mmproj_model: {mmproj_model}") response = mz_llama_cpp.llava_cpp_simple_interrogator( model_file=llama_cpp_model, mmproj_file=mmproj_model, image=image, options=options, system=system_prompt, question=question, ) response = response.strip() if response is not None and response != "": if args_dict.get("post_processing", False): # 双引号换成空格 response = response.replace("\"", " ") # 中括号换成空格 response = response.replace("[", " ") response = response.replace("]", " ") # 括号换成空格 response = response.replace("(", " ") response = response.replace(")", " ") # 去除多余空格 while response.find(" ") != -1: response = response.replace(" ", " ") # 从第一个为英文字母的地方开始截取 for i in range(len(response)): if response[i].isalpha(): response = response[i:] break response = response.strip() schema = get_schema_obj( keys_type={ "short_describes": get_schema_base_type("string"), "subject_tags": get_schema_array("string"), "action_tags": get_schema_array("string"), "light_tags": get_schema_array("string"), "scene_tags": get_schema_array("string"), "mood_tags": get_schema_array("string"), "style_tags": get_schema_array("string"), "object_tags": get_schema_array("string"), }, required=[ "short_describes", "subject_tags", "action_tags", "lights_tags", "scenes_tags", "moods_tags", "styles_tags", "objects_tags", ] ) response_json_str = mz_llama_cpp.llama_cpp_simple_interrogator_to_json( model_file=llama_cpp_model, system=mz_prompts.ImageCaptionerPostProcessing_System, question=f"Content: {response}", schema=schema, options=options, ) try: response_json = json.loads(response_json_str) except Exception as e: from . import half_json print("json.loads failed, try fix response_json_str: ", response_json_str) json_fixer = half_json.JSONFixer() fix_resp = json_fixer.fix(response_json_str) if fix_resp.success: print("fix success, use fixed response_json_str: ", fix_resp.line) response_json = json.loads(fix_resp.line) else: raise e responses = [] def pure_words(text: str) -> bool: number_of_spaces = text.count(" ") if number_of_spaces > 2: return False for c in text: if not c.isalpha() and c != "-" and c != "_" and c != " ": return False return True 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 != ""] # 去除带有空格和标点符号的字符串 value = [ v for v in value if pure_words(v)] # 空格换成下划线 value = [v.replace(" ", "_") for v in value] # 首字母小写 value = [v.lower() for v in value] if len(value) > 0: responses.append(f"{', '.join(value)}") description = response_json.get("short_describes", "") if description != "": responses.append(f"{description}") # 对response进行去重 response = ", ".join(responses) if keep_device is False: mz_llama_cpp.freed_gpu_memory(model_file=llama_cpp_model) # return response conditionings = None clip = args_dict.get("clip", None) if clip is not None: conditionings = mz_prompt_utils.Utils.a1111_clip_text_encode( clip, response, ) return {"ui": {"string": [mz_prompt_utils.Utils.to_debug_prompt(response),]}, "result": (response, conditionings)}