|
import discord |
|
import os |
|
import gradio as gr |
|
import asyncio |
|
import time |
|
from discord.ext import commands |
|
from gradio_client import Client |
|
from PIL import Image |
|
|
|
breaking bot |
|
|
|
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) |
|
|
|
MAX_WORKERS = 5 |
|
|
|
worker_semaphore = asyncio.Semaphore(MAX_WORKERS) |
|
|
|
async def jojo(ctx): |
|
start_time = time.time() |
|
style = 'JoJo' |
|
atchurl = 'https://cdn.discordapp.com/attachments/1100458786826747945/1111746037640601610/image.png' |
|
im = jojogan.predict(atchurl, style) |
|
end_time = time.time() |
|
generation_time = end_time - start_time |
|
await ctx.send(f"jojo image generated in {generation_time:.2f} seconds.") |
|
await ctx.send(file=discord.File(im)) |
|
|
|
async def process_requests(queue): |
|
while True: |
|
ctx = await queue.get() |
|
async with worker_semaphore: |
|
await jojo(ctx) |
|
queue.task_done() |
|
|
|
@bot.command() |
|
async def command(ctx, num_requests: int): |
|
start_time = time.time() |
|
queue = asyncio.Queue() |
|
|
|
|
|
worker_tasks = [] |
|
for _ in range(MAX_WORKERS): |
|
task = asyncio.create_task(process_requests(queue)) |
|
worker_tasks.append(task) |
|
|
|
|
|
for _ in range(num_requests): |
|
await queue.put(ctx) |
|
|
|
|
|
await queue.join() |
|
|
|
|
|
for task in worker_tasks: |
|
task.cancel() |
|
|
|
await ctx.send("Command executed.") |
|
end_time = time.time() |
|
generation_time = end_time - start_time |
|
await ctx.send(f"{style} everything done in {generation_time:.2f} seconds.") |
|
|
|
def run_bot(): |
|
bot.run(DISCORD_TOKEN) |
|
|
|
|
|
async def run_gradio_interface(): |
|
def greet(name): |
|
return "Hello " + name + "!" |
|
|
|
demo = gr.Interface(fn=greet, inputs="text", outputs="text") |
|
await demo.launch() |
|
|
|
|
|
async def main(): |
|
loop = asyncio.get_event_loop() |
|
tasks = [ |
|
loop.create_task(run_bot()), |
|
loop.create_task(run_gradio_interface()) |
|
] |
|
await asyncio.wait(tasks) |
|
|
|
|
|
asyncio.run(main()) |
|
|