SD web api

sd.py

import base64
import sys
from io import BytesIO

import requests
from PIL import Image


class Client:

  def __init__(self, base_url):
    self.base_url = base_url

  def call(self, op):
    r = requests.post(f'{self.base_url}/{op.path}', json=op.args)
    data = r.json()
    return op.transform(data)

class Txt2Img:
  path = 'sdapi/v1/txt2img'

  def __init__(self, **kwargs):
    kwargs = {key: value for key, value in kwargs.items() if value is not None}
    self.args = kwargs

  def transform(self, data):
    imgs = [
        base64.decodebytes(img_b64.encode()) for i, img_b64 in enumerate(data['images'])
    ]
    assert len(imgs) == 1
    return Image.open(BytesIO(imgs[0]))

class Img2Img:
  path = 'sdapi/v1/img2img'

  def __init__(self, **kwargs):
    defaults = dict(
        mask_blur=4,
        resize_mode=1,  # Just resize, Crop and resize, Resize and fill
        inpaint_full_res=False,
        inpaint_full_res_padding=32,
        inpaining_fill=1,  # Masked content: fill, original, latent noise, latent nothing
        sampler_index='Euler a',
        steps=20,
    )
    # Euler a, Euler, LMS, Heun,
    # DPM2, DPM2 a, DPM++ 2S a, DPM++ 2M, DPM fast, DPM adaptive,
    # LMS Karras, DPM2 Karras, DPM2 a Karras, DPM++ 2S a Karras,
    # DPM++ 2M Karras, DDIM
    kwargs = defaults | {key: value for key, value in kwargs.items() if value is not None}
    self.args = kwargs

  def transform(self, data):
    imgs = [
        base64.decodebytes(img_b64.encode()) for i, img_b64 in enumerate(data['images'])
    ]
    assert len(imgs) == 1
    return Image.open(BytesIO(imgs[0]))

def from_png_contents(contents):
  return 'data:image/png;base64, ' + base64.encodebytes(contents).decode()

def read_png(fname):
  assert fname.endswith('.png')
  return from_png_contents(open(fname, 'rb').read())

def gimp_to_encoded_png(width: int, height: int, src_bytes: bytes, mode: str):
  img = Image.frombytes(mode, (width, height), src_bytes)
  png = BytesIO()
  img.save(png, 'png')
  return from_png_contents(png.getvalue())

# Usage:
#
#   client = Client('http://127.0.0.1:7860')
#   img = client.call(
#       Img2Img(
#           seed=123,
#           init_images=[<encoded images>],
#           mask=<optional encoded mask>,
#           prompt='halloween cat in pumpkin',
#           negative_prompt='',
#       ))