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


DFIF_TOKEN = os.getenv('HF_TOKEN')
df = Client("huggingface-projects/IF", 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('👍')
        thread = await ctx.message.create_thread(name=f'{ctx.author} Image Upscaling Thread ')
        # create thread -> send new message inside thread + combined_image -> add reactions -> dfif2
        await thread.send(f'{ctx.author.mention}Generating images in thread, can take ~1 minute...')
        
        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)

            '''
        
        if png_files:
            for i, png_file in enumerate(png_files):
                png_file_path = os.path.join(stage_1_results, png_file)
                img = Image.open(png_file_path)
                image_path = os.path.join(stage_1_results, f'image{i+1}.png')
                img.save(image_path)
                with open(image_path, 'rb') as f:
                    await thread.send(f'{ctx.author.mention}Image {i+1}', file=discord.File(f, f'image{i+1}.png'))
                await asyncio.sleep(1)  # Add a delay between posting each image

        await thread.send(f'{ctx.author.mention}React with 👍 to the image you want to upscale!')         

    #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.')
        await ctx.message.add_reaction('❌')

#----------------------------------------------------------------------------------------------------------------------------
# Stage 2
async def dfif2(ctx, index: int, stage_1_result_path, thread):
    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 thread.send(f'{ctx.author.mention}Here is the upscaled image! :) ', file=discord.File(f, 'result.png'))
            #await ctx.reply('Here is the result of the second stage', file=discord.File(f, 'result.png'))
            await ctx.message.add_reaction('✔️')
            

        

    except Exception as e:
        print(f"Error: {e}")
        await ctx.reply('An error occurred while processing stage 2 upscaling. Please try again later.')
        await ctx.message.add_reaction('❌')
#----------------------------------------------------------------------------------------------------------------------------


        

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()