ComfyUIのプラグインを作る！

はじめに

ComfyUIはStableDiffusion等の画像生成AIをブラウザ上のノードベースUIで操作できるようにしたツールです。ComfyUIについては以前書いた記事もあるので良ければご覧ください。

今回はかなりマニアックですが、ComfyUI自体の裏側の仕組みやそのプラグイン作成に焦点を当てた解説があまり無かったので、そういったところを解説しています。また、ソースコードを眺めていて学びがあったのでそういった知見の共有として書いています。何かの参考になれば幸いです。

ComfyUIの大まかな構成

ComfyUIの構成はPythonで記述されたサーバー側とJavascriptで書かれたクライアント側に分けられます。サーバー側(py)では主にノードで組まれたワークフローの実行やノードの定義などを担っています。一方、クライアント側(js)ではブラウザに描画されるUIの部分を担っています。
(ちなみにノードUIの実装はlitegraph.jsというフレームワークを使っているみたいです。)

ComfyUIではサーバーとクライアントをインタラクティブに通信させるため、通常のHttp通信に加え、WebSocket通信を採用しています。それぞれの通信方法をどのように活用しているのか、具体例を見てみましょう。

Http通信

Http通信はクライアント側からサーバー側へリクエストを送信し、その返答を受信する一問一答のような形式でやり取りを行う通信手法です。

ComfyUIではワークフローがサーバー側で実行される故、各ノードのインプットやアウトプット、処理の内容などはサーバー側(nodes.py)で定義されています^[1]。このノード定義をノードUIとしてブラウザで描画するために、Http通信を利用してクライアント起動時にその情報をサーバー側にリクエストしています。

コード上の処理の流れ：1 -> 2 -> 3 -> 4

上図のように、単に何かの情報を取得したい等のケースではHttp通信のみを使えます。ただ、先述の通り、クライアントからサーバーへの単方向の一問一答形式なので、サーバーからクライアントに逐一何かを報告するみたいな使い方はできません。そこでWebSocket通信を使います。

Http ＋ WebSocket通信

WebSocket通信はサーバーとクライアントが接続している間、双方向のやり取りを可能にする通信手法です。（よくチャットツールとかで使われている）

ComfyUIではHttp通信でPOSTリクエストを行った後、サーバー側で行われている処理をリアルタイムでUIに反映させるため、WebSocket通信で逐一報告するようになっています。下図はQueue Promptボタンを押したときの処理を簡易的に表しています。

コード上の処理の流れ（ノードの実行まで）：1 -> 2 -> 3 -> 4 --> 5 -> 6

サーバー側からクライアントに通知を送るにはserver.pyで定義されたsend_sync関数を使います。例えばKSamplerノードの進捗（プログレスバーに使う）を通知しているのはここで行っています。サーバーからの通知があった際にはapi.jsのここの処理が実行されます。メッセージのタイプに応じて、予め設定しておいたイベントを発火させてることがわかると思います。

下のgifはQueue Promptを押した後のサーバーからのメッセージをそのままコンソールに表示しています。サーバー側からクライアントに通知が逐次渡っているのがわかります。

さて、ここで勘がいい方は気付くかと思いますが、クライアントはなにもComfyUIで用意されているWebUIでなくとも良いのです。つまり、サーバーに送るデータさえ適切に作ることが出来るのであれば、別のアプリケーションからでも画像生成できます。

既にそういったプラグインは存在していて、ComfyUI-BlenderAI-nodeではBlender側でノードを繋げてワークフローを作れるようにしているみたいです。また、Unreal EngineのプラグインComfyTexturesではあらかじめワークフローを作っておいて、入力する文字列や画像を変数化しているっぽいです。
https://github.com/AIGODLIKE/ComfyUI-BlenderAI-node?tab=readme-ov-file
https://github.com/AlexanderDzhoganov/ComfyTextures

ここまでComfyUIのざっくりとした仕組みについて語ってきましたが、ここからは本題で実際にプラグインを作ってみます。

Pythonでカスタムノードをつくる

まずは簡単なカスタムノードを作ってみましょう。
ComfyUIで特定の機能を持ったカスタムノードは、python側で定義してやるだけで作成できます。今回は文字列を入力する"TextInput"ノードとそれをコンソールにprintする"TextOutput"ノードを作ってみます。

1. ディレクトリ構成

まずはComfyUI/custom_nodes下に以下のようなフォルダ・ファイルを作成します。

MyCustomNodes
├── __init__.py
└── nodes.py

