deepfloydif.py

import discord
from discord import app_commands
import gradio as gr
from gradio_client import Client
import os
import threading 
import json
import random
from PIL import Image
import asyncio
import glob

MY_GUILD_ID = 1077674588122648679 if os.getenv("TEST_ENV", False) else 879548962464493619
MY_GUILD = discord.Object(id=MY_GUILD_ID)
HF_TOKEN = os.getenv('HF_TOKEN')
DISCORD_TOKEN = os.environ.get("DISCORD_TOKEN", Nonedeepfloyd_client = Client("huggingface-projects/IF", HF_TOKEN)

BOT_USER_ID = 1086256910572986469
DEEPFLOYD_CHANNEL_ID = 1121834257959092234

#-------------------------------------------------------------------------------------------------------------------------------------  
# deepfloydif stage 1 generation
def inference(prompt):
    negative_prompt = ''
    seed = random.randint(0, 1000)
    #seed = 1
    number_of_images = 4
    guidance_scale = 7
    custom_timesteps_1 = 'smart50'
    number_of_inference_steps = 50
    
    stage_1_results, stage_1_param_path, stage_1_result_path = df.predict(
        prompt, negative_prompt, seed, number_of_images, guidance_scale, custom_timesteps_1, number_of_inference_steps, api_name='/generate64')
    
    return [stage_1_results, stage_1_param_path, stage_1_result_path]        
#------------------------------------------------------------------------------------------------------------------------------------- 
async def try_deepfloydif(interaction, prompt):
    thread = None
    try:
        global BOT_USER_ID
        global DEEPFLOYD_CHANNEL_ID
        if interaction.user.id != BOT_USER_ID:
            if interaction.channel.id == DEEPFLOYD_CHANNEL_ID:
                await interaction.response.send_message(f"Working on it!")
                channel = interaction.channel
                message = await channel.send(f"DeepfloydIF Thread")
                await message.add_reaction('🔁')
                thread = await message.create_thread(name=f'{prompt}', auto_archive_duration=60) 
                await thread.send(f"[DISCLAIMER: HuggingBot is a **highly experimental** beta feature; Additional information" \
                f" on the DeepfloydIF model can be found here: https://huggingface.co/spaces/DeepFloyd/IF")

                dfif_command_message_id = message.id # used for updating the 'status' of our generations using reaction emojis
                
                await thread.send(f'{interaction.user.mention}Generating images in thread, can take ~1 minute...')

                # generation
                loop = asyncio.get_running_loop()
                result = await loop.run_in_executor(None, inference, prompt)  
                stage_1_results = result[0]
                stage_1_result_path = result[2]
                partialpath = stage_1_result_path[5:]
                png_files = list(glob.glob(f"{stage_1_results}/**/*.png"))

                img1 = None
                img2 = None
                img3 = None
                img4 = None

                if png_files:
                    png_file_index = 0
                    images = load_image(png_files, stage_1_results, png_file_index) 
                    img1 = images[0]
                    img2 = images[1]
                    img3 = images[2]
                    img4 = images[3]

                    combined_image = Image.new('RGB', (img1.width * 2, img1.height * 2))
                
                    combined_image.paste(img1, (0, 0))
                    combined_image.paste(img2, (img1.width, 0))
                    combined_image.paste(img3, (0, img1.height))
                    combined_image.paste(img4, (img1.width, img1.height))
                
                    combined_image_path = os.path.join(stage_1_results, f'{partialpath}{dfif_command_message_id}.png')
                    combined_image.save(combined_image_path)

                    with open(combined_image_path, 'rb') as f:
                        combined_image_dfif = await thread.send(f'{interaction.user.mention}React with the image number you want to upscale!', file=discord.File(
                            f, f'{partialpath}{dfif_command_message_id}.png')) # named something like: tmpgtv4qjix1111269940599738479.png                

                    emoji_list = ['↖️', '↗️', '↙️', '↘️']
                    await react1234(emoji_list, combined_image_dfif)
                        
                    await message.remove_reaction('🔁', client.user)
                    await message.add_reaction('✔️')

    except Exception as e:
        print(f"Error: {e}")
#-------------------------------------------------------------------------------------------------------------------------------------   
def load_image(png_files, stage_1_results, png_file_index):
    for file in png_files:
        png_file = png_files[png_file_index]
        png_path = os.path.join(stage_1_results, png_file)
        if png_file_index == 0:
            img1 = Image.open(png_path)
        if png_file_index == 1:
            img2 = Image.open(png_path)
        if png_file_index == 2:
            img3 = Image.open(png_path)
        if png_file_index == 3:
            img4 = Image.open(png_path)
        png_file_index = png_file_index + 1
    return [img1, img2, img3, img4]
#-------------------------------------------------------------------------------------------------------------------------------------