使用 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 



<p>FastAPI 是一个使用 Python 构建 Web API 的流行框架。在本例中，我们使用它来创建一个可以接受图像生成请求的服务器。我们还使用 GZip 中间件来压缩响应，这在以 Base64 格式发送回图像时特别有用。</p>

<blockquote>
<p>在生产环境中，您可能希望将生成的图像存储在 S3 存储桶或其他云存储中，并返回 URL 而不是 Base64 编码的字符串，以利用 CDN 和其他优化。</p>
</blockquote>

<h2>
  
  
  步骤 4：定义请求模型
</h2>

<p>我们现在需要为我们的 API 将接受的请求定义一个模型。<br>
</p>

<pre class="brush:php;toolbar:false">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 



<blockquote>
<p>此命令仅适用于安装了curl、jq 和base64 实用程序的基于UNIX 的系统。根据托管 FLUX 服务器的硬件，它也可能需要几分钟才能完成。</p>
</blockquote>

<h2>
  
  
  结论
</h2>

<p>恭喜！您已经使用 Python 成功创建了自己的 FLUX 服务器。此设置允许您通过简单的 API 根据文本提示生成图像。如果您对基本 FLUX 模型的结果不满意，您可以考虑微调模型，以便在特定用例上获得更好的性能。</p>

<h2>
  
  
  完整代码
</h2>

<p>您可以在下面找到本指南中使用的完整代码：<br>
</p>

<pre class="brush:php;toolbar:false">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中文网其他相关文章！