Skip to content

Commit 110c8dc

Browse files
KAJdevKAJdev
committed
convert to single custom node & bug fixes
1 parent 69d381a commit 110c8dc

File tree

10 files changed

+651
-606
lines changed

10 files changed

+651
-606
lines changed

‎packages/stablestudio-ui/src-tauri/resources/defaultGraph.js

Lines changed: 240 additions & 541 deletions
Large diffs are not rendered by default.

‎packages/stablestudio-ui/src-tauri/resources/stableStudioInputs.js

Lines changed: 286 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,286 @@
1+
/* eslint-disable @typescript-eslint/no-empty-function */
2+
/* eslint-disable @typescript-eslint/no-this-alias */
3+
import { app } from "../scripts/app.js";
4+
5+
function getWidgetType(config) {
6+
// Special handling for COMBO so we restrict links based on the entries
7+
let type = config[0];
8+
let linkType = type;
9+
if (type instanceof Array) {
10+
type = "COMBO";
11+
linkType = linkType.join(",");
12+
}
13+
return { type, linkType };
14+
}
15+
16+
app.registerExtension({
17+
name: "StableStudio.Inputs",
18+
registerCustomNodes() {
19+
class StableStudioNode {
20+
color = LGraphCanvas.node_colors.purple.color;
21+
bgcolor = LGraphCanvas.node_colors.purple.bgcolor;
22+
constructor() {
23+
this.addOutput("batch_size", "INT");
24+
this.addOutput("steps", "INT");
25+
this.addOutput("seed", "INT");
26+
this.addOutput("width", "INT");
27+
this.addOutput("height", "INT");
28+
this.addOutput("positive_prompt", "STRING");
29+
this.addOutput("negative_prompt", "STRING");
30+
this.addOutput("cfg", "FLOAT");
31+
this.addOutput(
32+
"sampler_name",
33+
"euler,euler_ancestral,heun,dpm_2,dpm_2_ancestral,lms,dpm_fast,dpm_adaptive,dpmpp_2s_ancestral,dpmpp_sde,dpmpp_sde_gpu,dpmpp_2m,dpmpp_2m_sde,dpmpp_2m_sde_gpu,ddim,uni_pc,uni_pc_bh2"
34+
);
35+
this.addOutput("ckpt_name", null);
36+
this.serialize_widgets = false;
37+
this.isVirtualNode = true;
38+
39+
this.updateLoop = setInterval(async () => {
40+
const resp = await fetch("/object_info", { cache: "no-cache" });
41+
const data = await resp.json();
42+
if (data) {
43+
const models =
44+
data.CheckpointLoader?.input?.required?.ckpt_name?.[0] ?? [];
45+
const samplers =
46+
data.KSampler?.input?.required?.sampler_name?.[0] ?? [];
47+
48+
if (!this.stableValues.ckpt_name) {
49+
this.stableValues.ckpt_name = models[0];
50+
}
51+
if (!this.stableValues.sampler_name) {
52+
this.stableValues.sampler_name = samplers[0];
53+
}
54+
55+
this.outputs[8].type = samplers.join(",");
56+
this.outputs[9].type = models.join(",");
57+
}
58+
}, 15000);
59+
60+
this.stableValues = {
61+
batch_size: 4,
62+
steps: 50,
63+
seed: 0,
64+
width: 512,
65+
height: 512,
66+
positive_prompt: "",
67+
negative_prompt: "",
68+
cfg: 7,
69+
sampler_name: "",
70+
ckpt_name: "",
71+
};
72+
}
73+
74+
applyToGraph() {
75+
if (!this.outputs.some((o) => o.links?.length)) return;
76+
77+
function get_links(node) {
78+
let links = [];
79+
for (const o of node.outputs) {
80+
if (!o.links?.length) continue;
81+
for (const l of o.links) {
82+
const linkInfo = app.graph.links[l];
83+
const n = node.graph.getNodeById(linkInfo.target_id);
84+
if (n.type == "Reroute") {
85+
links = links.concat(get_links(n));
86+
} else {
87+
links.push(l);
88+
}
89+
}
90+
}
91+
return links;
92+
}
93+
94+
let links = get_links(this);
95+
// For each output link copy our value over the original widget value
96+
for (const l of links) {
97+
const linkInfo = app.graph.links[l];
98+
99+
const node = this.graph.getNodeById(linkInfo.target_id);
100+
const input = node.inputs[linkInfo.target_slot];
101+
const widgetName = input.widget.name;
102+
if (widgetName) {
103+
const widget = node.widgets.find((w) => w.name === widgetName);
104+
if (widget) {
105+
let stableValue = this.outputs.find(
106+
(o) => o.slot_index === linkInfo.origin_slot
107+
)?.name;
108+
109+
if (!stableValue) {
110+
// try to match with name
111+
stableValue = this.outputs.find(
112+
(o) => o.name === input.name
113+
)?.name;
114+
}
115+
116+
if (!stableValue) continue;
117+
widget.value = this.stableValues[stableValue] ?? widget.value;
118+
if (widget.callback) {
119+
widget.callback(
120+
widget.value,
121+
app.canvas,
122+
node,
123+
app.canvas.graph_mouse,
124+
{}
125+
);
126+
}
127+
}
128+
}
129+
}
130+
}
131+
132+
onConnectionsChange(_, index, connected) {
133+
console.log("onConnectionsChange", index, connected);
134+
// if (connected) {
135+
// if (this.outputs[0].links?.length) {
136+
// if (!this.widgets?.length) {
137+
// this.#onFirstConnection();
138+
// }
139+
// if (!this.widgets?.length && this.outputs[0].widget) {
140+
// // On first load it often cant recreate the widget as the other node doesnt exist yet
141+
// // Manually recreate it from the output info
142+
// this.#createWidget(this.outputs[0].widget.config);
143+
// }
144+
// }
145+
// } else if (!this.outputs[0].links?.length) {
146+
// this.#onLastDisconnect();
147+
// }
148+
}
149+
150+
onConnectOutput(slot, type, input) {
151+
// Fires before the link is made allowing us to reject it if it isn't valid
152+
// No widget, we cant connect
153+
// if (!input.widget) {
154+
// if (!(input.type in ComfyWidgets)) return false;
155+
// }
156+
// if (this.outputs[slot].links?.length) {
157+
// return this.#isValidConnection(input);
158+
// }
159+
160+
// let a STRING connect to a ckpt_name
161+
if (
162+
this.outputs[slot].type === "STRING" &&
163+
input.name === "ckpt_name"
164+
) {
165+
return true;
166+
}
167+
}
168+
169+
#onFirstConnection() {
170+
// First connection can fire before the graph is ready on initial load so random things can be missing
171+
const linkId = this.outputs[0].links[0];
172+
const link = this.graph.links[linkId];
173+
if (!link) return;
174+
const theirNode = this.graph.getNodeById(link.target_id);
175+
if (!theirNode || !theirNode.inputs) return;
176+
const input = theirNode.inputs[link.target_slot];
177+
if (!input) return;
178+
var _widget;
179+
if (!input.widget) {
180+
if (!(input.type in ComfyWidgets)) return;
181+
_widget = { name: input.name, config: [input.type, {}] }; //fake widget
182+
} else {
183+
_widget = input.widget;
184+
}
185+
const widget = _widget;
186+
const { type, linkType } = getWidgetType(widget.config);
187+
// Update our output to restrict to the widget type
188+
this.outputs[0].type = linkType;
189+
this.outputs[0].name = type;
190+
this.outputs[0].widget = widget;
191+
this.#createWidget(widget.config, theirNode, widget.name);
192+
}
193+
194+
#createWidget(inputData, node, widgetName) {
195+
let type = inputData[0];
196+
197+
if (type instanceof Array) {
198+
type = "COMBO";
199+
}
200+
201+
let widget;
202+
if (type in ComfyWidgets) {
203+
widget = (ComfyWidgets[type](this, "value", inputData, app) || {})
204+
.widget;
205+
} else {
206+
widget = this.addWidget(type, "value", null, () => {}, {});
207+
}
208+
209+
if (node?.widgets && widget) {
210+
const theirWidget = node.widgets.find((w) => w.name === widgetName);
211+
if (theirWidget) {
212+
widget.value = theirWidget.value;
213+
}
214+
}
215+
216+
if (widget.type === "number" || widget.type === "combo") {
217+
addValueControlWidget(this, widget, "fixed");
218+
}
219+
220+
// When our value changes, update other widgets to reflect our changes
221+
// e.g. so LoadImage shows correct image
222+
const callback = widget.callback;
223+
const self = this;
224+
widget.callback = function () {
225+
const r = callback ? callback.apply(this, arguments) : undefined;
226+
self.applyToGraph();
227+
return r;
228+
};
229+
230+
// Grow our node if required
231+
const sz = this.computeSize();
232+
if (this.size[0] < sz[0]) {
233+
this.size[0] = sz[0];
234+
}
235+
if (this.size[1] < sz[1]) {
236+
this.size[1] = sz[1];
237+
}
238+
239+
requestAnimationFrame(() => {
240+
if (this.onResize) {
241+
this.onResize(this.size);
242+
}
243+
});
244+
}
245+
246+
#isValidConnection(input) {
247+
// Only allow connections where the configs match
248+
const config1 = this.outputs[0].widget.config;
249+
const config2 = input.widget.config;
250+
251+
if (config1[0] instanceof Array) {
252+
// These checks shouldnt actually be necessary as the types should match
253+
// but double checking doesn't hurt
254+
255+
// New input isnt a combo
256+
if (!(config2[0] instanceof Array)) return false;
257+
// New imput combo has a different size
258+
if (config1[0].length !== config2[0].length) return false;
259+
// New input combo has different elements
260+
if (config1[0].find((v, i) => config2[0][i] !== v)) return false;
261+
} else if (config1[0] !== config2[0]) {
262+
// Configs dont match
263+
return false;
264+
}
265+
266+
for (const k in config1[1]) {
267+
if (k !== "default") {
268+
if (config1[1][k] !== config2[1][k]) {
269+
return false;
270+
}
271+
}
272+
}
273+
274+
return true;
275+
}
276+
}
277+
278+
LiteGraph.registerNodeType(
279+
"StableStudioNode",
280+
Object.assign(StableStudioNode, {
281+
title: "StableStudio Input",
282+
})
283+
);
284+
StableStudioNode.category = "stablestudio";
285+
},
286+
});

