WSL2にUbuntu22.04とCUDAをインストールしてJupyterLab上でStable Diffusionを動かしたメモ
はじめに
話題のStable Diffusionをローカル環境で動かすまでに行った作業のメモです。ついでにJupyterLabでも使えるようにしました。
作業環境
- OS:Windows 11(21H2 Build 22000.856)
- GPU:RTX3070
手順
基本的にはこちらに書いてある内容と一緒です。
WSL2を有効にする
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でエラーが出るので実行しておく。
# 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」から最新のドライバをインストールする。
ドライバーのインストール(Ubuntu)
ここの「Option 1: Installation of Linux x86 CUDA Toolkit using WSL-Ubuntu Package – Recommended」という手順を実行すれば必要なドライバーやライブラリはほぼ入る。
$ 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のアカウント作成が必要） 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
これにて導入完了。
動作確認
モデルデータへのアクセストークンの取得
こちらを参考に、事前にモデルデータにはアクセスできるようにして、ユーザーアクセストークンを取得しておく。
一度コマンドラインで
$ 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