lambda_function_imageMake_and_storyMake.py

import base64
import boto3
import json
import os
import io

from PIL import Image
from botocore.exceptions import ClientError
import random
class ImageError(Exception):
    "Custom exception for errors returned by Amazon Titan Image Generator G1"

    def __init__(self, message):
        self.message = message





# AWS 서비스 클라이언트 설정
bedrock = boto3.client('bedrock-runtime', region_name='us-west-2')
s3 = boto3.client('s3')
# Claude 모델 ID
MODEL_ID = "anthropic.claude-3-haiku-20240307-v1:0"

# 사람의 이름
NAME = "신준혁"
hope = "소방관"
# 결과 파일 이름
def call_claude_haiku(base64_string):
    # 프롬프트 메시지에 NAME 변수를 삽입
    prompt = f"""이미지 속 인물의 직업은 {hope}입니다. 이미지 속 인물을 분석하고,가상의 인생 스토리를 만들어주세요.
    1. 특정 개인을 식별하는 것은 금지됩니다. 스토리의 주인공은 실존하지 않는 가상의 인물이어야 합니다.
    2. 스토리의 인물은 당신이라는 명칭으로 시작해야 합니다.
    3. 일대기 스토리의 글자는 150개로 제한해주세요.
    4. 이미지 속 인물, {NAME},의 그 직업에 대한 재미있는 일화를 만들어주세요.
    5. 마지막에 *를 적은 후에 이미지의 '헤어스타일', '성별','피부색'을 상세히 적어요.
    6. 인물의 어렸을 적 이야기는 들어가서는 안됩니다.
    7.  반드시 영어로 출력해야합니다.
    """
 
    prompt_config = {
        "anthropic_version": "bedrock-2023-05-31",
        "max_tokens": 1000,
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "source": {
                            "type": "base64",
                            "media_type": "image/jpeg",
                            "data": base64_string,
                        },
                    },
                    {"type": "text", "text": prompt},
                ],
            }
        ],
    }

    body = json.dumps(prompt_config)

    response = bedrock.invoke_model(
        body=body, modelId=MODEL_ID, accept="application/json", contentType="application/json"
    )
    response_body = json.loads(response.get("body").read())

    results = response_body.get("content")[0].get("text")
    return results

def generate_text_from_image(image_file):
    #image_data = image_file.read()
    base64_image = base64.b64encode(image_file).decode('utf-8')
    
    print("Base64 Image Length:", len(base64_image))
    
    text_result = call_claude_haiku(base64_image)
    return text_result

def generate_image( body):
    
    """
    Generate an image using Amazon Titan Image Generator G1 model on demand.
    Args:
        model_id (str): The model ID to use.
        body (str) : The request body to use.
    Returns:
        image_bytes (bytes): The image generated by the model.
    """
   
    model_id = 'amazon.titan-image-generator-v2:0'
    accept = "application/json"
    content_type = "application/json"

    try:
        response = bedrock.invoke_model(
            body=body,
            modelId=model_id,
            accept=accept,
            contentType=content_type
        )

        response_body = json.loads(response.get("body").read())

        base64_image = response_body.get("images")[0]
        base64_bytes = base64_image.encode('ascii')
        image_bytes = base64.b64decode(base64_bytes)

        finish_reason = response_body.get("error")

        if finish_reason is not None:
            raise ImageError(f"Image generation error. Error is {finish_reason}")

       
        return image_bytes

    except ClientError as err:
        message = err.response["Error"]["Message"]
        print("!!!!!!!!")
        raise


def lambda_handler(event, context):
    bucket_name = event['Records'][0]['s3']['bucket']['name']
    object_key = event['Records'][0]['s3']['object']['key']
    response = s3.get_object(Bucket=bucket_name, Key=object_key)
    image_data = response['Body'].read()
    try:
        generated_text = generate_text_from_image(image_data)
        generated_text = generated_text.split("*")
        hairstyle = generated_text[1]
        print(hairstyle)
        print('-------')
        
        """print(generated_text[0])
        story = generated_text[0]
        print(generated_text[1])
        style = generated_text[1]"""
        future_dream = "firefighter"
        random_seed = random.randint(0, 214783647)
        prompt =f"""
            A characteristic of person is {hairstyle}
            A person dedicatedly working on their {future_dream}.
            a background that matches their {future_dream},
            detaily background.
            Looking at viewer.
            Allowing the background to be an important element of the composition.
            A Whole Face.
        """
        print(prompt)
        body = json.dumps({
        "taskType": "TEXT_IMAGE",
        "textToImageParams": {
            "text": prompt.format(future_dream=future_dream)
        },
        "imageGenerationConfig": {
            "numberOfImages": 1,
            "height": 1024,
            "width": 1024,
            "cfgScale": 9.0,
            "seed": random_seed  # Random seed is used here
            }
        })
        image_bytes = generate_image(body=body)
        image = Image.open(io.BytesIO(image_bytes))
         # 처리된 이미지를 S3에 저장
        destination_bucket = bucket_name
        destination_key = object_key.replace('input/', 'output/')
        upload_image_to_s3(image, destination_bucket, destination_key)

    except Exception as e:
        print(f"Error: {e}")
    except ClientError as err:
        print(f"A client error occurred: {err}")
    else:
        print("야호!")
    return {
        'statusCode': 200,
        'body': json.dumps('Image processed successfully')
    }
    
    
def upload_image_to_s3(image, bucket_name, file_key):
    # 이미지를 바이트 스트림으로 변환
    image_bytes = io.BytesIO()
    image.save(image_bytes, format='PNG')
    image_bytes = image_bytes.getvalue()

    # S3에 이미지 업로드
    s3.put_object(Bucket=bucket_name, Key=file_key, Body=image_bytes, ContentType='image/png')