diff --git a/README.md b/README.md index 4819159..a9d5c31 100644 --- a/README.md +++ b/README.md @@ -50,57 +50,177 @@ This application uses environment variables for configuration. Below are the ava Models are automatically detected from the `MODEL_DIR`. Each subdirectory in `MODEL_DIR` is considered a model category. The application creates an enumeration of all files in each category, which can be used for validation in the application. -### Workflow Models - -The `WORKFLOW_MODELS` variable determines which workflow endpoints are available. -By default, it's set to "all", including all base model categories -If you want to only include models from a specific base model category, specify them in a comma separated list. -The available options are `sd1.5`, `sdxl`, and `flux`. -To specify stable diffusion 1.5 and stable diffusion xl workflows, you can set `WORKFLOW_MODELS` to `sd1.5,sdxl`. - ## Generating New Workflow Template Endpoints Since the ComfyUI prompt format is a little obtuse, it's common to wrap the workflow endpoints with a more user-friendly interface. -This can be done by following the pattern established in the `src/workflows` directory. +This can be done by adding conforming `.js` or `.ts` files to the `/workflows` directory in your dockerfile. + +Here is an example text-to-image workflow file. + +```typescript +import { z } from "zod"; +import { ComfyNode, Workflow } from "../types"; +import config from "../config"; + +let checkpoint: any = config.models.checkpoints.enum.optional(); +if (config.warmupCkpt) { + checkpoint = checkpoint.default(config.warmupCkpt); +} + +const RequestSchema = z.object({ + prompt: z.string().describe("The positive prompt for image generation"), + negative_prompt: z + .string() + .optional() + .default("text, watermark") + .describe("The negative prompt for image generation"), + width: z + .number() + .int() + .min(256) + .max(2048) + .optional() + .default(512) + .describe("Width of the generated image"), + height: z + .number() + .int() + .min(256) + .max(2048) + .optional() + .default(512) + .describe("Height of the generated image"), + seed: z + .number() + .int() + .optional() + .default(() => Math.floor(Math.random() * 100000000000)) + .describe("Seed for random number generation"), + steps: z + .number() + .int() + .min(1) + .max(100) + .optional() + .default(20) + .describe("Number of sampling steps"), + cfg_scale: z + .number() + .min(0) + .max(20) + .optional() + .default(8) + .describe("Classifier-free guidance scale"), + sampler_name: config.samplers + .optional() + .default("euler") + .describe("Name of the sampler to use"), + scheduler: config.schedulers + .optional() + .default("normal") + .describe("Type of scheduler to use"), + denoise: z + .number() + .min(0) + .max(1) + .optional() + .default(1) + .describe("Denoising strength"), + checkpoint, +}); + +type InputType = z.infer; + +function generateWorkflow(input: InputType): Record { + return { + "3": { + inputs: { + seed: input.seed, + steps: input.steps, + cfg: input.cfg_scale, + sampler_name: input.sampler_name, + scheduler: input.scheduler, + denoise: input.denoise, + model: ["4", 0], + positive: ["6", 0], + negative: ["7", 0], + latent_image: ["5", 0], + }, + class_type: "KSampler", + _meta: { + title: "KSampler", + }, + }, + "4": { + inputs: { + ckpt_name: input.checkpoint, + }, + class_type: "CheckpointLoaderSimple", + _meta: { + title: "Load Checkpoint", + }, + }, + "5": { + inputs: { + width: input.width, + height: input.height, + batch_size: 1, + }, + class_type: "EmptyLatentImage", + _meta: { + title: "Empty Latent Image", + }, + }, + "6": { + inputs: { + text: input.prompt, + clip: ["4", 1], + }, + class_type: "CLIPTextEncode", + _meta: { + title: "CLIP Text Encode (Prompt)", + }, + }, + "7": { + inputs: { + text: input.negative_prompt, + clip: ["4", 1], + }, + class_type: "CLIPTextEncode", + _meta: { + title: "CLIP Text Encode (Prompt)", + }, + }, + "8": { + inputs: { + samples: ["3", 0], + vae: ["4", 2], + }, + class_type: "VAEDecode", + _meta: { + title: "VAE Decode", + }, + }, + "9": { + inputs: { + filename_prefix: "ComfyUI", + images: ["8", 0], + }, + class_type: "SaveImage", + _meta: { + title: "Save Image", + }, + }, + }; +} + +const workflow: Workflow = { + RequestSchema, + generateWorkflow, +}; + +export default workflow; ``` -. -├── flux -│ ├── img2img.json -│ ├── img2img.ts -│ ├── txt2img.json -│ └── txt2img.ts -├── index.ts -├── sd1.5 -│ ├── img2img.json -│ ├── img2img.ts -│ ├── txt2img.json -│ └── txt2img.ts -└── sdxl - ├── img2img.json - ├── img2img.ts - ├── txt2img-with-refiner.json - ├── txt2img-with-refiner.ts - ├── txt2img.json - └── txt2img.ts - -3 directories, 15 files -``` - -Within the top level "workflows" directory, there are subdirectories for each base model category. -Within each base model category, there are JSON and TypeScript files for each workflow template. -The JSON files contain the original prompt format, and the TypeScript files contain the logic for converting a simpler input into the original prompt format. -The JSON files are for reference, and are not bundled into the final artifact. -Finally, the new workflow templates must be imported and added to the `workflows` object in `src/workflows/index.ts`. -From here they will be automatically added to the server, and have swagger docs generated. - -Producing these workflow templates can be fully automated using [Claude 3.5 Sonnet](https://www.anthropic.com/). -A script is provided to do this, `generateWorkflow.ts`. - -```bash -# First, compile the typescript -npm run build -# Then, run the script -node dist/generateWorkflow.js -``` \ No newline at end of file +Note your file MUST export a `Workflow` object, which contains a `RequestSchema` and a `generateWorkflow` function. The `RequestSchema` is a zod schema that describes the input to the workflow, and the `generateWorkflow` function takes the input and returns a ComfyUI API-format prompt. \ No newline at end of file diff --git a/package-lock.json b/package-lock.json index c522ead..47f43ff 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "comfyui-wrapper", - "version": "1.3.2", + "version": "1.4.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "comfyui-wrapper", - "version": "1.3.2", + "version": "1.4.0", "license": "MIT", "dependencies": { "@fastify/swagger": "^8.15.0", diff --git a/package.json b/package.json index f032a0e..4593071 100644 --- a/package.json +++ b/package.json @@ -1,12 +1,12 @@ { "name": "comfyui-wrapper", - "version": "1.3.2", - "description": "Wraps comfyui to make it easier to use as a web service", + "version": "1.4.0", + "description": "Wraps comfyui to make it easier to use as a stateless web service", "main": "dist/src/index.js", "scripts": { "test": "echo \"Error: no test specified\" && exit 1", "build": "tsc", - "build-binary": "tsc && pkg -t node18-linux-x64 --out-path bin ." + "build-binary": "tsc && pkg -t node18-linux-x64 --out-path bin --public --no-bytecode ." }, "author": "Shawn Rushefsky", "license": "MIT", @@ -15,8 +15,7 @@ "@types/chokidar": "^2.1.3", "@types/node": "^20.12.7", "minimist": "^1.2.8", - "pkg": "^5.8.1", - "typescript": "^5.4.5" + "pkg": "^5.8.1" }, "bin": { "comfyui-wrapper": "dist/src/index.js" @@ -27,6 +26,7 @@ "chokidar": "^3.6.0", "fastify": "^4.26.2", "fastify-type-provider-zod": "^2.0.0", - "zod": "^3.23.8" + "zod": "^3.23.8", + "typescript": "^5.4.5" } } diff --git a/src/config.ts b/src/config.ts index 7fb0a74..96fc41b 100644 --- a/src/config.ts +++ b/src/config.ts @@ -16,8 +16,11 @@ const { MODEL_DIR = "/opt/ComfyUI/models", WARMUP_PROMPT_FILE, WORKFLOW_MODELS = "all", + WORKFLOW_DIR = "/workflows", } = process.env; +fs.mkdirSync(WORKFLOW_DIR, { recursive: true }); + const comfyURL = `http://${DIRECT_ADDRESS}:${COMFYUI_PORT_HOST}`; const selfURL = `http://localhost:${PORT}`; const port = parseInt(PORT, 10); @@ -49,27 +52,6 @@ interface ComfyDescription { schedulers: string[]; } -function getPythonCommand(): string { - let pythonCommand = execSync( - "source /opt/ai-dock/etc/environment.sh && which python3", - { - encoding: "utf-8", - } - ).trim(); - if (!pythonCommand) { - pythonCommand = execSync( - "source /opt/ai-dock/etc/environment.sh && which python", - { - encoding: "utf-8", - } - ).trim(); - } - if (!pythonCommand) { - throw new Error("Python not found"); - } - return pythonCommand; -} - function getComfyUIDescription(): ComfyDescription { const temptComfyFilePath = path.join( "/opt/ComfyUI", @@ -136,6 +118,7 @@ const config = { startupCheckMaxTries, outputDir: OUTPUT_DIR, inputDir: INPUT_DIR, + workflowDir: WORKFLOW_DIR, warmupPrompt, warmupCkpt, samplers: z.enum(comfyDescription.samplers as [string, ...string[]]), diff --git a/src/server.ts b/src/server.ts index 5378634..12f10e6 100644 --- a/src/server.ts +++ b/src/server.ts @@ -18,14 +18,16 @@ import { launchComfyUI, shutdownComfyUI, processImage, + zodToMarkdownTable, } from "./utils"; import { PromptRequestSchema, PromptErrorResponseSchema, PromptResponseSchema, PromptRequest, - Workflow, WorkflowResponseSchema, + WorkflowTree, + isWorkflow, } from "./types"; import workflows from "./workflows"; import { z } from "zod"; @@ -265,61 +267,127 @@ server.after(() => { } ); - for (const baseModel in workflows) { - for (const workflowId in workflows[baseModel]) { - const workflow = workflows[baseModel][workflowId] as Workflow; - server.log.info(`Registering workflow ${baseModel}/${workflowId}`); - const BodySchema = z.object({ - id: z - .string() - .optional() - .default(() => randomUUID()), - input: workflow.RequestSchema, - webhook: z.string().optional(), - }); - - type BodyType = z.infer; - - app.post<{ - Body: BodyType; - }>( - `/workflow/${baseModel}/${workflowId}`, - { - schema: { - body: BodySchema, - response: { - 200: WorkflowResponseSchema, - 202: WorkflowResponseSchema, + // Recursively build the route tree from workflows + const walk = (tree: WorkflowTree, route = "/workflow") => { + for (const key in tree) { + const node = tree[key]; + if (isWorkflow(node)) { + const BodySchema = z.object({ + id: z + .string() + .optional() + .default(() => randomUUID()), + input: node.RequestSchema, + webhook: z.string().optional(), + }); + + type BodyType = z.infer; + + const description = zodToMarkdownTable(node.RequestSchema); + + app.post<{ + Body: BodyType; + }>( + `${route}/${key}`, + { + schema: { + description, + body: BodySchema, + response: { + 200: WorkflowResponseSchema, + 202: WorkflowResponseSchema, + }, }, }, - }, - async (request, reply) => { - const { id, input, webhook } = request.body; - const prompt = workflow.generateWorkflow(input); - - const resp = await fetch( - `http://localhost:${config.wrapperPort}/prompt`, - { - method: "POST", - headers: { - "Content-Type": "application/json", - }, - body: JSON.stringify({ prompt, id, webhook }), + async (request, reply) => { + const { id, input, webhook } = request.body; + const prompt = node.generateWorkflow(input); + + const resp = await fetch( + `http://localhost:${config.wrapperPort}/prompt`, + { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ prompt, id, webhook }), + } + ); + const body = await resp.json(); + if (!resp.ok) { + return reply.code(resp.status).send(body); } - ); - const body = await resp.json(); - if (!resp.ok) { + + body.input = input; + body.prompt = prompt; + return reply.code(resp.status).send(body); } + ); - body.input = input; - body.prompt = prompt; - - return reply.code(resp.status).send(body); - } - ); + server.log.info(`Registered workflow ${route}/${key}`); + } else { + walk(node as WorkflowTree, `${route}/${key}`); + } } - } + }; + walk(workflows); + + // for (const baseModel in workflows) { + // for (const workflowId in workflows[baseModel]) { + // const workflow = workflows[baseModel][workflowId] as Workflow; + // server.log.info(`Registering workflow ${baseModel}/${workflowId}`); + // const BodySchema = z.object({ + // id: z + // .string() + // .optional() + // .default(() => randomUUID()), + // input: workflow.RequestSchema, + // webhook: z.string().optional(), + // }); + + // type BodyType = z.infer; + + // app.post<{ + // Body: BodyType; + // }>( + // `/workflow/${baseModel}/${workflowId}`, + // { + // schema: { + // body: BodySchema, + // response: { + // 200: WorkflowResponseSchema, + // 202: WorkflowResponseSchema, + // }, + // }, + // }, + // async (request, reply) => { + // const { id, input, webhook } = request.body; + // const prompt = workflow.generateWorkflow(input); + + // const resp = await fetch( + // `http://localhost:${config.wrapperPort}/prompt`, + // { + // method: "POST", + // headers: { + // "Content-Type": "application/json", + // }, + // body: JSON.stringify({ prompt, id, webhook }), + // } + // ); + // const body = await resp.json(); + // if (!resp.ok) { + // return reply.code(resp.status).send(body); + // } + + // body.input = input; + // body.prompt = prompt; + + // return reply.code(resp.status).send(body); + // } + // ); + // } + // } }); let warm = false; @@ -343,9 +411,7 @@ export async function start() { // Start the server await server.listen({ port: config.wrapperPort, host: config.wrapperHost }); - server.log.info( - `ComfyUI API ${version} listening on ${server.server.address}` - ); + server.log.info(`ComfyUI API ${version} started.`); await warmupComfyUI(); warm = true; const warmupTime = Date.now() - start; diff --git a/src/types.ts b/src/types.ts index 4068cd6..2aeb918 100644 --- a/src/types.ts +++ b/src/types.ts @@ -48,6 +48,14 @@ export interface Workflow { generateWorkflow: (input: any) => Record; } +export function isWorkflow(obj: any): obj is Workflow { + return "RequestSchema" in obj && "generateWorkflow" in obj; +} + +export interface WorkflowTree { + [key: string]: WorkflowTree | Workflow; +} + export const WorkflowRequestSchema = z.object({ id: z .string() diff --git a/src/utils.ts b/src/utils.ts index 2acdbf1..df633f8 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -3,11 +3,10 @@ import { FastifyBaseLogger } from "fastify"; import { CommandExecutor } from "./commands"; import fs from "fs"; import fsPromises from "fs/promises"; -import { spawn } from "child_process"; -import * as readline from "readline"; import { Readable } from "stream"; import path from "path"; import { randomUUID } from "crypto"; +import { ZodObject, ZodRawShape, ZodTypeAny, ZodDefault } from "zod"; const commandExecutor = new CommandExecutor(); @@ -126,3 +125,93 @@ export async function processImage( } } } + +export function zodToMarkdownTable(schema: ZodObject): string { + const shape = schema.shape; + let markdownTable = "| Field | Type | Description | Default |\n|-|-|-|-|\n"; + + for (const [key, value] of Object.entries(shape)) { + const fieldName = key; + const { type: fieldType, isOptional } = getZodTypeName(value); + const fieldDescription = getZodDescription(value); + const defaultValue = getZodDefault(value); + + markdownTable += `| ${fieldName} | ${fieldType}${ + isOptional ? "" : "" + } | ${fieldDescription} | ${defaultValue || "**Required**"} |\n`; + } + + return markdownTable; +} + +function getZodTypeName(zodType: ZodTypeAny): { + type: string; + isOptional: boolean; +} { + let currentType = zodType; + let isOptional = false; + + while (currentType instanceof ZodDefault) { + currentType = currentType._def.innerType; + } + + if (currentType._def.typeName === "ZodOptional") { + isOptional = true; + currentType = currentType._def.innerType; + } + + let type: string; + switch (currentType._def.typeName) { + case "ZodString": + type = "string"; + break; + case "ZodNumber": + type = "number"; + break; + case "ZodBoolean": + type = "boolean"; + break; + case "ZodArray": + type = `${getZodTypeName(currentType._def.type).type}[]`; + break; + case "ZodObject": + type = "object"; + break; + case "ZodEnum": + type = `enum (${(currentType._def.values as string[]) + .map((val: string) => `\`${val}\``) + .join(", ")})`; + break; + case "ZodUnion": + type = currentType._def.options + .map((opt: any) => getZodTypeName(opt).type) + .join(", "); + break; + case "ZodLiteral": + type = `literal (${JSON.stringify(currentType._def.value)})`; + break; + default: + type = currentType._def.typeName.replace("Zod", "").toLowerCase(); + } + + return { type, isOptional }; +} + +function getZodDescription(zodType: ZodTypeAny): string { + let currentType: ZodTypeAny | undefined = zodType; + while (currentType) { + if (currentType.description) { + return currentType.description; + } + currentType = currentType._def.innerType; + } + return ""; +} + +function getZodDefault(zodType: ZodTypeAny): string { + if (zodType instanceof ZodDefault) { + const defaultValue = zodType._def.defaultValue(); + return JSON.stringify(defaultValue); + } + return "-"; +} diff --git a/src/workflows/index.ts b/src/workflows/index.ts index 11ef991..70f801a 100644 --- a/src/workflows/index.ts +++ b/src/workflows/index.ts @@ -1,40 +1,58 @@ import config from "../config"; -import fluxTxt2img from "../workflows/flux/txt2img"; -import fluxImg2img from "../workflows/flux/img2img"; -import sd15Txt2img from "../workflows/sd1.5/txt2img"; -import sd15Img2img from "../workflows/sd1.5/img2img"; -import sdxlTxt2img from "../workflows/sdxl/txt2img"; -import sdxlImg2img from "../workflows/sdxl/img2img"; -import sdxlTxt2imgWithRefiner from "../workflows/sdxl/txt2img-with-refiner"; -import { Workflow } from "../types"; +// import fluxTxt2img from "../workflows/flux/txt2img"; +// import fluxImg2img from "../workflows/flux/img2img"; +// import sd15Txt2img from "../workflows/sd1.5/txt2img"; +// import sd15Img2img from "../workflows/sd1.5/img2img"; +// import sdxlTxt2img from "../workflows/sdxl/txt2img"; +// import sdxlImg2img from "../workflows/sdxl/img2img"; +// import sdxlTxt2imgWithRefiner from "../workflows/sdxl/txt2img-with-refiner"; +import { WorkflowTree, isWorkflow } from "../types"; +import fs from "fs"; +import ts from "typescript"; -const workflows: { - [key: string]: { - [key: string]: Workflow; - }; -} = { - flux: { - txt2img: fluxTxt2img, - img2img: fluxImg2img, - }, - "sd1.5": { - txt2img: sd15Txt2img, - img2img: sd15Img2img, - }, - sdxl: { - txt2img: sdxlTxt2img, - img2img: sdxlImg2img, - "txt2img-with-refiner": sdxlTxt2imgWithRefiner, - }, -}; +const workflows: WorkflowTree = {}; + +const walk = (dir: string, tree: WorkflowTree) => { + const files = fs.readdirSync(dir); + files.forEach((file) => { + let filePath = `${dir}/${file}`; + if (fs.statSync(filePath).isDirectory()) { + tree[file] = {}; + walk(filePath, tree[file] as WorkflowTree); + } else { + // This is happening at runtime, so if it's .ts we need to compile it + const newPath = filePath.replace(".ts", ".js"); + if (file.endsWith(".ts")) { + const source = fs.readFileSync(filePath, "utf8"); + const result = ts.transpileModule(source, { + compilerOptions: { module: ts.ModuleKind.CommonJS }, + }); + // write it a sibling .js file + fs.writeFileSync(newPath, result.outputText); + } else if (file.endsWith(".js")) { + // fs.cpSync(filePath, newPath); + } else { + return; + } -if (config.workflowModels !== "all" && config.workflowModels !== "") { - const requestedModels = new Set(config.workflowModels.split(",")); - Object.keys(workflows).forEach((baseModel) => { - if (!requestedModels.has(baseModel)) { - delete workflows[baseModel]; + // Eval the file in the current context + console.log(`Evaluating ${newPath}`); + const workflow = eval(fs.readFileSync(newPath, "utf8")); + if (workflow && isWorkflow(workflow)) { + tree[file.replace(".js", "").replace(".ts", "")] = workflow; + } } }); -} +}; +walk(config.workflowDir, workflows); + +// if (config.workflowModels !== "all" && config.workflowModels !== "") { +// const requestedModels = new Set(config.workflowModels.split(",")); +// Object.keys(workflows).forEach((baseModel) => { +// if (!requestedModels.has(baseModel)) { +// delete workflows[baseModel]; +// } +// }); +// } export default workflows;