diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..9313599 --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/node.py b/node.py new file mode 100644 index 0000000..ce1e4e9 --- /dev/null +++ b/node.py @@ -0,0 +1,107 @@ +import folder_paths +import json +import requests +import torch +import re + +# 从文件中读取历史记录 +def load_prompt_history(file_path): + with open(file_path, 'r', encoding='utf-8') as file: + return json.load(file) + +# 调用DeepSeek API +def call_deepseek_api(api_key, user_input, prompt_history): + # 添加用户输入到历史记录 + prompt_history['messages'].append({"role": "user", "content": user_input}) + + # 设置API请求的URL和headers + url = "https://api.deepseek.com/v1/chat/completions" + headers = { + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json" + } + + # 设置请求体 + data = { + "model": "deepseek-chat", + "messages": prompt_history['messages'] + } + + # 发送请求 + response = requests.post(url, headers=headers, json=data) + + # 检查响应状态 + if response.status_code == 200: + return response.json() + else: + raise Exception(f"API request failed with status code {response.status_code}: {response.text}") + +# 提取 JSON 部分 +def extract_json_from_markdown(markdown_text): + # 使用正则表达式提取 ```json 和 ``` 之间的内容 + match = re.search(r'```json\n(.*?)\n```', markdown_text, re.DOTALL) + if match: + return match.group(1).strip() # 返回 JSON 部分 + else: + raise ValueError("No JSON content found in the markdown text.") + +class LLMProcessingNode: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "text": ("STRING", {"multiline": True}), + "api_key": ("STRING", {"default": ""}), + }, + } + + RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING", "STRING", "STRING", "STRING") + RETURN_NAMES = ("IS_NSFW", "角色头部以上服饰特征", "角色动作及表情", "角色上半身服饰特征", "角色下半身服饰特征", "其他", "NSFW") + FUNCTION = "process" + CATEGORY = "LLM Tag Classifier" + + def process(self, text, api_key): + # 从文件中加载历史记录 + prompt_history = load_prompt_history("custom_nodes\Comfyui-TagClassifier\prompt.json") + + # 调用DeepSeek API + try: + response = call_deepseek_api(api_key, text, prompt_history) + llm_output = response['choices'][0]['message']['content'] + + # 提取 JSON 部分 + json_content = extract_json_from_markdown(llm_output) + + # 打印提取的 JSON 内容，用于调试 + print("Extracted JSON Content:", json_content) + + # 解析 JSON + llm_output_json = json.loads(json_content) + + # 打印解析后的 JSON 内容，用于调试 + print("Parsed JSON Content:", llm_output_json) + + # 提取七个内容 + is_nsfw = llm_output_json.get("IS_NSFW", "") + head_features = llm_output_json.get("\u89d2\u8272\u5934\u90e8\u4ee5\u4e0a\u670d\u9970\u7279\u5f81", "") + action_expression = llm_output_json.get("\u89d2\u8272\u52a8\u4f5c\u53ca\u8868\u60c5", "") + upper_body_features = llm_output_json.get("\u89d2\u8272\u4e0a\u534a\u8eab\u670d\u9970\u7279\u5f81", "") + lower_body_features = llm_output_json.get("\u89d2\u8272\u4e0b\u534a\u8eab\u670d\u9970\u7279\u5f81", "") + other = llm_output_json.get("\u5176\u4ed6", "") + nsfw = llm_output_json.get("NSFW", "") + + return (is_nsfw, head_features, action_expression, upper_body_features, lower_body_features, other, nsfw) + except Exception as e: + raise Exception(f"Error processing LLM output: {e}") + +# 注册节点 +NODE_CLASS_MAPPINGS = { + "LLMProcessingNode": LLMProcessingNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "LLMProcessingNode": "LLM Tag Classifier", +} \ No newline at end of file diff --git a/prompt.json b/prompt.json new file mode 100644 index 0000000..f37b530 --- /dev/null +++ b/prompt.json @@ -0,0 +1 @@ +{"messages":[{"role":"system","content":"你是一个专用的danbooru tag分词器，你需要将用户输入的一串描述二次元角色插画的英文分成以下几个部分：\nIS_NSFW:True or False\n角色头部以上服饰特征：\n仅指主要角色（一般为女性）。包含不管角色在什么情况下都会有的外貌特征，包括角色名，使其具有泛用性，用户使用这一部分内容可以生成该角色的其他任意图像，\n\n角色动作及表情：\n正常情况下，要包含动作与面部表情，包括颜文字；当包含多人nsfw内容时，此条目输出为空\n\n角色上半身服饰特征：\n如果有nsfw内容，则仅保留与性无关的相关服饰特征\n\n角色下半身服饰特征：\n如果有nsfw内容，则仅保留与性无关的相关服饰特征\n\n其他：\n不确定是否属于角色本身特征，或者与任何人物都无关的tag\n\nNSFW：\n如果有两人及以上且有nsfw相关内容，则输出描述主角与其性爱对象的之间互动所有tag，示例\n1boy,1girl,head back,faceless male,leg grab,blush,nude,steaming body,lying,missionary,on back,on bed,sex,sheet grab,spread legs,sweat,motion lines,head back,arched back,cum,female ejaculation,overflow,female ejaculation,tongue out,\n如果为单人，如没有nsfw内容，该部分内容归于“角色动作及表情”，此类目输出为空\n如果为单人，如有nsfw内容（比如角色自慰），则仅输出那些nsfw相关词语如nipples,pussy,masturbation之类\n\n确保NSFW类目下的内容直接与另一个角色的角色tag合并时可以生成对应角色的相同体位图片而不干扰角色特征\n\n注意！！！NSFW与角色动作与表情两个类目互相冲突，不得同时有输出！！！\n\n用json的格式输出。不要回复任何多余的内容，不得重复输出，不得遗漏任何输入tag"},{"role":"user","content":"1girl, chen_qianyu_\\(arknights\\), 1boy, penis, hetero, nipples, uncensored, horns, black hair, dragon tail, breasts, finger in another's mouth, clothing aside, tail, dragon girl, dragon horns, full nelson, blush, open mouth, horn grab, looking at penis, thighhighs, panties aside, dark-skinned male, tongue, pussy, white panties, boots, one breast out, legs up, long hair, erection, teeth, imminent penetration, sweat, medium breasts, tongue out, solo focus, folded, thigh boots, underwear, panties, black thighhighs, veins, dark skin, upper teeth only, spread legs, veiny penis, reverse suspended congress, looking down, large penis, shirt, simple background, saliva, brown eyes"},{"role":"assistant","content":"```json\n{\n \"IS_NSFW\":\"True\",\n \"角色头部以上服饰特征\": \"1girl, chen_qianyu_(arknights), horns, black hair, dragon horns, dragon girl, long hair, brown eyes\",\n \"角色动作及表情\": \"\",\n \"角色上半身服饰特征\": \"shirt, one breast out, medium breasts\",\n \"角色下半身服饰特征\": \"thighhighs, black thighhighs, boots, thigh boots, white panties, underwear, panties,dragon tail,tail,\",\n \"其他\": \"simple background,solo focus\",\n \"NSFW\": \"1boy, penis, hetero, nipples, uncensored, breasts, finger in another's mouth, clothing aside, full nelson, blush, open mouth, horn grab, looking at penis, panties aside, dark-skinned male, tongue, pussy, legs up, erection, teeth, imminent penetration, sweat, tongue out, folded, veins, dark skin, upper teeth only, spread legs, veiny penis, reverse suspended congress, looking down, large penis, saliva\"\n}\n```"},{"role":"user","content":"tongue out, perlica_\\(arknights\\), blue eyes, tongue, pantyhose, food, ice cream, gloves, holding, multiple girls, ice cream cone, 2girls, dress, animal ears, chibi, cup, black gloves, long hair, jacket, holding cup, holding food, :q, grey hair, solo focus, sleeveless, open jacket, black pantyhose, black footwear, open clothes, blue skirt, long sleeves, white jacket, fingerless gloves, sleeveless dress, licking lips, grey dress, closed mouth, full body, mug, beer, puffy long sleeves, single glove, shoes, skirt, standing, white background, out of frame, plaid, beer mug, grey jacket, simple background, smile, alcohol, id card, :p, rabbit ears, thigh strap"},{"role":"assistant","content":"```json\n{\n \"IS_NSFW\": \"False\",\n \"角色头部以上服饰特征\": \" multiple girls, 2girls,perlica_(arknights), blue eyes, long hair, grey hair, rabbit ears, animal ears\",\n \"角色动作及表情\": \"smile, closed mouth, licking lips, :p, :q,holding, holding cup, holding food, standing, solo focus, full body\",\n \"角色上半身服饰特征\": \"jacket, white jacket, grey jacket, open jacket, open clothes, sleeveless, long sleeves, puffy long sleeves, fingerless gloves, black gloves, single glove, id card\",\n \"角色下半身服饰特征\": \"dress, grey dress, sleeveless dress, blue skirt, skirt, pantyhose, black pantyhose, thigh strap, shoes, black footwear\",\n \"其他\": \"food, ice cream, ice cream cone, cup, mug, beer, beer mug, alcohol, white background, simple background, out of frame, plaid, chibi\",\n \"NSFW\": \"\"\n}\n```"},{"role":"user","content":"\nsoyeonp19 kimetsu_no_yaiba kanroji_mitsuri kochou_shinobu bad_id bad_twitter_id commentary highres 2girls black_skirt braid breasts butterfly_hair_ornament cleavage couple gradient_hair green_eyes green_hair hair_ornament holding_hands looking_at_another mole mole_under_eye multicolored_hair multiple_girls pink_hair pleated_skirt purple_eyes shirt short_hair simple_background skirt smile twin_braids two-tone_hair white_background white_shirt yuri"},{"role":"assistant","content":"```json\n{\n \"IS_NSFW\": \"False\",\n \"角色头部以上服饰特征\": \"soyeonp19, kimetsu_no_yaiba, kanroji_mitsuri, kochou_shinobu, 2girls, multiple_girls, green_eyes, green_hair, pink_hair, purple_eyes, twin_braids, braid, short_hair, multicolored_hair, two-tone_hair, butterfly_hair_ornament, hair_ornament, mole, mole_under_eye\",\n \"角色动作及表情\": \"smile, looking_at_another, holding_hands, couple, yuri\",\n \"角色上半身服饰特征\": \"white_shirt, shirt, cleavage, breasts\",\n \"角色下半身服饰特征\": \"black_skirt, pleated_skirt, skirt\",\n \"其他\": \"commentary, highres, bad_id, bad_twitter_id, simple_background, white_background, gradient_hair\",\n \"NSFW\": \"\"\n}\n```"},{"role":"user","content":"\nmiruzawa_akechi extreme_hearts hayama_hiyori kodaka_saki commentary_request highres translated 2girls bare_shoulders barefoot black_camisole blue_eyes blue_shorts blush brown_hair camisole closed_eyes fingering hand_under_clothes heart long_hair looking_at_another multiple_girls off_shoulder open_mouth panties pink_camisole shirt short_hair shorts strap_slip toes underwear whispering white_panties white_shirt white_shorts yuri"},{"role":"assistant","content":"```json\n{\n \"IS_NSFW\": \"True\",\n \"角色头部以上服饰特征\": \"miruzawa_akechi, extreme_hearts, hayama_hiyori, kodaka_saki, 2girls, multiple_girls, blue_eyes, brown_hair, long_hair, short_hair\",\n \"角色动作及表情\": \"\",\n \"角色上半身服饰特征\": \"black_camisole, camisole, off_shoulder, pink_camisole, shirt, white_shirt, bare_shoulders, strap_slip\",\n \"角色下半身服饰特征\": \"blue_shorts, shorts, white_shorts, panties, white_panties, underwear, barefoot, toes\",\n \"其他\": \"commentary_request, highres, translated, heart, whispering\",\n \"NSFW\": \"blush, closed_eyes, fingering, hand_under_clothes, looking_at_another, open_mouth, yuri\"\n}\n```"}]} \ No newline at end of file