🚀

TensorRTを使ってStableDiffusionを爆速にする

2023/01/31に公開1

どうもだだっこぱんだです。今までずっとjs,tsとかを触ってましたが、最近はpythonかなり触ってます。
AIなんも知らない人間ですがAIについて書いてみました。

まずは結果を

右は普通の生成方法。Diffusersにattention_slicingを有効にして生成した結果です。
真ん中はxformersを有効にした場合。
左は今から紹介するTensorRTを使用する方法です。

はじめに

この記事では、技術関係なんも知らない!っという方向けのColabで試すだけの方法。
ある程度知識がある方向けのWindowsやLinuxで実際に動かす方法の二つを紹介します。
それいじょうのことはぼくはなにもわかりません🫠

Colabでお試し

ノートブックを開く

https://github.com/ddPn08/tensorrt-diffusion-colab

👆のリンクからノートブックのリポジトリにアクセスし、OPEN IN COLABボタンをクリックしてノートブックを開きます。

事前ビルド済みエンジンで実行

事前ビルド済みのエンジンを使って(詳しいことは後ほど)画像生成をしてみます。

ノートブックを開くとこのようなページに飛びます。
まずは上二つのセルを実行しましょう。(ドライブをマウントしなくて良い方はチェックボックスを外してください。)

二つのセルが正常に終了したら左の目次から"Inference with the tensorrt engine"まで飛びます。

そうしたら、一つ目のClone model repositoryを実行します。
engine_repositoryにHuggingFaceのリポジトリID入力すると他のモデルを指定できます(通常のチェックポイントファイルやDiffusersファイルとは形式が異なります。)

とりあえずデフォルトのままで正常に生成できるはずなのでやってみましょう

生成できました。

|------------|--------------|
|   Module   |   Latency    |
|------------|--------------|
|    CLIP    |     29.76 ms |
| UNet x 150 |  15593.99 ms |
|    VAE     |    125.54 ms |
|------------|--------------|
|  Pipeline  |  15756.69 ms |
|------------|--------------|

Colabの無料枠(T4)で150stepで15秒です。
えぐいですね

好きなモデルをビルドする

先ほどは事前に僕がビルドしたモデルを使って生成しました。
好きなモデルで高速生成するには各自でビルドしていただくことになります。

モデルをDiffusers形式に変換する (ckptしかない場合)

ckptのみ配布されているモデルは、まずDiffusers形式に変換しないといけません。
詳しい方法はkohyaさんのこちらの記事を見ていただければわかると思います
https://note.com/kohya_ss/n/n374f316fe4ad

ビルドの設定

左の目次からBuild tensorrt engineに飛びます。

各項目を説明していきます。

unet_pretrained_model_id UNetを読み込むDiffusersモデルのIDまたはパス
vae_pretrained_model_id VAEを読み込むDiffusersモデルのIDまたはパス
clip_pretrained_model_id CLIPのIDまたはパス
engine_dir Engineを保存するパス
onnx_dir ONNXモデルを保存するパス
hf_token HuggingFaceのトークン
denoising_prec fp16またはfp32を選択 ※1
height 生成時の縦解像度 ※2
width 生成時の横解像度 ※2

※1 fp16を選択した場合、モデルリポジトリのfp16ブランチが参照されます。
※2 ビルドしたEngineはおそらく解像度に依存しています(sda-nodeというソフトで解決されているかもしれません。)
 👆おそらく解像度に依存しているわけではない

これ以降の設定項目は自分もあまり理解していないので割愛します...

ビルドを実行

無料枠だとビルド中メモリ不足で止まる場合がありますが、気にせず再開して下さい。
最後の行にC^とかかれていたり、生成されたengineの数が足りなかったりしたら途中で中断されています。
まあまあ時間がかかるので画像生成でもしながら待ちましょう。

HuggingFaceにアップロード

ビルドしたエンジンをHuggingFaceにアップロードしておけます。

ローカルで実行

ローカルで実行する方法です。

Dockerで実行

こちら、Dockerfileでございます🤗
https://gist.github.com/ddPn08/6de89bf678040dbd38aaa6f093bcf5ef

普通に実行

Windowsの場合はWSLを利用してください。
Ubuntu前提で説明していきます。

CUDAをインストール

まずはCUDAをインストールします。今回僕は11.7を使います。

https://developer.nvidia.com/cuda-11-7-0-download-archive

👆のリンクからLinux > x86_64 > Ubuntu or WSL-Ubuntu
を選んでいきます。
今回は**deb (local)**の方法でインストールします。
方法は上のサイトに書いてあるので割愛!

cuDNNをインストール

cuDNN 8.6.0をインストールします。

https://developer.nvidia.com/rdp/cudnn-archive

debファイルをダウンロードしてcudaをインストールするサイト同じ要領で進めていきます。

TensorRTをインストール

TensorRT8.5をインストールします。
これもほぼ同じ要領です。

https://developer.nvidia.com/nvidia-tensorrt-8x-download

tensorrt tensorrt-dev tensorrt-libs

をインストールしてください。

nvinferのプラグインをビルド

tensorrtで画像生成するにはまず、TensorRTのプラグインをビルドする必要があります。
まずはTensorRT OSSをクローンします。

git clone https://github.com/NVIDIA/TensorRT
cd TensorRT
git submodule update --init --recursive
export TRT_OSSPATH=$PWD

次にmakeとかcmakeとかを実行してビルドします

cd $TRT_OSSPATH
mkdir -p build && cd build
cmake .. -DTRT_OUT_DIR=$PWD/out
cd plugin
make -j$(nproc)

export PLUGIN_LIBS="$TRT_OSSPATH/build/out/libnvinfer_plugin.so"

実行する

事前にpythonをインストールしておいてください。
$TRT_OSSPATH/demo/Diffusion に移動します。

次に、venvを作ってからモジュールをインストールします。

python -m venv venv
source venv/bin/activate
pip install --upgrade pip
pip install --upgrade tensorrt
pip install -r requirements.txt

とりあえずそのまま実行します。
HuggingFaceのトークンを取得しておいてください。

mkdir -p onnx engine output
export HF_TOKEN=<HuggingFaceのトークン>
LD_PRELOAD=${PLUGIN_LIBS} python3 demo-diffusion.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN -v

VRAMをまあまあ使います。
また、デフォルトでは"CompVis/stable-diffusion-v1-4"が使用されます。
別のモデルを使いたい場合は、models.pyCompVis/stable-diffusion-v1-4の部分を書き換えて使ってみてください。

AIものづくり研究会

Discussion

hotatehotate

とても勉強になる検証ありがとうございます。
TensorRTとstable-fastは併用可能なのでしょうか。またその際どれだけ早くなるのでしょうか
私のComfyUIの環境でどちらも上手く動かず、お手隙の際に検証していただけますと泣いて喜びます🙇