├── bunfig.toml ├── bun.lockb ├── changelog.json ├── .prettierrc ├── src ├── contansts.ts ├── types │ ├── tool.ts │ ├── sampler.ts │ ├── error.ts │ ├── manager.ts │ ├── event.ts │ └── api.ts ├── tools.ts ├── features │ ├── abstract.ts │ ├── monitoring.ts │ └── manager.ts ├── socket.ts ├── prompt-builder.ts ├── pool.ts ├── call-wrapper.ts └── client.ts ├── CHANGELOG.md ├── index.ts ├── tsconfig.json ├── LICENSE ├── .github └── workflows │ └── release.yml ├── test ├── tools.spec.ts ├── builder.spec.ts └── pool.spec.ts ├── package.json ├── examples ├── example-img2img-workflow.json ├── example-t2i.ts ├── example-txt2img-workflow.json ├── example-txt2img-upscaled-workflow.json ├── example-t2i-upscaled.ts ├── example-i2i.ts ├── example-pool.ts └── example-pool-basic-auth.ts ├── .gitignore └── README.md /bunfig.toml: -------------------------------------------------------------------------------- 1 | [install] 2 | registry = "https://registry.npmjs.org/" 3 | -------------------------------------------------------------------------------- /bun.lockb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/comfy-addons/comfyui-sdk/HEAD/bun.lockb -------------------------------------------------------------------------------- /changelog.json: -------------------------------------------------------------------------------- 1 | { 2 | "output": "CHANGELOG.md", 3 | "template": "keepachangelog" 4 | } 5 | -------------------------------------------------------------------------------- /.prettierrc: -------------------------------------------------------------------------------- 1 | { 2 | "trailingComma": "none", 3 | "semi": true, 4 | "singleQuote": false, 5 | "jsxSingleQuote": true, 6 | "printWidth": 120, 7 | "tabWidth": 2 8 | } 9 | -------------------------------------------------------------------------------- /src/contansts.ts: -------------------------------------------------------------------------------- 1 | export const LOAD_CHECKPOINTS_EXTENSION = "CheckpointLoaderSimple"; 2 | export const LOAD_LORAS_EXTENSION = "LoraLoader"; 3 | export const LOAD_KSAMPLER_EXTENSION = "KSampler"; 4 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | ### Changelog 2 | 3 | All notable changes to this project will be documented in this file. Dates are displayed in UTC. 4 | 5 | Generated by [`auto-changelog`](https://github.com/CookPete/auto-changelog). 6 | 7 | #### 0.2.49 8 | 9 | - feat(api): change uploadImage and uploadMask methods file type from B… [`#28`](https://github.com/comfy-addons/comfyui-sdk/pull/28) 10 | -------------------------------------------------------------------------------- /src/types/tool.ts: -------------------------------------------------------------------------------- 1 | export type FixArr = T extends readonly any[] ? Omit> : T; 2 | export type DropInitDot = T extends `.${infer U}` ? U : T; 3 | export type _DeepKeys = T extends object 4 | ? { 5 | [K in (string | number) & keyof T]: `${`.${K}`}${"" | _DeepKeys>}`; 6 | }[(string | number) & keyof T] 7 | : never; 8 | 9 | export type DeepKeys = DropInitDot<_DeepKeys>>; 10 | 11 | export type Simplify = { [K in keyof T]: T[K] } & {}; 12 | -------------------------------------------------------------------------------- /index.ts: -------------------------------------------------------------------------------- 1 | export { ComfyApi } from "./src/client"; 2 | export { CallWrapper } from "./src/call-wrapper"; 3 | export { ComfyPool, EQueueMode } from "./src/pool"; 4 | export { PromptBuilder } from "./src/prompt-builder"; 5 | 6 | export { TSamplerName, TSchedulerName } from "./src/types/sampler"; 7 | 8 | /** 9 | * Polyfill for CustomEvent in old NodeJS versions 10 | */ 11 | if (typeof CustomEvent === "undefined") { 12 | (global as any).CustomEvent = class CustomEvent extends Event { 13 | detail: any; 14 | constructor(event: any, params: any = {}) { 15 | super(event, params); 16 | this.detail = params.detail || null; 17 | } 18 | }; 19 | } 20 | -------------------------------------------------------------------------------- /src/tools.ts: -------------------------------------------------------------------------------- 1 | export const randomInt = (min: number, max: number) => { 2 | return Math.floor(Math.random() * (max - min + 1) + min); 3 | }; 4 | 5 | export const delay = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)); 6 | 7 | export const seed = () => randomInt(10000000000, 999999999999); 8 | 9 | /** 10 | * Encode POSIX path to NT path 11 | * 12 | * For example: `SDXL/realvisxlV40` -> `SDXL\\realvisxlV40` 13 | * 14 | * Useful for loading model with Windows's ComfyUI Client 15 | */ 16 | export const encodeNTPath = (path: string) => { 17 | return path.replace(/\//g, "\\"); 18 | }; 19 | 20 | export const encodePosixPath = (path: string) => { 21 | return path.replace(/\\/g, "/"); 22 | }; 23 | -------------------------------------------------------------------------------- /src/types/sampler.ts: -------------------------------------------------------------------------------- 1 | export type TSamplerName = 2 | | "euler" 3 | | "euler_cfg_pp" 4 | | "euler_ancestral" 5 | | "euler_ancestral_cfg_pp" 6 | | "heun" 7 | | "heunpp2" 8 | | "dpm_2" 9 | | "dpm_2_ancestral" 10 | | "lms" 11 | | "dpm_fast" 12 | | "dpm_adaptive" 13 | | "dpmpp_2s_ancestral" 14 | | "dpmpp_sde" 15 | | "dpmpp_sde_gpu" 16 | | "dpmpp_2m" 17 | | "dpmpp_2m_sde" 18 | | "dpmpp_2m_sde_gpu" 19 | | "dpmpp_3m_sde" 20 | | "dpmpp_3m_sde_gpu" 21 | | "ddpm" 22 | | "lcm" 23 | | "ipndm" 24 | | "ipndm_v" 25 | | "deis" 26 | | "ddim" 27 | | "uni_pc" 28 | | "uni_pc_bh2"; 29 | 30 | export type TSchedulerName = "normal" | "karras" | "exponential" | "sgm_uniform" | "simple" | "ddim_uniform"; 31 | -------------------------------------------------------------------------------- /tsconfig.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "outDir": "build", 4 | "module": "esnext", 5 | "target": "es6", 6 | "lib": ["es6", "dom", "es2016", "es2017", "es2019", "esnext"], 7 | "allowSyntheticDefaultImports": true, 8 | "sourceMap": true, 9 | "allowJs": false, 10 | "skipLibCheck": true, 11 | "jsx": "react", 12 | "declaration": true, 13 | "resolveJsonModule": true, 14 | "moduleResolution": "node", 15 | "forceConsistentCasingInFileNames": true, 16 | "noImplicitReturns": true, 17 | "noImplicitThis": true, 18 | "noImplicitAny": true, 19 | "strictNullChecks": true, 20 | "noUnusedLocals": true, 21 | "noUnusedParameters": false, 22 | "esModuleInterop": true, 23 | "baseUrl": ".", 24 | "typeRoots": ["./node_modules/@types"] 25 | }, 26 | "include": ["src/**/*", "index.ts", "test", "build.ts"], 27 | "exclude": ["node_modules", "build"] 28 | } 29 | -------------------------------------------------------------------------------- /src/features/abstract.ts: -------------------------------------------------------------------------------- 1 | import { ComfyApi } from "src/client"; 2 | 3 | export abstract class AbstractFeature extends EventTarget { 4 | protected client: ComfyApi; 5 | protected supported = false; 6 | 7 | constructor(client: ComfyApi) { 8 | super(); 9 | this.client = client; 10 | } 11 | 12 | get isSupported() { 13 | return this.supported; 14 | } 15 | 16 | public on(type: string, callback: (event: any) => void, options?: AddEventListenerOptions | boolean) { 17 | this.addEventListener(type, callback as any, options); 18 | return () => this.off(type, callback); 19 | } 20 | 21 | public off(type: string, callback: (event: any) => void, options?: EventListenerOptions | boolean): void { 22 | this.removeEventListener(type, callback as any, options); 23 | } 24 | 25 | abstract destroy(): void; 26 | 27 | /** 28 | * Check if this feature is supported by the current client 29 | */ 30 | abstract checkSupported(): Promise; 31 | } 32 | -------------------------------------------------------------------------------- /src/types/error.ts: -------------------------------------------------------------------------------- 1 | export class CallWrapperError extends Error { 2 | name = "CallWrapperError"; 3 | } 4 | 5 | export class WentMissingError extends CallWrapperError { 6 | name = "WentMissingError"; 7 | } 8 | 9 | export class FailedCacheError extends CallWrapperError { 10 | name = "FailedCacheError"; 11 | } 12 | 13 | export class EnqueueFailedError extends CallWrapperError { 14 | name = "EnqueueFailedError"; 15 | } 16 | 17 | export class DisconnectedError extends CallWrapperError { 18 | name = "DisconnectedError"; 19 | } 20 | 21 | export class ExecutionFailedError extends CallWrapperError { 22 | name = "ExecutionFailedError"; 23 | } 24 | 25 | export class CustomEventError extends CallWrapperError { 26 | name = "CustomEventError"; 27 | } 28 | 29 | export class ExecutionInterruptedError extends CallWrapperError { 30 | name = "ExecutionInterruptedError"; 31 | } 32 | 33 | export class MissingNodeError extends CallWrapperError { 34 | name = "MissingNodeError"; 35 | } 36 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 Saintno 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /.github/workflows/release.yml: -------------------------------------------------------------------------------- 1 | name: Release 2 | 3 | on: 4 | push: 5 | branches: 6 | - main 7 | 8 | jobs: 9 | release: 10 | runs-on: ubuntu-latest 11 | 12 | steps: 13 | - name: Checkout code 14 | uses: actions/checkout@v4 15 | 16 | - name: Set up Bun.js 17 | uses: oven-sh/setup-bun@v2 18 | 19 | - name: Install dependencies 20 | run: bun install 21 | 22 | - name: Build project 23 | run: bun run build 24 | 25 | - name: Set up Node.js 26 | uses: actions/setup-node@v4 27 | with: 28 | node-version: 22 29 | always-auth: true 30 | registry-url: "https://registry.npmjs.org" 31 | scope: "@saintno" 32 | 33 | - name: Bump version 34 | run: npm version patch --no-git-tag-version 35 | 36 | - name: Update changelog 37 | run: npm run version 38 | 39 | - name: Commit changes 40 | run: | 41 | git config --global user.name 'github-actions[bot]' 42 | git config --global user.email 'github-actions[bot]@users.noreply.github.com' 43 | git add . 44 | git commit -m 'ci: bump version and update changelog' 45 | git push 46 | 47 | - name: Publish to NPM 48 | run: npm publish --access public 49 | env: 50 | NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }} 51 | -------------------------------------------------------------------------------- /test/tools.spec.ts: -------------------------------------------------------------------------------- 1 | import { randomInt, delay, seed, encodeNTPath, encodePosixPath } from "src/tools"; 2 | 3 | describe("randomInt", () => { 4 | it("should generate a random integer within the specified range", () => { 5 | const min = 1; 6 | const max = 10; 7 | const result = randomInt(min, max); 8 | expect(result).toBeGreaterThanOrEqual(min); 9 | expect(result).toBeLessThanOrEqual(max); 10 | expect(Number.isInteger(result)).toBe(true); 11 | }); 12 | }); 13 | 14 | describe("delay", () => { 15 | it("should delay execution for the specified number of milliseconds", async () => { 16 | const ms = 1000; 17 | const start = Date.now(); 18 | await delay(ms); 19 | const end = Date.now(); 20 | expect(end - start).toBeGreaterThanOrEqual(ms); 21 | }); 22 | }); 23 | 24 | describe("seed", () => { 25 | it("should generate a random seed within the specified range", () => { 26 | const result = seed(); 27 | expect(result).toBeGreaterThanOrEqual(10000000000); 28 | expect(result).toBeLessThanOrEqual(999999999999); 29 | expect(Number.isInteger(result)).toBe(true); 30 | }); 31 | }); 32 | 33 | describe("encodeNTPath", () => { 34 | it("should encode a POSIX path to an NT path", () => { 35 | const path = "SDXL/realvisxlV40"; 36 | const result = encodeNTPath(path); 37 | expect(result).toBe("SDXL\\realvisxlV40"); 38 | }); 39 | }); 40 | 41 | describe("encodePosixPath", () => { 42 | it("should encode an NT path to a POSIX path", () => { 43 | const path = "SDXL\\realvisxlV40"; 44 | const result = encodePosixPath(path); 45 | expect(result).toBe("SDXL/realvisxlV40"); 46 | }); 47 | }); 48 | -------------------------------------------------------------------------------- /package.json: -------------------------------------------------------------------------------- 1 | { 2 | "name": "@saintno/comfyui-sdk", 3 | "version": "0.2.49", 4 | "description": "SDK for ComfyUI", 5 | "main": "build/index.esm.js", 6 | "typings": "build/index.d.ts", 7 | "type": "module", 8 | "exports": { 9 | ".": { 10 | "import": "./build/index.esm.js", 11 | "require": "./build/index.cjs", 12 | "types": "./build/index.d.ts", 13 | "default": "./build/index.esm.js" 14 | } 15 | }, 16 | "repository": { 17 | "type": "git", 18 | "url": "git+https://github.com/tctien342/comfyui-sdk.git" 19 | }, 20 | "author": "tctien342", 21 | "license": "MIT", 22 | "bugs": { 23 | "url": "https://github.com/tctien342/comfyui-sdk/issues" 24 | }, 25 | "homepage": "https://github.com/tctien342/comfyui-sdk#readme", 26 | "devDependencies": { 27 | "@types/bun": "^1.1.11", 28 | "@types/jest": "^29.5.12", 29 | "@types/node": "^22.2.0", 30 | "@types/ws": "^8.5.12", 31 | "auto-changelog": "^2.3.0", 32 | "commitizen": "^4.2.4", 33 | "cz-conventional-changelog": "^3.3.0", 34 | "dts-bundle-generator": "^9.5.1", 35 | "prettier": "^3.4.2" 36 | }, 37 | "peerDependencies": { 38 | "typescript": "^5.0.0" 39 | }, 40 | "keywords": [ 41 | "comfyui", 42 | "sdk", 43 | "typescript", 44 | "ts", 45 | "comfy", 46 | "stable-diffusion", 47 | "comfyui-api", 48 | "comfyui-sdk" 49 | ], 50 | "dependencies": { 51 | "ws": "^8.18.0" 52 | }, 53 | "scripts": { 54 | "version": "auto-changelog -p", 55 | "commit": "git-cz", 56 | "build": "bun build.ts", 57 | "format": "prettier --write ." 58 | }, 59 | "files": [ 60 | "build" 61 | ], 62 | "config": { 63 | "commitizen": { 64 | "path": "./node_modules/cz-conventional-changelog" 65 | } 66 | } 67 | } 68 | -------------------------------------------------------------------------------- /examples/example-img2img-workflow.json: -------------------------------------------------------------------------------- 1 | { 2 | "3": { 3 | "inputs": { 4 | "seed": 802716049688215, 5 | "steps": 20, 6 | "cfg": 8, 7 | "sampler_name": "euler", 8 | "scheduler": "normal", 9 | "denoise": 0.45, 10 | "model": ["4", 0], 11 | "positive": ["6", 0], 12 | "negative": ["7", 0], 13 | "latent_image": ["11", 0] 14 | }, 15 | "class_type": "KSampler", 16 | "_meta": { 17 | "title": "KSampler" 18 | } 19 | }, 20 | "4": { 21 | "inputs": { 22 | "ckpt_name": "SD15/counterfeitV30_v30.safetensors" 23 | }, 24 | "class_type": "CheckpointLoaderSimple", 25 | "_meta": { 26 | "title": "Load Checkpoint" 27 | } 28 | }, 29 | "6": { 30 | "inputs": { 31 | "text": "picture of tom and jerry", 32 | "clip": ["4", 1] 33 | }, 34 | "class_type": "CLIPTextEncode", 35 | "_meta": { 36 | "title": "CLIP Text Encode (Prompt)" 37 | } 38 | }, 39 | "7": { 40 | "inputs": { 41 | "text": "", 42 | "clip": ["4", 1] 43 | }, 44 | "class_type": "CLIPTextEncode", 45 | "_meta": { 46 | "title": "CLIP Text Encode (Prompt)" 47 | } 48 | }, 49 | "8": { 50 | "inputs": { 51 | "samples": ["3", 0], 52 | "vae": ["4", 2] 53 | }, 54 | "class_type": "VAEDecode", 55 | "_meta": { 56 | "title": "VAE Decode" 57 | } 58 | }, 59 | "9": { 60 | "inputs": { 61 | "filename_prefix": "ComfyUI", 62 | "images": ["8", 0] 63 | }, 64 | "class_type": "SaveImage", 65 | "_meta": { 66 | "title": "Save Image" 67 | } 68 | }, 69 | "10": { 70 | "inputs": { 71 | "image": "tom.jpg", 72 | "upload": "image" 73 | }, 74 | "class_type": "LoadImage", 75 | "_meta": { 76 | "title": "Load Image" 77 | } 78 | }, 79 | "11": { 80 | "inputs": { 81 | "pixels": ["12", 0], 82 | "vae": ["4", 2] 83 | }, 84 | "class_type": "VAEEncode", 85 | "_meta": { 86 | "title": "VAE Encode" 87 | } 88 | }, 89 | "12": { 90 | "inputs": { 91 | "width": 512, 92 | "height": 512, 93 | "upscale_method": "nearest-exact", 94 | "keep_proportion": true, 95 | "divisible_by": 2, 96 | "image": ["10", 0] 97 | }, 98 | "class_type": "ImageResizeKJ", 99 | "_meta": { 100 | "title": "Resize Image" 101 | } 102 | } 103 | } 104 | -------------------------------------------------------------------------------- /examples/example-t2i.ts: -------------------------------------------------------------------------------- 1 | import { CallWrapper } from "../src/call-wrapper"; 2 | import { ComfyApi } from "../src/client"; 3 | import { PromptBuilder } from "../src/prompt-builder"; 4 | import { seed } from "../src/tools"; 5 | import { TSamplerName, TSchedulerName } from "../src/types/sampler"; 6 | import ExampleTxt2ImgWorkflow from "./example-txt2img-workflow.json"; 7 | 8 | /** 9 | * Define a T2I (text to image) workflow task 10 | */ 11 | export const Txt2ImgPrompt = new PromptBuilder( 12 | ExampleTxt2ImgWorkflow, 13 | ["positive", "negative", "checkpoint", "seed", "batch", "step", "cfg", "sampler", "sheduler", "width", "height"], 14 | ["images"] 15 | ) 16 | .setInputNode("checkpoint", "4.inputs.ckpt_name") 17 | .setInputNode("seed", "3.inputs.seed") 18 | .setInputNode("batch", "5.inputs.batch_size") 19 | .setInputNode("negative", "7.inputs.text") 20 | .setInputNode("positive", "6.inputs.text") 21 | .setInputNode("cfg", "3.inputs.cfg") 22 | .setInputNode("sampler", "3.inputs.sampler_name") 23 | .setInputNode("sheduler", "3.inputs.scheduler") 24 | .setInputNode("step", "3.inputs.steps") 25 | .setInputNode("width", "5.inputs.width") 26 | .setInputNode("height", "5.inputs.height") 27 | .setOutputNode("images", "9"); 28 | 29 | /** 30 | * Initialize the client 31 | */ 32 | const api = new ComfyApi("http://localhost:8189").init(); 33 | 34 | /** 35 | * Set the workflow's input values 36 | */ 37 | const workflow = Txt2ImgPrompt.input( 38 | "checkpoint", 39 | "SDXL/realvisxlV40_v40LightningBakedvae.safetensors", 40 | /** 41 | * Use the client's osType to encode the path 42 | */ 43 | api.osType 44 | ) 45 | .input("seed", seed()) 46 | .input("step", 6) 47 | .input("cfg", 1) 48 | .input("sampler", "dpmpp_2m_sde_gpu") 49 | .input("sheduler", "sgm_uniform") 50 | .input("width", 1024) 51 | .input("height", 1024) 52 | .input("batch", 1) 53 | .input("positive", "A picture of cute dog on the street"); 54 | 55 | /** 56 | * Execute the workflow 57 | */ 58 | new CallWrapper(api, workflow) 59 | .onPending(() => console.log("Task is pending")) 60 | .onStart(() => console.log("Task is started")) 61 | .onPreview((blob) => console.log(blob)) 62 | .onFinished((data) => { 63 | console.log(data.images?.images.map((img: any) => api.getPathImage(img))); 64 | }) 65 | .onProgress((info) => console.log("Processing node", info.node, `${info.value}/${info.max}`)) 66 | .onFailed((err) => console.log("Task is failed", err)) 67 | .run(); 68 | -------------------------------------------------------------------------------- /src/socket.ts: -------------------------------------------------------------------------------- 1 | // src/WebSocketClient.ts 2 | 3 | import WebSocketLib from "ws"; 4 | 5 | // Define WebSocketInterface to allow for custom implementation 6 | export interface WebSocketInterface { 7 | new (url: string, protocols?: string | string[]): WebSocket; 8 | new (url: string, options?: any): WebSocket; 9 | CONNECTING: number; 10 | OPEN: number; 11 | CLOSING: number; 12 | CLOSED: number; 13 | } 14 | 15 | // Default WebSocket implementation based on environment 16 | let DefaultWebSocketImpl: WebSocketInterface; 17 | 18 | if (typeof window !== "undefined" && window.WebSocket) { 19 | // In a browser environment 20 | DefaultWebSocketImpl = window.WebSocket; 21 | } else { 22 | // In a Node.js environment 23 | DefaultWebSocketImpl = WebSocketLib as any; 24 | } 25 | 26 | export interface WebSocketClientOptions { 27 | headers?: { [key: string]: string }; 28 | customWebSocketImpl?: WebSocketInterface | null; 29 | } 30 | 31 | export class WebSocketClient { 32 | private socket: WebSocket; 33 | private readonly webSocketImpl: WebSocketInterface; 34 | 35 | constructor(url: string, options: WebSocketClientOptions = {}) { 36 | const { headers, customWebSocketImpl } = options; 37 | 38 | // Use custom WebSocket implementation if provided, otherwise use default 39 | this.webSocketImpl = customWebSocketImpl || DefaultWebSocketImpl; 40 | 41 | try { 42 | if (typeof window !== "undefined" && window.WebSocket) { 43 | // Browser environment - WebSocket does not support custom headers 44 | this.socket = new this.webSocketImpl(url); 45 | } else { 46 | // Node.js environment - using ws package, which supports custom headers 47 | const WebSocketConstructor = this.webSocketImpl as any; 48 | this.socket = new WebSocketConstructor(url, { headers }); 49 | } 50 | } catch (error) { 51 | console.error("WebSocket initialization failed:", error); 52 | throw new Error(`WebSocket initialization failed: ${error instanceof Error ? error.message : String(error)}`); 53 | } 54 | 55 | return this; 56 | } 57 | 58 | get client() { 59 | return this.socket; 60 | } 61 | 62 | public send(message: string) { 63 | if (this.socket && this.socket.readyState === this.webSocketImpl.OPEN) { 64 | this.socket.send(message); 65 | } else { 66 | console.error("WebSocket is not open or available"); 67 | } 68 | } 69 | 70 | public close() { 71 | if (this.socket) { 72 | this.socket.close(); 73 | } 74 | } 75 | } 76 | -------------------------------------------------------------------------------- /examples/example-txt2img-workflow.json: -------------------------------------------------------------------------------- 1 | { 2 | "3": { 3 | "inputs": { 4 | "seed": 509648683700218, 5 | "steps": 8, 6 | "cfg": 2, 7 | "sampler_name": "dpmpp_sde", 8 | "scheduler": "sgm_uniform", 9 | "denoise": 1, 10 | "model": ["4", 0], 11 | "positive": ["6", 0], 12 | "negative": ["7", 0], 13 | "latent_image": ["5", 0] 14 | }, 15 | "class_type": "KSampler", 16 | "_meta": { 17 | "title": "KSampler" 18 | } 19 | }, 20 | "4": { 21 | "inputs": { 22 | "ckpt_name": "SDXL/dreamshaperXL_v2TurboDpmppSDE.safetensors" 23 | }, 24 | "class_type": "CheckpointLoaderSimple", 25 | "_meta": { 26 | "title": "Load Checkpoint" 27 | } 28 | }, 29 | "5": { 30 | "inputs": { 31 | "width": 512, 32 | "height": 512, 33 | "batch_size": 1 34 | }, 35 | "class_type": "EmptyLatentImage", 36 | "_meta": { 37 | "title": "Empty Latent Image" 38 | } 39 | }, 40 | "6": { 41 | "inputs": { 42 | "text": "beautiful scenery nature glass bottle landscape", 43 | "clip": ["4", 1] 44 | }, 45 | "class_type": "CLIPTextEncode", 46 | "_meta": { 47 | "title": "CLIP Text Encode (Prompt)" 48 | } 49 | }, 50 | "7": { 51 | "inputs": { 52 | "text": "text, watermark", 53 | "clip": ["4", 1] 54 | }, 55 | "class_type": "CLIPTextEncode", 56 | "_meta": { 57 | "title": "CLIP Text Encode (Prompt)" 58 | } 59 | }, 60 | "8": { 61 | "inputs": { 62 | "samples": ["3", 0], 63 | "vae": ["4", 2] 64 | }, 65 | "class_type": "VAEDecode", 66 | "_meta": { 67 | "title": "VAE Decode" 68 | } 69 | }, 70 | "9": { 71 | "inputs": { 72 | "filename_prefix": "ComfyUI", 73 | "images": ["8", 0] 74 | }, 75 | "class_type": "SaveImage", 76 | "_meta": { 77 | "title": "Save Image" 78 | } 79 | }, 80 | "10": { 81 | "inputs": { 82 | "model_name": "4x-ClearRealityV1.pth" 83 | }, 84 | "class_type": "UpscaleModelLoader", 85 | "_meta": { 86 | "title": "Load Upscale Model" 87 | } 88 | }, 89 | "11": { 90 | "inputs": { 91 | "upscale_model": ["10", 0], 92 | "image": ["8", 0] 93 | }, 94 | "class_type": "ImageUpscaleWithModel", 95 | "_meta": { 96 | "title": "Upscale Image (using Model)" 97 | } 98 | }, 99 | "12": { 100 | "inputs": { 101 | "filename_prefix": "ComfyUI", 102 | "images": ["11", 0] 103 | }, 104 | "class_type": "SaveImage", 105 | "_meta": { 106 | "title": "Save Image" 107 | } 108 | } 109 | } 110 | -------------------------------------------------------------------------------- /examples/example-txt2img-upscaled-workflow.json: -------------------------------------------------------------------------------- 1 | { 2 | "3": { 3 | "inputs": { 4 | "seed": 509648683700218, 5 | "steps": 8, 6 | "cfg": 2, 7 | "sampler_name": "dpmpp_sde", 8 | "scheduler": "sgm_uniform", 9 | "denoise": 1, 10 | "model": ["4", 0], 11 | "positive": ["6", 0], 12 | "negative": ["7", 0], 13 | "latent_image": ["5", 0] 14 | }, 15 | "class_type": "KSampler", 16 | "_meta": { 17 | "title": "KSampler" 18 | } 19 | }, 20 | "4": { 21 | "inputs": { 22 | "ckpt_name": "SDXL/dreamshaperXL_v2TurboDpmppSDE.safetensors" 23 | }, 24 | "class_type": "CheckpointLoaderSimple", 25 | "_meta": { 26 | "title": "Load Checkpoint" 27 | } 28 | }, 29 | "5": { 30 | "inputs": { 31 | "width": 512, 32 | "height": 512, 33 | "batch_size": 1 34 | }, 35 | "class_type": "EmptyLatentImage", 36 | "_meta": { 37 | "title": "Empty Latent Image" 38 | } 39 | }, 40 | "6": { 41 | "inputs": { 42 | "text": "beautiful scenery nature glass bottle landscape", 43 | "clip": ["4", 1] 44 | }, 45 | "class_type": "CLIPTextEncode", 46 | "_meta": { 47 | "title": "CLIP Text Encode (Prompt)" 48 | } 49 | }, 50 | "7": { 51 | "inputs": { 52 | "text": "text, watermark", 53 | "clip": ["4", 1] 54 | }, 55 | "class_type": "CLIPTextEncode", 56 | "_meta": { 57 | "title": "CLIP Text Encode (Prompt)" 58 | } 59 | }, 60 | "8": { 61 | "inputs": { 62 | "samples": ["3", 0], 63 | "vae": ["4", 2] 64 | }, 65 | "class_type": "VAEDecode", 66 | "_meta": { 67 | "title": "VAE Decode" 68 | } 69 | }, 70 | "9": { 71 | "inputs": { 72 | "filename_prefix": "ComfyUI", 73 | "images": ["8", 0] 74 | }, 75 | "class_type": "SaveImage", 76 | "_meta": { 77 | "title": "Save Image" 78 | } 79 | }, 80 | "10": { 81 | "inputs": { 82 | "model_name": "4x-ClearRealityV1.pth" 83 | }, 84 | "class_type": "UpscaleModelLoader", 85 | "_meta": { 86 | "title": "Load Upscale Model" 87 | } 88 | }, 89 | "11": { 90 | "inputs": { 91 | "upscale_model": ["10", 0], 92 | "image": ["8", 0] 93 | }, 94 | "class_type": "ImageUpscaleWithModel", 95 | "_meta": { 96 | "title": "Upscale Image (using Model)" 97 | } 98 | }, 99 | "12": { 100 | "inputs": { 101 | "filename_prefix": "ComfyUI", 102 | "images": ["11", 0] 103 | }, 104 | "class_type": "SaveImage", 105 | "_meta": { 106 | "title": "Save Image" 107 | } 108 | } 109 | } 110 | -------------------------------------------------------------------------------- /examples/example-t2i-upscaled.ts: -------------------------------------------------------------------------------- 1 | import { CallWrapper } from "../src/call-wrapper"; 2 | import { ComfyApi } from "../src/client"; 3 | import { PromptBuilder } from "../src/prompt-builder"; 4 | import { seed } from "../src/tools"; 5 | import { TSamplerName, TSchedulerName } from "../src/types/sampler"; 6 | import ExampleTxt2ImgWorkflow from "./example-txt2img-upscaled-workflow.json"; 7 | 8 | /** 9 | * Define a T2I (text to image) with upscale step using model task 10 | */ 11 | export const Txt2ImgPrompt = new PromptBuilder( 12 | ExampleTxt2ImgWorkflow, 13 | ["positive", "negative", "checkpoint", "seed", "batch", "step", "cfg", "sampler", "sheduler", "width", "height"], 14 | ["images", "upscaled"] 15 | ) 16 | .setInputNode("checkpoint", "4.inputs.ckpt_name") 17 | .setInputNode("seed", "3.inputs.seed") 18 | .setInputNode("batch", "5.inputs.batch_size") 19 | .setInputNode("negative", "7.inputs.text") 20 | .setInputNode("positive", "6.inputs.text") 21 | .setInputNode("cfg", "3.inputs.cfg") 22 | .setInputNode("sampler", "3.inputs.sampler_name") 23 | .setInputNode("sheduler", "3.inputs.scheduler") 24 | .setInputNode("step", "3.inputs.steps") 25 | .setInputNode("width", "5.inputs.width") 26 | .setInputNode("height", "5.inputs.height") 27 | .setOutputNode("images", "9") 28 | .setOutputNode("upscaled", "12"); 29 | 30 | /** 31 | * Initialize the client 32 | */ 33 | const api = new ComfyApi("http://192.168.15.37:8189").init(); 34 | 35 | /** 36 | * Set the workflow's input values 37 | */ 38 | const workflow = Txt2ImgPrompt.input( 39 | "checkpoint", 40 | "SDXL/realvisxlV40_v40LightningBakedvae.safetensors", 41 | /** 42 | * Use the client's osType to encode the path 43 | */ 44 | api.osType 45 | ) 46 | .input("seed", seed()) 47 | .input("step", 6) 48 | .input("cfg", 1) 49 | .input("sampler", "dpmpp_2m_sde_gpu") 50 | .input("sheduler", "sgm_uniform") 51 | .input("width", 1024) 52 | .input("height", 1024) 53 | .input("batch", 1) 54 | .input("positive", "A picture of cute dog on the street"); 55 | 56 | /** 57 | * Execute the workflow 58 | */ 59 | new CallWrapper(api, workflow) 60 | .onPending(() => console.log("Task is pending")) 61 | .onStart(() => console.log("Task is started")) 62 | .onPreview((blob) => console.log(blob)) 63 | /** 64 | * Preview output of executed node 65 | */ 66 | .onOutput((outputName, outputVal) => console.log(`Output ${outputName} with value`, outputVal)) 67 | .onFinished((data) => { 68 | console.log("Final output", { 69 | images: data.images?.images.map((img: any) => api.getPathImage(img)), 70 | upscaled: data.upscaled?.images.map((img: any) => api.getPathImage(img)) 71 | }); 72 | }) 73 | .onProgress((info) => console.log("Processing node", info.node, `${info.value}/${info.max}`)) 74 | .onFailed((err) => console.log("Task is failed", err)) 75 | .run(); 76 | -------------------------------------------------------------------------------- /test/builder.spec.ts: -------------------------------------------------------------------------------- 1 | import { PromptBuilder } from "src/prompt-builder"; 2 | import Prompt from "../examples/example-txt2img-workflow.json"; 3 | 4 | import { describe, beforeEach, it, expect } from "bun:test"; 5 | 6 | describe("PromptBuilder with complex input", () => { 7 | let promptBuilder: PromptBuilder<"size", "images", typeof Prompt>; 8 | 9 | beforeEach(() => { 10 | promptBuilder = new PromptBuilder(Prompt, ["size"], ["images"]); 11 | }); 12 | 13 | it("should set and append input nodes correctly", () => { 14 | promptBuilder.setInputNode("size", "5.inputs.width"); 15 | promptBuilder.appendInputNode("size", "5.inputs.height"); 16 | 17 | expect(promptBuilder.mapInputKeys["size"]).toEqual(["5.inputs.width", "5.inputs.height"]); 18 | }); 19 | 20 | it("should set output nodes correctly", () => { 21 | promptBuilder.setOutputNode("images", "9"); 22 | 23 | expect(promptBuilder.mapOutputKeys["images"]).toBe("9"); 24 | }); 25 | 26 | it("should update input values correctly", () => { 27 | promptBuilder.setInputNode("size", "5.inputs.width"); 28 | promptBuilder.appendInputNode("size", "5.inputs.height"); 29 | 30 | const newPromptBuilder = promptBuilder 31 | .input("size", 1600) // Self update 32 | .clone() 33 | .input("size", 1500); // New instance update 34 | 35 | expect(promptBuilder.prompt["5"].inputs.width).toBe(1600); 36 | expect(newPromptBuilder.prompt["5"].inputs.width).toBe(1500); 37 | expect(newPromptBuilder.prompt["5"].inputs.height).toBe(1500); 38 | }); 39 | 40 | it("should have correct initial values for complex input structure", () => { 41 | expect(promptBuilder.prompt["3"].inputs.seed).toBe(509648683700218); 42 | expect(promptBuilder.prompt["4"].inputs.ckpt_name).toBe("SDXL/dreamshaperXL_v2TurboDpmppSDE.safetensors"); 43 | expect(promptBuilder.prompt["6"].inputs.text).toBe("beautiful scenery nature glass bottle landscape"); 44 | expect(promptBuilder.prompt["7"].inputs.text).toBe("text, watermark"); 45 | expect(promptBuilder.prompt["8"].inputs.samples).toEqual(["3", 0]); 46 | expect(promptBuilder.prompt["9"].inputs.filename_prefix).toBe("ComfyUI"); 47 | expect(promptBuilder.prompt["10"].inputs.model_name).toBe("4x-ClearRealityV1.pth"); 48 | expect(promptBuilder.prompt["11"].inputs.upscale_model).toEqual(["10", 0]); 49 | expect(promptBuilder.prompt["12"].inputs.filename_prefix).toBe("ComfyUI"); 50 | }); 51 | 52 | it("should clone the prompt builder correctly", () => { 53 | const clonedBuilder = promptBuilder.clone(); 54 | 55 | expect(clonedBuilder).not.toBe(promptBuilder); 56 | expect(clonedBuilder.prompt).toEqual(promptBuilder.prompt); 57 | expect(clonedBuilder.mapInputKeys).toEqual(promptBuilder.mapInputKeys); 58 | expect(clonedBuilder.mapOutputKeys).toEqual(promptBuilder.mapOutputKeys); 59 | }); 60 | 61 | it("should get the workflow correctly", () => { 62 | const workflow = promptBuilder.workflow; 63 | expect(workflow).toEqual(Prompt); 64 | }); 65 | 66 | it("should return the same instance when calling caller", () => { 67 | const callerInstance = promptBuilder.caller; 68 | expect(callerInstance).toBe(promptBuilder); 69 | }); 70 | }); 71 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Based on https://raw.githubusercontent.com/github/gitignore/main/Node.gitignore 2 | 3 | # Logs 4 | 5 | logs 6 | _.log 7 | npm-debug.log_ 8 | yarn-debug.log* 9 | yarn-error.log* 10 | lerna-debug.log* 11 | .pnpm-debug.log* 12 | 13 | # Caches 14 | 15 | .cache 16 | 17 | # Diagnostic reports (https://nodejs.org/api/report.html) 18 | 19 | report.[0-9]_.[0-9]_.[0-9]_.[0-9]_.json 20 | 21 | # Runtime data 22 | 23 | pids 24 | _.pid 25 | _.seed 26 | *.pid.lock 27 | 28 | # Directory for instrumented libs generated by jscoverage/JSCover 29 | 30 | lib-cov 31 | 32 | # Coverage directory used by tools like istanbul 33 | 34 | coverage 35 | *.lcov 36 | 37 | # nyc test coverage 38 | 39 | .nyc_output 40 | 41 | # Grunt intermediate storage (https://gruntjs.com/creating-plugins#storing-task-files) 42 | 43 | .grunt 44 | 45 | # Bower dependency directory (https://bower.io/) 46 | 47 | bower_components 48 | 49 | # node-waf configuration 50 | 51 | .lock-wscript 52 | 53 | # Compiled binary addons (https://nodejs.org/api/addons.html) 54 | 55 | build/Release 56 | 57 | # Dependency directories 58 | 59 | node_modules/ 60 | jspm_packages/ 61 | 62 | # Snowpack dependency directory (https://snowpack.dev/) 63 | 64 | web_modules/ 65 | 66 | # TypeScript cache 67 | 68 | *.tsbuildinfo 69 | 70 | # Optional npm cache directory 71 | 72 | .npm 73 | 74 | # Optional eslint cache 75 | 76 | .eslintcache 77 | 78 | # Optional stylelint cache 79 | 80 | .stylelintcache 81 | 82 | # Microbundle cache 83 | 84 | .rpt2_cache/ 85 | .rts2_cache_cjs/ 86 | .rts2_cache_es/ 87 | .rts2_cache_umd/ 88 | 89 | # Optional REPL history 90 | 91 | .node_repl_history 92 | 93 | # Output of 'npm pack' 94 | 95 | *.tgz 96 | 97 | # Yarn Integrity file 98 | 99 | .yarn-integrity 100 | 101 | # dotenv environment variable files 102 | 103 | .env 104 | .env.development.local 105 | .env.test.local 106 | .env.production.local 107 | .env.local 108 | 109 | # parcel-bundler cache (https://parceljs.org/) 110 | 111 | .parcel-cache 112 | 113 | # Next.js build output 114 | 115 | .next 116 | out 117 | 118 | # Nuxt.js build / generate output 119 | 120 | .nuxt 121 | dist 122 | 123 | # Gatsby files 124 | 125 | # Comment in the public line in if your project uses Gatsby and not Next.js 126 | 127 | # https://nextjs.org/blog/next-9-1#public-directory-support 128 | 129 | # public 130 | 131 | # vuepress build output 132 | 133 | .vuepress/dist 134 | 135 | # vuepress v2.x temp and cache directory 136 | 137 | .temp 138 | 139 | # Docusaurus cache and generated files 140 | 141 | .docusaurus 142 | 143 | # Serverless directories 144 | 145 | .serverless/ 146 | 147 | # FuseBox cache 148 | 149 | .fusebox/ 150 | 151 | # DynamoDB Local files 152 | 153 | .dynamodb/ 154 | 155 | # TernJS port file 156 | 157 | .tern-port 158 | 159 | # Stores VSCode versions used for testing VSCode extensions 160 | 161 | .vscode-test 162 | 163 | # yarn v2 164 | 165 | .yarn/cache 166 | .yarn/unplugged 167 | .yarn/build-state.yml 168 | .yarn/install-state.gz 169 | .pnp.* 170 | 171 | # IntelliJ based IDEs 172 | .idea 173 | 174 | # Finder (MacOS) folder config 175 | .DS_Store 176 | 177 | build 178 | 179 | test.ts 180 | test.js 181 | temp/ -------------------------------------------------------------------------------- /examples/example-i2i.ts: -------------------------------------------------------------------------------- 1 | import { ComfyApi } from "../src/client"; 2 | import { CallWrapper } from "../src/call-wrapper"; 3 | import { PromptBuilder } from "../src/prompt-builder"; 4 | import ExampleImg2ImgWorkflow from "./example-img2img-workflow.json"; 5 | import { TSamplerName, TSchedulerName } from "../src/types/sampler"; 6 | import { seed } from "../src/tools"; 7 | 8 | /** 9 | * Define a I2I (image to image) workflow task 10 | */ 11 | export const Img2ImgPrompt = new PromptBuilder( 12 | ExampleImg2ImgWorkflow, 13 | [ 14 | "sourceImg", 15 | "difference", 16 | "positive", 17 | "negative", 18 | "checkpoint", 19 | "cfg", 20 | "sampler", 21 | "sheduler", 22 | "seed", 23 | "step", 24 | "width", 25 | "height" 26 | ], 27 | ["images"] 28 | ) 29 | .setInputNode("sourceImg", "10.inputs.image") 30 | .setInputNode("checkpoint", "4.inputs.ckpt_name") 31 | .setInputNode("difference", "3.inputs.denoise") 32 | .setInputNode("seed", "3.inputs.seed") 33 | .setInputNode("negative", "7.inputs.text") 34 | .setInputNode("positive", "6.inputs.text") 35 | .setInputNode("cfg", "3.inputs.cfg") 36 | .setInputNode("sampler", "3.inputs.sampler_name") 37 | .setInputNode("sheduler", "3.inputs.scheduler") 38 | .setInputNode("step", "3.inputs.steps") 39 | .setInputNode("width", "12.inputs.width") 40 | .setInputNode("height", "12.inputs.height") 41 | .setOutputNode("images", "9"); 42 | 43 | /** 44 | * Initialize the client 45 | */ 46 | const api = new ComfyApi("http://localhost:8189").init(); 47 | 48 | /** 49 | * Prepare the image to be uploaded 50 | */ 51 | const exampleTomImg = "https://www.redwolf.in/image/cache/catalog/stickers/tom-face-sticker-india-600x800.jpg"; 52 | const downloadImg = await fetch(exampleTomImg); 53 | const imgBlob = await downloadImg.blob(); 54 | 55 | /** 56 | * Upload the source image to ComfyUI server 57 | */ 58 | const uploadedImg = await api.uploadImage(imgBlob, "tom-face-sticker-india.jpg"); 59 | if (!uploadedImg) { 60 | throw new Error("Failed to upload image"); 61 | } else { 62 | console.log("Uploaded source image", uploadedImg.url); 63 | } 64 | 65 | /** 66 | * Set the workflow's input values 67 | */ 68 | const workflow = Img2ImgPrompt.input( 69 | "checkpoint", 70 | "SDXL/realvisxlV40_v40LightningBakedvae.safetensors", 71 | /** 72 | * Use the client's osType to encode the path 73 | */ 74 | api.osType 75 | ) 76 | .input("sourceImg", uploadedImg.info.filename) 77 | .input("seed", seed()) 78 | .input("difference", 0.6) 79 | .input("step", 4) 80 | .input("cfg", 1) 81 | .input("sampler", "dpmpp_2m_sde_gpu") 82 | .input("sheduler", "sgm_uniform") 83 | .input("width", 768) 84 | .input("height", 768) 85 | .input("positive", "A picture of cute cat") 86 | .input("negative", "text, nsfw, blurry, bad draw, embeddings:easynegative"); 87 | 88 | new CallWrapper(api, workflow) 89 | .onPending(() => console.log("Task is pending")) 90 | .onStart(() => console.log("Task is started")) 91 | .onPreview((blob) => console.log(blob)) 92 | .onFinished((data) => { 93 | console.log(data.images?.images.map((img: any) => api.getPathImage(img))); 94 | }) 95 | .onProgress((info) => console.log("Processing node", info.node, `${info.value}/${info.max}`)) 96 | .onFailed((err) => console.log("Task is failed", err)) 97 | .run(); 98 | -------------------------------------------------------------------------------- /src/types/manager.ts: -------------------------------------------------------------------------------- 1 | export type TDefaultUI = "none" | "history" | "queue"; 2 | export type TExtensionActive = "Enabled" | "Disabled"; 3 | export type TPreviewMethod = "auto" | "latent2rgb" | "taesd" | "none"; 4 | export enum EInstallationState { 5 | NOT_INSTALLED = "not-installed", 6 | INSTALLED = "installed" 7 | } 8 | 9 | enum EModelType { 10 | CHECKPOINT = "checkpoint", 11 | UNCLIP = "unclip", 12 | CLIP = "clip", 13 | VAE = "VAE", 14 | LORA = "lora", 15 | T2I_ADAPTER = "T2I-Adapter", 16 | T2I_STYLE = "T2I-Style", 17 | CONTROLNET = "controlnet", 18 | CLIP_VISION = "clip_vision", 19 | GLIGEN = "gligen", 20 | UPSCALE = "upscale", 21 | EMBEDDINGS = "embeddings", 22 | ETC = "etc" 23 | } 24 | 25 | export enum EInstallType { 26 | GIT_CLONE = "git-clone", 27 | COPY = "copy", 28 | CNR = "cnr", 29 | UNZIP = "unzip" 30 | } 31 | 32 | export enum EExtensionUpdateCheckResult { 33 | NO_UPDATE = 0, 34 | UPDATE_AVAILABLE = 1, 35 | FAILED = 2 36 | } 37 | export enum EUpdateResult { 38 | UNCHANGED = 0, 39 | SUCCESS = 1, 40 | FAILED = 2 41 | } 42 | export type TExtensionNodeItem = { 43 | url: string; 44 | /** 45 | * Included nodes 46 | */ 47 | nodeNames: string[]; 48 | title_aux: string; 49 | title?: string; 50 | author?: string; 51 | description?: string; 52 | nickname?: string; 53 | }; 54 | 55 | export interface IExtensionInfo { 56 | author: string; 57 | title: string; 58 | id: string; 59 | reference: string; 60 | repository: string; 61 | files: string[]; 62 | install_type: EInstallType; 63 | description: string; 64 | stars: number; 65 | last_update: string; 66 | trust: boolean; 67 | state: EInstallationState; 68 | /** 69 | * @deprecated Use `state` instead 70 | */ 71 | installed: boolean; 72 | version: string; 73 | updatable: boolean; 74 | } 75 | 76 | export interface IExtensionBaseRequest { 77 | /** 78 | * Custom Node name 79 | */ 80 | title?: string; 81 | /** 82 | * Install method 83 | */ 84 | install_type: EInstallType; 85 | /** 86 | * Files to download, clone or copy (can be git url, file url or file path) 87 | */ 88 | files: string[]; 89 | } 90 | 91 | export interface IInstallExtensionRequest extends IExtensionBaseRequest { 92 | /** 93 | * Destination path for copying files when install_type is "copy", default is custom_node folder 94 | */ 95 | js_path?: string; 96 | /** 97 | * Python packages to be installed 98 | */ 99 | pip?: string[]; 100 | } 101 | 102 | export interface IExtensionUninstallRequest extends IExtensionBaseRequest { 103 | /** 104 | * Install method 105 | */ 106 | install_type: EInstallType.GIT_CLONE | EInstallType.COPY; 107 | /** 108 | * Destination path for remove files when install_type is "copy", default is custom_node folder 109 | */ 110 | js_path?: string; 111 | } 112 | 113 | export interface IExtensionUpdateRequest extends IExtensionBaseRequest { 114 | /** 115 | * Install method 116 | */ 117 | install_type: EInstallType.GIT_CLONE; 118 | } 119 | 120 | export interface IExtensionActiveRequest extends IExtensionBaseRequest { 121 | /** 122 | * Install method 123 | */ 124 | install_type: EInstallType.GIT_CLONE | EInstallType.COPY; 125 | /** 126 | * Active status 127 | */ 128 | installed: TExtensionActive; 129 | /** 130 | * Destination path of extension when install_type is "copy". Default is custom_node folder 131 | */ 132 | js_path?: string; 133 | } 134 | 135 | export interface IModelInstallRequest { 136 | /** 137 | * Model name 138 | */ 139 | name?: string; 140 | /** 141 | * Place to save the model, set to `default` to use type instead 142 | */ 143 | save_path: string; 144 | /** 145 | * Type of model 146 | */ 147 | type: EModelType; 148 | /** 149 | * Model filename 150 | */ 151 | filename: string; 152 | /** 153 | * Model url to be downloaded 154 | */ 155 | url: string; 156 | } 157 | 158 | export interface INodeMapItem { 159 | url: string; 160 | nodeNames: Array; 161 | title_aux: string; 162 | title?: string; 163 | author?: string; 164 | nickname?: string; 165 | description?: string; 166 | } 167 | -------------------------------------------------------------------------------- /src/types/event.ts: -------------------------------------------------------------------------------- 1 | import { EQueueMode } from "../pool"; 2 | import { ComfyApi } from "../client"; 3 | import { TMonitorEvent } from "../features/monitoring"; 4 | 5 | export type TEventStatus = { 6 | status: { 7 | exec_info: { 8 | queue_remaining: number; 9 | }; 10 | }; 11 | sid: string; 12 | }; 13 | 14 | export type TExecution = { 15 | prompt_id: string; 16 | }; 17 | 18 | export type TExecuting = TExecution & { 19 | node: string | null; 20 | }; 21 | 22 | export type TProgress = TExecuting & { 23 | value: number; 24 | max: number; 25 | }; 26 | 27 | export type TExecuted = TExecution & { 28 | node: string; 29 | output: T; 30 | }; 31 | 32 | export type TExecutionCached = TExecution & { 33 | nodes: string[]; 34 | }; 35 | 36 | export type TExecutionError = TExecution & { 37 | node_id: string; 38 | node_type: string; 39 | exception_message: string; 40 | exception_type: string; 41 | traceback: string[]; 42 | }; 43 | 44 | export type TExecutionInterrupted = TExecution & { 45 | node_id: string; 46 | node_type: string; 47 | executed: string[]; 48 | }; 49 | 50 | export type TEventKey = 51 | | "all" 52 | | "auth_error" 53 | | "connection_error" 54 | | "auth_success" 55 | | "status" 56 | | "progress" 57 | | "executing" 58 | | "executed" 59 | | "disconnected" 60 | | "execution_success" 61 | | "execution_start" 62 | | "execution_error" 63 | | "execution_cached" 64 | | "queue_error" 65 | | "reconnected" 66 | | "connected" 67 | | "log" 68 | | "terminal" 69 | | "reconnecting" 70 | | "b_preview"; 71 | 72 | export type TComfyAPIEventMap = { 73 | all: CustomEvent<{ type: string; data: any }>; 74 | auth_error: CustomEvent; 75 | auth_success: CustomEvent; 76 | connection_error: CustomEvent; 77 | execution_success: CustomEvent; 78 | status: CustomEvent; 79 | disconnected: CustomEvent; 80 | reconnecting: CustomEvent; 81 | connected: CustomEvent; 82 | reconnected: CustomEvent; 83 | b_preview: CustomEvent; 84 | log: CustomEvent<{ msg: string; data: any }>; 85 | terminal: CustomEvent<{ m: string; t: string }>; 86 | execution_start: CustomEvent; 87 | executing: CustomEvent; 88 | progress: CustomEvent; 89 | executed: CustomEvent; 90 | queue_error: CustomEvent; 91 | execution_error: CustomEvent; 92 | execution_interrupted: CustomEvent; 93 | execution_cached: CustomEvent; 94 | }; 95 | 96 | export type TComfyPoolEventKey = 97 | | "init" 98 | | "init_client" 99 | | "auth_error" 100 | | "connection_error" 101 | | "auth_success" 102 | | "added" 103 | | "removed" 104 | | "add_job" 105 | | "have_job" 106 | | "idle" 107 | | "terminal" 108 | | "ready" 109 | | "change_mode" 110 | | "connected" 111 | | "disconnected" 112 | | "reconnected" 113 | | "executing" 114 | | "executed" 115 | | "execution_interrupted" 116 | | "execution_error" 117 | | "system_monitor"; 118 | 119 | export type TComfyPoolEventMap = { 120 | init: CustomEvent; 121 | auth_error: CustomEvent<{ 122 | client: ComfyApi; 123 | clientIdx: number; 124 | res: Response; 125 | }>; 126 | connection_error: CustomEvent<{ 127 | client: ComfyApi; 128 | clientIdx: number; 129 | error: Error; 130 | }>; 131 | terminal: CustomEvent<{ clientIdx: number; m: string; t: string }>; 132 | ready: CustomEvent<{ client: ComfyApi; clientIdx: number }>; 133 | auth_success: CustomEvent<{ client: ComfyApi; clientIdx: number }>; 134 | loading_client: CustomEvent<{ client: ComfyApi; clientIdx: number }>; 135 | change_mode: CustomEvent<{ mode: EQueueMode }>; 136 | added: CustomEvent<{ client: ComfyApi; clientIdx: number }>; 137 | removed: CustomEvent<{ client: ComfyApi; clientIdx: number }>; 138 | connected: CustomEvent<{ client: ComfyApi; clientIdx: number }>; 139 | disconnected: CustomEvent<{ client: ComfyApi; clientIdx: number }>; 140 | reconnected: CustomEvent<{ client: ComfyApi; clientIdx: number }>; 141 | add_job: CustomEvent<{ jobIdx: number; weight: number }>; 142 | have_job: CustomEvent<{ client: ComfyApi; remain: number }>; 143 | idle: CustomEvent<{ client: ComfyApi }>; 144 | execution_interrupted: CustomEvent<{ client: ComfyApi; clientIdx: number }>; 145 | executing: CustomEvent<{ client: ComfyApi; clientIdx: number }>; 146 | executed: CustomEvent<{ client: ComfyApi; clientIdx: number }>; 147 | execution_error: CustomEvent<{ 148 | client: ComfyApi; 149 | clientIdx: number; 150 | error: Error; 151 | }>; 152 | system_monitor: CustomEvent<{ 153 | client: ComfyApi; 154 | clientIdx: number; 155 | data: TMonitorEvent; 156 | }>; 157 | }; 158 | -------------------------------------------------------------------------------- /examples/example-pool.ts: -------------------------------------------------------------------------------- 1 | import { ComfyApi } from "../src/client"; 2 | import { CallWrapper } from "../src/call-wrapper"; 3 | import { ComfyPool, EQueueMode } from "../src/pool"; 4 | import { PromptBuilder } from "../src/prompt-builder"; 5 | import ExampleTxt2ImgWorkflow from "./example-txt2img-workflow.json"; 6 | import { TSamplerName, TSchedulerName } from "../src/types/sampler"; 7 | import { seed } from "../src/tools"; 8 | 9 | /** 10 | * Define a T2I (text to image) workflow task 11 | */ 12 | export const Txt2ImgPrompt = new PromptBuilder( 13 | ExampleTxt2ImgWorkflow, // Get from `Save (API Format)` button in ComfyUI's website 14 | [ 15 | "positive", 16 | "negative", 17 | "checkpoint", 18 | "seed", 19 | "batch", 20 | "step", 21 | "cfg", 22 | "sampler", 23 | "sheduler", 24 | "width", 25 | "height", 26 | "sampler", 27 | "scheduler" 28 | ], 29 | ["images"] 30 | ) 31 | .setInputNode("checkpoint", "4.inputs.ckpt_name") 32 | .setInputNode("seed", "3.inputs.seed") 33 | .setInputNode("batch", "5.inputs.batch_size") 34 | .setInputNode("negative", "7.inputs.text") 35 | .setInputNode("positive", "6.inputs.text") 36 | .setInputNode("step", "3.inputs.steps") 37 | .setInputNode("width", "5.inputs.width") 38 | .setInputNode("height", "5.inputs.height") 39 | .setInputNode("cfg", "3.inputs.cfg") 40 | .setInputNode("sampler", "3.inputs.sampler_name") 41 | .setInputNode("scheduler", "3.inputs.scheduler") 42 | .setOutputNode("images", "9"); 43 | 44 | /** 45 | * Define a pool of ComfyApi 46 | */ 47 | const ApiPool = new ComfyPool( 48 | [ 49 | new ComfyApi("http://localhost:8188"), // Comfy Instance 1 50 | new ComfyApi("http://localhost:8189") // Comfy Instance 2 51 | ], 52 | EQueueMode.PICK_ZERO 53 | ) 54 | .on("init", () => console.log("Pool in initializing")) 55 | .on("loading_client", (ev) => console.log("Loading client", ev.detail.clientIdx)) 56 | .on("add_job", (ev) => console.log("Job added at index", ev.detail.jobIdx, "weight:", ev.detail.weight)) 57 | .on("added", (ev) => console.log("Client added", ev.detail.clientIdx)); 58 | 59 | /** 60 | * Define the generator function for all nodes 61 | */ 62 | const generateFn = async (api: ComfyApi, clientIdx?: number) => { 63 | return new Promise((resolve) => { 64 | /** 65 | * Set the workflow's input values 66 | */ 67 | const workflow = Txt2ImgPrompt.input( 68 | "checkpoint", 69 | "SDXL/realvisxlV40_v40LightningBakedvae.safetensors", 70 | /** 71 | * Use the client's osType to encode the path 72 | */ 73 | api.osType 74 | ) 75 | .input("seed", seed()) 76 | .input("step", 6) 77 | .input("width", 512) 78 | .input("height", 512) 79 | .input("batch", 2) 80 | .input("cfg", 1) 81 | .input("sampler", "dpmpp_2m_sde_gpu") 82 | .input("scheduler", "sgm_uniform") 83 | .input("positive", "A close up picture of cute Cat") 84 | .input("negative", "text, blurry, bad picture, nsfw"); 85 | 86 | new CallWrapper(api, workflow) 87 | .onPending((promptId) => 88 | console.log(`[${clientIdx}]`, `#${promptId}`, "Task is pending", { 89 | clientId: api.id, 90 | clientOs: api.osType 91 | }) 92 | ) 93 | .onStart((promptId) => console.log(`[${clientIdx}]`, `#${promptId}`, "Task is started")) 94 | .onPreview((blob, promptId) => console.log(`[${clientIdx}]`, `#${promptId}`, "Blob size", blob.size)) 95 | .onFinished((data, promptId) => { 96 | console.log(`[${clientIdx}]`, `#${promptId}`, "Task is finished"); 97 | const url = data.images?.images.map((img: any) => api.getPathImage(img)); 98 | resolve(url as string[]); 99 | }) 100 | .onProgress((info, promptId) => { 101 | console.log(`[${clientIdx}]`, `#${promptId}`, "Processing node", info.node, `${info.value}/${info.max}`); 102 | }) 103 | .onFailed((err, promptId) => { 104 | console.log(`[${clientIdx}]`, `#${promptId}`, "Task is failed", err); 105 | resolve([]); 106 | }) 107 | .run(); 108 | }); 109 | }; 110 | /** 111 | * Single shoot 112 | */ 113 | // const output = ApiPool.run(generateFn); 114 | 115 | /** 116 | * Multiple shoot using batch 117 | */ 118 | const jobA = ApiPool.batch( 119 | Array(5) 120 | .fill("") 121 | .map(() => generateFn), 122 | 10 // weight = 10, more weight = lower priority 123 | ).then((res) => { 124 | console.log("Batch A done"); 125 | return res.flat(); 126 | }); 127 | 128 | const jobB = ApiPool.batch( 129 | Array(5) 130 | .fill("") 131 | .map(() => generateFn), 132 | 0 // weight = 0, should be executed first 133 | ).then((res) => { 134 | console.log("Batch B done"); 135 | return res.flat(); 136 | }); 137 | 138 | console.log(await Promise.all([jobA, jobB]).then((res) => res.flat())); 139 | -------------------------------------------------------------------------------- /examples/example-pool-basic-auth.ts: -------------------------------------------------------------------------------- 1 | import { ComfyApi } from "../src/client"; 2 | import { CallWrapper } from "../src/call-wrapper"; 3 | import { ComfyPool, EQueueMode } from "../src/pool"; 4 | import { PromptBuilder } from "../src/prompt-builder"; 5 | import ExampleTxt2ImgWorkflow from "./example-txt2img-workflow.json"; 6 | import { seed } from "../src/tools"; 7 | import { BasicCredentials } from "../src/types/api"; 8 | import { TSamplerName, TSchedulerName } from "../src/types/sampler"; 9 | 10 | /** 11 | * Using with NginX basic auth 12 | */ 13 | const credentials: BasicCredentials = { 14 | type: "basic", 15 | username: "username", 16 | password: "password" 17 | }; 18 | 19 | /** 20 | * Define a T2I (text to image) workflow task 21 | */ 22 | export const Txt2ImgPrompt = new PromptBuilder( 23 | ExampleTxt2ImgWorkflow, // Get from `Save (API Format)` button in ComfyUI's website 24 | [ 25 | "positive", 26 | "negative", 27 | "checkpoint", 28 | "seed", 29 | "batch", 30 | "step", 31 | "cfg", 32 | "sampler", 33 | "sheduler", 34 | "width", 35 | "height", 36 | "sampler", 37 | "scheduler" 38 | ], 39 | ["images"] 40 | ) 41 | .setInputNode("checkpoint", "4.inputs.ckpt_name") 42 | .setInputNode("seed", "3.inputs.seed") 43 | .setInputNode("batch", "5.inputs.batch_size") 44 | .setInputNode("negative", "7.inputs.text") 45 | .setInputNode("positive", "6.inputs.text") 46 | .setInputNode("step", "3.inputs.steps") 47 | .setInputNode("width", "5.inputs.width") 48 | .setInputNode("height", "5.inputs.height") 49 | .setInputNode("cfg", "3.inputs.cfg") 50 | .setInputNode("sampler", "3.inputs.sampler_name") 51 | .setInputNode("scheduler", "3.inputs.scheduler") 52 | .setOutputNode("images", "9"); 53 | 54 | /** 55 | * Define a pool of ComfyApi 56 | */ 57 | const ApiPool = new ComfyPool( 58 | [ 59 | new ComfyApi("http://localhost:8188", "node-1", { 60 | credentials 61 | }), // Comfy Instance 1 62 | new ComfyApi("http://localhost:8189", "node-2", { 63 | credentials 64 | }) // Comfy Instance 2 65 | ], 66 | // "PICK_ZERO", Picks the client which has zero queue remaining. This is the default mode. (For who using along with ComfyUI web interface) 67 | // "PICK_LOWEST", Picks the client which has the lowest queue remaining. 68 | // "PICK_ROUTINE", Picks the client in a round-robin manner. 69 | EQueueMode.PICK_ZERO 70 | ) 71 | .on("init", () => console.log("Pool in initializing")) 72 | .on("loading_client", (ev) => console.log("Loading client", ev.detail.clientIdx)) 73 | .on("add_job", (ev) => console.log("Job added", ev.detail.jobIdx)) 74 | .on("added", (ev) => console.log("Client added", ev.detail.clientIdx)) 75 | .on("auth_success", (ev) => { 76 | console.info(`Client ${ev.detail.clientIdx} auth successfuly`); 77 | }) 78 | .on("auth_error", (ev) => { 79 | console.error(`Client ${ev.detail.clientIdx} auth failed`); 80 | }); 81 | 82 | /** 83 | * Define the generator function for all nodes 84 | */ 85 | const generateFn = async (api: ComfyApi, clientIdx?: number) => { 86 | return new Promise((resolve) => { 87 | /** 88 | * Set the workflow's input values 89 | */ 90 | const workflow = Txt2ImgPrompt.input( 91 | "checkpoint", 92 | "SDXL/realvisxlV40_v40LightningBakedvae.safetensors", 93 | /** 94 | * Use the client's osType to encode the path 95 | */ 96 | api.osType 97 | ) 98 | .input("seed", seed()) 99 | .input("step", 6) 100 | .input("width", 512) 101 | .input("height", 512) 102 | .input("batch", 2) 103 | .input("cfg", 1) 104 | .input("sampler", "dpmpp_2m_sde_gpu") 105 | .input("scheduler", "sgm_uniform") 106 | .input("positive", "A close up picture of cute Cat") 107 | .input("negative", "text, blurry, bad picture, nsfw"); 108 | 109 | new CallWrapper(api, workflow) 110 | .onPending((promptId) => 111 | console.log(`[${clientIdx}]`, `#${promptId}`, "Task is pending", { 112 | clientId: api.id, 113 | clientOs: api.osType 114 | }) 115 | ) 116 | .onStart((promptId) => console.log(`[${clientIdx}]`, `#${promptId}`, "Task is started")) 117 | .onPreview((blob, promptId) => console.log(`[${clientIdx}]`, `#${promptId}`, "Blob size", blob.size)) 118 | .onFinished(async (data, promptId) => { 119 | console.log(`[${clientIdx}]`, `#${promptId}`, "Task is finished"); 120 | /** 121 | * Use getImage instead of getURL because we use basic auth 122 | */ 123 | const url = await Promise.all(data.images?.images.map((img: any) => api.getImage(img))); 124 | resolve(url as string[]); 125 | }) 126 | .onProgress((info, promptId) => { 127 | console.log(`[${clientIdx}]`, `#${promptId}`, "Processing node", info.node, `${info.value}/${info.max}`); 128 | }) 129 | .onFailed((err, promptId) => { 130 | console.log(`[${clientIdx}]`, `#${promptId}`, "Task is failed", err); 131 | resolve([]); 132 | }) 133 | .run(); 134 | }); 135 | }; 136 | /** 137 | * Single shoot 138 | */ 139 | // const output = ApiPool.run(generateFn); 140 | 141 | /** 142 | * Multiple shoot using batch 143 | */ 144 | const output = await ApiPool.batch( 145 | Array(10) 146 | .fill("") 147 | .map(() => generateFn) 148 | ); 149 | 150 | console.log(output.flat()); 151 | -------------------------------------------------------------------------------- /src/types/api.ts: -------------------------------------------------------------------------------- 1 | export enum OSType { 2 | /** 3 | * Unix-like operating systems 4 | */ 5 | POSIX = "posix", 6 | /** 7 | * Windows operating systems 8 | */ 9 | NT = "nt", 10 | /** 11 | * Java virtual machine 12 | */ 13 | JAVA = "java" 14 | } 15 | 16 | export interface BasicCredentials { 17 | type: "basic"; 18 | username: string; 19 | password: string; 20 | } 21 | 22 | export interface BearerTokenCredentials { 23 | type: "bearer_token"; 24 | token: string; 25 | } 26 | 27 | export interface CustomCredentials { 28 | type: "custom"; 29 | headers: Record; 30 | } 31 | 32 | export interface HistoryResponse { 33 | [key: string]: HistoryEntry; 34 | } 35 | 36 | export interface HistoryEntry { 37 | prompt: PromptData; 38 | outputs: OutputData; 39 | status: StatusData; 40 | } 41 | 42 | export interface PromptData { 43 | [index: number]: number | string | NodeData | MetadataData; 44 | } 45 | 46 | export interface NodeData { 47 | [key: string]: { 48 | inputs: { [key: string]: any }; 49 | class_type: string; 50 | _meta: { title: string }; 51 | }; 52 | } 53 | 54 | export interface MetadataData { 55 | [key: string]: any; 56 | } 57 | 58 | export interface ImageInfo { 59 | filename: string; 60 | subfolder: string; 61 | type: string; 62 | } 63 | 64 | export interface OutputData { 65 | [key: string]: { 66 | width?: number[]; 67 | height?: number[]; 68 | ratio?: number[]; 69 | images?: ImageInfo[]; 70 | }; 71 | } 72 | 73 | export interface StatusData { 74 | status_str: string; 75 | completed: boolean; 76 | messages: [string, { [key: string]: any }][]; 77 | } 78 | 79 | export interface QueueResponse { 80 | queue_running: QueueItem[]; 81 | queue_pending: QueueItem[]; 82 | } 83 | 84 | export interface QueueItem { 85 | [index: number]: number | string | NodeData | MetadataData; 86 | } 87 | 88 | export interface QueuePromptResponse { 89 | prompt_id: string; 90 | number: number; 91 | node_errors: { [key: string]: any }; 92 | } 93 | 94 | export interface SystemStatsResponse { 95 | system: { 96 | os: OSType; 97 | python_version: string; 98 | embedded_python: boolean; 99 | }; 100 | devices: DeviceStats[]; 101 | } 102 | 103 | export interface DeviceStats { 104 | name: string; 105 | type: string; 106 | index: number; 107 | vram_total: number; 108 | vram_free: number; 109 | torch_vram_total: number; 110 | torch_vram_free: number; 111 | } 112 | 113 | export interface QueueStatus { 114 | exec_info: { queue_remaining: number }; 115 | } 116 | 117 | export interface NodeDefsResponse { 118 | [key: string]: NodeDef; 119 | } 120 | 121 | export interface NodeDef { 122 | input: { 123 | required: { 124 | [key: string]: 125 | | [string[], { tooltip?: string }] 126 | | [string, { tooltip?: string }] 127 | | TStringInput 128 | | TBoolInput 129 | | TNumberInput; 130 | }; 131 | optional?: { 132 | [key: string]: 133 | | [string[], { tooltip?: string }] 134 | | [string, { tooltip?: string }] 135 | | TStringInput 136 | | TBoolInput 137 | | TNumberInput; 138 | }; 139 | hidden: { 140 | [key: string]: string; 141 | }; 142 | }; 143 | input_order: { 144 | required: string[]; 145 | optional?: string[]; 146 | hidden: string[]; 147 | }; 148 | output: string[]; 149 | output_is_list: boolean[]; 150 | output_name: string[]; 151 | name: string; 152 | display_name: string; 153 | description: string; 154 | category: string; 155 | python_module: string; 156 | output_node: boolean; 157 | output_tooltips: string[]; 158 | } 159 | 160 | export interface NodeProgress { 161 | value: number; 162 | max: number; 163 | prompt_id: string; 164 | node: string; 165 | } 166 | 167 | export interface IInputNumberConfig { 168 | default: number; 169 | min: number; 170 | max: number; 171 | step?: number; 172 | round?: number; 173 | tooltip?: string; 174 | } 175 | export interface IInputStringConfig { 176 | default?: string; 177 | multiline?: boolean; 178 | dynamicPrompts?: boolean; 179 | tooltip?: string; 180 | } 181 | 182 | export type TStringInput = ["STRING", IInputStringConfig]; 183 | export type TBoolInput = ["BOOLEAN", { default: boolean; tooltip?: string }]; 184 | export type TNumberInput = ["INT" | "FLOAT", IInputNumberConfig]; 185 | 186 | /** 187 | * Represents a model folder in the ComfyUI system 188 | * @experimental API that may change in future versions 189 | */ 190 | export interface ModelFolder { 191 | name: string; 192 | folders: string[]; 193 | } 194 | 195 | /** 196 | * Represents a model file in the ComfyUI system 197 | * @experimental API that may change in future versions 198 | */ 199 | export interface ModelFile { 200 | name: string; 201 | pathIndex: number; 202 | } 203 | 204 | /** 205 | * Response format for model preview images 206 | * @experimental API that may change in future versions 207 | */ 208 | export interface ModelPreviewResponse { 209 | body: ArrayBuffer; 210 | contentType: string; 211 | } 212 | 213 | /** 214 | * Response format for a list of model files 215 | * @experimental API that may change in future versions 216 | */ 217 | export interface ModelFileListResponse { 218 | files: ModelFile[]; 219 | } 220 | 221 | /** 222 | * Response format for a list of model folders 223 | * @experimental API that may change in future versions 224 | */ 225 | export interface ModelFoldersResponse { 226 | folders: ModelFolder[]; 227 | } 228 | -------------------------------------------------------------------------------- /src/features/monitoring.ts: -------------------------------------------------------------------------------- 1 | import { AbstractFeature } from "./abstract"; 2 | import { FetchOptions } from "./manager"; 3 | 4 | const SYSTEM_MONITOR_EXTENSION = encodeURIComponent("Primitive boolean [Crystools]"); 5 | 6 | export type TMonitorEvent = { 7 | cpu_utilization: number; 8 | ram_total: number; 9 | ram_used: number; 10 | ram_used_percent: number; 11 | hdd_total: number; 12 | hdd_used: number; 13 | hdd_used_percent: number; 14 | device_type: "cuda"; 15 | gpus: Array<{ 16 | gpu_utilization: number; 17 | gpu_temperature: number; 18 | vram_total: number; 19 | vram_used: number; 20 | vram_used_percent: number; 21 | }>; 22 | }; 23 | 24 | export type TMonitorEventMap = { 25 | system_monitor: CustomEvent; 26 | }; 27 | 28 | export class MonitoringFeature extends AbstractFeature { 29 | private resources?: TMonitorEvent; 30 | private listeners: { 31 | event: keyof TMonitorEventMap; 32 | options?: AddEventListenerOptions | boolean; 33 | handler: (event: TMonitorEventMap[keyof TMonitorEventMap]) => void; 34 | }[] = []; 35 | private binded = false; 36 | 37 | async checkSupported() { 38 | const data = await this.client.getNodeDefs(SYSTEM_MONITOR_EXTENSION); 39 | if (data) { 40 | this.supported = true; 41 | this.bind(); 42 | } 43 | return this.supported; 44 | } 45 | 46 | public destroy(): void { 47 | this.listeners.forEach((listener) => { 48 | this.off(listener.event, listener.handler, listener.options); 49 | }); 50 | this.listeners = []; 51 | } 52 | 53 | private async fetchApi(path: string, options?: FetchOptions) { 54 | if (!this.supported) { 55 | return false; 56 | } 57 | return this.client.fetchApi(path, options); 58 | } 59 | 60 | public on( 61 | type: K, 62 | callback: (event: TMonitorEventMap[K]) => void, 63 | options?: AddEventListenerOptions | boolean 64 | ) { 65 | this.addEventListener(type, callback as any, options); 66 | this.listeners.push({ event: type, options, handler: callback }); 67 | return () => this.off(type, callback); 68 | } 69 | 70 | public off( 71 | type: K, 72 | callback: (event: TMonitorEventMap[K]) => void, 73 | options?: EventListenerOptions | boolean 74 | ): void { 75 | this.removeEventListener(type, callback as any, options); 76 | this.listeners = this.listeners.filter((listener) => listener.event !== type && listener.handler !== callback); 77 | } 78 | 79 | /** 80 | * Gets the monitor data. 81 | * 82 | * @returns The monitor data if supported, otherwise false. 83 | */ 84 | get monitorData() { 85 | if (!this.supported) { 86 | return false; 87 | } 88 | return this.resources; 89 | } 90 | 91 | /** 92 | * Sets the monitor configuration. 93 | */ 94 | async setConfig( 95 | config?: Partial<{ 96 | /** 97 | * Refresh per second (Default 0.5) 98 | */ 99 | rate: number; 100 | /** 101 | * Switch to enable/disable CPU monitoring 102 | */ 103 | switchCPU: boolean; 104 | /** 105 | * Switch to enable/disable GPU monitoring 106 | */ 107 | switchHDD: boolean; 108 | /** 109 | * Switch to enable/disable RAM monitoring 110 | */ 111 | switchRAM: boolean; 112 | /** 113 | * Path of HDD to monitor HDD usage (use getHddList to get the pick-able list) 114 | */ 115 | whichHDD: string; 116 | }> 117 | ) { 118 | if (!this.supported) { 119 | return false; 120 | } 121 | return this.fetchApi(`/api/crystools/monitor`, { 122 | method: "PATCH", 123 | body: JSON.stringify(config) 124 | }); 125 | } 126 | 127 | /** 128 | * Switches the monitor on or off. 129 | */ 130 | async switch(active: boolean) { 131 | if (!this.supported) { 132 | return false; 133 | } 134 | return this.fetchApi(`/api/crystools/monitor/switch`, { 135 | method: "POST", 136 | body: JSON.stringify({ monitor: active }) 137 | }); 138 | } 139 | 140 | /** 141 | * Gets the list of HDDs. 142 | */ 143 | async getHddList(): Promise> { 144 | if (!this.supported) { 145 | return null; 146 | } 147 | const data = await this.fetchApi(`/api/crystools/monitor/HDD`); 148 | if (data) { 149 | return data.json(); 150 | } 151 | return null; 152 | } 153 | 154 | /** 155 | * Gets the list of GPUs. 156 | */ 157 | async getGpuList(): Promise> { 158 | if (!this.supported) { 159 | return null; 160 | } 161 | const data = await this.fetchApi(`/api/crystools/monitor/GPU`); 162 | if (data) { 163 | return data.json(); 164 | } 165 | return null; 166 | } 167 | 168 | /** 169 | * Config gpu monitoring 170 | * @param index Index of the GPU 171 | * @param config Configuration of monitoring, set to `true` to enable monitoring 172 | */ 173 | async setGpuConfig(index: number, config: Partial<{ utilization: boolean; vram: boolean; temperature: boolean }>) { 174 | if (!this.supported) { 175 | return false; 176 | } 177 | return this.fetchApi(`/api/crystools/monitor/GPU/${index}`, { 178 | method: "PATCH", 179 | body: JSON.stringify(config) 180 | }); 181 | } 182 | 183 | private bind() { 184 | if (this.binded) { 185 | return; 186 | } else { 187 | this.binded = true; 188 | } 189 | this.client.on("all", (ev) => { 190 | const msg = ev.detail; 191 | if (msg.type === "crystools.monitor") { 192 | this.resources = msg.data; 193 | this.dispatchEvent(new CustomEvent("system_monitor", { detail: msg.data })); 194 | } 195 | }); 196 | } 197 | } 198 | -------------------------------------------------------------------------------- /test/pool.spec.ts: -------------------------------------------------------------------------------- 1 | import { ComfyPool, EQueueMode } from "../src/pool"; 2 | import { ComfyApi } from "../src/client"; 3 | import { test, describe, beforeEach, jest, expect } from "bun:test"; 4 | 5 | describe("ComfyPool", () => { 6 | let mockClient: ComfyApi; 7 | let comfyPool: ComfyPool; 8 | 9 | beforeEach(async () => { 10 | mockClient = { 11 | id: "client1", 12 | on: jest.fn((event, handler) => { 13 | if (event === "connected") { 14 | setTimeout(() => handler(), 0); // Simulate the connected event being dispatched 15 | } 16 | if (event === "status") { 17 | setTimeout( 18 | () => 19 | handler({ 20 | detail: { 21 | status: { exec_info: { queue_remaining: 0 } }, 22 | sid: "sid" 23 | } 24 | }), 25 | 0 26 | ); // Simulate the status event being dispatched 27 | } 28 | }), 29 | init: jest.fn().mockReturnValue({ 30 | waitForReady: jest.fn().mockResolvedValue(undefined) 31 | }), 32 | ext: { 33 | monitor: { 34 | isSupported: false, 35 | on: jest.fn() 36 | } 37 | }, 38 | destroy: () => { 39 | console.log("destroyed"); 40 | } 41 | } as unknown as ComfyApi; 42 | 43 | comfyPool = new ComfyPool([mockClient]); 44 | await new Promise((resolve) => setTimeout(resolve, 10)); 45 | }); 46 | 47 | test("should initialize clients and dispatch init event", async () => { 48 | const initListener = jest.fn(); 49 | 50 | comfyPool = new ComfyPool([mockClient]); 51 | comfyPool.addEventListener("init", initListener); 52 | await new Promise((resolve) => setTimeout(resolve, 10)); 53 | 54 | expect(initListener).toHaveBeenCalled(); 55 | expect(mockClient.on).toHaveBeenCalled(); 56 | expect(mockClient.init().waitForReady).toHaveBeenCalled(); 57 | }); 58 | 59 | test("should add a client to the pool", async () => { 60 | const newClient = { ...mockClient, id: "client2" } as ComfyApi; 61 | const addedListener = jest.fn(); 62 | comfyPool.addEventListener("added", addedListener); 63 | 64 | await comfyPool.addClient(newClient); 65 | 66 | expect(comfyPool.clients).toContain(newClient); 67 | expect(addedListener).toHaveBeenCalled(); 68 | }); 69 | 70 | test("should remove a client from the pool", () => { 71 | const removedListener = jest.fn(); 72 | comfyPool.addEventListener("removed", removedListener); 73 | 74 | comfyPool.removeClient(mockClient); 75 | 76 | expect(comfyPool.clients).not.toContain(mockClient); 77 | expect(removedListener).toHaveBeenCalledWith( 78 | expect.objectContaining({ detail: { client: mockClient, clientIdx: 0 } }) 79 | ); 80 | }); 81 | 82 | test("should remove a client from the pool by index", () => { 83 | const removedListener = jest.fn(); 84 | comfyPool.addEventListener("removed", removedListener); 85 | comfyPool.removeClientByIndex(0); 86 | expect(comfyPool.clients).not.toContain(mockClient); 87 | expect(removedListener).toHaveBeenCalledWith( 88 | expect.objectContaining({ detail: { client: mockClient, clientIdx: 0 } }) 89 | ); 90 | }); 91 | 92 | test("should change the mode of the queue", () => { 93 | const changeModeListener = jest.fn(); 94 | comfyPool.addEventListener("change_mode", changeModeListener); 95 | 96 | comfyPool.changeMode(EQueueMode.PICK_LOWEST); 97 | 98 | expect(comfyPool["mode"]).toBe(EQueueMode.PICK_LOWEST); 99 | expect(changeModeListener).toHaveBeenCalledWith( 100 | expect.objectContaining({ detail: { mode: EQueueMode.PICK_LOWEST } }) 101 | ); 102 | }); 103 | 104 | test("should pick a client by index", () => { 105 | const pickedClient = comfyPool.pick(0); 106 | expect(pickedClient).toBe(mockClient); 107 | }); 108 | 109 | test("should pick a client by ID", () => { 110 | const pickedClient = comfyPool.pickById("client1"); 111 | expect(pickedClient).toBe(mockClient); 112 | }); 113 | 114 | test("should run a job and dispatch events", async () => { 115 | const job = jest.fn().mockResolvedValue("result"); 116 | const executingListener = jest.fn(); 117 | const executedListener = jest.fn(); 118 | comfyPool.addEventListener("executing", executingListener); 119 | comfyPool.addEventListener("executed", executedListener); 120 | 121 | const result = await comfyPool.run(job); 122 | 123 | expect(job).toHaveBeenCalledWith(mockClient, 0); 124 | expect(executingListener).toHaveBeenCalledWith( 125 | expect.objectContaining({ detail: { client: mockClient, clientIdx: 0 } }) 126 | ); 127 | expect(executedListener).toHaveBeenCalledWith( 128 | expect.objectContaining({ detail: { client: mockClient, clientIdx: 0 } }) 129 | ); 130 | expect(result).toBe("result"); 131 | }); 132 | 133 | test("should handle client disconnection/reconnection", async () => { 134 | const disconnectedListener = jest.fn(); 135 | const reconnectedListener = jest.fn(); 136 | const tryClient = { 137 | ...mockClient, 138 | on: jest.fn((event, handler) => { 139 | if (event === "disconnected") { 140 | setTimeout(() => handler(), 0); // Simulate the disconnected event being dispatched 141 | } 142 | if (event === "reconnected") { 143 | setTimeout(() => handler(), 50); // Simulate the reconnected event being dispatched 144 | } 145 | }), 146 | id: "client2" 147 | } as any as ComfyApi; 148 | comfyPool = new ComfyPool([tryClient]); 149 | comfyPool.addEventListener("disconnected", disconnectedListener); 150 | comfyPool.addEventListener("reconnected", reconnectedListener); 151 | await new Promise((resolve) => setTimeout(resolve, 100)); 152 | 153 | expect(disconnectedListener).toHaveBeenCalled(); 154 | expect(reconnectedListener).toHaveBeenCalled(); 155 | }); 156 | 157 | test("should execute a batch of jobs", async () => { 158 | const job1 = jest.fn().mockResolvedValue("result1"); 159 | const job2 = jest.fn().mockResolvedValue("result2"); 160 | 161 | const newClient = { ...mockClient, id: "client2" } as ComfyApi; 162 | await comfyPool.addClient(newClient); 163 | 164 | const results = await comfyPool.batch([job1, job2]); 165 | 166 | expect(job1).toHaveBeenCalled(); 167 | expect(job2).toHaveBeenCalled(); 168 | expect(results).toEqual(["result1", "result2"]); 169 | }); 170 | }); 171 | -------------------------------------------------------------------------------- /src/prompt-builder.ts: -------------------------------------------------------------------------------- 1 | import { encodeNTPath, encodePosixPath } from "./tools"; 2 | import { NodeData, OSType } from "./types/api"; 3 | import { DeepKeys, Simplify } from "./types/tool"; 4 | 5 | export class PromptBuilder { 6 | prompt: T; 7 | mapInputKeys: Partial> = {}; 8 | mapOutputKeys: Partial> = {}; 9 | bypassNodes: (keyof T)[] = []; 10 | 11 | constructor(prompt: T, inputKeys: I[], outputKeys: O[]) { 12 | this.prompt = structuredClone(prompt); 13 | inputKeys.forEach((key) => { 14 | this.mapInputKeys[key] = undefined; 15 | }); 16 | outputKeys.forEach((key) => { 17 | this.mapOutputKeys[key] = undefined; 18 | }); 19 | return this; 20 | } 21 | 22 | /** 23 | * Creates a new instance of the PromptBuilder with the same prompt, input keys, and output keys. 24 | * 25 | * @returns A new instance of the PromptBuilder. 26 | */ 27 | clone(): PromptBuilder { 28 | const newBuilder = new PromptBuilder( 29 | this.prompt, 30 | Object.keys(this.mapInputKeys) as I[], 31 | Object.keys(this.mapOutputKeys) as O[] 32 | ); 33 | newBuilder.mapInputKeys = { ...this.mapInputKeys }; 34 | newBuilder.mapOutputKeys = { ...this.mapOutputKeys }; 35 | newBuilder.bypassNodes = [...this.bypassNodes]; 36 | return newBuilder; 37 | } 38 | 39 | /** 40 | * Marks a node to be bypassed at generation. 41 | * 42 | * @param node Node which will be bypassed. 43 | */ 44 | bypass(node: keyof T): PromptBuilder; 45 | 46 | /** 47 | * Marks multiple nodes to be bypassed at generation. 48 | * 49 | * @param nodes Array of nodes which will be bypassed. 50 | */ 51 | bypass(nodes: (keyof T)[]): PromptBuilder; 52 | 53 | bypass(nodes: keyof T | (keyof T)[]) { 54 | if (!Array.isArray(nodes)) { 55 | nodes = [nodes]; 56 | } 57 | const newBuilder = this.clone(); 58 | newBuilder.bypassNodes.push(...nodes); 59 | return newBuilder; 60 | } 61 | 62 | /** 63 | * Unmarks a node from bypass at generation. 64 | * 65 | * @param node Node to reverse bypass on. 66 | */ 67 | reinstate(node: keyof T): PromptBuilder; 68 | 69 | /** 70 | * Unmarks a collection of nodes from bypass at generation. 71 | * 72 | * @param nodes Array of nodes to reverse bypass on. 73 | */ 74 | reinstate(nodes: (keyof T)[]): PromptBuilder; 75 | 76 | reinstate(nodes: keyof T | (keyof T)[]) { 77 | if (!Array.isArray(nodes)) { 78 | nodes = [nodes]; 79 | } 80 | 81 | const newBuilder = this.clone(); 82 | for (const node of nodes) { 83 | newBuilder.bypassNodes.splice(newBuilder.bypassNodes.indexOf(node), 1); 84 | } 85 | return newBuilder; 86 | } 87 | 88 | /** 89 | * Sets the input node for a given key. Can be map multiple keys to the same input. 90 | * 91 | * @param input - The input node to set. 92 | * @param key - The key(s) to associate with the input node. Can be array of keys. 93 | * @returns This builder instance. 94 | */ 95 | setInputNode(input: I, key: DeepKeys | Array>) { 96 | return this.setRawInputNode(input, key); 97 | } 98 | 99 | /** 100 | * Sets the raw input node for the given input and key. This will bypass the typing check. Use for dynamic nodes. 101 | * 102 | * @param input - The input node to be set. 103 | * @param key - The key associated with the input node. 104 | * @returns The current instance for method chaining. 105 | */ 106 | setRawInputNode(input: I, key: string | string[]) { 107 | this.mapInputKeys[input] = key; 108 | return this.clone(); 109 | } 110 | 111 | /** 112 | * Appends raw input node keys to the map of input keys. This will bypass the typing check. Use for dynamic nodes. 113 | * 114 | * @param input - The input node to which the keys will be appended. 115 | * @param key - The key or array of keys to append to the input node. 116 | * @returns A clone of the current instance with the updated input keys. 117 | */ 118 | appendRawInputNode(input: I, key: string | string[]) { 119 | let keys = typeof key === "string" ? [key] : key; 120 | if (typeof this.mapInputKeys[input] === "string") { 121 | this.mapInputKeys[input] = [this.mapInputKeys[input] as string]; 122 | } 123 | this.mapInputKeys[input]?.push(...keys); 124 | return this.clone(); 125 | } 126 | 127 | /** 128 | * Appends mapped key into the input node. 129 | * 130 | * @param input - The input node to append. 131 | * @param key - The key(s) to associate with the input node. Can be array of keys. 132 | * @returns The updated prompt builder. 133 | */ 134 | appendInputNode(input: I, key: DeepKeys | Array>) { 135 | return this.appendRawInputNode(input, key); 136 | } 137 | 138 | /** 139 | * Sets the output node for a given key. This will bypass the typing check. Use for dynamic nodes. 140 | * 141 | * @param output - The output node to set. 142 | * @param key - The key to associate with the output node. 143 | * @returns This builder instance. 144 | */ 145 | setRawOutputNode(output: O, key: string) { 146 | this.mapOutputKeys[output] = key; 147 | return this.clone(); 148 | } 149 | 150 | /** 151 | * Sets the output node for a given key. 152 | * 153 | * @param output - The output node to set. 154 | * @param key - The key to associate with the output node. 155 | * @returns This builder instance. 156 | */ 157 | setOutputNode(output: O, key: DeepKeys) { 158 | return this.setRawOutputNode(output, key); 159 | } 160 | 161 | /** 162 | * Sets the value for a specific input key in the prompt builder. 163 | * 164 | * @template V - The type of the value being set. 165 | * @param {I} key - The input key. 166 | * @param {V} value - The value to set. 167 | * @param {OSType} [encodeOs] - The OS type to encode the path. 168 | * @returns A new prompt builder with the updated value. 169 | * @throws {Error} - If the key is not found. 170 | */ 171 | input(key: I, value: V, encodeOs?: OSType) { 172 | if (value !== undefined) { 173 | let valueToSet = value; 174 | /** 175 | * Handle encode path if needed, use for load models path 176 | */ 177 | if (encodeOs === OSType.NT && typeof valueToSet === "string") { 178 | valueToSet = encodeNTPath(valueToSet) as typeof valueToSet; 179 | } else if (encodeOs === OSType.POSIX && typeof valueToSet === "string") { 180 | valueToSet = encodePosixPath(valueToSet) as typeof valueToSet; 181 | } 182 | 183 | /** 184 | * Map the input key to the path in the prompt object 185 | */ 186 | let paths = this.mapInputKeys[key]; 187 | if (!paths) { 188 | throw new Error(`Key ${key} not found`); 189 | } 190 | if (typeof paths === "string") { 191 | paths = [paths]; 192 | } 193 | for (const path of paths as string[]) { 194 | const keys = path.split("."); 195 | let current = this.prompt as any; 196 | for (let i = 0; i < keys.length - 1; i++) { 197 | if (!current[keys[i]]) { 198 | current[keys[i]] = {}; // Alow to set value to undefined path 199 | } 200 | current = current[keys[i]]; 201 | } 202 | current[keys[keys.length - 1]] = valueToSet; 203 | } 204 | } 205 | return this as Simplify>; 206 | } 207 | 208 | /** 209 | * Sets the value for a any input key in the prompt builder. 210 | * 211 | * @template V - The type of the value being set. 212 | * @param {string} key - The input key. 213 | * @param {V} value - The value to set. 214 | * @param {OSType} [encodeOs] - The OS type to encode the path. 215 | * @returns A new prompt builder with the updated value. 216 | * @throws {Error} - If the key is not found. 217 | */ 218 | inputRaw(key: string, value: V, encodeOs?: OSType) { 219 | if (key === "__proto__" || key === "constructor" || key === "prototype") { 220 | throw new Error(`Invalid key: ${key}`); 221 | } 222 | if (value !== undefined) { 223 | let valueToSet = value; 224 | /** 225 | * Handle encode path if needed, use for load models path 226 | */ 227 | if (encodeOs === OSType.NT && typeof valueToSet === "string") { 228 | valueToSet = encodeNTPath(valueToSet) as typeof valueToSet; 229 | } else if (encodeOs === OSType.POSIX && typeof valueToSet === "string") { 230 | valueToSet = encodePosixPath(valueToSet) as typeof valueToSet; 231 | } 232 | 233 | const keys = key.split("."); 234 | let current = this.prompt as any; 235 | for (let i = 0; i < keys.length - 1; i++) { 236 | if (keys[i] === "__proto__" || keys[i] === "constructor") continue; 237 | if (!current[keys[i]]) { 238 | current[keys[i]] = {}; // Alow to set value to undefined path 239 | } 240 | current = current[keys[i]]; 241 | } 242 | if (keys[keys.length - 1] !== "__proto__" && keys[keys.length - 1] !== "constructor") { 243 | current[keys[keys.length - 1]] = valueToSet; 244 | } 245 | } 246 | return this as Simplify>; 247 | } 248 | 249 | /** 250 | * @deprecated Please call `input` directly instead 251 | */ 252 | get caller() { 253 | return this; 254 | } 255 | 256 | /** 257 | * Gets the workflow object of the prompt builder. 258 | */ 259 | get workflow() { 260 | return this.prompt as Simplify; 261 | } 262 | } 263 | -------------------------------------------------------------------------------- /src/features/manager.ts: -------------------------------------------------------------------------------- 1 | import { 2 | TDefaultUI, 3 | TExtensionNodeItem, 4 | EExtensionUpdateCheckResult, 5 | EUpdateResult, 6 | IExtensionInfo, 7 | TPreviewMethod, 8 | IInstallExtensionRequest, 9 | EInstallType, 10 | IExtensionUninstallRequest, 11 | IExtensionUpdateRequest, 12 | IExtensionActiveRequest, 13 | IModelInstallRequest, 14 | INodeMapItem 15 | } from "src/types/manager"; 16 | import { AbstractFeature } from "./abstract"; 17 | 18 | export interface FetchOptions extends RequestInit { 19 | headers?: { 20 | [key: string]: string; 21 | }; 22 | } 23 | 24 | export class ManagerFeature extends AbstractFeature { 25 | async checkSupported() { 26 | const data = await this.getVersion().catch(() => false); 27 | if (data !== false) { 28 | this.supported = true; 29 | } 30 | return this.supported; 31 | } 32 | 33 | public destroy(): void { 34 | this.supported = false; 35 | } 36 | 37 | private async fetchApi(path: string, options?: FetchOptions) { 38 | if (!this.supported) { 39 | return false; 40 | } 41 | return this.client.fetchApi(path, options); 42 | } 43 | 44 | /** 45 | * Set the default state to be displayed in the main menu when the browser starts. 46 | * 47 | * We use this api to checking if the manager feature is supported. 48 | * 49 | * Default will return the current state. 50 | * @deprecated Not working anymore 51 | */ 52 | async defaultUi(setUi?: TDefaultUI): Promise { 53 | return true; 54 | } 55 | 56 | async getVersion(): Promise { 57 | const callURL = "/manager/version"; 58 | const data = await this.client.fetchApi(callURL); 59 | if (data && data.ok) { 60 | return data.text() as Promise; 61 | } 62 | throw new Error("Failed to get version", { cause: data }); 63 | } 64 | 65 | /** 66 | * Retrieves a list of extension's nodes based on the specified mode. 67 | * 68 | * Usefull to find the node suitable for the current workflow. 69 | * 70 | * @param mode - The mode to determine the source of the nodes. Defaults to "local". 71 | * @returns A promise that resolves to an array of extension nodes. 72 | * @throws An error if the retrieval fails. 73 | */ 74 | async getNodeMapList(mode: "local" | "nickname" = "local"): Promise> { 75 | const listNodes: TExtensionNodeItem[] = []; 76 | const data = await this.fetchApi(`/customnode/getmappings?mode=${mode}`); 77 | if (data && data.ok) { 78 | const nodes: { [key: string]: [string[], any] } = await data.json(); 79 | for (const url in nodes) { 80 | const [nodeNames, nodeData] = nodes[url]; 81 | listNodes.push({ 82 | url, 83 | nodeNames, 84 | title_aux: nodeData.title_aux, 85 | title: nodeData.title, 86 | author: nodeData.author, 87 | nickname: nodeData.nickname, 88 | description: nodeData.description 89 | }); 90 | } 91 | return listNodes; 92 | } 93 | throw new Error("Failed to get node map list", { cause: data }); 94 | } 95 | 96 | /** 97 | * Checks for extension updates. 98 | * 99 | * @param mode - The mode to use for checking updates. Defaults to "local". 100 | * @returns The result of the extension update check. 101 | */ 102 | async checkExtensionUpdate(mode: "local" | "cache" = "local") { 103 | const data = await this.fetchApi(`/customnode/fetch_updates?mode=${mode}`); 104 | if (data && data.ok) { 105 | if (data.status === 201) { 106 | return EExtensionUpdateCheckResult.UPDATE_AVAILABLE; 107 | } 108 | return EExtensionUpdateCheckResult.NO_UPDATE; 109 | } 110 | return EExtensionUpdateCheckResult.FAILED; 111 | } 112 | 113 | /** 114 | * Updates all extensions. 115 | * @param mode - The update mode. Can be "local" or "cache". Defaults to "local". 116 | * @returns An object representing the result of the extension update. 117 | */ 118 | async updataAllExtensions(mode: "local" | "cache" = "local") { 119 | const data = await this.fetchApi(`/customnode/update_all?mode=${mode}`); 120 | if (data && data.ok) { 121 | if (data.status === 200) { 122 | return { type: EUpdateResult.UNCHANGED }; 123 | } 124 | return { 125 | type: EUpdateResult.SUCCESS, 126 | data: (await data.json()) as { updated: number; failed: number } 127 | } as const; 128 | } 129 | return { type: EUpdateResult.FAILED }; 130 | } 131 | 132 | /** 133 | * Updates the ComfyUI. 134 | * 135 | * @returns The result of the update operation. 136 | */ 137 | async updateComfyUI() { 138 | const data = await this.fetchApi("/comfyui_manager/update_comfyui"); 139 | if (data) { 140 | switch (data.status) { 141 | case 200: 142 | return EUpdateResult.UNCHANGED; 143 | case 201: 144 | return EUpdateResult.SUCCESS; 145 | default: 146 | return EUpdateResult.FAILED; 147 | } 148 | } 149 | return EUpdateResult.FAILED; 150 | } 151 | 152 | /** 153 | * Retrieves the list of extensions. 154 | * 155 | * @param mode - The mode to retrieve the extensions from. Can be "local" or "cache". Defaults to "local". 156 | * @param skipUpdate - Indicates whether to skip updating the extensions. Defaults to true. 157 | * @returns A promise that resolves to an object containing the channel and custom nodes, or false if the retrieval fails. 158 | * @throws An error if the retrieval fails. 159 | */ 160 | async getExtensionList( 161 | mode: "local" | "cache" = "local", 162 | skipUpdate: boolean = true 163 | ): Promise< 164 | | { 165 | channel: "local" | "default"; 166 | custom_nodes: IExtensionInfo[]; 167 | } 168 | | false 169 | > { 170 | const data = await this.fetchApi(`/customnode/getlist?mode=${mode}&skip_update=${skipUpdate}`); 171 | if (data && data.ok) { 172 | return data.json(); 173 | } 174 | throw new Error("Failed to get extension list", { cause: data }); 175 | } 176 | 177 | /** 178 | * Reboots the instance. 179 | * 180 | * @returns A promise that resolves to `true` if the instance was successfully rebooted, or `false` otherwise. 181 | */ 182 | async rebootInstance() { 183 | const data = await this.fetchApi("/manager/reboot").catch((e) => { 184 | return true; 185 | }); 186 | if (data !== true) return false; 187 | return true; 188 | } 189 | 190 | /** 191 | * Return the current preview method. Will set to `mode` if provided. 192 | * 193 | * @param mode - The preview method mode. 194 | * @returns The result of the preview method. 195 | * @throws An error if the preview method fails to set. 196 | */ 197 | async previewMethod(mode?: TPreviewMethod): Promise { 198 | let callURL = "/manager/preview_method"; 199 | if (mode) { 200 | callURL += `?value=${mode}`; 201 | } 202 | const data = await this.fetchApi(callURL); 203 | if (data && data.ok) { 204 | const result = await data.text(); 205 | if (!result) return mode; 206 | return result as TPreviewMethod; 207 | } 208 | throw new Error("Failed to set preview method", { cause: data }); 209 | } 210 | 211 | /** 212 | * Installs an extension based on the provided configuration. 213 | * 214 | * @param config - The configuration for the extension installation. 215 | * @returns A boolean indicating whether the installation was successful. 216 | * @throws An error if the installation fails. 217 | */ 218 | async installExtension(config: IInstallExtensionRequest) { 219 | const data = await this.fetchApi("/customnode/install", { 220 | method: "POST", 221 | body: JSON.stringify(config) 222 | }); 223 | if (data && data.ok) { 224 | return true; 225 | } 226 | throw new Error("Failed to install extension", { cause: data }); 227 | } 228 | 229 | /** 230 | * Try to fix installation of an extension by re-install it again with fixes. 231 | * 232 | * @param config - The configuration object for fixing the extension. 233 | * @returns A boolean indicating whether the extension was fixed successfully. 234 | * @throws An error if the fix fails. 235 | */ 236 | async fixInstallExtension( 237 | config: Omit & { 238 | install_type: EInstallType.GIT_CLONE; 239 | } 240 | ) { 241 | const data = await this.fetchApi("/customnode/fix", { 242 | method: "POST", 243 | body: JSON.stringify(config) 244 | }); 245 | if (data && data.ok) { 246 | return true; 247 | } 248 | throw new Error("Failed to fix extension installation", { cause: data }); 249 | } 250 | 251 | /** 252 | * Install an extension from a Git URL. 253 | * 254 | * @param url - The URL of the Git repository. 255 | * @returns A boolean indicating whether the installation was successful. 256 | * @throws An error if the installation fails. 257 | */ 258 | async installExtensionFromGit(url: string) { 259 | const data = await this.fetchApi("/customnode/install/git_url", { 260 | method: "POST", 261 | body: url 262 | }); 263 | if (data && data.ok) { 264 | return true; 265 | } 266 | throw new Error("Failed to install extension from git", { cause: data }); 267 | } 268 | 269 | /** 270 | * Installs pip packages. 271 | * 272 | * @param packages - An array of packages to install. 273 | * @returns A boolean indicating whether the installation was successful. 274 | * @throws An error if the installation fails. 275 | */ 276 | async installPipPackages(packages: string[]) { 277 | const data = await this.fetchApi("/customnode/install/pip", { 278 | method: "POST", 279 | body: packages.join(" ") 280 | }); 281 | if (data && data.ok) { 282 | return true; 283 | } 284 | throw new Error("Failed to install pip's packages", { cause: data }); 285 | } 286 | 287 | /** 288 | * Uninstalls an extension. 289 | * 290 | * @param config - The configuration for uninstalling the extension. 291 | * @returns A boolean indicating whether the uninstallation was successful. 292 | * @throws An error if the uninstallation fails. 293 | */ 294 | async uninstallExtension(config: IExtensionUninstallRequest) { 295 | const data = await this.fetchApi("/customnode/uninstall", { 296 | method: "POST", 297 | body: JSON.stringify(config) 298 | }); 299 | if (data && data.ok) { 300 | return true; 301 | } 302 | throw new Error("Failed to uninstall extension", { cause: data }); 303 | } 304 | 305 | /** 306 | * Updates the extension with the provided configuration. Only work with git-clone method 307 | * 308 | * @param config - The configuration object for the extension update. 309 | * @returns A boolean indicating whether the extension update was successful. 310 | * @throws An error if the extension update fails. 311 | */ 312 | async updateExtension(config: IExtensionUpdateRequest) { 313 | const data = await this.fetchApi("/customnode/update", { 314 | method: "POST", 315 | body: JSON.stringify(config) 316 | }); 317 | if (data && data.ok) { 318 | return true; 319 | } 320 | throw new Error("Failed to update extension", { cause: data }); 321 | } 322 | 323 | /** 324 | * Set the activation of extension. 325 | * 326 | * @param config - The configuration for the active extension. 327 | * @returns A boolean indicating whether the active extension was set successfully. 328 | * @throws An error if setting the active extension fails. 329 | */ 330 | async setActiveExtension(config: IExtensionActiveRequest) { 331 | const data = await this.fetchApi("/customnode/toggle_active", { 332 | method: "POST", 333 | body: JSON.stringify(config) 334 | }); 335 | if (data && data.ok) { 336 | return true; 337 | } 338 | throw new Error("Failed to set active extension", { cause: data }); 339 | } 340 | 341 | /** 342 | * Install a model from given info. 343 | * 344 | * @param info - The model installation request information. 345 | * @returns A boolean indicating whether the model installation was successful. 346 | * @throws An error if the model installation fails. 347 | */ 348 | async installModel(info: IModelInstallRequest) { 349 | const data = await this.fetchApi("/model/install", { 350 | method: "POST", 351 | body: JSON.stringify(info) 352 | }); 353 | if (data && data.ok) { 354 | return true; 355 | } 356 | throw new Error("Failed to install model", { cause: data }); 357 | } 358 | } 359 | -------------------------------------------------------------------------------- /src/pool.ts: -------------------------------------------------------------------------------- 1 | import { TComfyPoolEventMap } from "./types/event"; 2 | import { ComfyApi } from "./client"; 3 | import { delay } from "./tools"; 4 | 5 | interface JobItem { 6 | weight: number; 7 | /** 8 | * Only one of the following clientIds will be picked. 9 | */ 10 | includeClientIds?: string[]; 11 | /** 12 | * The following clientIds will be excluded from the picking list. 13 | */ 14 | excludeClientIds?: string[]; 15 | fn: (api: ComfyApi, clientIdx?: number) => Promise; 16 | } 17 | 18 | /** 19 | * Represents the mode for picking clients from a queue. 20 | * 21 | * - "PICK_ZERO": Picks the client which has zero queue remaining. This is the default mode. (For who using along with ComfyUI web interface) 22 | * - "PICK_LOWEST": Picks the client which has the lowest queue remaining. 23 | * - "PICK_ROUTINE": Picks the client in a round-robin manner. 24 | */ 25 | export enum EQueueMode { 26 | /** 27 | * Picks the client which has zero queue remaining. This is the default mode. (For who using along with ComfyUI web interface) 28 | */ 29 | "PICK_ZERO", 30 | /** 31 | * Picks the client which has the lowest queue remaining. 32 | */ 33 | "PICK_LOWEST", 34 | /** 35 | * Picks the client in a round-robin manner. 36 | */ 37 | "PICK_ROUTINE" 38 | } 39 | 40 | export class ComfyPool extends EventTarget { 41 | public clients: ComfyApi[] = []; 42 | private clientStates: Array<{ 43 | id: string; 44 | queueRemaining: number; 45 | locked: string | boolean; 46 | online: boolean; 47 | }> = []; 48 | 49 | private mode: EQueueMode = EQueueMode.PICK_ZERO; 50 | private jobQueue: Array = []; 51 | private routineIdx: number = 0; 52 | private listeners: { 53 | event: keyof TComfyPoolEventMap; 54 | options?: AddEventListenerOptions | boolean; 55 | handler: (event: TComfyPoolEventMap[keyof TComfyPoolEventMap]) => void; 56 | }[] = []; 57 | private maxQueueSize: number = 1000; 58 | 59 | constructor( 60 | clients: ComfyApi[], 61 | /** 62 | * The mode for picking clients from the queue. Defaults to "PICK_ZERO". 63 | */ 64 | mode: EQueueMode = EQueueMode.PICK_ZERO, 65 | opts?: { 66 | /** 67 | * The maximum size of the job queue. Defaults to 1000. 68 | */ 69 | maxQueueSize?: number; 70 | } 71 | ) { 72 | super(); 73 | this.mode = mode; 74 | if (opts?.maxQueueSize) { 75 | this.maxQueueSize = opts.maxQueueSize; 76 | } 77 | 78 | // Wait for event listeners to be attached before initializing the pool 79 | delay(1).then(() => { 80 | this.dispatchEvent(new CustomEvent("init")); 81 | clients.forEach((client) => { 82 | this.addClient(client); 83 | }); 84 | this.pickJob(); 85 | }); 86 | } 87 | 88 | public on( 89 | type: K, 90 | callback: (event: TComfyPoolEventMap[K]) => void, 91 | options?: AddEventListenerOptions | boolean 92 | ) { 93 | this.addEventListener(type, callback as any, options); 94 | this.listeners.push({ event: type, handler: callback, options }); 95 | return this; 96 | } 97 | 98 | public off( 99 | type: K, 100 | callback: (event: TComfyPoolEventMap[K]) => void, 101 | options?: EventListenerOptions | boolean 102 | ) { 103 | this.removeEventListener(type, callback as any, options); 104 | this.listeners = this.listeners.filter((listener) => listener.event !== type && listener.handler !== callback); 105 | return this; 106 | } 107 | 108 | /** 109 | * Removes all event listeners from the pool. 110 | */ 111 | public removeAllListeners() { 112 | this.listeners.forEach((listener) => { 113 | this.removeEventListener(listener.event, listener.handler, listener.options); 114 | }); 115 | this.listeners = []; 116 | } 117 | 118 | /** 119 | * Adds a client to the pool. 120 | * 121 | * @param client - The client to be added. 122 | * @returns Promise 123 | */ 124 | async addClient(client: ComfyApi) { 125 | const index = this.clients.push(client) - 1; 126 | this.clientStates.push({ 127 | id: client.id, 128 | queueRemaining: 0, 129 | locked: false, 130 | online: false 131 | }); 132 | await this.initializeClient(client, index); 133 | this.dispatchEvent(new CustomEvent("added", { detail: { client, clientIdx: index } })); 134 | } 135 | 136 | destroy() { 137 | this.clients.forEach((client) => client.destroy()); 138 | this.clients = []; 139 | this.clientStates = []; 140 | this.removeAllListeners(); 141 | } 142 | 143 | /** 144 | * Removes a client from the pool. 145 | * 146 | * @param client - The client to be removed. 147 | * @returns void 148 | */ 149 | removeClient(client: ComfyApi): void { 150 | const index = this.clients.indexOf(client); 151 | this.removeClientByIndex(index); 152 | } 153 | 154 | /** 155 | * Removes a client from the pool by its index. 156 | * 157 | * @param index - The index of the client to remove. 158 | * @returns void 159 | * @fires removed - Fires a "removed" event with the removed client and its index as detail. 160 | */ 161 | removeClientByIndex(index: number): void { 162 | if (index >= 0 && index < this.clients.length) { 163 | const client = this.clients.splice(index, 1)[0]; 164 | client.destroy(); 165 | this.clientStates.splice(index, 1); 166 | this.dispatchEvent(new CustomEvent("removed", { detail: { client, clientIdx: index } })); 167 | } 168 | } 169 | 170 | /** 171 | * Changes the mode of the queue. 172 | * 173 | * @param mode - The new mode to set for the queue. 174 | * @returns void 175 | */ 176 | changeMode(mode: EQueueMode): void { 177 | this.mode = mode; 178 | this.dispatchEvent(new CustomEvent("change_mode", { detail: { mode } })); 179 | } 180 | 181 | /** 182 | * Picks a ComfyApi client from the pool based on the given index. 183 | * 184 | * @param idx - The index of the client to pick. Defaults to 0 if not provided. 185 | * @returns The picked ComfyApi client. 186 | */ 187 | pick(idx: number = 0): ComfyApi { 188 | return this.clients[idx]; 189 | } 190 | 191 | /** 192 | * Retrieves a `ComfyApi` object from the pool based on the provided ID. 193 | * @param id - The ID of the `ComfyApi` object to retrieve. 194 | * @returns The `ComfyApi` object with the matching ID, or `undefined` if not found. 195 | */ 196 | pickById(id: string): ComfyApi | undefined { 197 | return this.clients.find((c) => c.id === id); 198 | } 199 | 200 | /** 201 | * Executes a job using the provided client and optional client index. 202 | * 203 | * @template T The type of the result returned by the job. 204 | * @param {Function} job The job to be executed. 205 | * @param {number} [weight] The weight of the job. 206 | * @param {Object} [clientFilter] An object containing client filtering options. 207 | * @returns {Promise} A promise that resolves with the result of the job. 208 | */ 209 | run( 210 | job: (client: ComfyApi, clientIdx?: number) => Promise, 211 | weight?: number, 212 | clientFilter?: { 213 | /** 214 | * Only one of the following clientIds will be picked. 215 | */ 216 | includeIds?: string[]; 217 | /** 218 | * The following clientIds will be excluded from the picking list. 219 | */ 220 | excludeIds?: string[]; 221 | } 222 | ): Promise { 223 | return new Promise((resolve, reject) => { 224 | const fn = async (client: ComfyApi, idx?: number) => { 225 | this.dispatchEvent(new CustomEvent("executing", { detail: { client, clientIdx: idx } })); 226 | try { 227 | resolve(await job(client, idx)); 228 | this.dispatchEvent(new CustomEvent("executed", { detail: { client, clientIdx: idx } })); 229 | } catch (e) { 230 | console.error(e); 231 | reject(e); 232 | this.dispatchEvent( 233 | new CustomEvent("execution_error", { 234 | detail: { client, clientIdx: idx, error: e } 235 | }) 236 | ); 237 | } 238 | }; 239 | this.claim(fn, weight, clientFilter); 240 | }); 241 | } 242 | 243 | /** 244 | * Executes a batch of asynchronous jobs concurrently and returns an array of results. 245 | * 246 | * @template T - The type of the result returned by each job. 247 | * @param jobs - An array of functions that represent the asynchronous jobs to be executed. 248 | * @param weight - An optional weight value to assign to each job. 249 | * @param clientFilter - An optional object containing client filtering options. 250 | * @returns A promise that resolves to an array of results, in the same order as the jobs array. 251 | */ 252 | batch( 253 | jobs: Array<(client: ComfyApi, clientIdx?: number) => Promise>, 254 | weight?: number, 255 | clientFilter?: { 256 | /** 257 | * Only one of the following clientIds will be picked. 258 | */ 259 | includeIds?: string[]; 260 | /** 261 | * The following clientIds will be excluded from the picking list. 262 | */ 263 | excludeIds?: string[]; 264 | } 265 | ): Promise { 266 | return Promise.all(jobs.map((task) => this.run(task, weight, clientFilter))); 267 | } 268 | 269 | private async initializeClient(client: ComfyApi, index: number) { 270 | this.dispatchEvent( 271 | new CustomEvent("loading_client", { 272 | detail: { client, clientIdx: index } 273 | }) 274 | ); 275 | const states = this.clientStates[index]; 276 | client.on("status", (ev) => { 277 | if (states.online === false) { 278 | this.dispatchEvent(new CustomEvent("connected", { detail: { client, clientIdx: index } })); 279 | } 280 | states.online = true; 281 | if (ev.detail.status.exec_info && ev.detail.status.exec_info.queue_remaining !== states.queueRemaining) { 282 | if (ev.detail.status.exec_info.queue_remaining > 0) { 283 | this.dispatchEvent( 284 | new CustomEvent("have_job", { 285 | detail: { client, remain: states.queueRemaining } 286 | }) 287 | ); 288 | } 289 | if (ev.detail.status.exec_info.queue_remaining === 0) { 290 | this.dispatchEvent(new CustomEvent("idle", { detail: { client } })); 291 | } 292 | } 293 | states.queueRemaining = ev.detail.status.exec_info.queue_remaining; 294 | if (this.mode !== EQueueMode.PICK_ZERO) { 295 | states.locked = false; 296 | } 297 | }); 298 | client.on("terminal", (ev) => { 299 | this.dispatchEvent( 300 | new CustomEvent("terminal", { 301 | detail: { 302 | clientIdx: index, 303 | ...ev.detail 304 | } 305 | }) 306 | ); 307 | }); 308 | client.on("disconnected", () => { 309 | states.online = false; 310 | states.locked = false; 311 | this.dispatchEvent( 312 | new CustomEvent("disconnected", { 313 | detail: { client, clientIdx: index } 314 | }) 315 | ); 316 | }); 317 | client.on("reconnected", () => { 318 | states.online = true; 319 | states.locked = false; 320 | this.dispatchEvent( 321 | new CustomEvent("reconnected", { 322 | detail: { client, clientIdx: index } 323 | }) 324 | ); 325 | }); 326 | client.on("execution_success", (ev) => { 327 | states.locked = false; 328 | }); 329 | client.on("execution_interrupted", (ev) => { 330 | states.locked = false; 331 | this.dispatchEvent( 332 | new CustomEvent("execution_interrupted", { 333 | detail: { 334 | client, 335 | clientIdx: index 336 | } 337 | }) 338 | ); 339 | }); 340 | client.on("execution_error", (ev) => { 341 | states.locked = false; 342 | this.dispatchEvent( 343 | new CustomEvent("execution_error", { 344 | detail: { 345 | client, 346 | clientIdx: index, 347 | error: new Error(ev.detail.exception_type, { cause: ev.detail }) 348 | } 349 | }) 350 | ); 351 | }); 352 | client.on("queue_error", (ev) => { 353 | states.locked = false; 354 | }); 355 | client.on("auth_error", (ev) => { 356 | this.dispatchEvent( 357 | new CustomEvent("auth_error", { 358 | detail: { client, clientIdx: index, res: ev.detail } 359 | }) 360 | ); 361 | }); 362 | client.on("auth_success", (ev) => { 363 | this.dispatchEvent( 364 | new CustomEvent("auth_success", { 365 | detail: { client, clientIdx: index } 366 | }) 367 | ); 368 | }); 369 | client.on("connection_error", (ev) => { 370 | this.dispatchEvent( 371 | new CustomEvent("connection_error", { 372 | detail: { client, clientIdx: index, res: ev.detail } 373 | }) 374 | ); 375 | }); 376 | /** 377 | * Wait for the client to be ready before start using it 378 | */ 379 | await client.init().waitForReady(); 380 | this.bindClientSystemMonitor(client, index); 381 | this.dispatchEvent(new CustomEvent("ready", { detail: { client, clientIdx: index } })); 382 | } 383 | 384 | private async bindClientSystemMonitor(client: ComfyApi, index: number) { 385 | if (client.ext.monitor.isSupported) { 386 | client.ext.monitor.on("system_monitor", (ev) => { 387 | this.dispatchEvent( 388 | new CustomEvent("system_monitor", { 389 | detail: { 390 | client, 391 | data: ev.detail, 392 | clientIdx: index 393 | } 394 | }) 395 | ); 396 | }); 397 | } 398 | } 399 | 400 | private pushJobByWeight(item: JobItem): number { 401 | const idx = this.jobQueue.findIndex((job) => job.weight > item.weight); 402 | if (idx === -1) { 403 | return this.jobQueue.push(item); 404 | } else { 405 | this.jobQueue.splice(idx, 0, item); 406 | return idx; 407 | } 408 | } 409 | 410 | private async claim( 411 | fn: (client: ComfyApi, clientIdx?: number) => Promise, 412 | weight?: number, 413 | clientFilter?: { 414 | includeIds?: string[]; 415 | excludeIds?: string[]; 416 | } 417 | ): Promise { 418 | if (this.jobQueue.length >= this.maxQueueSize) { 419 | throw new Error("Job queue limit reached"); 420 | } 421 | const inputWeight = weight === undefined ? this.jobQueue.length : weight; 422 | const idx = this.pushJobByWeight({ 423 | weight: inputWeight, 424 | fn, 425 | excludeClientIds: clientFilter?.excludeIds, 426 | includeClientIds: clientFilter?.includeIds 427 | }); 428 | this.dispatchEvent( 429 | new CustomEvent("add_job", { 430 | detail: { jobIdx: idx, weight: inputWeight } 431 | }) 432 | ); 433 | } 434 | 435 | private async getAvailableClient(includeIds?: string[], excludeIds?: string[], timeout = -1): Promise { 436 | let tries = 1; 437 | const start = Date.now(); 438 | while (true) { 439 | if (timeout > 0 && Date.now() - start > timeout) { 440 | throw new Error("Timeout waiting for an available client"); 441 | } 442 | if (tries < 100) tries++; 443 | let index = -1; 444 | const acceptedClients = this.clientStates.filter((c) => { 445 | if (!c.online) return false; 446 | if (includeIds && includeIds.length > 0) { 447 | return includeIds.includes(c.id); 448 | } 449 | if (excludeIds && excludeIds.length > 0) { 450 | return !excludeIds.includes(c.id); 451 | } 452 | return true; 453 | }); 454 | switch (this.mode) { 455 | case EQueueMode.PICK_ZERO: 456 | index = acceptedClients.findIndex((c) => c.queueRemaining === 0 && !c.locked && c.id); 457 | break; 458 | case EQueueMode.PICK_LOWEST: 459 | const queueSizes = acceptedClients.map((state) => 460 | state.online ? state.queueRemaining : Number.MAX_SAFE_INTEGER 461 | ); 462 | index = queueSizes.indexOf(Math.min(...queueSizes)); 463 | break; 464 | case EQueueMode.PICK_ROUTINE: 465 | index = this.routineIdx++ % acceptedClients.length; 466 | this.routineIdx = this.routineIdx % acceptedClients.length; 467 | break; 468 | } 469 | if (index !== -1 && acceptedClients[index]) { 470 | const trueIdx = this.clientStates.findIndex((c) => c.id === acceptedClients[index].id); 471 | this.clientStates[trueIdx].locked = true; 472 | const client = this.clients[trueIdx]; 473 | return client; 474 | } 475 | await delay(Math.min(tries * 10)); 476 | } 477 | } 478 | 479 | private async pickJob(): Promise { 480 | while (true) { 481 | if (this.jobQueue.length === 0) { 482 | await delay(100); 483 | continue; 484 | } 485 | const job = this.jobQueue.shift(); 486 | const client = await this.getAvailableClient(job?.includeClientIds, job?.excludeClientIds); 487 | const clientIdx = this.clients.indexOf(client); 488 | job?.fn?.(client, clientIdx); 489 | } 490 | } 491 | } 492 | -------------------------------------------------------------------------------- /src/call-wrapper.ts: -------------------------------------------------------------------------------- 1 | import { NodeData, NodeDef, NodeProgress } from "./types/api"; 2 | import { ComfyApi } from "./client"; 3 | import { PromptBuilder } from "./prompt-builder"; 4 | import { TExecutionCached } from "./types/event"; 5 | import { 6 | FailedCacheError, 7 | WentMissingError, 8 | EnqueueFailedError, 9 | DisconnectedError, 10 | CustomEventError, 11 | ExecutionFailedError, 12 | ExecutionInterruptedError, 13 | MissingNodeError 14 | } from "./types/error"; 15 | 16 | /** 17 | * Represents a wrapper class for making API calls using the ComfyApi client. 18 | * Provides methods for setting callback functions and executing the job. 19 | */ 20 | export class CallWrapper { 21 | private client: ComfyApi; 22 | private prompt: PromptBuilder; 23 | private started = false; 24 | private promptId?: string; 25 | private output: Record["mapOutputKeys"] | "_raw", any> = {} as any; 26 | 27 | private onPreviewFn?: (ev: Blob, promptId?: string) => void; 28 | private onPendingFn?: (promptId?: string) => void; 29 | private onStartFn?: (promptId?: string) => void; 30 | private onOutputFn?: ( 31 | key: keyof PromptBuilder["mapOutputKeys"] | "_raw", 32 | data: any, 33 | promptId?: string 34 | ) => void; 35 | private onFinishedFn?: ( 36 | data: Record["mapOutputKeys"], any> & { 37 | /** 38 | * The raw output data from the workflow execution. 39 | * Key is node_id, value is node output. 40 | */ 41 | _raw?: Record; 42 | }, 43 | promptId?: string 44 | ) => void; 45 | private onFailedFn?: (err: Error, promptId?: string) => void; 46 | private onProgressFn?: (info: NodeProgress, promptId?: string) => void; 47 | 48 | private onDisconnectedHandlerOffFn: any; 49 | private checkExecutingOffFn: any; 50 | private checkExecutedOffFn: any; 51 | private progressHandlerOffFn: any; 52 | private previewHandlerOffFn: any; 53 | private executionHandlerOffFn: any; 54 | private errorHandlerOffFn: any; 55 | private executionEndSuccessOffFn: any; 56 | private statusHandlerOffFn: any; 57 | private interruptionHandlerOffFn: any; 58 | 59 | /** 60 | * Constructs a new CallWrapper instance. 61 | * @param client The ComfyApi client. 62 | * @param workflow The workflow object. 63 | */ 64 | constructor(client: ComfyApi, workflow: PromptBuilder) { 65 | this.client = client; 66 | this.prompt = workflow; 67 | return this; 68 | } 69 | 70 | /** 71 | * Set the callback function to be called when a preview event occurs. 72 | * 73 | * @param fn - The callback function to be called. It receives a Blob object representing the event and an optional promptId string. 74 | * @returns The current instance of the CallWrapper. 75 | */ 76 | onPreview(fn: (ev: Blob, promptId?: string) => void) { 77 | this.onPreviewFn = fn; 78 | return this; 79 | } 80 | 81 | /** 82 | * Set a callback function to be executed when the job is queued. 83 | * @param {Function} fn - The callback function to be executed. 84 | * @returns The current instance of the CallWrapper. 85 | */ 86 | onPending(fn: (promptId?: string) => void) { 87 | this.onPendingFn = fn; 88 | return this; 89 | } 90 | 91 | /** 92 | * Set the callback function to be executed when the job start. 93 | * 94 | * @param fn - The callback function to be executed. It can optionally receive a `promptId` parameter. 95 | * @returns The current instance of the CallWrapper. 96 | */ 97 | onStart(fn: (promptId?: string) => void) { 98 | this.onStartFn = fn; 99 | return this; 100 | } 101 | 102 | /** 103 | * Sets the callback function to handle the output node when the workflow is executing. This is 104 | * useful when you want to handle the output of each nodes as they are being processed. 105 | * 106 | * All the nodes defined in the `mapOutputKeys` will be passed to this function when node is executed. 107 | * 108 | * @param fn - The callback function to handle the output. 109 | * @returns The current instance of the class. 110 | */ 111 | onOutput(fn: (key: keyof PromptBuilder["mapOutputKeys"] | "_raw", data: any, promptId?: string) => void) { 112 | this.onOutputFn = fn; 113 | return this; 114 | } 115 | 116 | /** 117 | * Set the callback function to be executed when the asynchronous operation is finished. 118 | * 119 | * @param fn - The callback function to be executed. It receives the data returned by the operation 120 | * and an optional promptId parameter. 121 | * @returns The current instance of the CallWrapper. 122 | */ 123 | onFinished( 124 | fn: ( 125 | data: Record["mapOutputKeys"], any> & { 126 | /** 127 | * The raw output data from the workflow execution. 128 | * Key is node_id, value is node output. 129 | */ 130 | _raw?: Record; 131 | }, 132 | promptId?: string 133 | ) => void 134 | ) { 135 | this.onFinishedFn = fn; 136 | return this; 137 | } 138 | 139 | /** 140 | * Set the callback function to be executed when the API call fails. 141 | * 142 | * @param fn - The callback function to be executed when the API call fails. 143 | * It receives an `Error` object as the first parameter and an optional `promptId` as the second parameter. 144 | * @returns The current instance of the CallWrapper. 145 | */ 146 | onFailed(fn: (err: Error, promptId?: string) => void) { 147 | this.onFailedFn = fn; 148 | return this; 149 | } 150 | 151 | /** 152 | * Set a callback function to be called when progress information is available. 153 | * @param fn - The callback function to be called with the progress information. 154 | * @returns The current instance of the CallWrapper. 155 | */ 156 | onProgress(fn: (info: NodeProgress, promptId?: string) => void) { 157 | this.onProgressFn = fn; 158 | return this; 159 | } 160 | 161 | /** 162 | * Run the call wrapper and returns the output of the executed job. 163 | * If the job is already cached, it returns the cached output. 164 | * If the job is not cached, it executes the job and returns the output. 165 | * 166 | * @returns A promise that resolves to the output of the executed job, 167 | * or `undefined` if the job is not found, 168 | * or `false` if the job execution fails. 169 | */ 170 | async run(): Promise["mapOutputKeys"] | "_raw", any> | undefined | false> { 171 | /** 172 | * Start the job execution. 173 | */ 174 | const job = await this.enqueueJob(); 175 | if (!job) { 176 | this.onFailedFn?.(new Error("Failed to queue prompt")); 177 | return false; 178 | } 179 | 180 | let promptLoadTrigger!: (value: boolean) => void; 181 | const promptLoadCached: Promise = new Promise((resolve) => { 182 | promptLoadTrigger = resolve; 183 | }); 184 | 185 | let jobDoneTrigger!: (value: Record["mapOutputKeys"] | "_raw", any> | false) => void; 186 | const jobDonePromise: Promise["mapOutputKeys"] | "_raw", any> | false> = 187 | new Promise((resolve) => { 188 | jobDoneTrigger = resolve; 189 | }); 190 | 191 | /** 192 | * Declare the function to check if the job is executing. 193 | */ 194 | const checkExecutingFn = (event: CustomEvent) => { 195 | if (event.detail && event.detail.prompt_id === job.prompt_id) { 196 | promptLoadTrigger(false); 197 | } 198 | }; 199 | /** 200 | * Declare the function to check if the job is cached. 201 | */ 202 | const checkExecutionCachedFn = (event: CustomEvent) => { 203 | const outputNodes = Object.values(this.prompt.mapOutputKeys).filter((n) => !!n) as string[]; 204 | if (event.detail.nodes.length > 0 && event.detail.prompt_id === job.prompt_id) { 205 | /** 206 | * Cached is true if all output nodes are included in the cached nodes. 207 | */ 208 | const cached = outputNodes.every((node) => event.detail.nodes.includes(node)); 209 | promptLoadTrigger(cached); 210 | } 211 | }; 212 | /** 213 | * Listen to the executing event. 214 | */ 215 | this.checkExecutingOffFn = this.client.on("executing", checkExecutingFn); 216 | this.checkExecutedOffFn = this.client.on("execution_cached", checkExecutionCachedFn); 217 | 218 | // race condition handling 219 | let wentMissing = false; 220 | let cachedOutputDone = false; 221 | let cachedOutputPromise: Promise< 222 | false | Record["mapOutputKeys"] | "_raw", any> | null 223 | > = Promise.resolve(null); 224 | 225 | const statusHandler = async () => { 226 | const queue = await this.client.getQueue(); 227 | const queueItems = [...queue.queue_pending, ...queue.queue_running]; 228 | for (const queueItem of queueItems) { 229 | if (queueItem[1] === job.prompt_id) { 230 | return; 231 | } 232 | } 233 | 234 | await cachedOutputPromise; 235 | if (cachedOutputDone) { 236 | return; 237 | } 238 | 239 | const output = await this.handleCachedOutput(job.prompt_id); 240 | 241 | wentMissing = true; 242 | 243 | if (output) { 244 | jobDoneTrigger(output); 245 | this.cleanupListeners(); 246 | return; 247 | } 248 | 249 | promptLoadTrigger(false); 250 | jobDoneTrigger(false); 251 | this.cleanupListeners(); 252 | this.onFailedFn?.(new WentMissingError("The job went missing!"), job.prompt_id); 253 | }; 254 | 255 | this.statusHandlerOffFn = this.client.on("status", statusHandler); 256 | 257 | await promptLoadCached; 258 | 259 | if (wentMissing) { 260 | return jobDonePromise; 261 | } 262 | 263 | cachedOutputPromise = this.handleCachedOutput(job.prompt_id); 264 | const output = await cachedOutputPromise; 265 | 266 | if (output) { 267 | cachedOutputDone = true; 268 | this.cleanupListeners(); 269 | jobDoneTrigger(output); 270 | return output; 271 | } 272 | if (output === false) { 273 | cachedOutputDone = true; 274 | this.cleanupListeners(); 275 | this.onFailedFn?.(new FailedCacheError("Failed to get cached output"), this.promptId); 276 | jobDoneTrigger(false); 277 | return false; 278 | } 279 | 280 | this.handleJobExecution(job.prompt_id, jobDoneTrigger); 281 | 282 | return jobDonePromise; 283 | } 284 | 285 | private async bypassWorkflowNodes(workflow: NodeData) { 286 | const nodeDefs: Record = {}; // cache node definitions 287 | 288 | for (const nodeId of this.prompt.bypassNodes) { 289 | if (!workflow[nodeId as string]) { 290 | throw new MissingNodeError(`Node ${nodeId.toString()} is missing from the workflow!`); 291 | } 292 | 293 | const classType = workflow[nodeId as string].class_type; 294 | 295 | const def = nodeDefs[classType] || (await this.client.getNodeDefs(classType))?.[classType]; 296 | if (!def) { 297 | throw new MissingNodeError(`Node type ${workflow[nodeId as string].class_type} is missing from server!`); 298 | } 299 | nodeDefs[classType] = def; 300 | 301 | const connections = new Map(); 302 | const connectedInputs: string[] = []; 303 | 304 | // connect output nodes to matching input nodes 305 | for (const [outputIdx, outputType] of def.output.entries()) { 306 | for (const [inputName, inputValue] of Object.entries(workflow[nodeId as string].inputs)) { 307 | if (connectedInputs.includes(inputName)) { 308 | continue; 309 | } 310 | 311 | if (def.input.required[inputName]?.[0] === outputType) { 312 | connections.set(outputIdx, inputValue); 313 | connectedInputs.push(inputName); 314 | break; 315 | } 316 | 317 | if (def.input.optional?.[inputName]?.[0] === outputType) { 318 | connections.set(outputIdx, inputValue); 319 | connectedInputs.push(inputName); 320 | break; 321 | } 322 | } 323 | } 324 | 325 | // search and replace all nodes' inputs referencing this node based on matching output type, or remove reference 326 | // if no matching output type was found 327 | for (const [conNodeId, conNode] of Object.entries(workflow)) { 328 | for (const [conInputName, conInputValue] of Object.entries(conNode.inputs)) { 329 | if (!Array.isArray(conInputValue) || conInputValue[0] !== nodeId) { 330 | continue; 331 | } 332 | 333 | if (connections.has(conInputValue[1])) { 334 | workflow[conNodeId].inputs[conInputName] = connections.get(conInputValue[1]); 335 | } else { 336 | delete workflow[conNodeId].inputs[conInputName]; 337 | } 338 | } 339 | } 340 | 341 | delete workflow[nodeId as string]; 342 | } 343 | 344 | return workflow; 345 | } 346 | 347 | private async enqueueJob() { 348 | let workflow = structuredClone(this.prompt.workflow) as NodeData; 349 | 350 | if (this.prompt.bypassNodes.length > 0) { 351 | try { 352 | workflow = await this.bypassWorkflowNodes(workflow); 353 | } catch (e) { 354 | if (e instanceof Response) { 355 | this.onFailedFn?.(new MissingNodeError("Failed to get workflow node definitions", { cause: await e.json() })); 356 | } else { 357 | this.onFailedFn?.(new MissingNodeError("There was a missing node in the workflow bypass.", { cause: e })); 358 | } 359 | return null; 360 | } 361 | } 362 | 363 | const job = await this.client.appendPrompt(workflow).catch(async (e) => { 364 | if (e instanceof Response) { 365 | this.onFailedFn?.(new EnqueueFailedError("Failed to queue prompt", { cause: await e.json() })); 366 | } else { 367 | this.onFailedFn?.(new EnqueueFailedError("Failed to queue prompt", { cause: e })); 368 | } 369 | return null; 370 | }); 371 | if (!job) { 372 | return; 373 | } 374 | 375 | this.promptId = job.prompt_id; 376 | this.onPendingFn?.(this.promptId); 377 | this.onDisconnectedHandlerOffFn = this.client.on("disconnected", () => 378 | this.onFailedFn?.(new DisconnectedError("Disconnected"), this.promptId) 379 | ); 380 | return job; 381 | } 382 | 383 | private async handleCachedOutput( 384 | promptId: string 385 | ): Promise["mapOutputKeys"] | "_raw", any> | false | null> { 386 | const hisData = await this.client.getHistory(promptId); 387 | if (hisData?.status?.completed) { 388 | const output = this.mapOutput(hisData.outputs); 389 | if (Object.values(output).some((v) => v !== undefined)) { 390 | this.onFinishedFn?.(output, this.promptId); 391 | return output; 392 | } else { 393 | return false; 394 | } 395 | } 396 | return null; 397 | } 398 | 399 | private mapOutput(outputNodes: any): Record["mapOutputKeys"] | "_raw", any> { 400 | const outputMapped = this.prompt.mapOutputKeys; 401 | const output: Record["mapOutputKeys"] | "_raw", any> = {} as any; 402 | 403 | for (const key in outputMapped) { 404 | const node = outputMapped[key]; 405 | if (node) { 406 | output[key as keyof PromptBuilder["mapOutputKeys"]] = outputNodes[node]; 407 | } else { 408 | if (!output._raw) { 409 | output._raw = {}; 410 | } 411 | output._raw[key] = outputNodes[key]; 412 | } 413 | } 414 | 415 | return output; 416 | } 417 | 418 | private handleJobExecution( 419 | promptId: string, 420 | jobDoneTrigger: (value: Record["mapOutputKeys"] | "_raw", any> | false) => void 421 | ): void { 422 | const reverseOutputMapped = this.reverseMapOutputKeys(); 423 | 424 | this.progressHandlerOffFn = this.client.on("progress", (ev) => this.handleProgress(ev, promptId)); 425 | this.previewHandlerOffFn = this.client.on("b_preview", (ev) => this.onPreviewFn?.(ev.detail, this.promptId)); 426 | 427 | const totalOutput = Object.keys(reverseOutputMapped).length; 428 | let remainingOutput = totalOutput; 429 | 430 | const executionHandler = (ev: CustomEvent) => { 431 | if (ev.detail.prompt_id !== promptId) return; 432 | 433 | const outputKey = reverseOutputMapped[ev.detail.node as keyof typeof this.prompt.mapOutputKeys]; 434 | if (outputKey) { 435 | this.output[outputKey as keyof PromptBuilder["mapOutputKeys"]] = ev.detail.output; 436 | this.onOutputFn?.(outputKey, ev.detail.output, this.promptId); 437 | remainingOutput--; 438 | } else { 439 | this.output._raw = this.output._raw || {}; 440 | this.output._raw[ev.detail.node as string] = ev.detail.output; 441 | this.onOutputFn?.(ev.detail.node as string, ev.detail.output, this.promptId); 442 | } 443 | 444 | if (remainingOutput === 0) { 445 | this.cleanupListeners(); 446 | this.onFinishedFn?.(this.output, this.promptId); 447 | jobDoneTrigger(this.output); 448 | } 449 | }; 450 | 451 | const executedEnd = async () => { 452 | if (remainingOutput !== 0) { 453 | // some cached output nodes might output after executedEnd, so check history data if an output is really missing 454 | const hisData = await this.client.getHistory(promptId); 455 | if (hisData?.status?.completed) { 456 | const outputCount = Object.keys(hisData.outputs).length; 457 | if (outputCount > 0 && outputCount - totalOutput === 0) { 458 | return; 459 | } 460 | } 461 | this.onFailedFn?.(new ExecutionFailedError("Execution failed"), this.promptId); 462 | this.cleanupListeners(); 463 | jobDoneTrigger(false); 464 | } 465 | }; 466 | 467 | this.executionEndSuccessOffFn = this.client.on("execution_success", executedEnd); 468 | this.executionHandlerOffFn = this.client.on("executed", executionHandler); 469 | this.errorHandlerOffFn = this.client.on("execution_error", (ev) => this.handleError(ev, promptId, jobDoneTrigger)); 470 | this.interruptionHandlerOffFn = this.client.on("execution_interrupted", (ev) => { 471 | if (ev.detail.prompt_id !== promptId) return; 472 | this.onFailedFn?.( 473 | new ExecutionInterruptedError("The execution was interrupted!", { cause: ev.detail }), 474 | ev.detail.prompt_id 475 | ); 476 | this.cleanupListeners(); 477 | jobDoneTrigger(false); 478 | }); 479 | } 480 | 481 | private reverseMapOutputKeys(): Record { 482 | const outputMapped: Partial> = this.prompt.mapOutputKeys; 483 | return Object.entries(outputMapped).reduce( 484 | (acc, [k, v]) => { 485 | if (v) acc[v] = k; 486 | return acc; 487 | }, 488 | {} as Record 489 | ); 490 | } 491 | 492 | private handleProgress(ev: CustomEvent, promptId: string) { 493 | if (ev.detail.prompt_id === promptId && !this.started) { 494 | this.started = true; 495 | this.onStartFn?.(this.promptId); 496 | } 497 | this.onProgressFn?.(ev.detail, this.promptId); 498 | } 499 | 500 | private handleError( 501 | ev: CustomEvent, 502 | promptId: string, 503 | resolve: (value: Record["mapOutputKeys"] | "_raw", any> | false) => void 504 | ) { 505 | if (ev.detail.prompt_id !== promptId) return; 506 | this.onFailedFn?.(new CustomEventError(ev.detail.exception_type, { cause: ev.detail }), ev.detail.prompt_id); 507 | this.cleanupListeners(); 508 | resolve(false); 509 | } 510 | 511 | private cleanupListeners() { 512 | this.onDisconnectedHandlerOffFn?.(); 513 | this.checkExecutingOffFn?.(); 514 | this.checkExecutedOffFn?.(); 515 | this.progressHandlerOffFn?.(); 516 | this.previewHandlerOffFn?.(); 517 | this.executionHandlerOffFn?.(); 518 | this.errorHandlerOffFn?.(); 519 | this.executionEndSuccessOffFn?.(); 520 | this.interruptionHandlerOffFn?.(); 521 | this.statusHandlerOffFn?.(); 522 | } 523 | } 524 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ✨ ComfyUI SDK ✨ 2 | 3 | [![NPM Version](https://img.shields.io/npm/v/@saintno/comfyui-sdk?style=flat-square)](https://www.npmjs.com/package/@saintno/comfyui-sdk) 4 | [![License](https://img.shields.io/npm/l/@saintno/comfyui-sdk?style=flat-square)](https://github.com/tctien342/comfyui-sdk/blob/main/LICENSE) 5 | ![CI](https://github.com/tctien342/comfyui-sdk/actions/workflows/release.yml/badge.svg) 6 | [![Buy Me a Coffee](https://img.shields.io/badge/Buy%20Me%20a%20Coffee-donate-yellow.svg)](https://www.buymeacoffee.com/tctien342) 7 | 8 | A robust and meticulously crafted TypeScript SDK 🚀 for seamless interaction with the [ComfyUI](https://github.com/comfyanonymous/ComfyUI) API. This SDK significantly simplifies the complexities of building, executing, and managing ComfyUI workflows, all while providing real-time updates and supporting multiple instances. 🖼️ 9 | 10 | ## 🧭 Table of Contents 11 | 12 | - [🌟 Key Features 🌟](#-key-features-) 13 | - [📦 Installation 📦](#-installation-) 14 | - [🚀 Getting Started 🚀](#-getting-started-) 15 | - [🎬 Basic Usage](#-basic-usage) 16 | - [🔄 Managing Multiple Instances with `ComfyPool`](#-managing-multiple-instances-with-comfypool) 17 | - [🔑 Authentication](#-authentication) 18 | - [📚 API Reference 📚](#-api-reference-) 19 | - [`ComfyApi`](#comfyapi) 20 | - [`CallWrapper`](#callwrapper) 21 | - [`PromptBuilder`](#promptbuilder) 22 | - [`ComfyPool`](#comfypool) 23 | - [🗂️ Enums](#-enums) 24 | - [🗄️ Types](#-types) 25 | - [🧩 Features](#-features) 26 | - [📂 Examples](#-examples) 27 | - [🤝 Contributing](#-contributing) 28 | - [📜 License](#-license) 29 | 30 | ## 🌟 Key Features 🌟 31 | 32 | - **💎 TypeScript Powered**: Enjoy a fully typed codebase, ensuring enhanced development, maintainability, and type safety. 🛡️ 33 | - **🏗️ Workflow Builder**: Construct and manipulate intricate ComfyUI workflows effortlessly using a fluent, intuitive builder pattern. 🧩 34 | - **🤹 Multi-Instance Management**: Handle a pool of ComfyUI instances with ease, employing flexible queueing strategies for optimal resource utilization. 🌐 35 | - **⚡ Real-Time Updates**: Subscribe to WebSocket events for live progress tracking, image previews, and error notifications. 🔔 36 | - **🔧 Custom WebSocket Support**: Supply your own WebSocket implementation for greater flexibility in different environments, with robust reconnection handling and fallback options. 🔄 37 | - **🔑 Authentication Flexibility**: Supports Basic Auth, Bearer Token, and Custom Authentication Headers, catering to diverse security requirements. 🔒 38 | - **🔌 Extension Support**: Seamlessly integrate with ComfyUI Manager and leverage system monitoring through the ComfyUI-Crystools extension. 🛠️ 39 | - **🔀 Flexible Node Bypassing**: Strategically bypass specific nodes in your workflows during generation, enabling advanced customization. ⏭️ 40 | - **📚 Comprehensive Examples**: Includes practical examples for Text-to-Image (T2I), Image-to-Image (I2I), and complex multi-node workflows. 📝 41 | - **🚨 Robust Error Handling**: Provides detailed error messages to facilitate debugging and graceful handling of API failures. 🐛 42 | - **📝 Automatic Changelog**: Automatically generates a changelog with each release, utilizing `auto-changelog` for transparent version tracking. 🔄 43 | 44 | ## 📦 Installation 📦 45 | 46 | ```bash 47 | bun add @saintno/comfyui-sdk 48 | ``` 49 | 50 | or 51 | 52 | ```bash 53 | npm i @saintno/comfyui-sdk 54 | ``` 55 | 56 | ## 🚀 Getting Started 🚀 57 | 58 | ### 🎬 Basic Usage 59 | 60 | Here's a simplified example to quickly get you started: 61 | 62 | ```typescript 63 | import { ComfyApi, CallWrapper, PromptBuilder, TSamplerName, TSchedulerName, seed } from "@saintno/comfyui-sdk"; 64 | import ExampleTxt2ImgWorkflow from "./example-txt2img-workflow.json"; 65 | 66 | const api = new ComfyApi("http://localhost:8189").init(); 67 | const workflow = new PromptBuilder( 68 | ExampleTxt2ImgWorkflow, 69 | ["positive", "negative", "checkpoint", "seed", "batch", "step", "cfg", "sampler", "sheduler", "width", "height"], 70 | ["images"] 71 | ) 72 | .setInputNode("checkpoint", "4.inputs.ckpt_name") 73 | .setInputNode("seed", "3.inputs.seed") 74 | .setInputNode("batch", "5.inputs.batch_size") 75 | .setInputNode("negative", "7.inputs.text") 76 | .setInputNode("positive", "6.inputs.text") 77 | .setInputNode("cfg", "3.inputs.cfg") 78 | .setInputNode("sampler", "3.inputs.sampler_name") 79 | .setInputNode("sheduler", "3.inputs.scheduler") 80 | .setInputNode("step", "3.inputs.steps") 81 | .setInputNode("width", "5.inputs.width") 82 | .setInputNode("height", "5.inputs.height") 83 | .setOutputNode("images", "9") 84 | .input("checkpoint", "SDXL/realvisxlV40_v40LightningBakedvae.safetensors", api.osType) 85 | .input("seed", seed()) 86 | .input("step", 6) 87 | .input("cfg", 1) 88 | .input("sampler", "dpmpp_2m_sde_gpu") 89 | .input("sheduler", "sgm_uniform") 90 | .input("width", 1024) 91 | .input("height", 1024) 92 | .input("batch", 1) 93 | .input("positive", "A picture of cute dog on the street"); 94 | 95 | new CallWrapper(api, workflow) 96 | .onFinished((data) => console.log(data.images?.images.map((img: any) => api.getPathImage(img)))) 97 | .run(); 98 | ``` 99 | 100 | #### 🔍 Breakdown 101 | 102 | - Import essential components from the SDK. 103 | - Create and initialize the `ComfyApi` instance. 104 | - Use `PromptBuilder` to define the workflow structure and set input nodes. 105 | - Set specific input values, including the checkpoint path, seed, and prompt. 106 | - Execute the workflow and log the generated image URLs using the `CallWrapper`. 107 | 108 | ### 🔄 Managing Multiple Instances with `ComfyPool` 109 | 110 | ```typescript 111 | import { 112 | ComfyApi, 113 | CallWrapper, 114 | ComfyPool, 115 | EQueueMode, 116 | PromptBuilder, 117 | seed, 118 | TSamplerName, 119 | TSchedulerName 120 | } from "@saintno/comfyui-sdk"; 121 | import ExampleTxt2ImgWorkflow from "./example-txt2img-workflow.json"; 122 | 123 | const ApiPool = new ComfyPool( 124 | [new ComfyApi("http://localhost:8188"), new ComfyApi("http://localhost:8189")], 125 | EQueueMode.PICK_ZERO 126 | ) 127 | .on("init", () => console.log("Pool in initializing")) 128 | .on("add_job", (ev) => console.log("Job added at index", ev.detail.jobIdx, "weight:", ev.detail.weight)) 129 | .on("added", (ev) => console.log("Client added", ev.detail.clientIdx)); 130 | 131 | const generateFn = async (api: ComfyApi, clientIdx?: number) => { 132 | const workflow = new PromptBuilder( 133 | ExampleTxt2ImgWorkflow, 134 | ["positive", "negative", "checkpoint", "seed", "batch", "step", "cfg", "sampler", "sheduler", "width", "height"], 135 | ["images"] 136 | ) 137 | .setInputNode("checkpoint", "4.inputs.ckpt_name") 138 | .setInputNode("seed", "3.inputs.seed") 139 | .setInputNode("batch", "5.inputs.batch_size") 140 | .setInputNode("negative", "7.inputs.text") 141 | .setInputNode("positive", "6.inputs.text") 142 | .setInputNode("step", "3.inputs.steps") 143 | .setInputNode("width", "5.inputs.width") 144 | .setInputNode("height", "5.inputs.height") 145 | .setInputNode("cfg", "3.inputs.cfg") 146 | .setInputNode("sampler", "3.inputs.sampler_name") 147 | .setInputNode("scheduler", "3.inputs.scheduler") 148 | .setOutputNode("images", "9") 149 | .input("checkpoint", "SDXL/realvisxlV40_v40LightningBakedvae.safetensors", api.osType) 150 | .input("seed", seed()) 151 | .input("step", 6) 152 | .input("width", 512) 153 | .input("height", 512) 154 | .input("batch", 2) 155 | .input("cfg", 1) 156 | .input("sampler", "dpmpp_2m_sde_gpu") 157 | .input("scheduler", "sgm_uniform") 158 | .input("positive", "A close up picture of cute Cat") 159 | .input("negative", "text, blurry, bad picture, nsfw"); 160 | 161 | return new Promise((resolve) => { 162 | new CallWrapper(api, workflow) 163 | .onFinished((data) => { 164 | const url = data.images?.images.map((img: any) => api.getPathImage(img)); 165 | resolve(url as string[]); 166 | }) 167 | .run(); 168 | }); 169 | }; 170 | 171 | const jobA = ApiPool.batch(Array(5).fill(generateFn), 10).then((res) => { 172 | console.log("Batch A done"); 173 | return res.flat(); 174 | }); 175 | 176 | const jobB = ApiPool.batch(Array(5).fill(generateFn), 0).then((res) => { 177 | console.log("Batch B done"); 178 | return res.flat(); 179 | }); 180 | 181 | console.log(await Promise.all([jobA, jobB]).then((res) => res.flat())); 182 | ``` 183 | 184 | #### 🔍 Breakdown 185 | 186 | - Create a `ComfyPool` with multiple `ComfyApi` instances. 187 | - Set up event listeners for pool initialization, job additions, and client connections. 188 | - Define an async function (`generateFn`) that creates a workflow, sets its inputs, and executes it with a `CallWrapper`. 189 | - Use `ApiPool.batch` to run multiple jobs and wait for all batches to complete. 190 | 191 | ### 🔑 Authentication 192 | 193 | ```typescript 194 | import { ComfyApi, BasicCredentials, BearerTokenCredentials, CustomCredentials } from "@saintno/comfyui-sdk"; 195 | 196 | // Basic Authentication 197 | const basicAuth = new ComfyApi("http://localhost:8189", "node-id", { 198 | credentials: { type: "basic", username: "username", password: "password" } as BasicCredentials 199 | }).init(); 200 | 201 | // Bearer Token Authentication 202 | const bearerAuth = new ComfyApi("http://localhost:8189", "node-id", { 203 | credentials: { type: "bearer_token", token: "your_bearer_token" } as BearerTokenCredentials 204 | }).init(); 205 | 206 | // Custom Header Authentication 207 | const customAuth = new ComfyApi("http://localhost:8189", "node-id", { 208 | credentials: { type: "custom", headers: { "X-Custom-Header": "your_custom_header" } } as CustomCredentials 209 | }).init(); 210 | ``` 211 | 212 | #### 🔍 Breakdown 213 | 214 | - Import the necessary types from the SDK. 215 | - Create `ComfyApi` instances using the corresponding credential types: `BasicCredentials`, `BearerTokenCredentials`, and `CustomCredentials`.. 216 | 217 | ### 🔌 Custom WebSocket Implementation 218 | 219 | ```typescript 220 | import { ComfyApi, WebSocketInterface } from "@saintno/comfyui-sdk"; 221 | import CustomWebSocket from "your-custom-websocket-library"; 222 | 223 | // Create a ComfyApi instance with a custom WebSocket implementation 224 | const api = new ComfyApi("http://localhost:8189", "node-id", { 225 | credentials: { type: "basic", username: "username", password: "password" }, 226 | customWebSocketImpl: CustomWebSocket as WebSocketInterface 227 | }).init(); 228 | ``` 229 | 230 | #### 🔍 Breakdown 231 | 232 | - Import the necessary types and your custom WebSocket implementation. 233 | - Pass the custom WebSocket implementation to the ComfyApi constructor using the `customWebSocketImpl` option. 234 | - The SDK will use your implementation instead of the default one, allowing for greater flexibility in environments where the standard WebSocket implementation might not be available or suitable. 235 | 236 | #### 📝 Benefits 237 | 238 | - **Environment Flexibility**: Run in environments where the standard WebSocket might not be available or optimal. 239 | - **Enhanced Stability**: The SDK includes robust reconnection logic with exponential backoff and jitter to handle connection issues gracefully. 240 | - **Fallback Mechanism**: Automatically falls back to HTTP polling if WebSocket connections fail, ensuring your application remains functional. 241 | - **Custom Protocol Support**: Implement custom protocols or security features through your WebSocket implementation. 242 | 243 | ## 📚 API Reference 📚 244 | 245 | ### `ComfyApi` 246 | 247 | #### 🏗️ Constructor 248 | 249 | ```typescript 250 | constructor(host: string, clientId: string, opts?: { forceWs?: boolean, wsTimeout?: number, credentials?: BasicCredentials | BearerTokenCredentials | CustomCredentials; }) 251 | ``` 252 | 253 | - `host`: The base URL of your ComfyUI server. 254 | - `clientId`: A unique ID for WebSocket communication (optional). Defaults to a generated ID. 255 | - `opts`: Optional settings: 256 | - `forceWs`: Boolean to force WebSocket usage. 257 | - `wsTimeout`: Timeout for WebSocket connections (milliseconds). 258 | - `credentials`: Optional authentication credentials. 259 | 260 | #### ⚙️ Methods 261 | 262 | - `init(maxTries?: number, delayTime?: number)`: Initializes the client and establishes connection. 263 | - `on(type: K, callback: (event: TComfyAPIEventMap[K]) => void, options?: AddEventListenerOptions | boolean)`: Attach an event listener. 264 | - `off(type: K, callback: (event: TComfyAPIEventMap[K]) => void, options?: EventListenerOptions | boolean)`: Detach an event listener. 265 | - `removeAllListeners()`: Detach all event listeners. 266 | - `fetchApi(route: string, options?: FetchOptions)`: Fetch data from the API endpoint. 267 | - `pollStatus(timeout?: number)`: Polls the ComfyUI server status. 268 | - `queuePrompt(number: number | null, workflow: object)`: Queues a prompt for processing. 269 | - `appendPrompt(workflow: object)`: Adds a prompt to the workflow queue. 270 | - `getQueue()`: Retrieves the current state of the queue. 271 | - `getHistories(maxItems?: number)`: Retrieves the prompt execution history. 272 | - `getHistory(promptId: string)`: Retrieves a specific history entry by ID. 273 | - `getSystemStats()`: Retrieves system and device statistics. 274 | - `getExtensions()`: Retrieves a list of installed extensions. 275 | - `getEmbeddings()`: Retrieves a list of available embeddings. 276 | - `getCheckpoints()`: Retrieves a list of available checkpoints. 277 | - `getLoras()`: Retrieves a list of available Loras. 278 | - `getSamplerInfo()`: Retrieves sampler and scheduler information. 279 | - `getNodeDefs(nodeName?: string)`: Retrieves node object definitions. 280 | - `getUserConfig()`: Get user configuration data. 281 | - `createUser(username: string)`: Create new user. 282 | - `getSettings()`: Get all setting values for the current user. 283 | - `getSetting(id: string)`: Get a specific setting for the current user. 284 | - `storeSettings(settings: Record)`: Store setting for the current user. 285 | - `storeSetting(id: string, value: unknown)`: Store a specific setting for the current user. 286 | - `uploadImage(file: Buffer | Blob, fileName: string, config?: { override?: boolean; subfolder?: string })`: Uploads an image file. 287 | - `uploadMask(file: Buffer | Blob, originalRef: ImageInfo)`: Uploads a mask file. 288 | - `freeMemory(unloadModels: boolean, freeMemory: boolean)`: Frees memory by unloading models. 289 | - `getPathImage(imageInfo: ImageInfo)`: Returns the path to an image. 290 | - `getImage(imageInfo: ImageInfo)`: Returns the blob data of image. 291 | - `getUserData(file: string)`: Get a user data file. 292 | - `storeUserData(file: string, data: unknown, options?: RequestInit & { overwrite?: boolean, stringify?: boolean, throwOnError?: boolean })`: Store a user data file. 293 | - `deleteUserData(file: string)`: Delete a user data file. 294 | - `moveUserData(source: string, dest: string, options?: RequestInit & { overwrite?: boolean })`: Move a user data file. 295 | - `listUserData(dir: string, recurse?: boolean, split?: boolean)`: List a user data file. 296 | - `interrupt()`: Interrupts the execution of the running prompt. 297 | - `reconnectWs(opened?: boolean)`: Reconnects to the WebSocket server. 298 | 299 | ### `CallWrapper` 300 | 301 | #### 🏗️ Constructor 302 | 303 | ```typescript 304 | constructor(client: ComfyApi, workflow: PromptBuilder) 305 | ``` 306 | 307 | - `client`: An instance of the `ComfyApi` client. 308 | - `workflow`: An instance of `PromptBuilder` defining the workflow. 309 | 310 | #### ⚙️ Methods 311 | 312 | - `onPreview(fn: (ev: Blob, promptId?: string) => void)`: Set callback for preview events. 313 | - `onPending(fn: (promptId?: string) => void)`: Set callback when job is queued. 314 | - `onStart(fn: (promptId?: string) => void)`: Set callback when the job is started. 315 | - `onOutput(fn: (key: keyof PromptBuilder["mapOutputKeys"], data: any, promptId?: string) => void)`: Sets a callback for when an output node is executed. 316 | - `onFinished(fn: (data: Record["mapOutputKeys"], any>, promptId?: string) => void)`: Set callback when the job is finished. 317 | - `onFailed(fn: (err: Error, promptId?: string) => void)`: Set callback when the job failed. 318 | - `onProgress(fn: (info: NodeProgress, promptId?: string) => void)`: Set callback for progress updates. 319 | - `run()`: Executes the workflow. 320 | 321 | ### `PromptBuilder` 322 | 323 | #### 🏗️ Constructor 324 | 325 | ```typescript 326 | constructor(prompt: T, inputKeys: I[], outputKeys: O[]) 327 | ``` 328 | 329 | - `prompt`: The initial workflow data object. 330 | - `inputKeys`: An array of input node keys. 331 | - `outputKeys`: An array of output node keys. 332 | 333 | #### ⚙️ Methods 334 | 335 | - `clone()`: Creates a new `PromptBuilder` instance with the same configuration. 336 | - `bypass(node: keyof T | (keyof T)[]): PromptBuilder`: Marks node(s) to be bypassed at generation. 337 | - `reinstate(node: keyof T | (keyof T)[]): PromptBuilder`: Unmarks node(s) from bypass at generation. 338 | - `setInputNode(input: I, key: DeepKeys | Array>)`: Sets input node path for a key. 339 | - `setRawInputNode(input: I, key: string | string[])`: Sets raw input node path for a key. 340 | - `appendInputNode(input: I, key: DeepKeys | Array>)`: Appends a node to the input node path. 341 | - `appendRawInputNode(input: I, key: string | string[])`: Appends a node to the raw input node path. 342 | - `setOutputNode(output: O, key: DeepKeys)`: Sets output node path for a key. 343 | - `setRawOutputNode(output: O, key: string)`: Sets raw output node path for a key. 344 | - `input(key: I, value: V, encodeOs?: OSType)`: Sets an input value. 345 | - `inputRaw(key: string, value: V, encodeOs?: OSType)`: Sets a raw input value with dynamic key. 346 | - `get workflow`: Retrieves the workflow object. 347 | - `get caller`: Retrieves current `PromptBuilder` object. 348 | 349 | ### `ComfyPool` 350 | 351 | #### 🏗️ Constructor 352 | 353 | ```typescript 354 | constructor(clients: ComfyApi[], mode: EQueueMode = EQueueMode.PICK_ZERO) 355 | ``` 356 | 357 | - `clients`: Array of `ComfyApi` instances. 358 | - `mode`: The queue mode using `EQueueMode` enum values. 359 | 360 | #### ⚙️ Methods 361 | 362 | - `on(type: K, callback: (event: TComfyPoolEventMap[K]) => void, options?: AddEventListenerOptions | boolean)`: Attach an event listener. 363 | - `off(type: K, callback: (event: TComfyPoolEventMap[K]) => void, options?: EventListenerOptions | boolean)`: Detach an event listener. 364 | - `addClient(client: ComfyApi)`: Adds a new client to the pool. 365 | - `removeClient(client: ComfyApi)`: Removes a client from the pool. 366 | - `removeClientByIndex(index: number)`: Removes a client by index. 367 | - `changeMode(mode: EQueueMode)`: Changes the queue mode. 368 | - `pick(idx?: number)`: Picks a client by index. 369 | - `pickById(id: string)`: Picks a client by ID. 370 | - `run(job: (client: ComfyApi, clientIdx?: number) => Promise, weight?: number, clientFilter?: { includeIds?: string[]; excludeIds?: string[] })`: Run a job with priority on an available client. 371 | - `batch(jobs: Array<(client: ComfyApi, clientIdx?: number) => Promise>, weight?: number, clientFilter?: { includeIds?: string[]; excludeIds?: string[] })`: Run multiple jobs concurrently. 372 | 373 | ### 🗂️ Enums 374 | 375 | - `EQueueMode`: 376 | - `PICK_ZERO`: Selects the client with zero remaining queue. 377 | - `PICK_LOWEST`: Selects the client with the lowest remaining queue. 378 | - `PICK_ROUTINE`: Selects clients in a round-robin manner. 379 | 380 | ### 🗄️ Types 381 | 382 | - `OSType`: 383 | - `POSIX`: For Unix-like systems. 384 | - `NT`: For Windows systems. 385 | - `JAVA`: For Java virtual machine. 386 | - `TSamplerName`: A union type of all available sampler names. 387 | - `TSchedulerName`: A union type of all available scheduler names. 388 | 389 | ### 🧩 Features 390 | 391 | - `ManagerFeature`: Provides methods to manage ComfyUI Manager Extension. 392 | 393 | ```typescript 394 | const api = new ComfyApi("http://localhost:8189").init(); 395 | await api.waitForReady(); 396 | 397 | if (api.ext.manager.isSupported) { 398 | await api.ext.manager.getExtensionList().then(console.log); 399 | // Check api.ext.manager for more methods 400 | } 401 | ``` 402 | 403 | - `MonitoringFeature`: Provides methods to monitor system resources using ComfyUI-Crystools Extension. 404 | 405 | ```typescript 406 | const api = new ComfyApi("http://localhost:8189").init(); 407 | await api.waitForReady(); 408 | 409 | if (api.ext.monitor.isSupported) { 410 | // For subscribing to system monitor events 411 | api.ext.monitor.on("system_monitor", (ev) => { 412 | console.log(ev.detail); 413 | }); 414 | 415 | // For getting current monitor data 416 | console.log(api.ext.monitor.monitorData); 417 | } 418 | ``` 419 | 420 | > Note: Features require respective extensions ([ComfyUI-Manager](https://github.com/ltdrdata/ComfyUI-Manager) and [ComfyUI-Crystools](https://github.com/crystian/ComfyUI-Crystools)) to be installed. 421 | 422 | ## 📂 Examples 423 | 424 | The `examples` directory contains practical demonstrations of SDK usage: 425 | 426 | - `example-i2i.ts`: Demonstrates image-to-image generation. 427 | - `example-pool.ts`: Demonstrates how to manage multiple ComfyUI instances using `ComfyPool`. 428 | - `example-pool-basic-auth.ts`: Demonstrates how to use `ComfyPool` with HTTP Basic Authentication. 429 | - `example-t2i.ts`: Demonstrates text-to-image generation. 430 | - `example-t2i-upscaled.ts`: Demonstrates text-to-image generation with upscaling. 431 | - `example-img2img-workflow.json`: Example workflow for image-to-image. 432 | - `example-txt2img-workflow.json`: Example workflow for text-to-image. 433 | - `example-txt2img-upscaled-workflow.json`: Example workflow for text-to-image with upscaling. 434 | 435 | ## 🤝 Contributing 436 | 437 | Contributions are always welcome! Feel free to submit pull requests or create issues for bug reports and feature enhancements. 🙏 438 | 439 | ## 📜 License 440 | 441 | This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for more details. 📄 442 | -------------------------------------------------------------------------------- /src/client.ts: -------------------------------------------------------------------------------- 1 | import { WebSocketClient, WebSocketInterface } from "./socket"; 2 | 3 | import { 4 | BasicCredentials, 5 | BearerTokenCredentials, 6 | CustomCredentials, 7 | HistoryEntry, 8 | HistoryResponse, 9 | ImageInfo, 10 | ModelFile, 11 | ModelFolder, 12 | ModelPreviewResponse, 13 | NodeDefsResponse, 14 | OSType, 15 | QueuePromptResponse, 16 | QueueResponse, 17 | QueueStatus, 18 | SystemStatsResponse 19 | } from "./types/api"; 20 | 21 | import { LOAD_CHECKPOINTS_EXTENSION, LOAD_KSAMPLER_EXTENSION, LOAD_LORAS_EXTENSION } from "./contansts"; 22 | import { TComfyAPIEventMap } from "./types/event"; 23 | import { delay } from "./tools"; 24 | import { ManagerFeature } from "./features/manager"; 25 | import { MonitoringFeature } from "./features/monitoring"; 26 | 27 | interface FetchOptions extends RequestInit { 28 | headers?: { [key: string]: string }; 29 | } 30 | 31 | export class ComfyApi extends EventTarget { 32 | public apiHost: string; 33 | public osType: OSType; 34 | public isReady: boolean = false; 35 | public listenTerminal: boolean = false; 36 | public lastActivity: number = Date.now(); 37 | 38 | private wsTimeout: number = 10000; 39 | private wsTimer: Timer | null = null; 40 | private _pollingTimer: NodeJS.Timeout | number | null = null; 41 | 42 | private apiBase: string; 43 | private clientId: string | null; 44 | private socket: WebSocketClient | null = null; 45 | private customWsImpl: WebSocketInterface | null = null; 46 | private listeners: { 47 | event: keyof TComfyAPIEventMap; 48 | options?: AddEventListenerOptions | boolean; 49 | handler: (event: TComfyAPIEventMap[keyof TComfyAPIEventMap]) => void; 50 | }[] = []; 51 | private credentials: BasicCredentials | BearerTokenCredentials | CustomCredentials | null = null; 52 | 53 | public ext = { 54 | /** 55 | * Interact with ComfyUI-Manager Extension 56 | */ 57 | manager: new ManagerFeature(this), 58 | /** 59 | * Interact with ComfyUI-Crystools Extension for track system resouces 60 | */ 61 | monitor: new MonitoringFeature(this) 62 | }; 63 | 64 | static generateId(): string { 65 | return "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx".replace(/[xy]/g, function (c) { 66 | const r = (Math.random() * 16) | 0; 67 | const v = c === "x" ? r : (r & 0x3) | 0x8; 68 | return v.toString(16); 69 | }); 70 | } 71 | 72 | public on( 73 | type: K, 74 | callback: (event: TComfyAPIEventMap[K]) => void, 75 | options?: AddEventListenerOptions | boolean 76 | ) { 77 | this.log("on", "Add listener", { type, callback, options }); 78 | this.addEventListener(type, callback as any, options); 79 | this.listeners.push({ event: type, handler: callback, options }); 80 | const clr = () => this.off(type, callback, options); 81 | return clr; 82 | } 83 | 84 | public off( 85 | type: K, 86 | callback: (event: TComfyAPIEventMap[K]) => void, 87 | options?: EventListenerOptions | boolean 88 | ): void { 89 | this.log("off", "Remove listener", { type, callback, options }); 90 | this.listeners = this.listeners.filter((listener) => listener.event !== type && listener.handler !== callback); 91 | this.removeEventListener(type, callback as any, options); 92 | } 93 | 94 | public removeAllListeners() { 95 | this.log("removeAllListeners", "Triggered"); 96 | this.listeners.forEach((listener) => { 97 | this.removeEventListener(listener.event, listener.handler, listener.options); 98 | }); 99 | this.listeners = []; 100 | } 101 | 102 | get id(): string { 103 | return this.clientId ?? this.apiBase; 104 | } 105 | 106 | /** 107 | * Retrieves the available features of the client. 108 | * 109 | * @returns An object containing the available features, where each feature is a key-value pair. 110 | */ 111 | get availableFeatures() { 112 | return Object.keys(this.ext).reduce( 113 | (acc, key) => ({ ...acc, [key]: this.ext[key as keyof typeof this.ext].isSupported }), 114 | {} 115 | ); 116 | } 117 | 118 | constructor( 119 | host: string, 120 | clientId: string = ComfyApi.generateId(), 121 | opts?: { 122 | /** 123 | * Do not fallback to HTTP if WebSocket is not available. 124 | * This will retry to connect to WebSocket on error. 125 | */ 126 | forceWs?: boolean; 127 | /** 128 | * Timeout for WebSocket connection. 129 | * Default is 10000ms. 130 | */ 131 | wsTimeout?: number; 132 | /** 133 | * Listen to terminal logs from the server. Default (false) 134 | */ 135 | listenTerminal?: boolean; 136 | /** 137 | * Custom WebSocket implementation. 138 | */ 139 | customWebSocketImpl?: WebSocketInterface; 140 | credentials?: BasicCredentials | BearerTokenCredentials | CustomCredentials; 141 | } 142 | ) { 143 | super(); 144 | this.apiHost = host; 145 | this.apiBase = host.split("://")[1]; 146 | this.clientId = clientId; 147 | if (opts?.credentials) { 148 | this.credentials = opts?.credentials; 149 | this.testCredentials(); 150 | } 151 | if (opts?.wsTimeout) { 152 | this.wsTimeout = opts.wsTimeout; 153 | } 154 | if (opts?.listenTerminal) { 155 | this.listenTerminal = opts.listenTerminal; 156 | } 157 | if (opts?.customWebSocketImpl) { 158 | this.customWsImpl = opts.customWebSocketImpl; 159 | } 160 | this.log("constructor", "Initialized", { host, clientId, opts }); 161 | return this; 162 | } 163 | 164 | /** 165 | * Destroys the client instance. 166 | */ 167 | destroy() { 168 | this.log("destroy", "Destroying client..."); 169 | // Clean up WebSocket timer 170 | if (this.wsTimer) clearInterval(this.wsTimer); 171 | // Clean up polling timer if exists 172 | if (this._pollingTimer) { 173 | clearInterval(this._pollingTimer as any); 174 | this._pollingTimer = null; 175 | } 176 | // Clean up socket event handlers 177 | if (this.socket?.client) { 178 | this.socket.client.onclose = null; 179 | this.socket.client.onerror = null; 180 | this.socket.client.onmessage = null; 181 | this.socket.client.onopen = null; 182 | this.socket.client.close(); 183 | } 184 | for (const ext in this.ext) { 185 | this.ext[ext as keyof typeof this.ext].destroy(); 186 | } 187 | this.socket?.close(); 188 | this.log("destroy", "Client destroyed"); 189 | this.removeAllListeners(); 190 | } 191 | 192 | private log(fnName: string, message: string, data?: any) { 193 | this.dispatchEvent(new CustomEvent("log", { detail: { fnName, message, data } })); 194 | } 195 | 196 | private apiURL(route: string): string { 197 | return `${this.apiHost}${route}`; 198 | } 199 | 200 | private getCredentialHeaders(): Record { 201 | if (!this.credentials) return {}; 202 | switch (this.credentials?.type) { 203 | case "basic": 204 | return { Authorization: `Basic ${btoa(`${this.credentials.username}:${this.credentials.password}`)}` }; 205 | case "bearer_token": 206 | return { Authorization: `Bearer ${this.credentials.token}` }; 207 | case "custom": 208 | return this.credentials.headers; 209 | default: 210 | return {}; 211 | } 212 | } 213 | 214 | private async testCredentials() { 215 | try { 216 | if (!this.credentials) return false; 217 | await this.pollStatus(2000); 218 | this.dispatchEvent(new CustomEvent("auth_success")); 219 | return true; 220 | } catch (e) { 221 | this.log("testCredentials", "Failed", e); 222 | if (e instanceof Response) { 223 | if (e.status === 401) { 224 | this.dispatchEvent(new CustomEvent("auth_error", { detail: e })); 225 | return; 226 | } 227 | } 228 | this.dispatchEvent(new CustomEvent("connection_error", { detail: e })); 229 | return false; 230 | } 231 | } 232 | 233 | private async testFeatures() { 234 | const exts = Object.values(this.ext); 235 | await Promise.all(exts.map((ext) => ext.checkSupported())); 236 | /** 237 | * Mark the client is ready to use the API. 238 | */ 239 | this.isReady = true; 240 | } 241 | 242 | /** 243 | * Fetches data from the API. 244 | * 245 | * @param route - The route to fetch data from. 246 | * @param options - The options for the fetch request. 247 | * @returns A promise that resolves to the response from the API. 248 | */ 249 | public async fetchApi(route: string, options?: FetchOptions): Promise { 250 | if (!options) { 251 | options = {}; 252 | } 253 | options.headers = { ...this.getCredentialHeaders() }; 254 | options.mode = "cors"; 255 | return fetch(this.apiURL(route), options); 256 | } 257 | 258 | /** 259 | * Polls the status for colab and other things that don't support websockets. 260 | * @returns {Promise} The status information. 261 | */ 262 | async pollStatus(timeout = 1000): Promise { 263 | const controller = new AbortController(); 264 | const timeoutId = setTimeout(() => controller.abort(), timeout); 265 | try { 266 | const response = await this.fetchApi("/prompt", { signal: controller.signal }); 267 | if (response.status === 200) { 268 | return response.json(); 269 | } else { 270 | throw response; 271 | } 272 | } catch (error: any) { 273 | this.log("pollStatus", "Failed", error); 274 | if (error.name === "AbortError") { 275 | throw new Error("Request timed out"); 276 | } 277 | throw error; 278 | } finally { 279 | clearTimeout(timeoutId); 280 | } 281 | } 282 | 283 | /** 284 | * Queues a prompt for processing. 285 | * @param {number} number The index at which to queue the prompt. using NULL will append to the end of the queue. 286 | * @param {object} workflow Additional workflow data. 287 | * @returns {Promise} The response from the API. 288 | */ 289 | async queuePrompt(number: number | null, workflow: object): Promise { 290 | const body = { client_id: this.clientId, prompt: workflow } as any; 291 | 292 | if (number !== null) { 293 | if (number === -1) { 294 | body["front"] = true; 295 | } else if (number !== 0) { 296 | body["number"] = number; 297 | } 298 | } 299 | 300 | try { 301 | const response = await this.fetchApi("/prompt", { 302 | method: "POST", 303 | headers: { "Content-Type": "application/json" }, 304 | body: JSON.stringify(body) 305 | }); 306 | 307 | if (response.status !== 200) { 308 | throw { response }; 309 | } 310 | 311 | return response.json(); 312 | } catch (e) { 313 | this.log("queuePrompt", "Can't queue prompt", e); 314 | throw e.response as Response; 315 | } 316 | } 317 | 318 | /** 319 | * Appends a prompt to the workflow queue. 320 | * 321 | * @param {object} workflow Additional workflow data. 322 | * @returns {Promise} The response from the API. 323 | */ 324 | async appendPrompt(workflow: object): Promise { 325 | return this.queuePrompt(null, workflow).catch((e) => { 326 | this.dispatchEvent(new CustomEvent("queue_error")); 327 | throw e; 328 | }); 329 | } 330 | 331 | /** 332 | * Retrieves the current state of the queue. 333 | * @returns {Promise} The queue state. 334 | */ 335 | async getQueue(): Promise { 336 | const response = await this.fetchApi("/queue"); 337 | return response.json(); 338 | } 339 | 340 | /** 341 | * Retrieves the prompt execution history. 342 | * @param {number} [maxItems=200] The maximum number of items to retrieve. 343 | * @returns {Promise} The prompt execution history. 344 | */ 345 | async getHistories(maxItems: number = 200): Promise { 346 | const response = await this.fetchApi(`/history?max_items=${maxItems}`); 347 | return response.json(); 348 | } 349 | 350 | /** 351 | * Retrieves the history entry for a given prompt ID. 352 | * @param promptId - The ID of the prompt. 353 | * @returns A Promise that resolves to the HistoryEntry object. 354 | */ 355 | async getHistory(promptId: string): Promise { 356 | const response = await this.fetchApi(`/history/${promptId}`); 357 | const history: HistoryResponse = await response.json(); 358 | return history[promptId]; 359 | } 360 | 361 | /** 362 | * Retrieves system and device stats. 363 | * @returns {Promise} The system stats. 364 | */ 365 | async getSystemStats(): Promise { 366 | const response = await this.fetchApi("/system_stats"); 367 | return response.json(); 368 | } 369 | 370 | /** 371 | * Retrieves the terminal logs from the server. 372 | */ 373 | async getTerminalLogs(): Promise<{ entries: Array<{ t: string; m: string }>; size: { cols: number; rows: number } }> { 374 | const response = await this.fetchApi("/internal/logs/raw"); 375 | return response.json(); 376 | } 377 | 378 | /** 379 | * Sets the terminal subscription status. 380 | * Enable will subscribe to terminal logs from websocket. 381 | */ 382 | async setTerminalSubscription(subscribe: boolean) { 383 | // Set the terminal subscription status again if call again 384 | this.listenTerminal = subscribe; 385 | // Send the request to the server 386 | await this.fetchApi("/internal/logs/subscribe", { 387 | method: "PATCH", 388 | headers: { "Content-Type": "application/json" }, 389 | body: JSON.stringify({ clientId: this.clientId, enabled: subscribe }) 390 | }); 391 | } 392 | 393 | /** 394 | * Retrieves a list of extension URLs. 395 | * @returns {Promise} A list of extension URLs. 396 | */ 397 | async getExtensions(): Promise { 398 | const response = await this.fetchApi("/extensions"); 399 | return response.json(); 400 | } 401 | 402 | /** 403 | * Retrieves a list of embedding names. 404 | * @returns {Promise} A list of embedding names. 405 | */ 406 | async getEmbeddings(): Promise { 407 | const response = await this.fetchApi("/embeddings"); 408 | return response.json(); 409 | } 410 | 411 | /** 412 | * Retrieves the checkpoints from the server. 413 | * @returns A promise that resolves to an array of strings representing the checkpoints. 414 | */ 415 | async getCheckpoints(): Promise { 416 | const nodeInfo = await this.getNodeDefs(LOAD_CHECKPOINTS_EXTENSION); 417 | if (!nodeInfo) return []; 418 | const output = nodeInfo[LOAD_CHECKPOINTS_EXTENSION].input.required?.ckpt_name?.[0]; 419 | if (!output) return []; 420 | return output as string[]; 421 | } 422 | 423 | /** 424 | * Retrieves the Loras from the node definitions. 425 | * @returns A Promise that resolves to an array of strings representing the Loras. 426 | */ 427 | async getLoras(): Promise { 428 | const nodeInfo = await this.getNodeDefs(LOAD_LORAS_EXTENSION); 429 | if (!nodeInfo) return []; 430 | const output = nodeInfo[LOAD_LORAS_EXTENSION].input.required?.lora_name?.[0]; 431 | if (!output) return []; 432 | return output as string[]; 433 | } 434 | 435 | /** 436 | * Retrieves the sampler information. 437 | * @returns An object containing the sampler and scheduler information. 438 | */ 439 | async getSamplerInfo() { 440 | const nodeInfo = await this.getNodeDefs(LOAD_KSAMPLER_EXTENSION); 441 | if (!nodeInfo) return {}; 442 | return { 443 | sampler: nodeInfo[LOAD_KSAMPLER_EXTENSION].input.required.sampler_name ?? [], 444 | scheduler: nodeInfo[LOAD_KSAMPLER_EXTENSION].input.required.scheduler ?? [] 445 | }; 446 | } 447 | 448 | /** 449 | * Retrieves node object definitions for the graph. 450 | * @returns {Promise} The node definitions. 451 | */ 452 | async getNodeDefs(nodeName?: string): Promise { 453 | const response = await this.fetchApi(`/object_info${nodeName ? `/${nodeName}` : ""}`); 454 | const result = await response.json(); 455 | if (Object.keys(result).length === 0) { 456 | return null; 457 | } 458 | return result; 459 | } 460 | 461 | /** 462 | * Retrieves user configuration data. 463 | * @returns {Promise} The user configuration data. 464 | */ 465 | async getUserConfig(): Promise { 466 | const response = await this.fetchApi("/users"); 467 | return response.json(); 468 | } 469 | 470 | /** 471 | * Creates a new user. 472 | * @param {string} username The username of the new user. 473 | * @returns {Promise} The response from the API. 474 | */ 475 | async createUser(username: string): Promise { 476 | const response = await this.fetchApi("/users", { 477 | method: "POST", 478 | headers: { "Content-Type": "application/json" }, 479 | body: JSON.stringify({ username }) 480 | }); 481 | return response; 482 | } 483 | 484 | /** 485 | * Retrieves all setting values for the current user. 486 | * @returns {Promise} A dictionary of setting id to value. 487 | */ 488 | async getSettings(): Promise { 489 | const response = await this.fetchApi("/settings"); 490 | return response.json(); 491 | } 492 | 493 | /** 494 | * Retrieves a specific setting for the current user. 495 | * @param {string} id The id of the setting to fetch. 496 | * @returns {Promise} The setting value. 497 | */ 498 | async getSetting(id: string): Promise { 499 | const response = await this.fetchApi(`/settings/${encodeURIComponent(id)}`); 500 | return response.json(); 501 | } 502 | 503 | /** 504 | * Stores a dictionary of settings for the current user. 505 | * @param {Record} settings Dictionary of setting id to value to save. 506 | * @returns {Promise} 507 | */ 508 | async storeSettings(settings: Record): Promise { 509 | await this.fetchApi(`/settings`, { 510 | method: "POST", 511 | headers: { "Content-Type": "application/json" }, 512 | body: JSON.stringify(settings) 513 | }); 514 | } 515 | 516 | /** 517 | * Stores a specific setting for the current user. 518 | * @param {string} id The id of the setting to update. 519 | * @param {unknown} value The value of the setting. 520 | * @returns {Promise} 521 | */ 522 | async storeSetting(id: string, value: unknown): Promise { 523 | await this.fetchApi(`/settings/${encodeURIComponent(id)}`, { 524 | method: "POST", 525 | headers: { "Content-Type": "application/json" }, 526 | body: JSON.stringify(value) 527 | }); 528 | } 529 | 530 | /** 531 | * Uploads an image file to the server. 532 | * @param file - The image file to upload. 533 | * @param fileName - The name of the image file. 534 | * @param override - Optional. Specifies whether to override an existing file with the same name. Default is true. 535 | * @returns A Promise that resolves to an object containing the image information and the URL of the uploaded image, 536 | * or false if the upload fails. 537 | */ 538 | async uploadImage( 539 | file: ArrayBuffer | Uint8Array | Blob, 540 | fileName: string, 541 | config?: { override?: boolean; subfolder?: string } 542 | ): Promise<{ info: ImageInfo; url: string } | false> { 543 | const formData = new FormData(); 544 | if (file instanceof ArrayBuffer || file instanceof Uint8Array) { 545 | formData.append("image", new Blob([file as BlobPart]), fileName); 546 | } else { 547 | formData.append("image", file as Blob, fileName); 548 | } 549 | formData.append("subfolder", config?.subfolder ?? ""); 550 | formData.append("overwrite", config?.override?.toString() ?? "false"); 551 | 552 | try { 553 | const response = await this.fetchApi("/upload/image", { method: "POST", body: formData }); 554 | const imgInfo = await response.json(); 555 | const mapped = { ...imgInfo, filename: imgInfo.name }; 556 | 557 | // Check if the response is successful 558 | if (!response.ok) { 559 | this.log("uploadImage", "Upload failed", response); 560 | return false; 561 | } 562 | 563 | return { info: mapped, url: this.getPathImage(mapped) }; 564 | } catch (e) { 565 | this.log("uploadImage", "Upload failed", e); 566 | return false; 567 | } 568 | } 569 | 570 | /** 571 | * Uploads a mask file to the server. 572 | * 573 | * @param file - The mask file to upload, can be an ArrayBuffer, Uint8Array, or Blob. 574 | * @param originalRef - The original reference information for the file. 575 | * @returns A Promise that resolves to an object containing the image info and URL if the upload is successful, or false if the upload fails. 576 | */ 577 | async uploadMask(file: ArrayBuffer | Uint8Array | Blob, originalRef: ImageInfo): Promise<{ info: ImageInfo; url: string } | false> { 578 | const formData = new FormData(); 579 | 580 | // Append the image file to the form data 581 | if (file instanceof ArrayBuffer || file instanceof Uint8Array) { 582 | formData.append("image", new Blob([file as BlobPart]), "mask.png"); 583 | } else { 584 | formData.append("image", file as Blob, "mask.png"); 585 | } 586 | 587 | // Append the original reference as a JSON string 588 | formData.append("original_ref", JSON.stringify(originalRef)); 589 | 590 | try { 591 | // Send the POST request to the /upload/mask endpoint 592 | const response = await this.fetchApi("/upload/mask", { method: "POST", body: formData }); 593 | 594 | // Check if the response is successful 595 | if (!response.ok) { 596 | this.log("uploadMask", "Upload failed", response); 597 | return false; 598 | } 599 | 600 | const imgInfo = await response.json(); 601 | const mapped = { ...imgInfo, filename: imgInfo.name }; 602 | return { info: mapped, url: this.getPathImage(mapped) }; 603 | } catch (error) { 604 | this.log("uploadMask", "Upload failed", error); 605 | return false; 606 | } 607 | } 608 | 609 | /** 610 | * Frees memory by unloading models and freeing memory. 611 | * 612 | * @param unloadModels - A boolean indicating whether to unload models. 613 | * @param freeMemory - A boolean indicating whether to free memory. 614 | * @returns A promise that resolves to a boolean indicating whether the memory was successfully freed. 615 | */ 616 | async freeMemory(unloadModels: boolean, freeMemory: boolean): Promise { 617 | const payload = { unload_models: unloadModels, free_memory: freeMemory }; 618 | 619 | try { 620 | const response = await this.fetchApi("/free", { 621 | method: "POST", 622 | headers: { "Content-Type": "application/json" }, 623 | body: JSON.stringify(payload) 624 | }); 625 | 626 | // Check if the response is successful 627 | if (!response.ok) { 628 | this.log("freeMemory", "Free memory failed", response); 629 | return false; 630 | } 631 | 632 | // Return the response object 633 | return true; 634 | } catch (error) { 635 | this.log("freeMemory", "Free memory failed", error); 636 | return false; 637 | } 638 | } 639 | 640 | /** 641 | * Returns the path to an image based on the provided image information. 642 | * @param imageInfo - The information of the image. 643 | * @returns The path to the image. 644 | */ 645 | getPathImage(imageInfo: ImageInfo): string { 646 | return this.apiURL( 647 | `/view?filename=${imageInfo.filename}&type=${imageInfo.type}&subfolder=${imageInfo.subfolder ?? ""}` 648 | ); 649 | } 650 | 651 | /** 652 | * Get blob of image based on the provided image information. Use when the server have credential. 653 | */ 654 | async getImage(imageInfo: ImageInfo): Promise { 655 | return this.fetchApi( 656 | `/view?filename=${imageInfo.filename}&type=${imageInfo.type}&subfolder=${imageInfo.subfolder ?? ""}` 657 | ).then((res) => res.blob()); 658 | } 659 | 660 | /** 661 | * Retrieves a user data file for the current user. 662 | * @param {string} file The name of the userdata file to load. 663 | * @returns {Promise} The fetch response object. 664 | */ 665 | async getUserData(file: string): Promise { 666 | return this.fetchApi(`/userdata/${encodeURIComponent(file)}`); 667 | } 668 | 669 | /** 670 | * Stores a user data file for the current user. 671 | * @param {string} file The name of the userdata file to save. 672 | * @param {unknown} data The data to save to the file. 673 | * @param {RequestInit & { overwrite?: boolean, stringify?: boolean, throwOnError?: boolean }} [options] Additional options for storing the file. 674 | * @returns {Promise} 675 | */ 676 | async storeUserData( 677 | file: string, 678 | data: unknown, 679 | options: RequestInit & { overwrite?: boolean; stringify?: boolean; throwOnError?: boolean } = { 680 | overwrite: true, 681 | stringify: true, 682 | throwOnError: true 683 | } 684 | ): Promise { 685 | const response = await this.fetchApi(`/userdata/${encodeURIComponent(file)}?overwrite=${options.overwrite}`, { 686 | method: "POST", 687 | headers: { "Content-Type": options.stringify ? "application/json" : "application/octet-stream" } as any, 688 | body: options.stringify ? JSON.stringify(data) : (data as any), 689 | ...options 690 | }); 691 | 692 | if (response.status !== 200 && options.throwOnError !== false) { 693 | this.log("storeUserData", "Error storing user data file", response); 694 | throw new Error(`Error storing user data file '${file}': ${response.status} ${response.statusText}`); 695 | } 696 | 697 | return response; 698 | } 699 | 700 | /** 701 | * Deletes a user data file for the current user. 702 | * @param {string} file The name of the userdata file to delete. 703 | * @returns {Promise} 704 | */ 705 | async deleteUserData(file: string): Promise { 706 | const response = await this.fetchApi(`/userdata/${encodeURIComponent(file)}`, { method: "DELETE" }); 707 | 708 | if (response.status !== 204) { 709 | this.log("deleteUserData", "Error deleting user data file", response); 710 | throw new Error(`Error removing user data file '${file}': ${response.status} ${response.statusText}`); 711 | } 712 | } 713 | 714 | /** 715 | * Moves a user data file for the current user. 716 | * @param {string} source The userdata file to move. 717 | * @param {string} dest The destination for the file. 718 | * @param {RequestInit & { overwrite?: boolean }} [options] Additional options for moving the file. 719 | * @returns {Promise} 720 | */ 721 | async moveUserData( 722 | source: string, 723 | dest: string, 724 | options: RequestInit & { overwrite?: boolean } = { overwrite: false } 725 | ): Promise { 726 | return this.fetchApi( 727 | `/userdata/${encodeURIComponent(source)}/move/${encodeURIComponent(dest)}?overwrite=${options.overwrite}`, 728 | { method: "POST" } 729 | ); 730 | } 731 | 732 | /** 733 | * Lists user data files for the current user. 734 | * @param {string} dir The directory in which to list files. 735 | * @param {boolean} [recurse] If the listing should be recursive. 736 | * @param {boolean} [split] If the paths should be split based on the OS path separator. 737 | * @returns {Promise} The list of files. 738 | */ 739 | async listUserData(dir: string, recurse?: boolean, split?: boolean): Promise { 740 | const response = await this.fetchApi( 741 | `/userdata?${new URLSearchParams({ dir, recurse: recurse?.toString() ?? "", split: split?.toString() ?? "" })}` 742 | ); 743 | 744 | if (response.status === 404) return []; 745 | if (response.status !== 200) { 746 | this.log("listUserData", "Error getting user data list", response); 747 | throw new Error(`Error getting user data list '${dir}': ${response.status} ${response.statusText}`); 748 | } 749 | 750 | return response.json(); 751 | } 752 | 753 | /** 754 | * Interrupts the execution of the running prompt. 755 | * @returns {Promise} 756 | */ 757 | async interrupt(): Promise { 758 | await this.fetchApi("/interrupt", { method: "POST" }); 759 | } 760 | 761 | /** 762 | * Initializes the client. 763 | * 764 | * @param maxTries - The maximum number of ping tries. 765 | * @param delayTime - The delay time between ping tries in milliseconds. 766 | * @returns The initialized client instance. 767 | */ 768 | init(maxTries = 10, delayTime = 1000) { 769 | this.pingSuccess(maxTries, delayTime) 770 | .then(() => { 771 | /** 772 | * Get system OS type on initialization. 773 | */ 774 | this.pullOsType(); 775 | /** 776 | * Test features on initialization. 777 | */ 778 | this.testFeatures(); 779 | /** 780 | * Create WebSocket connection on initialization. 781 | */ 782 | this.createSocket(); 783 | /** 784 | * Set terminal subscription on initialization. 785 | */ 786 | this.setTerminalSubscription(this.listenTerminal); 787 | }) 788 | .catch((e) => { 789 | this.log("init", "Failed", e); 790 | this.dispatchEvent(new CustomEvent("connection_error", { detail: e })); 791 | }); 792 | return this; 793 | } 794 | 795 | private async pingSuccess(maxTries = 10, delayTime = 1000) { 796 | let tries = 0; 797 | let ping = await this.ping(); 798 | while (!ping.status) { 799 | if (tries > maxTries) { 800 | throw new Error("Can't connect to the server"); 801 | } 802 | await delay(delayTime); // Wait for 1s before trying again 803 | ping = await this.ping(); 804 | tries++; 805 | } 806 | } 807 | 808 | async waitForReady() { 809 | while (!this.isReady) { 810 | await delay(100); 811 | } 812 | return this; 813 | } 814 | 815 | private async pullOsType() { 816 | this.getSystemStats().then((data) => { 817 | this.osType = data.system.os; 818 | }); 819 | } 820 | 821 | /** 822 | * Sends a ping request to the server and returns a boolean indicating whether the server is reachable. 823 | * @returns A promise that resolves to `true` if the server is reachable, or `false` otherwise. 824 | */ 825 | async ping() { 826 | const start = performance.now(); 827 | return this.pollStatus(5000) 828 | .then(() => { 829 | return { status: true, time: performance.now() - start } as const; 830 | }) 831 | .catch((error) => { 832 | this.log("ping", "Can't connect to the server", error); 833 | return { status: false } as const; 834 | }); 835 | } 836 | 837 | /** 838 | * Attempts to reconnect the WebSocket with an exponential backoff strategy 839 | * @param triggerEvent Whether to trigger disconnect/reconnect events 840 | */ 841 | public async reconnectWs(triggerEvent?: boolean) { 842 | if (triggerEvent) { 843 | this.dispatchEvent(new CustomEvent("disconnected")); 844 | this.dispatchEvent(new CustomEvent("reconnecting")); 845 | } 846 | 847 | // Maximum number of reconnection attempts 848 | const MAX_ATTEMPTS = 10; 849 | // Base delay in milliseconds 850 | const BASE_DELAY = 1000; 851 | // Maximum delay between attempts (15 seconds) 852 | const MAX_DELAY = 15000; 853 | 854 | let attempt = 0; 855 | 856 | const tryReconnect = () => { 857 | attempt++; 858 | this.log("socket", `WebSocket reconnection attempt #${attempt}`); 859 | 860 | // Clean up any existing socket 861 | if (this.socket?.client) { 862 | try { 863 | this.socket.close(); 864 | } catch (error) { 865 | this.log("socket", "Error while closing previous socket", error); 866 | } 867 | } 868 | 869 | this.socket = null; 870 | 871 | // Create new socket connection 872 | try { 873 | this.createSocket(true); 874 | } catch (error) { 875 | this.log("socket", "Error creating socket during reconnect", error); 876 | } 877 | 878 | // Calculate next retry delay with exponential backoff and jitter 879 | if (attempt < MAX_ATTEMPTS) { 880 | // Exponential backoff formula: baseDelay * 2^attempt + random jitter 881 | const exponentialDelay = Math.min(BASE_DELAY * Math.pow(2, attempt - 1), MAX_DELAY); 882 | 883 | // Add jitter (±30% of the delay) to prevent all clients reconnecting simultaneously 884 | const jitter = exponentialDelay * 0.3 * (Math.random() - 0.5); 885 | const delay = Math.max(1000, exponentialDelay + jitter); 886 | 887 | this.log("socket", `Will retry in ${Math.round(delay / 1000)} seconds`); 888 | 889 | // Check if the socket is reconnected within the timeout 890 | setTimeout(() => { 891 | if ( 892 | !this.socket?.client || 893 | (this.socket.client.readyState !== WebSocket.OPEN && this.socket.client.readyState !== WebSocket.CONNECTING) 894 | ) { 895 | this.log("socket", "Reconnection failed or timed out, retrying..."); 896 | tryReconnect(); // Retry if not connected 897 | } else { 898 | this.log("socket", "Reconnection successful"); 899 | } 900 | }, delay); 901 | } else { 902 | this.log("socket", `Maximum reconnection attempts (${MAX_ATTEMPTS}) reached.`); 903 | this.dispatchEvent(new CustomEvent("reconnection_failed")); 904 | } 905 | }; 906 | 907 | tryReconnect(); 908 | } 909 | 910 | private resetLastActivity() { 911 | this.lastActivity = Date.now(); 912 | } 913 | 914 | /** 915 | * Creates and connects a WebSocket for real-time updates. 916 | * @param {boolean} isReconnect If the socket connection is a reconnect attempt. 917 | */ 918 | /** 919 | * Creates and connects a WebSocket for real-time updates. 920 | * Falls back to polling if WebSocket is unavailable. 921 | * @param {boolean} isReconnect If the socket connection is a reconnect attempt. 922 | */ 923 | private createSocket(isReconnect: boolean = false) { 924 | let reconnecting = false; 925 | let usePolling = false; 926 | 927 | if (this.socket) { 928 | this.log("socket", "Socket already exists, skipping creation."); 929 | return; 930 | } 931 | 932 | const headers = { ...this.getCredentialHeaders() }; 933 | 934 | const existingSession = `?clientId=${this.clientId}`; 935 | const wsUrl = `ws${this.apiHost.includes("https:") ? "s" : ""}://${this.apiBase}/ws${existingSession}`; 936 | 937 | // Try to create WebSocket connection 938 | try { 939 | this.socket = new WebSocketClient(wsUrl, { headers, customWebSocketImpl: this.customWsImpl }); 940 | 941 | this.socket.client.onclose = () => { 942 | if (reconnecting || isReconnect) return; 943 | reconnecting = true; 944 | this.log("socket", "Socket closed -> Reconnecting"); 945 | this.reconnectWs(true); 946 | }; 947 | 948 | this.socket.client.onopen = () => { 949 | this.resetLastActivity(); 950 | reconnecting = false; 951 | usePolling = false; // Reset polling flag if we have an open connection 952 | this.log("socket", "Socket opened"); 953 | if (isReconnect) { 954 | this.dispatchEvent(new CustomEvent("reconnected")); 955 | } else { 956 | this.dispatchEvent(new CustomEvent("connected")); 957 | } 958 | }; 959 | } catch (error) { 960 | this.log("socket", "WebSocket creation failed, falling back to polling", error); 961 | this.socket = null; 962 | usePolling = true; 963 | this.dispatchEvent(new CustomEvent("websocket_unavailable", { detail: error })); 964 | 965 | // Set up polling mechanism 966 | this.setupPollingFallback(); 967 | } 968 | 969 | // Only continue with WebSocket setup if creation was successful 970 | if (this.socket?.client) { 971 | this.socket.client.onmessage = (event) => { 972 | this.resetLastActivity(); 973 | try { 974 | if (event.data instanceof ArrayBuffer) { 975 | const buffer = event.data; 976 | const view = new DataView(buffer); 977 | const eventType = view.getUint32(0); 978 | switch (eventType) { 979 | case 1: 980 | const imageType = view.getUint32(0); 981 | let imageMime; 982 | switch (imageType) { 983 | case 1: 984 | default: 985 | imageMime = "image/jpeg"; 986 | break; 987 | case 2: 988 | imageMime = "image/png"; 989 | } 990 | const imageBlob = new Blob([buffer.slice(8)], { type: imageMime }); 991 | this.dispatchEvent(new CustomEvent("b_preview", { detail: imageBlob })); 992 | break; 993 | default: 994 | throw new Error(`Unknown binary websocket message of type ${eventType}`); 995 | } 996 | } else if (typeof event.data === "string") { 997 | const msg = JSON.parse(event.data); 998 | if (!msg.data || !msg.type) return; 999 | this.dispatchEvent(new CustomEvent("all", { detail: msg })); 1000 | if (msg.type === "logs") { 1001 | this.dispatchEvent(new CustomEvent("terminal", { detail: msg.data.entries?.[0] || null })); 1002 | } else { 1003 | this.dispatchEvent(new CustomEvent(msg.type, { detail: msg.data })); 1004 | } 1005 | if (msg.data.sid) { 1006 | this.clientId = msg.data.sid; 1007 | } 1008 | } else { 1009 | this.log("socket", "Unhandled message", event); 1010 | } 1011 | } catch (error) { 1012 | this.log("socket", "Unhandled message", { event, error }); 1013 | } 1014 | }; 1015 | 1016 | this.socket.client.onerror = (e) => { 1017 | this.log("socket", "Socket error", e); 1018 | 1019 | // If this is the first error and we're not already in reconnect mode 1020 | if (!reconnecting && !usePolling) { 1021 | usePolling = true; 1022 | this.log("socket", "WebSocket error, will try polling as fallback"); 1023 | this.setupPollingFallback(); 1024 | } 1025 | }; 1026 | 1027 | if (!isReconnect) { 1028 | this.wsTimer = setInterval(() => { 1029 | if (reconnecting) return; 1030 | if (Date.now() - this.lastActivity > this.wsTimeout) { 1031 | reconnecting = true; 1032 | this.log("socket", "Connection timed out, reconnecting..."); 1033 | this.reconnectWs(true); 1034 | } 1035 | }, this.wsTimeout / 2); 1036 | } 1037 | } 1038 | } 1039 | 1040 | /** 1041 | * Sets up a polling mechanism as a fallback when WebSockets are unavailable 1042 | * Polls the server every 2 seconds for status updates 1043 | */ 1044 | /** 1045 | * Sets up a polling mechanism as a fallback when WebSockets are unavailable 1046 | * Polls the server every 2 seconds for status updates 1047 | */ 1048 | private setupPollingFallback() { 1049 | this.log("socket", "Setting up polling fallback mechanism"); 1050 | 1051 | // Clear any existing polling timer 1052 | if (this._pollingTimer) { 1053 | try { 1054 | clearInterval(this._pollingTimer as any); 1055 | this._pollingTimer = null; 1056 | } catch (e) { 1057 | this.log("socket", "Error clearing polling timer", e); 1058 | } 1059 | } 1060 | 1061 | // Poll every 2 seconds 1062 | const POLLING_INTERVAL = 2000; 1063 | 1064 | const pollFn = async () => { 1065 | try { 1066 | // Poll execution status 1067 | const status = await this.pollStatus(); 1068 | 1069 | // Simulate an event dispatch similar to WebSocket 1070 | this.dispatchEvent(new CustomEvent("status", { detail: status })); 1071 | 1072 | // Reset activity timestamp to prevent timeout 1073 | this.resetLastActivity(); 1074 | 1075 | // Try to re-establish WebSocket connection periodically 1076 | if (!this.socket || !this.socket.client || this.socket.client.readyState !== WebSocket.OPEN) { 1077 | this.log("socket", "Attempting to restore WebSocket connection"); 1078 | try { 1079 | this.createSocket(true); 1080 | } catch (error) { 1081 | // Continue with polling if WebSocket creation fails 1082 | this.log("socket", "WebSocket still unavailable, continuing with polling", error); 1083 | } 1084 | } else { 1085 | // WebSocket is back, we can stop polling 1086 | this.log("socket", "WebSocket connection restored, stopping polling"); 1087 | if (this._pollingTimer) { 1088 | clearInterval(this._pollingTimer as any); 1089 | this._pollingTimer = null; 1090 | } 1091 | } 1092 | } catch (error) { 1093 | this.log("socket", "Polling error", error); 1094 | } 1095 | }; 1096 | 1097 | // Using setInterval and casting to the expected type 1098 | this._pollingTimer = setInterval(pollFn, POLLING_INTERVAL) as any; 1099 | 1100 | this.log("socket", `Polling started with interval of ${POLLING_INTERVAL}ms`); 1101 | } 1102 | 1103 | /** 1104 | * Retrieves a list of all available model folders. 1105 | * @experimental API that may change in future versions 1106 | * @returns A promise that resolves to an array of ModelFolder objects. 1107 | */ 1108 | async getModelFolders(): Promise { 1109 | try { 1110 | const response = await this.fetchApi("/experiment/models"); 1111 | if (!response.ok) { 1112 | this.log("getModelFolders", "Failed to fetch model folders", response); 1113 | throw new Error(`Failed to fetch model folders: ${response.status} ${response.statusText}`); 1114 | } 1115 | return response.json(); 1116 | } catch (error) { 1117 | this.log("getModelFolders", "Error fetching model folders", error); 1118 | throw error; 1119 | } 1120 | } 1121 | 1122 | /** 1123 | * Retrieves a list of all model files in a specific folder. 1124 | * @experimental API that may change in future versions 1125 | * @param folder - The name of the model folder. 1126 | * @returns A promise that resolves to an array of ModelFile objects. 1127 | */ 1128 | async getModelFiles(folder: string): Promise { 1129 | try { 1130 | const response = await this.fetchApi(`/experiment/models/${encodeURIComponent(folder)}`); 1131 | if (!response.ok) { 1132 | this.log("getModelFiles", "Failed to fetch model files", { folder, response }); 1133 | throw new Error(`Failed to fetch model files: ${response.status} ${response.statusText}`); 1134 | } 1135 | return response.json(); 1136 | } catch (error) { 1137 | this.log("getModelFiles", "Error fetching model files", { folder, error }); 1138 | throw error; 1139 | } 1140 | } 1141 | 1142 | /** 1143 | * Retrieves a preview image for a specific model file. 1144 | * @experimental API that may change in future versions 1145 | * @param folder - The name of the model folder. 1146 | * @param pathIndex - The index of the folder path where the file is stored. 1147 | * @param filename - The name of the model file. 1148 | * @returns A promise that resolves to a ModelPreviewResponse object containing the preview image data. 1149 | */ 1150 | async getModelPreview(folder: string, pathIndex: number, filename: string): Promise { 1151 | try { 1152 | const response = await this.fetchApi( 1153 | `/experiment/models/preview/${encodeURIComponent(folder)}/${pathIndex}/${encodeURIComponent(filename)}` 1154 | ); 1155 | 1156 | if (!response.ok) { 1157 | this.log("getModelPreview", "Failed to fetch model preview", { folder, pathIndex, filename, response }); 1158 | throw new Error(`Failed to fetch model preview: ${response.status} ${response.statusText}`); 1159 | } 1160 | 1161 | const contentType = response.headers.get("content-type") || "image/webp"; 1162 | const body = await response.arrayBuffer(); 1163 | 1164 | return { body, contentType }; 1165 | } catch (error) { 1166 | this.log("getModelPreview", "Error fetching model preview", { folder, pathIndex, filename, error }); 1167 | throw error; 1168 | } 1169 | } 1170 | 1171 | /** 1172 | * Creates a URL for a model preview image. 1173 | * @experimental API that may change in future versions 1174 | * @param folder - The name of the model folder. 1175 | * @param pathIndex - The index of the folder path where the file is stored. 1176 | * @param filename - The name of the model file. 1177 | * @returns The URL string for the model preview. 1178 | */ 1179 | getModelPreviewUrl(folder: string, pathIndex: number, filename: string): string { 1180 | return this.apiURL( 1181 | `/experiment/models/preview/${encodeURIComponent(folder)}/${pathIndex}/${encodeURIComponent(filename)}` 1182 | ); 1183 | } 1184 | } 1185 | --------------------------------------------------------------------------------