Compare commits
No commits in common. "master" and "discord-bot-addition" have entirely different histories.
master
...
discord-bo
3
.gitignore
vendored
3
.gitignore
vendored
@ -1,5 +1,4 @@
|
||||
.vscode/
|
||||
venv/
|
||||
*.secret
|
||||
__pycache__/
|
||||
jmh/
|
||||
__pycache__/
|
||||
@ -7,4 +7,4 @@ This aims to be a discord bot that can take a prompt from a discord message and
|
||||
- discord.py: Python discord module
|
||||
|
||||
## Notebooks
|
||||
- testing.ipynb: Testing notebook with examples
|
||||
- testing.ipynb: Testing notebook with examples
|
||||
@ -1,219 +0,0 @@
|
||||
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))
|
||||
|
||||
|
||||
@ -1,171 +0,0 @@
|
||||
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))
|
||||
|
||||
BIN
images/404.jpg
BIN
images/404.jpg
Binary file not shown.
|
Before Width: | Height: | Size: 78 KiB |
@ -1,7 +1,5 @@
|
||||
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):
|
||||
@ -18,7 +16,7 @@ class DiffuseAPI():
|
||||
def set_steps(self, steps):
|
||||
try:
|
||||
new_steps = int(steps)
|
||||
if new_steps > 0 and new_steps <= 50:
|
||||
if 0 > new_steps < 50:
|
||||
self.num_steps = new_steps
|
||||
return True
|
||||
return False
|
||||
@ -42,7 +40,7 @@ class DiffuseAPI():
|
||||
return False
|
||||
except:
|
||||
return False
|
||||
|
||||
|
||||
def set_styles(self, styles):
|
||||
if type(styles) == list:
|
||||
self.styles = styles
|
||||
@ -66,43 +64,38 @@ class DiffuseAPI():
|
||||
return True
|
||||
case _:
|
||||
return False
|
||||
|
||||
|
||||
def get_orientation(self):
|
||||
if self.width == 512 and self.height == 512:
|
||||
return ("square", self.width, self.height)
|
||||
elif self.width == 1024:
|
||||
return ("landscape", self.width, self.height)
|
||||
return ("portrait", self.width, self.height)
|
||||
|
||||
|
||||
def set_nsfw_filter(self, filter_state):
|
||||
if type(filter_state) == bool:
|
||||
self.nsfw_enabled = filter_state
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def get_nsfw_filter(self):
|
||||
return self.nsfw_enabled
|
||||
|
||||
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
|
||||
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
|
||||
}
|
||||
|
||||
settings = {
|
||||
# "filter_nsfw": not self.nsfw_enabled,
|
||||
"enable_pnginfo": False,
|
||||
# "sd_hypernetwork": "anime_3"
|
||||
"filter_nsfw": not self.nsfw_enabled
|
||||
}
|
||||
|
||||
override_payload = {
|
||||
@ -110,37 +103,19 @@ class DiffuseAPI():
|
||||
}
|
||||
|
||||
payload.update(override_payload)
|
||||
|
||||
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/txt2img", json=payload) as image_json:
|
||||
image_data = await image_json.json()
|
||||
return image_data["images"][0]
|
||||
|
||||
async def generate_upscale(self, image):
|
||||
payload = {
|
||||
"resize_mode": 0,
|
||||
"show_extras_results": True,
|
||||
"gfpgan_visibility": 0,
|
||||
"codeformer_visibility": 0,
|
||||
"codeformer_weight": 0,
|
||||
"upscaling_resize": 2,
|
||||
"upscaling_resize_w": 512,
|
||||
"upscaling_resize_h": 1024,
|
||||
"upscaling_crop": True,
|
||||
"upscaler_1": "R-ESRGAN 4x+ Anime6B",
|
||||
"upscaler_2": "None",
|
||||
"extras_upscaler_2_visibility": 0,
|
||||
"upscale_first": False,
|
||||
"image": image
|
||||
}
|
||||
sess = aiohttp.ClientSession(self.url)
|
||||
alive = await sess.head('/')
|
||||
|
||||
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"]
|
||||
if alive.status != 200:
|
||||
alive.close()
|
||||
await sess.close()
|
||||
return None
|
||||
request = await sess.post("/sdapi/v1/txt2img", json=payload)
|
||||
try:
|
||||
req_json = await request.json()
|
||||
request.close()
|
||||
await sess.close()
|
||||
return req_json["images"][0]
|
||||
except:
|
||||
return None
|
||||
|
||||
@ -1,47 +0,0 @@
|
||||
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())
|
||||
File diff suppressed because one or more lines are too long
166
weeeabot.py
166
weeeabot.py
@ -1,11 +1,9 @@
|
||||
import discord
|
||||
from discord.ext import commands
|
||||
import aiohttp
|
||||
from base64 import b64decode, b64encode
|
||||
from base64 import b64decode
|
||||
from sys import exit
|
||||
from modules import diffuseapi
|
||||
import io
|
||||
import hashlib
|
||||
|
||||
try:
|
||||
with open("token.secret", 'r') as f:
|
||||
@ -17,39 +15,10 @@ except FileNotFoundError:
|
||||
intents = discord.Intents.default()
|
||||
intents.message_content = True
|
||||
|
||||
api = diffuseapi.DiffuseAPI("https://art.jurydoak.com", ["Bot"], True, 28)
|
||||
api = diffuseapi.DiffuseAPI("http://localhost:7860", ["default"], 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"]):
|
||||
|
||||
@ -61,38 +30,79 @@ class Settings:
|
||||
|
||||
|
||||
main_settings = Settings()
|
||||
activity = discord.Activity(type=discord.ActivityType.listening, name="!help main")
|
||||
activity = discord.Activity(type=discord.ActivityType.listening, name="!help")
|
||||
|
||||
|
||||
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)
|
||||
user_roles = [role.name for role in ctx.message.author.roles]
|
||||
await ctx.reply("Generating your image boo")
|
||||
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
|
||||
}
|
||||
|
||||
async with ctx.typing():
|
||||
if "Bot Ban" not in user_roles:
|
||||
# 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()
|
||||
settings = {
|
||||
"filter_nsfw": not main_settings.nsfw_enabled,
|
||||
"samples_save": True,
|
||||
}
|
||||
|
||||
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)
|
||||
else:
|
||||
await ctx.reply("NO!")
|
||||
override_payload = {
|
||||
"override_settings": settings
|
||||
}
|
||||
|
||||
payload.update(override_payload)
|
||||
image_data = await send_request(ctx, endpoint="/sdapi/v1/txt2img", payload=payload)
|
||||
|
||||
if image_data['images'][0] is None:
|
||||
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()
|
||||
@ -104,9 +114,9 @@ async def seed(ctx, arg):
|
||||
pass
|
||||
|
||||
# global ai_seed
|
||||
api.set_seed(arg)
|
||||
main_settings.ai_seed = arg
|
||||
|
||||
await ctx.reply(f"I have updated the seed to {api.seed} for you my master.")
|
||||
await ctx.reply(f"I have updated the seed to {main_settings.ai_seed} for you my master.")
|
||||
|
||||
@bot.command()
|
||||
async def steps(ctx, arg):
|
||||
@ -121,9 +131,9 @@ async def steps(ctx, arg):
|
||||
return
|
||||
|
||||
# global num_steps
|
||||
api.num_steps = arg
|
||||
main_settings.num_steps = arg
|
||||
|
||||
await ctx.reply(f"I have updated the steps to {api.num_steps} for you my master.")
|
||||
await ctx.reply(f"I have updated the steps to {main_settings.num_steps} for you my master.")
|
||||
|
||||
@bot.command()
|
||||
async def settings(ctx):
|
||||
@ -131,37 +141,37 @@ async def settings(ctx):
|
||||
global ai_seed, steps
|
||||
settings = f"""
|
||||
```
|
||||
seed: {api.seed}
|
||||
steps: {api.num_steps}
|
||||
seed: {main_settings.ai_seed}
|
||||
steps: {main_settings.num_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)
|
||||
if main_settings.url != "http://localhost:7860":
|
||||
main_settings.url = "http://localhost:7860"
|
||||
main_settings.styles = ["default"]
|
||||
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)
|
||||
main_settings.url = "https://art.jurydoak.com"
|
||||
main_settings.styles = ["Bot"]
|
||||
await ctx.reply("Set to main api")
|
||||
|
||||
@bot.command()
|
||||
async def test2(ctx):
|
||||
image_data = await api.generate_image(prompt="cute bot doing bot things")
|
||||
if image_data is None:
|
||||
return
|
||||
with open('/tmp/image.png', 'wb') as f:
|
||||
f.write(b64decode(image_data))
|
||||
|
||||
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)
|
||||
Loading…
x
Reference in New Issue
Block a user