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 # fix? # unstick
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')
DISCORD_TOKEN = os.environ.get("GRADIOTEST_TOKEN", None)

df = Client("huggingface-projects/IF", DFIF_TOKEN)
jojogan = Client("akhaliq/JoJoGAN", DFIF_TOKEN)


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

#---------------------------------------------------------------------------------------------------------------------------------------------- 
# jojo
@bot.command
async def jojo(ctx): # needs safety, easier to do safety function since they'll all have similar constraints
    
    await ctx.message.add_reaction('👍')
    thread = await ctx.message.create_thread(name=f'{ctx.author} Jojo Thread ')
    
    style = None
    style = 'JoJo'

    await thread.send(f'{ctx.author.mention}Generating JoJo image in thread, yare yare daze...')

    if style:
        if ctx.message.attachments:
            attachment = ctx.message.attachments[0]
            im = jojogan.predict(attachment.url, style)
            await ctx.message.reply(f'Here is the {style} version of it', file=discord.File(im))
        else:
            await ctx.message.channel.send("No attachments to be found...Can't animify dat! Try sending me an image 😉")
#----------------------------------------------------------------------------------------------------------------------------------------------     
# 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

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


            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'{i+1}{partialpath}.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'{i+1}{partialpath}.png'))
                    await asyncio.sleep(1) 
    
            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 -> posting images in thread')
            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') need to fix
        #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    

    emoji = reaction.emoji

    thread = reaction.message.channel
    tbd = reaction.message.content
    #index = 0 decided by index in front of partialpath 


    # magic begins
    attachment = reaction.message.attachments[0]
    image_name = attachment.filename
    # we know image_name will be something like 1tmpgtv4qjix.png
    # remove .png first
    indexpartialpath = image_name[:-4] # should be 1tmpgtv4qjix 
    # extract index as an integer (dfif2 needs integer)
    index = int(indexpartialpath[0]) - 1# should be 1
    # extract partialpath
    partialpath = indexpartialpath[1:] # should be tmpgtv4qjix
    # add /tmp/ to partialpath, save as new variable
    fullpath = "/tmp/" + partialpath # should be /tmp/tmpgtv4qjix
    
    if emoji == '👍':
        if reaction.message.attachments:
            if user.id == reaction.message.mentions[0].id:  # all we care about is upscaling whatever image this is 
                stage_1_result_path = fullpath
                index = index
                await dfif2(index, stage_1_result_path, thread)

        

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