index.js 2.54 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 } } config
8
9
10
 * @returns
 */
export async function start(config = undefined) {
pythongosssss's avatar
pythongosssss committed
11
12
13
14
15
16
	if(config?.resetEnv) {
		jest.resetModules();
		jest.resetAllMocks();
        lg.setup(global);
	}

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

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

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

/**
 *
 * @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 };
}