Home/weeabot
2
1
Fork 0
Code Issues 1 Pull Requests Projects Releases Wiki Activity

Added game

Browse Source
This commit is contained in:
Benjamyn Love 2023-08-21 11:14:45 +10:00
parent 233e8d1e47
commit f4718c83db
4 changed files with 310 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

228
commands/diffuse_cog.py
View File

@ -4,8 +4,69 @@ from discord import app_commands
from discord.ext import commands
from modules.diffuseapi import DiffuseAPI
from base64 import b64decode, b64encode
from typing import Literal
import io
class Game:
def __init__(self, player1, player2, turns_per_player, base_prompt):
self.player1 = player1
self.player2 = player2
self.current_turn = 1
self.player1_turn_count = 0
self.player2_turn_count = 0
self.turns_per_player = turns_per_player
self.prompt = base_prompt
# Adds new_prompt to self.prompt
def prompt_add(self, new_prompt):
self.prompt = f"{self.prompt}, {new_prompt}"
def get_settings(self):
data = {
"prompt": self.prompt,
"styles": ["Bot"],
"steps": 28,
"seed": -1,
"n_iter": 1,
"height": 768,
"width": 512,
"negative_prompts": "",
"cfg_scale": 12
}
return data
def get_players(self):
return [self.player1.id, self.player2.id]
def is_player_ingame(self, player_id):
if player_id in self.get_players():
return True
return False
# return True if player_id in self.get_players() else False
def take_turn(self, new_prompt):
match(self.current_turn):
case 0:
if self.player1_turn_count <= self.turns_per_player:
# Generate new prompt string
self.prompt_add(new_prompt)
# Increment turn and swap player
self.current_turn = 1
self.player1_turn_count += 1
return self.prompt
return None
case 1:
if self.player2_turn_count <= self.turns_per_player:
# Generate new prompt string
self.prompt_add(new_prompt)
# Increment turn and swap player
self.current_turn = 0
self.player2_turn_count += 1
return self.prompt
return None
class Confirm(discord.ui.View):
def __init__(self, api: DiffuseAPI):
super().__init__(timeout=None)
@ -33,15 +94,19 @@ class Confirm(discord.ui.View):
self.stop()
class DiffuseCog(commands.Cog):
def __init__(self, bot: commands.Bot):
self.bot = bot
self.api = DiffuseAPI("https://art.jurydoak.com", ["Bot"], True, 28)
self.user_settings = {}
self.prev_prompt = None
self.games = {}
self.users_with_active_games = []
def _get_user_settings(self, user_id: int):
user_data = self.user_settings.get(user_id)
if user_data == None:
if user_data is None:
self.user_settings[user_id] = {
"prompt": "",
"styles": ["Bot"],
@ -54,10 +119,136 @@ class DiffuseCog(commands.Cog):
"cfg_scale": 12
}
user_data = self.user_settings.get(user_id)
return user_data
@app_commands.command(name="startgame", description="Start a DEGEN GAME with another player")
async def start_game(self, interaction: discord.Integration, player: discord.Member, base_prompt: str):
"""/startgame @username (starts a game)"""
for game in self.games:
if self.games[game].is_player_ingame(interaction.user.id):
await interaction.response.send_message(f"{interaction.user.display_name} is still in a game")
if self.games[game].is_player_ingame(player.id):
await interaction.response.send_message(f"{player.display_name} is still in a game")
game_name = interaction.user.display_name + player.display_name
game = Game(interaction.user, player, 5, base_prompt)
self.games[game_name] = game
await interaction.response.defer()
print(f"SENDING PROMPT: {game.prompt} from User: {interaction.user.display_name}")
# Send the prompt + the users custom payload to the API
image_data = await self.api.generate_image(prompt=game.prompt, payload=game.get_settings())
# Verify that an image was returned
if image_data is None:
await interaction.followup.send("Failed to generate image, please bitch at admin :3")
# ### Bonk check
decoded_image = b64decode(image_data)
# Decode the image to bytes and send it in the response as a discord.File object
data = io.BytesIO(decoded_image)
await interaction.followup.send(f"Game started with {interaction.user.display_name} with {player.display_name} | 5 Turns | {game.prompt} \n Current player {game.player2.display_name}", file=discord.File(data, "prompt.png"))
await interaction.response.send_message(f"Game started with {interaction.user.display_name} with {player.display_name} | 5 Turns | {game.prompt}")
@app_commands.command(name="listgames", description="List running degen games")
async def list_games(self, interaction: discord.Interaction):
"""/listgames (Lists running games)"""
ret = ""
for game in self.games:
ret += f"{self.games[game].player1.display_name} vs. {self.games[game].player2.display_name}: {self.games[game].player1_turn_count + self.games[game].player2_turn_count} turns in\n"
await interaction.response.send_message(ret)
@app_commands.command(name="stopgame", description="Stop the current game you are in.")
async def stop_game(self, interaction: discord.Interaction):
# Make sure this user is in a game
for game in self.games:
if self.games[game].is_player_ingame(interaction.user.id):
del self.games[game]
break
await interaction.response.send_message(f"Game has been removed")
@app_commands.command(name="taketurn", description="Take your turn in degen game")
async def take_turn(self, interaction: discord.Interaction, prompt: str):
# Make sure this user is in a game
game_instance = None
for game in self.games:
if self.games[game].is_player_ingame(interaction.user.id):
game_instance = self.games[game]
break
if game_instance is None:
await interaction.response.send_message(f"{interaction.user.display_name} is not in a game")
# Get the current turn and check the appropriate players id
match(game_instance.current_turn):
case 0:
if interaction.user.id != game_instance.player1.id:
await interaction.response.send_message("It's not your turn.", ephemeral=True)
case 1:
if interaction.user.id != game_instance.player2.id:
await interaction.response.send_message("It's not your turn.")
# We are all gucci for the turn so we defer as we will be generating images now
await interaction.response.defer()
game_instance.take_turn(prompt)
print(f"SENDING PROMPT: {game_instance.prompt} from User: {interaction.user.display_name}")
# Send the prompt + the users custom payload to the API
image_data = await self.api.generate_image(prompt=game_instance.prompt, payload=game_instance.get_settings())
# Verify that an image was returned
if image_data is None:
await interaction.followup.send("Failed to generate image, please bitch at admin :3")
# ### Bonk check
decoded_image = b64decode(image_data)
# Decode the image to bytes and send it in the response as a discord.File object
data = io.BytesIO(decoded_image)
await interaction.followup.send(game_instance.prompt, file=discord.File(data, "prompt.png"))
@app_commands.command(name="pregen", description="Generate an image appending your prompt to the end")
async def pregenerate_image(self, interaction: discord.Interaction, prompt: str):
"""/pregen prompt (Generates an image using the prompt)"""
# This can take a while so we defer the response, this also tells the client the bot is thinking
await interaction.response.defer()
# Create an instance of the Confirm view so we can use it in our response
view = Confirm(self.api)
new_prompt = f"{self.prev_prompt}, {prompt}"
# Get any custom settings the user has entered
user_data = self._get_user_settings(interaction.user.id)
print(f"SENDING PROMPT: {new_prompt} from User: {interaction.user.display_name}")
# Send the prompt + the users custom payload to the API
image_data = await self.api.generate_image(prompt=new_prompt, payload=user_data)
# print(image_data)
# Verify that an image was returned
if image_data is None:
await interaction.followup.send("Failed to generate image, please bitch at admin :3")
# ### Bonk check
decoded_image = b64decode(image_data)
# h = hashlib.md5()
# h.update(decoded_image)
# digest = h.digest()
# return await self.bonk_check(interaction, digest)
### /Bonk check
# Decode the image to bytes and send it in the response as a discord.File object
data = io.BytesIO(decoded_image)
self.prev_prompt = new_prompt
await interaction.followup.send(new_prompt, file=discord.File(data, "prompt.png"), view=view)
@app_commands.command(name="gen", description="Generate an image with the following prompt")
async def generate_image(self, interaction: discord.Interaction, prompt: str):
"""/gen prompt (Generates an image using the prompt)"""
@ -75,26 +266,32 @@ class DiffuseCog(commands.Cog):
# Send the prompt + the users custom payload to the API
image_data = await self.api.generate_image(prompt=prompt, payload=user_data)
# print(image_data)
# Verify that an image was returned
if image_data is None:
await interaction.followup.send("Failed to generate image, please bitch at admin :3")
### Bonk check
# ### Bonk check
decoded_image = b64decode(image_data)
h = hashlib.md5()
h.update(decoded_image)
digest = h.digest()
# h = hashlib.md5()
# h.update(decoded_image)
# digest = h.digest()
if digest == b'i\xac`\xde\xbak\xba\xab{2Z\xcc\tK\xc2~':
await interaction.followup.send("Were no stranger to lewds, but you know the rules, and so do I", file=discord.File("images/404.jpg", "404.jpg"))
return
# return await self.bonk_check(interaction, digest)
### /Bonk check
# Decode the image to bytes and send it in the response as a discord.File object
data = io.BytesIO(decoded_image)
self.prev_prompt = prompt
await interaction.followup.send(prompt, file=discord.File(data, "prompt.png"), view=view)
# async def bonk_check(self, interaction, digest):
# digests = [b'i\xac`\xde\xbak\xba\xab{2Z\xcc\tK\xc2~']
# if digest in digests :
# await interaction.followup.send("Were no stranger to lewds, but you know the rules, and so do I", file=discord.File("images/404.jpg", "404.jpg"))
# return
@app_commands.command(name="steps", description="Set the amount of steps the bot will run")
async def set_steps(self, interaction: discord.Interaction, steps: int):
user_settings = self._get_user_settings(interaction.user.id)
@ -127,7 +324,8 @@ class DiffuseCog(commands.Cog):
await interaction.response.send_message("Set to main api")
@app_commands.command(name="orientation", description="Set the orientation of the generated image")
async def change_orientation(self, interaction: discord.Interaction, orientation: str):
# @app_commands.choices([app_commands.Choice("portrait"), app_commands.Choice("landscape"), app_commands.Choice("square")])
async def change_orientation(self, interaction: discord.Interaction, orientation: Literal["portrait", "landscape", "square"]):
if orientation.lower() not in ["portrait", "landscape", "square"]:
await interaction.response.send_message("You can only set the orientation to portrait, landscape or square")
return
@ -146,6 +344,16 @@ class DiffuseCog(commands.Cog):
await interaction.response.send_message(f"Set the orientation to {orientation}", ephemeral=True)
@app_commands.command(name="seed", description="Set a static seed (-1 is random)")
async def set_seed(self, interaction: discord.Interaction, seed: int):
try:
seed = int(seed)
user_settings = self._get_user_settings(interaction.user.id)
user_settings["seed"] = seed
await interaction.response.send_message(f"Set the seed to {seed}", ephemeral=True)
except:
await interaction.response.send_message("Failed to set the seed")
async def setup(bot: commands.Bot):
await bot.add_cog(DiffuseCog(bot), guild=discord.Object(id=775297479783874570))

