{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b2402bed",
   "metadata": {},
   "source": [
    "# 1. Hosting GenAI Models\n",
    "This Notebook is used to host 2 models to SageMaker Endpoints. The first model is Segment Anything Model (SAM) and the second model is the InPainting model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b8170ee",
   "metadata": {},
   "source": [
    "## 1.1 Install Dependencies and Import Libraries\n",
    "Here we are using `SageMaker, Boto3` libraries for Python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a641a723",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Import libraries\n",
    "import sagemaker, subprocess, boto3\n",
    "from datetime import datetime\n",
    "from sagemaker import s3, get_execution_role\n",
    "from sagemaker.pytorch import PyTorchModel\n",
    "from sagemaker.serializers import JSONSerializer\n",
    "from sagemaker.deserializers import JSONDeserializer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98e6da1e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get the right credentials, role and client for SageMaker\n",
    "sm_client = boto3.client(service_name=\"sagemaker\")\n",
    "role = get_execution_role()\n",
    "print(f'Role: {role}')\n",
    "\n",
    "INSTANCE_TYPE = 'ml.g5.2xlarge'\n",
    "\n",
    "bashCommand = \"tar -cpzf  code.tar.gz code/\"\n",
    "process = subprocess.Popen(bashCommand.split(), stdout=subprocess.PIPE)\n",
    "output, error = process.communicate()\n",
    "\n",
    "s3_client = boto3.client('s3')\n",
    "s3_resource = boto3.resource('s3')\n",
    "sts = boto3.client('sts')\n",
    "AWS_ACCOUNT_ID = sts.get_caller_identity()[\"Account\"]\n",
    "REGION = s3_client.meta.region_name\n",
    "\n",
    "bucket = f'gai-inpainting-s3'\n",
    "response = s3_client.list_buckets()\n",
    "for bucket in response['Buckets']:\n",
    "    if 'gai-inpainting-s3' in bucket[\"Name\"]:\n",
    "        bucket = bucket[\"Name\"]\n",
    "        break\n",
    "\n",
    "print(f'Bucket: {bucket}')\n",
    "sess = sagemaker.Session(default_bucket=bucket.split('s3://')[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6065a96",
   "metadata": {},
   "source": [
    "## 1.2 Model 1: SAM\n",
    "\n",
    "### Create and Host SAM Model\n",
    "In this section, we will demonstrate how to create and host the endpoint for SAM model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b205466",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "SAM_ENDPOINT_NAME = 'sam-pytorch-' + str(datetime.utcnow().strftime('%Y-%m-%d-%H-%M-%S-%f'))\n",
    "%store SAM_ENDPOINT_NAME\n",
    "\n",
    "prefix_sam = \"SAM/demo-custom-endpoint\"\n",
    "\n",
    "model_data_sam = s3.S3Uploader.upload(\"code.tar.gz\", f's3://{bucket}/{prefix_sam}')\n",
    "print(f'Model Data: {model_data_sam}')\n",
    "\n",
    "model_sam = PyTorchModel(entry_point='inference_sam.py',\n",
    "                     model_data=model_data_sam, \n",
    "                     framework_version='1.12', \n",
    "                     py_version='py38',\n",
    "                     role=role,\n",
    "                     env={'TS_MAX_RESPONSE_SIZE':'2000000000', 'SAGEMAKER_MODEL_SERVER_TIMEOUT' : '300'},\n",
    "                     sagemaker_session=sess,\n",
    "                     name='model-'+SAM_ENDPOINT_NAME)\n",
    "\n",
    "print(f'SAM Endpoint Name: {SAM_ENDPOINT_NAME}')\n",
    "\n",
    "predictor_sam = model_sam.deploy(initial_instance_count=1, \n",
    "                         instance_type=INSTANCE_TYPE,\n",
    "                         deserializers=JSONDeserializer(),\n",
    "                         endpoint_name=SAM_ENDPOINT_NAME)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15c3beda",
   "metadata": {},
   "source": [
    "## 1.3 Model 2: InPainting\n",
    "\n",
    "### Create and Host InPainting Model\n",
    "In this section, we will demonstrate how to create and host the endpoint for InPainting model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "202068c9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "INPAINTING_ENDPOINT_NAME = 'inpainting-pytorch-' + str(datetime.utcnow().strftime('%Y-%m-%d-%H-%M-%S-%f'))\n",
    "%store INPAINTING_ENDPOINT_NAME\n",
    "\n",
    "prefix_inpainting = \"InPainting/demo-custom-endpoint\"\n",
    "\n",
    "model_data_inpainting = s3.S3Uploader.upload(\"code.tar.gz\", f\"s3://{bucket}/{prefix_inpainting}\")\n",
    "print(f'Model Data: {model_data_inpainting}')\n",
    "\n",
    "model_inpainting = PyTorchModel(entry_point='inference_inpainting.py',\n",
    "                     model_data=model_data_inpainting, \n",
    "                     framework_version='1.12', \n",
    "                     py_version='py38',\n",
    "                     role=role,\n",
    "                     env={'TS_MAX_RESPONSE_SIZE':'2000000000', 'SAGEMAKER_MODEL_SERVER_TIMEOUT' : '300'},\n",
    "                     sagemaker_session=sess,\n",
    "                     name='model-'+INPAINTING_ENDPOINT_NAME)\n",
    "\n",
    "print(f'InPainting Endpoint Name: {INPAINTING_ENDPOINT_NAME}')\n",
    "\n",
    "predictor_inpainting = model_inpainting.deploy(initial_instance_count=1, \n",
    "                         instance_type=INSTANCE_TYPE,\n",
    "                         serializer=JSONSerializer(),\n",
    "                         deserializers=JSONDeserializer(),\n",
    "                         endpoint_name=INPAINTING_ENDPOINT_NAME,\n",
    "                         volume_size=128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9726d58f-159e-45dc-9fdc-d878ab059b06",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_python3"
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}