From 7c6d64bb0f94f204cf8b21882bf9f2217566e7fe Mon Sep 17 00:00:00 2001
From: jordan <jorraan@jorraan.au>
Date: Sat, 10 Dec 2022 22:32:57 +1100
Subject: [PATCH] Refactor

---
 DiffuseAPI.py | 41 ++++++++++++++++++++++-------------------
 1 file changed, 22 insertions(+), 19 deletions(-)

diff --git a/DiffuseAPI.py b/DiffuseAPI.py
index 51c2aac..506eba3 100644
--- a/DiffuseAPI.py
+++ b/DiffuseAPI.py
@@ -41,7 +41,7 @@ class DiffuseAPI():
             return False
         except:
             return False
-    
+
     def set_styles(self, styles):
         if type(styles) == list:
             self.styles = styles
@@ -65,20 +65,20 @@ class DiffuseAPI():
                 return True
             case _:
                 return False
-    
+
     def get_orientation(self):
         if self.width == 512 and self.height == 512:
             return ("square", self.width, self.height)
         elif self.width == 1024:
             return ("landscape", self.width, self.height)
         return ("portrait", self.width, self.height)
-    
+
     def set_nsfw_filter(self, filter_state):
         if type(filter_state) == bool:
             self.nsfw_enabled = filter_state
             return True
         return False
-    
+
     def get_nsfw_filter(self):
         return self.nsfw_enabled
 
@@ -87,7 +87,7 @@ class DiffuseAPI():
         "prompt": prompt,
         "styles": self.styles,
         "steps": self.num_steps,
-        "seed": self.ai_seed,
+        "seed": self.seed,
         "n_iter": 1,
         "height": self.height,
         "width": self.width,
@@ -112,17 +112,20 @@ class DiffuseAPI():
             alive.close()
             await sess.close()
             return None
-        request  = await sess.post("/sdapi/v1/txt2img", json=payload)
-        try:
-            req_json = await request.json()
-            request.close()
-            await sess.close()
-            return req_json["images", [0]]
-        except:
-            return None
-
-
-    
-
-    
-
+        async with aiohttp.ClientSession(self.url) as session:
+            async with session.head('/') as alive:
+                if alive.status != 200:
+                return None
+            async with session.post("/sdapi/v1/txt2img", json=payload) as image_json:
+                image_data = await image_json.json()
+                return image_data["images"][0]
+            
+            
+            
+            # try:
+            #     req_json = await request.json()
+            #     request.close()
+            #     await sess.close()
+            #     return req_json["images"][0]
+            # except:
+            #     return None