使用 FLUX、PyTorch 和 Streamlit 构建 AI 图像生成应用程序

# Web 开发 # 编程 # 教程 # AI

近年来，对 AI 生成图像的需求迅速增长。这些图像不仅用于艺术目的，还可用于各个行业的实际应用。例如，在时尚行业，AI 生成的图像可用于创建展示服装的虚拟模型。在汽车行业，AI 生成的图像可用于设计和测试新车型。最好的部分是什么？您现在可以在 Koyeb GPU 上运行自己的 AI 图像生成机器。

FLUX.1 [dev] 模型（由 BlackForestLabs 开发）是一种先进的 AI 图像生成模型，可产生出色的输出质量。它是一个 120 亿参数整流转换器，能够根据文本描述生成图像。它具有竞争性的提示跟随功能，并使用指导蒸馏进行训练。此外，生成的输出可用于个人、科学和非商业用途，如 FLUX.1 [dev] 非商业许可证中所述。

在本教程中，我们将学习如何设置 Streamlit 应用程序、集成 FLUX 模型进行实时图像生成，并使用 Docker 和 Koyeb 部署应用程序，确保可扩展的图像生成服务。

您可以使用下面的部署到 Koyeb 按钮部署本教程中构建的 FLUX 应用程序：

要求

要成功完成本教程，您需要满足以下条件：

Git 安装

Python 3.6+ 或更高版本

Hugging Face 账户

Koyeb 帐户

了解组件

文本到图像生成和 FLUX 模型

文本到图像的生成过程涉及模型解释输入文本并将其转换为视觉表示。此过程包括几个步骤：首先，模型将文本输入编码到潜在空间中，捕获单词的语义。接下来，它采用生成过程从该潜在空间中采样，生成与所述概念一致的图像。由 BlackForestLabs 开发的 FLUX 模型是一个最先进的 120 亿参数整流转换器，在包含图像对及其相应文本描述的大量数据集上进行训练。这种训练使模型能够学习语言和视觉内容之间的复杂关系。因此，用户可以输入详细的提示，FLUX 模型会生成准确（和创造性）反映这些提示的图像。

Streamlit

Streamlit 是一个开源 Python 库，旨在创建交互式数据应用程序（通常称为仪表板）。它使开发人员能够简单直观地构建和共享数据应用程序，而无需大量的 Web 开发专业知识。

Streamlit 应用程序以 Python 脚本的形式创建，然后在 Streamlit 环境中执行。该库提供了一组函数，可用于向应用程序添加交互元素，例如上传文件按钮。

步骤

设置环境：首先设置项目目录、安装必要的依赖项并配置环境变量。

设置 Streamlit：接下来，安装 Streamlit 并为您的应用程序创建初始用户界面。

使用FLUX模型生成AI图像：使用FLUX模型根据用户的提示和配置生成AI图像。

将 Streamlit 应用程序 Docker 化：创建一个 Dockerfile 来容器化您的应用程序，以实现一致的部署。

部署到 Koyeb：最后，在 Koyeb 平台上部署您的应用程序。

设置环境

让我们从创建一个新的 Streamlit 项目开始。为了保持 Python 依赖项井然有序，您应该创建一个虚拟环境。

首先，创建并导航到本地目录：

# Create and move to the new directory
mkdir example-koyeb-flux-images
cd example-koyeb-flux-images

之后，创建并激活一个新的虚拟环境：

# Create a virtual environment
python -m venv venv

# Active the virtual environment (Windows)
.\venv\Scripts\activate.bat

# Active the virtual environment (Linux)
source ./venv/bin/activate

现在，创建一个具有以下依赖项的“requirements.txt”文件：

streamlit
watchdog
diffusers
torch
torchvision
einops
huggingface_hub[hf_transfer]
safetensors
sentencepiece
transformers
tokenizers
protobuf
requests
invisible-watermark
accelerate
peft

