boto3 でAmazon Bedrock経由でStable Diffusionを叩いて、画像を保存する

何の記事？

• 2023/09/28にGAになったAmazon Bedrockに遅れて、boto3(python のAWS SDK)が対応した
  ので、boto3経由で叩く
  の続き
• 記事タイトルの通り

IAM Policy

• RecouseのARNはAPIサンプルのモデルIDがわかればOK
• モデルIDは、モデルのページ(例えばここ)に記載されている、APIRequestのmodelIdと一致する

policy.json

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "VisualEditor0",
            "Effect": "Allow",
            "Action": [
                "bedrock:InvokeModelWithResponseStream",
                "bedrock:InvokeModel"
            ],
            "Resource": [
                "arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-v2",
                "arn:aws:bedrock:us-east-1::foundation-model/stability.stable-diffusion-xl-v0"
            ]
        }
    ]
}

実行環境

前記事を参照いただきたい

コード

以下、公式サンプルを参考にpromptを渡す部分と画像を保存する部分を追加
https://docs.aws.amazon.com/bedrock/latest/userguide/api-methods-run-inference.html

invoke_stablediffusion.py

import boto3
import json
import base64
import sys
from datetime import datetime

sample_prompt = "A photograph of an dog on the top of a mountain covered in snow."
prompt_data = sys.argv[1] if len(sys.argv) > 0 else sample_prompt

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

body = json.dumps({
  "text_prompts": [
    { 
      "text": prompt_data 
    }
  ],
  "cfg_scale":10,
  "seed":20,
  "steps":50
})
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())
print(response_body['result'])

image_base64 = response_body.get("artifacts")[0].get("base64")

image_binary = base64.b64decode(image_base64)

file_name = "image{0}.png".format(datetime.now().strftime("%Y%m%d%H%M%S"))
print("saving file: {0}".format(file_name))
with open(file_name, "wb") as image_file:
    image_file.write(image_binary)
print("image is saved.")

実行結果

root@196c6f2f6568:~# python invoke_stablediffusion.py 'geek, monkey with glass is typing keyboards of his laptop.'
success
saving file: image20230930180344.png
image is saved.

保存された画像はこんな感じ