From 914d753fb64e1cb7645049398efde045cd2de3f0 Mon Sep 17 00:00:00 2001 From: Cyber Dick Lang <286878701@qq.com> Date: Tue, 15 Jul 2025 17:37:47 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20add=20LoraInfo=5FUTK=20node=20-=20?= =?UTF-8?q?=E7=A7=BB=E6=A4=8DLoRA=E4=BF=A1=E6=81=AF=E6=9F=A5=E8=AF=A2?= =?UTF-8?q?=E8=8A=82=E7=82=B9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 1 + __init__.py | 13 +++ nodes/tools/lora_info.py | 185 +++++++++++++++++++++++++++++++++++++++ web/lora_info.js | 47 ++++++++++ 4 files changed, 246 insertions(+) create mode 100644 nodes/tools/lora_info.py create mode 100644 web/lora_info.js diff --git a/README.md b/README.md index 90f22e0..d700b14 100644 --- a/README.md +++ b/README.md @@ -97,6 +97,7 @@ tqdm #### 系统工具 - **PurgeVRAM_UTK**：显存清理，支持选择性清理缓存和模型 +- **LoraInfo_UTK**：LoRA信息查询，获取触发词、示例提示词、基础模型等信息 #### 预设系统 - **LoadKontextPresets_UTK**：Kontext VLM系统预设，包含27种专业图像变换预设 diff --git a/__init__.py b/__init__.py index 74bff86..1d3535f 100644 --- a/__init__.py +++ b/__init__.py @@ -429,6 +429,16 @@ except ImportError as e: THINK_REMOVER_MAPPINGS = {} THINK_REMOVER_DISPLAY_MAPPINGS = {} +try: + from .nodes.tools.lora_info import \ + NODE_CLASS_MAPPINGS as LORA_INFO_MAPPINGS + from .nodes.tools.lora_info import \ + NODE_DISPLAY_NAME_MAPPINGS as LORA_INFO_DISPLAY_MAPPINGS +except ImportError as e: + print(f"导入错误: {e}") + LORA_INFO_MAPPINGS = {} + LORA_INFO_DISPLAY_MAPPINGS = {} + try: from .nodes.tools.kontext_presets import \ NODE_CLASS_MAPPINGS as KONTEXT_PRESETS_MAPPINGS @@ -467,6 +477,7 @@ NODE_CLASS_MAPPINGS.update(TEXTBOX_MAPPINGS) NODE_CLASS_MAPPINGS.update(TEXT_CONCATENATE_MAPPINGS) NODE_CLASS_MAPPINGS.update(MATH_EXPRESSION_MAPPINGS) NODE_CLASS_MAPPINGS.update(THINK_REMOVER_MAPPINGS) +NODE_CLASS_MAPPINGS.update(LORA_INFO_MAPPINGS) NODE_CLASS_MAPPINGS.update(KONTEXT_PRESETS_MAPPINGS) # 合并显示名称映射 @@ -497,6 +508,7 @@ NODE_DISPLAY_NAME_MAPPINGS.update(TEXTBOX_DISPLAY_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(TEXT_CONCATENATE_DISPLAY_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(MATH_EXPRESSION_DISPLAY_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(THINK_REMOVER_DISPLAY_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(LORA_INFO_DISPLAY_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(KONTEXT_PRESETS_DISPLAY_MAPPINGS) NODE_CATEGORIES = { @@ -529,6 +541,7 @@ NODE_CATEGORIES = { "TextConcatenate_UTK", "MathExpression_UTK", "ThinkRemover_UTK", + "LoraInfo_UTK", "LoadKontextPresets_UTK", ] } diff --git a/nodes/tools/lora_info.py b/nodes/tools/lora_info.py new file mode 100644 index 0000000..7d4617b --- /dev/null +++ b/nodes/tools/lora_info.py @@ -0,0 +1,185 @@ +""" +Lora Info Node +~~~~~~~~~~~~~ + +获取LoRA模型信息，包括触发词、示例提示词、基础模型等。 + +:copyright: (c) 2024 by May +:license: MIT, see LICENSE for more details. +""" + +import folder_paths +import hashlib +import requests +import json +import server +from aiohttp import web +import os + + +db_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'lora_info_db.json') + +def load_json_from_file(file_path): + try: + with open(file_path, 'r') as json_file: + data = json.load(json_file) + return data + except FileNotFoundError: + print(f"File not found: {file_path}") + return {} + except json.JSONDecodeError: + print(f"Error decoding JSON in file: {file_path}") + return {} + +def save_dict_to_json(data_dict, file_path): + try: + with open(file_path, 'w') as json_file: + json.dump(data_dict, json_file, indent=4) + print(f"Data saved to {file_path}") + except Exception as e: + print(f"Error saving JSON to file: {e}") + +def get_model_version_info(hash_value): + api_url = f"https://civitai.com/api/v1/model-versions/by-hash/{hash_value}" + response = requests.get(api_url) + + if response.status_code == 200: + return response.json() + else: + return {} + +def calculate_sha256(file_path): + sha256_hash = hashlib.sha256() + with open(file_path, "rb") as f: + for chunk in iter(lambda: f.read(4096), b""): + sha256_hash.update(chunk) + return sha256_hash.hexdigest() + +def get_lora_info(lora_name): + db = load_json_from_file(db_path) + output = None + examplePrompt = None + trainedWords = None + baseModel = None + + loraInfo = db.get(lora_name, {}) + + if isinstance(loraInfo, str): + loraInfo = {} + + output = loraInfo.get('output', None) + examplePrompt = loraInfo.get('examplePrompt', None) + trainedWords = loraInfo.get('trainedWords', None) + baseModel = loraInfo.get('baseModel', None) + + if output is None or baseModel is None: + output = "" + lora_path = folder_paths.get_full_path("loras", lora_name) + LORAsha256 = calculate_sha256(lora_path) + model_info = get_model_version_info(LORAsha256) + if model_info.get("trainedWords", None) is None: + trainedWords = "" + else: + trainedWords = ",".join(model_info.get("trainedWords")) + baseModel = model_info.get("baseModel", "") + images = model_info.get('images') + examplePrompt = None + modelID = model_info.get("modelId") + + if modelID: + output += f"URL: https://civitai.com/models/{modelID}\n" + if trainedWords: + output += "Triggers: " + trainedWords + output += "\n" + + if baseModel: + output += f"Base Model: {baseModel}\n" + if images: + output += "\nExamples:\n" + for image in images: + output += f"\nOutput: {image.get('url')}\n" + meta = image.get("meta") + if meta: + for key, value in meta.items(): + if examplePrompt is None and key == "prompt": + examplePrompt = value + output += f"{key}: {value}\n" + output += '\n' + + db[lora_name] = { + "output": output, + "trainedWords": trainedWords, + "examplePrompt": examplePrompt, + "baseModel": baseModel + } + save_dict_to_json(db, db_path) + + return (output, trainedWords, examplePrompt, baseModel) + + +@server.PromptServer.instance.routes.post('/lora_info_utk') +async def fetch_lora_info(request): + post = await request.post() + lora_name = post.get("lora_name") + (output, triggerWords, examplePrompt, baseModel) = get_lora_info(lora_name) + + return web.json_response({"output": output, "triggerWords": triggerWords, "examplePrompt": examplePrompt, "baseModel": baseModel}) + +class LoraInfo_UTK: + """ + LoRA信息节点 + + 获取LoRA模型的详细信息，包括： + - 触发词 (Trigger Words) + - 示例提示词 (Example Prompt) + - 基础模型 (Base Model) + - CivitAI链接 + - 示例图片 + """ + + @classmethod + def INPUT_TYPES(s): + LORA_LIST = sorted(folder_paths.get_filename_list("loras"), key=str.lower) + return { + "required": { + "lora_name": (LORA_LIST, {"default": LORA_LIST[0] if LORA_LIST else ""}) + }, + } + + RETURN_NAMES = ("lora_name", "trigger_words", "example_prompt", "info_text") + RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING") + FUNCTION = "lora_info" + OUTPUT_NODE = True + CATEGORY = "UniversalToolkit/Tools" + + def lora_info(self, lora_name): + (output, triggerWords, examplePrompt, baseModel) = get_lora_info(lora_name) + + # 构建信息文本 + info_text = f"LoRA: {lora_name}\n" + if baseModel: + info_text += f"Base Model: {baseModel}\n" + if triggerWords: + info_text += f"Trigger Words: {triggerWords}\n" + if examplePrompt: + info_text += f"Example Prompt: {examplePrompt}\n" + if output: + info_text += f"\n详细信息:\n{output}" + + return { + "ui": { + "text": (info_text,), + "model": (baseModel,) + }, + "result": (lora_name, triggerWords or "", examplePrompt or "", info_text) + } + + +# Node mappings +NODE_CLASS_MAPPINGS = { + "LoraInfo_UTK": LoraInfo_UTK, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "LoraInfo_UTK": "Lora Info (UTK)", +} \ No newline at end of file diff --git a/web/lora_info.js b/web/lora_info.js new file mode 100644 index 0000000..6e3632e --- /dev/null +++ b/web/lora_info.js @@ -0,0 +1,47 @@ +import { app } from "../../scripts/app.js"; +import { ComfyWidgets } from "../../scripts/widgets.js"; +import { api } from '../../scripts/api.js'; + +app.registerExtension({ + name: "LoraInfo_UTK", + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (nodeData.name === "LoraInfo_UTK") { + + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + onNodeCreated ? onNodeCreated.apply(this, []) : undefined; + + this.baseModelWidget = ComfyWidgets["STRING"](this, "Base Model", ["STRING", { multiline: false }], app).widget; + this.showValueWidget = ComfyWidgets["STRING"]( + this, + "output", + ["STRING", { multiline: true }], + app, + ).widget; + + const [loraNameWidget, baseModelWidget, outputWidget] = this.widgets; + + loraNameWidget.callback = () => { + const value = loraNameWidget.value; + + const body = new FormData(); + body.append("lora_name", value); + api + .fetchApi("/lora_info_utk", { method: "POST", body }) + .then((response) => response.json()) + .then((resp) => { + baseModelWidget.value = resp.baseModel; + outputWidget.value = resp.output; + }); + }; + } + + const onExecuted = nodeType.prototype.onExecuted; + nodeType.prototype.onExecuted = function (message) { + onExecuted?.apply(this, [message]); + this.showValueWidget.value = message.text[0]; + this.baseModelWidget.value = message.model[0]; + } + } + }, +}); \ No newline at end of file