ComfyUI APIでWebSocket進捗を確認しながらACE-Step音楽生成を行う

前回の記事では、ComfyUIのAPIを使って、ACE-Stepによる音楽生成をPythonから実行する方法を紹介しました。

今回は一歩進めて、WebSocketを使って生成の進捗状況をリアルタイムに受け取る方法を解説します。

この記事を読むことで、生成の進捗を把握し、安全に完了を検知できる実装パターンが理解できます。

ComfyUIのWebSocketでできること
WebSocketでの生成完了判定
TQDMと組み合わせた進捗表示の例
WebSocketを使う場面と注意点
まとめ

ComfyUIのWebSocketでできること

ComfyUIはHTTP APIに加えて、WebSocketによる双方向通信を提供しています。

WebSocketを使うことで、画像や音楽生成の進捗、内部ノードの実行状況、エラー通知などをリアルタイムで受信できます。

公式ドキュメントでも、HTTPのPOSTやGETよりも詳細な状態を取得できる手段として紹介されています。
https://docs.comfy.org/development/cloud/api-reference#websocket-for-real-time-progress

サーバーの標準ログ出力とほぼ同じ粒度でイベントが飛んでくるため、「処理が止まっていないか」「まだ計算中なのか」を確認したい場合に特に有効です。

WebSocketでの生成完了判定

WebSocketで流れてくるイベントにはいくつか種類があります。
https://docs.comfy.org/development/cloud/api-reference#websocket-message-types

進捗を表す progress イベントに加えて、実行完了を示すイベントが送られます。

一覧を見ると execution_success が完了のように見えますが、実際の挙動を確認すると、最後に届くのは executing でした。(公式の説明にあるようにexecutedで終わらない)

ワークフローやバージョンで挙動が変わる可能性もありますが、今回の例ではこのイベントを待つ形で動作しています。

公式ドキュメントでも、WebSocketは内部実行イベントをそのまま流す仕組みであると説明されています。イベント順序は将来変更される可能性があるため、WebSocketで送られてくるメッセージを実験して実装を判断するのが無難なようです。

TQDMと組み合わせた進捗表示の例

WebSocketの進捗情報は数値として受け取れるため、グラフィカルに表示することも可能です。

ここでは、PythonではおなじみのTQDMと組み合わせて進捗バーとして表示する例を試します。
以下は、進捗を待ち受ける処理部分を抜粋したサンプルです。

※事前に /prompt で生成ジョブを送信している前提での、WebSocket待機関数部分のみ。

import json
import time
import uuid
import websocket
from tqdm import tqdm

WS_URL = "ws://127.0.0.1:8188/ws"

def wait_and_show_progress(prompt_id: str, client_id: str, timeout: int = 1800):
    ws_url = f"{WS_URL}?clientId={client_id}"
    ws = websocket.create_connection(ws_url)
    start = time.time()
    bar = tqdm(total=100, desc="Comfy", unit="%")
    last = -1
    try:
        while True:
            if time.time() - start > timeout:
                raise TimeoutError("WebSocket timeout")
            raw = ws.recv()
            try:
                data = json.loads(raw)
            except Exception:
                continue
            mtype = data.get("type")
            pdata = data.get("data") or {}
            if pdata.get("prompt_id") != prompt_id:
                continue
            if mtype == "progress":
                v = pdata.get("value")
                maxv = pdata.get("max", 1)
                pct = int(v / maxv * 100)
                if pct != last:
                    bar.update(pct - bar.n)
                    last = pct
            if mtype == "execution_error":
                raise RuntimeError("execution_error from ComfyUI")
            if (mtype == "executing" and pid == prompt_id and pdata.get("node") is None):
                bar.update(100 - bar.n)
                bar.close()
                return
    finally:
        ws.close()

赤線部分は、実験に基づいて終了条件を調節しています。

実行すると、以下のような進捗表示が行われます。

submitted prompt_id: ...(省略)...
Comfy: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:14<00:00,  6.8%/s]

検証したCUDA環境では1分程度の音楽生成で14秒とすぐ終わりますが、長時間待つ生成や、リモートでComfyUIを動かしている場合には、「処理が進んでいる」ことを確認できる安心材料になります。