368 lines
10 KiB
JavaScript
368 lines
10 KiB
JavaScript
/** Split Primitives (misc) */
|
|
|
|
import { app } from "../../scripts/app.js"
|
|
import {
|
|
attachStabilizeHooks,
|
|
PRIMITIVES_TYPE,
|
|
syncPrimitivesLinkSlot,
|
|
COMBINE_PRIMITIVES_NODE_CLASS,
|
|
IoDirection,
|
|
PRIMITIVE_SLOT_TYPE,
|
|
SPLIT_PRIMITIVES_NODE_CLASS,
|
|
LENGTH_OUTPUT_NAME,
|
|
findLinkedCombineNode,
|
|
isPrimitiveRelayNode,
|
|
propagatePrimitiveSplitSync,
|
|
getStoredPrimitiveSlotTypes,
|
|
resolvePrimitiveSlotType,
|
|
reconcilePrimitiveSlotType,
|
|
scheduleStabilizeRetries,
|
|
syncNodeSizeToContent,
|
|
} from "./utils.js"
|
|
|
|
import { createTrailingOutputHelpers } from "./utils/trailing-output.js"
|
|
import {
|
|
formatTypedOutputName,
|
|
typeNameForPrimitiveSlotType,
|
|
listLinkedPrimitiveInputs,
|
|
pickDesiredDuringSync,
|
|
resolveDesiredPrimitiveSlots,
|
|
} from "./logic/split-primitives-names.js"
|
|
import { MISC_GRAPH_RESTORE_WINDOW_MS } from "./logic/restore-window.js"
|
|
|
|
const NODE_CLASS = SPLIT_PRIMITIVES_NODE_CLASS
|
|
const COMBINE_NODE_CLASS = COMBINE_PRIMITIVES_NODE_CLASS
|
|
|
|
const RELAY_PATCH_KEY = "__miscPrimitiveRelayPatch"
|
|
|
|
const lengthTrailing = createTrailingOutputHelpers(LENGTH_OUTPUT_NAME)
|
|
const countPrimitiveDataOutputs = lengthTrailing.countDataSlots
|
|
const ensureLengthOutput = lengthTrailing.ensureMeta
|
|
const isLengthOutputSlot = lengthTrailing.isMetaSlot
|
|
|
|
function patchRelayNodeConnectionCallbacks() {
|
|
if (LGraphNode.prototype[RELAY_PATCH_KEY]) {
|
|
return
|
|
}
|
|
LGraphNode.prototype[RELAY_PATCH_KEY] = true
|
|
|
|
const origChange = LGraphNode.prototype.onConnectionsChange
|
|
LGraphNode.prototype.onConnectionsChange = function (...args) {
|
|
const ret = origChange?.apply(this, args)
|
|
if (isPrimitiveRelayNode(this)) {
|
|
propagatePrimitiveSplitSync(this)
|
|
}
|
|
return ret
|
|
}
|
|
|
|
const origChain = LGraphNode.prototype.onConnectionsChainChange
|
|
if (origChain) {
|
|
LGraphNode.prototype.onConnectionsChainChange = function (...args) {
|
|
const ret = origChain.apply(this, args)
|
|
if (isPrimitiveRelayNode(this)) {
|
|
propagatePrimitiveSplitSync(this)
|
|
}
|
|
return ret
|
|
}
|
|
}
|
|
}
|
|
|
|
function addPrimitiveOutputBeforeLength(node, name, type) {
|
|
ensureLengthOutput(node)
|
|
if (isLengthOutputSlot(node)) {
|
|
const lengthSlot = node.outputs.pop()
|
|
node.addOutput(name, type)
|
|
node.outputs.push(lengthSlot)
|
|
} else {
|
|
node.addOutput(name, type)
|
|
ensureLengthOutput(node)
|
|
}
|
|
}
|
|
|
|
function collectExistingOutputSnapshot(splitNode) {
|
|
const end = countPrimitiveDataOutputs(splitNode)
|
|
const snapshot = []
|
|
for (let i = 0; i < end; i++) {
|
|
const out = splitNode.outputs?.[i]
|
|
if (!out?.name) {
|
|
continue
|
|
}
|
|
snapshot.push({ name: out.name, type: out.type })
|
|
}
|
|
return snapshot
|
|
}
|
|
|
|
function ensurePrimitiveOutputsCount(splitNode, count, defaultType = PRIMITIVE_SLOT_TYPE) {
|
|
ensureLengthOutput(splitNode)
|
|
while (countPrimitiveDataOutputs(splitNode) < count) {
|
|
const idx = countPrimitiveDataOutputs(splitNode)
|
|
addPrimitiveOutputBeforeLength(splitNode, `primitive_${String(idx + 1).padStart(2, "0")}`, defaultType)
|
|
}
|
|
}
|
|
|
|
function normalizeLengthOutput(splitNode) {
|
|
const lengthOut = splitNode.outputs[splitNode.outputs.length - 1]
|
|
if (!lengthOut) {
|
|
return
|
|
}
|
|
lengthOut.type = "INT"
|
|
lengthOut.name = LENGTH_OUTPUT_NAME
|
|
lengthOut.label = LENGTH_OUTPUT_NAME
|
|
}
|
|
|
|
function ensureSplitOutputLabels(splitNode) {
|
|
const end = countPrimitiveDataOutputs(splitNode)
|
|
for (let i = 0; i < end; i++) {
|
|
const out = splitNode.outputs?.[i]
|
|
if (!out?.name) {
|
|
continue
|
|
}
|
|
if (out.label !== out.name) {
|
|
out.label = out.name
|
|
}
|
|
}
|
|
normalizeLengthOutput(splitNode)
|
|
}
|
|
|
|
function setGenericPrimitiveOutputs(splitNode) {
|
|
const end = countPrimitiveDataOutputs(splitNode)
|
|
for (let i = 0; i < end; i++) {
|
|
const out = splitNode.outputs[i]
|
|
if (!out?.name?.includes("_")) {
|
|
continue
|
|
}
|
|
out.type = PRIMITIVE_SLOT_TYPE
|
|
out.label = out.name
|
|
}
|
|
}
|
|
|
|
function applySnapshotPlaceholders(splitNode, snapshot) {
|
|
if (!snapshot?.length) {
|
|
return
|
|
}
|
|
ensurePrimitiveOutputsCount(splitNode, snapshot.length, PRIMITIVE_SLOT_TYPE)
|
|
for (let i = 0; i < snapshot.length; i++) {
|
|
const out = splitNode.outputs?.[i]
|
|
const snap = snapshot[i]
|
|
if (!out || !snap) {
|
|
continue
|
|
}
|
|
out.name = snap.name
|
|
out.type = PRIMITIVE_SLOT_TYPE
|
|
out.label = snap.name
|
|
}
|
|
normalizeLengthOutput(splitNode)
|
|
}
|
|
|
|
function listLinkedWithResolvedTypes(combineNode) {
|
|
const linked = []
|
|
for (let i = 0; i < (combineNode.inputs ?? []).length; i++) {
|
|
const inp = combineNode.inputs[i]
|
|
if (!inp?.name?.startsWith("primitive_") || !inp.link) {
|
|
continue
|
|
}
|
|
linked.push({
|
|
name: inp.name,
|
|
type: resolvePrimitiveSlotType(combineNode, i, inp),
|
|
})
|
|
}
|
|
return linked
|
|
}
|
|
|
|
function getCombinedInput(splitNode) {
|
|
return splitNode.inputs?.find((inp) => inp.name === "combined")
|
|
}
|
|
|
|
function syncSplitFromCombine(splitNode, combineNode) {
|
|
ensureLengthOutput(splitNode)
|
|
|
|
const restoring = Date.now() < (splitNode._miscSplitRestoringUntil || 0)
|
|
const combinedInput = getCombinedInput(splitNode)
|
|
const hasCombinedLink = !!combinedInput?.link
|
|
|
|
if (!restoring && !combineNode && hasCombinedLink) {
|
|
ensurePrimitiveOutputsCount(splitNode, 1, PRIMITIVE_SLOT_TYPE)
|
|
normalizeLengthOutput(splitNode)
|
|
ensureSplitOutputLabels(splitNode)
|
|
return
|
|
}
|
|
|
|
if (restoring && !combineNode) {
|
|
setGenericPrimitiveOutputs(splitNode)
|
|
normalizeLengthOutput(splitNode)
|
|
return
|
|
}
|
|
|
|
const desired = pickDesiredDuringSync({
|
|
restoring,
|
|
combineNode,
|
|
hasCombinedLink,
|
|
cachedDesired: splitNode._miscSplitCachedDesired,
|
|
linked: combineNode ? listLinkedWithResolvedTypes(combineNode) : [],
|
|
stored: combineNode ? getStoredPrimitiveSlotTypes(combineNode) : null,
|
|
})
|
|
const desiredCount = desired.length
|
|
|
|
while (countPrimitiveDataOutputs(splitNode) < desiredCount) {
|
|
const idx = countPrimitiveDataOutputs(splitNode)
|
|
const slot = desired[idx]
|
|
if (!slot) {
|
|
break
|
|
}
|
|
addPrimitiveOutputBeforeLength(splitNode, formatTypedOutputName(slot.type, idx), slot.type)
|
|
}
|
|
|
|
while (countPrimitiveDataOutputs(splitNode) > desiredCount) {
|
|
const removeIndex = countPrimitiveDataOutputs(splitNode) - 1
|
|
const output = splitNode.outputs[removeIndex]
|
|
|
|
if (restoring) {
|
|
break
|
|
}
|
|
if (output?.links?.length) {
|
|
break
|
|
}
|
|
splitNode.removeOutput(removeIndex)
|
|
}
|
|
|
|
for (let i = 0; i < desiredCount; i++) {
|
|
const output = splitNode.outputs[i]
|
|
if (!output) {
|
|
break
|
|
}
|
|
|
|
const want = desired[i]
|
|
if (!want) {
|
|
continue
|
|
}
|
|
|
|
const typedName = formatTypedOutputName(want.type, i)
|
|
if (output.name !== typedName) {
|
|
output.name = typedName
|
|
}
|
|
|
|
if (restoring) {
|
|
output.type = PRIMITIVE_SLOT_TYPE
|
|
} else {
|
|
output.type = reconcilePrimitiveSlotType(output.type, want.type)
|
|
}
|
|
|
|
output.label = output.name
|
|
}
|
|
|
|
ensureSplitOutputLabels(splitNode)
|
|
syncNodeSizeToContent(splitNode)
|
|
}
|
|
|
|
function initSplitRestoreState(node) {
|
|
node._miscSplitGraphReady = true
|
|
node._miscSplitRestoringUntil = 0
|
|
node._miscSplitCachedOutputs = null
|
|
node._miscSplitCachedDesired = null
|
|
}
|
|
|
|
function setupSplitPrimitives(nodeType) {
|
|
const onDrawForeground = nodeType.prototype.onDrawForeground
|
|
|
|
attachStabilizeHooks(nodeType, {
|
|
configureStabilize: "none",
|
|
afterBind(node) {
|
|
initSplitRestoreState(node)
|
|
syncPrimitivesLinkSlot(getCombinedInput(node))
|
|
},
|
|
afterStabilizeOnCreate(node) {
|
|
scheduleStabilizeRetries(node, [80])
|
|
},
|
|
onConfigure(node) {
|
|
initSplitRestoreState(node)
|
|
node._miscSplitRestoringUntil = Date.now() + MISC_GRAPH_RESTORE_WINDOW_MS
|
|
node._miscSplitCachedOutputs = collectExistingOutputSnapshot(node)
|
|
node._miscSplitCachedDesired = node._miscSplitCachedOutputs
|
|
? node._miscSplitCachedOutputs.map((s) => ({ name: s.name, type: s.type }))
|
|
: null
|
|
applySnapshotPlaceholders(node, node._miscSplitCachedOutputs)
|
|
node.stabilize()
|
|
scheduleStabilizeRetries(node, [80, 200, 500, 1200], { doubleRaf: true })
|
|
},
|
|
onConnectionsChange(node, type, slotIndex, isConnected, _linkInfo, ioSlot) {
|
|
if (ioSlot && node.outputs?.includes(ioSlot)) {
|
|
return false
|
|
}
|
|
if (type === IoDirection.OUTPUT) {
|
|
return false
|
|
}
|
|
if (!node._miscSplitGraphReady) {
|
|
return false
|
|
}
|
|
const slot = ioSlot ?? node.inputs?.[slotIndex]
|
|
if (type === IoDirection.INPUT && isConnected && slot?.name === "combined") {
|
|
const combineNode = findLinkedCombineNode(node, COMBINE_NODE_CLASS)
|
|
combineNode?.stabilize?.()
|
|
node.stabilize()
|
|
scheduleStabilizeRetries(node, [80])
|
|
return false
|
|
}
|
|
},
|
|
onConnectionsChainChange(node) {
|
|
if (!node._miscSplitGraphReady) {
|
|
return false
|
|
}
|
|
},
|
|
})
|
|
|
|
nodeType.prototype.onDrawForeground = function (ctx) {
|
|
onDrawForeground?.call(this, ctx)
|
|
if (this._miscSplitGraphReady) {
|
|
ensureSplitOutputLabels(this)
|
|
}
|
|
}
|
|
|
|
nodeType.prototype.stabilize = function () {
|
|
if (!this._miscSplitGraphReady) {
|
|
return
|
|
}
|
|
|
|
syncPrimitivesLinkSlot(getCombinedInput(this))
|
|
|
|
const combineNode = findLinkedCombineNode(this, COMBINE_NODE_CLASS)
|
|
const restoring = Date.now() < (this._miscSplitRestoringUntil || 0)
|
|
syncSplitFromCombine(this, combineNode)
|
|
|
|
if (restoring && combineNode && this._miscSplitCachedOutputs) {
|
|
const current = collectExistingOutputSnapshot(this)
|
|
const cached = this._miscSplitCachedOutputs
|
|
const differs =
|
|
current.length !== cached.length ||
|
|
current.some((s, i) => s.name !== cached[i]?.name || s.type !== cached[i]?.type)
|
|
if (differs) {
|
|
requestAnimationFrame(() => this.scheduleStabilize(0))
|
|
}
|
|
}
|
|
|
|
syncNodeSizeToContent(this)
|
|
}
|
|
}
|
|
|
|
app.registerExtension({
|
|
name: "comfyui-misc.SplitPrimitives",
|
|
init() {
|
|
patchRelayNodeConnectionCallbacks()
|
|
},
|
|
async beforeRegisterNodeDef(nodeType, nodeData) {
|
|
if (nodeData.name !== NODE_CLASS) {
|
|
return
|
|
}
|
|
|
|
nodeData.input = nodeData.input ?? {}
|
|
nodeData.input.required = {
|
|
...(nodeData.input.required ?? {}),
|
|
combined: [PRIMITIVES_TYPE],
|
|
}
|
|
|
|
nodeData.output = [PRIMITIVE_SLOT_TYPE, "INT"]
|
|
nodeData.output_name = ["primitive_01", "length"]
|
|
|
|
setupSplitPrimitives(nodeType)
|
|
},
|
|
})
|