From d6a02a31a72c324199b772b089d760cd7514f8cc Mon Sep 17 00:00:00 2001 From: JeLuF Date: Sun, 17 Sep 2023 22:36:50 +0200 Subject: [PATCH] LoraManager: Implement 'Upload thumbnail' button --- ui/media/js/main.js | 1 + ui/plugins/ui/model-tools.plugin.js | 770 ++++++++++++++++++++++++++++ 2 files changed, 771 insertions(+) create mode 100644 ui/plugins/ui/model-tools.plugin.js diff --git a/ui/media/js/main.js b/ui/media/js/main.js index 7634f288..101a88e7 100644 --- a/ui/media/js/main.js +++ b/ui/media/js/main.js @@ -868,6 +868,7 @@ useAsThumbSaveBtn.addEventListener("click", (e) => { }) .then(() => { useAsThumbDialog.close() + document.dispatchEvent(new CustomEvent("saveThumb", { detail: useAsThumbSelect.selectedOptions })) }) .catch((error) => { console.error(error) diff --git a/ui/plugins/ui/model-tools.plugin.js b/ui/plugins/ui/model-tools.plugin.js new file mode 100644 index 00000000..09fd6499 --- /dev/null +++ b/ui/plugins/ui/model-tools.plugin.js @@ -0,0 +1,770 @@ +;(function() { + "use strict" + + let mergeCSS = ` + /*********** Main tab ***********/ + .tab-centered { + justify-content: center; + } + + #model-tool-tab-content { + background-color: var(--background-color3); + } + + #model-tool-tab-content .tab-content-inner { + text-align: initial; + } + + #model-tool-tab-bar .tab { + margin-bottom: 0px; + border-top-left-radius: var(--input-border-radius); + background-color: var(--background-color3); + padding: 6px 6px 0.8em 6px; + } + #tab-content-merge .tab-content-inner { + max-width: 100%; + padding: 10pt; + } + + /*********** Merge UI ***********/ + .merge-model-container { + margin-left: 15%; + margin-right: 15%; + text-align: left; + display: inline-grid; + grid-template-columns: 1fr 1fr; + grid-template-rows: auto auto auto; + gap: 0px 0px; + grid-auto-flow: row; + grid-template-areas: + "merge-input merge-config" + "merge-buttons merge-buttons"; + } + .merge-model-container p { + margin-top: 3pt; + margin-bottom: 3pt; + } + .merge-config .tab-content { + background: var(--background-color1); + border-radius: 3pt; + } + .merge-config .tab-content-inner { + text-align: left; + } + + .merge-input { + grid-area: merge-input; + padding-left:1em; + } + .merge-config { + grid-area: merge-config; + padding:1em; + } + .merge-config input { + margin-bottom: 3px; + } + .merge-config select { + margin-bottom: 3px; + } + .merge-buttons { + grid-area: merge-buttons; + padding:1em; + text-align: center; + } + #merge-button { + padding: 8px; + width:20em; + } + div#merge-log { + height:150px; + overflow-x:hidden; + overflow-y:scroll; + background:var(--background-color1); + border-radius: 3pt; + } + div#merge-log i { + color: hsl(var(--accent-hue), 100%, calc(2*var(--accent-lightness))); + font-family: monospace; + } + .disabled { + background: var(--background-color4); + color: var(--text-color); + } + #merge-type-tabs { + border-bottom: 1px solid black; + } + #merge-log-container { + display: none; + } + .merge-model-container #merge-warning { + color: var(--small-label-color); + } + + /*********** LORA UI ***********/ + .lora-manager-grid { + display: grid; + gap: 0px 8px; + grid-auto-flow: row; + } + + @media screen and (min-width: 1501px) { + .lora-manager-grid textarea { + height:350px; + } + + .lora-manager-grid { + grid-template-columns: auto 1fr 1fr; + grid-template-rows: auto 1fr; + grid-template-areas: + "selector selector selector" + "thumbnail keywords notes"; + } + } + + @media screen and (min-width: 1001px) and (max-width: 1500px) { + .lora-manager-grid textarea { + height:250px; + } + + .lora-manager-grid { + grid-template-columns: auto auto; + grid-template-rows: auto auto auto; + grid-template-areas: + "selector selector" + "thumbnail keywords" + "thumbnail notes"; + } + + } + + @media screen and (max-width: 1000px) { + .lora-manager-grid textarea { + height:200px; + } + + .lora-manager-grid { + grid-template-columns: auto; + grid-template-rows: auto auto auto auto; + grid-template-areas: + "selector" + "keywords" + "thumbnail" + "notes"; + } + + } + + .lora-manager-grid-selector { + grid-area: selector; + justify-self: start; + } + + .lora-manager-grid-thumbnail { + grid-area: thumbnail; + justify-self: center; + } + + .lora-manager-grid-keywords { + grid-area: keywords; + } + + .lora-manager-grid-notes { + grid-area: notes; + } + + .lora-manager-grid p { + margin-bottom: 2px; + } + + + ` + + let mergeUI = ` +
+
+

+ +

+ +

+

Important: Please merge models of similar type.
For e.g. SD 1.4 models with only SD 1.4/1.5 models,
SD 2.0 with SD 2.0-type, and SD 2.1 with SD 2.1-type models.

+
+ + + + + + + + + + + + + +
Base name of the output file.
Mix ratio and file suffix will be appended to this.
+ Image generation uses fp16, so it's a good choice.
Use fp32 if you want to use the result models for more mixes +
+
+
+
+

+
+
+
+
+
+ + Make a single file + + + Make multiple variations + +
+
+
+
+ Saves a single merged model file, at the specified merge ratio.

+ + + % + Model A's contribution to the mix. The rest will be from Model B. +
+
+
+
+ Saves multiple variations of the model, at different merge ratios.
Each variation will be saved as a separate file.

+ + + + + + + + + + + + + +
Number of models to create
% Smallest share of model A in the mix
% Share of model A added into the mix per step
Sigmoid function to be applied to the model share before mixing
+
+ Preview of variation ratios:
+ +
+
+
+
+
+ +
+
` + + + let loraUI=` +
+
+ + +
+
+

Thumbnail:

+
+ + +
+
+ + + +
+
+
+

Keywords: +

+ +

+ LORA model keywords can be used via the + Embeddings button. They get added to the embedding + keyword menu when the LORA has been selected in the image settings. +

+
+
+

Notes:

+ +

+ Civitai model page: + +

+
+
` + + let tabHTML=` +
+ + Lora Keywords + + + Merge Models + +
+
+
+
+ ${loraUI} +
+
+ +
+
+ ${mergeUI} +
+
+
` + + + ///////////////////// Function section + function smoothstep(x) { + return x * x * (3 - 2 * x) + } + + function smootherstep(x) { + return x * x * x * (x * (x * 6 - 15) + 10) + } + + function smootheststep(x) { + let y = -20 * Math.pow(x, 7) + y += 70 * Math.pow(x, 6) + y -= 84 * Math.pow(x, 5) + y += 35 * Math.pow(x, 4) + return y + } + function getCurrentTime() { + const now = new Date() + let hours = now.getHours() + let minutes = now.getMinutes() + let seconds = now.getSeconds() + + hours = hours < 10 ? `0${hours}` : hours + minutes = minutes < 10 ? `0${minutes}` : minutes + seconds = seconds < 10 ? `0${seconds}` : seconds + + return `${hours}:${minutes}:${seconds}` + } + + function addLogMessage(message) { + const logContainer = document.getElementById("merge-log") + logContainer.innerHTML += `${getCurrentTime()} ${message}
` + + // Scroll to the bottom of the log + logContainer.scrollTop = logContainer.scrollHeight + + document.querySelector("#merge-log-container").style.display = "block" + } + + function addLogSeparator() { + const logContainer = document.getElementById("merge-log") + logContainer.innerHTML += "
" + + logContainer.scrollTop = logContainer.scrollHeight + } + + function drawDiagram(fn) { + const SIZE = 300 + const canvas = document.getElementById("merge-canvas") + canvas.height = canvas.width = SIZE + const ctx = canvas.getContext("2d") + + // Draw coordinate system + ctx.scale(1, -1) + ctx.translate(0, -canvas.height) + ctx.lineWidth = 1 + ctx.beginPath() + + ctx.strokeStyle = "white" + ctx.moveTo(0, 0) + ctx.lineTo(0, SIZE) + ctx.lineTo(SIZE, SIZE) + ctx.lineTo(SIZE, 0) + ctx.lineTo(0, 0) + ctx.lineTo(SIZE, SIZE) + ctx.stroke() + ctx.beginPath() + ctx.setLineDash([1, 2]) + const n = SIZE / 10 + for (let i = n; i < SIZE; i += n) { + ctx.moveTo(0, i) + ctx.lineTo(SIZE, i) + ctx.moveTo(i, 0) + ctx.lineTo(i, SIZE) + } + ctx.stroke() + ctx.beginPath() + ctx.setLineDash([]) + ctx.beginPath() + ctx.strokeStyle = "black" + ctx.lineWidth = 3 + // Plot function + const numSamples = 20 + for (let i = 0; i <= numSamples; i++) { + const x = i / numSamples + const y = fn(x) + + const canvasX = x * SIZE + const canvasY = y * SIZE + + if (i === 0) { + ctx.moveTo(canvasX, canvasY) + } else { + ctx.lineTo(canvasX, canvasY) + } + } + ctx.stroke() + // Plot alpha values (yellow boxes) + let start = parseFloat(document.querySelector("#merge-start").value) + let step = parseFloat(document.querySelector("#merge-step").value) + let iterations = document.querySelector("#merge-count").value >> 0 + ctx.beginPath() + ctx.fillStyle = "yellow" + for (let i = 0; i < iterations; i++) { + const alpha = (start + i * step) / 100 + const x = alpha * SIZE + const y = fn(alpha) * SIZE + if (x <= SIZE) { + ctx.rect(x - 3, y - 3, 6, 6) + ctx.fill() + } else { + ctx.strokeStyle = "red" + ctx.moveTo(0, 0) + ctx.lineTo(0, SIZE) + ctx.lineTo(SIZE, SIZE) + ctx.lineTo(SIZE, 0) + ctx.lineTo(0, 0) + ctx.lineTo(SIZE, SIZE) + ctx.stroke() + addLogMessage("Warning: maximum ratio is ≥ 100%") + } + } + } + + function updateChart() { + let fn = (x) => x + switch (document.querySelector("#merge-interpolation").value) { + case "SmoothStep": + fn = smoothstep + break + case "SmootherStep": + fn = smootherstep + break + case "SmoothestStep": + fn = smootheststep + break + } + drawDiagram(fn) + } + + function initMergeUI() { + const tabSettingsSingle = document.querySelector("#tab-merge-opts-single") + const tabSettingsBatch = document.querySelector("#tab-merge-opts-batch") + linkTabContents(tabSettingsSingle) + linkTabContents(tabSettingsBatch) + + let mergeModelAField = new ModelDropdown(document.querySelector("#mergeModelA"), "stable-diffusion") + let mergeModelBField = new ModelDropdown(document.querySelector("#mergeModelB"), "stable-diffusion") + updateChart() + + // slider + const singleMergeRatioField = document.querySelector("#single-merge-ratio") + const singleMergeRatioSlider = document.querySelector("#single-merge-ratio-slider") + + function updateSingleMergeRatio() { + singleMergeRatioField.value = singleMergeRatioSlider.value / 10 + singleMergeRatioField.dispatchEvent(new Event("change")) + } + + function updateSingleMergeRatioSlider() { + if (singleMergeRatioField.value < 0) { + singleMergeRatioField.value = 0 + } else if (singleMergeRatioField.value > 100) { + singleMergeRatioField.value = 100 + } + + singleMergeRatioSlider.value = singleMergeRatioField.value * 10 + singleMergeRatioSlider.dispatchEvent(new Event("change")) + } + + singleMergeRatioSlider.addEventListener("input", updateSingleMergeRatio) + singleMergeRatioField.addEventListener("input", updateSingleMergeRatioSlider) + updateSingleMergeRatio() + + document.querySelector(".merge-config").addEventListener("change", updateChart) + + document.querySelector("#merge-button").addEventListener("click", async function(e) { + // Build request template + let model0 = mergeModelAField.value + let model1 = mergeModelBField.value + let request = { model0: model0, model1: model1 } + request["use_fp16"] = document.querySelector("#merge-fp").value == "fp16" + let iterations = document.querySelector("#merge-count").value >> 0 + let start = parseFloat(document.querySelector("#merge-start").value) + let step = parseFloat(document.querySelector("#merge-step").value) + + if (isTabActive(tabSettingsSingle)) { + start = parseFloat(singleMergeRatioField.value) + step = 0 + iterations = 1 + addLogMessage(`merge ratio = ${start}%`) + } else { + addLogMessage(`start = ${start}%`) + addLogMessage(`step = ${step}%`) + } + + if (start + (iterations - 1) * step >= 100) { + addLogMessage("Aborting: maximum ratio is ≥ 100%") + addLogMessage("Reduce the number of variations or the step size") + addLogSeparator() + document.querySelector("#merge-count").focus() + return + } + + if (document.querySelector("#merge-filename").value == "") { + addLogMessage("Aborting: No output file name specified") + addLogSeparator() + document.querySelector("#merge-filename").focus() + return + } + + // Disable merge button + e.target.disabled = true + e.target.classList.add("disabled") + let cursor = $("body").css("cursor") + let label = document.querySelector("#merge-button").innerHTML + $("body").css("cursor", "progress") + document.querySelector("#merge-button").innerHTML = "Merging models ..." + + addLogMessage("Merging models") + addLogMessage("Model A: " + model0) + addLogMessage("Model B: " + model1) + + // Batch main loop + for (let i = 0; i < iterations; i++) { + let alpha = (start + i * step) / 100 + + if (isTabActive(tabSettingsBatch)) { + switch (document.querySelector("#merge-interpolation").value) { + case "SmoothStep": + alpha = smoothstep(alpha) + break + case "SmootherStep": + alpha = smootherstep(alpha) + break + case "SmoothestStep": + alpha = smootheststep(alpha) + break + } + } + addLogMessage(`merging batch job ${i + 1}/${iterations}, alpha = ${alpha.toFixed(5)}...`) + + request["out_path"] = document.querySelector("#merge-filename").value + request["out_path"] += "-" + alpha.toFixed(5) + "." + document.querySelector("#merge-format").value + addLogMessage(`  filename: ${request["out_path"]}`) + + // sdkit documentation: "ratio - the ratio of the second model. 1 means only the second model will be used." + request["ratio"] = 1-alpha + let res = await fetch("/model/merge", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(request), + }) + const data = await res.json() + addLogMessage(JSON.stringify(data)) + } + addLogMessage( + "Done. The models have been saved to your models/stable-diffusion folder." + ) + addLogSeparator() + // Re-enable merge button + $("body").css("cursor", cursor) + document.querySelector("#merge-button").innerHTML = label + e.target.disabled = false + e.target.classList.remove("disabled") + + // Update model list + stableDiffusionModelField.innerHTML = "" + vaeModelField.innerHTML = "" + hypernetworkModelField.innerHTML = "" + await getModels() + }) + } + + const LoraUI = { + modelField: undefined, + keywordsField: undefined, + notesField: undefined, + civitaiImportBtn: undefined, + civitaiSecion: undefined, + civitaiAnchor: undefined, + image: undefined, + imagePlaceholder: undefined, + + init() { + LoraUI.modelField = new ModelDropdown(document.querySelector("#loraModel"), "lora") + LoraUI.keywordsField = document.querySelector("#lora-manager-keywords") + LoraUI.notesField = document.querySelector("#lora-manager-notes") + LoraUI.civitaiImportBtn = document.querySelector("#lora-keyword-from-civitai") + LoraUI.civitaiSection = document.querySelector("#civitai-section") + LoraUI.civitaiAnchor = document.querySelector("#civitai-model-page") + LoraUI.image = document.querySelector("#lora-manager-image") + LoraUI.imagePlaceholder = document.querySelector("#lora-manager-image-placeholder") + LoraUI.uploadBtn = document.querySelector("#lora-manager-upload-button") + LoraUI.uploadInput = document.querySelector("#lora-manager-upload-input") + + LoraUI.modelField.addEventListener("change", LoraUI.updateFields) + LoraUI.keywordsField.addEventListener("focusout", LoraUI.saveInfos) + LoraUI.notesField.addEventListener("focusout", LoraUI.saveInfos) + LoraUI.civitaiImportBtn.addEventListener("click", LoraUI.importFromCivitai) + + LoraUI.uploadBtn.addEventListener("click", (e) => LoraUI.uploadInput.click()) + LoraUI.uploadInput.addEventListener("change", LoraUI.uploadLoraThumb) + + document.addEventListener("saveThumb", LoraUI.updateFields) + + LoraUI.updateFields() + }, + + uploadLoraThumb(e) { + console.log(e) + if (LoraUI.uploadInput.files.length === 0) { + return + } + + let reader = new FileReader() + let file = LoraUI.uploadInput.files[0] + + reader.addEventListener("load", (event) => { + let img = document.createElement("img") + img.src = reader.result + onUseAsThumbnailClick( + { + use_lora_model: LoraUI.modelField.value, + }, + img + ) + }) + + if (file) { + reader.readAsDataURL(file) + } + }, + + updateFields() { + document.getElementById("civitai-section").classList.add("displayNone") + Bucket.retrieve(`modelinfo/lora/${LoraUI.modelField.value}`) + .then((info) => { + if (info == null) { + LoraUI.keywordsField.value = "" + LoraUI.notesField.value = "" + LoraUI.hideCivitaiLink() + } else { + LoraUI.keywordsField.value = info.keywords.join("\n") + LoraUI.notesField.value = info.notes + if ("civitai" in info && info["civitai"] != null) { + LoraUI.showCivitaiLink(info.civitai) + } else { + LoraUI.hideCivitaiLink() + } + } + }) + Bucket.getImageAsDataURL(`${profileNameField.value}/lora/${LoraUI.modelField.value}.png`) + .then((data) => { + LoraUI.image.src=data + LoraUI.image.classList.remove("displayNone") + LoraUI.imagePlaceholder.classList.add("displayNone") + }) + .catch((error) => { + LoraUI.image.classList.add("displayNone") + LoraUI.imagePlaceholder.classList.remove("displayNone") + }) + }, + + saveInfos() { + let info = { + keywords: LoraUI.keywordsField.value + .split("\n") + .filter((x) => (x != "")), + notes: LoraUI.notesField.value, + civitai: LoraUI.civitaiSection.checkVisibility() ? LoraUI.civitaiAnchor.href : null, + } + Bucket.store(`modelinfo/lora/${LoraUI.modelField.value}`, info) + }, + + importFromCivitai() { + document.body.style["cursor"] = "progress" + fetch("/sha256/lora/"+LoraUI.modelField.value) + .then((result) => result.json()) + .then((json) => fetch("https://civitai.com/api/v1/model-versions/by-hash/" + json.digest)) + .then((result) => result.json()) + .then((json) => { + document.body.style["cursor"] = "default" + if (json == null) { + return + } + if ("trainedWords" in json) { + LoraUI.keywordsField.value = json["trainedWords"].join("\n") + } else { + showToast("No keyword info found.") + } + if ("modelId" in json) { + LoraUI.showCivitaiLink("https://civitai.com/models/" + json.modelId) + } else { + LoraUI.hideCivitaiLink() + } + + LoraUI.saveInfos() + }) + }, + + showCivitaiLink(href) { + LoraUI.civitaiSection.classList.remove("displayNone") + LoraUI.civitaiAnchor.href = href + LoraUI.civitaiAnchor.innerHTML = LoraUI.civitaiAnchor.href + }, + + hideCivitaiLink() { + LoraUI.civitaiSection.classList.add("displayNone") + } + } + + createTab({ + id: "merge", + icon: "fa-toolbox", + label: "Model tools", + css: mergeCSS, + content: tabHTML, + onOpen: ({ firstOpen }) => { + if (!firstOpen) { + return + } + initMergeUI() + LoraUI.init() + const tabMergeUI = document.querySelector("#tab-model-mergeUI") + const tabLoraUI = document.querySelector("#tab-model-loraUI") + linkTabContents(tabMergeUI) + linkTabContents(tabLoraUI) + }, + }) +})() +async function getLoraKeywords(model) { + return Bucket.retrieve(`modelinfo/lora/${model}`) + .then((info) => info ? info.keywords : []) +}