使用 Amazon Bedrock 和 Langchain Agent 实现交互式文生图

{
	"Version": "2012-10-17",
	"Statement": [
		{
			"Sid": "VisualEditor0",
			"Effect": "Allow",
			"Action": [
				"s3:*",
				"bedrock:*"
			],
			"Resource": "*"
		}
	]
}

#requirements.txt
langchain
jupyter
boto3
Pillow
matplotlib
python-dotenv

#.env
profile=YOUR_AWS_PROFILE
bucket=YOUR_S3_BUCKET_TO_STORE_IMAGE

#python 3.10 环境
conda create -n py310 python=3.10 -y
conda activate py310
pip install -r requirements.txt
conda install -c conda-forge jupyterlab
jupyter lab

import base64
import io
import re
from PIL import Image
from langchain.chat_models import BedrockChat
from langchain.chains import LLMChain
from langchain.agents import initialize_agent
from langchain.agents import AgentType
from langchain.tools import tool
from langchain.prompts import PromptTemplate
import os
import boto3
import json
import matplotlib.pyplot as plt
from langchain.memory import ConversationBufferMemory
from dotenv import load_dotenv

# Load environment varibles including your aws profile and the s3 bucket to use. See https://docs.thinkwithwp.com/cli/latest/userguide/cli-configure-files.html
load_dotenv()

# Initialize the bedrock and s3 client
session = boto3.Session(profile_name=os.environ.get('profile'))

bedrock_client = session.client('bedrock-runtime')
s3_client = session.client('s3')

# Define the chat model. You can choose from anthropic.claude-instant-v1, anthropic.claude-v2 or other text models supported by Bedrock. We will use the chat model in both the prompt generation and the agent.
model_id="anthropic.claude-instant-v1"
# model_id="anthropic.claude-v2"
chat_model = BedrockChat(
    client=bedrock_client,
    model_id=model_id,
    model_kwargs = {"max_tokens_to_sample": 8000}
)

# Define the prompt generation tool
@tool(return_direct=True)
def generate_prompt_api(query: str) -> str:
    "Useful for generating a detailed prompt describing the scene."
    "Input: User's query."
    "Output: Stable diffusion prompts used to generate the image"
    prompt_template = """Give me one good and detailed prompt to generate the image using Stable Diffusion for {query}. The result should just include the prompt. No explanation. Example: <example>H: Give me one good and detailed prompt to generate the image using Stable Diffusion for an animal swimming in a lake. \n\n A: A photo of a brown bear swimming through a large, clear blue mountain lake surrounded by tall evergreen trees and snow capped mountains in the background. The bear is mid-stroke, with its head above the water and its front legs extended. Sparkling sunlight reflects off the smooth water. The image is sharply focused and highly detailed.</example>"""
    llm_chain = LLMChain(
            llm=chat_model,
            prompt=PromptTemplate.from_template(prompt_template),

        )

    result = llm_chain.predict(query=query)

    # The Claude model is a bit chatty, it might start the response with "Here is the prompt:". We need to strip that off.
    # print(result)
    striped = re.sub(r'^.*?:', '', result).strip()
    # print(striped)
    return striped

# Define the image generation tool. 
@tool(return_direct=True)
def text_to_image_api(query: str) -> str:
    "Useful for when you need to generate an image with a prompt."
    "Input: A detailed text-2-image prompt describing an image"
    "Output: Image url"
    # generate random integer values
    from random import randint
    body = json.dumps({
        "text_prompts": [
            { 
            "text": query 
            }
        ],
        "cfg_scale":10,
        "seed":randint(0, 100000),
        "steps":40,
        })
    # Here we use stable diffusion xl as our model
    modelId = "stability.stable-diffusion-xl-v0" 
    accept = "application/json"
    contentType = "application/json"

    response = bedrock_client.invoke_model(
        body=body, modelId=modelId, accept=accept, contentType=contentType
    )

    # Parsing the response (in base64) and plot the image
    response_body = json.loads(response.get("body").read())
    base_64_img_str = response_body["artifacts"][0].get("base64")
    img_bytes = base64.decodebytes(bytes(base_64_img_str, "utf-8"))
    img = Image.open(io.BytesIO(img_bytes))
    plt.imshow(img)
    plt.title(query[:80])

    # Save the image to s3 for later retrieval. 
    bucket=os.environ.get('bucket')
    key=f'img/{base_64_img_str[0:80]}.png'
    s3_client.put_object(Body=img_bytes, Bucket=bucket, Key=key)

    # Generate a presigned url for temporary public access
    generated_url = s3_client.generate_presigned_url(
        ClientMethod='get_object',
        Params={
            'Bucket': bucket,
            'Key': key
        },
        ExpiresIn=3600 # one hour in seconds, increase if needed
    )

    return generated_url

# Define a memory for the langchain agent so that it knows the whole context. The ai_prefix="Assistant" is needed as Claude uses this prefix for model generated content.
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True, ai_prefix="Assistant")

# Define the tools can be used by the agent
tools = [text_to_image_api, generate_prompt_api]

# Define the mrkl agent
mrkl = initialize_agent(
    tools, chat_model, agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION, verbose=True, max_iterations=5, memory=memory
)

# Let's see the prompt template it will use. You can also customize the prompt template if needed. See https://github.com/langchain-ai/langchain/issues/10721
print(mrkl.agent.llm_chain.prompt.messages)

output = mrkl.run("Generate some good prompt for an image with a cat playing football")

output = mrkl.run("Can you generate some better prompt?")

output = mrkl.run("Can you make the image more cyberpunk?")

output = mrkl.run("There isn't any football, plese regenerate prompt that the cat is playing football in the image")

output = mrkl.run("Fix the last prompt. the cat's color is black")

output = mrkl.run("生成水牛玩水的画的英文提示词")

output = mrkl.run("提示词中把水牛改成河马，其他不变")

output = mrkl.run("修改提示词，更加的油画风格")

output = mrkl.run("再加一些印象派风格")

output = mrkl.run("再加一些莫奈的风格")