Merge branch 'comfyanonymous:master' into batched_noise

README.md

@@ -45,6 +45,7 @@ Workflow examples can be found on the [Examples page](https://comfyanonymous.git
 |---------------------------|--------------------------------------------------------------------------------------------------------------------|
 | Ctrl + Enter              | Queue up current graph for generation                                                                              |
 | Ctrl + Shift + Enter      | Queue up current graph as first for generation                                                                     |
+| Ctrl + Z/Ctrl + Y         | Undo/Redo                                                                                                          |
 | Ctrl + S                  | Save workflow                                                                                                      |
 | Ctrl + O                  | Load workflow                                                                                                      |
 | Ctrl + A                  | Select all nodes                                                                                                   |
@@ -100,6 +101,7 @@ AMD users can install rocm and pytorch with pip if you don't have it already ins
 ```pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm5.6```
 
 This is the command to install the nightly with ROCm 5.7 that might have some performance improvements:
+
 ```pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm5.7```
 
 ### NVIDIA
@@ -192,7 +194,7 @@ To use a textual inversion concepts/embeddings in a text prompt put them in the
 
 Make sure you use the regular loaders/Load Checkpoint node to load checkpoints. It will auto pick the right settings depending on your GPU.
 
-You can set this command line setting to disable the upcasting to fp32 in some cross attention operations which will increase your speed. Note that this will very likely give you black images on SD2.x models. If you use xformers this option does not do anything.
+You can set this command line setting to disable the upcasting to fp32 in some cross attention operations which will increase your speed. Note that this will very likely give you black images on SD2.x models. If you use xformers or pytorch attention this option does not do anything.
 
 ```--dont-upcast-attention```
 

tests-ui/tests/extensions.test.js

@@ -0,0 +1,196 @@
+// @ts-check
+/// <reference path="../node_modules/@types/jest/index.d.ts" />
+const { start } = require("../utils");
+const lg = require("../utils/litegraph");
+
+describe("extensions", () => {
+	beforeEach(() => {
+		lg.setup(global);
+	});
+
+	afterEach(() => {
+		lg.teardown(global);
+	});
+
+	it("calls each extension hook", async () => {
+		const mockExtension = {
+			name: "TestExtension",
+			init: jest.fn(),
+			setup: jest.fn(),
+			addCustomNodeDefs: jest.fn(),
+			getCustomWidgets: jest.fn(),
+			beforeRegisterNodeDef: jest.fn(),
+			registerCustomNodes: jest.fn(),
+			loadedGraphNode: jest.fn(),
+			nodeCreated: jest.fn(),
+			beforeConfigureGraph: jest.fn(),
+			afterConfigureGraph: jest.fn(),
+		};
+
+		const { app, ez, graph } = await start({
+			async preSetup(app) {
+				app.registerExtension(mockExtension);
+			},
+		});
+
+		// Basic initialisation hooks should be called once, with app
+		expect(mockExtension.init).toHaveBeenCalledTimes(1);
+		expect(mockExtension.init).toHaveBeenCalledWith(app);
+
+		// Adding custom node defs should be passed the full list of nodes
+		expect(mockExtension.addCustomNodeDefs).toHaveBeenCalledTimes(1);
+		expect(mockExtension.addCustomNodeDefs.mock.calls[0][1]).toStrictEqual(app);
+		const defs = mockExtension.addCustomNodeDefs.mock.calls[0][0];
+		expect(defs).toHaveProperty("KSampler");
+		expect(defs).toHaveProperty("LoadImage");
+
+		// Get custom widgets is called once and should return new widget types
+		expect(mockExtension.getCustomWidgets).toHaveBeenCalledTimes(1);
+		expect(mockExtension.getCustomWidgets).toHaveBeenCalledWith(app);
+
+		// Before register node def will be called once per node type
+		const nodeNames = Object.keys(defs);
+		const nodeCount = nodeNames.length;
+		expect(mockExtension.beforeRegisterNodeDef).toHaveBeenCalledTimes(nodeCount);
+		for (let i = 0; i < nodeCount; i++) {
+			// It should be send the JS class and the original JSON definition
+			const nodeClass = mockExtension.beforeRegisterNodeDef.mock.calls[i][0];
+			const nodeDef = mockExtension.beforeRegisterNodeDef.mock.calls[i][1];
+
+			expect(nodeClass.name).toBe("ComfyNode");
+			expect(nodeClass.comfyClass).toBe(nodeNames[i]);
+			expect(nodeDef.name).toBe(nodeNames[i]);
+			expect(nodeDef).toHaveProperty("input");
+			expect(nodeDef).toHaveProperty("output");
+		}
+
+		// Register custom nodes is called once after registerNode defs to allow adding other frontend nodes
+		expect(mockExtension.registerCustomNodes).toHaveBeenCalledTimes(1);
+
+		// Before configure graph will be called here as the default graph is being loaded
+		expect(mockExtension.beforeConfigureGraph).toHaveBeenCalledTimes(1);
+		// it gets sent the graph data that is going to be loaded
+		const graphData = mockExtension.beforeConfigureGraph.mock.calls[0][0];
+
+		// A node created is fired for each node constructor that is called
+		expect(mockExtension.nodeCreated).toHaveBeenCalledTimes(graphData.nodes.length);
+		for (let i = 0; i < graphData.nodes.length; i++) {
+			expect(mockExtension.nodeCreated.mock.calls[i][0].type).toBe(graphData.nodes[i].type);
+		}
+
+		// Each node then calls loadedGraphNode to allow them to be updated
+		expect(mockExtension.loadedGraphNode).toHaveBeenCalledTimes(graphData.nodes.length);
+		for (let i = 0; i < graphData.nodes.length; i++) {
+			expect(mockExtension.loadedGraphNode.mock.calls[i][0].type).toBe(graphData.nodes[i].type);
+		}
+
+		// After configure is then called once all the setup is done
+		expect(mockExtension.afterConfigureGraph).toHaveBeenCalledTimes(1);
+
+		expect(mockExtension.setup).toHaveBeenCalledTimes(1);
+		expect(mockExtension.setup).toHaveBeenCalledWith(app);
+
+		// Ensure hooks are called in the correct order
+		const callOrder = [
+			"init",
+			"addCustomNodeDefs",
+			"getCustomWidgets",
+			"beforeRegisterNodeDef",
+			"registerCustomNodes",
+			"beforeConfigureGraph",
+			"nodeCreated",
+			"loadedGraphNode",
+			"afterConfigureGraph",
+			"setup",
+		];
+		for (let i = 1; i < callOrder.length; i++) {
+			const fn1 = mockExtension[callOrder[i - 1]];
+			const fn2 = mockExtension[callOrder[i]];
+			expect(fn1.mock.invocationCallOrder[0]).toBeLessThan(fn2.mock.invocationCallOrder[0]);
+		}
+
+		graph.clear();
+
+		// Ensure adding a new node calls the correct callback
+		ez.LoadImage();
+		expect(mockExtension.loadedGraphNode).toHaveBeenCalledTimes(graphData.nodes.length);
+		expect(mockExtension.nodeCreated).toHaveBeenCalledTimes(graphData.nodes.length + 1);
+		expect(mockExtension.nodeCreated.mock.lastCall[0].type).toBe("LoadImage");
+
+		// Reload the graph to ensure correct hooks are fired
+		await graph.reload();
+
+		// These hooks should not be fired again
+		expect(mockExtension.init).toHaveBeenCalledTimes(1);
+		expect(mockExtension.addCustomNodeDefs).toHaveBeenCalledTimes(1);
+		expect(mockExtension.getCustomWidgets).toHaveBeenCalledTimes(1);
+		expect(mockExtension.registerCustomNodes).toHaveBeenCalledTimes(1);
+		expect(mockExtension.beforeRegisterNodeDef).toHaveBeenCalledTimes(nodeCount);
+		expect(mockExtension.setup).toHaveBeenCalledTimes(1);
+
+		// These should be called again
+		expect(mockExtension.beforeConfigureGraph).toHaveBeenCalledTimes(2);
+		expect(mockExtension.nodeCreated).toHaveBeenCalledTimes(graphData.nodes.length + 2);
+		expect(mockExtension.loadedGraphNode).toHaveBeenCalledTimes(graphData.nodes.length + 1);
+		expect(mockExtension.afterConfigureGraph).toHaveBeenCalledTimes(2);
+	});
+
+	it("allows custom nodeDefs and widgets to be registered", async () => {
+		const widgetMock = jest.fn((node, inputName, inputData, app) => {
+			expect(node.constructor.comfyClass).toBe("TestNode");
+			expect(inputName).toBe("test_input");
+			expect(inputData[0]).toBe("CUSTOMWIDGET");
+			expect(inputData[1]?.hello).toBe("world");
+			expect(app).toStrictEqual(app);
+
+			return {
+				widget: node.addWidget("button", inputName, "hello", () => {}),
+			};
+		});
+
+		// Register our extension that adds a custom node + widget type
+		const mockExtension = {
+			name: "TestExtension",
+			addCustomNodeDefs: (nodeDefs) => {
+				nodeDefs["TestNode"] = {
+					output: [],
+					output_name: [],
+					output_is_list: [],
+					name: "TestNode",
+					display_name: "TestNode",
+					category: "Test",
+					input: {
+						required: {
+							test_input: ["CUSTOMWIDGET", { hello: "world" }],
+						},
+					},
+				};
+			},
+			getCustomWidgets: jest.fn(() => {
+				return {
+					CUSTOMWIDGET: widgetMock,
+				};
+			}),
+		};
+
+		const { graph, ez } = await start({
+			async preSetup(app) {
+				app.registerExtension(mockExtension);
+			},
+		});
+
+		expect(mockExtension.getCustomWidgets).toBeCalledTimes(1);
+
+		graph.clear();
+		expect(widgetMock).toBeCalledTimes(0);
+		const node = ez.TestNode();
+		expect(widgetMock).toBeCalledTimes(1);
+
+		// Ensure our custom widget is created
+		expect(node.inputs.length).toBe(0);
+		expect(node.widgets.length).toBe(1);
+		const w = node.widgets[0].widget;
+		expect(w.name).toBe("test_input");
+		expect(w.type).toBe("button");
+	});
+});

tests-ui/utils/index.js

@@ -4,11 +4,11 @@ const lg = require("./litegraph");
 
 /**
  *
- * @param { Parameters<mockApi>[0] & { resetEnv?: boolean } } config
+ * @param { Parameters<mockApi>[0] & { resetEnv?: boolean, preSetup?(app): Promise<void> } } config
  * @returns
  */
-export async function start(config = undefined) {
+export async function start(config = {}) {
-	if(config?.resetEnv) {
+	if(config.resetEnv) {
 		jest.resetModules();
 		jest.resetAllMocks();
         lg.setup(global);
@@ -16,6 +16,7 @@ export async function start(config = undefined) {
 
 	mockApi(config);
 	const { app } = require("../../web/scripts/app");
+	config.preSetup?.(app);
 	await app.setup();
 	return { ...Ez.graph(app, global["LiteGraph"], global["LGraphCanvas"]), app };
 }

web/extensions/core/maskeditor.js

@@ -42,7 +42,7 @@ async function uploadMask(filepath, formData) {
 	});
 
 	ComfyApp.clipspace.imgs[ComfyApp.clipspace['selectedIndex']] = new Image();
-	ComfyApp.clipspace.imgs[ComfyApp.clipspace['selectedIndex']].src = api.apiURL("/view?" + new URLSearchParams(filepath).toString() + app.getPreviewFormatParam());
+	ComfyApp.clipspace.imgs[ComfyApp.clipspace['selectedIndex']].src = api.apiURL("/view?" + new URLSearchParams(filepath).toString() + app.getPreviewFormatParam() + app.getRandParam());
 
 	if(ComfyApp.clipspace.images)
 		ComfyApp.clipspace.images[ComfyApp.clipspace['selectedIndex']] = filepath;

web/scripts/app.js

@@ -86,6 +86,10 @@ export class ComfyApp {
 			return "";
 	}
 
+	getRandParam() {
+		return "&rand=" + Math.random();
+	}
+
 	static isImageNode(node) {
 		return node.imgs || (node && node.widgets && node.widgets.findIndex(obj => obj.name === 'image') >= 0);
 	}
@@ -411,7 +415,7 @@ export class ComfyApp {
 		node.prototype.setSizeForImage = function (force) {
 			if(!force && this.animatedImages) return;
 
-			if (this.inputHeight) {
+			if (this.inputHeight || this.freeWidgetSpace > 210) {
 				this.setSize(this.size);
 				return;
 			}
@@ -437,7 +441,7 @@ export class ComfyApp {
 								return api.apiURL(
 									"/view?" +
 										new URLSearchParams(params).toString() +
-										(this.animatedImages ? "" : app.getPreviewFormatParam())
+										(this.animatedImages ? "" : app.getPreviewFormatParam()) + app.getRandParam()
 								);
 							})
 						);
@@ -1654,6 +1658,7 @@ export class ComfyApp {
 		if (missingNodeTypes.length) {
 			this.showMissingNodesError(missingNodeTypes);
 		}
+		await this.#invokeExtensionsAsync("afterConfigureGraph", missingNodeTypes);
 	}
 
 	/**

web/scripts/domWidget.js

@@ -120,6 +120,8 @@ function computeSize(size) {
 		freeSpace -= 220;
 	}
 
+	this.freeWidgetSpace = freeSpace;
+
 	if (freeSpace < 0) {
 		// Not enough space for all widgets so we need to grow
 		size[1] -= freeSpace;

web/scripts/widgets.js

@@ -310,7 +310,7 @@ export const ComfyWidgets = {
 				subfolder = name.substring(0, folder_separator);
 				name = name.substring(folder_separator + 1);
 			}
-			img.src = api.apiURL(`/view?filename=${encodeURIComponent(name)}&type=input&subfolder=${subfolder}${app.getPreviewFormatParam()}`);
+			img.src = api.apiURL(`/view?filename=${encodeURIComponent(name)}&type=input&subfolder=${subfolder}${app.getPreviewFormatParam()}${app.getRandParam()}`);
 			node.setSizeForImage?.();
 		}