From cce5ade578f9737b6d814ad7b0f7b0911012c3e3 Mon Sep 17 00:00:00 2001
From: AustinMroz <austinmroz@utexas.edu>
Date: Sun, 22 Dec 2024 15:18:31 -0600
Subject: [PATCH] Fix memory leak in preview code (#2012)

Co-authored-by: filtered <176114999+webfiltered@users.noreply.github.com>
---
 src/components/sidebar/tabs/queue/TaskItem.vue |  6 ++++++
 src/scripts/app.ts                             | 16 ++++++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/src/components/sidebar/tabs/queue/TaskItem.vue b/src/components/sidebar/tabs/queue/TaskItem.vue
index 70208d962b..9752d56eee 100644
--- a/src/components/sidebar/tabs/queue/TaskItem.vue
+++ b/src/components/sidebar/tabs/queue/TaskItem.vue
@@ -110,6 +110,9 @@ onMounted(() => {
 })
 
 onUnmounted(() => {
+  if (progressPreviewBlobUrl.value) {
+    URL.revokeObjectURL(progressPreviewBlobUrl.value)
+  }
   api.removeEventListener('b_preview', onProgressPreviewReceived)
 })
 
@@ -164,6 +167,9 @@ const formatTime = (time?: number) => {
 
 const onProgressPreviewReceived = async ({ detail }: CustomEvent) => {
   if (props.task.displayStatus === TaskItemDisplayStatus.Running) {
+    if (progressPreviewBlobUrl.value) {
+      URL.revokeObjectURL(progressPreviewBlobUrl.value)
+    }
     progressPreviewBlobUrl.value = URL.createObjectURL(detail)
   }
 }
diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index 1feac222ad..5464dc7fa4 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -1547,6 +1547,7 @@ export class ComfyApp {
     api.addEventListener('executing', ({ detail }) => {
       this.progress = null
       this.graph.setDirtyCanvas(true, false)
+      this.revokePreviews(this.runningNodeId)
       delete this.nodePreviewImages[this.runningNodeId]
     })
 
@@ -1589,6 +1590,8 @@ export class ComfyApp {
 
       const blob = detail
       const blobUrl = URL.createObjectURL(blob)
+      // Ensure clean up if `executing` event is missed.
+      this.revokePreviews(id)
       this.nodePreviewImages[id] = [blobUrl]
     })
 
@@ -2835,11 +2838,24 @@ export class ComfyApp {
     app.graph.setDirtyCanvas(true, true)
   }
 
+  /**
+   * Frees memory allocated to image preview blobs for a specific node, by revoking the URLs associated with them.
+   * @param nodeId ID of the node to revoke all preview images of
+   */
+  revokePreviews(nodeId: NodeId) {
+    if (!this.nodePreviewImages[nodeId]?.[Symbol.iterator]) return
+    for (const url of this.nodePreviewImages[nodeId]) {
+      URL.revokeObjectURL(url)
+    }
+  }
   /**
    * Clean current state
    */
   clean() {
     this.nodeOutputs = {}
+    for (const id of Object.keys(this.nodePreviewImages)) {
+      this.revokePreviews(id)
+    }
     this.nodePreviewImages = {}
     this.lastNodeErrors = null
     this.lastExecutionError = null