AWS 기술 블로그

Stable Diffusion을 Amazon SageMaker JumpStart로 편리하게 이용하기

Stable Diffusion 모델을 이용하면 텍스트를 이용하여 창조적인 이미지를 생성할 수 있습니다. AWS에서는 Amazon SageMaker JumpStart을 이용하여 기계 학습(ML)을 쉽게 사용할 수 있도록 사전 학습(pre-trained)된 모델을 제공하고 있는데, 2022년 10월 부터 Stable Diffusion 모델을 추가적으로 제공하고 있습니다. 이를 통해 Stable Diffusion 이미지를 쉽게 생성할 수 있으며, 즉시 Serving할 수 있도록 SageMaker Endpoint도 제공합니다. SageMaker Endpoint는 트래픽이 증가할 때는 자동으로 스케일 아웃(Scale-out)하므로, 트래픽 변동이 심할 때에도 효율적으로 인프라를 유지할 수 있으며 IAM 기반의 강화된 보안을 제공하고 있습니다.

SageMaker JumpStart 제공 Stable Diffusion Endpoint  고려사항

Amazon SageMaker Endpoint로 SageMaker JumpStart에서 제공한 Stable Diffusion 이미지 생성을 요청할 때 얻어진 응답(Response)은 아래와 같습니다. JSON 응답에는 “generated_image” 필드로 이미지의 RGB 정보를 전달합니다. 이를 클라이언트에서 활용하기 위해서는 이미지 포맷으로 변경하여야 합니다. 또한, SageMaker Endpoint로 Stable Diffusion 이미지 생성을 요청(Request)할 때에는 AWS IAM 인증을 하여야 하므로, 클라이언트는 민감한 정보인 IAM Credential을 가지고 있어야 하고, AWS SDK를 통해 API 요청을 수행하여야 합니다. 따라서 웹 브라우저 또는 모바일 앱에서는 IAM 인증 기반의 서비스를 제공하기 어렵습니다.

{
    "generated_image": [
        [[221,145,108],[237,141,98],[249,154,111],..]
        ...
    ],
    "prompt": "{
        predictions":[{
            "prompt": "astronaut on a horse", 
            "width": 768, 
            "height": 512,
            "num_images_per_prompt": 1, 
            "num_inference_steps": 50, 
            "guidance_scale": 7.5
        }]
    }
}

Stable Diffusion 모델 서빙 아키텍처

Stable Diffusion을 제공하기 위한 전체적인 아키텍처는 아래와 같습니다. SageMaker는 JumpStart로 제공되는 Stable Diffusion 모델을 가지고 있어서 입력된 텍스트로 부터 이미지를 생성할 수 있습니다. AWS Lambda는 IAM 인증을 통해 SageMaker Endpoint로 전달한 텍스트에 대한 응답으로 RGB배열 형태로 된 image map을 얻습니다.

이후 image map은 Amazon S3에 JPEG 파일로 저장하여 사용자가 쉽게 사용할 수 있도록 URL로 제공됩니다. Amazon API Gateway는 사용자의 요청을 Restful API로 받아서 Lambda 함수에 사용자의 요청을 전달하고, Lambda 함수가 생성한 URL 이미지 정보를 사용자에게 응답으로 전달합니다.

전체 서비스들의 배포는 AWS CDK를 이용하고, docker container 이미지는 ECR로 관리합니다. 이와 같이 SageMaker Endpoint에 대한 IAM 인증 및 이미지 파일 변환을 위해 API Gateway와 Lambda를 사용함으로써 웹 브라우저 또는 모바일 앱에서도 Open API를 통해 쉽고 편리하게 Stable Diffusion 이미지를 생성할 수 있습니다.

사용자가 Stable diffusion 이미지를 생성하는 순서를 정리하면 아래와 같습니다.

  • 단계 1: 사용자가 변환하고자 하는 text를 API Gateway에 전달합니다.
  • 단계 2: API Gateway는 전달받은 text를 Lambda 함수에 전달합니다.
  • 단계 3: Lambda는 IAM인증 후에 전달받은 text를 SageMaker Endpoint로 전달합니다.
  • 단계 4: SageMaker의 Stable Diffusion Model은 전달받은 text로 stable diffusion 이미지를 생성하고, 생성된 이미지 정보를 image map 형태로 Lambda 함수에 전달합니다.
  • 단계 5: Lambda 함수는 전달받은 image map으로 JPEG 이미지를 생성하여 S3에 저장합니다.
  • 단계 6: CloudFront의 도메인 주소와 저장된 이미지의 버킷 이름 및 키(key)를 조합하여, URL을 생성한 후에 API Gateway에 전달합니다.
  • 단계 7: API Gateway는 생성된 Stable diffusion 이미지에 대한 URL 정보를 사용자에게 전달합니다.
  • 단계 8: 사용자는 전달받은 URL을 이용하여 CloudFront로 부터 Stable diffusion 이미지를 다운로드 합니다.

