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

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

在本食谱中，我们将展示如何通过利用 Vertex AI 简化 PyTorch Stable Diffusion 模型的部署。

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

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

在 Vertex AI 端点上部署 Stable Diffusion 模型可以分为四个步骤

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

让我们更详细地了解每个步骤。您可以使用笔记本示例遵循并实现这些步骤。

注意：请记住，此食谱需要可计费的 Vertex AI，如笔记本示例中更详细地解释。

创建自定义 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,
)

如果您按照此笔记本操作，您还可以使用 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 模型部署到端点¶

更多资源¶

文档

教程

资源