# SageMaker JumpStart provides APIs as part of SageMaker SDK that allow you to deploy and fine-tune models in network isolation using scripts that SageMaker maintains.

import boto3
import base64
import io
import json
from PIL import Image
from sagemaker.jumpstart import utils
from sagemaker.jumpstart.model import JumpStartModel


model_id = "model-upscaling-stabilityai-stable-diffusion-x4-upscaler-fp16"
key_prefix = "stabilityai-metadata/assets"
low_res_img_file_name = "low_res_cat.jpg"
bucket_name = utils.get_jumpstart_content_bucket(boto3.Session().region_name)

s3 = boto3.client("s3")
s3.download_file(bucket_name, f"{key_prefix}/{low_res_img_file_name}", low_res_img_file_name)

with open(low_res_img_file_name,'rb') as f: low_res_image_bytes = f.read()
encoded_image = base64.b64encode(bytearray(low_res_image_bytes)).decode()

endpoint_input = { "prompt":"a cat", "image": encoded_image,  "num_inference_steps":50, "guidance_scale":7.5}

model = JumpStartModel(model_id=model_id)
predictor = model.deploy()

response = predictor.predict(endpoint_input)
decoded_image = base64.b64decode(response["generated_images"][0])
image = Image.open(io.BytesIO(decoded_image))
image.show()