import { subMinutes } from "date-fns"; import { Model } from "indexed_kv"; import createOpenApiFetch from "openapi_fetch"; import { info } from "std/log/mod.ts"; import { createEndpoint, createMethodFilter, createPathFilter } from "t_rest/server"; import { activeGenerationWorkers } from "../app/generationQueue.ts"; import { generationStore } from "../app/generationStore.ts"; import * as SdApi from "../app/sdApi.ts"; import { WorkerInstance, workerInstanceSchema, workerInstanceStore, } from "../app/workerInstanceStore.ts"; import { getAuthHeader } from "../utils/getAuthHeader.ts"; import { omitUndef } from "../utils/omitUndef.ts"; import { withAdmin } from "./withUser.ts"; export type WorkerData = Omit & { id: string; isActive: boolean; imagesPerMinute: number; stepsPerMinute: number; pixelsPerMinute: number; pixelStepsPerMinute: number; }; const STATS_INTERVAL_MIN = 10; async function getWorkerData(workerInstance: Model): Promise { const after = subMinutes(new Date(), STATS_INTERVAL_MIN); const generations = await generationStore.getBy("workerInstanceKey", { value: workerInstance.value.key, after: after, }); const imagesPerMinute = generations.length / STATS_INTERVAL_MIN; const stepsPerMinute = generations .map((generation) => generation.value.info?.steps ?? 0) .reduce((sum, steps) => sum + steps, 0) / STATS_INTERVAL_MIN; const pixelsPerMinute = generations .map((generation) => (generation.value.info?.width ?? 0) * (generation.value.info?.height ?? 0)) .reduce((sum, pixels) => sum + pixels, 0) / STATS_INTERVAL_MIN; const pixelStepsPerMinute = generations .map((generation) => (generation.value.info?.width ?? 0) * (generation.value.info?.height ?? 0) * (generation.value.info?.steps ?? 0) ) .reduce((sum, pixelSteps) => sum + pixelSteps, 0) / STATS_INTERVAL_MIN; return { id: workerInstance.id, key: workerInstance.value.key, name: workerInstance.value.name, gpu: workerInstance.value.gpu, lastError: workerInstance.value.lastError, lastOnlineTime: workerInstance.value.lastOnlineTime, isActive: activeGenerationWorkers.get(workerInstance.id)?.isProcessing ?? false, imagesPerMinute, stepsPerMinute, pixelsPerMinute, pixelStepsPerMinute, }; } export const workersRoute = createPathFilter({ "": createMethodFilter({ GET: createEndpoint( { query: null, body: null }, async () => { const workerInstances = await workerInstanceStore.getAll(); const workers = await Promise.all(workerInstances.map(getWorkerData)); return { status: 200, body: { type: "application/json", data: workers satisfies WorkerData[] }, }; }, ), POST: createEndpoint( { query: { sessionId: { type: "string" }, }, body: { type: "application/json", schema: workerInstanceSchema, }, }, async ({ query, body }) => { return withAdmin(query, async (user) => { const workerInstance = await workerInstanceStore.create(body.data); info(`User ${user.first_name} created worker ${workerInstance.value.name}`); return { status: 200, body: { type: "application/json", data: await getWorkerData(workerInstance) }, }; }); }, ), }), "{workerId}": createPathFilter({ "": createMethodFilter({ GET: createEndpoint( { query: null, body: null }, async ({ params }) => { const workerInstance = await workerInstanceStore.getById(params.workerId!); if (!workerInstance) { return { status: 404, body: { type: "text/plain", data: `Worker not found` } }; } return { status: 200, body: { type: "application/json", data: await getWorkerData(workerInstance) }, }; }, ), PATCH: createEndpoint( { query: { sessionId: { type: "string" }, }, body: { type: "application/json", schema: { ...workerInstanceSchema, required: [] }, }, }, async ({ params, query, body }) => { const workerInstance = await workerInstanceStore.getById(params.workerId!); if (!workerInstance) { return { status: 404, body: { type: "text/plain", data: `Worker not found` } }; } return withAdmin(query, async (user) => { info( `User ${user.first_name} updated worker ${workerInstance.value.name}: ${ JSON.stringify(body.data) }`, ); await workerInstance.update(omitUndef(body.data)); return { status: 200, body: { type: "application/json", data: await getWorkerData(workerInstance) }, }; }); }, ), DELETE: createEndpoint( { query: { sessionId: { type: "string" }, }, body: null, }, async ({ params, query }) => { const workerInstance = await workerInstanceStore.getById(params.workerId!); if (!workerInstance) { return { status: 404, body: { type: "text/plain", data: `Worker not found` } }; } return withAdmin(query, async (user) => { info(`User ${user.first_name} deleted worker ${workerInstance.value.name}`); await workerInstance.delete(); return { status: 200, body: { type: "application/json", data: null } }; }); }, ), }), "loras": createMethodFilter({ GET: createEndpoint( { query: null, body: null }, async ({ params }) => { const workerInstance = await workerInstanceStore.getById(params.workerId!); if (!workerInstance) { return { status: 404, body: { type: "text/plain", data: `Worker not found` } }; } const sdClient = createOpenApiFetch({ baseUrl: workerInstance.value.sdUrl, headers: getAuthHeader(workerInstance.value.sdAuth), }); const lorasResponse = await sdClient.GET("/sdapi/v1/loras", {}); 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 }, }; }, ), }), "models": createMethodFilter({ GET: createEndpoint( { query: null, body: null }, async ({ params }) => { const workerInstance = await workerInstanceStore.getById(params.workerId!); if (!workerInstance) { return { status: 404, body: { type: "text/plain", data: `Worker not found` } }; } const sdClient = createOpenApiFetch({ baseUrl: workerInstance.value.sdUrl, headers: getAuthHeader(workerInstance.value.sdAuth), }); const modelsResponse = await sdClient.GET("/sdapi/v1/sd-models", {}); 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; }