TensorRT で ComfyUI の画像生成を加速する [comfyanonymous/ComfyUI_TensorRT]

GeForce RTX 世代の GPU を使っていると、TensorRT で推論を加速することが可能です。ComfyUI でも TensorRT が利用可能になるカスタムノードを、以下の GitHub で公開してくれているので、こちらを利用してみます。

comfyanonymous/ComfyUI_TensorRT

なお、対応するモデルは上記 GitHub の Supports: 欄を参照してください。

ComfyUI 環境準備

Stable Diffusion 3 Medium を実行する環境を ComfyUI で用意した手順 を実行し、コンテナで ComfyUI を実行できるようにします。

起動 & Custom Node 導入

ComfyUI のディレクトリで以下コマンドを実行し、起動します。

python3 main.py

ブラウザで UI を開いたら、右下の Manager をクリックします。

ComfyUI Manager Menu が開くので、Custom Node Manager をクリックします。

上部の「Search」欄に tensorrt と入力します。すると TensorRT Node for ComfyUI があるので、そこの Install をクリックします。

しばらく待つと「Restart Required」と表示されるので、Restart ボタンをクリックします。あわせて、ブラウザの再読み込みを行い、画面を更新します。

TensorRT Engine を構築

生成する前に、既存のモデル (checkpoint) から TensorRT のエンジンを構築する必要があります。実際に生成する際は、モデルではなくこのエンジンを読み込む形になります。

まず、ComfyUI でワークフローをクリアし、以下のように Load Checkpoint と DYNAMIC TRT_MODEL CONVERSION を追加し、それぞれ MODEL をつなげます。

そのうえで、Load Checkpoint の ckpt_name を読み込みたいモデルにします。今回は Stable Diffusion 3 Medium のモデルを使います。

また、DYNAMIC TRT_MODEL_CONVERSION では filename_prefix を必要に応じて変更します。なお、この文字に続けてパラメータがファイル名に付加されるので、変更しなくても区別は可能です。

さらに、height_xx や width_xx の各パラメータも変更します。今回、VRAM 8GB の GeForce RTX 2070 を使用していますが、512 より大きな値にするとエンジンの構築に失敗しました。そのため、すべて 512 で統一しました。この値で指定した範囲内の解像度が利用できる形になります。

※後で気が付いたのですが、上のように min, max 解像度を同じ値にする場合、DYNAMIC TRT_MODEL_CONVERSION ではなく STATIC TRT_MODEL_CONVERSION を使用すればよかったようです。とりあえずこの記事では DYNAMIC TRT_MODEL_CONVERSION を使用して続行します。

これで Queue Prompt をクリックすると、エンジンが構築されます。ターミナルに処理状況が表示されます。

構築が完了すると、Prompt executed in xxx seconds と表示されます。今回は 331.47 秒かかりました。

ターミナルから ComfyUI/output/tensorrt ディレクトリを確認すると、エンジンが出力されていることが確認できます。

生成

Stable Diffusion 3 Medium を実行する環境を ComfyUI で用意した手順 と同様、Stable Diffusion 3 Medium を実行するワークフローを読み込ませます。そのうえで、モデルの読み込み部分を置き換えます。

まず、空白部分でダブルクリックし、検索画面を出したら tensor と検索します。すると TensorRT Loader があるので、それをクリックします。

TensrRT Loader が表示されるので、unet_name を構築したエンジンの名前に、model_type は今回は sd3 を選択します。

また、TensorRT Loader の MODEL を、ModelSamplingSD3 の model に接続します。

あとは EmptySD3LatentImage の witdh, height, batch_size を、エンジンを構築した際の値とマッチするように変更します。

これで Queue Prompt をクリックし、画像を生成します。すると、初回はエンジンのロード等で結構時間がかかります。今回は 43.92 秒かかっています。これだとあまり早く感じませんが、これは初回のみです。

同じ条件で再度生成してみると、以下の通り 4.28 秒で生成されます。これだとほとんど待たされる感じはしないですね。