AWS Lambdaによるサーバーレス画像⽣成AIの作成⽅法

2022年11月9日

森田和明

Kazuaki Morita

エリア事業本部西日本支社
インテグレーション＆ソリューション部
ＩＴアーキテクトグループ
主任 / フェロー

サーバーレス画像⽣成AIの構築に必要なLambdaの機能・アップデート
画像⽣成サーバーレスAPIの構築⼿順
まとめ
この記事の執筆者

AWS Lambdaは、サーバーの管理が不要で使⽤したミリ秒単位の時間課⾦で提供されるAWSのサーバーレスコンピューティングサービスです。

当初は負荷の軽いバッチ処理やAPIバックエンドを実現する環境として登場しましたが、数々の機能アップデートにより、画像⽣成AIを動作させることも可能となりました。AWS Lambdaの機能アップデートを振り返りながら、画像⽣成AIを動作させる⼿順をご紹介します。

※画像⽣成AIとしては、「Stable Diffusion」を使⽤します。AIモデルの利⽤にはライセンスへの同意が必要となります。詳細は公式サイトをご確認ください。

サーバーレス画像⽣成AIの構築に必要なLambdaの機能・アップデート

実⾏時間が5分から15分に延⻑

AWS Lambdaには実⾏時間に制限があり、当初は5分が上限で、上限に達した際は強制的に処理が終了されてしまいます。

https://aws.amazon.com/jp/about-aws/whats-new/2018/10/aws-lambda-supports-functions- that-can-run-up-to-15-minutes/

実⾏時間が15分間に延⻑されたことにより、画像⽣成など5分では終わらない処理もLambdaを使ってサーバーレスで実現できる可能性ができました。

バイナリサポート

バイナリサポートはAmazon API Gatewayのアップデートですが、API GatewayとAWS Lambdaを組み合わせてAPIを構成する場合に、バイナリデータを返却することが出来るようになりました。

https://aws.amazon.com/jp/about-aws/whats-new/2016/11/binary-data-now-supported-by-api- gateway/

Amazon API GatewayとAWS LambdaをLambda プロキシ統合で連携している場合は、Lambdaから以下のようなレスポンスを⽣成することで、バイナリデータを返却できます。

return {
    'headers': { "Content-Type": "image/png" },
    'statusCode': 200,
    'body': base64.b64encode(image).decode('utf-8'),
    'isBase64Encoded': True
}

メモリ上限が10GBに拡張

AWS Lambdaは設定により割り当てるメモリのサイズが指定できますが、この上限が現在は10GBまで拡張されています。

https://aws.amazon.com/jp/blogs/aws/new-for-aws-lambda-functions-with-up-to-10-gb-of- memory-and-6-vcpus/

画像⽣成などのAIの処理には多くのメモリが必要となりますので、メモリ上限の拡張も嬉しいアップデートでした。

コンテナイメージサポート

AWS Lambdaはプログラムのソースコードと依存関係のあるライブラリーをZIPファイルに纏めてアップロードしデプロイする形式となっています。ZIPファイルのサイズ上限が50MBのため、AIの動作に必要なライブラリーを使⽤することが現実的ではありませんでした。今回試した画像⽣成AIについてもライブラリーだけで1GB以上のサイズとなっており、Lambda上で動作させることができませんでしたが、ZIPファイルの代わりに10GBを上限としたコンテナイメージをデプロイすることが可能になり、AIの実⾏環境として利⽤出来るようになりました。

https://aws.amazon.com/jp/blogs/news/new-for-aws-lambda-container-image-support/

コンテナイメージがサポートされたことにより、デプロイサイズが増えたことに加え、AWS Lambdaで提供されていない開発⾔語であってもAWS Lambda上で動作させられるようになりました。

Amazon EFSサポート

