app.js 45.2 KB
Newer Older
pythongosssss's avatar
pythongosssss committed
1
import { ComfyLogging } from "./logging.js";
pythongosssss's avatar
pythongosssss committed
2
import { ComfyWidgets } from "./widgets.js";
3
import { ComfyUI, $el } from "./ui.js";
pythongosssss's avatar
pythongosssss committed
4
5
import { api } from "./api.js";
import { defaultGraph } from "./defaultGraph.js";
6
import { getPngMetadata, importA1111, getLatentMetadata } from "./pnginfo.js";
pythongosssss's avatar
pythongosssss committed
7

8
/**
9
10
11
12
13
14
15
 * @typedef {import("types/comfy").ComfyExtension} ComfyExtension
 */

export class ComfyApp {
	/**
	 * List of entries to queue
	 * @type {{number: number, batchCount: number}[]}
16
17
18
19
	 */
	#queueItems = [];
	/**
	 * If the queue is currently being processed
20
	 * @type {boolean}
21
22
23
	 */
	#processingQueue = false;

ltdrdata's avatar
ltdrdata committed
24
25
26
27
28
	/**
	 * Content Clipboard
	 * @type {serialized node object}
	 */
	static clipspace = null;
29
	static clipspace_invalidate_handler = null;
30
31
	static open_maskeditor = null;
	static clipspace_return_node = null;
ltdrdata's avatar
ltdrdata committed
32

pythongosssss's avatar
pythongosssss committed
33
34
	constructor() {
		this.ui = new ComfyUI(this);
pythongosssss's avatar
pythongosssss committed
35
		this.logging = new ComfyLogging(this);
36
37
38
39
40

		/**
		 * List of extensions that are registered with the app
		 * @type {ComfyExtension[]}
		 */
41
		this.extensions = [];
42
43
44
45
46

		/**
		 * Stores the execution output data for each node
		 * @type {Record<string, any>}
		 */
pythongosssss's avatar
pythongosssss committed
47
		this.nodeOutputs = {};
48

space-nuko's avatar
space-nuko committed
49
50
51
52
53
54
		/**
		 * Stores the preview image data for each node
		 * @type {Record<string, Image>}
		 */
		this.nodePreviewImages = {};

55
56
57
58
		/**
		 * If the shift key on the keyboard is pressed
		 * @type {boolean}
		 */
pythongosssss's avatar
pythongosssss committed
59
		this.shiftDown = false;
pythongosssss's avatar
pythongosssss committed
60
61
	}

62
63
64
65
66
67
68
69
	getPreviewFormatParam() {
		let preview_format = this.ui.settings.getSettingValue("Comfy.PreviewFormat");
		if(preview_format)
			return `&preview=${preview_format}`;
		else
			return "";
	}

70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
	static isImageNode(node) {
		return node.imgs || (node && node.widgets && node.widgets.findIndex(obj => obj.name === 'image') >= 0);
	}

	static onClipspaceEditorSave() {
		if(ComfyApp.clipspace_return_node) {
			ComfyApp.pasteFromClipspace(ComfyApp.clipspace_return_node);
		}
	}

	static onClipspaceEditorClosed() {
		ComfyApp.clipspace_return_node = null;
	}

	static copyToClipspace(node) {
		var widgets = null;
		if(node.widgets) {
			widgets = node.widgets.map(({ type, name, value }) => ({ type, name, value }));
		}

		var imgs = undefined;
		var orig_imgs = undefined;
		if(node.imgs != undefined) {
			imgs = [];
			orig_imgs = [];

			for (let i = 0; i < node.imgs.length; i++) {
				imgs[i] = new Image();
				imgs[i].src = node.imgs[i].src;
				orig_imgs[i] = imgs[i];
			}
		}

		var selectedIndex = 0;
		if(node.imageIndex) {
			selectedIndex = node.imageIndex;
		}

		ComfyApp.clipspace = {
			'widgets': widgets,
			'imgs': imgs,
			'original_imgs': orig_imgs,
			'images': node.images,
			'selectedIndex': selectedIndex,
			'img_paste_mode': 'selected' // reset to default im_paste_mode state on copy action
		};

		ComfyApp.clipspace_return_node = null;

		if(ComfyApp.clipspace_invalidate_handler) {
			ComfyApp.clipspace_invalidate_handler();
		}
	}

	static pasteFromClipspace(node) {
		if(ComfyApp.clipspace) {
			// image paste
			if(ComfyApp.clipspace.imgs && node.imgs) {
				if(node.images && ComfyApp.clipspace.images) {
					if(ComfyApp.clipspace['img_paste_mode'] == 'selected') {
130
						node.images = [ComfyApp.clipspace.images[ComfyApp.clipspace['selectedIndex']]];
131
					}
132
133
134
135
136
137
					else {
						node.images = ComfyApp.clipspace.images;
					}

					if(app.nodeOutputs[node.id + ""])
						app.nodeOutputs[node.id + ""].images = node.images;
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
				}

				if(ComfyApp.clipspace.imgs) {
					// deep-copy to cut link with clipspace
					if(ComfyApp.clipspace['img_paste_mode'] == 'selected') {
						const img = new Image();
						img.src = ComfyApp.clipspace.imgs[ComfyApp.clipspace['selectedIndex']].src;
						node.imgs = [img];
						node.imageIndex = 0;
					}
					else {
						const imgs = [];
						for(let i=0; i<ComfyApp.clipspace.imgs.length; i++) {
							imgs[i] = new Image();
							imgs[i].src = ComfyApp.clipspace.imgs[i].src;
							node.imgs = imgs;
						}
					}
				}
			}

			if(node.widgets) {
				if(ComfyApp.clipspace.images) {
					const clip_image = ComfyApp.clipspace.images[ComfyApp.clipspace['selectedIndex']];
					const index = node.widgets.findIndex(obj => obj.name === 'image');
					if(index >= 0) {
164
165
166
167
168
169
						if(node.widgets[index].type != 'image' && typeof node.widgets[index].value == "string" && clip_image.filename) {
							node.widgets[index].value = (clip_image.subfolder?clip_image.subfolder+'/':'') + clip_image.filename + (clip_image.type?` [${clip_image.type}]`:'');
						}
						else {
							node.widgets[index].value = clip_image;
						}
170
171
172
173
174
					}
				}
				if(ComfyApp.clipspace.widgets) {
					ComfyApp.clipspace.widgets.forEach(({ type, name, value }) => {
						const prop = Object.values(node.widgets).find(obj => obj.type === type && obj.name === name);
175
176
						if (prop && prop.type != 'button') {
							if(prop.type != 'image' && typeof prop.value == "string" && value.filename) {
Dr.Lt.Data's avatar
Dr.Lt.Data committed
177
178
179
180
181
182
183
								prop.value = (value.subfolder?value.subfolder+'/':'') + value.filename + (value.type?` [${value.type}]`:'');
							}
							else {
								prop.value = value;
								prop.callback(value);
							}
						}
184
185
186
187
188
189
190
191
					});
				}
			}

			app.graph.setDirtyCanvas(true);
		}
	}

pythongosssss's avatar
pythongosssss committed
192
193
	/**
	 * Invoke an extension callback
194
195
	 * @param {keyof ComfyExtension} method The extension callback to execute
	 * @param  {any[]} args Any arguments to pass to the callback
196
	 * @returns
pythongosssss's avatar
pythongosssss committed
197
	 */
pythongosssss's avatar
pythongosssss committed
198
199
200
201
202
203
204
	#invokeExtensions(method, ...args) {
		let results = [];
		for (const ext of this.extensions) {
			if (method in ext) {
				try {
					results.push(ext[method](...args, this));
				} catch (error) {
205
					console.error(
pythongosssss's avatar
pythongosssss committed
206
207
208
209
210
211
212
213
214
215
216
						`Error calling extension '${ext.name}' method '${method}'`,
						{ error },
						{ extension: ext },
						{ args }
					);
				}
			}
		}
		return results;
	}

pythongosssss's avatar
pythongosssss committed
217
218
219
220
221
	/**
	 * Invoke an async extension callback
	 * Each callback will be invoked concurrently
	 * @param {string} method The extension callback to execute
	 * @param  {...any} args Any arguments to pass to the callback
222
	 * @returns
pythongosssss's avatar
pythongosssss committed
223
	 */
