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": ["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) self.value = None self.api = api async def callback(self, interaction: discord.Interaction): await interaction.response.send_message("You wish to upscale this image!") # When the confirm function is run we will send the image to the upscale endpoint @discord.ui.button(label='Upscale', style=discord.ButtonStyle.green) async def confirm(self, interaction: discord.Interaction, button: discord.ui.Button): image_data = await interaction.message.attachments[0].read() await interaction.response.defer() upscaled_image = await self.api.generate_upscale(b64encode(image_data).decode('utf-8')) data = io.BytesIO(b64decode(upscaled_image)) await interaction.followup.send("", file=discord.File(data, "upscaled.png")) self.stop() # The cancel function will delete the message @discord.ui.button(label='Delete', style=discord.ButtonStyle.grey) async def cancel(self, interaction: discord.Interaction, button: discord.ui.Button): await interaction.message.delete() 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 is None: self.user_settings[user_id] = { "prompt": "", "styles": ["Bot"], "steps": 28, "seed": -1, "n_iter": 1, "height": 768, "width": 512, "negative_prompts": "", "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)""" # 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) # Get any custom settings the user has entered user_data = self._get_user_settings(interaction.user.id) print(f"SENDING PROMPT: {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=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 = 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) try: user_settings['steps'] = int(steps) await interaction.response.send_message("Set steps uWu", ephemeral=True) except Exception: await interaction.response.send_message("Failed to set steps, baka", ephemeral=True) @app_commands.command(name="api", description="Swaps between the main and test API", ) async def swap_api(self, interaction: discord.Interaction): user_settings = self._get_user_settings(interaction.user.id) '''Test function, currently changes the URL and bot settings''' if self.api.url != "http://localhost:7860": self.api.url = "http://localhost:7860" user_settings['style'] = ["default"] self.api.styles = ["default"] activity = discord.Activity(type=discord.ActivityType.listening, name="Waiting for your prompt :3 (Test API)") await self.bot.change_presence(activity=activity) await interaction.response.send_message("Set to fastboi") else: self.api.url = "https://art.jurydoak.com" self.api.styles = ["Bot"] user_settings['style'] = ["Bot"] activity = discord.Activity(type=discord.ActivityType.listening, name="Waiting for your prompt :3 (Main API)") await self.bot.change_presence(activity=activity) await interaction.response.send_message("Set to main api") @app_commands.command(name="orientation", description="Set the orientation of the generated image") # @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 user_data = self._get_user_settings(interaction.user.id) match orientation.lower(): case "portrait": user_data['height'] = 768 user_data['width'] = 512 case "landscape": user_data['height'] = 512 user_data['width'] = 768 case "square": user_data['height'] = 512 user_data['width'] = 512 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))