在 Vertex AI 上使用 PyTorch 推理 DLC 部署 FLUX

FLUX 是一款由 Black Forest Labs 创建的开源 120 亿参数修正流变换器，它可以根据文本描述生成图像，并突破了文本到图像生成的界限，非商业许可证使其广泛适用于探索和实验。此外，Google Cloud Vertex AI 是一个机器学习 (ML) 平台，可让您训练和部署 ML 模型和 AI 应用程序，以及自定义大型语言模型 (LLM) 以在您的 AI 驱动的应用程序中使用。

此示例展示了如何在 Vertex AI 上部署任何支持的 diffusers 文本到图像模型（在本例中为 black-forest-labs/FLUX.1-dev），方法是在 Google Cloud Platform (GCP) 中使用 Hugging Face PyTorch 推理 DLC（在 CPU 和 GPU 实例中都可用）。

设置/配置

首先，您需要在本地机器上安装 gcloud，它是 Google Cloud 的命令行工具，请按照 Cloud SDK 文档 - 安装 gcloud CLI 中的说明进行操作。

然后，您还需要安装 google-cloud-aiplatform Python SDK，它是以编程方式创建 Vertex AI 模型、注册模型、创建端点以及将其部署到 Vertex AI 所必需的。

!pip install --upgrade --quiet google-cloud-aiplatform

可选地，为了方便在本教程中使用命令，您需要为 GCP 设置以下环境变量

%env PROJECT_ID=your-project-id
%env LOCATION=your-location
%env CONTAINER_URI=us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-pytorch-inference-cu121.2-2.transformers.4-44.ubuntu2204.py311

然后，您需要登录您的 GCP 帐户并将项目 ID 设置为您要用来在 Vertex AI 上注册和部署模型的项目 ID。

!gcloud auth login
!gcloud auth application-default login  # For local development
!gcloud config set project $PROJECT_ID

登录后，您需要在 GCP 中启用必要的服务 API，例如 Vertex AI API、Compute Engine API 和 Google Container Registry 相关 API。

!gcloud services enable aiplatform.googleapis.com
!gcloud services enable compute.googleapis.com
!gcloud services enable container.googleapis.com
!gcloud services enable containerregistry.googleapis.com
!gcloud services enable containerfilesystem.googleapis.com

在 Vertex AI 中注册模型

一切设置完毕后，您就可以使用 google-cloud-aiplatform Python SDK 初始化 Vertex AI 会话，如下所示

import os
from google.cloud import aiplatform

aiplatform.init(
    project=os.getenv("PROJECT_ID"),
    location=os.getenv("LOCATION"),
)

由于 black-forest-labs/FLUX.1-dev 是一个受限制模型，因此您需要使用读访问令牌登录您的 Hugging Face Hub 帐户，该令牌要么是细粒度的，允许访问受限制模型，要么只是允许对您的帐户进行读访问。有关如何为 Hugging Face Hub 生成只读访问令牌的更多信息，请访问 https://huggingface.co/docs/hub/en/security-tokens。

!pip install --upgrade --quiet huggingface_hub

from huggingface_hub import interpreter_login

interpreter_login()

然后，您就可以“上传”模型，即在 Vertex AI 上注册模型。这不是真正的上传，因为该模型将在启动时通过 HF_MODEL_ID 环境变量从 Hugging Face Hub 中的 Hugging Face PyTorch 推理 DLC 中自动下载，因此上传的只是配置，而不是模型权重。

在深入代码之前，让我们快速回顾一下提供给 upload 方法的参数

