diff --git a/src/scripts/app.ts b/src/scripts/app.ts index 444a1fbe83..40db4ef4ea 100644 --- a/src/scripts/app.ts +++ b/src/scripts/app.ts @@ -63,6 +63,7 @@ import { useWidgetStore } from '@/stores/widgetStore' import { deserialiseAndCreate } from '@/extensions/core/vintageClipboard' import { st } from '@/i18n' import { normalizeI18nKey } from '@/utils/formatUtil' +import { useExtensionService } from '@/services/extensionService' export const ANIM_PREVIEW_WIDGET = '$$comfy_animation_preview' @@ -118,7 +119,6 @@ export class ComfyApp { vueAppReady: boolean api: ComfyApi ui: ComfyUI - extensions: ComfyExtension[] extensionManager: ExtensionManager _nodeOutputs: Record nodePreviewImages: Record @@ -184,6 +184,13 @@ export class ComfyApp { return false } + /** + * @deprecated Use useExtensionStore().extensions instead + */ + get extensions(): ComfyExtension[] { + return useExtensionStore().extensions + } + constructor() { this.vueAppReady = false this.ui = new ComfyUI(this) @@ -198,12 +205,6 @@ export class ComfyApp { this.menu = new ComfyAppMenu(this) this.bypassBgColor = '#FF00FF' - /** - * List of extensions that are registered with the app - * @type {ComfyExtension[]} - */ - this.extensions = [] - /** * Stores the execution output data for each node * @type {Record} @@ -223,7 +224,8 @@ export class ComfyApp { set nodeOutputs(value) { this._nodeOutputs = value - this.#invokeExtensions('onNodeOutputsUpdated', value) + if (this.vueAppReady) + useExtensionService().invokeExtensions('onNodeOutputsUpdated', value) } getPreviewFormatParam() { @@ -390,64 +392,6 @@ export class ComfyApp { } } - get enabledExtensions() { - if (!this.vueAppReady) { - return this.extensions - } - return useExtensionStore().enabledExtensions - } - - /** - * Invoke an extension callback - * @param {keyof ComfyExtension} method The extension callback to execute - * @param {any[]} args Any arguments to pass to the callback - * @returns - */ - #invokeExtensions(method, ...args) { - let results = [] - for (const ext of this.enabledExtensions) { - if (method in ext) { - try { - results.push(ext[method](...args, this)) - } catch (error) { - console.error( - `Error calling extension '${ext.name}' method '${method}'`, - { error }, - { extension: ext }, - { args } - ) - } - } - } - return results - } - - /** - * Invoke an async extension callback - * Each callback will be invoked concurrently - * @param {string} method The extension callback to execute - * @param {...any} args Any arguments to pass to the callback - * @returns - */ - async #invokeExtensionsAsync(method, ...args) { - return await Promise.all( - this.enabledExtensions.map(async (ext) => { - if (method in ext) { - try { - return await ext[method](...args, this) - } catch (error) { - console.error( - `Error calling extension '${ext.name}' method '${method}'`, - { error }, - { extension: ext }, - { args } - ) - } - } - }) - ) - } - #addRestoreWorkflowView() { const serialize = LGraph.prototype.serialize const self = this @@ -1636,32 +1580,6 @@ export class ComfyApp { } } - /** - * Loads all extensions from the API into the window in parallel - */ - async #loadExtensions() { - const extensionStore = useExtensionStore() - extensionStore.loadDisabledExtensionNames() - - const extensions = await api.getExtensions() - - // Need to load core extensions first as some custom extensions - // may depend on them. - await import('../extensions/core/index') - extensionStore.captureCoreExtensions() - await Promise.all( - extensions - .filter((extension) => !extension.includes('extensions/core')) - .map(async (ext) => { - try { - await import(/* @vite-ignore */ api.fileURL(ext)) - } catch (error) { - console.error('Error loading extension', ext, error) - } - }) - ) - } - /** * Set up the app on the page */ @@ -1676,7 +1594,7 @@ export class ComfyApp { useWorkspaceStore().workflow.syncWorkflows(), this.ui.settings.load() ]) - await this.#loadExtensions() + await useExtensionService().loadExtensions() this.#addProcessMouseHandler() this.#addProcessKeyHandler() @@ -1708,7 +1626,7 @@ export class ComfyApp { ro.observe(this.bodyRight) ro.observe(this.bodyBottom) - await this.#invokeExtensionsAsync('init') + await useExtensionService().invokeExtensionsAsync('init') await this.registerNodes() initWidgets(this) @@ -1745,7 +1663,7 @@ export class ComfyApp { this.#addCopyHandler() this.#addPasteHandler() - await this.#invokeExtensionsAsync('setup') + await useExtensionService().invokeExtensionsAsync('setup') } resizeCanvas() { @@ -1789,7 +1707,11 @@ export class ComfyApp { const nodeDefStore = useNodeDefStore() const nodeDefArray: ComfyNodeDef[] = Object.values(allNodeDefs) - this.#invokeExtensions('beforeRegisterVueAppNodeDefs', nodeDefArray, this) + useExtensionService().invokeExtensions( + 'beforeRegisterVueAppNodeDefs', + nodeDefArray, + this + ) nodeDefStore.updateNodeDefs(nodeDefArray) } @@ -1830,7 +1752,7 @@ export class ComfyApp { // Load node definitions from the backend const defs = await this.#getNodeDefs() await this.registerNodesFromDefs(defs) - await this.#invokeExtensionsAsync('registerCustomNodes') + await useExtensionService().invokeExtensionsAsync('registerCustomNodes') if (this.vueAppReady) { this.updateVueAppNodeDefs(defs) } @@ -1969,7 +1891,7 @@ export class ComfyApp { this.size = s this.serialize_widgets = true - app.#invokeExtensionsAsync('nodeCreated', this) + useExtensionService().invokeExtensionsAsync('nodeCreated', this) } configure(data: any) { @@ -2006,14 +1928,18 @@ export class ComfyApp { this.#addDrawBackgroundHandler(node) this.#addNodeKeyHandler(node) - await this.#invokeExtensionsAsync('beforeRegisterNodeDef', node, nodeData) + await useExtensionService().invokeExtensionsAsync( + 'beforeRegisterNodeDef', + node, + nodeData + ) LiteGraph.registerNodeType(nodeId, node) // Note: Do not move this to the class definition, it will be overwritten node.category = nodeData.category } async registerNodesFromDefs(defs: Record) { - await this.#invokeExtensionsAsync('addCustomNodeDefs', defs) + await useExtensionService().invokeExtensionsAsync('addCustomNodeDefs', defs) // Register a node for each definition for (const nodeId in defs) { @@ -2119,7 +2045,7 @@ export class ComfyApp { const missingNodeTypes: MissingNodeType[] = [] const missingModels = [] - await this.#invokeExtensionsAsync( + await useExtensionService().invokeExtensionsAsync( 'beforeConfigureGraph', graphData, missingNodeTypes @@ -2266,7 +2192,7 @@ export class ComfyApp { } } - this.#invokeExtensions('loadedGraphNode', node) + useExtensionService().invokeExtensions('loadedGraphNode', node) } // TODO: Properly handle if both nodes and models are missing (sequential dialogs?) @@ -2277,7 +2203,10 @@ export class ComfyApp { const paths = await api.getFolderPaths() this.#showMissingModelsError(missingModels, paths) } - await this.#invokeExtensionsAsync('afterConfigureGraph', missingNodeTypes) + await useExtensionService().invokeExtensionsAsync( + 'afterConfigureGraph', + missingNodeTypes + ) // @ts-expect-error zod types issue. Will be fixed after we enable ts-strict await workflowService.afterLoadNewGraph(workflow, this.graph.serialize()) requestAnimationFrame(() => { @@ -2778,14 +2707,10 @@ export class ComfyApp { /** * Registers a Comfy web extension with the app * @param {ComfyExtension} extension + * @deprecated Use useExtensionService().registerExtension instead */ registerExtension(extension: ComfyExtension) { - if (this.vueAppReady) { - useExtensionStore().registerExtension(extension) - } else { - // For jest testing. - this.extensions.push(extension) - } + useExtensionService().registerExtension(extension) } /** @@ -2824,7 +2749,10 @@ export class ComfyApp { } } - await this.#invokeExtensionsAsync('refreshComboInNodes', defs) + await useExtensionService().invokeExtensionsAsync( + 'refreshComboInNodes', + defs + ) if (this.vueAppReady) { this.updateVueAppNodeDefs(defs) diff --git a/src/services/extensionService.ts b/src/services/extensionService.ts new file mode 100644 index 0000000000..259d2a4e1f --- /dev/null +++ b/src/services/extensionService.ts @@ -0,0 +1,127 @@ +import { app } from '@/scripts/app' +import { api } from '@/scripts/api' +import { useCommandStore } from '@/stores/commandStore' +import { useExtensionStore } from '@/stores/extensionStore' +import { useKeybindingStore } from '@/stores/keybindingStore' +import { useMenuItemStore } from '@/stores/menuItemStore' +import { useSettingStore } from '@/stores/settingStore' +import { useWidgetStore } from '@/stores/widgetStore' +import { useBottomPanelStore } from '@/stores/workspace/bottomPanelStore' +import type { ComfyExtension } from '@/types/comfy' + +export const useExtensionService = () => { + const extensionStore = useExtensionStore() + const settingStore = useSettingStore() + + /** + * Loads all extensions from the API into the window in parallel + */ + const loadExtensions = async () => { + extensionStore.loadDisabledExtensionNames( + settingStore.get('Comfy.Extension.Disabled') + ) + + const extensions = await api.getExtensions() + + // Need to load core extensions first as some custom extensions + // may depend on them. + await import('../extensions/core/index') + extensionStore.captureCoreExtensions() + await Promise.all( + extensions + .filter((extension) => !extension.includes('extensions/core')) + .map(async (ext) => { + try { + await import(/* @vite-ignore */ api.fileURL(ext)) + } catch (error) { + console.error('Error loading extension', ext, error) + } + }) + ) + } + + /** + * Register an extension with the app + * @param extension The extension to register + */ + const registerExtension = (extension: ComfyExtension) => { + extensionStore.registerExtension(extension) + + useKeybindingStore().loadExtensionKeybindings(extension) + useCommandStore().loadExtensionCommands(extension) + useMenuItemStore().loadExtensionMenuCommands(extension) + settingStore.loadExtensionSettings(extension) + useBottomPanelStore().registerExtensionBottomPanelTabs(extension) + if (extension.getCustomWidgets) { + // TODO(huchenlei): We should deprecate the async return value of + // getCustomWidgets. + ;(async () => { + if (extension.getCustomWidgets) { + const widgets = await extension.getCustomWidgets(app) + useWidgetStore().registerCustomWidgets(widgets) + } + })() + } + } + + /** + * Invoke an extension callback + * @param {keyof ComfyExtension} method The extension callback to execute + * @param {any[]} args Any arguments to pass to the callback + * @returns + */ + const invokeExtensions = (method: keyof ComfyExtension, ...args: any[]) => { + const results: any[] = [] + for (const ext of extensionStore.enabledExtensions) { + if (method in ext) { + try { + results.push(ext[method](...args, app)) + } catch (error) { + console.error( + `Error calling extension '${ext.name}' method '${method}'`, + { error }, + { extension: ext }, + { args } + ) + } + } + } + return results + } + + /** + * Invoke an async extension callback + * Each callback will be invoked concurrently + * @param {string} method The extension callback to execute + * @param {...any} args Any arguments to pass to the callback + * @returns + */ + const invokeExtensionsAsync = async ( + method: keyof ComfyExtension, + ...args: any[] + ) => { + return await Promise.all( + extensionStore.enabledExtensions.map(async (ext) => { + if (method in ext) { + try { + return await ext[method](...args, app) + } catch (error) { + console.error( + `Error calling extension '${ext.name}' method '${method}'`, + { error }, + { extension: ext }, + { args } + ) + } + } + }) + ) + } + + return { + loadExtensions, + registerExtension, + invokeExtensions, + invokeExtensionsAsync + } +} diff --git a/src/stores/extensionStore.ts b/src/stores/extensionStore.ts index 860fe8fa23..b2da147319 100644 --- a/src/stores/extensionStore.ts +++ b/src/stores/extensionStore.ts @@ -1,13 +1,6 @@ import { ref, computed, markRaw } from 'vue' import { defineStore } from 'pinia' import type { ComfyExtension } from '@/types/comfy' -import { useKeybindingStore } from './keybindingStore' -import { useCommandStore } from './commandStore' -import { useSettingStore } from './settingStore' -import { app } from '@/scripts/app' -import { useMenuItemStore } from './menuItemStore' -import { useBottomPanelStore } from './workspace/bottomPanelStore' -import { useWidgetStore } from './widgetStore' /** * These extensions are always active, even if they are disabled in the setting. @@ -69,32 +62,10 @@ export const useExtensionStore = defineStore('extension', () => { } extensionByName.value[extension.name] = markRaw(extension) - useKeybindingStore().loadExtensionKeybindings(extension) - useCommandStore().loadExtensionCommands(extension) - useMenuItemStore().loadExtensionMenuCommands(extension) - useSettingStore().loadExtensionSettings(extension) - useBottomPanelStore().registerExtensionBottomPanelTabs(extension) - if (extension.getCustomWidgets) { - // TODO(huchenlei): We should deprecate the async return value of - // getCustomWidgets. - ;(async () => { - if (extension.getCustomWidgets) { - const widgets = await extension.getCustomWidgets(app) - useWidgetStore().registerCustomWidgets(widgets) - } - })() - } - /* - * Extensions are currently stored in both extensionStore and app.extensions. - * Legacy jest tests still depend on app.extensions being populated. - */ - app.extensions.push(extension) } - function loadDisabledExtensionNames() { - disabledExtensionNames.value = new Set( - useSettingStore().get('Comfy.Extension.Disabled') - ) + function loadDisabledExtensionNames(names: string[]) { + disabledExtensionNames.value = new Set(names) for (const name of ALWAYS_DISABLED_EXTENSIONS) { disabledExtensionNames.value.add(name) } @@ -109,7 +80,7 @@ export const useExtensionStore = defineStore('extension', () => { */ const coreExtensionNames = ref([]) function captureCoreExtensions() { - coreExtensionNames.value = app.extensions.map((ext) => ext.name) + coreExtensionNames.value = extensions.value.map((ext) => ext.name) } function isCoreExtension(name: string) { @@ -120,14 +91,6 @@ export const useExtensionStore = defineStore('extension', () => { return extensions.value.some((ext) => !isCoreExtension(ext.name)) }) - // Some core extensions are registered before the store is initialized, e.g. - // colorPalette. - // Register them manually here so the state of app.extensions and - // extensionByName are in sync. - for (const ext of app.extensions) { - extensionByName.value[ext.name] = markRaw(ext) - } - return { extensions, enabledExtensions,