diff --git a/README.md b/README.md index 63e958a..92c2a6b 100644 --- a/README.md +++ b/README.md @@ -15,7 +15,7 @@ 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`). -- Include models listed in ComfyUI's `extra_model_paths.yaml`. +- 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.) - Right, left, top and bottom toggleable sidebar modes. diff --git a/__init__.py b/__init__.py index 253dd6b..6b1e0e2 100644 --- a/__init__.py +++ b/__init__.py @@ -1,5 +1,6 @@ import os import sys +import copy import hashlib import importlib @@ -30,23 +31,38 @@ image_extensions = (".apng", ".gif", ".jpeg", ".jpg", ".png", ".webp") #hash_buffer_size = 4096 +_folder_names_and_paths: dict[str, tuple[list[str], list[str]]] = None +def folder_paths_folder_names_and_paths(refresh = False) -> dict[str, tuple[list[str], list[str]]]: + global _folder_names_and_paths + if refresh or _folder_names_and_paths is None: + _folder_names_and_paths = {} + for item_name in os.listdir(comfyui_model_uri): + item_path = os.path.join(comfyui_model_uri, item_name) + if not os.path.isdir(item_path): + 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: + dir_paths = [item_name] + extensions = [".ckpt", ".pt", ".bin", ".pth", ".safetensors"] # TODO: magic values + _folder_names_and_paths[item_name] = (dir_paths, extensions) + return _folder_names_and_paths -def folder_paths_get_folder_paths(folder_name): # API function crashes querying unknown model folder - paths = folder_paths.folder_names_and_paths +def folder_paths_get_folder_paths(folder_name, refresh = False) -> list[str]: # API function crashes querying unknown model folder + paths = folder_paths_folder_names_and_paths(refresh) if folder_name in paths: - return paths[folder_name][0][:] + return paths[folder_name][0] maybe_path = os.path.join(comfyui_model_uri, folder_name) if os.path.exists(maybe_path): return [maybe_path] return [] - -def folder_paths_get_supported_pt_extensions(folder_name): # Missing API function - paths = folder_paths.folder_names_and_paths +def folder_paths_get_supported_pt_extensions(folder_name, refresh = False) -> list[str]: # Missing API function + paths = folder_paths_folder_names_and_paths(refresh) if folder_name in paths: return paths[folder_name][1] - return set([".ckpt", ".pt", ".bin", ".pth", ".safetensors"]) + return set([".ckpt", ".pt", ".bin", ".pth", ".safetensors"]) # TODO: magic values def get_safetensor_header(path): diff --git a/web/model-manager.js b/web/model-manager.js index 190a05d..45f687e 100644 --- a/web/model-manager.js +++ b/web/model-manager.js @@ -744,6 +744,7 @@ class ModelManager extends ComfyDialog { #createModelTabHtml() { const modelGrid = $el("div.comfy-grid"); this.#el.modelGrid = modelGrid; + return [ $el("div.row.tab-header", [ $el("div.row.tab-header-flex-block", [ @@ -758,22 +759,7 @@ class ModelManager extends ComfyDialog { name: "model-type", onchange: () => this.#modelGridUpdate(), }, - [ // TODO: generate based on existing model folders - $el("option", ["checkpoints"]), - $el("option", ["clip"]), - $el("option", ["clip_vision"]), - $el("option", ["controlnet"]), - $el("option", ["diffusers"]), - $el("option", ["embeddings"]), - $el("option", ["gligen"]), - $el("option", ["hypernetworks"]), - $el("option", ["loras"]), - $el("option", ["style_models"]), - $el("option", ["unet"]), - $el("option", ["upscale_models"]), - $el("option", ["vae"]), - $el("option", ["vae_approx"]), - ] + [], ), ]), $el("div.row.tab-header-flex-block", [ @@ -794,16 +780,30 @@ 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 + } + + let modelTypeOptions = []; + for (const [key, value] of Object.entries(models)) { + const el = $el("option", [key]); + modelTypeOptions.push(el); + } + modelSelect.innerHTML = ""; + modelTypeOptions.forEach(option => modelSelect.add(option)); + modelSelect.value = modelType; + const searchAppend = this.#el.settings["model-search-always-append"].value; const searchText = this.#el.modelContentFilter.value + " " + searchAppend; - const modelType = this.#el.modelTypeSelect.value; - const models = this.#data.models; const modelList = ModelGrid.filter(models[modelType], searchText); const modelGrid = this.#el.modelGrid; - modelGrid.innerHTML = []; - const innerHTML = ModelGrid.generateInnerHtml(modelList, modelType, this.#el.settings); - modelGrid.append.apply(modelGrid, innerHTML); + modelGrid.innerHTML = ""; + const modelGridModels = ModelGrid.generateInnerHtml(modelList, modelType, this.#el.settings); + modelGrid.append.apply(modelGrid, modelGridModels); }; async #modelGridRefresh() {