Implementing a fast scaling and low cost Stable Diffusion inference solution with serverless and containers on AWS
import boto3
import gradio as gr
import requests
import json


# 设置AWS S3的访问密钥和密钥ID
#session = boto3.Session(profile_name='new')
#s3 = session.client('s3')
s3 = boto3.client('s3')
s3bucket =None
s3prefix =None


API_URL_INPUT = "https://XXX.execute-api.us-east-1.amazonaws.com/prod/"

with open('./payload1.json', 'r') as p1:
    payload1 = json.load(p1)
    formatted_payload1 = json.dumps(payload1, indent=4)
with open('./payload2.json', 'r') as p2:
    payload2 = json.load(p2)
    formatted_payload2 = json.dumps(payload2, indent=4)
with open('./payload3.json', 'r') as p3:
    payload3 = json.load(p3)
    formatted_payload3 = json.dumps(payload3, indent=4)
with open('./payload4.json', 'r') as p4:
    payload4 = json.load(p4)
    formatted_payload4 = json.dumps(payload4, indent=4)

# 创建请求 API

def api_request(api_url, json_input):
    global s3prefix
    global s3bucket
    #print (json_input)
    try:
        json_data = json.loads(json_input)  # 解析 JSON 数据
        s3prefix = json_data["alwayson_scripts"]["id_task"]
        #print ("s3prefix--->",s3prefix)
        response = requests.post(api_url, json=json_data)
        result = json.loads(response.text)
        obj_path  = result["output_location"]
        s3bucket = result["output_location"].split('/')[2]
        #print ("s3bucket------>",s3bucket)
        div = image_viewer(s3prefix)
        #print ("div------>",div)
        return div,obj_path
    except Exception as e:
        return {"error": str(e)}


# 获取特定前缀下的所有对象的列表

def list_objects(prefix):
    response = s3.list_objects(Bucket=s3bucket, Prefix=prefix)
    #print (response)
    while "Contents" in response:
        objects = response['Contents']
        return [obj['Key'] for obj in objects]

# 创建图像的HTML标记
def image_viewer(folder):
    image_url = []
    div = ""
    # print ("listobj--->",list_objects(folder))
    # if list_objects(folder) != None :
    #     prev_len = len(list_objects(folder))
    # else :
    #     prev_len = 0
    while True :
        objects = list_objects(folder)
        # cur_len = len(objects)
        # print ("objects ---->",objects)
        # print (prev_len,cur_len)
        while  objects:
            for object_key in objects:
                if object_key.endswith('.jpg') or object_key.endswith('.jpeg') or object_key.endswith('.png'):
                    image_url.append(s3.generate_presigned_url('get_object', Params={
                        'Bucket': s3bucket, 'Key': object_key}, ExpiresIn=3600))
            image_tags = ""
            for image in image_url:
                image_tags += f"<div style='padding: 5px'><a href='{image}' target='_blank'><img src='{image}'></a></div>"
                div = f"<div style='display: grid; grid-template-columns: repeat(5, 1fr); grid-gap: 10px'>{image_tags}</div>"
            return div
def choose_option(option):
    if option == "txt2img":
        return formatted_payload1
    elif option == "img2img":
        return formatted_payload2
    elif option == "txt2img with ControlNet":
        return formatted_payload3
    elif option == "txt2img with Lora & ControlNet":
        return formatted_payload4
    else:
        return ""
with gr.Blocks(title="SD-on-EKS Demo") as demo:
    gr.Markdown("# Stable diffusion on EKS API Demo")
    with gr.Tab("输入"):
        api_url =gr.Textbox(label="API 请求地址",value=API_URL_INPUT)
        object_path =  gr.Textbox(label="图片生成路径")
        with gr.Row():
            inputs = gr.Radio(["txt2img", "img2img", "txt2img with ControlNet", "txt2img with Lora & ControlNet"], label="选择一个请求")
        with gr.Row():
            payload = gr.Textbox(lines=10,label="Payload输入")
            picture_output = gr.HTML("图库")
        with gr.Row():
            request_button = gr.Button("提交")
            #total_timer =  gr.Textbox(label="图片生成时间(s)")
        inputs.change(fn=choose_option, inputs=inputs, outputs=payload)
    request_button.click(api_request, inputs=[api_url,payload],outputs=[picture_output,object_path])
demo.launch()

deploy/DemoUI/Demo2.py
aws-samples/stable-diffusion-on-eks
import boto3
import gradio as gr
import requests
import json
import time
import logging
import json
import botocore
import traceback


API_URL_INPUT = "https://XXX.execute-api.us-west-2.amazonaws.com/prod/"
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(process)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
from botocore.exceptions import ClientError

my_queue = 'output_queue'
# Create an SQS client object
sqs = boto3.client('sqs',region_name="us-west-2")

# Get the queue URL
sqsQueueUrl = sqs.get_queue_url(QueueName=my_queue)['QueueUrl']

