From 6a158d46b87f399bcc7b35af6b26c1ab53f1d1af Mon Sep 17 00:00:00 2001 From: "Alex \"mcmonkey\" Goodwin" <4000772+mcmonkey4eva@users.noreply.github.com> Date: Tue, 24 Sep 2024 09:48:15 +0900 Subject: [PATCH] [Draft] Model library sidebar tab (#837) * basic/empty model library sidebar tab in-progress * make it actually list out models * extremely primitive search impl * list out available folders (incomplete list atm) * load list dynamically * nice lil loading icon * that's not doing anything * run autoformatter * fix up some absolute vue shenanigans * swap to pi-box * is_fake_object * i think apply the tailwind thingo * trim '.safetensors' from end of display title * oop * after load, retain title if no new title is given * is_load_requested to prevent duplication * dirty initial model metadata load & preview based on node preview code * update model store tests * initial image icon for model lib * i hate this * better empty spacer * add api handler for '/models' * load model folders list instead of hardcoding * add a 'no content' placeholder for empty folders * autoformat * autoload model metadata * error handling on metadata loading * larger model icons * click a model to spawn a node for it * draggable model nodes * add a setting for whether to autoload or not * autoformat will be the death of me * cleanup promise code * make the model preview actually half-decent * revert bad unchecked change * put registration back --- .../common/TreeExplorerTreeNode.vue | 4 +- src/components/graph/GraphCanvas.vue | 20 +- .../sidebar/tabs/ModelLibrarySidebarTab.vue | 209 ++++++++++++++++++ .../tabs/modelLibrary/ModelPreview.vue | 89 ++++++++ .../tabs/modelLibrary/ModelTreeLeaf.vue | 130 +++++++++++ src/i18n.ts | 4 + src/scripts/api.ts | 29 ++- src/stores/coreSettings.ts | 8 + src/stores/modelStore.ts | 122 ++++++---- src/stores/modelToNodeStore.ts | 86 +++++++ src/views/GraphView.vue | 9 + tests-ui/tests/store/modelStore.test.ts | 3 +- 12 files changed, 663 insertions(+), 50 deletions(-) create mode 100644 src/components/sidebar/tabs/ModelLibrarySidebarTab.vue create mode 100644 src/components/sidebar/tabs/modelLibrary/ModelPreview.vue create mode 100644 src/components/sidebar/tabs/modelLibrary/ModelTreeLeaf.vue create mode 100644 src/stores/modelToNodeStore.ts diff --git a/src/components/common/TreeExplorerTreeNode.vue b/src/components/common/TreeExplorerTreeNode.vue index b8da42bb0..82e5ba04d 100644 --- a/src/components/common/TreeExplorerTreeNode.vue +++ b/src/components/common/TreeExplorerTreeNode.vue @@ -10,8 +10,8 @@ ]" ref="container" > -
- +
+ (null) @@ -49,7 +51,7 @@ const settingStore = useSettingStore() const nodeDefStore = useNodeDefStore() const workspaceStore = useWorkspaceStore() const canvasStore = useCanvasStore() - +const modelToNodeStore = useModelToNodeStore() const betaMenuEnabled = computed( () => settingStore.get('Comfy.UseNewMenu') !== 'Disabled' ) @@ -132,6 +134,22 @@ onMounted(async () => { loc.clientY ]) comfyApp.addNodeOnGraph(nodeDef, { pos }) + } else if (node.data instanceof ComfyModelDef) { + const model = node.data + const provider = modelToNodeStore.getNodeProvider(model.directory) + if (provider) { + const pos = comfyApp.clientPosToCanvasPos([ + loc.clientX - 20, + loc.clientY + ]) + const node = comfyApp.addNodeOnGraph(provider.nodeDef, { pos }) + const widget = node.widgets.find( + (widget) => widget.name === provider.key + ) + if (widget) { + widget.value = model.name + } + } } } } diff --git a/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue b/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue new file mode 100644 index 000000000..a0a24d3a0 --- /dev/null +++ b/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue @@ -0,0 +1,209 @@ + + + + + + + diff --git a/src/components/sidebar/tabs/modelLibrary/ModelPreview.vue b/src/components/sidebar/tabs/modelLibrary/ModelPreview.vue new file mode 100644 index 000000000..5b02843c8 --- /dev/null +++ b/src/components/sidebar/tabs/modelLibrary/ModelPreview.vue @@ -0,0 +1,89 @@ + + + + diff --git a/src/components/sidebar/tabs/modelLibrary/ModelTreeLeaf.vue b/src/components/sidebar/tabs/modelLibrary/ModelTreeLeaf.vue new file mode 100644 index 000000000..58957dda0 --- /dev/null +++ b/src/components/sidebar/tabs/modelLibrary/ModelTreeLeaf.vue @@ -0,0 +1,130 @@ + + + + + diff --git a/src/i18n.ts b/src/i18n.ts index 0ed171794..17bc1c4e1 100644 --- a/src/i18n.ts +++ b/src/i18n.ts @@ -21,6 +21,7 @@ const messages = { box: 'Box', briefcase: 'Briefcase', error: 'Error', + loading: 'Loading', findIssues: 'Find Issues', copyToClipboard: 'Copy to Clipboard', openNewIssue: 'Open New Issue', @@ -39,9 +40,11 @@ const messages = { searchWorkflows: 'Search Workflows', searchSettings: 'Search Settings', searchNodes: 'Search Nodes', + searchModels: 'Search Models', noResultsFound: 'No Results Found', searchFailedMessage: "We couldn't find any settings matching your search. Try adjusting your search terms.", + noContent: '(No Content)', noTasksFound: 'No Tasks Found', noTasksFoundMessage: 'There are no tasks in the queue.', newFolder: 'New Folder', @@ -54,6 +57,7 @@ const messages = { nodeLibraryTab: { sortOrder: 'Sort Order' }, + modelLibrary: 'Model Library', queueTab: { showFlatList: 'Show Flat List', backToAllTasks: 'Back to All Tasks', diff --git a/src/scripts/api.ts b/src/scripts/api.ts index f83cb488b..de250bc79 100644 --- a/src/scripts/api.ts +++ b/src/scripts/api.ts @@ -330,6 +330,18 @@ class ComfyApi extends EventTarget { return await res.json() } + /** + * Gets a list of model folder keys (eg ['checkpoints', 'loras', ...]) + * @returns The list of model folder keys + */ + async getModelFolders(): Promise { + const res = await this.fetchApi(`/models`) + if (res.status === 404) { + return null + } + return await res.json() + } + /** * Gets a list of models in the specified folder * @param {string} folder The folder to list models from, such as 'checkpoints' @@ -353,7 +365,22 @@ class ComfyApi extends EventTarget { const res = await this.fetchApi( `/view_metadata/${folder}?filename=${encodeURIComponent(model)}` ) - return await res.json() + const rawResponse = await res.text() + if (!rawResponse) { + return null + } + try { + return JSON.parse(rawResponse) + } catch (error) { + console.error( + 'Error viewing metadata', + res.status, + res.statusText, + rawResponse, + error + ) + return null + } } /** diff --git a/src/stores/coreSettings.ts b/src/stores/coreSettings.ts index 9270d0dd5..4f27c0aca 100644 --- a/src/stores/coreSettings.ts +++ b/src/stores/coreSettings.ts @@ -254,6 +254,14 @@ export const CORE_SETTINGS: SettingParams[] = [ step: 1 } }, + { + id: 'Comfy.ModelLibrary.AutoLoadAll', + name: 'Automatically load all model folders', + tooltip: + 'If true, all folders will load as soon as you open the model library (this may cause delays while it loads). If false, root level model folders will only load once you click on them.', + type: 'boolean', + defaultValue: false + }, { id: 'Comfy.Locale', name: 'Locale', diff --git a/src/stores/modelStore.ts b/src/stores/modelStore.ts index 490bad010..5b76c54d6 100644 --- a/src/stores/modelStore.ts +++ b/src/stores/modelStore.ts @@ -42,54 +42,71 @@ export class ComfyModelDef { image: string = '' /** Whether the model metadata has been loaded from the server, used for `load()` */ has_loaded_metadata: boolean = false + /** If true, a metadata load request has been triggered, but may or may not yet have finished loading */ + is_load_requested: boolean = false + /** If true, this is a fake model object used as a placeholder for something (eg a loading icon) */ + is_fake_object: boolean = false constructor(name: string, directory: string) { this.name = name - this.title = name + this.title = name.replaceAll('\\', '/').split('/').pop() + if (this.title.endsWith('.safetensors')) { + this.title = this.title.slice(0, -'.safetensors'.length) + } this.directory = directory } /** Loads the model metadata from the server, filling in this object if data is available */ async load(): Promise { - if (this.has_loaded_metadata) { + if (this.has_loaded_metadata || this.is_load_requested) { return } - const metadata = await api.viewMetadata(this.directory, this.name) - if (!metadata) { - return + this.is_load_requested = true + try { + const metadata = await api.viewMetadata(this.directory, this.name) + if (!metadata) { + return + } + this.title = + _findInMetadata( + metadata, + 'modelspec.title', + 'title', + 'display_name', + 'name' + ) || this.title + this.architecture_id = + _findInMetadata(metadata, 'modelspec.architecture', 'architecture') || + '' + this.author = + _findInMetadata(metadata, 'modelspec.author', 'author') || '' + this.description = + _findInMetadata(metadata, 'modelspec.description', 'description') || '' + this.resolution = + _findInMetadata(metadata, 'modelspec.resolution', 'resolution') || '' + this.usage_hint = + _findInMetadata(metadata, 'modelspec.usage_hint', 'usage_hint') || '' + this.trigger_phrase = + _findInMetadata( + metadata, + 'modelspec.trigger_phrase', + 'trigger_phrase' + ) || '' + this.image = + _findInMetadata( + metadata, + 'modelspec.thumbnail', + 'thumbnail', + 'image', + 'icon' + ) || '' + const tagsCommaSeparated = + _findInMetadata(metadata, 'modelspec.tags', 'tags') || '' + this.tags = tagsCommaSeparated.split(',').map((tag) => tag.trim()) + this.has_loaded_metadata = true + } catch (error) { + console.error('Error loading model metadata', this.name, this, error) } - this.title = - _findInMetadata( - metadata, - 'modelspec.title', - 'title', - 'display_name', - 'name' - ) || this.name - this.architecture_id = - _findInMetadata(metadata, 'modelspec.architecture', 'architecture') || '' - this.author = _findInMetadata(metadata, 'modelspec.author', 'author') || '' - this.description = - _findInMetadata(metadata, 'modelspec.description', 'description') || '' - this.resolution = - _findInMetadata(metadata, 'modelspec.resolution', 'resolution') || '' - this.usage_hint = - _findInMetadata(metadata, 'modelspec.usage_hint', 'usage_hint') || '' - this.trigger_phrase = - _findInMetadata(metadata, 'modelspec.trigger_phrase', 'trigger_phrase') || - '' - this.image = - _findInMetadata( - metadata, - 'modelspec.thumbnail', - 'thumbnail', - 'image', - 'icon' - ) || '' - const tagsCommaSeparated = - _findInMetadata(metadata, 'modelspec.tags', 'tags') || '' - this.tags = tagsCommaSeparated.split(',').map((tag) => tag.trim()) - this.has_loaded_metadata = true } } @@ -110,27 +127,42 @@ export class ModelStore { } } +const folderBlacklist = ['configs', 'custom_nodes'] + /** Model store handler, wraps individual per-folder model stores */ export const useModelStore = defineStore('modelStore', { state: () => ({ - modelStoreMap: {} as Record + modelStoreMap: {} as Record, + isLoading: {} as Record>, + modelFolders: [] as string[] }), actions: { async getModelsInFolderCached(folder: string): Promise { if (folder in this.modelStoreMap) { return this.modelStoreMap[folder] } - // TODO: needs a lock to avoid overlapping calls - const models = await api.getModels(folder) - if (!models) { - return null + if (this.isLoading[folder]) { + return this.isLoading[folder] } - const store = new ModelStore(folder, models) - this.modelStoreMap[folder] = store - return store + const promise = api.getModels(folder).then((models) => { + if (!models) { + return null + } + const store = new ModelStore(folder, models) + this.modelStoreMap[folder] = store + this.isLoading[folder] = false + return store + }) + this.isLoading[folder] = promise + return promise }, clearCache() { this.modelStoreMap = {} + }, + async getModelFolders() { + this.modelFolders = (await api.getModelFolders()).filter( + (folder) => !folderBlacklist.includes(folder) + ) } } }) diff --git a/src/stores/modelToNodeStore.ts b/src/stores/modelToNodeStore.ts new file mode 100644 index 000000000..038d5ec6a --- /dev/null +++ b/src/stores/modelToNodeStore.ts @@ -0,0 +1,86 @@ +import { ComfyNodeDefImpl } from '@/stores/nodeDefStore' +import { useNodeDefStore } from '@/stores/nodeDefStore' +import { defineStore } from 'pinia' +import { toRaw } from 'vue' + +/** Helper class that defines how to construct a node from a model. */ +export class ModelNodeProvider { + /** The node definition to use for this model. */ + public nodeDef: ComfyNodeDefImpl + + /** The node input key for where to inside the model name. */ + public key: string + + constructor(nodeDef: ComfyNodeDefImpl, key: string) { + this.nodeDef = nodeDef + this.key = key + } +} + +/** Service for mapping model types (by folder name) to nodes. */ +export const useModelToNodeStore = defineStore('modelToNode', { + state: () => ({ + modelToNodeMap: {} as Record, + nodeDefStore: useNodeDefStore(), + haveDefaultsLoaded: false + }), + actions: { + /** + * Get the node provider for the given model type name. + * @param modelType The name of the model type to get the node provider for. + * @returns The node provider for the given model type name. + */ + getNodeProvider(modelType: string): ModelNodeProvider { + this.registerDefaults() + return this.modelToNodeMap[modelType] + }, + + /** + * Register a node provider for the given model type name. + * @param modelType The name of the model type to register the node provider for. + * @param nodeProvider The node provider to register. + */ + registerNodeProvider(modelType: string, nodeProvider: ModelNodeProvider) { + this.registerDefaults() + this.modelToNodeMap[modelType] = nodeProvider + }, + + registerDefaults() { + if (this.haveDefaultsLoaded) { + return + } + if (Object.keys(this.nodeDefStore.nodeDefsByName).length === 0) { + return + } + this.haveDefaultsLoaded = true + this.registerNodeProvider( + 'checkpoints', + new ModelNodeProvider( + this.nodeDefStore.nodeDefsByName['CheckpointLoaderSimple'], + 'ckpt_name' + ) + ) + this.registerNodeProvider( + 'loras', + new ModelNodeProvider( + this.nodeDefStore.nodeDefsByName['LoraLoader'], + 'lora_name' + ) + ) + this.registerNodeProvider( + 'vae', + new ModelNodeProvider( + this.nodeDefStore.nodeDefsByName['VAELoader'], + 'vae_name' + ) + ) + this.registerNodeProvider( + 'controlnet', + new ModelNodeProvider( + this.nodeDefStore.nodeDefsByName['ControlNetLoader'], + 'control_net_name' + ) + ) + } + } +}) diff --git a/src/views/GraphView.vue b/src/views/GraphView.vue index 9fa92f20c..1a492e7d1 100644 --- a/src/views/GraphView.vue +++ b/src/views/GraphView.vue @@ -34,6 +34,7 @@ import { } from '@/stores/workflowStore' import QueueSidebarTab from '@/components/sidebar/tabs/QueueSidebarTab.vue' import NodeLibrarySidebarTab from '@/components/sidebar/tabs/NodeLibrarySidebarTab.vue' +import ModelLibrarySidebarTab from '@/components/sidebar/tabs/ModelLibrarySidebarTab.vue' import GlobalToast from '@/components/toast/GlobalToast.vue' import UnloadWindowConfirmDialog from '@/components/dialog/UnloadWindowConfirmDialog.vue' import BrowserTabTitle from '@/components/BrowserTabTitle.vue' @@ -126,6 +127,14 @@ const init = () => { component: markRaw(NodeLibrarySidebarTab), type: 'vue' }) + app.extensionManager.registerSidebarTab({ + id: 'model-library', + icon: 'pi pi-box', + title: t('sideToolbar.modelLibrary'), + tooltip: t('sideToolbar.modelLibrary'), + component: markRaw(ModelLibrarySidebarTab), + type: 'vue' + }) app.extensionManager.registerSidebarTab({ id: 'workflows', icon: 'pi pi-folder-open', diff --git a/tests-ui/tests/store/modelStore.test.ts b/tests-ui/tests/store/modelStore.test.ts index d488f2ddd..8d6d5c2ae 100644 --- a/tests-ui/tests/store/modelStore.test.ts +++ b/tests-ui/tests/store/modelStore.test.ts @@ -69,7 +69,8 @@ describe('useModelStore', () => { const folderStore = await store.getModelsInFolderCached('checkpoints') const model = folderStore.models['noinfo.safetensors'] await model.load() - expect(model.title).toBe('noinfo.safetensors') + expect(model.name).toBe('noinfo.safetensors') + expect(model.title).toBe('noinfo') expect(model.architecture_id).toBe('') expect(model.author).toBe('') expect(model.description).toBe('')