Skip to content

Commit 8da002b

Browse files
ZingzyZingzy
authored
Merge pull request #10 from Zingzy/9-feature-request-separate-configurations-into-a-toml-file
Separate Configurations into a TOML File
2 parents c97fb3f + c644dee commit 8da002b

15 files changed

+440
-179
lines changed

‎.env.example

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,2 +1 @@
1-
TOKEN=
2-
MONGODB_URI=
1+
TOKEN=

‎.gitignore

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -176,4 +176,7 @@ cython_debug/
176176
# PyPI configuration file
177177
.pypirc
178178

179+
# Configurations
180+
config.toml
181+
179182
todo.md

‎cogs/imagine_cog.py

Lines changed: 37 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,14 @@
22
import discord
33
from discord import app_commands
44
from discord.ext import commands
5+
import traceback
56

6-
from constants import MODELS
7+
from config import config
78
from utils.image_gen_utils import generate_image, validate_dimensions, validate_prompt
89
from utils.embed_utils import generate_pollinate_embed, generate_error_message
910
from utils.pollinate_utils import parse_url
1011
from utils.error_handler import send_error_embed
1112
from exceptions import DimensionTooSmallError, PromptTooLongError, APIError
12-
import traceback
1313

1414

1515
class ImagineButtonView(discord.ui.View):
@@ -20,7 +20,7 @@ def __init__(self) -> None:
2020
label="Regenerate",
2121
style=discord.ButtonStyle.secondary,
2222
custom_id="regenerate-button",
23-
emoji="<:redo:1187101382101180456>",
23+
emoji=f"<:redo:{config.bot.emojis['redo_emoji_id']}>",
2424
)
2525
async def regenerate(
2626
self, interaction: discord.Interaction, button: discord.ui.Button
@@ -29,7 +29,7 @@ async def regenerate(
2929
embed=discord.Embed(
3030
title="Regenerating Your Image",
3131
description="Please wait while we generate your image",
32-
color=discord.Color.blurple(),
32+
color=int(config.ui.colors.success, 16),
3333
),
3434
ephemeral=True,
3535
)
@@ -49,7 +49,7 @@ async def regenerate(
4949
embed=discord.Embed(
5050
title="Couldn't Generate the Requested Image 😔",
5151
description=f"```\n{e.message}\n```",
52-
color=discord.Color.red(),
52+
color=int(config.ui.colors.error, 16),
5353
),
5454
ephemeral=True,
5555
)
@@ -60,7 +60,7 @@ async def regenerate(
6060
embed=discord.Embed(
6161
title="Error",
6262
description=f"Error generating image : {e}",
63-
color=discord.Color.red(),
63+
color=int(config.ui.colors.error, 16),
6464
),
6565
ephemeral=True,
6666
)
@@ -86,7 +86,7 @@ async def regenerate(
8686
@discord.ui.button(
8787
style=discord.ButtonStyle.red,
8888
custom_id="delete-button",
89-
emoji="<:delete:1187102382312652800>",
89+
emoji=f"<:delete:{config.bot.emojis['delete_emoji_id']}>",
9090
)
9191
async def delete(self, interaction: discord.Interaction, button: discord.ui.Button):
9292
try:
@@ -99,8 +99,8 @@ async def delete(self, interaction: discord.Interaction, button: discord.ui.Butt
9999
await interaction.response.send_message(
100100
embed=discord.Embed(
101101
title="Error",
102-
description="You can only delete the images prompted by you",
103-
color=discord.Color.red(),
102+
description=config.ui.error_messages["delete_unauthorized"],
103+
color=int(config.ui.colors.error, 16),
104104
),
105105
ephemeral=True,
106106
)
@@ -114,7 +114,7 @@ async def delete(self, interaction: discord.Interaction, button: discord.ui.Butt
114114
embed=discord.Embed(
115115
title="Error Deleting the Image",
116116
description=f"{e}",
117-
color=discord.Color.red(),
117+
color=int(config.ui.colors.error, 16),
118118
),
119119
ephemeral=True,
120120
)
@@ -124,7 +124,7 @@ async def delete(self, interaction: discord.Interaction, button: discord.ui.Butt
124124
label="Bookmark",
125125
style=discord.ButtonStyle.secondary,
126126
custom_id="bookmark-button",
127-
emoji="<:save:1187101389822902344>",
127+
emoji=f"<:save:{config.bot.emojis['save_emoji_id']}>",
128128
)
129129
async def bookmark(
130130
self, interaction: discord.Interaction, button: discord.ui.Button
@@ -137,7 +137,7 @@ async def bookmark(
137137

138138
embed: discord.Embed = discord.Embed(
139139
description=f"**Prompt : {prompt}**",
140-
color=discord.Color.og_blurple(),
140+
color=int(config.ui.colors.success, 16),
141141
)
142142
embed.add_field(
143143
name="",
@@ -152,7 +152,7 @@ async def bookmark(
152152
embed=discord.Embed(
153153
title="Image Bookmarked",
154154
description="The image has been bookmarked and sent to your DMs",
155-
color=discord.Color.blurple(),
155+
color=int(config.ui.colors.success, 16),
156156
),
157157
ephemeral=True,
158158
)
@@ -164,7 +164,7 @@ async def bookmark(
164164
embed=discord.Embed(
165165
title="Error Bookmarking the Image",
166166
description=f"{e}",
167-
color=discord.Color.red(),
167+
color=int(config.ui.colors.error, 16),
168168
),
169169
ephemeral=True,
170170
)
@@ -174,17 +174,23 @@ async def bookmark(
174174
class Imagine(commands.Cog):
175175
def __init__(self, bot) -> None:
176176
self.bot = bot
177+
self.command_config = config.commands["pollinate"]
177178

178179
async def cog_load(self) -> None:
179180
await self.bot.wait_until_ready()
180181
self.bot.add_view(ImagineButtonView())
181182

182183
@app_commands.command(name="pollinate", description="Generate AI Images")
183184
@app_commands.choices(
184-
model=[app_commands.Choice(name=choice, value=choice) for choice in MODELS],
185+
model=[
186+
app_commands.Choice(name=choice, value=choice) for choice in config.MODELS
187+
],
185188
)
186189
@app_commands.guild_only()
187-
@app_commands.checks.cooldown(1, 10)
190+
@app_commands.checks.cooldown(
191+
config.commands["pollinate"].cooldown.rate,
192+
config.commands["pollinate"].cooldown.seconds,
193+
)
188194
@app_commands.describe(
189195
prompt="Prompt of the Image you want want to generate",
190196
height="Height of the Image",
@@ -200,22 +206,22 @@ async def imagine_command(
200206
self,
201207
interaction: discord.Interaction,
202208
prompt: str,
203-
width: int = 1000,
204-
height: int = 1000,
205-
model: app_commands.Choice[str] = MODELS[0],
206-
enhance: bool | None = None,
207-
safe: bool = False,
208-
cached: bool = False,
209-
nologo: bool = False,
210-
private: bool = False,
209+
width: int = config.commands["pollinate"].default_width,
210+
height: int = config.commands["pollinate"].default_height,
211+
model: app_commands.Choice[str] = config.MODELS[0],
212+
enhance: bool | None = config.image_generation.defaults.enhance,
213+
safe: bool = config.image_generation.defaults.safe,
214+
cached: bool = config.image_generation.defaults.cached,
215+
nologo: bool = config.image_generation.defaults.nologo,
216+
private: bool = config.image_generation.defaults.private,
211217
) -> None:
212218
validate_dimensions(width, height)
213219
validate_prompt(prompt)
214220

215221
await interaction.response.defer(thinking=True, ephemeral=private)
216222

217223
try:
218-
model = model.value
224+
model = model.value if model else None
219225
except Exception:
220226
pass
221227

@@ -250,7 +256,9 @@ async def imagine_command_error(
250256
embed: discord.Embed = await generate_error_message(
251257
interaction,
252258
error,
253-
cooldown_configuration=["- 1 time every 10 seconds"],
259+
cooldown_configuration=[
260+
f"- {self.command_config.cooldown.rate} time every {self.command_config.cooldown.seconds} seconds",
261+
],
254262
)
255263
return await interaction.response.send_message(embed=embed, ephemeral=True)
256264

@@ -277,7 +285,9 @@ async def imagine_command_error(
277285

278286
else:
279287
await send_error_embed(
280-
interaction, "An unexprected error occurred", f"```\n{str(error)}\n```"
288+
interaction,
289+
"An unexpected error occurred",
290+
f"```\n{str(error)}\n```",
281291
)
282292

283293

‎cogs/multi_pollinate_cog.py

Lines changed: 31 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import traceback
66
import asyncio
77

8+
from config import config
89
from utils.embed_utils import generate_error_message
910
from utils.image_gen_utils import generate_image, validate_dimensions, validate_prompt
1011
from utils.error_handler import send_error_embed
@@ -15,7 +16,6 @@
1516
DimensionTooSmallError,
1617
APIError,
1718
)
18-
from constants import MODELS
1919

2020

2121
class multiImagineButtonView(discord.ui.View):
@@ -39,7 +39,7 @@ def create_buttons(self) -> None:
3939
label="",
4040
style=discord.ButtonStyle.danger,
4141
custom_id="multiimagine_delete",
42-
emoji="<:delete:1187102382312652800>",
42+
emoji=f"<:delete:{config.bot.emojis['delete_emoji_id']}>",
4343
)
4444
)
4545

@@ -76,8 +76,8 @@ async def delete_image(self, interaction: discord.Interaction):
7676
await interaction.response.send_message(
7777
embed=discord.Embed(
7878
title="Error",
79-
description="You can only delete your own images",
80-
color=discord.Color.red(),
79+
description=config.ui.error_messages["delete_unauthorized"],
80+
color=int(config.ui.colors.error, 16),
8181
),
8282
ephemeral=True,
8383
)
@@ -91,7 +91,7 @@ async def delete_image(self, interaction: discord.Interaction):
9191
embed=discord.Embed(
9292
title="Error Deleting the Image",
9393
description=f"{e}",
94-
color=discord.Color.red(),
94+
color=int(config.ui.colors.error, 16),
9595
),
9696
ephemeral=True,
9797
)
@@ -101,6 +101,7 @@ async def delete_image(self, interaction: discord.Interaction):
101101
class Multi_pollinate(commands.Cog):
102102
def __init__(self, bot) -> None:
103103
self.bot = bot
104+
self.command_config = config.commands["multi_pollinate"]
104105

105106
async def cog_load(self) -> None:
106107
await self.bot.wait_until_ready()
@@ -112,7 +113,10 @@ async def get_info(interaction: discord.Interaction, index: int) -> None:
112113
@app_commands.command(
113114
name="multi-pollinate", description="Imagine multiple prompts"
114115
)
115-
@app_commands.checks.cooldown(1, 20)
116+
@app_commands.checks.cooldown(
117+
config.commands["multi_pollinate"].cooldown.rate,
118+
config.commands["multi_pollinate"].cooldown.seconds,
119+
)
116120
@app_commands.guild_only()
117121
@app_commands.describe(
118122
prompt="Prompt of the Image you want want to generate",
@@ -128,25 +132,25 @@ async def multiimagine_command(
128132
self,
129133
interaction: discord.Interaction,
130134
prompt: str,
131-
width: int = 1000,
132-
height: int = 1000,
133-
enhance: bool | None = None,
135+
width: int = config.commands["multi_pollinate"].default_width,
136+
height: int = config.commands["multi_pollinate"].default_height,
137+
enhance: bool | None = config.image_generation.defaults.enhance,
134138
negative: str | None = None,
135-
cached: bool = False,
136-
nologo: bool = False,
137-
private: bool = False,
139+
cached: bool = config.image_generation.defaults.cached,
140+
nologo: bool = config.image_generation.defaults.nologo,
141+
private: bool = config.image_generation.defaults.private,
138142
) -> None:
139143
validate_dimensions(width, height)
140144
validate_prompt(prompt)
141145

142-
total_models: int = len(MODELS)
146+
total_models: int = len(config.MODELS)
143147

144148
await interaction.response.send_message(
145149
embed=discord.Embed(
146150
title="Generating Image",
147151
description=f"Generating images across {total_models} models...\n"
148152
f"Completed: 0/{total_models} 0%",
149-
color=discord.Color.blurple(),
153+
color=int(config.ui.colors.success, 16),
150154
),
151155
ephemeral=private,
152156
)
@@ -178,12 +182,11 @@ async def update_progress() -> None:
178182
description=f"Generating images across {total_models} models...\n"
179183
f"Completed: {completed_count}/{total_models} "
180184
f"({(completed_count / total_models * 100):.2f}%)",
181-
color=discord.Color.blurple(),
185+
color=int(config.ui.colors.success, 16),
182186
)
183187
)
184188

185189
async def generate_for_model(i, model):
186-
"""Asynchronous function to generate an image for a specific model."""
187190
try:
188191
sub_start_time: datetime.datetime = datetime.datetime.now()
189192
dic, image = await generate_image(model=model, **command_args)
@@ -206,10 +209,13 @@ async def generate_for_model(i, model):
206209
try:
207210
results = await asyncio.wait_for(
208211
asyncio.gather(
209-
*[generate_for_model(i, model) for i, model in enumerate(MODELS)],
212+
*[
213+
generate_for_model(i, model)
214+
for i, model in enumerate(config.MODELS)
215+
],
210216
return_exceptions=True,
211217
),
212-
timeout=180,
218+
timeout=self.command_config.timeout_seconds,
213219
)
214220
except asyncio.TimeoutError:
215221
raise asyncio.TimeoutError
@@ -265,15 +271,17 @@ async def multiimagine_command_error(
265271
embed: discord.Embed = await generate_error_message(
266272
interaction,
267273
error,
268-
cooldown_configuration=["- 1 time every 20 seconds"],
274+
cooldown_configuration=[
275+
f"- {self.command_config.cooldown.rate} time every {self.command_config.cooldown.seconds} seconds",
276+
],
269277
)
270278
await interaction.response.send_message(embed=embed, ephemeral=True)
271279

272280
elif isinstance(error, asyncio.TimeoutError):
273281
await send_error_embed(
274282
interaction,
275283
"Timeout Error",
276-
"Image generation took too long and timed out. Please try again.",
284+
config.ui.error_messages["timeout"],
277285
)
278286

279287
elif isinstance(error, NoImagesGeneratedError):
@@ -306,7 +314,9 @@ async def multiimagine_command_error(
306314

307315
else:
308316
await send_error_embed(
309-
interaction, "An unexprected error occurred", f"```\n{str(error)}\n```"
317+
interaction,
318+
config.ui.error_messages["unknown"],
319+
f"```\n{str(error)}\n```",
310320
)
311321

312322

0 commit comments

Comments
 (0)