buttons class test

app.py

@@ -37,6 +37,42 @@ df = Client("huggingface-projects/IF", DFIF_TOKEN)
 #stable diffusion upscaler
 sdlu = Client("huggingface-projects/stable-diffusion-latent-upscaler", DFIF_TOKEN)
 
+#---------------------------------------------------------------------------------------------------------------------------------------------------
+class ButtonView(discord.ui.View):
+    def __init__(self, ctx, image_paths, stage_1_result_path):
+        super().__init__()
+        self.ctx = ctx
+        self.image_paths = image_paths
+        self.stage_1_result_path = stage_1_result_path
+
+    async def on_timeout(self):
+        for child in self.children:
+            child.disabled = True
+        self.stop()
+
+    @discord.ui.button(label='Image 1', style=discord.ButtonStyle.blurple)
+    async def image1_button(self, button: discord.ui.Button, interaction: discord.Interaction):
+        await self.ctx.invoke(self.ctx.bot.get_command('dfif2'), image_path=self.image_paths[0], stage_1_result_path=self.stage_1_result_path)
+        self.stop()
+
+    @discord.ui.button(label='Image 2', style=discord.ButtonStyle.blurple)
+    async def image2_button(self, button: discord.ui.Button, interaction: discord.Interaction):
+        await self.ctx.invoke(self.ctx.bot.get_command('dfif2'), image_path=self.image_paths[1], stage_1_result_path=self.stage_1_result_path)
+        self.stop()
+
+    @discord.ui.button(label='Image 3', style=discord.ButtonStyle.blurple)
+    async def image3_button(self, button: discord.ui.Button, interaction: discord.Interaction):
+        await self.ctx.invoke(self.ctx.bot.get_command('dfif2'), image_path=self.image_paths[2], stage_1_result_path=self.stage_1_result_path)
+        self.stop()
+
+    @discord.ui.button(label='Image 4', style=discord.ButtonStyle.blurple)
+    async def image4_button(self, button: discord.ui.Button, interaction: discord.Interaction):
+        await self.ctx.invoke(self.ctx.bot.get_command('dfif2'), image_path=self.image_paths[3], stage_1_result_path=self.stage_1_result_path)
+        self.stop()
+#---------------------------------------------------------------------------------------------------------------------------------------------------
+
+
+#---------------------------------------------------------------------------------------------------------------------------------------------------       
 # Set up discord bot 
 class MyClient(discord.Client):
     async def on_ready(self):
@@ -141,37 +177,29 @@ class MyClient(discord.Client):
                 combined_image_path = os.path.join(stage_1_results, 'combined_image.png')
                 combined_image.save(combined_image_path)
             
-                embed = discord.Embed(title="Combined Image")
-                embed.set_image(url="attachment://combined_image.png")
-
-                components = discord.ui.MessageComponents()
-                components.add_action_row(
-                    Button(style=ButtonStyle.blue, label="Image 1", custom_id="image_1"),
-                    Button(style=ButtonStyle.blue, label="Image 2", custom_id="image_2"),
-                    Button(style=ButtonStyle.blue, label="Image 3", custom_id="image_3"),
-                    Button(style=ButtonStyle.blue, label="Image 4", custom_id="image_4"),
-                )
-    
-                await message.channel.send(embed=embed, file=discord.File(combined_image_path, 'combined_image.png'), components=components)
-
-    
-    async def on_button_click(self, interaction):
-        if interaction.custom_id.startswith("image_"):
-            image_index = int(interaction.custom_id.split("_")[1])
-            await interaction.respond(content=f"You clicked image {image_index}")
-
-            #dfif stage 2
-            selected_index_for_stage_2 = image_index - 1
-            seed_2 = 0
-            guidance_scale_2 = 4
-            custom_timesteps_2 = 'smart50'
-            number_of_inference_steps_2 = 50
-            result_path = df.predict(stage_1_result_path, selected_index_for_stage_2, seed_2, guidance_scale_2, custom_timesteps_2, number_of_inference_steps_2, api_name='/upscale256')
-
-            # Send the upscaled image in a new message
-            upscale_image_path = os.path.join(result_path, "output.png")
-            with open(upscale_image_path, 'rb') as f:
-                await interaction.channel.send('Here is the upscaled image:', file=discord.File(f, 'upscaled_image.png'))       
+
+                # Send the combined image file as a discord attachment with the button view
+                with open(combined_image_path, 'rb') as f:
+                    view = ButtonView(message, [first_png_path, second_png_path, third_png_path, fourth_png_path], stage_1_result_path)
+                    await message.reply('Here is the combined image', file=discord.File(f, 'combined_image.png'), view=view)
+
+
+
+
+    #stage 2, can be stable diffusion too---------------------------------------------------------------------------------------------------------------------------------------------------
+    async def dfif2(self, ctx, image_path, stage_1_result_path):
+        selected_index_for_stage_2 = 0
+        seed_2 = 0
+        guidance_scale_2 = 4
+        custom_timesteps_2 = 'smart50'
+        number_of_inference_steps_2 = 50
+        result_path = df.predict(stage_1_result_path, selected_index_for_stage_2, seed_2, guidance_scale_2, custom_timesteps_2, number_of_inference_steps_2, api_name='/upscale256')
+        
+        with open(result_path, 'rb') as f:
+            await ctx.reply('Here is the result of the second stage', file=discord.File(f, 'result.png'))
+
+#---------------------------------------------------------------------------------------------------------------------------------------------------
+                
             
 
 DISCORD_TOKEN = os.environ.get("GRADIOTEST_TOKEN", None)