使用 FLUX、Python 和 Diffusers 創建人工智慧驅動的圖像生成 API 服務

FLUX（由 Black Forest Labs 開發）在過去幾個月席捲了 AI 影像生成領域。它不僅在許多基準測試中擊敗了 Stable Diffusion（先前的開源之王），還在某些指標上超越了 Dall-E 或 Midjourney 等專有模型。

但是您將如何在您的某個應用程式上使用 FLUX？人們可能會考慮使用 Replicate 等無伺服器主機，但這些主機很快就會變得非常昂貴，並且可能無法提供您所需的靈活性。這就是創建您自己的自訂 FLUX 伺服器派上用場的地方。

在本文中，我們將引導您使用 Python 建立自己的 FLUX 伺服器。該伺服器將允許您透過簡單的 API 根據文字提示產生圖像。無論您是執行此伺服器供個人使用還是將其部署為生產應用程式的一部分，本指南都將幫助您入門。

先決條件

在深入研究程式碼之前，讓我們確保您已設定必要的工具和函式庫：

• Python：您需要在電腦上安裝 Python 3，最好是 3.10 版本。
• torch：我們將用來運行 FLUX 的深度學習框架。
• 擴散器：提供對 FLUX 模型的存取。
• 變壓器：擴散器所需的依賴項。
• 句：運行 FLUX 分詞器所需
• protobuf：運行 FLUX 所需
• 加速：在某些情況下幫助更有效地載入 FLUX 模型。
• fastapi：用於建立可以接受影像產生請求的 Web 伺服器的框架。
• uvicorn：運行 FastAPI 伺服器所需。
• psutil：允許我們檢查我們的機器上有多少 RAM。

您可以透過執行以下指令來安裝所有函式庫：pip install torchifferstransformerssentpieceprotobufacceleratefastapiuvicorn。

如果您使用的是配備 M1 或 M2 晶片的 Mac，則應使用 Metal 設定 PyTorch 以獲得最佳效能。在繼續之前，請遵循官方 PyTorch with Metal 指南。

如果您打算在 GPU 裝置上執行 FLUX，您還需要確保至少有 12 GB 的 VRAM。或至少 12 GB RAM 用於在 CPU/MPS 上運行（這會更慢）。

第 1 步：設定環境

讓我們根據我們正在使用的硬體選擇正確的設備來運行推理來啟動腳本。

device = 'cuda' # can also be 'cpu' or 'mps'

import os

# MPS support in PyTorch is not yet fully implemented
if device == 'mps':
  os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"

import torch

if device == 'mps' and not torch.backends.mps.is_available():
      raise Exception("Device set to MPS, but MPS is not available")
elif device == 'cuda' and not torch.cuda.is_available():
      raise Exception("Device set to CUDA, but CUDA is not available")

您可以指定 cpu、cuda（用於 NVIDIA GPU）或 mps（對於 Apple 的 Metal Performance Shaders）。然後，該腳本檢查所選設備是否可用，如果不可用，則引發異常。

第 2 步：載入 FLUX 模型

接下來，我們載入 FLUX 模型。我們將以 fp16 精度加載模型，這將節省一些內存，而不會造成太大的品質損失。

此時，可能會要求您使用 HuggingFace 進行身份驗證，因為 FLUX 模型是門控的。為了成功進行身份驗證，您需要建立 HuggingFace 帳戶，前往模型頁面，接受條款，然後從您的帳戶設定建立 HuggingFace 令牌並將其作為 HF_TOKEN 環境變數新增至您的電腦。

device = 'cuda' # can also be 'cpu' or 'mps'

import os

# MPS support in PyTorch is not yet fully implemented
if device == 'mps':
  os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"

import torch

if device == 'mps' and not torch.backends.mps.is_available():
      raise Exception("Device set to MPS, but MPS is not available")
elif device == 'cuda' and not torch.cuda.is_available():
      raise Exception("Device set to CUDA, but CUDA is not available")

在這裡，我們使用擴散器庫來載入 FLUX 模型。我們使用的模型是 black-forest-labs/FLUX.1-dev，以 fp16 精度載入。

還有一個名為 FLUX Schnell 的時間步蒸餾模型，它具有更快的推理速度，但輸出的圖像細節較少，以及一個閉源的 FLUX Pro 模型。
我們將在這裡使用 Euler 調度程序，但您可以嘗試一下。您可以在此處閱讀有關調度程序的更多資訊。
由於影像生成可能會佔用大量資源，因此優化記憶體使用至關重要，尤其是在 CPU 或記憶體有限的裝置上運行時。

from diffusers import FlowMatchEulerDiscreteScheduler, FluxPipeline
import psutil

model_name = "black-forest-labs/FLUX.1-dev"

print(f"Loading {model_name} on {device}")

pipeline = FluxPipeline.from_pretrained(
      model_name,

      # Diffusion models are generally trained on fp32, but fp16
      # gets us 99% there in terms of quality, with just half the (V)RAM
      torch_dtype=torch.float16,

      # Ensure we don't load any dangerous binary code
      use_safetensors=True

      # We are using Euler here, but you can also use other samplers
      scheduler=FlowMatchEulerDiscreteScheduler()
).to(device)