2
modules/diffuseapi.py
View File

@ -100,7 +100,7 @@ class DiffuseAPI():
payload['prompt'] = prompt
settings = {
"filter_nsfw": not self.nsfw_enabled,
# "filter_nsfw": not self.nsfw_enabled,
"enable_pnginfo": False,
# "sd_hypernetwork": "anime_3"
}

86
modules/import requests.py Normal file
View File

@ -0,0 +1,86 @@
import requests
import io
import os.path
from PIL import Image
from config import Config
import uuid
import openai
from base64 import b64decode
cfg = Config()
working_directory = "auto_gpt_workspace"
def generate_image(prompt):
filename = str(uuid.uuid4()) + ".jpg"
# DALL-E
if cfg.image_provider == 'dalle':
openai.api_key = cfg.openai_api_key
response = openai.Image.create(
prompt=prompt,
n=1,
size="256x256",
response_format="b64_json",
)
print("Image Generated for prompt:" + prompt)
image_data = b64decode(response["data"][0]["b64_json"])
with open(working_directory + "/" + filename, mode="wb") as png:
png.write(image_data)
return "Saved to disk:" + filename
# STABLE DIFFUSION
elif cfg.image_provider == 'sd':
API_URL = "https://api-inference.huggingface.co/models/CompVis/stable-diffusion-v1-4"
headers = {"Authorization": "Bearer " + cfg.huggingface_api_token}
response = requests.post(API_URL, headers=headers, json={
"inputs": prompt,
})
image = Image.open(io.BytesIO(response.content))
print("Image Generated for prompt:" + prompt)
image.save(os.path.join(working_directory, filename))
return "Saved to disk:" + filename
elif cfg.image_provider == 'sdwu':
API_URL = "http://10.6.9.69:7860"
if payload is None:
payload = {
"prompt": prompt,
}
payload['prompt'] = prompt
settings = {
"enable_pnginfo": False,
}
override_payload = {
"override_settings": settings
}
payload.update(override_payload)
request = requests.post("/sdapi/v1/txt2img", json=payload)
image_data = request.json()
image = Image.open(io.BytesIO(image_data["images"][0]))
print("Image Generated for prompt:" + prompt)
image.save(os.path.join(working_directory, filename))
return "Saved to disk:" + filename
else:
return "No Image Provider Set"

4
slash_commands.py
View File

@ -16,7 +16,9 @@ except FileNotFoundError:
class aclient(commands.Bot):
def __init__(self):
activity = discord.Activity(type=discord.ActivityType.listening, name="Waiting for your prompt :3 (Main API)")
super().__init__(intents=discord.Intents.default(), command_prefix="!", activity=activity)
intents = discord.Intents.default()
intents.message_content = True
super().__init__(intents=intents, command_prefix="!", activity=activity, help_command=commands.DefaultHelpCommand())
self.synced = False