From d4ec0a9408cb60cdbfff44f2554d6fd325da5438 Mon Sep 17 00:00:00 2001
From: Luke Robles <lukelrobles@gmail.com>
Date: Thu, 14 Sep 2023 09:05:25 -0700
Subject: [PATCH] fix SD to use the latest api, but itll only work on my gaming
 computer for now

---
 app/cogs/stable_diffusion.py | 11 +++---
 app/stable_diffusion.py      | 70 +++++++-----------------------------
 2 files changed, 19 insertions(+), 62 deletions(-)

diff --git a/app/cogs/stable_diffusion.py b/app/cogs/stable_diffusion.py
index 50d794e6..76060476 100755
--- a/app/cogs/stable_diffusion.py
+++ b/app/cogs/stable_diffusion.py
@@ -35,25 +35,27 @@ class StableDiffusion(commands.Cog):
 
             port = "7860"
             ip = "192.168.1.80"
-            steps = 20
+            steps = 50
 
             # 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.19"
-                steps = 90
+                positive_prompt = (
+                    positive_prompt
+                    + " <lora:add_detail:0.3> <lora:bimbostyleThreeU:0.3> <lora:gigatitums:0.3> <lora:more_details:0.3>"
+                )
                 s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                 s.settimeout(1)
                 try:
                     s.connect((ip, int(port)))
                 except:
                     ip = "192.168.1.80"
-                    steps = 60
             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(
+                file_path = await stable_diffusion.generate_image(
                     ip=ip,
                     port=port,
                     positives=positive_prompt,
@@ -62,7 +64,6 @@ class StableDiffusion(commands.Cog):
                     steps=steps,
                 )
                 await original_message.edit(
-                    content=time_taken,
                     file=discord.File(
                         file_path,
                         filename="unknown.png",
diff --git a/app/stable_diffusion.py b/app/stable_diffusion.py
index d614efd4..04de643c 100755
--- a/app/stable_diffusion.py
+++ b/app/stable_diffusion.py
@@ -1,71 +1,27 @@
 from bs4 import BeautifulSoup
+import base64
 import httpx
 import tempfile
 
 
 async def generate_image(ip, port, positives, negatives, steps):
-    url = "http://" + ip + ":" + port + "/run/predict/"
-    headers = {
-        "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:108.0) Gecko/20100101 Firefox/108.0",
-        "Accept": "*/*",
-        "Accept-Language": "en-US,en;q=0.5",
-        "DNT": "1",
-        "Connection": "keep-alive",
-    }
+    url = "http://" + ip + ":" + port + "/sdapi/v1/txt2img"
 
     json_data = {
-        "fn_index": 52,
-        "data": [
-            positives,
-            negatives,
-            "None",
-            "None",
-            steps,
-            "Euler a",
-            True,
-            False,
-            1,
-            1,
-            7,
-            -1,
-            -1,
-            0,
-            0,
-            0,
-            False,
-            512,
-            512,
-            True,
-            0.7,
-            0,
-            0,
-            "None",
-            False,
-            False,
-            False,
-            False,
-            "",
-            "Seed",
-            "",
-            "Nothing",
-            "",
-            True,
-            False,
-            False,
-            None,
-        ],
+        "prompt": positives,
+        "negative_prompt": negatives,
+        "steps": steps,
+        "hr_upscale": 2,
+        "alwayson_scripts": {"ADetailer": {"args": [{"ad_model": "face_yolov8n.pt"}]}},
     }
-    client = httpx.AsyncClient()
-    response = await client.post(url, headers=headers, json=json_data, timeout=1200)
-    file_path = response.json()["data"][0][0]["name"]
 
-    imgdata = await client.get("http://%s:%s/file=%s" % (ip, port, file_path))
+    client = httpx.AsyncClient()
+
+    response = await client.post(url, json=json_data, timeout=1200)
+    b64_image = response.json()["images"][0]
     file, file_path = tempfile.mkstemp()
 
     with open(file_path, "wb") as f:
-        f.write(imgdata.content)
+        f.write(base64.b64decode(b64_image))
 
-    soup = BeautifulSoup(response.json()["data"][2], "html.parser")
-    time_taken = soup.find("p", class_="time").text
-
-    return file_path, time_taken
+    return file_path