はじめに

これはInfocom Advent Calendar 2024 25日目の記事です。

目次

ComfyUIとは？
- ComfyUIをちょっとやってみる
ComfyUIのAPIとは？
試してみた
さいごに
参考文献

ComfyUIとは？

ComfyUIとは、複雑な画像生成のワークフローを試したり作成したりするためのノード・グラフインターフェースのことです。
よく似たオープンソースとしてStable Diffusion WebUI（SDwebUI）が挙げられますが、ComfyUIの方がより複雑な処理をGUIで操作できると思っています。

ComfyUIをちょっとやってみる

gitHubから持ってきて、python main.pyすれば使えます。

画像生成してみる
試しにComfyUI-ManagerのWaifu Diffusion 1.5 Beta3 (fp16)を使ってみます。
ComfyUI-Managerには色々なカスタムノードがあります。
プロンプトに「apple」と入れて生成してみます。

赤っぽいりんごが出力されました。
ComfyUIは左から右にかけて実行されます。
このノードの流れのことを「ワークフロー」と呼び、jsonファイルで出力されます。
また、ワークフローは参照できるので他の人のワークフローも簡単に使えます。
左上の「ワークフロー」から「名前をつけて保存」でjson出力できます。

ちなみにワークフローの中身はこんな感じ

{
  "3": {
    "inputs": {
      "seed": 1074162188378014,
      "steps": 20,
      "cfg": 8,
      "sampler_name": "euler",
      "scheduler": "normal",
      "denoise": 1,
      "model": [
        "4",
        0
      ],
      "positive": [
        "6",
        0
      ],
      "negative": [
        "7",
        0
      ],
      "latent_image": [
        "5",
        0
      ]
    },
    "class_type": "KSampler",
    "_meta": {
      "title": "Kサンプラー"
    }
  }
  ......(続く)

ここの「3」という数字が冒頭のワークフローの画像で言う、Kサンプラーの「#3」になっていて、inputsにパラメータが入っています。
他にもComfyUIのノードにはLoRAやControlNetがあるので「自分が学習させたモデルを入れて、こういう構造で出したい」みたいなカスタマイズもできます。

ComfyUIのAPIとは？

ComfyUIをAPIサーバーとして、先ほど話したワークフローを自動実行できるようになります。
例えば「作ったワークフローのseed値やプロンプトを少しずつ変えて大量生成したい」が実現できます。
個人的に実験（ノーコード）→実用（API使うだけ）までが簡単に出来るのは魅力的だと感じます。（ただAPIを使う為のドキュメントは特になさそう）

試してみた

早速試します。今回はPythonでComfyUI APIを操作し、画像を取得するところまでやります。ComfyUIの例を参考に書きました。
注意：ComfyUIをAPIとして使うには、設定の「開発モードオプションを有効にする（API保存など）」をオンにして下さい

import sys
import json
import websocket
import uuid
import random
from PIL import Image
import io
from urllib.request import Request, urlopen
from setting import SERVER_ADDRESS, WORKFLOW_PATH

def run_workflow(workflow, client_id):
    """
        ワークフローをComfyUIに投げて実行する関数
    """
    data = json.dumps({"prompt": workflow, "client_id": client_id}).encode('utf-8')
    with urlopen(Request(f"http://{SERVER_ADDRESS}/prompt", data=data)) as response:
        return json.loads(response.read())

def get_history(prompt_id):
    with urlopen(f"http://{SERVER_ADDRESS}/history/{prompt_id}") as response:
        return json.load(response)

def create_image(ws, workflow, client_id):
    """
        ワークフローを投げて、生成した画像を生成
    """
    prompt_id = run_workflow(workflow, client_id)['prompt_id']

    # 画像が生成されたかをポーリングして確認
    # https://github.com/comfyanonymous/ComfyUI/blob/26e0ba8f8cf786575fc1324acb858ad81f3ef9d6/script_examples/websockets_api_example.py#L29
    while True:
        out = ws.recv()
        if isinstance(out, str):
            message = json.loads(out)
            if message['type'] == 'executing':
                data = message['data']
                if data['node'] is None and data['prompt_id'] == prompt_id:
                    break #Execution is done
        else:
            # If you want to be able to decode the binary stream for latent previews, here is how you can do it:
            # bytesIO = BytesIO(out[8:])
            # preview_image = Image.open(bytesIO) # This is your preview in PIL image format, store it in a global
            continue #previews are binary data
    
    return prompt_id

def get_image(file_data: dict):
    with urlopen(f'http://{SERVER_ADDRESS}/view?filename={file_data["filename"]}&type={file_data["type"]}&subfolder={file_data["subfolder"]}') as response:
        return response.read()

if __name__ == "__main__":

    # プロンプトをいじる
    args = sys.argv

    if len(args) > 1:
        prompt = args[1]
    else:
        # 特に指示がなければ緑のリンゴを出力
        prompt = "green apple"

    # JSONファイルを読み込み
    with open(WORKFLOW_PATH, 'r') as file:
        workflow = json.load(file)
    
    # ワークフロー内の#6の入力プロンプトを入れる
    workflow["6"]["inputs"]["text"] = prompt
    # 変更要素がないと実行されない
    workflow["3"]["inputs"]["seed"] = random.randint(1, 1000000)
    
    # 通信ソケットを使用して、画像生成状況を確認
    client_id = str(uuid.uuid4())
    ws = websocket.WebSocket()
    ws.connect(f"ws://{SERVER_ADDRESS}/ws?clientId={client_id}")
    prompt_id = create_image(ws, workflow, client_id)
    ws.close()
    
    # 出力したデータ情報を取得
    history = get_history(prompt_id)[prompt_id]

    # 画像データを取ってくる
    file_data = history['outputs']['9']['images'][0]
    image_data = get_image(file_data)
    image = Image.open(io.BytesIO(image_data))

    # 画像を保存
    image.save("test.png")

例えば、上記をtest.pyとして保存して、python test.py "red apple"と叩くと同じディレクトリにtest.pngとして赤いリンゴが出力されるはずです。

さいごに

調べて知りましたが、ComfyUIのPython用ライブラリを作ってくれた人がいるみたいです。これでもっと自動化が進みますね、どこかのタイミングで試したいです。

参考文献