ComfyUI_frontend/src/components/NodeSearchBoxPopover.vue

<template>
  <div>
    <Dialog
      v-model:visible="visible"
      pt:root="invisible-dialog-root"
      pt:mask="node-search-box-dialog-mask"
      modal
      :dismissable-mask="dismissable"
      @hide="clearFilters"
    >
      <template #container>
        <NodeSearchBox
          :filters="nodeFilters"
          @add-filter="addFilter"
          @remove-filter="removeFilter"
          @add-node="addNode"
        />
      </template>
    </Dialog>
  </div>
</template>

<script setup lang="ts">
import { app } from '@/scripts/app'
import { onMounted, onUnmounted, reactive, ref } from 'vue'
import NodeSearchBox from './NodeSearchBox.vue'
import Dialog from 'primevue/dialog'
import {
  INodeSlot,
  LiteGraphCanvasEvent,
  LGraphNode,
  LinkReleaseContext
} from '@comfyorg/litegraph'
import { FilterAndValue } from '@/services/nodeSearchService'
import { ComfyNodeDefImpl, useNodeDefStore } from '@/stores/nodeDefStore'

interface LiteGraphPointerEvent extends Event {
  canvasX: number
  canvasY: number
}

const visible = ref(false)
const dismissable = ref(true)
const triggerEvent = ref<LiteGraphCanvasEvent | null>(null)
const getNewNodeLocation = (): [number, number] => {
  if (triggerEvent.value === null) {
    return [100, 100]
  }

  const originalEvent = triggerEvent.value.detail
    .originalEvent as LiteGraphPointerEvent
  return [originalEvent.canvasX, originalEvent.canvasY]
}
const nodeFilters = reactive([])
const addFilter = (filter: FilterAndValue) => {
  nodeFilters.push(filter)
}
const removeFilter = (filter: FilterAndValue) => {
  const index = nodeFilters.findIndex((f) => f === filter)
  if (index !== -1) {
    nodeFilters.splice(index, 1)
  }
}
const clearFilters = () => {
  nodeFilters.splice(0, nodeFilters.length)
}
const closeDialog = () => {
  visible.value = false
}
const connectNodeOnLinkRelease = (
  node: LGraphNode,
  context: LinkReleaseContext
) => {
  const destIsInput = context.node_from !== undefined
  const srcNode = (
    destIsInput ? context.node_from : context.node_to
  ) as LGraphNode
  const srcSlotIndex: number = context.slot_from.slot_index
  const linkDataType = destIsInput
    ? context.type_filter_in
    : context.type_filter_out
  const destSlots = destIsInput ? node.inputs : node.outputs
  const destSlotIndex = destSlots.findIndex(
    (slot: INodeSlot) => slot.type === linkDataType
  )

  if (destSlotIndex === -1) {
    console.warn(
      `Could not find slot with type ${linkDataType} on node ${node.title}`
    )
    return
  }

  if (destIsInput) {
    srcNode.connect(srcSlotIndex, node, destSlotIndex)
  } else {
    node.connect(destSlotIndex, srcNode, srcSlotIndex)
  }
}
const addNode = (nodeDef: ComfyNodeDefImpl) => {
  closeDialog()

  const node = app.addNodeOnGraph(nodeDef, { pos: getNewNodeLocation() })

  const eventDetail = triggerEvent.value.detail
  if (eventDetail.subType === 'empty-release') {
    connectNodeOnLinkRelease(node, eventDetail.linkReleaseContext)
  }
}

const canvasEventHandler = (e: LiteGraphCanvasEvent) => {
  const shiftPressed = (e.detail.originalEvent as KeyboardEvent).shiftKey
  // Ignore empty releases unless shift is pressed
  // Empty release without shift is trigger right click menu
  if (e.detail.subType === 'empty-release' && !shiftPressed) {
    return
  }

  if (e.detail.subType === 'empty-release') {
    const destIsInput = e.detail.linkReleaseContext.node_from !== undefined
    const filter = useNodeDefStore().nodeSearchService.getFilterById(
      destIsInput ? 'input' : 'output'
    )

    const value = destIsInput
      ? e.detail.linkReleaseContext.type_filter_in
      : e.detail.linkReleaseContext.type_filter_out

    addFilter([filter, value])
  }
  triggerEvent.value = e
  visible.value = true
  // Prevent the dialog from being dismissed immediately
  dismissable.value = false
  setTimeout(() => {
    dismissable.value = true
  }, 300)
}

const handleEscapeKeyPress = (event) => {
  if (event.key === 'Escape') {
    closeDialog()
  }
}

onMounted(() => {
  document.addEventListener('litegraph:canvas', canvasEventHandler)
  document.addEventListener('keydown', handleEscapeKeyPress)
})

onUnmounted(() => {
  document.removeEventListener('litegraph:canvas', canvasEventHandler)
  document.removeEventListener('keydown', handleEscapeKeyPress)
})
</script>

<style>
.invisible-dialog-root {
  width: 30%;
  min-width: 24rem;
  max-width: 48rem;
  border: 0 !important;
  background-color: transparent !important;
  margin-top: 25vh;
}

.node-search-box-dialog-mask {
  align-items: flex-start !important;
}
</style>