slotDefaults.js 2.97 KB
Newer Older
1
import { app } from "/scripts/app.js";
2
import { ComfyWidgets } from "/scripts/widgets.js";
3
4
5
6
// Adds defaults for quickly adding nodes with middle click on the input/output

app.registerExtension({
	name: "Comfy.SlotDefaults",
7
	suggestionsNumber: null,
8
	init() {
pythongosssss's avatar
pythongosssss committed
9
		LiteGraph.search_filter_enabled = true;
10
		LiteGraph.middle_click_slot_add_default_node = true;
11
12
		this.suggestionsNumber = app.ui.settings.addSetting({
			id: "Comfy.NodeSuggestions.number",
reaper47's avatar
reaper47 committed
13
			name: "Number of nodes suggestions",
14
15
16
17
18
19
20
21
22
23
24
			type: "slider",
			attrs: {
				min: 1,
				max: 100,
				step: 1,
			},
			defaultValue: 5,
			onChange: (newVal, oldVal) => {
				this.setDefaults(newVal);
			}
		});
25
	},
26
27
	slot_types_default_out: {},
	slot_types_default_in: {},
28
	async beforeRegisterNodeDef(nodeType, nodeData, app) {
29
				var nodeId = nodeData.name;
30
		var inputs = [];
31
		inputs = nodeData["input"]["required"]; //only show required inputs to reduce the mess also not logical to create node with optional inputs
32
33
34
35
36
37
38
39
40
41
		for (const inputKey in inputs) {
			var input = (inputs[inputKey]);
			if (typeof input[0] !== "string") continue;

			var type = input[0]
			if (type in ComfyWidgets) {
				var customProperties = input[1]
				if (!(customProperties?.forceInput)) continue; //ignore widgets that don't force input
			}

42
43
			if (!(type in this.slot_types_default_out)) {
				this.slot_types_default_out[type] = ["Reroute"];
44
			}
45
46
			if (this.slot_types_default_out[type].includes(nodeId)) continue;
			this.slot_types_default_out[type].push(nodeId);
pythongosssss's avatar
pythongosssss committed
47
48
49
50
51
52
53
54

			// Input types have to be stored as lower case
			// Store each node that can handle this input type
			const lowerType = type.toLocaleLowerCase();
			if (!(lowerType in LiteGraph.registered_slot_in_types)) {
				LiteGraph.registered_slot_in_types[lowerType] = { nodes: [] };
			}
			LiteGraph.registered_slot_in_types[lowerType].nodes.push(nodeType.comfyClass);
55
		} 
56
57
58
59

		var outputs = nodeData["output"];
		for (const key in outputs) {
			var type = outputs[key];
60
61
			if (!(type in this.slot_types_default_in)) {
				this.slot_types_default_in[type] = ["Reroute"];// ["Reroute", "Primitive"];  primitive doesn't always work :'()
62
			}
63

64
			this.slot_types_default_in[type].push(nodeId);
pythongosssss's avatar
pythongosssss committed
65
66
67
68
69
70
71
72
73
74

			// Store each node that can handle this output type
			if (!(type in LiteGraph.registered_slot_out_types)) {
				LiteGraph.registered_slot_out_types[type] = { nodes: [] };
			}
			LiteGraph.registered_slot_out_types[type].nodes.push(nodeType.comfyClass);

			if(!LiteGraph.slot_types_out.includes(type)) {
				LiteGraph.slot_types_out.push(type);
			}
75
		}
comfyanonymous's avatar
comfyanonymous committed
76
		var maxNum = this.suggestionsNumber.value;
77
		this.setDefaults(maxNum);
78
	},
79
80
81
82
83
84
85
86
87
88
89
90
	setDefaults(maxNum) {

		LiteGraph.slot_types_default_out = {};
		LiteGraph.slot_types_default_in = {};

		for (const type in this.slot_types_default_out) {
			LiteGraph.slot_types_default_out[type] = this.slot_types_default_out[type].slice(0, maxNum);
		}
		for (const type in this.slot_types_default_in) {
			LiteGraph.slot_types_default_in[type] = this.slot_types_default_in[type].slice(0, maxNum);
		}
	}
91
});