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
#---------------------------------------------------------------------------------------------------------------------------------------------
async def safetychecks(ctx):
    if ctx.author.bot:
        return False

    # check if the bot is offline
    offline_bot_role_id = 1103676632667017266
    if offline_bot_role_id in bot.roles:
        return False

    # check if the command is in the allowed channel(s)
    channel = ctx.message.channel.id
    if channel != 1100458786826747945:
        return False
        
    # check if the user has the required role(s)   
    guild = 879548962464493619
    required_role_id = 900063512829755413  # @verified for now
    required_role = guild.get_role(required_role_id)
    if required_role not in ctx.author.roles:
        return False
    return True
#---------------------------------------------------------------------------------------------------------------------------------------------- 
# jojo
@bot.command()
async def jojo(ctx):
    # img + face    ✅
    # img + no face ✅
    # no image      ✅
    # no generation ✅
    # responsive?   ✅
    # ratelimits?   ❌
    # safety checks?❌
    # bot no crash  ❌
    
    if await safetychecks:
        try:    
            await ctx.message.add_reaction('🤖') 
            thread = await ctx.message.create_thread(name=f'{ctx.author} Jojo Thread')         
            if ctx.message.attachments:
                await thread.send(f'{ctx.author.mention}Generating images in thread, can take ~1 minute...yare yare, daze ...')  
                
                attachment = ctx.message.attachments[0]
                style = 'JoJo'
                im = jojogan.predict(attachment.url, style)
                #await ctx.message.reply(f'Here is the {style} version of it', file=discord.File(im))
        
                await thread.send(f'{ctx.author.mention}Here is the {style} version of it', file=discord.File(im))
                await ctx.message.add_reaction('✅') # img + face    
            else: # no image
                await thread.send(f"{ctx.author.mention}No attachments to be found...Can't animify dat! Try sending me an image 😉")
                await ctx.message.add_reaction('❌')             
        except Exception as e: # no generation / img + no face
            print(f"Error: {e}")
            await thread.send(f"{ctx.author.mention}Error: {e}")
            await ctx.message.add_reaction('❌')            
        
#---------------------------------------------------------------------------------------------------------------------------------------------- 
# Disney
@bot.command()
async def disney(ctx):
    
    if ctx.message.attachments:
        thread = await ctx.message.create_thread(name=f'{ctx.author} Disney Thread... ') 
        await thread.send(f'{ctx.author.mention}Generating images in thread, can take ~1 minute...')  
        
        attachment = ctx.message.attachments[0]
        style = 'Disney'
        im = jojogan.predict(attachment.url, style)

        await thread.send(f'Here is the {style} version of it', file=discord.File(im))
    else:
        await ctx.message.reply("No attachments to be found...Can't animify dat! Try sending me an image 😉") 
        
#---------------------------------------------------------------------------------------------------------------------------------------------- 
# Spider-Verse
@bot.command()
async def spiderverse(ctx):
    
    if ctx.message.attachments:
        thread = await ctx.message.create_thread(name=f'{ctx.author} Spider-Verse Thread... ') 
        await thread.send(f'{ctx.author.mention}Generating images in thread, can take ~1 minute...')  
        
        attachment = ctx.message.attachments[0]
        style = 'Spider-Verse'
        im = jojogan.predict(attachment.url, style)

        await thread.send(f'Here is the {style} version of it', file=discord.File(im))
    else:
        await ctx.message.reply("No attachments to be found...Can't animify dat! Try sending me an image 😉") 
        
#---------------------------------------------------------------------------------------------------------------------------------------------- 
# sketch
@bot.command()
async def sketch(ctx):
    
    if ctx.message.attachments:
        thread = await ctx.message.create_thread(name=f'{ctx.author} Sketch Thread... ') 
        await thread.send(f'{ctx.author.mention}Generating images in thread, can take ~1 minute...')  
        
        attachment = ctx.message.attachments[0]
        style = 'sketch'
        im = jojogan.predict(attachment.url, style)

        await thread.send(f'Here is the {style} version of it', file=discord.File(im))
    else:
        await ctx.message.reply("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 



    
    if emoji == '👍':
        if reaction.message.attachments:
            if user.id == reaction.message.mentions[0].id:  # all we care about is upscaling whatever image this is 
                
                # 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

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