text-generation-webui/extensions/openai/images.py

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