どうもだだっこぱんだです。今までずっとjs,tsとかを触ってましたが、最近はpythonかなり触ってます。
AIなんも知らない人間ですがAIについて書いてみました。
まずは結果を
右は普通の生成方法。Diffusersにattention_slicingを有効にして生成した結果です。
真ん中はxformersを有効にした場合。
左は今から紹介するTensorRTを使用する方法です。
はじめに
この記事では、技術関係なんも知らない！っという方向けのColabで試すだけの方法。
ある程度知識がある方向けのWindowsやLinuxで実際に動かす方法の二つを紹介します。
それいじょうのことはぼくはなにもわかりません🫠
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さんのこちらの記事を見ていただければわかると思います
ビルドの設定
左の目次から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でございます🤗
普通に実行
Windowsの場合はWSLを利用してください。
Ubuntu前提で説明していきます。
CUDAをインストール
まずはCUDAをインストールします。今回僕は11.7を使います。
👆のリンクからLinux > x86_64 > Ubuntu or WSL-Ubuntu
を選んでいきます。
今回は**deb (local)**の方法でインストールします。
方法は上のサイトに書いてあるので割愛！
cuDNNをインストール
cuDNN 8.6.0をインストールします。
debファイルをダウンロードしてcudaをインストールするサイト同じ要領で進めていきます。
TensorRTをインストール
TensorRT8.5をインストールします。
これもほぼ同じ要領です。
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.pyのCompVis/stable-diffusion-v1-4の部分を書き換えて使ってみてください。
Discussion