G5インスタンスをTerraformでつくって、10分でStableDiffusionを動かす

想定している人

• ローカルマシンで実行したくない/できない
• Pythonを動かすための環境構築は最低限の手間で済ませたい

最初に

StableDiffusionという、文字から画像を自動生成してくれるやつが何やらすごいという話を聞いたので週末で動かしてみた。

ローカル環境に機械学習をする環境は作ってないし、確実に大量のディスク容量が必要になる。CUDAとかのインストールも大変そうだったので、EC2インスタンスで最初から設定してあるAMIを使うことに決めた。

事前準備: GPUが利用できるインスタンスのQuotaを引き上げ

今回、GPUを利用するということでG5インスタンスを利用することにする。CPUしかないT3インスタンスだと、おそらくCUDAが動かないはず。まだGPUを利用できるインスタンスが利用できない場合は、Quotaを引き上げる必要がある。

AWSコンソールのQuotaページで、Quotaを確認できる。デフォルトでは0なので、SpotInstanceを利用したいか、On-Demandを利用したいかに合わせ制限を解除依頼する。

• All G and VT Spot Instance Requests
• Running On-Demand G and VT instances

クォータ値にはvCPUの値が入る。つまり、g5.xlarge を利用したい場合は4、g5.2xlargeを利用したい場合は8を入れる。

引き上げ依頼を出して、承認されるまでは待つ。これは10分の中には入らない。

EC2インスタンスの起動

Terraformを利用する。
Terraformコードは、昨日書いた｢セキュアなEC2インスタンスをTerraformで作成する｣のコードを利用して作成する。

$ git clone https://github.com/tomoyamachi/terraform-private-ec2-ssm.git
$ cd terraform-private-ec2-ssm

# Terraformに渡す変数をtfvarsファイルで設定する
$ vim terraform.tfvars

# 対象リージョンの Deep Learning AMI GPU PyTorchのAMIを設定する
image_id="ami-00b021cce8954bcab"
name="ssmtf"
instance="g5.xlarge"
block_volume_size="64"
session_role="SessionManager経由でログインするためのRole"
spot_instance=true

インスタンスにログイン後

インスタンスにログインすると、だいたいの環境がすでに整っているので、git-lfsだけインストールする。
StableDiffusionの公開されているモデルをGit経由でダウンロードするのだが、10GB以上あるのでgit-lfsが必要になる。

自分の場合は、遊んだらすぐにインスタンスを消すので、rootユーザで作業した。

$ sudo su -

# git-lfsのインストール
$ curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh | sudo bash
$ apt install git-lfs

# StableDiffusionのClone
$ git clone https://github.com/CompVis/stable-diffusion

# StableDiffusionのモデルをClone。 https://huggingface.co/ のログインIDとパスワードが必要になるので、事前にユーザ登録しておく。
$ git clone https://huggingface.co/CompVis/stable-diffusion-v-1-4-original

ダウンロードに時間がかかるので、別のセッションを開き、condaで環境構築をする。

$ cd stable-diffusion
# 環境構築
$ conda env create -f environment.yaml
$ conda init bash
# condaの環境を適用させるために、shellから一度ログアウト→ログインしなおす
$ exit
$ sudo su -
$ conda activate ldm

モデルのダウンロードが完了したら、モデルファイルをStableDiffusionのディレクトリに移動する。

# Git LFSが正しくインストールされていない場合、sd-v1-4.ckptのファイルサイズが小さいので、ファイルサイズを確認
$ ls -l stable-diffusion-v-1-4-original/
total 11688756
-rw-r--r-- 1 root root      11716 Aug 28 02:06 README.md
-rw-r--r-- 1 root root 7703807346 Aug 28 02:12 sd-v1-4-full-ema.ckpt
-rw-r--r-- 1 root root 4265380512 Aug 28 02:10 sd-v1-4.ckpt
-rw-r--r-- 1 root root      71237 Aug 28 02:06 v1-variants-scores.jpg

# モデルファイルをStableDiffusionのディレクトリに移動
$ mv stable-diffusion-v-1-4-original/sd-v1-4.ckpt stable-diffusion/
# モデルのサイズが大きいので、必要なものを移動したらモデルのディレクトリを削除
$ rm -fr stable-diffusion-v-1-4-original

あとは conda activate ldmした環境で、コマンドを実行すれば、outputs以下に画像が生成されている。

$ python scripts/txt2img.py --prompt "a photograph of an astronaut riding a horse" --plms --ckpt sd-v1-4.ckpt --n_samples 1

後片付け

とりあえず動いたのでひととおり遊んだら、インスタンスなどを削除する。

$ terraform destroy

最後に

StableDiffusionってなんなのか、どうやって動いているのか、みたいな話はまったくわからなかった。
ただ好奇心を満たしただけだった。