From 52298e6c356f54a212f39cac9c6e8bbee96645fc Mon Sep 17 00:00:00 2001 From: Maxwell Lang Date: Sun, 30 Oct 2022 17:03:41 -0500 Subject: [PATCH 1/4] Migrate over default QueueItem to QueueItem class, update processor, embeds, and imagine command to support QueueItem class. --- src/bot.ts | 24 ++-- src/embeds/addedToQueue.ts | 99 ++++++++------- src/embeds/imageResult.ts | 66 +++++----- src/embeds/imageSelectPrompt.ts | 4 +- src/embeds/processingPrompt.ts | 74 +++++------ src/embeds/queueDetails.ts | 20 +-- src/embeds/queueShow.ts | 8 +- src/events/commandImagine.ts | 29 ++--- src/events/commandQueueRemove.ts | 2 +- src/events/commandQueueShow.ts | 2 +- src/events/queueProcessor.ts | 20 +-- src/queueItems/QueueItem.ts | 207 +++++++++++++++++++++++++++++++ src/queueItems/index.ts | 26 ++++ src/types.ts | 40 ++---- 14 files changed, 414 insertions(+), 207 deletions(-) create mode 100644 src/queueItems/QueueItem.ts create mode 100644 src/queueItems/index.ts diff --git a/src/bot.ts b/src/bot.ts index 4b9f168..6f9ee4c 100644 --- a/src/bot.ts +++ b/src/bot.ts @@ -1,4 +1,4 @@ -import { BotCommand, BotEvent, QueueItem } from './types' +import { BotCommand, BotEvent, QueueItems } from './types' import { Client, GatewayIntentBits, Partials, REST, Routes } from 'discord.js' import { isEmpty, isNil } from 'ramda' @@ -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,14 @@ 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 addQueue = (queueItem: QueueItems.QueueItemInstances): number => this.queue.push(queueItem) /** * 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 removeQueue = (uuid: string): QueueItems.QueueItemInstances[] => { const queueIndex = this.queue.findIndex(queueItem => queueItem.uuid === uuid) if (queueIndex === -1) return [] @@ -92,7 +92,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 [] @@ -104,7 +104,7 @@ export class Bot extends Client { * @param uuid UUID of QueueItem. * @returns Array of found QueueItems. */ - public findQueue = (uuid: string): QueueItem | undefined => { + public findQueue = (uuid: string): QueueItems.QueueItemInstances | undefined => { return this.queue.find(queueItem => queueItem.uuid === uuid) } @@ -113,17 +113,17 @@ export class Bot extends Client { * @param uuid UUID of QueueItem. * @returns Array of found QueueItems. */ - 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..c0c2366 100644 --- a/src/embeds/addedToQueue.ts +++ b/src/embeds/addedToQueue.ts @@ -1,64 +1,67 @@ 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..fe067e0 100644 --- a/src/embeds/imageSelectPrompt.ts +++ b/src/embeds/imageSelectPrompt.ts @@ -1,7 +1,7 @@ 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 { return { embeds: [ new EmbedBuilder() 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/commandImagine.ts b/src/events/commandImagine.ts index abc0531..61e3e2a 100644 --- a/src/events/commandImagine.ts +++ b/src/events/commandImagine.ts @@ -1,7 +1,7 @@ import { Interaction } from 'discord.js' import { Bot } from '../bot' -import { BotEvent, QueueItem, QueueItemType } from '../types' -import addedToQueue, { QueueType } from '../embeds/addedToQueue' +import { BotEvent, QueueItems } from '../types' +import { addedToInstantQueue, addedToQueue, QueueType } from '../embeds/addedToQueue' import { v4 as uuidv4 } from 'uuid' import { getImageAttachmentURL, getRandomInt, validateHeight, validateWidth } from '../utils' @@ -23,36 +23,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.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) await interaction.editReply({ - embeds: addedToQueue(QueueType.Queued, queueItem, queuePos).embeds + embeds: addedToQueue(queueItem, queuePos).embeds }) } else { bot.addQueue(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..0a4639e 100644 --- a/src/events/commandQueueRemove.ts +++ b/src/events/commandQueueRemove.ts @@ -21,7 +21,7 @@ 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 } 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..8201bda 100644 --- a/src/events/queueProcessor.ts +++ b/src/events/queueProcessor.ts @@ -1,7 +1,7 @@ import { isEmpty, isNil } from 'ramda' import { Bot } from '../bot' import processingPrompt from '../embeds/processingPrompt' -import { BotEvent, QueueItemType } from '../types' +import { BotEvent, QueueItems } from '../types' import { AttachmentBuilder } from 'discord.js' import imageResult from '../embeds/imageResult' import sharp from 'sharp' @@ -26,14 +26,14 @@ 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 }) @@ -46,7 +46,7 @@ async function tick(bot: Bot) { } 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] @@ -125,8 +125,8 @@ async function tick(bot: Bot) { 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/queueItems/QueueItem.ts b/src/queueItems/QueueItem.ts new file mode 100644 index 0000000..56fa362 --- /dev/null +++ b/src/queueItems/QueueItem.ts @@ -0,0 +1,207 @@ +import { v4 as uuidv4 } from 'uuid' +import { getRandomInt } from '../utils' +import { ChatInputCommandInteraction, Snowflake } from "discord.js" +import { QueueItemTypes } from '.' +import { isEmpty, isNil } from 'ramda' + + +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 // 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 + } + + private _type: QueueItemTypes = QueueItemTypes.Default + public get type(): Readonly { + return this._type + } + + 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._discordCallerSnowflake = snowflake + } + + private _discordInteraction: ChatInputCommandInteraction + public get discordInteraction(): Readonly { + 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) + } + + /** + * 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/index.ts b/src/queueItems/index.ts new file mode 100644 index 0000000..f62da23 --- /dev/null +++ b/src/queueItems/index.ts @@ -0,0 +1,26 @@ +import * as QueueItem from './QueueItem' +// import * as QuickQueueItem from './QuickQueueItem' +// import * as RegeneratedQueueItem from './RegeneratedQueueItem' +// import * as VariantQueueItem from './VariantQueueItem' +// import * as UpscaledQueueItem from './UpscaledQueueItem' +// import * as ExtendedQueueItem from './ExtendedQueueItem' + +export enum QueueItemTypes { + Default = 'default', + Quick = 'quick', + Regenerated = 'regenerated', + Variant = 'variant', + Upscaled = 'upscaled', + Extended = 'extended' +} + +export type QueueItemInstances = QueueItem.QueueItem + +export { + QueueItem, + // QuickQueueItem, + // RegeneratedQueueItem, + // VariantQueueItem, + // UpscaledQueueItem, + // ExtendedQueueItem +} \ No newline at end of file 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 From 701bf692121632c4a5c04657d98ea3f278d768dc Mon Sep 17 00:00:00 2001 From: Maxwell Lang Date: Mon, 31 Oct 2022 00:14:07 -0500 Subject: [PATCH 2/4] commit --- src/embeds/addedToQueue.ts | 5 - src/events/commandImagine.ts | 5 +- src/events/commandImagineQuick.ts | 57 +++++---- src/queueItems/ExtendedQueueItem.ts | 0 src/queueItems/QueueItem.ts | 1 - src/queueItems/QuickQueueItem.ts | 18 +++ src/queueItems/RegeneratedQueueItem.ts | 0 src/queueItems/UpscaledQueueItem.ts | 0 src/queueItems/VariantQueueItem.ts | 157 +++++++++++++++++++++++++ src/queueItems/index.ts | 4 +- 10 files changed, 215 insertions(+), 32 deletions(-) create mode 100644 src/queueItems/ExtendedQueueItem.ts create mode 100644 src/queueItems/QuickQueueItem.ts create mode 100644 src/queueItems/RegeneratedQueueItem.ts create mode 100644 src/queueItems/UpscaledQueueItem.ts create mode 100644 src/queueItems/VariantQueueItem.ts diff --git a/src/embeds/addedToQueue.ts b/src/embeds/addedToQueue.ts index c0c2366..229be90 100644 --- a/src/embeds/addedToQueue.ts +++ b/src/embeds/addedToQueue.ts @@ -1,11 +1,6 @@ import { EmbedBuilder, codeBlock, ActionRowBuilder } from 'discord.js' import { BotEmbed, QueueItems } from '../types' -export enum QueueType { - Instant, - Queued -} - export function addedToInstantQueue(queueItem: QueueItems.QueueItemInstances): BotEmbed { let embeds: EmbedBuilder[] = [] let components: ActionRowBuilder[] = [] diff --git a/src/events/commandImagine.ts b/src/events/commandImagine.ts index 61e3e2a..416b8fe 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, QueueItems } from '../types' -import { addedToInstantQueue, 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 { getImageAttachmentURL } from '../utils' const botEvent: BotEvent = { name: 'Command Handler - Imagine', diff --git a/src/events/commandImagineQuick.ts b/src/events/commandImagineQuick.ts index 9b5e55d..5027df8 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 { BotEvent, QueueItems } from '../types' +import { addedToInstantQueue, addedToQueue } from '../embeds/addedToQueue' +import { getImageAttachmentURL, validateHeight, validateWidth } from '../utils' const botEvent: BotEvent = { name: 'Command Handler - Quick Imagine', @@ -26,36 +25,52 @@ 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.Quick, - prediction: { + // const queueItem: QueueItem = { + // discordCaller: interaction.user.id.toString(), + // seed: seed ?? getRandomInt(1, 99999999), + // uuid: uuidv4(), + // interaction, + // type: QueueItemType.Quick, + // prediction: { + // prompt, + // width, + // height, + // initImage, + // mask, + // promptStrength: promptStrength ?? 0.8, + // numOutputs: numOutputs ?? 1, + // numInferenceSteps: numInferenceSteps ?? 50, + // guidanceScale: guidanceScale ?? 7.5 + // } + // } + + const queueItem = new QueueItems.QuickQueueItem.QuickQueueItem({ + discordCallerSnowflake: interaction.user.id.toString(), + discordInteraction: interaction, + seed, prompt, - width, - height, + width: 512, + height: 512, initImage, mask, - promptStrength: promptStrength ?? 0.8, - numOutputs: numOutputs ?? 1, - numInferenceSteps: numInferenceSteps ?? 50, - guidanceScale: guidanceScale ?? 7.5 - } - } + promptStrength, + guidanceScale, + numInferenceSteps + }, { + test: true + }) if (bot.stableDiffusion.isProcessing() || bot.hasQueue()) { const queuePos = bot.addQueue(queueItem) await interaction.editReply({ - embeds: addedToQueue(QueueType.Queued, queueItem, queuePos).embeds + embeds: addedToQueue(queueItem, queuePos).embeds }) } else { bot.addQueue(queueItem) await interaction.editReply({ - embeds: addedToQueue(QueueType.Instant, queueItem).embeds + embeds: addedToInstantQueue(queueItem).embeds }) } } diff --git a/src/queueItems/ExtendedQueueItem.ts b/src/queueItems/ExtendedQueueItem.ts new file mode 100644 index 0000000..e69de29 diff --git a/src/queueItems/QueueItem.ts b/src/queueItems/QueueItem.ts index 56fa362..3dd4a07 100644 --- a/src/queueItems/QueueItem.ts +++ b/src/queueItems/QueueItem.ts @@ -4,7 +4,6 @@ import { ChatInputCommandInteraction, Snowflake } from "discord.js" import { QueueItemTypes } from '.' import { isEmpty, isNil } from 'ramda' - export type QueueItemBody = { input: { prompt?: string // Prompt text diff --git a/src/queueItems/QuickQueueItem.ts b/src/queueItems/QuickQueueItem.ts new file mode 100644 index 0000000..6f3233f --- /dev/null +++ b/src/queueItems/QuickQueueItem.ts @@ -0,0 +1,18 @@ +import { QueueItemTypes, QueueItem } from '.' +import { QueueItemConstructorInput } from './QueueItem' + +export class QuickQueueItem extends QueueItem.QueueItem { + constructor( + queueItemConstructorInput: QueueItemConstructorInput, + quickQueueItemConstructorInput: unknown /* QuickQueueItemConstructorInput */ + ) { + super(queueItemConstructorInput) + + // @ts-expect-error TODO: Figure out how to get this functional + this._type = QueueItemTypes.Quick // Set type to quick QueueItem + } + + + // QuickQueueItem specific properties + public test: string | undefined +} \ No newline at end of file diff --git a/src/queueItems/RegeneratedQueueItem.ts b/src/queueItems/RegeneratedQueueItem.ts new file mode 100644 index 0000000..e69de29 diff --git a/src/queueItems/UpscaledQueueItem.ts b/src/queueItems/UpscaledQueueItem.ts new file mode 100644 index 0000000..e69de29 diff --git a/src/queueItems/VariantQueueItem.ts b/src/queueItems/VariantQueueItem.ts new file mode 100644 index 0000000..c857f73 --- /dev/null +++ b/src/queueItems/VariantQueueItem.ts @@ -0,0 +1,157 @@ +import { QueueItem } from '.' +import { QueueItemConstructorInput } from './QueueItem' + +export type VariantQueueItemConstructorInput = { + // Whatever +} + +export class VariantQueueItem extends QueueItem.QueueItem { + constructor( + queueItemConstructorInput: QueueItemConstructorInput, + variantQueueItemConstructorInput: VariantQueueItemConstructorInput + ) { + super(queueItemConstructorInput) + + // Other stuff?!!?! + console.log('test') + } + + public test: string | undefined +} + + +// In variantQueueItem class, have a method that takes values from other queueItem types. +// Carry this over to remaining queueItem variants +// Ex: new VariantQueueItem().fromQueueItem(myQueueItem) + +// export type VariantQueueItemMetadata = { +// discordOwnerSnowflake: string +// messageIdSnowflake: string +// somethingVariantSpecific: string +// } + +// export type VariantQueueItemPrediction = { +// prompt?: string +// width: number +// height: number +// somethingVariantSpecific: number +// } + +// export class VariantQueueItem implements QueueItem { +// constructor( +// metadata: VariantQueueItemMetadata, +// prediction: VariantQueueItemPrediction +// ) { +// this.metadata = metadata +// this.prediction = prediction +// } + +// public type: QueueItemTypes = QueueItemTypes.Variant +// public metadata: VariantQueueItemMetadata +// public prediction: VariantQueueItemPrediction +// } + + + +// ========= + +// import { v4 as uuidv4 } from 'uuid' +// import { getRandomInt } from '../utils' +// import { ChatInputCommandInteraction } from "discord.js" +// import { QueueItemTypes } from '.' + + +// 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 = Omit & { +// initImage?: string +// promptStrength: number +// numOutputs: number +// numInferenceSteps: number +// guidanceScale: number + +// discordCallerSnowflake: string +// discordMessageSnowflake: string +// discordInteraction: ChatInputCommandInteraction +// } + +// export type QueueItemMutationInput = Partial + +// export class QueueItem { +// constructor(queueItemConstructorInput: QueueItemConstructorInput) { +// 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 ?? 0 +// this.height = queueItemConstructorInput.height ?? 0 +// this.initImage = queueItemConstructorInput.initImage ?? undefined +// this.mask = queueItemConstructorInput.mask ?? undefined +// this.promptStrength = queueItemConstructorInput.promptStrength ?? 0.8 +// this.numOutputs = 4 // This might always need to be 4 due to variant / quad preview support. +// this.numInferenceSteps = queueItemConstructorInput.numInferenceSteps ?? 50 +// this.guidanceScale = queueItemConstructorInput.guidanceScale ?? 7.5 +// } + +// // Metadata +// public readonly uuid: string +// public readonly type: QueueItemTypes = QueueItemTypes.Default +// public readonly imageData?: Buffer[] +// public readonly discordCallerSnowflake: string +// public readonly discordMessageSnowflake: string +// public readonly discordInteraction: ChatInputCommandInteraction + +// // Prediction +// public readonly seed: number +// public readonly prompt?: string // Prompt text +// public readonly width: number // Image size, min 64 +// public readonly height: number // Image size, min 64 +// public readonly initImage?: string // Starter image to generate from +// public readonly mask?: string // Mask image to generate with +// public readonly promptStrength: number // 0 - 1 float value, default 0.8 +// public readonly numOutputs: number // Number of images to generate +// public readonly numInferenceSteps: number // Number of steps, default 50 +// public readonly guidanceScale: number // Default 7.5 + + +// public mutateQueueItem(QueueItemMutationInput: QueueItemMutationInput) { +// // Takes current QueueItem and modifies any values provided in QueueItemMutationInput +// } + +// public createStableDiffusionRequestBody(): QueueItemBody { +// // This creates the stable diffusion request body that the +// // stable diffusion model takes as input +// 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/index.ts b/src/queueItems/index.ts index f62da23..a9ddf40 100644 --- a/src/queueItems/index.ts +++ b/src/queueItems/index.ts @@ -1,5 +1,5 @@ import * as QueueItem from './QueueItem' -// import * as QuickQueueItem from './QuickQueueItem' +import * as QuickQueueItem from './QuickQueueItem' // import * as RegeneratedQueueItem from './RegeneratedQueueItem' // import * as VariantQueueItem from './VariantQueueItem' // import * as UpscaledQueueItem from './UpscaledQueueItem' @@ -18,7 +18,7 @@ export type QueueItemInstances = QueueItem.QueueItem export { QueueItem, - // QuickQueueItem, + QuickQueueItem, // RegeneratedQueueItem, // VariantQueueItem, // UpscaledQueueItem, From 477d72cb8c68049aedc640bb03e697c93408598e Mon Sep 17 00:00:00 2001 From: Maxwell Lang Date: Wed, 23 Nov 2022 21:35:15 -0600 Subject: [PATCH 3/4] Implementation progress --- src/events/buttonImagineRegenerate.ts | 40 ++++++++++++++------------ src/events/buttonImagineUpscale.ts | 6 ++-- src/events/commandExtend.ts | 39 +++++++++++-------------- src/events/commandImagineQuick.ts | 26 ++--------------- src/queueItems/ExtendedQueueItem.ts | 14 +++++++++ src/queueItems/QueueItem.ts | 11 ++++--- src/queueItems/QuickQueueItem.ts | 15 +++++----- src/queueItems/RegeneratedQueueItem.ts | 15 ++++++++++ src/queueItems/UpscaledQueueItem.ts | 19 ++++++++++++ src/queueItems/index.ts | 8 +++--- 10 files changed, 109 insertions(+), 84 deletions(-) diff --git a/src/events/buttonImagineRegenerate.ts b/src/events/buttonImagineRegenerate.ts index 599ae9c..a1ba6ea 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 { BotEvent, QueueItems } from '../types' +import { addedToInstantQueue, addedToQueue } from '../embeds/addedToQueue' const botEvent: BotEvent = { name: 'Button Handler - Imagine Regenerate', @@ -14,7 +12,7 @@ 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, @@ -23,15 +21,21 @@ const botEvent: BotEvent = { return } - - const queueItem: QueueItem = { - ...referenceQueueItem, - type: QueueItemType.Regenerated, - discordCaller: interaction.user.id.toString(), - seed: getRandomInt(1, 99999999), - uuid: uuidv4() - } - + + const queueItem = new QueueItems.RegeneratedQueueItem.RegeneratedQueueItem({ + discordCallerSnowflake: interaction.user.id.toString(), + discordInteraction: interaction, + discordMessageSnowflake: interaction.message.id, + 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() @@ -43,15 +47,15 @@ const botEvent: BotEvent = { if (bot.stableDiffusion.isProcessing() || bot.hasQueue()) { const queuePos = bot.addQueue(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) - 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..cfca49d 100644 --- a/src/events/buttonImagineUpscale.ts +++ b/src/events/buttonImagineUpscale.ts @@ -1,7 +1,7 @@ import { AttachmentBuilder, Interaction } from "discord.js" import { isEmpty, isNil } from "ramda" import { Bot } from '../bot' -import { BotEvent, QueueItemType } from '../types' +import { BotEvent } from '../types' import imageResult from "../embeds/imageResult" import imageSelectPrompt from "../embeds/imageSelectPrompt" @@ -15,7 +15,7 @@ const botEvent: BotEvent = { if (interaction.customId === 'image-select-prompt-upscaled' && interaction.isSelectMenu()) { // TODO: Lock this action to user who triggered prompt? - const referenceQueueItem = bot.findLatestQueueItemReferenceByMessageID(interaction.message.id) + const referenceQueueItem = bot.findLatestQueueItemReferenceByMessageSnowflake(interaction.message.id) if (isNil(referenceQueueItem) || isNil(referenceQueueItem.imageData) || isEmpty(referenceQueueItem.imageData)) { await interaction.reply({ ephemeral: true, @@ -40,7 +40,7 @@ const botEvent: BotEvent = { files: [file] }) } else if (interaction.customId === 'button-imagine-result-upscale' && interaction.isButton()) { - const referenceQueueItem = bot.findLatestQueueItemReferenceByMessageID(interaction.message.id) + const referenceQueueItem = bot.findLatestQueueItemReferenceByMessageSnowflake(interaction.message.id) if (isNil(referenceQueueItem)) { await interaction.reply({ ephemeral: true, diff --git a/src/events/commandExtend.ts b/src/events/commandExtend.ts index f00ef4e..c4ddabd 100644 --- a/src/events/commandExtend.ts +++ b/src/events/commandExtend.ts @@ -1,7 +1,7 @@ import { Interaction } from 'discord.js' import { Bot } from '../bot' -import { BotEvent, QueueItem, QueueItemType } from '../types' -import addedToQueue, { QueueType } from '../embeds/addedToQueue' +import { BotEvent, QueueItems } from '../types' +import { addedToInstantQueue, addedToQueue } from '../embeds/addedToQueue' import { v4 as uuidv4 } from 'uuid' import { getImageAttachmentURL, getRandomInt } from '../utils' @@ -22,36 +22,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) await interaction.editReply({ - embeds: addedToQueue(QueueType.Queued, queueItem, queuePos).embeds + embeds: addedToQueue(queueItem, queuePos).embeds }) } else { bot.addQueue(queueItem) await interaction.editReply({ - embeds: addedToQueue(QueueType.Instant, queueItem).embeds + embeds: addedToInstantQueue(queueItem).embeds }) } } diff --git a/src/events/commandImagineQuick.ts b/src/events/commandImagineQuick.ts index 5027df8..a007bc8 100644 --- a/src/events/commandImagineQuick.ts +++ b/src/events/commandImagineQuick.ts @@ -20,44 +20,22 @@ 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: { - // prompt, - // width, - // height, - // initImage, - // mask, - // promptStrength: promptStrength ?? 0.8, - // numOutputs: numOutputs ?? 1, - // numInferenceSteps: numInferenceSteps ?? 50, - // guidanceScale: guidanceScale ?? 7.5 - // } - // } - const queueItem = new QueueItems.QuickQueueItem.QuickQueueItem({ discordCallerSnowflake: interaction.user.id.toString(), discordInteraction: interaction, seed, prompt, - width: 512, - height: 512, + width, + height, initImage, mask, promptStrength, guidanceScale, numInferenceSteps - }, { - test: true }) if (bot.stableDiffusion.isProcessing() || bot.hasQueue()) { diff --git a/src/queueItems/ExtendedQueueItem.ts b/src/queueItems/ExtendedQueueItem.ts index e69de29..abae92b 100644 --- a/src/queueItems/ExtendedQueueItem.ts +++ b/src/queueItems/ExtendedQueueItem.ts @@ -0,0 +1,14 @@ +import { QueueItemTypes, QueueItem } from '.' +import { QueueItemConstructorInput } from './QueueItem' + +export class ExtendedQueueItem extends QueueItem.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 index 3dd4a07..0c96a41 100644 --- a/src/queueItems/QueueItem.ts +++ b/src/queueItems/QueueItem.ts @@ -1,6 +1,6 @@ import { v4 as uuidv4 } from 'uuid' import { getRandomInt } from '../utils' -import { ChatInputCommandInteraction, Snowflake } from "discord.js" +import { ButtonInteraction, ChatInputCommandInteraction, Snowflake } from "discord.js" import { QueueItemTypes } from '.' import { isEmpty, isNil } from 'ramda' @@ -22,7 +22,7 @@ export type QueueItemBody = { export type QueueItemConstructorInput = { discordCallerSnowflake: Snowflake // Discord snowflake of invoking user discordMessageSnowflake?: Snowflake // Discord snowflake of Stable Confusion's related message embed - discordInteraction: ChatInputCommandInteraction // Discord interaction handle for QueueItem + discordInteraction: ChatInputCommandInteraction | ButtonInteraction // Discord interaction handle for QueueItem prompt?: string // Prompt text width: number // Image size, min 64 @@ -75,9 +75,8 @@ export class QueueItem { return this._uuid } - private _type: QueueItemTypes = QueueItemTypes.Default public get type(): Readonly { - return this._type + return QueueItemTypes.Default } private _imageData?: Buffer[] @@ -109,8 +108,8 @@ export class QueueItem { this._discordCallerSnowflake = snowflake } - private _discordInteraction: ChatInputCommandInteraction - public get discordInteraction(): Readonly { + private _discordInteraction: ChatInputCommandInteraction | ButtonInteraction + public get discordInteraction(): Readonly { return this._discordInteraction } diff --git a/src/queueItems/QuickQueueItem.ts b/src/queueItems/QuickQueueItem.ts index 6f3233f..7ab8462 100644 --- a/src/queueItems/QuickQueueItem.ts +++ b/src/queueItems/QuickQueueItem.ts @@ -3,16 +3,17 @@ import { QueueItemConstructorInput } from './QueueItem' export class QuickQueueItem extends QueueItem.QueueItem { constructor( - queueItemConstructorInput: QueueItemConstructorInput, - quickQueueItemConstructorInput: unknown /* QuickQueueItemConstructorInput */ + queueItemConstructorInput: QueueItemConstructorInput ) { super(queueItemConstructorInput) - - // @ts-expect-error TODO: Figure out how to get this functional - this._type = QueueItemTypes.Quick // Set type to quick QueueItem } + public get type(): Readonly { + return QueueItemTypes.Quick + } - // QuickQueueItem specific properties - public test: string | undefined + // 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 index e69de29..2ad45d0 100644 --- a/src/queueItems/RegeneratedQueueItem.ts +++ b/src/queueItems/RegeneratedQueueItem.ts @@ -0,0 +1,15 @@ +import { QueueItemTypes, QueueItem } from '.' +import { QueueItemConstructorInput } from './QueueItem' + +export class RegeneratedQueueItem extends QueueItem.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 index e69de29..466f602 100644 --- a/src/queueItems/UpscaledQueueItem.ts +++ b/src/queueItems/UpscaledQueueItem.ts @@ -0,0 +1,19 @@ +import { QueueItemTypes, QueueItem } from '.' +import { QueueItemConstructorInput } from './QueueItem' + +export class UpscaledQueueItem extends QueueItem.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/index.ts b/src/queueItems/index.ts index a9ddf40..ef80e26 100644 --- a/src/queueItems/index.ts +++ b/src/queueItems/index.ts @@ -1,9 +1,9 @@ import * as QueueItem from './QueueItem' import * as QuickQueueItem from './QuickQueueItem' -// import * as RegeneratedQueueItem from './RegeneratedQueueItem' +import * as RegeneratedQueueItem from './RegeneratedQueueItem' // import * as VariantQueueItem from './VariantQueueItem' // import * as UpscaledQueueItem from './UpscaledQueueItem' -// import * as ExtendedQueueItem from './ExtendedQueueItem' +import * as ExtendedQueueItem from './ExtendedQueueItem' export enum QueueItemTypes { Default = 'default', @@ -19,8 +19,8 @@ export type QueueItemInstances = QueueItem.QueueItem export { QueueItem, QuickQueueItem, - // RegeneratedQueueItem, + RegeneratedQueueItem, // VariantQueueItem, // UpscaledQueueItem, - // ExtendedQueueItem + ExtendedQueueItem } \ No newline at end of file From d38c14e471f4debdc95109beb534003a3b176fdd Mon Sep 17 00:00:00 2001 From: maxwlang Date: Thu, 24 Nov 2022 00:32:01 -0600 Subject: [PATCH 4/4] Additional implementation progress --- src/bot.ts | 58 +++++++-- src/embeds/imageSelectPrompt.ts | 3 +- src/events/buttonImagineRegenerate.ts | 25 ++-- src/events/buttonImagineUpscale.ts | 72 +++++++---- src/events/buttonImagineVariantOf.ts | 77 ++++++------ src/events/commandExtend.ts | 9 +- src/events/commandImagine.ts | 20 +++- src/events/commandImagineQuick.ts | 6 +- src/events/commandQueueRemove.ts | 2 +- src/events/queueProcessor.ts | 10 +- src/modules/stableDiffusion.ts | 48 ++++---- src/queueItems/ExtendedQueueItem.ts | 6 +- src/queueItems/QueueItem.ts | 35 ++++-- src/queueItems/QuickQueueItem.ts | 6 +- src/queueItems/RegeneratedQueueItem.ts | 6 +- src/queueItems/UpscaledQueueItem.ts | 6 +- src/queueItems/VariantQueueItem.ts | 159 ++----------------------- src/queueItems/index.ts | 20 ++-- 18 files changed, 268 insertions(+), 300 deletions(-) diff --git a/src/bot.ts b/src/bot.ts index 6f9ee4c..738661c 100644 --- a/src/bot.ts +++ b/src/bot.ts @@ -1,12 +1,12 @@ -import { BotCommand, BotEvent, QueueItems } 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}) { @@ -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: QueueItems.QueueItemInstances): 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): QueueItems.QueueItemInstances[] => { + public removeQueuedQueueItem = (uuid: string): QueueItems.QueueItemInstances[] => { const queueIndex = this.queue.findIndex(queueItem => queueItem.uuid === uuid) if (queueIndex === -1) return [] @@ -102,7 +142,7 @@ 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): QueueItems.QueueItemInstances | undefined => { return this.queue.find(queueItem => queueItem.uuid === uuid) @@ -111,7 +151,7 @@ export class Bot extends Client { /** * 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): QueueItems.QueueItemInstances | undefined => { return this.queueItemReferences.find(queueItem => queueItem.uuid === uuid) diff --git a/src/embeds/imageSelectPrompt.ts b/src/embeds/imageSelectPrompt.ts index fe067e0..865a1dd 100644 --- a/src/embeds/imageSelectPrompt.ts +++ b/src/embeds/imageSelectPrompt.ts @@ -2,6 +2,7 @@ import { ActionRowBuilder, EmbedBuilder, SelectMenuBuilder } from 'discord.js' import { BotEmbed, QueueItems } from '../types' 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: QueueItems.QueueItemInstances): 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/events/buttonImagineRegenerate.ts b/src/events/buttonImagineRegenerate.ts index a1ba6ea..b0a75cf 100644 --- a/src/events/buttonImagineRegenerate.ts +++ b/src/events/buttonImagineRegenerate.ts @@ -1,8 +1,8 @@ import { Interaction } from 'discord.js' import { isNil } from 'ramda' import { Bot } from '../bot' -import { BotEvent, QueueItems } from '../types' import { addedToInstantQueue, addedToQueue } from '../embeds/addedToQueue' +import { BotEvent, QueueItems } from '../types' const botEvent: BotEvent = { name: 'Button Handler - Imagine Regenerate', @@ -13,19 +13,29 @@ const botEvent: BotEvent = { if (interaction.customId !== 'button-imagine-result-regenerate') return 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 + } + + if (interaction.user.id !== referenceQueueItem.discordCallerSnowflake) { + await interaction.reply({ + ephemeral: true, + content: 'You can not modify content generated by other users.' + }) return } const queueItem = new QueueItems.RegeneratedQueueItem.RegeneratedQueueItem({ discordCallerSnowflake: interaction.user.id.toString(), - discordInteraction: interaction, - discordMessageSnowflake: interaction.message.id, + discordInteraction: referenceQueueItem.discordInteraction, + discordMessageSnowflake: referenceQueueItem.discordMessageSnowflake, prompt: referenceQueueItem.prompt, height: referenceQueueItem.height, width: referenceQueueItem.width, @@ -39,20 +49,15 @@ const botEvent: BotEvent = { // 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.' - }) - if (bot.stableDiffusion.isProcessing() || bot.hasQueue()) { - const queuePos = bot.addQueue(queueItem) + const queuePos = bot.addQueuedQueueItem(queueItem) await referenceQueueItem.discordInteraction.editReply({ embeds: addedToQueue(queueItem, queuePos).embeds, components: [] }) } else { - bot.addQueue(queueItem) + bot.addQueuedQueueItem(queueItem) await referenceQueueItem.discordInteraction.editReply({ embeds: addedToInstantQueue(queueItem).embeds, diff --git a/src/events/buttonImagineUpscale.ts b/src/events/buttonImagineUpscale.ts index cfca49d..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 } 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.findLatestQueueItemReferenceByMessageSnowflake(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.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 - } + // 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 c4ddabd..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, QueueItems } from '../types' import { addedToInstantQueue, addedToQueue } from '../embeds/addedToQueue' -import { v4 as uuidv4 } from 'uuid' -import { getImageAttachmentURL, getRandomInt } from '../utils' +import { BotEvent, QueueItems } from '../types' +import { getImageAttachmentURL } from '../utils' const botEvent: BotEvent = { name: 'Command Handler - Extend', @@ -37,13 +36,13 @@ const botEvent: BotEvent = { }) if (bot.stableDiffusion.isProcessing() || bot.hasQueue()) { - const queuePos = bot.addQueue(queueItem) + const queuePos = bot.addQueuedQueueItem(queueItem) await interaction.editReply({ embeds: addedToQueue(queueItem, queuePos).embeds }) } else { - bot.addQueue(queueItem) + bot.addQueuedQueueItem(queueItem) await interaction.editReply({ embeds: addedToInstantQueue(queueItem).embeds diff --git a/src/events/commandImagine.ts b/src/events/commandImagine.ts index 416b8fe..2e47b1d 100644 --- a/src/events/commandImagine.ts +++ b/src/events/commandImagine.ts @@ -1,7 +1,7 @@ import { Interaction } from 'discord.js' import { Bot } from '../bot' -import { BotEvent, QueueItems } from '../types' import { addedToInstantQueue, addedToQueue } from '../embeds/addedToQueue' +import { BotEvent, QueueItems } from '../types' import { getImageAttachmentURL } from '../utils' const botEvent: BotEvent = { @@ -37,17 +37,27 @@ const botEvent: BotEvent = { }) if (bot.stableDiffusion.isProcessing() || bot.hasQueue()) { - const queuePos = bot.addQueue(queueItem) + const queuePos = bot.addQueuedQueueItem(queueItem) - await interaction.editReply({ + 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({ + 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 a007bc8..4f7b988 100644 --- a/src/events/commandImagineQuick.ts +++ b/src/events/commandImagineQuick.ts @@ -1,7 +1,7 @@ import { Interaction } from 'discord.js' import { Bot } from '../bot' -import { BotEvent, QueueItems } from '../types' import { addedToInstantQueue, addedToQueue } from '../embeds/addedToQueue' +import { BotEvent, QueueItems } from '../types' import { getImageAttachmentURL, validateHeight, validateWidth } from '../utils' const botEvent: BotEvent = { @@ -39,13 +39,13 @@ const botEvent: BotEvent = { }) if (bot.stableDiffusion.isProcessing() || bot.hasQueue()) { - const queuePos = bot.addQueue(queueItem) + const queuePos = bot.addQueuedQueueItem(queueItem) await interaction.editReply({ embeds: addedToQueue(queueItem, queuePos).embeds }) } else { - bot.addQueue(queueItem) + bot.addQueuedQueueItem(queueItem) await interaction.editReply({ embeds: addedToInstantQueue(queueItem).embeds diff --git a/src/events/commandQueueRemove.ts b/src/events/commandQueueRemove.ts index 0a4639e..716551b 100644 --- a/src/events/commandQueueRemove.ts +++ b/src/events/commandQueueRemove.ts @@ -26,7 +26,7 @@ const botEvent: BotEvent = { return } - const deletedItems = bot.removeQueue(uuid) + const deletedItems = bot.removeQueuedQueueItem(uuid) const queueDetailsEmbed = queueDetails(deletedItems[0]) await interaction.reply({ diff --git a/src/events/queueProcessor.ts b/src/events/queueProcessor.ts index 8201bda..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 imageResult from '../embeds/imageResult' import processingPrompt from '../embeds/processingPrompt' import { BotEvent, QueueItems } from '../types' -import { AttachmentBuilder } from 'discord.js' -import imageResult from '../embeds/imageResult' -import sharp from 'sharp' const botEvent: BotEvent = { name: 'Queue Processor', @@ -40,7 +40,7 @@ async function tick(bot: Bot) { await message.suppressEmbeds() bot.log.debug('Processing failed') - bot.removeQueue(queueItem.uuid) + bot.removeQueuedQueueItem(queueItem.uuid) tickLock = false return } @@ -118,7 +118,7 @@ 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) { 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 index abae92b..e28d33f 100644 --- a/src/queueItems/ExtendedQueueItem.ts +++ b/src/queueItems/ExtendedQueueItem.ts @@ -1,7 +1,7 @@ -import { QueueItemTypes, QueueItem } from '.' -import { QueueItemConstructorInput } from './QueueItem' +import { QueueItemTypes } from '.' +import { QueueItem, QueueItemConstructorInput } from './QueueItem' -export class ExtendedQueueItem extends QueueItem.QueueItem { +export class ExtendedQueueItem extends QueueItem { constructor( queueItemConstructorInput: QueueItemConstructorInput ) { diff --git a/src/queueItems/QueueItem.ts b/src/queueItems/QueueItem.ts index 0c96a41..c1cc1be 100644 --- a/src/queueItems/QueueItem.ts +++ b/src/queueItems/QueueItem.ts @@ -1,8 +1,10 @@ -import { v4 as uuidv4 } from 'uuid' -import { getRandomInt } from '../utils' -import { ButtonInteraction, ChatInputCommandInteraction, Snowflake } from "discord.js" -import { QueueItemTypes } from '.' +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: { @@ -22,7 +24,7 @@ export type QueueItemBody = { export type QueueItemConstructorInput = { discordCallerSnowflake: Snowflake // Discord snowflake of invoking user discordMessageSnowflake?: Snowflake // Discord snowflake of Stable Confusion's related message embed - discordInteraction: ChatInputCommandInteraction | ButtonInteraction // Discord interaction handle for QueueItem + discordInteraction: ChatInputCommandInteraction | ButtonInteraction | SelectMenuInteraction // Discord interaction handle for QueueItem prompt?: string // Prompt text width: number // Image size, min 64 @@ -105,11 +107,11 @@ export class QueueItem { throw new Error('A snowflake must be provided.') } - this._discordCallerSnowflake = snowflake + this._discordMessageSnowflake = snowflake } - private _discordInteraction: ChatInputCommandInteraction | ButtonInteraction - public get discordInteraction(): Readonly { + private _discordInteraction: ChatInputCommandInteraction | ButtonInteraction | SelectMenuInteraction + public get discordInteraction(): ChatInputCommandInteraction | ButtonInteraction | SelectMenuInteraction { return this._discordInteraction } @@ -181,6 +183,23 @@ export class QueueItem { 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 diff --git a/src/queueItems/QuickQueueItem.ts b/src/queueItems/QuickQueueItem.ts index 7ab8462..581c8ef 100644 --- a/src/queueItems/QuickQueueItem.ts +++ b/src/queueItems/QuickQueueItem.ts @@ -1,7 +1,7 @@ -import { QueueItemTypes, QueueItem } from '.' -import { QueueItemConstructorInput } from './QueueItem' +import { QueueItemTypes } from '.' +import { QueueItem, QueueItemConstructorInput } from './QueueItem' -export class QuickQueueItem extends QueueItem.QueueItem { +export class QuickQueueItem extends QueueItem { constructor( queueItemConstructorInput: QueueItemConstructorInput ) { diff --git a/src/queueItems/RegeneratedQueueItem.ts b/src/queueItems/RegeneratedQueueItem.ts index 2ad45d0..439c7f1 100644 --- a/src/queueItems/RegeneratedQueueItem.ts +++ b/src/queueItems/RegeneratedQueueItem.ts @@ -1,7 +1,7 @@ -import { QueueItemTypes, QueueItem } from '.' -import { QueueItemConstructorInput } from './QueueItem' +import { QueueItemTypes } from '.' +import { QueueItem, QueueItemConstructorInput } from './QueueItem' -export class RegeneratedQueueItem extends QueueItem.QueueItem { +export class RegeneratedQueueItem extends QueueItem { constructor( queueItemConstructorInput: QueueItemConstructorInput ) { diff --git a/src/queueItems/UpscaledQueueItem.ts b/src/queueItems/UpscaledQueueItem.ts index 466f602..889b3b7 100644 --- a/src/queueItems/UpscaledQueueItem.ts +++ b/src/queueItems/UpscaledQueueItem.ts @@ -1,7 +1,7 @@ -import { QueueItemTypes, QueueItem } from '.' -import { QueueItemConstructorInput } from './QueueItem' +import { QueueItemTypes } from '.' +import { QueueItem, QueueItemConstructorInput } from './QueueItem' -export class UpscaledQueueItem extends QueueItem.QueueItem { +export class UpscaledQueueItem extends QueueItem { constructor( queueItemConstructorInput: QueueItemConstructorInput ) { diff --git a/src/queueItems/VariantQueueItem.ts b/src/queueItems/VariantQueueItem.ts index c857f73..075b499 100644 --- a/src/queueItems/VariantQueueItem.ts +++ b/src/queueItems/VariantQueueItem.ts @@ -1,157 +1,14 @@ -import { QueueItem } from '.' -import { QueueItemConstructorInput } from './QueueItem' +import { QueueItemTypes } from '.' +import { QueueItem, QueueItemConstructorInput } from './QueueItem' -export type VariantQueueItemConstructorInput = { - // Whatever -} - -export class VariantQueueItem extends QueueItem.QueueItem { +export class VariantQueueItem extends QueueItem { constructor( - queueItemConstructorInput: QueueItemConstructorInput, - variantQueueItemConstructorInput: VariantQueueItemConstructorInput + queueItemConstructorInput: QueueItemConstructorInput ) { super(queueItemConstructorInput) - - // Other stuff?!!?! - console.log('test') } - public test: string | undefined -} - - -// In variantQueueItem class, have a method that takes values from other queueItem types. -// Carry this over to remaining queueItem variants -// Ex: new VariantQueueItem().fromQueueItem(myQueueItem) - -// export type VariantQueueItemMetadata = { -// discordOwnerSnowflake: string -// messageIdSnowflake: string -// somethingVariantSpecific: string -// } - -// export type VariantQueueItemPrediction = { -// prompt?: string -// width: number -// height: number -// somethingVariantSpecific: number -// } - -// export class VariantQueueItem implements QueueItem { -// constructor( -// metadata: VariantQueueItemMetadata, -// prediction: VariantQueueItemPrediction -// ) { -// this.metadata = metadata -// this.prediction = prediction -// } - -// public type: QueueItemTypes = QueueItemTypes.Variant -// public metadata: VariantQueueItemMetadata -// public prediction: VariantQueueItemPrediction -// } - - - -// ========= - -// import { v4 as uuidv4 } from 'uuid' -// import { getRandomInt } from '../utils' -// import { ChatInputCommandInteraction } from "discord.js" -// import { QueueItemTypes } from '.' - - -// 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 = Omit & { -// initImage?: string -// promptStrength: number -// numOutputs: number -// numInferenceSteps: number -// guidanceScale: number - -// discordCallerSnowflake: string -// discordMessageSnowflake: string -// discordInteraction: ChatInputCommandInteraction -// } - -// export type QueueItemMutationInput = Partial - -// export class QueueItem { -// constructor(queueItemConstructorInput: QueueItemConstructorInput) { -// 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 ?? 0 -// this.height = queueItemConstructorInput.height ?? 0 -// this.initImage = queueItemConstructorInput.initImage ?? undefined -// this.mask = queueItemConstructorInput.mask ?? undefined -// this.promptStrength = queueItemConstructorInput.promptStrength ?? 0.8 -// this.numOutputs = 4 // This might always need to be 4 due to variant / quad preview support. -// this.numInferenceSteps = queueItemConstructorInput.numInferenceSteps ?? 50 -// this.guidanceScale = queueItemConstructorInput.guidanceScale ?? 7.5 -// } - -// // Metadata -// public readonly uuid: string -// public readonly type: QueueItemTypes = QueueItemTypes.Default -// public readonly imageData?: Buffer[] -// public readonly discordCallerSnowflake: string -// public readonly discordMessageSnowflake: string -// public readonly discordInteraction: ChatInputCommandInteraction - -// // Prediction -// public readonly seed: number -// public readonly prompt?: string // Prompt text -// public readonly width: number // Image size, min 64 -// public readonly height: number // Image size, min 64 -// public readonly initImage?: string // Starter image to generate from -// public readonly mask?: string // Mask image to generate with -// public readonly promptStrength: number // 0 - 1 float value, default 0.8 -// public readonly numOutputs: number // Number of images to generate -// public readonly numInferenceSteps: number // Number of steps, default 50 -// public readonly guidanceScale: number // Default 7.5 - - -// public mutateQueueItem(QueueItemMutationInput: QueueItemMutationInput) { -// // Takes current QueueItem and modifies any values provided in QueueItemMutationInput -// } - -// public createStableDiffusionRequestBody(): QueueItemBody { -// // This creates the stable diffusion request body that the -// // stable diffusion model takes as input -// 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 -// } -// } -// } -// } + 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 index ef80e26..7aff4ea 100644 --- a/src/queueItems/index.ts +++ b/src/queueItems/index.ts @@ -1,9 +1,9 @@ +import * as ExtendedQueueItem from './ExtendedQueueItem' import * as QueueItem from './QueueItem' import * as QuickQueueItem from './QuickQueueItem' import * as RegeneratedQueueItem from './RegeneratedQueueItem' -// import * as VariantQueueItem from './VariantQueueItem' -// import * as UpscaledQueueItem from './UpscaledQueueItem' -import * as ExtendedQueueItem from './ExtendedQueueItem' +import * as UpscaledQueueItem from './UpscaledQueueItem' +import * as VariantQueueItem from './VariantQueueItem' export enum QueueItemTypes { Default = 'default', @@ -14,13 +14,19 @@ export enum QueueItemTypes { Extended = 'extended' } -export type QueueItemInstances = QueueItem.QueueItem +export type QueueItemInstances = + QueueItem.QueueItem | + QuickQueueItem.QuickQueueItem | + RegeneratedQueueItem.RegeneratedQueueItem | + VariantQueueItem.VariantQueueItem | + UpscaledQueueItem.UpscaledQueueItem | + ExtendedQueueItem.ExtendedQueueItem export { QueueItem, QuickQueueItem, RegeneratedQueueItem, - // VariantQueueItem, - // UpscaledQueueItem, + VariantQueueItem, + UpscaledQueueItem, ExtendedQueueItem -} \ No newline at end of file +}