在依赖项方面，我们包含了用于用 Python 构建 Web 应用的 Streamlit、用于实时本地使用 FLUX 模型的 Hugging Face 以及用于监视文件系统事件的 watchdog。`accelerate` 包用于简化机器学习模型的训练和推理过程，使管理分布式训练和优化性能变得更加容易。`invisible-watermark` 包用于在生成的内容中嵌入不可见的水印，确保模型生成的输出的唯一性。

现在，您可以使用以下命令安装依赖项：

pip install -r requirements.txt

现在，让我们继续创建一个新的 Streamlit 项目。

设置 Streamlit

在此步骤中，您将设置 Streamlit UI，它将定义页面的视觉布局，以及用户加载 FLUX 模型以开始使用 AI 生成图像的能力。项目的所有逻辑都将驻留在此文件中，因此您可以先使用以下代码创建一个“app.py”文件：

# File: app.py

import os, torch, streamlit
from diffusers import FluxPipeline
from huggingface_hub import login, snapshot_download

# Log in to Hugging Face using the provided token from environment variables
login(token=os.getenv('HF_TOKEN'))

# Set the title of the Streamlit application
streamlit.title("AI Image Generation with FLUX.1-dev")

# Create a text input field for the user to enter a prompt for image generation
prompt = streamlit.text_input("Enter your prompt:", "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k")

# Create four columns for user inputs
col1, col2, col3, col4 = streamlit.columns(4)

# Column 1: Input for image width
with col1:
    num_width = streamlit.number_input("Width:", min_value=1, value=256)

# Column 2: Input for image height
with col2:
    num_height = streamlit.number_input("Height:", min_value=1, value=256)

# Column 3: Input for the number of images to generate
with col3:
    num_images = streamlit.number_input("Images:", min_value=1, value=3)

# Column 4: Input for the number of inference steps
with col4:
    num_inference_steps = streamlit.number_input("Steps:", min_value=1, value=3)

# Button to load the model from Hugging Face
if streamlit.button("Load Model"):
    snapshot_download(repo_id="black-forest-labs/FLUX.1-dev", cache_dir="./FLUX_1_dev")

上面的代码执行以下操作：

从导入 Streamlit 模块开始。

使用存储在环境变量中的令牌登录 Hugging Face，允许访问模型存储库。

定义每个步骤和页面标题及说明的文本值

创建四列来组织用户对图像尺寸、图像数量和推理步骤的输入。每列包含一个数字输入字段，其中指定了最小值和默认值。

创建一个“加载模型”按钮，用于从 Hugging Face 加载 FLUX 模型。单击后，它会使用 snapping_download 函数将模型快照下载到指定的缓存目录。

这样，您就设置了一个可以根据用户请求加载 FLUX 模型的 UI。现在，让我们继续根据用户提示和设置生成 AI 图像。

使用 FLUX 模型生成 AI 图像

在此步骤中，您将调用 FLUX 模型根据用户提示生成 AI 图像。根据输入中设置的默认值，模型将为用户的每个提示生成 3 张图像。在 `app.py` 文件中进行以下添加：

# File: app.py

# Existing code

# Button to generate images using the FLUX model
if streamlit.button("Generate Image"):
    # Load the FLUX model with specified parameters
    pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, cache_dir="./FLUX_1_dev")

    # Generate images based on the user input parameters
    images = pipe(prompt=prompt, num_inference_steps=num_inference_steps, height=num_height, width=num_width, num_images_per_prompt=num_images).images

    # Create three columns to display the generated images
    cols = streamlit.columns(3)

    # Loop through the generated images and display them in the columns
    for i, image in enumerate(images):
        cols[i % 3].image(image, caption=prompt)

上面的代码添加执行以下操作：

创建一个名为“生成图像”的按钮。单击后，它将实例化用于文本到图像生成的 Flux 管道。

使用用户配置的参数调用管道来生成所需的高度和宽度、要采取的推理步骤数以及要生成的图像数量。

创建三列来组织生成的图像。

现在，你可以使用以下方法运行 Streamlit 应用程序：

streamlit run ./app.py --server.port 8000

