Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Migrate QueueItems to class structure #1

Draft
wants to merge 4 commits into
base: main
Choose a base branch
from
Draft

Migrate QueueItems to class structure #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
52298e6
Migrate over default QueueItem to QueueItem class, update processor, …
maxwlang Oct 30, 2022
701bf69
commit
maxwlang Oct 31, 2022
477d72c
Implementation progress
maxwlang Nov 24, 2022
d38c14e
Additional implementation progress
maxwlang Nov 24, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 58 additions & 18 deletions src/bot.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
import { BotCommand, BotEvent, QueueItem } from './types'
import bPromise from 'bluebird'
import { Client, GatewayIntentBits, Partials, REST, Routes } from 'discord.js'
import fs from 'fs'
import { isEmpty, isNil } from 'ramda'

import Config from './config/bot'
import { Logger } from 'winston'
import Config from './config/bot'
import StableDiffusion from './modules/stableDiffusion'
import bPromise from 'bluebird'
import fs from 'fs'
import { BotCommand, BotEvent, QueueItems } from './types'


export class Bot extends Client {
constructor(config: typeof Config, logger: Logger, options: {[k: string]: string}) {
Expand All @@ -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
Expand Down Expand Up @@ -71,14 +71,54 @@ export class Bot extends Client {
* @param queueItem QueueItem to add to queue.
* @returns Number of items in queue.
*/
public addQueue = (queueItem: QueueItem): number => this.queue.push(queueItem)
public addQueuedQueueItem = (queueItem: QueueItems.QueueItemInstances): number => this.queue.push(queueItem)

/**
* Updates content in a queue or reference QueueItem
* @param queueItem QueueItem to add to queue.
* @returns Dictionary of updated QueueItems.
*/
public updateQueueItem = (
updateFunction: (queueItem: QueueItems.QueueItemInstances) => QueueItems.QueueItemInstances,
uuid: string
): {
queue: QueueItems.QueueItemInstances | undefined,
referenceQueue: QueueItems.QueueItemInstances | undefined
} => {
const queueIndex = this.queue.findIndex(queueItem => queueItem.uuid === uuid)
const referenceQueueIndex = this.queueItemReferences.findIndex(queueItem => queueItem.uuid === uuid)

if (queueIndex === -1 && referenceQueueIndex === -1) throw new Error('No QueueItem in queue or queue references to update.')

const updated: {
queue: QueueItems.QueueItemInstances | undefined,
referenceQueue: QueueItems.QueueItemInstances | undefined
} = {
queue: undefined,
referenceQueue: undefined
}

if (queueIndex !== -1) {
const queueItem = this.queue[queueIndex]
this.queue[queueIndex] = updateFunction(queueItem)
updated.queue = this.queue[queueIndex]
}

if (referenceQueueIndex !== -1) {
const referenceQueueItem = this.queueItemReferences[referenceQueueIndex]
this.queueItemReferences[referenceQueueIndex] = updateFunction(referenceQueueItem)
updated.referenceQueue = this.queueItemReferences[referenceQueueIndex]
}

return updated
}

/**
* Deletes a QueueItem from queue, adds it to queueItemReferences. Throws if it can't find a QueueItem for uuid.
* @param uuid UUID of QueueItem.
* @returns Array of deleted QueueItems.
*/
public removeQueue = (uuid: string): QueueItem[] => {
public removeQueuedQueueItem = (uuid: string): QueueItems.QueueItemInstances[] => {
const queueIndex = this.queue.findIndex(queueItem => queueItem.uuid === uuid)

if (queueIndex === -1) return []
Expand All @@ -92,7 +132,7 @@ export class Bot extends Client {
* @param uuid UUID of QueueItem.
* @returns Array of deleted QueueItems.
*/
public removeQueueItemReference = (uuid: string): QueueItem[] => {
public removeQueueItemReference = (uuid: string): QueueItems.QueueItemInstances[] => {
const queueIndex = this.queueItemReferences.findIndex(queueItem => queueItem.uuid === uuid)

if (queueIndex === -1) return []
Expand All @@ -102,28 +142,28 @@ export class Bot extends Client {
/**
* Finds a QueueItem in queue storage by uuid.
* @param uuid UUID of QueueItem.
* @returns Array of found QueueItems.
* @returns Found QueueItem.
*/
public findQueue = (uuid: string): QueueItem | undefined => {
public findQueue = (uuid: string): QueueItems.QueueItemInstances | undefined => {
return this.queue.find(queueItem => queueItem.uuid === uuid)
}

/**
* Finds a QueueItem in reference storage by uuid.
* @param uuid UUID of QueueItem.
* @returns Array of found QueueItems.
* @returns Found QueueItem.
*/
public findQueueItemReference = (uuid: string): QueueItem | undefined => {
public findQueueItemReference = (uuid: string): QueueItems.QueueItemInstances | undefined => {
return this.queueItemReferences.find(queueItem => queueItem.uuid === uuid)
}

/**
* Finds the latest QueueItem in reference storage by message ID.
* Finds the latest QueueItem in reference storage by message snowflake.
* @param uuid UUID of QueueItem.
* @returns Array of found QueueItems.
*/
public findLatestQueueItemReferenceByMessageID = (messageId: string): QueueItem | undefined => {
const queueItems = this.queueItemReferences.filter(queueItem => queueItem.messageId === messageId)
public findLatestQueueItemReferenceByMessageSnowflake = (snowflake: string): QueueItems.QueueItemInstances | undefined => {
const queueItems = this.queueItemReferences.filter(queueItem => queueItem.discordMessageSnowflake === snowflake)
return queueItems[queueItems.length - 1]
}
}
Expand Down
104 changes: 51 additions & 53 deletions src/embeds/addedToQueue.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,64 +1,62 @@
import { EmbedBuilder, codeBlock, ActionRowBuilder } from 'discord.js'
import { BotEmbed, QueueItem } from '../types'
import { BotEmbed, QueueItems } from '../types'

export enum QueueType {
Instant,
Queued
}

export default function(queueType: QueueType, queueItem: QueueItem, queuePosition?: number): BotEmbed {
export function addedToInstantQueue(queueItem: QueueItems.QueueItemInstances): BotEmbed {
let embeds: EmbedBuilder[] = []
let components: ActionRowBuilder<any>[] = []

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<any>[] = []

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
Expand Down
Loading