app.py

import discord
import os
import threading
import gradio as gr
import requests
import json
import random
import time
import re
from discord import Embed, Color
from discord.ext import commands
# test
from gradio_client import Client
from PIL import Image
#from ratelimiter import RateLimiter
#
from datetime import datetime
from pytz import timezone
#
import asyncio



zurich_tz = timezone("Europe/Zurich")

def convert_to_timezone(dt, tz):
    return dt.astimezone(tz).strftime("%Y-%m-%d %H:%M:%S %Z")

DFIF_TOKEN = os.getenv('HF_TOKEN')
df = Client("huggingface-projects/IF", DFIF_TOKEN)
sdlu = Client("huggingface-projects/stable-diffusion-latent-upscaler", DFIF_TOKEN)

DISCORD_TOKEN = os.environ.get("GRADIOTEST_TOKEN", None)
intents = discord.Intents.default()
intents.message_content = True

bot = commands.Bot(command_prefix='!', intents=intents)


#----------------------------------------------------------------------------------------------------------------------------------------------

@bot.event
async def on_ready():
    print('Logged on as', bot.user)
    bot.log_channel = bot.get_channel(1100458786826747945) # 1100458786826747945 = bot-test, 1107006391547342910 = lunarbot server

#---------------------------------------------------------------------------------------------------------------------------------------------- 
# Stage 1
@bot.command()
async def deepfloydif(ctx, *, prompt: str):
    try:
        prompt = prompt.strip()[:100]  # Limit the prompt length to 100 characters
        prompt = re.sub(r'[^\w\s]', '', prompt)  # Remove special characters

        def check_reaction(reaction, user):
            return user == ctx.author and str(reaction.emoji) in ['1️⃣', '2️⃣', '3️⃣', '4️⃣']   

        await ctx.message.add_reaction('👍')
        
        number_of_images = 4
        current_time = int(time.time())
        random.seed(current_time)
        seed = random.randint(0, 2**32 - 1)
        stage_1_results, stage_1_param_path, stage_1_result_path = df.predict(prompt, "blur", seed, number_of_images, 7.0, 'smart100', 50, api_name="/generate64")
        png_files = [f for f in os.listdir(stage_1_results) if f.endswith('.png')]

        if png_files:
            first_png = png_files[0]
            second_png = png_files[1]
            third_png = png_files[2]
            fourth_png = png_files[3]

            first_png_path = os.path.join(stage_1_results, first_png)
            second_png_path = os.path.join(stage_1_results, second_png)
            third_png_path = os.path.join(stage_1_results, third_png)
            fourth_png_path = os.path.join(stage_1_results, fourth_png)

            img1 = Image.open(first_png_path)
            img2 = Image.open(second_png_path)
            img3 = Image.open(third_png_path)
            img4 = Image.open(fourth_png_path)

            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, 'combined_image.png')
            combined_image.save(combined_image_path)

            # Trigger the second stage prediction
            #await dfif2(ctx, stage_1_result_path)


        #1
        sent_message = None

        #2
        # create thread -> send new message inside thread + combined_image -> add reactions -> dfif2
        sent_message = await ctx.reply('Creating thread...')

        #3
        # Create a public thread based on the sent message
        # could replace sent_message with initial user command message? try later
        thread = await sent_message.create_thread(name='Image Upscaling Thread')

        #4
        with open(combined_image_path, 'rb') as f:
            threadmsg = await thread.send('Here is the combined image. React with the image number you want to upscale!', file=discord.File(f, 'combined_image.png'))          
            
        # bot reacts with appropriate emojis to the post, both showing the user what options they have,
        # as well as showing which post to react to.
        for emoji in ['1️⃣', '2️⃣', '3️⃣', '4️⃣']:
            await threadmsg.add_reaction(emoji)            

        reaction, user = await bot.wait_for('reaction_add', check=check_reaction)
        if str(reaction.emoji) == '1️⃣':
            await ctx.reply("Upscaling the first image...")
            index = 0
            await dfif2(ctx, index, stage_1_result_path)
        elif str(reaction.emoji) == '2️⃣':
            await ctx.reply("Upscaling the second image...")
            index = 1
            await dfif2(ctx, index, stage_1_result_path)            
        elif str(reaction.emoji) == '3️⃣':
            await ctx.reply("Upscaling the third image...")
            index = 2
            await dfif2(ctx, index, stage_1_result_path)            
        elif str(reaction.emoji) == '4️⃣':
            await ctx.reply("Upscaling the fourth image...")
            index = 3
            await dfif2(ctx, index, stage_1_result_path)            

    #deepfloydif try/except
    except Exception as e:
        print(f"Error: {e}")
        await ctx.reply('An error occurred while processing your request. Please wait 5 seconds before retrying.')

#----------------------------------------------------------------------------------------------------------------------------
# Stage 2
async def dfif2(ctx, index: int, stage_1_result_path):
    try:
        selected_index_for_stage_2 = index
        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'))
    
    except Exception as e:
        print(f"Error: {e}")
        await ctx.reply('An error occurred while processing stage 2 upscaling. Please try again later.')
#----------------------------------------------------------------------------------------------------------------------------


        

def run_bot():
    bot.run(DISCORD_TOKEN)

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

def greet(name):
    return "Hello " + name + "!"

demo = gr.Interface(fn=greet, inputs="text", outputs="text")
demo.launch()