wraith-executioner-year-2-comfyui-misc/web/any_split.js

/** Any Split (misc) */

import { app } from "../../scripts/app.js"
import {
  attachStabilizeHooks,
  IoDirection,
  LENGTH_OUTPUT_NAME,
  countPrimitiveDataOutputs,
  ensureLengthOutput,
  followConnectionUntilType,
  formatDataOutputName,
  pinWildcardSlot,
  scheduleStabilizeRetries,
  slotLabelForLinkType,
  syncNodeSizeToContent,
} from "./utils.js"
import { computeDesiredDataOutputCount } from "./logic/any-split-layout.js"
import { MISC_GRAPH_RESTORE_WINDOW_MS } from "./logic/restore-window.js"
import { getGraphLink } from "./utils/graph-links.js"

const NODE_CLASS = "Any Split (misc)"
const MIN_DATA_OUTPUTS = 1

const lengthTrailing = {
  isMeta: (node) => node.outputs[node.outputs.length - 1]?.name === LENGTH_OUTPUT_NAME,
  countData: countPrimitiveDataOutputs,
  ensureMeta: ensureLengthOutput,
}

function getSplitInput(node) {
  return node.inputs?.find((inp) => inp.name === "*")
}

function addDataOutputBeforeLength(node, num = 1, nodeType = "*") {
  lengthTrailing.ensureMeta(node)
  for (let n = 0; n < num; n++) {
    if (!lengthTrailing.isMeta(node)) {
      lengthTrailing.ensureMeta(node)
    }

    const indexOutput = node.outputs[node.outputs.length - 1]
    const savedLengthLinkIds = [...(indexOutput?.links || [])]
    node.outputs.pop()

    const slotIndex = node.outputs.length + 1
    node.addOutput(formatDataOutputName("any", slotIndex), nodeType || "*")
    node.addOutput(LENGTH_OUTPUT_NAME, "INT")

    const newLengthPos = node.outputs.length - 1
    for (const linkId of savedLengthLinkIds) {
      const link = getGraphLink(app.graph, linkId)
      if (link && link.origin_id === node.id) {
        link.origin_slot = newLengthPos
      }
    }
    node.outputs[newLengthPos].type = "INT"
  }
}

function renumberDataOutputs(node) {
  const end = lengthTrailing.countData(node)
  for (let i = 0; i < end; i++) {
    node.outputs[i].name = formatDataOutputName("any", i + 1)
  }
}

function ensureMinimumDataOutputs(node, minCount, nodeType) {
  lengthTrailing.ensureMeta(node)
  while (lengthTrailing.countData(node) < minCount) {
    addDataOutputBeforeLength(node, 1, nodeType)
  }
  renumberDataOutputs(node)
}

function removeUnusedDataOutputsFromEnd(node, minCount) {
  if (node.removed) {
    return
  }
  lengthTrailing.ensureMeta(node)
  const lengthPos = node.outputs.length - 1
  for (let i = lengthPos - 1; i >= minCount; i--) {
    if (!node.outputs[i]?.links?.length) {
      node.removeOutput(i)
      continue
    }
    break
  }
  renumberDataOutputs(node)
}

function normalizeLengthOutput(node) {
  const lengthOut = node.outputs[node.outputs.length - 1]
  if (!lengthOut) {
    return
  }
  lengthOut.type = "INT"
  lengthOut.name = LENGTH_OUTPUT_NAME
  lengthOut.label = LENGTH_OUTPUT_NAME
}

function resolveElementType(node) {
  const splitInput = getSplitInput(node)
  const inputIndex = splitInput ? node.inputs.indexOf(splitInput) : 0
  const connectedType = followConnectionUntilType(node, IoDirection.INPUT, inputIndex, true)
  return connectedType?.type || "*"
}

function setupAnySplit(nodeType) {
  nodeType.prototype.serialize_widgets = true

  attachStabilizeHooks(nodeType, {
    configureStabilize: "none",
    afterBind(node) {
      node._miscSplitRestoringUntil = 0
    },
    onConfigure(node) {
      node._miscSplitRestoringUntil = Date.now() + MISC_GRAPH_RESTORE_WINDOW_MS
      node.stabilize()
      scheduleStabilizeRetries(node, [80, 200, 500])
    },
    onConnectionsChange(node, type, _slotIndex, _isConnected, _linkInfo, ioSlot) {
      if (ioSlot && node.outputs?.includes(ioSlot)) {
        const metaIndex = node.outputs.length - 1
        if (node.outputs.indexOf(ioSlot) === metaIndex && ioSlot.name === LENGTH_OUTPUT_NAME) {
          return false
        }
      }
      if (type === IoDirection.OUTPUT && ioSlot?.name === LENGTH_OUTPUT_NAME) {
        return false
      }
    },
  })

  nodeType.prototype.stabilize = function () {
    if (this.removed) {
      return
    }

    const splitInput = getSplitInput(this)
    pinWildcardSlot(splitInput)

    this.nodeType = resolveElementType(this)

    lengthTrailing.ensureMeta(this)

    const restoring = Date.now() < (this._miscSplitRestoringUntil || 0)
    if (!restoring) {
      removeUnusedDataOutputsFromEnd(this, MIN_DATA_OUTPUTS)
    }

    const dataOutputs = this.outputs.slice(0, lengthTrailing.countData(this))
    const desired = computeDesiredDataOutputCount(dataOutputs)
    ensureMinimumDataOutputs(this, desired, this.nodeType)

    const dataCount = lengthTrailing.countData(this)
    const primitivesLabel = slotLabelForLinkType(this.nodeType)
    for (let i = 0; i < dataCount; i++) {
      const output = this.outputs[i]
      output.type = this.nodeType
      output.label = primitivesLabel ?? output.name
    }

    normalizeLengthOutput(this)
    syncNodeSizeToContent(this)
    this.graph?.setDirtyCanvas?.(true, false)
  }
}

app.registerExtension({
  name: "comfyui-misc.AnySplit",
  async beforeRegisterNodeDef(nodeType, nodeData) {
    if (nodeData.name !== NODE_CLASS) {
      return
    }

    nodeData.input = nodeData.input ?? {}
    nodeData.input.required = {
      "*": ["*"],
      ...(nodeData.input.required ?? {}),
    }

    nodeData.output = ["*", "INT"]
    nodeData.output_name = ["any_01", "length"]

    setupAnySplit(nodeType)
  },
})