Add previewMediaType flag for simpler node preview rendering (#2694)

Co-authored-by: huchenlei <huchenlei@proton.me>

src/composables/useNodeImage.ts

@@ -10,6 +10,7 @@ const VIDEO_DEFAULT_OPTIONS = {
 } as const
 const MEDIA_LOAD_TIMEOUT = 8192 as const
 const MAX_RETRIES = 1 as const
+const VIDEO_PIXEL_OFFSET = 64 as const
 
 type MediaElement = HTMLImageElement | HTMLVideoElement
 
@@ -54,11 +55,6 @@ export const useNodePreview = <T extends MediaElement>(
   const loadElements = async (urls: string[]) =>
     Promise.all(urls.map((url) => loadElementWithTimeout(url)))
 
-  const render = () => {
-    node.setSizeForImage?.()
-    node.graph?.setDirtyCanvas(true)
-  }
-
   /**
    * Displays media element(s) on the node.
    */
@@ -77,7 +73,7 @@ export const useNodePreview = <T extends MediaElement>(
         )
         if (validElements.length) {
           onLoaded?.(validElements)
-          render()
+          node.graph?.setDirtyCanvas(true)
         }
       })
       .catch(() => {
@@ -97,6 +93,8 @@ export const useNodePreview = <T extends MediaElement>(
  * Attaches a preview image to a node.
  */
 export const useNodeImage = (node: LGraphNode) => {
+  node.previewMediaType = 'image'
+
   const loadElement = (url: string): Promise<HTMLImageElement | null> =>
     new Promise((resolve) => {
       const img = new Image()
@@ -108,6 +106,7 @@ export const useNodeImage = (node: LGraphNode) => {
   const onLoaded = (elements: HTMLImageElement[]) => {
     node.imageIndex = null
     node.imgs = elements
+    node.setSizeForImage?.()
   }
 
   return useNodePreview(node, {
@@ -123,6 +122,8 @@ export const useNodeImage = (node: LGraphNode) => {
  * Attaches a preview video to a node.
  */
 export const useNodeVideo = (node: LGraphNode) => {
+  node.previewMediaType = 'video'
+
   const loadElement = (url: string): Promise<HTMLVideoElement | null> =>
     new Promise((resolve) => {
       const video = document.createElement('video')
@@ -152,7 +153,8 @@ export const useNodeVideo = (node: LGraphNode) => {
     }
 
     node.videoContainer.replaceChildren(videoElement)
-    node.imageOffset = 64
+    node.imageOffset = VIDEO_PIXEL_OFFSET
+    node.setSizeForImage?.(true)
   }
 
   return useNodePreview(node, {

src/extensions/core/uploadAudio.ts

@@ -199,6 +199,8 @@ app.registerExtension({
         )
         uploadWidget.label = 'choose file to upload'
 
+        node.previewMediaType = 'audio'
+
         return { widget: uploadWidget }
       }
     }

src/types/litegraph-augmentation.d.ts

@@ -129,6 +129,8 @@ declare module '@comfyorg/litegraph' {
     videoContainer?: HTMLElement
     /** Whether the node's preview media is loading */
     isLoading?: boolean
+    /** The content type of the node's preview media */
+    previewMediaType?: 'image' | 'video' | 'audio' | 'model'
 
     preview: string[]
     /** Index of the currently selected image on a multi-image node such as Preview Image */
@@ -141,7 +143,7 @@ declare module '@comfyorg/litegraph' {
     /** @deprecated Unused */
     inputHeight?: unknown
 
-    /** @deprecated Unused */
+    /** The y offset of the image preview to the top of the node body. */
     imageOffset?: number
     /** Callback for pasting an image file into the node */
     pasteFile?(file: File): void

src/utils/litegraphUtil.ts

@@ -8,67 +8,23 @@ import {
 import type { IComboWidget } from '@comfyorg/litegraph/dist/types/widgets'
 import _ from 'lodash'
 
-import { ComfyInputsSpec, ComfyNodeDef, InputSpec } from '@/types/apiTypes'
-
-const IMAGE_NODE_PROPERTY = 'image_upload'
-const VIDEO_NODE_PROPERTY = 'video_upload'
-
-const getNodeData = (node: LGraphNode): ComfyNodeDef | undefined =>
-  node.constructor?.nodeData as ComfyNodeDef | undefined
-
-const getInputSpecsFromData = (
-  inputData: ComfyInputsSpec | undefined
-): InputSpec[] => {
-  if (!inputData) return []
-
-  const { required, optional } = inputData
-  const inputSpecs: InputSpec[] = []
-  if (required) {
-    for (const value of Object.values(required)) {
-      inputSpecs.push(value)
-    }
-  }
-  if (optional) {
-    for (const value of Object.values(optional)) {
-      inputSpecs.push(value)
-    }
-  }
-  return inputSpecs
+type ImageNode = LGraphNode & { imgs: HTMLImageElement[] | undefined }
+type VideoNode = LGraphNode & {
+  videoContainer: HTMLElement | undefined
+  imgs: HTMLVideoElement[] | undefined
 }
 
-const hasImageElements = (imgs: unknown[]): boolean =>
-  Array.isArray(imgs) &&
-  imgs.some((img): img is HTMLImageElement => img instanceof HTMLImageElement)
-
-const hasInputProperty = (
-  node: LGraphNode | undefined,
-  property: string
-): boolean => {
-  if (!node) return false
-  const nodeData = getNodeData(node)
-  if (!nodeData?.input) return false
-
-  const inputs = getInputSpecsFromData(nodeData.input)
-  return inputs.some((input) => input?.[1]?.[property])
-}
-
-type ImageNode = LGraphNode & { imgs: HTMLImageElement[] }
-type VideoNode = LGraphNode & { videoContainer: HTMLElement }
-
 export function isImageNode(node: LGraphNode | undefined): node is ImageNode {
   if (!node) return false
-  if (node.imgs?.length && hasImageElements(node.imgs)) return true
-  if (!node.widgets) return false
-
-  return hasInputProperty(node, IMAGE_NODE_PROPERTY)
+  return (
+    node.previewMediaType === 'image' ||
+    (node.previewMediaType !== 'video' && !!node.imgs?.length)
+  )
 }
 
 export function isVideoNode(node: LGraphNode | undefined): node is VideoNode {
   if (!node) return false
-  if (node.videoContainer) return true
-  if (!node.widgets) return false
-
-  return hasInputProperty(node, VIDEO_NODE_PROPERTY)
+  return node.previewMediaType === 'video' || !!node.videoContainer
 }
 
 export function addToComboValues(widget: IComboWidget, value: string) {