Home/weeabot
2
1
Fork 0
Code Issues 1 Pull Requests Projects Releases Wiki Activity

Refactor

Browse Source
This commit is contained in:
jordan 2022-12-10 22:32:57 +11:00
parent 8a3c228f92
commit 7c6d64bb0f
1 changed files with 22 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

41
DiffuseAPI.py
View File

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