2. nodes.pyに書く

nodes.pyにカスタムノードの定義をクラスとして記述します。このクラスの書き方は本家の方のnodes.pyが参考になると思います。

nodes.py

class TextInput:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"text": ("STRING", {"multiline": True})}}
    RETURN_TYPES = ("STRING",)
    FUNCTION = "run"
    CATEGORY = "MyCustomNodes"

    def run(self, text, seed = None):
        return (text,)

class TextOutput:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"text": ("STRING", {"forceInput": True})}}
    OUTPUT_NODE = True
    RETURN_TYPES = ()
    FUNCTION  = "run"
    CATEGORY = "MyCustomNodes"

    def run(self, text):
        print(text)
        return ()

NODE_CLASS_MAPPINGS = {
    "TextInput": TextInput,
    "TextOutput": TextOutput,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "TextInput": "TextInput",
    "TextOutput": "TextOutput",
}

3. __init__.pyに書く

NODE_CLASS_MAPPINGSとNODE_DISPLAY_NAME_MAPPINGSを明示的に公開させるため、__init__.pyの方で以下のように記述します。

__init__.py

from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']

ちなみにカスタムノードがどこでどうやってロードされているかというのが、本家nodes.pyのここの辺りでcustom_nodesディレクトリを走査して__init__.pyを見つけ、動的にモジュールをインポートしているのが分かります。

4. 動作確認

ComfyUIを起動し、"TextInput"、"TextOutput"ノードが追加されているのを確認します。
下図のようにQueuePromptを押しコンソールにテキストがprintされればokです。

5. nodes.pyを少し修正

しかし、このままでは一つ問題があって、テキストを編集した後の最初の一回は正常にprintされるのですが、連続で実行するとなにもprintされません。これはComfyUIの仕様で、入力が以前実行したときと同じ場合、そのノードは実行しないようになっています。

これを解決するには、実行時に毎回値が変化するシード値を"TextInput"の入力に設けます。この入力タイプはINT:seedという名前としてデフォルトで用意されてます。

nodes.py

class TextInput:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
                    "text": ("STRING", {"multiline": True}),
+                    "seed": ("INT:seed", {}),
                }}

Javascriptでノードの見た目を変える

ただ、新しく追加されたseedに関する入力は直接ユーザーが触る必要はないので、UIとして表示させない方が良さそうです。そこでJavascriptの方でノードの見た目を変えてみます。

1. web.jsを追加

Javascriptを記述するファイルとして、web.jsを所定のディレクトリ下に配置します。

MyCustomNodes
├── js
│   └── web.js  <-これを追加
├── __init__.py
└── nodes.py

2. __init__.pyを修正

追加したjsファイルを認識できるように__init__.pyを以下のように修正します。

__init__.py

from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+ WEB_DIRECTORY = "./js"
- __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
+ __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY']

3. web.jsで特定のインプットを非表示にする

ここでは"TextInput"ノードのseedとそれに関連するインプット(widget)を非表示にしています。
beforeRegisterNodeDefはノード定義が登録される前に呼び出される関数です。

app.registerExtensionの引数で指定するのがComfyExtensionという型のオブジェクトなのですが、その型定義がcomfy.d.tsにあります。それを見れば他にどのような関数が定義できるかがわかると思います。

js/web.js

import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.js";

app.registerExtension({
    name: "MyCustomNodes",
    async beforeRegisterNodeDef(nodeType, nodeData, app) {
        if (nodeData.name === "TextInput") {
            const origOnNodeCreated = nodeType.prototype.onNodeCreated;
            nodeType.prototype.onNodeCreated = function () {
                const r = origOnNodeCreated ? origOnNodeCreated.apply(this) : undefined;
                for (const w of this.widgets) {
                    if (w.name === "seed") {
                        w.type = "converted-widget";
                        if (!w.linkedWidgets) continue;
                        for (const lw of w.linkedWidgets) {
                            lw.type = "converted-widget";
                        }
                    }
                }
                return r;
            }
        }
    }
})

ComfyUIを再起動してseed関連のUIが消えていればokです。

4. web.jsでノードの色を変える

同じような具合で、ノードの色も変えてみます。

ちなみにnodeTypeはLGraphNodeという型で、その定義はlitegraph.d.tsにあります。

js/web.js

