test_stable_image_upscale.py

import base64
import json

import requests
from PIL import Image
from os import getenv

SSC_HOSTING_PORT = getenv("SSC_HOSTING_PORT", "8080")

# Set up the request details
url = f"http://localhost:{SSC_HOSTING_PORT}/invocations"

# define image assets
init_image = "assets/dapper_cat.jpg"

b64_encoded_image = base64.b64encode(open(init_image, "rb").read()).decode("utf-8")

workflows = [
    'upscale-fast',
    'upscale-conservative',
    'upscale-creative'
]

workflow_request_dict = {
    'upscale-fast': {
            "image":  b64_encoded_image,
            "output_format": "png",
        },
    'upscale-conservative': {
            "image": b64_encoded_image,
            "prompt": "a handsome cat",
            "negative_prompt": "realistic",
            "seed": 123,
            "creativity": 0.5,
            "output_format": "png",
        },
    'upscale-creative': {
            "image": b64_encoded_image,
            "prompt": "a distinguished jaguar",
            "negative_prompt": "cartoon",
            "creativity": 0.35,
            "seed": 123,
            "output_format": "png",
            "style_preset": "modeling-compound"
        }
}

workflow_request_tuples = [(workflow, workflow_request_dict[workflow]) 
                          for workflow in workflows 
                          if workflow in workflow_request_dict]

for workflow, request_dict in workflow_request_tuples:
    assert type(workflow) == str
    assert type(request_dict) == dict
    custom_attributes = {"model": workflow}
    json_attributes = json.dumps(custom_attributes)
    encoded_attributes = base64.b64encode(json_attributes.encode('utf-8')).decode('utf-8')


    headers = {
        "Content-Type": "application/json",
        "accept": "application/json",
        "x-amzn-sagemaker-custom-attributes": encoded_attributes  
    }

    # Send the request
    print(f"Sending request for workflow: {workflow}")
    import time
    t0 = time.time()
    response = requests.post(url, headers=headers, data=json.dumps(request_dict))
    print(f"\t\t{workflow} Response status code: {response.status_code}\n")
    # Check if the response is successful
    if response.status_code == 200:
        t1 = time.time()
        print(f"latency = {t1-t0}")
        # Parse the JSON response

        response_json = response.json()
        # Extract the first element from the "images" list
        image_base64 = response_json['body']['images'][0]
        
        if image_base64:
            # Decode the base64 image data
            image_data = base64.b64decode(image_base64)
            
            # Save the decoded image as a JPEG file
            fname = f"output_{workflow}.jpeg"
            with open(fname, "wb") as image_file:
                image_file.write(image_data)
            
            print(f"Image saved as {fname}")
        else:
            print("No image data found in the response.")
    else:
        
        print(f"Request failed with status code {response.status_code}: {response.text}")