weeabot/modules/diffuseapi.py_ref
2022-12-10 22:38:37 +11:00

42 lines
1.2 KiB
Plaintext

import aiohttp
import asyncio
class DiffuseAPI:
def __init__(self, url, nsfw_enabled, styles, steps):
self.url = url
self.nsfw_enabled = nsfw_enabled
self.styles = styles
self.steps = steps
self.seed = -1
def _generate_payload(self):
payload = {
"prompt": "",
"negative_prompt": "",
"steps": self.steps,
"seed": self.seed,
"styles": self.styles,
"height": 1024,
"width": 512
}
settings = {
"filter_nsfw": not self.nsfw_enabled,
"samples_save": True,
}
override_payload = {
"override_settings": settings
}
payload.update(override_payload)
return payload
async def generate_image(self, prompt, negative_prompt=""):
async with aiohttp.ClientSession(self.url) as session:
payload = self._generate_payload()
payload.update({"prompt": prompt, "negative_prompt": negative_prompt})
print(payload)
async with session.post('/sdapi/v1/txt2img', json=payload) as image_handler:
image_data = await image_handler.json()
print(image_data)