View models in sub-directories of ComfyUI/models.

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.

__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):

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() {