AWS LambdaでStableDiffusionを動かしたい

(更新) img2img対応版の記事をアップしました。

2022/9/21 別記事をアップしました。主な内容はimgimgへの対応です。以降は↓を参照ください。

はじめに

StableDiffusionが衝撃的です。
使用するにはGPUが必要という要件があったためAWS AWS SageMakerで動かしていたのですが、色々と試すのにGPU付きのインスタンスは費用が若干気になるところ。
ですが先日、CPUのみで動作する「stable_diffusion.openvino」が公開されました。
動作要件を見ると「AWS Lambdaでも動きそう？」→「動いた！」ので、一先ずここに共有します。

本記事は以下に示す「stable_diffusion.openvino」をAWS Lambdaに移植したものです。

作成した画像サンプル

「墨田川」と「桜」のキーワードで作成した画像です。
AWS Lambda上で作成しましたが、1枚作成するのに4分前後かかっています。

費用に関する注意事項
本記事ではLambda関数に設定可能なメモリ上限の割り当てを行うなど費用が高額になる可能性のある設定を行っています。
内容を検証する場合はご注意ください。

AWS Lambdaで動かす場合のポイント

AWS Lambdaへ動かす際に修正が必要となった主なポイントを示します。
AWS Lambdaの設定値に関しては適宜変更してください。

• 画像を生成するための文字列(prompt)などのパラメータはLambda呼び出し時の値を使用するように修正
• OpenVINO 2022.1.0と必要なライブラリを個別にインストール
• 必要に応じてキャッシュディレクトリを/tmpに指定
• 生成した画像はS3へ格納
• AWS Lambdaの設定値の変更→メモリ：最大値の10240MBを指定、エフェメラルストレージ：5120MB、タイムアウト：5分

開発環境など

• AWS Cloud9を使用
• ストレージはCloud9のデフォルトのサイズでは不足するため、50 GiBに拡張

本記事ではコンテナ形式のAWS Lambdaで動かしています。
デプロイやテストなどの手順については本記事には記載していないため(本記事内ではコードのみ記載)、作業の流れなどを知りたい方は以下の記事を参照してください。

ファイル構成と概要

最終的には以下のような構成になりました。

ファイル構成

├ lambda-Stable-Diffusion/　・・・作業用のプロジェクトディレクトリ。任意の名称でOKです。
│　├ app.py　・・・StableDiffusionに関するメイン処理を実装するPythonファイルです。
│　├ Dockerfile　・・・コンテナイメージを作成するためのファイルです。
|  ├ glibc-2.27/　・・・GLIBC 2.27 を求められるため、ダウンロードして展開した状態です。
|  ├ openvino-2022.repo　・・・OpenVINO 2022.1.0をインストールするための設定ファイルです。
|  ├ requirements.txt　・・・パッケージリストです。

それでは、上から順番に紹介していきます。

app.py

• StableDiffusionに関するメインの処理を、および、Lambda関数(ハンドラ)を実装しています。
• BUCKET_NAMEに生成した生成した画像を保存するS3のバケット名を設定(バケットがなければ新規に作成)してください。
• 記事執筆時点(2022/9/12)で動作確認をしていますが、もしエラーが出た場合は、OpenVINOのバージョンとインストールディレクトリ、hf_hub_downloadに指定している各filenameの値を確認してみてください。

app.py

# -- coding: utf-8 --`
import argparse
import os
import inspect
import numpy as np

# openvino
import sys
sys.path.append('/opt/intel/openvino_2022.1.0.643/python/python3.9/')
from openvino.runtime import Core

# tokenizer
from transformers import CLIPTokenizer
# scheduler
from diffusers import LMSDiscreteScheduler
# utils
from tqdm import tqdm
import cv2
from huggingface_hub import hf_hub_download

from datetime import datetime
import boto3
s3 = boto3.resource('s3')
# 生成した画像ファイルを格納するバケット名を設定します
BUCKET_NAME = 'YOUR BUCKET NAME'