该应用程序现在已在 http://localhost:8000 上准备就绪。通过更改默认输入值和您自己的自定义提示来生成 AI 图像，从而测试该应用程序。

现在，让我们将应用程序docker化，以确保多个部署之间的一致性。

将 Streamlit 应用程序 Docker 化

Docker 化部署有助于创建一致且可重复的环境，确保应用程序在任何系统上以相同的方式运行。它简化了依赖项管理并增强了可扩展性，使部署更加高效和可靠。要进行 docker 化，请使用以下代码在项目的根目录中创建一个“Dockerfile”：

FROM python:3.12 AS runner

WORKDIR /app

COPY app.py requirements.txt .

RUN pip install -r requirements.txt --root-user-action=ignore
RUN pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121 --root-user-action=ignore # [!code ++]

RUN apt-get update 
RUN apt-get install -y libsm6 libxext6 git git-lfs # [!code ++]
RUN rm -rf /var/lib/apt/lists/*

EXPOSE 8000 # [!code ++]

ENV HF_HUB_ENABLE_HF_TRANSFER=1

ARG HF_TOKEN # [!code ++]
CMD ["streamlit", "run", "./app.py", "--server.port", "8000"] # [!code ++]

除了用于部署 Python 应用程序的常用 Dockerfile 之外，还对此代码进行了以下调整/添加：

pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121 --root-user-action=ignore 安装支持 CUDA 的 PyTorch 以实现 GPU 加速。

RUN apt-get update && apt-get install -y libsm6 libxext6 git git-lfs && rm -rf /var/lib/apt/lists/* 用于安装 git-lfs 和 git，然后清理包列表以减小镜像大小。

EXPOSE 8000 用于指定 Streamlit 应用程序将运行的端口。

CMD ["streamlit", "run", "./app.py", "--server.port", "8000"] 用于定义在端口 8000 上启动 Streamlit 应用程序的命令。

完成所有配置后，让我们继续将应用程序部署到 Koyeb。

部署至 Koyeb

现在您已经在本地运行该应用程序，您还可以将其部署在 Koyeb 上并使其在互联网上可用。

在您的 GitHub 帐户上创建一个新的存储库，以便您可以推送您的代码。

您可以从 GitHub 下载 Python 的标准 .gitignore 文件，以排除将某些目录和文件推送到存储库：

curl -L https://raw.githubusercontent.com/github/gitignore/main/Python.gitignore -o .gitignore

在终端中运行以下命令来提交并将代码推送到存储库：

git init
git add .
git commit -m "first commit"
git branch -M main
git remote add origin [Your GitHub repository URL]
git push -u origin main

现在，您应该已将所有本地代码保存在远程存储库中。现在是时候部署应用程序了。

在 Koyeb 控制面板中的“概述”选项卡上，单击“创建 Web 服务”以启动应用程序创建和部署过程。

选择 GitHub 作为部署方法。

从菜单中选择您的存储库。或者，通过在公共存储库字段中输入 https://github.com/koyeb/example-flux-1-pytorch 来从与本教程相关联的示例存储库进行部署。

在实例选择中，选择一个 GPU 实例。

在 HF_TOKEN 环境变量中设置您的 HuggingFace 访问令牌。

最后，单击部署按钮。

应用程序部署完成后，您可以访问 Koyeb 服务 URL（以 `.koyeb.app` 结尾）来访问 Streamlit 应用程序。

结论

在本教程中，您使用带有 Streamlit 框架的 FLUX dev 模型构建了一个 AI 图像生成应用程序。在此过程中，您学习了如何在 Python 中调用 Flux 管道来随时生成 AI 图像，使用用户提供的提示和设置进行自定义，并使用 Streamlit 框架快速制作用户界面原型。

鉴于应用程序是使用 Git 部署选项部署的，随后将代码推送到已部署的分支将自动为您的应用程序启动新的构建。部署成功后，对应用程序的更改将生效。如果部署失败，Koyeb 会保留最后一次可运行的生产部署，确保您的应用程序不间断运行。

CLIS.CC