From eba7bab15ef9d24da354f624695eb11c0d50cb01 Mon Sep 17 00:00:00 2001
From: cmdr2 <secondary.cmdr2@gmail.com>
Date: Tue, 1 Aug 2023 16:16:38 +0530
Subject: [PATCH] Allow named models in the dropdown

---
 ui/easydiffusion/model_manager.py | 20 ++++++++++++++------
 ui/media/js/searchable-models.js  | 28 +++++++++-------------------
 2 files changed, 23 insertions(+), 25 deletions(-)

diff --git a/ui/easydiffusion/model_manager.py b/ui/easydiffusion/model_manager.py
index 2e7fef67..7a4ddde3 100644
--- a/ui/easydiffusion/model_manager.py
+++ b/ui/easydiffusion/model_manager.py
@@ -290,11 +290,11 @@ def is_malicious_model(file_path):
 def getModels(scan_for_malicious: bool = True):
     models = {
         "options": {
-            "stable-diffusion": ["sd-v1-4"],
+            "stable-diffusion": [{"sd-v1-4": "SD 1.4"}],
             "vae": [],
             "hypernetwork": [],
             "lora": [],
-            "codeformer": ["codeformer"],
+            "codeformer": [{"codeformer": "CodeFormer"}],
             "embeddings": [],
             "controlnet": [],
         },
@@ -305,9 +305,9 @@ def getModels(scan_for_malicious: bool = True):
     class MaliciousModelException(Exception):
         "Raised when picklescan reports a problem with a model"
 
-    def scan_directory(directory, suffixes, directoriesFirst: bool = True):
+    def scan_directory(directory, suffixes, directoriesFirst: bool = True, default_entries=[]):
+        tree = list(default_entries)
         nonlocal models_scanned
-        tree = []
         for entry in sorted(
             os.scandir(directory),
             key=lambda entry: (entry.is_file() == directoriesFirst, entry.name.lower()),
@@ -326,7 +326,14 @@ def getModels(scan_for_malicious: bool = True):
                         raise MaliciousModelException(entry.path)
                 if scan_for_malicious:
                     known_models[entry.path] = mtime
-                tree.append(entry.name[: -len(matching_suffix)])
+                model_id = entry.name[: -len(matching_suffix)]
+                model_exists = False
+                for m in tree:  # allows default "named" models, like CodeFormer and known ControlNet models
+                    if (isinstance(m, str) and model_id == m) or (isinstance(m, dict) and model_id in m):
+                        model_exists = True
+                        break
+                if not model_exists:
+                    tree.append(model_id)
             elif entry.is_dir():
                 scan = scan_directory(entry.path, suffixes, directoriesFirst=False)
 
@@ -343,7 +350,8 @@ def getModels(scan_for_malicious: bool = True):
             os.makedirs(models_dir)
 
         try:
-            models["options"][model_type] = scan_directory(models_dir, model_extensions)
+            default_tree = models["options"].get(model_type, [])
+            models["options"][model_type] = scan_directory(models_dir, model_extensions, default_entries=default_tree)
         except MaliciousModelException as e:
             models["scan-error"] = str(e)
 
diff --git a/ui/media/js/searchable-models.js b/ui/media/js/searchable-models.js
index 174faf77..299c60dc 100644
--- a/ui/media/js/searchable-models.js
+++ b/ui/media/js/searchable-models.js
@@ -552,17 +552,23 @@ class ModelDropdown {
                     this.createModelNodeList(`${folderName || ""}/${childFolderName}`, childModels, false)
                 )
             } else {
+                let modelId = model
+                let modelName = model
+                if (typeof model === "object") {
+                    modelId = Object.keys(model)[0]
+                    modelName = model[modelId]
+                }
                 const classes = ["model-file"]
                 if (isRootFolder) {
                     classes.push("in-root-folder")
                 }
                 // Remove the leading slash from the model path
-                const fullPath = folderName ? `${folderName.substring(1)}/${model}` : model
+                const fullPath = folderName ? `${folderName.substring(1)}/${modelId}` : modelId
                 modelsMap.set(
-                    model,
+                    modelId,
                     createElement("li", { "data-path": fullPath }, classes, [
                         createElement("i", undefined, ["fa-regular", "fa-file", "icon"]),
-                        model,
+                        modelName,
                     ])
                 )
             }
@@ -643,22 +649,6 @@ async function getModels(scanForMalicious = true) {
             makeImageBtn.disabled = true
         }
 
-        /* This code should no longer be needed. Commenting out for now, will cleanup later.
-        const sd_model_setting_key = "stable_diffusion_model"
-        const vae_model_setting_key = "vae_model"
-        const hypernetwork_model_key = "hypernetwork_model"
-
-        const stableDiffusionOptions = modelsOptions['stable-diffusion']
-        const vaeOptions = modelsOptions['vae']
-        const hypernetworkOptions = modelsOptions['hypernetwork']
-
-        // TODO: set default for model here too
-        SETTINGS[sd_model_setting_key].default = stableDiffusionOptions[0]
-        if (getSetting(sd_model_setting_key) == '' || SETTINGS[sd_model_setting_key].value == '') {
-            setSetting(sd_model_setting_key, stableDiffusionOptions[0])
-        }
-        */
-
         // notify ModelDropdown objects to refresh
         document.dispatchEvent(new Event("refreshModels"))
     } catch (e) {