Merge pull request #4 from SaladTechnologies/schema-improvements

0.4.1 - Swagger docs updates
SaladTechnologies · Sep 25, 2024 · db5b23c · db5b23c
2 parents 36bf388 + 171c36b
commit db5b23c
Show file tree

Hide file tree

Showing 9 changed files with 53 additions and 63 deletions.
diff --git a/docker-compose.yml b/docker-compose.yml
@@ -0,0 +1,16 @@
+services:
+  comfyui:
+    image: saladtechnologies/comfyui:comfy0.2.0-api1.4.0-base
+    volumes:
+      - ./bin:/app/bin
+    command: ["./app/bin/comfyui-api"]
+    ports:
+      - "3000:3000"
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              capabilities: [ gpu ]
+              count: all
+
diff --git a/example-workflows/flux/img2img.ts b/example-workflows/flux/img2img.ts
@@ -198,6 +198,8 @@ function generateWorkflow(input: InputType): Record<string, ComfyNode> {
 const workflow: Workflow = {
   RequestSchema,
   generateWorkflow,
+  summary: "Image-to-Image",
+  description: "Text-guided Image-to-Image generation",
 };
 
 export default workflow;
diff --git a/example-workflows/sd1.5/img2img.ts b/example-workflows/sd1.5/img2img.ts
@@ -204,6 +204,8 @@ function generateWorkflow(input: InputType): Record<string, ComfyNode> {
 const workflow: Workflow = {
   RequestSchema,
   generateWorkflow,
+  summary: "Image-to-Image",
+  description: "Text-guided Image-to-Image generation",
 };
 
 export default workflow;
diff --git a/example-workflows/sdxl/img2img.ts b/example-workflows/sdxl/img2img.ts
@@ -216,6 +216,8 @@ function generateWorkflow(input: InputType): Record<string, ComfyNode> {
 const workflow: Workflow = {
   RequestSchema,
   generateWorkflow,
+  summary: "Image-to-Image",
+  description: "Text-guided Image-to-Image generation",
 };
 
 export default workflow;
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "comfyui-api",
-  "version": "1.4.0",
+  "version": "1.4.1",
   "description": "Wraps comfyui to make it easier to use as a stateless web service",
   "main": "dist/src/index.js",
   "scripts": {

diff --git a/src/config.ts b/src/config.ts
@@ -17,6 +17,7 @@ const {
   WARMUP_PROMPT_FILE,
   WORKFLOW_MODELS = "all",
   WORKFLOW_DIR = "/workflows",
+  MARKDOWN_SCHEMA_DESCRIPTIONS = "true",
 } = process.env;
 
 fs.mkdirSync(WORKFLOW_DIR, { recursive: true });
@@ -132,6 +133,7 @@ const config = {
     }
   >,
   workflowModels: WORKFLOW_MODELS,
+  markdownSchemaDescriptions: MARKDOWN_SCHEMA_DESCRIPTIONS === "true",
 };
 
 const model_dirs = fs.readdirSync(MODEL_DIR);

diff --git a/src/server.ts b/src/server.ts
@@ -78,6 +78,8 @@ server.after(() => {
     "/health",
     {
       schema: {
+        summary: "Health Probe",
+        description: "Check if the server is healthy",
         response: {
           200: z.object({
             version: z.literal(version),
@@ -103,6 +105,8 @@ server.after(() => {
     "/ready",
     {
       schema: {
+        summary: "Readiness Probe",
+        description: "Check if the server is ready to serve traffic",
         response: {
           200: z.object({
             version: z.literal(version),
@@ -127,6 +131,9 @@ server.after(() => {
     "/models",
     {
       schema: {
+        summary: "List Models",
+        description:
+          "List all available models. This is from the contents of the models directory.",
         response: {
           200: ModelResponseSchema,
         },
@@ -143,6 +150,8 @@ server.after(() => {
     "/prompt",
     {
       schema: {
+        summary: "Submit Prompt",
+        description: "Submit an API-formatted ComfyUI prompt.",
         body: PromptRequestSchema,
         response: {
           200: PromptResponseSchema,
@@ -283,14 +292,25 @@ server.after(() => {
 
         type BodyType = z.infer<typeof BodySchema>;
 
-        const description = zodToMarkdownTable(node.RequestSchema);
+        let description = "";
+        if (config.markdownSchemaDescriptions) {
+          description = zodToMarkdownTable(node.RequestSchema);
+        } else if (node.description) {
+          description = node.description;
+        }
+
+        let summary = key;
+        if (node.summary) {
+          summary = node.summary;
+        }
 
         app.post<{
           Body: BodyType;
         }>(
           `${route}/${key}`,
           {
             schema: {
+              summary,
               description,
               body: BodySchema,
               response: {
@@ -332,62 +352,6 @@ server.after(() => {
     }
   };
   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<typeof BodySchema>;
-
-  //     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;

diff --git a/src/types.ts b/src/types.ts
@@ -46,6 +46,8 @@ export const WorkflowSchema = z.object({
 export interface Workflow {
   RequestSchema: z.ZodObject<any, any>;
   generateWorkflow: (input: any) => Record<string, ComfyNode>;
+  description?: string;
+  summary?: string;
 }
 
 export function isWorkflow(obj: any): obj is Workflow {