利用Stable Diffusion Web UI API实现高级图像生成

下面代码是一个 Python 脚本,用于与 Stable Diffusion 模型的 Web UI 服务器进行交互,实现文本到图像(txt2img)和图像到图像(img2img)的转换。这个脚本展示了如何通过编程方式使用 API 来生成和修改图像,这在自动化图像生成和处理的应用中非常有价值。

from datetime import datetime
import urllib.request
import base64
import json
import time
import os
# 文档https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/API

webui_server_url = ''  # 输入api

out_dir = 'api_out'
out_dir_t2i = os.path.join(out_dir, 'txt2img')
out_dir_i2i = os.path.join(out_dir, 'img2img')
os.makedirs(out_dir_t2i, exist_ok=True)
os.makedirs(out_dir_i2i, exist_ok=True)


def timestamp():
    return datetime.fromtimestamp(time.time()).strftime("%Y%m%d-%H%M%S")


def encode_file_to_base64(path):
    with open(path, 'rb') as file:
        return base64.b64encode(file.read()).decode('utf-8')


def decode_and_save_base64(base64_str, save_path):
    with open(save_path, "wb") as file:
        file.write(base64.b64decode(base64_str))


def call_api(api_endpoint, **payload):
    data = json.dumps(payload).encode('utf-8')
    request = urllib.request.Request(
        f'{webui_server_url}/{api_endpoint}',
        headers={'Content-Type': 'application/json'},
        data=data,
    )
    response = urllib.request.urlopen(request)
    return json.loads(response.read().decode('utf-8'))


def call_txt2img_api(**payload):
    response = call_api('sdapi/v1/txt2img', **payload)
    for index, image in enumerate(response.get('images')):
        save_path = os.path.join(out_dir_t2i, f'txt2img-{timestamp()}-{index}.png')
        decode_and_save_base64(image, save_path)


def call_img2img_api(**payload):
    response = call_api('sdapi/v1/img2img', **payload)
    for index, image in enumerate(response.get('images')):
        save_path = os.path.join(out_dir_i2i, f'img2img-{timestamp()}-{index}.png')
        decode_and_save_base64(image, save_path)


if __name__ == '__main__':
    payload = {
        "prompt": "masterpiece, (best quality:1.1), 1girl <lora:lora_model:1>",  # extra networks also in prompts
        "negative_prompt": "",
        "seed": 1,
        "steps": 20,
        "width": 512,
        "height": 512,
        "cfg_scale": 7,
        "sampler_name": "DPM++ 2M Karras",
        "n_iter": 1,
        "batch_size": 1,

        # x/y/z图的参数示例
        # "script_name": "x/y/z plot",
        # "script_args": [
        #     1,
        #     "10,20",
        #     [],
        #     0,
        #     "",
        #     [],
        #     0,
        #     "",
        #     [],
        #     True,
        #     True,
        #     False,
        #     False,
        #     0,
        #     False
        # ],

        # Refiner和ControlNet的示例参数
        # "alwayson_scripts": {
        #     "ControlNet": {
        #         "args": [
        #             {
        #                 "batch_images": "",
        #                 "control_mode": "Balanced",
        #                 "enabled": True,
        #                 "guidance_end": 1,
        #                 "guidance_start": 0,
        #                 "image": {
        #                     "image": encode_file_to_base64(r"B:\path\to\control\img.png"),
        #                     "mask": None  # base64, None when not need
        #                 },
        #                 "input_mode": "simple",
        #                 "is_ui": True,
        #                 "loopback": False,
        #                 "low_vram": False,
        #                 "model": "control_v11p_sd15_canny [d14c016b]",
        #                 "module": "canny",
        #                 "output_dir": "",
        #                 "pixel_perfect": False,
        #                 "processor_res": 512,
        #                 "resize_mode": "Crop and Resize",
        #                 "threshold_a": 100,
        #                 "threshold_b": 200,
        #                 "weight": 1
        #             }
        #         ]
        #     },
        #     "Refiner": {
        #         "args": [
        #             True,
        #             "sd_xl_refiner_1.0",
        #             0.5
        #         ]
        #     }
        # },
        # "enable_hr": True,
        # "hr_upscaler": "R-ESRGAN 4x+ Anime6B",
        # "hr_scale": 2,
        # "denoising_strength": 0.5,
        # "styles": ['style 1', 'style 2'],
        # "override_settings": {
        #     'sd_model_checkpoint': "sd_xl_base_1.0",  # this can use to switch sd model
        # },
    }
    call_txt2img_api(**payload)

    init_images = [
        encode_file_to_base64("output.png"),
        # encode_file_to_base64(r"B:\path\to\img_2.png"),
        # "https://image.can/also/be/a/http/url.png",
    ]

    batch_size = 2
    payload = {
        "prompt": "1girl, blue hair",
        "seed": 1,
        "steps": 20,
        "width": 512,
        "height": 512,
        "denoising_strength": 0.5,
        "n_iter": 1,
        "init_images": init_images,
        "batch_size": batch_size if len(init_images) == 1 else len(init_images),
        # "mask": encode_file_to_base64(r"B:\path\to\mask.png")
    }
    # if len(init_images) > 1 then batch_size should be == len(init_images)
    # else if len(init_images) == 1 then batch_size can be any value int >= 1
    call_img2img_api(**payload)

    # there exist a useful extension that allows converting of webui calls to api payload
    # particularly useful when you wish setup arguments of extensions and scripts
    # https://github.com/huchenlei/sd-webui-api-payload-display

代码解析

  1. 导入必要的库: 使用datetime、urllib.request、base64、json、time和os等库来处理日期时间、网络请求、文件编码、JSON处理、时间和文件系统操作。
  2. 初始化输出目录: 创建两个输出目录out_dir_t2i和out_dir_i2i,分别用于存储文本到图像和图像到图像的结果。
  3. 定义辅助函数: timestamp(): 生成当前时间戳。 encode_file_to_base64(path): 将文件编码为Base64字符串。 decode_and_save_base64(base64_str, save_path): 将Base64字符串解码并保存为文件。
  4. 定义API调用函数: call_api(api_endpoint, **payload): 发送请求到指定的API端点。 call_txt2img_api(**payload): 调用文本到图像的API。 call_img2img_api(**payload): 调用图像到图像的API。
  5. 实际的API调用示例: 在脚本的主体部分,提供了用于调用文本到图像和图像到图像API的示例payload。

代码的实际应用

这个脚本在AI图像生成和处理的领域具有广泛的应用,例如:

  • 创意艺术和设计:艺术家和设计师可以使用这个脚本来自动化地生成创意图像。
  • 数据可视化:数据科学家可以利用图像到图像的功能来增强数据可视化效果。
  • 媒体和娱乐:媒体公司可以用它来快速生成或修改图像,用于社交媒体、广告等。

应用场景

  • 自动化内容生成:自动创建图像内容,用于社交媒体、网站等。
  • 图像编辑和增强:修改已有图像,加入新元素或改变风格。