Feat: TorchServe SageMaker MME GPU example (#4116)

* init torchserve-mme-gpu

* add test data

* add mme notebook

* Update torchserve_multi_model_endpoint.ipynb

* add dockerfile

* fmt notebook

* fmt

* chmod 755 for build_and_push.sh

* fix typo

* model config

* fix grammer

* fix grammer

* fix grammer

* fix grammer

* fix grammer

* fix grammer

Author: lxning

inference/torchserve/mme-gpu/torchserve_multi_model_endpoint.ipynb

@@ -0,0 +1,22 @@
+segment-anything-py==1.0
+opencv-python-headless==4.7.0.68
+transformers==4.28.1
+ftfy
+diffusers
+xformers
+tqdm
+#easydict==1.9.0
+#scikit-image==0.17.2
+#scikit-learn==0.24.2
+easydict
+scikit-image
+tensorflow
+joblib
+matplotlib
+albumentations==0.5.2
+hydra-core==1.1.0
+pytorch-lightning
+tabulate
+kornia==0.5.0
+webdataset
+omegaconf==2.1.2

inference/torchserve/mme-gpu/workspace/code/requirements.txt

@@ -0,0 +1,26 @@
+ARG BASE_IMAGE
+
+FROM $BASE_IMAGE
+
+#Install any additional libraries
+RUN pip install segment-anything-py==1.0
+RUN pip install opencv-python-headless==4.7.0.68
+RUN pip install matplotlib==3.6.3
+RUN pip install diffusers
+RUN pip install tqdm
+RUN pip install easydict
+RUN pip install scikit-image
+RUN pip install xformers
+RUN pip install tensorflow
+RUN pip install joblib
+RUN pip install matplotlib
+RUN pip install albumentations==0.5.2
+RUN pip install hydra-core==1.1.0
+RUN pip install pytorch-lightning
+RUN pip install tabulate
+RUN pip install kornia==0.5.0
+RUN pip install webdataset
+RUN pip install omegaconf==2.1.2
+RUN pip install transformers==4.28.1
+RUN pip install accelerate
+RUN pip install ftfy

inference/torchserve/mme-gpu/workspace/docker/Dockerfile

@@ -0,0 +1,50 @@
+#!/usr/bin/env bash
+
+# This script shows how to build the Docker image and push it to ECR to be ready for use
+# by SageMaker.
+
+# The argument to this script is the image name. This will be used as the image on the local
+# machine and combined with the account and region to form the repository name for ECR.
+reponame=$1
+versiontag=$2
+baseimage=$3
+regionname=$4
+account=$5
+
+if [ "$reponame" == "" ] || [ "$versiontag" == "" ]  || [ "$baseimage" == "" ] || [ "$regionname" == "" ] || [ "$account" == "" ]
+then
+    echo "Usage: $0 <repo-name> <version-tag> <base-image> <region> <account>"
+    exit 1
+fi
+
+if [ $? -ne 0 ]
+then
+    exit 255
+fi
+
+fullname="${account}.dkr.ecr.${regionname}.amazonaws.com/${reponame}:${versiontag}"
+
+# If the repository doesn't exist in ECR, create it.
+aws ecr describe-repositories --repository-names "${reponame}" > /dev/null 2>&1
+
+if [ $? -ne 0 ]
+then
+    aws ecr create-repository --repository-name "${reponame}" > /dev/null
+fi
+
+# Get the login command from ECR in order to pull down the SageMaker PyTorch image
+aws ecr get-login-password --region $regionname | docker login --username AWS --password-stdin ${baseimage} 
+
+aws ecr get-login-password --region $regionname | docker login --username AWS --password-stdin ${account}.dkr.ecr."${regionname}".amazonaws.com
+
+# Build the docker image locally with the image name and then push it to ECR
+# with the full name.
+docker build  -t ${reponame} . --build-arg BASE_IMAGE=${baseimage}
+docker tag ${reponame} ${fullname}
+
+# Get the login command from ECR in order to pull down the SageMaker PyTorch image
+aws ecr get-login-password --region $regionname | docker login --username AWS --password-stdin ${account}.dkr.ecr."${regionname}".amazonaws.com
+
+docker push ${fullname}
+echo "${fullname}"
+

inference/torchserve/mme-gpu/workspace/docker/build_and_push.sh

@@ -0,0 +1,39 @@
+import base64
+import json
+import io
+import numpy as np
+from PIL import Image
+import httpx
+from io import BytesIO
+
+def encode_image(img):
+    
+    # Convert the image to bytes
+    with io.BytesIO() as output:
+        img.save(output, format="JPEG")
+        img_bytes = output.getvalue()
+    
+    return base64.b64encode(img_bytes).decode('utf8')
+
+img_file = 'sample1.png'
+img_bytes = None
+with Image.open(img_file) as f:
+    img_bytes = encode_image(f)
+
+mask_file = 'sample1_mask.jpg'
+mask = Image.open(mask_file)
+mask_bytes = encode_image(mask)
+
+payload = {
+        "image": img_bytes, 
+        "mask_image": mask_bytes,
+        }
+
+url="http://127.0.0.1:8080/predictions/lama"
+response = httpx.post(url, json=payload, timeout=None)
+encoded_masks_string = response.json()['generated_image']
+base64_bytes_masks = base64.b64decode(encoded_masks_string)
+print(base64_bytes_masks)
+with Image.open(io.BytesIO(base64_bytes_masks)) as f:
+    generated_image_rgb=f.convert("RGB")
+    generated_image_rgb.show()

inference/torchserve/mme-gpu/workspace/lama/client.py

@@ -0,0 +1,160 @@
+import os
+import sys
+import numpy as np
+import torch
+import yaml
+import glob
+import argparse
+from PIL import Image
+from omegaconf import OmegaConf
+from pathlib import Path
+import json
+import base64
+
+from abc import ABC
+from io import BytesIO
+from ts.context import Context
+from ts.torch_handler.base_handler import BaseHandler
+
+from saicinpainting.evaluation.utils import move_to_device
+from saicinpainting.training.trainers import load_checkpoint
+from saicinpainting.evaluation.data import pad_tensor_to_modulo
+
+os.environ['OMP_NUM_THREADS'] = '1'
+os.environ['OPENBLAS_NUM_THREADS'] = '1'
+os.environ['MKL_NUM_THREADS'] = '1'
+os.environ['VECLIB_MAXIMUM_THREADS'] = '1'
+os.environ['NUMEXPR_NUM_THREADS'] = '1'
+
+class LamaHandler(BaseHandler, ABC):
+
+    def __init__(self):
+        super().__init__()
+        self.initialized = False
+
+    def initialize(self, ctx: Context):
+        properties = ctx.system_properties
+        model_dir = properties.get("model_dir")
+
+        self.device = torch.device(
+            "cuda:" + str(properties.get("gpu_id"))
+            if torch.cuda.is_available() and properties.get("gpu_id") is not None
+            else "cpu"
+        )
+
+        predict_config = OmegaConf.load(f'{model_dir}/configs/prediction/default.yaml')
+        predict_config.model.path = f'{model_dir}/big-lama'
+        with open(f'{predict_config.model.path}/config.yaml', 'r') as f:
+            train_config = OmegaConf.create(yaml.safe_load(f))
+
+        train_config.training_model.predict_only = True
+        train_config.visualizer.kind = 'noop'
+
+        checkpoint_path = os.path.join(
+            predict_config.model.path, 'models',
+            predict_config.model.checkpoint
+        )
+        self.model = load_checkpoint(
+            train_config,
+            checkpoint_path,
+            strict=False,
+            map_location='cpu')
+        self.model.freeze()
+        self.model.to(self.device)
+
+        self.initialized = True
+
+    def preprocess(self, data):
+
+        requests = []
+        for row in data:
+            request = row.get("data") or row.get("body")
+
+            if isinstance(request, (bytearray, bytes)):
+                request = json.loads(request.decode('utf-8'))            
+
+            if isinstance(request, dict) and \
+                    "image" in request and \
+                    "mask_image" in request:
+                img = request["image"]
+                if isinstance(img, str):
+                    img = base64.b64decode(img)
+
+                with Image.open(BytesIO(img)) as f:
+                    img_rgb = f.convert("RGB")
+                    img_np_array = np.array(img_rgb)
+                    request["image"] = img_np_array
+
+                mask_img = request["mask_image"]
+                if isinstance(mask_img, str):
+                    mask_img = base64.b64decode(mask_img)
+
+                with Image.open(BytesIO(mask_img)) as f:
+                    mask_img_rgb = f.convert("L")
+                    mask_img_np_array = np.array(mask_img_rgb)
+                    request["mask_image"] = mask_img_np_array
+
+                requests.append(request)
+            else:
+                raise RuntimeError("Dict request must include image and mask_image")
+
+        return requests
+
+    def inference(self, data):
+
+        responses = []
+        for request in data:
+            mod = 8
+            img = torch.from_numpy(request["image"]).float().div(255.)
+            mask = torch.from_numpy(request["mask_image"]).float()
+
+            batch = {}
+            batch['image'] = img.permute(2, 0, 1).unsqueeze(0)
+            batch['mask'] = mask[None, None]
+            unpad_to_size = [batch['image'].shape[2], batch['image'].shape[3]]
+            batch['image'] = pad_tensor_to_modulo(batch['image'], mod)
+            batch['mask'] = pad_tensor_to_modulo(batch['mask'], mod)
+            #batch = move_to_device(batch, 'cuda')
+            batch = move_to_device(batch, self.device)
+            batch['mask'] = (batch['mask'] > 0) * 1
+
+            batch = self.model(batch)
+            cur_res = batch['inpainted'][0].permute(1, 2, 0)
+            cur_res = cur_res.detach().cpu().numpy()
+
+            if unpad_to_size is not None:
+                orig_height, orig_width = unpad_to_size
+                cur_res = cur_res[:orig_height, :orig_width]
+
+            cur_res = np.clip(cur_res * 255, 0, 255).astype('uint8')
+
+            output_img = Image.fromarray(cur_res)
+            output_img_bytes = self.encode_image(output_img)
+
+            print(f'output_img_bytes:{output_img_bytes}')
+            responses.append({"generated_image": output_img_bytes})
+
+        return responses
+
+    def handle(self, data, context):
+        requests = self.preprocess(data)
+        responses = self.inference(requests)
+
+        return responses
+
+    def dilate_mask(self, mask, dilate_factor=15):
+        mask = mask.astype(np.uint8)
+        mask = cv2.dilate(
+            mask,
+            np.ones((dilate_factor, dilate_factor), np.uint8),
+            iterations=1
+        )
+        return mask
+
+    def encode_image(self, img):
+        # Convert the image to bytes
+        with BytesIO() as output:
+            img.save(output, format="JPEG")
+            img_bytes = output.getvalue()
+
+        return base64.b64encode(img_bytes).decode("utf-8")

inference/torchserve/mme-gpu/workspace/lama/custom_handler.py

@@ -0,0 +1,5 @@
+minWorkers: 4 
+maxWorkers: 4
+batchSize: 1
+maxBatchDelay: 200
+responseTimeout: 300

inference/torchserve/mme-gpu/workspace/lama/model-config.yaml

@@ -0,0 +1,20 @@
+opencv-python
+pyyaml
+tqdm
+numpy
+easydict==1.9.0
+scikit-image==0.17.2
+scikit-learn==0.24.2
+tensorflow
+joblib
+matplotlib
+pandas
+albumentations==0.5.2
+hydra-core==1.1.0
+#pytorch-lightning==1.2.9
+pytorch-lightning
+tabulate
+kornia==0.5.0
+webdataset
+packaging
+omegaconf==2.1.2

inference/torchserve/mme-gpu/workspace/lama/requirements.txt

@@ -0,0 +1,36 @@
+import base64
+import json
+from json import JSONEncoder
+import io
+import numpy as np
+from PIL import Image
+import httpx
+
+def encode_image(img):
+    
+    # Convert the image to bytes
+    with io.BytesIO() as output:
+        img.save(output, format="JPEG")
+        img_bytes = output.getvalue()
+    
+    return base64.b64encode(img_bytes).decode('utf8')
+
+img_file = 'sample1.png'
+img_bytes = None
+with Image.open(img_file) as f:
+    img_bytes = encode_image(f)
+
+gen_args = json.dumps(dict(point_coords=[750, 500], point_labels=1, dilate_kernel_size=15))
+
+payload = {
+        "image": img_bytes, 
+        "gen_args": gen_args
+        }
+
+url="http://127.0.0.1:8080/predictions/sam"
+response = httpx.post(url, json=payload, timeout=None)
+encoded_masks_string = response.json()['generated_image']
+base64_bytes_masks = base64.b64decode(encoded_masks_string)
+with Image.open(io.BytesIO(base64_bytes_masks)) as f:
+    generated_image_rgb=f.convert("RGB")
+    generated_image_rgb.show()