使用TorchServe在GPU上运行多个生成式AI模型，并使用Amazon SageMaker多模型端点，可节省高达75%的推理成本

使用TorchServe在GPU上运行多个生成式AI模型，使用Amazon SageMaker多模型端点，可节省高达75%的推理成本

多模型端点（MME）是Amazon SageMaker的一个强大功能，旨在简化机器学习（ML）模型的部署和运行。通过MME，您可以在单个服务容器上托管多个模型，并将所有模型都放在一个端点后面。SageMaker平台会自动管理模型的加载和卸载，并根据流量模式扩展资源，减少管理大量模型的操作负担。这个功能对于需要加速计算的深度学习和生成AI模型特别有益。通过资源共享和简化的模型管理实现的成本节约使SageMaker MME成为您在AWS上托管大规模模型的绝佳选择。

最近，生成AI应用引起了广泛的关注和想象力。客户希望在GPU上部署生成AI模型，但同时又关注成本。SageMaker MME支持GPU实例，对于这类应用来说是一个很好的选择。今天，我们很高兴地宣布TorchServe支持SageMaker MME。这种新的模型服务器支持使您在使用TorchServe客户最熟悉的服务堆栈的同时，享受MME的所有好处。在本篇文章中，我们演示了如何使用TorchServe在SageMaker MME上托管生成AI模型，如稳定扩散和任意对象分割模型，并构建了一个语言引导的编辑解决方案，可以帮助艺术家和内容创作者更快地开发和迭代他们的艺术作品。

解决方案概述

语言引导的编辑是一个常见的跨行业生成AI用例。它可以帮助艺术家和内容创作者更加高效地满足内容需求，通过自动化重复任务、优化活动和为最终客户提供超个性化的体验。企业可以从增加内容输出、节约成本、改善个性化和提升客户体验中获益。在本篇文章中，我们演示了如何使用MME TorchServe构建语言辅助编辑功能，允许您通过提供文本指令来擦除图像中的任何不需要的对象，并修改或替换图像中的任何对象。

每个用例的用户体验流程如下：

• 在Amazon SageMaker JumpStart上对Llama 2进行微调，用于文本生成
• “通往量子互联网的新途径”
• 当给气味命名时，AI可以媲美人类的嗅觉

• 要删除一个不需要的对象，首先选择图像中的对象以突出显示。此操作将像素坐标和原始图像发送给生成AI模型，该模型为对象生成分割掩码。在确认正确的对象选择之后，您可以将原始图像和掩码图像发送给第二个模型进行删除操作。下面是该用户流程的详细说明。

步骤1：从图像中选择一个对象（“狗”）	步骤2：确认正确的对象已突出显示	步骤3：从图像中擦除对象

• 要修改或替换一个对象，首先选择并突出显示所需对象，然后按照上述相同的过程操作。一旦确认正确的对象选择，您可以通过提供原始图像、掩码和文本提示来修改对象。模型将根据提供的指令更改突出显示的对象。下面是该第二个用户流程的详细说明。

步骤1：从图像中选择一个对象（“花瓶”）	步骤2：确认正确的对象已突出显示	步骤3：提供文本提示（“未来主义花瓶”）来修改对象

为了支持这个解决方案，我们使用了三个生成式AI模型：Segment Anything Model (SAM)、Large Mask Inpainting Model (LaMa)和Stable Diffusion Inpaint (SD)。以下是这些模型在用户体验工作流中的使用方式：

移除不需要的对象	修改或替换对象

1. Segment Anything Model (SAM) 用于生成感兴趣对象的分割掩码。由 Meta Research 开发的 SAM 是一个开源模型，可以对图像中的任何对象进行分割。该模型在一个名为 SA-1B 的大型数据集上进行训练，该数据集包含超过 1100 万张图像和 11 亿个分割掩码。有关 SAM 的更多信息，请参阅其网站和研究论文。
2. LaMa 用于从图像中移除任何不需要的对象。LaMa 是一个生成对抗网络 (GAN) 模型，专门用于使用不规则掩码填充图像的缺失部分。该模型采用图像全局上下文和单步架构，使用 Fourier 卷积，使其能够以更快的速度实现最先进的结果。有关 LaMa 的更多详细信息，请访问其网站和研究论文。
3. SD 2 inpaint 模型来自 Stability AI，用于修改或替换图像中的对象。该模型允许我们通过提供文本提示来编辑掩码区域中的对象。inpaint 模型基于文本到图像的 SD 模型，可以使用简单的文本提示创建高质量图像。它提供了额外的参数，如原始图像和掩码图像，可以快速修改和恢复现有内容。有关在 AWS 上使用 Stable Diffusion 模型的更多信息，请参阅使用 Stable Diffusion 模型创建高质量图像，并在 Amazon SageMaker 上以低成本部署。

