From e3339da557a006e513275af069c21b9433583602 Mon Sep 17 00:00:00 2001
From: Luke Robles <98352913+lrobles-iterable@users.noreply.github.com>
Date: Fri, 14 Oct 2022 09:40:44 -0700
Subject: [PATCH] Adding time taken to stable diffusion renders

---
 app/cogs/actual_utils.py |  4 ++--
 app/stable_diffusion.py  | 11 ++++++++---
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/app/cogs/actual_utils.py b/app/cogs/actual_utils.py
index 63b80171..2aade524 100644
--- a/app/cogs/actual_utils.py
+++ b/app/cogs/actual_utils.py
@@ -121,7 +121,7 @@ class ActualUtils(commands.Cog):
                 original_message = await ctx.followup.send(
                     "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,
                     port=port,
                     positives=positive_prompt,
@@ -129,7 +129,7 @@ class ActualUtils(commands.Cog):
                     steps=steps,
                 )
                 await original_message.edit(
-                    content="",
+                    content=time_taken,
                     file=discord.File(
                         file_path,
                         filename="unknown.png",
diff --git a/app/stable_diffusion.py b/app/stable_diffusion.py
index 82faabe5..b7984fe9 100644
--- a/app/stable_diffusion.py
+++ b/app/stable_diffusion.py
@@ -1,6 +1,8 @@
+from bs4 import BeautifulSoup
 import base64
-import json
 import httpx
+import json
+import requests
 import socket
 import tempfile
 
@@ -59,10 +61,13 @@ async def generate_image(ip, port, positives, negatives, steps):
     client = httpx.AsyncClient()
     response = await client.post(url, headers=headers, json=json_data, timeout=1200)
     imageb64 = response.json()["data"][0][0].split("data:image/png;base64,")[1]
-
     imgdata = base64.b64decode(imageb64)
     file, file_path = tempfile.mkstemp()
 
     with open(file_path, "wb") as f:
         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