Zack
2 years ago
parent
3bc51e4046
commit
49350f308c
@ -0,0 +1,197 @@
|
|||||||
|
import asyncio
|
||||||
|
import argparse
|
||||||
|
from collections import Counter
|
||||||
|
import json
|
||||||
|
import pathlib
|
||||||
|
import re
|
||||||
|
|
||||||
|
|
||||||
|
import discord
|
||||||
|
from discord.ext import commands
|
||||||
|
import gradio as gr
|
||||||
|
from gradio import utils
|
||||||
|
import requests
|
||||||
|
|
||||||
|
from typing import Dict, List
|
||||||
|
|
||||||
|
from utils import *
|
||||||
|
|
||||||
|
|
||||||
|
lock = asyncio.Lock()
|
||||||
|
|
||||||
|
bot = commands.Bot("", intents=discord.Intents(messages=True, guilds=True))
|
||||||
|
|
||||||
|
|
||||||
|
GUILD_SPACES_FILE = "guild_spaces.pkl"
|
||||||
|
|
||||||
|
|
||||||
|
if pathlib.Path(GUILD_SPACES_FILE).exists():
|
||||||
|
guild_spaces = read_pickle_file(GUILD_SPACES_FILE)
|
||||||
|
assert isinstance(guild_spaces, dict), f"{GUILD_SPACES_FILE} in invalid format."
|
||||||
|
guild_blocks = {}
|
||||||
|
delete_keys = []
|
||||||
|
for k, v in guild_spaces.items():
|
||||||
|
try:
|
||||||
|
guild_blocks[k] = gr.Interface.load(v, src="spaces")
|
||||||
|
except ValueError:
|
||||||
|
delete_keys.append(k)
|
||||||
|
for k in delete_keys:
|
||||||
|
del guild_spaces[k]
|
||||||
|
else:
|
||||||
|
guild_spaces: Dict[int, str] = {}
|
||||||
|
guild_blocks: Dict[int, gr.Blocks] = {}
|
||||||
|
|
||||||
|
|
||||||
|
HASHED_USERS_FILE = "users.pkl"
|
||||||
|
|
||||||
|
if pathlib.Path(HASHED_USERS_FILE).exists():
|
||||||
|
hashed_users = read_pickle_file(HASHED_USERS_FILE)
|
||||||
|
assert isinstance(hashed_users, list), f"{HASHED_USERS_FILE} in invalid format."
|
||||||
|
else:
|
||||||
|
hashed_users: List[str] = []
|
||||||
|
|
||||||
|
|
||||||
|
@bot.event
|
||||||
|
async def on_ready():
|
||||||
|
print(f"Logged in as {bot.user}")
|
||||||
|
print(f"Running in {len(bot.guilds)} servers...")
|
||||||
|
|
||||||
|
|
||||||
|
async def run_prediction(space: gr.Blocks, *inputs):
|
||||||
|
inputs = list(inputs)
|
||||||
|
fn_index = 0
|
||||||
|
processed_inputs = space.serialize_data(fn_index=fn_index, inputs=inputs)
|
||||||
|
batch = space.dependencies[fn_index]["batch"]
|
||||||
|
|
||||||
|
if batch:
|
||||||
|
processed_inputs = [[inp] for inp in processed_inputs]
|
||||||
|
|
||||||
|
outputs = await space.process_api(
|
||||||
|
fn_index=fn_index, inputs=processed_inputs, request=None, state={}
|
||||||
|
)
|
||||||
|
outputs = outputs["data"]
|
||||||
|
|
||||||
|
if batch:
|
||||||
|
outputs = [out[0] for out in outputs]
|
||||||
|
|
||||||
|
processed_outputs = space.deserialize_data(fn_index, outputs)
|
||||||
|
processed_outputs = utils.resolve_singleton(processed_outputs)
|
||||||
|
|
||||||
|
return processed_outputs
|
||||||
|
|
||||||
|
|
||||||
|
async def display_stats(message: discord.Message):
|
||||||
|
await message.channel.send(
|
||||||
|
f"Running in {len(bot.guilds)} servers\n"
|
||||||
|
f"Total # of users: {len(hashed_users)}\n"
|
||||||
|
f"------------------"
|
||||||
|
)
|
||||||
|
await message.channel.send(f"Most popular spaces:")
|
||||||
|
# display the top 10 most frequently occurring strings and their counts
|
||||||
|
spaces = guild_spaces.values()
|
||||||
|
counts = Counter(spaces)
|
||||||
|
for space, count in counts.most_common(10):
|
||||||
|
await message.channel.send(f"- {space}: {count}")
|
||||||
|
|
||||||
|
|
||||||
|
async def load_space(guild: discord.Guild, message: discord.Message, content: str):
|
||||||
|
iframe_url = (
|
||||||
|
requests.get(f"https://huggingface.co/api/spaces/{content}/host")
|
||||||
|
.json()
|
||||||
|
.get("host")
|
||||||
|
)
|
||||||
|
if iframe_url is None:
|
||||||
|
return await message.channel.send(
|
||||||
|
f"Space: {content} not found. If you'd like to make a prediction, enclose the inputs in quotation marks."
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
await message.channel.send(
|
||||||
|
f"Loading Space: https://huggingface.co/spaces/{content}..."
|
||||||
|
)
|
||||||
|
interface = gr.Interface.load(content, src="spaces")
|
||||||
|
guild_spaces[guild.id] = content
|
||||||
|
guild_blocks[guild.id] = interface
|
||||||
|
asyncio.create_task(update_pickle_file(guild_spaces, GUILD_SPACES_FILE))
|
||||||
|
if len(content) > 32 - len(f"{bot.name} []"): # type: ignore
|
||||||
|
nickname = content[: 32 - len(f"{bot.name} []") - 3] + "..." # type: ignore
|
||||||
|
else:
|
||||||
|
nickname = content
|
||||||
|
nickname = f"{bot.name} [{nickname}]" # type: ignore
|
||||||
|
await guild.me.edit(nick=nickname)
|
||||||
|
await message.channel.send(
|
||||||
|
"Ready to make predictions! Type in your inputs and enclose them in quotation marks."
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def disconnect_space(bot: commands.Bot, guild: discord.Guild):
|
||||||
|
guild_spaces.pop(guild.id, None)
|
||||||
|
guild_blocks.pop(guild.id, None)
|
||||||
|
asyncio.create_task(update_pickle_file(guild_spaces, GUILD_SPACES_FILE))
|
||||||
|
await guild.me.edit(nick=bot.name) # type: ignore
|
||||||
|
|
||||||
|
|
||||||
|
async def make_prediction(guild: discord.Guild, message: discord.Message, content: str):
|
||||||
|
if guild.id in guild_spaces:
|
||||||
|
params = re.split(r' (?=")', content)
|
||||||
|
params = [p.strip("'\"") for p in params]
|
||||||
|
space = guild_blocks[guild.id]
|
||||||
|
predictions = await run_prediction(space, *params)
|
||||||
|
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
|
||||||
|
else:
|
||||||
|
await message.channel.send(
|
||||||
|
"No Space is currently running. Please type in the name of a Hugging Face Space name first, e.g. abidlabs/en2fr"
|
||||||
|
)
|
||||||
|
await guild.me.edit(nick=bot.name) # type: ignore
|
||||||
|
|
||||||
|
|
||||||
|
@bot.event
|
||||||
|
async def on_message(message: discord.Message):
|
||||||
|
if message.author == bot.user:
|
||||||
|
return
|
||||||
|
h = hash_user_id(message.author.id)
|
||||||
|
if h not in hashed_users:
|
||||||
|
hashed_users.append(h)
|
||||||
|
asyncio.create_task(update_pickle_file(hashed_users, HASHED_USERS_FILE))
|
||||||
|
else:
|
||||||
|
if message.content:
|
||||||
|
content = remove_tags(message.content)
|
||||||
|
guild = message.channel.guild
|
||||||
|
assert guild, "Message not sent in a guild."
|
||||||
|
|
||||||
|
if content.strip() == "exit":
|
||||||
|
await disconnect_space(bot, guild)
|
||||||
|
elif content.strip() == "stats":
|
||||||
|
await display_stats(message)
|
||||||
|
elif content.startswith('"') or content.startswith("'"):
|
||||||
|
await make_prediction(guild, message, content)
|
||||||
|
else:
|
||||||
|
await load_space(guild, 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="",
|
||||||
|
)
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
if args.token.strip():
|
||||||
|
discord_token = args.token
|
||||||
|
bot.env = "staging" # type: ignore
|
||||||
|
bot.name = "StagingBot" # type: ignore
|
||||||
|
else:
|
||||||
|
with open("secrets.json") as fp:
|
||||||
|
discord_token = json.load(fp)["discord_token"]
|
||||||
|
bot.env = "prod" # type: ignore
|
||||||
|
bot.name = "GradioBot" # type: ignore
|
||||||
|
|
||||||
|
bot.run(discord_token)
|
||||||
@ -0,0 +1,41 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import pickle
|
||||||
|
import hashlib
|
||||||
|
import pathlib
|
||||||
|
from typing import Dict, List
|
||||||
|
|
||||||
|
import discord
|
||||||
|
|
||||||
|
lock = asyncio.Lock()
|
||||||
|
|
||||||
|
|
||||||
|
async def update_pickle_file(data: Dict | List, file_path: str):
|
||||||
|
async with lock:
|
||||||
|
with open(file_path, "wb") as fp:
|
||||||
|
pickle.dump(data, fp)
|
||||||
|
|
||||||
|
|
||||||
|
def read_pickle_file(file_path: str):
|
||||||
|
with open(file_path, "rb") as fp:
|
||||||
|
return pickle.load(fp)
|
||||||
|
|
||||||
|
|
||||||
|
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)
|
||||||
|
|
||||||
|
|
||||||
|
def remove_tags(content: str) -> str:
|
||||||
|
content = content.replace("<@1040198143695933501>", "")
|
||||||
|
content = content.replace("<@1057338428938788884>", "")
|
||||||
|
return content.strip()
|
||||||
|
|
||||||
|
|
||||||
|
def hash_user_id(user_id: int) -> str:
|
||||||
|
return hashlib.sha256(str(user_id).encode("utf-8")).hexdigest()
|
||||||
@ -1,21 +0,0 @@
|
|||||||
MIT License
|
|
||||||
|
|
||||||
Copyright (c) 2023 pliny
|
|
||||||
|
|
||||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
||||||
of this software and associated documentation files (the "Software"), to deal
|
|
||||||
in the Software without restriction, including without limitation the rights
|
|
||||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
||||||
copies of the Software, and to permit persons to whom the Software is
|
|
||||||
furnished to do so, subject to the following conditions:
|
|
||||||
|
|
||||||
The above copyright notice and this permission notice shall be included in all
|
|
||||||
copies or substantial portions of the Software.
|
|
||||||
|
|
||||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
||||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
||||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
||||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
||||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
||||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
||||||
SOFTWARE.
|
|
||||||
@ -1,62 +0,0 @@
|
|||||||
from flask import Flask, request, jsonify
|
|
||||||
import openai
|
|
||||||
import logging
|
|
||||||
from dalle3 import Dalle
|
|
||||||
import os
|
|
||||||
import gradio as gr
|
|
||||||
import requests
|
|
||||||
from PIL import Image
|
|
||||||
from io import BytesIO
|
|
||||||
import dotenv
|
|
||||||
|
|
||||||
app = Flask(__name__)
|
|
||||||
|
|
||||||
dotenv.load_dotenv(".env")
|
|
||||||
|
|
||||||
# Initialize OpenAI API, INPUT YOUR OWN OPENAI KEY
|
|
||||||
openai.api_key = os.getenv("OPENAI_API_KEY")
|
|
||||||
|
|
||||||
# Initialize DALLE3 API, INPUT YOUR OWN COOKIE
|
|
||||||
cookie = os.getenv("DALLE_COOKIE")
|
|
||||||
dalle = Dalle(cookie)
|
|
||||||
|
|
||||||
|
|
||||||
def interpret_text_with_gpt(text):
|
|
||||||
model_engine = "text-davinci-002"
|
|
||||||
panel_instructions = "Create a comic panel where"
|
|
||||||
refined_prompt = f"{panel_instructions} {text}"
|
|
||||||
|
|
||||||
response = openai.Completion.create(
|
|
||||||
engine=model_engine,
|
|
||||||
prompt=refined_prompt,
|
|
||||||
max_tokens=100
|
|
||||||
)
|
|
||||||
|
|
||||||
final_prompt = response.choices[0].text.strip()
|
|
||||||
return final_prompt
|
|
||||||
|
|
||||||
def generate_images_with_dalle(refined_prompt):
|
|
||||||
dalle.create(refined_prompt)
|
|
||||||
urls = dalle.get_urls()
|
|
||||||
return urls
|
|
||||||
|
|
||||||
def gradio_interface(text):
|
|
||||||
refined_prompt = interpret_text_with_gpt(text)
|
|
||||||
comic_panel_urls = generate_images_with_dalle(refined_prompt)
|
|
||||||
|
|
||||||
output = []
|
|
||||||
for i, url in enumerate(comic_panel_urls):
|
|
||||||
response = requests.get(url)
|
|
||||||
img = Image.open(BytesIO(response.content))
|
|
||||||
caption = f"Caption for panel {i+1}"
|
|
||||||
output.append((img, caption))
|
|
||||||
|
|
||||||
return output
|
|
||||||
|
|
||||||
iface = gr.Interface(
|
|
||||||
fn=gradio_interface,
|
|
||||||
inputs=["text"],
|
|
||||||
outputs=[gr.outputs.Image(type="pil", label="Comic Panels"), "text"]
|
|
||||||
)
|
|
||||||
|
|
||||||
iface.launch()
|
|
||||||
Loading…
Reference in new issue