@artifyfun/comfy-ui-client
Version:
WebSockets API client for ComfyUI
399 lines (396 loc) • 12 kB
JavaScript
;
var __create = Object.create;
var __defProp = Object.defineProperty;
var __getOwnPropDesc = Object.getOwnPropertyDescriptor;
var __getOwnPropNames = Object.getOwnPropertyNames;
var __getProtoOf = Object.getPrototypeOf;
var __hasOwnProp = Object.prototype.hasOwnProperty;
var __export = (target, all) => {
for (var name in all)
__defProp(target, name, { get: all[name], enumerable: true });
};
var __copyProps = (to, from, except, desc) => {
if (from && typeof from === "object" || typeof from === "function") {
for (let key of __getOwnPropNames(from))
if (!__hasOwnProp.call(to, key) && key !== except)
__defProp(to, key, { get: () => from[key], enumerable: !(desc = __getOwnPropDesc(from, key)) || desc.enumerable });
}
return to;
};
var __toESM = (mod, isNodeMode, target) => (target = mod != null ? __create(__getProtoOf(mod)) : {}, __copyProps(
// If the importer is in node compatibility mode or this is not an ESM
// file that has been converted to a CommonJS file using a Babel-
// compatible transform (i.e. "__esModule" has not been set), then set
// "default" to the CommonJS "module.exports" for node compatibility.
isNodeMode || !mod || !mod.__esModule ? __defProp(target, "default", { value: mod, enumerable: true }) : target,
mod
));
var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: true }), mod);
// src/index.ts
var src_exports = {};
__export(src_exports, {
ComfyUIClient: () => ComfyUIClient
});
module.exports = __toCommonJS(src_exports);
// src/client.ts
var import_pino = __toESM(require("pino"));
var import_isomorphic_ws = __toESM(require("isomorphic-ws"));
var logger = (0, import_pino.default)({
level: "info"
});
var ComfyUIClient = class {
host;
clientId;
historyResult = {};
eventEmitter = () => {
};
handlers;
ws;
constructor(host, clientId, eventEmitter) {
this.host = host;
this.clientId = clientId;
this.eventEmitter = eventEmitter || (() => {
});
this.handlers = {
open: [],
close: [],
error: [],
message: []
};
}
connect() {
return new Promise(async (resolve, reject) => {
if (this.ws) {
await this.disconnect();
}
const url = `${this.host.replace("http", "ws")}/ws?clientId=${this.clientId}`;
logger.info(`Connecting to url: ${url}`);
this.ws = new import_isomorphic_ws.default(url);
if (typeof window !== "undefined") {
this.ws.onopen = (event) => {
this.handlers.open.forEach((cb) => cb(event));
};
this.ws.onclose = (event) => {
this.handlers.close.forEach((cb) => cb(event));
};
this.ws.onerror = (event) => {
this.handlers.error.forEach((cb) => cb(event));
};
this.ws.onmessage = (event) => {
this.handlers.message.forEach((cb) => {
cb(event.data, event.data instanceof Blob);
});
};
this.ws.on = (event, callback) => {
if (this.handlers[event]) {
this.handlers[event].push(callback);
} else {
console.error(`Unknown event type: ${event}`);
}
};
this.ws.off = (event, callback) => {
if (this.handlers[event]) {
this.handlers[event] = this.handlers[event].filter(
(cb) => cb !== callback
);
}
};
}
this.ws.on("open", () => {
logger.info("Connection open");
resolve();
});
this.ws.on("close", () => {
logger.info("Connection closed");
});
this.ws.on("error", (err) => {
logger.error({ err }, "WebSockets error");
reject(err);
this.eventEmitter("error", err);
});
this.ws.on("message", (data, isBinary) => {
if (isBinary) {
logger.debug("Received binary data");
} else {
logger.debug("Received data: %s", data.toString());
this.eventEmitter("message", data);
}
});
});
}
async disconnect() {
if (this.ws) {
this.ws.close();
this.ws = void 0;
}
}
async getEmbeddings() {
const res = await fetch(`${this.host}/embeddings`);
const json = await res.json();
if ("error" in json) {
throw new Error(JSON.stringify(json));
}
return json;
}
async getExtensions() {
const res = await fetch(`${this.host}/extensions`);
const json = await res.json();
if ("error" in json) {
throw new Error(JSON.stringify(json));
}
return json;
}
async queuePrompt(prompt) {
const res = await fetch(`${this.host}/prompt`, {
method: "POST",
headers: {
Accept: "application/json",
"Content-Type": "application/json"
},
body: JSON.stringify({
prompt,
client_id: this.clientId
})
});
const json = await res.json();
if ("error" in json) {
throw new Error(JSON.stringify(json));
}
return json;
}
interrupt() {
return fetch(`${this.host}/interrupt`, {
method: "POST",
headers: {
Accept: "application/json",
"Content-Type": "application/json"
}
});
}
async editHistory(params) {
const res = await fetch(`${this.host}/history`, {
method: "POST",
headers: {
Accept: "application/json",
"Content-Type": "application/json"
},
body: JSON.stringify(params)
});
const json = await res.json();
if ("error" in json) {
throw new Error(JSON.stringify(json));
}
}
async uploadImage(image, filename, overwrite) {
const formData = new FormData();
formData.append("image", new Blob([image]), filename);
if (overwrite !== void 0) {
formData.append("overwrite", overwrite.toString());
}
const res = await fetch(`${this.host}/upload/image`, {
method: "POST",
body: formData
});
const json = await res.json();
if ("error" in json) {
throw new Error(JSON.stringify(json));
}
return json;
}
async uploadMask(image, filename, originalRef, overwrite) {
const formData = new FormData();
formData.append("image", new Blob([image]), filename);
formData.append("originalRef", JSON.stringify(originalRef));
if (overwrite !== void 0) {
formData.append("overwrite", overwrite.toString());
}
const res = await fetch(`${this.host}/upload/mask`, {
method: "POST",
body: formData
});
const json = await res.json();
if ("error" in json) {
throw new Error(JSON.stringify(json));
}
return json;
}
async getImage(filename, subfolder, type) {
const res = await fetch(
`${this.host}/view?` + new URLSearchParams({
filename,
subfolder,
type
})
);
const blob = await res.blob();
return blob;
}
async viewMetadata(folderName, filename) {
const res = await fetch(
`${this.host}/view_metadata/${folderName}?filename=${filename}`
);
const json = await res.json();
if ("error" in json) {
throw new Error(JSON.stringify(json));
}
return json;
}
async getSystemStats() {
const res = await fetch(`${this.host}/system_stats`);
const json = await res.json();
if ("error" in json) {
throw new Error(JSON.stringify(json));
}
return json;
}
async getPrompt() {
const res = await fetch(`${this.host}/prompt`);
const json = await res.json();
if ("error" in json) {
throw new Error(JSON.stringify(json));
}
return json;
}
async getObjectInfo(nodeClass) {
const res = await fetch(
`${this.host}/object_info` + (nodeClass ? `/${nodeClass}` : "")
);
const json = await res.json();
if ("error" in json) {
throw new Error(JSON.stringify(json));
}
return json;
}
async getHistory(fetchOptionOrPromptId, promptId) {
let fetchOption;
let actualPromptId;
if (typeof fetchOptionOrPromptId === "string" && promptId === void 0) {
actualPromptId = fetchOptionOrPromptId;
fetchOption = void 0;
} else {
fetchOption = fetchOptionOrPromptId;
actualPromptId = promptId;
}
const host = fetchOption ? fetchOption.host : this.host;
const method = fetchOption ? fetchOption.method : "get";
const res = await fetch(
`${host}/history` + (actualPromptId ? `/${actualPromptId}` : ""),
{
method
}
);
const json = await res.json();
if ("error" in json) {
throw new Error(JSON.stringify(json));
}
this.historyResult = json;
return json;
}
async getQueue(fetchOption) {
const host = fetchOption ? fetchOption.host : this.host;
const method = fetchOption ? fetchOption.method : "get";
const res = await fetch(`${host}/queue`, {
method
});
const json = await res.json();
if ("error" in json) {
throw new Error(JSON.stringify(json));
}
return json;
}
async deleteQueue(id) {
const res = await fetch(`${this.host}/queue`, {
method: "POST",
headers: {
Accept: "application/json",
"Content-Type": "application/json"
},
body: JSON.stringify({
delete: id
})
});
const json = await res.json();
if ("error" in json) {
throw new Error(JSON.stringify(json));
}
return json;
}
// async saveImages(response: ImagesResponse, outputDir: string) {
// for (const nodeId of Object.keys(response)) {
// for (const img of response[nodeId]) {
// const arrayBuffer = await img.blob.arrayBuffer();
// const outputPath = join(outputDir, img.image.filename);
// // @ts-ignore
// await writeFile(outputPath, Buffer.from(arrayBuffer));
// }
// }
// }
async getResult(fetchOptionOrPrompt, promptParam) {
let fetchOption;
let prompt;
if (promptParam === void 0) {
prompt = fetchOptionOrPrompt;
fetchOption = void 0;
} else {
fetchOption = fetchOptionOrPrompt;
prompt = promptParam;
}
const queue = await this.queuePrompt(prompt);
const promptId = queue.prompt_id;
return new Promise((resolve, reject) => {
const onMessage = async (data, isBinary) => {
if (isBinary) {
return;
}
try {
const message = JSON.parse(data.toString());
if (message.type === "executing") {
const messageData = message.data;
if (!messageData.node) {
const donePromptId = messageData.prompt_id;
logger.info(`Done executing prompt (ID: ${donePromptId})`);
if (messageData.prompt_id === promptId) {
const historyRes = await this.getHistory(fetchOption, promptId);
const history = historyRes[promptId];
this.ws?.off("message", onMessage);
return resolve(history);
}
}
}
} catch (err) {
return reject(err);
}
};
this.ws?.on("message", onMessage);
});
}
async getImages(prompt) {
return new Promise(async (resolve, reject) => {
try {
const outputImages = {};
const history = await this.getResult(prompt);
for (const nodeId of Object.keys(history.outputs)) {
const nodeOutput = history.outputs[nodeId];
if (nodeOutput.images) {
const imagesOutput = [];
for (const image of nodeOutput.images) {
const blob = await this.getImage(
image.filename,
image.subfolder,
image.type
);
imagesOutput.push({
blob,
image
});
}
outputImages[nodeId] = imagesOutput;
}
}
resolve(outputImages);
} catch (err) {
return reject(err);
}
});
}
};