sqsRes = boto3.resource('sqs',region_name="us-west-2")
queue = sqsRes.Queue(sqsQueueUrl)

s3 = boto3.client('s3')
start_time = None
end_time = None
total_time = None

with open('./payload1.json', 'r') as p1:
    payload1 = json.load(p1)
    formatted_payload1 = json.dumps(payload1, indent=4)
with open('./payload2.json', 'r') as p2:
    payload2 = json.load(p2)
    formatted_payload2 = json.dumps(payload2, indent=4)
with open('./payload3.json', 'r') as p3:
    payload3 = json.load(p3)
    formatted_payload3 = json.dumps(payload3, indent=4)
with open('./payload4.json', 'r') as p4:
    payload4 = json.load(p4)
    formatted_payload4 = json.dumps(payload4, indent=4)


def choose_option(option):
    if option == "txt2img":
        return formatted_payload1
    elif option == "img2img":
        return formatted_payload2
    elif option == "txt2img with ControlNet":
        return formatted_payload3
    elif option == "txt2img with Lora & ControlNet":
        return formatted_payload4
    else:
        return ""
    
# 生成预签名 URL 的函数
def generate_presigned_url(s3_uri):
    try:
        # 解析 S3 URI 获取桶名和对象键
        s3_uri_parts = s3_uri.replace('s3://', '').split('/')
        bucket_name = s3_uri_parts[0]
        object_key = '/'.join(s3_uri_parts[1:])
        
        # 生成预签名 URL
        presigned_url = s3.generate_presigned_url('get_object',
                                                  Params={'Bucket': bucket_name, 'Key': object_key},
                                                  ExpiresIn=3600)  # 预签名 URL 的有效期（以秒为单位）
        return presigned_url
    except botocore.exceptions.ParamValidationError as e:
        print("生成预签名 URL 时出错：", e)
        return None


def receive_messages(queue, max_number, wait_time):
    """
    Receive a batch of messages in a single request from an SQS queue.

    :param queue: The queue from which to receive messages.
    :param max_number: The maximum number of messages to receive. The actual number
                       of messages received might be less.
    :param wait_time: The maximum time to wait (in seconds) before returning. When
                      this number is greater than zero, long polling is used. This
                      can result in reduced costs and fewer false empty responses.
    :return: The list of Message objects received. These each contain the body
             of the message and metadata and custom attributes.
    """
    try:
        messages = queue.receive_messages(
            MaxNumberOfMessages=max_number,
            WaitTimeSeconds=wait_time,
            AttributeNames=['All'],
            MessageAttributeNames=['All']
        )
    except Exception as error:
        logger.error(f"Error receiving messages: {error}")
    else:
        return messages

def delete_message(message):
    """
    Delete a message from a queue. Clients must delete messages after they
    are received and processed to remove them from the queue.

    :param message: The message to delete. The message's queue URL is contained in
                    the message's metadata.
    :return: None
    """
    try:
        message.delete()
    except ClientError as error:
        traceback.print_exc()
        raise error    
    

def api_request(api_url, json_input):
    global s3prefix
    global start_time
    print ("收到请求,开始计时")
    start_time = time.time()
    #print (json_input)
    try:
        json_data = json.loads(json_input)  # 解析 JSON 数据
        s3prefix = json_data["alwayson_scripts"]["id_task"]
        print ("s3prefix--->",s3prefix)
        response = requests.post(api_url, json=json_data)
        result = json.loads(response.text)
        obj_path  = result["output_location"]
        div = image_gallery(s3prefix)
        #print ("div------>",div)
        return div,obj_path
    except Exception as e:
        return {"error": str(e)}



def image_gallery (task_id):
    global end_time
    global total_time
    while True:
        try:
            received_messages = receive_messages(queue, 1, 20)
            img_urls =[]
            print('received_messages',received_messages)
            for message in received_messages:
                end_time = time.time()
                total_time = end_time - start_time
                print ("任务结束,停止计时,总耗时",total_time,"秒")
                Payload = json.loads(message.body)
                taskid_output = json.loads(Payload['Message'])['parameters']['id_task']
                if int(taskid_output) == int(task_id):
                    print("Processing message with id_task:", taskid_output)
                    url_output = json.loads(Payload['Message'])['parameters']['image_url']
                    url_output_dict = [item.strip('\'') for item in url_output.split(',')]
                    for s3uri in url_output_dict:
                        presigned_url = generate_presigned_url(s3uri)
                        img_urls.append(presigned_url)
                        #print ("presigned_url---->",(presigned_url))
                    #print (taskid_output)
                    print ("图片队列长度--->",len(img_urls))
                delete_message(message)
            #break
        except Exception as e:
            print(e)
            continue
        return (img_urls)


