API Extras example for stable-diffusion-webui

extras.py

#! /usr/bin/python3
"""
extras.py: Upscale PNG images in DIR_IN into DIR_OUT.
Usage: Set API_URL, DIR_IN, and DIR_OUT. Then run ./extras.py
For API documentation see: http://localhost:7860/docs#/
"""

import sys
import io
import os
import base64
import json
from io import BytesIO
import requests
from PIL import Image, PngImagePlugin

API_URL = "http://localhost:7860/sdapi/v1/extra-batch-images"
DIR_IN = "extras-in/"
DIR_OUT = "extras-out/"


def pil_to_base64(pil_image):
    """Encode PIL Image to base64 string."""
    with BytesIO() as stream:
        meta = PngImagePlugin.PngInfo()
        for k, v in pil_image.info.items():
            if isinstance(k, str) and isinstance(v, str):
                meta.add_text(k, v)
        pil_image.save(stream, "PNG", pnginfo=meta)
        base64_str = str(base64.b64encode(stream.getvalue()), "utf-8")
        return "data:image/png;base64," + base64_str


# Get files without directories or subdirectories
image_names = [f for f in os.listdir(DIR_IN) if os.path.isfile(os.path.join(DIR_IN, f))]
print(f"extras.py: Upscaling {len(image_names)} images...")

image_list = []
for i in image_names:
    image_list.append({"data": pil_to_base64(Image.open(DIR_IN + i)), "name": i})

payload = {
    "resize_mode": 0,
    "gfpgan_visibility": 0,
    "codeformer_visibility": 0,
    "codeformer_weight": 0,
    "upscaling_resize": 2,
    "upscaler_1": "R-ESRGAN 4x+",
    "upscaler_2": "None",
    "extras_upscaler_2_visibility": 0,
    "imageList": image_list,
    "upscale_first": False,  # upscale before restoring faces (not
    # applied if visibility = 0?)
}

payloadJson = json.dumps(payload)
resp = requests.post(url=API_URL, data=payloadJson).json()

if resp.get("images") is None:
    print("extras.py: Error, Post response:")
    print(resp)
    sys.exit(1)

else:
    index = 0
    for i in resp["images"]:
        index = index + 1
        metadata = PngImagePlugin.PngInfo()
        img = Image.open(io.BytesIO(base64.b64decode(i)))
        for key, value in img.info.items():
            if isinstance(key, str) and isinstance(value, str):
                metadata.add_text(key, value)
        img.save(
            DIR_OUT + image_names[index - 1] + "-" + str(index) + ".png",
            "PNG",
            pnginfo=metadata,
        )
    print("extras.py: Done")

curl -X 'POST' \
  'http://54.189.78.1:7860/sdapi/v1/extra-single-image' \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -d '{
  "resize_mode": 0,
  "show_extras_results": true,
  "gfpgan_visibility": 0,
  "codeformer_visibility": 0,
  "codeformer_weight": 0,
  "upscaling_resize": 2,
  "upscaling_crop": true,
  "upscaler_1": "None",
  "upscaler_2": "None",
  "extras_upscaler_2_visibility": 0,
  "upscale_first": false,
  "image": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAA3UAAAG..."
}'