Add the edit feature to the random_image command too

Author: Zingzy

cogs/cross_pollinate_cog.py

@@ -35,9 +35,9 @@ async def on_submit(self, interaction: discord.Interaction):
         try:
             # Validate the prompt
             validate_prompt(self.edit_prompt.value)
-            
+
             await interaction.response.defer(thinking=True)
-            
+
             start = datetime.datetime.now()
             discord_logger.log_image_generation(
                 action="edit_start",
@@ -52,9 +52,9 @@ async def on_submit(self, interaction: discord.Interaction):
             dic, edited_image = await generate_cross_pollinate(
                 prompt=self.edit_prompt.value,
                 image_url=self.original_image_url,
-                nologo=config.image_generation.defaults.nologo
+                nologo=config.image_generation.defaults.nologo,
             )
-            
+
             time_taken = (datetime.datetime.now() - start).total_seconds()
             discord_logger.log_image_generation(
                 action="edit_complete",
@@ -72,7 +72,13 @@ async def on_submit(self, interaction: discord.Interaction):
 
             time_taken_delta = datetime.datetime.now() - start
             embed = await generate_cross_pollinate_embed(
-                interaction, False, dic, time_taken_delta, self.edit_prompt.value, self.original_image_url, image_file.filename
+                interaction,
+                False,
+                dic,
+                time_taken_delta,
+                self.edit_prompt.value,
+                self.original_image_url,
+                image_file.filename,
             )
 
             # Send the edited image with cross-pollinate buttons
@@ -153,7 +159,10 @@ async def edit(self, interaction: discord.Interaction, button: discord.ui.Button
                 return
 
             # Get the original image URL from the message embed
-            if not interaction.message.embeds or not interaction.message.embeds[0].image:
+            if (
+                not interaction.message.embeds
+                or not interaction.message.embeds[0].image
+            ):
                 await interaction.response.send_message(
                     embed=SafeEmbed(
                         title="🎨 No Image Found",
@@ -163,27 +172,27 @@ async def edit(self, interaction: discord.Interaction, button: discord.ui.Button
                     ephemeral=True,
                 )
                 return
-            
+
             original_image_url = interaction.message.embeds[0].image.url
-            
+
             # Get the original prompt based on the embed fields
             interaction_data: dict = interaction.message.embeds[0].to_dict()
-            
+
             # Check if it's a cross-pollinate embed or regular pollinate embed
             prompt_field = None
             for field in interaction_data.get("fields", []):
                 if field["name"] in ["Cross-Pollinate Prompt 🐝", "Prompt"]:
                     prompt_field = field
                     break
-            
+
             original_prompt: str = (
                 prompt_field["value"][3:-3] if prompt_field else "Unknown prompt"
             )
-            
+
             # Show the edit modal
             modal = EditImageModal(original_image_url, original_prompt)
             await interaction.response.send_modal(modal)
-            
+
         except Exception as e:
             discord_logger.log_error(
                 error_type="edit_button_error",
@@ -552,7 +561,13 @@ async def cross_pollinate_command(
 
             time_taken_delta: datetime.timedelta = datetime.datetime.now() - start
             embed: SafeEmbed = await generate_cross_pollinate_embed(
-                interaction, private, dic, time_taken_delta, prompt, image_url, image_file.filename
+                interaction,
+                private,
+                dic,
+                time_taken_delta,
+                prompt,
+                image_url,
+                image_file.filename,
             )
 
             if private:

cogs/imagine_cog.py

@@ -150,7 +150,10 @@ async def edit(self, interaction: discord.Interaction, button: discord.ui.Button
                 return
 
             # Get the original image URL from the message embed
-            if not interaction.message.embeds or not interaction.message.embeds[0].image:
+            if (
+                not interaction.message.embeds
+                or not interaction.message.embeds[0].image
+            ):
                 await interaction.response.send_message(
                     embed=SafeEmbed(
                         title="🎨 No Image Found",
@@ -160,17 +163,17 @@ async def edit(self, interaction: discord.Interaction, button: discord.ui.Button
                     ephemeral=True,
                 )
                 return
-            
+
             original_image_url = interaction.message.embeds[0].image.url
-            
+
             # Get the original prompt
             interaction_data: dict = interaction.message.embeds[0].to_dict()
             original_prompt: str = interaction_data["fields"][0]["value"][3:-3]
-            
+
             # Show the edit modal
             modal = EditImageModal(original_image_url, original_prompt)
             await interaction.response.send_modal(modal)
-            
+
         except Exception as e:
             discord_logger.log_error(
                 error_type="edit_button_error",

cogs/multi_pollinate_cog.py

@@ -259,7 +259,8 @@ async def generate_for_model(i, model):
 
         embeds[0].add_field(name="Prompt", value=f"```{prompt}```", inline=False)
         embeds[0].add_field(
-            name="Total Processing Time", value=f"```{time_taken.total_seconds():.2f}s```"
+            name="Total Processing Time",
+            value=f"```{time_taken.total_seconds():.2f}s```",
         )
 
         embeds[0].set_footer(text=f"Generated by {interaction.user}")

cogs/random_cog.py

@@ -7,14 +7,25 @@
 from utils.image_gen_utils import generate_image, validate_dimensions
 from utils.embed_utils import generate_error_message, SafeEmbed
 from utils.error_handler import send_error_embed
+from utils.logger import discord_logger
 from exceptions import DimensionTooSmallError, APIError
 
+# Import cross-pollinate button view for edit functionality
+from cogs.cross_pollinate_cog import CrossPollinateButtonView
+
 
 class RandomImage(commands.Cog):
     def __init__(self, bot) -> None:
         self.bot = bot
         self.command_config = config.commands["random"]
 
+    async def cog_load(self) -> None:
+        await self.bot.wait_until_ready()
+        self.bot.add_view(CrossPollinateButtonView())
+        discord_logger.log_bot_event(
+            action="cog_load", status="success", details={"cog": "RandomImage"}
+        )
+
     @app_commands.command(name="random", description="Generate Random AI Images")
     @app_commands.choices(
         model=[
@@ -81,7 +92,6 @@ 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)
@@ -100,7 +110,9 @@ async def random_image_command(
             await interaction.followup.send(embed=embed, ephemeral=True)
             return
         else:
-            await interaction.followup.send(embed=embed, file=image_file)
+            # Use CrossPollinateButtonView for public images to get edit, delete, and bookmark functionality
+            view = CrossPollinateButtonView()
+            await interaction.followup.send(embed=embed, view=view, file=image_file)
 
     @random_image_command.error
     async def random_image_command_error(
@@ -140,4 +152,6 @@ async def random_image_command_error(
 
 async def setup(bot) -> None:
     await bot.add_cog(RandomImage(bot))
-    print("Random Image cog loaded")
+    discord_logger.log_bot_event(
+        action="cog_setup", status="success", details={"cog": "RandomImage"}
+    )