added output and base model label

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];
       }
     }
   },

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"

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)