Spaces:
Runtime error
Runtime error
import asyncio | |
import argparse | |
import pathlib | |
import os | |
import discord | |
from discord.ext import commands | |
import gradio as gr | |
import anyio | |
captioner = gr.Interface.load("olivierdehaene/git-large-coco", src="spaces") | |
lock = asyncio.Lock() | |
bot = commands.Bot("", intents=discord.Intents(messages=True, guilds=True)) | |
async def on_ready(): | |
print(f"Logged in as {bot.user}") | |
print(f"Running in {len(bot.guilds)} servers...") | |
async def run_prediction(img): | |
output = await anyio.to_thread.run_sync(captioner, img) | |
return output | |
def remove_tags(content: str) -> str: | |
content = content.replace("<@1040198143695933501>", "") | |
content = content.replace("<@1057338428938788884>", "") | |
return content.strip() | |
async def send_file_or_text(channel, file_or_text: str): | |
# if the file exists, send as a file | |
if pathlib.Path(str(file_or_text)).exists(): | |
with open(file_or_text, "rb") as f: | |
return await channel.send(file=discord.File(f)) | |
else: | |
return await channel.send(file_or_text) | |
async def make_prediction(message: discord.Message, content: str): | |
predictions = await run_prediction(message.attachments[0].url) | |
print(message.attachments[0].url) | |
if isinstance(predictions, (tuple, list)): | |
for p in predictions: | |
await send_file_or_text(message.channel, p) | |
else: | |
await send_file_or_text(message.channel, predictions) | |
return | |
async def on_message(message: discord.Message): | |
if message.author == bot.user: | |
return | |
if message.content: | |
content = remove_tags(message.content) | |
await make_prediction(message, content) | |
if __name__ == "__main__": | |
parser = argparse.ArgumentParser() | |
parser.add_argument( | |
"--token", | |
type=str, | |
help="API key for the Discord bot. You can set this to your Discord token if you'd like to make your own clone of the Gradio Bot.", | |
required=False, | |
default=os.getenv("DISCORD_TOKEN", ""), | |
) | |
args = parser.parse_args() | |
bot.run(args.token) |