diff --git a/bot/img2imgCommand.ts b/bot/img2imgCommand.ts index e73ac1b..0376a65 100644 --- a/bot/img2imgCommand.ts +++ b/bot/img2imgCommand.ts @@ -1,7 +1,7 @@ import { Collections, Grammy, GrammyStatelessQ } from "../deps.ts"; -import { formatUserChat } from "../utils.ts"; +import { formatUserChat } from "../common/utils.ts"; import { jobStore } from "../db/jobStore.ts"; -import { parsePngInfo, PngInfo } from "../sd.ts"; +import { parsePngInfo, PngInfo } from "../common/parsePngInfo.ts"; import { Context, logger } from "./mod.ts"; export const img2imgQuestion = new GrammyStatelessQ.StatelessQuestion( diff --git a/bot/mod.ts b/bot/mod.ts index e95545e..ca816ad 100644 --- a/bot/mod.ts +++ b/bot/mod.ts @@ -1,5 +1,5 @@ import { Grammy, GrammyAutoQuote, GrammyFiles, GrammyParseMode, Log } from "../deps.ts"; -import { formatUserChat } from "../utils.ts"; +import { formatUserChat } from "../common/utils.ts"; import { session, SessionFlavor } from "./session.ts"; import { queueCommand } from "./queueCommand.ts"; import { txt2imgCommand, txt2imgQuestion } from "./txt2imgCommand.ts"; diff --git a/bot/pnginfoCommand.ts b/bot/pnginfoCommand.ts index e83bbe7..b9600b6 100644 --- a/bot/pnginfoCommand.ts +++ b/bot/pnginfoCommand.ts @@ -1,7 +1,6 @@ import { Grammy, GrammyParseMode, GrammyStatelessQ } from "../deps.ts"; -import { fmt } from "../utils.ts"; - -import { getPngInfo, parsePngInfo } from "../sd.ts"; +import { fmt } from "../common/utils.ts"; +import { getPngInfo, parsePngInfo } from "../common/parsePngInfo.ts"; import { Context } from "./mod.ts"; export const pnginfoQuestion = new GrammyStatelessQ.StatelessQuestion( diff --git a/bot/queueCommand.ts b/bot/queueCommand.ts index 000b22a..09d7648 100644 --- a/bot/queueCommand.ts +++ b/bot/queueCommand.ts @@ -1,8 +1,9 @@ import { Grammy, GrammyParseMode } from "../deps.ts"; -import { fmt, getFlagEmoji } from "../utils.ts"; +import { fmt } from "../common/utils.ts"; import { runningWorkers } from "../tasks/pingWorkers.ts"; import { jobStore } from "../db/jobStore.ts"; import { Context, logger } from "./mod.ts"; +import { getFlagEmoji } from "../common/getFlagEmoji.ts"; export async function queueCommand(ctx: Grammy.CommandContext) { let formattedMessage = await getMessageText(); diff --git a/bot/session.ts b/bot/session.ts index 0a769ef..d18c97d 100644 --- a/bot/session.ts +++ b/bot/session.ts @@ -1,6 +1,6 @@ import { db } from "../db/db.ts"; import { Grammy, GrammyKvStorage } from "../deps.ts"; -import { SdApi, SdTxt2ImgRequest } from "../sd.ts"; +import { SdApi, SdTxt2ImgRequest } from "../common/sdApi.ts"; export type SessionFlavor = Grammy.SessionFlavor; diff --git a/bot/txt2imgCommand.ts b/bot/txt2imgCommand.ts index 9b4e48b..0f088a5 100644 --- a/bot/txt2imgCommand.ts +++ b/bot/txt2imgCommand.ts @@ -1,7 +1,7 @@ import { Grammy, GrammyStatelessQ } from "../deps.ts"; -import { formatUserChat } from "../utils.ts"; +import { formatUserChat } from "../common/utils.ts"; import { jobStore } from "../db/jobStore.ts"; -import { getPngInfo, parsePngInfo, PngInfo } from "../sd.ts"; +import { getPngInfo, parsePngInfo, PngInfo } from "../common/parsePngInfo.ts"; import { Context, logger } from "./mod.ts"; export const txt2imgQuestion = new GrammyStatelessQ.StatelessQuestion( diff --git a/common/getFlagEmoji.ts b/common/getFlagEmoji.ts new file mode 100644 index 0000000..6d7a18a --- /dev/null +++ b/common/getFlagEmoji.ts @@ -0,0 +1,51 @@ +/** Language to biggest country emoji map */ +const languageToFlagMap: Record = { + "en": "🇺🇸", + "zh": "🇨🇳", + "es": "🇪🇸", + "hi": "🇮🇳", + "ar": "🇪🇬", + "pt": "🇧🇷", + "bn": "🇧🇩", + "ru": "🇷🇺", + "ja": "🇯🇵", + "pa": "🇮🇳", + "de": "🇩🇪", + "ko": "🇰🇷", + "fr": "🇫🇷", + "tr": "🇹🇷", + "ur": "🇵🇰", + "it": "🇮🇹", + "th": "🇹🇭", + "vi": "🇻🇳", + "pl": "🇵🇱", + "uk": "🇺🇦", + "uz": "🇺🇿", + "su": "🇮🇩", + "sw": "🇹🇿", + "nl": "🇳🇱", + "fi": "🇫🇮", + "el": "🇬🇷", + "da": "🇩🇰", + "cs": "🇨🇿", + "sk": "🇸🇰", + "bg": "🇧🇬", + "sv": "🇸🇪", + "be": "🇧🇾", + "hu": "🇭🇺", + "lt": "🇱🇹", + "lv": "🇱🇻", + "et": "🇪🇪", + "sl": "🇸🇮", + "hr": "🇭🇷", + "zu": "🇿🇦", + "id": "🇮🇩", + "is": "🇮🇸", + "lb": "🇱🇺", // Luxembourgish - Luxembourg +}; + +export function getFlagEmoji(languageCode?: string): string | undefined { + const language = languageCode?.split("-").pop()?.toLowerCase(); + if (!language) return; + return languageToFlagMap[language]; +} diff --git a/pnginfo.test.ts b/common/parsePngInfo.test.ts similarity index 98% rename from pnginfo.test.ts rename to common/parsePngInfo.test.ts index cd820ca..26e22e3 100644 --- a/pnginfo.test.ts +++ b/common/parsePngInfo.test.ts @@ -3,7 +3,7 @@ import { assertEquals, assertMatch, } from "https://deno.land/std@0.135.0/testing/asserts.ts"; -import { parsePngInfo } from "./sd.ts"; +import { parsePngInfo } from "./parsePngInfo.ts"; Deno.test("parses pnginfo", async (t) => { await t.step("1", () => { diff --git a/common/parsePngInfo.ts b/common/parsePngInfo.ts new file mode 100644 index 0000000..b324b58 --- /dev/null +++ b/common/parsePngInfo.ts @@ -0,0 +1,117 @@ +import { PngChunksExtract, PngChunkText } from "../deps.ts"; + +export function getPngInfo(pngData: Uint8Array): string | undefined { + return PngChunksExtract.default(pngData) + .filter((chunk) => chunk.name === "tEXt") + .map((chunk) => PngChunkText.decode(chunk.data)) + .find((textChunk) => textChunk.keyword === "parameters") + ?.text; +} + +export interface PngInfo { + prompt: string; + negative_prompt: string; + steps: number; + cfg_scale: number; + width: number; + height: number; + sampler_name: string; + seed: number; + denoising_strength: number; +} + +export function parsePngInfo(pngInfo: string): Partial { + const tags = pngInfo.split(/[,;]+|\.+\s|\n/u); + let part: "prompt" | "negative_prompt" | "params" = "prompt"; + const params: Partial = {}; + const prompt: string[] = []; + const negativePrompt: string[] = []; + for (const tag of tags) { + const paramValuePair = tag.trim().match(/^(\w+\s*\w*):\s+(.*)$/u); + if (paramValuePair) { + const [, param, value] = paramValuePair; + switch (param.replace(/\s+/u, "").toLowerCase()) { + case "positiveprompt": + case "positive": + case "prompt": + case "pos": + part = "prompt"; + prompt.push(value.trim()); + break; + case "negativeprompt": + case "negative": + case "neg": + part = "negative_prompt"; + negativePrompt.push(value.trim()); + break; + case "steps": + case "cycles": { + part = "params"; + const steps = Number(value.trim()); + if (steps > 0) params.steps = Math.min(steps, 50); + break; + } + case "cfgscale": + case "cfg": + case "detail": { + part = "params"; + const cfgScale = Number(value.trim()); + if (cfgScale > 0) params.cfg_scale = Math.min(cfgScale, 20); + break; + } + case "size": + case "resolution": { + part = "params"; + const [width, height] = value.trim() + .split(/\s*[x,]\s*/u, 2) + .map((v) => v.trim()) + .map(Number); + if (width > 0 && height > 0) { + params.width = Math.min(width, 2048); + params.height = Math.min(height, 2048); + } + break; + } + case "denoisingstrength": + case "denoising": + case "denoise": { + part = "params"; + // allow percent or decimal + let denoisingStrength: number; + if (value.trim().endsWith("%")) { + denoisingStrength = Number(value.trim().slice(0, -1).trim()) / 100; + } else { + denoisingStrength = Number(value.trim()); + } + denoisingStrength = Math.min(Math.max(denoisingStrength, 0), 1); + params.denoising_strength = denoisingStrength; + break; + } + case "seed": + case "model": + case "modelhash": + case "modelname": + case "sampler": + part = "params"; + // ignore for now + break; + default: + break; + } + } else if (tag.trim().length > 0) { + switch (part) { + case "prompt": + prompt.push(tag.trim()); + break; + case "negative_prompt": + negativePrompt.push(tag.trim()); + break; + default: + break; + } + } + } + if (prompt.length > 0) params.prompt = prompt.join(", "); + if (negativePrompt.length > 0) params.negative_prompt = negativePrompt.join(", "); + return params; +} diff --git a/sd.ts b/common/sdApi.ts similarity index 69% rename from sd.ts rename to common/sdApi.ts index 8f00bef..56d10a6 100644 --- a/sd.ts +++ b/common/sdApi.ts @@ -1,4 +1,4 @@ -import { Async, AsyncX, PngChunksExtract, PngChunkText } from "./deps.ts"; +import { Async, AsyncX } from "../deps.ts"; export interface SdApi { url: string; @@ -298,119 +298,3 @@ export class SdApiError extends Error { super(message); } } - -export function getPngInfo(pngData: Uint8Array): string | undefined { - return PngChunksExtract.default(pngData) - .filter((chunk) => chunk.name === "tEXt") - .map((chunk) => PngChunkText.decode(chunk.data)) - .find((textChunk) => textChunk.keyword === "parameters") - ?.text; -} - -export interface PngInfo { - prompt: string; - negative_prompt: string; - steps: number; - cfg_scale: number; - width: number; - height: number; - sampler_name: string; - seed: number; - denoising_strength: number; -} - -export function parsePngInfo(pngInfo: string): Partial { - const tags = pngInfo.split(/[,;]+|\.+\s|\n/u); - let part: "prompt" | "negative_prompt" | "params" = "prompt"; - const params: Partial = {}; - const prompt: string[] = []; - const negativePrompt: string[] = []; - for (const tag of tags) { - const paramValuePair = tag.trim().match(/^(\w+\s*\w*):\s+(.*)$/u); - if (paramValuePair) { - const [, param, value] = paramValuePair; - switch (param.replace(/\s+/u, "").toLowerCase()) { - case "positiveprompt": - case "positive": - case "prompt": - case "pos": - part = "prompt"; - prompt.push(value.trim()); - break; - case "negativeprompt": - case "negative": - case "neg": - part = "negative_prompt"; - negativePrompt.push(value.trim()); - break; - case "steps": - case "cycles": { - part = "params"; - const steps = Number(value.trim()); - if (steps > 0) params.steps = Math.min(steps, 50); - break; - } - case "cfgscale": - case "cfg": - case "detail": { - part = "params"; - const cfgScale = Number(value.trim()); - if (cfgScale > 0) params.cfg_scale = Math.min(cfgScale, 20); - break; - } - case "size": - case "resolution": { - part = "params"; - const [width, height] = value.trim() - .split(/\s*[x,]\s*/u, 2) - .map((v) => v.trim()) - .map(Number); - if (width > 0 && height > 0) { - params.width = Math.min(width, 2048); - params.height = Math.min(height, 2048); - } - break; - } - case "denoisingstrength": - case "denoising": - case "denoise": { - part = "params"; - // allow percent or decimal - let denoisingStrength: number; - if (value.trim().endsWith("%")) { - denoisingStrength = Number(value.trim().slice(0, -1).trim()) / 100; - } else { - denoisingStrength = Number(value.trim()); - } - denoisingStrength = Math.min(Math.max(denoisingStrength, 0), 1); - params.denoising_strength = denoisingStrength; - break; - } - case "seed": - case "model": - case "modelhash": - case "modelname": - case "sampler": - part = "params"; - // ignore for now - break; - default: - break; - } - } else if (tag.trim().length > 0) { - switch (part) { - case "prompt": - prompt.push(tag.trim()); - break; - case "negative_prompt": - negativePrompt.push(tag.trim()); - break; - default: - break; - } - } - } - if (prompt.length > 0) params.prompt = prompt.join(", "); - if (negativePrompt.length > 0) params.negative_prompt = negativePrompt.join(", "); - return params; -} diff --git a/common/utils.ts b/common/utils.ts new file mode 100644 index 0000000..3362e39 --- /dev/null +++ b/common/utils.ts @@ -0,0 +1,60 @@ +import { GrammyParseMode, GrammyTypes } from "../deps.ts"; + +export function formatOrdinal(n: number) { + if (n % 100 === 11 || n % 100 === 12 || n % 100 === 13) return `${n}th`; + if (n % 10 === 1) return `${n}st`; + if (n % 10 === 2) return `${n}nd`; + if (n % 10 === 3) return `${n}rd`; + return `${n}th`; +} + +export const fmt = ( + rawStringParts: TemplateStringsArray | GrammyParseMode.Stringable[], + ...stringLikes: GrammyParseMode.Stringable[] +): GrammyParseMode.FormattedString => { + let text = ""; + const entities: GrammyTypes.MessageEntity[] = []; + + const length = Math.max(rawStringParts.length, stringLikes.length); + for (let i = 0; i < length; i++) { + for (const stringLike of [rawStringParts[i], stringLikes[i]]) { + if (stringLike instanceof GrammyParseMode.FormattedString) { + entities.push( + ...stringLike.entities.map((e) => ({ + ...e, + offset: e.offset + text.length, + })), + ); + } + if (stringLike != null) text += stringLike.toString(); + } + } + return new GrammyParseMode.FormattedString(text, entities); +}; + +export function formatUserChat(ctx: { from?: GrammyTypes.User; chat?: GrammyTypes.Chat }) { + const msg: string[] = []; + if (ctx.from) { + msg.push(ctx.from.first_name); + if (ctx.from.last_name) msg.push(ctx.from.last_name); + if (ctx.from.username) msg.push(`(@${ctx.from.username})`); + if (ctx.from.language_code) msg.push(`(${ctx.from.language_code.toUpperCase()})`); + } + if (ctx.chat) { + if ( + ctx.chat.type === "group" || + ctx.chat.type === "supergroup" || + ctx.chat.type === "channel" + ) { + msg.push("in"); + msg.push(ctx.chat.title); + if ( + (ctx.chat.type === "supergroup" || ctx.chat.type === "channel") && + ctx.chat.username + ) { + msg.push(`(@${ctx.chat.username})`); + } + } + } + return msg.join(" "); +} diff --git a/db/jobStore.ts b/db/jobStore.ts index b1e71ad..3f76a9a 100644 --- a/db/jobStore.ts +++ b/db/jobStore.ts @@ -1,5 +1,6 @@ import { GrammyTypes, IKV } from "../deps.ts"; -import { PngInfo, SdTxt2ImgInfo } from "../sd.ts"; +import { SdTxt2ImgInfo } from "../common/sdApi.ts"; +import { PngInfo } from "../common/parsePngInfo.ts"; import { db } from "./db.ts"; export interface JobSchema { diff --git a/tasks/pingWorkers.ts b/tasks/pingWorkers.ts index ae445c5..4cfcd2d 100644 --- a/tasks/pingWorkers.ts +++ b/tasks/pingWorkers.ts @@ -1,6 +1,6 @@ import { Async, Log } from "../deps.ts"; import { getGlobalSession } from "../bot/session.ts"; -import { sdGetConfig } from "../sd.ts"; +import { sdGetConfig } from "../common/sdApi.ts"; const logger = () => Log.getLogger(); diff --git a/tasks/processJobs.ts b/tasks/processJobs.ts index 43f8b05..b8a70fc 100644 --- a/tasks/processJobs.ts +++ b/tasks/processJobs.ts @@ -11,8 +11,14 @@ import { } from "../deps.ts"; import { bot } from "../bot/mod.ts"; import { getGlobalSession, GlobalData, WorkerData } from "../bot/session.ts"; -import { fmt, formatUserChat } from "../utils.ts"; -import { SdApiError, sdImg2Img, SdProgressResponse, SdResponse, sdTxt2Img } from "../sd.ts"; +import { fmt, formatUserChat } from "../common/utils.ts"; +import { + SdApiError, + sdImg2Img, + SdProgressResponse, + SdResponse, + sdTxt2Img, +} from "../common/sdApi.ts"; import { JobSchema, jobStore } from "../db/jobStore.ts"; import { runningWorkers } from "./pingWorkers.ts"; diff --git a/tasks/returnHangedJobs.ts b/tasks/returnHangedJobs.ts index a45e148..8b8a868 100644 --- a/tasks/returnHangedJobs.ts +++ b/tasks/returnHangedJobs.ts @@ -1,5 +1,5 @@ import { FmtDuration, Log } from "../deps.ts"; -import { formatUserChat } from "../utils.ts"; +import { formatUserChat } from "../common/utils.ts"; import { jobStore } from "../db/jobStore.ts"; const logger = () => Log.getLogger(); diff --git a/tasks/updateJobStatusMsgs.ts b/tasks/updateJobStatusMsgs.ts index 74c8eb6..fa2bb51 100644 --- a/tasks/updateJobStatusMsgs.ts +++ b/tasks/updateJobStatusMsgs.ts @@ -1,6 +1,6 @@ import { Log } from "../deps.ts"; import { bot } from "../bot/mod.ts"; -import { formatOrdinal } from "../utils.ts"; +import { formatOrdinal } from "../common/utils.ts"; import { jobStore } from "../db/jobStore.ts"; const logger = () => Log.getLogger(); diff --git a/utils.ts b/utils.ts deleted file mode 100644 index abbd8ee..0000000 --- a/utils.ts +++ /dev/null @@ -1,112 +0,0 @@ -import { GrammyParseMode, GrammyTypes } from "./deps.ts"; - -export function formatOrdinal(n: number) { - if (n % 100 === 11 || n % 100 === 12 || n % 100 === 13) return `${n}th`; - if (n % 10 === 1) return `${n}st`; - if (n % 10 === 2) return `${n}nd`; - if (n % 10 === 3) return `${n}rd`; - return `${n}th`; -} - -export const fmt = ( - rawStringParts: TemplateStringsArray | GrammyParseMode.Stringable[], - ...stringLikes: GrammyParseMode.Stringable[] -): GrammyParseMode.FormattedString => { - let text = ""; - const entities: GrammyTypes.MessageEntity[] = []; - - const length = Math.max(rawStringParts.length, stringLikes.length); - for (let i = 0; i < length; i++) { - for (const stringLike of [rawStringParts[i], stringLikes[i]]) { - if (stringLike instanceof GrammyParseMode.FormattedString) { - entities.push( - ...stringLike.entities.map((e) => ({ - ...e, - offset: e.offset + text.length, - })), - ); - } - if (stringLike != null) text += stringLike.toString(); - } - } - return new GrammyParseMode.FormattedString(text, entities); -}; - -export function formatUserChat(ctx: { from?: GrammyTypes.User; chat?: GrammyTypes.Chat }) { - const msg: string[] = []; - if (ctx.from) { - msg.push(ctx.from.first_name); - if (ctx.from.last_name) msg.push(ctx.from.last_name); - if (ctx.from.username) msg.push(`(@${ctx.from.username})`); - if (ctx.from.language_code) msg.push(`(${ctx.from.language_code.toUpperCase()})`); - } - if (ctx.chat) { - if ( - ctx.chat.type === "group" || - ctx.chat.type === "supergroup" || - ctx.chat.type === "channel" - ) { - msg.push("in"); - msg.push(ctx.chat.title); - if ( - (ctx.chat.type === "supergroup" || ctx.chat.type === "channel") && - ctx.chat.username - ) { - msg.push(`(@${ctx.chat.username})`); - } - } - } - return msg.join(" "); -} - -/** Language to biggest country emoji map */ -const languageToFlagMap: Record = { - "en": "🇺🇸", // English - United States - "zh": "🇨🇳", // Chinese - China - "es": "🇪🇸", // Spanish - Spain - "hi": "🇮🇳", // Hindi - India - "ar": "🇪🇬", // Arabic - Egypt - "pt": "🇧🇷", // Portuguese - Brazil - "bn": "🇧🇩", // Bengali - Bangladesh - "ru": "🇷🇺", // Russian - Russia - "ja": "🇯🇵", // Japanese - Japan - "pa": "🇮🇳", // Punjabi - India - "de": "🇩🇪", // German - Germany - "ko": "🇰🇷", // Korean - South Korea - "fr": "🇫🇷", // French - France - "tr": "🇹🇷", // Turkish - Turkey - "ur": "🇵🇰", // Urdu - Pakistan - "it": "🇮🇹", // Italian - Italy - "th": "🇹🇭", // Thai - Thailand - "vi": "🇻🇳", // Vietnamese - Vietnam - "pl": "🇵🇱", // Polish - Poland - "uk": "🇺🇦", // Ukrainian - Ukraine - "uz": "🇺🇿", // Uzbek - Uzbekistan - "su": "🇮🇩", // Sundanese - Indonesia - "sw": "🇹🇿", // Swahili - Tanzania - "nl": "🇳🇱", // Dutch - Netherlands - "fi": "🇫🇮", // Finnish - Finland - "el": "🇬🇷", // Greek - Greece - "da": "🇩🇰", // Danish - Denmark - "cs": "🇨🇿", // Czech - Czech Republic - "sk": "🇸🇰", // Slovak - Slovakia - "bg": "🇧🇬", // Bulgarian - Bulgaria - "sv": "🇸🇪", // Swedish - Sweden - "be": "🇧🇾", // Belarusian - Belarus - "hu": "🇭🇺", // Hungarian - Hungary - "lt": "🇱🇹", // Lithuanian - Lithuania - "lv": "🇱🇻", // Latvian - Latvia - "et": "🇪🇪", // Estonian - Estonia - "sl": "🇸🇮", // Slovenian - Slovenia - "hr": "🇭🇷", // Croatian - Croatia - "zu": "🇿🇦", // Zulu - South Africa - "id": "🇮🇩", // Indonesian - Indonesia - "is": "🇮🇸", // Icelandic - Iceland - "lb": "🇱🇺", // Luxembourgish - Luxembourg -}; - -export function getFlagEmoji(languageCode?: string): string | undefined { - const language = languageCode?.split("-").pop()?.toLowerCase(); - if (!language) return; - return languageToFlagMap[language]; -}