SageMaker Endpoint로 추론(Inference) 요청 방법

Lambda 함수에서 Sagemaker Endpoint로 추론(Inference) 요청시에 아래와 같이 “ContentType“과 “Accept“을 지정하여야 합니다.

"ContentType": " application/json",
"Accept": "application/json",

이때 RequestBody에는 아래 포맷으로 Stable Diffusion에 필요한 정보를 전달합니다. width, height 값은 8로 나누어지는 정수로서 이미지의 크기를 지정합니다. num_images_per_prompt은 한번에 생성되는 이미지의 개수이고, num_inference_steps는 이미지 생성시 denoising 단계를 의미하며 숫자를 높이면 더 높은 품질의 이미지를 얻을 수 있습니다. guidance_scale은 prompt에 가까운 정도를 표현합니다.

{
    predictions":[{
        "prompt": "astronaut on a horse",
        "width": 768,
        "height": 512,
        "num_images_per_prompt": 1,
        "num_inference_steps": 50,
        "guidance_scale": 7.5
    }]
}

이와 같이 lambda_function.py에서는 Python의 boto3을 이용해 SageMaker Endpoint로 요청(request)을 전달하는데, ContentType은 “application/json“을, Accept 헤더로는 “Accept='application/json” 또는 “Accept='application/json;jpeg“을 사용하여야 합니다.

import boto3

runtime = boto3.Session().client('sagemaker-runtime')
response = runtime.invoke_endpoint(
	EndpointName=endpoint, 
	ContentType='application/json', 
	Accept='application/json;jpeg', 
	Body=json.dumps(payload))

RGB 이미지 데이터를 변환하여 Amazon S3에 업로드 하는 경우

SageMaker Endpoint로 query시에 Accept헤더를 “application/json“로 하는 경우에, RGB로 된 text데이터가 내려옵니다. 이미지 데이터는 JSON의 “Body“와 “generated_image“로 부터 추출한 후에, PIL(pillow)numpy 라이브러리를 사용하여 S3에 저장할 수 있는 바이너리 이미지 데이터로 변환합니다. 이때lambda_function.py의 코드는 아래와 같습니다.

from PIL import Image
import numpy as np

def parse_response(query_response):
    response_dict = json.loads(query_response)
    return response_dict["generated_images"], response_dict["prompt"]
    
response_payload = response['Body'].read().decode('utf-8')
generated_images, prompt = parse_response(response_payload)
        
image = Image.fromarray(np.uint8(generated_images[0]))
buffer = io.BytesIO()
image.save(buffer, "jpeg")
buffer.seek(0)
            
s3 = boto3.client('s3')
s3.upload_fileobj(buffer, mybucket, mykey, ExtraArgs={"ContentType": "image/jpeg"})

그런데, Lambda에서 pillow, numpy 라이브러리를 “pip install --target=[lambda 폴더] pillow numpy“와 같이 설치한 후 압축해서 올리려면 layer를 추가하는 문제가 있으므로, docker container를 이용하여 pillow, numpy와 같은 라이브러리를 사용할 수 있도록 했습니다. 이때의 Dockerfile의 예는 아래와 같습니다.

FROM amazon/aws-lambda-python:3.8

RUN pip3 install --upgrade pip
RUN python -m pip install joblib awsiotsdk
RUN pip install numpy pillow
WORKDIR /var/task/lambda
COPY lambda_function.py /var/task
COPY . .

CMD ["lambda_function.lambda_handler"]

JPEG로 encoding된 이미지를 S3에 업로드 하는 경우

SageMaker Endpoint로 query시에 Accept헤더를 “application/json;jpeg“로 설정하면 SageMaker Endpoint가 base64로 encoding된 JPEG 이미지를 전달합니다. JPEG로 된 이미지를 이용하면, RGB 배열 형태보다 작은 크기로 이미지를 전달받을 수 있습니다. 아래에서는 전달받은 이미지를 base64 decoding 후에 인메모리 바이너리 스트림으로 변경하여 S3로 업로드합니다.

response_payload = response['Body'].read().decode('utf-8')
generated_images, prompt = parse_response(response_payload)