pythongosssss's avatar
pythongosssss committed
224
225
226
227
228
229
230
	async #invokeExtensionsAsync(method, ...args) {
		return await Promise.all(
			this.extensions.map(async (ext) => {
				if (method in ext) {
					try {
						return await ext[method](...args, this);
					} catch (error) {
231
						console.error(
pythongosssss's avatar
pythongosssss committed
232
233
234
235
236
237
238
239
240
241
242
							`Error calling extension '${ext.name}' method '${method}'`,
							{ error },
							{ extension: ext },
							{ args }
						);
					}
				}
			})
		);
	}

pythongosssss's avatar
pythongosssss committed
243
244
245
246
247
	/**
	 * Adds special context menu handling for nodes
	 * e.g. this adds Open Image functionality for nodes that show images
	 * @param {*} node The node to add the menu handler
	 */
pythongosssss's avatar
pythongosssss committed
248
249
250
251
252
253
254
255
256
257
258
259
260
	#addNodeContextMenuHandler(node) {
		node.prototype.getExtraMenuOptions = function (_, options) {
			if (this.imgs) {
				// If this node has images then we add an open in new tab item
				let img;
				if (this.imageIndex != null) {
					// An image is selected so select that
					img = this.imgs[this.imageIndex];
				} else if (this.overIndex != null) {
					// No image is selected but one is hovered
					img = this.imgs[this.overIndex];
				}
				if (img) {
pythongosssss's avatar
pythongosssss committed
261
262
263
					options.unshift(
						{
							content: "Open Image",
264
265
266
267
268
							callback: () => {
								let url = new URL(img.src);
								url.searchParams.delete('preview');
								window.open(url, "_blank")
							},
pythongosssss's avatar
pythongosssss committed
269
270
271
272
273
						},
						{
							content: "Save Image",
							callback: () => {
								const a = document.createElement("a");
274
275
276
277
								let url = new URL(img.src);
								url.searchParams.delete('preview');
								a.href = url;
								a.setAttribute("download", new URLSearchParams(url.search).get("filename"));
pythongosssss's avatar
pythongosssss committed
278
279
280
281
282
283
								document.body.append(a);
								a.click();
								requestAnimationFrame(() => a.remove());
							},
						}
					);
pythongosssss's avatar
pythongosssss committed
284
285
				}
			}
ltdrdata's avatar
ltdrdata committed
286

287
288
289
290
291
			options.push({
					content: "Bypass",
					callback: (obj) => { if (this.mode === 4) this.mode = 0; else this.mode = 4; this.graph.change(); }
				});

292
293
294
295
296
297
			// prevent conflict of clipspace content
			if(!ComfyApp.clipspace_return_node) {
				options.push({
						content: "Copy (Clipspace)",
						callback: (obj) => { ComfyApp.copyToClipspace(this); }
					});
ltdrdata's avatar
ltdrdata committed
298

299
300
301
302
303
304
				if(ComfyApp.clipspace != null) {
					options.push({
							content: "Paste (Clipspace)",
							callback: () => { ComfyApp.pasteFromClipspace(this); }
						});
				}
305

306
307
308
309
310
311
312
313
314
315
				if(ComfyApp.isImageNode(this)) {
					options.push({
							content: "Open in MaskEditor",
							callback: (obj) => {
								ComfyApp.copyToClipspace(this);
								ComfyApp.clipspace_return_node = this;
								ComfyApp.open_maskeditor();
							}
						});
				}
ltdrdata's avatar
ltdrdata committed
316
			}
pythongosssss's avatar
pythongosssss committed
317
318
319
		};
	}

320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
	#addNodeKeyHandler(node) {
		const app = this;
		const origNodeOnKeyDown = node.prototype.onKeyDown;

		node.prototype.onKeyDown = function(e) {
			if (origNodeOnKeyDown && origNodeOnKeyDown.apply(this, e) === false) {
				return false;
			}

			if (this.flags.collapsed || !this.imgs || this.imageIndex === null) {
				return;
			}

			let handled = false;

			if (e.key === "ArrowLeft" || e.key === "ArrowRight") {
				if (e.key === "ArrowLeft") {
					this.imageIndex -= 1;
				} else if (e.key === "ArrowRight") {
					this.imageIndex += 1;
				}
				this.imageIndex %= this.imgs.length;

				if (this.imageIndex < 0) {
					this.imageIndex = this.imgs.length + this.imageIndex;
				}
				handled = true;
			} else if (e.key === "Escape") {
				this.imageIndex = null;
				handled = true;
			}

			if (handled === true) {
				e.preventDefault();
				e.stopImmediatePropagation();
				return false;
			}
		}
	}

pythongosssss's avatar
pythongosssss committed
360
361
362
363
364
	/**
	 * Adds Custom drawing logic for nodes
	 * e.g. Draws images and handles thumbnail navigation on nodes that output images
	 * @param {*} node The node to add the draw handler
	 */
pythongosssss's avatar
pythongosssss committed
365
366
	#addDrawBackgroundHandler(node) {
		const app = this;
367
368
369
370
371
372
373
374
375
376
377

		function getImageTop(node) {
			let shiftY;
			if (node.imageOffset != null) {
				shiftY = node.imageOffset;
			} else {
				if (node.widgets?.length) {
					const w = node.widgets[node.widgets.length - 1];
					shiftY = w.last_y;
					if (w.computeSize) {
						shiftY += w.computeSize()[1] + 4;
378
379
380
381
382
					}
					else if(w.computedHeight) {
						shiftY += w.computedHeight;
					}
					else {
383
384
385
386
387
388
389
390
391
392
						shiftY += LiteGraph.NODE_WIDGET_HEIGHT + 4;
					}
				} else {
					shiftY = node.computeSize()[1];
				}
			}
			return shiftY;
		}

		node.prototype.setSizeForImage = function () {
393
394
395
396
			if (this.inputHeight) {
				this.setSize(this.size);
				return;
			}
397
398
399
400
401
402
			const minHeight = getImageTop(this) + 220;
			if (this.size[1] < minHeight) {
				this.setSize([this.size[0], minHeight]);
			}
		};

