UNPKG

@artifyfun/comfy-ui-client

Version:
399 lines (396 loc) 12 kB
"use strict"; 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); } }); } };