with gr.Blocks(title="SD-on-EKS Demo") as demo:
    gr.Markdown("# Stable diffusion on EKS API Demo")
    with gr.Tab("输入"):
        api_url =gr.Textbox(label="API 请求地址",value=API_URL_INPUT)
        object_path =  gr.Textbox(label="图片生成路径")
        with gr.Row():
            inputs = gr.Radio(["txt2img", "img2img", "txt2img with ControlNet", "txt2img with Lora & ControlNet"], label="选择一个请求")
        with gr.Row():
            payload = gr.Textbox(lines=10,label="Payload输入")
            gallery = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery", columns=[5], rows=[5], object_fit="contain", height=400)
        with gr.Row():
            request_button = gr.Button("提交")
            total_timer =  gr.Textbox(label="图片生成时间(s)",value=total_time)
        inputs.change(fn=choose_option, inputs=inputs, outputs=payload)
    request_button.click(api_request, inputs=[api_url,payload],outputs=[gallery,object_path])
demo.launch()
deploy/DemoUI/Demo3.py
aws-samples/stable-diffusion-on-eks
import boto3
import gradio as gr
import requests


# 设置AWS S3的访问密钥和密钥ID
#session = boto3.Session(profile_name='ue1')
#s3 = session.client('s3')
s3 = boto3.client('s3')


# 获取特定前缀下的所有对象的列表

def list_objects(Bucket_name,prefix):
    response = s3.list_objects(Bucket=Bucket_name, Prefix=prefix)
    objects = response['Contents']
    return [obj['Key'] for obj in objects]



# 创建图像的HTML标记
def gallery_viewer(Bucket_name,folder):
    objects = list_objects(Bucket_name,folder)
    image_url = []
    for object_key in objects:
        if object_key.endswith('.jpg') or object_key.endswith('.jpeg') or object_key.endswith('.png'):
            image_url.append(s3.generate_presigned_url('get_object', Params={
                'Bucket': Bucket_name, 'Key': object_key}, ExpiresIn=3600))
            #print (image_url)
    return image_url


with gr.Blocks(title="SD-on-EKS Demo") as demo:
    gr.Markdown("# Stable diffusion on EKS API Demo")
    with gr.Tab("输入"):
        with gr.Row():
            Bucket_name = gr.Textbox(label="存储桶名称",value = "sdoneksdataplanestack-outputs3bucket9fe85b9f-pzlnjnu6xsl7")
            folder = gr.Textbox(label="文件夹名称")
            request_button = gr.Button("提交")
        gallery = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery", columns=[5], rows=[5],object_fit="contain", height=400)        
    request_button.click(gallery_viewer, inputs=[Bucket_name,folder],outputs=gallery)
demo.launch(share=True)


deploy/DemoUI/S3Broswer.py
aws-samples/stable-diffusion-on-eks
import argparse
import boto3
import urllib3
import csv
from tqdm import tqdm
import os

# 初始化S3客户端
#session = boto3.Session(profile_name='ue1')
#s3 = session.client('s3')
s3 = boto3.client('s3')

# 读取CSV文件
def read_csv(csv_file):
    urls = []
    with open(csv_file, 'r') as file:
        reader = csv.DictReader(file)
        for row in reader:
            urls.append(row)
    return urls

# 下载文件并上传到S3
def download_and_upload_to_s3(url, bucket, key, index, total):
    http = urllib3.PoolManager()
    response = http.request('GET', url, preload_content=False)

    total_size = int(response.headers['Content-Length'])
    chunk_size = 1024 * 1024  # 1 MB chunks (adjust as needed)
    
    # 使用 tqdm 创建下载进度条
    progress_bar = tqdm(total=total_size, unit='B', unit_scale=True, desc=f"Downloading {url} ({index}/{total})")

    # 打开本地临时文件用于存储下载内容
    temp_file_path = '/tmp/temp_file'
    with open(temp_file_path, 'wb') as temp_file:
        for chunk in response.stream(chunk_size):
            temp_file.write(chunk)
            progress_bar.update(len(chunk))

    progress_bar.close()
    
    # 上传到S3并监控上传进度
    s3.upload_file(
        temp_file_path,
        bucket,
        key,
        Callback=ProgressPercentage(progress_bar, total_size)
    )
    
    os.remove(temp_file_path)  # 删除临时文件

    print(f"正在从 {url} 下载到S3{bucket}/{key} ({index}/{total})")

# 用于监控上传进度的回调函数
class ProgressPercentage:
    def __init__(self, progress_bar, total_size):
        self.progress_bar = progress_bar
        self.total_size = total_size
        self.uploaded = 0

    def __call__(self, bytes_amount):
        self.uploaded += bytes_amount
        self.progress_bar.update(bytes_amount)

