Merge pull request #1362 from JeLuF/textualinv

Basic embeddings support

ui/easydiffusion/model_manager.py

@@ -27,6 +27,7 @@ MODEL_EXTENSIONS = {
     "realesrgan": [".pth"],
     "lora": [".ckpt", ".safetensors"],
     "codeformer": [".pth"],
+    "embeddings": [".pt", ".bin", ".safetensors"],
 }
 DEFAULT_MODELS = {
     "stable-diffusion": [
@@ -57,6 +58,9 @@ def init():
 def load_default_models(context: Context):
     set_vram_optimizations(context)
 
+    config = app.getConfig()
+    context.embeddings_path = os.path.join(app.MODELS_DIR, "embeddings")
+
     # init default model paths
     for model_type in MODELS_TO_LOAD_ON_START:
         context.model_paths[model_type] = resolve_model_to_use(model_type=model_type, fail_if_not_found=False)
@@ -317,6 +321,7 @@ def getModels(scan_for_malicious: bool = True):
             "hypernetwork": [],
             "lora": [],
             "codeformer": ["codeformer"],
+            "embeddings": [],
         },
     }
 
@@ -375,6 +380,7 @@ def getModels(scan_for_malicious: bool = True):
     listModels(model_type="hypernetwork")
     listModels(model_type="gfpgan")
     listModels(model_type="lora")
+    listModels(model_type="embeddings")
 
     if scan_for_malicious and models_scanned > 0:
         log.info(f"[green]Scanned {models_scanned} models. Nothing infected[/]")

ui/index.html

@@ -244,6 +244,9 @@
                         <td><label for="hypernetwork_strength_slider">Hypernetwork Strength:</label></td>
                         <td> <input id="hypernetwork_strength_slider" name="hypernetwork_strength_slider" class="editor-slider" value="100" type="range" min="0" max="100"> <input id="hypernetwork_strength" name="hypernetwork_strength" size="4" pattern="^[0-9\.]+$" onkeypress="preventNonNumericalInput(event)"><br/></td>
                     </tr>
+                    <tr id="embeddings-container" class="pl-5 displayNone"><td><label for="embeddings-button">Embedding:</label></td><td>
+                        <button id="embeddings-button" class="tertiaryButton">Add embedding to prompt</button>
+                    </td></tr>
                     <tr id="tiling_container" class="pl-5"><td><label for="tiling">Seamless Tiling:</label></td><td>
                         <select id="tiling" name="tiling">
                             <option value="none" selected>None</option>
@@ -562,6 +565,26 @@
         </div>
     </dialog>
 
+    <dialog id="embeddings-dialog">
+        <div id="embeddings-dialog-header" class="dialog-header">
+            <div id="embeddings-dialog-header-left" class="dialog-header-left">
+                <h4>Embeddings</h4>
+                <span>Add embeddings to the prompt (click) or negative prompt (shift-click)</span>
+            </div>
+            <div id="embeddings-dialog-header-right">
+                <i id="embeddings-dialog-close-button" class="fa-solid fa-xmark fa-lg"></i>
+            </div>
+        </div>
+        <div>
+            <i class="fa-solid fa-magnifying-glass"></i>
+            <input id="embeddings-search-box" type="text" spellcheck="false" autocomplete="off" placeholder="Search...">
+            <span style="float:right;"><label>Mode:</label>&nbsp;<select id="embeddings-mode"><option value="insert">Insert at cursor position</option><option value="append">Append at the end</option></select>
+        </div>
+        <div id="embeddings-list">
+        </div>
+        </div>
+    </dialog>
+
     <div id="image-editor" class="popup image-editor-popup">
         <div>
             <i class="close-button fa-solid fa-xmark"></i>

ui/media/css/main.css

@@ -1650,3 +1650,29 @@ body.wait-pause {
     }
 }
 
+#embeddings-button {
+    background-color: var(--background-color3);
+}
+
+#embeddings-button:hover {
+    background-color: var(--button-hover-background);
+}
+
+#embeddings-dialog {
+    overflow: clip;
+}
+
+#embeddings-list {
+    height: 70vH;
+    width: 40vW;
+    overflow-y: scroll;
+}
+
+#embeddings-list button {
+    margin-top: 2px;
+    margin-bottom: 2px;
+}
+
+#embeddings-list::-webkit-scrollbar-thumb {
+  background: var(--background-color3);
+}

ui/media/js/main.js