import base64
img_str = base64.b64decode(generated_images[0])
buffer = io.BytesIO(img_str)  
s3.upload_fileobj(buffer, mybucket, mykey, ExtraArgs={"ContentType": "image/jpeg"})

AWS CDK 내부 리소스 생성 코드 준비

본 게시글은 Typescript를 이용하여 CDK 배포 준비를 합니다. S3 bucket을 아래와 같이 생성합니다. bucketName으로 Bucket 이름을 지정할 수 있습니다. 아래처럼 지정하지 않으면 cdk로 시작하는 bucket을 생성합니다.

const s3Bucket = new s3.Bucket(this, "gg-depolyment-storage",{
    // bucketName: bucketName,
    blockPublicAccess: s3.BlockPublicAccess.BLOCK_ALL,
    removalPolicy: cdk.RemovalPolicy.DESTROY,
    autoDeleteObjects: true,
    publicReadAccess: false,
    versioned: false,
});

CloudFront를 생성합니다. 여기서 CloudFront의 Origin은 생성한 S3 Bucket로 지정합니다.

const distribution = new cloudFront.Distribution(this, 'cloudfront', {
    defaultBehavior: {
      origin: new origins.S3Origin(s3Bucket),
      allowedMethods: cloudFront.AllowedMethods.ALLOW_ALL,
      cachePolicy: cloudFront.CachePolicy.CACHING_DISABLED,
      viewerProtocolPolicy: cloudFront.ViewerProtocolPolicy.REDIRECT_TO_HTTPS,
    },
    priceClass: cloudFront.PriceClass.PRICE_CLASS_200,  
});

lambda-stable-diffusion” 이름을 가지는 Lambda를 docker container 환경으로 생성합니다. Lambda가 S3 Bucket에 대한 Read/Write 속성 및 SageMaker에 대한 권한을 가지도록 설정합니다.

const mlLambda = new lambda.DockerImageFunction(this, "lambda-stable-diffusion", {
    description: 'lambda function for stable diffusion',
    functionName: 'lambda-stable-diffusion',
    memorySize: 512,
    code: lambda.DockerImageCode.fromImageAsset(path.join(__dirname, '../../lambda')),
    timeout: cdk.Duration.seconds(60),
    environment: {
        bucket: s3Bucket.bucketName,
        endpoint: endpoint,
        domain: distribution.domainName
    }
}); 
s3Bucket.grantReadWrite(mlLambda);
const SageMakerPolicy = new iam.PolicyStatement({
    actions: ['sagemaker:*'],
    resources: ['*'],
});
mlLambda.role?.attachInlinePolicy(
    new iam.Policy(this, 'sagemaker-policy', {
        statements: [SageMakerPolicy],
    }),
);

API Gateway와 리소스로 “text2image”를 생성한 후에 POST method를 설정합니다.

const api = new apiGateway.RestApi(this, 'api-stable-diffusion', {
    endpointTypes: [apiGateway.EndpointType.REGIONAL],
    deployOptions: {
      stageName: stage,
    },
});
const text2image = api.root.addResource('text2image');
text2image.addMethod('POST', new apiGateway.LambdaIntegration(mlLambda, {
    passthroughBehavior: apiGateway.PassthroughBehavior.WHEN_NO_TEMPLATES,
    integrationResponses: [{
      statusCode: '200',
    }], 
    proxy:false, 
}), {
    methodResponses: [{
        statusCode: '200',
        responseModels: {
          'application/json': apiGateway.Model.EMPTY_MODEL,
        }, 
      }
    ]
});

API Gateway의 Invoke URL을 아래와 같이 확인합니다. 이 정보는 API Gateway Console에서도 확인할 수 있습니다.

new cdk.CfnOutput(this, 'apiUrl', {
    value: api.url,
    description: 'The url of API Gateway',
});

직접 실습 해보기

사전 준비 사항

이 솔루션을 사용하기 위해서는 사전에 아래와 같은 준비가 되어야 합니다.

SageMaker Studio 셋업

SageMaker Studio를 설치하여 SageMaker JumpStart를 사용하고자 합니다. SageMaker Console에 접속하여 [Create domain]을 선택합니다. 편의상 서울 리전에서 모든 동작을 수행합니다. 이후 아래처럼 Domain name을 입력하고 User profile과 Execution role은 기본값을 유지한 상태에서 [Submit]을 선택합니다. 여기서는 Domain name을 “MyStableDiffusion”으로 입력하였습니다.

이후 아래처럼 VPC와 Subnet을 생성합니다.

