[Refactor] Use LGraphNode.progress API (#3281)

src/components/BrowserTabTitle.vue

@@ -17,7 +17,9 @@ const TITLE_SUFFIX = ' - ComfyUI'
 
 const executionStore = useExecutionStore()
 const executionText = computed(() =>
-  executionStore.isIdle ? '' : `[${executionStore.executionProgress}%]`
+  executionStore.isIdle
+    ? ''
+    : `[${Math.round(executionStore.executionProgress * 100)}%]`
 )
 
 const settingStore = useSettingStore()
@@ -41,7 +43,7 @@ const workflowNameText = computed(() => {
 
 const nodeExecutionTitle = computed(() =>
   executionStore.executingNode && executionStore.executingNodeProgress
-    ? `${executionText.value}[${executionStore.executingNodeProgress}%] ${executionStore.executingNode.type}`
+    ? `${executionText.value}[${Math.round(executionStore.executingNodeProgress * 100)}%] ${executionStore.executingNode.type}`
     : ''
 )
 

src/components/graph/GraphCanvas.vue

@@ -68,6 +68,7 @@ import { IS_CONTROL_WIDGET, updateControlWidgetLabel } from '@/scripts/widgets'
 import { useColorPaletteService } from '@/services/colorPaletteService'
 import { useWorkflowService } from '@/services/workflowService'
 import { useCommandStore } from '@/stores/commandStore'
+import { useExecutionStore } from '@/stores/executionStore'
 import { useCanvasStore } from '@/stores/graphStore'
 import { useNodeDefStore } from '@/stores/nodeDefStore'
 import { useSettingStore } from '@/stores/settingStore'
@@ -80,6 +81,7 @@ const settingStore = useSettingStore()
 const nodeDefStore = useNodeDefStore()
 const workspaceStore = useWorkspaceStore()
 const canvasStore = useCanvasStore()
+const executionStore = useExecutionStore()
 const betaMenuEnabled = computed(
   () => settingStore.get('Comfy.UseNewMenu') !== 'Disabled'
 )
@@ -158,6 +160,25 @@ watch(
   }
 )
 
+// Update the progress of the executing node
+watch(
+  () =>
+    [executionStore.executingNodeId, executionStore.executingNodeProgress] as [
+      string | null,
+      number | null
+    ],
+  ([executingNodeId, executingNodeProgress]) => {
+    if (!executingNodeId) return
+    for (const node of comfyApp.graph.nodes) {
+      if (node.id == executingNodeId) {
+        node.progress = executingNodeProgress ?? undefined
+      } else {
+        node.progress = undefined
+      }
+    }
+  }
+)
+
 const loadCustomNodesI18n = async () => {
   try {
     const i18nData = await api.getCustomNodesI18n()

src/scripts/app.ts

@@ -562,27 +562,15 @@ export class ComfyApp {
     LGraphCanvas.prototype.drawNodeShape = function (
       node,
       ctx,
-      size,
+      _size,
       _fgcolor,
-      bgcolor
+      _bgcolor
     ) {
       // @ts-expect-error fixme ts strict error
       const res = origDrawNodeShape.apply(this, arguments)
 
       const nodeErrors = self.lastNodeErrors?.[node.id]
 
-      // @ts-expect-error fixme ts strict error
-      if (self.progress && node.id === +self.runningNodeId) {
-        ctx.fillStyle = 'green'
-        ctx.fillRect(
-          0,
-          0,
-          size[0] * (self.progress.value / self.progress.max),
-          6
-        )
-        ctx.fillStyle = bgcolor
-      }
-
       // Highlight inputs that failed validation
       if (nodeErrors) {
         ctx.lineWidth = 2

src/stores/executionStore.ts

@@ -10,21 +10,29 @@ import type {
 } from '@/schemas/apiSchema'
 import type {
   ComfyNode,
-  ComfyWorkflowJSON
+  ComfyWorkflowJSON,
+  NodeId
 } from '@/schemas/comfyWorkflowSchema'
 import { api } from '@/scripts/api'
 
 import { ComfyWorkflow } from './workflowStore'
 
 export interface QueuedPrompt {
-  nodes: Record<string, boolean>
+  /**
+   * The nodes that are queued to be executed. The key is the node id and the
+   * value is a boolean indicating if the node has been executed.
+   */
+  nodes: Record<NodeId, boolean>
+  /**
+   * The workflow that is queued to be executed
+   */
   workflow?: ComfyWorkflow
 }
 
 export const useExecutionStore = defineStore('execution', () => {
   const clientId = ref<string | null>(null)
   const activePromptId = ref<string | null>(null)
-  const queuedPrompts = ref<Record<string, QueuedPrompt>>({})
+  const queuedPrompts = ref<Record<NodeId, QueuedPrompt>>({})
   const executingNodeId = ref<string | null>(null)
   const executingNode = computed<ComfyNode | null>(() => {
     if (!executingNodeId.value) return null
@@ -47,11 +55,7 @@ export const useExecutionStore = defineStore('execution', () => {
   const _executingNodeProgress = ref<ProgressWsMessage | null>(null)
   const executingNodeProgress = computed(() =>
     _executingNodeProgress.value
-      ? Math.round(
-          (_executingNodeProgress.value.value /
-            _executingNodeProgress.value.max) *
-            100
-        )
+      ? _executingNodeProgress.value.value / _executingNodeProgress.value.max
       : null
   )
 
@@ -75,7 +79,7 @@ export const useExecutionStore = defineStore('execution', () => {
     if (!activePrompt.value) return 0
     const total = totalNodesToExecute.value
     const done = nodesExecuted.value
-    return Math.round((done / total) * 100)
+    return done / total
   })
 
   function bindExecutionEvents() {
@@ -185,14 +189,41 @@ export const useExecutionStore = defineStore('execution', () => {
   return {
     isIdle,
     clientId,
+    /**
+     * The id of the prompt that is currently being executed
+     */
     activePromptId,
+    /**
+     * The queued prompts
+     */
     queuedPrompts,
+    /**
+     * The id of the node that is currently being executed
+     */
     executingNodeId,
+    /**
+     * The prompt that is currently being executed
+     */
     activePrompt,
+    /**
+     * The total number of nodes to execute
+     */
     totalNodesToExecute,
+    /**
+     * The number of nodes that have been executed
+     */
     nodesExecuted,
+    /**
+     * The progress of the execution
+     */
     executionProgress,
+    /**
+     * The node that is currently being executed
+     */
     executingNode,
+    /**
+     * The progress of the executing node (if the node reports progress)
+     */
     executingNodeProgress,
     bindExecutionEvents,
     unbindExecutionEvents,