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

import shutil # for doing image movement magic



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:
        try: 
            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

            #current_time = int(time.time())
            #random.seed(current_time)

            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
            
        except Exception as e:
            print(f"Error: {e}")
            await ctx.reply('stage 1 error -> pre generation')
            await ctx.message.add_reaction('❌')

        
        try:
            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')
        except Exception as e:
            print(f"Error: {e}")
            await ctx.reply('stage 1 error -> during generation')
            await ctx.message.add_reaction('❌')            
            
        try:
            await thread.send(f'{ctx.author.mention}Generating images in thread, can take ~1 minute...')

            png_files = [f for f in os.listdir(stage_1_results) if f.endswith('.png')]
            
            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!')             
        except Exception as e:
            print(f"Error: {e}")
            await ctx.reply('stage 1 error -> threads')
            await ctx.message.add_reaction('❌')
            
        

    #deepfloydif try/except
    except Exception as e:
        print(f"Error: {e}")
        await ctx.reply('An error occurred in stage 1')
        await ctx.message.add_reaction('❌')

#----------------------------------------------------------------------------------------------------------------------------
# Stage 2
async def dfif2(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'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('✅') need to fix this
            

        

    except Exception as e:
        print(f"Error: {e}")
        await ctx.reply('An error occured in stage 2')
        await ctx.message.add_reaction('❌')
#----------------------------------------------------------------------------------------------------------------------------
@bot.event
async def on_reaction_add(reaction, user):
    # safety checks first
    if user.bot:
        return    
        
    # do they have the required (verified) role?
    guild = reaction.message.guild

    # Check if the user has the desired role
    required_role_id = 897376942817419265  # huggingfolks for now
    required_role = guild.get_role(required_role_id)

    if required_role not in user.roles:
        return    

    attachment = reaction.message.attachments[0]
    emoji = reaction.emoji
    
    if emoji == '👍':
        if message.attachments:
            if user.id == message.mentions[0].id:  # all we care about is upscaling whatever image this is 
                stage_1_result_path = "/tmp/s1rp"
                attachment_path = os.path.join(stage_1_result_path, attachment.filename)
                await attachment.save(attachment_path)
                await dfif2(0, attachment_path, thread=reaction.message.reference.resolved)

        

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