Jetson nanoでStableDiffusionを動かすメモ
本当に取り急ぎのメモ。Jetson nanoもしくはJetson Orin NXが手に入ったら追記。たぶん。
そもそもJetson nanoでStableDiffusionが動くのか
どうやら動くらしい
9月くらいにJetsonの存在を知り、実際に動かしている人を探し続けてようやく発見
jetson nanoでstable diffusion動いたけどやはりメモリが課題で画像サイズを下げないと厳しい
— モンステラ (@momonstera) November 2, 2022
環境構築できただけで割と満足だけどメモリ節約もうちょい頑張ってみようかな
どのように環境構築されたのか聞いてみた
とても丁寧に書いてくださりました、感謝しかない…正直に言うと「WebUIがそのまま動かないかな」とか思ってたので戦々恐々としている。ちゃんと勉強しなきゃ…
まずjetson nanoに対応してるjetpackはpython3.6だったので最新のdiffusersが対応してなさそうと思い、ubuntu20.04 python3.8 cuda対応pytorch1.12の環境をdockerで作りました。jetpackmlの環境ではpipから最新のdiffusersがインストールできませんでしたがgithubからなら動いたりしますかね?🤔
— モンステラ (@momonstera) November 13, 2022
jetson nanoのcudaが動くdocker構築参考サイトhttps://t.co/aaedsLvQww
— モンステラ (@momonstera) November 13, 2022
python3.8でcuda対応のpytorchが公式で配布されてなさそうだったため、個人でコンパイルする必要がありそう
jetson nanoでpython3.8に対応したpytorchの配布&ゼロから作る方法が載ってるサイトhttps://t.co/xKnssBGCnq
diffusersでSDを普通に動かすとメモリが足りなくてKilledされてしまうのと、swap領域を作れば一応動くがめちゃくちゃ遅くなるので工夫が必要
— モンステラ (@momonstera) November 13, 2022
sliced attentionとcpu offload、euler schedulerとか試しましたhttps://t.co/3By9zJGycH
cpu offloadを使う場合pipe = pipe . to("cuda")は必要なさそう
あと、runwayml/stable-diffusion-v1-5を使う場合、v1-5-pruned-emaonly.ckptの方を使ってくれてるのかよく分からなかったため、ローカルにダウンロードしてckptからdiffusers用モデルに変換して動かしてます
— モンステラ (@momonstera) November 13, 2022
@Kaitu_Ai
— モンステラ (@momonstera) November 13, 2022
こんな感じですかね？それで速度とクオリティですが、速度はswap領域を使うか使わないかにかかってて、めちゃくちゃぶれます。512×512画素、推論25ステップで約10分〜1時間でした。
256×256にすると最大で30分くらいでしたかね？
dancing cat, concept art pic.twitter.com/bsUfuG5NxF
メモリ節約などの改善について
軽量なデスクトップ使ったり、デスクトップ使わずsshで動かしたりするとメモリが節約されて早くなりそうと思ってます
— モンステラ (@momonstera) November 13, 2022
あと、SDを構成するモデルを全部メモリに乗せるよりも1つ1つ読み込ませて、推論して、解放させるってしたほうが速いかもと考え、今度試して見ようと思ってます
作成されたDockerfileのスクリーンショット
色々雑で恥ずかしいですが、Dockerfileのスクショも置いておきますね pic.twitter.com/noCEW3l7LP
— モンステラ (@momonstera) November 13, 2022
補足
Jetson AGX Xavier 512x512 30Steps での生成時間が約25秒らしい
メモリの差があるものの、Jetson Orin NX 16GBなら実行時間が1/4くらいになるかも？