@@ -116,6 +116,12 @@ let streamImageProgressField = document.querySelector("#stream_image_progress")
 let thumbnailSizeField = document.querySelector("#thumbnail_size-input")
 let autoscrollBtn = document.querySelector("#auto_scroll_btn")
 let autoScroll = document.querySelector("#auto_scroll")
+let embeddingsButton = document.querySelector("#embeddings-button")
+let embeddingsDialog = document.querySelector("#embeddings-dialog")
+let embeddingsDialogCloseBtn = embeddingsDialog.querySelector("#embeddings-dialog-close-button")
+let embeddingsSearchBox = document.querySelector("#embeddings-search-box")
+let embeddingsList = document.querySelector("#embeddings-list")
+let embeddingsModeField = document.querySelector("#embeddings-mode")
 
 let makeImageBtn = document.querySelector("#makeImage")
 let stopImageBtn = document.querySelector("#stopImage")
@@ -2131,6 +2137,79 @@ document.getElementById("toggle-cloudflare-tunnel").addEventListener("click", as
     console.log(`Cloudflare tunnel ${command} result:`, res)
 })
 
+/* Embeddings */
+
+function updateEmbeddingsList(filter="") {
+    function html(model, prefix="", filter="") {
+        filter = filter.toLowerCase()
+        let toplevel=""
+        let folders=""
+       
+        model?.forEach( m => {
+            if (typeof(m) == "string") {
+                if (m.toLowerCase().search(filter)!=-1) {
+                    toplevel += `<button data-embedding="${m}">${m}</button> `
+                }
+            } else {
+                let subdir = html(m[1], prefix+m[0]+"/", filter)
+                if (subdir != "") {
+                   folders += `<h4>${prefix}${m[0]}</h4>` + subdir
+                }
+            }
+        })
+        return toplevel + folders
+    }
+
+    function onButtonClick(e) {
+        let text = e.target.dataset["embedding"]
+        console.log(e.shiftKey, text)
+
+        if (embeddingsModeField.value == "insert") {
+            if (e.shiftKey) {
+                insertAtCursor(negativePromptField, text)
+            } else {
+                insertAtCursor(promptField, text)
+            }
+        } else {
+            let pad=""
+            if (e.shiftKey) {
+                if (!negativePromptField.value.endsWith(" ")) {
+                    pad = " "
+                }
+                negativePromptField.value += pad + text
+            } else {
+                if (!promptField.value.endsWith(" ")) {
+                    pad = " "
+                }
+                promptField.value += pad + text
+            }
+        }
+    }
+
+    embeddingsList.innerHTML = html(modelsOptions.embeddings, "", filter)
+    embeddingsList.querySelectorAll("button").forEach( (b) => { b.addEventListener("click", onButtonClick)})
+}
+
+embeddingsButton.addEventListener("click", () => { 
+    updateEmbeddingsList()
+    embeddingsSearchBox.value=""
+    embeddingsDialog.showModal() 
+})
+embeddingsDialogCloseBtn.addEventListener("click", (e) => {
+    embeddingsDialog.close()
+})
+embeddingsSearchBox.addEventListener("input", (e) => {
+    updateEmbeddingsList(embeddingsSearchBox.value)
+})
+
+modalDialogCloseOnBackdropClick(embeddingsDialog)
+makeDialogDraggable(embeddingsDialog)
+
+
+if (testDiffusers.checked) {
+    document.getElementById("embeddings-container").classList.remove("displayNone")
+}
+
 /* Pause function */
 document.querySelectorAll(".tab").forEach(linkTabContents)
 

ui/media/js/parameters.js

@@ -441,6 +441,7 @@ async function getAppConfig() {
                 option.disabled = true
             })
             document.querySelector("#clip_skip_config").classList.remove("displayNone")
+            document.querySelector("#embeddings-container").classList.remove("displayNone")
         }
 
         console.log("get config status response", config)

ui/media/js/utils.js

@@ -997,6 +997,22 @@ async function getStorageData(key) {
     });
 }
 
+function insertAtCursor(field, text) {
+    if (field.selectionStart || field.selectionStart == "0") {
+        var startPos = field.selectionStart
+        var endPos = field.selectionEnd
+        var before = field.value.substring(0, startPos)
+        var after = field.value.substring(endPos, field.value.length)
+
+        if (!before.endsWith(" ")) { before += " " }
+        if (!after.startsWith(" ")) { after = " "+after }
+
+        field.value = before + text + after
+    } else {
+        field.value += text
+    }
+}
+
 // indexedDB debug functions
 async function getAllKeys() {
     return openDB().then(db => {