app.py

import discord
import gradio_client
from gradio_client import Client
import gradio as gr
import os
import threading

#for deepfloydif
import requests
import json
import random
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib.image as mpimg




#todos
#alert

# Get Gradio client
jojogan = gradio_client.Client("akhaliq/JoJoGAN")

#token update
DFIF_TOKEN = os.getenv('DFIF_TOKEN')
#deepfloydif
#df = Client("DeepFloyd/IF") #not reliable at the moment
df = Client("huggingface-projects/IF", hf_token=DFIF_TOKEN)


# Set up discord bot 
class MyClient(discord.Client):
    async def on_ready(self):
        print('Logged on as', self.user)

    #----------------------------------------------------------------------------------------------------
    async def on_message(self, message):

        #tldr, bot should run if
        #1) it does not have @offline role
        #2) user has @verified role
        #3) bot is in #bot-test channel
        

        # if the bot has this role, it won't run
        OFFLINE_ROLE_ID = 1103676632667017266  # 1103676632667017266 = @offline-bot
        guild = message.guild
        bot_member = guild.get_member(self.user.id)
        if any(role.id == OFFLINE_ROLE_ID for role in bot_member.roles):
            return
        
        # don't respond to ourselves
        if message.author == self.user:
            return

        # if the message author doesn't have this role, the bot won't run
        REQUIRED_ROLE_ID = 900063512829755413 # 900063512829755413 = @verified
        if not any(role.id == REQUIRED_ROLE_ID for role in message.author.roles):
            return            

        # channels where bot will accept commands
        ALLOWED_CHANNEL_IDS = [1100458786826747945] # 1100458786826747945 = #bot-test
        if message.channel.id not in ALLOWED_CHANNEL_IDS:
            return

        #----------------------------------------------------------------------------------------------------
        #jojogan
        
        if message.content.find("!help") != -1:
            await message.reply("Use !jojo !disney !spidey or !sketch. Have fun!", mention_author=True)

        style = None
        if message.content.startswith('!jojo'):
            style = 'JoJo'
        if message.content.startswith('!disney'):
            style = 'Disney'
        if message.content.startswith('!spidey'):
            style = 'Spider-Verse'
        if message.content.startswith('!sketch'):
            style = 'sketch'

        if style:
            if message.attachments:
                attachment = message.attachments[0]

                # jojogan = gradio_client.Client("akhaliq/JoJoGAN")
                # predict(img, model, api_name="/predict") -> output
                # im = jojogan.predict(img, model) -> output
                # im = jojogan.predict(attachment.url, style) -> output
                
                im = jojogan.predict(attachment.url, style)
                await message.reply(f'Here is the {style} version of it', file=discord.File(im))
            else:
                await message.channel.send("No attachments to be found...Can't animify dat! Try sending me an image 😉")

        #----------------------------------------------------------------------------------------------------
        #deepfloydif
        '''
        if message.content.startswith('!deepfloydif'):
            text_input = message.content[12:].strip()
            if text_input:
                im = deepfloydif.predict(text_input)
                im_bytes = BytesIO()
                im.save(im_bytes, 'PNG')
                im_bytes.seek(0)
                await message.reply(f'Here is the image generated for the input "{text_input}"', file=discord.File(im_bytes, 'output.png'))
            else:
                await message.channel.send("No text input provided. Please provide some text after the !deepfloydif command.")        
        
        '''
        #deepfloydif
        # start with preset values for now
        if message.content.startswith('!deepfloydif'):

            #predict
            stage_1_results, stage_1_param_path, stage_1_result_path = df.predict("gradio written on a wall", "blur", 1,4,7.0, 'smart100',50, api_name="/generate64")
            
            # Assuming stage_1_results contains the path to the directory
            png_files = [f for f in os.listdir(stage_1_results) if f.endswith('.png')]

            # Assuming png_files contains the list of all PNG files in the directory
            if png_files:
                first_png = png_files[0]
                first_png_path = os.path.join(stage_1_results, first_png)

            # Send the image file as a Discord attachment
            with open(first_png_path, 'rb') as f:
                await message.reply(f'Here is the first generated image', file=discord.File(f, 'first_png.png'))   








            '''
            # stage 1
            prompt = "llama"
            negative_prompt = "blue"
            seed = 7  # maybe random seed each time?    seed = 7       seed = random.randint(0, 2**32 - 1)
            number_of_images = 4
            guidance_scale = 7
            custom_timesteps_1 = 'smart100'
            number_of_inference_steps = 100

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

            # stage 2
            selected_index_for_stage_2 = -1
            custom_timesteps_2 = 'smart100' # could reset to smart50 if index was the issue
            seed = 362572064 # again, could randomize this        seed = 362572064      seed = random.randint(0, 2**32 - 1)
            
            # predict(stage_1_result_path, selected_index_for_stage_2, seed, guidance_scale, custom_timesteps_2, number_of_inference_steps, api_name="/upscale256") -> result
            
            img = df.predict(stage_1_result_path, selected_index_for_stage_2, seed, guidance_scale, custom_timesteps_2, number_of_inference_steps, api_name="/upscale256")
    
            # Save the generated image to a file
            img_path = "/tmp/generated_image.png"
            img.save(img_path)
            
            # Send the image file as a Discord attachment
            with open(img_path, 'rb') as f:
                await message.reply(f'Here is the generated image', file=discord.File(f, 'generated_image.png'))                
            
            
            
            
            
            '''
       
            

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

def run_bot():
  client.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()