将 PyTorch Stable Diffusion 模型部署为 Vertex AI 端点¶

创建日期：2023 年 11 月 14 日 | 最后更新：2023 年 11 月 14 日 | 最后验证：未验证

部署大型模型，例如 Stable Diffusion，可能具有挑战性且耗时。

在这个 Recipe 中，我们将展示如何利用 Vertex AI 简化 PyTorch Stable Diffusion 模型的部署。

PyTorch 是 Stability AI 在 Stable Diffusion v1.5 上使用的框架。Vertex AI 是一个全托管的机器学习平台，提供工具和基础设施，旨在帮助 ML 从业者利用 PyTorch 等开源框架，加速和扩展生产环境中的 ML。

通过四个步骤，您可以部署一个 PyTorch Stable Diffusion 模型 (v1.5)。

通过四个步骤即可将您的 Stable Diffusion 模型部署到 Vertex AI 端点

创建自定义 TorchServe 处理程序。
将模型工件上传到 Google Cloud Storage (GCS)。
使用模型工件和预构建的 PyTorch 容器镜像创建 Vertex AI 模型。
将 Vertex AI 模型部署到端点。

让我们更详细地看看每个步骤。您可以使用 Notebook 示例来遵循和实现这些步骤。

注意：请记住，此 Recipe 需要付费的 Vertex AI，Notebook 示例中对此有更详细的说明。

创建自定义 TorchServe 处理程序¶

TorchServe 是一个简单灵活的工具，用于服务 PyTorch 模型。部署到 Vertex AI 的模型使用 TorchServe 处理请求并返回模型的响应。您必须创建一个自定义 TorchServe 处理程序，以便包含在上传到 Vertex AI 的模型工件中。将处理程序文件包含在与其他模型工件相同的目录中，例如：model_artifacts/handler.py。

创建处理程序文件后，必须将处理程序打包为模型归档文件 (MAR)。输出文件必须命名为 model.mar。

!torch-model-archiver \
-f \
--model-name <your_model_name> \
--version 1.0 \
 --handler model_artifacts/handler.py \
--export-path model_artifacts

将模型工件上传到 Google Cloud Storage (GCS)¶

在此步骤中，我们将模型工件（例如模型文件或处理程序）上传到 GCS。将工件存储在 GCS 上的优点是您可以在一个中心存储桶中跟踪这些工件。

BUCKET_NAME = "your-bucket-name-unique"  # @param {type:"string"}
BUCKET_URI = f"gs://{BUCKET_NAME}/"

# Will copy the artifacts into the bucket
!gsutil cp -r model_artifacts $BUCKET_URI

使用模型工件和预构建的 PyTorch 容器镜像创建 Vertex AI 模型¶

将模型工件上传到 GCS 存储桶后，您可以将 PyTorch 模型上传到 Vertex AI 模型注册表。通过 Vertex AI 模型注册表，您可以概览您的模型，从而更好地组织、跟踪和训练新版本。为此，您可以使用 Vertex AI SDK 和这个预构建的 PyTorch 容器。

from google.cloud import aiplatform as vertexai
PYTORCH_PREDICTION_IMAGE_URI = (
    "us-docker.pkg.dev/vertex-ai/prediction/pytorch-gpu.1-12:latest"
)
MODEL_DISPLAY_NAME = "stable_diffusion_1_5-unique"
MODEL_DESCRIPTION = "stable_diffusion_1_5 container"

vertexai.init(project='your_project', location='us-central1', staging_bucket=BUCKET_NAME)

model = aiplatform.Model.upload(
    display_name=MODEL_DISPLAY_NAME,
    description=MODEL_DESCRIPTION,
    serving_container_image_uri=PYTORCH_PREDICTION_IMAGE_URI,
    artifact_uri=BUCKET_URI,
)

将 Vertex AI 模型部署到端点¶

将模型上传到 Vertex AI 模型注册表后，您可以将其部署到 Vertex AI 端点。为此，您可以使用控制台或 Vertex AI SDK。在此示例中，您将模型部署到 NVIDIA Tesla P100 GPU 和 n1-standard-8 机器。您可以指定您的机器类型。

endpoint = aiplatform.Endpoint.create(display_name=ENDPOINT_DISPLAY_NAME)

model.deploy(
    endpoint=endpoint,
    deployed_model_display_name=MODEL_DISPLAY_NAME,
    machine_type="n1-standard-8",
    accelerator_type="NVIDIA_TESLA_P100",
    accelerator_count=1,
    traffic_percentage=100,
    deploy_request_timeout=1200,
    sync=True,
)

如果您遵循此notebook，您还可以使用 Vertex AI SDK 获取在线预测，如下面的代码片段所示。

instances = [{"prompt": "An examplePup dog with a baseball jersey."}]
response = endpoint.predict(instances=instances)

with open("img.jpg", "wb") as g:
    g.write(base64.b64decode(response.predictions[0]))

display.Image("img.jpg")

使用模型工件和预构建的 PyTorch 容器镜像创建 Vertex AI 模型

将 PyTorch Stable Diffusion 模型部署为 Vertex AI 端点¶

创建自定义 TorchServe 处理程序¶

将模型工件上传到 Google Cloud Storage (GCS)¶

使用模型工件和预构建的 PyTorch 容器镜像创建 Vertex AI 模型¶

将 Vertex AI 模型部署到端点¶

更多资源¶

文档

教程

资源