Synced repo using 'sync_with_huggingface' Github Action

app.py

@@ -4,7 +4,7 @@ import threading
 import discord
 import gradio as gr
 from audioldm2 import audioldm2_create
-from deepfloydif import deepfloydif_stage_1, deepfloydif_stage_2_react_check
+from deepfloydif import deepfloydif_generate64
 from discord import app_commands
 from discord.ext import commands
 from falcon import continue_falcon, try_falcon
@@ -38,7 +38,7 @@ async def on_ready():
 
 
 @client.hybrid_command(
-    name="testtext",
+    name="falcon",
     with_app_command=True,
     description="Enter some text to chat with the bot! Like this: /falcon Hello, how are you?",
 )
@@ -84,7 +84,7 @@ async def on_message(message):
 async def deepfloydif(ctx, prompt: str):
     """DeepfloydIF stage 1 generation"""
     try:
-        await deepfloydif_stage_1(ctx, prompt, client)
+        await deepfloydif_generate64(ctx, prompt, client)
     except Exception as e:
         print(f"Error: {e}")
 
@@ -113,15 +113,6 @@ async def audioldm2(ctx, prompt: str):
         print(f"Error: (app.py){e}")
 
 
-@client.event
-async def on_reaction_add(reaction, user):
-    """Checks for a reaction in order to call dfif2"""
-    try:
-        await deepfloydif_stage_2_react_check(reaction, user)
-    except Exception as e:
-        print(f"Error: {e} (known error, does not cause issues, low priority)")
-
-
 def run_bot():
     client.run(DISCORD_TOKEN)
 

deepfloydif.py

@@ -8,6 +8,8 @@ import discord
 from gradio_client import Client
 from PIL import Image
 
+from discord.ui import Button, View
+
 HF_TOKEN = os.getenv("HF_TOKEN")
 deepfloydif_client = Client("huggingface-projects/IF", HF_TOKEN)
 
@@ -15,8 +17,8 @@ BOT_USER_ID = 1086256910572986469 if os.getenv("TEST_ENV", False) else 110223665
 DEEPFLOYDIF_CHANNEL_ID = 1121834257959092234 if os.getenv("TEST_ENV", False) else 1119313215675973714
 
 
-def deepfloydif_stage_1_inference(prompt):
-    """Generates an image based on a prompt"""
+def deepfloydif_generate64_inference(prompt):
+    """Generates four images based on a prompt"""
     negative_prompt = ""
     seed = random.randint(0, 1000)
     number_of_images = 4
@@ -26,7 +28,7 @@ def deepfloydif_stage_1_inference(prompt):
     (
         stage_1_images,
         stage_1_param_path,
-        path_for_stage_2_upscaling,
+        path_for_upscale256_upscaling,
     ) = deepfloydif_client.predict(
         prompt,
         negative_prompt,
@@ -37,19 +39,19 @@ def deepfloydif_stage_1_inference(prompt):
         number_of_inference_steps,
         api_name="/generate64",
     )
