commit PromptEmbellish node

2024-03-17 20:30:16 +08:00
parent e982d8a503
commit a6ebe99785
6 changed files with 133 additions and 1 deletions
@@ -37,6 +37,7 @@ git clone https://github.com/chflame163/ComfyUI_LayerStyle.git
 ## Update
 <font size="4">**If the dependency package error after updating, please reinstall the relevant dependency packages. for details, please refer to [here](https://github.com/chflame163/ComfyUI_LayerStyle/issues/5). </font><br />    

+* Commit [PromptEmbellish](#PromptEmbellish) node. it output polished prompt words, and support inputting images as references.
 * Ultra nodes have been fully upgraded to V2 version, with the addition of VITMatte edge processing method, which is suitable for handling semi transparent areas. Include [MaskEdgeUltraDetailV2](#MaskEdgeUltraDetailV2), [SegmentAnythingUltraV2](#SegmentAnythingUltraV2), [RmBgUltraV2](#RmBgUltraV2) and [PersonMaskUltraV2](#PersonMaskUltraV2) nodes.
 * Commit [Color of Shadow & Highlight](#Highlight) node, it can adjust the color of the dark and bright parts separately. Commit [Shadow & Highlight Mask](#Shadow) node, it can output mask for dark and bright areas.
 * Commit [CropByMaskV2](#CropByMaskV2) node, On the basis of the original node, it supports ```crop_box``` input, making it convenient to cut layers of the same size.
@@ -518,6 +519,19 @@ Node options:
 * replace_with_word: That word will replace the exclude_word.


+### <a id="table1">PromptEmbellish</a>
+Enter simple prompt words, output polished prompt words, and support inputting images as references. This node currently uses Google Gemini API as the backend service. Please ensure that the network environment can use Gemini normally.
+Please apply for your API key on [Google AI Studio](https://makersuite.google.com/app/apikey),  And fill it in ```api_key.ini```, This file is located in the root directory of the plugin. Open it using text editing software, fill in your API key after ```google_api_key=``` and save it.
+![image](image/prompt_embellish_example.png)
+
+Node options:   
+![image](image/prompt_embellish_node.png)
+
+* image: Optional, input image as a reference for prompt words.
+* aip: The Api used. At present, there is only one option, "google-gemini".
+* token_limit: The maximum token limit for generating prompt words.
+* discribe: Enter a simple description here. supports Chinese text input.
+
 ### <a id="table1">ImageShift</a>
 Shift the image. this node supports the output of displacement seam masks, making it convenient to create continuous textures.
 ![image](image/image_shift_example.png)
@@ -38,6 +38,7 @@ git clone https://github.com/chflame163/ComfyUI_LayerStyle.git
 ## 更新说明
 <font size="4">**如果本插件更新后出现依赖包错误，请重新安装相关依赖包。详情见[这里](https://github.com/chflame163/ComfyUI_LayerStyle/issues/5)。  </font><br /> 

+* 添加 [PromptEmbellish](#PromptEmbellish) 节点, 对简单的提示词润色，支持图片输入参考，支持中文输入。
 * Ultra 节点全面升级到V2版本，增加了VITMatte边缘处理方法，此方法适合处理半透明区域。包括 [MaskEdgeUltraDetailV2](#MaskEdgeUltraDetailV2), [SegmentAnythingUltraV2](#SegmentAnythingUltraV2), [RmBgUltraV2](#RmBgUltraV2) 以及 [PersonMaskUltraV2](#PersonMaskUltraV2) 节点。
 * 添加 [Color of Shadow & Highlight](#Highlight) 节点，可对暗部和亮部分别进行色彩调整。添加 [Shadow & Highlight Mask](#Shadow) 节点, 可输出暗部和亮部的遮罩。
 * 添加 [CropByMaskV2](#CropByMaskV2) 节点，在原节点基础上支持```crop_box```输入，方便裁切相同尺寸的图层。
@@ -514,6 +515,19 @@ ImageScaleByAspectRatio的V2升级版
 * exclude_word: 需要排除的关键词。
 * replace_with_word: 替换exclude_word的关键词。

+### <a id="table1">PromptEmbellish</a>
+输入简单的提示词，输出经过润色的提示词，支持输入图片作为参考。这个节点目前使用Google Gemini API作为后端服务，请确保网络环境可以正常使用Gemini。
+请在[Google AI Studio](https://makersuite.google.com/app/apikey)申请你的API key, 并将其填到```api_key.ini```, 这个文件位于插件根目录。用文本编辑软件打开，在```google_api_key=```后面填入你的API key并保存。
+![image](image/prompt_embellish_example.png)
+
+节点选项说明:   
+![image](image/prompt_embellish_node.png)
+
+* image: 可选项，输入图像作为提示词参考。
+* aip: 使用的Api。目前只有"google-gemini"一个选项。
+* token_limit: 生成提示词的最大token限制。
+* discribe: 在这里输入简单的描述。支持中文。
+

 ### <a id="table1">ImageShift</a>
 使图片产生位移。此节点支持位移接缝遮罩的输出，方便制作连续贴图。
@@ -1387,7 +1387,12 @@ def check_image_file(file_name:str, interval:int) -> object:
            break
        time.sleep(interval / 1000)

-
+# 判断字符串是否包含中文
+def is_contain_chinese(check_str:str) -> bool:
+    for ch in check_str:
+        if u'\u4e00' <= ch <= u'\u9fff':
+            return True
+    return False

 '''CLASS'''

@@ -0,0 +1,99 @@
+from .imagefunc import *
+
+NODE_NAME = 'PromptEmbellish'
+
+class PromptEmbellish:
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(self):
+        api_list = ['google-gemini']
+        return {
+            "required": {
+                "api": (api_list,),
+                "token_limit": ("INT", {"default": 40, "min": 2, "max": 1024, "step": 1}),
+                "describe": ("STRING", {"default": "", "multiline": True}),
+            },
+            "optional": {
+                "image": ("IMAGE",),
+            }
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("text",)
+    FUNCTION = 'prompt_embellish'
+    CATEGORY = '😺dzNodes/LayerUtility'
+    OUTPUT_NODE = True
+
+    def prompt_embellish(self, api, token_limit, describe, image=None):
+        if describe == "" and image is None:
+            return ("",)
+        import google.generativeai as genai
+        ret_text = ""
+        first_step_prompt = (f"You are creating a prompt for Stable Diffusion to generate an image. "
+                             f"First step:Using '{describe}' as the basic content, "
+                             f"polish and embellish it to describe into text, keep it on {token_limit} tokens."
+                             f"Second step: Generate a Stable Diffusion text prompt for based on first step in at least {token_limit} words."
+                             f"Only respond with the prompt itself, but embellish it."
+                             )
+
+        genai.configure(api_key=get_api_key('google_api_key'), transport='rest')
+        if describe != "":
+            model = genai.GenerativeModel('gemini-pro')
+            log(f"{NODE_NAME}: Request to gemini-pro...")
+            response = model.generate_content(first_step_prompt)
+            print(response)
+            ret_text = response.text
+            ret_text = ret_text[ret_text.rfind(':') + 1:]
+            ret_text = ret_text[ret_text.rfind('\n') + 1:]
+            # log(f"{NODE_NAME}: Text2Image Prompt is:\n\033[1;36m{ret_text}\033[m")
+            if is_contain_chinese(describe):
+                translate_prompt = (f"Please translate the text in parentheses into English:({describe})"
+                                    )
+                response = model.generate_content(translate_prompt)
+                print(response)
+                ret_discribe = response.text
+            else:
+                ret_discribe = describe
+
+        if image is not None:
+            if describe != "":
+                second_step_prompt = (f"You are creating a prompt for Stable Diffusion to generate an image. "
+                                      f"First step:Modify and polish the content in parentheses to match this photo,"
+                                      f"but must keep '{describe}': ({ret_text}) "
+                                      f"Second step: Find objects that is similar in parentheses from the content of the first step"
+                                      f" and replace it with the content in parentheses: ({describe})"
+                                      f"Third step: Generate a Stable Diffusion text prompt for based on second step in at least {token_limit} words."                                 
+                                      f"Only respond with the prompt itself, but embellish it."
+                                      )
+            else:
+                second_step_prompt = (f"You are creating a prompt for Stable Diffusion to generate an image. "
+                                      f"First step: describe this image, "
+                                      f"polish and embellish it into text, discrete it in {token_limit} tokens."
+                                      f"Second step: Generate a Stable Diffusion text prompt for based on first step in at least {token_limit} words."
+                                      f"Only respond with the prompt itself, but embellish it."
+                                      )
+            _image = tensor2pil(image).convert('RGB')
+            model = genai.GenerativeModel('gemini-pro-vision')
+            log(f"{NODE_NAME}: Request to gemini-pro-vision...")
+            response = model.generate_content([second_step_prompt, _image])
+            print(response)
+            ret_text = response.text
+            ret_text = ret_text[ret_text.rfind(':') + 1:]
+            ret_text = ret_text.replace('(','').replace(')','')
+            if describe != "":
+                ret_text = f"((({ret_discribe}))), {ret_text}"
+            # log(f"{NODE_NAME}: Text2Image by ImageRefrence Prompt is:\n\033[1;36m{ret_text}\033[m")
+        log(f"{NODE_NAME}: Prompt is:\n\033[1;36m{ret_text}\033[m")
+        log(f"{NODE_NAME} Processed.", message_type='finish')
+        return (ret_text,)
+
+NODE_CLASS_MAPPINGS = {
+    "LayerUtility: PromptEmbellish": PromptEmbellish
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "LayerUtility: PromptEmbellish": "LayerUtility: PromptEmbellish"
+}