Extend command, disable dm renders

maxwlang · maxwlang · commit 68fc82f41fa2 · 2022-10-05T02:01:28.000-05:00
diff --git a/src/commands/extend.ts b/src/commands/extend.ts
@@ -0,0 +1,40 @@
+import { SlashCommandBuilder } from 'discord.js';
+import { BotCommand } from '../types';
+export default function(): BotCommand {
+    const command = new SlashCommandBuilder()
+        .setName('extend')
+        .setDMPermission(false)
+        .setDescription('Generate images from a base image without a prompt')
+        .addAttachmentOption(option =>
+            option.setName('image')
+                .setDescription('Inital jpg or png image to generate variations of. Will be resized to the specified width and height')
+                .setRequired(true))
+        .addAttachmentOption(option =>
+            option.setName('mask')
+                .setDescription('Black and white jpg or png image to use as mask for inpainting over init_image'))
+        .addNumberOption(option =>
+            option.setName('pstrength')
+                .setDescription('Prompt strength when using init image. 1.0 corresponds to full destruction of information in image')
+                .setMinValue(0)
+                .setMaxValue(1))
+        .addIntegerOption(option =>
+            option.setName('numsteps')
+            .setDescription('Number of denoising steps. Default 50')
+            .setMinValue(1)
+            .setMaxValue(500))
+        .addNumberOption(option =>
+            option.setName('guidescale')
+            .setDescription('Scale for classifier-free guidance. Default 7.5')
+            .setMinValue(1)
+            .setMaxValue(20))
+        .addIntegerOption(option =>
+            option.setName('seed')
+            .setDescription('Generation seed. Randomized if not provided')
+            .setMinValue(0)
+            .setMaxValue(999999999))
+
+    return {
+        command,
+        commandJson: command.toJSON()
+    }
+}
diff --git a/src/commands/imagine.ts b/src/commands/imagine.ts
@@ -3,6 +3,7 @@ import { BotCommand } from '../types';
 export default function(): BotCommand {
     const command = new SlashCommandBuilder()
         .setName('imagine')
+        .setDMPermission(false)
         .setDescription('Generate images from prompts')
         .addStringOption(option =>
             option.setName('prompt')
diff --git a/src/commands/imagineQuick.ts b/src/commands/imagineQuick.ts
@@ -3,6 +3,7 @@ import { BotCommand } from '../types';
 export default function(): BotCommand {
     const command = new SlashCommandBuilder()
         .setName('imagine-quick')
+        .setDMPermission(false)
         .setDescription('Generate a single image based on a prompt. Faster.')
         .addStringOption(option =>
             option.setName('prompt')
diff --git a/src/commands/queue.ts b/src/commands/queue.ts
@@ -3,6 +3,7 @@ import { BotCommand } from '../types';
 export default function(): BotCommand {
     const command =  new SlashCommandBuilder()
         .setName('queue')
+        .setDMPermission(false)
         .setDescription('Prompt queue')
             .addSubcommand(subcommand =>
                 subcommand
diff --git a/src/embeds/addedToQueue.ts b/src/embeds/addedToQueue.ts
@@ -20,7 +20,7 @@ export default function(queueType: QueueType, queueItem: QueueItem, queuePositio
                     .setThumbnail('https://i.imgur.com/pXmPAAG.gif')
                     .setTimestamp()
                     .addFields([
-                        {name: 'Prompt', value: codeBlock(queueItem.prediction.prompt), inline: false},
+                        {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},
@@ -44,7 +44,7 @@ export default function(queueType: QueueType, queueItem: QueueItem, queuePositio
                     .setThumbnail('https://i.imgur.com/pXmPAAG.gif')
                     .setTimestamp()
                     .addFields([
-                        {name: 'Prompt', value: codeBlock(queueItem.prediction.prompt), inline: false},
+                        {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},
diff --git a/src/embeds/imageResult.ts b/src/embeds/imageResult.ts
@@ -14,7 +14,7 @@ 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: codeBlock(queueItem.prediction.prompt), inline: false},
+                        {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},
@@ -37,7 +37,7 @@ export default function(queueItem: QueueItem): BotEmbed {
                     .setDescription('Result for upscale')
                     .setTimestamp()
                     .addFields([
-                        {name: 'Prompt', value: codeBlock(queueItem.prediction.prompt), inline: false},
+                        {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},
@@ -58,7 +58,7 @@ export default function(queueItem: QueueItem): BotEmbed {
                     .setDescription('Variant results are displayed below')
                     .setTimestamp()
                     .addFields([
-                        {name: 'Prompt', value: codeBlock(queueItem.prediction.prompt), inline: false},
+                        {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},
@@ -98,7 +98,7 @@ export default function(queueItem: QueueItem): BotEmbed {
                     .setDescription('Result for parameters')
                     .setTimestamp()
                     .addFields([
-                        {name: 'Prompt', value: codeBlock(queueItem.prediction.prompt), inline: false},
+                        {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},
diff --git a/src/embeds/processingPrompt.ts b/src/embeds/processingPrompt.ts
@@ -15,7 +15,7 @@ export default function(queueItem: QueueItem): BotEmbed {
                     .setThumbnail('https://i.gifer.com/2RNf.gif')
                     .setTimestamp()
                     .addFields([
-                        {name: 'Prompt', value: codeBlock(queueItem.prediction.prompt), inline: false},
+                        {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},
@@ -36,7 +36,7 @@ export default function(queueItem: QueueItem): BotEmbed {
                     .setThumbnail('https://i.gifer.com/2RNf.gif')
                     .setTimestamp()
                     .addFields([
-                        {name: 'Prompt', value: codeBlock(queueItem.prediction.prompt), inline: false},
+                        {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},
@@ -57,7 +57,7 @@ export default function(queueItem: QueueItem): BotEmbed {
                     .setThumbnail('https://i.gifer.com/2RNf.gif')
                     .setTimestamp()
                     .addFields([
-                        {name: 'Prompt', value: codeBlock(queueItem.prediction.prompt), inline: false},
+                        {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},
@@ -78,7 +78,7 @@ export default function(queueItem: QueueItem): BotEmbed {
                     .setThumbnail('https://i.gifer.com/2RNf.gif')
                     .setTimestamp()
                     .addFields([
-                        {name: 'Prompt', value: codeBlock(queueItem.prediction.prompt), inline: false},
+                        {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},
@@ -100,7 +100,7 @@ export default function(queueItem: QueueItem): BotEmbed {
                     .setThumbnail('https://i.gifer.com/2RNf.gif')
                     .setTimestamp()
                     .addFields([
-                        {name: 'Prompt', value: codeBlock(queueItem.prediction.prompt), inline: false},
+                        {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},
diff --git a/src/embeds/queueDetails.ts b/src/embeds/queueDetails.ts
@@ -10,7 +10,7 @@ export default function(queueItem: QueueItem): BotEmbed {
                 .setThumbnail('https://i.imgur.com/wj1PRKP.png')
                 .setTimestamp()
                 .addFields([
-                    {name: 'Prompt', value: codeBlock(queueItem.prediction.prompt), inline: false},
+                    {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},
diff --git a/src/embeds/queueShow.ts b/src/embeds/queueShow.ts
@@ -15,7 +15,7 @@ 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: codeBlock(queueItem.prediction.prompt), inline: true},
+            {name: 'Prompt', value: queueItem.prediction.prompt ? codeBlock(queueItem.prediction.prompt) : 'Not Supplied', inline: true},
             {name: 'Owner', value: userMention(queueItem.discordCaller), inline: true},
             {name: 'Prompt ID', value: codeBlock(queueItem.uuid), inline: true},
         ]
diff --git a/src/events/commandExtend.ts b/src/events/commandExtend.ts
@@ -0,0 +1,60 @@
+import { Interaction } from 'discord.js'
+import { Bot } from '../bot'
+import { BotEvent, QueueItem, QueueItemType } from '../types'
+import addedToQueue, { QueueType } from '../embeds/addedToQueue'
+import { v4 as uuidv4 } from 'uuid'
+import { getImageAttachmentURL, getRandomInt } from '../utils'
+
+const botEvent: BotEvent = {
+    name: 'Command Handler - Extend',
+    event: 'interactionCreate',
+    once: false,
+    async execute(bot: Bot, interaction: Interaction) {
+        if (!interaction.isChatInputCommand()) return
+        if (interaction.commandName !== 'extend') return
+
+        await interaction.deferReply()
+
+        const initImage = getImageAttachmentURL(interaction.options.getAttachment('image'))
+        const mask = getImageAttachmentURL(interaction.options.getAttachment('mask'))
+        const promptStrength = interaction.options.getNumber('pstrength')
+        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.Extended,
+          prediction: {
+            prompt: undefined,
+            width: 512,
+            height: 512,
+            initImage,
+            mask,
+            promptStrength: promptStrength ?? 0.8,
+            numOutputs: 4,
+            numInferenceSteps: numInferenceSteps ?? 50,
+            guidanceScale: guidanceScale ?? 7.5
+          }
+        }
+
+        if (bot.stableDiffusion.isProcessing() || bot.hasQueue()) {
+          const queuePos = bot.addQueue(queueItem)
+
+          await interaction.editReply({
+            embeds: addedToQueue(QueueType.Queued, queueItem, queuePos).embeds
+          })
+        } else {
+          bot.addQueue(queueItem)
+
+          await interaction.editReply({
+            embeds: addedToQueue(QueueType.Instant, queueItem).embeds
+          })
+        }
+    }
+}
+
+export default botEvent
diff --git a/src/modules/stableDiffusion.ts b/src/modules/stableDiffusion.ts
@@ -18,6 +18,7 @@ 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)
+        if (queueItem.type === QueueItemType.Extended) return this.processExtendRequest(queueItem) // ToDo: Can we reduce duplicate code?
 
         this.abortController = new AbortController()
         this.processing = true
@@ -75,6 +76,32 @@ export default class StableDiffusion {
         return results.data.output as string[]
     }
 
+    private async processExtendRequest(queueItem: QueueItem): Promise<false | string[]> {
+        if (this.processing) return false
+        this.abortController = new AbortController()
+        this.processing = true
+
+        const 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
+            }
+        }
+
+        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
@@ -27,9 +27,12 @@ export enum QueueItemType {
     Quick = 'quick',
     Regenerated = 'regenerated',
     Variant = 'variant',
-    Upscaled = 'upscaled'
+    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
@@ -39,7 +42,7 @@ export interface QueueItem {
     interaction: ChatInputCommandInteraction
     discordCaller: string // Snowflake of caller for command
     prediction: {
-        prompt: string // Prompt text
+        prompt?: string // Prompt text
         width: number // Image size, min 64
         height: number // Image size, min 64
         initImage: string | undefined // Starter image to generate from

Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,7 @@ export default function(args: QueueShowArgs): BotEmbed {`
`15`	`15`
`16`	`16`	`const queueField = [`
`17`	`17`	{name: position === 0 ? 'Next Up' : 'Queue Position', value: position === 0 ? 'This will be generated next' : `#${position}`, inline: false},
`18`		`- {name: 'Prompt', value: codeBlock(queueItem.prediction.prompt), inline: true},`
	`18`	`+ {name: 'Prompt', value: queueItem.prediction.prompt ? codeBlock(queueItem.prediction.prompt) : 'Not Supplied', inline: true},`
`19`	`19`	`{name: 'Owner', value: userMention(queueItem.discordCaller), inline: true},`
`20`	`20`	`{name: 'Prompt ID', value: codeBlock(queueItem.uuid), inline: true},`
`21`	`21`	`]`