161 lines
5.1 KiB
Python
161 lines
5.1 KiB
Python
import discord
|
|
from discord.ext import commands
|
|
import aiohttp
|
|
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:
|
|
discord_client_token = f.read()
|
|
except FileNotFoundError:
|
|
print("Cannot locate token.secret please generate a token")
|
|
exit(69)
|
|
|
|
intents = discord.Intents.default()
|
|
intents.message_content = True
|
|
|
|
api = diffuseapi.DiffuseAPI("https://art.jurydoak.com", ["Bot"], False, 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"]):
|
|
|
|
self.nsfw_enabled = nsfw_enabled
|
|
self.num_steps = num_steps
|
|
self.ai_seed = ai_seed
|
|
self.url = url
|
|
self.styles = styles
|
|
|
|
|
|
main_settings = Settings()
|
|
activity = discord.Activity(type=discord.ActivityType.listening, name="!help main")
|
|
|
|
|
|
bot = commands.Bot(intents=intents, command_prefix="!", activity=activity, help_command=commands.DefaultHelpCommand())
|
|
|
|
@bot.command()
|
|
async def prompt(ctx, *args):
|
|
'''Generate an image with the provided prompt'''
|
|
prompt = " ".join(args)
|
|
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()
|
|
|
|
if digest == b'i\xac`\xde\xbak\xba\xab{2Z\xcc\tK\xc2~':
|
|
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())
|
|
data = io.BytesIO(decoded_image)
|
|
await ctx.reply("", file=discord.File(data, "_".join(args) + ".png"), view=view)
|
|
|
|
|
|
@bot.command()
|
|
async def seed(ctx, arg):
|
|
'''Set the seed for the image generation'''
|
|
try:
|
|
arg = int(arg)
|
|
except:
|
|
pass
|
|
|
|
# global ai_seed
|
|
api.set_seed(arg)
|
|
|
|
await ctx.reply(f"I have updated the seed to {api.seed} for you my master.")
|
|
|
|
@bot.command()
|
|
async def steps(ctx, arg):
|
|
'''Set how many steps the AI will run (max 50)'''
|
|
try:
|
|
arg = int(arg)
|
|
except:
|
|
pass
|
|
|
|
if arg > 50:
|
|
await ctx.reply("I'm sorry Dave, I can't do that")
|
|
return
|
|
|
|
# global num_steps
|
|
api.set_steps(arg)
|
|
|
|
await ctx.reply(f"I have updated the steps to {api.num_steps} for you my master.")
|
|
|
|
@bot.command()
|
|
async def settings(ctx):
|
|
'''See the currently configured settings (BROKEN)'''
|
|
global ai_seed, steps
|
|
settings = f"""
|
|
```
|
|
seed: {api.seed}
|
|
steps: {api.steps}
|
|
```
|
|
"""
|
|
await ctx.message.channel.send(settings)
|
|
|
|
@bot.command()
|
|
async def upscale(ctx):
|
|
"""Upscale the attached image"""
|
|
orig_image_data = await ctx.message.attachments[0].read()
|
|
new_data = await api.generate_upscale(b64encode(orig_image_data).decode('utf-8'))
|
|
data = io.BytesIO(b64decode(new_data))
|
|
await ctx.reply("", file=discord.File(data, "upscaled_img.png"))
|
|
|
|
@bot.command()
|
|
async def test(ctx):
|
|
'''Test function, currently changes the URL and bot settings'''
|
|
if api.url != "http://localhost:7860":
|
|
api.url = "http://localhost:7860"
|
|
api.styles = ["default"]
|
|
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")
|
|
|
|
|
|
|
|
|
|
bot.run(discord_client_token) |