Dom widget store (#2899)

src/components/graph/DomWidgets.vue

@@ -0,0 +1,22 @@
+<template>
+  <div>
+    <DomWidget v-for="widget in widgets" :key="widget.id" :widget="widget" />
+  </div>
+</template>
+
+<script setup lang="ts">
+import { computed } from 'vue'
+
+import DomWidget from '@/components/graph/widgets/DomWidget.vue'
+import { DOMWidget } from '@/scripts/domWidget'
+import { useDomWidgetStore } from '@/stores/domWidgetStore'
+
+const domWidgetStore = useDomWidgetStore()
+const widgets = computed(() =>
+  Array.from(
+    domWidgetStore.widgetInstances.values() as Iterable<
+      DOMWidget<HTMLElement, object | string>
+    >
+  )
+)
+</script>

src/components/graph/GraphCanvas.vue

@@ -33,6 +33,7 @@
   </SelectionOverlay>
   <NodeTooltip v-if="tooltipEnabled" />
   <NodeBadge />
+  <DomWidgets />
 </template>
 
 <script setup lang="ts">
@@ -40,6 +41,7 @@ import { computed, onMounted, ref, watch, watchEffect } from 'vue'
 
 import LiteGraphCanvasSplitterOverlay from '@/components/LiteGraphCanvasSplitterOverlay.vue'
 import BottomPanel from '@/components/bottomPanel/BottomPanel.vue'
+import DomWidgets from '@/components/graph/DomWidgets.vue'
 import GraphCanvasMenu from '@/components/graph/GraphCanvasMenu.vue'
 import NodeBadge from '@/components/graph/NodeBadge.vue'
 import NodeTooltip from '@/components/graph/NodeTooltip.vue'

src/components/graph/widgets/DomWidget.vue

@@ -0,0 +1,19 @@
+<template>
+  <div ref="widgetElement" />
+</template>
+
+<script setup lang="ts">
+import { onMounted, ref } from 'vue'
+
+import type { DOMWidget } from '@/scripts/domWidget'
+
+const { widget } = defineProps<{
+  widget: DOMWidget<HTMLElement, any>
+}>()
+
+const widgetElement = ref<HTMLElement>()
+
+onMounted(() => {
+  widgetElement.value.appendChild(widget.element)
+})
+</script>

src/scripts/domWidget.ts

@@ -1,5 +1,5 @@
 import { LGraphCanvas, LGraphNode } from '@comfyorg/litegraph'
-import type { Size, Vector4 } from '@comfyorg/litegraph'
+import type { Vector4 } from '@comfyorg/litegraph'
 import type { ISerialisedNode } from '@comfyorg/litegraph/dist/types/serialisation'
 import type {
   ICustomWidget,
@@ -8,7 +8,9 @@ import type {
 
 import { useChainCallback } from '@/composables/functional/useChainCallback'
 import { app } from '@/scripts/app'
+import { useDomWidgetStore } from '@/stores/domWidgetStore'
 import { useSettingStore } from '@/stores/settingStore'
+import { generateRandomSuffix } from '@/utils/formatUtil'
 
 interface Rect {
   height: number
@@ -19,6 +21,7 @@ interface Rect {
 
 export interface DOMWidget<T extends HTMLElement, V extends object | string>
   extends ICustomWidget<T> {
+  // ICustomWidget properties
   type: 'custom'
   element: T
   options: DOMWidgetOptions<T, V>
@@ -30,6 +33,9 @@ export interface DOMWidget<T extends HTMLElement, V extends object | string>
    */
   inputEl?: T
   callback?: (value: V) => void
+  // DOMWidget properties
+  /** The unique ID of the widget. */
+  id: string
 }
 
 export interface DOMWidgetOptions<
@@ -146,30 +152,35 @@ LGraphCanvas.prototype.computeVisibleNodes = function (
 export class DOMWidgetImpl<T extends HTMLElement, V extends object | string>
   implements DOMWidget<T, V>
 {
-  type: 'custom'
-  name: string
-  element: T
-  options: DOMWidgetOptions<T, V>
+  readonly type: 'custom'
+  readonly name: string
+  readonly element: T
+  readonly options: DOMWidgetOptions<T, V>
   computedHeight?: number
   callback?: (value: V) => void
-  private mouseDownHandler?: (event: MouseEvent) => void
 
-  constructor(
-    name: string,
-    type: string,
-    element: T,
-    options: DOMWidgetOptions<T, V> = {}
-  ) {
+  readonly id: string
+  mouseDownHandler?: (event: MouseEvent) => void
+
+  constructor(obj: {
+    id: string
+    name: string
+    type: string
+    element: T
+    options: DOMWidgetOptions<T, V>
+  }) {
     // @ts-expect-error custom widget type
-    this.type = type
-    this.name = name
-    this.element = element
-    this.options = options
+    this.type = obj.type
+    this.name = obj.name
+    this.element = obj.element
+    this.options = obj.options
 
-    if (element.blur) {
+    this.id = obj.id
+
+    if (this.element.blur) {
       this.mouseDownHandler = (event) => {
-        if (!element.contains(event.target as HTMLElement)) {
-          element.blur()
+        if (!this.element.contains(event.target as HTMLElement)) {
+          this.element.blur()
         }
       }
       document.addEventListener('mousedown', this.mouseDownHandler)
@@ -304,9 +315,6 @@ LGraphNode.prototype.addDOMWidget = function <
 ): DOMWidget<T, V> {
   options = { hideOnZoom: true, selectOn: ['focus', 'click'], ...options }
 
-  if (!element.parentElement) {
-    app.canvasContainer.append(element)
-  }
   element.hidden = true
   element.style.display = 'none'
 
@@ -317,7 +325,14 @@ LGraphNode.prototype.addDOMWidget = function <
     element.title = tooltip
   }
 
-  const widget = new DOMWidgetImpl(name, type, element, options)
+  const widget = new DOMWidgetImpl({
+    id: `${this.id}:${name}:${generateRandomSuffix()}`,
+    name,
+    type,
+    element,
+    options
+  })
+
   // Workaround for https://github.com/Comfy-Org/ComfyUI_frontend/issues/2493
   // Some custom nodes are explicitly expecting getter and setter of `value`
   // property to be on instance instead of prototype.
@@ -374,5 +389,7 @@ LGraphNode.prototype.addDOMWidget = function <
     options.afterResize?.call(widget, this)
   })
 
+  useDomWidgetStore().registerWidget(widget)
+
   return widget
 }

src/services/litegraphService.ts

@@ -424,6 +424,7 @@ export const useLitegraphService = () => {
         } else {
           const host = createImageHost(this)
           this.setSizeForImage(true)
+          // @ts-expect-error host is not a standard DOM widget option.
           const widget = this.addDOMWidget(
             ANIM_PREVIEW_WIDGET,
             'img',

src/stores/domWidgetStore.ts

@@ -0,0 +1,36 @@
+/**
+ * Stores all DOM widgets that are used in the canvas.
+ */
+import { defineStore } from 'pinia'
+import { markRaw, ref } from 'vue'
+
+import type { DOMWidget } from '@/scripts/domWidget'
+
+export const useDomWidgetStore = defineStore('domWidget', () => {
+  // Map to reference actual widget instances
+  // Widgets are stored as raw values to avoid reactivity issues
+  const widgetInstances = ref(
+    new Map<string, DOMWidget<HTMLElement, object | string>>()
+  )
+
+  // Register a widget with the store
+  const registerWidget = <T extends HTMLElement, V extends object | string>(
+    widget: DOMWidget<T, V>
+  ) => {
+    widgetInstances.value.set(
+      widget.id,
+      markRaw(widget as unknown as DOMWidget<HTMLElement, object | string>)
+    )
+  }
+
+  // Unregister a widget from the store
+  const unregisterWidget = (widgetId: string) => {
+    widgetInstances.value.delete(widgetId)
+  }
+
+  return {
+    widgetInstances,
+    registerWidget,
+    unregisterWidget
+  }
+})

src/utils/formatUtil.ts

@@ -312,3 +312,9 @@ export const paramsToCacheKey = (params: unknown): string => {
 
   return String(params)
 }
+
+/**
+ * Generates a random 4-character string to use as a unique suffix
+ */
+export const generateRandomSuffix = (): string =>
+  Math.random().toString(36).substring(2, 6)