Stable DiffusionをEC2で動かしてHTTPリクエストで画像生成させる

Introduction

最近盛り上がっている画像生成AIですが、
その中でもStable Diffusionはオープンソースであり、
好きな環境にインストールできて生成した画像もある程度自由に使えます。

今回はEC2でStable Diffusionを構築し、
ついでにHTTPリクエストで画像生成して返すサンプルをつくってみました。

Setup

EC２を使うのでAWSアカウントは必要です。
けっこう大きめのインスタンスを使うので、実際にやってみる人は止め忘れに注意。

必要なアカウント作成など

ではまず、Stable Diffusion構築の際に必要になるアカウントのの作成や
必要ファイルのダウンロードをします。

ここで、HuggingFaceのユーザー登録と、
NVIDIA DeveloperでNVIDIAのユーザー登録を行います。
それぞれ認証確認用メールがきているので、指示にしたがって認証しておきましょう。

次に、NVIDIA cuDNNライブラリのダウンロードを行います。
ここにあるダウンロード用画面を開き、
"I Agree To the Terms of the cuDNN Software License Agreement" にチェック、
"Download cuDNN v8.5.0 (August 8th, 2022), for CUDA 11.x" をクリックしてから
"Local Installer for Linux x86_64 (Tar)" をダウンロードします。
このファイルはあとで使います。

EC2インスタンスの起動

※基本的にここにある手順でやってます

では、AWSコンソールからEC2インスタンスを起動しましょう。
gとかpとかのデカめのインスタンスが必要らしいので、今回はp3を使ってます。
↓の設定でインスタンスを起動。

• インスタンスタイプ: p3.2xlarge
• AMI: "Amazon Linux" 　＞ "Deep Learning AMI GPU TensorFlow 2.9.1 (Amazon Linux 2) 20220830"
• ストレージ: 汎用SSD(GP3) 100GB

SSHと3000番ポート(Nodeアプリ用)をアクセス可能に設定。

SSHログインしてセットアップ

EC2インスタンスが起動したら、SSHログインします。

% ssh -i <KeyPairファイル> ec2-user@＜IPアドレス＞
　・・・

CUDA用GPUの存在を確認。

% lspci | grep -i nvidia
00:1e.0 3D controller: NVIDIA Corporation GV100GL [Tesla V100 SXM2 16GB] (rev a1)

必要なライブラリのダウンロード。

% sudo amazon-linux-extras install epel
% sudo yum install kernel-devel-$(uname -r) kernel-headers-$(uname -r)

CUDAリポジトリを追加してドライバをインストール。

% ARCH=$( /bin/arch )
% echo $ARCH
x86_64

% distribution=rhel7
% sudo yum-config-manager --add-repo http://developer.download.nvidia.com/compute/cuda/repos/$distribution/${ARCH}/cuda-$distribution.repo
% sudo yum install -y nvidia-driver-latest-dkms cuda cuda-drivers

インストールOK。

% nvidia-smi -q

==============NVSMI LOG==============

Timestamp                                 : Wed Sep  7 08:26:49 2022
Driver Version                            : 515.65.01
CUDA Version                              : 11.7

Attached GPUs                             : 1
GPU 00000000:00:1E.0
    Product Name                          : Tesla V100-SXM2-16GB
    Product Brand                         : Tesla
    Product Architecture                  : Volta

~/.bashrcに環境変数を設定しておきます。

export CUDA_HOME=/usr/local/cuda-11.0
export PATH=${CUDA_HOME}/bin${PATH:+:${PATH}}

CUDA Toolkitをインストールします。

% mkdir ~/install && cd ~/install
% wget https://developer.download.nvidia.com/compute/cuda/11.3.1/local_installers/cuda_11.3.1_465.19.01_linux.run
% sudo sh cuda_11.3.1_465.19.01_linux.run

画像のように、Acceptを選択し、Driverのチェックははずしてtoolkitをインストール。
途中でYES/NOとかきかれたらYESを選択しましょう。

さきほどダウンロードしたファイルを、ローカルからscpでEC2にアップロードします。

% scp -i <KeyPairファイル>　＜さっきダウンロードしたcudnn-linux-x86_64-8.5.0.96_cuda11-archive.tar.xz＞ ec2-user@＜IPアドレス＞:/home/ec2-user/install

tarを解凍して必要ファイルをコピーしましょう。

% cd ~/install
% tar Jxfv cudnn-linux-x86_64-8.5.0.96_cuda11-archive.tar.xz
% cd cudnn-linux-x86_64-8.5.0.96_cuda11-archive/

% sudo cp include/cudnn*.h /usr/local/cuda/include
% sudo cp -P lib/libcudnn* /usr/local/cuda/lib64
% sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*

Annacondaをインストール

Anacondaがはいってなかったのでインストールします。

% wget https://repo.anaconda.com/archive/Anaconda3-2022.05-Linux-x86_64.sh
% sh Anaconda3-2022.05-Linux-x86_64.sh
% exec $SHELL -l

Stable Diffusion本体のインストール

StableDiffusionをインストールします。
私の場合、けっこう時間かかりました。

% cd ~/
% git clone https://github.com/CompVis/stable-diffusion

# /home/ec2-user/conda/bin/conda
% conda init bash
% conda env create -f environment.yaml

次にモデルのダウンロードをするのですが、
その前にここの画面でAccess Repositoryしておかないと
git clone時に403になるので注意しましょう。　　

#HuggingFaceアカウント情報を入力
% cd ~/
% git clone https://huggingface.co/CompVis/stable-diffusion-v-1-4-original

データは数GBあるのでしばらく時間がかります。
そしてデータの移動。

