ComfyUI_frontend/src/composables/usePaste.ts

import { useEventListener } from '@vueuse/core'

import type { LGraphCanvas, LGraphNode } from '@/lib/litegraph/src/litegraph'
import { LiteGraph } from '@/lib/litegraph/src/litegraph'
import type { ComfyWorkflowJSON } from '@/platform/workflow/validation/schemas/workflowSchema'
import { useCanvasStore } from '@/renderer/core/canvas/canvasStore'
import { app } from '@/scripts/app'
import { useWorkspaceStore } from '@/stores/workspaceStore'
import { isAudioNode, isImageNode, isVideoNode } from '@/utils/litegraphUtil'
import { shouldIgnoreCopyPaste } from '@/workbench/eventHelpers'

/** A node that supports pasting files */
interface PasteableNode {
  pasteFile?(file: File): void
  pasteFiles?(files: File[]): void
}

function pasteClipboardItems(data: DataTransfer): boolean {
  const rawData = data.getData('text/html')
  const match = rawData.match(/data-metadata="([A-Za-z0-9+/=]+)"/)?.[1]
  if (!match) return false
  try {
    // Decode UTF-8 safe base64
    const binaryString = atob(match)
    const bytes = Uint8Array.from(binaryString, (c) => c.charCodeAt(0))
    const decodedData = new TextDecoder().decode(bytes)
    useCanvasStore().getCanvas()._deserializeItems(JSON.parse(decodedData), {})
    return true
  } catch (err) {
    console.error(err)
  }
  return false
}

function pasteItemsOnNode(
  items: DataTransferItemList,
  node: PasteableNode | null,
  contentType: string
): void {
  if (!node) return

  const filteredItems = Array.from(items).filter((item) =>
    item.type.startsWith(contentType)
  )

  const blob = filteredItems[0]?.getAsFile()
  if (!blob) return

  node.pasteFile?.(blob)
  node.pasteFiles?.(
    Array.from(filteredItems)
      .map((i) => i.getAsFile())
      .filter((f) => f !== null)
  )
}

export function pasteImageNode(
  canvas: LGraphCanvas,
  items: DataTransferItemList,
  imageNode: PasteableNode | null = null
): void {
  const {
    graph,
    graph_mouse: [posX, posY]
  } = canvas

  if (!imageNode) {
    // No image node selected: add a new one
    const newNode = LiteGraph.createNode('LoadImage')
    if (newNode) {
      newNode.pos = [posX, posY]
      imageNode = graph?.add(newNode) ?? null
    }
    graph?.change()
  }

  pasteItemsOnNode(items, imageNode, 'image')
}

/**
 * Adds a handler on paste that extracts and loads images or workflows from pasted JSON data
 */
export const usePaste = () => {
  const workspaceStore = useWorkspaceStore()
  const canvasStore = useCanvasStore()

  useEventListener(document, 'paste', async (e) => {
    if (shouldIgnoreCopyPaste(e.target)) {
      // Default system copy
      return
    }
    // ctrl+shift+v is used to paste nodes with connections
    // this is handled by litegraph
    if (workspaceStore.shiftDown) return

    const { canvas } = canvasStore
    if (!canvas) return

    const { graph } = canvas
    let data: DataTransfer | string | null = e.clipboardData
    if (!data) throw new Error('No clipboard data on clipboard event')

    const { items } = data

    const currentNode = canvas.current_node as LGraphNode
    const isNodeSelected = currentNode?.is_selected

    const isImageNodeSelected = isNodeSelected && isImageNode(currentNode)
    const isVideoNodeSelected = isNodeSelected && isVideoNode(currentNode)
    const isAudioNodeSelected = isNodeSelected && isAudioNode(currentNode)

    let audioNode: LGraphNode | null = isAudioNodeSelected ? currentNode : null
    const 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/')) {
        pasteImageNode(canvas as LGraphCanvas, 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 {
          pasteItemsOnNode(items, videoNode, 'video')
          return
        }
      } else if (item.type.startsWith('audio/')) {
        if (!audioNode) {
          // No audio node selected: add a new one
          const newNode = LiteGraph.createNode('LoadAudio')
          if (newNode) {
            newNode.pos = [canvas.graph_mouse[0], canvas.graph_mouse[1]]
            audioNode = graph?.add(newNode) ?? null
          }
          graph?.change()
        }
        pasteItemsOnNode(items, audioNode, 'audio')
        return
      }
    }
    if (pasteClipboardItems(data)) return

    // No image found. Look for node data
    data = data.getData('text/plain')
    let workflow: ComfyWorkflowJSON | null = null
    try {
      data = data.slice(data.indexOf('{'))
      workflow = JSON.parse(data)
    } catch (err) {
      try {
        data = data.slice(data.indexOf('workflow\n'))
        data = data.slice(data.indexOf('{'))
        workflow = JSON.parse(data)
      } catch (error) {
        workflow = null
      }
    }

    if (workflow && workflow.version && workflow.nodes && workflow.extra) {
      await app.loadGraphData(workflow)
    } else {
      if (
        (e.target instanceof HTMLTextAreaElement &&
          e.target.type === 'textarea') ||
        (e.target instanceof HTMLInputElement && e.target.type === 'text')
      ) {
        return
      }

      // Litegraph default paste
      canvas.pasteFromClipboard()
    }
  })
}