app.py

import discord
import os
import threading
from threading import active_count
import gradio as gr 
import requests
import json
import random
import time
import re
from discord import Embed, Color
from discord.ext import commands
from gradio_client import Client
from PIL import Image
import asyncio
import shutil # for doing image movement magic 
import concurrent.futures
# reboot
import multiprocessing


#semaphore = asyncio.Semaphore(2)

DFIF_TOKEN = os.getenv('HF_TOKEN')
DISCORD_TOKEN = os.environ.get("GRADIOTEST_TOKEN", None)

jojogan = Client("akhaliq/JoJoGAN", DFIF_TOKEN)

intents = discord.Intents.default()
intents.message_content = True

bot = commands.Bot(command_prefix='!', intents=intents)


@bot.command()
async def active(ctx):
    # get the number of active threads
    activethreadcount = active_count()
    await ctx.message.reply(f'active threads: {activethreadcount}')
    
@bot.command()
async def jojo(ctx):

    # get the number of active threads
    activethreadcount = active_count()
    # report the number of active threads
    print(activethreadcount)
    start_time = time.time()
    attachment = await discord_before(ctx)
    if attachment:
        asyncio.create_task(generation(ctx, attachment, start_time))
        #await generation(ctx, attachment)
        
async def discord_before(ctx):
    if ctx.message.attachments: 
        await asyncio.gather(
            ctx.message.add_reaction('🤖'),
            asyncio.sleep(0.5),
            ctx.message.add_reaction('⌛')  
        )
        attachment = ctx.message.attachments[0]

        # get the number of active threads
        activethreadcount = active_count()
        # report the number of active threads
        print(activethreadcount)
        
        return attachment
    else:
        await asyncio.gather(
            ctx.send(f"{ctx.author.mention} No attachments to be found... Can't animify dat! Try sending me an image 😉"),
            asyncio.sleep(0.5),
            ctx.message.add_reaction('❌') 
        )
        return None

async def generation(ctx, attachment, start_time):
    if attachment:
        style = 'JoJo'

        loop = asyncio.get_running_loop()
        im = loop.run_in_executor(None, jojogan.predict, attachment.url, style)
        #im = await asyncio.get_running_loop().run_in_executor(None, jojogan.predict, attachment.url, style)
        #semaphore.release()

        #im = await asyncio.to_thread(jojogan.predict, attachment.url, style)
        #im = jojogan.predict(attachment.url, style)

        # get the number of active threads
        activethreadcount = active_count()
        # report the number of active threads
        print(activethreadcount)
        
        end_time = time.time()
        total_time = end_time - start_time
        
        await asyncio.gather(
            ctx.send(f'{ctx.author.mention} Jojo image generated in {total_time:.2f} seconds.', file=discord.File(im)),
            asyncio.sleep(0.5),
            ctx.message.remove_reaction('⌛', bot.user),
            asyncio.sleep(0.5),
            ctx.message.add_reaction('✅'),
            asyncio.sleep(0.5)
        )   
        


 


#--------------------------------------------------------
def get_total_threads():
    return multiprocessing.cpu_count()

total_threads = get_total_threads()
print("Total cpus:", total_threads)

# get the number of active threads
activethreadcount = active_count()
# report the number of active threads
print(activethreadcount)
    
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()
#--------------------------------------------------------



'''
 async def discord_after(ctx, im):
    await asyncio.gather(
        ctx.send(f'{ctx.author.mention} Here is the jojo version of it'),
        ctx.send(file=discord.File(im)),
        ctx.message.remove_reaction('⌛', bot.user),
        ctx.message.add_reaction('✅')
    )  

'''


'''
def blocking_io():
    #generate
    #return image

async def main():
    loop = asyncio.get_running_loop()


    
    exeresult = await loop.run_in_executor(
        None, blocking_io)
    print('default thread pool', exeresult)

'''