{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8af3794b",
   "metadata": {},
   "source": [
    "# Image Captioning: Run EXAONE Atelier from AWS Marketplace"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7ea02ec",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "In this demo notebook, we demonstrate how to use the SageMaker Python SDK to deploy EXAONE Atelier Image Captioning Model from AWS Marketplace.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35642ab2",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Setup\n",
    "\n",
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88e9674d-e2b9-4f45-8cc1-3b9fab569e58",
   "metadata": {},
   "source": [
    "\n",
    "To subscribe to the model package:\n",
    "\n",
    "1. Open the model package **listing page**\n",
    "1. On the AWS Marketplace listing, click on the **Continue to Subscribe** button.\n",
    "1. On the **Subscribe to this software** page, review and click on **\"Accept Offer\"** if you and your organization agrees with EULA, pricing, and support terms.\n",
    "1. Once you click on **Continue to configuration** button and then choose a **region**, you will see a **Product Arn** displayed. This is the model package ARN that you need to specify while creating a deployable model using Boto3. Copy the ARN corresponding to your region and specify the same in the following cell.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b55e677-3429-4668-b100-bd63d2a4c401",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%pip install --upgrade --quiet sagemaker\n",
    "%pip install --upgrade --quiet sagemaker accelerate datasets tritonclient[all]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "014424a8-7f8f-46a7-8963-2c3d454878b8",
   "metadata": {
    "jumpStartAlterations": [
     "modelIdVersion"
    ],
    "tags": []
   },
   "outputs": [],
   "source": [
    "import boto3\n",
    "\n",
    "#model_package = \"exaone-atelier-i2t-limited-85f3f0d181593a10b7aef9bea522a333\" # EXAONE Atelier - Image to Text - Limited\n",
    "model_package = \"exaone-atelier-i2t-76c77246a8343a23a36b2ce80c06f4f6\" # EXAONE Atelier - Image to Text\n",
    "\n",
    "model_package_map = {\n",
    "    \"us-east-1\": f\"arn:aws:sagemaker:us-east-1:865070037744:model-package/{model_package}\",\n",
    "    \"us-east-2\": f\"arn:aws:sagemaker:us-east-2:057799348421:model-package/{model_package}\",\n",
    "    \"us-west-1\": f\"arn:aws:sagemaker:us-west-1:382657785993:model-package/{model_package}\",\n",
    "    \"us-west-2\": f\"arn:aws:sagemaker:us-west-2:594846645681:model-package/{model_package}\",\n",
    "    \"ca-central-1\": f\"arn:aws:sagemaker:ca-central-1:470592106596:model-package/{model_package}\",\n",
    "    \"eu-central-1\": f\"arn:aws:sagemaker:eu-central-1:446921602837:model-package/{model_package}\",\n",
    "    \"eu-west-1\": f\"arn:aws:sagemaker:eu-west-1:985815980388:model-package/{model_package}\",\n",
    "    \"eu-west-2\": f\"arn:aws:sagemaker:eu-west-2:856760150666:model-package/{model_package}\",\n",
    "    \"eu-west-3\": f\"arn:aws:sagemaker:eu-west-3:843114510376:model-package/{model_package}\",\n",
    "    \"eu-north-1\": f\"arn:aws:sagemaker:eu-north-1:136758871317:model-package/{model_package}\",\n",
    "    \"ap-southeast-1\": f\"arn:aws:sagemaker:ap-southeast-1:192199979996:model-package/{model_package}\",\n",
    "    \"ap-southeast-2\": f\"arn:aws:sagemaker:ap-southeast-2:666831318237:model-package/{model_package}\",\n",
    "    \"ap-northeast-2\": f\"arn:aws:sagemaker:ap-northeast-2:745090734665:model-package/{model_package}\",\n",
    "    \"ap-northeast-1\": f\"arn:aws:sagemaker:ap-northeast-1:977537786026:model-package/{model_package}\",\n",
    "    \"ap-south-1\": f\"arn:aws:sagemaker:ap-south-1:077584701553:model-package/{model_package}\",\n",
    "    \"sa-east-1\": f\"arn:aws:sagemaker:sa-east-1:270155090741:model-package/{model_package}\",\n",
    "}\n",
    "region = boto3.Session().region_name\n",
    "if region not in model_package_map.keys():\n",
    "    raise Exception(f\"Current boto3 session region {region} is not supported.\")\n",
    "\n",
    "model_package_arn = model_package_map[region]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01729938-4f1c-4ddf-bef4-63a314909eb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import strftime\n",
    "create_date = strftime(\"%m%d-%H%M%s\") \n",
    "model_name = f'exaone-atelier-i2t-{create_date}'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fbfbdc2",
   "metadata": {},
   "source": [
    "### Changing instance type\n",
    "---\n",
    "\n",
    "\n",
    "Model are supported on the following instance types:\n",
    " - Exaone Atelier - Image to Text - Limited: `ml.g5.12xlarge`\n",
    " - Exaone Atelier - Image to Text: `ml.g5.xlarge`, `ml.g5.12xlarge`, `ml.g5.48xlarge`, `ml.p4d.24xlarge`\n",
    "\n",
    "Exaone Atelier - Image to Text - Limited offers 5 days of free-trial.\n",
    "\n",
    "Below are average inference times to process a single image request for each supported instance type. The actual response time may differ due to various reasons including network condition.\n",
    "\n",
    "\n",
    "|Instance Type|Inference Time (sec)|\n",
    "|---|---|\n",
    "|ml.p4d.24xlarge|1.85|\n",
    "|ml.g5.48xlarge|3.68|\n",
    "|ml.g5.12xlarge|6.59|\n",
    "|ml.g5.xlarge|25.85|\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e52afae-868d-4736-881f-7180f393003a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from sagemaker import ModelPackage\n",
    "from sagemaker import get_execution_role\n",
    "import sagemaker\n",
    "\n",
    "role = get_execution_role()\n",
    "sagemaker_session = sagemaker.Session()\n",
    "\n",
    "\n",
    "model = ModelPackage(\n",
    "    role=role,\n",
    "    model_package_arn=model_package_arn,\n",
    "    sagemaker_session=sagemaker_session,\n",
    ")\n",
    "\n",
    "model.deploy(\n",
    "    initial_instance_count=1,\n",
    "    instance_type='ml.g5.12xlarge', # choose preferred instance type\n",
    "    endpoint_name=model_name,\n",
    "    container_startup_health_check_timeout=3600,\n",
    ")\n",
    "model.endpoint_name"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ef7207e-01ba-4ac2-b4a9-c8f6f0e1c498",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Invoke the Endpoint\n",
    "---\n",
    "\n",
    "In this notebook, we show how you can encode an image into bytes, send a request, and decode the response.\n",
    "\n",
    "***\n",
    "### Notes\n",
    "- This model receives an image in byte and returns 4 captions in total along with the confidence score for each caption.\n",
    "- The confidence score indicates how sure the model is that each caption well describes the given image.\n",
    "***\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71b1a232",
   "metadata": {},
   "outputs": [],
   "source": [
    "import base64\n",
    "from PIL import Image\n",
    "from io import BytesIO\n",
    "import numpy as np\n",
    "import tritonclient.http as httpclient\n",
    "import requests\n",
    "import json\n",
    "import boto3\n",
    "\n",
    "smr_client = boto3.client(\"sagemaker-runtime\")\n",
    "\n",
    "def encode_image(image):\n",
    "    buffer = BytesIO()\n",
    "    image.save(buffer, format=\"JPEG\")\n",
    "    img_str = base64.b64encode(buffer.getvalue())\n",
    "    return img_str\n",
    "\n",
    "def get_sample_binary(payload):\n",
    "\n",
    "    inputs = []\n",
    "    outputs = []\n",
    "    for idx, dic in enumerate(payload[\"inputs\"]):\n",
    "        input_name = dic[\"name\"]\n",
    "        input_value = dic[\"data\"][0]\n",
    "\n",
    "        input_value = np.array([input_value.encode('utf-8')], dtype=np.object_)\n",
    "\n",
    "        input_value = np.expand_dims(input_value, axis=0)\n",
    "        inputs.append(httpclient.InferInput(input_name, [1, 1], \"BYTES\"))\n",
    "        inputs[idx].set_data_from_numpy(input_value)\n",
    "\n",
    "    outputs.append(httpclient.InferRequestedOutput(\"generated_caption\", binary_data=True))\n",
    "\n",
    "    request_body, header_length = httpclient.InferenceServerClient.generate_request_body(\n",
    "        inputs, outputs=outputs\n",
    "    )\n",
    "    return request_body, header_length\n",
    "\n",
    "\n",
    "def invoke_endpoint(endpoint_name, payload):\n",
    "    import re\n",
    "    request_body, header_length = get_sample_binary(payload)\n",
    "    response = smr_client.invoke_endpoint(\n",
    "        EndpointName=endpoint_name,\n",
    "        ContentType=\"application/vnd.sagemaker-triton.binary+json;json-header-size={}\".format(\n",
    "            header_length\n",
    "        ),\n",
    "        Body=request_body\n",
    "    )\n",
    "    data = response[\"Body\"].read()\n",
    "    ptn = re.compile(rb'\\{\"binary_data_size\":[0-9]*\\}')\n",
    "    match = json.loads(ptn.search(data).group().decode('utf-8'))\n",
    "    binary_data_size = match['binary_data_size']\n",
    "    binary_data = data[len(data)-binary_data_size+1:]\n",
    "    binary_data = binary_data.replace(b'\\x00', b'')\n",
    "    binary_data = binary_data.replace(b'\\x01', b'').decode('utf-8')    \n",
    "\n",
    "    return eval(binary_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c53a7e37",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Prepare image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f874bb7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_image(image, keep_ratio=False):\n",
    "    # By default, EXAONE Atelier Image to Text internally center crops the image \n",
    "    # by min(width, height). Find your best option between keep_ratio=True or keep_ratio=False.\n",
    "    width, height = image.size\n",
    "    if keep_ratio:\n",
    "        if width < 256 or height < 256:\n",
    "            if width > height:\n",
    "                ratio = width / height        \n",
    "                image = image.resize((int( 256* ratio), 256)) \n",
    "            else:\n",
    "                ratio = height / width\n",
    "                image = image.resize((256, int( 256 * ratio)))               \n",
    "        elif width > 4096 or height > 4096:       \n",
    "            if width > height:\n",
    "                ratio = height / width     \n",
    "                image = image.resize((4096, int( 4096 * ratio)))\n",
    "            else:\n",
    "                ratio = width / height\n",
    "                image = image.resize((int( 4096 * ratio), 4096))\n",
    "    else:\n",
    "        if width < 256 or height < 256:       \n",
    "            image = image.resize((256, 256)) \n",
    "        elif width > 4096 or height > 4096:       \n",
    "            image = image.resize((4096, 4096)) \n",
    "        else:\n",
    "            resize_dim = min(width, height)\n",
    "            image = image.resize((resize_dim, resize_dim))  \n",
    "    return image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9085e0f1-bbf2-4c36-8eb3-f9dd8d59f0dd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def generate_input(image):\n",
    "    input_image = encode_image(image)\n",
    "    inputs = dict(\n",
    "        image=input_image,\n",
    "    )\n",
    "\n",
    "    payload = {\n",
    "        \"inputs\": [\n",
    "            {\"name\": name, \"shape\": [1, -1], \"datatype\": \"BYTES\", \"data\": [data.decode('utf8')]}\n",
    "            for name, data in inputs.items()\n",
    "        ]\n",
    "    }\n",
    "    return payload"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9d37ede",
   "metadata": {},
   "source": [
    "### Example 1 (single image from web)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ce614f6-ca67-44f6-ae1e-e74fa3cbeffc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "url = 'https://github.com/LGAI-Research/EXAONE-Atelier/blob/main/images/example_1.png?raw=true'\n",
    "image = Image.open(requests.get(url, stream=True).raw).convert('RGB')  \n",
    "# prepare image to avoid image size out of range error (256 x 256 ~ 4096 x 4096)\n",
    "input_image = prepare_image(image)\n",
    "display(input_image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cbde5e7-1068-41f9-999a-70ef04e1cbbb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "endpoint_name=model.endpoint_name\n",
    "start_time = time.time()\n",
    "captions = invoke_endpoint(endpoint_name, generate_input(input_image))\n",
    "print('{:.2f} sec'.format(time.time()-start_time))\n",
    "print(captions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd5db6b6",
   "metadata": {},
   "source": [
    "### Example 2 (batch inference on images in local directory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9059366",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import time\n",
    "\n",
    "image_dir = 'image_data/'\n",
    "\n",
    "path = Path(image_dir)\n",
    "\n",
    "image_files = [\n",
    "*path.glob('**/*.png'), *path.glob('**/*.jpg'),\n",
    "*path.glob('**/*.jpeg')\n",
    "    ]\n",
    "\n",
    "endpoint_name=model.endpoint_name\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09455e42",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, file in enumerate(image_files): \n",
    "    image = Image.open(file).convert('RGB')\n",
    "    image = prepare_image(image)\n",
    "    display(image)\n",
    "    input_image = encode_image(image)\n",
    "        \n",
    "    start_time = time.time()\n",
    "    captions = invoke_endpoint(endpoint_name, generate_input(input_image))\n",
    "    print('{:.2f} sec'.format(time.time()-start_time))\n",
    "    print(captions)    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e062d29",
   "metadata": {},
   "source": [
    "## Clean Up the Endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24cc5560",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Delete the SageMaker endpoint\n",
    "model.sagemaker_session.delete_endpoint(model.endpoint_name)\n",
    "model.sagemaker_session.delete_endpoint_config(model.endpoint_name)\n",
    "model.delete_model()"
   ]
  }
 ],
 "metadata": {
  "instance_type": "ml.t3.medium",
  "kernelspec": {
   "display_name": "conda_pytorch_p310",
   "language": "python",
   "name": "conda_pytorch_p310"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}