Amazon EFSはAWSマネージドなファイルストレージサービスです。AWS Lambdaは基本的にファイルを永続化する機能は提供されておらず、必要な場合はAmazon S3と連携する必要がありました。AIフレームワークが Amazon S3から直接モデルをロードできないなど課題がありましたが、Amazon EFSがサポートされることで、AI⽤途に利⽤しやすくなりました。

https://aws.amazon.com/jp/about-aws/whats-new/2020/06/aws-lambda-support-for-amazon- elastic-file-system-now-generally-/

今回の検証ではコンテナイメージ内にAIのモデルデータを含める形としましたが、EFSを使うとモデルの更新時にLambdaの再デプロイが不要になったり、複数モデルを切り替えて利⽤するなど活⽤の幅が広がります。

Function URLs

個⼈的に2022年⼀番のアップデートと感じているのがこのFunction URLs（関数 URL）です。Function URLsは、API Gatewayを⽤意しなくてもAWS Lambda単体でHTTPSエンドポイントが提供されるというものです。API Gatewayが提供する豊富な機能は不要だが、HTTPSエンドポイントを実現したい場合に活⽤できます。

https://aws.amazon.com/jp/blogs/news/announcing-aws-lambda-function-urls-built-in-https- endpoints-for-single-function-microservices/

画像⽣成AIでの活⽤の観点として重要なのがAmazon API Gatewayの29秒制限の回避に使⽤できる点です。Amazon API Gatewayには統合タイムアウトがあり上限の引き上げも不可となっています。そのため、AWS Lambdaが15分間処理可能なのですが、HTTPSエンドポイントとしては29秒以内にレスポンスを返却する必要がありました。Function URLsを使⽤するとこの29秒の上限がなくなるため、15分以内にレスポンスを返却する APIを作成することが可能となりました。

画像⽣成サーバーレスAPIの構築⼿順

それでは、画像⽣成AIを構築する⼿順をご紹介します。

SAM プロジェクトの作成

AWS SAM(サーバーレスアプリケーションモデル)を使ってプロジェクトを作成します。

sam init --name "stable-diffusion-lambda" \
    --package-type Image \
    --base-image "amazon/python3.9-base" \
    --app-template "hello-world-lambda-image" \
    --tracing

これ以降はstable-diffusion-lambda内で⾏います。

画像⽣成AIのモデルを取得

Stable Diffusionの学習済みモデルを取得します。

ライブラリーのインストール

pip install --no-cache-dir diffusers transformers scipy torch --extra-index-url https://download.pytorch.org/whl/cpu

Hugging Faceへログイン

モデルの取得にはライセンスに同意しておく必要があります。Webサイトでライセンスに同意後、APIキーを⽣成します。

APIキーを使⽤してログインを⾏います。

huggingface-cli login

モデルの取得

以下のPythonスクリプトでモデルを取得します。

from diffusers import StableDiffusionPipeline

model_id = 'CompVis/stable-diffusion-v1-4'

model = StableDiffusionPipeline.from_pretrained(model_id)
model.save_pretrained('hello_world/model')

AWS Lambdaのソースを修正

app.py

Lambda関数のプログラムです。Lambdaの場合は/tmp領域のみ書き込みができるため、⽣成した画像は/tmpに出⼒しています。

出⼒した画像をBase64形式に変換しレスポンスとして返却します。

import base64
import json
import os

from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained('/model')
pipe = pipe.to('cpu')

def lambda_handler(event, context):

    width = int(os.environ['WIDTH'])
    height = int(os.environ['HEIGHT'])

    try:
        os.remove('/tmp/image.png')
    except:
        pass

    body = json.loads(event['body'])
    prompt = body['prompt']
    print(prompt)

    image = pipe(prompt, width=width, height=height, guidance_scale=7.5).images[0]
    image.save('/tmp/image.png')

    with open('/tmp/image.png', 'rb') as f:
        base64_img = base64.b64encode(f.read()).decode('utf-8')

    return {
        "statusCode": 200,
        "headers": { "Content-Type": "image/png" },
        "body": base64_img,
        "isBase64Encoded": True
    }

