From f94831d0544efaab431e3f1198b2c810c35a317a Mon Sep 17 00:00:00 2001 From: bymyself Date: Sat, 22 Feb 2025 16:37:42 -0700 Subject: [PATCH] Add node video previews (#2635) --- src/assets/css/style.css | 6 + src/composables/useNodeImage.ts | 162 ++++++++++++++++-- src/composables/useNodeImageUpload.ts | 10 +- src/composables/usePaste.ts | 61 ++++--- .../widgets/useImageUploadWidget.ts | 21 ++- src/extensions/core/saveImageExtraOutput.ts | 6 +- src/extensions/core/uploadImage.ts | 17 +- src/services/litegraphService.ts | 74 ++------ src/stores/imagePreviewStore.ts | 55 +++--- src/types/litegraph-augmentation.d.ts | 4 + src/utils/formatUtil.ts | 10 +- src/utils/litegraphUtil.ts | 97 ++++++++++- 12 files changed, 374 insertions(+), 149 deletions(-) diff --git a/src/assets/css/style.css b/src/assets/css/style.css index baf0d0b7f..4bb3ea08f 100644 --- a/src/assets/css/style.css +++ b/src/assets/css/style.css @@ -532,6 +532,12 @@ dialog::backdrop { height: var(--comfy-img-preview-height); } +.comfy-img-preview video { + object-fit: contain; + height: 100%; + width: 100%; +} + .comfy-missing-nodes li button { font-size: 12px; margin-left: 5px; diff --git a/src/composables/useNodeImage.ts b/src/composables/useNodeImage.ts index a1a1b84bf..0919535ac 100644 --- a/src/composables/useNodeImage.ts +++ b/src/composables/useNodeImage.ts @@ -2,21 +2,163 @@ import type { LGraphNode } from '@comfyorg/litegraph' import { useNodeOutputStore } from '@/stores/imagePreviewStore' +const VIDEO_WIDGET_NAME = 'video-preview' as const +const VIDEO_DEFAULT_OPTIONS = { + playsInline: true, + controls: true, + loop: true +} as const +const MEDIA_LOAD_TIMEOUT = 8192 as const +const MAX_RETRIES = 1 as const + +type MediaElement = HTMLImageElement | HTMLVideoElement + +interface NodePreviewOptions { + loadElement: (url: string) => Promise + onLoaded?: (elements: T[]) => void + onFailedLoading?: () => void +} + +const createContainer = () => { + const container = document.createElement('div') + container.classList.add('comfy-img-preview') + return container +} + +const createTimeout = (ms: number) => + new Promise((resolve) => setTimeout(() => resolve(null), ms)) + +export const useNodePreview = ( + node: LGraphNode, + options: NodePreviewOptions +) => { + const { loadElement, onLoaded, onFailedLoading } = options + const nodeOutputStore = useNodeOutputStore() + + const loadElementWithTimeout = async ( + url: string, + retryCount = 0 + ): Promise => { + const result = await Promise.race([ + loadElement(url), + createTimeout(MEDIA_LOAD_TIMEOUT) + ]) + + if (result === null && retryCount < MAX_RETRIES) { + return loadElementWithTimeout(url, retryCount + 1) + } + + return result + } + + 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. + */ + function showPreview() { + if (node.isLoading) return + + const outputUrls = nodeOutputStore.getNodeImageUrls(node) + if (!outputUrls?.length) return + + node.isLoading = true + + loadElements(outputUrls) + .then((elements) => { + const validElements = elements.filter( + (el): el is NonNullable> => el !== null + ) + if (validElements.length) { + onLoaded?.(validElements) + render() + } + }) + .catch(() => { + onFailedLoading?.() + }) + .finally(() => { + node.isLoading = false + }) + } + + return { + showPreview + } +} + /** * Attaches a preview image to a node. */ export const useNodeImage = (node: LGraphNode) => { - const nodeOutputStore = useNodeOutputStore() + const loadElement = (url: string): Promise => + new Promise((resolve) => { + const img = new Image() + img.onload = () => resolve(img) + img.onerror = () => resolve(null) + img.src = url + }) - /** Displays output image(s) on the node. */ - function showImage(output: string | string[]) { - if (!output) return - nodeOutputStore.setNodeOutputs(node, output) - node.setSizeForImage?.() - node.graph?.setDirtyCanvas(true) + const onLoaded = (elements: HTMLImageElement[]) => { + node.imageIndex = null + node.imgs = elements } - return { - showImage - } + return useNodePreview(node, { + loadElement, + onLoaded, + onFailedLoading: () => { + node.imgs = undefined + } + }) +} + +/** + * Attaches a preview video to a node. + */ +export const useNodeVideo = (node: LGraphNode) => { + const loadElement = (url: string): Promise => + new Promise((resolve) => { + const video = document.createElement('video') + Object.assign(video, VIDEO_DEFAULT_OPTIONS) + video.onloadeddata = () => resolve(video) + video.onerror = () => resolve(null) + video.src = url + }) + + const addVideoDomWidget = (container: HTMLElement) => { + const hasWidget = node.widgets?.some((w) => w.name === VIDEO_WIDGET_NAME) + if (!hasWidget) { + node.addDOMWidget(VIDEO_WIDGET_NAME, 'video', container, { + hideOnZoom: false + }) + } + } + + const onLoaded = (videoElements: HTMLVideoElement[]) => { + const videoElement = videoElements[0] + if (!videoElement) return + + if (!node.videoContainer) { + node.videoContainer = createContainer() + addVideoDomWidget(node.videoContainer) + } + + node.videoContainer.replaceChildren(videoElement) + node.imageOffset = 64 + } + + return useNodePreview(node, { + loadElement, + onLoaded, + onFailedLoading: () => { + node.videoContainer = undefined + } + }) } diff --git a/src/composables/useNodeImageUpload.ts b/src/composables/useNodeImageUpload.ts index a353f536a..8e3b0fe81 100644 --- a/src/composables/useNodeImageUpload.ts +++ b/src/composables/useNodeImageUpload.ts @@ -6,7 +6,6 @@ import { useNodePaste } from '@/composables/useNodePaste' import { api } from '@/scripts/api' import { useToastStore } from '@/stores/toastStore' -const ACCEPTED_IMAGE_TYPES = 'image/jpeg,image/png,image/webp' const PASTED_IMAGE_EXPIRY_MS = 2000 const uploadFile = async (file: File, isPasted: boolean) => { @@ -32,6 +31,11 @@ interface ImageUploadOptions { fileFilter?: (file: File) => boolean onUploadComplete: (paths: string[]) => void allow_batch?: boolean + /** + * The file types to accept. + * @example 'image/png,image/jpeg,image/webp,video/webm,video/mp4' + */ + accept?: string } /** @@ -41,7 +45,7 @@ export const useNodeImageUpload = ( node: LGraphNode, options: ImageUploadOptions ) => { - const { fileFilter, onUploadComplete, allow_batch } = options + const { fileFilter, onUploadComplete, allow_batch, accept } = options const isPastedFile = (file: File): boolean => file.name === 'image.png' && @@ -81,7 +85,7 @@ export const useNodeImageUpload = ( const { openFileSelection } = useNodeFileInput({ fileFilter, allow_batch, - accept: ACCEPTED_IMAGE_TYPES, + accept, onSelect: handleUploadBatch }) diff --git a/src/composables/usePaste.ts b/src/composables/usePaste.ts index 9871daa8e..558f77a80 100644 --- a/src/composables/usePaste.ts +++ b/src/composables/usePaste.ts @@ -6,7 +6,7 @@ import { app } from '@/scripts/app' import { useCanvasStore } from '@/stores/graphStore' import { useWorkspaceStore } from '@/stores/workspaceStore' import { ComfyWorkflowJSON } from '@/types/comfyWorkflow' -import { isImageNode } from '@/utils/litegraphUtil' +import { isImageNode, isVideoNode } from '@/utils/litegraphUtil' /** * Adds a handler on paste that extracts and loads images or workflows from pasted JSON data @@ -15,6 +15,23 @@ export const usePaste = () => { const workspaceStore = useWorkspaceStore() const canvasStore = useCanvasStore() + const pasteItemOnNode = ( + items: DataTransferItemList, + node: LGraphNode | null + ) => { + if (!node) return + + const blob = items[0]?.getAsFile() + if (!blob) return + + node.pasteFile?.(blob) + node.pasteFiles?.( + Array.from(items) + .map((i) => i.getAsFile()) + .filter((f) => f !== null) + ) + } + useEventListener(document, 'paste', async (e: ClipboardEvent) => { // ctrl+shift+v is used to paste nodes with connections // this is handled by litegraph @@ -30,38 +47,38 @@ export const usePaste = () => { let data = e.clipboardData || window.clipboardData const items: DataTransferItemList = data.items + const currentNode = canvas.current_node as LGraphNode + const isNodeSelected = currentNode?.is_selected + + const isImageNodeSelected = isNodeSelected && isImageNode(currentNode) + const isVideoNodeSelected = isNodeSelected && isVideoNode(currentNode) + + let imageNode: LGraphNode | null = isImageNodeSelected ? currentNode : null + const videoNode: LGraphNode | null = isVideoNodeSelected + ? currentNode + : null + // Look for image paste data for (const item of items) { if (item.type.startsWith('image/')) { - let imageNode: LGraphNode | null = null - - // If an image node is selected, paste into it - const currentNode = canvas.current_node as LGraphNode - if ( - currentNode && - currentNode.is_selected && - isImageNode(currentNode) - ) { - imageNode = currentNode - } - - // No image node selected: add a new one if (!imageNode) { + // No image node selected: add a new one const newNode = LiteGraph.createNode('LoadImage') // @ts-expect-error array to Float32Array newNode.pos = [...canvas.graph_mouse] imageNode = graph.add(newNode) ?? null graph.change() } - const blob = item.getAsFile() - if (blob) imageNode?.pasteFile?.(blob) - - imageNode?.pasteFiles?.( - Array.from(items) - .map((i) => i.getAsFile()) - .filter((f) => f !== null) - ) + pasteItemOnNode(items, imageNode) return + } else if (item.type.startsWith('video/')) { + if (!videoNode) { + // No video node selected: add a new one + // TODO: when video node exists + } else { + pasteItemOnNode(items, videoNode) + return + } } } diff --git a/src/composables/widgets/useImageUploadWidget.ts b/src/composables/widgets/useImageUploadWidget.ts index 60b54027d..f76238895 100644 --- a/src/composables/widgets/useImageUploadWidget.ts +++ b/src/composables/widgets/useImageUploadWidget.ts @@ -1,20 +1,25 @@ import type { LGraphNode } from '@comfyorg/litegraph' import { IComboWidget } from '@comfyorg/litegraph/dist/types/widgets' -import { useNodeImage } from '@/composables/useNodeImage' +import { useNodeImage, useNodeVideo } from '@/composables/useNodeImage' import { useNodeImageUpload } from '@/composables/useNodeImageUpload' import { useValueTransform } from '@/composables/useValueTransform' import type { ComfyWidgetConstructor } from '@/scripts/widgets' +import { useNodeOutputStore } from '@/stores/imagePreviewStore' import type { ComfyApp } from '@/types' import type { InputSpec, ResultItem } from '@/types/apiTypes' import { createAnnotatedPath } from '@/utils/formatUtil' import { addToComboValues } from '@/utils/litegraphUtil' +const ACCEPTED_IMAGE_TYPES = 'image/png,image/jpeg,image/webp' +const ACCEPTED_VIDEO_TYPES = 'video/webm,video/mp4' + type InternalFile = string | ResultItem type InternalValue = InternalFile | InternalFile[] type ExposedValue = string | string[] const isImageFile = (file: File) => file.type.startsWith('image/') +const isVideoFile = (file: File) => file.type.startsWith('video/') const findFileComboWidget = (node: LGraphNode, inputName: string) => node.widgets!.find((w) => w.name === inputName) as IComboWidget & { @@ -30,9 +35,13 @@ export const useImageUploadWidget = () => { ) => { const inputOptions = inputData[1] ?? {} const { imageInputName, allow_batch, image_folder = 'input' } = inputOptions + const nodeOutputStore = useNodeOutputStore() - const { showImage } = useNodeImage(node) + const isVideo = !!inputOptions.video_upload + const accept = isVideo ? ACCEPTED_VIDEO_TYPES : ACCEPTED_IMAGE_TYPES + const { showPreview } = isVideo ? useNodeVideo(node) : useNodeImage(node) + const fileFilter = isVideo ? isVideoFile : isImageFile const fileComboWidget = findFileComboWidget(node, imageInputName) const initialFile = `${fileComboWidget.value}` const formatPath = (value: InternalFile) => @@ -56,7 +65,8 @@ export const useImageUploadWidget = () => { // Setup file upload handling const { openFileSelection } = useNodeImageUpload(node, { allow_batch, - fileFilter: isImageFile, + fileFilter, + accept, onUploadComplete: (output) => { output.forEach((path) => addToComboValues(fileComboWidget, path)) // @ts-expect-error litegraph combo value type does not support arrays yet @@ -78,7 +88,7 @@ export const useImageUploadWidget = () => { // Add our own callback to the combo widget to render an image when it changes const cb = node.callback fileComboWidget.callback = function (...args) { - showImage(fileComboWidget.value) + nodeOutputStore.setNodeOutputs(node, fileComboWidget.value) if (cb) return cb.apply(this, args) } @@ -86,7 +96,8 @@ export const useImageUploadWidget = () => { // The value isnt set immediately so we need to wait a moment // No change callbacks seem to be fired on initial setting of the value requestAnimationFrame(() => { - showImage(fileComboWidget.value) + nodeOutputStore.setNodeOutputs(node, fileComboWidget.value) + showPreview() }) return { widget: uploadWidget } diff --git a/src/extensions/core/saveImageExtraOutput.ts b/src/extensions/core/saveImageExtraOutput.ts index 3ebc5a7ed..5b7a8e14e 100644 --- a/src/extensions/core/saveImageExtraOutput.ts +++ b/src/extensions/core/saveImageExtraOutput.ts @@ -7,7 +7,11 @@ import { applyTextReplacements } from '../../scripts/utils' app.registerExtension({ name: 'Comfy.SaveImageExtraOutput', async beforeRegisterNodeDef(nodeType, nodeData, app) { - if (nodeData.name === 'SaveImage' || nodeData.name === 'SaveAnimatedWEBP') { + if ( + nodeData.name === 'SaveImage' || + nodeData.name === 'SaveAnimatedWEBP' || + nodeData.name === 'SaveAnimatedWEBM' + ) { const onNodeCreated = nodeType.prototype.onNodeCreated // When the SaveImage node is created we want to override the serialization of the output name widget to run our S&R nodeType.prototype.onNodeCreated = function () { diff --git a/src/extensions/core/uploadImage.ts b/src/extensions/core/uploadImage.ts index e3bb7c12d..3b6ab441c 100644 --- a/src/extensions/core/uploadImage.ts +++ b/src/extensions/core/uploadImage.ts @@ -4,10 +4,17 @@ import { app } from '../../scripts/app' // Adds an upload button to the nodes -const isImageComboInput = (inputSpec: InputSpec) => { +const isMediaUploadComboInput = (inputSpec: InputSpec) => { const [inputName, inputOptions] = inputSpec - if (!inputOptions || inputOptions['image_upload'] !== true) return false - return isComboInputSpecV1(inputSpec) || inputName === 'COMBO' + if (!inputOptions) return false + + const isUploadInput = + inputOptions['image_upload'] === true || + inputOptions['video_upload'] === true + + return ( + isUploadInput && (isComboInputSpecV1(inputSpec) || inputName === 'COMBO') + ) } const createUploadInput = ( @@ -29,10 +36,10 @@ app.registerExtension({ if (!required) return const found = Object.entries(required).find(([_, input]) => - isImageComboInput(input) + isMediaUploadComboInput(input) ) - // If image combo input found, attach upload input + // If media combo input found, attach upload input if (found) { const [inputName, inputSpec] = found required.upload = createUploadInput(inputName, inputSpec) diff --git a/src/services/litegraphService.ts b/src/services/litegraphService.ts index da132ec1e..d7c56b183 100644 --- a/src/services/litegraphService.ts +++ b/src/services/litegraphService.ts @@ -9,6 +9,7 @@ import { import { Vector2 } from '@comfyorg/litegraph' import { IBaseWidget, IWidget } from '@comfyorg/litegraph/dist/types/widgets' +import { useNodeImage, useNodeVideo } from '@/composables/useNodeImage' import { st } from '@/i18n' import { ANIM_PREVIEW_WIDGET, ComfyApp, app } from '@/scripts/app' import { $el } from '@/scripts/ui' @@ -20,7 +21,7 @@ import { ComfyNodeDef } from '@/types/apiTypes' import type { NodeId } from '@/types/comfyWorkflow' import { normalizeI18nKey } from '@/utils/formatUtil' import { is_all_same_aspect_ratio } from '@/utils/imageUtil' -import { isImageNode } from '@/utils/litegraphUtil' +import { getImageTop, isImageNode, isVideoNode } from '@/utils/litegraphUtil' import { useExtensionService } from './extensionService' @@ -363,26 +364,6 @@ export const useLitegraphService = () => { * @param {*} node The node to add the draw handler */ function addDrawBackgroundHandler(node: typeof LGraphNode) { - function getImageTop(node: LGraphNode) { - let shiftY: number - if (node.imageOffset != null) { - return node.imageOffset - } else if (node.widgets?.length) { - const w = node.widgets[node.widgets.length - 1] - shiftY = w.last_y - if (w.computeSize) { - shiftY += w.computeSize()[1] + 4 - } else if (w.computedHeight) { - shiftY += w.computedHeight - } else { - shiftY += LiteGraph.NODE_WIDGET_HEIGHT + 4 - } - } else { - return node.computeSize()[1] - } - return shiftY - } - node.prototype.setSizeForImage = function ( this: LGraphNode, force: boolean @@ -405,49 +386,24 @@ export const useLitegraphService = () => { ) { if (this.flags.collapsed) return - let imgURLs: string[] = [] const nodeOutputStore = useNodeOutputStore() + const output = nodeOutputStore.getNodeOutputs(this) - let imagesChanged = nodeOutputStore.isImagesChanged(this) - if (output && imagesChanged) { - this.animatedImages = output.animated?.find(Boolean) - this.images = output.images - imgURLs = nodeOutputStore.getNodeImageUrls(this) - } - const preview = nodeOutputStore.getNodePreviews(this) - if (this.preview !== preview) { - this.preview = preview - imagesChanged = true - if (preview != null) { - imgURLs.push(...preview) - } - } - if (imagesChanged) { - this.imageIndex = null - if (imgURLs.length > 0) { - Promise.all( - imgURLs.flat().map((src) => { - return new Promise((r) => { - const img = new Image() - img.onload = () => r(img) - img.onerror = () => r(null) - img.src = src - }) - }) - ).then((imgs) => { - if ( - (!output || this.images === output.images) && - (!preview || this.preview === preview) - ) { - this.imgs = imgs.filter(Boolean) - this.setSizeForImage?.() - app.graph.setDirtyCanvas(true) - } - }) + const isNewOutput = output && this.images !== output.images + const isNewPreview = preview && this.preview !== preview + + if (isNewPreview) this.preview = preview + if (isNewOutput) this.images = output.images + + if (isNewOutput || isNewPreview) { + this.animatedImages = output?.animated?.find(Boolean) + + if (this.animatedImages || isVideoNode(this)) { + useNodeVideo(this).showPreview() } else { - this.imgs = undefined + useNodeImage(this).showPreview() } } diff --git a/src/stores/imagePreviewStore.ts b/src/stores/imagePreviewStore.ts index 449c9fbc4..a4336bd1b 100644 --- a/src/stores/imagePreviewStore.ts +++ b/src/stores/imagePreviewStore.ts @@ -5,7 +5,7 @@ import { api } from '@/scripts/api' import { ExecutedWsMessage, ResultItem } from '@/types/apiTypes' import { parseFilePath } from '@/utils/formatUtil' -const toOutputs = ( +const createOutputs = ( filenames: string[], type: string ): ExecutedWsMessage['output'] => { @@ -14,62 +14,56 @@ const toOutputs = ( } } -const getPreviewParam = (node: LGraphNode) => { +const getPreviewParam = (node: LGraphNode): string => { if (node.animatedImages) return '' return app.getPreviewFormatParam() } export const useNodeOutputStore = defineStore('nodeOutput', () => { - function getNodeOutputs(node: LGraphNode): ExecutedWsMessage['output'] { - return app.nodeOutputs[node.id + ''] + const getNodeId = (node: LGraphNode): string => node.id.toString() + + function getNodeOutputs( + node: LGraphNode + ): ExecutedWsMessage['output'] | undefined { + return app.nodeOutputs[getNodeId(node)] } - function getNodePreviews(node: LGraphNode): string[] { - return app.nodePreviewImages[node.id + ''] + function getNodePreviews(node: LGraphNode): string[] | undefined { + return app.nodePreviewImages[getNodeId(node)] } - function getNodeImageUrls(node: LGraphNode): string[] { + function getNodeImageUrls(node: LGraphNode): string[] | undefined { + const previews = getNodePreviews(node) + if (previews?.length) return previews + const outputs = getNodeOutputs(node) - if (!outputs?.images?.length) return [] + if (!outputs?.images?.length) return + + const rand = app.getRandParam() + const previewParam = getPreviewParam(node) return outputs.images.map((image) => { const imgUrlPart = new URLSearchParams(image) - const rand = app.getRandParam() - const previewParam = getPreviewParam(node) return api.apiURL(`/view?${imgUrlPart}${previewParam}${rand}`) }) } - /** - * Checks if the node's images have changed from what's stored - * @returns true if images have changed, false otherwise - */ - function isImagesChanged(node: LGraphNode): boolean { - const currentImages = node.images || [] - const { images: newImages } = getNodeOutputs(node) ?? {} - if (!newImages?.length) return false - - return currentImages !== newImages - } - function setNodeOutputs( node: LGraphNode, filenames: string | string[] | ResultItem, - options: { folder?: string } = {} + { folder = 'input' }: { folder?: string } = {} ) { - if (!filenames) return + if (!filenames || !node) return - const { folder = 'input' } = options - const nodeId = node.id + '' + const nodeId = getNodeId(node) if (typeof filenames === 'string') { - app.nodeOutputs[nodeId] = toOutputs([filenames], folder) + app.nodeOutputs[nodeId] = createOutputs([filenames], folder) } else if (!Array.isArray(filenames)) { app.nodeOutputs[nodeId] = filenames } else { - const resultItems = toOutputs(filenames, folder) + const resultItems = createOutputs(filenames, folder) if (!resultItems?.images?.length) return - app.nodeOutputs[nodeId] = resultItems } } @@ -78,7 +72,6 @@ export const useNodeOutputStore = defineStore('nodeOutput', () => { getNodeOutputs, getNodeImageUrls, getNodePreviews, - setNodeOutputs, - isImagesChanged + setNodeOutputs } }) diff --git a/src/types/litegraph-augmentation.d.ts b/src/types/litegraph-augmentation.d.ts index 4ee151dbe..51f50b9d6 100644 --- a/src/types/litegraph-augmentation.d.ts +++ b/src/types/litegraph-augmentation.d.ts @@ -111,6 +111,10 @@ declare module '@comfyorg/litegraph' { animatedImages?: boolean imgs?: HTMLImageElement[] images?: ExecutedWsMessage['output'] + /** Container for the node's video preview */ + videoContainer?: HTMLElement + /** Whether the node's preview media is loading */ + isLoading?: boolean preview: string[] /** Index of the currently selected image on a multi-image node such as Preview Image */ diff --git a/src/utils/formatUtil.ts b/src/utils/formatUtil.ts index fe7cc81ad..71344fed4 100644 --- a/src/utils/formatUtil.ts +++ b/src/utils/formatUtil.ts @@ -215,9 +215,11 @@ export function isValidUrl(url: string): boolean { return false } } +const hasAnnotation = (filepath: string): boolean => + /\[(input|output|temp)\]/i.test(filepath) -const createAnnotation = (rootFolder = 'input'): string => - rootFolder !== 'input' ? ` [${rootFolder}]` : '' +const createAnnotation = (filepath: string, rootFolder = 'input'): string => + !hasAnnotation(filepath) && rootFolder !== 'input' ? ` [${rootFolder}]` : '' const createPath = (filename: string, subfolder = ''): string => subfolder ? `${subfolder}/${filename}` : filename @@ -229,8 +231,8 @@ export function createAnnotatedPath( ): string { const { rootFolder = 'input', subfolder } = options if (typeof item === 'string') - return `${createPath(item, subfolder)}${createAnnotation(rootFolder)}` - return `${createPath(item.filename ?? '', item.subfolder)}${createAnnotation(item.type)}` + return `${createPath(item, subfolder)}${createAnnotation(item, rootFolder)}` + return `${createPath(item.filename ?? '', item.subfolder)}${item.type ? createAnnotation(item.type, rootFolder) : ''}` } /** diff --git a/src/utils/litegraphUtil.ts b/src/utils/litegraphUtil.ts index 187293b8d..4fedcf6cd 100644 --- a/src/utils/litegraphUtil.ts +++ b/src/utils/litegraphUtil.ts @@ -1,15 +1,74 @@ -import type { ColorOption, IWidget } from '@comfyorg/litegraph' -import { LGraphGroup, LGraphNode, isColorable } from '@comfyorg/litegraph' +import type { ColorOption } from '@comfyorg/litegraph' +import { + LGraphGroup, + LGraphNode, + LiteGraph, + isColorable +} from '@comfyorg/litegraph' import type { IComboWidget } from '@comfyorg/litegraph/dist/types/widgets' import _ from 'lodash' -export function isImageNode(node: LGraphNode) { - return ( - node.imgs || - (node && - node.widgets && - node.widgets.findIndex((obj: IWidget) => obj.name === 'image') >= 0) - ) +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 +} + +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) +} + +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) } export function addToComboValues(widget: IComboWidget, value: string) { @@ -56,3 +115,23 @@ export function executeWidgetsCallback( } } } + +export function getImageTop(node: LGraphNode) { + let shiftY: number + if (node.imageOffset != null) { + return node.imageOffset + } else if (node.widgets?.length) { + const w = node.widgets[node.widgets.length - 1] + shiftY = w.last_y ?? 0 + if (w.computeSize) { + shiftY += w.computeSize()[1] + 4 + } else if (w.computedHeight) { + shiftY += w.computedHeight + } else { + shiftY += LiteGraph.NODE_WIDGET_HEIGHT + 4 + } + } else { + return node.computeSize()[1] + } + return shiftY +}