-    return [stage_1_images, stage_1_param_path, path_for_stage_2_upscaling]
+    return [stage_1_images, stage_1_param_path, path_for_upscale256_upscaling]
 
 
-def deepfloydif_stage_2_inference(index, path_for_stage_2_upscaling):
-    """Upscales one of the images from deepfloydif_stage_1_inference based on the chosen index"""
-    selected_index_for_stage_2 = index
+def deepfloydif_upscale256_inference(index, path_for_upscale256_upscaling):
+    """Upscales one of the images from deepfloydif_generate64_inference based on the chosen index"""
+    selected_index_for_upscale256 = index
     seed_2 = 0
     guidance_scale_2 = 4
     custom_timesteps_2 = "smart50"
     number_of_inference_steps_2 = 50
     result_path = deepfloydif_client.predict(
-        path_for_stage_2_upscaling,
-        selected_index_for_stage_2,
+        path_for_upscale256_upscaling,
+        selected_index_for_upscale256,
         seed_2,
         guidance_scale_2,
         custom_timesteps_2,
@@ -59,10 +61,34 @@ def deepfloydif_stage_2_inference(index, path_for_stage_2_upscaling):
     return result_path
 
 
-async def react_1234(reaction_emojis, combined_image_dfif):
-    """Sets up 4 reaction emojis so the user can choose an image to upscale for deepfloydif"""
-    for emoji in reaction_emojis:
-        await combined_image_dfif.add_reaction(emoji)
+def deepfloydif_upscale1024_inference(index, path_for_upscale256_upscaling, prompt):
+    """Upscales to stage 2, then stage 3"""
+    selected_index_for_upscale256 = index
+    seed_2 = 0  # default seed for stage 2 256 upscaling
+    guidance_scale_2 = 4  # default for stage 2
+    custom_timesteps_2 = "smart50"  # default for stage 2
+    number_of_inference_steps_2 = 50  # default for stage 2
+    negative_prompt = ""  # empty (not used, could add in the future)
+
+    seed_3 = 0  # default for stage 3 1024 upscaling
+    guidance_scale_3 = 9  # default for stage 3
+    number_of_inference_steps_3 = 40  # default for stage 3
+
+    result_path = deepfloydif_client.predict(
+        path_for_upscale256_upscaling,
+        selected_index_for_upscale256,
+        seed_2,
+        guidance_scale_2,
+        custom_timesteps_2,
+        number_of_inference_steps_2,
+        prompt,
+        negative_prompt,
+        seed_3,
+        guidance_scale_3,
+        number_of_inference_steps_3,
+        api_name="/upscale1024",
+    )
+    return result_path
 
 
 def load_image(png_files, stage_1_images):
@@ -76,7 +102,7 @@ def load_image(png_files, stage_1_images):
 
 def combine_images(png_files, stage_1_images, partial_path):
     if os.environ.get("TEST_ENV") == "True":
-        print("Combining images for deepfloydif_stage_1")
+        print("Combining images for deepfloydif_generate64")
     images = load_image(png_files, stage_1_images)
     combined_image = Image.new("RGB", (images[0].width * 2, images[0].height * 2))
     combined_image.paste(images[0], (0, 0))
@@ -88,108 +114,132 @@ def combine_images(png_files, stage_1_images, partial_path):
     return combined_image_path
 
 
-async def deepfloydif_stage_1(ctx, prompt, client):
+async def deepfloydif_generate64(ctx, prompt, client):
     """DeepfloydIF command (generate images with realistic text using slash commands)"""
     try:
         if ctx.author.id != BOT_USER_ID:
             if ctx.channel.id == DEEPFLOYDIF_CHANNEL_ID:
                 if os.environ.get("TEST_ENV") == "True":
-                    print("Safety checks passed for deepfloydif_stage_1")
+                    print("Safety checks passed for deepfloydif_generate64")
+                channel = client.get_channel(DEEPFLOYDIF_CHANNEL_ID)
                 # interaction.response message can't be used to create a thread, so we create another message
-                message = await ctx.send(f"**{prompt}** - {ctx.author.mention}")
-                if len(prompt) > 99:
-                    small_prompt = prompt[:99]
-                else:
-                    small_prompt = prompt
-                thread = await message.create_thread(name=f"{small_prompt}", auto_archive_duration=60)
-                await thread.send(f"{ctx.author.mention} Generating images in thread, can take ~1 minute...")
+                message = await ctx.send(f"**{prompt}** - {ctx.author.mention} <a:loading:1114111677990981692>")
 
                 loop = asyncio.get_running_loop()
-                result = await loop.run_in_executor(None, deepfloydif_stage_1_inference, prompt)
+                result = await loop.run_in_executor(None, deepfloydif_generate64_inference, prompt)
                 stage_1_images = result[0]
-                path_for_stage_2_upscaling = result[2]
+                path_for_upscale256_upscaling = result[2]
 
-                partial_path = pathlib.Path(path_for_stage_2_upscaling).name
+                partial_path = pathlib.Path(path_for_upscale256_upscaling).name
                 png_files = list(glob.glob(f"{stage_1_images}/**/*.png"))
 
                 if png_files:
+                    await message.delete()
                     combined_image_path = combine_images(png_files, stage_1_images, partial_path)
                     if os.environ.get("TEST_ENV") == "True":
-                        print("Images combined for deepfloydif_stage_1")
-                    with open(combined_image_path, "rb") as f:
-                        combined_image_dfif = await thread.send(
-                            f"{ctx.author.mention} React with the image number you want to upscale!",
+                        print("Images combined for deepfloydif_generate64")
+
+                    with Image.open(combined_image_path) as img:
+                        width, height = img.size
+                        new_width = width * 3
+                        new_height = height * 3
+                        resized_img = img.resize((new_width, new_height))
+                        x2_combined_image_path = combined_image_path
+                        resized_img.save(x2_combined_image_path)
+
+                    # making image bigger, more readable
+                    with open(x2_combined_image_path, "rb") as f:  # was combined_image_path
+                        button1 = Button(custom_id="0", emoji="↖")
+                        button2 = Button(custom_id="1", emoji="↗")
+                        button3 = Button(custom_id="2", emoji="↙")
+                        button4 = Button(custom_id="3", emoji="↘")
+
+                        async def button_callback(interaction):
+                            index = int(interaction.data["custom_id"])  # 0,1,2,3
+
+                            await interaction.response.send_message(
+                                f"{ctx.author.mention} <a:loading:1114111677990981692>", ephemeral=True
+                            )
+                            result_path = await deepfloydif_upscale256(index, path_for_upscale256_upscaling)
+
+                            # create and use upscale 1024 button
+                            with open(result_path, "rb") as f:
+                                upscale1024 = Button(
+                                    label="High-quality upscale (x4)", custom_id=str(index)
+                                )  # "0", "1" etc
+                                upscale1024.callback = upscale1024_callback
+                                view = View(timeout=None)
+                                view.add_item(upscale1024)
+
+                                await interaction.delete_original_response()
+                                await channel.send(
+                                    content=(
+                                        f"{ctx.author.mention} Here is the upscaled image! Click to upscale even more!"
+                                    ),
+                                    file=discord.File(f, f"{prompt}.png"),
+                                    view=view,
+                                )
+
+                        async def upscale1024_callback(interaction):
+                            index = int(interaction.data["custom_id"])
+
+                            await interaction.response.send_message(
+                                f"{ctx.author.mention} <a:loading:1114111677990981692>", ephemeral=True
+                            )
+                            result_path = await deepfloydif_upscale1024(index, path_for_upscale256_upscaling, prompt)
+
+                            with open(result_path, "rb") as f:
+                                await interaction.delete_original_response()
+                                await channel.send(
+                                    content=f"{ctx.author.mention} Here's your high-quality x16 image!",
+                                    file=discord.File(f, f"{prompt}.png"),
+                                )
+
+                        button1.callback = button_callback
+                        button2.callback = button_callback
+                        button3.callback = button_callback
+                        button4.callback = button_callback
+
+                        view = View(timeout=None)
+                        view.add_item(button1)
+                        view.add_item(button2)
+                        view.add_item(button3)
+                        view.add_item(button4)
+
+                        # could store this message as combined_image_dfif in case it's useful for future testing
+                        await ctx.send(
+                            f"{ctx.author.mention} Click a button to upscale! (make larger + enhance quality)",
                             file=discord.File(f, f"{partial_path}.png"),
+                            view=view,
                         )
-                    emoji_list = ["↖️", "↗️", "↙️", "↘️"]
-                    await react_1234(emoji_list, combined_image_dfif)
                 else:
-                    await thread.send(f"{ctx.author.mention} No PNG files were found, cannot post them!")
+                    await ctx.send(f"{ctx.author.mention} No PNG files were found, cannot post them!")
 
     except Exception as e:
         print(f"Error: {e}")
 
 
-async def deepfloydif_stage_2_react_check(reaction, user):
-    """Checks for a reaction in order to call dfif2"""
+async def deepfloydif_upscale256(index: int, path_for_upscale256_upscaling):
+    """upscaling function for images generated using /deepfloydif"""
     try:
-        if os.environ.get("TEST_ENV") == "True":
-            print("Running deepfloydif_stage_2_react_check")
-        global BOT_USER_ID
-        global DEEPFLOYDIF_CHANNEL_ID
-        if user.id != BOT_USER_ID:
-            thread = reaction.message.channel
-            thread_parent_id = thread.parent.id
-            if thread_parent_id == DEEPFLOYDIF_CHANNEL_ID:
-                if reaction.message.attachments:
-                    if user.id == reaction.message.mentions[0].id:
-                        attachment = reaction.message.attachments[0]
-                        image_name = attachment.filename
-                        partial_path = image_name[:-4]
-                        full_path = "/tmp/" + partial_path
-                        emoji = reaction.emoji
-                        if emoji == "↖️":
-                            index = 0
-                        elif emoji == "↗️":
-                            index = 1
-                        elif emoji == "↙️":
-                            index = 2
-                        elif emoji == "↘️":
-                            index = 3
-                        path_for_stage_2_upscaling = full_path
-                        thread = reaction.message.channel
-                        await deepfloydif_stage_2(
-                            index,
-                            path_for_stage_2_upscaling,
-                            thread,
-                        )
+        loop = asyncio.get_running_loop()
+        result_path = await loop.run_in_executor(
+            None, deepfloydif_upscale256_inference, index, path_for_upscale256_upscaling
+        )
+        return result_path
+
     except Exception as e:
-        print(f"Error: {e} (known error, does not cause issues, low priority)")
+        print(f"Error: {e}")
 
 
-async def deepfloydif_stage_2(index: int, path_for_stage_2_upscaling, thread):
+async def deepfloydif_upscale1024(index: int, path_for_upscale256_upscaling, prompt):
     """upscaling function for images generated using /deepfloydif"""
     try:
-        if os.environ.get("TEST_ENV") == "True":
-            print("Running deepfloydif_stage_2")
-        if index == 0:
-            position = "top left"
-        elif index == 1:
-            position = "top right"
-        elif index == 2:
-            position = "bottom left"
-        elif index == 3:
-            position = "bottom right"
-        await thread.send(f"Upscaling the {position} image...")
-
-        # run blocking function in executor
         loop = asyncio.get_running_loop()
         result_path = await loop.run_in_executor(
-            None, deepfloydif_stage_2_inference, index, path_for_stage_2_upscaling
+            None, deepfloydif_upscale1024_inference, index, path_for_upscale256_upscaling, prompt
         )
+        return result_path
 
-        with open(result_path, "rb") as f:
-            await thread.send("Here is the upscaled image!", file=discord.File(f, "result.png"))
-        await thread.edit(archived=True)
     except Exception as e:
         print(f"Error: {e}")