西川和久の不定期コラム
Stable Diffusion高速化技術続々登場！TensorRT、SSD-1B、LCMなどを試してみる
2023年11月9日 06:06
今年7月27日に次世代Stable DiffusionのSDXL 1.0が登場してから数カ月経ったが、ここに来てSD(XL)を高速化する技術がいくつか出てきたので紹介したい。
VRAMオフロードのオン/オフが可能になったNVIDIAドライバ546.01
高速化技術をご紹介する前に以前から指摘していたNVIDIAドライバのVRAMオフロードによるAUTOMATIC1111などのスローダウン。待つこと半年近くでやっとオン/オフできるドライバの546.01が登場。設定方法などをご紹介したい。
VRAMオフロードとは、VRAMが不足した時、メインメモリを使う手法だ。これによって、たとえばVRAMが8GBでも、メインメモリが許す限りいくらでも見かけ上のVRAMを増やすことができる。一般的にアプリはCUDA経由でアクセスするため、コードを1文字も変更することなく、これに対応可能だ。
「それはGood! VRAMが8GBでもメインメモリ32GBなら24GB程度は楽勝では!? 」と思うだろうが、そうはいかない理由がある。
LLMだと130億パラメータ以上になると、とにかくVRAMが多くないとModel自体をロードできず動かない。どちらかと言えばGPUの速度よりVRAM容量が肝となる。そしてこのケースではVRAMオフロードは有効。実際筆者はcodellama-13b-instruct.Q4_K_M.ggufをChatGPT替わりに日頃使っている。VRAM+メインメモリで39.6GB中35.7GB(内VRAM 24GB)ほどが必要となるもののレスポンスは問題ないレベルだ。
ところがStable DiffusionでUpscaleやControlNetなどVRAMを多く使う処理の場合、VRAM不足分をメインメモリで補うと処理速度がガクンと落ち(10倍前後の時間がかかる)、いくら待っても処理が終わらないと言うスローダウンが発生する。
さらに、たとえばこれまでVRAM 12GB内で処理できていたものが、VRAMオフロードに切り替えるタイミングが速いのか、すべてVRAMでまかなわずメインメモリも使用しだすため、やはりスローダウンが発生する。正直この状態では何倍も待つため使い物にならない。
従って筆者は結構長い間、VRAMオフロードを搭載していないドライバ、531.61を使ってきた経緯がある。
「nvidia-smiかNVIDIAコントロールパネルでオン/オフできれば……」と思っていたが、10月31日に対応した546.01がリリースされた。
手順はNVIDIAコントロールパネルのプロラム設定でVRAMオフロードをオフにしたいpython.exeを指定し、CUDA - Sysmem Fallback Policy / Prefer No Sysmem Fallbackに変更する(戻す/オン時はDriver Default)。またpython.exeを使うアプリは再起動する必要がある。
注意する点としては、pythonのバージョンを変更できるminicondaなどを使っていると、動作しているpython.exeは環境ごとにpathが変わる。これを探すにはタスクマネージャーを使いpythonでプロセスを検索、出た一覧からプロパティでpathを見れば場所が分かる。
これでやっと最新ドライバが使えるようになり一安心。VRAMオフロードをオフにした結果、AUTOMATIC1111などがスローダウンせず動作している。
参考までにLinux版のドライバは今のところVRAMオフロード未対応(535.129.03)。もちろんVRAM容量が足らないと、CUDAで指定容量が確保できずエラーとなるが、これは以前からそうなので問題ない。
TensorRT
10月18日、NVIDIAは最新ドライバでStable Diffusionを倍速にするTensorRT(Tensor Real-Time)に対応。AUTOMATIC1111の拡張機能、Stable-Diffusion-WebUI-TensorRTをリリースした。
動作原理をザックリ説明すると、TensorRT用にModelを変換し、GPU上で推論、高速化する仕掛けだ。設定方法は、
- Extensionsでhttps://github.com/NVIDIA/Stable-Diffusion-WebUI-TensorRTをインストール
- 再起動後、TensorRTの項目が増えているのを確認
- Settings > User interface でQuicksettings listにsd_unetを追加。再起動。上部にSD Unetが増えているのを確認
こんな感じで一般的な拡張機能をインストールするのと(3)の手順が増えている以外は同じ。
実際の使用方法は、
- 使用したいModelを事前に選択しておき(ただしSD 1.5用)、TensorRT Exporterをクリック。次に[Export Default Engine]をクリック
- 何分(GPU性能による)か待つとonnx形式とtrt形式ファイル(どちらもGB級)ができあがる
- 変換したModelを選び、SD Unetで対になる名前を選ぶ
- 以降は普通にいつも通り生成する
と、ざっくり言えばこんな感じだ。タブにあるTensorRT LoRAは、使いたいLoRAを指定し、これも変換する。保存場所は[AUTOMATIC111フォルダ]/models/Unet-trtとUnet-onnxになる。
ただしこの状態ではADetailer(顔を解像度に合わせて書き直す機能)は使えるものの、UpscaleとControlNet、FreeUは動作しない(ほかもあるかもしれないが)。
まずControlNetは何をしても今のところ使えない。FreeUも同様。これはあきらめるしかない。Upscaleは、TensorRT自体がModelに加え解像度にも依存するため一工夫する必要がある(token数、batch数にも依存するがここは必要に応じて)。
たとえばSD 1.5で一般的な解像度だと、512x512、512x768(768x512)などがある。そしてこれをフルHD相当にUpscaleした場合、1,920x1,920、1,280x1,920(1,920x1,280)となるだろうか。TensorRTは、これらすべての解像度を含むtrt形式ファイルを用意しなければならないのだ(onnx形式は各Model1つのみ)。
上記のTensorRT Exporterの一覧を見ると、
512x512 | Batch Size 1 (Static)
768x768 | Batch Size 1 (Static)
1024x1024 | Batch Size 1 (Static)
256x256 - 512x512 | Batch Size 1-4 (Dynamic)
512x512 - 768x768 | Batch Size 1-4 (Dynamic)
768x768 - 1024x1024 | Batch Size 1-4 (Dynamic)
このような項目が並ぶ。Staticは記述の解像度固定、Dynamicは記述の解像度から解像度までどんな解像度でもOKとなる(ただし64の倍数)。
見るとこの中にはUpscale用の解像度がない。この場合、Advanced Settingsを開き、任意の幅x高さの解像度を入力してtrt形式ファイルを作ることになる。この時、Staticの何かを選んでAdvanced Settingsを開くのと、Dynamicの何かを選んでAdvanced Settingsを開くのでは出てくる項目の内容が異なる(前者の方が少ない)。
StaticならOptimal height/Optimal width。DynamicならMin height/Max height、Min width/Max widthを変更する。日頃縦位置のポートレートが多いならStaticの場合、
- 512x768
- 1,280x1,920
この2つのtrt形式ファイルを作りSD Unetで切り替えながら運用するか、Dynamicだと、
- 512x512-1,920x1,920
としておけば縦位置でも横位置でも、そしてUpscaleにも1つで対応できる。
いくつも解像度の異なるtrt形式ファイルを作ってしまうと、該当する解像度ごとにファイル切り替えが発生し、ここで時間がかかり、倍速の威力が逆にマイナスになってしまうため要注意。
筆者はこの理由で512x512-1,920x1,920(Dynamic)を1本だけ作り処理している。これなら縦位置でも横位置でもUpscaleでも何でもこなせる。が、1点問題があるとすればVRAMをより多く消費することだろうか。
さて、本当に速いのか? 確認したところ……
Prompt: (best quality),portrait photo of a japanese woman sitting on a car,clear eyes,middle breasts,t-shirt,shorts,35mm photograph,film,bokeh,professional,4k,highly detailed,
Negative prompt: (worst quality),llustration,3d,2d,painting,cartoons,(deformed|distorted|disfigured:1.2),(mutated hands AND fingers:1.2),3fingers,4fingers,long fingers,
Steps: 20, Sampler: DPM++ 2M Karras, CFG scale: 5, Seed: 4053779178, Size: 512x768, Model: beautifulRealistic_v7
- 512x768
00:01, 19.87it/s vs 00:02, 11.37it/s - 1,280x1,920(R-ESRGAN 4x+ Anime6B)
00:12, 1.54it/s vs 00:15, 1.19it/s
※GeForce RTX 3090をThunderbolt 3で接続
512x768は確かにほぼ2倍速い。Upscale時はUpscalerがCPUで処理するためここまでの差は出ないが、それでも少し速い。LoRAは試したところSD UnetでLoRAの項目は1つしか選べず、複数LoRAが使えずこれは痛い。
ここまではSD 1.5の話で通常のAUTOMATIC1111で動作する。SDXLは、(執筆時)AUTOMATIC1111をdevへ切り替えなければ使えない(git switch dev)。dev、つまり開発版はリリース版より安定していないこともあり常用は避けたいところ。またLoRAは変換時エラーで落ち、変換できない=使えない。
速度比較すると(Promptは上記と解像度以外同じ。もちろん832x1,216などSDXL固有の解像度で別途onnx/trt形式ファイルを使用)、下記のようにSD 1.5ほどではないものの、それでも気持ち速くなる。ただしUpscaleはCPU側がより多く動くので差がなくなる。
- 832x1,216
00:05, 5.73it/s vs 00:06, 3.42it/s - 1,280x1,920(R-ESRGAN 4x+ Anime6B)
00:20, 1.66it/s vs 00:19, 1.29it/s
以上のように、SD 1.5の時、本当に倍速になるTensorRTだが、欠点も説明した通り山盛りだ。万人におすすめできる機能ではない。またこの拡張機能の副作用なのかバグなのか、img2imgのinterrogate CLIPがエラーで落ちる。
従って使い方としては「いつも使うModelが決まっている」、「LoRAは1つ(SD 1.5のみ)」、UpscaleはOKな写真を後でまとめて行なう……という、(あるかどうか分からないが)業務用途的な感じだろうか。
できれば通常版でのSDXL対応や、ControlNet/FeeU対応、そしてimg2imgのinterrogate CLIPエラーも修正してほしいところだが、リリース直後何度か更新がありバタバタしたが、その後まったく動きがない。やりました! 的なポーズでないことを祈りたい。
SSD-1B
SSD-1B(Segmind Stable Diffusion Model) は、SDXLを50%小型化し、txt2imgの生成機能を維持しながら60%の高速化を実現したModelだ。詳細はsegmind/ SSD-1Bを参考にしてほしいが、実際画像を生成するには、以下のようなコードとなる。
from diffusers import StableDiffusionXLPipeline
import torch
import torchvision
pipe = StableDiffusionXLPipeline.from_pretrained("segmind/SSD-1B", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
#pipe = StableDiffusionXLPipeline.from_pretrained("furusu/SSD-1B-anime", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
pipe.to("cuda")
# if using torch < 2.0
# pipe.enable_xformers_memory_efficient_attention()
# Prompt here
prompt = "professional portrait photo of a woman, , 20 y.o., sitting, looking at viewer, in cafe"
# Negative prompt
neg_prompt = "(worst quality),llustration,3d,2d,painting,cartoons,(deformed|distorted|disfigured:1.2),(mutated hands AND fingers:1.2),"
image = pipe(prompt=prompt, negative_prompt=neg_prompt).images[0]
# 画像を保存
tensor_image = torchvision.transforms.ToTensor()(image)
torchvision.utils.save_image(tensor_image, "output.jpg")
上の方にあるpipeの1つ(下側)がコメントアウト(#)になっているのはfurusu / SSD-1B-animeと言うアニメ版だ。上側をコメントアウトして下側の#を外せばアニメ画が生成される。
Modelのファイルサイズは4.47GB(SDXLは6.5GB)、生成速度は、
- 50/50 [00:04<00:00, 12.32it/s] vs 50/50 [00:07<00:00, 6.98steps/s]
※GeForce RTX 4090をPCIeで接続
なので確かに倍近く速くなっている。とは言えアニメ版はともかくとして、リアル版は個人的な好みとしてはイマイチ。掲載したサンプルが西洋系の顔なのでまだ見れるが、東洋系の顔だと残念な状態になる……と、直ちに乗り換えられる状況ではない。
LCM(Latent Consistency Model for Stable Diffusion)
筆者はロジックまでは理解していないので、詳細はここを参考にしてほしいが、試すにはAUTOMATIC1111の拡張機能sd-webui-lcmがあるので簡単にできる。Promptは先のTensorRTと同じものを使用。結果はこんな感じだ。
512x512(Guidance scale for base: 8, Number of inference steps for base: 4)の画像4枚を”LCM inference time: 0.42435789108276367 seconds”で生成できる(GeForce RTX 4090をPCIeで接続)。同じ環境でSD 1.5 512x512(batch:4)だと2秒なので約5倍速いことになる。
1枚0.1秒だとリアルタイムのPrompt入力にも対応できそうだ……と言うことで開発されたのがこのReal-Time-Latent-Consistency-Model。使い方はサイトの通りだが、Windows環境での使用時は一点注意。それはtorchがcompileに対応していないため、app-img2img.pyの81行目の、
pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead",fullgraph=True)
ここをコメントアウト(文頭に#を付ける)しなければならない。このcompileはModelを最適化するものだが、Linux環境では使えるものの、変換にものすごく時間がかかる上にVRAM 24GBでも途中でVRAM不足エラーとなり、結果どうなのか不明のままだ。
試すと確かにリアルタイムでPrompt入力に対応し、画像がバンバン出るのでおもしろいが、単語の途中でも反応するため、妙な絵も多く出る。ある意味先のAUTOMATIC1111の拡張機能の方が実用的だろう。とは言え、出てくる絵はSSD-1B同様、ご覧のようにまだまだ。
以上、Stable Diffusionの高速化技術、TensorRT、SSD-1B、LCMを試してみた。すべて標準の状態より速く生成できるものの、絵のクオリティで同レベルはTensorRTのみ。そのTensorRTも制限がかなり多い。SSD-1BとLCMはこれからと言った感じか。
だから駄目と言うつもりはなく(SDXLも出だしは結構ひどかった)、今後に期待したいところ。ソフトウェアで倍速(以上)になるならこれほどうれしいことはない。