这三个模型都托管在 SageMaker MME 上，从而减少了管理多个端点的操作负担。除此之外，使用 MME 还消除了由于资源共享而导致某些模型被低效利用的担忧。您可以通过改善实例饱和度的好处来节省成本。下面的架构图说明了如何使用 TorchServe 在 SageMaker MME 上提供这三个模型。

我们已经在 GitHub 存储库中发布了实现此解决方案架构的代码。要跟随本文的其余部分，使用 notebook 文件。建议在 SageMaker 笔记本实例上使用 conda_python3 (Python 3.10.10) 内核运行此示例。

扩展 TorchServe 容器

第一步是准备模型托管容器。SageMaker 提供了一个托管的 PyTorch 深度学习容器 (DLC)，您可以使用以下代码片段检索它：

# 使用 SageMaker PyTorch DLC 作为基础镜像
baseimage = sagemaker.image_uris.retrieve(
    framework="pytorch",
    region=region,
    py_version="py310",
    image_scope="inference",
    version="2.0.0",
    instance_type="ml.g5.2xlarge",
)
print(baseimage)

因为模型需要资源和额外的软件包，而这些在基本的 PyTorch DLC 上不存在，所以您需要构建一个 Docker 镜像。然后将该镜像上传到 Amazon Elastic Container Registry (Amazon ECR)，以便我们可以直接从 SageMaker 访问。在 Docker 文件中列出了自定义安装的库：

ARG BASE_IMAGE

FROM $BASE_IMAGE

# 安装其他库
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

运行Shell命令文件以在本地构建自定义镜像并将其推送到Amazon ECR：

%%capture build_output

reponame = "torchserve-mme-demo"
versiontag = "genai-0.1"

# 构建自己的Docker镜像
!cd workspace/docker && ./build_and_push.sh {reponame} {versiontag} {baseimage} {region} {account}

准备模型文件

新的MME与TorchServe支持的主要区别在于如何准备模型文件。代码仓库为每个模型（models文件夹）提供了一个骨架文件夹，用于存放TorchServe所需的文件。我们按照相同的四步骤来准备每个模型的.tar文件。下面的代码是SD模型的骨架文件夹的示例：

workspace
|--sd
   |-- custom_handler.py
   |-- model-config.yaml

第一步是在models文件夹中下载预训练模型检查点：

import diffusers
import torch
import transformers

pipeline = diffusers.StableDiffusionInpaintPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2-inpainting", torch_dtype=torch.float16
)

sd_dir = "workspace/sd/model"
pipeline.save_pretrained(sd_dir)

下一步是定义一个custom_handler.py文件。这是为了定义模型接收到请求时的行为，例如加载模型、预处理输入和后处理输出。handle方法是请求的主要入口点，它接受一个请求对象并返回一个响应对象。它加载预训练模型检查点，并对输入和输出数据应用预处理和后处理方法。下面的代码片段示例了custom_handler.py文件的简单结构。有关详细信息，请参阅TorchServe处理程序API。

def initialize(self, ctx: Context):

def preprocess(self, data):

def inference(self, data):

def handle(self, data, context):
    requests = self.preprocess(data)
    responses = self.inference(requests)

    return responses

TorchServe所需的最后一个文件是model-config.yaml。该文件定义了模型服务器的配置，例如工作器数量和批量大小。配置是在模型级别上的，下面的代码展示了一个示例配置文件。有关参数的完整列表，请参阅GitHub仓库。

minWorkers: 1
maxWorkers: 1
batchSize: 1
maxBatchDelay: 200
responseTimeout: 300

最后一步是使用torch-model-archiver模块将所有模型文件打包成一个.tar.gz文件：

!torch-model-archiver --model-name sd --version 1.0 --handler workspace/sd/custom_handler.py --extra-files workspace/sd/model --config-file workspace/sam/model-config.yaml --archive-format no-archive!cd sd && tar cvzf sd.tar.gz .

创建多模型端点

创建SageMaker MME的步骤与以前相同。在这个特定的示例中，您使用SageMaker SDK启动一个端点。首先定义Amazon Simple Storage Service（Amazon S3）位置和托管容器。这个S3位置是SageMaker根据调用模式动态加载模型的位置。托管容器是您在之前的步骤中构建并推送到Amazon ECR的自定义容器。请参阅以下代码：

# 这是我们的MME将从S3上读取模型的位置。
multi_model_s3uri = output_path

然后您需要定义一个MulitDataModel，它包含所有属性，如模型位置、托管容器和权限访问：

print(multi_model_s3uri)
model = Model(
    model_data=f"{multi_model_s3uri}/sam.tar.gz",
    image_uri=container,
    role=role,
    sagemaker_session=smsess,
    env={"TF_ENABLE_ONEDNN_OPTS": "0"},
)

