diff --git a/src/config.ts b/src/config.ts index 435e1b1..4786ab6 100644 --- a/src/config.ts +++ b/src/config.ts @@ -1,5 +1,7 @@ import assert from "node:assert"; import fs from "node:fs"; +import path from "node:path"; +import { z } from "zod"; const { CMD = "init.sh", HOST = "::", @@ -10,11 +12,13 @@ const { STARTUP_CHECK_MAX_TRIES = "10", OUTPUT_DIR = "/opt/ComfyUI/output", INPUT_DIR = "/opt/ComfyUI/input", - CKPT_DIR = "/opt/ComfyUI/models/checkpoints", + MODEL_DIR = "/opt/ComfyUI/models", WARMUP_PROMPT_FILE, + WORKFLOW_MODELS = "all", } = process.env; const comfyURL = `http://${DIRECT_ADDRESS}:${COMFYUI_PORT_HOST}`; +const selfURL = `http://localhost:${PORT}`; const port = parseInt(PORT, 10); const startupCheckInterval = parseInt(STARTUP_CHECK_INTERVAL_S, 10) * 1000; const startupCheckMaxTries = parseInt(STARTUP_CHECK_MAX_TRIES, 10); @@ -39,12 +43,11 @@ if (WARMUP_PROMPT_FILE) { } } -const allCheckpoints = fs.readdirSync(CKPT_DIR); - const config = { comfyLaunchCmd: CMD, wrapperHost: HOST, wrapperPort: port, + selfURL, comfyHost: DIRECT_ADDRESS, comfyPort: COMFYUI_PORT_HOST, comfyURL, @@ -54,7 +57,28 @@ const config = { inputDir: INPUT_DIR, warmupPrompt, warmupCkpt, - checkpoints: allCheckpoints, + models: {} as Record< + string, + { + dir: string; + all: string[]; + enum: z.ZodEnum<[string, ...string[]]>; + } + >, + workflowModels: WORKFLOW_MODELS, }; +const model_dirs = fs.readdirSync(MODEL_DIR); +for (const model_dir of model_dirs) { + const model_path = path.join(MODEL_DIR, model_dir); + if (fs.statSync(model_path).isDirectory()) { + const all = fs.readdirSync(model_path); + config.models[model_dir] = { + dir: model_path, + all, + enum: z.enum(all as [string, ...string[]]), + }; + } +} + export default config; diff --git a/src/server.ts b/src/server.ts index 58931ff..1ac3f06 100644 --- a/src/server.ts +++ b/src/server.ts @@ -27,7 +27,7 @@ import { Workflow, WorkflowResponseSchema, } from "./types"; -import { workflows } from "./workflows"; +import workflows from "./workflows"; import { z } from "zod"; import { randomUUID } from "crypto"; diff --git a/src/types.ts b/src/types.ts index 3dbdc85..a478e3a 100644 --- a/src/types.ts +++ b/src/types.ts @@ -66,7 +66,3 @@ export const WorkflowResponseSchema = z.object({ webhook: z.string().optional(), status: z.enum(["ok"]).optional(), }); - -export const AvailableCheckpoints = z.enum( - config.checkpoints as unknown as readonly [string, ...string[]] -); diff --git a/src/workflows/flux/img2img.json b/src/workflows/flux/img2img.json new file mode 100644 index 0000000..980fa64 --- /dev/null +++ b/src/workflows/flux/img2img.json @@ -0,0 +1,151 @@ +{ + "6": { + "inputs": { + "text": "A noble wolf stands by a raging river in the style of a japanese scroll", + "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": "ComfyUI", + "images": [ + "8", + 0 + ] + }, + "class_type": "SaveImage", + "_meta": { + "title": "Save Image" + } + }, + "27": { + "inputs": { + "width": 1024, + "height": 1024, + "batch_size": 1 + }, + "class_type": "EmptySD3LatentImage", + "_meta": { + "title": "EmptySD3LatentImage" + } + }, + "30": { + "inputs": { + "ckpt_name": "flux1-schnell-fp8.safetensors" + }, + "class_type": "CheckpointLoaderSimple", + "_meta": { + "title": "Load Checkpoint" + } + }, + "31": { + "inputs": { + "seed": 226018262510838, + "steps": 4, + "cfg": 1, + "sampler_name": "euler", + "scheduler": "simple", + "denoise": 0.8, + "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": "IMG_0655.JPG", + "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": 1024, + "height": 1024, + "interpolation": "nearest", + "method": "fill / crop", + "condition": "always", + "multiple_of": 8, + "image": [ + "37", + 0 + ] + }, + "class_type": "ImageResize+", + "_meta": { + "title": "🔧 Image Resize" + } + } +} \ No newline at end of file diff --git a/src/workflows/flux/img2img.ts b/src/workflows/flux/img2img.ts index 40a85aa..d067029 100644 --- a/src/workflows/flux/img2img.ts +++ b/src/workflows/flux/img2img.ts @@ -1,28 +1,95 @@ import { z } from "zod"; -import { ComfyNode, Workflow, AvailableCheckpoints } from "../../types"; +import { ComfyNode, Workflow } from "../../types"; import config from "../../config"; -let checkpoint: any = AvailableCheckpoints.optional(); +let checkpoint: any = config.models.checkpoints.enum.optional(); if (config.warmupCkpt) { - checkpoint = AvailableCheckpoints.default(config.warmupCkpt); + checkpoint = checkpoint.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), + prompt: z.string().describe("The positive prompt for image generation"), + negative_prompt: z + .string() + .optional() + .describe("The negative prompt for image generation"), + width: z + .number() + .int() + .min(256) + .max(2048) + .optional() + .default(1024) + .describe("Width of the generated image"), + height: z + .number() + .int() + .min(256) + .max(2048) + .optional() + .default(1024) + .describe("Height of the generated image"), 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(), + .default(() => Math.floor(Math.random() * 1000000000000000)) + .describe("Seed for random number generation"), + steps: z + .number() + .int() + .min(1) + .max(100) + .optional() + .default(4) + .describe("Number of sampling steps"), + cfg_scale: z + .number() + .min(0) + .max(20) + .optional() + .default(1) + .describe("Classifier-free guidance scale"), + sampler_name: z + .enum(["euler"]) + .optional() + .default("euler") + .describe("Name of the sampler to use"), + scheduler: z + .enum(["simple"]) + .optional() + .default("simple") + .describe("Type of scheduler to use"), + denoise: z + .number() + .min(0) + .max(1) + .optional() + .default(0.8) + .describe("Denoising strength"), checkpoint, + image: z.string().describe("Input image for img2img"), + interpolation: z + .enum(["nearest"]) + .optional() + .default("nearest") + .describe("Interpolation method for image resizing"), + resize_method: z + .enum(["fill / crop"]) + .optional() + .default("fill / crop") + .describe("Method for resizing the image"), + resize_condition: z + .enum(["always"]) + .optional() + .default("always") + .describe("Condition for when to resize the image"), + multiple_of: z + .number() + .int() + .optional() + .default(8) + .describe("Ensure image dimensions are multiples of this value"), }); type InputType = z.infer; @@ -51,7 +118,7 @@ function generateWorkflow(input: InputType): Record { }, "9": { inputs: { - filename_prefix: "", + filename_prefix: "ComfyUI", images: ["8", 0], }, class_type: "SaveImage", @@ -59,17 +126,6 @@ function generateWorkflow(input: InputType): Record { 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, @@ -83,8 +139,8 @@ function generateWorkflow(input: InputType): Record { inputs: { seed: input.seed, steps: input.steps, - cfg: input.cfg, - sampler_name: input.sampler, + cfg: input.cfg_scale, + sampler_name: input.sampler_name, scheduler: input.scheduler, denoise: input.denoise, model: ["30", 0], @@ -99,7 +155,7 @@ function generateWorkflow(input: InputType): Record { }, "33": { inputs: { - text: "", + text: input.negative_prompt, clip: ["30", 1], }, class_type: "CLIPTextEncode", @@ -131,10 +187,10 @@ function generateWorkflow(input: InputType): Record { inputs: { width: input.width, height: input.height, - interpolation: "nearest", - method: "fill / crop", - condition: "always", - multiple_of: 8, + interpolation: input.interpolation, + method: input.resize_method, + condition: input.resize_condition, + multiple_of: input.multiple_of, image: ["37", 0], }, class_type: "ImageResize+", diff --git a/src/workflows/flux/txt2img.json b/src/workflows/flux/txt2img.json new file mode 100644 index 0000000..52916d4 --- /dev/null +++ b/src/workflows/flux/txt2img.json @@ -0,0 +1,107 @@ +{ + "6": { + "inputs": { + "text": "a bottle with a beautiful rainbow galaxy inside it on top of a wooden table in the middle of a modern kitchen beside a plate of vegetables and mushrooms and a wine glasse that contains a planet earth with a plate with a half eaten apple pie on it", + "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": "Flux", + "images": [ + "8", + 0 + ] + }, + "class_type": "SaveImage", + "_meta": { + "title": "Save Image" + } + }, + "27": { + "inputs": { + "width": 1024, + "height": 1024, + "batch_size": 1 + }, + "class_type": "EmptySD3LatentImage", + "_meta": { + "title": "EmptySD3LatentImage" + } + }, + "30": { + "inputs": { + "ckpt_name": "flux1-schnell-fp8.safetensors" + }, + "class_type": "CheckpointLoaderSimple", + "_meta": { + "title": "Load Checkpoint" + } + }, + "31": { + "inputs": { + "seed": 1030319533692526, + "steps": 4, + "cfg": 1, + "sampler_name": "euler", + "scheduler": "simple", + "denoise": 1, + "model": [ + "30", + 0 + ], + "positive": [ + "6", + 0 + ], + "negative": [ + "33", + 0 + ], + "latent_image": [ + "27", + 0 + ] + }, + "class_type": "KSampler", + "_meta": { + "title": "KSampler" + } + }, + "33": { + "inputs": { + "text": "", + "clip": [ + "30", + 1 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "CLIP Text Encode (Negative Prompt)" + } + } +} \ No newline at end of file diff --git a/src/workflows/flux/txt2img.ts b/src/workflows/flux/txt2img.ts index a283b19..30d1bf2 100644 --- a/src/workflows/flux/txt2img.ts +++ b/src/workflows/flux/txt2img.ts @@ -1,24 +1,72 @@ import { z } from "zod"; -import { ComfyNode, Workflow, AvailableCheckpoints } from "../../types"; +import { ComfyNode, Workflow } from "../../types"; import config from "../../config"; -let checkpoint: any = AvailableCheckpoints.optional(); +let checkpoint: any = config.models.checkpoints.enum.optional(); if (config.warmupCkpt) { - checkpoint = AvailableCheckpoints.default(config.warmupCkpt); + checkpoint = checkpoint.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), + prompt: z.string().describe("The positive prompt for image generation"), + negative_prompt: z + .string() + .optional() + .describe("The negative prompt for image generation"), + width: z + .number() + .int() + .min(256) + .max(2048) + .optional() + .default(1024) + .describe("Width of the generated image"), + height: z + .number() + .int() + .min(256) + .max(2048) + .optional() + .default(1024) + .describe("Height of the generated image"), seed: z .number() .int() .optional() - .default(() => Math.floor(Math.random() * 100000)), - 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 + .default(() => Math.floor(Math.random() * 1000000000000000)) + .describe("Seed for random number generation"), + steps: z + .number() + .int() + .min(1) + .max(100) + .optional() + .default(4) + .describe("Number of sampling steps"), + cfg_scale: z + .number() + .min(0) + .max(20) + .optional() + .default(1) + .describe("Classifier-free guidance scale"), + sampler_name: z + .enum(["euler"]) + .optional() + .default("euler") + .describe("Name of the sampler to use"), + scheduler: z + .enum(["simple"]) + .optional() + .default("simple") + .describe("Type of scheduler to use"), + denoise: z + .number() + .min(0) + .max(1) + .optional() + .default(1) + .describe("Denoising strength"), checkpoint, }); @@ -48,7 +96,7 @@ function generateWorkflow(input: InputType): Record { }, "9": { inputs: { - filename_prefix: "", + filename_prefix: "Flux", images: ["8", 0], }, class_type: "SaveImage", @@ -80,10 +128,10 @@ function generateWorkflow(input: InputType): Record { inputs: { seed: input.seed, steps: input.steps, - cfg: 1.0, - sampler_name: input.sampler, + cfg: input.cfg_scale, + sampler_name: input.sampler_name, scheduler: input.scheduler, - denoise: 1, + denoise: input.denoise, model: ["30", 0], positive: ["6", 0], negative: ["33", 0], @@ -96,7 +144,7 @@ function generateWorkflow(input: InputType): Record { }, "33": { inputs: { - text: "", + text: input.negative_prompt || "", clip: ["30", 1], }, class_type: "CLIPTextEncode", diff --git a/src/workflows/index.ts b/src/workflows/index.ts index 6c4d6b4..a94e0f2 100644 --- a/src/workflows/index.ts +++ b/src/workflows/index.ts @@ -1,9 +1,15 @@ +import config from "../config"; import fluxTxt2img from "../workflows/flux/txt2img"; import fluxImg2img from "../workflows/flux/img2img"; import sd15Txt2img from "../workflows/sd1.5/txt2img"; import sdxlImg2img from "../workflows/sdxl/img2img"; +import { Workflow } from "../types"; -export const workflows: any = { +const workflows: { + [key: string]: { + [key: string]: Workflow; + }; +} = { flux: { txt2img: fluxTxt2img, img2img: fluxImg2img, @@ -15,3 +21,14 @@ export const workflows: any = { img2img: sdxlImg2img, }, }; + +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; diff --git a/src/workflows/sd1.5/txt2img.json b/src/workflows/sd1.5/txt2img.json new file mode 100644 index 0000000..4317ccd --- /dev/null +++ b/src/workflows/sd1.5/txt2img.json @@ -0,0 +1,107 @@ +{ + "3": { + "inputs": { + "seed": 712610403220747, + "steps": 20, + "cfg": 8, + "sampler_name": "euler", + "scheduler": "normal", + "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": "dreamshaper_8.safetensors" + }, + "class_type": "CheckpointLoaderSimple", + "_meta": { + "title": "Load Checkpoint" + } + }, + "5": { + "inputs": { + "width": 512, + "height": 512, + "batch_size": 1 + }, + "class_type": "EmptyLatentImage", + "_meta": { + "title": "Empty Latent Image" + } + }, + "6": { + "inputs": { + "text": "beautiful scenery nature glass bottle landscape, , purple galaxy bottle,", + "clip": [ + "4", + 1 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "CLIP Text Encode (Prompt)" + } + }, + "7": { + "inputs": { + "text": "text, watermark", + "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" + } + } +} \ No newline at end of file diff --git a/src/workflows/sd1.5/txt2img.ts b/src/workflows/sd1.5/txt2img.ts index 4c12fce..5d4331b 100644 --- a/src/workflows/sd1.5/txt2img.ts +++ b/src/workflows/sd1.5/txt2img.ts @@ -1,26 +1,72 @@ import { z } from "zod"; -import { ComfyNode, Workflow, AvailableCheckpoints } from "../../types"; +import { ComfyNode, Workflow } from "../../types"; import config from "../../config"; -let checkpoint: any = AvailableCheckpoints.optional(); +let checkpoint: any = config.models.checkpoints.enum.optional(); if (config.warmupCkpt) { - checkpoint = AvailableCheckpoints.default(config.warmupCkpt); + checkpoint = checkpoint.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), + prompt: z.string().describe("The positive prompt for image generation"), + negative_prompt: z + .string() + .optional() + .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() * 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"), + .default(712610403220747) + .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: z + .enum(["euler"]) + .optional() + .default("euler") + .describe("Name of the sampler to use"), + scheduler: z + .enum(["normal"]) + .optional() + .default("normal") + .describe("Type of scheduler to use"), + denoise: z + .number() + .min(0) + .max(1) + .optional() + .default(1) + .describe("Denoising strength"), checkpoint, }); @@ -32,10 +78,10 @@ function generateWorkflow(input: InputType): Record { inputs: { seed: input.seed, steps: input.steps, - cfg: input.cfg, + cfg: input.cfg_scale, sampler_name: input.sampler_name, scheduler: input.scheduler, - denoise: 1, + denoise: input.denoise, model: ["4", 0], positive: ["6", 0], negative: ["7", 0], diff --git a/src/workflows/sdxl/img2img.json b/src/workflows/sdxl/img2img.json new file mode 100644 index 0000000..1a6f8d5 --- /dev/null +++ b/src/workflows/sdxl/img2img.json @@ -0,0 +1,152 @@ +{ + "8": { + "inputs": { + "samples": [ + "36", + 0 + ], + "vae": [ + "14", + 2 + ] + }, + "class_type": "VAEDecode", + "_meta": { + "title": "VAE Decode" + } + }, + "9": { + "inputs": { + "filename_prefix": "img2img", + "images": [ + "8", + 0 + ] + }, + "class_type": "SaveImage", + "_meta": { + "title": "Save Image" + } + }, + "14": { + "inputs": { + "ckpt_name": "sd_xl_base_1.0.safetensors" + }, + "class_type": "CheckpointLoaderSimple", + "_meta": { + "title": "Load Checkpoint Base" + } + }, + "16": { + "inputs": { + "width": 4096, + "height": 4096, + "crop_w": 0, + "crop_h": 0, + "target_width": 4096, + "target_height": 4096, + "text_g": "a professional photo of a young man smiling\n\nhigh resolution, highly detailed, 4k", + "text_l": "a professional photo of a young man smiling\n\nhigh resolution, highly detailed, 4k", + "clip": [ + "14", + 1 + ] + }, + "class_type": "CLIPTextEncodeSDXL", + "_meta": { + "title": "CLIPTextEncodeSDXL" + } + }, + "19": { + "inputs": { + "width": 4096, + "height": 4096, + "crop_w": 0, + "crop_h": 0, + "target_width": 4096, + "target_height": 4096, + "text_g": "blurry, horror, rendering, illustration, drawing, painting", + "text_l": "blurry, horror, rendering, illustration, drawing, painting", + "clip": [ + "14", + 1 + ] + }, + "class_type": "CLIPTextEncodeSDXL", + "_meta": { + "title": "CLIPTextEncodeSDXL" + } + }, + "36": { + "inputs": { + "seed": 887855663168366, + "steps": 20, + "cfg": 5.5, + "sampler_name": "dpmpp_2m_sde_gpu", + "scheduler": "exponential", + "denoise": 0.75, + "model": [ + "14", + 0 + ], + "positive": [ + "16", + 0 + ], + "negative": [ + "19", + 0 + ], + "latent_image": [ + "39", + 0 + ] + }, + "class_type": "KSampler", + "_meta": { + "title": "KSampler" + } + }, + "38": { + "inputs": { + "image": "ComfyUI_00376_.png", + "upload": "image" + }, + "class_type": "LoadImage", + "_meta": { + "title": "Load Image" + } + }, + "39": { + "inputs": { + "pixels": [ + "40", + 0 + ], + "vae": [ + "14", + 2 + ] + }, + "class_type": "VAEEncode", + "_meta": { + "title": "VAE Encode" + } + }, + "40": { + "inputs": { + "upscale_method": "nearest-exact", + "width": 1024, + "height": 1024, + "crop": "center", + "image": [ + "38", + 0 + ] + }, + "class_type": "ImageScale", + "_meta": { + "title": "Upscale Image" + } + } +} \ No newline at end of file diff --git a/src/workflows/sdxl/img2img.ts b/src/workflows/sdxl/img2img.ts index fc4a68a..683b14d 100644 --- a/src/workflows/sdxl/img2img.ts +++ b/src/workflows/sdxl/img2img.ts @@ -1,29 +1,97 @@ import { z } from "zod"; -import { ComfyNode, Workflow, AvailableCheckpoints } from "../../types"; +import { ComfyNode, Workflow } from "../../types"; import config from "../../config"; -let checkpoint: any = AvailableCheckpoints.optional(); +let checkpoint: any = config.models.checkpoints.enum.optional(); if (config.warmupCkpt) { - checkpoint = AvailableCheckpoints.default(config.warmupCkpt); + checkpoint = checkpoint.default(config.warmupCkpt); } const RequestSchema = z.object({ - prompt: z.string(), - negativePrompt: z.string().optional().default(""), - width: z.number().int().min(256).max(2048).optional().default(1024), - height: z.number().int().min(256).max(2048).optional().default(1024), + prompt: z.string().describe("The positive prompt for image generation"), + negative_prompt: z + .string() + .optional() + .describe("The negative prompt for image generation"), + width: z + .number() + .int() + .min(256) + .max(4096) + .optional() + .default(4096) + .describe("Width of the generated image"), + height: z + .number() + .int() + .min(256) + .max(4096) + .optional() + .default(4096) + .describe("Height of the generated image"), 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(0).max(20).optional().default(5.5), - denoise: z.number().min(0).max(1).optional().default(0.75), - sampler: z.enum(["dpmpp_2m_sde_gpu"]).optional().default("dpmpp_2m_sde_gpu"), - scheduler: z.enum(["exponential"]).optional().default("exponential"), + .default(() => Math.floor(Math.random() * 1000000000000000)) + .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(5.5) + .describe("Classifier-free guidance scale"), + sampler_name: z + .enum(["dpmpp_2m_sde_gpu"]) + .optional() + .default("dpmpp_2m_sde_gpu") + .describe("Name of the sampler to use"), + scheduler: z + .enum(["exponential"]) + .optional() + .default("exponential") + .describe("Type of scheduler to use"), + denoise: z + .number() + .min(0) + .max(1) + .optional() + .default(0.75) + .describe("Denoising strength"), checkpoint, - image: z.string(), + image: z.string().describe("Input image for img2img"), + upscale_method: z + .enum(["nearest-exact"]) + .optional() + .default("nearest-exact") + .describe( + "Method used for upscaling if input image is smaller than target size" + ), + target_width: z + .number() + .int() + .min(256) + .max(4096) + .optional() + .default(1024) + .describe("Target width for upscaling"), + target_height: z + .number() + .int() + .min(256) + .max(4096) + .optional() + .default(1024) + .describe("Target height for upscaling"), }); type InputType = z.infer; @@ -84,8 +152,8 @@ function generateWorkflow(input: InputType): Record { crop_h: 0, target_width: input.width, target_height: input.height, - text_g: input.negativePrompt, - text_l: input.negativePrompt, + text_g: input.negative_prompt, + text_l: input.negative_prompt, clip: ["14", 1], }, class_type: "CLIPTextEncodeSDXL", @@ -97,8 +165,8 @@ function generateWorkflow(input: InputType): Record { inputs: { seed: input.seed, steps: input.steps, - cfg: input.cfg, - sampler_name: input.sampler, + cfg: input.cfg_scale, + sampler_name: input.sampler_name, scheduler: input.scheduler, denoise: input.denoise, model: ["14", 0], @@ -133,9 +201,9 @@ function generateWorkflow(input: InputType): Record { }, "40": { inputs: { - upscale_method: "nearest-exact", - width: 1024, - height: 1024, + upscale_method: input.upscale_method, + width: input.target_width, + height: input.target_height, crop: "center", image: ["38", 0], },