From 18a88505ab4ef0c914701fae7c17be5eb46e209a Mon Sep 17 00:00:00 2001 From: Zane <90309290+Zingzy@users.noreply.github.com> Date: Thu, 21 Dec 2023 18:45:20 +0530 Subject: [PATCH] Refactor API and utils modules, add cooldown checks --- .github/workflows/format.yaml | 12 +++++ api.py | 4 +- cogs/imagine_cog.py | 51 +++++++++++++++++++++ main.py | 85 ++++++++++++++++++++++++----------- utils.py | 23 ++++++++-- 5 files changed, 144 insertions(+), 31 deletions(-) create mode 100644 .github/workflows/format.yaml diff --git a/.github/workflows/format.yaml b/.github/workflows/format.yaml new file mode 100644 index 0000000..36e5332 --- /dev/null +++ b/.github/workflows/format.yaml @@ -0,0 +1,12 @@ +on: [push, pull_request] + +jobs: + python-black: + name: Python Black + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - name: Python Black + uses: cytopia/docker-black@0.8 + with: + path: '.' \ No newline at end of file diff --git a/api.py b/api.py index 61d9c9f..0b25120 100644 --- a/api.py +++ b/api.py @@ -8,14 +8,16 @@ def home(): return "Hello, I am alive" + @app.route("/health") def health(): return jsonify({"status": "ok", "message": "I am alive"}), 200 + def run(): app.run(host="0.0.0.0", port=8080) def keep_alive(): t = Thread(target=run) - t.start() \ No newline at end of file + t.start() diff --git a/cogs/imagine_cog.py b/cogs/imagine_cog.py index 5b4d604..bede15d 100644 --- a/cogs/imagine_cog.py +++ b/cogs/imagine_cog.py @@ -150,6 +150,7 @@ async def model_autocomplete(self, @app_commands.command(name="imagine", description="Imagine a prompt") @app_commands.autocomplete(model=model_autocomplete) + @app_commands.checks.cooldown(1, 15) @app_commands.describe(prompt="Imagine a prompt", height="Height of the image", width="Width of the image", negative="The things not to include in the image", cached="Removes the image seed", nologo="Remove the logo", enhance="Disables Prompt enhancing if set to False", private="Only you can see the generated Image if set to True") async def imagine_command(self, interaction, prompt:str, model: str = "Dreamshaper", width:int = 1000, height:int = 1000, negative:str|None = None, cached:bool = False, nologo:bool = False, enhance:bool = True, private:bool = False): await interaction.response.send_message(embed=discord.Embed(title="Generating Image", description="Please wait while we generate your image", color=discord.Color.blurple()), ephemeral=True) @@ -192,6 +193,7 @@ async def imagine_command(self, interaction, prompt:str, model: str = "Dreamshap return @app_commands.command(name="multi-imagine", description="Imagine multiple prompts") + @app_commands.checks.cooldown(1, 30) @app_commands.describe(prompt="Imagine a prompt", height="Height of the image", width="Width of the image", negative="The things not to include in the image", cached="Removes the image seed", nologo="Remove the logo", enhance="Disables Prompt enhancing if set to False", private="Only you can see the generated Image if set to True") async def multiimagine_command(self, interaction, prompt:str, width:int = 1000, height:int = 1000, negative:str|None = None, cached:bool = False, nologo:bool = False, enhance:bool = True, private:bool = False): @@ -241,6 +243,55 @@ async def multiimagine_command(self, interaction, prompt:str, width:int = 1000, return + @imagine_command.error + async def imagine_command_error( + self, interaction: discord.Interaction, error: app_commands.AppCommandError + ): + if isinstance(error, app_commands.CommandOnCooldown): + end_time = datetime.datetime.now() + datetime.timedelta( + seconds=error.retry_after + ) + end_time_str = end_time.strftime("%Y-%m-%d %H:%M:%S UTC") + end_time_ts = f"" + + hours, remainder = divmod(error.retry_after, 3600) + minutes, seconds = divmod(remainder, 60) + seconds = round(seconds) + time_left = f"{ hours + ' hour, ' if not hours<1 else ''}{int(minutes)} minute{'s' if minutes != 1 else ''} and {seconds} second{'s' if seconds != 1 else ''}" + + embed = discord.Embed( + title="⏳ Cooldown", + description=f"You have to wait until **{end_time_ts}** ({time_left}) before using the command again.", + color=discord.Color.red(), + ) + + await interaction.response.send_message(embed=embed, ephemeral=True) + + @multiimagine_command.error + async def multiimagine_command_error( + self, interaction: discord.Interaction, error: app_commands.AppCommandError + ): + if isinstance(error, app_commands.CommandOnCooldown): + end_time = datetime.datetime.now() + datetime.timedelta( + seconds=error.retry_after + ) + end_time_str = end_time.strftime("%Y-%m-%d %H:%M:%S UTC") + end_time_ts = f"" + + hours, remainder = divmod(error.retry_after, 3600) + minutes, seconds = divmod(remainder, 60) + seconds = round(seconds) + time_left = f"{ hours + ' hour, ' if not hours<1 else ''}{int(minutes)} minute{'s' if minutes != 1 else ''} and {seconds} second{'s' if seconds != 1 else ''}" + + embed = discord.Embed( + title="⏳ Cooldown", + description=f"You have to wait until **{end_time_ts}** ({time_left}) before using the again.", + color=discord.Color.red(), + ) + + await interaction.response.send_message(embed=embed, ephemeral=True) + + async def setup(bot): await bot.add_cog(Imagine(bot)) print("Imagine cog loaded") \ No newline at end of file diff --git a/main.py b/main.py index 83f3966..127dbdc 100644 --- a/main.py +++ b/main.py @@ -20,6 +20,31 @@ start_time = None latencies = [] +commands_ = { + " 🎨": """Generates AI Images based on your prompts +- **prompt** 🗣️ : Your prompt for the Image to be generated +- **model** 🤖 : The model to be used for generating the Image +- **width** ↔️ : The width of your prompted Image +- **height** ↕️ : The height of your prompted Image +- **cached** : specifies whether to return a cached image +- **negative** ❎ : Specifies what not to be in the generated images +- **nologo** 🚫 : Specifies whether to remove the logo from the generated images (deafault False) +- **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 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 +- **cached** : specifies whether to return a cached image +- **negative** ❎ : Specifies what not to be in the generated images +- **nologo** 🚫 : Specifies whether to remove the logo from the generated images (deafault False) +- **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 +""", + " ❓": "Displays this", +} + class pollinationsBot(commands.Bot): def __init__(self): @@ -53,15 +78,24 @@ async def on_ready(self): async def load(): - for filename in os.listdir('./cogs'): - if filename.endswith('.py'): - await bot.load_extension(f'cogs.{filename[:-3]}') + for filename in os.listdir("./cogs"): + if filename.endswith(".py"): + await bot.load_extension(f"cogs.{filename[:-3]}") + @bot.event async def on_message(message): if message.author == bot.user: return + if bot.user in message.mentions: + embed = discord.Embed( + description="Hello, I am the Pollinations.ai Bot. I am here to help you with your AI needs. Type `!help` or click to get started.", + color=discord.Color.og_blurple(), + ) + + await message.reply(embed=embed) + await bot.process_commands(message) @@ -102,16 +136,16 @@ async def ping(ctx): latency = (end - ctx.start) * 1000 + 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) - embed.add_field( - name="Websocket Latency", value=f"{bot.latency * 1000:.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( - name="Average Ping", value=f"{average_ping:.2f} ms", inline=False + name="Average Message Latency", + value=f"{average_ping:.2f} ms", + inline=False, ) global start_time @@ -141,31 +175,30 @@ async def ping(ctx): print(e, file=sys.stdout) -# @bot.hybrid_command(name="help", description="View the various commands of this server") -# async def help(ctx): -# user = bot.get_user(1168874960081649684) -# profilePicture = user.avatar.url +@bot.hybrid_command(name="help", description="View the various commands of this server") +async def help(ctx): + user = bot.get_user(1123551005993357342) + profilePicture = user.avatar.url -# embed = discord.Embed( -# title="SecretSanctuary Bot Commands", -# url=APP_URI, -# description="Here is the list of the available commands:", -# color=discord.Color.og_blurple(), -# ) + embed = discord.Embed( + title="Pollinations.ai Bot Commands", + url=APP_URI, + description="Here is the list of the available commands:", + color=discord.Color.og_blurple(), + ) -# embed.set_thumbnail(url=profilePicture) -# for i in commands_.keys(): -# embed.add_field(name=i, value=commands_[i], inline=False) + embed.set_thumbnail(url=profilePicture) + for i in commands_.keys(): + embed.add_field(name=i, value=commands_[i], inline=False) -# embed.set_footer( -# text="Information requested by: {}".format(ctx.author.name), -# icon_url=ctx.author.avatar.url, -# ) + embed.set_footer( + text="Information requested by: {}".format(ctx.author.name), + icon_url=ctx.author.avatar.url, + ) -# await ctx.send(embed=embed) + await ctx.send(embed=embed) if __name__ == "__main__": - # run_server() keep_alive() bot.run(token=TOKEN) diff --git a/utils.py b/utils.py index df2fb95..1173655 100644 --- a/utils.py +++ b/utils.py @@ -18,6 +18,7 @@ db = client["pollinations"] collection = db["prompts"] + def get_prompt_data(message_id: int): try: return collection.find_one({"_id": message_id}) @@ -25,29 +26,43 @@ def get_prompt_data(message_id: int): print(e) return None + def save_prompt_data(message_id: int, data: dict): try: collection.insert_one(data) except Exception as e: print(e) + def update_prompt_data(message_id: int, data: dict): try: collection.update_one({"_id": message_id}, {"$set": data}) except Exception as e: print(e) + def delete_prompt_data(message_id: int): try: collection.delete_one({"_id": message_id}) except Exception as e: print(e) -async def generate_image(prompt: str, width: int = 500, height: int = 500, model: str = "turbo", negative: str|None = None, cached:bool = False, nologo:bool = False, enhance:bool = True): - + +async def generate_image( + prompt: str, + width: int = 500, + height: int = 500, + model: str = "turbo", + negative: str | None = None, + cached: bool = False, + nologo: bool = False, + enhance: bool = True, +): model = model.lower() - print(f"Generating image with prompt: {prompt}, width: {width}, height: {height}, model: {model}, negative: {negative}, cached: {cached}, nologo: {nologo}, enhance: {enhance}") + print( + f"Generating image with prompt: {prompt}, width: {width}, height: {height}, model: {model}, negative: {negative}, cached: {cached}, nologo: {nologo}, enhance: {enhance}" + ) seed = str(random.randint(0, 1000000000)) @@ -69,7 +84,7 @@ async def generate_image(prompt: str, width: int = 500, height: int = 500, model "cached": cached, "nologo": nologo, "enhance": enhance, - "bookmark_url": quote(url, safe=':/&=?'), + "bookmark_url": quote(url, safe=":/&=?"), } dic["seed"] = seed if not cached else None