diff --git a/src/config.ts b/src/config.ts index 03902c4..435e1b1 100644 --- a/src/config.ts +++ b/src/config.ts @@ -10,6 +10,7 @@ const { STARTUP_CHECK_MAX_TRIES = "10", OUTPUT_DIR = "/opt/ComfyUI/output", INPUT_DIR = "/opt/ComfyUI/input", + CKPT_DIR = "/opt/ComfyUI/models/checkpoints", WARMUP_PROMPT_FILE, } = process.env; @@ -18,18 +19,28 @@ const port = parseInt(PORT, 10); const startupCheckInterval = parseInt(STARTUP_CHECK_INTERVAL_S, 10) * 1000; const startupCheckMaxTries = parseInt(STARTUP_CHECK_MAX_TRIES, 10); -let warmupPrompt: string | undefined; +let warmupPrompt: any | undefined; +let warmupCkpt: string | undefined; if (WARMUP_PROMPT_FILE) { assert(fs.existsSync(WARMUP_PROMPT_FILE), "Warmup prompt file not found"); try { warmupPrompt = JSON.parse( fs.readFileSync(WARMUP_PROMPT_FILE, { encoding: "utf-8" }) ); + for (const nodeId in warmupPrompt) { + const node = warmupPrompt[nodeId]; + if (node.class_type === "CheckpointLoaderSimple") { + warmupCkpt = node.inputs.ckpt_name; + break; + } + } } catch (e: any) { throw new Error(`Failed to parse warmup prompt: ${e.message}`); } } +const allCheckpoints = fs.readdirSync(CKPT_DIR); + const config = { comfyLaunchCmd: CMD, wrapperHost: HOST, @@ -42,6 +53,8 @@ const config = { outputDir: OUTPUT_DIR, inputDir: INPUT_DIR, warmupPrompt, + warmupCkpt, + checkpoints: allCheckpoints, }; export default config; diff --git a/src/server.ts b/src/server.ts index e099ff2..966d177 100644 --- a/src/server.ts +++ b/src/server.ts @@ -264,7 +264,14 @@ server.after(() => { body: JSON.stringify({ prompt, id, webhook }), } ); - return reply.code(resp.status).send(await resp.json()); + const body = await resp.json(); + if (!resp.ok) { + return reply.code(resp.status).send(body); + } + + body.prompt = prompt; + + return reply.code(resp.status).send(body); } ); } diff --git a/src/types.ts b/src/types.ts index 2a127da..bc9ed92 100644 --- a/src/types.ts +++ b/src/types.ts @@ -1,5 +1,6 @@ -import { z } from "zod"; +import { string, z } from "zod"; import { randomUUID } from "crypto"; +import config from "./config"; export const ComfyNodeSchema = z.object({ inputs: z.any(), @@ -42,10 +43,10 @@ export const WorkflowSchema = z.object({ generateWorkflow: z.function(), }); -export type Workflow = { +export interface Workflow { RequestSchema: z.ZodObject; generateWorkflow: (input: any) => Record; -}; +} export const WorkflowRequestSchema = z.object({ id: z @@ -57,3 +58,7 @@ export const WorkflowRequestSchema = z.object({ }); export type WorkflowRequest = z.infer; + +export const AvailableCheckpoints = z.enum( + config.checkpoints as unknown as readonly [string, ...string[]] +); diff --git a/src/workflows/flux/img2img.ts b/src/workflows/flux/img2img.ts new file mode 100644 index 0000000..40a85aa --- /dev/null +++ b/src/workflows/flux/img2img.ts @@ -0,0 +1,153 @@ +import { z } from "zod"; +import { ComfyNode, Workflow, AvailableCheckpoints } from "../../types"; +import config from "../../config"; + +let checkpoint: any = AvailableCheckpoints.optional(); +if (config.warmupCkpt) { + checkpoint = AvailableCheckpoints.default(config.warmupCkpt); +} + +const RequestSchema = z.object({ + prompt: z.string(), + width: z.number().int().min(256).max(1024).optional().default(1024), + height: z.number().int().min(256).max(1024).optional().default(1024), + seed: z + .number() + .int() + .optional() + .default(() => Math.floor(Math.random() * 1000000000000000)), + steps: z.number().int().min(1).max(10).optional().default(2), + sampler: z.enum(["euler"]).optional().default("euler"), + scheduler: z.enum(["simple"]).optional().default("simple"), + denoise: z.number().min(0).max(1).optional().default(0.8), + cfg: z.number().min(1).max(30).optional().default(1), + image: z.string(), + checkpoint, +}); + +type InputType = z.infer; + +function generateWorkflow(input: InputType): Record { + return { + "6": { + inputs: { + text: input.prompt, + clip: ["30", 1], + }, + class_type: "CLIPTextEncode", + _meta: { + title: "CLIP Text Encode (Positive Prompt)", + }, + }, + "8": { + inputs: { + samples: ["31", 0], + vae: ["30", 2], + }, + class_type: "VAEDecode", + _meta: { + title: "VAE Decode", + }, + }, + "9": { + inputs: { + filename_prefix: "", + images: ["8", 0], + }, + class_type: "SaveImage", + _meta: { + title: "Save Image", + }, + }, + "27": { + inputs: { + width: input.width, + height: input.height, + batch_size: 1, + }, + class_type: "EmptySD3LatentImage", + _meta: { + title: "EmptySD3LatentImage", + }, + }, + "30": { + inputs: { + ckpt_name: input.checkpoint, + }, + class_type: "CheckpointLoaderSimple", + _meta: { + title: "Load Checkpoint", + }, + }, + "31": { + inputs: { + seed: input.seed, + steps: input.steps, + cfg: input.cfg, + sampler_name: input.sampler, + scheduler: input.scheduler, + denoise: input.denoise, + model: ["30", 0], + positive: ["6", 0], + negative: ["33", 0], + latent_image: ["38", 0], + }, + class_type: "KSampler", + _meta: { + title: "KSampler", + }, + }, + "33": { + inputs: { + text: "", + clip: ["30", 1], + }, + class_type: "CLIPTextEncode", + _meta: { + title: "CLIP Text Encode (Negative Prompt)", + }, + }, + "37": { + inputs: { + image: input.image, + upload: "image", + }, + class_type: "LoadImage", + _meta: { + title: "Load Image", + }, + }, + "38": { + inputs: { + pixels: ["40", 0], + vae: ["30", 2], + }, + class_type: "VAEEncode", + _meta: { + title: "VAE Encode", + }, + }, + "40": { + inputs: { + width: input.width, + height: input.height, + interpolation: "nearest", + method: "fill / crop", + condition: "always", + multiple_of: 8, + image: ["37", 0], + }, + class_type: "ImageResize+", + _meta: { + title: "🔧 Image Resize", + }, + }, + }; +} + +const workflow: Workflow = { + RequestSchema, + generateWorkflow, +}; + +export default workflow; diff --git a/src/workflows/flux/txt2img.ts b/src/workflows/flux/txt2img.ts index e54c28d..a283b19 100644 --- a/src/workflows/flux/txt2img.ts +++ b/src/workflows/flux/txt2img.ts @@ -1,5 +1,11 @@ import { z } from "zod"; -import { ComfyNode, Workflow } from "../../types"; +import { ComfyNode, Workflow, AvailableCheckpoints } from "../../types"; +import config from "../../config"; + +let checkpoint: any = AvailableCheckpoints.optional(); +if (config.warmupCkpt) { + checkpoint = AvailableCheckpoints.default(config.warmupCkpt); +} const RequestSchema = z.object({ prompt: z.string(), @@ -13,15 +19,12 @@ const RequestSchema = z.object({ steps: z.number().int().min(1).max(10).optional().default(4), sampler: z.enum(["euler"]).optional().default("euler"), // This may need to be expanded with more options scheduler: z.enum(["simple"]).optional().default("simple"), // This may need to be expanded with more options - checkpoint: z - .enum(["flux1-schnell-fp8.safetensors"]) - .optional() - .default("flux1-schnell-fp8.safetensors"), // This may need to be expanded with more options + checkpoint, }); -type Input = z.infer; +type InputType = z.infer; -export function generateWorkflow(input: Input): Record { +function generateWorkflow(input: InputType): Record { return { "6": { inputs: { diff --git a/src/workflows/index.ts b/src/workflows/index.ts index 7f8fe56..d3e0a04 100644 --- a/src/workflows/index.ts +++ b/src/workflows/index.ts @@ -1,7 +1,13 @@ import fluxTxt2img from "../workflows/flux/txt2img"; +import fluxImg2img from "../workflows/flux/img2img"; +import sd15Txt2img from "../workflows/sd1.5/txt2img"; export const workflows: any = { flux: { txt2img: fluxTxt2img, + img2img: fluxImg2img, + }, + "sd1.5": { + txt2img: sd15Txt2img, }, }; diff --git a/src/workflows/sd1.5/txt2img.ts b/src/workflows/sd1.5/txt2img.ts new file mode 100644 index 0000000..4c12fce --- /dev/null +++ b/src/workflows/sd1.5/txt2img.ts @@ -0,0 +1,117 @@ +import { z } from "zod"; +import { ComfyNode, Workflow, AvailableCheckpoints } from "../../types"; +import config from "../../config"; + +let checkpoint: any = AvailableCheckpoints.optional(); +if (config.warmupCkpt) { + checkpoint = AvailableCheckpoints.default(config.warmupCkpt); +} + +const RequestSchema = z.object({ + prompt: z.string(), + negative_prompt: z.string().optional().default("text, watermark"), + width: z.number().int().min(64).max(2048).optional().default(512), + height: z.number().int().min(64).max(2048).optional().default(512), + seed: z + .number() + .int() + .optional() + .default(() => Math.floor(Math.random() * 1000000000000000)), + steps: z.number().int().min(1).max(100).optional().default(20), + cfg: z.number().min(1).max(20).optional().default(8), + sampler_name: z.enum(["euler"]).optional().default("euler"), + scheduler: z.enum(["normal"]).optional().default("normal"), + checkpoint, +}); + +type InputType = z.infer; + +function generateWorkflow(input: InputType): Record { + return { + "3": { + inputs: { + seed: input.seed, + steps: input.steps, + cfg: input.cfg, + sampler_name: input.sampler_name, + scheduler: input.scheduler, + denoise: 1, + 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;