From 37d4cc974b03886ad6d9c7d5864ceb1b37373361 Mon Sep 17 00:00:00 2001
From: bymyself <cbyrne@comfy.org>
Date: Mon, 24 Feb 2025 08:45:27 -0700
Subject: [PATCH] Cleanup node's file input elements on node removal (#2703)

Co-authored-by: huchenlei <huchenlei@proton.me>
---
 src/composables/useNodeFileInput.ts   | 28 +++++++++++++--------------
 src/composables/useNodeImageUpload.ts |  2 +-
 src/extensions/core/uploadAudio.ts    | 23 +++++++++++-----------
 3 files changed, 26 insertions(+), 27 deletions(-)

diff --git a/src/composables/useNodeFileInput.ts b/src/composables/useNodeFileInput.ts
index 6902bd5b1..5b8982b94 100644
--- a/src/composables/useNodeFileInput.ts
+++ b/src/composables/useNodeFileInput.ts
@@ -1,3 +1,7 @@
+import type { LGraphNode } from '@comfyorg/litegraph'
+
+import { useChainCallback } from './functional/useChainCallback'
+
 interface FileInputOptions {
   accept?: string
   allow_batch?: boolean
@@ -8,7 +12,7 @@ interface FileInputOptions {
 /**
  * Creates a file input for a node.
  */
-export function useNodeFileInput(options: FileInputOptions) {
+export function useNodeFileInput(node: LGraphNode, options: FileInputOptions) {
   const {
     accept,
     allow_batch = false,
@@ -16,30 +20,26 @@ export function useNodeFileInput(options: FileInputOptions) {
     onSelect
   } = options
 
-  const fileInput = document.createElement('input')
+  let fileInput: HTMLInputElement | null = document.createElement('input')
   fileInput.type = 'file'
   fileInput.accept = accept ?? '*'
   fileInput.multiple = allow_batch
-  fileInput.style.visibility = 'hidden'
 
   fileInput.onchange = () => {
-    if (fileInput.files?.length) {
+    if (fileInput?.files?.length) {
       const files = Array.from(fileInput.files).filter(fileFilter)
       if (files.length) onSelect(files)
     }
   }
 
-  document.body.append(fileInput)
-
-  /**
-   * Shows the system file picker dialog for selecting files.
-   */
-  function openFileSelection() {
-    fileInput.click()
-  }
+  node.onRemoved = useChainCallback(node.onRemoved, () => {
+    if (fileInput) {
+      fileInput.onchange = null
+      fileInput = null
+    }
+  })
 
   return {
-    fileInput,
-    openFileSelection
+    openFileSelection: () => fileInput?.click()
   }
 }
diff --git a/src/composables/useNodeImageUpload.ts b/src/composables/useNodeImageUpload.ts
index 8e3b0fe81..47e8ae198 100644
--- a/src/composables/useNodeImageUpload.ts
+++ b/src/composables/useNodeImageUpload.ts
@@ -82,7 +82,7 @@ export const useNodeImageUpload = (
   })
 
   // Handle file input
-  const { openFileSelection } = useNodeFileInput({
+  const { openFileSelection } = useNodeFileInput(node, {
     fileFilter,
     allow_batch,
     accept,
diff --git a/src/extensions/core/uploadAudio.ts b/src/extensions/core/uploadAudio.ts
index 985dcb82b..fd903bcc2 100644
--- a/src/extensions/core/uploadAudio.ts
+++ b/src/extensions/core/uploadAudio.ts
@@ -2,6 +2,7 @@
 import type { IWidget } from '@comfyorg/litegraph'
 import type { IStringWidget } from '@comfyorg/litegraph/dist/types/widgets'
 
+import { useNodeFileInput } from '@/composables/useNodeFileInput'
 import type { DOMWidget } from '@/scripts/domWidget'
 import { useToastStore } from '@/stores/toastStore'
 import { ComfyNodeDef } from '@/types/apiTypes'
@@ -179,23 +180,21 @@ app.registerExtension({
           }
         }
 
-        const fileInput = document.createElement('input')
-        fileInput.type = 'file'
-        fileInput.accept = 'audio/*'
-        fileInput.style.display = 'none'
-        fileInput.onchange = () => {
-          if (fileInput.files.length) {
-            uploadFile(audioWidget, audioUIWidget, fileInput.files[0], true)
+        const { openFileSelection } = useNodeFileInput(node, {
+          accept: 'audio/*',
+          onSelect: (files) => {
+            if (files?.length) {
+              uploadFile(audioWidget, audioUIWidget, files[0], true)
+            }
           }
-        }
+        })
+
         // The widget to pop up the upload dialog.
         const uploadWidget = node.addWidget(
           'button',
           inputName,
-          /* value=*/ '',
-          () => {
-            fileInput.click()
-          },
+          '',
+          openFileSelection,
           { serialize: false }
         )
         uploadWidget.label = 'choose file to upload'