biiig changes
This commit is contained in:
parent
90663d9182
commit
aa6d0bee6e
153
commands/diffuse_cog.py
Normal file
153
commands/diffuse_cog.py
Normal file
@ -0,0 +1,153 @@
|
||||
import discord
|
||||
import hashlib
|
||||
from discord import app_commands
|
||||
from discord.ext import commands
|
||||
from modules.diffuseapi import DiffuseAPI
|
||||
from base64 import b64decode, b64encode
|
||||
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", ["Bot"], True, 28)
|
||||
self.user_settings = {}
|
||||
|
||||
def _get_user_settings(self, user_id: int):
|
||||
user_data = self.user_settings.get(user_id)
|
||||
if user_data == 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="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)
|
||||
|
||||
# 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()
|
||||
|
||||
if digest == b'i\xac`\xde\xbak\xba\xab{2Z\xcc\tK\xc2~':
|
||||
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)
|
||||
await interaction.followup.send(prompt, file=discord.File(data, "prompt.png"), view=view)
|
||||
|
||||
@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")
|
||||
async def change_orientation(self, interaction: discord.Interaction, orientation: str):
|
||||
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)
|
||||
|
||||
|
||||
async def setup(bot: commands.Bot):
|
||||
await bot.add_cog(DiffuseCog(bot), guild=discord.Object(id=775297479783874570))
|
||||
|
||||
|
||||
@ -18,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
|
||||
@ -83,22 +83,26 @@ class DiffuseAPI():
|
||||
def get_nsfw_filter(self):
|
||||
return self.nsfw_enabled
|
||||
|
||||
async def generate_image(self, prompt, neg_prompt=""):
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"styles": self.styles,
|
||||
"steps": self.num_steps,
|
||||
"seed": self.seed,
|
||||
"n_iter": 1,
|
||||
"height": self.height,
|
||||
"width": self.width,
|
||||
"negative_prompts": neg_prompt,
|
||||
"cfg_scale": self.cfg_scale
|
||||
}
|
||||
async def generate_image(self, prompt, neg_prompt="", payload=None):
|
||||
if payload is None:
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"styles": self.styles,
|
||||
"steps": self.num_steps,
|
||||
"seed": self.seed,
|
||||
"n_iter": 1,
|
||||
"height": self.height,
|
||||
"width": self.width,
|
||||
"negative_prompts": neg_prompt,
|
||||
"cfg_scale": self.cfg_scale
|
||||
}
|
||||
|
||||
payload['prompt'] = prompt
|
||||
|
||||
settings = {
|
||||
"filter_nsfw": not self.nsfw_enabled,
|
||||
"enable_pnginfo": False
|
||||
"enable_pnginfo": False,
|
||||
# "sd_hypernetwork": "anime_3"
|
||||
}
|
||||
|
||||
override_payload = {
|
||||
@ -122,7 +126,7 @@ class DiffuseAPI():
|
||||
"gfpgan_visibility": 0,
|
||||
"codeformer_visibility": 0,
|
||||
"codeformer_weight": 0,
|
||||
"upscaling_resize": 4,
|
||||
"upscaling_resize": 2,
|
||||
"upscaling_resize_w": 512,
|
||||
"upscaling_resize_h": 1024,
|
||||
"upscaling_crop": True,
|
||||
|
||||
45
slash_commands.py
Normal file
45
slash_commands.py
Normal file
@ -0,0 +1,45 @@
|
||||
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)")
|
||||
super().__init__(intents=discord.Intents.default(), command_prefix="!", activity=activity)
|
||||
|
||||
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())
|
||||
Loading…
x
Reference in New Issue
Block a user