phixxy/sparkytron3000
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

updated stable_diffusion plugin to save images and a log file

Browse source
This commit is contained in:
phixxy 2024-01-17 23:28:40 -08:00
parent e838005102
commit b460bd4821
1 changed files with 7 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

16
plugins/stable_diffusion.py
View file

@ -270,13 +270,6 @@ async def imagine(ctx):
await handle_error(error) await handle_error(error)
for i in r['images']: 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]))) image = Image.open(io.BytesIO(base64.b64decode(i.split(",", 1)[0])))
png_payload = {"image": "data:image/png;base64," + i} png_payload = {"image": "data:image/png;base64," + i}
@ -291,9 +284,9 @@ async def imagine(ctx):
pnginfo.add_text("parameters", response2.get("info")) pnginfo.add_text("parameters", response2.get("info"))
try: try:
if ctx.channel.is_nsfw(): if ctx.channel.is_nsfw():
folder = "users/" + str(ctx.author.id) + '/nsfw/' folder = "tmp/"
else: else:
folder = "users/" + str(ctx.author.id) + '/sfw/' folder = "tmp/sfw/"
except: except:
folder = "users/" + str(ctx.author.id) + '/' folder = "users/" + str(ctx.author.id) + '/'
my_filename = folder + str(time.time_ns()) + ".png" my_filename = folder + str(time.time_ns()) + ".png"
@ -301,7 +294,10 @@ async def imagine(ctx):
with open(my_filename, "rb") as fh: with open(my_filename, "rb") as fh:
f = discord.File(fh, filename=my_filename) 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() await http_session.close()