% cd ~/stable-diffusion/models/ldm
% mkdir stable-diffusion-v1
% mv ../../../stable-diffusion-v-1-4-original/sd-v1-4.ckpt stable-diffusion-v1/model.ckpt

ここでシェルを再起動して、↓のように(base)となればOKです。

% exec $SHELL -l
(base) [ec2-user@ip-172-31-95-72 ldm]$

サンプルスクリプトで動作確認

実行準備ができたので、スクリプト（txt2img.py）を動かしてテキストから画像を生成してみます。

% conda activate ldm
% cd ~/stable-diffusion
% python3 scripts/txt2img.py --prompt "an apple on the desk" --plms

生成画像をみるためにscpでダウンロードします。

% zip archive -r ./outputs/txt2img-samples/

#ローカルからscp実行してダウンロード
% scp -i <KeyPairファイル>  ec2-user@＜IPアドレス＞:/home/ec2-user/stable-diffusion/archive.zip ~/path/my/images

画像を開くと、リンゴの画像が生成されていました。
しっかり動いていそうです。

Image generate by HTTP request

では次に、自分で記述したプログラムから画像生成をおこなってみます。
まずはプログラム実行のためにアクセストークンが必要なので、 ここで発行します。

そして、↓のようなpythonファイル(my_scripts.py)をscriptディレクトリに作成します。

# my_scripts.py
import sys
from torch import autocast
from diffusers import StableDiffusionPipeline

args = sys.argv

file_name =  args[1]
prompt = args[2]

TOKEN="<ACCESSS TOKEN>"

pipe = StableDiffusionPipeline.from_pretrained(
	"CompVis/stable-diffusion-v1-4", 
	use_auth_token=TOKEN
).to("cuda")

#prompt = "a photo of an astronaut riding a horse on mars"

with autocast("cuda"):
    image = pipe(prompt)["sample"][0]  

# 画像を保存(ディレクトリは作成しておく)    
image.save("/path/your/my_outdir/" + file_name)

TOKEN変数にさきほど発行したアクセストークンを設定し、StableDiffusionPipelineを使って
画像生成します。

my_scripts.pyを実行してみます。
さきほどと同じく、画像が生成されます。（保存しているのは１つだけ）

% python3 /path/your/my_scripts.py foo.png "a photo of an astronaut riding a horse on mars"

実行時に403エラーがでたら、huggingface-cliコマンドで
ログインしておく必要があるかもしれません。　　

% huggingface-cli login

Node.jsからStable Diffusionを実行

StableDiffusionPipelineをつかって画像生成ができたので、
次はHTTPリクエストをうけて画像生成をしてみます。

FlaskとかつかってStableDiffusionPipelineで画像生成すれば
シンプルにできそうな気がするのですが、あまり詳しくないので、
Node.jsからexecSyncでシェルスクリプトを実行するという強引な方法をやります。

では、nvmとNode(v16)をインストールします。

% curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.34.0/install.sh | bash
% nvm install v16.17.0

HTTPサーバはfastifyを使うのでnpmでインストール。

% cd ~/
% mkdir node-app && cd node-app
% npm install fastify --save

HTTPリクエストをうけとるjsファイルを作成。
3000番ポートでサーバが起動し、
POSTでファイル名とプロンプト（画像生成のための文字列）を受け取ります。

//server.js

const fastify = require('fastify')({ logger:{level:'info'}})
const fs = require('fs')
const { execSync } = require('child_process')

// Declare a route
fastify.post('/generate-image', async (request, reply) => {
  const file_name = request.body.file_name;
  const prompt = request.body.prompt;
  const exec_command = '/path/your/my_scripts.sh ' + file_name +  ' "' + prompt + '"';
  console.log("exec_command : " + exec_command);

  const result = execSync(exec_command);
  const buffer = fs.readFileSync('/path/your/my_outdir/' + file_name)
  reply.type('image/png')
  reply.send(buffer)
})

// Run the server!
const start = async () => {
  try {
    await fastify.listen({port: 3000, host: '0.0.0.0'})
  } catch (err) {
    fastify.log.error(err)
    process.exit(1)
  }
}
start()

server.jsから実行するシェルを作成。
pythonスクリプト実行前にいろいろ(conda activate ldmとか)
やりたいかもしれないからつくったけど、必要ない気がする。

#!/bin/sh

file_name=$1
prompt=$2
echo $file_name
echo $prompt

py_command="python3 /path/your/my_scripts.py ${file_name} \"${prompt}\""
eval $py_command

HTTPサーバを起動しましょう。

% cd ~/node-app
% node server.js
{"level":30,"time":1662619564856,"pid":5714,"hostname":"ip-xxx-xxx-xxx-xxx.ec2.internal","msg":"Server listening at http://0.0.0.0:3000"}

ローカルからcurlでHTTPリクエストをPOST投げます。
Bodyのpromptに画像生成する文字列、file_nameにファイル名を設定します。
outputオプションでローカル保存用のファイル名を指定して
レスポンスをそのままpngで保存します。

% curl -H "Content-Type:application/json" -X POST http://＜IPアドレス＞:3000/generate-image -d '{"file_name":"foo.png","prompt":"River and Cherry Trees"}' --output example.png

% open example.png

これで（やっつけだけど）「HTTPリクエストで画像生成」ができました。
しかし、同時に実行したらたぶん落ちるので注意。

Summary

今回はStable DiffusionのセットアップからHTTPリクエストで画像生成をやってみました。
セットアップできれば使い方はシンプルなので、
いろいろな使いかたができそうです。

References

• Stable DiffusionをAWSのEC2インスタンスで動かすまでの手順
• Stable Diffusion を Colab で Web アプリ化する
• Stable Diffusion API 開発
• EC2 G4インスタンスのAmazon Linux 2にNVIDIA CUDAをインストールしてみた