From f69e5ffa7b0e4ae490c755e0bfa51ff4e57f9f97 Mon Sep 17 00:00:00 2001 From: pythongosssss <125205205+pythongosssss@users.noreply.github.com> Date: Sun, 27 Aug 2023 13:40:11 +0100 Subject: [PATCH] Refactor lora info, add checkpoint info --- py/{lora_info.py => model_info.py} | 6 +- web/js/loraInfo.js | 397 --------------------- web/js/{loraInfo.css => modelInfo.css} | 35 +- web/js/modelInfo.js | 457 +++++++++++++++++++++++++ 4 files changed, 482 insertions(+), 413 deletions(-) rename py/{lora_info.py => model_info.py} (95%) delete mode 100644 web/js/loraInfo.js rename web/js/{loraInfo.css => modelInfo.css} (64%) create mode 100644 web/js/modelInfo.js diff --git a/py/lora_info.py b/py/model_info.py similarity index 95% rename from py/lora_info.py rename to py/model_info.py index 898b384..30cd4ad 100644 --- a/py/lora_info.py +++ b/py/model_info.py @@ -35,7 +35,11 @@ async def load_metadata(request): if not file_path: return web.Response(status=404) - meta = get_metadata(file_path) + try: + meta = get_metadata(file_path) + except: + meta = None + if meta is None: meta = {} diff --git a/web/js/loraInfo.js b/web/js/loraInfo.js deleted file mode 100644 index ef84845..0000000 --- a/web/js/loraInfo.js +++ /dev/null @@ -1,397 +0,0 @@ -import { app } from "../../../scripts/app.js"; -import { $el, ComfyDialog } from "../../../scripts/ui.js"; -import { api } from "../../../scripts/api.js"; -import { addStylesheet, getUrl } from "./common/utils.js"; - -addStylesheet(getUrl("loraInfo.css", import.meta.url)); - -const MAX_TAGS = 500; - -class LoraMetadataDialog extends ComfyDialog { - constructor(name, metadata) { - super(); - - this.element.classList.add("pysssss-lora-metadata"); - } - - show(metadata) { - super.show( - $el( - "div", - Object.keys(metadata).map((k) => - $el("div", [$el("label", { textContent: k }), $el("span", { textContent: metadata[k] })]) - ) - ) - ); - } -} - -class LoraInfoDialog extends ComfyDialog { - #metadata; - - get tagFrequency() { - if (!this.#metadata.ss_tag_frequency) return []; - - const datasets = JSON.parse(this.#metadata.ss_tag_frequency); - const tags = {}; - for (const setName in datasets) { - const set = datasets[setName]; - for (const t in set) { - if (t in tags) { - tags[t] += set[t]; - } else { - tags[t] = set[t]; - } - } - } - - return Object.entries(tags).sort((a, b) => b[1] - a[1]); - } - - get resolutions() { - let res = []; - if (this.#metadata.ss_bucket_info) { - const parsed = JSON.parse(this.#metadata.ss_bucket_info); - if (parsed?.buckets) { - for (const { resolution, count } of Object.values(parsed.buckets)) { - res.push([count, `${resolution.join("x")} * ${count}`]); - } - } - } - res = res.sort((a, b) => b[0] - a[0]).map((a) => a[1]); - let r = this.#metadata.ss_resolution; - if (r) { - const s = r.split(","); - const w = s[0].replace("(", ""); - const h = s[1].replace(")", ""); - res.push(`${w.trim()}x${h.trim()} (Base res)`); - } else if ((r = this.#metadata["modelspec.resolution"])) { - res.push(r + " (Base res"); - } - if (!res.length) { - res.push("⚠️ Unknown"); - } - return res; - } - - get customNotes() { - return this.#metadata["pysssss.notes"]; - } - - get hash() { - return this.#metadata["pysssss.sha256"]; - } - - getTagList(tags) { - return tags.map((t) => - $el( - "li.pysssss-lora-tag", - { - dataset: { - tag: t[0], - }, - $: (el) => { - el.onclick = () => { - el.classList.toggle("pysssss-lora-tag--selected"); - }; - }, - }, - [ - $el("p", { - textContent: t[0], - }), - $el("span", { - textContent: t[1], - }), - ] - ) - ); - } - - constructor(name, metadata) { - super(); - - this.element.classList.add("pysssss-lora-info"); - this.#metadata = metadata; - - let tags = this.tagFrequency; - let hasMore; - if (tags?.length) { - const c = tags.length; - let list; - if (c > MAX_TAGS) { - tags = tags.slice(0, MAX_TAGS); - hasMore = $el("p", [ - $el("span", { textContent: `⚠️ Only showing first ${MAX_TAGS} tags ` }), - $el("a", { - href: "#", - textContent: `Show all ${c}`, - onclick: () => { - list.replaceChildren(...this.getTagList(this.tagFrequency)); - hasMore.remove(); - }, - }), - ]); - } - list = $el("ol.pysssss-lora-tags-list", this.getTagList(tags)); - this.tags = $el("div", [list]); - } else { - this.tags = $el("p", { textContent: "⚠️ No tag frequency metadata found" }); - } - - const resolutions = $el("label", { textContent: "Resolution:" }, [ - $el( - "select", - this.resolutions.map((r) => $el("option", { textContent: r })) - ), - ]); - - let notes = []; - - if (this.customNotes) { - const r = new RegExp("(\\bhttps?:\\/\\/[^\\s]+)", "g"); - let end = 0; - let m; - do { - m = r.exec(this.customNotes); - let pos; - let fin = 0; - if (m) { - pos = m.index; - fin = m.index + m[0].length; - } else { - pos = this.customNotes.length; - } - - let pre = this.customNotes.substring(end, pos); - if (pre) { - pre = pre.replaceAll("\n", "
"); - notes.push( - $el("span", { - innerHTML: pre, - }) - ); - } - if (m) { - notes.push( - $el("a", { - href: m[0], - textContent: m[0], - target: "_blank", - }) - ); - } - - end = fin; - } while (m); - } - - if (notes.length) { - notes = $el("p", { textContent: "Notes: " }, notes); - } else { - notes = $el("p", { - textContent: "Notes: " + (this.customNotes ?? `Add custom notes in ${name.split(".")[0] + ".txt"}`), - }); - } - - const img = $el("img", { - style: { - display: "hidden", - maxHeight: "300px", - marginLeft: "10px", - }, - }); - - let triggers = $el("div"); - - let main; - const civitaiInfo = this.createCivitaiInfo((images, words, description) => { - const image = images.find((img) => img.nsfw === "None") || images[0]; - img.src = image.url; - img.style.display = "block"; - triggers.append(words); - triggers.append(description); - }); - - main = $el( - "main", - { - style: { - display: "flex", - }, - }, - [ - $el("div", [ - $el("p", { - textContent: "Output Name: " + (metadata.ss_output_name || "⚠️ Unknown"), - }), - $el("p", { - textContent: "Base Model: " + (metadata.ss_sd_model_name || "⚠️ Unknown"), - }), - $el("p", { - textContent: "Clip Skip: " + (metadata.ss_clip_skip || "⚠️ Unknown"), - }), - resolutions, - notes, - civitaiInfo, - ]), - img, - ] - ); - - this.content = $el( - "div.pysssss-lora-content", - [$el("h2", { textContent: name }), main, triggers, this.tags, hasMore].filter(Boolean) - ); - } - - createCivitaiInfo(loadCallback) { - 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", - }), - $el("br"), - $el("br") - ); - - loadCallback( - res.images, - $el("p", { textContent: "Trained Words: " }, [ - $el("pre", { - textContent: res.trainedWords.join(", "), - style: { - whiteSpace: "pre-wrap", - margin: "10px 0", - background: "#222", - padding: "5px", - borderRadius: "5px", - maxHeight: "250px", - overflow: "auto", - }, - }), - ]), - $el("div", { innerHTML: res.description, style: { maxHeight: "250px", overflow: "auto" } }) - ); - } 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(); - - function copyTags(e, tags) { - const textarea = $el("textarea", { - parent: document.body, - style: { - position: "fixed", - }, - textContent: tags.map((el) => el.dataset.tag).join(", "), - }); - textarea.select(); - try { - document.execCommand("copy"); - if (!e.target.dataset.text) { - e.target.dataset.text = e.target.textContent; - } - e.target.textContent = "Copied " + tags.length + " tags"; - setTimeout(() => { - e.target.textContent = e.target.dataset.text; - }, 1000); - } catch (ex) { - prompt("Copy to clipboard: Ctrl+C, Enter", text); - } finally { - document.body.removeChild(textarea); - } - } - - btns.unshift( - $el("button", { - type: "button", - textContent: "Copy Selected", - onclick: (e) => { - copyTags(e, [...this.tags.querySelectorAll(".pysssss-lora-tag--selected")]); - }, - }), - $el("button", { - type: "button", - textContent: "Copy All", - onclick: (e) => { - copyTags(e, [...this.tags.querySelectorAll(".pysssss-lora-tag")]); - }, - }), - $el("button", { - type: "button", - textContent: "View raw metadata", - onclick: (e) => { - new LoraMetadataDialog().show(this.#metadata); - }, - }) - ); - return btns; - } - - show() { - super.show(this.content); - } -} - -app.registerExtension({ - name: "pysssss.LoraInfo", - beforeRegisterNodeDef(nodeType, nodeData, app) { - if (nodeType.comfyClass === "LoraLoader" || nodeType.comfyClass === "LoraLoader|pysssss") { - const getExtraMenuOptions = nodeType.prototype.getExtraMenuOptions; - nodeType.prototype.getExtraMenuOptions = function (_, options) { - let value = this.widgets[0].value; - if (!value) { - return; - } - if (value.content) { - value = value.content; - } - options.unshift({ - content: "View info...", - callback: async () => { - const meta = await (await api.fetchApi("/pysssss/metadata/" + encodeURIComponent(`loras/${value}`))).json(); - new LoraInfoDialog(value, meta).show(); - }, - }); - - return getExtraMenuOptions?.apply(this, arguments); - }; - } - }, -}); diff --git a/web/js/loraInfo.css b/web/js/modelInfo.css similarity index 64% rename from web/js/loraInfo.css rename to web/js/modelInfo.css index dc88dba..9bfbf82 100644 --- a/web/js/loraInfo.css +++ b/web/js/modelInfo.css @@ -1,27 +1,32 @@ -.pysssss-lora-info { +.pysssss-model-info { color: white; font-family: sans-serif; max-width: 90vw; } -.pysssss-lora-content { +.pysssss-model-content { display: flex; flex-direction: column; overflow: hidden; } -.pysssss-lora-info h2 { +.pysssss-model-info h2 { text-align: center; margin: 0 0 10px 0; } -.pysssss-lora-info p { +.pysssss-model-info main img { + max-height: 300px; + margin-left: 10px; + max-width: 50%; +} +.pysssss-model-info p { margin: 5px 0; } -.pysssss-lora-info a { +.pysssss-model-info a { color: dodgerblue; } -.pysssss-lora-info a:hover { +.pysssss-model-info a:hover { text-decoration: underline; } -.pysssss-lora-tags-list { +.pysssss-model-tags-list { display: flex; flex-wrap: wrap; list-style: none; @@ -31,7 +36,7 @@ margin: 10px 0; padding: 0; } -.pysssss-lora-tag { +.pysssss-model-tag { background-color: rgb(128, 213, 247); color: #000; display: flex; @@ -41,7 +46,7 @@ padding: 2px 5px; cursor: pointer; } -.pysssss-lora-tag--selected span::before { +.pysssss-model-tag--selected span::before { content: "✅"; position: absolute; background-color: dodgerblue; @@ -51,13 +56,13 @@ bottom: 0; text-align: center; } -.pysssss-lora-tag:hover { +.pysssss-model-tag:hover { outline: 2px solid dodgerblue; } -.pysssss-lora-tag p { +.pysssss-model-tag p { margin: 0; } -.pysssss-lora-tag span { +.pysssss-model-tag span { text-align: center; border-radius: 5px; background-color: dodgerblue; @@ -68,14 +73,14 @@ overflow: hidden; } -.pysssss-lora-metadata .comfy-modal-content { +.pysssss-model-metadata .comfy-modal-content { max-width: 100%; } -.pysssss-lora-metadata label { +.pysssss-model-metadata label { margin-right: 1ch; color: #ccc; } -.pysssss-lora-metadata span { +.pysssss-model-metadata span { color: dodgerblue; } diff --git a/web/js/modelInfo.js b/web/js/modelInfo.js new file mode 100644 index 0000000..dd372a7 --- /dev/null +++ b/web/js/modelInfo.js @@ -0,0 +1,457 @@ +import { app } from "../../../scripts/app.js"; +import { $el, ComfyDialog } from "../../../scripts/ui.js"; +import { api } from "../../../scripts/api.js"; +import { addStylesheet, getUrl } from "./common/utils.js"; + +addStylesheet(getUrl("modelInfo.css", import.meta.url)); + +const MAX_TAGS = 500; + +class MetadataDialog extends ComfyDialog { + constructor() { + super(); + + this.element.classList.add("pysssss-model-metadata"); + } + show(metadata) { + super.show( + $el( + "div", + Object.keys(metadata).map((k) => + $el("div", [$el("label", { textContent: k }), $el("span", { textContent: metadata[k] })]) + ) + ) + ); + } +} + +class ModelInfoDialog extends ComfyDialog { + constructor(name) { + super(); + this.name = name; + this.element.classList.add("pysssss-model-info"); + } + + get customNotes() { + return this.metadata["pysssss.notes"]; + } + + get hash() { + return this.metadata["pysssss.sha256"]; + } + + async show(req) { + this.info = $el("div"); + this.img = $el("img", { style: { display: "none" } }); + this.main = $el("main", { style: { display: "flex" } }, [this.info, this.img]); + this.content = $el("div.pysssss-model-content", [$el("h2", { textContent: this.name }), this.main]); + + const loading = $el("div", { textContent: "ℹ️ Loading...", parent: this.content }); + + super.show(this.content); + + this.metadata = await (await req).json(); + this.viewMetadata.style.cursor = this.viewMetadata.style.opacity = ""; + this.viewMetadata.removeAttribute("disabled"); + + loading.remove(); + this.addInfo(); + } + + createButtons() { + const btns = super.createButtons(); + this.viewMetadata = $el("button", { + type: "button", + textContent: "View raw metadata", + disabled: "disabled", + style: { + opacity: 0.5, + cursor: "not-allowed", + }, + onclick: (e) => { + if (this.metadata) { + new MetadataDialog().show(this.metadata); + } + }, + }); + + btns.unshift(this.viewMetadata); + return btns; + } + + getNoteInfo() { + if (this.customNotes) { + let notes = []; + // Extract links from notes + const r = new RegExp("(\\bhttps?:\\/\\/[^\\s]+)", "g"); + let end = 0; + let m; + do { + m = r.exec(this.customNotes); + let pos; + let fin = 0; + if (m) { + pos = m.index; + fin = m.index + m[0].length; + } else { + pos = this.customNotes.length; + } + + let pre = this.customNotes.substring(end, pos); + if (pre) { + pre = pre.replaceAll("\n", "
"); + notes.push( + $el("span", { + innerHTML: pre, + }) + ); + } + if (m) { + notes.push( + $el("a", { + href: m[0], + textContent: m[0], + target: "_blank", + }) + ); + } + + end = fin; + } while (m); + return $el("span", notes); + } else { + return `Add custom notes in ${this.name.split(".")[0] + ".txt"}`; + } + } + + addInfo() { + this.addInfoEntry("Notes", this.getNoteInfo()); + } + + addInfoEntry(name, value) { + return $el( + "p", + { + parent: this.info, + }, + [ + typeof name === "string" ? $el("label", { textContent: name + ": " }) : name, + typeof value === "string" ? $el("span", { textContent: value }) : value, + ] + ); + } + + async getCivitaiDetails() { + const req = await fetch("https://civitai.com/api/v1/model-versions/by-hash/" + this.hash); + if (req.status === 200) { + return await req.json(); + } else if (req.status === 404) { + throw new Error("Model not found"); + } else { + throw new Error(`Error loading info (${req.status}) ${req.statusText}`); + } + } + + addCivitaiInfo() { + const promise = this.getCivitaiDetails(); + const content = $el("span", { textContent: "ℹ️ Loading..." }); + + this.addInfoEntry( + $el("label", [ + $el("img", { + style: { + width: "18px", + position: "relative", + top: "3px", + margin: "0 5px 0 0", + }, + src: "https://civitai.com/favicon.ico", + }), + $el("span", { textContent: "Civitai: " }), + ]), + content + ); + + return promise + .then((info) => { + content.replaceChildren( + $el("a", { + href: "https://civitai.com/models/" + info.modelId, + textContent: "View " + info.model.name, + target: "_blank", + }) + ); + + if (info.images?.length) { + this.img.src = info.images[0].url; + this.img.style.display = ""; + } + + return info; + }) + .catch((err) => { + content.textContent = "⚠️ " + err.message; + }); + } +} + +class LoraInfoDialog extends ModelInfoDialog { + getTagFrequency() { + if (!this.metadata.ss_tag_frequency) return []; + + const datasets = JSON.parse(this.metadata.ss_tag_frequency); + const tags = {}; + for (const setName in datasets) { + const set = datasets[setName]; + for (const t in set) { + if (t in tags) { + tags[t] += set[t]; + } else { + tags[t] = set[t]; + } + } + } + + return Object.entries(tags).sort((a, b) => b[1] - a[1]); + } + + getResolutions() { + let res = []; + if (this.metadata.ss_bucket_info) { + const parsed = JSON.parse(this.metadata.ss_bucket_info); + if (parsed?.buckets) { + for (const { resolution, count } of Object.values(parsed.buckets)) { + res.push([count, `${resolution.join("x")} * ${count}`]); + } + } + } + res = res.sort((a, b) => b[0] - a[0]).map((a) => a[1]); + let r = this.metadata.ss_resolution; + if (r) { + const s = r.split(","); + const w = s[0].replace("(", ""); + const h = s[1].replace(")", ""); + res.push(`${w.trim()}x${h.trim()} (Base res)`); + } else if ((r = this.metadata["modelspec.resolution"])) { + res.push(r + " (Base res"); + } + if (!res.length) { + res.push("⚠️ Unknown"); + } + return res; + } + + getTagList(tags) { + return tags.map((t) => + $el( + "li.pysssss-model-tag", + { + dataset: { + tag: t[0], + }, + $: (el) => { + el.onclick = () => { + el.classList.toggle("pysssss-model-tag--selected"); + }; + }, + }, + [ + $el("p", { + textContent: t[0], + }), + $el("span", { + textContent: t[1], + }), + ] + ) + ); + } + + addTags() { + let tags = this.getTagFrequency(); + let hasMore; + if (tags?.length) { + const c = tags.length; + let list; + if (c > MAX_TAGS) { + tags = tags.slice(0, MAX_TAGS); + hasMore = $el("p", [ + $el("span", { textContent: `⚠️ Only showing first ${MAX_TAGS} tags ` }), + $el("a", { + href: "#", + textContent: `Show all ${c}`, + onclick: () => { + list.replaceChildren(...this.getTagList(this.getTagFrequency())); + hasMore.remove(); + }, + }), + ]); + } + list = $el("ol.pysssss-model-tags-list", this.getTagList(tags)); + this.tags = $el("div", [list]); + } else { + this.tags = $el("p", { textContent: "⚠️ No tag frequency metadata found" }); + } + + this.content.append(this.tags); + + if (hasMore) { + this.content.append(hasMore); + } + } + + async addInfo() { + this.addInfoEntry("Name", this.metadata.ss_output_name || "⚠️ Unknown"); + this.addInfoEntry("Base Model", this.metadata.ss_sd_model_name || "⚠️ Unknown"); + this.addInfoEntry("Clip Skip", this.metadata.ss_clip_skip || "⚠️ Unknown"); + + this.addInfoEntry( + "Resolution", + $el( + "select", + this.getResolutions().map((r) => $el("option", { textContent: r })) + ) + ); + + super.addInfo(); + const p = this.addCivitaiInfo(); + this.addTags(); + + const info = await p; + if (info) { + $el( + "p", + { + parent: this.content, + textContent: "Trained Words: ", + }, + [ + $el("pre", { + textContent: info.trainedWords.join(", "), + style: { + whiteSpace: "pre-wrap", + margin: "10px 0", + background: "#222", + padding: "5px", + borderRadius: "5px", + maxHeight: "250px", + overflow: "auto", + }, + }), + ] + ); + $el("div", { + parent: this.content, + innerHTML: info.description, + style: { + maxHeight: "250px", + overflow: "auto", + }, + }); + } + } + + createButtons() { + const btns = super.createButtons(); + + function copyTags(e, tags) { + const textarea = $el("textarea", { + parent: document.body, + style: { + position: "fixed", + }, + textContent: tags.map((el) => el.dataset.tag).join(", "), + }); + textarea.select(); + try { + document.execCommand("copy"); + if (!e.target.dataset.text) { + e.target.dataset.text = e.target.textContent; + } + e.target.textContent = "Copied " + tags.length + " tags"; + setTimeout(() => { + e.target.textContent = e.target.dataset.text; + }, 1000); + } catch (ex) { + prompt("Copy to clipboard: Ctrl+C, Enter", text); + } finally { + document.body.removeChild(textarea); + } + } + + btns.unshift( + $el("button", { + type: "button", + textContent: "Copy Selected", + onclick: (e) => { + copyTags(e, [...this.tags.querySelectorAll(".pysssss-model-tag--selected")]); + }, + }), + $el("button", { + type: "button", + textContent: "Copy All", + onclick: (e) => { + copyTags(e, [...this.tags.querySelectorAll(".pysssss-model-tag")]); + }, + }) + ); + + return btns; + } +} + +class CheckpointInfoDialog extends ModelInfoDialog { + async addInfo() { + super.addInfo(); + const info = await this.addCivitaiInfo(); + if (info) { + this.addInfoEntry("Base Model", info.baseModel || "⚠️ Unknown"); + + $el("div", { + parent: this.content, + innerHTML: info.description, + style: { + maxHeight: "250px", + overflow: "auto", + }, + }); + } + } +} + +const infoHandler = { + LoraLoader: "loras", + "LoraLoader|pysssss": "loras", + CheckpointLoader: "checkpoints", + CheckpointLoaderSimple: "checkpoints", + "CheckpointLoader|pysssss": "checkpoints", +}; + +app.registerExtension({ + name: "pysssss.ModelInfo", + beforeRegisterNodeDef(nodeType) { + const type = infoHandler[nodeType.comfyClass]; + + if (type) { + const cls = type === "loras" ? LoraInfoDialog : CheckpointInfoDialog; + const getExtraMenuOptions = nodeType.prototype.getExtraMenuOptions; + nodeType.prototype.getExtraMenuOptions = function (_, options) { + let value = this.widgets[0].value; + if (!value) { + return; + } + if (value.content) { + value = value.content; + } + options.unshift({ + content: "View info...", + callback: async () => { + new cls(value).show(api.fetchApi("/pysssss/metadata/" + encodeURIComponent(`${type}/${value}`))); + }, + }); + + return getExtraMenuOptions?.apply(this, arguments); + }; + } + }, +});