Fix bug with components staying when they shouldn't, introduce queueitem types, variant support

maxwlang · maxwlang · commit 184aa7ce91e7 · 2022-09-24T22:35:30.000-05:00
diff --git a/src/embeds/imageSelectPrompt.ts b/src/embeds/imageSelectPrompt.ts
@@ -15,7 +15,7 @@ export default function(queueItem: QueueItem): BotEmbed {
             new ActionRowBuilder()
                 .addComponents([
                     new SelectMenuBuilder()
-                        .setCustomId('image-select-prompt')
+                        .setCustomId(`image-select-prompt-${queueItem.type}`)
                         .setPlaceholder('Select an image')
                         .addOptions(
                             {
diff --git a/src/events/buttonImagineRegenerate.ts b/src/events/buttonImagineRegenerate.ts
@@ -1,72 +1,58 @@
 import { Interaction } from 'discord.js'
 import { isNil } from 'ramda'
 import { Bot } from '../bot'
-import { BotEvent, QueueItem } from '../types'
+import { BotEvent, QueueItem, QueueItemType } from '../types'
 import { getRandomInt } from '../utils'
 import { v4 as uuidv4 } from 'uuid'
-import addedToQueue from '../embeds/addedToQueue'
-import addedToQueueNoWait from '../embeds/addedToQueueNoWait'
+import addedToQueue, { QueueType } from '../embeds/addedToQueue'
 
 const botEvent: BotEvent = {
-    name: 'Command Handler - Ping',
+    name: 'Button Handler - Imagine Regenerate',
     event: 'interactionCreate',
     once: false,
     async execute(bot: Bot, interaction: Interaction) {
         if (!interaction.isButton()) return
         if (interaction.customId !== 'button-imagine-result-regenerate') return
 
-        const referenceQueueItem = bot.findQueueItemReferenceByMessageID(interaction.message.id)
+        const referenceQueueItem = bot.findLatestQueueItemReferenceByMessageID(interaction.message.id)
         if (isNil(referenceQueueItem)) {
             await interaction.reply({
                 ephemeral: true,
                 content: 'There is no reference of this image generation. Most likely, this generation is too old to modify.'
-            }).then(message => {
-                // TODO: Remove message after 5 seconds
-                // setTimeout(async () => {
-                //     if (!message.interaction.isRepliable) return
-                //     const msg = await message.awaitMessageComponent()
-                //     await msg.deleteReply()
-                // }, 5000)
             })
+
             return
         }
         
         const queueItem: QueueItem = {
             ...referenceQueueItem,
+            type: QueueItemType.Regenerated,
             discordCaller: interaction.user.id.toString(),
             seed: getRandomInt(1, 99999999),
             uuid: uuidv4()
         }
         
-        bot.log.debug(`isProcessing: ${bot.stableDiffusion.isProcessing()} hasQueue: ${bot.hasQueue()} queueLength: ${bot.queue.length}`)
-        
         // 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.'
-        }).then(message => {
-            // TODO: Remove message after 5 seconds
-            // setTimeout(async () => {
-            //     if (!message.interaction.isRepliable) return
-            //     const msg = await message.awaitMessageComponent()
-            //     await msg.deleteReply()
-            // }, 5000)
         })
 
-
         if (bot.stableDiffusion.isProcessing() || bot.hasQueue()) {
             const queuePos = bot.addQueue(queueItem)
 
             await referenceQueueItem.interaction.editReply({
-                embeds: addedToQueue(queuePos, queueItem).embeds
+                embeds: addedToQueue(QueueType.Queued, queueItem, queuePos).embeds,
+                components: []
             })
         } else {
             bot.addQueue(queueItem)
 
             await referenceQueueItem.interaction.editReply({
-                embeds: addedToQueueNoWait(queueItem).embeds
+                embeds: addedToQueue(QueueType.Instant, queueItem).embeds,
+                components: []
             })
         }
     }
diff --git a/src/events/buttonImagineVariantOf.ts b/src/events/buttonImagineVariantOf.ts
@@ -0,0 +1,89 @@
+import { Interaction } from "discord.js"
+import { isEmpty, isNil } from "ramda"
+import { Bot } from '../bot'
+import { BotEvent, QueueItem, QueueItemType } from '../types'
+import imageSelectPrompt from '../embeds/imageSelectPrompt'
+import { v4 as uuidv4 } from 'uuid'
+import addedToQueue, { QueueType } from "../embeds/addedToQueue"
+
+const botEvent: BotEvent = {
+    name: 'Button Handler - Imagine Variant Of',
+    event: 'interactionCreate',
+    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)
+            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
+            }
+
+            // Ensure PNG
+            const selectedImageB64 = `data:image/png;base64,${referenceQueueItem.imageData[0].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()
+            }
+            
+            // 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,
+                    components: []
+                })
+            } else {
+                bot.addQueue(queueItem)
+                await referenceQueueItem.interaction.editReply({
+                    embeds: addedToQueue(QueueType.Instant, queueItem).embeds,
+                    components: []
+                })
+            }
+        } else if (interaction.customId === 'button-imagine-result-variant-of' && interaction.isButton()) {
+            const referenceQueueItem = bot.findLatestQueueItemReferenceByMessageID(interaction.message.id)
+            if (isNil(referenceQueueItem)) {
+                await interaction.reply({
+                    ephemeral: true,
+                    content: 'There is no reference of this image generation. Most likely, this generation is too old to modify.'
+                })
+                return
+            }
+
+            const imageSelectPromptEmbed = imageSelectPrompt({...referenceQueueItem, type: QueueItemType.Variant})
+
+            await referenceQueueItem.interaction.editReply({
+                embeds: imageSelectPromptEmbed.embeds,
+                components: imageSelectPromptEmbed.components
+            })
+        }
+        return
+    }
+}
+
+export default botEvent
diff --git a/src/events/queueProcessor.ts b/src/events/queueProcessor.ts
@@ -21,7 +21,7 @@ async function tick(bot: Bot) {
     if (bot.stableDiffusion.isProcessing() || tickLock) return
     if (isNil(bot.queue) || isEmpty(bot.queue)) return
     tickLock = true
-    bot.log.debug(`isProcessing ; queueLength: ${bot.queue.length}`)
+    bot.log.debug(`isProcessing ; queueLength: ${bot.queue.length - 1}`)
     const queueItem = bot.queue[0]
 
     try {
@@ -53,6 +53,8 @@ async function tick(bot: Bot) {
                 const file = new AttachmentBuilder(buf, {
                     name: `stable-confusion_${queueItem.uuid}.jpeg`
                 })
+
+                queueItem.imageData = [buf]
         
                 await queueItem.interaction.editReply({
                     embeds: imageResultEmbed.embeds,
@@ -105,6 +107,7 @@ async function tick(bot: Bot) {
                     name: `stable-confusion_${queueItem.uuid}_collage.png`
                 })
         
+                queueItem.imageData = imageBuffers
         
                 await queueItem.interaction.editReply({
                     embeds: imageResultEmbed.embeds,
diff --git a/src/modules/stableDiffusion.ts b/src/modules/stableDiffusion.ts
@@ -1,4 +1,4 @@
-import { QueueItem, } from "../types"
+import { QueueItem, QueueItemType } from "../types"
 import axios from 'axios'
 import { isNil } from "ramda"
 
@@ -17,6 +17,8 @@ export default class StableDiffusion {
 
     public async processRequest(queueItem: QueueItem): Promise<false | string[]> {
         if (this.processing) return false
+        if (queueItem.type === QueueItemType.Variant) return this.processVariantRequest(queueItem)
+
         this.abortController = new AbortController()
         this.processing = true
 
@@ -46,6 +48,33 @@ export default class StableDiffusion {
         return results.data.output as string[]
     }
 
+    private async processVariantRequest(queueItem: QueueItem): Promise<false | string[]> {
+        if (this.processing) return false
+        this.abortController = new AbortController()
+        this.processing = true
+
+        const 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
+            }
+        }
+
+        const results = await axios.post(`${this.host}:${this.port ?? 5000}/predictions`, body, {
+            headers: {'Content-Type': 'application/json'}
+        }).catch(e => {
+            return null
+        })
+
+        this.processing = false
+        if (isNil(results) || isNil(results.data)) return false
+        return results.data.output as string[]
+    }
+
     public cancelRequest(): boolean | undefined {
         if (!this.abortController) throw new Error('No active abort controller')
         if (!this.processing) return false
diff --git a/src/types.ts b/src/types.ts
@@ -23,17 +23,18 @@ export interface BotEvent {
 }
 
 export enum QueueItemType {
-    Default,
-    Quick,
-    Regenerated,
-    Variant,
-    Upscaled
+    Default = 'default',
+    Quick = 'quick',
+    Regenerated = 'regenerated',
+    Variant = 'variant',
+    Upscaled = 'upscaled'
 }
 
 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
@@ -44,8 +45,8 @@ export interface QueueItem {
         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
+        numOutputs: number // Number of images to generate, default 1
         numInferenceSteps: number // Number of steps, default 50
         guidanceScale: number // Default 7.5
     }
-}
+}
diff --git a/src/utils.ts b/src/utils.ts
@@ -22,4 +22,4 @@ export const getImageAttachmentURL = (attachment: Attachment | null): string | u
     if (isNil(attachment) || isNil(attachment.contentType)) return
     if (!imageMimes.includes(attachment.contentType)) return
     return attachment.url
-}
+}

Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,7 @@ export default function(queueItem: QueueItem): BotEmbed {`
`15`	`15`	`new ActionRowBuilder()`
`16`	`16`	`.addComponents([`
`17`	`17`	`new SelectMenuBuilder()`
`18`		`- .setCustomId('image-select-prompt')`
	`18`	+ .setCustomId(`image-select-prompt-${queueItem.type}`)
`19`	`19`	`.setPlaceholder('Select an image')`
`20`	`20`	`.addOptions(`
`21`	`21`	`{`