pythongosssss's avatar
pythongosssss committed
403
404
		node.prototype.onDrawBackground = function (ctx) {
			if (!this.flags.collapsed) {
space-nuko's avatar
space-nuko committed
405
406
407
				let imgURLs = []
				let imagesChanged = false

pythongosssss's avatar
pythongosssss committed
408
409
410
411
				const output = app.nodeOutputs[this.id + ""];
				if (output && output.images) {
					if (this.images !== output.images) {
						this.images = output.images;
space-nuko's avatar
space-nuko committed
412
413
						imagesChanged = true;
						imgURLs = imgURLs.concat(output.images.map(params => {
414
							return api.apiURL("/view?" + new URLSearchParams(params).toString() + app.getPreviewFormatParam());
space-nuko's avatar
space-nuko committed
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
						}))
					}
				}

				const preview = app.nodePreviewImages[this.id + ""]
				if (this.preview !== preview) {
					this.preview = preview
					imagesChanged = true;
					if (preview != null) {
						imgURLs.push(preview);
					}
				}

				if (imagesChanged) {
					this.imageIndex = null;
					if (imgURLs.length > 0) {
pythongosssss's avatar
pythongosssss committed
431
						Promise.all(
space-nuko's avatar
space-nuko committed
432
							imgURLs.map((src) => {
pythongosssss's avatar
pythongosssss committed
433
434
435
								return new Promise((r) => {
									const img = new Image();
									img.onload = () => r(img);
comfyanonymous's avatar
comfyanonymous committed
436
									img.onerror = () => r(null);
space-nuko's avatar
space-nuko committed
437
									img.src = src
pythongosssss's avatar
pythongosssss committed
438
439
440
								});
							})
						).then((imgs) => {
space-nuko's avatar
space-nuko committed
441
							if ((!output || this.images === output.images) && (!preview || this.preview === preview)) {
pythongosssss's avatar
pythongosssss committed
442
								this.imgs = imgs.filter(Boolean);
443
								this.setSizeForImage?.();
pythongosssss's avatar
pythongosssss committed
444
445
446
447
								app.graph.setDirtyCanvas(true);
							}
						});
					}
space-nuko's avatar
space-nuko committed
448
449
450
					else {
						this.imgs = null;
					}
pythongosssss's avatar
pythongosssss committed
451
				}
pythongosssss's avatar
pythongosssss committed
452

pythongosssss's avatar
pythongosssss committed
453
454
455
456
457
458
				if (this.imgs && this.imgs.length) {
					const canvas = graph.list_of_graphcanvas[0];
					const mouse = canvas.graph_mouse;
					if (!canvas.pointer_is_down && this.pointerDown) {
						if (mouse[0] === this.pointerDown.pos[0] && mouse[1] === this.pointerDown.pos[1]) {
							this.imageIndex = this.pointerDown.index;
pythongosssss's avatar
pythongosssss committed
459
						}
pythongosssss's avatar
pythongosssss committed
460
461
						this.pointerDown = null;
					}
pythongosssss's avatar
pythongosssss committed
462

pythongosssss's avatar
pythongosssss committed
463
464
465
466
467
468
469
					let w = this.imgs[0].naturalWidth;
					let h = this.imgs[0].naturalHeight;
					let imageIndex = this.imageIndex;
					const numImages = this.imgs.length;
					if (numImages === 1 && !imageIndex) {
						this.imageIndex = imageIndex = 0;
					}
pythongosssss's avatar
pythongosssss committed
470

471
					const shiftY = getImageTop(this);
pythongosssss's avatar
pythongosssss committed
472

pythongosssss's avatar
pythongosssss committed
473
474
475
					let dw = this.size[0];
					let dh = this.size[1];
					dh -= shiftY;
pythongosssss's avatar
pythongosssss committed
476

pythongosssss's avatar
pythongosssss committed
477
478
479
480
481
482
483
484
485
486
487
488
					if (imageIndex == null) {
						let best = 0;
						let cellWidth;
						let cellHeight;
						let cols = 0;
						let shiftX = 0;
						for (let c = 1; c <= numImages; c++) {
							const rows = Math.ceil(numImages / c);
							const cW = dw / c;
							const cH = dh / rows;
							const scaleX = cW / w;
							const scaleY = cH / h;
pythongosssss's avatar
pythongosssss committed
489

pythongosssss's avatar
pythongosssss committed
490
491
492
493
							const scale = Math.min(scaleX, scaleY, 1);
							const imageW = w * scale;
							const imageH = h * scale;
							const area = imageW * imageH * numImages;
pythongosssss's avatar
pythongosssss committed
494

pythongosssss's avatar
pythongosssss committed
495
496
497
498
499
500
							if (area > best) {
								best = area;
								cellWidth = imageW;
								cellHeight = imageH;
								cols = c;
								shiftX = c * ((cW - imageW) / 2);
pythongosssss's avatar
pythongosssss committed
501
							}
pythongosssss's avatar
pythongosssss committed
502
						}
pythongosssss's avatar
pythongosssss committed
503

pythongosssss's avatar
pythongosssss committed
504
505
506
507
508
509
510
511
						let anyHovered = false;
						this.imageRects = [];
						for (let i = 0; i < numImages; i++) {
							const img = this.imgs[i];
							const row = Math.floor(i / cols);
							const col = i % cols;
							const x = col * cellWidth + shiftX;
							const y = row * cellHeight + shiftY;
pythongosssss's avatar
pythongosssss committed
512
							if (!anyHovered) {
pythongosssss's avatar
pythongosssss committed
513
								anyHovered = LiteGraph.isInsideRectangle(
pythongosssss's avatar
pythongosssss committed
514
515
516
517
									mouse[0],
									mouse[1],
									x + this.pos[0],
									y + this.pos[1],
pythongosssss's avatar
pythongosssss committed
518
519
									cellWidth,
									cellHeight
pythongosssss's avatar
pythongosssss committed
520
								);
pythongosssss's avatar
pythongosssss committed
521
522
523
								if (anyHovered) {
									this.overIndex = i;
									let value = 110;
pythongosssss's avatar
pythongosssss committed
524
									if (canvas.pointer_is_down) {
pythongosssss's avatar
pythongosssss committed
525
526
527
528
										if (!this.pointerDown || this.pointerDown.index !== i) {
											this.pointerDown = { index: i, pos: [...mouse] };
										}
										value = 125;
pythongosssss's avatar
pythongosssss committed
529
									}
pythongosssss's avatar
pythongosssss committed
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
									ctx.filter = `contrast(${value}%) brightness(${value}%)`;
									canvas.canvas.style.cursor = "pointer";
								}
							}
							this.imageRects.push([x, y, cellWidth, cellHeight]);
							ctx.drawImage(img, x, y, cellWidth, cellHeight);
							ctx.filter = "none";
						}

						if (!anyHovered) {
							this.pointerDown = null;
							this.overIndex = null;
						}
					} else {
						// Draw individual
						const scaleX = dw / w;
						const scaleY = dh / h;
						const scale = Math.min(scaleX, scaleY, 1);

						w *= scale;
						h *= scale;

						let x = (dw - w) / 2;
						let y = (dh - h) / 2 + shiftY;
						ctx.drawImage(this.imgs[imageIndex], x, y, w, h);

						const drawButton = (x, y, sz, text) => {
							const hovered = LiteGraph.isInsideRectangle(mouse[0], mouse[1], x + this.pos[0], y + this.pos[1], sz, sz);
							let fill = "#333";
							let textFill = "#fff";
							let isClicking = false;
							if (hovered) {
								canvas.canvas.style.cursor = "pointer";
								if (canvas.pointer_is_down) {
									fill = "#1e90ff";
									isClicking = true;
pythongosssss's avatar
pythongosssss committed
566
								} else {
pythongosssss's avatar
pythongosssss committed
567
568
									fill = "#eee";
									textFill = "#000";
pythongosssss's avatar
pythongosssss committed
569
								}
pythongosssss's avatar
pythongosssss committed
570
571
572
							} else {
								this.pointerWasDown = null;
							}
pythongosssss's avatar
pythongosssss committed
573

pythongosssss's avatar
pythongosssss committed
574
575
576
577
578
579
580
581
							ctx.fillStyle = fill;
							ctx.beginPath();
							ctx.roundRect(x, y, sz, sz, [4]);
							ctx.fill();
							ctx.fillStyle = textFill;
							ctx.font = "12px Arial";
							ctx.textAlign = "center";
							ctx.fillText(text, x + 15, y + 20);
pythongosssss's avatar
pythongosssss committed
582

pythongosssss's avatar
pythongosssss committed
583
584
							return isClicking;
						};
pythongosssss's avatar
pythongosssss committed
585

pythongosssss's avatar
pythongosssss committed
586
587
588
589
590
						if (numImages > 1) {
							if (drawButton(x + w - 35, y + h - 35, 30, `${this.imageIndex + 1}/${numImages}`)) {
								let i = this.imageIndex + 1 >= numImages ? 0 : this.imageIndex + 1;
								if (!this.pointerDown || !this.pointerDown.index === i) {
									this.pointerDown = { index: i, pos: [...mouse] };
pythongosssss's avatar
pythongosssss committed
591
								}
pythongosssss's avatar
pythongosssss committed
592
							}
pythongosssss's avatar
pythongosssss committed
593

pythongosssss's avatar
pythongosssss committed
594
595
596
							if (drawButton(x + w - 35, y + 5, 30, `x`)) {
								if (!this.pointerDown || !this.pointerDown.index === null) {
									this.pointerDown = { index: null, pos: [...mouse] };
pythongosssss's avatar
pythongosssss committed
597
598
599
600
601
602
603
604
605
								}
							}
						}
					}
				}
			}
		};
	}

pythongosssss's avatar
pythongosssss committed
606
607
608
	/**
	 * Adds a handler allowing drag+drop of files onto the window to load workflows
	 */
pythongosssss's avatar
pythongosssss committed
609
610
611
612
613
	#addDropHandler() {
		// Get prompt from dropped PNG or json
		document.addEventListener("drop", async (event) => {
			event.preventDefault();
			event.stopPropagation();
614
615
616
617
618

			const n = this.dragOverNode;
			this.dragOverNode = null;
			// Node handles file drop, we dont use the built in onDropFile handler as its buggy
			// If you drag multiple files it will call it multiple times with the same file
619
			if (n && n.onDragDrop && (await n.onDragDrop(event))) {
620
621
				return;
			}
622
623
			// Dragging from Chrome->Firefox there is a file but its a bmp, so ignore that
			if (event.dataTransfer.files.length && event.dataTransfer.files[0].type !== "image/bmp") {
624
			await this.handleFile(event.dataTransfer.files[0]);
625
626
627
628
629
630
631
632
633
634
635
			} else {
				// Try loading the first URI in the transfer list
				const validTypes = ["text/uri-list", "text/x-moz-url"];
				const match = [...event.dataTransfer.types].find((t) => validTypes.find(v => t === v));
				if (match) {
					const uri = event.dataTransfer.getData(match)?.split("\n")?.[0];
					if (uri) {
						await this.handleFile(await (await fetch(uri)).blob());
					}
				}
			}
pythongosssss's avatar
pythongosssss committed
636
		});
