Push me

And then just touch me
So I can get my
Satisfaction

modules/diffuseapi.py

@@ -1,5 +1,7 @@
 import aiohttp
-
+from base64 import b64decode
+from sys import exit
+from pprint import pprint
 
 class DiffuseAPI():
     def __init__(self, url, styles, nsfw_enabled=True, num_steps=28):
@@ -40,7 +42,7 @@ class DiffuseAPI():
             return False
         except:
             return False
-    
+
     def set_styles(self, styles):
         if type(styles) == list:
             self.styles = styles
@@ -64,20 +66,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
 
@@ -95,7 +97,8 @@ class DiffuseAPI():
         }
 
         settings = {
-        "filter_nsfw": not self.nsfw_enabled
+        "filter_nsfw": not self.nsfw_enabled,
+        "enable_pnginfo": False
         }
 
         override_payload = {
@@ -104,18 +107,36 @@ class DiffuseAPI():
 
         payload.update(override_payload)
 
-        sess = aiohttp.ClientSession(self.url)
-        alive = await sess.head('/')
+        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]
 
-        if alive.status != 200:
-            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 def generate_upscale(self, image):
+        payload = {
+            "resize_mode": 0,
+            "show_extras_results": True,
+            "gfpgan_visibility": 0,
+            "codeformer_visibility": 0,
+            "codeformer_weight": 0,
+            "upscaling_resize": 4,
+            "upscaling_resize_w": 512,
+            "upscaling_resize_h": 1024,
+            "upscaling_crop": True,
+            "upscaler_1": "R-ESRGAN 4x+ Anime6B",
+            "upscaler_2": "None",
+            "extras_upscaler_2_visibility": 0,
+            "upscale_first": False,
+            "image": image
+        }
+
+        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/extra-single-image", json=payload) as image_json:
+                image_data = await image_json.json()
+                return image_data["image"]