Deploy Endpoint#

Follow THIS to deploy a syn endpoint on sagemaker

Invoke Endpoint#

Let invoke the deployed endpoint

import boto3
import json
import base64
import io
from PIL import Image


# create sagemaker runtime client
sm_client = boto3.client("sagemaker-runtime")


# endpoint name
endpoint_name = "sdxl-1-0-jumpstart-2024-01-24-09-54-52-906"

Let create a request body with prompt and invoke the endpoint

# payload request
payload = json.dumps({
  "text_prompts": [
    {
    "text": "wife and baby",
    "weight": 1.0
    }
  ],
  "cfg_scale": 15,
  "samples": 1,
  "seed": 3,
  "style_preset": "anime",
  "num_inference_steps": 30,
  "height": 1024,
  "width": 1024
})


# invoke endpoint
response = sm_client.invoke_endpoint(
    EndpointName=endpoint_name,
    ContentType="application/json",
    Body=payload
)

Finally parse response, decode base64 response and save image

data = response['Body'].read().decode('utf-8')
data = json.loads(data)
imgstring = data['artifacts'][0]['base64']


# decode image
imgdata = base64.b64decode(imgstring)


# save image
filename = 'some_image.jpg'
with open(filename, 'wb') as f:
    f.write(imgdata)


# open image in notebook
image = Image.open(io.BytesIO(imgdata))
display(image)

Lambda Function#

Let create a lambda function which invoke the endpoint and save image to S3. Project structure

|--Dockerfile
|--requirements.txt
|--index.py

Handler code

"""
lambda call sagemaker diffusion model
"""


import sagemaker
from stability_sdk_sagemaker.predictor import StabilityPredictor
from stability_sdk.api import GenerationRequest, GenerationResponse, TextPrompt
from PIL import Image
import io
import base64
import boto3
import datetime
import json
import os 


#  


STYLE = "anime"


s3_client = boto3.client("s3")


sagemaker_session = sagemaker.Session()


deployed_model = StabilityPredictor(
    endpoint_name=os.environ["ENDPOINT_NAME"], sagemaker_session=sagemaker_session
)




def decode_and_show(model_response: GenerationResponse) -> None:
    """
    Decodes and displays an image from SDXL output


    Args:
        model_response (GenerationResponse): The response object from the deployed SDXL model.


    Returns:
        None
    """
    # file name
    name = datetime.datetime.now().strftime("%m-%d-%Y-%H-%M-%S")
    # key
    key = f"diffuision/{name}.png"
    # image
    image = model_response.artifacts[0].base64
    image_data = base64.b64decode(image.encode())
    s3_client.upload_fileobj(io.BytesIO(image_data), os.environ["BUCKET_NAME"], key)
    # signed url
    sign_url = s3_client.generate_presigned_url(
        "get_object",
        Params={"Bucket": os.environ["BUCKET_NAME"], "Key": key},
        ExpiresIn=3600,
    )
    # image = Image.open(io.BytesIO(image_data))
    # image.save("hehe.png")
    return sign_url




def handler(event, context):
    """
    handler
    """


    # parse prompt
    try:
        prompt = event["queryStringParameters"]["prompt"]
    except:
        prompt = "fish"
    
    try:
        style = event["queryStringParameters"]["style"]
    except:
        style = "anime"


    # call model
    output = deployed_model.predict(
        GenerationRequest(
            text_prompts=[TextPrompt(text=prompt)],
            style_preset=style,
            seed=3,
            height=1024,
            width=1024,
        )
    )


    # save image to s3
    try: 
        url = decode_and_show(output)
    except:
        url = "ERROR"


    # return
    return {
        "statusCode": 200,
        "headers": {
            "Access-Control-Allow-Origin": "*",
            "Access-Control-Allow-Headers": "Content-Type",
            "Access-Control-Allow-Methods": "OPTIONS,GET",
        },
        "body": json.dumps(
            {
                "url": url,
            }
        ),
    }

Here is Dockerfile

FROM public.ecr.aws/lambda/python:3.9


# create code dir inside container
RUN mkdir ${LAMBDA_TASK_ROOT}/source


# copy code to container
COPY "requirements.txt" ${LAMBDA_TASK_ROOT}/source


# copy handler function to container
COPY ./index.py ${LAMBDA_TASK_ROOT}


# install dependencies for running time environment
COPY ./package/ ${LAMBDA_TASK_ROOT}


# set the CMD to your handler
CMD [ "index.handler" ]

And requirements.txt

sagemaker
stability-sdk[sagemaker] @ git+https://github.com/Stability-AI/stability-sdk.git@sagemaker

Reference#