ldooks/dragon-bot
ldooks/dragon-bot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Adding time taken to stable diffusion renders

Browse Source
This commit is contained in:
Luke Robles 2022-10-14 09:40:44 -07:00
parent 2381957cd1
commit e3339da557
2 changed files with 10 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
app/cogs/actual_utils.py
View File

@ -121,7 +121,7 @@ class ActualUtils(commands.Cog):
original_message = await ctx.followup.send( original_message = await ctx.followup.send(
"Please be patient, I'm generating your image" "Please be patient, I'm generating your image"
) )
file_path = await stable_diffusion.generate_image( file_path, time_taken = await stable_diffusion.generate_image(
ip=ip, ip=ip,
port=port, port=port,
positives=positive_prompt, positives=positive_prompt,
@ -129,7 +129,7 @@ class ActualUtils(commands.Cog):
steps=steps, steps=steps,
) )
await original_message.edit( await original_message.edit(
content="", content=time_taken,
file=discord.File( file=discord.File(
file_path, file_path,
filename="unknown.png", filename="unknown.png",

11
app/stable_diffusion.py
View File

@ -1,6 +1,8 @@
from bs4 import BeautifulSoup
import base64 import base64
import json
import httpx import httpx
import json
import requests
import socket import socket
import tempfile import tempfile
@ -59,10 +61,13 @@ async def generate_image(ip, port, positives, negatives, steps):
client = httpx.AsyncClient() client = httpx.AsyncClient()
response = await client.post(url, headers=headers, json=json_data, timeout=1200) response = await client.post(url, headers=headers, json=json_data, timeout=1200)
imageb64 = response.json()["data"][0][0].split("data:image/png;base64,")[1] imageb64 = response.json()["data"][0][0].split("data:image/png;base64,")[1]
imgdata = base64.b64decode(imageb64) imgdata = base64.b64decode(imageb64)
file, file_path = tempfile.mkstemp() file, file_path = tempfile.mkstemp()
with open(file_path, "wb") as f: with open(file_path, "wb") as f:
f.write(imgdata) f.write(imgdata)
return file_path
soup = BeautifulSoup(response.json()["data"][2], "html.parser")
time_taken = soup.find("p", class_="time").text
return file_path, time_taken