637

638
639
640
641
642
643
		// Always clear over node on drag leave
		this.canvasEl.addEventListener("dragleave", async () => {
			if (this.dragOverNode) {
				this.dragOverNode = null;
				this.graph.setDirtyCanvas(false, true);
			}
pythongosssss's avatar
pythongosssss committed
644
		});
645

646
647
648
649
650
651
652
653
654
		// Add handler for dropping onto a specific node
		this.canvasEl.addEventListener(
			"dragover",
			(e) => {
				this.canvas.adjustMouseEvent(e);
				const node = this.graph.getNodeOnPos(e.canvasX, e.canvasY);
				if (node) {
					if (node.onDragOver && node.onDragOver(e)) {
						this.dragOverNode = node;
655
656

						// dragover event is fired very frequently, run this on an animation frame
657
658
659
660
661
662
663
664
665
666
						requestAnimationFrame(() => {
							this.graph.setDirtyCanvas(false, true);
						});
						return;
					}
				}
				this.dragOverNode = null;
			},
			false
		);
pythongosssss's avatar
pythongosssss committed
667
668
	}

pythongosssss's avatar
pythongosssss committed
669
	/**
670
	 * Adds a handler on paste that extracts and loads images or workflows from pasted JSON data
pythongosssss's avatar
pythongosssss committed
671
	 */
pythongosssss's avatar
pythongosssss committed
672
673
	#addPasteHandler() {
		document.addEventListener("paste", (e) => {
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
			let data = (e.clipboardData || window.clipboardData);
			const items = data.items;

			// Look for image paste data
			for (const item of items) {
				if (item.type.startsWith('image/')) {
					var imageNode = null;

					// If an image node is selected, paste into it
					if (this.canvas.current_node &&
						this.canvas.current_node.is_selected &&
						ComfyApp.isImageNode(this.canvas.current_node)) {
						imageNode = this.canvas.current_node;
					}

					// No image node selected: add a new one
					if (!imageNode) {
						const newNode = LiteGraph.createNode("LoadImage");
						newNode.pos = [...this.canvas.graph_mouse];
						imageNode = this.graph.add(newNode);
						this.graph.change();
					}
					const blob = item.getAsFile();
					imageNode.pasteFile(blob);
					return;
				}
			}

			// No image found. Look for node data
			data = data.getData("text/plain");
pythongosssss's avatar
pythongosssss committed
704
705
706
707
708
709
710
711
712
713
714
			let workflow;
			try {
				data = data.slice(data.indexOf("{"));
				workflow = JSON.parse(data);
			} catch (err) {
				try {
					data = data.slice(data.indexOf("workflow\n"));
					data = data.slice(data.indexOf("{"));
					workflow = JSON.parse(data);
				} catch (error) {}
			}
pythongosssss's avatar
pythongosssss committed
715

pythongosssss's avatar
pythongosssss committed
716
717
718
			if (workflow && workflow.version && workflow.nodes && workflow.extra) {
				this.loadGraphData(workflow);
			}
719
			else {
720
721
722
723
				if (e.target.type === "text" || e.target.type === "textarea") {
					return;
				}

724
				// Litegraph default paste
725
				this.canvas.pasteFromClipboard();
726
727
728
729
730
731
732
733
734
735
736
737
			}


		});
	}


	/**
	 * Adds a handler on copy that serializes selected nodes to JSON
	 */
	#addCopyHandler() {
		document.addEventListener("copy", (e) => {
Michael Abrahams's avatar
Michael Abrahams committed
738
739
740
741
742
			if (e.target.type === "text" || e.target.type === "textarea") {
				// Default system copy
				return;
			}
			// copy nodes and clear clipboard
743
			if (this.canvas.selected_nodes) {
Michael Abrahams's avatar
Michael Abrahams committed
744
745
746
747
748
				this.canvas.copyToClipboard();
				e.clipboardData.clearData();
				e.preventDefault();
				e.stopImmediatePropagation();
				return false;
749
			}
pythongosssss's avatar
pythongosssss committed
750
751
752
		});
	}

753

Jairo Correa's avatar
Jairo Correa committed
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
	/**
	 * Handle mouse
	 *
	 * Move group by header
	 */
	#addProcessMouseHandler() {
		const self = this;

		const origProcessMouseDown = LGraphCanvas.prototype.processMouseDown;
		LGraphCanvas.prototype.processMouseDown = function(e) {
			const res = origProcessMouseDown.apply(this, arguments);

			this.selected_group_moving = false;

			if (this.selected_group && !this.selected_group_resizing) {
				var font_size =
					this.selected_group.font_size || LiteGraph.DEFAULT_GROUP_FONT_SIZE;
				var height = font_size * 1.4;

				// Move group by header
				if (LiteGraph.isInsideRectangle(e.canvasX, e.canvasY, this.selected_group.pos[0], this.selected_group.pos[1], this.selected_group.size[0], height)) {
					this.selected_group_moving = true;
				}
			}

			return res;
		}

		const origProcessMouseMove = LGraphCanvas.prototype.processMouseMove;
		LGraphCanvas.prototype.processMouseMove = function(e) {
			const orig_selected_group = this.selected_group;

			if (this.selected_group && !this.selected_group_resizing && !this.selected_group_moving) {
				this.selected_group = null;
			}

			const res = origProcessMouseMove.apply(this, arguments);

			if (orig_selected_group && !this.selected_group_resizing && !this.selected_group_moving) {
				this.selected_group = orig_selected_group;
			}

			return res;
		};
	}

800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
	/**
	 * Handle keypress
	 *
	 * Ctrl + M mute/unmute selected nodes
	 */
	#addProcessKeyHandler() {
		const self = this;
		const origProcessKey = LGraphCanvas.prototype.processKey;
		LGraphCanvas.prototype.processKey = function(e) {
			if (!this.graph) {
				return;
			}

			var block_default = false;

			if (e.target.localName == "input") {
				return;
			}

819
820
			if (e.type == "keydown" && !e.repeat) {

821
				// Ctrl + M mute/unmute
822
				if (e.key === 'm' && e.ctrlKey) {
823
824
825
826
827
828
829
830
831
832
833
					if (this.selected_nodes) {
						for (var i in this.selected_nodes) {
							if (this.selected_nodes[i].mode === 2) { // never
								this.selected_nodes[i].mode = 0; // always
							} else {
								this.selected_nodes[i].mode = 2; // never
							}
						}
					}
					block_default = true;
				}
834

835
836
				// Ctrl + B bypass
				if (e.key === 'b' && e.ctrlKey) {
837
838
839
840
841
842
843
844
845
846
847
					if (this.selected_nodes) {
						for (var i in this.selected_nodes) {
							if (this.selected_nodes[i].mode === 4) { // never
								this.selected_nodes[i].mode = 0; // always
							} else {
								this.selected_nodes[i].mode = 4; // never
							}
						}
					}
					block_default = true;
				}
848
849
850
851
852

				// Ctrl+C Copy
				if ((e.key === 'c') && (e.metaKey || e.ctrlKey)) {
					if (e.shiftKey) {
						this.copyToClipboard(true);
Michael Abrahams's avatar
Michael Abrahams committed
853
						e.clipboardData.clearData();
854
855
						block_default = true;
					}
Michael Abrahams's avatar
Michael Abrahams committed
856
857
858
859
					else {
						// Trigger onCopy
						return true;
					}
860
861
862
863
864
865
866
867
868
				}

				// Ctrl+V Paste
				if ((e.key === 'v') && (e.metaKey || e.ctrlKey)) {
					if (e.shiftKey) {
						this.pasteFromClipboard(true);
						block_default = true;
					}
					else {
Michael Abrahams's avatar
Michael Abrahams committed
869
						// Trigger onPaste
870
871
872
						return true;
					}
				}
873
874
875
876
877
878
879
880
881
882
			}

			this.graph.change();

			if (block_default) {
				e.preventDefault();
				e.stopImmediatePropagation();
				return false;
			}

883
884
			// Fall through to Litegraph defaults
			return origProcessKey.apply(this, arguments);
885
886
887
		};
	}