mme = MultiDataModel(
    name="torchserve-mme-genai-" + datetime.now().strftime("%Y-%m-%d-%H-%M-%S"),
    model_data_prefix=multi_model_s3uri,
    model=model,
    sagemaker_session=smsess,
)
print(mme)

deploy()函数创建一个端点配置并托管端点：

mme.deploy(
    initial_instance_count=1,
    instance_type="ml.g5.2xlarge",
    serializer=sagemaker.serializers.JSONSerializer(),
    deserializer=sagemaker.deserializers.JSONDeserializer(),
)

在我们提供的示例中，我们还展示了如何使用SDK列出模型并动态添加新模型。 add_model()函数将您的本地模型.tar文件复制到MME S3位置：

# 只有sam.tar.gz可见！
list(mme.list_models())

models = ["sd/sd.tar.gz", "lama/lama.tar.gz"]
for model in models:
    mme.add_model(model_data_source=model)

调用模型

现在我们在一个MME上托管了三个模型，我们可以按顺序调用每个模型来构建我们的语言辅助编辑功能。要调用每个模型，在predictor.predict()函数中提供一个target_model参数。模型名称只是我们上传的模型.tar文件的名称。以下是一个示例代码片段，用于SAM模型，它接受像素坐标、点标签和膨胀核大小，并生成像素位置处对象的分割掩码：

img_file = "workspace/test_data/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 = json.dumps({"image": img_bytes, "gen_args": gen_args}).encode("utf-8")

response = predictor.predict(data=payload, target_model="/sam.tar.gz")
encoded_masks_string = json.loads(response.decode("utf-8"))["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()

要从图像中移除不需要的对象，请使用从SAM生成的分割掩码，将其与原始图像一起输入LaMa模型。以下是一个示例。

示例图像	来自SAM的分割掩码	使用LaMa擦除狗

要使用文本提示修改或替换图像中的任何对象，请将SAM生成的分割掩码与原始图像和文本提示一起输入SD模型，如下例所示。

!		!
示例图像	来自SAM的分割掩码	使用文本提示的SD模型替换 “在长凳上的仓鼠”

成本节省

SageMaker MME的好处随着模型合并规模的增加而增加。下表显示了本文中三个模型的GPU内存使用情况。它们使用一个SageMaker MME部署在一个g5.2xlarge实例上。

模型	GPU内存（MiB）
分割任意模型	3,362
稳定扩散绘画	3,910
Lama	852

当使用一个端点托管这三个模型时，您可以看到成本节省的效果。对于具有数百或数千个模型的用例，节省的成本会更大。

例如，考虑100个稳定扩散模型。每个模型本身可以通过一个ml.g5.2xlarge端点（4 GiB内存）提供服务，在美国东部（弗吉尼亚北部）地区的成本为每小时1.52美元。如果使用各自的端点提供所有100个模型，每月的费用将达到218,880美元。通过使用SageMaker MME，使用ml.g5.2xlarge实例可以同时托管四个模型的单一端点。这将使生产推理成本降低75%，仅为每月54,720美元。下表总结了单模型端点和多模型端点在此示例中的差异。在加载完所有模型后，给定具有足够内存的端点配置，稳定状态调用延迟将类似于单模型端点的延迟。

.	单模型端点	多模型端点
每月总端点价格	218,880美元	54,720美元
端点实例类型	ml.g5.2xlarge	ml.g5.2xlarge
CPU内存容量（GiB）	32	32
GPU内存容量（GiB）	24	24
每小时端点价格	1.52美元	1.52美元
每个端点的实例数量	2	2
100个模型所需的端点数	100	25

清理

完成后，请按照笔记本中的清理部分中的说明删除本文中提供的资源，以避免产生不必要的费用。有关推理实例成本的详细信息，请参阅Amazon SageMaker Pricing。

结论

本文演示了通过在SageMaker MME上托管的TorchServe上使用生成式AI模型实现的语言辅助编辑功能。我们分享的示例说明了我们如何使用SageMaker MME的资源共享和简化模型管理，同时仍然利用TorchServe作为我们的模型服务堆栈。我们使用了三个深度学习基础模型：SAM、SD 2 Inpainting和LaMa。这些模型使我们能够构建强大的功能，例如从图像中擦除任何不需要的对象，并通过提供文本指令修改或替换图像中的任何对象。这些功能可以帮助艺术家和内容创作者更高效地工作，并通过自动化重复任务、优化活动和提供超个性化体验来满足其内容需求。我们邀请您探索本文中提供的示例，并使用SageMaker MME上的TorchServe构建自己的UI体验。

要开始，请参阅支持的算法、框架和使用 GPU 支持实例的多模型端点。