diff --git a/.env.example b/.env.example index 84db104..f74c661 100644 --- a/.env.example +++ b/.env.example @@ -1,2 +1 @@ -TOKEN= -MONGODB_URI= \ No newline at end of file +TOKEN= \ No newline at end of file diff --git a/.gitignore b/.gitignore index ad0cd5b..02db827 100644 --- a/.gitignore +++ b/.gitignore @@ -176,4 +176,7 @@ cython_debug/ # PyPI configuration file .pypirc +# Configurations +config.toml + todo.md \ No newline at end of file diff --git a/cogs/imagine_cog.py b/cogs/imagine_cog.py index ec7e453..5d3a6df 100644 --- a/cogs/imagine_cog.py +++ b/cogs/imagine_cog.py @@ -2,14 +2,14 @@ import discord from discord import app_commands from discord.ext import commands +import traceback -from constants import MODELS +from config import config from utils.image_gen_utils import generate_image, validate_dimensions, validate_prompt from utils.embed_utils import generate_pollinate_embed, generate_error_message from utils.pollinate_utils import parse_url from utils.error_handler import send_error_embed from exceptions import DimensionTooSmallError, PromptTooLongError, APIError -import traceback class ImagineButtonView(discord.ui.View): @@ -20,7 +20,7 @@ def __init__(self) -> None: label="Regenerate", style=discord.ButtonStyle.secondary, custom_id="regenerate-button", - emoji="<:redo:1187101382101180456>", + emoji=f"<:redo:{config.bot.emojis['redo_emoji_id']}>", ) async def regenerate( self, interaction: discord.Interaction, button: discord.ui.Button @@ -29,7 +29,7 @@ async def regenerate( embed=discord.Embed( title="Regenerating Your Image", description="Please wait while we generate your image", - color=discord.Color.blurple(), + color=int(config.ui.colors.success, 16), ), ephemeral=True, ) @@ -49,7 +49,7 @@ async def regenerate( embed=discord.Embed( title="Couldn't Generate the Requested Image 😔", description=f"```\n{e.message}\n```", - color=discord.Color.red(), + color=int(config.ui.colors.error, 16), ), ephemeral=True, ) @@ -60,7 +60,7 @@ async def regenerate( embed=discord.Embed( title="Error", description=f"Error generating image : {e}", - color=discord.Color.red(), + color=int(config.ui.colors.error, 16), ), ephemeral=True, ) @@ -86,7 +86,7 @@ async def regenerate( @discord.ui.button( style=discord.ButtonStyle.red, custom_id="delete-button", - emoji="<:delete:1187102382312652800>", + emoji=f"<:delete:{config.bot.emojis['delete_emoji_id']}>", ) async def delete(self, interaction: discord.Interaction, button: discord.ui.Button): try: @@ -99,8 +99,8 @@ async def delete(self, interaction: discord.Interaction, button: discord.ui.Butt await interaction.response.send_message( embed=discord.Embed( title="Error", - description="You can only delete the images prompted by you", - color=discord.Color.red(), + description=config.ui.error_messages["delete_unauthorized"], + color=int(config.ui.colors.error, 16), ), ephemeral=True, ) @@ -114,7 +114,7 @@ async def delete(self, interaction: discord.Interaction, button: discord.ui.Butt embed=discord.Embed( title="Error Deleting the Image", description=f"{e}", - color=discord.Color.red(), + color=int(config.ui.colors.error, 16), ), ephemeral=True, ) @@ -124,7 +124,7 @@ async def delete(self, interaction: discord.Interaction, button: discord.ui.Butt label="Bookmark", style=discord.ButtonStyle.secondary, custom_id="bookmark-button", - emoji="<:save:1187101389822902344>", + emoji=f"<:save:{config.bot.emojis['save_emoji_id']}>", ) async def bookmark( self, interaction: discord.Interaction, button: discord.ui.Button @@ -137,7 +137,7 @@ async def bookmark( embed: discord.Embed = discord.Embed( description=f"**Prompt : {prompt}**", - color=discord.Color.og_blurple(), + color=int(config.ui.colors.success, 16), ) embed.add_field( name="", @@ -152,7 +152,7 @@ async def bookmark( embed=discord.Embed( title="Image Bookmarked", description="The image has been bookmarked and sent to your DMs", - color=discord.Color.blurple(), + color=int(config.ui.colors.success, 16), ), ephemeral=True, ) @@ -164,7 +164,7 @@ async def bookmark( embed=discord.Embed( title="Error Bookmarking the Image", description=f"{e}", - color=discord.Color.red(), + color=int(config.ui.colors.error, 16), ), ephemeral=True, ) @@ -174,6 +174,7 @@ async def bookmark( class Imagine(commands.Cog): def __init__(self, bot) -> None: self.bot = bot + self.command_config = config.commands["pollinate"] async def cog_load(self) -> None: await self.bot.wait_until_ready() @@ -181,10 +182,15 @@ async def cog_load(self) -> None: @app_commands.command(name="pollinate", description="Generate AI Images") @app_commands.choices( - model=[app_commands.Choice(name=choice, value=choice) for choice in MODELS], + model=[ + app_commands.Choice(name=choice, value=choice) for choice in config.MODELS + ], ) @app_commands.guild_only() - @app_commands.checks.cooldown(1, 10) + @app_commands.checks.cooldown( + config.commands["pollinate"].cooldown.rate, + config.commands["pollinate"].cooldown.seconds, + ) @app_commands.describe( prompt="Prompt of the Image you want want to generate", height="Height of the Image", @@ -200,14 +206,14 @@ async def imagine_command( self, interaction: discord.Interaction, prompt: str, - width: int = 1000, - height: int = 1000, - model: app_commands.Choice[str] = MODELS[0], - enhance: bool | None = None, - safe: bool = False, - cached: bool = False, - nologo: bool = False, - private: bool = False, + width: int = config.commands["pollinate"].default_width, + height: int = config.commands["pollinate"].default_height, + model: app_commands.Choice[str] = config.MODELS[0], + enhance: bool | None = config.image_generation.defaults.enhance, + safe: bool = config.image_generation.defaults.safe, + cached: bool = config.image_generation.defaults.cached, + nologo: bool = config.image_generation.defaults.nologo, + private: bool = config.image_generation.defaults.private, ) -> None: validate_dimensions(width, height) validate_prompt(prompt) @@ -215,7 +221,7 @@ async def imagine_command( await interaction.response.defer(thinking=True, ephemeral=private) try: - model = model.value + model = model.value if model else None except Exception: pass @@ -250,7 +256,9 @@ async def imagine_command_error( embed: discord.Embed = await generate_error_message( interaction, error, - cooldown_configuration=["- 1 time every 10 seconds"], + cooldown_configuration=[ + f"- {self.command_config.cooldown.rate} time every {self.command_config.cooldown.seconds} seconds", + ], ) return await interaction.response.send_message(embed=embed, ephemeral=True) @@ -277,7 +285,9 @@ async def imagine_command_error( else: await send_error_embed( - interaction, "An unexprected error occurred", f"```\n{str(error)}\n```" + interaction, + "An unexpected error occurred", + f"```\n{str(error)}\n```", ) diff --git a/cogs/multi_pollinate_cog.py b/cogs/multi_pollinate_cog.py index ec7be91..c33c31f 100644 --- a/cogs/multi_pollinate_cog.py +++ b/cogs/multi_pollinate_cog.py @@ -5,6 +5,7 @@ import traceback import asyncio +from config import config from utils.embed_utils import generate_error_message from utils.image_gen_utils import generate_image, validate_dimensions, validate_prompt from utils.error_handler import send_error_embed @@ -15,7 +16,6 @@ DimensionTooSmallError, APIError, ) -from constants import MODELS class multiImagineButtonView(discord.ui.View): @@ -39,7 +39,7 @@ def create_buttons(self) -> None: label="", style=discord.ButtonStyle.danger, custom_id="multiimagine_delete", - emoji="<:delete:1187102382312652800>", + emoji=f"<:delete:{config.bot.emojis['delete_emoji_id']}>", ) ) @@ -76,8 +76,8 @@ async def delete_image(self, interaction: discord.Interaction): await interaction.response.send_message( embed=discord.Embed( title="Error", - description="You can only delete your own images", - color=discord.Color.red(), + description=config.ui.error_messages["delete_unauthorized"], + color=int(config.ui.colors.error, 16), ), ephemeral=True, ) @@ -91,7 +91,7 @@ async def delete_image(self, interaction: discord.Interaction): embed=discord.Embed( title="Error Deleting the Image", description=f"{e}", - color=discord.Color.red(), + color=int(config.ui.colors.error, 16), ), ephemeral=True, ) @@ -101,6 +101,7 @@ async def delete_image(self, interaction: discord.Interaction): class Multi_pollinate(commands.Cog): def __init__(self, bot) -> None: self.bot = bot + self.command_config = config.commands["multi_pollinate"] async def cog_load(self) -> None: await self.bot.wait_until_ready() @@ -112,7 +113,10 @@ async def get_info(interaction: discord.Interaction, index: int) -> None: @app_commands.command( name="multi-pollinate", description="Imagine multiple prompts" ) - @app_commands.checks.cooldown(1, 20) + @app_commands.checks.cooldown( + config.commands["multi_pollinate"].cooldown.rate, + config.commands["multi_pollinate"].cooldown.seconds, + ) @app_commands.guild_only() @app_commands.describe( prompt="Prompt of the Image you want want to generate", @@ -128,25 +132,25 @@ async def multiimagine_command( self, interaction: discord.Interaction, prompt: str, - width: int = 1000, - height: int = 1000, - enhance: bool | None = None, + width: int = config.commands["multi_pollinate"].default_width, + height: int = config.commands["multi_pollinate"].default_height, + enhance: bool | None = config.image_generation.defaults.enhance, negative: str | None = None, - cached: bool = False, - nologo: bool = False, - private: bool = False, + cached: bool = config.image_generation.defaults.cached, + nologo: bool = config.image_generation.defaults.nologo, + private: bool = config.image_generation.defaults.private, ) -> None: validate_dimensions(width, height) validate_prompt(prompt) - total_models: int = len(MODELS) + total_models: int = len(config.MODELS) await interaction.response.send_message( embed=discord.Embed( title="Generating Image", description=f"Generating images across {total_models} models...\n" f"Completed: 0/{total_models} 0%", - color=discord.Color.blurple(), + color=int(config.ui.colors.success, 16), ), ephemeral=private, ) @@ -178,12 +182,11 @@ async def update_progress() -> None: description=f"Generating images across {total_models} models...\n" f"Completed: {completed_count}/{total_models} " f"({(completed_count / total_models * 100):.2f}%)", - color=discord.Color.blurple(), + color=int(config.ui.colors.success, 16), ) ) async def generate_for_model(i, model): - """Asynchronous function to generate an image for a specific model.""" try: sub_start_time: datetime.datetime = datetime.datetime.now() dic, image = await generate_image(model=model, **command_args) @@ -206,10 +209,13 @@ async def generate_for_model(i, model): try: results = await asyncio.wait_for( asyncio.gather( - *[generate_for_model(i, model) for i, model in enumerate(MODELS)], + *[ + generate_for_model(i, model) + for i, model in enumerate(config.MODELS) + ], return_exceptions=True, ), - timeout=180, + timeout=self.command_config.timeout_seconds, ) except asyncio.TimeoutError: raise asyncio.TimeoutError @@ -265,7 +271,9 @@ async def multiimagine_command_error( embed: discord.Embed = await generate_error_message( interaction, error, - cooldown_configuration=["- 1 time every 20 seconds"], + cooldown_configuration=[ + f"- {self.command_config.cooldown.rate} time every {self.command_config.cooldown.seconds} seconds", + ], ) await interaction.response.send_message(embed=embed, ephemeral=True) @@ -273,7 +281,7 @@ async def multiimagine_command_error( await send_error_embed( interaction, "Timeout Error", - "Image generation took too long and timed out. Please try again.", + config.ui.error_messages["timeout"], ) elif isinstance(error, NoImagesGeneratedError): @@ -306,7 +314,9 @@ async def multiimagine_command_error( else: await send_error_embed( - interaction, "An unexprected error occurred", f"```\n{str(error)}\n```" + interaction, + config.ui.error_messages["unknown"], + f"```\n{str(error)}\n```", ) diff --git a/cogs/random_cog.py b/cogs/random_cog.py index 757d138..c8e660e 100644 --- a/cogs/random_cog.py +++ b/cogs/random_cog.py @@ -3,23 +3,29 @@ from discord import app_commands from discord.ext import commands +from config import config from utils.image_gen_utils import generate_image, validate_dimensions from utils.embed_utils import generate_error_message from utils.error_handler import send_error_embed from exceptions import DimensionTooSmallError, APIError -from constants import MODELS class RandomImage(commands.Cog): def __init__(self, bot) -> None: self.bot = bot + self.command_config = config.commands["random"] @app_commands.command(name="random", description="Generate Random AI Images") @app_commands.choices( - model=[app_commands.Choice(name=choice, value=choice) for choice in MODELS], + model=[ + app_commands.Choice(name=choice, value=choice) for choice in config.MODELS + ], ) @app_commands.guild_only() - @app_commands.checks.cooldown(1, 10) + @app_commands.checks.cooldown( + config.commands["random"].cooldown.rate, + config.commands["random"].cooldown.seconds, + ) @app_commands.describe( height="Height of the image", width="Width of the image", @@ -31,19 +37,19 @@ def __init__(self, bot) -> None: async def random_image_command( self, interaction: discord.Interaction, - width: int = 1000, - height: int = 1000, - model: app_commands.Choice[str] = MODELS[0], + width: int = config.commands["random"].default_width, + height: int = config.commands["random"].default_height, + model: app_commands.Choice[str] = config.MODELS[0], negative: str | None = None, - nologo: bool = False, - private: bool = False, + nologo: bool = config.image_generation.defaults.nologo, + private: bool = config.image_generation.defaults.private, ) -> None: validate_dimensions(width, height) await interaction.response.defer(thinking=True, ephemeral=private) try: - model = model.value + model = model.value if model else None except Exception: pass @@ -75,6 +81,7 @@ async def random_image_command( else "", timestamp=datetime.datetime.now(datetime.timezone.utc), url=dic["url"], + color=int(config.ui.colors.success, 16), ) embed.add_field(name="Seed", value=f"```{dic['seed']}```", inline=True) @@ -103,7 +110,9 @@ async def random_image_command_error( embed: discord.Embed = await generate_error_message( interaction, error, - cooldown_configuration=["- 1 time every 10 seconds"], + cooldown_configuration=[ + f"- {self.command_config.cooldown.rate} time every {self.command_config.cooldown.seconds} seconds", + ], ) await interaction.response.send_message(embed=embed, ephemeral=True) @@ -123,7 +132,9 @@ async def random_image_command_error( else: await send_error_embed( - interaction, "An unexprected error occurred", f"```\n{str(error)}\n```" + interaction, + config.ui.error_messages["unknown"], + f"```\n{str(error)}\n```", ) diff --git a/config.py b/config.py new file mode 100644 index 0000000..5c0b927 --- /dev/null +++ b/config.py @@ -0,0 +1,131 @@ +from pydantic import BaseModel, model_validator +import tomli +from typing import List, Dict, Optional +from pathlib import Path +import sys + + +class BotConfig(BaseModel): + command_prefix: str + bot_id: str + avatar_url: str + commands: Dict[str, str] + emojis: Dict[str, str] + + +class APIConfig(BaseModel): + models_list_endpoint: str + image_gen_endpoint: str + models_refresh_interval_minutes: int + max_timeout_seconds: int + + +class ImageGenerationDefaults(BaseModel): + width: int + height: int + safe: bool + cached: bool + nologo: bool + enhance: bool + private: bool + + +class ImageGenerationValidation(BaseModel): + min_width: int + min_height: int + max_prompt_length: int + max_enhanced_prompt_length: int = 80 + + +class CommandCooldown(BaseModel): + rate: int + seconds: int + per_minute: Optional[int] = None + per_day: Optional[int] = None + + +class CommandConfig(BaseModel): + cooldown: CommandCooldown + default_width: int + default_height: int + timeout_seconds: Optional[int] = None + max_prompt_length: Optional[int] = None + + +class UIColors(BaseModel): + success: str + error: str + warning: str + + +class UIConfig(BaseModel): + bot_invite_url: str + support_server_url: str + github_repo_url: str + api_provider_url: str + bot_creator_avatar: str + colors: UIColors + error_messages: Dict[str, str] + + +class ResourcesConfig(BaseModel): + waiting_gifs: List[str] + + +class ImageGenerationConfig(BaseModel): + referer: str + fallback_model: str + defaults: ImageGenerationDefaults + validation: ImageGenerationValidation + + +class Config(BaseModel): + bot: BotConfig + api: APIConfig + image_generation: ImageGenerationConfig + commands: Dict[str, CommandConfig] + ui: UIConfig + resources: ResourcesConfig + MODELS: List[str] = [] # Initialize with empty list as default + + @model_validator(mode="after") + def validate_structure(self): + required_commands = {"pollinate", "multi_pollinate", "random"} + if not all(cmd in self.commands for cmd in required_commands): + missing = required_commands - self.commands.keys() + raise ValueError(f"Missing required commands: {missing}") + return self + + +def load_config(path: str = "config.toml") -> Config: + """Load and validate config from TOML file""" + config_path = Path(path) + if not config_path.exists(): + raise FileNotFoundError(f"Config file not found at {config_path}") + + with open(config_path, "rb") as f: + config_data = tomli.load(f) + + return Config(**config_data) + + +def initialize_models(config_instance: Config) -> List[str]: + """Pre-initialize models list by fetching from the API""" + import requests + + try: + response = requests.get(config_instance.api.models_list_endpoint) + if response.ok: + return response.json() + except Exception as e: + print(f"Error pre-initializing models: {e}", file=sys.stderr) + return [config_instance.image_generation.fallback_model] + + +# Load config on import +try: + config: Config = load_config() + # Pre-initialize models list + config.MODELS = initialize_models(config) +except Exception as e: + raise RuntimeError(f"Failed to load config: {e}") from e diff --git a/config.template.toml b/config.template.toml new file mode 100644 index 0000000..392f43a --- /dev/null +++ b/config.template.toml @@ -0,0 +1,100 @@ +[bot] +command_prefix = "!" +bot_id = "" +avatar_url = "" + +[bot.commands] # Command IDs of the bot (generated after commands are registered) +pollinate_id = "" +multi_pollinate_id = "" +random_id = "" +help_id = "" +invite_id = "" +about_id = "" + +[bot.emojis] # emojis used in the bot +github_emoji_id = "" +redo_emoji_id = "" +delete_emoji_id = "" +save_emoji_id = "" + +[api] # API endpoints and configurations +models_list_endpoint = "https://image.pollinations.ai/models" +image_gen_endpoint = "https://image.pollinations.ai/prompt" +models_refresh_interval_minutes = 5 +max_timeout_seconds = 120 + +[image_generation] +referer = "discordbot" +fallback_model = "flux" + +[image_generation.defaults] +width = 1000 +height = 1000 +safe = false +cached = false +nologo = false +enhance = true +private = false + +[image_generation.validation] +min_width = 16 +min_height = 16 +max_prompt_length = 2000 +max_enhanced_prompt_length = 80 + +[commands.pollinate] +default_width = 1000 +default_height = 1000 + +[commands.pollinate.cooldown] +rate = 1 +seconds = 10 + +[commands.multi_pollinate] +default_width = 1000 +default_height = 1000 +timeout_seconds = 180 +max_prompt_length = 4000 + +[commands.multi_pollinate.cooldown] +rate = 1 +seconds = 20 + +[commands.random] +default_width = 1000 +default_height = 1000 + +[commands.random.cooldown] +rate = 1 +seconds = 10 + +[ui] +bot_invite_url = "" +support_server_url = "" +github_repo_url = "" +api_provider_url = "https://pollinations.ai/" +bot_creator_avatar = "" + +[ui.colors] +success = "0x7289da" # og_blurple +error = "0xe74c3c" # red +warning = "0xf1c40f" # yellow + +[ui.error_messages] # Default error messages +dimension_too_small = "Width and Height must be greater than 16" +prompt_too_long = "Prompt must be less than 2000 characters" +delete_unauthorized = "You can only delete your own images" +rate_limit = "Rate limit exceeded. Please try again later" +resource_not_found = "The requested resource was not found" +timeout = "Image generation took too long and timed out. Please try again." +unknown = "An unexpected error occurred" + +[resources] +waiting_gifs = [ + "https://media3.giphy.com/media/l0HlBO7eyXzSZkJri/giphy.gif?cid=ecf05e475p246q1gdcu96b5mkqlqvuapb7xay2hywmki7f5q&ep=v1_gifs_search&rid=giphy.gif&ct=g", + "https://media2.giphy.com/media/QBd2kLB5qDmysEXre9/giphy.gif?cid=ecf05e47ha6xwa7rq38dcst49nefabwwrods631hvz67ptfg&ep=v1_gifs_search&rid=giphy.gif&ct=g", + "https://media2.giphy.com/media/ZgqJGwh2tLj5C/giphy.gif?cid=ecf05e47gflyso481izbdcrw7y8okfkgdxgc7zoh34q9rxim&ep=v1_gifs_search&rid=giphy.gif&ct=g", + "https://media0.giphy.com/media/EWhLjxjiqdZjW/giphy.gif?cid=ecf05e473fifxe2bg4act0zq73nkyjw0h69fxi52t8jt37lf&ep=v1_gifs_search&rid=giphy.gif&ct=g", + "https://i.giphy.com/26BRuo6sLetdllPAQ.webp", + "https://i.giphy.com/tXL4FHPSnVJ0A.gif" +] \ No newline at end of file diff --git a/constants.py b/constants.py deleted file mode 100644 index 6682e2a..0000000 --- a/constants.py +++ /dev/null @@ -1,21 +0,0 @@ -import os -from dotenv import load_dotenv -import requests -import json - -load_dotenv(override=True) - -TOKEN: str = os.environ["TOKEN"] -MONGODB_URI: str = os.environ["MONGODB_URI"] - -r: requests.Response = requests.get("https://image.pollinations.ai/models") -MODELS: list[str] = json.loads(r.text) - -WAITING_GIFS: list[str] = [ - "https://media3.giphy.com/media/l0HlBO7eyXzSZkJri/giphy.gif?cid=ecf05e475p246q1gdcu96b5mkqlqvuapb7xay2hywmki7f5q&ep=v1_gifs_search&rid=giphy.gif&ct=g", - "https://media2.giphy.com/media/QBd2kLB5qDmysEXre9/giphy.gif?cid=ecf05e47ha6xwa7rq38dcst49nefabwwrods631hvz67ptfg&ep=v1_gifs_search&rid=giphy.gif&ct=g", - "https://media2.giphy.com/media/ZgqJGwh2tLj5C/giphy.gif?cid=ecf05e47gflyso481izbdcrw7y8okfkgdxgc7zoh34q9rxim&ep=v1_gifs_search&rid=giphy.gif&ct=g", - "https://media0.giphy.com/media/EWhLjxjiqdZjW/giphy.gif?cid=ecf05e473fifxe2bg4act0zq73nkyjw0h69fxi52t8jt37lf&ep=v1_gifs_search&rid=giphy.gif&ct=g", - "https://i.giphy.com/26BRuo6sLetdllPAQ.webp", - "https://i.giphy.com/tXL4FHPSnVJ0A.gif", -] diff --git a/main.py b/main.py index f4f54c6..4fd2255 100644 --- a/main.py +++ b/main.py @@ -6,17 +6,17 @@ import statistics import time import sys -import requests -import json -from constants import MODELS, TOKEN +import aiohttp +from config import config -load_dotenv() +load_dotenv(override=True) +TOKEN: str = os.environ["TOKEN"] start_time = None latencies: list = [] commands_: dict[str, str] = { - " 🎨": """Generates AI Images based on your prompts + f" 🎨": """Generates AI Images based on your prompts - **prompt** đŸ—Ŗī¸ : Your prompt for the Image to be generated - **width** ↔ī¸ : The width of your prompted Image - **height** ↕ī¸ : The height of your prompted Image @@ -26,7 +26,7 @@ - **nologo** đŸšĢ : Specifies whether to remove the logo from the generated images (deafault False) - **private** 🔒 : when set to True the generated Image will only be visible to you """, - " 🎨": """Generates AI Images using all available models + f" 🎨": """Generates AI Images using all available models - **prompt** đŸ—Ŗī¸ : Your prompt for the Image to be generated - **width** ↔ī¸ : The width of your prompted Image - **height** ↕ī¸ : The height of your prompted Image @@ -36,17 +36,16 @@ - **enhance** đŸ–ŧī¸ : Specifies whether to enhance the image prompt or not (default True) - **private** 🔒 : when set to True the generated Image will only be visible to you """, - " 🎨": """Generates Random AI Images + f" 🎨": """Generates Random AI Images - **width** ↔ī¸ : The width of your prompted Image - **height** ↕ī¸ : The height of your prompted Image - **negative** ❎ : Specifies what not to be in the generated images - **nologo** đŸšĢ : Specifies whether to remove the logo from the generated images (deafault False) - **private** 🔒 : when set to True the generated Image will only be visible to you """, - " 🏆": "Shows the Global Leaderboard", - " ❓": "Displays this", - " 📨": "Invite the bot to your server", - " ℹī¸": "About the bot", + f" ❓": "Displays this", + f" 📨": "Invite the bot to your server", + f" ℹī¸": "About the bot", } @@ -56,19 +55,22 @@ def __init__(self) -> None: intents.messages = True intents.message_content = True - super().__init__(command_prefix="!", intents=intents, help_command=None) + super().__init__( + command_prefix=config.bot.command_prefix, intents=intents, help_command=None + ) self.synced = False - @tasks.loop(minutes=5) + @tasks.loop(minutes=config.api.models_refresh_interval_minutes) async def refresh_models(self) -> None: try: - response = requests.get("https://image.pollinations.ai/models") - if response.ok: - global MODELS - MODELS.clear() - MODELS.extend(json.loads(response.text)) - print(f"Models refreshed: {MODELS}") + async with aiohttp.ClientSession() as session: + async with session.get(config.api.models_list_endpoint) as response: + if response.ok: + config.MODELS.clear() + config.MODELS.extend(await response.json()) + print(f"Models refreshed: {config.MODELS}") except Exception as e: + config.MODELS = [config.image_generation.fallback_model] print(f"Error refreshing models: {e}", file=sys.stdout) async def on_ready(self) -> None: @@ -91,7 +93,7 @@ async def on_ready(self) -> None: print(f"Logged in as {self.user.name} (ID: {self.user.id})") print(f"Connected to {len(self.guilds)} guilds") - print(f"Available MODELS: {MODELS}") + print(f"Available MODELS: {config.MODELS}") bot = pollinationsBot() @@ -108,14 +110,13 @@ async def on_message(message) -> None: if message.author == bot.user: return - if bot.user in message.mentions: - if message.type is not discord.MessageType.reply: - embed = discord.Embed( - description="Hello, I am the Pollinations.ai Bot. I am here to help you with your AI needs. **To Generate Images click or or type `/help` for more commands**.", - color=discord.Color.og_blurple(), - ) + if bot.user in message.mentions and message.type is not discord.MessageType.reply: + embed = discord.Embed( + description=f"Hello, I am the Pollinations.ai Bot. I am here to help you with your AI needs. **To Generate Images click or or type `/help` for more commands**.", + color=int(config.ui.colors.success, 16), + ) - await message.reply(embed=embed) + await message.reply(embed=embed) await bot.process_commands(message) @@ -150,7 +151,7 @@ async def before_invoke(ctx) -> None: @bot.command() async def ping(ctx) -> None: try: - embed = discord.Embed(title="Pong!", color=discord.Color.green()) + embed = discord.Embed(title="Pong!", color=int(config.ui.colors.success, 16)) message = await ctx.send(embed=embed) end: float = time.perf_counter() @@ -159,7 +160,6 @@ async def ping(ctx) -> None: embed.add_field(name="Ping", value=f"{bot.latency * 1000:.2f} ms", inline=False) embed.add_field(name="Message Latency", value=f"{latency:.2f} ms", inline=False) - # Calculate the average ping of the bot in the last 10 minutes if latencies: average_ping = statistics.mean(latencies) embed.add_field( @@ -181,12 +181,10 @@ async def ping(ctx) -> None: inline=False, ) embed.set_footer( - text="Information requested by: {}".format(ctx.author.name), + text=f"Information requested by: {ctx.author.name}", icon_url=ctx.author.avatar.url, ) - embed.set_thumbnail( - url="https://uploads.poxipage.com/7q5iw7dwl5jc3zdjaergjhpat27tws8bkr9fgy45_938843265627717703-webp" - ) + embed.set_thumbnail(url=config.bot_avatar_url) await message.edit(embed=embed) @@ -196,13 +194,16 @@ async def ping(ctx) -> None: @bot.hybrid_command(name="help", description="View the various commands of this server") async def help(ctx) -> None: - user: discord.User | None = bot.get_user(1123551005993357342) - profilePicture: str = user.avatar.url + user: discord.User | None = bot.get_user(int(config.bot.bot_id)) + try: + profilePicture: str = user.avatar.url + except AttributeError: + profilePicture = config.bot.avatar_url embed = discord.Embed( title="Pollinations.ai Bot Commands", description="Here is the list of the available commands:", - color=discord.Color.og_blurple(), + color=int(config.ui.colors.success, 16), ) embed.set_thumbnail(url=profilePicture) @@ -210,7 +211,7 @@ async def help(ctx) -> None: embed.add_field(name=i, value=commands_[i], inline=False) embed.set_footer( - text="Information requested by: {}".format(ctx.author.name), + text=f"Information requested by: {ctx.author.name}", icon_url=ctx.author.avatar.url, ) @@ -221,13 +222,13 @@ async def help(ctx) -> None: async def invite(ctx) -> None: embed = discord.Embed( title="Invite the bot to your server", - url="https://discord.com/api/oauth2/authorize?client_id=1123551005993357342&permissions=534791060544&scope=bot%20applications.commands", + url=config.ui.bot_invite_url, description="Click the link above to invite the bot to your server", - color=discord.Color.og_blurple(), + color=int(config.ui.colors.success, 16), ) embed.set_footer( - text="Information requested by: {}".format(ctx.author.name), + text=f"Information requested by: {ctx.author.name}", icon_url=ctx.author.avatar.url, ) @@ -236,18 +237,21 @@ async def invite(ctx) -> None: @bot.hybrid_command(name="about", description="About the bot") async def about(ctx) -> None: - user: discord.User | None = bot.get_user(1123551005993357342) - profilePicture: str = user.avatar.url + user: discord.User | None = bot.get_user(int(config.bot.bot_id)) + try: + profilePicture: str = user.avatar.url + except AttributeError: + profilePicture = config.bot.avatar_url embed = discord.Embed( title="About Pollinations.ai Bot 🙌", - url="https://pollinations.ai/", + url=config.ui.api_provider_url, description="I am the official Pollinations.ai Bot. I can generate AI Images from your prompts ✨.", - color=discord.Color.og_blurple(), + color=int(config.ui.colors.success, 16), ) github_emoji: discord.Emoji | None = discord.utils.get( - bot.emojis, id=1187437992093155338, name="github" + bot.emojis, id=int(config.bot.emojis["github_emoji_id"]), name="github" ) embed.set_thumbnail(url=profilePicture) @@ -263,17 +267,17 @@ async def about(ctx) -> None: ) embed.add_field( name="How do I use this bot? 🤔", - value="You can use this bot by typing `/help` or clicking to get started.", + value=f"You can use this bot by typing `/help` or clicking to get started.", inline=False, ) embed.add_field( name="How do I report a bug? đŸĒ˛", - value="You can report a bug by joining our [Discord Server](https://discord.gg/SFasNG4n6b).", + value=f"You can report a bug by joining our [Discord Server]({config.ui.support_server_url}).", inline=False, ) embed.add_field( name=f"How do I contribute to this project? {str(github_emoji)}", - value="This project is open source. You can contribute to this project by visiting our [GitHub Repository](https://github.com/zingzy/pollinations.ai-bot).", + value=f"This project is open source. You can contribute to this project by visiting our [GitHub Repository]({config.ui.github_repo_url}).", inline=False, ) @@ -294,7 +298,7 @@ async def about(ctx) -> None: embed.set_footer( text="Bot created by Zngzy", - icon_url="https://i.ibb.co/6Pb7XG9/18622ff1cc55d7dca730d1ac246b6192.png", + icon_url=config.ui.bot_creator_avatar, ) await ctx.send(embed=embed) diff --git a/pyproject.toml b/pyproject.toml index 95df6aa..e8feff4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,8 +9,10 @@ dependencies = [ "discord>=2.3.2", "discord-py>=2.4.0", "pillow>=11.0.0", + "pydantic>=2.10.6", "python-dotenv>=1.0.1", "requests>=2.32.3", "ruff>=0.9.4", + "tomli>=2.2.1", "typing-extensions>=4.12.2", ] diff --git a/requirements.txt b/requirements.txt index ef9bdf3..cfbc591 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,17 +1,21 @@ aiohappyeyeballs==2.4.4 aiohttp==3.11.11 aiosignal==1.3.2 +annotated-types==0.7.0 attrs==24.3.0 certifi==2024.12.14 charset-normalizer==3.4.1 -discord-py==2.4.0 +discord.py==2.4.0 frozenlist==1.5.0 idna==3.10 multidict==6.1.0 pillow==11.0.0 propcache==0.2.1 +pydantic==2.10.6 +pydantic_core==2.27.2 python-dotenv==1.0.1 requests==2.32.3 -typing-extensions==4.12.2 +tomli==2.2.1 +typing_extensions==4.12.2 urllib3==2.3.0 yarl==1.18.3 diff --git a/utils/embed_utils.py b/utils/embed_utils.py index 85de833..6c4f6ba 100644 --- a/utils/embed_utils.py +++ b/utils/embed_utils.py @@ -3,7 +3,7 @@ import datetime from discord import Embed import random -from constants import WAITING_GIFS +from config import config __all__: list[str] = ("generate_pollinate_embed", "generate_error_message") @@ -15,7 +15,10 @@ async def generate_pollinate_embed( time_taken: datetime.timedelta, ) -> Embed: embed = discord.Embed( - title="", timestamp=datetime.datetime.now(datetime.timezone.utc), url=dic["url"] + title="", + timestamp=datetime.datetime.now(datetime.timezone.utc), + url=dic["url"], + color=int(config.ui.colors.success, 16), ) embed.add_field( @@ -24,7 +27,11 @@ async def generate_pollinate_embed( inline=False, ) - if len(dic["prompt"]) < 80 or dic["enhance"]: + if ( + len(dic["prompt"]) + < config.image_generation.validation.max_enhanced_prompt_length + or dic["enhance"] + ): if "enhanced_prompt" in dic and dic["enhanced_prompt"] is not None: embed.add_field( name="Enhanced Prompt", @@ -64,8 +71,6 @@ async def generate_error_message( if cooldown_configuration is None: cooldown_configuration: list[str] = [ "- 1 time every 10 seconds", - "- 5 times every 60 seconds", - "- 200 times every 24 hours", ] end_time = datetime.datetime.now() + datetime.timedelta(seconds=error.retry_after) @@ -74,10 +79,10 @@ async def generate_error_message( embed = discord.Embed( title="âŗ Cooldown", description=f"### You can use this command again ", - color=discord.Color.red(), + color=int(config.ui.colors.error, 16), timestamp=interaction.created_at, ) - embed.set_image(url=random.choice(WAITING_GIFS)) + embed.set_image(url=random.choice(config.resources.waiting_gifs)) embed.add_field( name="How many times can I use this command?", diff --git a/utils/error_handler.py b/utils/error_handler.py index 2474aeb..87e79e7 100644 --- a/utils/error_handler.py +++ b/utils/error_handler.py @@ -1,22 +1,18 @@ -import discord +from discord import Embed, Interaction +from config import config -# Create and sends an error embed async def send_error_embed( - interaction: discord.Interaction, title: str, description: str -): - embed = discord.Embed( - title=title, - description=description, - color=discord.Color.red(), + interaction: Interaction, title: str, description: str = None +) -> None: + if not description: + description = config.ui.error_messages["unknown"] + + embed = Embed( + title=title, description=description, color=int(config.ui.colors.error, 16) ) - # Try different methods to send the error message - for send_method in [ - lambda: interaction.response.send_message(embed=embed, ephemeral=True), - lambda: interaction.edit_original_response(embed=embed), - lambda: interaction.followup.send(embed=embed, ephemeral=True), - ]: - try: - return await send_method() - except Exception: - continue + + if not interaction.response.is_done(): + await interaction.response.send_message(embed=embed, ephemeral=True) + else: + await interaction.followup.send(embed=embed, ephemeral=True) diff --git a/utils/image_gen_utils.py b/utils/image_gen_utils.py index 4fdbd52..bf9465a 100644 --- a/utils/image_gen_utils.py +++ b/utils/image_gen_utils.py @@ -1,5 +1,4 @@ import random -from constants import MODELS import aiohttp import io from urllib.parse import quote @@ -7,30 +6,34 @@ import json from PIL import Image from exceptions import PromptTooLongError, DimensionTooSmallError, APIError +from config import config __all__: list[str] = ("generate_image", "validate_prompt", "validate_dimensions") def validate_prompt(prompt) -> None: - if len(prompt) > 2000: - raise PromptTooLongError("Prompt must be less than 2000 characters") + if len(prompt) > config.image_generation.validation.max_prompt_length: + raise PromptTooLongError(config.ui.error_messages["prompt_too_long"]) def validate_dimensions(width, height) -> None: - if width < 16 or height < 16: - raise DimensionTooSmallError("Width and Height must be greater than 16") + if ( + width < config.image_generation.validation.min_width + or height < config.image_generation.validation.min_height + ): + raise DimensionTooSmallError(config.ui.error_messages["dimension_too_small"]) async def generate_image( prompt: str = None, - width: int = 800, - height: int = 800, - model: str = f"{MODELS[0]}", - safe: bool = False, - cached: bool = False, - nologo: bool = False, - enhance: bool = False, - private: bool = False, + width: int = config.image_generation.defaults.width, + height: int = config.image_generation.defaults.height, + model: str = config.MODELS[0], + safe: bool = config.image_generation.defaults.safe, + cached: bool = config.image_generation.defaults.cached, + nologo: bool = config.image_generation.defaults.nologo, + enhance: bool = config.image_generation.defaults.enhance, + private: bool = config.image_generation.defaults.private, **kwargs, ): print( @@ -40,16 +43,16 @@ async def generate_image( seed = str(random.randint(0, 1000000000)) - url: str = f"https://image.pollinations.ai/prompt/{prompt}" + url: str = f"{config.api.image_gen_endpoint}/{prompt}" url += "" if cached else f"?seed={seed}" url += f"&width={width}" url += f"&height={height}" - url += f"&model={model}" + url += f"&model={model}" if model else "" url += f"&safe={safe}" if safe else "" url += f"&nologo={nologo}" if nologo else "" url += f"&enhance={enhance}" if enhance else "" url += f"&nofeed={private}" if private else "" - url += "&referer=discordbot" + url += f"&referer={config.image_generation.referer}" dic = { "prompt": prompt, @@ -73,9 +76,9 @@ async def generate_image( f"Server error occurred while generating image with status code: {response.status}\nPlease try again later" ) elif response.status == 429: - raise APIError("Rate limit exceeded. Please try again later") + raise APIError(config.ui.error_messages["rate_limit"]) elif response.status == 404: - raise APIError("The requested resource was not found") + raise APIError(config.ui.error_messages["resource_not_found"]) elif response.status != 200: raise APIError( f"API request failed with status code: {response.status}", @@ -95,7 +98,11 @@ async def generate_image( try: dic["nsfw"] = user_comment["has_nsfw_concept"] - if enhance or len(prompt) < 80: + if ( + enhance + or len(prompt) + < config.image_generation.validation.max_enhanced_prompt_length + ): enhance_prompt = user_comment["prompt"] if enhance_prompt == prompt: dic["enhanced_prompt"] = None diff --git a/utils/pollinate_utils.py b/utils/pollinate_utils.py index 8fa6218..bc6bd15 100644 --- a/utils/pollinate_utils.py +++ b/utils/pollinate_utils.py @@ -1,5 +1,5 @@ from urllib.parse import urlparse, parse_qs -from constants import MODELS +from config import config def parse_url(url: str) -> dict: @@ -9,7 +9,7 @@ def parse_url(url: str) -> dict: data = { "width": int(params.get("width", [1000])[0]), "height": int(params.get("height", [1000])[0]), - "model": params.get("model", [MODELS[0]])[0], + "model": params.get("model", [config.MODELS[0]])[0], "safe": True if params.get("safe", [False])[0] == "True" else False, "cached": True if "seed" not in params else False, "nologo": True if params.get("nologo", [False])[0] == "True" else False,