🐈

WSL2にUbuntu22.04とCUDAをインストールしてJupyterLab上でStable Diffusionを動かしたメモ

2022/08/27に公開約5,900字

はじめに

話題のStable Diffusionをローカル環境で動かすまでに行った作業のメモです。ついでにJupyterLabでも使えるようにしました。

作業環境

  • OS:Windows 11(21H2 Build 22000.856)
  • GPU:RTX3070

手順

基本的にはこちらに書いてある内容と一緒です。

https://qiita.com/ryohassay/items/15fab8a72847be13d16d

WSL2を有効にする

https://docs.microsoft.com/ja-jp/windows/wsl/install

Ubuntuを入れる

Microsoft StoreからWSL用のUbuntuをインストール。Ubuntu 22.04.4 LTSを入れた。
インストールしたあとはとりあえず一度起動しておく。

初期状態でバックアップとっておく

やってもやらなくてもよいが、まっさらなOSを保存しておくとあとでなにかと便利。

wsl.exe --export Ubuntu-22.04 Ubuntu-22.04_original.tar

などとしてバックアップ。これから作る環境用にインポート。

wsl.exe --import sd_ubuntu-22.04 D:\sd .\Ubuntu-22.04_original.tar

Ubuntu初期設定

WSL2のイメージをImportするとログインユーザがrootになってしまうようなので修正。ついでにaptの問題も修正しておく。

cat << EOF > /etc/wsl.conf
[user]
default=your-user-name
EOF

aptでエラーが出るので実行しておく。

https://level69.net/archives/30464
# apt purge needrestart
# apt update
# apt upgrade
# apt autoremove

ここで一度Ubuntuを閉じ、念のためWSLを再起動。
再度開くと一般ユーザになっているはず。

wsl -t sd_ubuntu-22.04

ドライバーの更新(Windows)

WSL2でGPUを使うために、WSL2側のUbuntuにドライバーをインストールする必要はない。
下記のアドレスにアクセスして、「Get CUDA Driver」から最新のドライバをインストールする。

https://developer.nvidia.com/cuda/wsl

ドライバーのインストール(Ubuntu)

ここの「Option 1: Installation of Linux x86 CUDA Toolkit using WSL-Ubuntu Package – Recommended」という手順を実行すれば必要なドライバーやライブラリはほぼ入る。

https://docs.nvidia.com/cuda/wsl-user-guide/index.html#cuda-support-for-wsl2
$ wget https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-wsl-ubuntu.pin
$ sudo mv cuda-wsl-ubuntu.pin /etc/apt/preferences.d/cuda-repository-pin-600
$ wget https://developer.download.nvidia.com/compute/cuda/11.7.0/local_installers/cuda-repo-wsl-ubuntu-11-7-local_11.7.0-1_amd64.deb
$ sudo dpkg -i cuda-repo-wsl-ubuntu-11-7-local_11.7.0-1_amd64.deb
$ sudo cp /var/cuda-repo-wsl-ubuntu-11-7-local/cuda-*-keyring.gpg /usr/share/keyrings/
$ sudo apt-get update
$ sudo apt-get -y install cuda
$ rm cuda-repo-wsl-ubuntu-11-7-local_11.7.0-1_amd64.deb

ここまででいったんWSL2を再起動。

wsl -t sd_ubuntu-22.04

nvidia-smiコマンドでGPUが認識されていればたぶんOK

$ nvidia-smi
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 515.65.01    Driver Version: 516.94       CUDA Version: 11.7     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  On   | 00000000:01:00.0  On |                  N/A |
|  0%   42C    P8    20W / 156W |   1404MiB /  8192MiB |      1%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

CuDNNインストール

下記からCuDNNのパッケージをダウンロード(NVIDIAのアカウント作成が必要)

https://developer.nvidia.com/cuDNN
wgetなどのコマンドでダウンロードできないので、WSL環境にコピーしてください。
$ sudo dpkg -i cudnn-local-repo-ubuntu2204-8.5.0.96_1.0-1_amd64.deb
$ sudo cp /var/cudnn-local-repo-ubuntu2204-8.5.0.96/cudnn-local-*-keyring.gpg /usr/share/keyrings/
$ sudo apt update
$ sudo apt install libcudnn8
$ sudo ln -s /usr/lib/wsl/lib/libcuda.so.1 /usr/local/cuda/lib64/libcuda.so
$ export LD_LIBRARY_PATH="/usr/local/cuda/lib64:$LD_LIBRARY_PATH"

minicondaインストール

Anacondaはデカいのでminicondaをインストール。

$ wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
$ bash Miniconda3-latest-Linux-x86_64.sh
[Enter]
[yes(ライセンス同意)]
[Enter]
$ export PATH=~/miniconda3/bin:$PATH >> ~/.bashrc
$ source .bashrc
$ rm Miniconda3-latest-Linux-x86_64.sh

環境準備

minicondaの環境を作成する。

$ mkdir StableDiffusion
$ cd StableDiffusion
$ conda create -n StableDiffusion python=3.8.5
$ conda activate StableDiffusion
PyTorchがらみのパッケージ
$ conda install pytorch-gpu=1.11.0 torchvision=0.12.0 cudatoolkit=11.7 -c pytorch -c conda-forge
Stable Diffusionがらみのパッケージ
$ conda install diffusers transformers scipy ftfy -c conda-forge

次のコマンドを打って”True"が返ってくれば無事にGPUが動作しています。

$ python -c "import torch; print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0));"
True
NVIDIA GeForce RTX 3070

Jupyter Lab導入

Jupyter Labをインストールする。そのままだと進捗表示が出なかったので導入。最新版では表示が崩れたのでipywidgetはバージョン指定が必要だった。

$ pip install jupyterlab ipywidgets==7.7.2

これにて導入完了。

動作確認

モデルデータへのアクセストークンの取得

こちらを参考に、事前にモデルデータにはアクセスできるようにして、ユーザーアクセストークンを取得しておく。

https://note.com/npaka/n/ndd549d2ce556

一度コマンドラインで

$ huggingface-cli login

を実行してアクセストークンを一度入れておいたほうがいいらしい。

JupyterLab上での動作確認

JupyterLabを起動して、ブラウザからアクセスする。IPアドレスは"ip a"コマンド等で事前に調べておく。

$ jupyter lab --NotebookApp.token='' --ip 0.0.0.0 

Notebookでの実行

自分の環境ではGPUのメモリが8GBしかないため、精度を落としたモデルを利用する必要がある。メモリが潤沢な場合はfp32のデータを利用しても良い。

from diffusers import StableDiffusionPipeline
from torch import autocast
import torch

YOUR_TOKEN="<TOKEN_HERE>"

pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16, use_auth_token=YOUR_TOKEN)
pipe.to("cuda")

画像生成

prompt = "a photo of an astronaut riding a horse on mars"
with autocast("cuda"):
    image = pipe(prompt)["sample"][0]
    image.show()

Enjoy!

Discussion

ログインするとコメントできます