Dockerfile

コンテナとしてLambda関数を作成します。AIモデルもコンテナ内に含める形にしました。

⽣成されるコンテナイメージは約7GB程度で、デプロイ可能な10GB以内に収まりました。

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

RUN python3.9 -m pip install --no-cache-dir diffusers transformers scipy torch --extra-index-url https://download.pytorch.org/whl/cpu

RUN yum install -y <span style="font-family: arial;">\</span>
        mesa-libGL-devel <span style="font-family: arial;">\</span>
        mesa-libGLU-devel <span style="font-family: arial;">\</span>
        libpng-devel &amp;&amp; <span style="font-family: arial;">\</span>
    yum clean all

COPY ./model/ /model/
COPY app.py ./

# Command can be overwritten by providing a different command in the template directly.
CMD ["app.lambda_handler"]

template.yaml

タイムアウト時間を900秒（15分）、メモリサイズを10GBにセットします。また、Function URLsを有効化します。認証なしの形式と試しましたが、必要に応じて認証を有効化してください。

AWSTemplateFormatVersion: '2010-09-09'
Transform: AWS::Serverless-2016-10-31
Description: >
  python3.9

  Sample SAM Template for stable-diffusion-lambda

# More info about Globals: https://github.com/awslabs/serverless-application-model/blob/master/docs/globals.rst
Globals:
  Function:
    Timeout: 900
    MemorySize: 10240
    Tracing: Active

Resources:
  HelloWorldFunction:
    Type: AWS::Serverless::Function # More info about Function Resource: https://github.com/awslabs/serverless-application-model/blob/master/versions/2016-10-31.md#awsserverlessfunction
    Properties:
      PackageType: Image
      Architectures:
        - x86_64
      FunctionUrlConfig:
        AuthType: NONE
      Environment:
        Variables:
          WIDTH: 512
          HEIGHT: 512
    Metadata:
      Dockerfile: Dockerfile
      DockerContext: ./hello_world
      DockerTag: python3.9-v1

Outputs:
  HelloWorldFunctionUrlEndpoint:
      Description: "HelloWorldFunction URL Endpoint"
      Value:
        Fn::GetAtt: HelloWorldFunctionUrl.FunctionUrl

ビルドとデプロイ

プログラムが揃ったので、ビルドとデプロイを⾏います。

ビルド

sam build

テスト実⾏

ローカル環境でLambdaを呼び出します。

echo '{"body": "{<span style="font-family: arial;">\</span>"prompt<span style="font-family: arial;">\</span>": <span style="font-family: arial;">\</span>"a photo of an astronaut riding a horse on mars<span style="font-family: arial;">\</span>"}"}' | sam local invoke --event -

レスポンスはJSONで返却されます。

結果からがPNG画像を⽣成するにはこのようになります。

echo '{"body": "{<span style="font-family: arial;">\</span>"prompt<span style="font-family: arial;">\</span>": <span style="font-family: arial;">\</span>"a photo of an astronaut riding a horse on mars<span style="font-family: arial;">\</span>"}"}' |  sam local invoke --event - &gt; response.json
cat response.json | jq -r .body | base64 -d &gt; output.png

デプロイ

ローカル環境でうまく動作したらデプロイしましょう。

sam deploy --guided

CloudFormationのアウトプットとしてFunction URLsのエンドポイントが出⼒されます。

cURLによるAPI呼び出し

curl -X POST <span style="font-family: arial;">\</span>
-H 'Content-Type: application/json' <span style="font-family: arial;">\</span>
-d '{"prompt": "a photo of an astronaut riding a horse on mars"}' <span style="font-family: arial;">\</span>
[Function URLsのエンドポイント] <span style="font-family: arial;">\</span>
--output image.png

Lambdaで画像が⽣成できました。