From 3a2fb9a3355e13ce3a657bd733fb0f42a788c89d Mon Sep 17 00:00:00 2001 From: Benjamyn Love Date: Wed, 23 Aug 2023 15:06:10 +1000 Subject: [PATCH] Moved game to separate cog --- commands/diffuse_cog.py | 150 +---------------------------------- commands/game_cog.py | 171 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 175 insertions(+), 146 deletions(-) create mode 100644 commands/game_cog.py diff --git a/commands/diffuse_cog.py b/commands/diffuse_cog.py index d76a06c..090ca7d 100644 --- a/commands/diffuse_cog.py +++ b/commands/diffuse_cog.py @@ -8,63 +8,7 @@ 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): @@ -98,18 +42,17 @@ class Confirm(discord.ui.View): class DiffuseCog(commands.Cog): def __init__(self, bot: commands.Bot): self.bot = bot - self.api = DiffuseAPI("https://art.jurydoak.com", ["Bot"], True, 28) + self.api = DiffuseAPI("https://art.jurydoak.com", ["default"], 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 is None: self.user_settings[user_id] = { "prompt": "", - "styles": ["Bot"], + "styles": ["default"], "steps": 28, "seed": -1, "n_iter": 1, @@ -123,92 +66,7 @@ class DiffuseCog(commands.Cog): 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")) + diff --git a/commands/game_cog.py b/commands/game_cog.py new file mode 100644 index 0000000..f6fb5b4 --- /dev/null +++ b/commands/game_cog.py @@ -0,0 +1,171 @@ +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)) +