Add back ability to use hashes in black/whitelist

They are displayed in the UI after all, just not in the dropdown but at the bottom
DominikDoom · Jan 14, 2023 · 7fdad1b · 7fdad1b
1 parent a91a098
commit 7fdad1b
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 3 deletions.
diff --git a/javascript/tagAutocomplete.js b/javascript/tagAutocomplete.js
@@ -266,6 +266,7 @@ function hideResults(textArea) {
     selectedTag = null;
 }
 
+var currentModelHash = "";
 var currentModelName = "";
 // Function to check activation criteria
 function isEnabled() {
@@ -275,13 +276,14 @@ function isEnabled() {
             .map(x => x.trim())
             .filter(x => x.length > 0);
 
+        let shortHash = currentModelHash.substring(0, 10);
         if (CFG.activeIn.modelListMode.toLowerCase() === "blacklist") {
             // If the current model is in the blacklist, disable
-            return !modelList.includes(currentModelName);
+            return modelList.filter(x => x === currentModelName || x === currentModelHash || x === shortHash).length === 0;
         } else {
             // If the current model is in the whitelist, enable.
             // An empty whitelist is ignored.
-            return modelList.length === 0 || modelList.includes(currentModelName);
+            return modelList.length === 0 || modelList.filter(x => x === currentModelName || x === currentModelHash || x === shortHash).length > 0;
         }
     } else {
         return false;
@@ -1100,6 +1102,7 @@ async function setup() {
             }, 500);
         });
     });
+
     // Add change listener to model dropdown to react to model changes
     let modelDropdown = gradioApp().querySelector("#setting_sd_model_checkpoint select");
     currentModelName = modelDropdown.value;
@@ -1108,6 +1111,19 @@ async function setup() {
             currentModelName = modelDropdown.value;
         }, 100);
     });
+    // Add mutation observer for the model hash text to also allow hash-based blacklist again
+    let modelHashText = gradioApp().querySelector("#sd_checkpoint_hash");
+    if (modelHashText) {
+        currentModelHash = modelHashText.title
+        let modelHashObserver = new MutationObserver((mutationList, observer) => {
+            for (const mutation of mutationList) {
+                if (mutation.type === "attributes" && mutation.attributeName === "title") {
+                    currentModelHash = mutation.target.title;
+                }
+            }
+        });
+        modelHashObserver.observe(modelHashText, { attributes: true });
+    }
 
     // Not found, we're on a page without prompt textareas
     if (textAreas.every(v => v === null || v === undefined)) return;

diff --git a/scripts/tag_autocomplete_helper.py b/scripts/tag_autocomplete_helper.py
@@ -215,7 +215,7 @@ def on_ui_settings():
     shared.opts.add_option("tac_activeIn.img2img", shared.OptionInfo(True, "Active in img2img (Requires restart)", section=TAC_SECTION))
     shared.opts.add_option("tac_activeIn.negativePrompts", shared.OptionInfo(True, "Active in negative prompts (Requires restart)", section=TAC_SECTION))
     shared.opts.add_option("tac_activeIn.thirdParty", shared.OptionInfo(True, "Active in third party textboxes [Dataset Tag Editor] (Requires restart)", section=TAC_SECTION))
-    shared.opts.add_option("tac_activeIn.modelList", shared.OptionInfo("", "List of model names (with file extension) to use as black/whitelist, separated by commas.", section=TAC_SECTION))
+    shared.opts.add_option("tac_activeIn.modelList", shared.OptionInfo("", "List of model names (with file extension) or their hashes to use as black/whitelist, separated by commas.", section=TAC_SECTION))
     shared.opts.add_option("tac_activeIn.modelListMode", shared.OptionInfo("Blacklist", "Mode to use for model list", gr.Dropdown, lambda: {"choices": ["Blacklist","Whitelist"]}, section=TAC_SECTION))
     # Results related settings
     shared.opts.add_option("tac_maxResults", shared.OptionInfo(5, "Maximum results", section=TAC_SECTION))