From d0cbb4536bcbda8199d8f4b883fc2819b117b5fa Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Mon, 21 Aug 2023 13:58:37 +0100
Subject: [PATCH] Add custom notes and civitai link

---
 js/loraInfo.js  | 69 +++++++++++++++++++++++++++++++++++++++++++++++++
 py/lora_info.py | 24 +++++++++++++++--
 2 files changed, 91 insertions(+), 2 deletions(-)

diff --git a/js/loraInfo.js b/js/loraInfo.js
index 7997bac..8c81815 100644
--- a/js/loraInfo.js
+++ b/js/loraInfo.js
@@ -72,6 +72,14 @@ class LoraInfoDialog extends ComfyDialog {
 		return res;
 	}
 
+	get customNotes() {
+		return this.#metadata["pysssss.notes"];
+	}
+
+	get hash() {
+		return this.#metadata["pysssss.sha256"];
+	}
+
 	getTagList(tags) {
 		return tags.map((t) =>
 			$el(
@@ -136,6 +144,21 @@ class LoraInfoDialog extends ComfyDialog {
 			),
 		]);
 
+		let notes = null;
+		if (this.customNotes?.startsWith("http")) {
+			notes = $el("p", { textContent: "Notes: " }, [
+				$el("a", {
+					href: this.customNotes,
+					textContent: this.customNotes,
+					target: "_blank",
+				}),
+			]);
+		} else {
+			notes = $el("p", {
+				textContent: "Notes: " + (this.customNotes ?? `Add custom notes in ${name.split(".")[0] + ".txt"}`),
+			});
+		}
+
 		this.content = $el(
 			"div.pysssss-lora-content",
 			[
@@ -150,12 +173,58 @@ class LoraInfoDialog extends ComfyDialog {
 					textContent: "Clip Skip: " + (metadata.ss_clip_skip || "⚠️ Unknown"),
 				}),
 				resolutions,
+				notes,
+				this.createCivitaiInfo(),
 				this.tags,
 				hasMore,
 			].filter(Boolean)
 		);
 	}
 
+	createCivitaiInfo() {
+		if (!this.hash) return;
+
+		const info = $el("span", { textContent: "ℹ️ Loading..." });
+		const el = $el("p", [
+			$el("img", {
+				style: {
+					width: "18px",
+					position: "relative",
+					top: "3px",
+					marginRight: "5px",
+				},
+				src: "https://civitai.com/favicon.ico",
+			}),
+			$el("span", { textContent: "Civitai: " }),
+			info,
+		]);
+
+		(async () => {
+			try {
+				const req = await fetch("https://civitai.com/api/v1/model-versions/by-hash/" + this.hash);
+				if (req.status === 200) {
+					const res = await req.json();
+					info.replaceChildren(
+						$el("a", {
+							href: "https://civitai.com/models/" + res.modelId,
+							textContent: "View " + res.model.name,
+							target: "_blank",
+						})
+					);
+				} else if (req.status === 404) {
+					info.textContent = "⚠️ Model not found";
+				} else {
+					info.textContent = `⚠️ Error loading info (${req.status}) ` + req.statusText;
+				}
+			} catch (error) {
+				console.error(error);
+				info.textContent = "⚠️ Error loading info";
+			}
+		})();
+
+		return el;
+	}
+
 	createButtons() {
 		const btns = super.createButtons();
 
diff --git a/py/lora_info.py b/py/lora_info.py
index b93e6fb..8dcc2f7 100644
--- a/py/lora_info.py
+++ b/py/lora_info.py
@@ -1,8 +1,9 @@
+import hashlib
 import json
 from aiohttp import web
 from server import PromptServer
 import folder_paths
-
+import os
 
 def get_metadata(filepath):
     with open(filepath, "rb") as file:
@@ -32,6 +33,25 @@ async def load_metadata(request):
         type, name)
     if not file_path:
         return web.Response(status=404)
-
+    
     meta = get_metadata(file_path)
+
+    file_no_ext = os.path.splitext(file_path)[0]
+
+    info_file = file_no_ext + ".txt"
+    if os.path.isfile(info_file):
+        with open(info_file, "r") as f:
+            meta["pysssss.notes"] = f.read()
+
+    hash_file = file_no_ext + ".sha256"
+    if os.path.isfile(hash_file):
+        with open(hash_file, "rt") as f:
+            meta["pysssss.sha256"] = f.read()
+    else:
+        with open(file_path, "rb") as f:
+            meta["pysssss.sha256"] = hashlib.sha256(f.read()).hexdigest()
+        with open(hash_file, "wt") as f:
+            f.write(meta["pysssss.sha256"])
+            
+
     return web.json_response(meta)