Jairo Correa's avatar
Jairo Correa committed
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
	/**
	 * Draws group header bar
	 */
	#addDrawGroupsHandler() {
		const self = this;

		const origDrawGroups = LGraphCanvas.prototype.drawGroups;
		LGraphCanvas.prototype.drawGroups = function(canvas, ctx) {
			if (!this.graph) {
				return;
			}

			var groups = this.graph._groups;

			ctx.save();
			ctx.globalAlpha = 0.7 * this.editor_alpha;

			for (var i = 0; i < groups.length; ++i) {
				var group = groups[i];

				if (!LiteGraph.overlapBounding(this.visible_area, group._bounding)) {
					continue;
				} //out of the visible area

				ctx.fillStyle = group.color || "#335";
				ctx.strokeStyle = group.color || "#335";
				var pos = group._pos;
				var size = group._size;
				ctx.globalAlpha = 0.25 * this.editor_alpha;
				ctx.beginPath();
				var font_size =
					group.font_size || LiteGraph.DEFAULT_GROUP_FONT_SIZE;
				ctx.rect(pos[0] + 0.5, pos[1] + 0.5, size[0], font_size * 1.4);
				ctx.fill();
				ctx.globalAlpha = this.editor_alpha;
			}

			ctx.restore();

			const res = origDrawGroups.apply(this, arguments);
			return res;
		}
	}

pythongosssss's avatar
pythongosssss committed
932
	/**
pythongosssss's avatar
pythongosssss committed
933
	 * Draws node highlights (executing, drag drop) and progress bar
pythongosssss's avatar
pythongosssss committed
934
	 */
935
	#addDrawNodeHandler() {
936
		const origDrawNodeShape = LGraphCanvas.prototype.drawNodeShape;
pythongosssss's avatar
pythongosssss committed
937
		const self = this;
938

pythongosssss's avatar
pythongosssss committed
939
		LGraphCanvas.prototype.drawNodeShape = function (node, ctx, size, fgcolor, bgcolor, selected, mouse_over) {
940
			const res = origDrawNodeShape.apply(this, arguments);
pythongosssss's avatar
pythongosssss committed
941

942
			const nodeErrors = self.lastNodeErrors?.[node.id];
943

944
			let color = null;
945
			let lineWidth = 1;
946
947
948
949
950
			if (node.id === +self.runningNodeId) {
				color = "#0f0";
			} else if (self.dragOverNode && node.id === self.dragOverNode.id) {
				color = "dodgerblue";
			}
951
			else if (nodeErrors?.errors) {
952
953
954
				color = "red";
				lineWidth = 2;
			}
955
956
957
958
			else if (self.lastExecutionError && +self.lastExecutionError.node_id === node.id) {
				color = "#f0f";
				lineWidth = 2;
			}
959
960

			if (color) {
pythongosssss's avatar
pythongosssss committed
961
				const shape = node._shape || node.constructor.shape || LiteGraph.ROUND_SHAPE;
962
				ctx.lineWidth = lineWidth;
pythongosssss's avatar
pythongosssss committed
963
964
965
				ctx.globalAlpha = 0.8;
				ctx.beginPath();
				if (shape == LiteGraph.BOX_SHAPE)
pythongosssss's avatar
pythongosssss committed
966
					ctx.rect(-6, -6 - LiteGraph.NODE_TITLE_HEIGHT, 12 + size[0] + 1, 12 + size[1] + LiteGraph.NODE_TITLE_HEIGHT);
pythongosssss's avatar
pythongosssss committed
967
968
969
970
971
972
973
974
975
976
977
				else if (shape == LiteGraph.ROUND_SHAPE || (shape == LiteGraph.CARD_SHAPE && node.flags.collapsed))
					ctx.roundRect(
						-6,
						-6 - LiteGraph.NODE_TITLE_HEIGHT,
						12 + size[0] + 1,
						12 + size[1] + LiteGraph.NODE_TITLE_HEIGHT,
						this.round_radius * 2
					);
				else if (shape == LiteGraph.CARD_SHAPE)
					ctx.roundRect(
						-6,
pythongosssss's avatar
pythongosssss committed
978
						-6 - LiteGraph.NODE_TITLE_HEIGHT,
pythongosssss's avatar
pythongosssss committed
979
980
						12 + size[0] + 1,
						12 + size[1] + LiteGraph.NODE_TITLE_HEIGHT,
pythongosssss's avatar
pythongosssss committed
981
						[this.round_radius * 2, this.round_radius * 2, 2, 2]
pythongosssss's avatar
pythongosssss committed
982
				);
pythongosssss's avatar
pythongosssss committed
983
984
				else if (shape == LiteGraph.CIRCLE_SHAPE)
					ctx.arc(size[0] * 0.5, size[1] * 0.5, size[0] * 0.5 + 6, 0, Math.PI * 2);
985
				ctx.strokeStyle = color;
pythongosssss's avatar
pythongosssss committed
986
987
988
				ctx.stroke();
				ctx.strokeStyle = fgcolor;
				ctx.globalAlpha = 1;
989
990
991
992
993
994
995
			}

			if (self.progress && node.id === +self.runningNodeId) {
				ctx.fillStyle = "green";
				ctx.fillRect(0, 0, size[0] * (self.progress.value / self.progress.max), 6);
				ctx.fillStyle = bgcolor;
			}
pythongosssss's avatar
pythongosssss committed
996

997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
			// Highlight inputs that failed validation
			if (nodeErrors) {
				ctx.lineWidth = 2;
				ctx.strokeStyle = "red";
				for (const error of nodeErrors.errors) {
					if (error.extra_info && error.extra_info.input_name) {
						const inputIndex = node.findInputSlot(error.extra_info.input_name)
						if (inputIndex !== -1) {
							let pos = node.getConnectionPos(true, inputIndex);
							ctx.beginPath();
							ctx.arc(pos[0] - node.pos[0], pos[1] - node.pos[1], 12, 0, 2 * Math.PI, false)
							ctx.stroke();
						}
					}
pythongosssss's avatar
pythongosssss committed
1011
1012
1013
1014
1015
				}
			}

			return res;
		};
1016
1017
1018
1019

		const origDrawNode = LGraphCanvas.prototype.drawNode;
		LGraphCanvas.prototype.drawNode = function (node, ctx) {
			var editor_alpha = this.editor_alpha;
1020
			var old_color = node.bgcolor;
1021
1022
1023
1024
1025

			if (node.mode === 2) { // never
				this.editor_alpha = 0.4;
			}

1026
1027
1028
1029
1030
			if (node.mode === 4) { // never
				node.bgcolor = "#FF00FF";
				this.editor_alpha = 0.2;
			}

1031
1032
1033
			const res = origDrawNode.apply(this, arguments);

			this.editor_alpha = editor_alpha;
1034
			node.bgcolor = old_color;
1035
1036
1037

			return res;
		};
pythongosssss's avatar
pythongosssss committed
1038
1039
	}

pythongosssss's avatar
pythongosssss committed
1040
1041
1042
	/**
	 * Handles updates from the API socket
	 */
pythongosssss's avatar
pythongosssss committed
1043
	#addApiUpdateHandlers() {
pythongosssss's avatar
pythongosssss committed
1044
1045
		api.addEventListener("status", ({ detail }) => {
			this.ui.setStatus(detail);
pythongosssss's avatar
pythongosssss committed
1046
1047
		});

pythongosssss's avatar
pythongosssss committed
1048
1049
1050
		api.addEventListener("reconnecting", () => {
			this.ui.dialog.show("Reconnecting...");
		});
pythongosssss's avatar
pythongosssss committed
1051

pythongosssss's avatar
pythongosssss committed
1052
1053
1054
		api.addEventListener("reconnected", () => {
			this.ui.dialog.close();
		});
pythongosssss's avatar
pythongosssss committed
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064

		api.addEventListener("progress", ({ detail }) => {
			this.progress = detail;
			this.graph.setDirtyCanvas(true, false);
		});

		api.addEventListener("executing", ({ detail }) => {
			this.progress = null;
			this.runningNodeId = detail;
			this.graph.setDirtyCanvas(true, false);
space-nuko's avatar
space-nuko committed
1065
			delete this.nodePreviewImages[this.runningNodeId]
pythongosssss's avatar
pythongosssss committed
1066
1067
		});

pythongosssss's avatar
pythongosssss committed
1068
1069
		api.addEventListener("executed", ({ detail }) => {
			this.nodeOutputs[detail.node] = detail.output;
1070
			const node = this.graph.getNodeById(detail.node);
space-nuko's avatar
space-nuko committed
1071
1072
1073
			if (node) {
				if (node.onExecuted)
					node.onExecuted(detail.output);
1074
			}
pythongosssss's avatar
pythongosssss committed
1075
		});
