diff --git a/comfy_extras/nodes_subflow.py b/comfy_extras/nodes_subflow.py index c1f6c9833..71616ce56 100644 --- a/comfy_extras/nodes_subflow.py +++ b/comfy_extras/nodes_subflow.py @@ -1,24 +1,13 @@ import folder_paths -import json -import os.path as osp class Subflow: @classmethod def INPUT_TYPES(s): return {"required": { "subflow_name": (folder_paths.get_filename_list("subflows"), ), }} RETURN_TYPES = () - FUNCTION = "exec_subflow" + FUNCTION = "" CATEGORY = "loaders" - - def exec_subflow(self, subflow_name): - subflow_path = folder_paths.get_full_path("subflows", subflow_name) - with open(subflow_path) as f: - if osp.splitext(subflow_path)[1] == ".json": - subflow_data = json.load(f) - return subflow_data - - return None NODE_CLASS_MAPPINGS = { "Subflow": Subflow, diff --git a/web/extensions/core/subflow.js b/web/extensions/core/subflow.js index 1c1973c2d..363b7d8ab 100644 --- a/web/extensions/core/subflow.js +++ b/web/extensions/core/subflow.js @@ -1,10 +1,80 @@ import { app } from "../../scripts/app.js"; import { api } from "../../scripts/api.js"; +class InMemorySubflow { + constructor() {} + + onConfigure() { + console.log(this); + } + + getExportedOutput(slot) { + return this.subflow.extras.outputSlots[slot]; + }; + + getExportedInput(slot) { + return this.subflow.extras.inputSlots[slot]; + }; + + async refreshNode(subflow) { + if (!subflow) return; + this.updateSubflowPrompt(subflow); + this.refreshPins(subflow); + } + + refreshPins(subflow) { + if(!subflow) + return; + + subflow.extras = { inputSlots: [], outputSlots: [] }; + const { inputSlots } = subflow.extras; + const { outputSlots } = subflow.extras; + + // remove all existing pins + const numInputs = this.inputs?.length ?? 0; + const numOutputs = this.outputs?.length ?? 0; + for(let i = numInputs-1; i > -1; i--) { + this.removeInput(i); + } + for(let i = numOutputs-1; i > -1; i--) { + this.removeOutput(i); + } + + const subflowNodes = subflow.nodes; + // add the new pins and keep track of where the exported vars go to within the inner nodes + for (const subflowNode of subflowNodes) { + const exports = subflowNode.properties?.exports; + if (exports) { + let pinNum = 0; + for (const inputRef of exports.inputs) { + const input = subflowNode.inputs.find(q => q.name === inputRef); + if (!input) continue; + this.addInput(input.name, input.type); + inputSlots.push([subflowNode.id, pinNum]); + pinNum++; + } + pinNum = 0; + for (const outputRef of exports.outputs) { + const output = subflowNode.outputs.find(q => q.name === outputRef); + if (!output) continue; + this.addOutput(output.name, output.type); + outputSlots.push([subflowNode.id, pinNum]); + pinNum++; + } + } + } + + this.size[0] = 180; + }; + + updateSubflowPrompt(subflow) { + this.subflow = subflow; + }; +} + app.registerExtension({ name: "Comfy.Subflow", async nodeCreated(node) { - if (!node.widgets) return; if (node.widgets[0].name !== "subflow_name") return; @@ -68,8 +138,6 @@ app.registerExtension({ return Math.max(100, LiteGraph.NODE_TEXT_SIZE * subflowName.length * 0.45 + 160); }; - node.onAdded = () => updateSubflowPrompt(node.widgets[0].value); - node.onConfigure = () => updateSubflowPrompt(node.widgets[0].value); node.widgets[0].callback = (subflowName) => refreshNode(subflowName); node.getExportedOutput = (slot) => { @@ -79,7 +147,13 @@ app.registerExtension({ node.getExportedInput = (slot) => { return inputSlots[slot]; }; - - refreshNode(node.widgets[0].value); - } + }, + registerCustomNodes() { + LiteGraph.registerNodeType( + "InMemorySubflow", + Object.assign(InMemorySubflow, { + title: "InMemorySubflow", + }) + ); + }, }); diff --git a/web/lib/litegraph.core.js b/web/lib/litegraph.core.js index c14f839ad..959bbf343 100644 --- a/web/lib/litegraph.core.js +++ b/web/lib/litegraph.core.js @@ -13023,6 +13023,63 @@ LGraphNode.prototype.executeAction = function(action) return false; }; + // LGraphCanvas.onDecoupleSubflow = function(value, options, e, menu, node) { + // const { subflow } = node; + // if (!subflow) return; + + // const subgraph = new LGraph( subflow ); + + // for (const node of subgraph._nodes) { + // console.log(node); + // graph.add(node); + // } + // }; + + LGraphCanvas.onCollapseToSubflow = function(value, options, e, menu, node) { + const calculatePosition = (nodes) => { + const sum = [0, 0]; + const N = nodes.length; + for (const node of nodes) { + sum[0] += node.pos[0]; + sum[1] += node.pos[1]; + } + return [Math.round(sum[0])/N, Math.round(sum[1]/N)]; + }; + + var graphcanvas = LGraphCanvas.active_canvas; + if (graphcanvas.selected_nodes && Object.keys(graphcanvas.selected_nodes).length > 1){ + const nodes = Object.values(graphcanvas.selected_nodes); + + + // New subflow is the current graph - non included nodes + const subgraph = new LGraph(graph.serialize()); + const subgraphNodes = subgraph._nodes.map(n => n.id); + const nodesToKeep = nodes.map(n => n.id); + + for (const n of subgraphNodes) { + if(!nodesToKeep.includes(n)) { + subgraph.remove(subgraph.getNodeById(n)); + } + } + const subflow = subgraph.serialize(); + const subflowNode = LiteGraph.createNode("InMemorySubflow"); + subflowNode.updateSubflowPrompt(subflow); + subflowNode.pos = calculatePosition(subflow.nodes); + subflowNode.size[0] = 180; + subflowNode.refreshNode(subflow); + graph.add(subflowNode); + + // remove selected nodes + for (const node of nodes) { + graph.remove(node); + } + } + + console.log(graph._nodes); + + return false; + }; + LGraphCanvas.onMenuNodeRemove = function(value, options, e, menu, node) { if (!node) { throw "no node passed"; @@ -13281,9 +13338,20 @@ LGraphNode.prototype.executeAction = function(action) content: "Align Selected To", has_submenu: true, callback: LGraphCanvas.onNodeAlign, - }) + }); + options.push({ + content: "Collapse to Subflow", + callback: LGraphCanvas.onCollapseToSubflow + }); } + // if (Object.keys(this.selected_nodes).length == 1 && node.subflow) { + // options.push({ + // content: "Decouple Subflow", + // callback: LGraphCanvas.onDecoupleSubflow, + // }); + // } + options.push(null, { content: "Remove", disabled: !(node.removable !== false && !node.block_delete ), diff --git a/web/scripts/app.js b/web/scripts/app.js index 92032608d..eb7c981d0 100644 --- a/web/scripts/app.js +++ b/web/scripts/app.js @@ -1533,7 +1533,7 @@ export class ComfyApp { } const getInputRef = (inputNode, inputSlot) => { - if (inputNode.type == "Subflow") { + if (inputNode.subflow) { // input should be mapped to inner node const [ localOriginId, originSlot ] = inputNode.getExportedOutput(inputSlot); const originId = subflowIdMapping[inputNode.id] + localOriginId; @@ -1582,7 +1582,7 @@ export class ComfyApp { } if (link) { - if (node.type == "Subflow") { + if (node.subflow) { // inner node's input should be used const [ localTargetId, targetSlot ] = node.getExportedInput(link.target_slot); const targetId = subflowIdMapping[node.id] + localTargetId; @@ -1594,7 +1594,7 @@ export class ComfyApp { } } - if (node.type != "Subflow") { + if (!node.subflow) { output[String(node.id)] = { inputs, class_type: node.comfyClass,