add Leaderboard command, View for Multi-Imagine

cogs/imagine_cog.py

@@ -252,10 +252,10 @@ async def imagine_command(self, interaction, prompt:str, model: str = "Dreamshap
         view = self.ImagineButtonView(link=dic["bookmark_url"])
 
         if private:
-            response = await interaction.followup.send(f"## {prompt} - {interaction.user.mention}", file=image_file, ephemeral=True)
+            response = await interaction.followup.send(f"## {prompt} - {interaction.user.mention}\n### Model - `{model}`", file=image_file, ephemeral=True)
             return
         else:
-            response = await interaction.channel.send(f"## {prompt} - {interaction.user.mention}", file=image_file, view=view)
+            response = await interaction.channel.send(f"## {prompt} - {interaction.user.mention}\n### Model - `{model}`", file=image_file, view=view)
 
         message_id = response.id
         dic["_id"] = message_id

cogs/leaderboard_cog.py

@@ -0,0 +1,50 @@
+import datetime
+from discord import app_commands, ui
+from discord.ext import commands
+import discord
+from utils import *
+from constants import *
+
+class leaderboard(commands.Cog):
+    def __init__(self, bot):
+        self.bot = bot
+
+    @app_commands.command(name="leaderboard", description="Shows the leaderboard of the server")
+    @app_commands.guild_only()
+    async def leaderboard_command(self, interaction):
+        await interaction.response.defer()
+
+        leaderboard = generate_global_leaderboard()
+
+        leaderboard_ = {}
+
+        for i in list(leaderboard.keys())[1:]:
+            user = await self.bot.fetch_user(i)
+            leaderboard_[i] = {"name": user.name, "points": leaderboard[i]}
+
+        top_user = await self.bot.fetch_user(list(leaderboard.keys())[0])
+        top_user_id = top_user.id
+
+        try:
+            top_user_avatar = top_user.avatar.url
+        except Exception as e:
+            top_user_avatar = top_user.default_avatar.url
+
+        embed = discord.Embed(title="Top 10 prompters", color=discord.Color.gold(), description="This Shows the top 10 bot users across all of the servers the bot is in.", timestamp=datetime.datetime.utcnow())
+        embed.set_thumbnail(url=top_user_avatar)
+
+        embed.add_field(name=f"{NUMBER_EMOJIES[1]}  {top_user.name} - {leaderboard[top_user_id]} points", value="** **")
+
+        for i, user in enumerate(leaderboard_):
+            embed.add_field(name=f"{NUMBER_EMOJIES[i+2]}  {leaderboard_[user]['name']} - {leaderboard_[user]['points']} points", inline=False, value="** **")
+
+        try:
+            embed.set_footer(text=f"Requested by {interaction.user.name}", icon_url=interaction.user.avatar.url)
+        except Exception as e:
+            embed.set_footer(text=f"Requested by {interaction.user.name}", icon_url=interaction.user.default_avatar.url)
+
+        await interaction.followup.send(embed=embed)
+
+async def setup(bot):
+    await bot.add_cog(leaderboard(bot))
+    print("leaderboard cog loaded")

cogs/multi_imagine_cog.py

@@ -7,9 +7,109 @@
 from constants import *
 
 class Multi_imagine(commands.Cog):
-    def __init(self, bot):
+    def __init__(self, bot):
         self.bot = bot
 
+    async def cog_load(self):
+        await self.bot.wait_until_ready()
+        self.bot.add_view(self.multiImagineButtonView())
+
+    async def regenerate(interaction: discord.Interaction, button: discord.ui.Button, data: dict, model_no: int):
+        data["model"] = MODELS[model_no]
+
+        dic, img = await generate_image(**data)
+
+        image_file = discord.File(img, "image.png")
+        if data["nsfw"]:
+            image_file.filename = f"SPOILER_{image_file.filename}"
+
+        response = await interaction.channel.send(f"## {data['prompt']} - {interaction.user.mention}\n### Model - `{MODELS[model_no]}`", file=image_file)
+
+        message_id = response.id
+        dic["_id"] = message_id
+        dic["channel_id"] = interaction.channel.id
+        dic["user_id"] = interaction.user.id
+        dic["guild_id"] = interaction.guild.id
+        dic["bookmarks"] = []
+        dic["author"] = interaction.user.id
+        dic["likes"] = []
+
+        user_data = get_user_data(interaction.user.id)
+        if user_data is None:
+            user_data = {
+                "_id": interaction.user.id,
+                "bookmarks": [],
+                "likes": [],
+                "prompts": [],
+                "last_prompt": None,
+            }
+            save_user_data(interaction.user.id, user_data)
+
+        user_data["prompts"].append(message_id)
+        user_data["last_prompt"] = message_id
+
+        update_user_data(interaction.user.id, user_data)
+        save_prompt_data(message_id, dic)
+
+    class multiImagineButtonView(discord.ui.View):
+        def __init__(self):
+            super().__init__(timeout=None)
+
+        @discord.ui.button(label="V1", style=discord.ButtonStyle.secondary, custom_id="v1")
+        async def regerate_1(self, interaction: discord.Interaction, button: discord.ui.Button):
+            await interaction.response.send_message(embed=discord.Embed(title="Regenerating Image 1", description="Please wait while we regenerate your image", color=discord.Color.blurple()), ephemeral=True)
+            data = get_multi_imagined_prompt_data(interaction.message.id)
+
+            await Multi_imagine.regenerate(interaction, button, data, 0)
+
+        @discord.ui.button(label="V2", style=discord.ButtonStyle.secondary, custom_id="v2")
+        async def regerate_2(self, interaction: discord.Interaction, button: discord.ui.Button):
+            await interaction.response.send_message(embed=discord.Embed(title="Regenerating Image 2", description="Please wait while we regenerate your image", color=discord.Color.blurple()), ephemeral=True)
+            data = get_multi_imagined_prompt_data(interaction.message.id)
+
+            await Multi_imagine.regenerate(interaction, button, data, 1)
+
+        @discord.ui.button(label="V3", style=discord.ButtonStyle.secondary, custom_id="v3")
+        async def regerate_3(self, interaction: discord.Interaction, button: discord.ui.Button):
+            await interaction.response.send_message(embed=discord.Embed(title="Regenerating Image 3", description="Please wait while we regenerate your image", color=discord.Color.blurple()), ephemeral=True)
+            data = get_multi_imagined_prompt_data(interaction.message.id)
+
+            await Multi_imagine.regenerate(interaction, button, data, 2)
+
+        @discord.ui.button(label="V4", style=discord.ButtonStyle.secondary, custom_id="v4")
+        async def regerate_4(self, interaction: discord.Interaction, button: discord.ui.Button):
+            await interaction.response.send_message(embed=discord.Embed(title="Regenerating Image 4", description="Please wait while we regenerate your image", color=discord.Color.blurple()), ephemeral=True)
+            data = get_multi_imagined_prompt_data(interaction.message.id)
+
+            await Multi_imagine.regenerate(interaction, button, data, 3)
+
+        @discord.ui.button(label="V5", style=discord.ButtonStyle.secondary, custom_id="v5")
+        async def regerate_5(self, interaction: discord.Interaction, button: discord.ui.Button):
+            await interaction.response.send_message(embed=discord.Embed(title="Regenerating Image 5", description="Please wait while we regenerate your image", color=discord.Color.blurple()), ephemeral=True)
+            data = get_multi_imagined_prompt_data(interaction.message.id)
+
+            await Multi_imagine.regenerate(interaction, button, data, 4)
+
+        @discord.ui.button(label="V6", style=discord.ButtonStyle.secondary, custom_id="v6")
+        async def regerate_6(self, interaction: discord.Interaction, button: discord.ui.Button):
+            await interaction.response.send_message(embed=discord.Embed(title="Regenerating Image 6", description="Please wait while we regenerate your image", color=discord.Color.blurple()), ephemeral=True)
+            data = get_multi_imagined_prompt_data(interaction.message.id)
+
+            await Multi_imagine.regenerate(interaction, button, data, 5)
+
+        @discord.ui.button(label="V7", style=discord.ButtonStyle.secondary, custom_id="v7")
+        async def regerate_7(self, interaction: discord.Interaction, button: discord.ui.Button):
+            await interaction.response.send_message(embed=discord.Embed(title="Regenerating Image 7", description="Please wait while we regenerate your image", color=discord.Color.blurple()), ephemeral=True)
+            data = get_multi_imagined_prompt_data(interaction.message.id)
+
+            await Multi_imagine.regenerate(interaction, button, data, 6)
+
+        @discord.ui.button(style=discord.ButtonStyle.red, custom_id="multiimagine_delete", emoji="<:delete:1187102382312652800>")
+        async def delete(self, interaction: discord.Interaction, button: discord.ui.Button):
+            delete_multi_imagined_prompt_data(interaction.message.id)
+
+            await interaction.message.delete()
+
     @app_commands.command(name="multi-imagine", description="Imagine multiple prompts")
     @app_commands.checks.cooldown(1, 30)
     @app_commands.guild_only()
@@ -18,8 +118,6 @@ async def multiimagine_command(self, interaction, prompt:str, width:int = 1000,
 
         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)
 
-        models = ['Deliberate', 'Playground', 'Pixart', 'Dreamshaper', 'Turbo', 'Formulaxl', 'Dpo']
-
         if len(prompt) > 1500:
             await interaction.channel.send(embed=discord.Embed(title="Error", description="Prompt must be less than 1500 characters", color=discord.Color.red()))
             return
@@ -29,17 +127,23 @@ async def multiimagine_command(self, interaction, prompt:str, width:int = 1000,
             return
 
         images = []
+        image_urls = {}
         description = ""
         counter = 1
 
-        for i in models:
+        for i in MODELS:
             try:
                 time = datetime.datetime.now()
                 dic, image = await generate_image(prompt, width, height, i, negative, cached, nologo, enhance)
+
+                image_urls[i] = dic["bookmark_url"]
+
                 time_taken = datetime.datetime.now() - time
                 await interaction.followup.send(f"Generated `{i} model` Image in `{round(time_taken.total_seconds(), 2)}` seconds ✅", ephemeral=True)
+
                 description += f"Image {counter} model :  `{i}`\n"
                 counter += 1
+
                 images.append(image)
             except Exception as e:
                 print(e)
@@ -56,15 +160,32 @@ async def multiimagine_command(self, interaction, prompt:str, width:int = 1000,
             file_name = f"{prompt}_{idx}.png" if not is_nsfw else f"SPOILER_{prompt}_{idx}.png"
             files.append(discord.File(img, file_name))
 
+        multi_imagine_view = self.multiImagineButtonView()
+
         if not len(files) == 0:
             if private:
                 response = await interaction.followup.send(f'## {prompt} - {interaction.user.mention}\n{description}', files=files, ephemeral= True)
             else:
-                response = await interaction.channel.send(f'## {prompt} - {interaction.user.mention}\n{description}', files=files)
+                response = await interaction.channel.send(f'## {prompt} - {interaction.user.mention}\n{description}', files=files, view=multi_imagine_view)
         else:
             await interaction.followup.send(embed=discord.Embed(title="Error", description="No images were generated", color=discord.Color.red()), ephemeral=True)
             return
 
+        message_id = response.id
+        dic["_id"] = message_id
+        dic["channel_id"] = interaction.channel.id
+        dic["user_id"] = interaction.user.id
+        dic["guild_id"] = interaction.guild.id
+        dic["urls"] = image_urls
+        dic["author"] = interaction.user.id
+        dic["nsfw"] = is_nsfw
+
+        del dic["bookmark_url"]
+        del dic["seed"]
+        del dic["model"]
+
+        save_multi_imagined_prompt_data(message_id, dic)
+
     @multiimagine_command.error
     async def multiimagine_command_error(
         self, interaction: discord.Interaction, error: app_commands.AppCommandError

constants.py

@@ -6,6 +6,7 @@
 TOKEN = os.environ["TOKEN"]
 MONGODB_URI = os.environ["MONGODB_URI"]
 APP_URI = os.environ["APP_URI"]
+MODELS = ['Deliberate', 'Playground', 'Pixart', 'Dreamshaper', 'Turbo', 'Formulaxl', 'Dpo']
 
 with open("nsfw.txt", "r") as r:
     nsfw = r.read().split("\n")

utils.py

@@ -5,6 +5,7 @@
 from urllib.parse import quote
 from pymongo import MongoClient
 import sys
+import itertools
 
 client = MongoClient(MONGODB_URI)
 
@@ -17,6 +18,20 @@
 db = client["pollinations"]
 prompts = db["prompts"]
 users = db["users"]
+multi_prompts = db["multi_prompts"]
+
+NUMBER_EMOJIES = {
+    1:"🥇",
+    2:"🥈",
+    3:"🥉",
+    4:"4️⃣",
+    5:"5️⃣",
+    6:"6️⃣",
+    7:"7️⃣",
+    8:"8️⃣",
+    9:"9️⃣",
+    10:"1️⃣0️⃣"
+}
 
 
 def get_prompt_data(message_id: int):
@@ -48,6 +63,35 @@ def delete_prompt_data(message_id: int):
         print(e)
 
 
+def get_multi_imagined_prompt_data(message_id: int):
+    try:
+        return multi_prompts.find_one({"_id": message_id})
+    except Exception as e:
+        print(e)
+        return None
+
+
+def save_multi_imagined_prompt_data(message_id: int, data: dict):
+    try:
+        multi_prompts.insert_one(data)
+    except Exception as e:
+        print(e)
+
+
+def update_multi_imagined_prompt_data(message_id: int, data: dict):
+    try:
+        multi_prompts.update_one({"_id": message_id}, {"$set": data})
+    except Exception as e:
+        print(e)
+
+
+def delete_multi_imagined_prompt_data(message_id: int):
+    try:
+        multi_prompts.delete_one({"_id": message_id})
+    except Exception as e:
+        print(e)
+
+
 def get_prompts_counts():
     try:
         return prompts.count_documents({})
@@ -78,6 +122,21 @@ def update_user_data(user_id: int, data: dict):
         print(e)
 
 
+def generate_global_leaderboard():
+    try:
+        documents = users.find()
+
+        data = {doc["_id"]: len(doc["prompts"]) for doc in documents}
+
+        sorted_data = dict(sorted(data.items(), key=lambda item: item[1], reverse=True))
+
+        top_10_users = dict(itertools.islice(sorted_data.items(), 10))
+        return top_10_users
+
+    except Exception as e:
+        print(e)
+        return None
+
 async def generate_image(
     prompt: str,
     width: int = 500,
@@ -87,6 +146,7 @@ async def generate_image(
     cached: bool = False,
     nologo: bool = False,
     enhance: bool = True,
+    **kwargs
 ):
     model = model.lower()