class StableDiffusion:
    def __init__(
            self,
            scheduler,
            model="bes-dev/stable-diffusion-v1-4-openvino",
            tokenizer="openai/clip-vit-large-patch14",
            device="CPU"
    ):
        self.tokenizer = CLIPTokenizer.from_pretrained(tokenizer, cache_dir='/tmp')
        self.scheduler = scheduler
        # models
        self.core = Core()
        # text features
        self._text_encoder = self.core.read_model(
            hf_hub_download(repo_id=model, filename="text_encoder.xml", cache_dir='/tmp'),
            hf_hub_download(repo_id=model, filename="text_encoder.bin", cache_dir='/tmp')
        )
        self.text_encoder = self.core.compile_model(self._text_encoder, device)
        # diffusion
        self._unet = self.core.read_model(
            hf_hub_download(repo_id=model, filename="unet.xml", cache_dir='/tmp'),
            hf_hub_download(repo_id=model, filename="unet.bin", cache_dir='/tmp')
        )
        self.unet = self.core.compile_model(self._unet, device)
        self.latent_shape = tuple(self._unet.inputs[0].shape)[1:]
        # decoder
        self._vae = self.core.read_model(
            hf_hub_download(repo_id=model, filename='vae_decoder.xml', cache_dir='/tmp'),
            hf_hub_download(repo_id=model, filename='vae_decoder.bin', cache_dir='/tmp')
        )
        self.vae = self.core.compile_model(self._vae, device)

    def __call__(self, prompt, num_inference_steps = 32, guidance_scale = 7.5, eta = 0.0):
        result = lambda var: next(iter(var.values()))

        # extract condition
        tokens = self.tokenizer(
            prompt,
            padding="max_length",
            max_length=self.tokenizer.model_max_length,
            truncation=True
        ).input_ids
        text_embeddings = result(
            self.text_encoder.infer_new_request({"tokens": np.array([tokens])})
        )

        # do classifier free guidance
        if guidance_scale > 1.0:
            tokens_uncond = self.tokenizer(
                "",
                padding="max_length",
                max_length=self.tokenizer.model_max_length,
                truncation=True
            ).input_ids
            uncond_embeddings = result(
                self.text_encoder.infer_new_request({"tokens": np.array([tokens_uncond])})
            )
            text_embeddings = np.concatenate((uncond_embeddings, text_embeddings), axis=0)

        # make noise
        latents = np.random.randn(*self.latent_shape)

        # set timesteps
        accepts_offset = "offset" in set(inspect.signature(self.scheduler.set_timesteps).parameters.keys())
        extra_set_kwargs = {}
        if accepts_offset:
            extra_set_kwargs["offset"] = 1

        self.scheduler.set_timesteps(num_inference_steps, **extra_set_kwargs)

        # if we use LMSDiscreteScheduler, let's make sure latents are mulitplied by sigmas
        if isinstance(self.scheduler, LMSDiscreteScheduler):
            latents = latents * self.scheduler.sigmas[0]

        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
        # and should be between [0, 1]
        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
        extra_step_kwargs = {}
        if accepts_eta:
            extra_step_kwargs["eta"] = eta

        for i, t in tqdm(enumerate(self.scheduler.timesteps)):
            # expand the latents if we are doing classifier free guidance
            latent_model_input = np.stack([latents, latents], 0) if guidance_scale > 1.0 else latents
            if isinstance(self.scheduler, LMSDiscreteScheduler):
                sigma = self.scheduler.sigmas[i]
                latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5)

            # predict the noise residual
            noise_pred = result(self.unet.infer_new_request({
                "latent_model_input": latent_model_input,
                "t": t,
                "encoder_hidden_states": text_embeddings
            }))

            # perform guidance
            if guidance_scale > 1.0:
                noise_pred = noise_pred[0] + guidance_scale * (noise_pred[1] - noise_pred[0])

            # compute the previous noisy sample x_t -> x_t-1
            if isinstance(self.scheduler, LMSDiscreteScheduler):
                latents = self.scheduler.step(noise_pred, i, latents, **extra_step_kwargs)["prev_sample"]
            else:
                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs)["prev_sample"]

        image = result(self.vae.infer_new_request({
            "latents": np.expand_dims(latents, 0)
        }))

        # convert tensor to opencv's image format
        image = (image / 2 + 0.5).clip(0, 1)
        image = (image[0].transpose(1, 2, 0)[:, :, ::-1] * 255).astype(np.uint8)
        
        return image