SageMaker Studio Console에서 기 생성한 “MyStableDiffusion”을 선택한 후에 [Create user profile]을 선택하고 이후 아래처럼 기본값으로 [Next]를 선택합니다. 이후 [Jupyter Lab 3.0]등을 기본값으로 선택 후에 [Submit]을 선택하여 user profile를 생성합니다.

Stable Diffusion 모델 서빙을 위한 SageMaker Endpoint 생성

SageMaker Studio Console에서 Studio와 User profile을 고른 후에 [Open Studio]를 선택합니다. 이후 아래와 같이 [Home] 화면에서 [JumpStart]를 선택합니다.

[SageMaker JumpStart]에서 [Stable Diffusion 2.1 base]를 선택합니다. 만약 [Stabel Diffusion 2.1]이 화면에 보이지 않는다면, 오른쪽 상단의 [Search]에서 “Stable Diffusion”을 입력하여 검색한 후에 선택합니다.

아래처럼 [Stable Diffusion 2.1 base]에서 오른쪽의 [Open notebook]을 선택합니다.

“Introduction to JumpStart – Text to Image”이 오픈되면 아래로 스크롤해서 “Clean up the endpoint”를 아래처럼 주석처리 합니다.

“Introduction to JumpStart – Text to Image”에서 상단 메뉴의 [Run]을 선택한 후에, [Run All Cells]을 선택하여 노트북을 실행합니다. 십분에서 수 십분 정도가 지나면 노트북에 있는 모든 Cell이 수행됩니다.

노트북 실행이 완료되면 왼쪽 메뉴의 [Deployment] – [Endpoint]를 선택하여, 생성된 Endpoint 정보를 아래처럼 확인합니다.

생성된 Endpoint를 선택하여 들어간 후에 아래와 같은 [ENDPOINT DETAILS]에서 Endpoint 이름을 복사합니다. 여기서는 “jumpstart-example-infer-model-txt2img-s-2023-02-08-13-53-49-534″을 복사하였습니다.

AWS Cloud9 생성

추론을 위한 인프라에는 API Gateway, S3, Lambda, CloudFront가 있으며, AWS CDK로 배포합니다. 상세한 배포정보는 cdk-stable-diffusion-stack.ts을 참조합니다. AWS Cloud9을 생성하기 위하여 Cloud9 console에서 [Create environment]를 선택한 후에 아래처럼 Name을 입력합니다. 여기서는 “Stabel Diffusion”이라고 입력하였습니다. 이후 나머지는 기본값을 유지하고 [Create]를 선택합니다.

Cloud9 인스턴스가 생성된 후에 [Open]을 선택하여 진입한 후 아래처럼 터미널을 실행합니다.

AWS CDK로 솔루션 배포

아래와 같이 관련 코드를 다운로드 합니다.

curl https://aws-korea-tech-blog-public.s3.ap-northeast-2.amazonaws.com/20230216-Stable-Diffustion-CDK/stable-diffusion-api-server.zip -o stable-diffusion-api-server.zip

이후 압축을 풀고 이동합니다.

unzip stable-diffusion-api-server.zip && cd stable-diffusion-api-server

인프라 생성시 SageMaker의 Endpoint 정보가 필요하므로, 아래와 같이 좌측 파일탐색기에서 “stable-diffusion-api-server/cdk-stable-diffusion/lib/cdk-stable-diffusion-stack.ts“를 선택하여 이전 단계에서 복사한 Endpoint의 이름을 수정하고 저장합니다.

CDK 폴더(cdk-stable-diffusion)로 이동하여 “aws-cdk-lib"와 “path” 라이브러리를 npm으로 설치합니다. 여기서, “aws-cdk-lib”은 CDK 2.0 라이브러리로서 v2.64.0을 사용하였습니다.

cd cdk-stable-diffusion && npm install aws-cdk-lib@2.64.0 path

아래 명령어로 전체 인프라를 설치합니다.

cdk deploy

CDK로 인프라 설치가 완료되면 아래와 같이 설치된 인프라의 정보를 알 수 있습니다. “appUrl”은 API Gateway의 Invoke URL로서 클라이언트가 Stable Diffusion을 요청할 때 필요합니다. 아래 케이스에서는 “appUrl”이 “https://734ury6k98.execute-api.ap-northeast-2.amazonaws.com/dev” 입니다. API Gateway의 API 리소스(Resource)이름이 “text2image”이므로, 클라이언트에서 API Gateway로 Stable Diffusion을 요청할 때 사용하는 URL은 https://734ury6k98.execute-api.ap-northeast-2.amazonaws.com/dev/text2image임을 알 수 있습니다.

