mirror of
https://github.com/oobabooga/text-generation-webui.git
synced 2024-11-23 08:28:21 +01:00
48 lines
2.1 KiB
Python
48 lines
2.1 KiB
Python
|
import os
|
||
|
|
import time
|
||
|
|
import requests
|
||
|
|
from extensions.openai.errors import *
|
||
|
|
|
||
|
|
def generations(prompt: str, size: str, response_format: str, n: int):
|
||
|
|
# Stable Diffusion callout wrapper for txt2img
|
||
|
|
# Low effort implementation for compatibility. With only "prompt" being passed and assuming DALL-E
|
||
|
|
# the results will be limited and likely poor. SD has hundreds of models and dozens of settings.
|
||
|
|
# If you want high quality tailored results you should just use the Stable Diffusion API directly.
|
||
|
|
# it's too general an API to try and shape the result with specific tags like "masterpiece", etc,
|
||
|
|
# Will probably work best with the stock SD models.
|
||
|
|
# SD configuration is beyond the scope of this API.
|
||
|
|
# At this point I will not add the edits and variations endpoints (ie. img2img) because they
|
||
|
|
# require changing the form data handling to accept multipart form data, also to properly support
|
||
|
|
# url return types will require file management and a web serving files... Perhaps later!
|
||
|
|
|
||
|
|
width, height = [ int(x) for x in size.split('x') ] # ignore the restrictions on size
|
||
|
|
|
||
|
|
# to hack on better generation, edit default payload.
|
||
|
|
payload = {
|
||
|
|
'prompt': prompt, # ignore prompt limit of 1000 characters
|
||
|
|
'width': width,
|
||
|
|
'height': height,
|
||
|
|
'batch_size': n,
|
||
|
|
'restore_faces': True, # slightly less horrible
|
||
|
|
}
|
||
|
|
|
||
|
|
resp = {
|
||
|
|
'created': int(time.time()),
|
||
|
|
'data': []
|
||
|
|
}
|
||
|
|
|
||
|
|
# TODO: support SD_WEBUI_AUTH username:password pair.
|
||
|
|
sd_url = f"{os.environ['SD_WEBUI_URL']}/sdapi/v1/txt2img"
|
||
|
|
|
||
|
|
response = requests.post(url=sd_url, json=payload)
|
||
|
|
r = response.json()
|
||
|
|
if response.status_code != 200 or 'images' not in r:
|
||
|
|
raise ServiceUnavailableError(r.get('detail', [{'msg': 'Unknown error calling Stable Diffusion'}])[0]['msg'], code = response.status_code)
|
||
|
|
# r['parameters']...
|
||
|
|
for b64_json in r['images']:
|
||
|
|
if response_format == 'b64_json':
|
||
|
|
resp['data'].extend([{'b64_json': b64_json}])
|
||
|
|
else:
|
||
|
|
resp['data'].extend([{'url': f'data:image/png;base64,{b64_json}'}]) # yeah it's lazy. requests.get() will not work with this
|
||
|
|
|
||
|
|
return resp
|