diff --git a/src/bot.ts b/src/bot.ts index 4b9f168..738661c 100644 --- a/src/bot.ts +++ b/src/bot.ts @@ -1,12 +1,12 @@ -import { BotCommand, BotEvent, QueueItem } from './types' +import bPromise from 'bluebird' import { Client, GatewayIntentBits, Partials, REST, Routes } from 'discord.js' +import fs from 'fs' import { isEmpty, isNil } from 'ramda' - -import Config from './config/bot' import { Logger } from 'winston' +import Config from './config/bot' import StableDiffusion from './modules/stableDiffusion' -import bPromise from 'bluebird' -import fs from 'fs' +import { BotCommand, BotEvent, QueueItems } from './types' + export class Bot extends Client { constructor(config: typeof Config, logger: Logger, options: {[k: string]: string}) { @@ -28,18 +28,18 @@ export class Bot extends Client { /** * Array of QueueItems awaiting processing. */ - public queue: QueueItem[] + public queue: QueueItems.QueueItemInstances[] /** * Array of previously used QueueItems. */ - public queueItemReferences: QueueItem[] + public queueItemReferences: QueueItems.QueueItemInstances[] /** * Current QueueItem processing. * TODO: Support multi-processing */ - public processing: QueueItem | null + public processing: QueueItems.QueueItemInstances | null /** * The bot's running config @@ -71,14 +71,54 @@ export class Bot extends Client { * @param queueItem QueueItem to add to queue. * @returns Number of items in queue. */ - public addQueue = (queueItem: QueueItem): number => this.queue.push(queueItem) + public addQueuedQueueItem = (queueItem: QueueItems.QueueItemInstances): number => this.queue.push(queueItem) + + /** + * Updates content in a queue or reference QueueItem + * @param queueItem QueueItem to add to queue. + * @returns Dictionary of updated QueueItems. + */ + public updateQueueItem = ( + updateFunction: (queueItem: QueueItems.QueueItemInstances) => QueueItems.QueueItemInstances, + uuid: string + ): { + queue: QueueItems.QueueItemInstances | undefined, + referenceQueue: QueueItems.QueueItemInstances | undefined + } => { + const queueIndex = this.queue.findIndex(queueItem => queueItem.uuid === uuid) + const referenceQueueIndex = this.queueItemReferences.findIndex(queueItem => queueItem.uuid === uuid) + + if (queueIndex === -1 && referenceQueueIndex === -1) throw new Error('No QueueItem in queue or queue references to update.') + + const updated: { + queue: QueueItems.QueueItemInstances | undefined, + referenceQueue: QueueItems.QueueItemInstances | undefined + } = { + queue: undefined, + referenceQueue: undefined + } + + if (queueIndex !== -1) { + const queueItem = this.queue[queueIndex] + this.queue[queueIndex] = updateFunction(queueItem) + updated.queue = this.queue[queueIndex] + } + + if (referenceQueueIndex !== -1) { + const referenceQueueItem = this.queueItemReferences[referenceQueueIndex] + this.queueItemReferences[referenceQueueIndex] = updateFunction(referenceQueueItem) + updated.referenceQueue = this.queueItemReferences[referenceQueueIndex] + } + + return updated + } /** * Deletes a QueueItem from queue, adds it to queueItemReferences. Throws if it can't find a QueueItem for uuid. * @param uuid UUID of QueueItem. * @returns Array of deleted QueueItems. */ - public removeQueue = (uuid: string): QueueItem[] => { + public removeQueuedQueueItem = (uuid: string): QueueItems.QueueItemInstances[] => { const queueIndex = this.queue.findIndex(queueItem => queueItem.uuid === uuid) if (queueIndex === -1) return [] @@ -92,7 +132,7 @@ export class Bot extends Client { * @param uuid UUID of QueueItem. * @returns Array of deleted QueueItems. */ - public removeQueueItemReference = (uuid: string): QueueItem[] => { + public removeQueueItemReference = (uuid: string): QueueItems.QueueItemInstances[] => { const queueIndex = this.queueItemReferences.findIndex(queueItem => queueItem.uuid === uuid) if (queueIndex === -1) return [] @@ -102,28 +142,28 @@ export class Bot extends Client { /** * Finds a QueueItem in queue storage by uuid. * @param uuid UUID of QueueItem. - * @returns Array of found QueueItems. + * @returns Found QueueItem. */ - public findQueue = (uuid: string): QueueItem | undefined => { + public findQueue = (uuid: string): QueueItems.QueueItemInstances | undefined => { return this.queue.find(queueItem => queueItem.uuid === uuid) } /** * Finds a QueueItem in reference storage by uuid. * @param uuid UUID of QueueItem. - * @returns Array of found QueueItems. + * @returns Found QueueItem. */ - public findQueueItemReference = (uuid: string): QueueItem | undefined => { + public findQueueItemReference = (uuid: string): QueueItems.QueueItemInstances | undefined => { return this.queueItemReferences.find(queueItem => queueItem.uuid === uuid) } /** - * Finds the latest QueueItem in reference storage by message ID. + * Finds the latest QueueItem in reference storage by message snowflake. * @param uuid UUID of QueueItem. * @returns Array of found QueueItems. */ - public findLatestQueueItemReferenceByMessageID = (messageId: string): QueueItem | undefined => { - const queueItems = this.queueItemReferences.filter(queueItem => queueItem.messageId === messageId) + public findLatestQueueItemReferenceByMessageSnowflake = (snowflake: string): QueueItems.QueueItemInstances | undefined => { + const queueItems = this.queueItemReferences.filter(queueItem => queueItem.discordMessageSnowflake === snowflake) return queueItems[queueItems.length - 1] } } diff --git a/src/embeds/addedToQueue.ts b/src/embeds/addedToQueue.ts index f4f9292..229be90 100644 --- a/src/embeds/addedToQueue.ts +++ b/src/embeds/addedToQueue.ts @@ -1,64 +1,62 @@ import { EmbedBuilder, codeBlock, ActionRowBuilder } from 'discord.js' -import { BotEmbed, QueueItem } from '../types' +import { BotEmbed, QueueItems } from '../types' -export enum QueueType { - Instant, - Queued -} - -export default function(queueType: QueueType, queueItem: QueueItem, queuePosition?: number): BotEmbed { +export function addedToInstantQueue(queueItem: QueueItems.QueueItemInstances): BotEmbed { let embeds: EmbedBuilder[] = [] let components: ActionRowBuilder[] = [] - switch (queueType) { - case QueueType.Queued: - embeds = [ - new EmbedBuilder() - .setColor('#030354') - .setTitle('Prompt added to queue') - .setDescription(`Your prompt has been added to the image generation queue and is currently in position number ${queuePosition}.`) - .setThumbnail('https://i.imgur.com/pXmPAAG.gif') - .setTimestamp() - .addFields([ - {name: 'Prompt', value: queueItem.prediction.prompt ? codeBlock(queueItem.prediction.prompt) : 'Not Supplied', inline: false}, - {name: 'Width', value: queueItem.prediction.width.toString(), inline: true}, - {name: 'Height', value: queueItem.prediction.height.toString(), inline: true}, - {name: 'Prompt Strength', value: queueItem.prediction.promptStrength.toString(), inline: true}, - {name: 'Steps', value: queueItem.prediction.numInferenceSteps.toString(), inline: true}, - {name: 'Guidance Scale', value: queueItem.prediction.guidanceScale.toString(), inline: true}, - {name: 'Has Mask', value: typeof queueItem.prediction.mask === 'undefined' ? 'No' : 'Yes', inline: true}, - {name: 'Has Base Image', value: typeof queueItem.prediction.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, - {name: 'Seed', value: codeBlock(queueItem.seed.toString()), inline: false}, - {name: 'Prompt ID', value: codeBlock(queueItem.uuid), inline: false} - ]) - ] - break + embeds = [ + new EmbedBuilder() + .setColor('#030354') + .setTitle('Prompt will generate shortly') + .setDescription(`Your prompt is currently first in the queue and will begin generating shortly.`) + .setThumbnail('https://i.imgur.com/pXmPAAG.gif') + .setTimestamp() + .addFields([ + {name: 'Prompt', value: queueItem.prompt ? codeBlock(queueItem.prompt) : 'Not Supplied', inline: false}, + {name: 'Width', value: queueItem.width.toString(), inline: true}, + {name: 'Height', value: queueItem.height.toString(), inline: true}, + {name: 'Prompt Strength', value: queueItem.promptStrength.toString(), inline: true}, + {name: 'Steps', value: queueItem.numInferenceSteps.toString(), inline: true}, + {name: 'Guidance Scale', value: queueItem.guidanceScale.toString(), inline: true}, + {name: 'Has Mask', value: typeof queueItem.mask === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Has Base Image', value: typeof queueItem.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Seed', value: codeBlock(queueItem.seed.toString()), inline: false}, + {name: 'Prompt ID', value: codeBlock(queueItem.uuid), inline: false} + ]) + ] - case QueueType.Instant: - default: - embeds = [ - new EmbedBuilder() - .setColor('#030354') - .setTitle('Prompt will generate shortly') - .setDescription(`Your prompt is currently first in the queue and will begin generating shortly.`) - .setThumbnail('https://i.imgur.com/pXmPAAG.gif') - .setTimestamp() - .addFields([ - {name: 'Prompt', value: queueItem.prediction.prompt ? codeBlock(queueItem.prediction.prompt) : 'Not Supplied', inline: false}, - {name: 'Width', value: queueItem.prediction.width.toString(), inline: true}, - {name: 'Height', value: queueItem.prediction.height.toString(), inline: true}, - {name: 'Prompt Strength', value: queueItem.prediction.promptStrength.toString(), inline: true}, - {name: 'Steps', value: queueItem.prediction.numInferenceSteps.toString(), inline: true}, - {name: 'Guidance Scale', value: queueItem.prediction.guidanceScale.toString(), inline: true}, - {name: 'Has Mask', value: typeof queueItem.prediction.mask === 'undefined' ? 'No' : 'Yes', inline: true}, - {name: 'Has Base Image', value: typeof queueItem.prediction.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, - {name: 'Seed', value: codeBlock(queueItem.seed.toString()), inline: false}, - {name: 'Prompt ID', value: codeBlock(queueItem.uuid), inline: false} - ]) - ] + return { + embeds, + components } +} + +export function addedToQueue(queueItem: QueueItems.QueueItemInstances, queuePosition: number): BotEmbed { + let embeds: EmbedBuilder[] = [] + let components: ActionRowBuilder[] = [] + + embeds = [ + new EmbedBuilder() + .setColor('#030354') + .setTitle('Prompt added to queue') + .setDescription(`Your prompt has been added to the image generation queue and is currently in position number ${queuePosition}.`) + .setThumbnail('https://i.imgur.com/pXmPAAG.gif') + .setTimestamp() + .addFields([ + {name: 'Prompt', value: queueItem.prompt ? codeBlock(queueItem.prompt) : 'Not Supplied', inline: false}, + {name: 'Width', value: queueItem.width.toString(), inline: true}, + {name: 'Height', value: queueItem.height.toString(), inline: true}, + {name: 'Prompt Strength', value: queueItem.promptStrength.toString(), inline: true}, + {name: 'Steps', value: queueItem.numInferenceSteps.toString(), inline: true}, + {name: 'Guidance Scale', value: queueItem.guidanceScale.toString(), inline: true}, + {name: 'Has Mask', value: typeof queueItem.mask === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Has Base Image', value: typeof queueItem.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Seed', value: codeBlock(queueItem.seed.toString()), inline: false}, + {name: 'Prompt ID', value: codeBlock(queueItem.uuid), inline: false} + ]) + ] - return { embeds, components diff --git a/src/embeds/imageResult.ts b/src/embeds/imageResult.ts index 035317c..ea9e5a5 100644 --- a/src/embeds/imageResult.ts +++ b/src/embeds/imageResult.ts @@ -1,12 +1,12 @@ import { EmbedBuilder, ActionRowBuilder, ButtonBuilder, ButtonStyle, codeBlock } from 'discord.js' -import { BotEmbed, QueueItem, QueueItemType } from '../types' +import { BotEmbed, QueueItems } from '../types' -export default function(queueItem: QueueItem): BotEmbed { +export default function(queueItem: QueueItems.QueueItemInstances): BotEmbed { let embeds: EmbedBuilder[] = [] let components: ActionRowBuilder[] = [] switch (queueItem.type) { - case QueueItemType.Quick: + case QueueItems.QueueItemTypes.Quick: embeds = [ new EmbedBuilder() .setColor('#6aaa64') @@ -14,14 +14,14 @@ export default function(queueItem: QueueItem): BotEmbed { .setDescription('Result for parameters\n\nGenerated using the quick imagine command. Check out `/imagine` for new features!') .setTimestamp() .addFields([ - {name: 'Prompt', value: queueItem.prediction.prompt ? codeBlock(queueItem.prediction.prompt) : 'Not Supplied', inline: false}, - {name: 'Width', value: queueItem.prediction.width.toString(), inline: true}, - {name: 'Height', value: queueItem.prediction.height.toString(), inline: true}, - {name: 'Prompt Strength', value: queueItem.prediction.promptStrength.toString(), inline: true}, - {name: 'Steps', value: queueItem.prediction.numInferenceSteps.toString(), inline: true}, - {name: 'Guidance Scale', value: queueItem.prediction.guidanceScale.toString(), inline: true}, - {name: 'Has Mask', value: typeof queueItem.prediction.mask === 'undefined' ? 'No' : 'Yes', inline: true}, - {name: 'Has Base Image', value: typeof queueItem.prediction.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Prompt', value: queueItem.prompt ? codeBlock(queueItem.prompt) : 'Not Supplied', inline: false}, + {name: 'Width', value: queueItem.width.toString(), inline: true}, + {name: 'Height', value: queueItem.height.toString(), inline: true}, + {name: 'Prompt Strength', value: queueItem.promptStrength.toString(), inline: true}, + {name: 'Steps', value: queueItem.numInferenceSteps.toString(), inline: true}, + {name: 'Guidance Scale', value: queueItem.guidanceScale.toString(), inline: true}, + {name: 'Has Mask', value: typeof queueItem.mask === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Has Base Image', value: typeof queueItem.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, {name: 'Seed', value: codeBlock(queueItem.seed.toString()), inline: false}, {name: 'Prompt ID', value: codeBlock(queueItem.uuid), inline: false} ]) @@ -29,7 +29,7 @@ export default function(queueItem: QueueItem): BotEmbed { ] break - case QueueItemType.Upscaled: + case QueueItems.QueueItemTypes.Upscaled: embeds = [ new EmbedBuilder() .setColor('#6aaa64') @@ -37,12 +37,12 @@ export default function(queueItem: QueueItem): BotEmbed { .setDescription('Result for upscale') .setTimestamp() .addFields([ - {name: 'Prompt', value: queueItem.prediction.prompt ? codeBlock(queueItem.prediction.prompt) : 'Not Supplied', inline: false}, - {name: 'Prompt Strength', value: queueItem.prediction.promptStrength.toString(), inline: true}, - {name: 'Steps', value: queueItem.prediction.numInferenceSteps.toString(), inline: true}, - {name: 'Guidance Scale', value: queueItem.prediction.guidanceScale.toString(), inline: true}, - {name: 'Has Mask', value: typeof queueItem.prediction.mask === 'undefined' ? 'No' : 'Yes', inline: true}, - {name: 'Has Base Image', value: typeof queueItem.prediction.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Prompt', value: queueItem.prompt ? codeBlock(queueItem.prompt) : 'Not Supplied', inline: false}, + {name: 'Prompt Strength', value: queueItem.promptStrength.toString(), inline: true}, + {name: 'Steps', value: queueItem.numInferenceSteps.toString(), inline: true}, + {name: 'Guidance Scale', value: queueItem.guidanceScale.toString(), inline: true}, + {name: 'Has Mask', value: typeof queueItem.mask === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Has Base Image', value: typeof queueItem.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, {name: 'Seed', value: codeBlock(queueItem.seed.toString()), inline: false}, {name: 'Prompt ID', value: codeBlock(queueItem.uuid), inline: false} ]) @@ -50,7 +50,7 @@ export default function(queueItem: QueueItem): BotEmbed { ] break - case QueueItemType.Variant: + case QueueItems.QueueItemTypes.Variant: embeds = [ new EmbedBuilder() .setColor('#6aaa64') @@ -58,12 +58,12 @@ export default function(queueItem: QueueItem): BotEmbed { .setDescription('Variant results are displayed below') .setTimestamp() .addFields([ - {name: 'Prompt', value: queueItem.prediction.prompt ? codeBlock(queueItem.prediction.prompt) : 'Not Supplied', inline: false}, - {name: 'Prompt Strength', value: queueItem.prediction.promptStrength.toString(), inline: true}, - {name: 'Steps', value: queueItem.prediction.numInferenceSteps.toString(), inline: true}, - {name: 'Guidance Scale', value: queueItem.prediction.guidanceScale.toString(), inline: true}, - {name: 'Has Mask', value: typeof queueItem.prediction.mask === 'undefined' ? 'No' : 'Yes', inline: true}, - {name: 'Has Base Image', value: typeof queueItem.prediction.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Prompt', value: queueItem.prompt ? codeBlock(queueItem.prompt) : 'Not Supplied', inline: false}, + {name: 'Prompt Strength', value: queueItem.promptStrength.toString(), inline: true}, + {name: 'Steps', value: queueItem.numInferenceSteps.toString(), inline: true}, + {name: 'Guidance Scale', value: queueItem.guidanceScale.toString(), inline: true}, + {name: 'Has Mask', value: typeof queueItem.mask === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Has Base Image', value: typeof queueItem.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, {name: 'Seed', value: codeBlock(queueItem.seed.toString()), inline: false}, {name: 'Prompt ID', value: codeBlock(queueItem.uuid), inline: false} ]) @@ -88,8 +88,8 @@ export default function(queueItem: QueueItem): BotEmbed { ] break - case QueueItemType.Regenerated: - case QueueItemType.Default: + case QueueItems.QueueItemTypes.Regenerated: + case QueueItems.QueueItemTypes.Default: default: embeds = [ new EmbedBuilder() @@ -98,12 +98,12 @@ export default function(queueItem: QueueItem): BotEmbed { .setDescription('Result for parameters') .setTimestamp() .addFields([ - {name: 'Prompt', value: queueItem.prediction.prompt ? codeBlock(queueItem.prediction.prompt) : 'Not Supplied', inline: false}, - {name: 'Prompt Strength', value: queueItem.prediction.promptStrength.toString(), inline: true}, - {name: 'Steps', value: queueItem.prediction.numInferenceSteps.toString(), inline: true}, - {name: 'Guidance Scale', value: queueItem.prediction.guidanceScale.toString(), inline: true}, - {name: 'Has Mask', value: typeof queueItem.prediction.mask === 'undefined' ? 'No' : 'Yes', inline: true}, - {name: 'Has Base Image', value: typeof queueItem.prediction.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Prompt', value: queueItem.prompt ? codeBlock(queueItem.prompt) : 'Not Supplied', inline: false}, + {name: 'Prompt Strength', value: queueItem.promptStrength.toString(), inline: true}, + {name: 'Steps', value: queueItem.numInferenceSteps.toString(), inline: true}, + {name: 'Guidance Scale', value: queueItem.guidanceScale.toString(), inline: true}, + {name: 'Has Mask', value: typeof queueItem.mask === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Has Base Image', value: typeof queueItem.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, {name: 'Seed', value: codeBlock(queueItem.seed.toString()), inline: false}, {name: 'Prompt ID', value: codeBlock(queueItem.uuid), inline: false} ]) diff --git a/src/embeds/imageSelectPrompt.ts b/src/embeds/imageSelectPrompt.ts index 78df5b8..865a1dd 100644 --- a/src/embeds/imageSelectPrompt.ts +++ b/src/embeds/imageSelectPrompt.ts @@ -1,7 +1,8 @@ import { ActionRowBuilder, EmbedBuilder, SelectMenuBuilder } from 'discord.js' -import { BotEmbed, QueueItem } from '../types' +import { BotEmbed, QueueItems } from '../types' -export default function(queueItem: QueueItem): BotEmbed { +export default function(queueItem: QueueItems.QueueItemInstances): BotEmbed { + console.log(`Trigger: image-select-prompt-${queueItem.type}-${queueItem.uuid}`) return { embeds: [ new EmbedBuilder() @@ -15,7 +16,7 @@ export default function(queueItem: QueueItem): BotEmbed { new ActionRowBuilder() .addComponents([ new SelectMenuBuilder() - .setCustomId(`image-select-prompt-${queueItem.type}`) + .setCustomId(`image-select-prompt-${queueItem.type}-${queueItem.uuid}`) .setPlaceholder('Select an image') .addOptions( { diff --git a/src/embeds/processingPrompt.ts b/src/embeds/processingPrompt.ts index c6a594b..c9d980a 100644 --- a/src/embeds/processingPrompt.ts +++ b/src/embeds/processingPrompt.ts @@ -1,12 +1,12 @@ import { EmbedBuilder, codeBlock, ActionRowBuilder } from 'discord.js' -import { BotEmbed, QueueItem, QueueItemType } from '../types' +import { BotEmbed, QueueItems } from '../types' -export default function(queueItem: QueueItem): BotEmbed { +export default function(queueItem: QueueItems.QueueItemInstances): BotEmbed { let embeds: EmbedBuilder[] = [] let components: ActionRowBuilder[] = [] switch (queueItem.type) { - case QueueItemType.Quick: + case QueueItems.QueueItemTypes.Quick: embeds = [ new EmbedBuilder() .setColor('#030354') @@ -15,19 +15,19 @@ export default function(queueItem: QueueItem): BotEmbed { .setThumbnail('https://i.gifer.com/2RNf.gif') .setTimestamp() .addFields([ - {name: 'Prompt', value: queueItem.prediction.prompt ? codeBlock(queueItem.prediction.prompt) : 'Not Supplied', inline: false}, - {name: 'Prompt Strength', value: queueItem.prediction.promptStrength.toString(), inline: true}, - {name: 'Steps', value: queueItem.prediction.numInferenceSteps.toString(), inline: true}, - {name: 'Guidance Scale', value: queueItem.prediction.guidanceScale.toString(), inline: true}, - {name: 'Has Mask', value: typeof queueItem.prediction.mask === 'undefined' ? 'No' : 'Yes', inline: true}, - {name: 'Has Base Image', value: typeof queueItem.prediction.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Prompt', value: queueItem.prompt ? codeBlock(queueItem.prompt) : 'Not Supplied', inline: false}, + {name: 'Prompt Strength', value: queueItem.promptStrength.toString(), inline: true}, + {name: 'Steps', value: queueItem.numInferenceSteps.toString(), inline: true}, + {name: 'Guidance Scale', value: queueItem.guidanceScale.toString(), inline: true}, + {name: 'Has Mask', value: typeof queueItem.mask === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Has Base Image', value: typeof queueItem.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, {name: 'Seed', value: codeBlock(queueItem.seed.toString()), inline: false}, {name: 'Prompt ID', value: codeBlock(queueItem.uuid), inline: false} ]) ] break - case QueueItemType.Regenerated: + case QueueItems.QueueItemTypes.Regenerated: embeds = [ new EmbedBuilder() .setColor('#030354') @@ -36,19 +36,19 @@ export default function(queueItem: QueueItem): BotEmbed { .setThumbnail('https://i.gifer.com/2RNf.gif') .setTimestamp() .addFields([ - {name: 'Prompt', value: queueItem.prediction.prompt ? codeBlock(queueItem.prediction.prompt) : 'Not Supplied', inline: false}, - {name: 'Prompt Strength', value: queueItem.prediction.promptStrength.toString(), inline: true}, - {name: 'Steps', value: queueItem.prediction.numInferenceSteps.toString(), inline: true}, - {name: 'Guidance Scale', value: queueItem.prediction.guidanceScale.toString(), inline: true}, - {name: 'Has Mask', value: typeof queueItem.prediction.mask === 'undefined' ? 'No' : 'Yes', inline: true}, - {name: 'Has Base Image', value: typeof queueItem.prediction.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Prompt', value: queueItem.prompt ? codeBlock(queueItem.prompt) : 'Not Supplied', inline: false}, + {name: 'Prompt Strength', value: queueItem.promptStrength.toString(), inline: true}, + {name: 'Steps', value: queueItem.numInferenceSteps.toString(), inline: true}, + {name: 'Guidance Scale', value: queueItem.guidanceScale.toString(), inline: true}, + {name: 'Has Mask', value: typeof queueItem.mask === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Has Base Image', value: typeof queueItem.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, {name: 'Seed', value: codeBlock(queueItem.seed.toString()), inline: false}, {name: 'Prompt ID', value: codeBlock(queueItem.uuid), inline: false} ]) ] break - case QueueItemType.Variant: + case QueueItems.QueueItemTypes.Variant: embeds = [ new EmbedBuilder() .setColor('#030354') @@ -57,19 +57,19 @@ export default function(queueItem: QueueItem): BotEmbed { .setThumbnail('https://i.gifer.com/2RNf.gif') .setTimestamp() .addFields([ - {name: 'Prompt', value: queueItem.prediction.prompt ? codeBlock(queueItem.prediction.prompt) : 'Not Supplied', inline: false}, - {name: 'Prompt Strength', value: queueItem.prediction.promptStrength.toString(), inline: true}, - {name: 'Steps', value: queueItem.prediction.numInferenceSteps.toString(), inline: true}, - {name: 'Guidance Scale', value: queueItem.prediction.guidanceScale.toString(), inline: true}, - {name: 'Has Mask', value: typeof queueItem.prediction.mask === 'undefined' ? 'No' : 'Yes', inline: true}, - {name: 'Has Base Image', value: typeof queueItem.prediction.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Prompt', value: queueItem.prompt ? codeBlock(queueItem.prompt) : 'Not Supplied', inline: false}, + {name: 'Prompt Strength', value: queueItem.promptStrength.toString(), inline: true}, + {name: 'Steps', value: queueItem.numInferenceSteps.toString(), inline: true}, + {name: 'Guidance Scale', value: queueItem.guidanceScale.toString(), inline: true}, + {name: 'Has Mask', value: typeof queueItem.mask === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Has Base Image', value: typeof queueItem.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, {name: 'Seed', value: codeBlock(queueItem.seed.toString()), inline: false}, {name: 'Prompt ID', value: codeBlock(queueItem.uuid), inline: false} ]) ] break - case QueueItemType.Upscaled: + case QueueItems.QueueItemTypes.Upscaled: embeds = [ new EmbedBuilder() .setColor('#030354') @@ -78,19 +78,19 @@ export default function(queueItem: QueueItem): BotEmbed { .setThumbnail('https://i.gifer.com/2RNf.gif') .setTimestamp() .addFields([ - {name: 'Prompt', value: queueItem.prediction.prompt ? codeBlock(queueItem.prediction.prompt) : 'Not Supplied', inline: false}, - {name: 'Prompt Strength', value: queueItem.prediction.promptStrength.toString(), inline: true}, - {name: 'Steps', value: queueItem.prediction.numInferenceSteps.toString(), inline: true}, - {name: 'Guidance Scale', value: queueItem.prediction.guidanceScale.toString(), inline: true}, - {name: 'Has Mask', value: typeof queueItem.prediction.mask === 'undefined' ? 'No' : 'Yes', inline: true}, - {name: 'Has Base Image', value: typeof queueItem.prediction.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Prompt', value: queueItem.prompt ? codeBlock(queueItem.prompt) : 'Not Supplied', inline: false}, + {name: 'Prompt Strength', value: queueItem.promptStrength.toString(), inline: true}, + {name: 'Steps', value: queueItem.numInferenceSteps.toString(), inline: true}, + {name: 'Guidance Scale', value: queueItem.guidanceScale.toString(), inline: true}, + {name: 'Has Mask', value: typeof queueItem.mask === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Has Base Image', value: typeof queueItem.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, {name: 'Seed', value: codeBlock(queueItem.seed.toString()), inline: false}, {name: 'Prompt ID', value: codeBlock(queueItem.uuid), inline: false} ]) ] break - case QueueItemType.Default: + case QueueItems.QueueItemTypes.Default: default: embeds = [ new EmbedBuilder() @@ -100,12 +100,12 @@ export default function(queueItem: QueueItem): BotEmbed { .setThumbnail('https://i.gifer.com/2RNf.gif') .setTimestamp() .addFields([ - {name: 'Prompt', value: queueItem.prediction.prompt ? codeBlock(queueItem.prediction.prompt) : 'Not Supplied', inline: false}, - {name: 'Prompt Strength', value: queueItem.prediction.promptStrength.toString(), inline: true}, - {name: 'Steps', value: queueItem.prediction.numInferenceSteps.toString(), inline: true}, - {name: 'Guidance Scale', value: queueItem.prediction.guidanceScale.toString(), inline: true}, - {name: 'Has Mask', value: typeof queueItem.prediction.mask === 'undefined' ? 'No' : 'Yes', inline: true}, - {name: 'Has Base Image', value: typeof queueItem.prediction.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Prompt', value: queueItem.prompt ? codeBlock(queueItem.prompt) : 'Not Supplied', inline: false}, + {name: 'Prompt Strength', value: queueItem.promptStrength.toString(), inline: true}, + {name: 'Steps', value: queueItem.numInferenceSteps.toString(), inline: true}, + {name: 'Guidance Scale', value: queueItem.guidanceScale.toString(), inline: true}, + {name: 'Has Mask', value: typeof queueItem.mask === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Has Base Image', value: typeof queueItem.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, {name: 'Seed', value: codeBlock(queueItem.seed.toString()), inline: false}, {name: 'Prompt ID', value: codeBlock(queueItem.uuid), inline: false} ]) diff --git a/src/embeds/queueDetails.ts b/src/embeds/queueDetails.ts index 43f285c..115e7bc 100644 --- a/src/embeds/queueDetails.ts +++ b/src/embeds/queueDetails.ts @@ -1,7 +1,7 @@ import { EmbedBuilder, codeBlock } from 'discord.js' -import { BotEmbed, QueueItem } from '../types' +import { BotEmbed, QueueItems } from '../types' -export default function(queueItem: QueueItem): BotEmbed { +export default function(queueItem: QueueItems.QueueItemInstances): BotEmbed { return { embeds: [ new EmbedBuilder() @@ -10,14 +10,14 @@ export default function(queueItem: QueueItem): BotEmbed { .setThumbnail('https://i.imgur.com/wj1PRKP.png') .setTimestamp() .addFields([ - {name: 'Prompt', value: queueItem.prediction.prompt ? codeBlock(queueItem.prediction.prompt) : 'Not Supplied', inline: false}, - {name: 'Width', value: queueItem.prediction.width.toString(), inline: true}, - {name: 'Height', value: queueItem.prediction.height.toString(), inline: true}, - {name: 'Prompt Strength', value: queueItem.prediction.promptStrength.toString(), inline: true}, - {name: 'Steps', value: queueItem.prediction.numInferenceSteps.toString(), inline: true}, - {name: 'Guidance Scale', value: queueItem.prediction.guidanceScale.toString(), inline: true}, - {name: 'Has Mask', value: typeof queueItem.prediction.mask === 'undefined' ? 'No' : 'Yes', inline: true}, - {name: 'Has Base Image', value: typeof queueItem.prediction.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Prompt', value: queueItem.prompt ? codeBlock(queueItem.prompt) : 'Not Supplied', inline: false}, + {name: 'Width', value: queueItem.width.toString(), inline: true}, + {name: 'Height', value: queueItem.height.toString(), inline: true}, + {name: 'Prompt Strength', value: queueItem.promptStrength.toString(), inline: true}, + {name: 'Steps', value: queueItem.numInferenceSteps.toString(), inline: true}, + {name: 'Guidance Scale', value: queueItem.guidanceScale.toString(), inline: true}, + {name: 'Has Mask', value: typeof queueItem.mask === 'undefined' ? 'No' : 'Yes', inline: true}, + {name: 'Has Base Image', value: typeof queueItem.initImage === 'undefined' ? 'No' : 'Yes', inline: true}, {name: 'Seed', value: codeBlock(queueItem.seed.toString()), inline: false}, {name: 'Prompt ID', value: codeBlock(queueItem.uuid), inline: false} ]) diff --git a/src/embeds/queueShow.ts b/src/embeds/queueShow.ts index d3a6b0d..dafb75a 100644 --- a/src/embeds/queueShow.ts +++ b/src/embeds/queueShow.ts @@ -1,8 +1,8 @@ import { EmbedBuilder, codeBlock, ActionRowBuilder, userMention, ButtonBuilder, ButtonStyle } from 'discord.js' -import { BotEmbed, QueueItem } from '../types' +import { BotEmbed, QueueItems } from '../types' interface QueueShowArgs { - queueItems: QueueItem[] + queueItems: QueueItems.QueueItemInstances[] positions: number[] } @@ -15,8 +15,8 @@ export default function(args: QueueShowArgs): BotEmbed { const queueField = [ {name: position === 0 ? 'Next Up' : 'Queue Position', value: position === 0 ? 'This will be generated next' : `#${position}`, inline: false}, - {name: 'Prompt', value: queueItem.prediction.prompt ? codeBlock(queueItem.prediction.prompt) : 'Not Supplied', inline: true}, - {name: 'Owner', value: userMention(queueItem.discordCaller), inline: true}, + {name: 'Prompt', value: queueItem.prompt ? codeBlock(queueItem.prompt) : 'Not Supplied', inline: true}, + {name: 'Owner', value: userMention(queueItem.discordCallerSnowflake), inline: true}, {name: 'Prompt ID', value: codeBlock(queueItem.uuid), inline: true}, ] diff --git a/src/events/buttonImagineRegenerate.ts b/src/events/buttonImagineRegenerate.ts index 599ae9c..b0a75cf 100644 --- a/src/events/buttonImagineRegenerate.ts +++ b/src/events/buttonImagineRegenerate.ts @@ -1,10 +1,8 @@ import { Interaction } from 'discord.js' import { isNil } from 'ramda' import { Bot } from '../bot' -import { BotEvent, QueueItem, QueueItemType } from '../types' -import { getRandomInt } from '../utils' -import { v4 as uuidv4 } from 'uuid' -import addedToQueue, { QueueType } from '../embeds/addedToQueue' +import { addedToInstantQueue, addedToQueue } from '../embeds/addedToQueue' +import { BotEvent, QueueItems } from '../types' const botEvent: BotEvent = { name: 'Button Handler - Imagine Regenerate', @@ -14,44 +12,55 @@ const botEvent: BotEvent = { if (!interaction.isButton()) return if (interaction.customId !== 'button-imagine-result-regenerate') return - const referenceQueueItem = bot.findLatestQueueItemReferenceByMessageID(interaction.message.id) + const referenceQueueItem = bot.findLatestQueueItemReferenceByMessageSnowflake(interaction.message.id) + if (isNil(referenceQueueItem)) { await interaction.reply({ ephemeral: true, content: 'There is no reference of this image generation. Most likely, this generation is too old to modify.' }) - + return } - - const queueItem: QueueItem = { - ...referenceQueueItem, - type: QueueItemType.Regenerated, - discordCaller: interaction.user.id.toString(), - seed: getRandomInt(1, 99999999), - uuid: uuidv4() + + if (interaction.user.id !== referenceQueueItem.discordCallerSnowflake) { + await interaction.reply({ + ephemeral: true, + content: 'You can not modify content generated by other users.' + }) + + return } - - // Remove old attachment - await interaction.message.removeAttachments() - await interaction.reply({ - ephemeral: true, - content: 'Your images are being regenerated. The original message will be updated once complete.' + const queueItem = new QueueItems.RegeneratedQueueItem.RegeneratedQueueItem({ + discordCallerSnowflake: interaction.user.id.toString(), + discordInteraction: referenceQueueItem.discordInteraction, + discordMessageSnowflake: referenceQueueItem.discordMessageSnowflake, + prompt: referenceQueueItem.prompt, + height: referenceQueueItem.height, + width: referenceQueueItem.width, + initImage: referenceQueueItem.initImage, + mask: referenceQueueItem.mask, + promptStrength: referenceQueueItem.promptStrength, + guidanceScale: referenceQueueItem.guidanceScale, + numInferenceSteps: referenceQueueItem.numInferenceSteps }) + // Remove old attachment + await interaction.message.removeAttachments() + if (bot.stableDiffusion.isProcessing() || bot.hasQueue()) { - const queuePos = bot.addQueue(queueItem) + const queuePos = bot.addQueuedQueueItem(queueItem) - await referenceQueueItem.interaction.editReply({ - embeds: addedToQueue(QueueType.Queued, queueItem, queuePos).embeds, + await referenceQueueItem.discordInteraction.editReply({ + embeds: addedToQueue(queueItem, queuePos).embeds, components: [] }) } else { - bot.addQueue(queueItem) + bot.addQueuedQueueItem(queueItem) - await referenceQueueItem.interaction.editReply({ - embeds: addedToQueue(QueueType.Instant, queueItem).embeds, + await referenceQueueItem.discordInteraction.editReply({ + embeds: addedToInstantQueue(queueItem).embeds, components: [] }) } diff --git a/src/events/buttonImagineUpscale.ts b/src/events/buttonImagineUpscale.ts index 2b9a2cb..128b1e1 100644 --- a/src/events/buttonImagineUpscale.ts +++ b/src/events/buttonImagineUpscale.ts @@ -1,10 +1,10 @@ import { AttachmentBuilder, Interaction } from "discord.js" import { isEmpty, isNil } from "ramda" import { Bot } from '../bot' -import { BotEvent, QueueItemType } from '../types' - import imageResult from "../embeds/imageResult" import imageSelectPrompt from "../embeds/imageSelectPrompt" +import { BotEvent, QueueItems } from '../types' + const botEvent: BotEvent = { name: 'Button Handler - Imagine Upscale', @@ -12,48 +12,78 @@ const botEvent: BotEvent = { once: false, async execute(bot: Bot, interaction: Interaction) { if (!interaction.isButton() && !interaction.isSelectMenu()) return - if (interaction.customId === 'image-select-prompt-upscaled' && interaction.isSelectMenu()) { - // TODO: Lock this action to user who triggered prompt? + if (interaction.customId !== 'button-imagine-result-upscale' && !interaction.customId.startsWith('image-select-prompt-upscaled')) return + + const referenceQueueItem = bot.findLatestQueueItemReferenceByMessageSnowflake(interaction.message.id) + + if (isNil(referenceQueueItem) || isNil(referenceQueueItem.imageData) || isEmpty(referenceQueueItem.imageData)) { + await interaction.reply({ + ephemeral: true, + content: 'There is no reference of this image generation. Most likely, this generation is too old to modify.' + }) + return + } + + // Handle image select prompt result + if (interaction.customId === `image-select-prompt-upscaled-${referenceQueueItem.uuid}` && interaction.isSelectMenu()) { - const referenceQueueItem = bot.findLatestQueueItemReferenceByMessageID(interaction.message.id) - if (isNil(referenceQueueItem) || isNil(referenceQueueItem.imageData) || isEmpty(referenceQueueItem.imageData)) { + if (interaction.user.id !== referenceQueueItem.discordCallerSnowflake) { await interaction.reply({ ephemeral: true, - content: 'There is no reference of this image generation. Most likely, this generation is too old to modify.' + content: 'You can not modify content generated by other users.' }) + return } // We can cheat here, we don't really need to generate anything. We just take the image out of the collage. // Maybe we add ai upscaling here in the future? const [selectedImage] = interaction.values - const imageResultEmbed = imageResult({...referenceQueueItem, type: QueueItemType.Upscaled}) + const imageResultEmbed = imageResult(referenceQueueItem) const file = new AttachmentBuilder(referenceQueueItem.imageData[+selectedImage], { name: `stable-confusion_${referenceQueueItem.uuid}_upscaled.jpeg`, - description: referenceQueueItem.prediction.prompt + description: referenceQueueItem.prompt }) - await referenceQueueItem.interaction.editReply({ + await referenceQueueItem.discordInteraction.editReply({ embeds: imageResultEmbed.embeds, components: imageResultEmbed.components, files: [file] }) - } else if (interaction.customId === 'button-imagine-result-upscale' && interaction.isButton()) { - const referenceQueueItem = bot.findLatestQueueItemReferenceByMessageID(interaction.message.id) - if (isNil(referenceQueueItem)) { - await interaction.reply({ - ephemeral: true, - content: 'There is no reference of this image generation. Most likely, this generation is too old to modify.' - }) - return - } + // Handle button invoke + } else if (interaction.customId === 'button-imagine-result-upscale' && interaction.isButton()) { await interaction.deferUpdate() // We defer the update here because we'll update the queueItem interaction instead. - const imageSelectPromptEmbed = imageSelectPrompt({...referenceQueueItem, type: QueueItemType.Upscaled}) + // Create an UpscaledQueueItem + const upscaledQueueItem = new QueueItems.UpscaledQueueItem.UpscaledQueueItem({ + discordCallerSnowflake: interaction.user.id.toString(), + discordInteraction: interaction, + discordMessageSnowflake: interaction.message.id, + height: referenceQueueItem.height, + width: referenceQueueItem.width, + initImage: referenceQueueItem.initImage, + mask: referenceQueueItem.mask, + guidanceScale: referenceQueueItem.guidanceScale, + numInferenceSteps: referenceQueueItem.numInferenceSteps, + prompt: referenceQueueItem.prompt, + promptStrength: referenceQueueItem.promptStrength, + seed: referenceQueueItem.seed + }) + + upscaledQueueItem.imageData = referenceQueueItem.imageData + + // Convert existing queueitem to UpscaledQueueItem + bot.updateQueueItem(queueItem => { + upscaledQueueItem.inheritUUIDFromQueueItem(queueItem, bot) + upscaledQueueItem.discordMessageSnowflake = queueItem.discordMessageSnowflake + return upscaledQueueItem + }, referenceQueueItem.uuid) - await referenceQueueItem.interaction.editReply({ + // Send image select prompt + const imageSelectPromptEmbed = imageSelectPrompt(upscaledQueueItem) + await referenceQueueItem.discordInteraction.editReply({ embeds: imageSelectPromptEmbed.embeds, components: imageSelectPromptEmbed.components }) diff --git a/src/events/buttonImagineVariantOf.ts b/src/events/buttonImagineVariantOf.ts index 4bd6c91..4a3a379 100644 --- a/src/events/buttonImagineVariantOf.ts +++ b/src/events/buttonImagineVariantOf.ts @@ -1,10 +1,9 @@ import { Interaction } from "discord.js" import { isEmpty, isNil } from "ramda" import { Bot } from '../bot' -import { BotEvent, QueueItem, QueueItemType } from '../types' +import { addedToInstantQueue, addedToQueue } from "../embeds/addedToQueue" import imageSelectPrompt from '../embeds/imageSelectPrompt' -import { v4 as uuidv4 } from 'uuid' -import addedToQueue, { QueueType } from "../embeds/addedToQueue" +import { BotEvent, QueueItems } from '../types' const botEvent: BotEvent = { name: 'Button Handler - Imagine Variant Of', @@ -12,10 +11,10 @@ const botEvent: BotEvent = { once: false, async execute(bot: Bot, interaction: Interaction) { if (!interaction.isButton() && !interaction.isSelectMenu()) return - if (interaction.customId === 'image-select-prompt-variant' && interaction.isSelectMenu()) { - // TODO: Lock this action to user who triggered prompt? - const referenceQueueItem = bot.findLatestQueueItemReferenceByMessageID(interaction.message.id) + // Handle selection + if (interaction.customId === 'image-select-prompt-variant' && interaction.isSelectMenu()) { + const referenceQueueItem = bot.findLatestQueueItemReferenceByMessageSnowflake(interaction.message.id) if (isNil(referenceQueueItem) || isNil(referenceQueueItem.imageData) || isEmpty(referenceQueueItem.imageData)) { await interaction.reply({ ephemeral: true, @@ -23,52 +22,54 @@ const botEvent: BotEvent = { }) return } + + if (interaction.user.id !== referenceQueueItem.discordCallerSnowflake) { + await interaction.reply({ + ephemeral: true, + content: 'You can not modify content generated by other users.' + }) + + return + } const [selectedImage] = interaction.values // Ensure PNG const selectedImageB64 = `data:image/png;base64,${referenceQueueItem.imageData[+selectedImage].toString('base64')}` - const queueItem: QueueItem = { - ...referenceQueueItem, - type: QueueItemType.Variant, - prediction: { - prompt: referenceQueueItem.prediction.prompt, - promptStrength: 0.8, - initImage: selectedImageB64, - numOutputs: 4, - height: referenceQueueItem.prediction.height, - width: referenceQueueItem.prediction.width, - guidanceScale: referenceQueueItem.prediction.guidanceScale, - mask: referenceQueueItem.prediction.mask, - numInferenceSteps: referenceQueueItem.prediction.numInferenceSteps - }, - discordCaller: interaction.user.id.toString(), - uuid: uuidv4() - } + + const queueItem = new QueueItems.VariantQueueItem.VariantQueueItem({ + discordCallerSnowflake: interaction.user.id.toString(), + discordInteraction: interaction, + discordMessageSnowflake: interaction.message.id, + height: referenceQueueItem.height, + width: referenceQueueItem.width, + initImage: selectedImageB64, + mask: referenceQueueItem.mask, + guidanceScale: referenceQueueItem.guidanceScale, + numInferenceSteps: referenceQueueItem.numInferenceSteps, + prompt: referenceQueueItem.prompt, + }) // Remove old attachment await interaction.message.removeAttachments() - - await interaction.reply({ - ephemeral: true, - content: 'Your image variants are being generated. The original message will be updated once complete.' - }) - + if (bot.stableDiffusion.isProcessing() || bot.hasQueue()) { - const queuePos = bot.addQueue(queueItem) - await referenceQueueItem.interaction.editReply({ - embeds: addedToQueue(QueueType.Queued, queueItem, queuePos).embeds, + const queuePos = bot.addQueuedQueueItem(queueItem) + await queueItem.discordInteraction.editReply({ + embeds: addedToQueue(queueItem, queuePos).embeds, components: [] }) } else { - bot.addQueue(queueItem) - await referenceQueueItem.interaction.editReply({ - embeds: addedToQueue(QueueType.Instant, queueItem).embeds, + bot.addQueuedQueueItem(queueItem) + await queueItem.discordInteraction.editReply({ + embeds: addedToInstantQueue(queueItem).embeds, components: [] }) } + + // Handle button invoke } else if (interaction.customId === 'button-imagine-result-variant-of' && interaction.isButton()) { - const referenceQueueItem = bot.findLatestQueueItemReferenceByMessageID(interaction.message.id) + const referenceQueueItem = bot.findLatestQueueItemReferenceByMessageSnowflake(interaction.message.id) if (isNil(referenceQueueItem)) { await interaction.reply({ ephemeral: true, @@ -79,9 +80,9 @@ const botEvent: BotEvent = { await interaction.deferUpdate() // We defer the update here because we'll update the queueItem interaction instead. - const imageSelectPromptEmbed = imageSelectPrompt({...referenceQueueItem, type: QueueItemType.Variant}) + const imageSelectPromptEmbed = imageSelectPrompt(referenceQueueItem) - await referenceQueueItem.interaction.editReply({ + await referenceQueueItem.discordInteraction.editReply({ embeds: imageSelectPromptEmbed.embeds, components: imageSelectPromptEmbed.components }) diff --git a/src/events/commandExtend.ts b/src/events/commandExtend.ts index f00ef4e..27f6ff9 100644 --- a/src/events/commandExtend.ts +++ b/src/events/commandExtend.ts @@ -1,9 +1,8 @@ import { Interaction } from 'discord.js' import { Bot } from '../bot' -import { BotEvent, QueueItem, QueueItemType } from '../types' -import addedToQueue, { QueueType } from '../embeds/addedToQueue' -import { v4 as uuidv4 } from 'uuid' -import { getImageAttachmentURL, getRandomInt } from '../utils' +import { addedToInstantQueue, addedToQueue } from '../embeds/addedToQueue' +import { BotEvent, QueueItems } from '../types' +import { getImageAttachmentURL } from '../utils' const botEvent: BotEvent = { name: 'Command Handler - Extend', @@ -22,36 +21,31 @@ const botEvent: BotEvent = { const guidanceScale = interaction.options.getNumber('guidancescale') const seed = interaction.options.getInteger('seed') - const queueItem: QueueItem = { - discordCaller: interaction.user.id.toString(), - seed: seed ?? getRandomInt(1, 99999999), - uuid: uuidv4(), - interaction, - type: QueueItemType.Extended, - prediction: { - prompt: undefined, - width: 512, - height: 512, - initImage, - mask, - promptStrength: promptStrength ?? 0.8, - numOutputs: 4, - numInferenceSteps: numInferenceSteps ?? 50, - guidanceScale: guidanceScale ?? 7.5 - } - } + const queueItem = new QueueItems.ExtendedQueueItem.ExtendedQueueItem({ + discordCallerSnowflake: interaction.user.id.toString(), + discordInteraction: interaction, + seed, + prompt: undefined, + width: 512, + height: 512, + initImage, + mask, + promptStrength, + guidanceScale, + numInferenceSteps + }) if (bot.stableDiffusion.isProcessing() || bot.hasQueue()) { - const queuePos = bot.addQueue(queueItem) + const queuePos = bot.addQueuedQueueItem(queueItem) await interaction.editReply({ - embeds: addedToQueue(QueueType.Queued, queueItem, queuePos).embeds + embeds: addedToQueue(queueItem, queuePos).embeds }) } else { - bot.addQueue(queueItem) + bot.addQueuedQueueItem(queueItem) await interaction.editReply({ - embeds: addedToQueue(QueueType.Instant, queueItem).embeds + embeds: addedToInstantQueue(queueItem).embeds }) } } diff --git a/src/events/commandImagine.ts b/src/events/commandImagine.ts index abc0531..2e47b1d 100644 --- a/src/events/commandImagine.ts +++ b/src/events/commandImagine.ts @@ -1,9 +1,8 @@ import { Interaction } from 'discord.js' import { Bot } from '../bot' -import { BotEvent, QueueItem, QueueItemType } from '../types' -import addedToQueue, { QueueType } from '../embeds/addedToQueue' -import { v4 as uuidv4 } from 'uuid' -import { getImageAttachmentURL, getRandomInt, validateHeight, validateWidth } from '../utils' +import { addedToInstantQueue, addedToQueue } from '../embeds/addedToQueue' +import { BotEvent, QueueItems } from '../types' +import { getImageAttachmentURL } from '../utils' const botEvent: BotEvent = { name: 'Command Handler - Imagine', @@ -23,37 +22,42 @@ const botEvent: BotEvent = { const guidanceScale = interaction.options.getNumber('guidancescale') const seed = interaction.options.getInteger('seed') - const queueItem: QueueItem = { - discordCaller: interaction.user.id.toString(), - seed: seed ?? getRandomInt(1, 99999999), - uuid: uuidv4(), - interaction, - type: QueueItemType.Default, - prediction: { + const queueItem = new QueueItems.QueueItem.QueueItem({ + discordCallerSnowflake: interaction.user.id.toString(), + discordInteraction: interaction, + seed, prompt, width: 512, height: 512, initImage, mask, - promptStrength: promptStrength ?? 0.8, - numOutputs: 4, - numInferenceSteps: numInferenceSteps ?? 50, - guidanceScale: guidanceScale ?? 7.5 - } - } + promptStrength, + guidanceScale, + numInferenceSteps + }) if (bot.stableDiffusion.isProcessing() || bot.hasQueue()) { - const queuePos = bot.addQueue(queueItem) + const queuePos = bot.addQueuedQueueItem(queueItem) - await interaction.editReply({ - embeds: addedToQueue(QueueType.Queued, queueItem, queuePos).embeds + const message = await interaction.editReply({ + embeds: addedToQueue(queueItem, queuePos).embeds }) + + bot.updateQueueItem(queueItem => { + queueItem.discordMessageSnowflake = message.id + return queueItem + }, queueItem.uuid) } else { - bot.addQueue(queueItem) + bot.addQueuedQueueItem(queueItem) - await interaction.editReply({ - embeds: addedToQueue(QueueType.Instant, queueItem).embeds + const message = await interaction.editReply({ + embeds: addedToInstantQueue(queueItem).embeds }) + + bot.updateQueueItem(queueItem => { + queueItem.discordMessageSnowflake = message.id + return queueItem + }, queueItem.uuid) } } } diff --git a/src/events/commandImagineQuick.ts b/src/events/commandImagineQuick.ts index 9b5e55d..4f7b988 100644 --- a/src/events/commandImagineQuick.ts +++ b/src/events/commandImagineQuick.ts @@ -1,9 +1,8 @@ import { Interaction } from 'discord.js' import { Bot } from '../bot' -import { BotEvent, QueueItem, QueueItemType } from '../types' -import addedToQueue, { QueueType } from '../embeds/addedToQueue' -import { v4 as uuidv4 } from 'uuid' -import { getImageAttachmentURL, getRandomInt, validateHeight, validateWidth } from '../utils' +import { addedToInstantQueue, addedToQueue } from '../embeds/addedToQueue' +import { BotEvent, QueueItems } from '../types' +import { getImageAttachmentURL, validateHeight, validateWidth } from '../utils' const botEvent: BotEvent = { name: 'Command Handler - Quick Imagine', @@ -21,41 +20,35 @@ const botEvent: BotEvent = { const initImage = getImageAttachmentURL(interaction.options.getAttachment('image')) const mask = getImageAttachmentURL(interaction.options.getAttachment('mask')) const promptStrength = interaction.options.getNumber('pstrength') - const numOutputs = interaction.options.getInteger('numout') const numInferenceSteps = interaction.options.getInteger('numsteps') const guidanceScale = interaction.options.getNumber('guidancescale') const seed = interaction.options.getInteger('seed') - const queueItem: QueueItem = { - discordCaller: interaction.user.id.toString(), - seed: seed ?? getRandomInt(1, 99999999), - uuid: uuidv4(), - interaction, - type: QueueItemType.Quick, - prediction: { + const queueItem = new QueueItems.QuickQueueItem.QuickQueueItem({ + discordCallerSnowflake: interaction.user.id.toString(), + discordInteraction: interaction, + seed, prompt, width, height, initImage, mask, - promptStrength: promptStrength ?? 0.8, - numOutputs: numOutputs ?? 1, - numInferenceSteps: numInferenceSteps ?? 50, - guidanceScale: guidanceScale ?? 7.5 - } - } + promptStrength, + guidanceScale, + numInferenceSteps + }) if (bot.stableDiffusion.isProcessing() || bot.hasQueue()) { - const queuePos = bot.addQueue(queueItem) + const queuePos = bot.addQueuedQueueItem(queueItem) await interaction.editReply({ - embeds: addedToQueue(QueueType.Queued, queueItem, queuePos).embeds + embeds: addedToQueue(queueItem, queuePos).embeds }) } else { - bot.addQueue(queueItem) + bot.addQueuedQueueItem(queueItem) await interaction.editReply({ - embeds: addedToQueue(QueueType.Instant, queueItem).embeds + embeds: addedToInstantQueue(queueItem).embeds }) } } diff --git a/src/events/commandQueueRemove.ts b/src/events/commandQueueRemove.ts index b925514..716551b 100644 --- a/src/events/commandQueueRemove.ts +++ b/src/events/commandQueueRemove.ts @@ -21,12 +21,12 @@ const botEvent: BotEvent = { return } - if (interaction.user.id !== queueItem.discordCaller && !isAdminRequest) { + if (interaction.user.id !== queueItem.discordCallerSnowflake && !isAdminRequest) { await interaction.reply('You do not own this queue item.') return } - const deletedItems = bot.removeQueue(uuid) + const deletedItems = bot.removeQueuedQueueItem(uuid) const queueDetailsEmbed = queueDetails(deletedItems[0]) await interaction.reply({ diff --git a/src/events/commandQueueShow.ts b/src/events/commandQueueShow.ts index 63c53f2..d0b8f4a 100644 --- a/src/events/commandQueueShow.ts +++ b/src/events/commandQueueShow.ts @@ -18,7 +18,7 @@ const botEvent: BotEvent = { let queueItems = bot.queue const userFilter = interaction.options.getUser('user') if (!isNil(userFilter)) { - queueItems = bot.queue.filter(queueItem => queueItem.discordCaller === userFilter.id) + queueItems = bot.queue.filter(queueItem => queueItem.discordCallerSnowflake === userFilter.id) if (isEmpty(queueItems)) { await interaction.reply('This user has no queue items.') diff --git a/src/events/queueProcessor.ts b/src/events/queueProcessor.ts index a30f7ec..6771f6f 100644 --- a/src/events/queueProcessor.ts +++ b/src/events/queueProcessor.ts @@ -1,10 +1,10 @@ +import { AttachmentBuilder } from 'discord.js' import { isEmpty, isNil } from 'ramda' +import sharp from 'sharp' import { Bot } from '../bot' -import processingPrompt from '../embeds/processingPrompt' -import { BotEvent, QueueItemType } from '../types' -import { AttachmentBuilder } from 'discord.js' import imageResult from '../embeds/imageResult' -import sharp from 'sharp' +import processingPrompt from '../embeds/processingPrompt' +import { BotEvent, QueueItems } from '../types' const botEvent: BotEvent = { name: 'Queue Processor', @@ -26,27 +26,27 @@ async function tick(bot: Bot) { try { const processingPromptEmbed = processingPrompt(queueItem) - const message = await queueItem.interaction.editReply({ + const message = await queueItem.discordInteraction.editReply({ embeds: processingPromptEmbed.embeds }) - queueItem.messageId = message.id + queueItem.discordMessageSnowflake = message.id const processResults = await bot.stableDiffusion.processRequest(queueItem) if (processResults === false) { - const message = await queueItem.interaction.editReply({ + const message = await queueItem.discordInteraction.editReply({ content: 'Processing for this prompt has failed.', embeds: undefined }) await message.suppressEmbeds() bot.log.debug('Processing failed') - bot.removeQueue(queueItem.uuid) + bot.removeQueuedQueueItem(queueItem.uuid) tickLock = false return } switch(queueItem.type) { - case QueueItemType.Quick: { + case QueueItems.QueueItemTypes.Quick: { const imageResultEmbed = imageResult(queueItem) const data: string = processResults[0].split(',')[1] const buf = Buffer.from(data, 'base64') @@ -56,7 +56,7 @@ async function tick(bot: Bot) { queueItem.imageData = [buf] - await queueItem.interaction.editReply({ + await queueItem.discordInteraction.editReply({ embeds: imageResultEmbed.embeds, components: imageResultEmbed.components, files: [file] @@ -64,7 +64,7 @@ async function tick(bot: Bot) { break } - case QueueItemType.Default: + case QueueItems.QueueItemTypes.Default: default: { const imageResultEmbed = imageResult(queueItem) const imageData: string[] = processResults.map(base64 => base64.split(',')[1]) @@ -109,7 +109,7 @@ async function tick(bot: Bot) { queueItem.imageData = imageBuffers - await queueItem.interaction.editReply({ + await queueItem.discordInteraction.editReply({ embeds: imageResultEmbed.embeds, components: imageResultEmbed.components, files: [file] @@ -118,15 +118,15 @@ async function tick(bot: Bot) { } bot.log.debug('Processing complete') - bot.removeQueue(queueItem.uuid) + bot.removeQueuedQueueItem(queueItem.uuid) tickLock = false // @ts-ignore no-implicit-any } catch(e: any) { bot.log.error(e.toString()) bot.log.error(e.stack) tickLock = false - await queueItem.interaction.deleteReply().catch() - await queueItem.interaction.followUp('An error occured with your request.').catch() + await queueItem.discordInteraction.deleteReply().catch() + await queueItem.discordInteraction.followUp('An error occured with your request.').catch() } } diff --git a/src/modules/stableDiffusion.ts b/src/modules/stableDiffusion.ts index 870c5ac..d9ad3d2 100644 --- a/src/modules/stableDiffusion.ts +++ b/src/modules/stableDiffusion.ts @@ -1,6 +1,6 @@ -import { QueueItem, QueueItemType } from "../types" import axios from 'axios' import { isNil } from "ramda" +import { QueueItems } from '../types' export default class StableDiffusion { constructor(host: string, port: number) { @@ -15,7 +15,7 @@ export default class StableDiffusion { public isProcessing = () => this.processing - public async processRequest(queueItem: QueueItem): Promise { + public async processRequest(queueItem: QueueItems.QueueItemInstances): Promise { if (this.processing) return false this.abortController = new AbortController() this.processing = true @@ -23,28 +23,28 @@ export default class StableDiffusion { let body switch (queueItem.type) { - case QueueItemType.Extended: { + case QueueItems.QueueItemTypes.Extended: { body = { input: { - width: queueItem.prediction.width, - height: queueItem.prediction.height, - "init_image": queueItem.prediction.initImage, - "prompt_strength": queueItem.prediction.promptStrength, - "num_outputs": queueItem.prediction.numOutputs + width: queueItem.width, + height: queueItem.height, + "init_image": queueItem.initImage, + "prompt_strength": queueItem.promptStrength, + "num_outputs": queueItem.numOutputs } } break } - case QueueItemType.Variant: { + case QueueItems.QueueItemTypes.Variant: { body = { input: { - prompt: queueItem.prediction.prompt, - width: queueItem.prediction.width, - height: queueItem.prediction.height, - "init_image": queueItem.prediction.initImage, - "prompt_strength": queueItem.prediction.promptStrength, - "num_outputs": queueItem.prediction.numOutputs + prompt: queueItem.prompt, + width: queueItem.width, + height: queueItem.height, + "init_image": queueItem.initImage, + "prompt_strength": queueItem.promptStrength, + "num_outputs": queueItem.numOutputs } } break @@ -53,15 +53,15 @@ export default class StableDiffusion { default: { body = { input: { - prompt: queueItem.prediction.prompt, - width: queueItem.prediction.width, - height: queueItem.prediction.height, - "init_image": queueItem.prediction.initImage, - mask: queueItem.prediction.mask, - "prompt_strength": queueItem.prediction.promptStrength, - "num_outputs": queueItem.prediction.numOutputs, - "num_inference_steps": queueItem.prediction.numInferenceSteps, - "guidance_scale": queueItem.prediction.guidanceScale, + prompt: queueItem.prompt, + width: queueItem.width, + height: queueItem.height, + "init_image": queueItem.initImage, + mask: queueItem.mask, + "prompt_strength": queueItem.promptStrength, + "num_outputs": queueItem.numOutputs, + "num_inference_steps": queueItem.numInferenceSteps, + "guidance_scale": queueItem.guidanceScale, seed: queueItem.seed } } diff --git a/src/queueItems/ExtendedQueueItem.ts b/src/queueItems/ExtendedQueueItem.ts new file mode 100644 index 0000000..e28d33f --- /dev/null +++ b/src/queueItems/ExtendedQueueItem.ts @@ -0,0 +1,14 @@ +import { QueueItemTypes } from '.' +import { QueueItem, QueueItemConstructorInput } from './QueueItem' + +export class ExtendedQueueItem extends QueueItem { + constructor( + queueItemConstructorInput: QueueItemConstructorInput + ) { + super(queueItemConstructorInput) + } + + public get type(): Readonly { + return QueueItemTypes.Extended + } +} \ No newline at end of file diff --git a/src/queueItems/QueueItem.ts b/src/queueItems/QueueItem.ts new file mode 100644 index 0000000..c1cc1be --- /dev/null +++ b/src/queueItems/QueueItem.ts @@ -0,0 +1,224 @@ +import { ButtonInteraction, ChatInputCommandInteraction, SelectMenuInteraction, Snowflake } from "discord.js" +import { isEmpty, isNil } from 'ramda' +import { v4 as uuidv4, validate as validateUuid } from 'uuid' +import { QueueItemTypes } from '.' +import { Bot } from '../bot' +import { QueueItems } from '../types' +import { getRandomInt } from '../utils' + +export type QueueItemBody = { + input: { + prompt?: string // Prompt text + width: number // Image size, min 64 + height: number // Image size, min 64 + "init_image": string | undefined // Starter image to generate from + mask: string | undefined // Mask image to generate with + "prompt_strength": number // 0 - 1 float value, default 0.8 + "num_outputs": number // Number of images to generate, default 1 + "num_inference_steps": number // Number of steps, default 50 + "guidance_scale": number // Default 7.5 + seed: number // RNG numeric seed + } +} + +export type QueueItemConstructorInput = { + discordCallerSnowflake: Snowflake // Discord snowflake of invoking user + discordMessageSnowflake?: Snowflake // Discord snowflake of Stable Confusion's related message embed + discordInteraction: ChatInputCommandInteraction | ButtonInteraction | SelectMenuInteraction // Discord interaction handle for QueueItem + + prompt?: string // Prompt text + width: number // Image size, min 64 + height: number // Image size, min 64 + initImage: string | undefined // Starter image to generate from + mask: string | undefined // Mask image to generate with + promptStrength?: number | null // 0 - 1 float value, default 0.8 + numInferenceSteps?: number | null // Number of steps, default 50 + guidanceScale?: number | null // Default 7.5 + seed?: number | null // RNG numeric seed +} + +export type QueueItemMutationInput = Partial + +export class QueueItem { + constructor(queueItemConstructorInput: QueueItemConstructorInput) { + if (queueItemConstructorInput.height && (queueItemConstructorInput.height > 2048 || queueItemConstructorInput.height < 64)) { + throw new Error('Invalid QueueItem height provided. Must be >= 64 || <= 2048.') + } + + if (queueItemConstructorInput.width && (queueItemConstructorInput.width > 2048 || queueItemConstructorInput.height < 64)) { + throw new Error('Invalid QueueItem width provided. Must be >= 64 || <= 2048.') + } + + if (queueItemConstructorInput.promptStrength && (queueItemConstructorInput.promptStrength > 1 || queueItemConstructorInput.height < 0)) { + throw new Error('Invalid QueueItem promptStrength provided. Must be >= 0 || <= 1.') + } + + this._uuid = uuidv4() // Request identifier + this._seed = queueItemConstructorInput.seed ?? getRandomInt(1, 99999999) // Generation seed + this._discordCallerSnowflake = queueItemConstructorInput.discordCallerSnowflake + this._discordMessageSnowflake = queueItemConstructorInput.discordMessageSnowflake + this._discordInteraction = queueItemConstructorInput.discordInteraction + + this._prompt = queueItemConstructorInput.prompt ?? undefined + this._width = queueItemConstructorInput.width ?? 64 + this._height = queueItemConstructorInput.height ?? 64 + this._initImage = queueItemConstructorInput.initImage ?? undefined + this._mask = queueItemConstructorInput.mask ?? undefined + this._promptStrength = queueItemConstructorInput.promptStrength ?? 0.8 + this._numOutputs = 4 // This should always be 4 here due to variant / quad preview support. + this._numInferenceSteps = queueItemConstructorInput.numInferenceSteps ?? 50 + this._guidanceScale = queueItemConstructorInput.guidanceScale ?? 7.5 + } + + /* == Metadata == */ + + private _uuid: string + public get uuid(): Readonly { + return this._uuid + } + + public get type(): Readonly { + return QueueItemTypes.Default + } + + private _imageData?: Buffer[] + public get imageData(): Buffer[] | undefined { + return this._imageData + } + public set imageData(imageData: Buffer[] | undefined) { + if (isNil(imageData) || isEmpty(imageData)) { + throw new Error('A valid populated buffer array must be provided.') + } + + this._imageData = imageData + } + + private _discordCallerSnowflake: string + public get discordCallerSnowflake(): Readonly { + return this._discordCallerSnowflake + } + + private _discordMessageSnowflake?: string + public get discordMessageSnowflake(): Snowflake | undefined { + return this._discordMessageSnowflake + } + public set discordMessageSnowflake(snowflake: Snowflake | undefined) { + if (isNil(snowflake)) { + throw new Error('A snowflake must be provided.') + } + + this._discordMessageSnowflake = snowflake + } + + private _discordInteraction: ChatInputCommandInteraction | ButtonInteraction | SelectMenuInteraction + public get discordInteraction(): ChatInputCommandInteraction | ButtonInteraction | SelectMenuInteraction { + return this._discordInteraction + } + + /* == Prediction == */ + + private _seed: number + public get seed(): Readonly { + return this._seed + } + + // Prompt text + private _prompt?: string + public get prompt(): Readonly { + return this._prompt + } + + // Image size, min 64 + private _width: number + public get width(): Readonly { + return this._width + } + + // Image size, min 64 + private _height: number + public get height(): Readonly { + return this._height + } + + // Starter image to generate from + private _initImage?: string + public get initImage(): Readonly { + return this._initImage + } + + // Mask image to generate with + private _mask?: string + public get mask(): Readonly { + return this._mask + } + + // 0 - 1 float value, default 0.8 + private _promptStrength: number + public get promptStrength(): Readonly { + return this._promptStrength + } + + // Number of images to generate + private _numOutputs: number + public get numOutputs(): Readonly { + return this._numOutputs + } + + // Number of steps, default 50 + private _numInferenceSteps: number + public get numInferenceSteps(): Readonly { + return this._numInferenceSteps + } + + // Default 7.5 + private _guidanceScale: number + public get guidanceScale(): Readonly { + return this._guidanceScale + } + + /** + * Creates a new random seed and assigns to QueueItem + */ + public shuffleSeed() { + this._seed = getRandomInt(1, 99999999) + } + + /** + * Inherits a UUID from a passed QueueItem + * @param queueItem The QueueItem to inherit from. + * @returns UUID of QueueItem + */ + public inheritUUIDFromQueueItem(queueItem: QueueItems.QueueItemInstances, bot: Bot) { + if (isNil(queueItem.uuid) || isEmpty(queueItem.uuid)) { + throw new Error('UUID must exist.') + } + + if (!validateUuid(queueItem.uuid)) { + throw new Error('Invalid UUID.') + } + + this._uuid = queueItem.uuid + } + + /** + * Ceates the Stable Diffusion request body that the Stable Diffusion module sends to the model + * @returns A QueueItem body formatted for Stable Diffusion input + */ + public createStableDiffusionRequestBody(): QueueItemBody { + + return { + input: { + prompt: this._prompt, + width: this._width, + height: this._height, + "init_image": this._initImage, + mask: this._mask, + "prompt_strength": this._promptStrength, + "num_outputs": this._numOutputs, + "num_inference_steps": this._numInferenceSteps, + "guidance_scale": this._guidanceScale, + seed: this._seed + } + } + } +} diff --git a/src/queueItems/QuickQueueItem.ts b/src/queueItems/QuickQueueItem.ts new file mode 100644 index 0000000..581c8ef --- /dev/null +++ b/src/queueItems/QuickQueueItem.ts @@ -0,0 +1,19 @@ +import { QueueItemTypes } from '.' +import { QueueItem, QueueItemConstructorInput } from './QueueItem' + +export class QuickQueueItem extends QueueItem { + constructor( + queueItemConstructorInput: QueueItemConstructorInput + ) { + super(queueItemConstructorInput) + } + + public get type(): Readonly { + return QueueItemTypes.Quick + } + + // Number of images to generate + public get numOutputs(): Readonly { + return 1 + } +} \ No newline at end of file diff --git a/src/queueItems/RegeneratedQueueItem.ts b/src/queueItems/RegeneratedQueueItem.ts new file mode 100644 index 0000000..439c7f1 --- /dev/null +++ b/src/queueItems/RegeneratedQueueItem.ts @@ -0,0 +1,15 @@ +import { QueueItemTypes } from '.' +import { QueueItem, QueueItemConstructorInput } from './QueueItem' + +export class RegeneratedQueueItem extends QueueItem { + constructor( + queueItemConstructorInput: QueueItemConstructorInput + ) { + super(queueItemConstructorInput) + this.shuffleSeed() + } + + public get type(): Readonly { + return QueueItemTypes.Regenerated + } +} \ No newline at end of file diff --git a/src/queueItems/UpscaledQueueItem.ts b/src/queueItems/UpscaledQueueItem.ts new file mode 100644 index 0000000..889b3b7 --- /dev/null +++ b/src/queueItems/UpscaledQueueItem.ts @@ -0,0 +1,19 @@ +import { QueueItemTypes } from '.' +import { QueueItem, QueueItemConstructorInput } from './QueueItem' + +export class UpscaledQueueItem extends QueueItem { + constructor( + queueItemConstructorInput: QueueItemConstructorInput + ) { + super(queueItemConstructorInput) + } + + public get type(): Readonly { + return QueueItemTypes.Upscaled + } + + // Number of images to generate + public get numOutputs(): Readonly { + return 1 + } +} \ No newline at end of file diff --git a/src/queueItems/VariantQueueItem.ts b/src/queueItems/VariantQueueItem.ts new file mode 100644 index 0000000..075b499 --- /dev/null +++ b/src/queueItems/VariantQueueItem.ts @@ -0,0 +1,14 @@ +import { QueueItemTypes } from '.' +import { QueueItem, QueueItemConstructorInput } from './QueueItem' + +export class VariantQueueItem extends QueueItem { + constructor( + queueItemConstructorInput: QueueItemConstructorInput + ) { + super(queueItemConstructorInput) + } + + public get type(): Readonly { + return QueueItemTypes.Variant + } +} \ No newline at end of file diff --git a/src/queueItems/index.ts b/src/queueItems/index.ts new file mode 100644 index 0000000..7aff4ea --- /dev/null +++ b/src/queueItems/index.ts @@ -0,0 +1,32 @@ +import * as ExtendedQueueItem from './ExtendedQueueItem' +import * as QueueItem from './QueueItem' +import * as QuickQueueItem from './QuickQueueItem' +import * as RegeneratedQueueItem from './RegeneratedQueueItem' +import * as UpscaledQueueItem from './UpscaledQueueItem' +import * as VariantQueueItem from './VariantQueueItem' + +export enum QueueItemTypes { + Default = 'default', + Quick = 'quick', + Regenerated = 'regenerated', + Variant = 'variant', + Upscaled = 'upscaled', + Extended = 'extended' +} + +export type QueueItemInstances = + QueueItem.QueueItem | + QuickQueueItem.QuickQueueItem | + RegeneratedQueueItem.RegeneratedQueueItem | + VariantQueueItem.VariantQueueItem | + UpscaledQueueItem.UpscaledQueueItem | + ExtendedQueueItem.ExtendedQueueItem + +export { + QueueItem, + QuickQueueItem, + RegeneratedQueueItem, + VariantQueueItem, + UpscaledQueueItem, + ExtendedQueueItem +} diff --git a/src/types.ts b/src/types.ts index afbdc2d..2dcffbb 100644 --- a/src/types.ts +++ b/src/types.ts @@ -1,5 +1,11 @@ -import { ActionRowBuilder, ChatInputCommandInteraction, EmbedBuilder, RESTPostAPIApplicationCommandsJSONBody, SlashCommandBuilder } from 'discord.js' +import { ActionRowBuilder, EmbedBuilder, RESTPostAPIApplicationCommandsJSONBody } from 'discord.js' import { Bot } from './bot' +import * as QueueItems from './queueItems' + +// https://stackoverflow.com/a/40076355 +export type DeepPartial = { + [P in keyof T]?: T[P] extends object ? DeepPartial : T[P] +} export interface BotCommand { command: any // Should be a variant of SlashCommandBuilder @@ -22,34 +28,4 @@ export interface BotEvent { ) => void | Promise } -export enum QueueItemType { - Default = 'default', - Quick = 'quick', - Regenerated = 'regenerated', - Variant = 'variant', - Upscaled = 'upscaled', - Extended = 'extended' -} - -// TODO: Create a base queueitem, extend for each queueitemtype. Introduce metadata object. - -export interface QueueItem { - uuid: string // Unique queue uuid - seed: number // RNG - type: QueueItemType - imageData?: Buffer[] // Array of generated image buffers for request - messageId?: string // Discord main message snowflake ID. - interaction: ChatInputCommandInteraction - discordCaller: string // Snowflake of caller for command - prediction: { - prompt?: string // Prompt text - width: number // Image size, min 64 - height: number // Image size, min 64 - initImage: string | undefined // Starter image to generate from - mask: string | undefined // Mask image to generate with - promptStrength: number // 0 - 1 float value, default 0.8 - numOutputs: number // Number of images to generate, default 1 - numInferenceSteps: number // Number of steps, default 50 - guidanceScale: number // Default 7.5 - } -} +export { QueueItems } \ No newline at end of file