部署 PyTorch 稳定扩散模型作为 Vertex AI 端点

创建时间：2025 年 4 月 1 日 | 最后更新时间：2025 年 4 月 1 日 | 最后验证：未验证

部署大型模型，如稳定扩散，可能具有挑战性和耗时。

在本菜谱中，我们将展示如何通过利用 Vertex AI 来简化 PyTorch 稳定扩散模型的部署。

PyTorch 是 Stability AI 在稳定扩散 v1.5 中使用的框架。Vertex AI 是一个全托管的机器学习平台，它提供了工具和基础设施，旨在帮助机器学习从业者加速和扩展生产中的机器学习，同时受益于像 PyTorch 这样的开源框架。

在四个步骤中，您可以部署 PyTorch 稳定扩散模型（v1.5）。

您可以在四个步骤中将您的稳定扩散模型部署到 Vertex AI 端点。

创建自定义 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 稳定扩散模型作为 Vertex AI 端点

创建自定义 TorchServe 处理器

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

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

将 Vertex AI 模型部署到端点 ¶

更多资源 ¶

文档

教程

资源