81 lines
2.5 KiB
Python
Executable File
81 lines
2.5 KiB
Python
Executable File
from discord import option
|
|
from discord.ext import commands
|
|
import core_utils
|
|
import discord
|
|
import os
|
|
|
|
|
|
class StableDiffusion(commands.Cog):
|
|
def __init__(self, bot):
|
|
self.bot = bot
|
|
|
|
@commands.has_role("stable-diffuser")
|
|
@commands.slash_command(
|
|
guild_ids=core_utils.my_guilds,
|
|
name="sd",
|
|
description="Pass a prompt and optional negative prompt to stable diffusion",
|
|
)
|
|
@option(
|
|
"positive_prompt",
|
|
description="The positive prompt to pass to stable diffusion",
|
|
input_type="str",
|
|
requred=True,
|
|
)
|
|
@option(
|
|
"negative_prompt",
|
|
description="An optional set of negatives you want to pass to stable diffusion",
|
|
required=False,
|
|
)
|
|
async def sd(
|
|
self, ctx: commands.Context, positive_prompt: str, negative_prompt: str
|
|
):
|
|
|
|
if ctx.channel.name == "stable-diffusion" or ctx.channel.name == "bot-testing":
|
|
import socket
|
|
import stable_diffusion
|
|
|
|
port = "7860"
|
|
ip = "192.168.1.80"
|
|
steps = 20
|
|
|
|
# Send my requests to my gaming computer with the 3080 (if its up)
|
|
if ctx.author.id == core_utils.my_id:
|
|
ip = "192.168.1.29"
|
|
steps = 60
|
|
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
s.settimeout(1)
|
|
try:
|
|
s.connect((ip, int(port)))
|
|
except:
|
|
ip = "192.168.1.80"
|
|
try:
|
|
await ctx.defer()
|
|
original_message = await ctx.followup.send(
|
|
"Please be patient, I'm generating your image"
|
|
)
|
|
file_path, time_taken = await stable_diffusion.generate_image(
|
|
ip=ip,
|
|
port=port,
|
|
positives=positive_prompt,
|
|
negatives=negative_prompt or None,
|
|
steps=steps,
|
|
)
|
|
await original_message.edit(
|
|
content=time_taken,
|
|
file=discord.File(
|
|
file_path,
|
|
filename="unknown.png",
|
|
),
|
|
)
|
|
os.remove(file_path)
|
|
except Exception as e:
|
|
await ctx.respond(
|
|
"Stable diffusion isnt running right now, sorry.\n%s" % e,
|
|
)
|
|
else:
|
|
await ctx.respond("You can only do that in the stable-diffusion channel")
|
|
|
|
|
|
def setup(bot):
|
|
bot.add_cog(StableDiffusion(bot))
|