diff --git a/package-lock.json b/package-lock.json index f111cb5..2295a56 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "comfyui-wrapper", - "version": "1.2.0", + "version": "1.3.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "comfyui-wrapper", - "version": "1.2.0", + "version": "1.3.0", "license": "MIT", "dependencies": { "@fastify/swagger": "^8.15.0", @@ -1611,9 +1611,9 @@ } }, "node_modules/micromatch": { - "version": "4.0.7", - "resolved": "https://registry.npmjs.org/micromatch/-/micromatch-4.0.7.tgz", - "integrity": "sha512-LPP/3KorzCwBxfeUuZmaR6bG2kdeHSbe0P2tY3FLRU4vYrjYz5hI4QZwV0njUx3jeuKe67YukQ1LSPZBKDqO/Q==", + "version": "4.0.8", + "resolved": "https://registry.npmjs.org/micromatch/-/micromatch-4.0.8.tgz", + "integrity": "sha512-PXwfBhYu0hBCPw8Dn0E+WDYb7af3dSLVWKi3HGv84IdF4TyFoC0ysxFd0Goxw7nSv4T/PzEJQxsYsEiFCKo2BA==", "dev": true, "dependencies": { "braces": "^3.0.3", diff --git a/package.json b/package.json index ff7f5f6..24a3d6f 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "comfyui-wrapper", - "version": "1.2.0", + "version": "1.3.0", "description": "Wraps comfyui to make it easier to use as a web service", "main": "dist/src/index.js", "scripts": { diff --git a/src/config.ts b/src/config.ts index 4786ab6..4ac552d 100644 --- a/src/config.ts +++ b/src/config.ts @@ -72,7 +72,9 @@ 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); + const all = fs + .readdirSync(model_path) + .filter((f) => !(f.startsWith("put_") && f.endsWith("_here"))); config.models[model_dir] = { dir: model_path, all, diff --git a/src/server.ts b/src/server.ts index 1ac3f06..abfe200 100644 --- a/src/server.ts +++ b/src/server.ts @@ -40,6 +40,19 @@ const server = Fastify({ server.setValidatorCompiler(validatorCompiler); server.setSerializerCompiler(serializerCompiler); +const modelSchema: any = {}; +for (const modelType in config.models) { + modelSchema[modelType] = z.string().array(); +} + +const ModelResponseSchema = z.object(modelSchema); +type ModelResponse = z.infer; + +const modelResponse: ModelResponse = {}; +for (const modelType in config.models) { + modelResponse[modelType] = config.models[modelType].all; +} + server.register(fastifySwagger, { openapi: { openapi: "3.0.0", @@ -47,12 +60,6 @@ server.register(fastifySwagger, { title: "Comfy Wrapper API", version, }, - servers: [ - { - url: `http://localhost:${config.wrapperPort}`, - description: "Local server", - }, - ], }, transform: jsonSchemaTransform, }); @@ -62,36 +69,68 @@ server.register(fastifySwaggerUI, { server.after(() => { const app = server.withTypeProvider(); - app.route({ - method: "GET", - url: "/health", - schema: { - response: { - 200: z.object({ - version: z.string(), - status: z.literal("healthy"), - }), - 500: z.object({ - version: z.string(), - status: z.literal("not healthy"), - }), + app.get( + "/health", + { + schema: { + response: { + 200: z.object({ + version: z.literal(version), + status: z.literal("healthy"), + }), + 500: z.object({ + version: z.literal(version), + status: z.literal("not healthy"), + }), + }, }, }, - handler: async (request, reply) => { + async (request, reply) => { // 200 if ready, 500 if not if (warm) { return reply.code(200).send({ version, status: "healthy" }); } return reply.code(500).send({ version, status: "not healthy" }); + } + ); + + app.get( + "/ready", + { + schema: { + response: { + 200: z.object({ + version: z.literal(version), + status: z.literal("ready"), + }), + 500: z.object({ + version: z.literal(version), + status: z.literal("not ready"), + }), + }, + }, }, - }); + async (request, reply) => { + if (warm) { + return reply.code(200).send({ version, status: "ready" }); + } + return reply.code(500).send({ version, status: "not ready" }); + } + ); - app.get("/ready", async (request, reply) => { - if (warm) { - return reply.code(200).send({ version, status: "ready" }); + app.get( + "/models", + { + schema: { + response: { + 200: ModelResponseSchema, + }, + }, + }, + async (request, reply) => { + return modelResponse; } - return reply.code(500).send({ version, status: "not ready" }); - }); + ); app.post<{ Body: PromptRequest;