172 lines
7.4 KiB
Python
172 lines
7.4 KiB
Python
import discord
|
|
import hashlib
|
|
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": ["default"],
|
|
"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 GameCog(commands.Cog):
|
|
def __init__(self, bot: commands.Bot):
|
|
self.bot = bot
|
|
self.api = DiffuseAPI("https://art.jurydoak.com", ["default"], True, 28)
|
|
self.games = {}
|
|
|
|
|
|
@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.mention}", 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
|
|
channel = interaction.channel
|
|
await interaction.response.defer(ephemeral=True)
|
|
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"), ephemeral=True)
|
|
if game_instance.current_turn == 0:
|
|
player_name = game_instance.player1.mention
|
|
else:
|
|
player_name = game_instance.player2.mention
|
|
await channel.send(f"Turn taken: {player_name}'s turn")
|
|
|
|
async def setup(bot: commands.Bot):
|
|
await bot.add_cog(GameCog(bot), guild=discord.Object(id=775297479783874570))
|
|
|