From f6c65d3fe7f9ac7dc3a0dc0364f603eb3e359d48 Mon Sep 17 00:00:00 2001 From: Jacob Segal Date: Mon, 7 Jul 2025 17:16:38 -0700 Subject: [PATCH] Progress bars work in subgraphs --- src/components/graph/GraphCanvas.vue | 38 ++++------ src/services/litegraphService.ts | 10 ++- src/stores/executionStore.ts | 103 +++++++++++++++++++++------ 3 files changed, 102 insertions(+), 49 deletions(-) diff --git a/src/components/graph/GraphCanvas.vue b/src/components/graph/GraphCanvas.vue index 362802a8c..4189c02a2 100644 --- a/src/components/graph/GraphCanvas.vue +++ b/src/components/graph/GraphCanvas.vue @@ -192,36 +192,22 @@ watch( // Update the progress of executing nodes watch( - () => executionStore.nodeProgressStates, - (nodeProgressStates) => { - // Clear progress for all nodes first - for (const node of comfyApp.graph.nodes) { - 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) { - node.progress = progressState.value / progressState.max - } else { - // Update progress if it was already set - node.progress = Math.min( - node.progress, - progressState.value / progressState.max - ) - } - } + () => + [executionStore.nodeLocationProgressStates, canvasStore.canvas] as const, + ([nodeLocationProgressStates, canvas]) => { + if (!canvas?.graph) return + for (const node of canvas.graph.nodes) { + const nodeLocatorId = useWorkflowStore().nodeIdToNodeLocatorId(node.id) + const progressState = nodeLocationProgressStates[nodeLocatorId] + if (progressState && progressState.state === 'running') { + node.progress = progressState.value / progressState.max + } else { + node.progress = undefined } } // Force canvas redraw to ensure progress updates are visible - comfyApp.graph.setDirtyCanvas(true, false) + canvas.graph.setDirtyCanvas(true, false) }, { deep: true } ) diff --git a/src/services/litegraphService.ts b/src/services/litegraphService.ts index c2715a593..cb6efd627 100644 --- a/src/services/litegraphService.ts +++ b/src/services/litegraphService.ts @@ -108,7 +108,11 @@ export const useLitegraphService = () => { */ #setupStrokeStyles() { this.strokeStyles['running'] = function (this: LGraphNode) { - if (this.id == app.runningNodeId) { + const nodeId = String(this.id) + const nodeLocatorId = useWorkflowStore().nodeIdToNodeLocatorId(nodeId) + const state = + useExecutionStore().nodeLocationProgressStates[nodeLocatorId]?.state + if (state === 'running') { return { color: '#0f0' } } } @@ -364,7 +368,9 @@ export const useLitegraphService = () => { #setupStrokeStyles() { this.strokeStyles['running'] = function (this: LGraphNode) { const nodeId = String(this.id) - const state = useExecutionStore().nodeProgressStates[nodeId]?.state + const nodeLocatorId = useWorkflowStore().nodeIdToNodeLocatorId(nodeId) + const state = + useExecutionStore().nodeLocationProgressStates[nodeLocatorId]?.state if (state === 'running') { return { color: '#0f0' } } diff --git a/src/stores/executionStore.ts b/src/stores/executionStore.ts index 9182d04f3..306df2dfa 100644 --- a/src/stores/executionStore.ts +++ b/src/stores/executionStore.ts @@ -25,6 +25,7 @@ import type { import { api } from '@/scripts/api' import { app } from '@/scripts/app' import type { NodeLocatorId } from '@/types/nodeIdentification' +import { createNodeLocatorId } from '@/types/nodeIdentification' import { useCanvasStore } from './graphStore' import { ComfyWorkflow, useWorkflowStore } from './workflowStore' @@ -53,6 +54,86 @@ export const useExecutionStore = defineStore('execution', () => { // This is the progress of all nodes in the currently executing workflow const nodeProgressStates = ref>({}) + /** + * Convert execution context node IDs to NodeLocatorIds + * @param nodeId The node ID from execution context (could be hierarchical) + * @returns The NodeLocatorId + */ + const executionIdToNodeLocatorId = ( + nodeId: string | number + ): NodeLocatorId => { + const nodeIdStr = String(nodeId) + + if (!nodeIdStr.includes(':')) { + // It's a top-level node ID + return nodeIdStr as NodeLocatorId + } + + // It's a hierarchical node ID + const parts = nodeIdStr.split(':') + const localNodeId = parts[parts.length - 1] + const subgraphs = getSubgraphsFromInstanceIds(app.graph, parts) + const nodeLocatorId = createNodeLocatorId(subgraphs.at(-1)!.id, localNodeId) + return nodeLocatorId + } + + const mergeHierarchicalProgressStates = ( + currentState: NodeProgressState | undefined, + newState: NodeProgressState + ): NodeProgressState => { + if (currentState === undefined) { + return newState + } + + const mergedState = { ...currentState } + if (mergedState.state === 'error') { + return mergedState + } else if (newState.state === 'running') { + const newPerc = newState.max > 0 ? newState.value / newState.max : 0.0 + if (mergedState.state === 'running') { + const oldPerc = + mergedState.max > 0 ? mergedState.value / mergedState.max : 0.0 + if (oldPerc === 0.0) { + mergedState.value = newState.value + mergedState.max = newState.max + } else if (newPerc < oldPerc) { + mergedState.value = newState.value + mergedState.max = newState.max + } + } else { + mergedState.value = newState.value + mergedState.max = newState.max + } + mergedState.state = 'running' + } + + return mergedState + } + + const nodeLocationProgressStates = computed< + Record + >(() => { + const result: Record = {} + + const states = nodeProgressStates.value // Apparently doing this inside `Object.entries` causes issues + for (const [_, state] of Object.entries(states)) { + // Convert the node ID to a NodeLocatorId + const parts = String(state.display_node_id).split(':') + for (let i = 0; i < parts.length; i++) { + const executionId = parts.slice(0, i + 1).join(':') + const locatorId = executionIdToNodeLocatorId(executionId) + if (!locatorId) continue + + result[locatorId] = mergeHierarchicalProgressStates( + result[locatorId], + state + ) + } + } + + return result + }) + // Easily access all currently executing node IDs const executingNodeIds = computed(() => { return Object.entries(nodeProgressStates) @@ -310,27 +391,6 @@ export const useExecutionStore = defineStore('execution', () => { ) } - /** - * Convert execution context node IDs to NodeLocatorIds - * @param nodeId The node ID from execution context (could be hierarchical) - * @returns The NodeLocatorId - */ - const executionIdToNodeLocatorId = ( - nodeId: string | number - ): NodeLocatorId => { - const nodeIdStr = String(nodeId) - - // If it's a hierarchical ID, use the workflow store's conversion - if (nodeIdStr.includes(':')) { - const result = workflowStore.hierarchicalIdToNodeLocatorId(nodeIdStr) - // If conversion fails, return the original ID as-is - return result ?? (nodeIdStr as NodeLocatorId) - } - - // For simple node IDs, we need the active subgraph context - return workflowStore.nodeIdToNodeLocatorId(nodeIdStr) - } - /** * Convert a NodeLocatorId to an execution context ID (hierarchical ID) * @param locatorId The NodeLocatorId @@ -399,6 +459,7 @@ export const useExecutionStore = defineStore('execution', () => { * All node progress states from progress_state events */ nodeProgressStates, + nodeLocationProgressStates, bindExecutionEvents, unbindExecutionEvents, storePrompt,