index.js 3.12 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
7
const fs = require("fs");
const path = require("path");

const html = fs.readFileSync(path.resolve(__dirname, "../../web/index.html"))
8
9
10

/**
 *
11
12
13
14
15
 * @param { Parameters<typeof mockApi>[0] & { 
 * 	resetEnv?: boolean, 
 * 	preSetup?(app): Promise<void>,
 *  localStorage?: Record<string, string> 
 * } } config
16
17
 * @returns
 */
pythongosssss's avatar
pythongosssss committed
18
19
export async function start(config = {}) {
	if(config.resetEnv) {
pythongosssss's avatar
pythongosssss committed
20
21
22
		jest.resetModules();
		jest.resetAllMocks();
        lg.setup(global);
23
24
		localStorage.clear();
		sessionStorage.clear();
pythongosssss's avatar
pythongosssss committed
25
26
	}

27
28
29
	Object.assign(localStorage, config.localStorage ?? {});
	document.body.innerHTML = html;

30
31
	mockApi(config);
	const { app } = require("../../web/scripts/app");
pythongosssss's avatar
pythongosssss committed
32
	config.preSetup?.(app);
33
	await app.setup();
34

pythongosssss's avatar
pythongosssss committed
35
	return { ...Ez.graph(app, global["LiteGraph"], global["LGraphCanvas"]), app };
36
37
38
}

/**
pythongosssss's avatar
pythongosssss committed
39
40
 * @param { ReturnType<Ez["graph"]>["graph"] } graph
 * @param { (hasReloaded: boolean) => (Promise<void> | void) } cb
41
42
43
44
45
46
47
48
 */
export async function checkBeforeAndAfterReload(graph, cb) {
	await cb(false);
	await graph.reload();
	await cb(true);
}

/**
pythongosssss's avatar
pythongosssss committed
49
50
 * @param { string } name
 * @param { Record<string, string | [string | string[], any]> } input
51
 * @param { (string | string[])[] | Record<string, string | string[]> } output
pythongosssss's avatar
pythongosssss committed
52
 * @returns { Record<string, import("../../web/types/comfy").ComfyObjectInfo> }
53
54
55
56
57
58
59
60
61
 */
export function makeNodeDef(name, input, output = {}) {
	const nodeDef = {
		name,
		category: "test",
		output: [],
		output_name: [],
		output_is_list: [],
		input: {
pythongosssss's avatar
pythongosssss committed
62
			required: {},
63
64
		},
	};
pythongosssss's avatar
pythongosssss committed
65
	for (const k in input) {
66
67
		nodeDef.input.required[k] = typeof input[k] === "string" ? [input[k], {}] : [...input[k]];
	}
pythongosssss's avatar
pythongosssss committed
68
	if (output instanceof Array) {
69
70
71
		output = output.reduce((p, c) => {
			p[c] = c;
			return p;
pythongosssss's avatar
pythongosssss committed
72
		}, {});
73
	}
pythongosssss's avatar
pythongosssss committed
74
	for (const k in output) {
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
		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
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
}

/**
 *
 * @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
121
122
123
124
125
126
127
128
129

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

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