pythongosssss's avatar
pythongosssss committed
1076

1077
		api.addEventListener("execution_start", ({ detail }) => {
space-nuko's avatar
space-nuko committed
1078
			this.runningNodeId = null;
1079
			this.lastExecutionError = null
Chris's avatar
Chris committed
1080
1081
1082
1083
			this.graph._nodes.forEach((node) => {
				if (node.onExecutionStart)
					node.onExecutionStart()
			})
1084
1085
1086
1087
1088
1089
1090
1091
1092
		});

		api.addEventListener("execution_error", ({ detail }) => {
			this.lastExecutionError = detail;
			const formattedError = this.#formatExecutionError(detail);
			this.ui.dialog.show(formattedError);
			this.canvas.draw(true, true);
		});

space-nuko's avatar
space-nuko committed
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
		api.addEventListener("b_preview", ({ detail }) => {
			const id = this.runningNodeId
			if (id == null)
				return;

			const blob = detail
			const blobUrl = URL.createObjectURL(blob)
			this.nodePreviewImages[id] = [blobUrl]
		});

pythongosssss's avatar
pythongosssss committed
1103
1104
1105
		api.init();
	}

1106
1107
	#addKeyboardHandler() {
		window.addEventListener("keydown", (e) => {
pythongosssss's avatar
pythongosssss committed
1108
			this.shiftDown = e.shiftKey;
1109
		});
pythongosssss's avatar
pythongosssss committed
1110
1111
1112
		window.addEventListener("keyup", (e) => {
			this.shiftDown = e.shiftKey;
		});
1113
1114
	}

1115
	/**
1116
	 * Loads all extensions from the API into the window in parallel
1117
1118
	 */
	async #loadExtensions() {
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
	    const extensions = await api.getExtensions();
	    this.logging.addEntry("Comfy.App", "debug", { Extensions: extensions });
	
	    const extensionPromises = extensions.map(async ext => {
	        try {
	            await import(api.apiURL(ext));
	        } catch (error) {
	            console.error("Error loading extension", ext, error);
	        }
	    });
	
	    await Promise.all(extensionPromises);
1131
1132
	}

pythongosssss's avatar
pythongosssss committed
1133
1134
1135
1136
	/**
	 * Set up the app on the page
	 */
	async setup() {
1137
1138
		await this.#loadExtensions();

pythongosssss's avatar
pythongosssss committed
1139
		// Create and mount the LiteGraph in the DOM
ltdrdata's avatar
ltdrdata committed
1140
1141
1142
		const mainCanvas = document.createElement("canvas")
		mainCanvas.style.touchAction = "none"
		const canvasEl = (this.canvasEl = Object.assign(mainCanvas, { id: "graph-canvas" }));
1143
		canvasEl.tabIndex = "1";
pythongosssss's avatar
pythongosssss committed
1144
1145
		document.body.prepend(canvasEl);

Jairo Correa's avatar
Jairo Correa committed
1146
		this.#addProcessMouseHandler();
1147
		this.#addProcessKeyHandler();
Jairo Correa's avatar
Jairo Correa committed
1148

pythongosssss's avatar
pythongosssss committed
1149
1150
1151
1152
		this.graph = new LGraph();
		const canvas = (this.canvas = new LGraphCanvas(canvasEl, this.graph));
		this.ctx = canvasEl.getContext("2d");

1153
1154
1155
		LiteGraph.release_link_on_empty_shows_menu = true;
		LiteGraph.alt_drag_do_clone_nodes = true;

pythongosssss's avatar
pythongosssss committed
1156
1157
1158
		this.graph.start();

		function resizeCanvas() {
1159
1160
1161
1162
1163
1164
			// Limit minimal scale to 1, see https://github.com/comfyanonymous/ComfyUI/pull/845
			const scale = Math.max(window.devicePixelRatio, 1);
			const { width, height } = canvasEl.getBoundingClientRect();
			canvasEl.width = Math.round(width * scale);
			canvasEl.height = Math.round(height * scale);
			canvasEl.getContext("2d").scale(scale, scale);
pythongosssss's avatar
pythongosssss committed
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
			canvas.draw(true, true);
		}

		// Ensure the canvas fills the window
		resizeCanvas();
		window.addEventListener("resize", resizeCanvas);

		await this.#invokeExtensionsAsync("init");
		await this.registerNodes();

		// Load previous workflow
		let restored = false;
		try {
			const json = localStorage.getItem("workflow");
			if (json) {
				const workflow = JSON.parse(json);
				this.loadGraphData(workflow);
				restored = true;
			}
1184
1185
1186
		} catch (err) {
			console.error("Error loading previous workflow", err);
		}
pythongosssss's avatar
pythongosssss committed
1187
1188
1189

		// We failed to restore a workflow so load the default
		if (!restored) {
pythongosssss's avatar
pythongosssss committed
1190
			this.loadGraphData();
pythongosssss's avatar
pythongosssss committed
1191
1192
1193
1194
1195
		}

		// Save current workflow automatically
		setInterval(() => localStorage.setItem("workflow", JSON.stringify(this.graph.serialize())), 1000);

1196
		this.#addDrawNodeHandler();
Jairo Correa's avatar
Jairo Correa committed
1197
		this.#addDrawGroupsHandler();
pythongosssss's avatar
pythongosssss committed
1198
1199
		this.#addApiUpdateHandlers();
		this.#addDropHandler();
1200
		this.#addCopyHandler();
pythongosssss's avatar
pythongosssss committed
1201
		this.#addPasteHandler();
1202
		this.#addKeyboardHandler();
pythongosssss's avatar
pythongosssss committed
1203

pythongosssss's avatar
pythongosssss committed
1204
1205
1206
		await this.#invokeExtensionsAsync("setup");
	}

pythongosssss's avatar
pythongosssss committed
1207
1208
1209
	/**
	 * Registers nodes with the graph
	 */
