ComfyUI APIを試してみた



# はじめに
これは[Infocom Advent Calendar 2024](https://qiita.com/advent-calendar/2024/infocom) 25日目の記事です。


# 目次
1. ComfyUIとは？
    * ComfyUIをちょっとやってみる
1. ComfyUIのAPIとは？
1. 試してみた
1. さいごに
1. 参考文献


# ComfyUIとは？
[ComfyUI](https://github.com/comfyanonymous/ComfyUI)とは、複雑な画像生成のワークフローを試したり作成したりするためのノード・グラフインターフェースのことです。
よく似たオープンソースとして[Stable Diffusion WebUI](https://github.com/AUTOMATIC1111/stable-diffusion-webui)（SDwebUI）が挙げられますが、ComfyUIの方がより複雑な処理をGUIで操作できると思っています。

## ComfyUIをちょっとやってみる
gitHubから持ってきて、`python main.py`すれば使えます。
![image1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2995327/43cf2276-ac73-a40e-e504-44f37c814792.png)

* 画像生成してみる  
試しに[ComfyUI-Manager](https://github.com/ltdrdata/ComfyUI-Manager)の[Waifu Diffusion 1.5 Beta3 (fp16)](https://huggingface.co/waifu-diffusion/wd-1-5-beta3)を使ってみます。
ComfyUI-Managerには色々なカスタムノードがあります。
プロンプトに「apple」と入れて生成してみます。
![image2.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2995327/1efef066-a1f9-08eb-86a3-134d7b1ba0f3.png)


赤っぽいりんごが出力されました 。 
ComfyUIは左から右にかけて実行されます。
このノードの流れのことを「ワークフロー」と呼び、jsonファイルで出力されます。
また、ワークフローは参照できるので他の人のワークフローも簡単に使えます。
左上の「ワークフロー」から「名前をつけて保存」でjson出力できます。
![image3.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2995327/6386cd4c-9863-7da3-4ff1-626b71dc7cb2.png)


ちなみにワークフローの中身はこんな感じ
```{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の例](https://github.com/comfyanonymous/ComfyUI/blob/master/script_examples/websockets_api_example.py)を参考に書きました。  
注意：ComfyUIをAPIとして使うには、設定の「開発モードオプションを有効にする（API保存など）」をオンにして下さい

```python
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用ライブラリ](https://github.com/pydn/ComfyUI-to-Python-Extension)を作ってくれた人がいるみたいです。これでもっと自動化が進みますね、どこかのタイミングで試したいです。

# 参考文献
* https://ja.aicu.ai/2024/09/29/comfymaster21-20240929/
* https://zenn.dev/sunwood_ai_labs/articles/12085336e9d1ab
* https://github.com/comfyanonymous/ComfyUI/blob/master/script_examples/websockets_api_example.py