app.registerExtension({
    name: "MyCustomNodes",
    async beforeRegisterNodeDef(nodeType, nodeData, app) {
        if (nodeData.name === "TextInput") {
            //中略...
+            nodeType.prototype.color = LGraphCanvas.node_colors.green.color;
+            nodeType.prototype.bgcolor = LGraphCanvas.node_colors.green.bgcolor;
        }
+        else if (nodeData.name === "TextOutput") {
+            nodeType.prototype.color = LGraphCanvas.node_colors.green.color;
+            nodeType.prototype.bgcolor = LGraphCanvas.node_colors.green.bgcolor;
+        }
    }
})

外部アプリと連携させてみる

ここでは簡単な例として、C#のコンソールアプリを新しいクライアントとしてサーバーと接続してみます。「コンソールアプリからサーバーにテキストを送信　→　jsのクライアントにテキストを送信　→　TextInputノードにセット　→　ワークフローを実行　→　結果をコンソールアプリに送信」という一連の流れを実装します。

1. ディレクトリ構成

新しいプラグインとしてMyCustomClientというフォルダを作ります。
そのディレクトリ構成は以下のようにします。

MyCustomClient
├── console
│   ├── console.csproj
│   └── Program.cs
├── js
│   └── web.js
├── __init__.py
└── nodes.py

コンソールアプリの方は.Netが必要です。
.Netをインストール後、MyCustomClient下でdotnet new console -n consoleでプロジェクトを作ると楽です。
また、csprojの方でRestSharpとNewtonsoft.Jsonを使えるように設定しておいてください。

csprojの中身

console.csproj

<Project Sdk="Microsoft.NET.Sdk">

  <PropertyGroup>
    <OutputType>Exe</OutputType>
    <TargetFramework>net7.0</TargetFramework>
    <ImplicitUsings>enable</ImplicitUsings>
    <Nullable>enable</Nullable>
  </PropertyGroup>

  <ItemGroup>
    <PackageReference Include="Newtonsoft.Json" Version="13.0.3" />
    <PackageReference Include="RestSharp" Version="110.2.0" />
  </ItemGroup>

</Project>

2. __init__.pyに書く

前半の部分は先ほどと同じです。
後半にコンソールアプリからHttpリクエストできるエンドポイントを新たに追加します。このエンドポイントでは、WebSocket通信で接続しているjs側のクライアントに"TextInput"ノードのテキストを更新する＆queuePromptを実行するためのメッセージを送っています。

__init__py

import os
import server
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
from aiohttp import web

WEB_DIRECTORY = "./js"
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY']

# テキストを更新し、queuePromptを実行するエンドポイントを追加
@server.PromptServer.instance.routes.post('/my_custom_client/update_text')
async def update_text(request):
    data = await request.json()
    text = data["text"]
    # クライアントでテキストを更新するメッセージを送信
    server.PromptServer.instance.send_sync("update_text", data)
    # クライアントでqueuePromptを呼び出すメッセージを送信
    server.PromptServer.instance.send_sync("run_workflow", {})
    return web.Response()

3. nodes.pyに書く

新たに"TextSend"というカスタムノードを追加します。
このノードは入力されたテキストをWebSocket接続しているコンソールアプリに送信します。

nodes.py

import server

class TextSend:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"text": ("STRING", {"forceInput": True})}}
    OUTPUT_NODE = True
    RETURN_TYPES = ()
    FUNCTION = "run"
    CATEGORY = "MyCustomClient"

    def run(self, text):
        # テキストをクライアントに送信する（コンソールアプリの方で受け取る）
        text = text + " (from TextSend node)" 
        server.PromptServer.instance.send_sync("send_text", {"text": text})
        return ()


NODE_CLASS_MAPPINGS = {
    "TextSend": TextSend,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "TextSend": "TextSend",
}

4. web.jsに書く

jsのクライアント側にWebSocket通信のメッセージが送られてきたときのイベントリスナーを登録します。
apiインスタンスのaddEventListener関数で第一引数にメッセージのタイプと第二引数にイベント発火時の関数を記述します。

js/web.js

import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.js";

app.registerExtension({
    name: "MyCustomClient",

    async setup() {
        // queuePromptを呼び出すイベントリスナーを追加
        api.addEventListener("run_workflow", ({ detail }) => {
            app.queuePrompt(0);
        });

        // TextInputノードのテキストを更新するイベントリスナーを追加
        api.addEventListener("update_text", ({ detail }) => {
            const nodes = app.graph.findNodesByType("TextInput");
            for (const node of nodes) {
                const text = detail["text"];
                node.widgets[0].value = text;
            }
        });
    }
})

5. Program.csに書く

