From 090fda2f223f9229299c8c4989915f6d0553a8b6 Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Sun, 25 Aug 2024 21:53:58 -0400 Subject: [PATCH] Node library custom bookmark folder (#631) * Add new folder button * Add tree util test * nit * Support empty folder in node library * Drag to bookmark folder * Use bookmark icon for bookmark folder * Highlight on dragover * nit * Auto-expand on item added * Extract bookmark system as store * Add context menu on bookmark folder * Add editable text * Fix reactivity * Plumb editable text * refactor * Rename node * Fix focus * Prevent name collision * nit * Add new folder * nested folder support * Change drag behavior * Add basic playwright tests * nit * Target tree-node-content instead of tree-node --- browser_tests/ComfyPage.ts | 24 +- browser_tests/menu.spec.ts | 113 +++++++++- src/components/common/EditableText.vue | 81 +++++++ .../sidebar/tabs/NodeLibrarySidebarTab.vue | 208 +++++++++++++----- .../tabs/nodeLibrary/NodeTreeFolder.vue | 90 ++++++++ src/i18n.ts | 8 + src/stores/nodeBookmarkStore.ts | 137 ++++++++++++ src/stores/nodeDefStore.ts | 25 ++- src/utils/treeUtil.ts | 25 ++- tests-ui/tests/utils/treeUtilTest.test.ts | 63 ++++++ 10 files changed, 692 insertions(+), 82 deletions(-) create mode 100644 src/components/common/EditableText.vue create mode 100644 src/components/sidebar/tabs/nodeLibrary/NodeTreeFolder.vue create mode 100644 src/stores/nodeBookmarkStore.ts create mode 100644 tests-ui/tests/utils/treeUtilTest.test.ts diff --git a/browser_tests/ComfyPage.ts b/browser_tests/ComfyPage.ts index 1e8772ef..6413edf7 100644 --- a/browser_tests/ComfyPage.ts +++ b/browser_tests/ComfyPage.ts @@ -65,6 +65,14 @@ class NodeLibrarySidebarTab { return this.page.locator('.node-lib-node-preview') } + get tabContainer() { + return this.page.locator('.sidebar-content-container') + } + + get newFolderButton() { + return this.tabContainer.locator('.new-folder-button') + } + async open() { if (await this.selectedTabButton.isVisible()) { return @@ -74,16 +82,20 @@ class NodeLibrarySidebarTab { await this.nodeLibraryTree.waitFor({ state: 'visible' }) } + folderSelector(folderName: string) { + return `.p-tree-node-content:has(> .node-lib-tree-node-label:has(.folder-label:has-text("${folderName}")))` + } + getFolder(folderName: string) { - return this.page.locator( - `.p-tree-node-content:has(> .node-lib-tree-node-label:has(.folder-label:has-text("${folderName}")))` - ) + return this.page.locator(this.folderSelector(folderName)) + } + + nodeSelector(nodeName: string) { + return `.p-tree-node-content:has(> .node-lib-tree-node-label:has(.node-label:has-text("${nodeName}")))` } getNode(nodeName: string) { - return this.page.locator( - `.p-tree-node-content:has(> .node-lib-tree-node-label:has(.node-label:has-text("${nodeName}")))` - ) + return this.page.locator(this.nodeSelector(nodeName)) } } diff --git a/browser_tests/menu.spec.ts b/browser_tests/menu.spec.ts index a3eba0ae..3048694f 100644 --- a/browser_tests/menu.spec.ts +++ b/browser_tests/menu.spec.ts @@ -60,10 +60,15 @@ test.describe('Menu', () => { }) test.describe('Node library sidebar', () => { - test('Node preview and drag to canvas', async ({ comfyPage }) => { + test.beforeEach(async ({ comfyPage }) => { + await comfyPage.setSetting('Comfy.NodeLibrary.Bookmarks', []) // Open the sidebar const tab = comfyPage.menu.nodeLibraryTab await tab.open() + }) + + test('Node preview and drag to canvas', async ({ comfyPage }) => { + const tab = comfyPage.menu.nodeLibraryTab await tab.getFolder('sampling').click() // Hover over a node to display the preview @@ -86,11 +91,7 @@ test.describe('Menu', () => { }) test('Bookmark node', async ({ comfyPage }) => { - await comfyPage.setSetting('Comfy.NodeLibrary.Bookmarks', []) - - // Open the sidebar const tab = comfyPage.menu.nodeLibraryTab - await tab.open() await tab.getFolder('sampling').click() // Bookmark the node @@ -116,13 +117,107 @@ test.describe('Menu', () => { test('Ignores unrecognized node', async ({ comfyPage }) => { await comfyPage.setSetting('Comfy.NodeLibrary.Bookmarks', ['foo']) - // Open the sidebar const tab = comfyPage.menu.nodeLibraryTab - await tab.open() - expect(await tab.getFolder('sampling').count()).toBe(1) expect(await tab.getNode('foo').count()).toBe(0) - await comfyPage.setSetting('Comfy.NodeLibrary.Bookmarks', []) + }) + + test('Displays empty bookmarks folder', async ({ comfyPage }) => { + await comfyPage.setSetting('Comfy.NodeLibrary.Bookmarks', ['foo/']) + const tab = comfyPage.menu.nodeLibraryTab + expect(await tab.getFolder('foo').count()).toBe(1) + }) + + test('Can add new bookmark folder', async ({ comfyPage }) => { + const tab = comfyPage.menu.nodeLibraryTab + await tab.newFolderButton.click() + await comfyPage.page.keyboard.press('Enter') + expect(await tab.getFolder('New Folder').count()).toBe(1) + expect(await comfyPage.getSetting('Comfy.NodeLibrary.Bookmarks')).toEqual( + ['New Folder/'] + ) + }) + + test('Can add nested bookmark folder', async ({ comfyPage }) => { + await comfyPage.setSetting('Comfy.NodeLibrary.Bookmarks', ['foo/']) + const tab = comfyPage.menu.nodeLibraryTab + + await tab.getFolder('foo').click({ button: 'right' }) + await comfyPage.page.getByLabel('New Folder').click() + await comfyPage.page.keyboard.press('Enter') + + expect(await comfyPage.getSetting('Comfy.NodeLibrary.Bookmarks')).toEqual( + ['foo/', 'foo/New Folder/'] + ) + }) + + test('Can delete bookmark folder', async ({ comfyPage }) => { + await comfyPage.setSetting('Comfy.NodeLibrary.Bookmarks', ['foo/']) + const tab = comfyPage.menu.nodeLibraryTab + + await tab.getFolder('foo').click({ button: 'right' }) + await comfyPage.page.getByLabel('Delete').click() + + expect(await comfyPage.getSetting('Comfy.NodeLibrary.Bookmarks')).toEqual( + [] + ) + }) + + test('Can rename bookmark folder', async ({ comfyPage }) => { + await comfyPage.setSetting('Comfy.NodeLibrary.Bookmarks', ['foo/']) + const tab = comfyPage.menu.nodeLibraryTab + + await tab.getFolder('foo').click({ button: 'right' }) + await comfyPage.page.getByLabel('Rename').click() + await comfyPage.page.keyboard.insertText('bar') + await comfyPage.page.keyboard.press('Enter') + + expect(await comfyPage.getSetting('Comfy.NodeLibrary.Bookmarks')).toEqual( + ['bar/'] + ) + }) + + test('Can add bookmark by dragging node to bookmark folder', async ({ + comfyPage + }) => { + await comfyPage.setSetting('Comfy.NodeLibrary.Bookmarks', ['foo/']) + const tab = comfyPage.menu.nodeLibraryTab + await tab.getFolder('sampling').click() + await comfyPage.page.dragAndDrop( + tab.nodeSelector('KSampler (Advanced)'), + tab.folderSelector('foo') + ) + expect(await comfyPage.getSetting('Comfy.NodeLibrary.Bookmarks')).toEqual( + ['foo/', 'foo/KSampler (Advanced)'] + ) + }) + + test('Can add bookmark by clicking bookmark button', async ({ + comfyPage + }) => { + const tab = comfyPage.menu.nodeLibraryTab + await tab.getFolder('sampling').click() + await tab + .getNode('KSampler (Advanced)') + .locator('.bookmark-button') + .click() + expect(await comfyPage.getSetting('Comfy.NodeLibrary.Bookmarks')).toEqual( + ['KSampler (Advanced)'] + ) + }) + + test('Can unbookmark node', async ({ comfyPage }) => { + await comfyPage.setSetting('Comfy.NodeLibrary.Bookmarks', [ + 'KSampler (Advanced)' + ]) + const tab = comfyPage.menu.nodeLibraryTab + await tab + .getNode('KSampler (Advanced)') + .locator('.bookmark-button') + .click() + expect(await comfyPage.getSetting('Comfy.NodeLibrary.Bookmarks')).toEqual( + [] + ) }) }) diff --git a/src/components/common/EditableText.vue b/src/components/common/EditableText.vue new file mode 100644 index 00000000..942c5607 --- /dev/null +++ b/src/components/common/EditableText.vue @@ -0,0 +1,81 @@ + + + + + diff --git a/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue b/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue index ed85fb23..ab05d6e2 100644 --- a/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue +++ b/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue @@ -2,6 +2,15 @@ + diff --git a/src/components/sidebar/tabs/nodeLibrary/NodeTreeFolder.vue b/src/components/sidebar/tabs/nodeLibrary/NodeTreeFolder.vue new file mode 100644 index 00000000..bd91566d --- /dev/null +++ b/src/components/sidebar/tabs/nodeLibrary/NodeTreeFolder.vue @@ -0,0 +1,90 @@ + + + diff --git a/src/i18n.ts b/src/i18n.ts index ba62f5c1..2ff2ff61 100644 --- a/src/i18n.ts +++ b/src/i18n.ts @@ -2,6 +2,7 @@ import { createI18n } from 'vue-i18n' const messages = { en: { + error: 'Error', findIssues: 'Find Issues', copyToClipboard: 'Copy to Clipboard', openNewIssue: 'Open New Issue', @@ -10,6 +11,8 @@ const messages = { reconnecting: 'Reconnecting', reconnected: 'Reconnected', delete: 'Delete', + rename: 'Rename', + customize: 'Customize', experimental: 'BETA', deprecated: 'DEPR', loadWorkflow: 'Load Workflow', @@ -21,6 +24,7 @@ const messages = { "We couldn't find any settings matching your search. Try adjusting your search terms.", noTasksFound: 'No Tasks Found', noTasksFoundMessage: 'There are no tasks in the queue.', + newFolder: 'New Folder', sideToolbar: { themeToggle: 'Toggle Theme', queue: 'Queue', @@ -35,11 +39,14 @@ const messages = { } }, zh: { + error: '错误', showReport: '显示报告', imageFailedToLoad: '图像加载失败', reconnecting: '重新连接中', reconnected: '已重新连接', delete: '删除', + rename: '重命名', + customize: '定制', loadWorkflow: '加载工作流', settings: '设置', searchSettings: '搜索设置', @@ -49,6 +56,7 @@ const messages = { noTasksFoundMessage: '队列中没有任务。', searchFailedMessage: '我们找不到与您的搜索匹配的任何设置。请尝试调整搜索条件。', + newFolder: '新建文件夹', sideToolbar: { themeToggle: '主题切换', queue: '队列', diff --git a/src/stores/nodeBookmarkStore.ts b/src/stores/nodeBookmarkStore.ts new file mode 100644 index 00000000..58d9a614 --- /dev/null +++ b/src/stores/nodeBookmarkStore.ts @@ -0,0 +1,137 @@ +import { defineStore } from 'pinia' +import { computed } from 'vue' +import { useSettingStore } from './settingStore' +import { useNodeDefStore } from './nodeDefStore' +import { ComfyNodeDefImpl, createDummyFolderNodeDef } from './nodeDefStore' +import { buildNodeDefTree } from './nodeDefStore' +import type { TreeNode } from 'primevue/treenode' +import _ from 'lodash' + +export const useNodeBookmarkStore = defineStore('nodeBookmark', () => { + const settingStore = useSettingStore() + const nodeDefStore = useNodeDefStore() + + const bookmarks = computed(() => + settingStore.get('Comfy.NodeLibrary.Bookmarks') + ) + + const bookmarksSet = computed>(() => new Set(bookmarks.value)) + + const bookmarkedRoot = computed(() => + buildBookmarkTree(bookmarks.value) + ) + + // For a node in custom bookmark folders, check if its nodePath is in bookmarksSet + // For a node in the nodeDefStore, check if its name is bookmarked at top level + const isBookmarked = (node: ComfyNodeDefImpl) => + bookmarksSet.value.has(node.nodePath) || + bookmarksSet.value.has(node.display_name) + + const toggleBookmark = (node: ComfyNodeDefImpl) => { + if (isBookmarked(node)) { + deleteBookmark(node.nodePath) + } else { + addBookmark(node.display_name) + } + } + + const buildBookmarkTree = (bookmarks: string[]) => { + const bookmarkNodes = bookmarks + .map((bookmark: string) => { + if (bookmark.endsWith('/')) return createDummyFolderNodeDef(bookmark) + + const parts = bookmark.split('/') + const displayName = parts.pop() + const category = parts.join('/') + const srcNodeDef = nodeDefStore.nodeDefsByDisplayName[displayName] + if (!srcNodeDef) { + return null + } + const nodeDef = _.clone(srcNodeDef) + nodeDef.category = category + return nodeDef + }) + .filter((nodeDef) => nodeDef !== null) + return buildNodeDefTree(bookmarkNodes) + } + + const addBookmark = (nodePath: string) => { + settingStore.set('Comfy.NodeLibrary.Bookmarks', [ + ...bookmarks.value, + nodePath + ]) + } + + const deleteBookmark = (nodePath: string) => { + settingStore.set( + 'Comfy.NodeLibrary.Bookmarks', + bookmarks.value.filter((b: string) => b !== nodePath) + ) + } + + const addNewBookmarkFolder = (parent?: ComfyNodeDefImpl) => { + const parentPath = parent ? parent.nodePath : '' + let newFolderPath = parentPath + 'New Folder/' + let suffix = 1 + while (bookmarks.value.some((b: string) => b.startsWith(newFolderPath))) { + newFolderPath = parentPath + `New Folder ${suffix}/` + suffix++ + } + addBookmark(newFolderPath) + return newFolderPath + } + + const renameBookmarkFolder = ( + folderNode: ComfyNodeDefImpl, + newName: string + ) => { + if (!folderNode.isDummyFolder) { + throw new Error('Cannot rename non-folder node') + } + + const newNodePath = + folderNode.category.split('/').slice(0, -1).concat(newName).join('/') + + '/' + + if (newNodePath === folderNode.nodePath) { + return + } + + if (bookmarks.value.some((b: string) => b.startsWith(newNodePath))) { + throw new Error(`Folder name "${newNodePath}" already exists`) + } + + settingStore.set( + 'Comfy.NodeLibrary.Bookmarks', + bookmarks.value.map((b: string) => + b.startsWith(folderNode.nodePath) + ? b.replace(folderNode.nodePath, newNodePath) + : b + ) + ) + } + + const deleteBookmarkFolder = (folderNode: ComfyNodeDefImpl) => { + if (!folderNode.isDummyFolder) { + throw new Error('Cannot delete non-folder node') + } + settingStore.set( + 'Comfy.NodeLibrary.Bookmarks', + bookmarks.value.filter( + (b: string) => + b !== folderNode.nodePath && !b.startsWith(folderNode.nodePath) + ) + ) + } + + return { + bookmarks, + bookmarkedRoot, + isBookmarked, + toggleBookmark, + addBookmark, + addNewBookmarkFolder, + renameBookmarkFolder, + deleteBookmarkFolder + } +}) diff --git a/src/stores/nodeDefStore.ts b/src/stores/nodeDefStore.ts index 602d7eba..2165a35b 100644 --- a/src/stores/nodeDefStore.ts +++ b/src/stores/nodeDefStore.ts @@ -205,6 +205,14 @@ export class ComfyNodeDefImpl { }) return new ComfyOutputsSpec(result) } + + get nodePath(): string { + return (this.category ? this.category + '/' : '') + this.display_name + } + + get isDummyFolder(): boolean { + return this.name === '' + } } export const SYSTEM_NODE_DEFS: Record = { @@ -244,10 +252,19 @@ export const SYSTEM_NODE_DEFS: Record = { } export function buildNodeDefTree(nodeDefs: ComfyNodeDefImpl[]): TreeNode { - return buildTree(nodeDefs, (nodeDef: ComfyNodeDefImpl) => [ - ...nodeDef.category.split('/').filter((s) => s !== ''), - nodeDef.display_name - ]) + return buildTree(nodeDefs, (nodeDef: ComfyNodeDefImpl) => + nodeDef.nodePath.split('/') + ) +} + +export function createDummyFolderNodeDef(folderPath: string): ComfyNodeDefImpl { + return plainToClass(ComfyNodeDefImpl, { + name: '', + display_name: '', + category: folderPath.endsWith('/') ? folderPath.slice(0, -1) : folderPath, + python_module: 'nodes', + description: 'Dummy Folder Node (User should never see this string)' + }) } interface State { diff --git a/src/utils/treeUtil.ts b/src/utils/treeUtil.ts index 3269a1ff..50bb8c8d 100644 --- a/src/utils/treeUtil.ts +++ b/src/utils/treeUtil.ts @@ -17,7 +17,12 @@ export function buildTree( for (const item of items) { const keys = typeof key === 'string' ? item[key] : key(item) let parent = root - for (const k of keys) { + for (let i = 0; i < keys.length; i++) { + const k = keys[i] + // 'a/b/c/' represents an empty folder 'c' in folder 'b' in folder 'a' + // 'a/b/c/' is split into ['a', 'b', 'c', ''] + if (k === '' && i === keys.length - 1) break + const id = parent.key + '/' + k if (!map[id]) { const node: TreeNode = { @@ -31,7 +36,7 @@ export function buildTree( } parent = map[id] } - parent.leaf = true + parent.leaf = keys[keys.length - 1] !== '' parent.data = item } return root @@ -68,3 +73,19 @@ export function sortedTree(node: TreeNode): TreeNode { return newNode } + +export const findNodeByKey = (root: TreeNode, key: string): TreeNode | null => { + if (root.key === key) { + return root + } + if (!root.children) { + return null + } + for (const child of root.children) { + const result = findNodeByKey(child, key) + if (result) { + return result + } + } + return null +} diff --git a/tests-ui/tests/utils/treeUtilTest.test.ts b/tests-ui/tests/utils/treeUtilTest.test.ts new file mode 100644 index 00000000..c3adf14e --- /dev/null +++ b/tests-ui/tests/utils/treeUtilTest.test.ts @@ -0,0 +1,63 @@ +import { buildTree } from '@/utils/treeUtil' + +describe('buildTree', () => { + it('should handle empty folder items correctly', () => { + const items = [ + { path: 'a/b/c/' }, + { path: 'a/b/d.txt' }, + { path: 'a/e/' }, + { path: 'f.txt' } + ] + + const tree = buildTree(items, (item) => item.path.split('/')) + + expect(tree).toEqual({ + key: 'root', + label: 'root', + children: [ + { + key: 'root/a', + label: 'a', + leaf: false, + children: [ + { + key: 'root/a/b', + label: 'b', + leaf: false, + children: [ + { + key: 'root/a/b/c', + label: 'c', + leaf: false, + children: [], + data: { path: 'a/b/c/' } + }, + { + key: 'root/a/b/d.txt', + label: 'd.txt', + leaf: true, + children: [], + data: { path: 'a/b/d.txt' } + } + ] + }, + { + key: 'root/a/e', + label: 'e', + leaf: false, + children: [], + data: { path: 'a/e/' } + } + ] + }, + { + key: 'root/f.txt', + label: 'f.txt', + leaf: true, + children: [], + data: { path: 'f.txt' } + } + ] + }) + }) +})