def main(csv_file):
    # 读取CSV文件
    urls = read_csv(csv_file)
    total_files = len(urls)
    
    for index, row in enumerate(urls, start=1):
        url = row['url']
        bucket = row['bucket']
        key = row['key']
        
        try:
            download_and_upload_to_s3(url, bucket, key, index, total_files)
        except Exception as e:
            print(f"Failed to download and upload {url} to S3 bucket {bucket} with key {key}: {str(e)}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Download and upload files from a CSV to S3')
    parser.add_argument('csv_file', type=str, help='Path to the CSV file')

    args = parser.parse_args()
    main(args.csv_file)

deploy/tools/S3uploader/s3uploader.py
aws-samples/stable-diffusion-on-eks
import boto3
import json
import requests
import logging
import sys
import argparse
import time
import os
logger = logging.getLogger(__name__)

pid = os.getpid()


logging.basicConfig(level=logging.INFO, format='From Client: %(asctime)s - %(process)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

url = '<api gw url>'


json_data = {
    "alwayson_scripts": {
        "task": "text-to-image",
        "sd_model_checkpoint": "135643.safetensors",
        "id_task": "16869",
        "uid": "123",
        "save_dir": "outputs"
    },
    "prompt": "A dog",
    "steps": 20,
    "width": 768,
    "height": 768
}
    
headers = {'Content-Type': 'application/json','x-api-key': '<api key>'}
json_data = json.dumps(json_data)
response = requests.post(url, data=json_data,headers = headers)
logger.info(f"response for request {response}")

    




load_test/multi-process-client.py
aws-samples/stable-diffusion-on-eks
# from multiprocessing import Process
import time
import sys
import logging
import subprocess
import boto3
import traceback
from PIL import Image
from datetime import datetime
import io
import json
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(process)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
args = sys.argv
# print("Command-line arguments:")
# for arg in args:
#     print(arg)
from botocore.exceptions import ClientError

num_of_requests = int(args[1])
interval = int(args[2])
sqsQueueUrl = '<sqs url>'
procs = []
normals = []
abnormals = []
SQS_WAIT_TIME_SECONDS = 1

s3_resource = boto3.resource("s3")
def get_bucket_and_key(s3uri):
    pos = s3uri.find("/", 5)
    bucket = s3uri[5:pos]
    key = s3uri[pos + 1 :]
    return bucket, key

def receive_messages(queue, max_number, wait_time):
    """
    Receive a batch of messages in a single request from an SQS queue.

    :param queue: The queue from which to receive messages.
    :param max_number: The maximum number of messages to receive. The actual number
                       of messages received might be less.
    :param wait_time: The maximum time to wait (in seconds) before returning. When
                      this number is greater than zero, long polling is used. This
                      can result in reduced costs and fewer false empty responses.
    :return: The list of Message objects received. These each contain the body
             of the message and metadata and custom attributes.
    """
    try:
        messages = queue.receive_messages(
            MaxNumberOfMessages=max_number,
            WaitTimeSeconds=wait_time,
            AttributeNames=['All'],
            MessageAttributeNames=['All']
        )
    except Exception as error:
        logger.error(f"Error receiving messages: {error}")
    else:
        return messages

def delete_message(message):
    """
    Delete a message from a queue. Clients must delete messages after they
    are received and processed to remove them from the queue.

    :param message: The message to delete. The message's queue URL is contained in
                    the message's metadata.
    :return: None
    """
    try:
        message.delete()
    except ClientError as error:
        traceback.print_exc()
        raise error

# send request
i = 0
for i in range(num_of_requests):
    proc = subprocess.Popen(['python', 'multi-process-client.py'], stdout=subprocess.PIPE)
    procs.append(proc)
    time.sleep(interval)


sqsRes = boto3.resource('sqs')
queue = sqsRes.Queue(sqsQueueUrl)

time_lst = []

j = 0
for j in range(num_of_requests):
    
    return_code = procs[j].wait()
    # logger.info(f"Request number {j} finished with return code {return_code}.")
    if return_code == 0:
        normals.append(return_code)

        while True:
            try:
                start_time = time.perf_counter()
                received_messages = receive_messages(queue, 1, 20)
                if len(received_messages)==0:
                    print('wait again')
                    received_messages = receive_messages(queue, 1, 20)

                for message in received_messages:
                    Payload = json.loads(message.body)
                    output = json.loads(Payload['Message'])['parameters']['image_url']
                    if len(output)>0:
                        while True:
                            try:
                                image_bucket, image_key = get_bucket_and_key(output)
                                image_object = s3_resource.Object(image_bucket, image_key)
                                image = image_object.get()["Body"].read()
                                break
                            except Exception as e:
                                print(e)
                                continue
                    else:
                        print('missing output')
                        print('debug',Payload)
                        continue

                    print('time spent', j, time.perf_counter() - start_time)
                    time_lst.append(time.perf_counter() - start_time)
                    delete_message(message)
                break
            except Exception as e:
                print(e)
                continue
            
        

print('time lag log',str(time_lst))
print('time_mean',str(sum(time_lst)/len(time_lst)))
print('time_max',str(max(time_lst)))



load_test/multi-process-master.py
aws-samples/stable-diffusion-on-eks
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

import os
import io
import asyncio
import base64
import json
import time
import datetime
import requests
import boto3
import logging
import traceback
import uuid
import aioboto3
import difflib
import signal
from aiohttp_client_cache import CachedSession, CacheBackend
from botocore.exceptions import ClientError
from PIL import PngImagePlugin, Image
from requests.adapters import HTTPAdapter, Retry
from aws_xray_sdk.core import xray_recorder, patch_all
from aws_xray_sdk.core.models.trace_header import TraceHeader

patch_all()

logging.getLogger("aws_xray_sdk").setLevel(logging.ERROR)

aws_default_region = os.getenv("AWS_DEFAULT_REGION")
sqs_queue_url = os.getenv("SQS_QUEUE_URL")
sns_topic_arn = os.getenv("SNS_TOPIC_ARN")
s3_bucket = os.getenv("S3_BUCKET")
dynamic_sd_model = os.getenv("DYNAMIC_SD_MODEL")

current_model_name = ''
sqsRes = boto3.resource('sqs')
snsRes = boto3.resource('sns')
s3Res = boto3.resource('s3')
ab3_session = aioboto3.Session()

apiBaseUrl = "http://localhost:8080/sdapi/v1/"
apiClient = requests.Session()
retries = Retry(
    total=3,
    connect=100,
    backoff_factor=0.1,
    allowed_methods=["GET", "POST"])
apiClient.mount('http://', HTTPAdapter(max_retries=retries))
REQUESTS_TIMEOUT_SECONDS = 60

cache = CacheBackend(
    cache_name='memory-cache',
    expire_after=600
)

ALWAYSON_SCRIPTS_EXCLUDE_KEYS = ['task', 'id_task', 'uid',
                                 'sd_model_checkpoint', 'image_link', 'save_dir', 'sd_vae', 'override_settings']

shutdown = False


def main():
    # Initialization:
    # 1. Environment parameters;
    # 2. AWS services resources(sqs/sns/s3);
    # 3. SD API readiness check, current checkpoint cached;
    print_env()

    queue = sqsRes.Queue(sqs_queue_url)
    SQS_WAIT_TIME_SECONDS = 20
    topic = snsRes.Topic(sns_topic_arn)

    check_readiness()

    # main loop
    # 1. Pull msg from sqs;
    # 2. Translate parameteres;
    # 3. (opt)Switch model;
    # 4. (opt)Prepare inputs for image downloading and encoding;
    # 5. Call SD API;
    # 6. Prepare outputs for decoding, uploading and notifying;
    # 7. Delete msg;
    while True:
        if shutdown:
            print('Received SIGTERM, shutting down...')
            break

        received_messages = receive_messages(queue, 1, SQS_WAIT_TIME_SECONDS)

        for message in received_messages:
            with xray_recorder.in_segment('Queue-Agent') as segment:
                # Retrieve x-ray trace header from SQS message
                traceHeaderStr = message.attributes['AWSTraceHeader']
                sqsTraceHeader = TraceHeader.from_header_str(traceHeaderStr)
                # Update current segment to link with SQS
                segment.trace_id = sqsTraceHeader.root
                segment.parent_id = sqsTraceHeader.parent
                segment.sampled = sqsTraceHeader.sampled

                try:
                    snsPayload = json.loads(message.body)
                    payload = json.loads(snsPayload['Message'])
                    taskHeader = payload['alwayson_scripts']
                    taskType = taskHeader['task'] if 'task' in taskHeader else None
                    taskId = taskHeader['id_task'] if 'id_task' in taskHeader else None
                    folder = get_prefix(
                        payload['s3_output_path']) if 's3_output_path' in payload else None
                    print(
                        f"Start process {taskType} task with ID: {taskId}")

                    if dynamic_sd_model == 'true' and taskHeader['sd_model_checkpoint']:
                        switch_model(taskHeader['sd_model_checkpoint'])
                        print(f'Current model is: {current_model_name}.')

                    if taskType == 'text-to-image':
                        r = invoke_txt2img(payload, taskHeader)
                    elif taskType == 'image-to-image':
                        r = invoke_img2img(payload, taskHeader)
                    else:
                        raise RuntimeError(
                            f'Unsupported task type: {taskType}')

                    imgOutputs = post_invocations(folder, r, 80)

                except Exception as e:
                    content = json.dumps(failed(taskHeader, repr(e)))
                    traceback.print_exc()
                else:
                    content = json.dumps(succeed(imgOutputs, r, taskHeader))
                finally:
                    upload_outputs(content, folder)
                    publish_message(topic, content)
                    delete_message(message)
                    print(
                        f"End process {taskType} task with ID: {taskId}")


def print_env():
    print(f'AWS_DEFAULT_REGION={aws_default_region}')
    print(f'SQS_QUEUE_URL={sqs_queue_url}')
    print(f'SNS_TOPIC_ARN={sns_topic_arn}')
    print(f'S3_BUCKET={s3_bucket}')
    print(f'DYNAMIC_SD_MODEL={dynamic_sd_model}')


def check_readiness():
    while True:
        try:
            print('Checking service readiness...')
            # checking with options "sd_model_checkpoint" also for caching current model
            opts = invoke_get_options()
            print('Service is ready.')
            if "sd_model_checkpoint" in opts:
                global current_model_name
                current_model_name = opts['sd_model_checkpoint']
                print(f'Init model is: {current_model_name}.')
            break
        except Exception as e:
            print(repr(e))
            time.sleep(1)


def signalHandler(signum, frame):
    global shutdown
    shutdown = True


def get_time(f):

    def inner(*arg, **kwarg):
        s_time = time.time()
        res = f(*arg, **kwarg)
        e_time = time.time()
        print('Used: {:.4f} seconds on api: {}.'.format(
            e_time - s_time, arg[0]))
        return res
    return inner


def receive_messages(queue, max_number, wait_time):
    """
    Receive a batch of messages in a single request from an SQS queue.

    :param queue: The queue from which to receive messages.
    :param max_number: The maximum number of messages to receive. The actual number
                       of messages received might be less.
    :param wait_time: The maximum time to wait (in seconds) before returning. When
                      this number is greater than zero, long polling is used. This
                      can result in reduced costs and fewer false empty responses.
    :return: The list of Message objects received. These each contain the body
             of the message and metadata and custom attributes.
    """
    try:
        messages = queue.receive_messages(
            MaxNumberOfMessages=max_number,
            WaitTimeSeconds=wait_time,
            AttributeNames=['All'],
            MessageAttributeNames=['All']
        )
    except ClientError as error:
        traceback.print_exc()
        raise error
    else:
        return messages


def publish_message(topic, message):
    try:
        response = topic.publish(Message=message)
        message_id = response['MessageId']
    except ClientError as error:
        traceback.print_exc()
        raise error
    else:
        return message_id


def delete_message(message):
    """
    Delete a message from a queue. Clients must delete messages after they
    are received and processed to remove them from the queue.

    :param message: The message to delete. The message's queue URL is contained in
                    the message's metadata.
    :return: None
    """
    try:
        message.delete()
    except ClientError as error:
        traceback.print_exc()
        raise error


def get_bucket_and_key(s3uri):
    pos = s3uri.find('/', 5)
    bucket = s3uri[5: pos]
    key = s3uri[pos + 1:]
    return bucket, key


def get_prefix(path):
    pos = path.find('/')
    return path[pos + 1:]


def str_simularity(a, b):
    return difflib.SequenceMatcher(None, a, b).ratio()


def encode_to_base64(buffer):
    return str(base64.b64encode(buffer))[2:-1]


def decode_to_image(encoding):
    image = None
    try:
        image = Image.open(io.BytesIO(base64.b64decode(encoding)))
    except Exception as e:
        raise e
    return image


def export_pil_to_bytes(image, quality):
    with io.BytesIO() as output_bytes:

        use_metadata = False
        metadata = PngImagePlugin.PngInfo()
        for key, value in image.info.items():
            if isinstance(key, str) and isinstance(value, str):
                metadata.add_text(key, value)
                use_metadata = True
        image.save(output_bytes, format="PNG", pnginfo=(
            metadata if use_metadata else None), quality=quality if quality else 80)

        bytes_data = output_bytes.getvalue()

    return bytes_data


@xray_recorder.capture('text-to-image')
def invoke_txt2img(body, header):
    return do_invocations(apiBaseUrl+"txt2img", prepare_payload(body, header))


@xray_recorder.capture('image-to-image')
def invoke_img2img(body, header):
    return do_invocations(apiBaseUrl+"img2img", prepare_payload(body, header))


def invoke_set_options(options):
    return do_invocations(apiBaseUrl+"options", options)


def invoke_get_options():
    return do_invocations(apiBaseUrl+"options")


def invoke_get_model_names():
    return sorted([x["title"] for x in do_invocations(apiBaseUrl+"sd-models")])


def invoke_refresh_checkpoints():
    return do_invocations(apiBaseUrl+"refresh-checkpoints", {})


def switch_model(name, find_closest=True):
    global current_model_name
    # check current model
    if find_closest:
        name = name.lower()
        current = current_model_name.lower()

    if name in current:
        return current_model_name

    # refresh then check from model list
    invoke_refresh_checkpoints()
    models = invoke_get_model_names()
    found_model = None

    # exact matching
    if name in models:
        found_model = name
    # find closest
    elif find_closest:
        max_sim = 0.0
        max_model = None
        for model in models:
            sim = str_simularity(name, model.lower())
            if sim > max_sim:
                max_sim = sim
                max_model = model
        found_model = max_model

    if not found_model:
        raise RuntimeError(f'Model not found: {name}')
    elif found_model != current_model_name:
        options = {}
        options["sd_model_checkpoint"] = found_model
        invoke_set_options(options)
        current_model_name = found_model

    return current_model_name


def succeed(images, response, header):
    n_iter = response['parameters']['n_iter']
    batch_size = response['parameters']['batch_size']
    parameters = {}
    parameters['id_task'] = header['id_task']
    parameters['status'] = 1
    parameters['image_url'] = ','.join(
        images[: n_iter * batch_size])
    parameters['seed'] = ','.join(
        str(x) for x in json.loads(response['info'])['all_seeds'])
    parameters['error_msg'] = ''
    parameters['image_mask_url'] = ','.join(
        images[n_iter * batch_size:])
    return {
        'images': [''],
        'parameters': parameters,
        'info': ''
    }


def failed(header, message):
    parameters = {}
    parameters['id_task'] = header['id_task']
    parameters['status'] = 0
    parameters['image_url'] = ''
    parameters['seed'] = []
    parameters['error_msg'] = message
    parameters['image_mask_url'] = ''
    return {
        'images': [''],
        'parameters': parameters,
        'info': ''
    }


@get_time
def do_invocations(url, body=None):
    if body is None:
        response = apiClient.get(
            url=url, timeout=(1, REQUESTS_TIMEOUT_SECONDS))
    else:
        response = apiClient.post(
            url=url, json=body, timeout=(1, REQUESTS_TIMEOUT_SECONDS))
    response.raise_for_status()
    return response.json()


def post_invocations(folder, response, quality):
    defaultFolder = datetime.date.today().strftime("%Y-%m-%d")
    if not folder:
        folder = defaultFolder
    images = []
    results = []
    if "images" in response.keys():
        images = [export_pil_to_bytes(decode_to_image(i), quality)
                  for i in response["images"]]
    elif "image" in response.keys():
        images = [export_pil_to_bytes(
            decode_to_image(response["image"]), quality)]

    if len(images) > 0:
        loop = asyncio.get_event_loop()
        tasks = [loop.create_task(async_upload(i, folder)) for i in images]
        results = loop.run_until_complete(asyncio.gather(*tasks))

    return results


def upload_outputs(content, folder):
    try:
        defaultFolder = datetime.date.today().strftime("%Y-%m-%d")
        if not folder:
            folder = defaultFolder
        suffix = 'out'
        content_type = f'application/json'
        file_name = f"response-{uuid.uuid4()}"

        bucket = s3Res.Bucket(s3_bucket)
        bucket.put_object(
            Body=content, Key=f'{folder}/{file_name}.{suffix}', ContentType=content_type)
    except Exception as error:
        traceback.print_exc()
        raise error


async def async_get(url):
    try:
        if url.startswith("http://") or url.startswith("https://"):
            async with CachedSession(cache=cache) as session:
                async with session.get(url) as res:
                    res.raise_for_status()
                    # todo: need a counter to delete expired responses
                    # await session.delete_expired_responses()
                    # print(res.from_cache, res.created_at, res.expires, res.is_expired)
                    return await res.read()
        elif url.startswith("s3://"):
            bucket, key = get_bucket_and_key(url)
            async with ab3_session.resource("s3") as s3:
                obj = await s3.Object(bucket, key)
                res = await obj.get()
                return await res['Body'].read()
    except Exception as e:
        raise e


async def async_upload(object_bytes, folder, file_name=None, suffix=None):
    try:
        async with ab3_session.resource("s3") as s3:
            if suffix == 'out':
                content_type = f'application/json'
                if file_name is None:
                    file_name = f"response-{uuid.uuid4()}"
            else:
                suffix = 'png'
                content_type = f'image/{suffix}'
                if file_name is None:
                    file_name = datetime.datetime.now().strftime(
                        f"%Y%m%d%H%M%S-{uuid.uuid4()}")

            bucket = await s3.Bucket(s3_bucket)
            await bucket.put_object(
                Body=object_bytes, Key=f'{folder}/{file_name}.{suffix}', ContentType=content_type)

            return f's3://{s3_bucket}/{folder}/{file_name}.{suffix}'
    except Exception as e:
        raise e


async def async_publish_message(content):
    try:
        async with ab3_session.resource("sns") as sns:
            topic = await sns.Topic(sns_topic_arn)
            response = await topic.publish(Message=content)
            return response['MessageId']
    except Exception as e:
        raise e


def exclude_keys(dictionary, keys):
    key_set = set(dictionary.keys()) - set(keys)
    return {key: dictionary[key] for key in key_set}


def prepare_payload(body, header):
    try:
        urls = []
        offset = 0
        # img2img image link
        if 'image_link' in header:
            urls.extend(header['image_link'].split(','))
            offset = len(urls)
        # ControlNet image link
        if 'controlnet' in header:
            for x in header['controlnet']['args']:
                if 'image_link' in x:
                    urls.append(x['image_link'])
        # Generate payload including ControlNet units
        if len(urls) > 0:
            loop = asyncio.get_event_loop()
            tasks = [loop.create_task(async_get(u)) for u in urls]
            results = loop.run_until_complete(asyncio.gather(*tasks))
            if offset > 0:
                init_images = [encode_to_base64(x) for x in results[:offset]]
                body.update({"init_images": init_images})

            if 'controlnet' in header:
                for x in header['controlnet']['args']:
                    if 'image_link' in x:
                        x['input_image'] = encode_to_base64(results[offset])
                        offset += 1

        # dbt compatible for override_settings
        override_settings = {}
        if 'sd_vae' in header:
            override_settings.update({'sd_vae': header['sd_vae']})
        if 'override_settings' in header:
            override_settings.update(header['override_settings'])
        if override_settings:
            if 'override_settings' in body:
                body['override_settings'].update(override_settings)
            else:
                body.update({'override_settings': override_settings})

        # dbt compatible for alwayson_scripts
        body.update({'alwayson_scripts': exclude_keys(
            header, ALWAYSON_SCRIPTS_EXCLUDE_KEYS)})
    except Exception as e:
        raise e

    return body


if __name__ == '__main__':
    for sig in [signal.SIGINT, signal.SIGHUP, signal.SIGTERM]:
        signal.signal(sig, signalHandler)
    main()

src/backend/queue_agent/main.py
aws-samples/stable-diffusion-on-eks
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

import os
import boto3
import json
import traceback

sns_client = boto3.client('sns')


def lambda_handler(event, context):
    if event['httpMethod'] == 'POST':
        try:
            payload = json.loads(event['body'])

            if ('alwayson_scripts' not in payload or 'task' not in payload['alwayson_scripts'] or 'sd_model_checkpoint' not in payload['alwayson_scripts'] or 'id_task' not in payload['alwayson_scripts']):
                return {
                    'statusCode': 400,
                    'body': 'Incorrect payload structure'
                }

            task = payload['alwayson_scripts']['task']
            id_task = payload['alwayson_scripts']['id_task']
            sd_model_checkpoint = payload['alwayson_scripts']['sd_model_checkpoint']
            s3_output_path = f"{os.environ['S3_OUTPUT_BUCKET']}/{payload['alwayson_scripts']['id_task']}"

            payload['s3_output_path'] = s3_output_path

            print(event['headers'])
            print(event['queryStringParameters'])

            sns_client.publish(
                TargetArn=os.environ['SNS_TOPIC_ARN'],
                Message=json.dumps({"default": json.dumps(payload)}),
                MessageStructure='json',
                MessageAttributes={
                    'sd_model_checkpoint': {
                        'DataType': 'String',
                        'StringValue': sd_model_checkpoint
                    }
                }
            )

            return {
                'statusCode': 200,
                'body': json.dumps({
                    "id_task": id_task,
                    "task": task,
                    "sd_model_checkpoint": sd_model_checkpoint,
                    "output_location": f"s3://{s3_output_path}"
                })
            }

        except Exception as e:
            traceback.print_exc()
            return {
                'statusCode': 400,
                'body': str(e)
            }
    else:
        return {
            'statusCode': 400,
            'body': "Unsupported HTTP method"
        }

src/frontend/input_function/app.py
aws-samples/stable-diffusion-on-eks
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

import boto3, os

def get_ebs_volume_id(instance_id):
    ec2 = boto3.client('ec2')
    try:
        response = ec2.describe_instances(InstanceIds=[instance_id])
        volume_id = response['Reservations'][0]['Instances'][0]['BlockDeviceMappings'][-1]['Ebs']['VolumeId']
        return volume_id
    except Exception as e:
        print("Error:", e)
        return None

def get_instance_tag(instance, key):
    for tag in instance.get('Tags', []):
        if tag['Key'] == key:
            return tag['Value']
    return None

def modify_ebs_throughput_and_iops(volume_id, throughput, iops):
    ec2 = boto3.client('ec2')
    try:
        response = ec2.modify_volume(VolumeId=volume_id, Throughput=throughput, Iops=iops)
        print("Successfully modified EBS throughput of " + volume_id)
    except Exception as e:
        print("Error:", e)

# Entrypoint
def lambda_handler(event, context):
    ec2 = boto3.client('ec2')
    instance_id = event['detail']['instance-id']
    print (f"Process with instance " + instance_id)
    target_ec2_tag_key = os.environ['TARGET_EC2_TAG_KEY']
    target_ec2_tag_value = os.environ['TARGET_EC2_TAG_VALUE']
    # Get the instance tag
    response = ec2.describe_instances(InstanceIds=[instance_id])
    instance = response['Reservations'][0]['Instances'][0]
    instance_tag_value = get_instance_tag(instance, target_ec2_tag_key)
    volume_id = get_ebs_volume_id(instance_id)
    if instance_tag_value:
        # Determine if EBS throughput needs to be modified based on instance name
        if target_ec2_tag_value in instance_tag_value:
            print (f"Got matching tag from " + instance_id + " ,processing...")
            throughput_value = int(os.environ['THROUGHPUT_VALUE'])  # Modify throughput
            IOPS_value = int(os.environ['IOPS_VALUE']) # Modify IOPS
            modify_ebs_throughput_and_iops(volume_id, throughput_value, IOPS_value)
        else:
            print (f"Skipped " + instance_id)

src/tools/ebs_throughput_tuner/app.py
aws-samples/stable-diffusion-on-eks