pythongosssss's avatar
pythongosssss committed
1210
1211
1212
1213
	async registerNodes() {
		const app = this;
		// Load node definitions from the backend
		const defs = await api.getNodeDefs();
ltdrdata's avatar
ltdrdata committed
1214
		await this.registerNodesFromDefs(defs);
1215
1216
1217
1218
		await this.#invokeExtensionsAsync("registerCustomNodes");
	}

    async registerNodesFromDefs(defs) {
pythongosssss's avatar
pythongosssss committed
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
		await this.#invokeExtensionsAsync("addCustomNodeDefs", defs);

		// Generate list of known widgets
		const widgets = Object.assign(
			{},
			ComfyWidgets,
			...(await this.#invokeExtensionsAsync("getCustomWidgets")).filter(Boolean)
		);

		// Register a node for each definition
		for (const nodeId in defs) {
			const nodeData = defs[nodeId];
			const node = Object.assign(
				function ComfyNode() {
1233
1234
1235
1236
					var inputs = nodeData["input"]["required"];
					if (nodeData["input"]["optional"] != undefined){
					    inputs = Object.assign({}, nodeData["input"]["required"], nodeData["input"]["optional"])
					}
pythongosssss's avatar
pythongosssss committed
1237
1238
1239
1240
1241
					const config = { minWidth: 1, minHeight: 1 };
					for (const inputName in inputs) {
						const inputData = inputs[inputName];
						const type = inputData[0];

1242
						let widgetCreated = true;
1243
1244
1245
1246
1247
1248
1249
1250
1251
						if (Array.isArray(type)) {
							// Enums
							Object.assign(config, widgets.COMBO(this, inputName, inputData, app) || {});
						} else if (`${type}:${inputName}` in widgets) {
							// Support custom widgets by Type:Name
							Object.assign(config, widgets[`${type}:${inputName}`](this, inputName, inputData, app) || {});
						} else if (type in widgets) {
							// Standard type widgets
							Object.assign(config, widgets[type](this, inputName, inputData, app) || {});
1252
						} else {
1253
1254
							// Node connection inputs
							this.addInput(inputName, type);
1255
							widgetCreated = false;
1256
						}
1257
1258

						if(widgetCreated && inputData[1]?.forceInput && config?.widget) {
1259
							if (!config.widget.options) config.widget.options = {};
1260
							config.widget.options.forceInput = inputData[1].forceInput;
pythongosssss's avatar
pythongosssss committed
1261
1262
1263
						}
					}

1264
1265
1266
					for (const o in nodeData["output"]) {
						const output = nodeData["output"][o];
						const outputName = nodeData["output_name"][o] || output;
1267
1268
						const outputShape = nodeData["output_is_list"][o] ? LiteGraph.GRID_SHAPE : LiteGraph.CIRCLE_SHAPE ;
						this.addOutput(outputName, output, { shape: outputShape });
1269
1270
					}

pythongosssss's avatar
pythongosssss committed
1271
1272
1273
1274
1275
					const s = this.computeSize();
					s[0] = Math.max(config.minWidth, s[0] * 1.5);
					s[1] = Math.max(config.minHeight, s[1]);
					this.size = s;
					this.serialize_widgets = true;
1276
1277

					app.#invokeExtensionsAsync("nodeCreated", this);
pythongosssss's avatar
pythongosssss committed
1278
1279
				},
				{
City's avatar
City committed
1280
					title: nodeData.display_name || nodeData.name,
pythongosssss's avatar
pythongosssss committed
1281
1282
1283
1284
1285
1286
1287
					comfyClass: nodeData.name,
				}
			);
			node.prototype.comfyClass = nodeData.name;

			this.#addNodeContextMenuHandler(node);
			this.#addDrawBackgroundHandler(node, app);
1288
			this.#addNodeKeyHandler(node);
pythongosssss's avatar
pythongosssss committed
1289

1290
			await this.#invokeExtensionsAsync("beforeRegisterNodeDef", node, nodeData);
pythongosssss's avatar
pythongosssss committed
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
			LiteGraph.registerNodeType(nodeId, node);
			node.category = nodeData.category;
		}
	}

	/**
	 * Populates the graph with the specified workflow data
	 * @param {*} graphData A serialized graph object
	 */
	loadGraphData(graphData) {
1301
1302
		this.clean();

1303
		let reset_invalid_values = false;
pythongosssss's avatar
pythongosssss committed
1304
		if (!graphData) {
1305
			graphData = structuredClone(defaultGraph);
1306
			reset_invalid_values = true;
pythongosssss's avatar
pythongosssss committed
1307
		}
1308

pythongosssss's avatar
const  
pythongosssss committed
1309
		const missingNodeTypes = [];
1310
		for (let n of graphData.nodes) {
1311
			// Patch T2IAdapterLoader to ControlNetLoader since they are the same node now
1312
			if (n.type == "T2IAdapterLoader") n.type = "ControlNetLoader";
1313
1314
1315
1316
1317

			// Find missing node types
			if (!(n.type in LiteGraph.registered_node_types)) {
				missingNodeTypes.push(n.type);
			}
1318
1319
		}

1320
1321
1322
		try {
			this.graph.configure(graphData);
		} catch (error) {
pythongosssss's avatar
pythongosssss committed
1323
			let errorHint = [];
pythongosssss's avatar
pythongosssss committed
1324
1325
			// Try extracting filename to see if it was caused by an extension script
			const filename = error.fileName || (error.stack || "").match(/(\/extensions\/.*\.js)/)?.[1];
1326
1327
			const pos = (filename || "").indexOf("/extensions/");
			if (pos > -1) {
pythongosssss's avatar
pythongosssss committed
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
				errorHint.push(
					$el("span", { textContent: "This may be due to the following script:" }),
					$el("br"),
					$el("span", {
						style: {
							fontWeight: "bold",
						},
						textContent: filename.substring(pos),
					})
				);
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
			}

			// Show dialog to let the user know something went wrong loading the data
			this.ui.dialog.show(
				$el("div", [
					$el("p", { textContent: "Loading aborted due to error reloading workflow data" }),
					$el("pre", {
						style: { padding: "5px", backgroundColor: "rgba(255,0,0,0.2)" },
						textContent: error.toString(),
					}),
					$el("pre", {
						style: {
							padding: "5px",
							color: "#ccc",
							fontSize: "10px",
							maxHeight: "50vh",
							overflow: "auto",
pythongosssss's avatar
pythongosssss committed
1355
							backgroundColor: "rgba(0,0,0,0.2)",
1356
1357
1358
						},
						textContent: error.stack || "No stacktrace available",
					}),
pythongosssss's avatar
pythongosssss committed
1359
					...errorHint,
1360
1361
				]).outerHTML
			);
pythongosssss's avatar
pythongosssss committed
1362

1363
1364
			return;
		}
pythongosssss's avatar
pythongosssss committed
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378

		for (const node of this.graph._nodes) {
			const size = node.computeSize();
			size[0] = Math.max(node.size[0], size[0]);
			size[1] = Math.max(node.size[1], size[1]);
			node.size = size;

			if (node.widgets) {
				// If you break something in the backend and want to patch workflows in the frontend
				// This is the place to do this
				for (let widget of node.widgets) {
					if (node.type == "KSampler" || node.type == "KSamplerAdvanced") {
						if (widget.name == "sampler_name") {
							if (widget.value.startsWith("sample_")) {
pythongosssss's avatar
pythongosssss committed
1379
								widget.value = widget.value.slice(7);
pythongosssss's avatar
pythongosssss committed
1380
1381
							}
						}
1382
1383
					}
					if (node.type == "KSampler" || node.type == "KSamplerAdvanced" || node.type == "PrimitiveNode") {
1384
						if (widget.name == "control_after_generate") {
1385
							if (widget.value === true) {
1386
								widget.value = "randomize";
1387
1388
							} else if (widget.value === false) {
								widget.value = "fixed";
1389
1390
							}
						}
pythongosssss's avatar
pythongosssss committed
1391
					}
1392
1393
1394
1395
1396
1397
1398
					if (reset_invalid_values) {
						if (widget.type == "combo") {
							if (!widget.options.values.includes(widget.value) && widget.options.values.length > 0) {
								widget.value = widget.options.values[0];
							}
						}
					}
pythongosssss's avatar
pythongosssss committed
1399
1400
1401
1402
1403
				}
			}

			this.#invokeExtensions("loadedGraphNode", node);
		}
1404
1405
1406

		if (missingNodeTypes.length) {
			this.ui.dialog.show(
pythongosssss's avatar
pythongosssss committed
1407
				`When loading the graph, the following node types were not found: <ul>${Array.from(new Set(missingNodeTypes)).map(
1408
					(t) => `<li>${t}</li>`
pythongosssss's avatar
pythongosssss committed
1409
				).join("")}</ul>Nodes that have failed to load will show as red on the graph.`
1410
			);
pythongosssss's avatar
pythongosssss committed
1411
			this.logging.addEntry("Comfy.App", "warn", {
1412
				MissingNodes: missingNodeTypes,
pythongosssss's avatar
pythongosssss committed
1413
			});
1414
		}
pythongosssss's avatar
pythongosssss committed
1415
1416
	}

pythongosssss's avatar
pythongosssss committed
1417
1418
1419
1420
	/**
	 * Converts the current graph workflow for sending to the API
	 * @returns The workflow and node links
	 */
1421
	async graphToPrompt() {
pythongosssss's avatar
pythongosssss committed
1422
1423
		const workflow = this.graph.serialize();
		const output = {};
1424
1425
1426
		// Process nodes in order of execution
		for (const node of this.graph.computeExecutionOrder(false)) {
			const n = workflow.nodes.find((n) => n.id === node.id);
1427
1428

			if (node.isVirtualNode) {
1429
1430
1431
1432
				// Don't serialize frontend only nodes but let them make changes
				if (node.applyToGraph) {
					node.applyToGraph(workflow);
				}
1433
1434
1435
				continue;
			}

1436
			if (node.mode === 2 || node.mode === 4) {
1437
1438
1439
1440
				// Don't serialize muted nodes
				continue;
			}

1441
			const inputs = {};
pythongosssss's avatar
pythongosssss committed
1442
1443
1444
1445
			const widgets = node.widgets;

			// Store all widget values
			if (widgets) {
1446
1447
				for (const i in widgets) {
					const widget = widgets[i];
1448
					if (!widget.options || widget.options.serialize !== false) {
1449
						inputs[widget.name] = widget.serializeValue ? await widget.serializeValue(n, i) : widget.value;
pythongosssss's avatar
pythongosssss committed
1450
1451
1452
1453
1454
1455
					}
				}
			}

			// Store all node links
			for (let i in node.inputs) {
1456
1457
				let parent = node.getInputNode(i);
				if (parent) {
1458
					let link = node.getInputLink(i);
1459
					while (parent.mode === 4 || parent.isVirtualNode) {
1460
						let found = false;
1461
1462
1463
						if (parent.isVirtualNode) {
							link = parent.getInputLink(link.origin_slot);
							if (link) {
1464
								parent = parent.getInputNode(link.target_slot);
1465
1466
1467
1468
1469
								if (parent) {
									found = true;
								}
							}
						} else if (link && parent.mode === 4) {
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
							let all_inputs = [link.origin_slot];
							if (parent.inputs) {
								all_inputs = all_inputs.concat(Object.keys(parent.inputs))
								for (let parent_input in all_inputs) {
									parent_input = all_inputs[parent_input];
									if (parent.inputs[parent_input].type === node.inputs[i].type) {
										link = parent.getInputLink(parent_input);
										if (link) {
											parent = parent.getInputNode(parent_input);
										}
										found = true;
										break;
1482
1483
1484
1485
									}
								}
							}
						}
1486

1487
1488
1489
1490
1491
						if (!found) {
							break;
						}
					}

1492
1493
1494
					if (link) {
						inputs[node.inputs[i].name] = [String(link.origin_id), parseInt(link.origin_slot)];
					}
pythongosssss's avatar
pythongosssss committed
1495
1496
1497
1498
1499
1500
1501
1502
1503
				}
			}

			output[String(node.id)] = {
				inputs,
				class_type: node.comfyClass,
			};
		}

1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
		// Remove inputs connected to removed nodes

		for (const o in output) {
			for (const i in output[o].inputs) {
				if (Array.isArray(output[o].inputs[i])
					&& output[o].inputs[i].length === 2
					&& !output[output[o].inputs[i][0]]) {
					delete output[o].inputs[i];
				}
			}
		}

pythongosssss's avatar
pythongosssss committed
1516
1517
1518
		return { workflow, output };
	}

1519
	#formatPromptError(error) {
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
		if (error == null) {
			return "(unknown error)"
		}
		else if (typeof error === "string") {
			return error;
		}
		else if (error.stack && error.message) {
			return error.toString()
		}
		else if (error.response) {
			let message = error.response.error.message;
			if (error.response.error.details)
			message += ": " + error.response.error.details;
			for (const [nodeID, nodeError] of Object.entries(error.response.node_errors)) {
			message += "\n" + nodeError.class_type + ":"
				for (const errorReason of nodeError.errors) {
					message += "\n    - " + errorReason.message + ": " + errorReason.details
				}
			}
			return message
		}
		return "(unknown error)"
	}

