Skip to content

Commit 280e256

Browse files
p1atdevp1atdev
committed
feat: refresh command
1 parent e83d412 commit 280e256

File tree

5 files changed

+241
-122
lines changed

5 files changed

+241
-122
lines changed

‎bot.ts

Lines changed: 46 additions & 122 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,7 @@
11
import { AUTO1111, ImagineOptions } from "./auto1111.ts"
2-
import {
3-
createBot,
4-
Intents,
5-
startBot,
6-
CreateSlashApplicationCommand,
7-
InteractionResponseTypes,
8-
ApplicationCommandOptionTypes,
9-
} from "./deps.ts"
2+
import { createCommands, registerCommands } from "./command.ts"
3+
import { createBot, Intents, startBot, InteractionResponseTypes } from "./deps.ts"
4+
import { log } from "./log.ts"
105
import { Secret } from "./secret.ts"
116
import { base64ToBlob } from "./utils.ts"
127

@@ -16,12 +11,6 @@ const client = new AUTO1111({
1611
host: Secret.AUTO1111_Host,
1712
})
1813

19-
const [sdModels, samplers, promptStyles] = await Promise.all([
20-
client.sdModels(),
21-
client.samplers(),
22-
client.promptStyles(),
23-
])
24-
2514
const bot = createBot({
2615
token: DISCORD_TOKEN,
2716
intents: Intents.Guilds | Intents.GuildMessages | Intents.MessageContent | Intents.GuildMembers,
@@ -32,118 +21,24 @@ const bot = createBot({
3221
},
3322
})
3423

35-
const switchModelCommand: CreateSlashApplicationCommand = {
36-
name: "switch",
37-
description: "Switch stable diffusion model",
38-
options: [
39-
{
40-
type: ApplicationCommandOptionTypes.String,
41-
name: "name",
42-
description: "Stable diffusion model name",
43-
choices: (sdModels.length > 25 ? sdModels.slice(0, 25) : sdModels).map((model) => ({
44-
name: model.model_name,
45-
value: model.title,
46-
})),
47-
required: true,
48-
},
49-
],
50-
}
24+
const refreshParameters = async () => {
25+
await client.refreshCheckpoints()
5126

52-
const imagineCommand: CreateSlashApplicationCommand = {
53-
name: "imagine",
54-
description: "Thinking to the world of dreams...",
55-
options: [
56-
{
57-
type: ApplicationCommandOptionTypes.String,
58-
name: "prompt",
59-
description: "Positive prompt",
60-
required: true,
61-
},
62-
{
63-
type: ApplicationCommandOptionTypes.String,
64-
name: "negative",
65-
description: "Negative prompt",
66-
required: false,
67-
},
68-
{
69-
type: ApplicationCommandOptionTypes.String,
70-
name: "prompt-style",
71-
description: "Prompt style",
72-
choices: promptStyles.map((style) => ({
73-
name: style.name,
74-
value: style.name,
75-
})),
76-
required: false,
77-
},
78-
{
79-
type: ApplicationCommandOptionTypes.String,
80-
name: "aspect",
81-
description: "Aspect ratio",
82-
choices: [
83-
{
84-
name: "2:3",
85-
value: "2:3",
86-
},
87-
{
88-
name: "1:1",
89-
value: "1:1",
90-
},
91-
{
92-
name: "3:2",
93-
value: "3:2",
94-
},
95-
],
96-
required: false,
97-
},
98-
{
99-
type: ApplicationCommandOptionTypes.Integer,
100-
name: "seed",
101-
description: "Like finding a single shining star in the vastness of space",
102-
required: false,
103-
},
104-
{
105-
type: ApplicationCommandOptionTypes.String,
106-
name: "sampler",
107-
description: "Sampling method",
108-
choices: (samplers.length > 25 ? samplers.slice(0, 25) : samplers).map((sampler) => ({
109-
name: sampler.name,
110-
value: sampler.name,
111-
})),
112-
required: false,
113-
},
114-
{
115-
type: ApplicationCommandOptionTypes.Integer,
116-
name: "steps",
117-
description: "Number of steps",
118-
required: false,
119-
},
120-
{
121-
type: ApplicationCommandOptionTypes.Number,
122-
name: "scale",
123-
description: "CFG scale",
124-
required: false,
125-
},
126-
{
127-
type: ApplicationCommandOptionTypes.Integer,
128-
name: "count",
129-
description: "Number of images to generate",
130-
required: false,
131-
minValue: 1,
132-
maxValue: 4,
133-
},
134-
],
27+
return await Promise.all([client.sdModels(), client.samplers(), client.promptStyles()])
13528
}
13629

137-
const commands = [imagineCommand, switchModelCommand]
30+
const [_sdModels, _samplers, promptStyles] = await refreshParameters()
31+
32+
const refreshCommands = async () => {
33+
const commands = createCommands(...(await refreshParameters()))
34+
await registerCommands(bot, commands, false)
35+
log.info("Commands refreshed")
36+
}
13837

139-
await Promise.all(
140-
commands.map((command) => {
141-
bot.helpers.createGuildApplicationCommand(command, Secret.GUILD_ID)
142-
// bot.helpers.createGlobalApplicationCommand(command)
143-
})
144-
)
145-
await bot.helpers.upsertGuildApplicationCommands(Secret.GUILD_ID, commands)
146-
// await bot.helpers.upsertGlobalApplicationCommands(commands)
38+
bot.events.ready = async () => {
39+
log.success("Successfully connected to gateway")
40+
await refreshCommands()
41+
}
14742

14843
bot.events.interactionCreate = async (b, interaction) => {
14944
switch (interaction.data?.name) {
@@ -163,6 +58,8 @@ bot.events.interactionCreate = async (b, interaction) => {
16358
return
16459
}
16560

61+
log.info(`Switching model from ${auto1111options.sd_model_checkpoint} to ${model?.value}...`)
62+
16663
await b.helpers.sendInteractionResponse(interaction.id, interaction.token, {
16764
type: InteractionResponseTypes.ChannelMessageWithSource,
16865
data: {
@@ -176,6 +73,8 @@ bot.events.interactionCreate = async (b, interaction) => {
17673
content: `✅ Model switched to **${model?.value}** successfully!`,
17774
})
17875

76+
log.info(`Model switched to ${model?.value} successfully!`)
77+
17978
break
18079
}
18180
case "imagine": {
@@ -213,6 +112,8 @@ bot.events.interactionCreate = async (b, interaction) => {
213112
count: interaction.data.options?.find((o) => o.name === "count")?.value as number,
214113
}
215114

115+
log.info("Imagine:", options)
116+
216117
const promptStyleName = interaction.data.options?.find((o) => o.name === "prompt-style")?.value
217118
if (typeof promptStyleName === "string") {
218119
const promptStyle = promptStyles.find((style) => style.name === promptStyleName)
@@ -237,6 +138,9 @@ bot.events.interactionCreate = async (b, interaction) => {
237138
while (!finished) {
238139
await new Promise((resolve) => setTimeout(resolve, 500))
239140
const progress = await client.progress()
141+
if (progress.progress === 0) {
142+
continue
143+
}
240144
await b.helpers.editOriginalInteractionResponse(interaction.token, {
241145
content: `${paramerters}\n${(progress.progress * 100).toFixed(
242146
1
@@ -259,6 +163,26 @@ bot.events.interactionCreate = async (b, interaction) => {
259163

260164
break
261165
}
166+
case "refresh": {
167+
log.info("Refreshing...")
168+
169+
await b.helpers.sendInteractionResponse(interaction.id, interaction.token, {
170+
type: InteractionResponseTypes.ChannelMessageWithSource,
171+
data: {
172+
content: "Refreshing...",
173+
},
174+
})
175+
176+
await refreshCommands()
177+
178+
await b.helpers.editOriginalInteractionResponse(interaction.token, {
179+
content: "✅ Refreshed!",
180+
})
181+
182+
log.info("Refreshed")
183+
184+
break
185+
}
262186
default: {
263187
break
264188
}

‎command.ts

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
1+
import { Bot, CreateSlashApplicationCommand, ApplicationCommandOptionTypes } from "./deps.ts"
2+
import { PromptStyle } from "./types/promptStyle.ts"
3+
import { Sampler } from "./types/sampler.ts"
4+
import { SDModel } from "./types/sdModel.ts"
5+
import { Secret } from "./secret.ts"
6+
7+
const createSwitchModelCommand = (sdModels: SDModel[]) => {
8+
const switchModelCommand: CreateSlashApplicationCommand = {
9+
name: "switch",
10+
description: "Switch stable diffusion model",
11+
options: [
12+
{
13+
type: ApplicationCommandOptionTypes.String,
14+
name: "name",
15+
description: "Stable diffusion model name",
16+
choices: (sdModels.length > 25 ? sdModels.slice(0, 25) : sdModels).map((model) => ({
17+
name: model.model_name,
18+
value: model.title,
19+
})),
20+
required: true,
21+
},
22+
],
23+
}
24+
25+
return switchModelCommand
26+
}
27+
28+
const createImagineCommand = (samplers: Sampler[], promptStyles: PromptStyle[]) => {
29+
const imagineCommand: CreateSlashApplicationCommand = {
30+
name: "imagine",
31+
description: "Thinking to the world of dreams...",
32+
options: [
33+
{
34+
type: ApplicationCommandOptionTypes.String,
35+
name: "prompt",
36+
description: "Positive prompt",
37+
required: true,
38+
},
39+
{
40+
type: ApplicationCommandOptionTypes.String,
41+
name: "negative",
42+
description: "Negative prompt",
43+
required: false,
44+
},
45+
{
46+
type: ApplicationCommandOptionTypes.String,
47+
name: "prompt-style",
48+
description: "Prompt style",
49+
choices: promptStyles.map((style) => ({
50+
name: style.name,
51+
value: style.name,
52+
})),
53+
required: false,
54+
},
55+
{
56+
type: ApplicationCommandOptionTypes.String,
57+
name: "aspect",
58+
description: "Aspect ratio",
59+
choices: [
60+
{
61+
name: "2:3",
62+
value: "2:3",
63+
},
64+
{
65+
name: "1:1",
66+
value: "1:1",
67+
},
68+
{
69+
name: "3:2",
70+
value: "3:2",
71+
},
72+
],
73+
required: false,
74+
},
75+
{
76+
type: ApplicationCommandOptionTypes.Integer,
77+
name: "seed",
78+
description: "Like finding a single shining star in the vastness of space",
79+
required: false,
80+
},
81+
{
82+
type: ApplicationCommandOptionTypes.String,
83+
name: "sampler",
84+
description: "Sampling method",
85+
choices: (samplers.length > 25 ? samplers.slice(0, 25) : samplers).map((sampler) => ({
86+
name: sampler.name,
87+
value: sampler.name,
88+
})),
89+
required: false,
90+
},
91+
{
92+
type: ApplicationCommandOptionTypes.Integer,
93+
name: "steps",
94+
description: "Number of steps",
95+
required: false,
96+
},
97+
{
98+
type: ApplicationCommandOptionTypes.Number,
99+
name: "scale",
100+
description: "CFG scale",
101+
required: false,
102+
},
103+
{
104+
type: ApplicationCommandOptionTypes.Integer,
105+
name: "count",
106+
description: "Number of images to generate",
107+
required: false,
108+
minValue: 1,
109+
maxValue: 4,
110+
},
111+
],
112+
}
113+
114+
return imagineCommand
115+
}
116+
117+
const createRefreshCommand = () => {
118+
const refreshCommand: CreateSlashApplicationCommand = {
119+
name: "refresh",
120+
description: "Refresh stable diffusion model and other parameters",
121+
}
122+
123+
return refreshCommand
124+
}
125+
126+
export const createCommands = (sdModels: SDModel[], samplers: Sampler[], promptStyles: PromptStyle[]) => {
127+
const commands = [
128+
createSwitchModelCommand(sdModels),
129+
createImagineCommand(samplers, promptStyles),
130+
createRefreshCommand(),
131+
]
132+
133+
return commands
134+
}
135+
136+
export const registerCommands = async (bot: Bot, commands: CreateSlashApplicationCommand[], global: boolean) => {
137+
const current = global
138+
? await bot.helpers.getGlobalApplicationCommands()
139+
: await bot.helpers.getGuildApplicationCommands(Secret.GUILD_ID)
140+
current.forEach(async (command) => {
141+
if (global) {
142+
await bot.helpers.deleteGlobalApplicationCommand(command.id)
143+
} else {
144+
await bot.helpers.deleteGuildApplicationCommand(command.id, Secret.GUILD_ID)
145+
}
146+
})
147+
148+
await Promise.all(
149+
commands.map(async (command) => {
150+
if (global) {
151+
// await bot.helpers.deleteGlobalApplicationCommand(command.name)
152+
await bot.helpers.createGlobalApplicationCommand(command)
153+
} else {
154+
// await bot.helpers.deleteGuildApplicationCommand(command.name, Secret.GUILD_ID)
155+
await bot.helpers.createGuildApplicationCommand(command, Secret.GUILD_ID)
156+
}
157+
})
158+
)
159+
160+
if (global) {
161+
await bot.helpers.upsertGlobalApplicationCommands(commands)
162+
} else {
163+
await bot.helpers.upsertGuildApplicationCommands(Secret.GUILD_ID, commands)
164+
}
165+
}

0 commit comments

Comments
 (0)