eris/api/workersRoute.ts

244 lines
8.0 KiB
TypeScript
Raw Normal View History

2023-10-15 19:13:38 +00:00
import { subMinutes } from "date-fns";
import { Model } from "indexed_kv";
import createOpenApiFetch from "openapi_fetch";
import { info } from "std/log/mod.ts";
2023-10-11 01:59:52 +00:00
import { createEndpoint, createMethodFilter, createPathFilter } from "t_rest/server";
2023-10-15 19:13:38 +00:00
import { activeGenerationWorkers } from "../app/generationQueue.ts";
import { generationStore } from "../app/generationStore.ts";
2023-10-11 01:59:52 +00:00
import * as SdApi from "../app/sdApi.ts";
2023-10-17 13:03:14 +00:00
import {
WorkerInstance,
workerInstanceSchema,
workerInstanceStore,
} from "../app/workerInstanceStore.ts";
2023-10-13 11:47:57 +00:00
import { getAuthHeader } from "../utils/getAuthHeader.ts";
2023-10-19 21:37:03 +00:00
import { omitUndef } from "../utils/omitUndef.ts";
2023-10-23 00:39:01 +00:00
import { withAdmin } from "./withUser.ts";
2023-10-13 11:47:57 +00:00
export type WorkerData = Omit<WorkerInstance, "sdUrl" | "sdAuth"> & {
id: string;
isActive: boolean;
imagesPerMinute: number;
stepsPerMinute: number;
pixelsPerMinute: number;
pixelStepsPerMinute: number;
};
const STATS_INTERVAL_MIN = 10;
async function getWorkerData(workerInstance: Model<WorkerInstance>): Promise<WorkerData> {
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,
2023-11-10 02:50:02 +00:00
gpu: workerInstance.value.gpu,
2023-10-13 11:47:57 +00:00
lastError: workerInstance.value.lastError,
lastOnlineTime: workerInstance.value.lastOnlineTime,
isActive: activeGenerationWorkers.get(workerInstance.id)?.isProcessing ?? false,
imagesPerMinute,
stepsPerMinute,
pixelsPerMinute,
pixelStepsPerMinute,
};
}
2023-10-11 01:59:52 +00:00
export const workersRoute = createPathFilter({
"": createMethodFilter({
2023-10-13 11:47:57 +00:00
GET: createEndpoint(
2023-10-11 01:59:52 +00:00
{ query: null, body: null },
async () => {
2023-10-13 11:47:57 +00:00
const workerInstances = await workerInstanceStore.getAll();
const workers = await Promise.all(workerInstances.map(getWorkerData));
2023-10-11 01:59:52 +00:00
return {
status: 200,
2023-10-13 11:47:57 +00:00
body: { type: "application/json", data: workers satisfies WorkerData[] },
2023-10-11 01:59:52 +00:00
};
},
),
2023-10-13 11:47:57 +00:00
POST: createEndpoint(
{
query: {
sessionId: { type: "string" },
},
body: {
type: "application/json",
2023-10-17 13:03:14 +00:00
schema: workerInstanceSchema,
2023-10-13 11:47:57 +00:00
},
},
async ({ query, body }) => {
2023-10-23 00:39:01 +00:00
return withAdmin(query, async (user) => {
2023-10-17 13:03:14 +00:00
const workerInstance = await workerInstanceStore.create(body.data);
2023-10-23 00:39:01 +00:00
info(`User ${user.first_name} created worker ${workerInstance.value.name}`);
2023-10-17 13:03:14 +00:00
return {
status: 200,
body: { type: "application/json", data: await getWorkerData(workerInstance) },
};
2023-10-23 00:39:01 +00:00
});
2023-10-11 01:59:52 +00:00
},
),
}),
2023-10-13 11:47:57 +00:00
"{workerId}": createPathFilter({
"": createMethodFilter({
GET: createEndpoint(
{ query: null, body: null },
async ({ params }) => {
2023-10-19 21:37:03 +00:00
const workerInstance = await workerInstanceStore.getById(params.workerId!);
2023-10-13 11:47:57 +00:00
if (!workerInstance) {
return { status: 404, body: { type: "text/plain", data: `Worker not found` } };
}
2023-10-11 01:59:52 +00:00
return {
2023-10-13 11:47:57 +00:00
status: 200,
2023-10-17 13:03:14 +00:00
body: { type: "application/json", data: await getWorkerData(workerInstance) },
2023-10-11 01:59:52 +00:00
};
2023-10-13 11:47:57 +00:00
},
),
PATCH: createEndpoint(
{
query: {
sessionId: { type: "string" },
},
body: {
type: "application/json",
2023-10-17 13:03:14 +00:00
schema: { ...workerInstanceSchema, required: [] },
2023-10-13 11:47:57 +00:00
},
},
async ({ params, query, body }) => {
2023-10-19 21:37:03 +00:00
const workerInstance = await workerInstanceStore.getById(params.workerId!);
2023-10-13 11:47:57 +00:00
if (!workerInstance) {
return { status: 404, body: { type: "text/plain", data: `Worker not found` } };
}
2023-10-23 00:39:01 +00:00
return withAdmin(query, async (user) => {
2023-10-17 13:03:14 +00:00
info(
2023-10-23 00:39:01 +00:00
`User ${user.first_name} updated worker ${workerInstance.value.name}: ${
2023-10-17 13:03:14 +00:00
JSON.stringify(body.data)
}`,
);
2023-10-19 21:37:03 +00:00
await workerInstance.update(omitUndef(body.data));
2023-10-17 13:03:14 +00:00
return {
status: 200,
body: { type: "application/json", data: await getWorkerData(workerInstance) },
};
2023-10-23 00:39:01 +00:00
});
2023-10-13 11:47:57 +00:00
},
),
DELETE: createEndpoint(
{
query: {
sessionId: { type: "string" },
},
body: null,
},
async ({ params, query }) => {
2023-10-19 21:37:03 +00:00
const workerInstance = await workerInstanceStore.getById(params.workerId!);
2023-10-13 11:47:57 +00:00
if (!workerInstance) {
return { status: 404, body: { type: "text/plain", data: `Worker not found` } };
}
2023-10-23 00:39:01 +00:00
return withAdmin(query, async (user) => {
info(`User ${user.first_name} deleted worker ${workerInstance.value.name}`);
2023-10-17 13:03:14 +00:00
await workerInstance.delete();
return { status: 200, body: { type: "application/json", data: null } };
2023-10-23 00:39:01 +00:00
});
2023-10-13 11:47:57 +00:00
},
),
}),
"loras": createMethodFilter({
GET: createEndpoint(
{ query: null, body: null },
async ({ params }) => {
2023-10-19 21:37:03 +00:00
const workerInstance = await workerInstanceStore.getById(params.workerId!);
2023-10-13 11:47:57 +00:00
if (!workerInstance) {
return { status: 404, body: { type: "text/plain", data: `Worker not found` } };
}
const sdClient = createOpenApiFetch<SdApi.paths>({
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 }) => {
2023-10-19 21:37:03 +00:00
const workerInstance = await workerInstanceStore.getById(params.workerId!);
2023-10-13 11:47:57 +00:00
if (!workerInstance) {
return { status: 404, body: { type: "text/plain", data: `Worker not found` } };
}
const sdClient = createOpenApiFetch<SdApi.paths>({
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 },
};
},
),
}),
2023-10-11 01:59:52 +00:00
}),
});
export interface Lora {
name: string;
alias: string | null;
metadata: object;
}