From 6313c5d67e91de63668ae0b530a08036cea75667 Mon Sep 17 00:00:00 2001
From: pinks <lisq@cock.li>
Date: Wed, 11 Oct 2023 03:59:52 +0200
Subject: [PATCH] feat: workers api

---
 api/serveApi.ts     |   2 +
 api/workersRoute.ts | 103 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 105 insertions(+)
 create mode 100644 api/workersRoute.ts

diff --git a/api/serveApi.ts b/api/serveApi.ts
index 9373870..9781194 100644
--- a/api/serveApi.ts
+++ b/api/serveApi.ts
@@ -4,6 +4,7 @@ import { paramsRoute } from "./paramsRoute.ts";
 import { sessionsRoute } from "./sessionsRoute.ts";
 import { statsRoute } from "./statsRoute.ts";
 import { usersRoute } from "./usersRoute.ts";
+import { workersRoute } from "./workersRoute.ts";
 
 export const serveApi = createLoggerMiddleware(
   createPathFilter({
@@ -12,6 +13,7 @@ export const serveApi = createLoggerMiddleware(
     "users": usersRoute,
     "settings/params": paramsRoute,
     "stats": statsRoute,
+    "workers": workersRoute,
   }),
   { filterStatus: (status) => status >= 400 },
 );
diff --git a/api/workersRoute.ts b/api/workersRoute.ts
new file mode 100644
index 0000000..a1227cc
--- /dev/null
+++ b/api/workersRoute.ts
@@ -0,0 +1,103 @@
+import { createEndpoint, createMethodFilter, createPathFilter } from "t_rest/server";
+import { activeGenerationWorkers } from "../app/generationQueue.ts";
+import { getConfig } from "../app/config.ts";
+import * as SdApi from "../app/sdApi.ts";
+import createOpenApiFetch from "openapi_fetch";
+
+export const workersRoute = createPathFilter({
+  "": createMethodFilter({
+    "GET": createEndpoint(
+      { query: null, body: null },
+      async () => {
+        const activeWorkers = activeGenerationWorkers;
+        const { sdInstances } = await getConfig();
+
+        const workers = Object.entries(sdInstances).map(([sdInstanceId, sdInstance]) => ({
+          id: sdInstanceId,
+          name: sdInstance.name ?? sdInstanceId,
+          maxResolution: sdInstance.maxResolution,
+          active: activeWorkers.has(sdInstanceId),
+          lastOnline: null,
+          imagesPerMinute: null,
+          pixelsPerSecond: null,
+          pixelStepsPerSecond: null,
+        }));
+
+        return {
+          status: 200,
+          body: { type: "application/json", data: workers },
+        };
+      },
+    ),
+  }),
+
+  "{workerId}/loras": createMethodFilter({
+    GET: createEndpoint(
+      { query: null, body: null },
+      async ({ params }) => {
+        const { sdInstances } = await getConfig();
+        const sdInstance = sdInstances[params.workerId];
+        if (!sdInstance) {
+          return { status: 404, body: { type: "text/plain", data: `Worker not found` } };
+        }
+        const sdClient = createOpenApiFetch<SdApi.paths>({ baseUrl: sdInstance.api.url });
+        const lorasResponse = await sdClient.GET("/sdapi/v1/loras", {
+          headers: sdInstance.api.auth ? { Authorization: sdInstance.api.auth } : undefined,
+        });
+        if (lorasResponse.error) {
+          return {
+            status: 500,
+            body: { type: "text/plain", data: `Loras request failed: ${lorasResponse["error"]}` },
+          };
+        }
+        const loras = (lorasResponse.data as Lora[]).map((lora) => ({
+          name: lora.name,
+          alias: lora.alias ?? null,
+        }));
+        return {
+          status: 200,
+          body: { type: "application/json", data: loras },
+        };
+      },
+    ),
+  }),
+
+  "{workerId}/models": createMethodFilter({
+    GET: createEndpoint(
+      { query: null, body: null },
+      async ({ params }) => {
+        const { sdInstances } = await getConfig();
+        const sdInstance = sdInstances[params.workerId];
+        if (!sdInstance) {
+          return { status: 404, body: { type: "text/plain", data: `Worker not found` } };
+        }
+        const sdClient = createOpenApiFetch<SdApi.paths>({ baseUrl: sdInstance.api.url });
+        const modelsResponse = await sdClient.GET("/sdapi/v1/sd-models", {
+          headers: sdInstance.api.auth ? { Authorization: sdInstance.api.auth } : undefined,
+        });
+        if (modelsResponse.error) {
+          return {
+            status: 500,
+            body: { type: "text/plain", data: `Models request failed: ${modelsResponse["error"]}` },
+          };
+        }
+        const models = modelsResponse.data.map((model) => ({
+          title: model.title,
+          modelName: model.model_name,
+          hash: model.hash,
+          sha256: model.sha256,
+        }));
+        return {
+          status: 200,
+          body: { type: "application/json", data: models },
+        };
+      },
+    ),
+  }),
+});
+
+export interface Lora {
+  name: string;
+  alias: string | null;
+  metadata: object;
+}