Fix execution output & previews not displayed (#4506)

Author: webfiltered

src/stores/executionStore.ts

@@ -1,4 +1,4 @@
-import type { LGraph, Subgraph } from '@comfyorg/litegraph'
+import type { LGraph, LGraphNode, Subgraph } from '@comfyorg/litegraph'
 import { defineStore } from 'pinia'
 import { computed, ref } from 'vue'
 
@@ -24,7 +24,7 @@ import type {
 } from '@/schemas/comfyWorkflowSchema'
 import { api } from '@/scripts/api'
 import { app } from '@/scripts/app'
-import type { NodeLocatorId } from '@/types/nodeIdentification'
+import type { NodeExecutionId, NodeLocatorId } from '@/types/nodeIdentification'
 import { createNodeLocatorId } from '@/types/nodeIdentification'
 
 import { useCanvasStore } from './graphStore'
@@ -54,6 +54,22 @@ export const useExecutionStore = defineStore('execution', () => {
   // This is the progress of all nodes in the currently executing workflow
   const nodeProgressStates = ref<Record<string, NodeProgressState>>({})
 
+  /**
+   * @deprecated Workaround for Subgraph phase 1 - execution IDs may share locator IDs.
+   * The most recently executed execution ID for each node locator ID.
+   */
+  const locatorIdToExecutionIdMap = new Map<NodeLocatorId, NodeExecutionId>()
+
+  /**
+   * Get the NodeLocatorId for a node.
+   * @param node The node
+   * @returns The NodeLocatorId
+   */
+  const getNodeLocatorId = (node: LGraphNode): NodeLocatorId =>
+    node.graph?.isRootGraph
+      ? node.id.toString()
+      : [node.graph?.id, node.id].join(':')
+
   /**
    * Convert execution context node IDs to NodeLocatorIds
    * @param nodeId The node ID from execution context (could be execution ID)
@@ -468,6 +484,8 @@ export const useExecutionStore = defineStore('execution', () => {
     _executingNodeProgress,
     // NodeLocatorId conversion helpers
     executionIdToNodeLocatorId,
-    nodeLocatorIdToExecutionId
+    nodeLocatorIdToExecutionId,
+    locatorIdToExecutionIdMap,
+    getNodeLocatorId
   }
 })

src/stores/imagePreviewStore.ts

@@ -8,9 +8,12 @@ import {
 } from '@/schemas/apiSchema'
 import { api } from '@/scripts/api'
 import { app } from '@/scripts/app'
+import type { NodeExecutionId } from '@/types'
 import { parseFilePath } from '@/utils/formatUtil'
 import { isVideoNode } from '@/utils/litegraphUtil'
 
+import { useExecutionStore } from './executionStore'
+
 const createOutputs = (
   filenames: string[],
   type: ResultItemType,
@@ -23,16 +26,20 @@ const createOutputs = (
 }
 
 export const useNodeOutputStore = defineStore('nodeOutput', () => {
-  const getNodeId = (node: LGraphNode): string => node.id.toString()
+  const executionStore = useExecutionStore()
+  const getMostRecentExecutionId = (node: LGraphNode): NodeExecutionId =>
+    executionStore.locatorIdToExecutionIdMap.get(
+      executionStore.getNodeLocatorId(node)
+    ) ?? node.id.toString()
 
   function getNodeOutputs(
     node: LGraphNode
   ): ExecutedWsMessage['output'] | undefined {
-    return app.nodeOutputs[getNodeId(node)]
+    return app.nodeOutputs[getMostRecentExecutionId(node)]
   }
 
   function getNodePreviews(node: LGraphNode): string[] | undefined {
-    return app.nodePreviewImages[getNodeId(node)]
+    return app.nodePreviewImages[getMostRecentExecutionId(node)]
   }
 
   /**
@@ -96,7 +103,7 @@ export const useNodeOutputStore = defineStore('nodeOutput', () => {
   ) {
     if (!filenames || !node) return
 
-    const nodeId = getNodeId(node)
+    const nodeId = getMostRecentExecutionId(node)
 
     if (typeof filenames === 'string') {
       app.nodeOutputs[nodeId] = createOutputs([filenames], folder, isAnimated)

src/utils/executionUtil.ts

@@ -14,6 +14,7 @@ import type {
   ComfyApiWorkflow,
   ComfyWorkflowJSON
 } from '@/schemas/comfyWorkflowSchema'
+import { useExecutionStore } from '@/stores/executionStore'
 
 import { ExecutableGroupNodeDTO, isGroupNode } from './executableGroupNodeDto'
 import { compressWidgetInputSlots } from './litegraphUtil'
@@ -58,6 +59,8 @@ export const graphToPrompt = async (
   options: { sortNodes?: boolean; queueNodeIds?: NodeId[] } = {}
 ): Promise<{ workflow: ComfyWorkflowJSON; output: ComfyApiWorkflow }> => {
   const { sortNodes = false, queueNodeIds } = options
+  const executionStore = useExecutionStore()
+  executionStore.locatorIdToExecutionIdMap.clear()
 
   for (const node of graph.computeExecutionOrder(false)) {
     const innerNodes = node.getInnerNodes
@@ -159,6 +162,13 @@ export const graphToPrompt = async (
       ]
     }
 
+    // Temporary workaround for Subgraph phase 1. Overwrites the ID, but keeps the image.
+    const baseNodeId = node.id.split(':').at(-1)
+    executionStore.locatorIdToExecutionIdMap.set(
+      `${node.subgraphId}:${baseNodeId}`,
+      node.id
+    )
+
     output[String(node.id)] = {
       inputs,
       // TODO(huchenlei): Filter out all nodes that cannot be mapped to a