def handler(event, context):
    if event.setdefault('seed', None) is not None:
        np.random.seed(event['seed'])
    scheduler = LMSDiscreteScheduler(
        beta_start=event.setdefault('beta_start', 0.00085),
        beta_end=event.setdefault('beta_end', 0.012),
        beta_schedule=event.setdefault('beta_schedule', 'scaled_linear'),
        tensor_format='np'
    )
    stable_diffusion = StableDiffusion(
        model = event.setdefault('model', 'bes-dev/stable-diffusion-v1-4-openvino'),
        scheduler = scheduler,
        tokenizer = event.setdefault('tokenizer', 'openai/clip-vit-large-patch14')
    )
    image = stable_diffusion(
        prompt = event.setdefault('prompt', 'Cherry blossoms in Tokyo, Sumida River, blue sky, cartoon'),
        num_inference_steps = event.setdefault('num_inference_steps', 32),
        guidance_scale = event.setdefault('guidance_scale', 7.5),
        eta = event.setdefault('eta', 0.0)
    )
    
    output_img = event.setdefault('output', 'sd_' + datetime.now().strftime('%Y-%m-%d-%H-%M-%S') + '.png') 
    cv2.imwrite('/tmp/' + output_img , image)

    bucket = s3.Bucket(BUCKET_NAME)
    bucket.upload_file('/tmp/' + output_img, output_img)

Dockerfile

• OpenVINO 2022.1.0とGLIBC 2.27はpipでインストールできないため個別にインストールしています。

Dockerfile

FROM public.ecr.aws/lambda/python:3.9

COPY app.py ./
COPY requirements.txt  ./
RUN pip install -r requirements.txt

RUN yum -y update
COPY openvino-2022.repo /etc/yum.repos.d
RUN yum -y install openvino-2022.1.0
RUN cp -pa /opt/intel/openvino_2022.1.0.643/runtime/lib/intel64/* /var/task/
RUN cp -pa /opt/intel/openvino_2022.1.0.643/runtime/3rdparty/tbb/lib/* /var/task/

RUN yum -y install gcc make gcc-c++ zlib-devel bison bison-devel gzip glibc-static mesa-libGL-devel
COPY glibc-2.27/ /opt/glibc-2.27/
WORKDIR /opt/glibc-2.27/build
RUN /opt/glibc-2.27/configure --prefix=/var/task
RUN make
RUN make install
RUN cp -pa /var/task/lib/libm.so.6 /lib64/

CMD [ "app.handler" ]

glibc-2.27

• GLIBC 2.27が必要なライブラリがあるため(ない場合、"version GLIBC 2.27 Not found"エラーとなる)。

CMD

$ wget https://ftp.gnu.org/gnu/glibc/glibc-2.27.tar.gz
$ tar zxvf glibc-2.27.tar.gz

openvino-2022.repo

• OpenVINO 2022インストール用のyumリポジトリを追加するための設定ファイルです。

openvino-2022.repo

[OpenVINO]
name=Intel(R) Distribution of OpenVINO 2022
baseurl=https://yum.repos.intel.com/openvino/2022
enabled=1
gpgcheck=1
repo_gpgcheck=1
gpgkey=https://yum.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB

requirements.txt

• 「stable_diffusion.openvino」のオリジナルの"requirements.txt"ではここにopenvinoが記述されています。

requirements.txt

numpy==1.19.5
opencv-python==4.5.5.64
transformers==4.16.2
diffusers==0.2.4
tqdm==4.64.0
huggingface_hub==0.9.0
scipy==1.9.0

コード等の準備が終わったら

上述のコードや設定ファイルなどの準備が終わった後の手順の概略です。

• Amazon ECRコンソールでリポジトリを用意
• Cloud9に戻り、docker buildコマンドでビルドし、用意したECRリポジトリへpush
  (このビルドからpushについてはECRでリポジトリを作成すると実行すべきコマンドが表示されるようになるため、その手順に従ってください)
• AWS Lambdaコンソールで関数を作成し、ビルドしたイメージの指定、および、設定値を変更
• テストを実施(promptなどは以下のように指定します)

イベントJSON

{
  "prompt": "Cherry blossoms in Tokyo, Sumida River, blue sky"
}

以上です。

おわりに

AWSでStableDiffusionを試してみたいという場合は簡単で速いためSageMakerがよいと思います。
一方でCPUのみで使えるという選択肢は用途が広がりそうな期待感があります。
現状ではまだ「stable_diffusion.openvino」をAWS Lambdaに移植しただけのため色々と改善点のあるコードですが、少しでもどなたかの参考になれば幸いです。