removed import requests.py

commands/diffuse_cog.py

@@ -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 = hashlib.md5()
-        h.update(decoded_image)
+        # h.update(decoded_image)
-        digest = h.digest()
+        # digest = h.digest()
 
-        if digest == b'i\xac`\xde\xbak\xba\xab{2Z\xcc\tK\xc2~':
+        # return await self.bonk_check(interaction, digest)
-            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
         ### /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))

modules/diffuseapi.py

@@ -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"
         }

slash_commands.py

@@ -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