diff --git a/js/index.js b/js/index.js index eece2d2..dec06c3 100644 --- a/js/index.js +++ b/js/index.js @@ -10,6 +10,7 @@ app.registerExtension({ 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", @@ -22,6 +23,7 @@ app.registerExtension({ nodeType.prototype.onExecuted = function (message) { onExecuted?.apply(this, [message]); this.showValueWidget.value = message.text[0]; + this.baseModelWidget.value = message.model[0]; } } }, diff --git a/lora_info.py b/lora_info.py index 18d41ad..3c2a332 100644 --- a/lora_info.py +++ b/lora_info.py @@ -30,7 +30,7 @@ def get_model_version_info(hash_value): if response.status_code == 200: return response.json() else: - return None + return {} def calculate_sha256(file_path): sha256_hash = hashlib.sha256() @@ -39,70 +39,94 @@ def calculate_sha256(file_path): sha256_hash.update(chunk) return sha256_hash.hexdigest() +def get_lora_info(lora_name): + db = load_json_from_file('./custom_nodes/lora-info/db.json') + 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) + trainedWords = ",".join(model_info.get("trainedWords")) + baseModel = model_info.get("baseModel") + images = model_info.get('images') + examplePrompt = None + + 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, './custom_nodes/lora-info/db.json') + + return (output, trainedWords, examplePrompt, baseModel) + +LORA_LIST = sorted(folder_paths.get_filename_list("loras"), key=str.lower) + class LoraInfo: def __init__(self): pass @classmethod def INPUT_TYPES(s): - LORA_LIST = sorted(folder_paths.get_filename_list("loras"), key=str.lower) - return { "required": { "lora_name": (LORA_LIST, ) }, } - RETURN_TYPES = () + RETURN_NAMES = ("trigger_words", "example_prompt") + RETURN_TYPES = ("STRING", "STRING") FUNCTION = "lora_info" OUTPUT_NODE = True CATEGORY = "jitcoder" def lora_info(self, lora_name): - lora_tags = load_json_from_file('./custom_nodes/lora-info/db.json') - output = lora_tags.get(lora_name, None) if lora_tags is not None else None - - if output 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: - trainedWords = model_info.get("trainedWords") - baseModel = model_info.get("baseModel") - images = model_info.get('images') - - if trainedWords: - output += "Triggers: " + ",".join(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(): - output += f"{key}: {value}\n" - output += '\n' - lora_tags[lora_name] = output - save_dict_to_json(lora_tags, './custom_nodes/lora-info/db.json') - else: - output = "LoRA model not found on CivitAI" - return {"ui": {"text": (output,)}} - + (output, triggerWords, examplePrompt, baseModel) = get_lora_info(lora_name) + return {"ui": {"text": (output,), "model": (baseModel,)}, "result": (triggerWords, examplePrompt)} # A dictionary that contains all nodes you want to export with their names # NOTE: names should be globally unique NODE_CLASS_MAPPINGS = { - "LoraInfo": LoraInfo + "LoraInfo": LoraInfo, } # A dictionary that contains the friendly/humanly readable titles for the nodes NODE_DISPLAY_NAME_MAPPINGS = { - "LoraInfo": "Lora Info" + "LoraInfo": "Lora Info", } WEB_DIRECTORY = "./js" \ No newline at end of file diff --git a/preview.png b/preview.png index b6cfae0..4aa6c1f 100644 Binary files a/preview.png and b/preview.png differ diff --git a/readme.md b/readme.md index 70eca00..cd23f88 100644 --- a/readme.md +++ b/readme.md @@ -2,7 +2,10 @@ Shows Lora Base Model, Trigger Words and Examples -Special thanks to `badjeff` for doing all the actual hard work (https://github.com/badjeff/comfyui_lora_tag_loader) +Special thanks to: + +- `badjeff` for doing all the actual hard work (https://github.com/badjeff/comfyui_lora_tag_loader) +- `alessandroperilli` for all the feedback. ![image](preview.png)