{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Create Stable Diffusion endpoint from custom pix2pix model\n",
    "look for detailed instruction here: [Huggingface_deploy_instructpix2pix](https://github.com/aws/amazon-sagemaker-examples/tree/main/advanced_functionality/huggingface_deploy_instructpix2pix)\n",
    "\n",
    "- **Step 1** - Create environment for custom model compilation\n",
    "- **Step 2** - Create endpoint files\n",
    "- **Step 3** - Pack the model and upload to S3\n",
    "- **Step 4** - Create sagemaker model and run inference endpoint\n",
    "- **Step 5** - Test image generation\n",
    "- **Step 6** - clean up Sagemaker endpoints to avoid unexpected charges\n",
    "\n",
    "## Endpoint parameters:\n",
    "**Input:**\n",
    "    data structure with:\n",
    "      - `prompt` - positive text prompt\n",
    "      - `image_url` - URL to image\n",
    "**Output:**\n",
    "      - generated image as string.\n"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "45ce2bc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Previous session cleean up\n",
    "#!rm -rf ./efs"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Step 1 - create environment for custom model compilation\n",
    "\n",
    "Tip: you need to create S3 bucket to upload compiled model to Sagemaker\n",
    "Put the name of this bucket to `s3_path_to_models`.\n"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8345ba24",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "local_package_dir = \"./efs/pix2pix_model\"  # Directory where local files created\n",
    "local_model_upload = os.path.join (local_package_dir, \"model\") # directory where uploaded file structure created\n",
    "# please create private S3 bucket for compiling model and put it's name here\n",
    "s3_path_to_models = \"s3://miro-integration-ml-model-store\"\n",
    "# create and archive model file\n",
    "my_model_name = \"pix2pix-model-0001.tar.gz\"\n",
    "endpoint_name = \"huggingface-example-model-pix2pix-demo-on-miro-1\"\n",
    "\n",
    "# create dirs\n",
    "!mkdir ./efs\n",
    "!mkdir $local_package_dir\n",
    "!mkdir $local_model_upload\n",
    "!mkdir $local_model_upload/code/"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Step 2 - create endpoint files\n",
    "- `requirements.txt` - custom requirements.txt\n",
    "- `inference.py` - custom inference"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "66930bcf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing ./efs/pix2pix_model/model/code/requirements.txt\n"
     ]
    }
   ],
   "source": [
    "%%writefile $local_model_upload/code/requirements.txt\n",
    "pillow\n",
    "diffusers==0.12.1\n",
    "transformers==4.26.1\n",
    "accelerate==0.16.0\n",
    "safetensors==0.2.8\n",
    "scipy==1.7.3\n",
    "--extra-index-url https://download.pytorch.org/whl/cu113\n",
    "torch==1.11.0+cu113\n",
    "--extra-index-url https://download.pytorch.org/whl/cu113\n",
    "torchaudio==0.11.0+cu113\n",
    "--extra-index-url https://download.pytorch.org/whl/cu113\n",
    "torchvision==0.12.0+cu113\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3439c3dd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing ./efs/pix2pix_model/model/code/inference.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile $local_model_upload/code/inference.py\n",
    "\n",
    "import os, json, random\n",
    "import pkg_resources\n",
    "from PIL import Image, ImageOps\n",
    "from io import BytesIO\n",
    "import base64, requests\n",
    "import torch\n",
    "\n",
    "\n",
    "# Initializations section\n",
    "from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler\n",
    "\n",
    "\n",
    "# Helpers: All helpers going here\n",
    "\n",
    "# Model\n",
    "# model_fn(model_dir) overrides the default method for loading a model. \n",
    "# The return value model will be used in thepredict_fn for predictions.\n",
    "#     model_dir is the the path to your unzipped model.tar.gz\n",
    "def model_fn(model_dir):\n",
    "    # Log start dir\n",
    "    # Load and return model from dir\n",
    "    # \n",
    "    print (\"----- Custom pix2pix inference on line: model load. Let's look around first.\", )\n",
    "    # Use the DPMSolverMultistepScheduler (DPM-Solver++) scheduler here instead\n",
    "    model_id = \"timbrooks/instruct-pix2pix\"\n",
    "    pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, torch_dtype=torch.float16, safety_checker=None)\n",
    "    pipe.to(\"cuda\")\n",
    "    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)\n",
    "\n",
    "    # You can uncomment this to look installed python packages in log\n",
    "    #installed_packages = pkg_resources.working_set\n",
    "    #installed_packages_list = sorted([\"%s==%s\" % (i.key, i.version)\n",
    "    #   for i in installed_packages])\n",
    "    #print (\"Checking installed packages... \\nPackages list: \", installed_packages_list)\n",
    "    \n",
    "    \n",
    "    print (\"----- Custom model found in my model directory: \", model_dir)\n",
    "    \n",
    "    print (\"----- Custom model loaded, see you on inference side ----\")\n",
    "    return pipe\n",
    "\n",
    "# Data-pre processing\n",
    "# The return value data will be used in predict_fn for predictions. The inputs are:\n",
    "#    input_data is the raw body of your request.\n",
    "#    content_type is the content type from the request header.\n",
    "#\n",
    "# def input_fn(input_data, content_type):\n",
    "    \n",
    "# Image download helper\n",
    "def download_image(url):\n",
    "    image = Image.open(requests.get(url, stream=True).raw)\n",
    "    image = ImageOps.exif_transpose(image)\n",
    "    image = image.convert(\"RGB\")\n",
    "    return image\n",
    "\n",
    "\n",
    "# Predictor\n",
    "# predict_fn(processed_data, model) overrides the default method for predictions. \n",
    "# The return value predictions will be used in output_fn.\n",
    "#\n",
    "#    model returned value from model_fn methond\n",
    "#    processed_data returned value from input_fn method\n",
    "#\n",
    "\n",
    "def predict_fn(data, model): # predict_fn(processed_data, model) # if you have several assets\n",
    "    # destruct model and tokenizer\n",
    "    # \n",
    "    \n",
    "    print (\"----- Custom inference online: predictor\")\n",
    "    print (\"Request: \", data)\n",
    "    prompt = data[\"prompt\"]\n",
    "    img_url = data[\"image_url\"]\n",
    "    \n",
    "    image = download_image(img_url)\n",
    "    image = model(prompt, image=image, num_inference_steps=10, image_guidance_scale=1).images[0]\n",
    "    bytes_io = BytesIO()\n",
    "    image.save(bytes_io, \"JPEG\")\n",
    "\n",
    "        \n",
    "    result_string = base64.encodebytes(bytes_io.getvalue()).decode(\"utf-8\")\n",
    "    \n",
    "    return result_string\n",
    "\n",
    "# Data post-processing\n",
    "# The return value result will be the response to your request (e.g.JSON). The inputs are:\n",
    "#\n",
    "#    predictions is the result from predict_fn.\n",
    "#    accept is the return accept type from the HTTP Request, e.g. application/json.\n",
    "#\n",
    "#"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Step 3 - pack the model and upload to S3"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b11e782d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./\n",
      "./code/\n",
      "./code/requirements.txt\n",
      "./code/inference.py\n",
      "upload: efs/pix2pix_model/pix2pix-model-0001.tar.gz to s3://miro-integration-ml-model-store/pix2pix-model-0001.tar.gz\n",
      "CPU times: user 20.6 ms, sys: 7.84 ms, total: 28.5 ms\n",
      "Wall time: 1.2 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# Archive model files & prepare to upload\n",
    "!tar -cvzf $local_package_dir/$my_model_name --exclude=\".ipynb_*\" -C $local_model_upload .\n",
    "\n",
    "# Upload our model to S3 bucket\n",
    "import os\n",
    "\n",
    "full_model_path = os.path.join(s3_path_to_models, my_model_name)\n",
    "!aws s3 cp $local_package_dir/$my_model_name $full_model_path"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "# --- From here we create SageMaker model and inference endpoint ---\n",
    "## Step 4 - Create sagemaker model and run inference endpoint"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6489513c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model created:  <sagemaker.huggingface.model.HuggingFaceModel object at 0x7fb33d02b460>\n",
      "-----------!Endpoint name:  huggingface-example-model-pix2pix-demo-on-miro-1\n",
      "CPU times: user 1.34 s, sys: 160 ms, total: 1.5 s\n",
      "Wall time: 6min 4s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "import sagemaker, boto3, os\n",
    "from sagemaker.huggingface.model import HuggingFaceModel\n",
    "\n",
    "full_model_path = os.path.join(s3_path_to_models, my_model_name)\n",
    "\n",
    "# prepare model execution role\n",
    "try:\n",
    "    role = sagemaker.get_execution_role()\n",
    "except ValueError:\n",
    "    iam = boto3.client('iam')\n",
    "    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']\n",
    "\n",
    "# create model\n",
    "# create Hugging Face Model Class\n",
    "huggingface_model = HuggingFaceModel(\n",
    "   model_data=full_model_path,       # path to your model and script\n",
    "   role=role,                    # iam role with permissions to create an Endpoint\n",
    "   transformers_version=\"4.12\",  # transformers version used\n",
    "   pytorch_version=\"1.9\",        # pytorch version used\n",
    "   py_version='py38',            # python version used\n",
    "   name=endpoint_name\n",
    ")\n",
    "\n",
    "print (\"Model created: \", huggingface_model)\n",
    "\n",
    "# create inference endpoint\n",
    "# deploy the endpoint endpoint\n",
    "# Inference types: https://aws.amazon.com/sagemaker/pricing/\n",
    "instance_type = \"ml.g4dn.2xlarge\"\n",
    "\n",
    "predictor = huggingface_model.deploy(\n",
    "    initial_instance_count=1,\n",
    "    instance_type=instance_type,\n",
    "    endpoint_name=endpoint_name\n",
    "    )\n",
    "\n",
    "print(\"Endpoint name: \", endpoint_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Step 5 - Test image generation"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "id": "13c753d3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4.59 ms, sys: 0 ns, total: 4.59 ms\n",
      "Wall time: 3.07 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "data = {\n",
    "    \"prompt\" : \"turn him into cyborg\",\n",
    "    \"image_url\" : \"https://raw.githubusercontent.com/timothybrooks/instruct-pix2pix/main/imgs/example.jpg\"\n",
    "}\n",
    "\n",
    "result = predictor.predict(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0e1aeb9",
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "# Create image from result\n",
    "from io import BytesIO\n",
    "import base64\n",
    "from PIL import Image\n",
    "\n",
    "# What to do with string to make it bytes again\n",
    "new_byte_io = BytesIO(base64.decodebytes(result.encode(\"utf-8\")))\n",
    "\n",
    "new_image = Image.open(new_byte_io)\n",
    "new_image.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "22976151",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Call endpoint:  huggingface-example-model-pix2pix-demo-on-miro-1\n",
      "With parameters:  b'{    \"action\": \"modify\", \"prompt\": \"turn him into cyborg\", \"image_url\": \"https://raw.githubusercontent.com/timothybrooks/instruct-pix2pix/main/imgs/example.jpg\"}'\n",
      "Received reply from endpoint, len:  4\n",
      "CPU times: user 58.8 ms, sys: 7.57 ms, total: 66.4 ms\n",
      "Wall time: 3.21 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "import json, base64\n",
    "from io import BytesIO\n",
    "from PIL import Image\n",
    "\n",
    "runtime = boto3.client('runtime.sagemaker')\n",
    "parameters = b'{    \"action\": \"modify\", \"prompt\": \"turn him into cyborg\", \"image_url\": \"https://raw.githubusercontent.com/timothybrooks/instruct-pix2pix/main/imgs/example.jpg\"}'\n",
    "#\n",
    "print (\"Call endpoint: \", endpoint_name)\n",
    "print (\"With parameters: \", parameters)\n",
    "response = runtime.invoke_endpoint(EndpointName=endpoint_name,\n",
    "                                   ContentType='application/json',\n",
    "                                   Body=parameters)\n",
    "print (\"Received reply from endpoint, len: \", len(response))\n",
    "# What to do with string to make it bytes again\n",
    "response_image = response[\"Body\"]\n",
    "stream = response_image.read()\n",
    "data = json.loads(stream)\n",
    "# base64.encodebytes(byte_io.getvalue()).decode(\"utf-8\")\n",
    "new_byte_io = BytesIO(base64.decodebytes(data.encode(\"utf-8\")))\n",
    "new_image = Image.open(new_byte_io)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6111ffe",
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "new_image"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Step 6 - After successful demo, please don't forget to clean up the endpoints to avoid unexpected charges"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "069aff90",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "predictor = sagemaker.predictor.Predictor(endpoint_name)\n",
    "predictor.delete_model()\n",
    "predictor.delete_endpoint()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ba0506b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_tensorflow2_p310",
   "language": "python",
   "name": "conda_tensorflow2_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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}