deepfloydif.py

import asyncio
import glob
import os
import pathlib
import random
import threading

import gradio as gr
import discord
from gradio_client import Client
from PIL import Image
from discord.ext import commands
from discord.ui import Button, View

HF_TOKEN = os.getenv("HF_TOKEN")
deepfloydif_client = Client("huggingface-projects/IF", HF_TOKEN)
DISCORD_TOKEN = os.getenv("DISCORD_TOKEN")
intents = discord.Intents.all()
bot = commands.Bot(command_prefix="/", intents=intents)


@bot.event
async def on_ready():
    print(f"Logged in as {bot.user} (ID: {bot.user.id})")
    synced = await bot.tree.sync()
    print(f"Synced commands: {', '.join([s.name for s in synced])}.")
    print("------")


@bot.hybrid_command(
    name="deepfloydif",
    description="Enter a prompt to generate an image! Can generate realistic text, too!",
)
async def deepfloydif(ctx, prompt: str):
    """DeepfloydIF stage 1 generation"""
    try:
        await deepfloydif_generate64(ctx, prompt)
    except Exception as e:
        print(f"Error: {e}")


def deepfloydif_generate64_inference(prompt):
    """Generates four images based on a prompt"""
    negative_prompt = ""
    seed = random.randint(0, 1000)
    number_of_images = 4
    guidance_scale = 7
    custom_timesteps_1 = "smart50"
    number_of_inference_steps = 50
    (
        stage_1_images,
        stage_1_param_path,
        path_for_upscale256_upscaling,
    ) = deepfloydif_client.predict(
        prompt,
        negative_prompt,
        seed,
        number_of_images,
        guidance_scale,
        custom_timesteps_1,
        number_of_inference_steps,
        api_name="/generate64",
    )
    return [stage_1_images, stage_1_param_path, path_for_upscale256_upscaling]


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_upscale256_upscaling,
        selected_index_for_upscale256,
        seed_2,
        guidance_scale_2,
        custom_timesteps_2,
        number_of_inference_steps_2,
        api_name="/upscale256",
    )
    return result_path


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):
    """Opens images as variables so we can combine them later"""
    results = []
    for file in png_files:
        png_path = os.path.join(stage_1_images, file)
        results.append(Image.open(png_path))
    return results


def combine_images(png_files, stage_1_images, partial_path):
    if os.environ.get("TEST_ENV") == "True":
        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))
    combined_image.paste(images[1], (images[0].width, 0))
    combined_image.paste(images[2], (0, images[0].height))
    combined_image.paste(images[3], (images[0].width, images[0].height))
    combined_image_path = os.path.join(stage_1_images, f"{partial_path}.png")
    combined_image.save(combined_image_path)
    return combined_image_path


async def deepfloydif_generate64(ctx, prompt):
    """DeepfloydIF command (generate images with realistic text using slash commands)"""
    try:
        if ctx.guild.id == 879548962464493619:
            if ctx.channel.id != 1119313215675973714:
                return
        channel = ctx.channel
        # 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} (generating...)")

        loop = asyncio.get_running_loop()
        result = await loop.run_in_executor(None, deepfloydif_generate64_inference, prompt)
        stage_1_images = result[0]
        path_for_upscale256_upscaling = result[2]

        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_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"{interaction.user.mention} (upscaling...)", 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))
                        upscale1024.callback = upscale1024_callback
                        view = View(timeout=None)
                        view.add_item(upscale1024)

                        await interaction.delete_original_response()
                        await channel.send(
                            content=(
                                f"{interaction.user.mention} Here is the upscaled image! Click the button"
                                " 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"{interaction.user.mention} (upscaling...)", 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"{interaction.user.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 channel.send(
                    f"**{prompt}** - {ctx.author.mention} Click a button to upscale! (make larger + enhance quality)",
                    file=discord.File(f, f"{partial_path}.png"),
                    view=view,
                )
        else:
            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_upscale256(index: int, path_for_upscale256_upscaling):
    """upscaling function for images generated using /deepfloydif"""
    try:
        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}")


async def deepfloydif_upscale1024(index: int, path_for_upscale256_upscaling, prompt):
    """upscaling function for images generated using /deepfloydif"""
    try:
        loop = asyncio.get_running_loop()
        result_path = await loop.run_in_executor(
            None, deepfloydif_upscale1024_inference, index, path_for_upscale256_upscaling, prompt
        )
        return result_path

    except Exception as e:
        print(f"Error: {e}")


def run_bot():
    bot.run(DISCORD_TOKEN)


threading.Thread(target=run_bot).start()


welcome_message = """
## Add this bot to your server by clicking this link: 

https://discord.com/api/oauth2/authorize?client_id=1154395078735953930&permissions=51200&scope=bot

## How to use it?

The bot can be triggered via `/deepfloydif` followed by your text prompt.

This will generate images based on the text prompt. You can upscale the images using the buttons up to 16x!

⚠️ Note ⚠️: Please make sure this bot's command does have the same name as another command in your server.

⚠️ Note ⚠️: Bot commands do not work in DMs with the bot as of now.
"""


with gr.Blocks() as demo:
    gr.Markdown(f"""
    # Discord bot of https://huggingface.co/spaces/DeepFloyd/IF
    {welcome_message}
    """)


demo.queue(concurrency_count=100)
demo.queue(max_size=100)
demo.launch()