improved automated endpoints

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;

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";
 

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[]]
-);

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"
+    }
+  }
+}