diff --git a/src/components/graph/GraphCanvas.vue b/src/components/graph/GraphCanvas.vue index 8ef41fb34..f1e024b15 100644 --- a/src/components/graph/GraphCanvas.vue +++ b/src/components/graph/GraphCanvas.vue @@ -72,7 +72,6 @@ import { useWorkflowAutoSave } from '@/composables/useWorkflowAutoSave' import { useWorkflowPersistence } from '@/composables/useWorkflowPersistence' import { CORE_SETTINGS } from '@/constants/coreSettings' import { i18n, t } from '@/i18n' -import type { NodeId } from '@/schemas/comfyWorkflowSchema' import { UnauthorizedError, api } from '@/scripts/api' import { app as comfyApp } from '@/scripts/app' import { ChangeTracker } from '@/scripts/changeTracker' @@ -191,22 +190,46 @@ watch( } ) -// Update the progress of the executing node +// Update the progress of executing nodes watch( - () => - [ - executionStore.executingNodeId, - executionStore.executingNodeProgress - ] satisfies [NodeId | null, number | null], - ([executingNodeId, executingNodeProgress]) => { + () => executionStore.nodeProgressStates, + (nodeProgressStates) => { + // Clear progress for all nodes first for (const node of comfyApp.graph.nodes) { - if (node.id == executingNodeId) { - node.progress = executingNodeProgress ?? undefined - } else { - node.progress = undefined + node.progress = undefined + } + + // Then set progress for nodes with progress states + for (const nodeId in nodeProgressStates) { + const progressState = nodeProgressStates[nodeId] + const node = comfyApp.graph.getNodeById(progressState.display_node_id) + + if (node && progressState) { + // Only show progress for running nodes + if (progressState.state === 'running') { + if (node.progress === undefined || node.progress === 0.0) { + console.log( + `${Date.now()} Setting progress for node ${node.id} to ${progressState.value / progressState.max}=${progressState.value}/${progressState.max} due to ${nodeId}` + ) + node.progress = progressState.value / progressState.max + } else { + // Update progress if it was already set + console.log( + `${Date.now()} Setting progress for node ${node.id} to Math.min(${node.progress}, ${progressState.value / progressState.max}=${progressState.value}/${progressState.max}) due to ${nodeId}` + ) + node.progress = Math.min( + node.progress, + progressState.value / progressState.max + ) + } + } } } - } + + // TODO - Do we need to force canvas redraw here? + // comfyApp.graph.setDirtyCanvas(true, true) + }, + { deep: true } ) // Update node slot errors diff --git a/src/components/graph/widgets/TextPreviewWidget.vue b/src/components/graph/widgets/TextPreviewWidget.vue index 8d161b8e3..7244634a4 100644 --- a/src/components/graph/widgets/TextPreviewWidget.vue +++ b/src/components/graph/widgets/TextPreviewWidget.vue @@ -20,33 +20,44 @@ import { useExecutionStore } from '@/stores/executionStore' import { linkifyHtml, nl2br } from '@/utils/formatUtil' const modelValue = defineModel({ required: true }) -defineProps<{ +const props = defineProps<{ widget?: object + nodeId: NodeId }>() const executionStore = useExecutionStore() const isParentNodeExecuting = ref(true) const formattedText = computed(() => nl2br(linkifyHtml(modelValue.value))) -let executingNodeId: NodeId | null = null +let parentNodeId: NodeId | null = null onMounted(() => { - executingNodeId = executionStore.executingNodeId + // Get the parent node ID from props if provided + // For backward compatibility, fall back to the first executing node + parentNodeId = props.nodeId }) // Watch for either a new node has starting execution or overall execution ending const stopWatching = watch( - [() => executionStore.executingNode, () => executionStore.isIdle], + [() => executionStore.executingNodeIds, () => executionStore.isIdle], () => { + if (executionStore.isIdle) { + isParentNodeExecuting.value = false + stopWatching() + return + } + + // Check if parent node is no longer in the executing nodes list if ( - executionStore.isIdle || - (executionStore.executingNode && - executionStore.executingNode.id !== executingNodeId) + parentNodeId && + !executionStore.executingNodeIds.includes(parentNodeId) ) { isParentNodeExecuting.value = false stopWatching() } - if (!executingNodeId) { - executingNodeId = executionStore.executingNodeId + + // Set parent node ID if not set yet + if (!parentNodeId && executionStore.executingNodeIds.length > 0) { + parentNodeId = executionStore.executingNodeIds[0] } } ) diff --git a/src/composables/useBrowserTabTitle.ts b/src/composables/useBrowserTabTitle.ts index 95d60752d..6f0b85a82 100644 --- a/src/composables/useBrowserTabTitle.ts +++ b/src/composables/useBrowserTabTitle.ts @@ -1,6 +1,7 @@ import { useTitle } from '@vueuse/core' import { computed } from 'vue' +import { t } from '@/i18n' import { useExecutionStore } from '@/stores/executionStore' import { useSettingStore } from '@/stores/settingStore' import { useWorkflowStore } from '@/stores/workflowStore' @@ -36,11 +37,35 @@ export const useBrowserTabTitle = () => { : DEFAULT_TITLE }) - const nodeExecutionTitle = computed(() => - executionStore.executingNode && executionStore.executingNodeProgress - ? `${executionText.value}[${Math.round(executionStore.executingNodeProgress * 100)}%] ${executionStore.executingNode.type}` - : '' - ) + const nodeExecutionTitle = computed(() => { + // Check if any nodes are in progress + const nodeProgressEntries = Object.entries( + executionStore.nodeProgressStates + ) + const runningNodes = nodeProgressEntries.filter( + ([_, state]) => state.state === 'running' + ) + + if (runningNodes.length > 0) { + // If multiple nodes are running + if (runningNodes.length > 1) { + return `${executionText.value}[${runningNodes.length} ${t('g.nodesRunning', 'nodes running')}]` + } + // If only one node is running + else { + const [nodeId, state] = runningNodes[0] + const progress = Math.round((state.value / state.max) * 100) + const nodeType = + executionStore.activePrompt?.workflow?.changeTracker?.activeState?.nodes.find( + (n) => String(n.id) === nodeId + )?.type || 'Node' + + return `${executionText.value}[${progress}%] ${nodeType}` + } + } + + return '' + }) const workflowTitle = computed( () => diff --git a/src/composables/widgets/useProgressTextWidget.ts b/src/composables/widgets/useProgressTextWidget.ts index 625f0f9cf..078c461aa 100644 --- a/src/composables/widgets/useProgressTextWidget.ts +++ b/src/composables/widgets/useProgressTextWidget.ts @@ -23,6 +23,9 @@ export const useTextPreviewWidget = ( name: inputSpec.name, component: TextPreviewWidget, inputSpec, + componentProps: { + nodeId: node.id + }, options: { getValue: () => widgetValue.value, setValue: (value: string | object) => { diff --git a/src/extensions/core/groupNode.ts b/src/extensions/core/groupNode.ts index ccc973390..4ddf60e38 100644 --- a/src/extensions/core/groupNode.ts +++ b/src/extensions/core/groupNode.ts @@ -9,6 +9,7 @@ import { } from '@/schemas/comfyWorkflowSchema' import type { ComfyNodeDef } from '@/schemas/nodeDefSchema' import { useDialogService } from '@/services/dialogService' +import { useExecutionStore } from '@/stores/executionStore' import { useNodeDefStore } from '@/stores/nodeDefStore' import { useToastStore } from '@/stores/toastStore' import { useWidgetStore } from '@/stores/widgetStore' @@ -1178,9 +1179,10 @@ export class GroupNodeHandler { node.onDrawForeground = function (ctx) { // @ts-expect-error fixme ts strict error onDrawForeground?.apply?.(this, arguments) + const executionStore = useExecutionStore() if ( - // @ts-expect-error fixme ts strict error - +app.runningNodeId === this.id && + executionStore.nodeProgressStates[this.id] && + executionStore.nodeProgressStates[this.id].state === 'running' && this.runningInternalNodeId !== null ) { // @ts-expect-error fixme ts strict error @@ -1275,6 +1277,45 @@ export class GroupNodeHandler { // @ts-expect-error fixme ts strict error (_, id) => id ) + /* + // Handle progress_state events for multiple executing nodes + const progress_state = handleEvent.call( + this, + 'progress_state', + (d) => { + // Check if any of our inner nodes are in this progress state update + for (const nodeId in d.nodes) { + const innerNodeIndex = this.innerNodes?.findIndex((n) => n.id == nodeId); + if (innerNodeIndex > -1) return nodeId; + } + return null; + }, + (d, id, node) => { + // Create a new progress_state event with just our group node + const newProgressState = { ...d }; + newProgressState.nodes = { [id]: { + node: id, + state: 'running', + value: 0, + max: 1, + prompt_id: d.prompt_id + }}; + + // If we have a specific running internal node, update its state + if (node.runningInternalNodeId !== null) { + const innerNodeId = this.innerNodes[node.runningInternalNodeId].id; + if (d.nodes[innerNodeId]) { + newProgressState.nodes[id] = { + ...d.nodes[innerNodeId], + node: id + }; + } + } + + return newProgressState; + } + ); + */ const executed = handleEvent.call( this, @@ -1294,6 +1335,7 @@ export class GroupNodeHandler { this.node.onRemoved = function () { // @ts-expect-error fixme ts strict error onRemoved?.apply(this, arguments) + // api.removeEventListener('progress_state', progress_state) api.removeEventListener('executing', executing) api.removeEventListener('executed', executed) } diff --git a/src/locales/en/main.json b/src/locales/en/main.json index ebfa79a73..2aed21615 100644 --- a/src/locales/en/main.json +++ b/src/locales/en/main.json @@ -133,7 +133,8 @@ "copyURL": "Copy URL", "releaseTitle": "{package} {version} Release", "progressCountOf": "of", - "keybindingAlreadyExists": "Keybinding already exists on" + "keybindingAlreadyExists": "Keybinding already exists on", + "nodesRunning": "nodes running" }, "manager": { "title": "Custom Nodes Manager", @@ -1490,4 +1491,4 @@ "whatsNewPopup": { "learnMore": "Learn more" } -} \ No newline at end of file +} diff --git a/src/schemas/apiSchema.ts b/src/schemas/apiSchema.ts index da637612e..c5979c6bd 100644 --- a/src/schemas/apiSchema.ts +++ b/src/schemas/apiSchema.ts @@ -48,6 +48,22 @@ const zProgressWsMessage = z.object({ node: zNodeId }) +const zNodeProgressState = z.object({ + value: z.number(), + max: z.number(), + state: z.enum(['pending', 'running', 'finished', 'error']), + node_id: zNodeId, + prompt_id: zPromptId, + display_node_id: zNodeId.optional(), + parent_node_id: zNodeId.optional(), + real_node_id: zNodeId.optional() +}) + +const zProgressStateWsMessage = z.object({ + prompt_id: zPromptId, + nodes: z.record(zNodeId, zNodeProgressState) +}) + const zExecutingWsMessage = z.object({ node: zNodeId, display_node: zNodeId, @@ -132,6 +148,8 @@ export type ProgressTextWsMessage = z.infer export type DisplayComponentWsMessage = z.infer< typeof zDisplayComponentWsMessage > +export type NodeProgressState = z.infer +export type ProgressStateWsMessage = z.infer // End of ws messages const zPromptInputItem = z.object({ diff --git a/src/scripts/api.ts b/src/scripts/api.ts index 85316a74a..84048c906 100644 --- a/src/scripts/api.ts +++ b/src/scripts/api.ts @@ -15,6 +15,7 @@ import type { LogsRawResponse, LogsWsMessage, PendingTaskItem, + ProgressStateWsMessage, ProgressTextWsMessage, ProgressWsMessage, PromptResponse, @@ -103,7 +104,17 @@ interface BackendApiCalls { logs: LogsWsMessage /** Binary preview/progress data */ b_preview: Blob + /** Binary preview with metadata (node_id, prompt_id) */ + b_preview_with_metadata: { + blob: Blob + nodeId: string + parentNodeId: string + displayNodeId: string + realNodeId: string + promptId: string + } progress_text: ProgressTextWsMessage + progress_state: ProgressStateWsMessage display_component: DisplayComponentWsMessage } @@ -432,6 +443,33 @@ export class ComfyApi extends EventTarget { }) this.dispatchCustomEvent('b_preview', imageBlob) break + case 4: + // PREVIEW_IMAGE_WITH_METADATA + const decoder4 = new TextDecoder() + const metadataLength = view.getUint32(4) + const metadataBytes = event.data.slice(8, 8 + metadataLength) + const metadata = JSON.parse(decoder4.decode(metadataBytes)) + const imageData4 = event.data.slice(8 + metadataLength) + + let imageMime4 = metadata.image_type + + const imageBlob4 = new Blob([imageData4], { + type: imageMime4 + }) + + // Dispatch enhanced preview event with metadata + this.dispatchCustomEvent('b_preview_with_metadata', { + blob: imageBlob4, + nodeId: metadata.node_id, + displayNodeId: metadata.display_node_id, + parentNodeId: metadata.parent_node_id, + realNodeId: metadata.real_node_id, + promptId: metadata.prompt_id + }) + + // Also dispatch legacy b_preview for backward compatibility + this.dispatchCustomEvent('b_preview', imageBlob4) + break default: throw new Error( `Unknown binary websocket message of type ${eventType}` @@ -461,6 +499,7 @@ export class ComfyApi extends EventTarget { case 'execution_cached': case 'execution_success': case 'progress': + case 'progress_state': case 'executed': case 'graphChanged': case 'promptQueued': diff --git a/src/scripts/app.ts b/src/scripts/app.ts index 7744ae757..c1474de52 100644 --- a/src/scripts/app.ts +++ b/src/scripts/app.ts @@ -194,6 +194,8 @@ export class ComfyApp { /** * @deprecated Use useExecutionStore().executingNodeId instead + * TODO: Update to support multiple executing nodes. This getter returns only the first executing node. + * Consider updating consumers to handle multiple nodes or use executingNodeIds array. */ get runningNodeId(): NodeId | null { return useExecutionStore().executingNodeId @@ -635,10 +637,6 @@ export class ComfyApp { api.addEventListener('executing', () => { this.graph.setDirtyCanvas(true, false) - // @ts-expect-error fixme ts strict error - this.revokePreviews(this.runningNodeId) - // @ts-expect-error fixme ts strict error - delete this.nodePreviewImages[this.runningNodeId] }) api.addEventListener('executed', ({ detail }) => { @@ -689,15 +687,13 @@ export class ComfyApp { this.canvas.draw(true, true) }) - api.addEventListener('b_preview', ({ detail }) => { - const id = this.runningNodeId - if (id == null) return - - const blob = detail + api.addEventListener('b_preview_with_metadata', ({ detail }) => { + // Enhanced preview with explicit node context + const { blob, displayNodeId } = detail + this.revokePreviews(displayNodeId) const blobUrl = URL.createObjectURL(blob) - // Ensure clean up if `executing` event is missed. - this.revokePreviews(id) - this.nodePreviewImages[id] = [blobUrl] + // Preview cleanup is now handled in progress_state event to support multiple concurrent previews + this.nodePreviewImages[displayNodeId] = [blobUrl] }) api.init() diff --git a/src/scripts/domWidget.ts b/src/scripts/domWidget.ts index 9846234fc..ae237752d 100644 --- a/src/scripts/domWidget.ts +++ b/src/scripts/domWidget.ts @@ -237,6 +237,7 @@ export class ComponentWidgetImpl< component: Component inputSpec: InputSpec props?: P + componentProps?: Record options: DOMWidgetOptions }) { super({ @@ -245,7 +246,9 @@ export class ComponentWidgetImpl< }) this.component = obj.component this.inputSpec = obj.inputSpec - this.props = obj.props + this.props = obj.componentProps + ? ({ ...obj.props, ...obj.componentProps } as P) + : obj.props } override computeLayoutSize() { diff --git a/src/services/litegraphService.ts b/src/services/litegraphService.ts index afd31d83b..04addfde1 100644 --- a/src/services/litegraphService.ts +++ b/src/services/litegraphService.ts @@ -33,6 +33,7 @@ import type { import type { ComfyNodeDef as ComfyNodeDefV1 } from '@/schemas/nodeDefSchema' import { ComfyApp, app } from '@/scripts/app' import { $el } from '@/scripts/ui' +import { useExecutionStore } from '@/stores/executionStore' import { useCanvasStore } from '@/stores/graphStore' import { useNodeOutputStore } from '@/stores/imagePreviewStore' import { ComfyNodeDefImpl } from '@/stores/nodeDefStore' @@ -362,7 +363,12 @@ export const useLitegraphService = () => { */ #setupStrokeStyles() { this.strokeStyles['running'] = function (this: LGraphNode) { - if (this.id == app.runningNodeId) { + const nodeProgressStates = useExecutionStore().nodeProgressStates + const nodeId = String(this.id) + if ( + nodeProgressStates[nodeId] && + nodeProgressStates[nodeId].state === 'running' + ) { return { color: '#0f0' } } } diff --git a/src/stores/executionStore.ts b/src/stores/executionStore.ts index a8667da0d..c2a5cc595 100644 --- a/src/stores/executionStore.ts +++ b/src/stores/executionStore.ts @@ -12,6 +12,8 @@ import type { ExecutionErrorWsMessage, ExecutionStartWsMessage, NodeError, + NodeProgressState, + ProgressStateWsMessage, ProgressTextWsMessage, ProgressWsMessage } from '@/schemas/apiSchema' @@ -46,7 +48,22 @@ export const useExecutionStore = defineStore('execution', () => { const queuedPrompts = ref>({}) const lastNodeErrors = ref | null>(null) const lastExecutionError = ref(null) - const executingNodeId = ref(null) + // This is the progress of all nodes in the currently executing workflow + const nodeProgressStates = ref>({}) + + // Easily access all currently executing node IDs + const executingNodeIds = computed(() => { + return Object.entries(nodeProgressStates) + .filter(([_, state]) => state.state === 'running') + .map(([nodeId, _]) => nodeId) + }) + + // For backward compatibility - stores the primary executing node ID + const executingNodeId = computed(() => { + return executingNodeIds.value.length > 0 ? executingNodeIds.value[0] : null + }) + + // For backward compatibility - returns the primary executing node const executingNode = computed(() => { if (!executingNodeId.value) return null @@ -116,7 +133,7 @@ export const useExecutionStore = defineStore('execution', () => { } } - // This is the progress of the currently executing node, if any + // This is the progress of the currently executing node (for backward compatibility) const _executingNodeProgress = ref(null) const executingNodeProgress = computed(() => _executingNodeProgress.value @@ -153,6 +170,7 @@ export const useExecutionStore = defineStore('execution', () => { api.addEventListener('executed', handleExecuted) api.addEventListener('executing', handleExecuting) api.addEventListener('progress', handleProgress) + api.addEventListener('progress_state', handleProgressState) api.addEventListener('status', handleStatus) api.addEventListener('execution_error', handleExecutionError) } @@ -165,6 +183,7 @@ export const useExecutionStore = defineStore('execution', () => { api.removeEventListener('executed', handleExecuted) api.removeEventListener('executing', handleExecuting) api.removeEventListener('progress', handleProgress) + api.removeEventListener('progress_state', handleProgressState) api.removeEventListener('status', handleStatus) api.removeEventListener('execution_error', handleExecutionError) api.removeEventListener('progress_text', handleProgressText) @@ -194,19 +213,41 @@ export const useExecutionStore = defineStore('execution', () => { if (!activePrompt.value) return - if (executingNodeId.value && activePrompt.value) { - // Seems sometimes nodes that are cached fire executing but not executed - activePrompt.value.nodes[executingNodeId.value] = true + // Update the executing nodes list + if (e.detail === null) { + if (activePromptId.value) { + delete queuedPrompts.value[activePromptId.value] + } } - if (typeof e.detail === 'string') { - executingNodeId.value = executionIdToCurrentId(e.detail) ?? null - } else { - executingNodeId.value = e.detail - if (executingNodeId.value === null) { - if (activePromptId.value) { - delete queuedPrompts.value[activePromptId.value] - } - activePromptId.value = null + } + + function handleProgressState(e: CustomEvent) { + const { nodes } = e.detail + + // Revoke previews for nodes that are starting to execute + for (const nodeId in nodes) { + const nodeState = nodes[nodeId] + if (nodeState.state === 'running' && !nodeProgressStates.value[nodeId]) { + // This node just started executing, revoke its previews + // Note that we're doing the *actual* node id instead of the display node id + // here intentionally. That way, we don't clear the preview every time a new node + // within an expanded graph starts executing. + app.revokePreviews(nodeId) + delete app.nodePreviewImages[nodeId] + } + } + + // Update the progress states for all nodes + nodeProgressStates.value = nodes + + // If we have progress for the currently executing node, update it for backwards compatibility + if (executingNodeId.value && nodes[executingNodeId.value]) { + const nodeState = nodes[executingNodeId.value] + _executingNodeProgress.value = { + value: nodeState.value, + max: nodeState.max, + prompt_id: nodeState.prompt_id, + node: nodeState.display_node_id || nodeState.node_id } } } @@ -310,9 +351,13 @@ export const useExecutionStore = defineStore('execution', () => { */ lastExecutionError, /** - * The id of the node that is currently being executed + * The id of the node that is currently being executed (backward compatibility) */ executingNodeId, + /** + * The list of all nodes that are currently executing + */ + executingNodeIds, /** * The prompt that is currently being executed */ @@ -330,13 +375,17 @@ export const useExecutionStore = defineStore('execution', () => { */ executionProgress, /** - * The node that is currently being executed + * The node that is currently being executed (backward compatibility) */ executingNode, /** - * The progress of the executing node (if the node reports progress) + * The progress of the executing node (backward compatibility) */ executingNodeProgress, + /** + * All node progress states from progress_state events + */ + nodeProgressStates, bindExecutionEvents, unbindExecutionEvents, storePrompt, diff --git a/tests-ui/tests/composables/BrowserTabTitle.spec.ts b/tests-ui/tests/composables/BrowserTabTitle.spec.ts index 328f41e80..dd88fb1cd 100644 --- a/tests-ui/tests/composables/BrowserTabTitle.spec.ts +++ b/tests-ui/tests/composables/BrowserTabTitle.spec.ts @@ -3,12 +3,20 @@ import { nextTick, reactive } from 'vue' import { useBrowserTabTitle } from '@/composables/useBrowserTabTitle' +// Mock i18n module +vi.mock('@/i18n', () => ({ + t: (key: string, fallback: string) => + key === 'g.nodesRunning' ? 'nodes running' : fallback +})) + // Mock the execution store const executionStore = reactive({ isIdle: true, executionProgress: 0, executingNode: null as any, - executingNodeProgress: 0 + executingNodeProgress: 0, + nodeProgressStates: {} as any, + activePrompt: null as any }) vi.mock('@/stores/executionStore', () => ({ useExecutionStore: () => executionStore @@ -37,6 +45,8 @@ describe('useBrowserTabTitle', () => { executionStore.executionProgress = 0 executionStore.executingNode = null as any executionStore.executingNodeProgress = 0 + executionStore.nodeProgressStates = {} + executionStore.activePrompt = null // reset setting and workflow stores ;(settingStore.get as any).mockReturnValue('Enabled') @@ -97,13 +107,41 @@ describe('useBrowserTabTitle', () => { expect(document.title).toBe('[30%]ComfyUI') }) - it('shows node execution title when executing a node', async () => { + it('shows node execution title when executing a node using nodeProgressStates', async () => { executionStore.isIdle = false executionStore.executionProgress = 0.4 - executionStore.executingNodeProgress = 0.5 - executionStore.executingNode = { type: 'Foo' } + executionStore.nodeProgressStates = { + '1': { state: 'running', value: 5, max: 10, node: '1', prompt_id: 'test' } + } + executionStore.activePrompt = { + workflow: { + changeTracker: { + activeState: { + nodes: [{ id: 1, type: 'Foo' }] + } + } + } + } useBrowserTabTitle() await nextTick() expect(document.title).toBe('[40%][50%] Foo') }) + + it('shows multiple nodes running when multiple nodes are executing', async () => { + executionStore.isIdle = false + executionStore.executionProgress = 0.4 + executionStore.nodeProgressStates = { + '1': { + state: 'running', + value: 5, + max: 10, + node: '1', + prompt_id: 'test' + }, + '2': { state: 'running', value: 8, max: 10, node: '2', prompt_id: 'test' } + } + useBrowserTabTitle() + await nextTick() + expect(document.title).toBe('[40%][2 nodes running]') + }) })