コンソールアプリの実装です。
WebSocket通信でサーバーと接続し、Http通信でテキストを送信します。特定のタイプのメッセージが受信するまでWebSocket接続を維持しています。
コードを書いたらcsprojがある階層でdotnet buildします。

console/Program.cs

using System;
using System.Net.WebSockets;
using System.Text;
using Newtonsoft.Json;
using RestSharp;

namespace MyCustomClient
{
    class Program
    {
        // サーバーアドレスを静的変数で定義
        private static readonly string SERVER_ADDRESS = "127.0.0.1:8188";

        static async Task Main(string[] args)
        {
            // ユーザー入力を受け取る
            Console.WriteLine("Enter message: ");
            string? input = Console.ReadLine();
            if (string.IsNullOrEmpty(input))
            {
                Console.WriteLine("Input is empty. Exiting...");
                return;
            }

            // メッセージ送信を試みる
            try
            {
                await SendMessage(input);
            }
            catch (Exception e)
            {
                Console.WriteLine(e.Message);
            }
        }

        // サーバーにメッセージを送信するメソッド
        static async Task SendMessage(string message)
        {
            // 送信するメッセージを表示
            Console.WriteLine("Sending message: " + message);
            using (ClientWebSocket client = new ClientWebSocket())
            {
                // WebSocketサーバーへ接続
                Uri serverUri = new Uri($"ws://{SERVER_ADDRESS}/ws");
                await client.ConnectAsync(serverUri, CancellationToken.None);

                // RESTクライアントを作成
                RestClient restClient = new RestClient($"http://{SERVER_ADDRESS}");

                // 送信メッセージを格納
                var dic = new Dictionary<string, string>
                {
                    { "text", message }
                };

                // REST APIを通じてメッセージ送信
                RestRequest restRequest = new RestRequest("/my_custom_client/update_text", Method.Post);
                string jsonData = JsonConvert.SerializeObject(dic);
                restRequest.AddParameter("application/json", jsonData, ParameterType.RequestBody);
                await restClient.ExecuteAsync(restRequest);

                // サーバーからの応答を待ち続ける
                while (client.State == WebSocketState.Open)
                {
                    var receiveBuffer = new byte[1024];
                    var result = await client.ReceiveAsync(new ArraySegment<byte>(receiveBuffer), CancellationToken.None);

                    // 受信データをJSONに変換
                    var json = Encoding.UTF8.GetString(receiveBuffer, 0, result.Count);
                    var comfyReceiveObject = JsonConvert.DeserializeObject<ComfyReceivedObject>(json);

                    bool isClose = false;

                    // 応答タイプによって処理を分岐
                    switch (comfyReceiveObject?.Type)
                    {
                        case "send_text":
                            Console.WriteLine("Received send_text: " + comfyReceiveObject.Data?["text"]);
                            isClose = true;
                            break;
                    }

                    // 応答処理後に接続を閉じる
                    if (isClose)
                    {
                        break;
                    }
                }
            }
        }
    }

    // サーバーからの応答を格納するクラス
    public class ComfyReceivedObject
    {
        [JsonProperty("type")]
        public string? Type { get; set; }

        [JsonProperty("data")]
        public Dictionary<string, object>? Data { get; set; }
    }
}

6. 動作確認

ComfyUIを起動し、"TextInput"と"TextSend"を繋げます。
その後、ビルドしたコンソールアプリ(console.exe)を実行します。任意の文字列を送信後、QueuePromptが実行されコンソール側にテキストが返ってきたらokです。

今回は非常に地味ですが、これを応用することで、任意のアプリケーションとComfyUIを連携させることが可能になります。
(以前作成したGrasshopperと連携するプラグインComfyGHはこれとほぼ同じ原理です)

まとめ

今回はComfyUIの裏側の仕組みとプラグインの作成というだいぶマニアックな解説をしました。
個人的な感想としては、ComfyUIは外部からカスタマイズし易いように設計されていて、そういった要素も人気の要因の一つになっているのだと感じました。また、任意のアプリケーション上でAIを使いたいとなったとき、そのアプリ上で実行する環境を一から作成するのではなく、こういったカスタマイズ性に富んだフレームワークに乗っからせる方が色々と恩恵が得られそうだなと思いました。
今回作成したコードは以下のgithubリポジトリに置いてあります。
https://github.com/4kk11/MyConfyPlugin

脚注

1. PrimitiveやNoteなどの「値をインプットにセットする」、「UIで表示するだけ」みたいなノードはjs側のみで作られています