‎packages/stablestudio-ui/src-tauri/src/main.rs

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -96,17 +96,24 @@ fn extract_comfy(
9696
// dont block the main thread
9797
extract_thread.join().unwrap();
9898

99-
// move resources/defaultGraph.js to APPDATA/ComfyUI/ComfyUI/web/scripts/defaultGraph.js
100-
let resource_path = handle
99+
let default_graph = handle
101100
.path_resolver()
102101
.resolve_resource("resources/defaultGraph.js")
103102
.expect("failed to resolve resource");
103+
let stable_inputs = handle
104+
.path_resolver()
105+
.resolve_resource("resources/stableStudioInputs.js")
106+
.expect("failed to resolve resource");
104107

105-
let mut dest = std::path::PathBuf::from(target_dir);
108+
let mut dest = std::path::PathBuf::from(&target_dir);
106109
dest.push("ComfyUI/ComfyUI/web/scripts/defaultGraph.js");
110+
std::fs::create_dir_all(dest.parent().unwrap()).unwrap();
111+
std::fs::copy(default_graph, dest).unwrap();
107112

113+
let mut dest = std::path::PathBuf::from(&target_dir);
114+
dest.push("ComfyUI/ComfyUI/web/extensions/stableStudioInputs.js");
108115
std::fs::create_dir_all(dest.parent().unwrap()).unwrap();
109-
std::fs::copy(resource_path, dest).unwrap();
116+
std::fs::copy(stable_inputs, dest).unwrap();
110117

111118
println!("extracted zip");
112119
Ok("completed".to_string())

‎packages/stablestudio-ui/src/App/index.tsx

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,7 @@ export namespace App {
134134

135135
if (!comfyExists) {
136136
setIsSetup(SetupState.NotStarted);
137-
setMessage("Installing ComfyUI...");
137+
setMessage("Downloading ComfyUI...");
138138

139139
// delete the old comfyui zip if it exists
140140
if (

‎packages/stablestudio-ui/src/Comfy/index.tsx

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ export type ComfyApp = {
2121
_nodes: {
2222
title: string;
2323
type: string;
24+
stableValues?: any;
2425
widgets: {
2526
name: string;
2627
type: string;

0 commit comments

Comments
 (0)