From 5a7d645ee2ae43904c78d6096a98c87144629b3a Mon Sep 17 00:00:00 2001 From: Christian Bastian Date: Sun, 21 Jan 2024 07:08:15 -0500 Subject: [PATCH] Added basic auto-suggest dropdown for model directories. --- README.md | 4 +- __init__.py | 63 +++++++++++++++- web/model-manager.css | 25 +++++++ web/model-manager.js | 166 ++++++++++++++++++++++++++++++++++++++---- 4 files changed, 239 insertions(+), 19 deletions(-) diff --git a/README.md b/README.md index 92c2a6b..7115976 100644 --- a/README.md +++ b/README.md @@ -14,7 +14,8 @@ Currently it is still missing some features it should have. - Search bar in models tab. - Advanced keyword search using `"multiple words in quotes"` or a minus sign to `-exclude`. -- Search `/`subdirectories of model directories based on your file structure (for example, `/styles/clothing`). +- Search `/`subdirectories of model directories based on your file structure (for example, `/0/1.5/styles/clothing`). + - Add `/` at the start of the search bar to see auto-complete suggestions. - Include models listed in ComfyUI's `extra_model_paths.yaml` or added in `ComfyUI/models`. - Button to copy a model to the ComfyUI clipboard or embedding to system clipboard. (Embedding copying requires secure http connection.) - Button to add model to ComfyUI graph or embedding to selected nodes. (For small screens/low resolution.) @@ -58,7 +59,6 @@ Currently it is still missing some features it should have. ### Search filtering and sort -- ☐ Add auto-suggest paths in search - ☐ Filters dropdown - ☐ Stable Diffusion model version/Clip/Upscale/? - ☐ Favorites diff --git a/__init__.py b/__init__.py index 6b1e0e2..850df94 100644 --- a/__init__.py +++ b/__init__.py @@ -40,6 +40,8 @@ def folder_paths_folder_names_and_paths(refresh = False) -> dict[str, tuple[list item_path = os.path.join(comfyui_model_uri, item_name) if not os.path.isdir(item_path): continue + if item_name == "configs": + continue if item_name in folder_paths.folder_names_and_paths: dir_paths, extensions = copy.deepcopy(folder_paths.folder_names_and_paths[item_name]) else: @@ -294,7 +296,7 @@ async def load_download_models(request): # TODO: Stop sending redundant information item = { "name": model, - "search-path": os.path.join(model_type, rel_path, model).replace(os.path.sep, "/"), # TODO: Remove hack + "search-path": "/" + os.path.join(model_type, str(base_path_index), rel_path, model).replace(os.path.sep, "/"), # TODO: Remove hack "path": os.path.join(rel_path, model), } if image is not None: @@ -307,6 +309,65 @@ async def load_download_models(request): return web.json_response(models) +def linear_directory_list(refresh = False) -> dict[str, list]: + model_paths = folder_paths_folder_names_and_paths(refresh) + dir_list = [] + dir_list.append({ "name": "", "childIndex": 1, "childCount": len(model_paths) }) + for model_dir_name, (model_dirs, _) in model_paths.items(): + dir_list.append({ "name": model_dir_name, "childIndex": None, "childCount": len(model_dirs) }) + for model_dir_index, (_, (model_dirs, extension_whitelist)) in enumerate(model_paths.items()): + model_dir_child_index = len(dir_list) + dir_list[model_dir_index + 1]["childIndex"] = model_dir_child_index + for dir_path_index, dir_path in enumerate(model_dirs): + dir_list.append({ "name": str(dir_path_index), "childIndex": None, "childCount": None }) + for dir_path_index, dir_path in enumerate(model_dirs): + if not os.path.exists(dir_path) or os.path.isfile(dir_path): + continue + + #dir_list.append({ "name": str(dir_path_index), "childIndex": None, "childCount": 0 }) + dir_stack = [(dir_path, model_dir_child_index + dir_path_index)] + while len(dir_stack) > 0: # DEPTH-FIRST + dir_path, dir_index = dir_stack.pop() + + dir_items = os.listdir(dir_path) + dir_items = sorted(dir_items, key=str.casefold) + + dir_list[dir_index]["childIndex"] = len(dir_list) + dir_child_count = 0 + + # TODO: sort content of directory: alphabetically + # TODO: sort content of directory: files first + + subdirs = [] + for item_name in dir_items: # BREADTH-FIRST + item_path = os.path.join(dir_path, item_name) + if os.path.isdir(item_path): + # dir + subdir_index = len(dir_list) # this must be done BEFORE `dir_list.append` + subdirs.append((item_path, subdir_index)) + dir_list.append({ "name": item_name, "childIndex": None, "childCount": 0 }) + dir_child_count += 1 + else: + # file + _, file_extension = os.path.splitext(item_name) + if extension_whitelist is None or file_extension in extension_whitelist: + dir_list.append({ "name": item_name }) + dir_child_count += 1 + dir_list[dir_index]["childCount"] = dir_child_count + subdirs.reverse() + for dir_path, subdir_index in subdirs: + dir_stack.append((dir_path, subdir_index)) + return dir_list + + +@server.PromptServer.instance.routes.get("/model-manager/directory-list") +async def directory_list(request): + #body = await request.json() + dir_list = linear_directory_list(True) + #json.dump(dir_list, sys.stdout, indent=4) + return web.json_response(dir_list) + + def_headers = { "User-Agent": "Mozilla/5.0 (iPad; CPU OS 12_2 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Mobile/15E148" } diff --git a/web/model-manager.css b/web/model-manager.css index fd9fc87..168e8c3 100644 --- a/web/model-manager.css +++ b/web/model-manager.css @@ -332,12 +332,37 @@ width: 100%; } +.model-manager .search-models { + display: flex; + flex-direction: row; + flex: 1; +} + .model-manager .search-text-area, .model-manager .source-text-area, .model-manager .model-type-dropdown { flex: 1; } +.model-manager .search-dropdown { + position: absolute; + background-color: var(--bg-color); + border: 2px var(--border-color) solid; + top: 45px; + width: 94.7%; + max-height: 30vh; + overflow: auto; + border-radius: 10px; +} + +.search-dropdown:empty { + display: none; +} + +.model-manager .search-dropdown > p { + margin-left: 20px; +} + .model-manager .button-success { color: green; border-color: green; diff --git a/web/model-manager.js b/web/model-manager.js index 45f687e..52ce1a1 100644 --- a/web/model-manager.js +++ b/web/model-manager.js @@ -36,7 +36,10 @@ function modelNodeType(modelType) { else if (modelType === "upscale_models") { return "UpscaleModelLoader"; } else if (modelType === "vae") { return "VAELoader"; } else if (modelType === "vae_approx") { return undefined; } - else { console.warn(`ModelType ${modelType} unrecognized.`); return undefined; } + else { + //console.warn(`ModelType ${modelType} unrecognized.`); + return undefined; + } } function modelWidgetIndex(nodeType) { @@ -424,8 +427,11 @@ class ModelGrid { } static generateInnerHtml(models, modelType, settingsElements) { - const showAddButton = settingsElements["model-show-add-button"].checked; - const showCopyButton = settingsElements["model-show-copy-button"].checked; + // TODO: seperate text and model logic; getting too messy + // TODO: fallback on button failure to copy text? + const canShowButtons = modelNodeType(modelType) !== undefined; + const showAddButton = canShowButtons && settingsElements["model-show-add-button"].checked; + const showCopyButton = canShowButtons && settingsElements["model-show-copy-button"].checked; const strictDragToAdd = settingsElements["model-add-drag-strict-on-field"].checked; const addOffset = parseInt(settingsElements["model-add-offset"].value); const showModelExtension = settingsElements["model-show-label-extensions"].checked; @@ -532,6 +538,7 @@ class ModelManager extends ComfyDialog { modelGrid: null, modelTypeSelect: null, + modelDirectorySearchOptions: null, modelContentFilter: null, sidebarButtons: null, @@ -555,6 +562,8 @@ class ModelManager extends ComfyDialog { #data = { sources: [], models: {}, + modelDirectories: null, + previousModelDirectoryFilter: "", }; /** @type {SourceList} */ @@ -745,6 +754,11 @@ class ModelManager extends ComfyDialog { const modelGrid = $el("div.comfy-grid"); this.#el.modelGrid = modelGrid; + const searchDropdown = $el("div.search-dropdown", { + $: (el) => (this.#el.modelDirectorySearchOptions = el), + style: { display: "none" }, + }); + return [ $el("div.row.tab-header", [ $el("div.row.tab-header-flex-block", [ @@ -753,21 +767,31 @@ class ModelManager extends ComfyDialog { textContent: "⟳", onclick: () => this.#modelGridRefresh(), }), - $el("select.model-type-dropdown", - { - $: (el) => (this.#el.modelTypeSelect = el), - name: "model-type", - onchange: () => this.#modelGridUpdate(), - }, - [], - ), + $el("select.model-type-dropdown", { + $: (el) => (this.#el.modelTypeSelect = el), + name: "model-type", + onchange: () => this.#modelGridUpdate(), + }), ]), $el("div.row.tab-header-flex-block", [ - $el("input.search-text-area", { - $: (el) => (this.#el.modelContentFilter = el), - placeholder: "example: styles/clothing -.pt", - onkeyup: (e) => e.key === "Enter" && this.#modelGridUpdate(), - }), + $el("div.search-models", [ + $el("input.search-text-area", { + $: (el) => (this.#el.modelContentFilter = el), + placeholder: "example: /0/1.5/styles/clothing -.pt", + onkeyup: (e) => e.key === "Enter" && this.#modelGridUpdate(), + oninput: () => this.#updateSearchDropdown(), + onfocus: () => { + if (searchDropdown.innerHTML === "") { + searchDropdown.style.display = "none"; + } + else { + searchDropdown.style.display = "block"; + } + }, + onblur: () => { searchDropdown.style.display = "none"; }, + }), + searchDropdown, + ]), $el("button.icon-button", { type: "button", textContent: "🔍︎", @@ -782,6 +806,7 @@ class ModelManager extends ComfyDialog { #modelGridUpdate() { const models = this.#data.models; const modelSelect = this.#el.modelTypeSelect; + let modelType = modelSelect.value; if (models[modelType] === undefined) { modelType = "checkpoints"; // TODO: magic value @@ -808,6 +833,7 @@ class ModelManager extends ComfyDialog { async #modelGridRefresh() { this.#data.models = await request("/model-manager/models"); + this.#data.modelDirectories = await request("/model-manager/directory-list"); this.#modelGridUpdate(); }; @@ -1003,6 +1029,114 @@ class ModelManager extends ComfyDialog { this.#el.settingsTab = settingsTab; return [settingsTab]; } + + #getFilterDirectory(filter, directory, sep, cwd = 0) { + // TODO: directories === undefined + let filterIndex0 = 1; + while (true) { + const filterIndex1 = filter.indexOf(sep, filterIndex0); + if (filterIndex1 === -1) { + // end of filter + break; + } + + const item = directory[cwd]; + if (item["childCount"] === undefined) { + // file + break; + } + + const childCount = item["childCount"]; + if (childCount === 0) { + // directory is empty + cwd = null; + break; + } + const childIndex = item["childIndex"]; + const items = directory.slice(childIndex, childIndex + childCount); + + const word = filter.substring(filterIndex0, filterIndex1); + cwd = null; + for (let itemIndex = 0; itemIndex < items.length; itemIndex++) { + const itemName = items[itemIndex]["name"]; + if (itemName === word) { + // directory exists + cwd = childIndex + itemIndex; + break; + } + } + if (cwd === null) { + // directory does not exist + break; + } + filterIndex0 = filterIndex1 + 1; + } + return [filterIndex0, cwd]; + } + + async #updateSearchDropdown() { + const directories = this.#data.modelDirectories; + const previousFilter = this.#data.previousModelDirectoryFilter; + const searchDropdown = this.#el.modelDirectorySearchOptions; + const filter = this.#el.modelContentFilter.value; + const modelType = this.#el.modelTypeSelect.value; + + if (previousFilter !== filter) { + let options = []; + const sep = "/"; + if (filter[0] === sep) { + let initCwd = null; + const root = directories[0]; + const rootChildIndex = root["childIndex"]; + const rootChildCount = root["childCount"]; + for (let i = rootChildIndex; i < rootChildIndex + rootChildCount; i++) { + const modelDir = directories[i]; + if (modelDir["name"] === modelType) { + initCwd = i; + break; + } + } + const [filterIndex0, cwd] = this.#getFilterDirectory( + filter, + directories, + sep, + initCwd + ); + if (cwd !== null) { + const lastWord = filter.substring(filterIndex0); + const item = directories[cwd]; + if (item["childIndex"] !== undefined) { + const childIndex = item["childIndex"]; + const childCount = item["childCount"]; + const items = directories.slice(childIndex, childIndex + childCount); + for (let i = 0; i < items.length; i++) { + const itemName = items[i]["name"]; + if (itemName.startsWith(lastWord)) { + options.push(itemName); + } + } + } + else { + const filename = item["name"]; + if (filename.startsWith(lastWord)) { + options.push(filename); + } + } + } + } + + const innerHtml = options.map((text) => { + const el = document.createElement("p"); + el.innerHTML = text; + return el; + }); + searchDropdown.innerHTML = ""; + searchDropdown.append.apply(searchDropdown, innerHtml); + searchDropdown.style.display = options.length === 0 ? "none" : "block"; + } + + this.#data.previousModelDirectoryFilter = filter; + } } let instance;