1544
1545
1546
1547
1548
1549
1550
1551
1552
	#formatExecutionError(error) {
		if (error == null) {
			return "(unknown error)"
		}

		const traceback = error.traceback.join("")
		const nodeId = error.node_id
		const nodeType = error.node_type

1553
		return `Error occurred when executing ${nodeType}:\n\n${error.exception_message}\n\n${traceback}`
1554
1555
	}

m957ymj75urz's avatar
m957ymj75urz committed
1556
	async queuePrompt(number, batchCount = 1) {
1557
		this.#queueItems.push({ number, batchCount });
pythongosssss's avatar
pythongosssss committed
1558

1559
1560
1561
1562
		// Only have one action process the items so each one gets a unique seed correctly
		if (this.#processingQueue) {
			return;
		}
1563

1564
		this.#processingQueue = true;
1565
		this.lastNodeErrors = null;
1566

1567
1568
1569
		try {
			while (this.#queueItems.length) {
				({ number, batchCount } = this.#queueItems.pop());
pythongosssss's avatar
pythongosssss committed
1570

1571
1572
1573
1574
				for (let i = 0; i < batchCount; i++) {
					const p = await this.graphToPrompt();

					try {
1575
1576
1577
1578
1579
						const res = await api.queuePrompt(number, p);
						this.lastNodeErrors = res.node_errors;
						if (this.lastNodeErrors.length > 0) {
							this.canvas.draw(true, true);
						}
1580
					} catch (error) {
1581
						const formattedError = this.#formatPromptError(error)
1582
1583
						this.ui.dialog.show(formattedError);
						if (error.response) {
1584
							this.lastNodeErrors = error.response.node_errors;
1585
1586
							this.canvas.draw(true, true);
						}
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
						break;
					}

					for (const n of p.workflow.nodes) {
						const node = graph.getNodeById(n.id);
						if (node.widgets) {
							for (const widget of node.widgets) {
								// Allow widgets to run callbacks after a prompt has been queued
								// e.g. random seed after every gen
								if (widget.afterQueued) {
									widget.afterQueued();
								}
							}
m957ymj75urz's avatar
m957ymj75urz committed
1600
						}
pythongosssss's avatar
pythongosssss committed
1601
					}
1602
1603
1604

					this.canvas.draw(true, true);
					await this.ui.queue.update();
pythongosssss's avatar
pythongosssss committed
1605
1606
				}
			}
1607
1608
		} finally {
			this.#processingQueue = false;
m957ymj75urz's avatar
m957ymj75urz committed
1609
		}
pythongosssss's avatar
pythongosssss committed
1610
	}
pythongosssss's avatar
pythongosssss committed
1611

pythongosssss's avatar
pythongosssss committed
1612
1613
	/**
	 * Loads workflow data from the specified file
1614
	 * @param {File} file
pythongosssss's avatar
pythongosssss committed
1615
	 */
pythongosssss's avatar
pythongosssss committed
1616
1617
1618
	async handleFile(file) {
		if (file.type === "image/png") {
			const pngInfo = await getPngMetadata(file);
1619
1620
1621
1622
1623
1624
			if (pngInfo) {
				if (pngInfo.workflow) {
					this.loadGraphData(JSON.parse(pngInfo.workflow));
				} else if (pngInfo.parameters) {
					importA1111(this.graph, pngInfo.parameters);
				}
pythongosssss's avatar
pythongosssss committed
1625
			}
1626
		} else if (file.type === "application/json" || file.name?.endsWith(".json")) {
pythongosssss's avatar
pythongosssss committed
1627
1628
1629
1630
1631
			const reader = new FileReader();
			reader.onload = () => {
				this.loadGraphData(JSON.parse(reader.result));
			};
			reader.readAsText(file);
1632
		} else if (file.name?.endsWith(".latent") || file.name?.endsWith(".safetensors")) {
1633
1634
1635
1636
			const info = await getLatentMetadata(file);
			if (info.workflow) {
				this.loadGraphData(JSON.parse(info.workflow));
			}
pythongosssss's avatar
pythongosssss committed
1637
1638
		}
	}
1639

1640
1641
1642
1643
	/**
	 * Registers a Comfy web extension with the app
	 * @param {ComfyExtension} extension
	 */
1644
1645
1646
1647
1648
1649
1650
1651
1652
	registerExtension(extension) {
		if (!extension.name) {
			throw new Error("Extensions must have a 'name' property.");
		}
		if (this.extensions.find((ext) => ext.name === extension.name)) {
			throw new Error(`Extension named '${extension.name}' already registered.`);
		}
		this.extensions.push(extension);
	}
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664

	/**
	 * Refresh combo list on whole nodes
	 */
	async refreshComboInNodes() {
		const defs = await api.getNodeDefs();

		for(let nodeNum in this.graph._nodes) {
			const node = this.graph._nodes[nodeNum];

			const def = defs[node.type];

1665
1666
1667
1668
1669
			// HOTFIX: The current patch is designed to prevent the rest of the code from breaking due to primitive nodes,
			//         and additional work is needed to consider the primitive logic in the refresh logic.
			if(!def)
				continue;

1670
1671
1672
			for(const widgetNum in node.widgets) {
				const widget = node.widgets[widgetNum]
				if(widget.type == "combo" && def["input"]["required"][widget.name] !== undefined) {
comfyanonymous's avatar
comfyanonymous committed
1673
					widget.options.values = def["input"]["required"][widget.name][0];
1674

1675
					if(widget.name != 'image' && !widget.options.values.includes(widget.value)) {
1676
						widget.value = widget.options.values[0];
1677
						widget.callback(widget.value);
1678
1679
1680
1681
1682
					}
				}
			}
		}
	}
1683
1684
1685
1686
1687
1688

	/**
	 * Clean current state
	 */
	clean() {
		this.nodeOutputs = {};
space-nuko's avatar
space-nuko committed
1689
		this.nodePreviewImages = {}
1690
		this.lastNodeErrors = null;
1691
		this.lastExecutionError = null;
space-nuko's avatar
space-nuko committed
1692
		this.runningNodeId = null;
1693
	}
pythongosssss's avatar
pythongosssss committed
1694
1695
1696
}

export const app = new ComfyApp();