Amazon BedrockのAPIとGradioで簡単なUIを作って画像生成してみた

はじめに

こんにちは。NTTドコモサービスイノベーション部の中村圭佑です。普段の業務では画像認識に関する研究開発を行っています。具体的には画像生成AI、農業AI等の研究開発に取り組んでいます。今回はAWSの生成系AIサービスである「Amazon Bedrock」と機械学習モデルのデモを行うWebアプリを作ることができるpythonのライブラリ「Gradio」を使ってUI上で画像生成してみたので紹介しようと思います。

Amazon Bedrockとは？

Amazon BedrockはAnthropic、Cohere、Meta、Stability AI、Amazon などを始めとした企業が提供する基盤モデル(Foundation model)を単一のAPIで動かすことができるフルマネージド型サービスです。Bedrockは現在(※2023年11月15日)General Availabilityの状態ですが、画像生成モデルであるStable DiffusionはPublic Previewとなっております。そのため、us-east-1とus-west-2のリージョンのみ利用可能となっています。モデルは多岐に渡り、画像生成、対話型のLLMによる検索拡張生成・テキスト生成・要約・コード生成など様々なユースケースに適用できるサービスとなっています。

Gradioとは？

Gradioは、機械学習のデモをWebアプリケーション上で行うためのオープンソースのPythonライブラリです。Gradioの利点としては、簡単かつ短時間でオレンジ基調のわかりやすいユーザーインターフェースを作成できることです。有名なstable-diffusion-webuiでも使われており、画像生成界隈では馴染み深いUIとなっています。

事前準備

動かすにあたってはswitchroleや最小権限を持つIAMユーザー、一時的なクレデンシャルを駆使しましょう。

今回使うモデルはPublic Preview中のためus-east-1かus-west-2のリージョンを指定して、BedrockのページからModel accessに飛んで使うモデルにアクセスリクエストを飛ばしましょう。画像生成モデルであるStable Diffusion XLを有効化します。右上のEditを押してモデルを選択後Save changesを押せば完了です。数分でAccess grantedとなります。これをやらないとモデルを呼び出す権限がないのでAPI叩いてもエラーが出ます。Stable Diffusionについては私が解説記事を過去に掲載していますので、ご興味ある方はぜひ確認してみてください。

APIを叩く時にBoto3(AWS SDK for Python)を使うので実行環境にインストールしていきます。Bedrockが対応しているのはboto3のバージョン1.28.57以上となっています。

pip install boto3

今回はawscliを使ってクレデンシャルを設定するのでインストールします。

pip install awscli

aws configureコマンドを実行して、クレデンシャルを設定します。

aws configure

これにより、以下の情報を順番に入力するプロンプトが表示されます。 - AWS Access Key ID
- AWS Secret Access Key
- Default region name
- Default output format

Default output formatはjsonにしています。入力した情報は、~/.aws/credentialsと~/.aws/configに保存されます。

注意: クレデンシャル情報は機密情報なので、不要な共有や公開を避けるよう注意してください。また、可能であればIAMロールや一時的なセキュリティクレデンシャルを使用することをおすすめします。

プログラム

公式のモデル呼び出しサンプルを参考に入力と出力を設定して作っていきます。
入力は
- text_prompts(text): プロンプトです。ここで入力する内容に沿って画像が生成されます。
- cfg_scale: 生成する画像にプロンプトの影響をどの程度反映するか決める値です。
- seed: 初期値であるランダムノイズ画像を生み出すための値です。seed値を変えると初期ノイズも異なるため、同じプロンプトでも異なる画像が出力されます。
- steps: 反復回数=デノイジング回数です。

出力は画像です。

以上の入力と出力の表示をGradioで行えるようプログラムを書きます。

#ライブラリ
import gradio as gr
import json
import requests
import io
import base64
from PIL import Image
import boto3
import sys
from datetime import datetime
import numpy as np

bedrock = boto3.client(service_name='bedrock-runtime', region_name="us-east-1")

# プロンプト入力、画像生成
def GenImg(prompt_data, steps, seed, cfg_scale):
    #payload
    body = json.dumps({
        "text_prompts": [
        { 
            "text": prompt_data
        }
    ],
        "cfg_scale":int(cfg_scale),
        "seed":int(seed),
        "steps":int(steps)
    })

    # APIに投げて画像生成
    modelId = "stability.stable-diffusion-xl-v0" 
    accept = "application/json"
    contentType = "application/json"

    response = bedrock.invoke_model(
        body=body, modelId=modelId, accept=accept, contentType=contentType
    )
    response_body = json.loads(response.get("body").read())
    image_base64 = response_body.get("artifacts")[0].get("base64")
    output_img = Image.open(io.BytesIO(base64.b64decode(image_base64)))
    return output_img


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=1, min_width=600):
            prompt_data = gr.Textbox(label="prompt_data")
            steps = gr.Textbox(label="Steps", value="20")
            cfg_scale =  gr.Textbox(label="cfg_scale", value="10")
        with gr.Column(scale=2, min_width=300):
            text_button = gr.Button("Generate image")
            output_img = gr.Image()
    with gr.Row(scale=1, min_width=600): 
            seed = gr.Textbox(label="seed", value="1")       
            text_button.click(GenImg, inputs=[prompt_data, steps, seed,  cfg_scale], outputs=output_img)

# 起動
demo.launch()

このコードを実行してweb UIを起動して画像生成してみましょう。デフォルトではローカルコンピュータの http://127.0.0.1:7860 でlistenしているので、このURLに接続します。プロンプトを打ってGenerate imageボタンを押すとBedrockのAPIが裏側で叩かれてweb ui上で画像が表示されます。

おわりに

これでBedrockとGradioを連携させることができました。Amazon Bedrockは画像生成の他にも自然言語モデルの提供も行っています。これらのAPIを使い基盤モデルを利用&運用することで様々な社会課題の解決につながるかもしれません。Diffusionモデルをベースとしたtxt2imgの画像生成に関しても、発展が目覚ましいので将来が非常に楽しみです。

またNTTドコモでは画像生成・物体検出・姿勢推定・一般物体認識・特定物体認識・文字認識・類似画像検索等、様々なAI技術の研究開発をしています。インターンや新卒・キャリア採用を積極的に実施しているため、気になる方は以下のリンクをご参考にしてください。