Spaces:
Sleeping
Sleeping
| import os | |
| import time | |
| import requests | |
| from extensions.openai.errors import ServiceUnavailableError | |
| 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 negative prompts | |
| # or "masterpiece", etc. 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! | |
| base_model_size = 512 if 'SD_BASE_MODEL_SIZE' not in os.environ else int(os.environ.get('SD_BASE_MODEL_SIZE', 512)) | |
| sd_defaults = { | |
| 'sampler_name': 'DPM++ 2M Karras', # vast improvement | |
| 'steps': 30, | |
| } | |
| 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, | |
| } | |
| payload.update(sd_defaults) | |
| scale = min(width, height) / base_model_size | |
| if scale >= 1.2: | |
| # for better performance with the default size (1024), and larger res. | |
| scaler = { | |
| 'width': width // scale, | |
| 'height': height // scale, | |
| 'hr_scale': scale, | |
| 'enable_hr': True, | |
| 'hr_upscaler': 'Latent', | |
| 'denoising_strength': 0.68, | |
| } | |
| payload.update(scaler) | |
| resp = { | |
| 'created': int(time.time()), | |
| 'data': [] | |
| } | |
| from extensions.openai.script import params | |
| # TODO: support SD_WEBUI_AUTH username:password pair. | |
| sd_url = f"{os.environ.get('SD_WEBUI_URL', params.get('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: | |
| print(r) | |
| raise ServiceUnavailableError(r.get('error', 'Unknown error calling Stable Diffusion'), code=response.status_code, internal_message=r.get('errors', None)) | |
| # 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 | |