updated stable_diffusion plugin to save images and a log file

This commit is contained in:
phixxy 2024-01-17 23:28:40 -08:00
parent e838005102
commit b460bd4821

View file

@ -270,13 +270,6 @@ async def imagine(ctx):
await handle_error(error)
for i in r['images']:
if not os.path.isdir("users/" + str(ctx.author.id)):
os.makedirs("users/" + str(ctx.author.id))
if not os.path.isdir("users/" + str(ctx.author.id) + '/nsfw/'):
os.makedirs("users/" + str(ctx.author.id) + '/nsfw/')
if not os.path.isdir("users/" + str(ctx.author.id) + '/sfw/'):
os.makedirs("users/" + str(ctx.author.id) + '/sfw/')
image = Image.open(io.BytesIO(base64.b64decode(i.split(",", 1)[0])))
png_payload = {"image": "data:image/png;base64," + i}
@ -291,9 +284,9 @@ async def imagine(ctx):
pnginfo.add_text("parameters", response2.get("info"))
try:
if ctx.channel.is_nsfw():
folder = "users/" + str(ctx.author.id) + '/nsfw/'
folder = "tmp/"
else:
folder = "users/" + str(ctx.author.id) + '/sfw/'
folder = "tmp/sfw/"
except:
folder = "users/" + str(ctx.author.id) + '/'
my_filename = folder + str(time.time_ns()) + ".png"
@ -301,10 +294,13 @@ async def imagine(ctx):
with open(my_filename, "rb") as fh:
f = discord.File(fh, filename=my_filename)
await ctx.send(file=f)
log_data = {'Author': ctx.author.name, 'Prompt': prompt, "Filename":my_filename}
with open("databases/stable_diffusion.log", 'a') as log_file:
log_file.writelines(log_data)
await ctx.send(f'Generated by {ctx.author.name}\n{prompt}', file=f)
await http_session.close()
@commands.command(
description="Describe",