feat: manage workers via webui

This commit is contained in:
pinks 2023-10-13 13:47:57 +02:00
parent a251d5e965
commit 0b85c99c92
25 changed files with 974 additions and 256 deletions

View File

@ -13,10 +13,7 @@ You can put these in `.env` file or pass them as environment variables.
- `TG_BOT_TOKEN` - Telegram bot token. Get yours from [@BotFather](https://t.me/BotFather). - `TG_BOT_TOKEN` - Telegram bot token. Get yours from [@BotFather](https://t.me/BotFather).
Required. Required.
- `SD_API_URL` - URL to Stable Diffusion API. Only used on first run. Default:
`http://127.0.0.1:7860/`
- `TG_ADMIN_USERNAMES` - Comma separated list of usernames of users that can use admin commands. - `TG_ADMIN_USERNAMES` - Comma separated list of usernames of users that can use admin commands.
Only used on first run. Optional.
## Running ## Running

View File

@ -30,18 +30,15 @@ export const paramsRoute = createMethodFilter({
return { status: 401, body: { type: "text/plain", data: "Must be logged in" } }; return { status: 401, body: { type: "text/plain", data: "Must be logged in" } };
} }
const chat = await bot.api.getChat(session.userId); const chat = await bot.api.getChat(session.userId);
if (chat.type !== "private") { if (chat.type !== "private") throw new Error("Chat is not private");
throw new Error("Chat is not private"); if (!chat.username) {
}
const userName = chat.username;
if (!userName) {
return { status: 403, body: { type: "text/plain", data: "Must have a username" } }; return { status: 403, body: { type: "text/plain", data: "Must have a username" } };
} }
const config = await getConfig(); const config = await getConfig();
if (!config?.adminUsernames?.includes(userName)) { if (!config?.adminUsernames?.includes(chat.username)) {
return { status: 403, body: { type: "text/plain", data: "Must be an admin" } }; return { status: 403, body: { type: "text/plain", data: "Must be an admin" } };
} }
logger().info(`User ${userName} updated default params: ${JSON.stringify(body.data)}`); logger().info(`User ${chat.username} updated default params: ${JSON.stringify(body.data)}`);
const defaultParams = deepMerge(config.defaultParams ?? {}, body.data); const defaultParams = deepMerge(config.defaultParams ?? {}, body.data);
await setConfig({ defaultParams }); await setConfig({ defaultParams });
return { status: 200, body: { type: "application/json", data: config.defaultParams } }; return { status: 200, body: { type: "application/json", data: config.defaultParams } };

View File

@ -1,25 +1,54 @@
// deno-lint-ignore-file require-await
import { createEndpoint, createMethodFilter, createPathFilter } from "t_rest/server"; import { createEndpoint, createMethodFilter, createPathFilter } from "t_rest/server";
import { liveGlobalStats } from "../app/globalStatsStore.ts"; import { globalStats } from "../app/globalStats.ts";
import { getDailyStats } from "../app/dailyStatsStore.ts"; import { getDailyStats } from "../app/dailyStatsStore.ts";
import { getUserStats } from "../app/userStatsStore.ts"; import { getUserStats } from "../app/userStatsStore.ts";
import { getUserDailyStats } from "../app/userDailyStatsStore.ts"; import { getUserDailyStats } from "../app/userDailyStatsStore.ts";
import { generationStore } from "../app/generationStore.ts";
import { subMinutes } from "date-fns";
const STATS_INTERVAL_MIN = 3;
export const statsRoute = createPathFilter({ export const statsRoute = createPathFilter({
"": createMethodFilter({ "": createMethodFilter({
GET: createEndpoint( GET: createEndpoint(
{ query: null, body: null }, { query: null, body: null },
async () => { async () => {
const stats = liveGlobalStats; const after = subMinutes(new Date(), STATS_INTERVAL_MIN);
const generations = await generationStore.getAll({ 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 { return {
status: 200, status: 200,
body: { body: {
type: "application/json", type: "application/json",
data: { data: {
imageCount: stats.imageCount, imageCount: globalStats.imageCount,
pixelCount: stats.pixelCount, stepCount: globalStats.stepCount,
userCount: stats.userIds.length, pixelCount: globalStats.pixelCount,
timestamp: stats.timestamp, pixelStepCount: globalStats.pixelStepCount,
userCount: globalStats.userIds.length,
imagesPerMinute,
stepsPerMinute,
pixelsPerMinute,
pixelStepsPerMinute,
}, },
}, },
}; };

View File

@ -3,26 +3,6 @@ import { getConfig } from "../app/config.ts";
import { bot } from "../bot/mod.ts"; import { bot } from "../bot/mod.ts";
export const usersRoute = createPathFilter({ export const usersRoute = createPathFilter({
"{userId}": createMethodFilter({
GET: createEndpoint(
{ query: null, body: null },
async ({ params }) => {
const chat = await bot.api.getChat(params.userId);
if (chat.type !== "private") {
throw new Error("Chat is not private");
}
const config = await getConfig();
const isAdmin = chat.username && config?.adminUsernames?.includes(chat.username);
return {
status: 200,
body: {
type: "application/json",
data: { ...chat, isAdmin },
},
};
},
),
}),
"{userId}/photo": createMethodFilter({ "{userId}/photo": createMethodFilter({
GET: createEndpoint( GET: createEndpoint(
{ query: null, body: null }, { query: null, body: null },
@ -51,4 +31,25 @@ export const usersRoute = createPathFilter({
}, },
), ),
}), }),
"{userId}": createMethodFilter({
GET: createEndpoint(
{ query: null, body: null },
async ({ params }) => {
const chat = await bot.api.getChat(params.userId);
if (chat.type !== "private") {
throw new Error("Chat is not private");
}
const config = await getConfig();
const isAdmin = chat.username && config?.adminUsernames?.includes(chat.username);
return {
status: 200,
body: {
type: "application/json",
data: { ...chat, isAdmin },
},
};
},
),
}),
}); });

View File

@ -3,47 +3,261 @@ import { activeGenerationWorkers } from "../app/generationQueue.ts";
import { getConfig } from "../app/config.ts"; import { getConfig } from "../app/config.ts";
import * as SdApi from "../app/sdApi.ts"; import * as SdApi from "../app/sdApi.ts";
import createOpenApiFetch from "openapi_fetch"; import createOpenApiFetch from "openapi_fetch";
import { sessions } from "./sessionsRoute.ts";
import { bot } from "../bot/mod.ts";
import { getLogger } from "std/log/mod.ts";
import { WorkerInstance, workerInstanceStore } from "../app/workerInstanceStore.ts";
import { getAuthHeader } from "../utils/getAuthHeader.ts";
import { Model } from "indexed_kv";
import { generationStore } from "../app/generationStore.ts";
import { subMinutes } from "date-fns";
const logger = () => getLogger();
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,
lastError: workerInstance.value.lastError,
lastOnlineTime: workerInstance.value.lastOnlineTime,
isActive: activeGenerationWorkers.get(workerInstance.id)?.isProcessing ?? false,
imagesPerMinute,
stepsPerMinute,
pixelsPerMinute,
pixelStepsPerMinute,
};
}
export const workersRoute = createPathFilter({ export const workersRoute = createPathFilter({
"": createMethodFilter({ "": createMethodFilter({
"GET": createEndpoint( GET: createEndpoint(
{ query: null, body: null }, { query: null, body: null },
async () => { async () => {
const activeWorkers = activeGenerationWorkers; const workerInstances = await workerInstanceStore.getAll();
const { sdInstances } = await getConfig(); const workers = await Promise.all(workerInstances.map(getWorkerData));
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 { return {
status: 200, status: 200,
body: { type: "application/json", data: workers }, body: { type: "application/json", data: workers satisfies WorkerData[] },
};
},
),
POST: createEndpoint(
{
query: {
sessionId: { type: "string" },
},
body: {
type: "application/json",
schema: {
type: "object",
properties: {
key: { type: "string" },
name: { type: ["string", "null"] },
sdUrl: { type: "string" },
sdAuth: {
type: ["object", "null"],
properties: {
user: { type: "string" },
password: { type: "string" },
},
required: ["user", "password"],
},
},
required: ["key", "name", "sdUrl", "sdAuth"],
},
},
},
async ({ query, body }) => {
const session = sessions.get(query.sessionId);
if (!session?.userId) {
return { status: 401, body: { type: "text/plain", data: "Must be logged in" } };
}
const chat = await bot.api.getChat(session.userId);
if (chat.type !== "private") throw new Error("Chat is not private");
if (!chat.username) {
return { status: 403, body: { type: "text/plain", data: "Must have a username" } };
}
const config = await getConfig();
if (!config?.adminUsernames?.includes(chat.username)) {
return { status: 403, body: { type: "text/plain", data: "Must be an admin" } };
}
const workerInstance = await workerInstanceStore.create({
key: body.data.key,
name: body.data.name,
sdUrl: body.data.sdUrl,
sdAuth: body.data.sdAuth,
});
logger().info(`User ${chat.username} created worker ${workerInstance.id}`);
const worker = await getWorkerData(workerInstance);
return {
status: 200,
body: { type: "application/json", data: worker satisfies WorkerData },
}; };
}, },
), ),
}), }),
"{workerId}/loras": createMethodFilter({ "{workerId}": createPathFilter({
"": createMethodFilter({
GET: createEndpoint( GET: createEndpoint(
{ query: null, body: null }, { query: null, body: null },
async ({ params }) => { async ({ params }) => {
const { sdInstances } = await getConfig(); const workerInstance = await workerInstanceStore.getById(params.workerId);
const sdInstance = sdInstances[params.workerId]; if (!workerInstance) {
if (!sdInstance) {
return { status: 404, body: { type: "text/plain", data: `Worker not found` } }; return { status: 404, body: { type: "text/plain", data: `Worker not found` } };
} }
const sdClient = createOpenApiFetch<SdApi.paths>({ baseUrl: sdInstance.api.url }); const worker: WorkerData = await getWorkerData(workerInstance);
const lorasResponse = await sdClient.GET("/sdapi/v1/loras", { return {
headers: sdInstance.api.auth ? { Authorization: sdInstance.api.auth } : undefined, status: 200,
body: { type: "application/json", data: worker satisfies WorkerData },
};
},
),
PATCH: createEndpoint(
{
query: {
sessionId: { type: "string" },
},
body: {
type: "application/json",
schema: {
type: "object",
properties: {
key: { type: "string" },
name: { type: ["string", "null"] },
sdUrl: { type: "string" },
auth: {
type: ["object", "null"],
properties: {
user: { type: "string" },
password: { type: "string" },
},
required: ["user", "password"],
},
},
},
},
},
async ({ params, query, body }) => {
const workerInstance = await workerInstanceStore.getById(params.workerId);
if (!workerInstance) {
return { status: 404, body: { type: "text/plain", data: `Worker not found` } };
}
const session = sessions.get(query.sessionId);
if (!session?.userId) {
return { status: 401, body: { type: "text/plain", data: "Must be logged in" } };
}
const chat = await bot.api.getChat(session.userId);
if (chat.type !== "private") throw new Error("Chat is not private");
if (!chat.username) {
return { status: 403, body: { type: "text/plain", data: "Must have a username" } };
}
const config = await getConfig();
if (!config?.adminUsernames?.includes(chat.username)) {
return { status: 403, body: { type: "text/plain", data: "Must be an admin" } };
}
if (body.data.name !== undefined) {
workerInstance.value.name = body.data.name;
}
if (body.data.sdUrl !== undefined) {
workerInstance.value.sdUrl = body.data.sdUrl;
}
if (body.data.auth !== undefined) {
workerInstance.value.sdAuth = body.data.auth;
}
logger().info(
`User ${chat.username} updated worker ${params.workerId}: ${JSON.stringify(body.data)}`,
);
await workerInstance.update();
const worker = await getWorkerData(workerInstance);
return {
status: 200,
body: { type: "application/json", data: worker satisfies WorkerData },
};
},
),
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` } };
}
const session = sessions.get(query.sessionId);
if (!session?.userId) {
return { status: 401, body: { type: "text/plain", data: "Must be logged in" } };
}
const chat = await bot.api.getChat(session.userId);
if (chat.type !== "private") throw new Error("Chat is not private");
if (!chat.username) {
return { status: 403, body: { type: "text/plain", data: "Must have a username" } };
}
const config = await getConfig();
if (!config?.adminUsernames?.includes(chat.username)) {
return { status: 403, body: { type: "text/plain", data: "Must be an admin" } };
}
logger().info(`User ${chat.username} deleted worker ${params.workerId}`);
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<SdApi.paths>({
baseUrl: workerInstance.value.sdUrl,
headers: getAuthHeader(workerInstance.value.sdAuth),
}); });
const lorasResponse = await sdClient.GET("/sdapi/v1/loras", {});
if (lorasResponse.error) { if (lorasResponse.error) {
return { return {
status: 500, status: 500,
@ -62,23 +276,26 @@ export const workersRoute = createPathFilter({
), ),
}), }),
"{workerId}/models": createMethodFilter({ "models": createMethodFilter({
GET: createEndpoint( GET: createEndpoint(
{ query: null, body: null }, { query: null, body: null },
async ({ params }) => { async ({ params }) => {
const { sdInstances } = await getConfig(); const workerInstance = await workerInstanceStore.getById(params.workerId);
const sdInstance = sdInstances[params.workerId]; if (!workerInstance) {
if (!sdInstance) {
return { status: 404, body: { type: "text/plain", data: `Worker not found` } }; return { status: 404, body: { type: "text/plain", data: `Worker not found` } };
} }
const sdClient = createOpenApiFetch<SdApi.paths>({ baseUrl: sdInstance.api.url }); const sdClient = createOpenApiFetch<SdApi.paths>({
const modelsResponse = await sdClient.GET("/sdapi/v1/sd-models", { baseUrl: workerInstance.value.sdUrl,
headers: sdInstance.api.auth ? { Authorization: sdInstance.api.auth } : undefined, headers: getAuthHeader(workerInstance.value.sdAuth),
}); });
const modelsResponse = await sdClient.GET("/sdapi/v1/sd-models", {});
if (modelsResponse.error) { if (modelsResponse.error) {
return { return {
status: 500, status: 500,
body: { type: "text/plain", data: `Models request failed: ${modelsResponse["error"]}` }, body: {
type: "text/plain",
data: `Models request failed: ${modelsResponse["error"]}`,
},
}; };
} }
const models = modelsResponse.data.map((model) => ({ const models = modelsResponse.data.map((model) => ({
@ -94,6 +311,7 @@ export const workersRoute = createPathFilter({
}, },
), ),
}), }),
}),
}); });
export interface Lora { export interface Lora {

View File

@ -21,27 +21,8 @@ export const configSchema = {
negative_prompt: { type: "string" }, negative_prompt: { type: "string" },
}, },
}, },
sdInstances: {
type: "object",
additionalProperties: {
type: "object",
properties: {
name: { type: "string" },
api: {
type: "object",
properties: {
url: { type: "string" },
auth: { type: "string" },
}, },
required: ["url"], required: ["adminUsernames", "maxUserJobs", "maxJobs", "defaultParams"],
},
maxResolution: { type: "number" },
},
required: ["api", "maxResolution"],
},
},
},
required: ["adminUsernames", "maxUserJobs", "maxJobs", "defaultParams", "sdInstances"],
} as const satisfies JsonSchema; } as const satisfies JsonSchema;
export type Config = jsonType<typeof configSchema>; export type Config = jsonType<typeof configSchema>;
@ -55,13 +36,6 @@ export async function getConfig(): Promise<Config> {
maxUserJobs: config?.maxUserJobs ?? Infinity, maxUserJobs: config?.maxUserJobs ?? Infinity,
maxJobs: config?.maxJobs ?? Infinity, maxJobs: config?.maxJobs ?? Infinity,
defaultParams: config?.defaultParams ?? {}, defaultParams: config?.defaultParams ?? {},
sdInstances: config?.sdInstances ??
{
"local": {
api: { url: Deno.env.get("SD_API_URL") ?? "http://127.0.0.1:7860/" },
maxResolution: 1024 * 1024,
},
},
}; };
} }
@ -73,7 +47,6 @@ export async function setConfig(newConfig: Partial<Config>): Promise<void> {
maxUserJobs: newConfig.maxUserJobs ?? oldConfig.maxUserJobs, maxUserJobs: newConfig.maxUserJobs ?? oldConfig.maxUserJobs,
maxJobs: newConfig.maxJobs ?? oldConfig.maxJobs, maxJobs: newConfig.maxJobs ?? oldConfig.maxJobs,
defaultParams: newConfig.defaultParams ?? oldConfig.defaultParams, defaultParams: newConfig.defaultParams ?? oldConfig.defaultParams,
sdInstances: newConfig.sdInstances ?? oldConfig.sdInstances,
}; };
await db.set(["config"], config); await db.set(["config"], config);
} }

View File

@ -13,10 +13,12 @@ export const dailyStatsSchema = {
properties: { properties: {
userIds: { type: "array", items: { type: "number" } }, userIds: { type: "array", items: { type: "number" } },
imageCount: { type: "number" }, imageCount: { type: "number" },
stepCount: { type: "number" },
pixelCount: { type: "number" }, pixelCount: { type: "number" },
pixelStepCount: { type: "number" },
timestamp: { type: "number" }, timestamp: { type: "number" },
}, },
required: ["userIds", "imageCount", "pixelCount", "timestamp"], required: ["userIds", "imageCount", "stepCount", "pixelCount", "pixelStepCount", "timestamp"],
} as const satisfies JsonSchema; } as const satisfies JsonSchema;
export type DailyStats = jsonType<typeof dailyStatsSchema>; export type DailyStats = jsonType<typeof dailyStatsSchema>;
@ -27,7 +29,9 @@ export const getDailyStats = kvMemoize(
async (year: number, month: number, day: number): Promise<DailyStats> => { async (year: number, month: number, day: number): Promise<DailyStats> => {
const userIdSet = new Set<number>(); const userIdSet = new Set<number>();
let imageCount = 0; let imageCount = 0;
let stepCount = 0;
let pixelCount = 0; let pixelCount = 0;
let pixelStepCount = 0;
const after = new Date(Date.UTC(year, month - 1, day)); const after = new Date(Date.UTC(year, month - 1, day));
const before = new Date(Date.UTC(year, month - 1, day + 1)); const before = new Date(Date.UTC(year, month - 1, day + 1));
@ -39,13 +43,19 @@ export const getDailyStats = kvMemoize(
) { ) {
userIdSet.add(generation.value.from.id); userIdSet.add(generation.value.from.id);
imageCount++; imageCount++;
stepCount += generation.value.info?.steps ?? 0;
pixelCount += (generation.value.info?.width ?? 0) * (generation.value.info?.height ?? 0); pixelCount += (generation.value.info?.width ?? 0) * (generation.value.info?.height ?? 0);
pixelStepCount += (generation.value.info?.width ?? 0) *
(generation.value.info?.height ?? 0) *
(generation.value.info?.steps ?? 0);
} }
return { return {
userIds: [...userIdSet], userIds: [...userIdSet],
imageCount, imageCount,
stepCount,
pixelCount, pixelCount,
pixelStepCount,
timestamp: Date.now(), timestamp: Date.now(),
}; };
}, },

View File

@ -10,12 +10,14 @@ import { bot } from "../bot/mod.ts";
import { PngInfo } from "../bot/parsePngInfo.ts"; import { PngInfo } from "../bot/parsePngInfo.ts";
import { formatOrdinal } from "../utils/formatOrdinal.ts"; import { formatOrdinal } from "../utils/formatOrdinal.ts";
import { formatUserChat } from "../utils/formatUserChat.ts"; import { formatUserChat } from "../utils/formatUserChat.ts";
import { getAuthHeader } from "../utils/getAuthHeader.ts";
import { SdError } from "./SdError.ts"; import { SdError } from "./SdError.ts";
import { getConfig } from "./config.ts"; import { getConfig } from "./config.ts";
import { db, fs } from "./db.ts"; import { db, fs } from "./db.ts";
import { SdGenerationInfo } from "./generationStore.ts"; import { SdGenerationInfo } from "./generationStore.ts";
import * as SdApi from "./sdApi.ts"; import * as SdApi from "./sdApi.ts";
import { uploadQueue } from "./uploadQueue.ts"; import { uploadQueue } from "./uploadQueue.ts";
import { workerInstanceStore } from "./workerInstanceStore.ts";
const logger = () => getLogger(); const logger = () => getLogger();
@ -34,7 +36,7 @@ interface GenerationJob {
chat: Chat; chat: Chat;
requestMessage: Message; requestMessage: Message;
replyMessage: Message; replyMessage: Message;
sdInstanceId?: string; workerInstanceKey?: string;
progress?: number; progress?: number;
} }
@ -47,18 +49,17 @@ export const activeGenerationWorkers = new Map<string, Worker<GenerationJob>>();
*/ */
export async function processGenerationQueue() { export async function processGenerationQueue() {
while (true) { while (true) {
const config = await getConfig(); for await (const workerInstance of workerInstanceStore.listAll()) {
const activeWorker = activeGenerationWorkers.get(workerInstance.id);
for (const [sdInstanceId, sdInstance] of Object.entries(config?.sdInstances ?? {})) {
const activeWorker = activeGenerationWorkers.get(sdInstanceId);
if (activeWorker?.isProcessing) { if (activeWorker?.isProcessing) {
continue; continue;
} }
const workerSdClient = createOpenApiClient<SdApi.paths>({ const workerSdClient = createOpenApiClient<SdApi.paths>({
baseUrl: sdInstance.api.url, baseUrl: workerInstance.value.sdUrl,
headers: { "Authorization": sdInstance.api.auth }, headers: getAuthHeader(workerInstance.value.sdAuth),
}); });
// check if worker is up // check if worker is up
const activeWorkerStatus = await workerSdClient.GET("/sdapi/v1/memory", { const activeWorkerStatus = await workerSdClient.GET("/sdapi/v1/memory", {
signal: AbortSignal.timeout(10_000), signal: AbortSignal.timeout(10_000),
@ -70,16 +71,20 @@ export async function processGenerationQueue() {
return response; return response;
}) })
.catch((error) => { .catch((error) => {
logger().debug(`Worker ${sdInstanceId} is down: ${error}`); workerInstance.update({ lastError: { message: error.message, time: Date.now() } })
.catch(() => undefined);
logger().debug(`Worker ${workerInstance.value.key} is down: ${error}`);
}); });
if (!activeWorkerStatus?.data) { if (!activeWorkerStatus?.data) {
continue; continue;
} }
// create worker // create worker
const newWorker = generationQueue.createWorker(async ({ state }, updateJob) => { const newWorker = generationQueue.createWorker(async ({ state }, updateJob) => {
await processGenerationJob(state, updateJob, sdInstanceId); await processGenerationJob(state, updateJob, workerInstance.id);
}); });
newWorker.addEventListener("error", (e) => { newWorker.addEventListener("error", (e) => {
logger().error( logger().error(
`Generation failed for ${formatUserChat(e.detail.job.state)}: ${e.detail.error}`, `Generation failed for ${formatUserChat(e.detail.job.state)}: ${e.detail.error}`,
@ -96,11 +101,19 @@ export async function processGenerationQueue() {
}, },
).catch(() => undefined); ).catch(() => undefined);
newWorker.stopProcessing(); newWorker.stopProcessing();
logger().info(`Stopped worker ${sdInstanceId}`); workerInstance.update({ lastError: { message: e.detail.error.message, time: Date.now() } })
.catch(() => undefined);
logger().info(`Stopped worker ${workerInstance.value.key}`);
}); });
newWorker.addEventListener("complete", () => {
workerInstance.update({ lastOnlineTime: Date.now() }).catch(() => undefined);
});
await workerInstance.update({ lastOnlineTime: Date.now() });
newWorker.processJobs(); newWorker.processJobs();
activeGenerationWorkers.set(sdInstanceId, newWorker); activeGenerationWorkers.set(workerInstance.id, newWorker);
logger().info(`Started worker ${sdInstanceId}`); logger().info(`Started worker ${workerInstance.value.key}`);
} }
await delay(60_000); await delay(60_000);
} }
@ -112,19 +125,19 @@ export async function processGenerationQueue() {
async function processGenerationJob( async function processGenerationJob(
state: GenerationJob, state: GenerationJob,
updateJob: (job: Partial<JobData<GenerationJob>>) => Promise<void>, updateJob: (job: Partial<JobData<GenerationJob>>) => Promise<void>,
sdInstanceId: string, workerInstanceId: string,
) { ) {
const startDate = new Date(); const startDate = new Date();
const config = await getConfig(); const config = await getConfig();
const sdInstance = config?.sdInstances?.[sdInstanceId]; const workerInstance = await workerInstanceStore.getById(workerInstanceId);
if (!sdInstance) { if (!workerInstance) {
throw new Error(`Unknown sdInstanceId: ${sdInstanceId}`); throw new Error(`Unknown workerInstanceId: ${workerInstanceId}`);
} }
const workerSdClient = createOpenApiClient<SdApi.paths>({ const workerSdClient = createOpenApiClient<SdApi.paths>({
baseUrl: sdInstance.api.url, baseUrl: workerInstance.value.sdUrl,
headers: { "Authorization": sdInstance.api.auth }, headers: getAuthHeader(workerInstance.value.sdAuth),
}); });
state.sdInstanceId = sdInstanceId; state.workerInstanceKey = workerInstance.value.key;
state.progress = 0; state.progress = 0;
logger().debug(`Generation started for ${formatUserChat(state)}`); logger().debug(`Generation started for ${formatUserChat(state)}`);
await updateJob({ state: state }); await updateJob({ state: state });
@ -142,14 +155,16 @@ async function processGenerationJob(
await bot.api.editMessageText( await bot.api.editMessageText(
state.replyMessage.chat.id, state.replyMessage.chat.id,
state.replyMessage.message_id, state.replyMessage.message_id,
`Generating your prompt now... 0% using ${sdInstance.name || sdInstanceId}`, `Generating your prompt now... 0% using ${
workerInstance.value.name || workerInstance.value.key
}`,
{ maxAttempts: 1 }, { maxAttempts: 1 },
).catch(() => undefined); ).catch(() => undefined);
// reduce size if worker can't handle the resolution // reduce size if worker can't handle the resolution
const size = limitSize( const size = limitSize(
{ ...config.defaultParams, ...state.task.params }, { ...config.defaultParams, ...state.task.params },
sdInstance.maxResolution, 1024 * 1024,
); );
function limitSize( function limitSize(
{ width, height }: { width?: number; height?: number }, { width, height }: { width?: number; height?: number },
@ -233,7 +248,7 @@ async function processGenerationJob(
state.replyMessage.message_id, state.replyMessage.message_id,
`Generating your prompt now... ${ `Generating your prompt now... ${
(progressResponse.data.progress * 100).toFixed(0) (progressResponse.data.progress * 100).toFixed(0)
}% using ${sdInstance.name || sdInstanceId}`, }% using ${workerInstance.value.name || workerInstance.value.key}`,
{ maxAttempts: 1 }, { maxAttempts: 1 },
).catch(() => undefined); ).catch(() => undefined);
} }
@ -268,7 +283,7 @@ async function processGenerationJob(
from: state.from, from: state.from,
requestMessage: state.requestMessage, requestMessage: state.requestMessage,
replyMessage: state.replyMessage, replyMessage: state.replyMessage,
sdInstanceId: sdInstanceId, workerInstanceKey: workerInstance.value.key,
startDate, startDate,
endDate: new Date(), endDate: new Date(),
imageKeys, imageKeys,

View File

@ -5,7 +5,7 @@ import { db } from "./db.ts";
export interface GenerationSchema { export interface GenerationSchema {
from: User; from: User;
chat: Chat; chat: Chat;
sdInstanceId?: string; sdInstanceId?: string; // TODO: change to workerInstanceKey
info?: SdGenerationInfo; info?: SdGenerationInfo;
startDate?: Date; startDate?: Date;
endDate?: Date; endDate?: Date;
@ -48,6 +48,7 @@ export interface SdGenerationInfo {
type GenerationIndices = { type GenerationIndices = {
fromId: number; fromId: number;
chatId: number; chatId: number;
workerInstanceKey: string;
}; };
export const generationStore = new Store<GenerationSchema, GenerationIndices>( export const generationStore = new Store<GenerationSchema, GenerationIndices>(
@ -57,6 +58,7 @@ export const generationStore = new Store<GenerationSchema, GenerationIndices>(
indices: { indices: {
fromId: { getValue: (item) => item.from.id }, fromId: { getValue: (item) => item.from.id },
chatId: { getValue: (item) => item.chat.id }, chatId: { getValue: (item) => item.chat.id },
workerInstanceKey: { getValue: (item) => item.sdInstanceId ?? "" },
}, },
}, },
); );

View File

@ -9,17 +9,19 @@ export const globalStatsSchema = {
properties: { properties: {
userIds: { type: "array", items: { type: "number" } }, userIds: { type: "array", items: { type: "number" } },
imageCount: { type: "number" }, imageCount: { type: "number" },
stepCount: { type: "number" },
pixelCount: { type: "number" }, pixelCount: { type: "number" },
pixelStepCount: { type: "number" },
timestamp: { type: "number" }, timestamp: { type: "number" },
}, },
required: ["userIds", "imageCount", "pixelCount", "timestamp"], required: ["userIds", "imageCount", "stepCount", "pixelCount", "pixelStepCount", "timestamp"],
} as const satisfies JsonSchema; } as const satisfies JsonSchema;
export type GlobalStats = jsonType<typeof globalStatsSchema>; export type GlobalStats = jsonType<typeof globalStatsSchema>;
export const liveGlobalStats: GlobalStats = await getGlobalStats(); export const globalStats: GlobalStats = await getGlobalStats();
export async function getGlobalStats(): Promise<GlobalStats> { async function getGlobalStats(): Promise<GlobalStats> {
// find the year/month/day of the first generation // find the year/month/day of the first generation
const startDate = await generationStore.getAll({}, { limit: 1 }) const startDate = await generationStore.getAll({}, { limit: 1 })
.then((generations) => generations[0]?.id) .then((generations) => generations[0]?.id)
@ -28,7 +30,9 @@ export async function getGlobalStats(): Promise<GlobalStats> {
// iterate to today and sum up stats // iterate to today and sum up stats
const userIdSet = new Set<number>(); const userIdSet = new Set<number>();
let imageCount = 0; let imageCount = 0;
let stepCount = 0;
let pixelCount = 0; let pixelCount = 0;
let pixelStepCount = 0;
const tomorrow = addDays(new Date(), 1); const tomorrow = addDays(new Date(), 1);
@ -44,13 +48,17 @@ export async function getGlobalStats(): Promise<GlobalStats> {
); );
for (const userId of dailyStats.userIds) userIdSet.add(userId); for (const userId of dailyStats.userIds) userIdSet.add(userId);
imageCount += dailyStats.imageCount; imageCount += dailyStats.imageCount;
stepCount += dailyStats.stepCount;
pixelCount += dailyStats.pixelCount; pixelCount += dailyStats.pixelCount;
pixelStepCount += dailyStats.pixelStepCount;
} }
return { return {
userIds: [...userIdSet], userIds: [...userIdSet],
imageCount, imageCount,
stepCount,
pixelCount, pixelCount,
pixelStepCount,
timestamp: Date.now(), timestamp: Date.now(),
}; };
} }

View File

@ -9,7 +9,7 @@ import { bot } from "../bot/mod.ts";
import { formatUserChat } from "../utils/formatUserChat.ts"; import { formatUserChat } from "../utils/formatUserChat.ts";
import { db, fs } from "./db.ts"; import { db, fs } from "./db.ts";
import { generationStore, SdGenerationInfo } from "./generationStore.ts"; import { generationStore, SdGenerationInfo } from "./generationStore.ts";
import { liveGlobalStats } from "./globalStatsStore.ts"; import { globalStats } from "./globalStats.ts";
const logger = () => getLogger(); const logger = () => getLogger();
@ -18,7 +18,7 @@ interface UploadJob {
chat: Chat; chat: Chat;
requestMessage: Message; requestMessage: Message;
replyMessage: Message; replyMessage: Message;
sdInstanceId: string; workerInstanceKey?: string;
startDate: Date; startDate: Date;
endDate: Date; endDate: Date;
imageKeys: Deno.KvKey[]; imageKeys: Deno.KvKey[];
@ -56,7 +56,7 @@ export async function processUploadQueue() {
fmt`${bold("CFG scale:")} ${state.info.cfg_scale}, `, fmt`${bold("CFG scale:")} ${state.info.cfg_scale}, `,
fmt`${bold("Seed:")} ${state.info.seed}, `, fmt`${bold("Seed:")} ${state.info.seed}, `,
fmt`${bold("Size")}: ${state.info.width}x${state.info.height}, `, fmt`${bold("Size")}: ${state.info.width}x${state.info.height}, `,
fmt`${bold("Worker")}: ${state.sdInstanceId}, `, state.workerInstanceKey ? fmt`${bold("Worker")}: ${state.workerInstanceKey}, ` : "",
fmt`${bold("Time taken")}: ${format(jobDurationMs, { ignoreZero: true })}`, fmt`${bold("Time taken")}: ${format(jobDurationMs, { ignoreZero: true })}`,
] ]
: [], : [],
@ -104,7 +104,7 @@ export async function processUploadQueue() {
await generationStore.create({ await generationStore.create({
from: state.from, from: state.from,
chat: state.chat, chat: state.chat,
sdInstanceId: state.sdInstanceId, sdInstanceId: state.workerInstanceKey,
startDate: state.startDate, startDate: state.startDate,
endDate: new Date(), endDate: new Date(),
info: state.info, info: state.info,
@ -112,11 +112,13 @@ export async function processUploadQueue() {
// update live stats // update live stats
{ {
liveGlobalStats.imageCount++; globalStats.imageCount++;
liveGlobalStats.pixelCount += state.info.width * state.info.height; globalStats.stepCount += state.info.steps;
const userIdSet = new Set(liveGlobalStats.userIds); globalStats.pixelCount += state.info.width * state.info.height;
globalStats.pixelStepCount += state.info.width * state.info.height * state.info.steps;
const userIdSet = new Set(globalStats.userIds);
userIdSet.add(state.from.id); userIdSet.add(state.from.id);
liveGlobalStats.userIds = [...userIdSet]; globalStats.userIds = [...userIdSet];
} }
// delete the status message // delete the status message

View File

@ -13,14 +13,24 @@ export const userStatsSchema = {
properties: { properties: {
userId: { type: "number" }, userId: { type: "number" },
imageCount: { type: "number" }, imageCount: { type: "number" },
stepCount: { type: "number" },
pixelCount: { type: "number" }, pixelCount: { type: "number" },
pixelStepCount: { type: "number" },
tagCountMap: { tagCountMap: {
type: "object", type: "object",
additionalProperties: { type: "number" }, additionalProperties: { type: "number" },
}, },
timestamp: { type: "number" }, timestamp: { type: "number" },
}, },
required: ["userId", "imageCount", "pixelCount", "tagCountMap", "timestamp"], required: [
"userId",
"imageCount",
"stepCount",
"pixelCount",
"pixelStepCount",
"tagCountMap",
"timestamp",
],
} as const satisfies JsonSchema; } as const satisfies JsonSchema;
export type UserStats = jsonType<typeof userStatsSchema>; export type UserStats = jsonType<typeof userStatsSchema>;
@ -48,7 +58,9 @@ export const getUserStats = kvMemoize(
["userStats"], ["userStats"],
async (userId: number): Promise<UserStats> => { async (userId: number): Promise<UserStats> => {
let imageCount = 0; let imageCount = 0;
let stepCount = 0;
let pixelCount = 0; let pixelCount = 0;
let pixelStepCount = 0;
const tagCountMap: Record<string, number> = {}; const tagCountMap: Record<string, number> = {};
logger().info(`Calculating user stats for ${userId}`); logger().info(`Calculating user stats for ${userId}`);
@ -57,7 +69,11 @@ export const getUserStats = kvMemoize(
const generation of generationStore.listBy("fromId", { value: userId }) const generation of generationStore.listBy("fromId", { value: userId })
) { ) {
imageCount++; imageCount++;
stepCount += generation.value.info?.steps ?? 0;
pixelCount += (generation.value.info?.width ?? 0) * (generation.value.info?.height ?? 0); pixelCount += (generation.value.info?.width ?? 0) * (generation.value.info?.height ?? 0);
pixelStepCount += (generation.value.info?.width ?? 0) *
(generation.value.info?.height ?? 0) *
(generation.value.info?.steps ?? 0);
const tags = generation.value.info?.prompt const tags = generation.value.info?.prompt
// split on punctuation and newlines // split on punctuation and newlines
@ -85,7 +101,9 @@ export const getUserStats = kvMemoize(
return { return {
userId, userId,
imageCount, imageCount,
stepCount,
pixelCount, pixelCount,
pixelStepCount,
tagCountMap, tagCountMap,
timestamp: Date.now(), timestamp: Date.now(),
}; };

View File

@ -0,0 +1,38 @@
import { Store } from "indexed_kv";
import { JsonSchema, jsonType } from "t_rest/server";
import { db } from "./db.ts";
export const workerInstanceSchema = {
type: "object",
properties: {
// used for counting stats
key: { type: "string" },
// used for display
name: { type: ["string", "null"] },
sdUrl: { type: "string" },
sdAuth: {
type: ["object", "null"],
properties: {
user: { type: "string" },
password: { type: "string" },
},
required: ["user", "password"],
},
lastOnlineTime: { type: "number" },
lastError: {
type: "object",
properties: {
message: { type: "string" },
time: { type: "number" },
},
required: ["message", "time"],
},
},
required: ["key", "name", "sdUrl", "sdAuth"],
} as const satisfies JsonSchema;
export type WorkerInstance = jsonType<typeof workerInstanceSchema>;
export const workerInstanceStore = new Store<WorkerInstance>(db, "workerInstances", {
indices: {},
});

View File

@ -1,9 +1,9 @@
import { CommandContext } from "grammy"; import { CommandContext } from "grammy";
import { bold, fmt } from "grammy_parse_mode"; import { bold, fmt } from "grammy_parse_mode";
import { getConfig } from "../app/config.ts";
import { activeGenerationWorkers, generationQueue } from "../app/generationQueue.ts"; import { activeGenerationWorkers, generationQueue } from "../app/generationQueue.ts";
import { getFlagEmoji } from "../utils/getFlagEmoji.ts"; import { getFlagEmoji } from "../utils/getFlagEmoji.ts";
import { ErisContext } from "./mod.ts"; import { ErisContext } from "./mod.ts";
import { workerInstanceStore } from "../app/workerInstanceStore.ts";
export async function queueCommand(ctx: CommandContext<ErisContext>) { export async function queueCommand(ctx: CommandContext<ErisContext>) {
let formattedMessage = await getMessageText(); let formattedMessage = await getMessageText();
@ -14,8 +14,8 @@ export async function queueCommand(ctx: CommandContext<ErisContext>) {
handleFutureUpdates().catch(() => undefined); handleFutureUpdates().catch(() => undefined);
async function getMessageText() { async function getMessageText() {
const config = await getConfig();
const allJobs = await generationQueue.getAllJobs(); const allJobs = await generationQueue.getAllJobs();
const workerInstances = await workerInstanceStore.getAll();
const processingJobs = allJobs const processingJobs = allJobs
.filter((job) => job.lockUntil > new Date()).map((job) => ({ ...job, index: 0 })); .filter((job) => job.lockUntil > new Date()).map((job) => ({ ...job, index: 0 }));
const waitingJobs = allJobs const waitingJobs = allJobs
@ -30,24 +30,17 @@ export async function queueCommand(ctx: CommandContext<ErisContext>) {
`${job.index}. `, `${job.index}. `,
fmt`${bold(job.state.from.first_name)} `, fmt`${bold(job.state.from.first_name)} `,
job.state.from.last_name ? fmt`${bold(job.state.from.last_name)} ` : "", job.state.from.last_name ? fmt`${bold(job.state.from.last_name)} ` : "",
job.state.from.username ? `(@${job.state.from.username}) ` : "",
getFlagEmoji(job.state.from.language_code) ?? "", getFlagEmoji(job.state.from.language_code) ?? "",
job.state.chat.type === "private" ? " in private chat " : ` in ${job.state.chat.title} `, job.index === 0 && job.state.progress && job.state.workerInstanceKey
job.state.chat.type !== "private" && job.state.chat.type !== "group" && ? `(${(job.state.progress * 100).toFixed(0)}% using ${job.state.workerInstanceKey}) `
job.state.chat.username
? `(@${job.state.chat.username}) `
: "",
job.index === 0 && job.state.progress && job.state.sdInstanceId
? `(${(job.state.progress * 100).toFixed(0)}% using ${job.state.sdInstanceId}) `
: "", : "",
"\n", "\n",
]) ])
: ["Queue is empty.\n"], : ["Queue is empty.\n"],
"\nActive workers:\n", "\nActive workers:\n",
...Object.entries(config.sdInstances).flatMap(([sdInstanceId, sdInstance]) => [ ...workerInstances.flatMap((workerInstace) => [
activeGenerationWorkers.get(sdInstanceId)?.isProcessing ? "✅ " : "☠️ ", activeGenerationWorkers.get(workerInstace.id)?.isProcessing ? "✅ " : "☠️ ",
fmt`${bold(sdInstance.name || sdInstanceId)} `, fmt`${bold(workerInstace.value.name || workerInstace.value.key)} `,
`(max ${(sdInstance.maxResolution / 1000000).toFixed(1)} Mpx) `,
"\n", "\n",
]), ]),
]); ]);

View File

@ -42,6 +42,7 @@
"react": "https://esm.sh/react@18.2.0?dev", "react": "https://esm.sh/react@18.2.0?dev",
"react-dom/client": "https://esm.sh/react-dom@18.2.0/client?dev", "react-dom/client": "https://esm.sh/react-dom@18.2.0/client?dev",
"react-router-dom": "https://esm.sh/react-router-dom@6.16.0?dev", "react-router-dom": "https://esm.sh/react-router-dom@6.16.0?dev",
"react-intl": "https://esm.sh/react-intl@6.4.7?external=react&alias=@types/react:react&dev",
"swr": "https://esm.sh/swr@2.2.4?dev", "swr": "https://esm.sh/swr@2.2.4?dev",
"swr/mutation": "https://esm.sh/swr@2.2.4/mutation?dev", "swr/mutation": "https://esm.sh/swr@2.2.4/mutation?dev",
"use-local-storage": "https://esm.sh/use-local-storage@3.0.0?dev", "use-local-storage": "https://esm.sh/use-local-storage@3.0.0?dev",

View File

@ -4,8 +4,9 @@ import useLocalStorage from "use-local-storage";
import { AppHeader } from "./AppHeader.tsx"; import { AppHeader } from "./AppHeader.tsx";
import { QueuePage } from "./QueuePage.tsx"; import { QueuePage } from "./QueuePage.tsx";
import { SettingsPage } from "./SettingsPage.tsx"; import { SettingsPage } from "./SettingsPage.tsx";
import { StatsPage } from "./StatsPage.tsx";
import { WorkersPage } from "./WorkersPage.tsx";
import { fetchApi, handleResponse } from "./apiClient.tsx"; import { fetchApi, handleResponse } from "./apiClient.tsx";
import { HomePage } from "./HomePage.tsx";
export function App() { export function App() {
// store session ID in the local storage // store session ID in the local storage
@ -30,7 +31,8 @@ export function App() {
/> />
<div className="self-center w-full max-w-screen-md flex flex-col items-stretch gap-4 p-4"> <div className="self-center w-full max-w-screen-md flex flex-col items-stretch gap-4 p-4">
<Routes> <Routes>
<Route path="/" element={<HomePage />} /> <Route path="/" element={<StatsPage />} />
<Route path="/workers" element={<WorkersPage sessionId={sessionId} />} />
<Route path="/queue" element={<QueuePage />} /> <Route path="/queue" element={<QueuePage />} />
<Route path="/settings" element={<SettingsPage sessionId={sessionId} />} /> <Route path="/settings" element={<SettingsPage sessionId={sessionId} />} />
</Routes> </Routes>

View File

@ -57,6 +57,9 @@ export function AppHeader(
<NavTab to="/"> <NavTab to="/">
Stats Stats
</NavTab> </NavTab>
<NavTab to="/workers">
Workers
</NavTab>
<NavTab to="/queue"> <NavTab to="/queue">
Queue Queue
</NavTab> </NavTab>

View File

@ -2,18 +2,18 @@ import { cx } from "@twind/core";
import React from "react"; import React from "react";
function CounterDigit(props: { value: number; transitionDurationMs?: number }) { function CounterDigit(props: { value: number; transitionDurationMs?: number }) {
const { value, transitionDurationMs = 1000 } = props; const { value, transitionDurationMs = 1500 } = props;
const rads = -(Math.floor(value) % 1_000_000) * 2 * Math.PI * 0.1; const rads = -(Math.floor(value) % 1_000_000) * 2 * Math.PI * 0.1;
return ( return (
<span className="w-[1em] h-[1.3em] relative"> <span className="w-[1em] relative">
{Array.from({ length: 10 }).map((_, i) => ( {Array.from({ length: 10 }).map((_, i) => (
<span <span
key={i} key={i}
className="absolute inset-0 grid place-items-center transition-transform duration-1000 ease-in-out [backface-visibility:hidden]" className="absolute inset-0 grid place-items-center transition-transform duration-1000 ease-in-out [backface-visibility:hidden]"
style={{ style={{
transform: `rotateX(${rads + i * 2 * Math.PI * 0.1}rad)`, transform: `rotateX(${rads + i * 2 * Math.PI * 0.1}rad)`,
transformOrigin: `center center 1.8em`, transformOrigin: `center center 2em`,
transitionDuration: `${transitionDurationMs}ms`, transitionDuration: `${transitionDurationMs}ms`,
}} }}
> >
@ -24,31 +24,35 @@ function CounterDigit(props: { value: number; transitionDurationMs?: number }) {
); );
} }
const Spacer = () =>
<span className="border-l-1 mx-[0.1em] border-zinc-200 dark:border-zinc-700" />;
const CounterText = (props: { children: React.ReactNode }) => (
<span className="self-center px-[0.1em]">
{props.children}
</span>
);
export function Counter(props: { export function Counter(props: {
value: number; value: number;
digits: number; digits: number;
fractionDigits?: number;
transitionDurationMs?: number; transitionDurationMs?: number;
className?: string; className?: string;
postfix?: string;
}) { }) {
const { value, digits, transitionDurationMs, className } = props; const { value, digits, fractionDigits = 0, transitionDurationMs, className, postfix } = props;
return ( return (
<span <span
className={cx( className={cx(
"inline-flex items-stretch border-1 border-zinc-300 dark:border-zinc-700 shadow-inner rounded-md overflow-hidden", "inline-flex h-[1.5em] items-stretch border-1 border-zinc-300 dark:border-zinc-700 shadow-inner rounded-md overflow-hidden",
className, className,
)} )}
> >
{Array.from({ length: digits }) {Array.from({ length: digits })
.flatMap((_, i) => [ .flatMap((_, i) => [
i > 0 && i % 3 === 0 i > 0 && i % 3 === 0 ? <Spacer key={`spacer${i}`} /> : null,
? (
<span
key={`spacer${i}`}
className="border-l-1 mx-[0.1em] border-zinc-200 dark:border-zinc-700"
/>
)
: null,
<CounterDigit <CounterDigit
key={`digit${i}`} key={`digit${i}`}
value={value / 10 ** i} value={value / 10 ** i}
@ -56,6 +60,30 @@ export function Counter(props: {
/>, />,
]) ])
.reverse()} .reverse()}
{fractionDigits > 0 && (
<>
<Spacer />
<CounterText>.</CounterText>
<Spacer />
{Array.from({ length: fractionDigits })
.flatMap((_, i) => [
i > 0 && i % 3 === 0 ? <Spacer key={`fractionSpacer${i}`} /> : null,
<CounterDigit
key={`fractionDigit${i}`}
value={value * 10 ** (i + 1)}
transitionDurationMs={transitionDurationMs}
/>,
])}
</>
)}
{postfix && (
<>
<Spacer />
<CounterText>{postfix}</CounterText>
</>
)}
</span> </span>
); );
} }

View File

@ -1,48 +0,0 @@
import React from "react";
import { fetchApi, handleResponse } from "./apiClient.tsx";
import useSWR from "swr";
import { eachDayOfInterval, endOfMonth, startOfMonth, subMonths } from "date-fns";
import { UTCDateMini } from "@date-fns/utc";
import { Counter } from "./Counter.tsx";
export function HomePage() {
const globalStats = useSWR(
["stats", "GET", {}] as const,
(args) => fetchApi(...args).then(handleResponse),
{ refreshInterval: 2_000 },
);
return (
<div className="my-16 flex flex-col gap-16 text-zinc-600 dark:text-zinc-400">
<p className="flex flex-col items-center gap-2 text-2xl sm:text-3xl">
<span>Pixels painted</span>
<Counter
className="font-bold text-zinc-700 dark:text-zinc-300"
value={globalStats.data?.pixelCount ?? 0}
digits={12}
transitionDurationMs={3_000}
/>
</p>
<div className="flex flex-col md:flex-row gap-16 md:gap-8">
<p className="flex-grow flex flex-col items-center gap-2 text-2xl">
<span>Images generated</span>
<Counter
className="font-bold text-zinc-700 dark:text-zinc-300"
value={globalStats.data?.imageCount ?? 0}
digits={9}
transitionDurationMs={1_500}
/>
</p>
<p className="flex-grow flex flex-col items-center gap-2 text-2xl">
<span>Unique users</span>
<Counter
className="font-bold text-zinc-700 dark:text-zinc-300"
value={globalStats.data?.userCount ?? 0}
digits={6}
transitionDurationMs={1_500}
/>
</p>
</div>
</div>
);
}

View File

@ -41,7 +41,7 @@ export function QueuePage() {
<Progress className="w-full h-full" value={job.state.progress} />} <Progress className="w-full h-full" value={job.state.progress} />}
</span> </span>
<span className="text-sm text-zinc-500 dark:text-zinc-400"> <span className="text-sm text-zinc-500 dark:text-zinc-400">
{job.state.sdInstanceId} {job.state.workerInstanceKey}
</span> </span>
</li> </li>
))} ))}

94
ui/StatsPage.tsx Normal file
View File

@ -0,0 +1,94 @@
import React from "react";
import { fetchApi, handleResponse } from "./apiClient.tsx";
import useSWR from "swr";
import { Counter } from "./Counter.tsx";
export function StatsPage() {
const globalStats = useSWR(
["stats", "GET", {}] as const,
(args) => fetchApi(...args).then(handleResponse),
{ refreshInterval: 2_000 },
);
return (
<div className="my-16 flex flex-col gap-16 text-zinc-600 dark:text-zinc-400">
<p className="flex flex-col items-center gap-2 text-2xl sm:text-3xl">
<span>Pixelsteps diffused</span>
<Counter
className="font-bold text-zinc-700 dark:text-zinc-300"
value={globalStats.data?.pixelStepCount ?? 0}
digits={15}
transitionDurationMs={3_000}
/>
<Counter
className="text-base"
value={(globalStats.data?.pixelStepsPerMinute ?? 0) / 60}
digits={9}
transitionDurationMs={2_000}
postfix="/s"
/>
</p>
<p className="flex flex-col items-center gap-2 text-2xl sm:text-3xl">
<span>Pixels painted</span>
<Counter
className="font-bold text-zinc-700 dark:text-zinc-300"
value={globalStats.data?.pixelCount ?? 0}
digits={15}
transitionDurationMs={3_000}
/>
<Counter
className="text-base"
value={(globalStats.data?.pixelsPerMinute ?? 0) / 60}
digits={9}
transitionDurationMs={2_000}
postfix="/s"
/>
</p>
<div className="flex flex-col md:flex-row gap-16 md:gap-8">
<p className="flex-grow flex flex-col items-center gap-2 text-2xl">
<span>Steps processed</span>
<Counter
className="font-bold text-zinc-700 dark:text-zinc-300"
value={globalStats.data?.stepCount ?? 0}
digits={9}
transitionDurationMs={3_000}
/>
<Counter
className="text-base"
value={(globalStats.data?.stepsPerMinute ?? 0) / 60}
digits={3}
fractionDigits={3}
transitionDurationMs={2_000}
postfix="/s"
/>
</p>
<p className="flex-grow flex flex-col items-center gap-2 text-2xl">
<span>Images generated</span>
<Counter
className="font-bold text-zinc-700 dark:text-zinc-300"
value={globalStats.data?.imageCount ?? 0}
digits={9}
transitionDurationMs={3_000}
/>
<Counter
className="text-base"
value={(globalStats.data?.imagesPerMinute ?? 0) / 60}
digits={3}
fractionDigits={3}
transitionDurationMs={2_000}
postfix="/s"
/>
</p>
</div>
<p className="flex-grow flex flex-col items-center gap-2 text-2xl">
<span>Unique users</span>
<Counter
className="font-bold text-zinc-700 dark:text-zinc-300"
value={globalStats.data?.userCount ?? 0}
digits={6}
transitionDurationMs={1_500}
/>
</p>
</div>
);
}

330
ui/WorkersPage.tsx Normal file
View File

@ -0,0 +1,330 @@
import React, { useRef } from "react";
import { FormattedRelativeTime } from "react-intl";
import useSWR, { useSWRConfig } from "swr";
import { WorkerData } from "../api/workersRoute.ts";
import { Counter } from "./Counter.tsx";
import { fetchApi, handleResponse } from "./apiClient.tsx";
export function WorkersPage(props: { sessionId?: string }) {
const { sessionId } = props;
const createWorkerModalRef = useRef<HTMLDialogElement>(null);
const session = useSWR(
sessionId ? ["sessions/{sessionId}", "GET", { params: { sessionId } }] as const : null,
(args) => fetchApi(...args).then(handleResponse),
);
const user = useSWR(
session.data?.userId
? ["users/{userId}", "GET", { params: { userId: String(session.data.userId) } }] as const
: null,
(args) => fetchApi(...args).then(handleResponse),
);
const workers = useSWR(
["workers", "GET", {}] as const,
(args) => fetchApi(...args).then(handleResponse),
{ refreshInterval: 5000 },
);
return (
<>
<ul className="my-4 flex flex-col gap-2">
{workers.data?.map((worker) => (
<WorkerListItem key={worker.id} worker={worker} sessionId={sessionId} />
))}
</ul>
{user.data?.isAdmin && (
<button
className="button-filled"
onClick={() => createWorkerModalRef.current?.showModal()}
>
Add worker
</button>
)}
<dialog
className="dialog"
ref={createWorkerModalRef}
>
<form
method="dialog"
className="flex flex-col gap-4 p-4"
onSubmit={(e) => {
e.preventDefault();
const data = new FormData(e.target as HTMLFormElement);
const key = data.get("key") as string;
const name = data.get("name") as string;
const sdUrl = data.get("url") as string;
const user = data.get("user") as string;
const password = data.get("password") as string;
console.log(key, name, user, password);
workers.mutate(async () => {
const worker = await fetchApi("workers", "POST", {
query: { sessionId: sessionId! },
body: {
type: "application/json",
data: {
key,
name: name || null,
sdUrl,
sdAuth: user && password ? { user, password } : null,
},
},
}).then(handleResponse);
return [...(workers.data ?? []), worker];
});
createWorkerModalRef.current?.close();
}}
>
<div className="flex gap-4">
<label className="flex flex-1 flex-col items-stretch gap-1">
<span className="text-sm">
Key
</span>
<input
className="input-text"
type="text"
name="key"
required
/>
<span className="text-sm text-zinc-500">
Used for counting statistics
</span>
</label>
<label className="flex flex-1 flex-col items-stretch gap-1">
<span className="text-sm">
Name
</span>
<input
className="input-text"
type="text"
name="name"
/>
<span className="text-sm text-zinc-500">
Used for display
</span>
</label>
</div>
<label className="flex flex-col items-stretch gap-1">
<span className="text-sm">
URL
</span>
<input
className="input-text"
type="url"
name="url"
required
pattern="https?://.*"
/>
</label>
<div className="flex gap-4">
<label className="flex flex-1 flex-col items-stretch gap-1">
<span className="text-sm">
User
</span>
<input
className="input-text"
type="text"
name="user"
/>
</label>
<label className="flex flex-1 flex-col items-stretch gap-1">
<span className="text-sm">
Password
</span>
<input
className="input-text"
type="password"
name="password"
/>
</label>
</div>
<div className="flex gap-2 items-center justify-end">
<button
type="button"
className="button-outlined"
onClick={() => createWorkerModalRef.current?.close()}
>
Close
</button>
<button type="submit" disabled={!sessionId} className="button-filled">
Save
</button>
</div>
</form>
</dialog>
</>
);
}
function WorkerListItem(props: { worker: WorkerData; sessionId?: string }) {
const { worker, sessionId } = props;
const editWorkerModalRef = useRef<HTMLDialogElement>(null);
const deleteWorkerModalRef = useRef<HTMLDialogElement>(null);
const session = useSWR(
sessionId ? ["sessions/{sessionId}", "GET", { params: { sessionId } }] as const : null,
(args) => fetchApi(...args).then(handleResponse),
);
const user = useSWR(
session.data?.userId
? ["users/{userId}", "GET", { params: { userId: String(session.data.userId) } }] as const
: null,
(args) => fetchApi(...args).then(handleResponse),
);
const { mutate } = useSWRConfig();
return (
<li className="flex flex-col gap-2 rounded-md bg-zinc-100 dark:bg-zinc-800 p-2">
<p className="font-bold">
{worker.name ?? worker.key}
</p>
{worker.isActive ? <p> Active</p> : (
<>
<p>
Last seen {worker.lastOnlineTime
? (
<FormattedRelativeTime
value={(worker.lastOnlineTime - Date.now()) / 1000}
numeric="auto"
updateIntervalInSeconds={1}
/>
)
: "never"}
</p>
{worker.lastError && (
<p className="text-red-500">
{worker.lastError.message} (
<FormattedRelativeTime
value={(worker.lastError.time - Date.now()) / 1000}
numeric="auto"
updateIntervalInSeconds={1}
/>)
</p>
)}
</>
)}
<p className="flex gap-1">
<Counter value={worker.imagesPerMinute} digits={2} fractionDigits={1} /> images per minute,
{" "}
<Counter value={worker.stepsPerMinute} digits={3} /> steps per minute
</p>
<p className="flex gap-2">
{user.data?.isAdmin && (
<>
<button
className="button-outlined"
onClick={() => editWorkerModalRef.current?.showModal()}
>
Settings
</button>
<button
className="button-outlined"
onClick={() => deleteWorkerModalRef.current?.showModal()}
>
Delete
</button>
</>
)}
</p>
<dialog
className="dialog"
ref={editWorkerModalRef}
>
<form
method="dialog"
className="flex flex-col gap-4 p-4"
onSubmit={(e) => {
e.preventDefault();
const data = new FormData(e.target as HTMLFormElement);
const user = data.get("user") as string;
const password = data.get("password") as string;
console.log(user, password);
fetchApi("workers/{workerId}", "PATCH", {
params: { workerId: worker.id },
query: { sessionId: sessionId! },
body: {
type: "application/json",
data: {
auth: user && password ? { user, password } : null,
},
},
});
editWorkerModalRef.current?.close();
}}
>
<div className="flex gap-4">
<label className="flex flex-1 flex-col items-stretch gap-1">
<span className="text-sm">
User
</span>
<input
className="input-text"
type="text"
name="user"
/>
</label>
<label className="flex flex-1 flex-col items-stretch gap-1">
<span className="text-sm">
Password
</span>
<input
className="input-text"
type="password"
name="password"
/>
</label>
</div>
<div className="flex gap-2 items-center justify-end">
<button
type="button"
className="button-outlined"
onClick={() => editWorkerModalRef.current?.close()}
>
Close
</button>
<button type="submit" disabled={!sessionId} className="button-filled">
Save
</button>
</div>
</form>
</dialog>
<dialog
className="dialog"
ref={deleteWorkerModalRef}
>
<form
method="dialog"
className="flex flex-col gap-4 p-4"
onSubmit={(e) => {
e.preventDefault();
fetchApi("workers/{workerId}", "DELETE", {
params: { workerId: worker.id },
query: { sessionId: sessionId! },
}).then(handleResponse).then(() => mutate(["workers", "GET", {}]));
deleteWorkerModalRef.current?.close();
}}
>
<p>
Are you sure you want to delete this worker?
</p>
<div className="flex gap-2 items-center justify-end">
<button
type="button"
className="button-outlined"
onClick={() => deleteWorkerModalRef.current?.close()}
>
Close
</button>
<button type="submit" disabled={!sessionId} className="button-filled">
Delete
</button>
</div>
</form>
</dialog>
</li>
);
}

View File

@ -4,9 +4,12 @@ import React from "react";
import { App } from "./App.tsx"; import { App } from "./App.tsx";
import "./twind.ts"; import "./twind.ts";
import { BrowserRouter } from "react-router-dom"; import { BrowserRouter } from "react-router-dom";
import { IntlProvider } from "react-intl";
createRoot(document.body).render( createRoot(document.body).render(
<BrowserRouter> <BrowserRouter>
<IntlProvider locale="en">
<App /> <App />
</IntlProvider>
</BrowserRouter>, </BrowserRouter>,
); );

View File

@ -1,7 +1,7 @@
import { Chat, User } from "grammy_types"; import { Chat, User } from "grammy_types";
export function formatUserChat( export function formatUserChat(
ctx: { from?: User; chat?: Chat; sdInstanceId?: string }, ctx: { from?: User; chat?: Chat; workerInstanceKey?: string },
) { ) {
const msg: string[] = []; const msg: string[] = [];
if (ctx.from) { if (ctx.from) {
@ -26,8 +26,8 @@ export function formatUserChat(
} }
} }
} }
if (ctx.sdInstanceId) { if (ctx.workerInstanceKey) {
msg.push(`using ${ctx.sdInstanceId}`); msg.push(`using ${ctx.workerInstanceKey}`);
} }
return msg.join(" "); return msg.join(" ");
} }

4
utils/getAuthHeader.ts Normal file
View File

@ -0,0 +1,4 @@
export function getAuthHeader(auth: { user: string; password: string } | null) {
if (!auth) return {};
return { "Authorization": `Basic ${btoa(`${auth.user}:${auth.password}`)}` };
}