• display_name 是将在 Vertex AI 模型注册表中显示的名称。
• serving_container_image_uri 是用于提供模型服务的 Hugging Face PyTorch 推理 DLC 的位置。
• serving_container_environment_variables 是在容器运行时使用的环境变量，因此这些变量与 huggingface-inference-toolkit Python SDK 定义的环境变量保持一致，该 SDK 公开了以下一些环境变量
  • HF_MODEL_ID 是 Hugging Face Hub 中模型的标识符。要浏览所有支持的模型，请查看 https://huggingface.co/models?sort=trending，并按您要使用的任务（例如 text-classification）进行筛选。
  • HF_TASK 是 Hugging Face Hub 中的任务标识符。要查看所有支持的任务，请查看 [https://huggingface.co/docs/transformers/en/task_summary#natural-language-processing。
  • HUGGING_FACE_HUB_TOKEN 是 Hugging Face Hub 令牌，因为 black-forest-labs/FLUX.1-dev 是一个受限制模型，因此需要使用该令牌。

有关支持的 aiplatform.Model.upload 参数的更多信息，请查看其 Python 引用，网址为 https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform.Model#google_cloud_aiplatform_Model_upload。

from huggingface_hub import get_token

model = aiplatform.Model.upload(
    display_name="black-forest-labs--FLUX.1-dev",
    serving_container_image_uri=os.getenv("CONTAINER_URI"),
    serving_container_environment_variables={
        "HF_MODEL_ID": "black-forest-labs/FLUX.1-dev",
        "HF_TASK": "text-to-image",
        "HF_TOKEN": get_token(),
    },
)
model.wait()

在 Vertex AI 中部署模型

模型注册到 Vertex AI 后，您需要定义要部署模型的目标端点，然后将模型部署链接到该端点资源。

为此，您需要调用 aiplatform.Endpoint.create 方法以创建一个新的 Vertex AI 端点资源（尚未链接到模型或任何可使用的东西）。

endpoint = aiplatform.Endpoint.create(display_name="black-forest-labs--FLUX.1-dev-endpoint")

现在您可以在 Vertex AI 的端点上部署注册的模型。

deploy 方法将先前创建的端点资源链接到包含服务容器配置的模型，然后它将在 Vertex AI 的指定实例中部署模型。

在进入代码之前，让我们快速回顾一下提供给 deploy 方法的参数

• endpoint 是要将模型部署到的端点，它是可选的，默认情况下将设置为模型显示名称加上 _endpoint 后缀。
• machine_type、accelerator_type 和 accelerator_count 是定义要使用的实例以及要使用的加速器和加速器数量的参数。machine_type 和 accelerator_type 是绑定的，因此您需要选择支持您正在使用的加速器的实例，反之亦然。有关不同实例的更多信息，请访问 Compute Engine 文档 - GPU 机器类型，以及有关 accelerator_type 命名方式的更多信息，请访问 Vertex AI 文档 - MachineSpec。
• enable_access_logging 是一个参数，用于启用端点访问日志记录，即记录 Google Cloud Logging 中对端点发出的请求信息。

有关支持的 aiplatform.Model.deploy 参数的更多信息，您可以在 https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform.Model#google_cloud_aiplatform_Model_deploy 上查看其 Python 参考。

%%time
deployed_model = model.deploy(
    endpoint=endpoint,
    machine_type="g2-standard-48",
    accelerator_type="NVIDIA_L4",
    accelerator_count=4,
    enable_access_logging=True,
)

在 Vertex AI 上进行在线预测

最后，您可以使用 predict 方法在 Vertex AI 上运行在线预测，该方法将请求发送到运行中的端点，该端点位于容器内指定的 /predict 路由中，并遵循 Vertex AI I/O 负载格式。

%%time
output = deployed_model.predict(
    instances=["a photo of an astronaut riding a horse on mars"],
    parameters={
        "width": 512,
        "height": 512,
        "num_inference_steps": 8,
        "guidance_scale": 3.5,
    },
)

由于部署的模型是文本到图像模型，输出负载将包含为提供的文本描述生成的图像，但以 base64 编码；这意味着您需要对其进行解码并使用 Pillow Python 库加载它。

!pip install --upgrade --quiet Pillow

import base64
import io

from IPython.display import display
from PIL import Image

image = Image.open(io.BytesIO(base64.b64decode(output.predictions[0])))
display(image)

警告 目前，Google Cloud 上的 Hugging Face DLC 具有 60 秒的请求超时，这意味着如果预测请求花费的时间超过 60 秒，将引发 HTTP 503 错误。为了防止这种情况，您可以根据 Vertex AI 文档 - 从自定义训练的模型获取在线预测 提高默认预测超时，方法是提交支持工单或联系您的 Google Cloud 代表。无论如何，对于公开发布的 Hugging Face DLC，Google Cloud 团队目前正在努力将它们列入白名单，以将预测超时提高到 600 秒。

资源清理

最后，您可以在与以下代码相同的 Python 会话中以编程方式释放资源

• deployed_model.undeploy_all 从所有端点取消部署模型。
• deployed_model.delete 在 undeploy_all 之后优雅地删除模型部署到的端点。
• model.delete 从注册表中删除模型。

deployed_model.undeploy_all()
deployed_model.delete()
model.delete()