Merge pull request 'stable-diffuse-api' (#4) from stable-diffuse-api into master

Reviewed-on: #4

.gitignore

@@ -2,3 +2,4 @@
 venv/
 *.secret
 __pycache__/
+jmh/

commands/diffuse_cog.py

@@ -0,0 +1,219 @@
+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 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", ["default"], True, 28)
+        self.user_settings = {}
+        self.prev_prompt = None
+
+
+    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": ["default"],
+                "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="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))
+
+

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))
+

modules/diffuseapi.py

@@ -1,5 +1,7 @@
 import aiohttp
-
+from base64 import b64decode
+from sys import exit
+from pprint import pprint
 
 class DiffuseAPI():
     def __init__(self, url, styles, nsfw_enabled=True, num_steps=28):
@@ -16,7 +18,7 @@ class DiffuseAPI():
     def set_steps(self, steps):
         try:
             new_steps = int(steps)
-            if 0 > new_steps < 50:
+            if new_steps > 0 and new_steps <= 50:
                 self.num_steps = new_steps
                 return True
             return False
@@ -81,21 +83,26 @@ class DiffuseAPI():
     def get_nsfw_filter(self):
         return self.nsfw_enabled
 
-    async def generate_image(self, prompt, neg_prompt=""):
+    async def generate_image(self, prompt, neg_prompt="", payload=None):
-        payload = {
+        if payload is None:
-        "prompt": prompt,
+            payload = {
-        "styles": self.styles,
+            "prompt": prompt,
-        "steps": self.num_steps,
+            "styles": self.styles,
-        "seed": self.seed,
+            "steps": self.num_steps,
-        "n_iter": 1,
+            "seed": self.seed,
-        "height": self.height,
+            "n_iter": 1,
-        "width": self.width,
+            "height": self.height,
-        "negative_prompts": neg_prompt,
+            "width": self.width,
-        "cfg_scale": self.cfg_scale
+            "negative_prompts": neg_prompt,
-        }
+            "cfg_scale": self.cfg_scale
+            }
+
+        payload['prompt'] = prompt
 
         settings = {
-        "filter_nsfw": not self.nsfw_enabled
+#        "filter_nsfw": not self.nsfw_enabled,
+        "enable_pnginfo": False,
+#        "sd_hypernetwork": "anime_3"
         }
 
         override_payload = {
@@ -104,18 +111,36 @@ class DiffuseAPI():
 
         payload.update(override_payload)
         
-        sess = aiohttp.ClientSession(self.url)
+        async with aiohttp.ClientSession(self.url) as session:
-        alive = await sess.head('/')
+            async with session.head('/') as alive:
+                if alive.status != 200:
+                    return None
+            async with session.post("/sdapi/v1/txt2img", json=payload) as image_json:
+                image_data = await image_json.json()
+                return image_data["images"][0]
 
-        if alive.status != 200:
+    async def generate_upscale(self, image):
-            alive.close()
+        payload = {
-            await sess.close()
+            "resize_mode": 0,
-            return None
+            "show_extras_results": True,
-        request  = await sess.post("/sdapi/v1/txt2img", json=payload)
+            "gfpgan_visibility": 0,
-        try:
+            "codeformer_visibility": 0,
-            req_json = await request.json()
+            "codeformer_weight": 0,
-            request.close()
+            "upscaling_resize": 2,
-            await sess.close()
+            "upscaling_resize_w": 512,
-            return req_json["images"][0]
+            "upscaling_resize_h": 1024,
-        except:
+            "upscaling_crop": True,
-            return None
+            "upscaler_1": "R-ESRGAN 4x+ Anime6B",
+            "upscaler_2": "None",
+            "extras_upscaler_2_visibility": 0,
+            "upscale_first": False,
+            "image": image
+        }
+
+        async with aiohttp.ClientSession(self.url) as session:
+            async with session.head('/') as alive:
+                if alive.status != 200:
+                    return None
+            async with session.post("/sdapi/v1/extra-single-image", json=payload) as image_json:
+                image_data = await image_json.json()
+                return image_data["image"]

slash_commands.py

@@ -0,0 +1,47 @@
+import discord
+from discord import app_commands
+from discord.ext import commands
+import os
+import asyncio
+
+
+try:
+    with open("token.secret", 'r') as f:
+        discord_client_token = f.read()
+except FileNotFoundError:
+    print("Cannot locate token.secret please generate a token")
+    exit(69)
+
+
+class aclient(commands.Bot):
+    def __init__(self):
+        activity = discord.Activity(type=discord.ActivityType.listening, name="Waiting for your prompt :3 (Main API)")
+        intents = discord.Intents.default()
+        intents.message_content = True
+        super().__init__(intents=intents, command_prefix="!", activity=activity, help_command=commands.DefaultHelpCommand())
+        
+        self.synced = False
+
+    async def load(self):
+        for filename in os.listdir("./commands"):
+            if filename.endswith(".py"):
+                print(f"Loading commands.{filename[:-3]}")
+                await self.load_extension(f"commands.{filename[:-3]}")
+
+    async def on_ready(self):
+        await self.wait_until_ready()
+        if not self.synced:
+            await tree.sync(guild=discord.Object(id=775297479783874570))
+            self.synced = True
+        print(f"We have logged in as {self.user}.")
+        pass
+
+client = aclient()
+tree = client.tree
+
+async def main():
+    await client.load()
+    await client.start(discord_client_token)
+
+# client.run(discord_client_token)
+asyncio.run(main())

weeeabot.py

@@ -1,9 +1,11 @@
 import discord
 from discord.ext import commands
 import aiohttp
-from base64 import b64decode
+from base64 import b64decode, b64encode
 from sys import exit
 from modules import diffuseapi
+import io
+import hashlib
 
 try:
     with open("token.secret", 'r') as f:
@@ -15,10 +17,39 @@ except FileNotFoundError:
 intents = discord.Intents.default()
 intents.message_content = True
 
-api = diffuseapi.DiffuseAPI("http://localhost:7860", ["default"], True, 28)
+api = diffuseapi.DiffuseAPI("https://art.jurydoak.com", ["Bot"], True, 28)
 
 # client = discord.Client(intents=intents)
 
+class Confirm(discord.ui.View):
+    def __init__(self):
+        super().__init__()
+        self.value = None
+
+    async def callback(self, interaction: discord.Interaction):
+        await interaction.response.send_message("You wish to upscale this image!")
+
+    # When the confirm button is pressed, set the inner value to `True` and
+    # stop the View from listening to more input.
+    # We also send the user an ephemeral message that we're confirming their choice.
+    @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 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()
+
+
+    # This one is similar to the confirmation button except sets the inner value to `False`
+    @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 Settings:
     def __init__(self, nsfw_enabled=True, num_steps=28, ai_seed=-1, url="https://art.jurydoak.com", styles=["Bot"]):
 
@@ -30,79 +61,38 @@ class Settings:
 
 
 main_settings = Settings()
-activity = discord.Activity(type=discord.ActivityType.listening, name="!help")
+activity = discord.Activity(type=discord.ActivityType.listening, name="!help main")
 
 
 bot = commands.Bot(intents=intents, command_prefix="!", activity=activity, help_command=commands.DefaultHelpCommand())
 
-async def send_request(ctx, endpoint, payload):
-    global num_steps
-    global ai_seed
-    sess = aiohttp.ClientSession(main_settings.url)
-    alive = await sess.head('/')
-
-    if alive.status != 200:
-        await ctx.reply("I'm sorry but it appears my brain is offline")
-        alive.close()
-        await sess.close()
-        return {
-            'images': [None]
-        }
-
-    print(payload)
-    request  = await sess.post(endpoint, json=payload)
-    try:
-        req_json = await request.json()
-    except:
-        print(await request.text())
-        return {
-            'images': [None]
-        }
-    request.close()
-    await sess.close()
-    return req_json
-
 @bot.command()
 async def prompt(ctx, *args):
     '''Generate an image with the provided prompt'''
     prompt = " ".join(args)
-    await ctx.reply("Generating your image boo")
+    user_roles = [role.name for role in ctx.message.author.roles]
-    payload = {
-        "prompt": prompt,
-        "styles": main_settings.styles,
-        "steps": main_settings.num_steps,
-        "seed": main_settings.ai_seed,
-        "n_iter": 1,
-        "height": 1024,
-        "negative_prompts": "nsfw, not safe for work, nudity, multiple keyboards",
-        "cfg_scale": 12
-    }
 
-    settings = {
+    async with ctx.typing():
-        "filter_nsfw": not main_settings.nsfw_enabled,
+        if "Bot Ban" not in user_roles:
-        "samples_save": True,
+            # await ctx.reply("Generating your image boo")
-    }
+            image_data = await api.generate_image(prompt)
+            view = Confirm()
+            if image_data is None:
+                await ctx.reply("Something went wrong, please report this to the admin so it can be ignored")
+                return
+            decoded_image = b64decode(image_data)
+            h = hashlib.md5()
+            h.update(decoded_image)
+            digest = h.digest()
 
-    override_payload = {
+            if digest == b'i\xac`\xde\xbak\xba\xab{2Z\xcc\tK\xc2~':
-        "override_settings": settings
+                await ctx.reply("Were no stranger to lewds, but you know the rules, and so do I", file=discord.File("images/404.jpg", "404.jpg"))
-    }
+                return
-
+            print(h.digest())
-    payload.update(override_payload)
+            data = io.BytesIO(decoded_image)
-    image_data = await send_request(ctx, endpoint="/sdapi/v1/txt2img", payload=payload)
+            await ctx.reply("", file=discord.File(data, "_".join(args) + ".png"), view=view)
-
+        else:
-    if image_data['images'][0] is None:
+            await ctx.reply("NO!")
-        await ctx.reply("Something went wrong, please report this to the admin so it can be ignored")
-        return
-
-    with open('/tmp/image.png', 'wb') as f:
-        f.write(b64decode(image_data['images'][0]))
-
-    embed = discord.Embed()
-    upload_file = discord.File("/tmp/image.png", filename="image.png")
-    embed.set_image(url="attachment://image.png")
-    # embed.title = prompt
-
-    await ctx.reply("", file=upload_file, embed=embed)
 
 
 @bot.command()
@@ -114,9 +104,9 @@ async def seed(ctx, arg):
         pass
 
     # global ai_seed
-    main_settings.ai_seed = arg
+    api.set_seed(arg)
 
-    await ctx.reply(f"I have updated the seed to {main_settings.ai_seed} for you my master.")
+    await ctx.reply(f"I have updated the seed to {api.seed} for you my master.")
 
 @bot.command()
 async def steps(ctx, arg):
@@ -131,9 +121,9 @@ async def steps(ctx, arg):
         return
 
     # global num_steps
-    main_settings.num_steps = arg
+    api.num_steps = arg
 
-    await ctx.reply(f"I have updated the steps to {main_settings.num_steps} for you my master.")
+    await ctx.reply(f"I have updated the steps to {api.num_steps} for you my master.")
 
 @bot.command()
 async def settings(ctx):
@@ -141,37 +131,37 @@ async def settings(ctx):
     global ai_seed, steps
     settings = f"""
     ```
-    seed: {main_settings.ai_seed}
+    seed: {api.seed}
-    steps: {main_settings.num_steps}
+    steps: {api.num_steps}
     ```
     """
     await ctx.message.channel.send(settings)
 
 @bot.command()
-async def test(ctx):
+async def upscale(ctx):
-    '''Test function, currently changes the URL and bot settings'''
+    """Upscale the attached image"""
-    if main_settings.url != "http://localhost:7860":
+    orig_image_data = await ctx.message.attachments[0].read()
-        main_settings.url = "http://localhost:7860"
+    new_data = await api.generate_upscale(b64encode(orig_image_data).decode('utf-8'))
-        main_settings.styles = ["default"]
+    data = io.BytesIO(b64decode(new_data))
-        await ctx.reply("Set to fastboi")
+    await ctx.reply("", file=discord.File(data, "upscaled_img.png"))
-    else:
-        main_settings.url = "https://art.jurydoak.com"
-        main_settings.styles = ["Bot"]
-        await ctx.reply("Set to main api")
 
 @bot.command()
-async def test2(ctx):
+async def test(ctx):
-    image_data = await api.generate_image(prompt="cute bot doing bot things")
+    '''Test function, currently changes the URL and bot settings'''
-    if image_data is None:
+    if api.url != "http://localhost:7860":
-        return
+        api.url = "http://localhost:7860"
-    with open('/tmp/image.png', 'wb') as f:
+        api.styles = ["default"]
-        f.write(b64decode(image_data))
+        activity = discord.Activity(type=discord.ActivityType.listening, name="!help fast")
+        await bot.change_presence(activity=activity)
+        await ctx.reply("Set to fastboi")
+    else:
+        api.url = "https://art.jurydoak.com"
+        api.styles = ["Bot"]
+        activity = discord.Activity(type=discord.ActivityType.listening, name="!help main")
+        await bot.change_presence(activity=activity)
+        await ctx.reply("Set to main api")
+
 
-    embed = discord.Embed()
-    upload_file = discord.File("/tmp/image.png", filename="image.png")
-    embed.set_image(url="attachment://image.png")
-    # embed.title = prompt
 
-    await ctx.reply("", file=upload_file, embed=embed)
 
 bot.run(discord_client_token)