使用 BentoML 部署 DeepFloyd IF

TL;DR

Hugging Face 已詳細說明使用 BentoML（一個開源模型服務框架）部署 DeepFloyd IF 的工作流程。此整合使開發者能夠透過在不同 GPU 資源上獨立擴展推理運行器來管理 DeepFloyd IF 複雜的多階段架構，以優化生產環境的性能。

DeepFloyd IF 架構與功能

DeepFloyd IF 是一種直接在像素空間運作的開源文字到圖像模型，這使其與如 Stable Diffusion 這類潛在擴散模型有所區別。它採用由凍結的文字編碼器和三個級聯的像素擴散模組組成的模組化結構：

  • Stage 1: 生成 64x64 像素的基礎圖像。
  • Stage 2 和 3: 逐步將圖像上採樣至最終解析度 1024x1024 像素。

為了增強對複雜提示的語言理解，DeepFloyd IF 將 T5-XXL-1.1 大型語言模型 (LLM) 作為其文字編碼器。

使用 BentoML 進行生產部署

BentoML 提供了一個統一的框架，用於將模型從 Hugging Face Hub 過渡到可投入生產的 AI 應用程式。部署流程遵循五個步驟的生命週期：

  1. Define a model: 使用在 PyTorch 或 TensorFlow 等庫中訓練的模型。
  2. Save the model: 將模型存儲在 BentoML 本地 Model Store 中，以便管理和服務訪問。
  3. Create a BentoML Service: 使用 service.py 檔案封裝模型，並透過 Runners 定義服務邏輯，Runners 能夠處理大規模模型推理並公開 API 以進行輸入/輸出處理。
  4. Build a Bento: 將服務和模型打包成一個「Bento」——一個可部署的 artefact，透過 YAML 配置檔案包含所有必要的程式碼與依賴。
  5. Deploy the Bento: 將 Bento 容器化為 Docker 镜像，以進行 Kubernetes 部署，或使用 Yatai 在 Kubernetes 上實現自動擴展。

技術實作與資源管理

部署 DeepFloyd IF 需要大量的計算資源。建議的硬體配置包括至少 2 個 16GB VRAM 的 GPU 或 1 個 40GB VRAM 的 GPU。對於生產級別的服務，單一 Tesla T4 不建議使用。

GPU 配置策略

BentoML 允許為 DeepFloyd IF 管線的每個階段獨立擴展 Runners，使開發者能夠根據每個階段的具體需求分配資源：

  • High VRAM (40GB+): 所有模型可以在單一 GPU 上運行。
  • Dual Tesla T4 (15GB each): 第階段 1 可以分配給第一個 GPU，而階段 2 和 3 則分配給第二個 GPU。
  • Mixed GPU Setup: 階段 1 可以分配給主要 GPU（例如 T4），而階段 2 和 3 則分佈在兩個額外的 VRAM 較小的 GPU 上。

部署工作流程

要實施此設置，開發者使用 import_models.py 將所需階段下載到 BentoML Model Store，並使用 service.py 定義接受 JSON 提示和負面提示以返回生成圖像的 API。最終產物透過 bentoml build 打包，可透過 bentoml serve 在本地提供服務，或使用 bentoml containerize 部署到雲端。

Sources