66 lines
2.1 KiB
TypeScript
66 lines
2.1 KiB
TypeScript
import type { LGraphNode, IWidget } from "@comfyorg/litegraph";
|
|
import type { ComfyExtension } from "@comfyorg/comfyui-frontend-types";
|
|
|
|
const extension: ComfyExtension = {
|
|
name: "dynamicInputs",
|
|
|
|
beforeRegisterNodeDef(nodeType, nodeData, app) {
|
|
if ((nodeType as any).comfyClass === "DrawThingsLoRA") {
|
|
// update inputs when a new node is added
|
|
setCallback(nodeType.prototype, "onAdded", function (this: any, graph: any) {
|
|
updateInputs(this)
|
|
})
|
|
|
|
// or when the model changes
|
|
setCallback(nodeType.prototype, "onWidgetChanged", function (this: any, name: string, value: any, old_Value: any, widget: any) {
|
|
updateInputs(this)
|
|
})
|
|
}
|
|
},
|
|
}
|
|
|
|
export default extension;
|
|
|
|
function updateInputs(node: LGraphNode) {
|
|
const s = saveSize(node)
|
|
|
|
// the first input should always be lora
|
|
if (node.inputs.length < 1 || node.inputs.length > 2 || node.inputs[0].type !== "DT_LORA") {
|
|
for (let i = node.inputs.length - 1; i >= 0; i--) {
|
|
node.removeInput(i)
|
|
}
|
|
node.addInput("lora", "DT_LORA")
|
|
}
|
|
|
|
const modelWidget = node?.widgets?.find((w) => (w.options as any)?.modelType)
|
|
if (!modelWidget) return
|
|
const modifier = (modelWidget.value as any)?.value?.modifier
|
|
|
|
if (modifier) {
|
|
if (node.inputs.length < 2 || node.inputs[1]?.type !== "IMAGE") {
|
|
node.removeInput(1)
|
|
node.addInput("control_image", "IMAGE")
|
|
}
|
|
} else {
|
|
while (node.inputs.length > 1) {
|
|
node.removeInput(1)
|
|
}
|
|
}
|
|
s()
|
|
}
|
|
|
|
function saveSize(node: LGraphNode) {
|
|
const width = node.width
|
|
return () => node.setSize([width, node.computeSize()[1]])
|
|
}
|
|
|
|
export function setCallback<T extends object, K extends keyof T>(target: T, callbackName: K, callback: T[K]) {
|
|
const originalCallback = target[callbackName]
|
|
target[callbackName] = function (this: any, ...args: any[]) {
|
|
const r = (originalCallback as any)?.apply(this, args)
|
|
;(callback as any)?.apply(this, args)
|
|
return r
|
|
} as any
|
|
}
|
|
|