stable-diffusion-webui-reForge(または1111)のAPIを使った画像生成スクリプト

import requests
import json
import os
import io
import base64
from pydantic import BaseModel, Field
from PIL import Image, PngImagePlugin
from datetime import datetime, timezone, timedelta

# 日本標準時
JST = timezone(timedelta(hours=9))

# APIエンドポイント
SD_SERVER_URL = "http://127.0.0.1:7860"  # LAN内のサーバーURLに変更してください

# 入力データを管理するクラス
class ImageGenerationParams(BaseModel):
    prompt: str = Field(..., example="1girl", description="生成したい画像のプロンプト")
    negative_prompt: str = Field(
        "worst quality", example="worst quality", description="否定プロンプト"
    )
    steps: int = Field(25, example=25, description="生成ステップ数")
    cfg_scale: float = Field(7.0, example=7.0, description="CFGスケールの値")
    width: int = Field(512, example=512, description="生成画像の幅")
    height: int = Field(512, example=512, description="生成画像の高さ")
    sampler_index: str = Field(
        "Euler a", example="Euler a", description="サンプラーの種類"
    )
    batch_count: int = Field(1, example=1, description="バッチ生成の数")
    batch_size: int = Field(1, example=1, description="1バッチあたりの画像数")
    enable_hr: bool = Field(
        False, description="高解像度リサイズモードを有効にするかどうか"
    )
    denoising_strength: float = Field(
        0.4, example=0.4, description="高解像度リサイズ時のノイズ低減強度"
    )
    hr_scale: float = Field(1.5, example=1.5, description="高解像度リサイズのスケール")
    hr_steps: int = Field(
        15, example=15, description="高解像度リサイズ時の追加ステップ数"
    )
    hr_upscaler: str = Field(
        "R-ESRGAN 4x+ Anime6B", example="R-ESRGAN 4x+ Anime6B", description="高解像度リサイズのアップスケーラー"
    )

def text_to_image(params: ImageGenerationParams, output_dir: str, unique_id: str):
    """
    Stable Diffusion API (txt2img) を使用して画像を生成し、PNGとして保存する。

    Args:
        params: ImageGenerationParamsオブジェクト（画像生成パラメータを含む）
        output_dir: 出力先ディレクトリ
        unique_id: ファイル名に使用する一意の識別子
    """
    # APIへのリクエストペイロード
    payload = {
        "prompt": params.prompt,
        "negative_prompt": params.negative_prompt,
        "steps": params.steps,
        "cfg_scale": params.cfg_scale,
        "width": params.width,
        "height": params.height,
        "sampler_index": params.sampler_index,
        "batch_count": params.batch_count,
        "batch_size": params.batch_size,
        "enable_hr": params.enable_hr,
        "denoising_strength": params.denoising_strength,
        "hr_scale": params.hr_scale,
        "hr_upscaler": params.hr_upscaler,
        "hr_second_pass_steps": params.hr_steps,
    }

    headers = {"Content-type": "application/json"}
    url = f"{SD_SERVER_URL}/sdapi/v1/txt2img"

    # APIリクエスト送信
    response = requests.post(url, json=payload, headers=headers, timeout=300)
    response.raise_for_status()

    # レスポンス処理
    result = response.json()
    meta = json.loads(result["info"])["infotexts"][0]

    # 出力ディレクトリの作成
    os.makedirs(output_dir, exist_ok=True)

    # 画像の保存
    for i, image_data in enumerate(result["images"]):
        image = Image.open(io.BytesIO(base64.b64decode(image_data.split(",", 1)[0])))
        pnginfo = PngImagePlugin.PngInfo()
        pnginfo.add_text("parameters", meta)
        now = datetime.now(JST).strftime("%Y%m%d%H%M%S")
        file_path = os.path.join(output_dir, f"{now}_{unique_id}_{i+1}.png")
        print(f"Saving image to {file_path}")
        image.save(file_path, pnginfo=pnginfo)

# 生成パラメータの設定
params = ImageGenerationParams(
    prompt="a futuristic cityscape",
    negative_prompt="low quality, blurry",
    steps=50,
    cfg_scale=7.5,
    width=512,
    height=512,
    sampler_index="Euler a",
    batch_count=1,
    batch_size=1,
    enable_hr=True,
    denoising_strength=0.4,
    hr_scale=2.0,
    hr_steps=20,
    hr_upscaler="R-ESRGAN 4x+ Anime6B",
)

# 出力ディレクトリと一意の識別子
output_dir = "./generated_images"
unique_id = "cityscape"

# 画像生成
text_to_image(params, output_dir, unique_id)