index.js 2.79 KB
Newer Older
1
2
const { mockApi } = require("./setup");
const { Ez } = require("./ezgraph");
pythongosssss's avatar
pythongosssss committed
3
const lg = require("./litegraph");
4
5
6

/**
 *
pythongosssss's avatar
pythongosssss committed
7
 * @param { Parameters<mockApi>[0] & { resetEnv?: boolean, preSetup?(app): Promise<void> } } config
8
9
 * @returns
 */
pythongosssss's avatar
pythongosssss committed
10
11
export async function start(config = {}) {
	if(config.resetEnv) {
pythongosssss's avatar
pythongosssss committed
12
13
14
15
16
		jest.resetModules();
		jest.resetAllMocks();
        lg.setup(global);
	}

17
18
	mockApi(config);
	const { app } = require("../../web/scripts/app");
pythongosssss's avatar
pythongosssss committed
19
	config.preSetup?.(app);
20
	await app.setup();
pythongosssss's avatar
pythongosssss committed
21
	return { ...Ez.graph(app, global["LiteGraph"], global["LGraphCanvas"]), app };
22
23
24
}

/**
pythongosssss's avatar
pythongosssss committed
25
26
 * @param { ReturnType<Ez["graph"]>["graph"] } graph
 * @param { (hasReloaded: boolean) => (Promise<void> | void) } cb
27
28
29
30
31
32
33
34
 */
export async function checkBeforeAndAfterReload(graph, cb) {
	await cb(false);
	await graph.reload();
	await cb(true);
}

/**
pythongosssss's avatar
pythongosssss committed
35
36
 * @param { string } name
 * @param { Record<string, string | [string | string[], any]> } input
37
 * @param { (string | string[])[] | Record<string, string | string[]> } output
pythongosssss's avatar
pythongosssss committed
38
 * @returns { Record<string, import("../../web/types/comfy").ComfyObjectInfo> }
39
40
41
42
43
44
45
46
47
 */
export function makeNodeDef(name, input, output = {}) {
	const nodeDef = {
		name,
		category: "test",
		output: [],
		output_name: [],
		output_is_list: [],
		input: {
pythongosssss's avatar
pythongosssss committed
48
			required: {},
49
50
		},
	};
pythongosssss's avatar
pythongosssss committed
51
	for (const k in input) {
52
53
		nodeDef.input.required[k] = typeof input[k] === "string" ? [input[k], {}] : [...input[k]];
	}
pythongosssss's avatar
pythongosssss committed
54
	if (output instanceof Array) {
55
56
57
		output = output.reduce((p, c) => {
			p[c] = c;
			return p;
pythongosssss's avatar
pythongosssss committed
58
		}, {});
59
	}
pythongosssss's avatar
pythongosssss committed
60
	for (const k in output) {
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
		nodeDef.output.push(output[k]);
		nodeDef.output_name.push(k);
		nodeDef.output_is_list.push(false);
	}

	return { [name]: nodeDef };
}

/**
/**
 * @template { any } T
 * @param { T } x
 * @returns { x is Exclude<T, null | undefined> }
 */
export function assertNotNullOrUndefined(x) {
	expect(x).not.toEqual(null);
	expect(x).not.toEqual(undefined);
	return true;
pythongosssss's avatar
pythongosssss committed
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
}

/**
 *
 * @param { ReturnType<Ez["graph"]>["ez"] } ez
 * @param { ReturnType<Ez["graph"]>["graph"] } graph
 */
export function createDefaultWorkflow(ez, graph) {
	graph.clear();
	const ckpt = ez.CheckpointLoaderSimple();

	const pos = ez.CLIPTextEncode(ckpt.outputs.CLIP, { text: "positive" });
	const neg = ez.CLIPTextEncode(ckpt.outputs.CLIP, { text: "negative" });

	const empty = ez.EmptyLatentImage();
	const sampler = ez.KSampler(
		ckpt.outputs.MODEL,
		pos.outputs.CONDITIONING,
		neg.outputs.CONDITIONING,
		empty.outputs.LATENT
	);

	const decode = ez.VAEDecode(sampler.outputs.LATENT, ckpt.outputs.VAE);
	const save = ez.SaveImage(decode.outputs.IMAGE);
	graph.arrange();

	return { ckpt, pos, neg, empty, sampler, decode, save };
}
pythongosssss's avatar
pythongosssss committed
107
108
109
110
111
112
113
114
115

export async function getNodeDefs() {
	const { api } = require("../../web/scripts/api");
	return api.getNodeDefs();
}

export async function getNodeDef(nodeId) {
	return (await getNodeDefs())[nodeId];
}