122 lines
3.2 KiB
Python
122 lines
3.2 KiB
Python
import aiohttp
|
|
|
|
|
|
class DiffuseAPI():
|
|
def __init__(self, url, styles, nsfw_enabled=True, num_steps=28):
|
|
|
|
self.nsfw_enabled = nsfw_enabled
|
|
self.num_steps = num_steps
|
|
self.url = url
|
|
self.styles = styles
|
|
self.seed = -1
|
|
self.width = 512
|
|
self.height = 1024
|
|
self.cfg_scale = 12
|
|
|
|
def set_steps(self, steps):
|
|
try:
|
|
new_steps = int(steps)
|
|
if 0 > new_steps < 50:
|
|
self.num_steps = new_steps
|
|
return True
|
|
return False
|
|
except:
|
|
return False
|
|
|
|
def set_seed(self, seed):
|
|
try:
|
|
new_seed = int(seed)
|
|
self.seed = new_seed
|
|
return True
|
|
except:
|
|
return False
|
|
|
|
def set_cfg_scale(self, scale):
|
|
try:
|
|
new_scale = int(scale)
|
|
if 0 > new_scale < 30:
|
|
self.cfg_scale = new_scale
|
|
return True
|
|
return False
|
|
except:
|
|
return False
|
|
|
|
def set_styles(self, styles):
|
|
if type(styles) == list:
|
|
self.styles = styles
|
|
return True
|
|
return False
|
|
|
|
def set_orientation(self, orientation):
|
|
new_orientation = str(orientation)
|
|
match new_orientation:
|
|
case "portrait":
|
|
self.width = 512
|
|
self.height = 1024
|
|
return True
|
|
case "landscape":
|
|
self.width = 1024
|
|
self.height = 512
|
|
return True
|
|
case "square":
|
|
self.width = 512
|
|
self.height = 512
|
|
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
|
|
|
|
async def generate_image(self, prompt, neg_prompt=""):
|
|
payload = {
|
|
"prompt": prompt,
|
|
"styles": self.styles,
|
|
"steps": self.num_steps,
|
|
"seed": self.seed,
|
|
"n_iter": 1,
|
|
"height": self.height,
|
|
"width": self.width,
|
|
"negative_prompts": neg_prompt,
|
|
"cfg_scale": self.cfg_scale
|
|
}
|
|
|
|
settings = {
|
|
"filter_nsfw": not self.nsfw_enabled
|
|
}
|
|
|
|
override_payload = {
|
|
"override_settings": settings
|
|
}
|
|
|
|
payload.update(override_payload)
|
|
|
|
sess = aiohttp.ClientSession(self.url)
|
|
alive = await sess.head('/')
|
|
|
|
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
|