mz_llama_core_nodes.py

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)}