此程式碼檢查總可用內存，並在系統 RAM 小於 64 GB 時啟用注意力切片。注意力切片可減少影像產生過程中的記憶體使用，這對於資源有限的裝置至關重要。

第 3 步：使用 FastAPI 建立 API

接下來，我們將設定 FastAPI 伺服器，它將提供用於產生映像的 API。

# Recommended if running on MPS or CPU with < 64 GB of RAM
total_memory = psutil.virtual_memory().total
total_memory_gb = total_memory / (1024 ** 3)
if (device == 'cpu' or device == 'mps') and total_memory_gb < 64:
      print("Enabling attention slicing")
      pipeline.enable_attention_slicing()

FastAPI 是一個使用 Python 建立 Web API 的熱門框架。在本例中，我們使用它來建立一個可以接受影像產生請求的伺服器。我們還使用 GZip 中間件來壓縮回應，這在以 Base64 格式傳回圖片時特別有用。

在生產環境中，您可能想要將產生的映像儲存在 S3 儲存桶或其他雲端儲存中，並傳回 URL 而不是 Base64 編碼的字串，以利用 CDN 和其他最佳化。

步驟 4：定義請求模型

我們現在需要為我們的 API 將接受的請求定義一個模型。

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field, conint, confloat
from fastapi.middleware.gzip import GZipMiddleware
from io import BytesIO
import base64

app = FastAPI()

# We will be returning the image as a base64 encoded string
# which we will want compressed
app.add_middleware(GZipMiddleware, minimum_size=1000, compresslevel=7)

此GenerateRequest模型定義了產生影像所需的參數。提示欄位是您要建立的圖像的文字描述。其他欄位包括圖像尺寸、推理步驟數和批次大小。

步驟5：建立影像產生端點

現在，讓我們建立處理影像產生請求的端點。

device = 'cuda' # can also be 'cpu' or 'mps'

import os

# MPS support in PyTorch is not yet fully implemented
if device == 'mps':
  os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"

import torch

if device == 'mps' and not torch.backends.mps.is_available():
      raise Exception("Device set to MPS, but MPS is not available")
elif device == 'cuda' and not torch.cuda.is_available():
      raise Exception("Device set to CUDA, but CUDA is not available")

此端點處理影像產生過程。它首先根據 FLUX 的要求驗證高度和寬度是否為 8 的倍數。然後，它根據提供的提示生成圖像，並將它們作為 base64 編碼的字串傳回。

第6步：啟動伺服器

最後，讓我們添加一些程式碼以在腳本運行時啟動伺服器。

from diffusers import FlowMatchEulerDiscreteScheduler, FluxPipeline
import psutil

model_name = "black-forest-labs/FLUX.1-dev"

print(f"Loading {model_name} on {device}")

pipeline = FluxPipeline.from_pretrained(
      model_name,

      # Diffusion models are generally trained on fp32, but fp16
      # gets us 99% there in terms of quality, with just half the (V)RAM
      torch_dtype=torch.float16,

      # Ensure we don't load any dangerous binary code
      use_safetensors=True

      # We are using Euler here, but you can also use other samplers
      scheduler=FlowMatchEulerDiscreteScheduler()
).to(device)

此程式碼在連接埠 8000 上啟動 FastAPI 伺服器，由於 0.0.0.0 綁定，不僅可以從 http://localhost:8000 存取它，還可以使用主機的 IP 位址從同一網路上的其他裝置存取它。

第 7 步：本地測試您的伺服器

現在您的 FLUX 伺服器已啟動並運行，是時候對其進行測試了。您可以使用curl（一種用於發出HTTP請求的命令列工具）與您的伺服器互動：

# Recommended if running on MPS or CPU with < 64 GB of RAM
total_memory = psutil.virtual_memory().total
total_memory_gb = total_memory / (1024 ** 3)
if (device == 'cpu' or device == 'mps') and total_memory_gb < 64:
      print("Enabling attention slicing")
      pipeline.enable_attention_slicing()

此指令僅適用於安裝了curl、jq 和base64 公用程式的基於UNIX 的系統。根據託管 FLUX 伺服器的硬件，它也可能需要幾分鐘才能完成。

結論

恭喜！您已經使用 Python 成功創建了自己的 FLUX 伺服器。此設定可讓您透過簡單的 API 根據文字提示產生圖像。如果您對基本 FLUX 模型的結果不滿意，您可以考慮微調模型，以便在特定用例上獲得更好的效能。

完整程式碼

您可以在下面找到本指南中使用的完整程式碼：

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field, conint, confloat
from fastapi.middleware.gzip import GZipMiddleware
from io import BytesIO
import base64

app = FastAPI()

# We will be returning the image as a base64 encoded string
# which we will want compressed
app.add_middleware(GZipMiddleware, minimum_size=1000, compresslevel=7)

以上是使用 FLUX、Python 和 Diffusers 創建人工智慧驅動的圖像生成 API 服務的詳細內容。更多資訊請關注PHP中文網其他相關文章！