diff --git a/src/modules/sd-images/SDImagesBotBase.ts b/src/modules/sd-images/SDImagesBotBase.ts new file mode 100644 index 00000000..8d55586a --- /dev/null +++ b/src/modules/sd-images/SDImagesBotBase.ts @@ -0,0 +1,65 @@ +import { SDNodeApi, getModelByParam, IModel } from "./api"; +import { OnMessageContext, OnCallBackQueryData } from "../types"; +import { sleep, uuidv4 } from "./utils"; +import { InlineKeyboard, InputFile } from "grammy"; + +export class SDImagesBotBase { + sdNodeApi: SDNodeApi; + + queue: string[] = []; + + constructor() { + this.sdNodeApi = new SDNodeApi(); + } + + waitingQueue = async (uuid: string, ctx: OnMessageContext | OnCallBackQueryData,) => { + this.queue.push(uuid); + + let idx = this.queue.findIndex((v) => v === uuid); + + if (idx !== 0) { + ctx.reply( + `You are ${idx + 1}/${this.queue.length + }, wait about ${idx * 30} seconds` + ); + } + + // waiting queue + while (idx !== 0) { + await sleep(3000 * this.queue.findIndex((v) => v === uuid)); + idx = this.queue.findIndex((v) => v === uuid); + } + } + + generateImage = async ( + ctx: OnMessageContext | OnCallBackQueryData, + refundCallback: (reason?: string) => void, + prompt: string, + model: IModel, + seed?: number + ) => { + const uuid = uuidv4(); + + try { + await this.waitingQueue(uuid, ctx); + + ctx.chatAction = "upload_photo"; + + const imageBuffer = await this.sdNodeApi.generateImage( + prompt, + model, + seed + ); + + await ctx.replyWithPhoto(new InputFile(imageBuffer), { + caption: `/${model.aliases[0]} ${prompt}`, + }); + } catch (e) { + console.error(e); + ctx.reply(`Error: something went wrong... Refunding payments`); + refundCallback(); + } + + this.queue = this.queue.filter((v) => v !== uuid); + } +} \ No newline at end of file diff --git a/src/modules/sd-images/api/helpers.ts b/src/modules/sd-images/api/helpers.ts new file mode 100644 index 00000000..a94ebc11 --- /dev/null +++ b/src/modules/sd-images/api/helpers.ts @@ -0,0 +1,8 @@ +export const waitingExecute = (fn: () => Promise, ms: number) => new Promise((resolve, reject) => { + const timeoutId = setTimeout(() => { + console.error('SD images Error: waitingExecute time is up'); + reject('Error: waitingExecute time is up'); + }, ms); + + fn().then(resolve).catch(reject).finally(() => clearTimeout(timeoutId)) + }); \ No newline at end of file diff --git a/src/modules/sd-images/api/index.ts b/src/modules/sd-images/api/index.ts new file mode 100644 index 00000000..cd737d9f --- /dev/null +++ b/src/modules/sd-images/api/index.ts @@ -0,0 +1,2 @@ +export * from './sd-node-api'; +export * from './models-config'; \ No newline at end of file diff --git a/src/modules/sd-images/api/models-config.ts b/src/modules/sd-images/api/models-config.ts new file mode 100644 index 00000000..5dd9aa0a --- /dev/null +++ b/src/modules/sd-images/api/models-config.ts @@ -0,0 +1,135 @@ +export interface IModel { + path: string; + name: string; + id: string; + hash: string; + shortName: string; + link: string; + baseModel: 'SD 1.5' | 'SDXL 1.0'; + aliases: string[]; + defaultPrompt: string; +} + +export const MODELS_CONFIGS: IModel[] = [ + { + path: "deliberate_v2.safetensors", + name: "Deliberate", + id: '4823', + hash: '9ABA26ABDF', + shortName: 'deliberate', + link: 'https://civitai.com/models/4823/deliberate', + baseModel: 'SD 1.5', + aliases: ['del', '4823', 'h9ab'], + defaultPrompt: 'a cute kitten made out of metal, (cyborg:1.1), ([tail | detailed wire]:1.3), (intricate details), hdr, (intricate details, hyperdetailed:1.2), cinematic shot, vignette, centered', + }, + { + path: "dreamshaper_8.safetensors", + name: "DreamShaper", + id: '4384', + hash: '879DB523C3', + shortName: 'dreamshaper', + link: 'https://civitai.com/models/4384/dreamshaper', + baseModel: 'SD 1.5', + aliases: ['dream', '4384', 'h879'], + defaultPrompt: '8k portrait of beautiful cyborg with brown hair, intricate, elegant, highly detailed, majestic, digital photography, art by artgerm and ruan jia and greg rutkowski surreal painting gold butterfly filigree, broken glass, (masterpiece, sidelighting, finely detailed beautiful eyes: 1.2), hdr, ' + }, + { + path: "majicmixRealistic_betterV2V25.safetensors", + name: 'majicMIX realistic 麦橘写实', + id: '43331', + hash: 'D7E2AC2F4A', + shortName: 'majicmix-realistic', + link: 'https://civitai.com/models/43331/majicmix-realistic', + baseModel: 'SD 1.5', + aliases: ['maji', '4333', 'h7e2'], + defaultPrompt: '1girl,sitting on a cozy couch,crossing legs,soft light' + }, + { + path: "revAnimated_v122.safetensors", + name: 'ReV Animated', + id: '7371', + hash: '4199BCDD14', + shortName: 'rev-animated', + link: 'https://civitai.com/models/7371/rev-animated', + baseModel: 'SD 1.5', + aliases: ['rev', '7371', 'h419'], + defaultPrompt: '((best quality)), ((masterpiece)), (detailed), alluring succubus, ethereal beauty, perched on a cloud, (fantasy illustration:1.3), enchanting gaze, captivating pose, delicate wings, otherworldly charm, mystical sky, (Luis Royo:1.2), (Yoshitaka Amano:1.1), moonlit night, soft colors, (detailed cloudscape:1.3), (high-resolution:1.2)' + }, + // { + // path: "v1-5-pruned-emaonly.safetensors", + // name: '', + // id: '', + // hash: '', + // shortName: '', + // link: '', + // baseModel: 'SD 1.5' + // }, + { + path: "animePastelDream_softBakedVae.safetensors", + name: 'Anime Pastel Dream', + id: '23521', + hash: '4BE38C1A17', + shortName: 'anime-pastel-dream', + link: 'https://civitai.com/models/23521/anime-pastel-dream', + baseModel: 'SD 1.5', + aliases: ['anime', '2352', 'h4be'], + defaultPrompt: 'masterpiece, best quality, ultra-detailed, illustration,(1girl),beautiful detailed eyes, looking at viewer, close up, (breast focus), pink hair, shy, cat ears' + }, + { + path: 'cyberrealistic_v33.safetensors', + name: 'CyberRealistic', + id: '15003', + hash: '7A4DBBA12F', + shortName: 'cyberrealistic', + link: 'https://civitai.com/models/15003/cyberrealistic', + baseModel: 'SD 1.5', + aliases: ['cyber', '1500', 'h7a4'], + defaultPrompt: ' (8k, RAW photo, highest quality), beautiful girl, close up, dress, (detailed eyes:0.8), defiance512, (looking at the camera:1.4), (highest quality), (best shadow), intricate details, interior, ginger hair:1.3, dark studio, muted colors, freckles ', + }, + { + path: 'dreamshaperXL10_alpha2Xl10.safetensors', + name: 'DreamShaper XL1.0', + id: '112902', + hash: '0F1B80CFE8', + shortName: 'dreamshaper-xl10', + link: 'https://civitai.com/models/112902/dreamshaper-xl10', + baseModel: 'SDXL 1.0', + aliases: ['xl_dream', '1129', 'h0f1'], + defaultPrompt: 'barry allen the flash on wheelchair moving at supersonic speed creating flame trails, speed trails, motion blur, electricity speed outdoor, realistic highly detailed cinematic cinematography, movie shots footage,', + }, + { + path: 'sdXL_v10VAEFix.safetensors', + name: 'SD XL', + id: '101055', + hash: 'E6BB9EA85B', + shortName: 'sd-xl', + link: 'https://civitai.com/models/101055/sd-xl', + baseModel: 'SDXL 1.0', + aliases: ['xl', '1010', 'he6b'], + defaultPrompt: 'A cat is sitting in a kimono, in the style of renaissance - inspired chiaroscuro, hyper - realistic portraiture, nicolas mignard, old master influenced fantasy, portraitures with hidden meanings, dom qwek, art of burma' + }, + { + path: 'sdxlUnstableDiffusers_v5UnchainedSlayer.safetensors', + name: 'SDXL Unstable Diffusers ☛ YamerMIX', + id: '84040', + hash: 'EF924AAE79', + shortName: 'sdxl-unstable-diffusers-yamermix', + link: 'https://civitai.com/models/84040/sdxl-unstable-diffusers-yamermix', + baseModel: 'SDXL 1.0', + aliases: ['xl_dif', '8404', 'hef9'], + defaultPrompt: 'pastel color, from above, upper body, depth of field, masterpiece, best quality, best quality, 1girl sitting on a swing, school uniform, black hair, blue eyes, autumn, park' + } +]; + +export const getModelByParam = (param: string) => { + const model = MODELS_CONFIGS.find(m => + m.id === param || + m.hash === param || + m.shortName === param || + m.aliases.includes(param) + ); + + return model; +} + +export const modelsAliases = []; \ No newline at end of file diff --git a/src/modules/sd-images/api/sd-node-api.ts b/src/modules/sd-images/api/sd-node-api.ts new file mode 100644 index 00000000..a4bf30ee --- /dev/null +++ b/src/modules/sd-images/api/sd-node-api.ts @@ -0,0 +1,55 @@ +import { Client } from './sd-node-client' +import { MODELS_CONFIGS, getModelByParam, IModel } from './models-config'; + +const NEGATIVE_PROMPT = '(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation'; + +export class SDNodeApi { + client: Client; + + constructor() { + this.client = new Client() + } + + generateImage = async (prompt: string, model: IModel, seed?: number) => { + const { images } = await this.client.txt2img({ + prompt, + negativePrompt: NEGATIVE_PROMPT, + width: model.baseModel === 'SDXL 1.0' ? 1024 : 512, + height: model.baseModel === 'SDXL 1.0' ? 1024 : 768, + steps: 26, + batchSize: 1, + cfgScale: 7, + seed, + model: model.path + }) + + return images[0]; + } + + generateImagesPreviews = async (prompt: string, model: IModel) => { + const params = { + prompt, + negativePrompt: NEGATIVE_PROMPT, + width: model.baseModel === 'SDXL 1.0' ? 1024 : 512, + height: model.baseModel === 'SDXL 1.0' ? 1024 : 768, + steps: 15, + batchSize: 1, + cfgScale: 7, + model: model.path + }; + + const res = await Promise.all([ + this.client.txt2img(params), + this.client.txt2img(params), + this.client.txt2img(params), + this.client.txt2img(params) + ]); + + return { + images: res.map(r => r.images[0]), + parameters: {}, + all_seeds: res.map(r => r.all_seeds[0]), + info: '' + }; + } +} \ No newline at end of file diff --git a/src/modules/sd-images/sd-node-client.ts b/src/modules/sd-images/api/sd-node-client.ts similarity index 79% rename from src/modules/sd-images/sd-node-client.ts rename to src/modules/sd-images/api/sd-node-client.ts index 953bcd83..ccb34c1e 100644 --- a/src/modules/sd-images/sd-node-client.ts +++ b/src/modules/sd-images/api/sd-node-client.ts @@ -1,7 +1,9 @@ -import { ComfyClient } from '../qrcode/comfy/ComfyClient'; -import config from "../../config"; -import { sleep, waitingExecute } from './utils'; +import { ComfyClient } from '../../qrcode/comfy/ComfyClient'; +import config from "../../../config"; +import { sleep } from '../utils'; import { buildImgPrompt } from './text_to_img_config'; +import { MODELS_CONFIGS } from './models-config'; +import { waitingExecute } from './helpers'; export type Txt2ImgOptions = { hires?: { @@ -31,7 +33,7 @@ export type Txt2ImgOptions = { name: string args?: string[] } - model?: MODELS + model?: string; } export type Txt2ImgResponse = { @@ -41,19 +43,16 @@ export type Txt2ImgResponse = { info: string } -export enum MODELS { - "XL_BASE_1.0" = "sd_xl_base_1.0.safetensors", -} - const getRandomSeed = () => Math.round(Math.random() * 1e15); export class Client { constructor() { } - txt2img = async (options: Txt2ImgOptions): Promise => { + txt2img = async (options: Txt2ImgOptions, serverConfig?: { host: string, wsHost: string }): Promise => { const comfyClient = new ComfyClient({ host: config.comfyHost, - wsHost: config.comfyWsHost + wsHost: config.comfyWsHost, + ...serverConfig }); try { @@ -70,12 +69,11 @@ export class Client { ...options, seed, clientId: comfyClient.clientId, - model: MODELS['XL_BASE_1.0'] }); const r = await comfyClient.queuePrompt(prompt); - const promptResult = await waitingExecute(() => comfyClient.waitingPromptExecution(r.prompt_id), 1000 * 120); + const promptResult = await waitingExecute(() => comfyClient.waitingPromptExecution(r.prompt_id), 1000 * 180); const history = await comfyClient.history(r.prompt_id); @@ -88,7 +86,7 @@ export class Client { return { images, parameters: {}, - all_seeds: history.outputs['9'].images.map((i, idx) => String(seed + idx)), + all_seeds: [String(seed)], info: '' } as Txt2ImgResponse; } catch (e) { diff --git a/src/modules/sd-images/text_to_img_config.ts b/src/modules/sd-images/api/text_to_img_config.ts similarity index 94% rename from src/modules/sd-images/text_to_img_config.ts rename to src/modules/sd-images/api/text_to_img_config.ts index 26a95bad..8bacc2f1 100644 --- a/src/modules/sd-images/text_to_img_config.ts +++ b/src/modules/sd-images/api/text_to_img_config.ts @@ -9,7 +9,7 @@ export function buildImgPrompt(options: Txt2ImgOptions & { clientId: string }) { seed: options.seed, steps: options.steps, cfg: 8, - sampler_name: "euler_ancestral", + sampler_name: "dpmpp_2m", scheduler: "karras", denoise: 1, model: ["4", 0], @@ -25,8 +25,8 @@ export function buildImgPrompt(options: Txt2ImgOptions & { clientId: string }) { }, 5: { inputs: { - width: 1024, - height: 1024, + width: options.width || 1024, + height: options.height || 1024, batch_size: options.batchSize || 1, }, class_type: "EmptyLatentImage", @@ -35,7 +35,7 @@ export function buildImgPrompt(options: Txt2ImgOptions & { clientId: string }) { inputs: { text: options.prompt, clip: ["4", 1] }, class_type: "CLIPTextEncode", }, - 7: { inputs: { text: "", clip: ["4", 1] }, class_type: "CLIPTextEncode" }, + 7: { inputs: { text: options.negativePrompt, clip: ["4", 1] }, class_type: "CLIPTextEncode" }, 8: { inputs: { samples: ["3", 0], vae: ["4", 2] }, class_type: "VAEDecode", @@ -136,7 +136,7 @@ export function buildImgPrompt(options: Txt2ImgOptions & { clientId: string }) { { name: "LATENT", type: "LATENT", links: [2], slot_index: 0 }, ], properties: { "Node name for S&R": "EmptyLatentImage" }, - widgets_values: [1024, 1024, 1], + widgets_values: [options.width || 1024, options.height || 1024, 1], }, { id: 9, @@ -169,11 +169,11 @@ export function buildImgPrompt(options: Txt2ImgOptions & { clientId: string }) { ], properties: { "Node name for S&R": "KSampler" }, widgets_values: [ - 877217196051697, + options.seed, "randomize", - 30, + options.steps, 8, - "euler_ancestral", + "dpmpp_2m", "karras", 1, ], diff --git a/src/modules/sd-images/helpers.ts b/src/modules/sd-images/helpers.ts new file mode 100644 index 00000000..13558455 --- /dev/null +++ b/src/modules/sd-images/helpers.ts @@ -0,0 +1,135 @@ +import config from "../../config"; +import { OnCallBackQueryData, OnMessageContext } from "../types"; +import { getModelByParam, IModel, MODELS_CONFIGS } from "./api"; + +export enum COMMAND { + TEXT_TO_IMAGE = 'image', + TEXT_TO_IMAGES = 'images', + CONSTRUCTOR = 'constructor', + HELP = 'help' +} + +export interface IOperation { + command: COMMAND; + prompt: string; + model: IModel; +} + +const removeSpaceFromBegin = (text: string) => { + if (!text) return ''; + + let idx = 0; + + // const regex = /^[a-zA-Z\d]$/; + + while (!!text[idx] && text[idx] === ' ') idx++; + + return text.slice(idx); +} + +const parsePrompts = (fullText: string): { modelId: string, prompt: string } => { + let modelId = ''; + let prompt: any; + + let text = fullText; + + if (text.startsWith('/') || text.startsWith(',')) { + const startIdx = text.indexOf(' '); + text = startIdx > -1 ? text.slice(startIdx) : ''; + } + + try { + const startIdx = text.indexOf('--model='); + const endIdx = text.indexOf(' ', startIdx); + + if (startIdx > -1) { + prompt = text.split(''); + const modelParamStr = prompt.splice(startIdx, endIdx - startIdx); + + prompt = prompt.join(''); + prompt = removeSpaceFromBegin(prompt); + + modelId = modelParamStr.join('').split('=')[1].replace(/[^a-zA-Z\-\_\d]/g, ""); + } else { + prompt = removeSpaceFromBegin(text); + } + } catch (e) { + console.log('Warning: sd images parse prompts', e); + } + + // console.log({ modelId, prompt }); + + return { modelId, prompt }; +} + +type Context = OnMessageContext | OnCallBackQueryData; + +export const parseCtx = (ctx: Context): IOperation | false => { + try { + if (!ctx.message?.text) { + return false; + } + + let { + modelId, + prompt + } = parsePrompts(ctx.message?.text); + + let model = getModelByParam(modelId); + let command; + + if (ctx.hasCommand('image')) { + command = COMMAND.TEXT_TO_IMAGE; + } + + if (ctx.hasCommand('images')) { + command = COMMAND.TEXT_TO_IMAGES; + } + + if (ctx.hasCommand('all')) { + command = COMMAND.CONSTRUCTOR; + } + + if (ctx.hasCommand('SD')) { + command = COMMAND.HELP; + } + + const startWithCmdSymbol = !!ctx.message?.text?.startsWith('/'); + + if (startWithCmdSymbol) { + const cmd = String(ctx.message?.text?.slice(1).split(' ')[0]); + const modelFromCmd = getModelByParam(cmd); + + if (modelFromCmd) { + command = COMMAND.TEXT_TO_IMAGE; + model = modelFromCmd; + } + } + + const startWithSpecialSymbol = !!ctx.message?.text?.startsWith(','); + + if (startWithSpecialSymbol) { + command = COMMAND.TEXT_TO_IMAGE; + } + + if (!model) { + model = MODELS_CONFIGS[0]; + } + + if (!prompt) { + prompt = model.defaultPrompt; + } + + if (command) { + return { + command, + model, + prompt + } + } + } catch (e) { + console.log('Error: SD images parse prompts', e); + } + + return false; +} \ No newline at end of file diff --git a/src/modules/sd-images/index.ts b/src/modules/sd-images/index.ts index 90cdbe69..266ca3af 100644 --- a/src/modules/sd-images/index.ts +++ b/src/modules/sd-images/index.ts @@ -1,49 +1,30 @@ -import { SDNodeApi } from "./sd-node-api"; -import config from "../../config"; import { InlineKeyboard, InputFile } from "grammy"; import { OnMessageContext, OnCallBackQueryData } from "../types"; -import { sleep, uuidv4 } from "./utils"; -import { showcasePrompts } from "./showcase"; -import { AbortController, AbortSignal } from "grammy/out/shim.node"; - -enum SupportedCommands { - IMAGE = "image", - IMAGES = "images", - // SHOWCASE = 'image_example', -} - -enum SESSION_STEP { - IMAGE_SELECT = "IMAGE_SELECT", - IMAGE_GENERATED = "IMAGE_GENERATED", -} +import { SDImagesBotBase } from './SDImagesBotBase'; +import { COMMAND, parseCtx } from './helpers'; +import { getModelByParam, IModel, MODELS_CONFIGS } from "./api"; +import { uuidv4 } from "./utils"; interface ISession { id: string; author: string; - step: SESSION_STEP; prompt: string; + model: IModel; all_seeds: string[]; + command: COMMAND; } -export class SDImagesBot { - sdNodeApi: SDNodeApi; - - private queue: string[] = []; +export class SDImagesBot extends SDImagesBotBase { private sessions: ISession[] = []; - private showcaseCount = 0; - - constructor() { - this.sdNodeApi = new SDNodeApi(); - } public isSupportedEvent( ctx: OnMessageContext | OnCallBackQueryData ): boolean { - const hasCommand = ctx.hasCommand(Object.values(SupportedCommands)); + const operation = !!parseCtx(ctx); const hasCallbackQuery = this.isSupportedCallbackQuery(ctx); - return hasCallbackQuery || hasCommand; + return hasCallbackQuery || operation; } public getEstimatedPrice(ctx: any) { @@ -66,136 +47,85 @@ export class SDImagesBot { ctx: OnMessageContext | OnCallBackQueryData, refundCallback: (reason?: string) => void ) { - if (!this.isSupportedEvent(ctx)) { - console.log(`### unsupported command ${ctx.message?.text}`); - refundCallback("Unsupported command"); - await ctx.reply("### unsupported command"); + if (this.isSupportedCallbackQuery(ctx)) { + this.onImgSelected(ctx, refundCallback); return; } - if (ctx.hasCommand(SupportedCommands.IMAGE)) { - return this.onImageCmd(ctx, refundCallback); - } - - if (ctx.hasCommand(SupportedCommands.IMAGES)) { - return this.onImagesCmd(ctx, refundCallback); - } - - // if (ctx.hasCommand(SupportedCommands.SHOWCASE)) { - // this.onShowcaseCmd(ctx); - // return; - // } + const operation = parseCtx(ctx); - if (this.isSupportedCallbackQuery(ctx)) { - return this.onImgSelected(ctx, refundCallback); + if (!operation) { + console.log(`### unsupported command ${ctx.message?.text}`); + ctx.reply("### unsupported command"); + return refundCallback("Unsupported command"); } - console.log(`### unsupported command`); - await ctx.reply("### unsupported command"); - } - - onImageCmd = async ( - ctx: OnMessageContext | OnCallBackQueryData, - refundCallback: (reason?: string) => void - ) => { - const uuid = uuidv4(); - - // /qr s.country/ai astronaut, exuberant, anime girl, smile, sky, colorful - try { - const prompt: any = ctx.match - ? ctx.match - : config.stableDiffusion.imageDefaultMessage; - - const authorObj = await ctx.getAuthor(); - const author = `@${authorObj.user.username}`; - - if (!prompt) { - refundCallback("Wrong prompts"); - await ctx.reply(`${author} please add prompt to your message`); + switch (operation.command) { + case COMMAND.TEXT_TO_IMAGE: + this.generateImage( + ctx, + refundCallback, + operation.prompt, + operation.model + ); return; - } - ctx.chatAction = "upload_photo"; - this.queue.push(uuid); - let idx = this.queue.findIndex((v) => v === uuid); + case COMMAND.TEXT_TO_IMAGES: + this.onImagesCmd( + ctx, + refundCallback, + operation.prompt, + operation.model + ); + return; - if (idx >= 0) { - await ctx.reply( - `You are #${idx + 1} in line, wait ~${idx * 30 + 30} seconds` + case COMMAND.CONSTRUCTOR: + this.onConstructorCmd( + ctx, + refundCallback, + operation.prompt, + operation.model ); - } + return; - // waiting queue - while (idx !== 0) { - await sleep(3000 * this.queue.findIndex((v) => v === uuid)); - idx = this.queue.findIndex((v) => v === uuid); - } + case COMMAND.HELP: + await ctx.reply('Stable Diffusion Models: \n'); - const imageBuffer = await this.sdNodeApi.generateImage(prompt); - await ctx.replyWithPhoto(new InputFile(imageBuffer), { - caption: `/image ${prompt}`, - }); + for (let i = 0; i < MODELS_CONFIGS.length; i++) { + const model = MODELS_CONFIGS[i]; - // await ctx.reply(`/image ${prompt}`); - } catch (e: any) { - console.log(e); - this.queue = this.queue.filter((v) => v !== uuid); - ctx.chatAction = null; - refundCallback(e); - await ctx.reply(`Error: something went wrong...`); + await ctx.reply(`${model.name}: ${model.link} \n \nUsing: /${model.aliases[0]} /${model.aliases[1]} /${model.aliases[2]} \n`); + } + return; } - this.queue = this.queue.filter((v) => v !== uuid); - }; + console.log(`### unsupported command`); + ctx.reply("### unsupported command"); + } onImagesCmd = async ( ctx: OnMessageContext | OnCallBackQueryData, - refundCallback: (reason?: string) => void + refundCallback: (reason?: string) => void, + prompt: string, + model: IModel ) => { const uuid = uuidv4(); try { - const prompt: any = ctx.match - ? ctx.match - : config.stableDiffusion.imagesDefaultMessage; - const authorObj = await ctx.getAuthor(); const author = `@${authorObj.user.username}`; - if (!prompt) { - refundCallback("Wrong prompts"); - await ctx.reply(`${author} please add prompt to your message`); - return; - } - ctx.chatAction = "upload_photo"; - this.queue.push(uuid); + await this.waitingQueue(uuid, ctx); - let idx = this.queue.findIndex((v) => v === uuid); - - if (idx >= 0) { - await ctx.reply( - `You are #${idx + 1} in line, wait ~${idx * 30 + 30} seconds` - ); - } - - // waiting queue - while (idx !== 0) { - await sleep(3000 * this.queue.findIndex((v) => v === uuid)); - - idx = this.queue.findIndex((v) => v === uuid); - } - - // ctx.reply(`${author} starting to generate your images`); - const res = await this.sdNodeApi.generateImagesPreviews(prompt); - - // res.images.map(img => new InputFile(Buffer.from(img, 'base64'))); + const res = await this.sdNodeApi.generateImagesPreviews(prompt, model); const newSession: ISession = { id: uuidv4(), author, - prompt: String(prompt), - step: SESSION_STEP.IMAGE_SELECT, + prompt: prompt, all_seeds: res.all_seeds, + model, + command: COMMAND.TEXT_TO_IMAGES }; this.sessions.push(newSession); @@ -208,23 +138,19 @@ export class SDImagesBot { })) ); - // await ctx.reply("Please choose 1 of 4 images for next high quality generation", { - // parse_mode: "HTML", - // reply_markup: new InlineKeyboard() - // .text("1", `${newSession.id}_1`) - // .text("2", `${newSession.id}_2`) - // .text("3", `${newSession.id}_3`) - // .text("4", `${newSession.id}_4`) - // .row() - // }); + await ctx.reply("Please choose 1 of 4 images for next high quality generation", { + parse_mode: "HTML", + reply_markup: new InlineKeyboard() + .text("1", `${newSession.id}_1`) + .text("2", `${newSession.id}_2`) + .text("3", `${newSession.id}_3`) + .text("4", `${newSession.id}_4`) + .row() + }); } catch (e: any) { - ctx.chatAction = null; - console.log(e); - this.queue = this.queue.filter((v) => v !== uuid); - ctx.chatAction = null; + ctx.reply(`Error: something went wrong...`); refundCallback(e.message); - await ctx.reply(`Error: something went wrong...`); } this.queue = this.queue.filter((v) => v !== uuid); @@ -244,9 +170,11 @@ export class SDImagesBot { return; } - const [sessionId, imageNumber] = ctx.callbackQuery.data.split("_"); + const [sessionId, ...paramsArray] = ctx.callbackQuery.data.split("_"); + + const params = paramsArray.join('_'); - if (!sessionId || !imageNumber) { + if (!sessionId || !params) { refundCallback("Wrong params"); return; } @@ -257,44 +185,92 @@ export class SDImagesBot { refundCallback("Wrong author"); return; } - ctx.chatAction = "upload_photo"; - // ctx.reply(`${author} starting to generate your image ${imageNumber} in high quality`); - const imageBuffer = await this.sdNodeApi.generateImageFull( - session.prompt, - +session.all_seeds[+imageNumber - 1] - ); - await ctx.replyWithPhoto(new InputFile(imageBuffer), { - caption: `/image ${session.prompt}`, - }); + let model; + + if (session.command === COMMAND.CONSTRUCTOR) { + model = getModelByParam(params); + + if (!model) { + console.log("wrong model"); + refundCallback("Wrong callbackQuery"); + return; + } + + this.generateImage( + ctx, + refundCallback, + session.prompt, + model + ); + + return; + } + + if (session.command === COMMAND.TEXT_TO_IMAGES) { + this.generateImage( + ctx, + refundCallback, + session.prompt, + session.model, + Number(session.all_seeds[+params - 1]) + ); + + return; + } } catch (e: any) { - ctx.chatAction = null; console.log(e); + ctx.reply(`Error: something went wrong...`); refundCallback(e.message); - await ctx.reply(`Error: something went wrong...`); } } - onShowcaseCmd = async (ctx: OnMessageContext | OnCallBackQueryData) => { - const uuid = uuidv4(); - + onConstructorCmd = async ( + ctx: OnMessageContext | OnCallBackQueryData, + refundCallback: (reason?: string) => void, + prompt: string, + model: IModel + ) => { try { - if (this.showcaseCount >= showcasePrompts.length) { - this.showcaseCount = 0; - } + const authorObj = await ctx.getAuthor(); + const author = `@${authorObj.user.username}`; - const prompt = showcasePrompts[this.showcaseCount++]; + const newSession: ISession = { + id: uuidv4(), + author, + prompt: String(prompt), + all_seeds: [], + model, + command: COMMAND.CONSTRUCTOR + }; + + this.sessions.push(newSession); - const imageBuffer = await this.sdNodeApi.generateImage(prompt); - ctx.chatAction = "upload_photo"; - await ctx.replyWithPhoto(new InputFile(imageBuffer)); + const buttonsPerRow = 2; + let rowCount = buttonsPerRow; + const keyboard = new InlineKeyboard(); - await ctx.reply(`/image ${prompt}`); + for (let i = 0; i < MODELS_CONFIGS.length; i++) { + keyboard.text(MODELS_CONFIGS[i].name, `${newSession.id}_${MODELS_CONFIGS[i].hash}`); + + rowCount--; + + if (!rowCount) { + keyboard.row(); + rowCount = buttonsPerRow; + } + } + + keyboard.row(); + + await ctx.reply(prompt, { + parse_mode: "HTML", + reply_markup: keyboard + }); } catch (e: any) { - ctx.chatAction = null; console.log(e); - await ctx.reply(`Error: something went wrong...`); - // throw new Error(e?.message); + ctx.reply(`Error: something went wrong...`); + refundCallback(e); } }; -} +} \ No newline at end of file diff --git a/src/modules/sd-images/sd-node-api.ts b/src/modules/sd-images/sd-node-api.ts deleted file mode 100644 index 7165a2dd..00000000 --- a/src/modules/sd-images/sd-node-api.ts +++ /dev/null @@ -1,54 +0,0 @@ -// import sdwebui, { Client, SamplingMethod } from 'node-sd-webui' -import { Client } from './sd-node-client' - -const NEGATIVE_PROMPT = 'ugly, deformed, watermark'; - -export class SDNodeApi { - client: Client; - - constructor() { - this.client = new Client() - } - - generateImage = async (prompt: string) => { - const { images } = await this.client.txt2img({ - prompt, - negativePrompt: NEGATIVE_PROMPT, - width: 1024, - height: 1024, - steps: 30, - batchSize: 1, - }) - - return images[0]; // Buffer.from(images[0], 'base64'); - } - - generateImageFull = async (prompt: string, seed: number) => { - const { images } = await this.client.txt2img({ - prompt, - negativePrompt: NEGATIVE_PROMPT, - width: 1024, - height: 1024, - steps: 30, - batchSize: 1, - cfgScale: 7, - seed - }) - - return images[0]; - } - - generateImagesPreviews = async (prompt: string) => { - const res = await this.client.txt2img({ - prompt, - negativePrompt: NEGATIVE_PROMPT, - width: 1024, - height: 1024, - steps: 15, - batchSize: 4, - cfgScale: 10, - }) - - return res; - } -} \ No newline at end of file diff --git a/src/modules/sd-images/utils.ts b/src/modules/sd-images/utils.ts index dff7a6e6..89fc4789 100644 --- a/src/modules/sd-images/utils.ts +++ b/src/modules/sd-images/utils.ts @@ -6,13 +6,4 @@ export const uuidv4 = () => { return [1, 2, 3, 4].map(() => rand()).join('-'); }; -export const sleep = (ms: number) => new Promise(res => setTimeout(res, ms)); - -export const waitingExecute = (fn: () => Promise, ms: number) => new Promise((resolve, reject) => { - const timeoutId = setTimeout(() => { - console.error('SD images Error: waitingExecute time is up'); - reject('Error: waitingExecute time is up'); - }, ms); - - fn().then(resolve).catch(reject).finally(() => clearTimeout(timeoutId)) -}); \ No newline at end of file +export const sleep = (ms: number) => new Promise(res => setTimeout(res, ms)); \ No newline at end of file