fred-dev/comfy_ui_ali

1import { app } from "../../scripts/app.js";2import { api } from "../../scripts/api.js";3 4function addMenuHandler(nodeType, cb) {5	const getOpts = nodeType.prototype.getExtraMenuOptions;6	nodeType.prototype.getExtraMenuOptions = function () {7		const r = getOpts.apply(this, arguments);8		cb.apply(this, arguments);9		return r;10	};11}12 13function distance(node1, node2) {14	let dx = (node1.pos[0] + node1.size[0]/2) - (node2.pos[0] + node2.size[0]/2);15	let dy = (node1.pos[1] + node1.size[1]/2) - (node2.pos[1] + node2.size[1]/2);16	return Math.sqrt(dx * dx + dy * dy);17}18 19function lookup_nearest_nodes(node) {20	let nearest_distance = Infinity;21	let nearest_node = null;22	for(let other of app.graph._nodes) {23		if(other === node)24			continue;25 26		let dist = distance(node, other);27		if (dist < nearest_distance && dist < 1000) {28			nearest_distance = dist;29			nearest_node = other;30		}31	}32 33	return nearest_node;34}35 36function lookup_nearest_inputs(node) {37	let input_map = {};38 39	for(let i in node.inputs) {40		let input = node.inputs[i];41 42		if(input.link || input_map[input.type])43			continue;44 45		input_map[input.type] = {distance: Infinity, input_name: input.name, node: null, slot: null};46	}47 48	let x = node.pos[0];49	let y = node.pos[1] + node.size[1]/2;50 51	for(let other of app.graph._nodes) {52		if(other === node || !other.outputs)53			continue;54 55		let dx = x - (other.pos[0] + other.size[0]);56		let dy = y - (other.pos[1] + other.size[1]/2);57 58		if(dx < 0)59			continue;60 61		let dist = Math.sqrt(dx * dx + dy * dy);62 63		for(let input_type in input_map) {64			for(let j in other.outputs) {65				let output = other.outputs[j];66				if(output.type == input_type) {67					if(input_map[input_type].distance > dist) {68						input_map[input_type].distance = dist;69						input_map[input_type].node = other;70						input_map[input_type].slot = parseInt(j);71					}72				}73			}74		}75	}76 77	let res = {};78	for (let i in input_map) {79		if (input_map[i].node) {80			res[i] = input_map[i];81		}82	}83 84	return res;85}86 87function connect_inputs(nearest_inputs, node) {88	for(let i in nearest_inputs) {89		let info = nearest_inputs[i];90		info.node.connect(info.slot, node.id, info.input_name);91	}92}93 94function node_info_copy(src, dest, connect_both, copy_shape) {95	// copy input connections96	for(let i in src.inputs) {97		let input = src.inputs[i];98		if (input.widget !== undefined) {99			const destWidget = dest.widgets.find(x => x.name === input.widget.name);100			dest.convertWidgetToInput(destWidget);101		}102		if(input.link) {103			let link = app.graph.links[input.link];104			let src_node = app.graph.getNodeById(link.origin_id);105			src_node.connect(link.origin_slot, dest.id, input.name);106		}107	}108 109	// copy output connections110	if(connect_both) {111		let output_links = {};112		for(let i in src.outputs) {113			let output = src.outputs[i];114			if(output.links) {115				let links = [];116				for(let j in output.links) {117					links.push(app.graph.links[output.links[j]]);118				}119				output_links[output.name] = links;120			}121		}122 123		for(let i in dest.outputs) {124			let links = output_links[dest.outputs[i].name];125			if(links) {126				for(let j in links) {127					let link = links[j];128					let target_node = app.graph.getNodeById(link.target_id);129					dest.connect(parseInt(i), target_node, link.target_slot);130				}131			}132		}133	}134 135	if(copy_shape) {136		dest.color = src.color;137		dest.bgcolor = src.bgcolor;138		dest.size = max(src.size, dest.size);139	}140 141	app.graph.afterChange();142}143 144app.registerExtension({145	name: "Comfy.Manager.NodeFixer",146	beforeRegisterNodeDef(nodeType, nodeData, app) {147		addMenuHandler(nodeType, function (_, options) {148			options.push({149				content: "Fix node (recreate)",150				callback: () => {151					let new_node = LiteGraph.createNode(nodeType.comfyClass);152					new_node.pos = [this.pos[0], this.pos[1]];153					app.canvas.graph.add(new_node, false);154					node_info_copy(this, new_node, true);155					app.canvas.graph.remove(this);156				},157			});158		});159	}160});161