Amazon Nova Canvasを使った画像生成チャットアプリを作ってみた #AWS – Qiita

import base64
import json
import os
import random
import boto3
import streamlit as st

REGION = "us-east-1"
IMG_MODEL_ID = "amazon.nova-canvas-v1:0"
TXT_MODEL_ID = "anthropic.claude-3-haiku-20240307-v1:0"

@st.cache_resource
def get_bedrock_client():
    return boto3.client(service_name="bedrock-runtime", region_name=REGION)

def generate_image(native_message, image_size, image_num, system_prompt, ng_text):
    """画像生成を行う
    """
    message = system_prompt + native_message
    print(f'textToImageParams: {message}')
    
    seed = random.randint(0, 858993460)
    native_request = {
        "taskType": "TEXT_IMAGE",
        "textToImageParams": {
            "text": message,
            "negativeText": ng_text
        },
        "imageGenerationConfig": {
            "seed": seed,
            "quality": "standard",
            "height": image_size,
            "width": image_size,
            "cfgScale": 10,
            "numberOfImages": image_num,
        },
    }
    request = json.dumps(native_request)
    
    bedrock_client = get_bedrock_client()
    response = bedrock_client.invoke_model(modelId=IMG_MODEL_ID, body=request)
    model_response = json.loads(response["body"].read())
    
    image_path_list = []
    for base64_image_data in model_response["images"]:

        # ローカルフォルダに画像イメージを保存
        i, output_dir = 1, "output"
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        while os.path.exists(os.path.join(output_dir, f"nova_canvas_{i}.png")):
            i += 1
        
        image_data = base64.b64decode(base64_image_data)        
        image_path = os.path.join(output_dir, f"nova_canvas_{i}.png")
        with open(image_path, "wb") as file:
            file.write(image_data)
        
        print(f"The generated image has been saved to {image_path}")
        image_path_list.append(image_path)
        
    return image_path_list


def textract_from_input(message):
    """チャットで指示された日本を英語に翻訳し、単語の切り出しを行う
    本処理はNova Canvasではなく、HaikuやSonnetなどの文章処理の基盤モデルを用いる
    """
    bedrock_client = get_bedrock_client()
    system_prompt = "あなたは日本語から英語への翻訳をし、"\
    "翻訳した文の中で重要な意味を持つ英単語を10個未満で抽出してください。"\
    "英単語は意味のある3単語未満でカンマで区切ってください。"\
    "抽出した英単語は、カンマ区切りで出力してください。"\
    "抽出する単語は前置詞や冠詞は含んではいけません。"\
    "かわいい猫と女の子が楽しく遊んでいる"\
    "cute cat, girl, play, joyful"
    
    message = {
        "role": "user",
        "content": [{"text": message}]
    }
    messages = [message]
    system_prompts = [{"text" : system_prompt}]
    
    inference_config = {"temperature": 0.5}
    additional_model_fields = {"top_k": 200}
    
    response = bedrock_client.converse(
        modelId=TXT_MODEL_ID,
        messages=messages,
        system=system_prompts,
        inferenceConfig=inference_config,
        additionalModelRequestFields=additional_model_fields
    )
    words = response["output"]["message"]["content"][0]["text"]
    return words


def display_history(messages):
    """チャットの履歴表示
    """
    for message in messages:
        display_img_content(message)


def display_img_content(message):
    """メッセージと画像を表示
    Novaで複数画像を生成した場合にも対応
    """
    contents = message["content"]
    print(f'message内容: {contents}')
    with st.chat_message(message["role"]):
        for content in contents:
            if content.get('text', None) != None:
                st.write(content["text"])
            else:
                st.image(content["image"])


def sidebar():
    """サイドバーを表示
    """
    with st.sidebar:
        st.sidebar.title("画像設定")
        
        # 画像サイズ
        image_size = st.sidebar.slider(
            "画像サイズ",
            min_value=320,
            max_value=960,
            step=64,
            value=320
        )
        
        # 画像数
        image_number = st.sidebar.slider(
            "画像数",
            min_value=1,
            max_value=5,
            step=1,
            value=1
        )
    
        # システムプロンプト
        system_prompt = st.sidebar.text_area(
            "システムプロンプト",
            value="minimalist icon, simple, flat design, dual-line design, 1.5px stroke weight line, "\
            "solid background, monochromatic, 2D, ",
            height=200
        )
        
        # 負のプロンプト
        negative_text = st.sidebar.text_area(
            "ネガティブプロンプト",
            value="3D, color, photo, multiple images",
            height=100
        )
        
        return image_size, image_number, system_prompt, negative_text


def main():
    """メイン処理
    """
    st.title("Simple Icon Generator by Nova Canvas")
    img_size, img_num, system_prompt, negative_text = sidebar()

    if "messages" not in st.session_state:
        st.session_state.messages = []

    display_history(st.session_state.messages)

    if prompt := st.chat_input("What's up?"):
        input_msg = {"role": "user", "content": [{"text": prompt}]}
        display_img_content(input_msg)
        st.session_state.messages.append(input_msg)
        print(f'st.session_state.messages: {st.session_state.messages}')
        
        # 入力された内容の拡張
        print(f'input_msg: {input_msg}')
        textract_resp = textract_from_input(input_msg["content"][0]["text"])
        print(f'texttract: {textract_resp}')

        # 画像生成
        resp_img_path_list = generate_image(
            textract_resp, img_size, img_num, system_prompt, negative_text
        )
        
        resp_img_contents = []
        for resp_img_path in resp_img_path_list:
            resp_img_contents.append({"text": f"Image created! filepath: {resp_img_path}"})
            resp_img_contents.append({"image": resp_img_path})
        resp_msg = {
            "role": "assistant",
            "content": resp_img_contents
        }
        display_img_content(resp_msg)
        st.session_state.messages.append(resp_msg)


if __name__ == "__main__":
    main()