Curl 명령어로 추론

Curl 명령어로 아래와 같이 실행할 수 있습니다. 여기서 curl이 접속하는 URL은 CDK Outputs의 “appUrl”과 리소스명인 “text2image”를 조합한 “https://734ury6k98.execute-api.ap-northeast-2.amazonaws.com/dev/text2image”입니다.

curl -X POST https://734ury6k98.execute-api.ap-northeast-2.amazonaws.com/dev/text2image -H "Content-Type: application/json" -d '{"text":"astronaut on a horse"}'

아래는 curl로 요청한 추론에 대한 결과의 예입니다. “body”에 추론의 결과로 생성된 이미지의 URL이 있습니다.

{"statusCode": 200, "body": "https://d283dvdglbetjo.cloudfront.net/img_20230208-014926"}

상기 Curl 명령어에서 text의 내용을 바꾸면, 아래와 같은 Stable Diffusion 이미지를 얻을 수 있습니다.

“I see trees of green Red roses too. I see them bloom. For me and you. And I think to myself. What a wonderful world” (Louis Armstrong의 What a Wonderful World)

ukrainian girl with blue and yellow clothes near big ruined building, concept art, trending on artstation, highly detailed, intricate, sharp focus, digital art

a portrait of a korean woman that is a representation of Korean culture, buenos aires, fantasy, intricate, highly detailed, digital painting, artstation, concept art, smooth, sharp focus, illustration, art by artgerm and greg rutkowski and alphonse mucha

a young blonde male jedi with short hair standing still looking at the sunset concept art by Doug Chiang cinematic, realistic painting, high definition, concept art, portait image, path tracing, serene landscape, high quality, highly detailed, 8K, soft colors, warm colors, turbulent sea, high coherence, anatomically correct, hyperrealistic, concept art, defined face, five fingers, symmetrical

the eye of the storm, atmospheric, hyper realistic, 8k, epic composition, cinematic, octane render, artstation landscape vista photography by Carr Clifton & Galen Rowell, 16K resolution, Landscape veduta photo by Dustin Lefevre & tdraw, 8k resolution, detailed landscape painting by Ivan Shishkin, DeviantArt, Flickr, rendered in Enscape, Miyazaki, Nausicaa Ghibli, Breath of The Wild, 4k detailed post processing, artstation, rendering by octane

리소스 정리하기

API Gateway, S3, Lambda, CloudFront는 아래와 같이 정리합니다.

cdk destroy

SageMaker Models, SageMaker Endpoints, SageMaker Endpoint configuration에서 생성했던 “jumpstart-example-infer-model-txt2img-s-2023-02-08-13-53-49-534”에 대한 model, endpoint, endpoint configuration을 삭제합니다.

Cloud9 console에서 배포에 사용되었던 Cloud9을 삭제합니다.

결론

Amazon SageMaker JumpStart에서 제공하고 있는 Stable Diffusion 모델을 활용하여 클라이언트에서 쉽고 편리하게 Stable Diffusion 추론을 수행하는 방법을 설명하였습니다. 이를 통해 Stable Diffusion 이미지를 웹 브라우저 또는 모바일 앱에서도 쉽게 생성하여 활용할 수 있습니다. 또한, AWS CDK로 Stable Diffusion를 위한 인프라를 구축하는 방법을 예제와 함께 설명하였습니다. Stable Diffusion은 계속 발전하고 있고 이미 다양한 용도로 활용되고 있어서 향후 인공 지능 아트분야 뿐 아니라 산업 전반에서 많은 영향을 줄 수 있을 것으로 기대합니다.

실습 코드 및 도움 자료

아래의 링크에서 실습 소스 파일 및 기계 학습과 관련된 자료를 확인하실 수 있습니다.

Kyoung-Su Park

Kyoung-Su Park

박경수 솔루션즈 아키텍트는 다양한 워크로드에 대한 개발 경험을 바탕으로 고객이 최적의 솔루션을 선택하여 비즈니스 성과를 달성할 수 있도록 고객과 함께 효율적인 아키텍처를 구성하는 역할을 수행하고 있습니다. 현재 AWS의 Machine Learning, IoT, Analytics TFC에서 활동하고 있습니다.

Joonyong Park

Joonyong Park

박준용 파트너 솔루션스 아키텍트는 파트너 역량강화를 위하여 AWS Ambassador, BlackBelt, APN Immersion Day, JetPack 등의 프로그램을 수행하고 있으며 파트너사에 Data Analytics와 Machine Learning 분야 Enablement를 준비 및 제공하고 있습니다.