我在我的脸上训练了 FLUX.1 (你也可以这样做)

使用 Python 和 Replicate 的简易微调指南

虽然大型语言模型 (LLMs) 最近备受关注，但图像生成模型同样快速发展。目前的尖端技术是 FLUX.1，这是来自黑森林实验室 (Stable Diffusion 团队的一个分支) 的图像模型。在本文中，我分享了我用来微调这个模型的完整过程，以便生成我自己的无限高质量照片。

微调模型输出 (采用了丰厚的头发和体态调整)。

FLUX.1 是一个可以处理文本和图像的模型。更具体地说，它 接受文本和图像输入并生成图像输出。这使得您可以仅通过文本提示生成图像，或者使用文本编辑图像！

该模型有三个版本：Pro、Dev 和 Schnell (德语中的“快”)。Pro 是表现最佳的模型，而 Schnell 是最快和最便宜的。

虽然这些模型可以做惊人的事情，但它们 无法生成特定人物的照片 (除非这些人是知名公共人物)。这就是微调能为我们提供帮助的地方。

微调 FLUX.1

微调 是指通过 额外训练使模型适应特定用例。这很有价值，因为它通常能产生比更大（更昂贵）的模型更好的效果。

油管

高效微调

虽然微调比从头开始训练模型要简单和便宜得多，但在处理像 FLUX.1 这样具有 120 亿参数的模型时，仍然是一项重大任务。作为参考，120 亿参数的模型在 FP32 表示下，将需要 48GB 的内存进行推理 和 超过 200GB 用于完整微调！

由于这需要专用的 GPU 集群，让我们采取更聪明的方法。换句话说，我们可以使用一种称为 LoRA（低秩适应） 的技术，它 在保持基础参数不变的情况下，向模型添加相对少量新的可训练参数，而不是重新训练所有模型参数。

关于 LoRA 的更多背景知识，请查看这个链接。

示例：使用 Python 和 Replicate 微调 FLUX.1

在了解了背景之后，让我们通过 Python 代码来逐步演示这个过程。这基本上遵循了 Replicate 的博客上的指南[1]。另外，感谢 Jarret Burkett，他通过在 Replicate 上发布 训练端点 简化了这个过程[2]。

油管中可以免费获得。

需求：

• Python 3.8+
• 一个 Replicate API 密钥
• 10–20 张高质量的个人照片 (至少 1024 x 1024)

导入

首先，我们需要导入几个 Python 包。关键的是 replicate，我们将使用它来生成图像描述，微调 FLUX，以及使用微调后的模型。

from PIL import Image
import os
import shutil

import replicate
from dotenv import load_dotenv

接下来，我从一个秘密的 .env 文件中导入我的 replicate API 密钥。或者，您可以将此变量添加到全局环境中（或硬编码它）。

# load vars from .env
load_dotenv()

# set replicate api key
replicate = replicate.Client(api_token=os.getenv("REPLICATE_API_KEY"))

筛选训练图像

在我的训练运行过程中，我将 20 张我在各种场景下的图片放在名为 raw/ 的文件夹中。以下是其中一些图像。

用于训练的原始照片。图片来自作者。

请注意，这些图片的大小和分辨率各不相同。在用这些图片训练 FLUX 之前，我们需要进行一些关键的预处理步骤。具体来说，我们需要使每张图片变为正方形、将尺寸减少到 1024 x 1024 (它们的质量都比这更高)，并以 .png (或 .jpeg) 格式保存。

# Set input and output directories
input_dir = "raw/"  # 将其更改为您的文件夹
output_dir = "data/"

# Ensure output directory exists
os.makedirs(output_dir, exist_ok=True)

# image counter
i = 0

# Process each image
for filename in os.listdir(input_dir):
    # Check for PNG and JPEG
    if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
        img_path = os.path.join(input_dir, filename)
        img = Image.open(img_path)

        # Get original dimensions
        width, height = img.size
        new_size = min(width, height)  # Smallest dimension for square crop

        # Calculate cropping box (centered)
        left = (width - new_size) // 2
        top = (height - new_size) // 2
        right = left + new_size
        bottom = top + new_size

        # Crop the image
        img_cropped = img.crop((left, top, right, bottom))

        # Resize to 1024x1024 using high-quality resampling
        img_resized = img_cropped.resize((1024, 1024), Image.LANCZOS)

        # Crop and save the image
        output_path = os.path.join(output_dir, f"img-{i}.png")
        img_resized.save(output_path, format="PNG")
        i += 1

下面是处理后这些图像的样子。

用于训练的图片 (处理后)。图片来自作者。

生成描述

现在我们已经准备好训练的图像，接下来需要做的是为它们撰写描述。

这一步的一个关键点是要在每个描述中包含一个 “触发标记”。这个是 模型可以学习和我的面部关联的唯一字符序列。在这里，我使用的触发标记是 shw-tlb。

一个示例描述可以是这样的：

一张 shw-tlb 的照片。一个光头并留有胡须的男人穿着黑色衬衫，坐在白色砖墙旁。

虽然我可以手动生成这些描述，但我使用了一个多模态 LLM (llava-13b) 通过 replicate 的 API 为我撰写这些描述。

for filename in os.listdir(output_dir):
    if filename.lower().endswith('.png'):
        img_path = os.path.join(output_dir, filename)
        image = open(img_path, "rb")

        # generate caption
        output = replicate.run(
            "yorickvp/llava-13b:a0fdc44e4f2e1f20f2bb4e27846899953ac8e66c5886c5878fa1d6b73ce009e5",
            input={
                "image": image,
                "prompt": "请为这张图片写描述。"
            }
        )
        # extract text from output
        response = " ".join(list(output))

        # add trigger token "shw-tlb"
        caption = f"一张 shw-tlb 的照片。{response}"
     
        # save caption as text file
        caption_filename = img_path.split('.')[0] + '.txt'
        with open(caption_filename, "w") as file:
            print(caption)
            file.write(caption)

到最后，我们有一个名为 data/ 的文件夹，其中包含 20 张预处理过的我的图片和对应的描述。

作为一个文件步骤，我们可以将这个文件夹压缩成一个 .zip 文件。

# inputs: name of .zip file, compression format, directory to compress
shutil.make_archive("data", 'zip', "data/")

微调模型

接下来，我们可以微调我们的模型。这实际上是过程中的最简单部分。我们甚至不需要写代码，因为有一个 web 界面 可以完成我即将编程进行的所有工作 (如果您这么做，可以跳到下一部分)。

不过，通过 Python 的过程仍然很简单。首先，我们需要在 Replicate 上 创建一个新模型。

model = replicate.models.create(
    owner="shawhint",
    name="flux-shaw-ft",
    visibility="public",  # 或者选择 "private" 如果你更喜欢
    hardware="gpu-t4",  # Replicate 会为微调模型覆盖此设置
    description="FLUX.1 基于我的照片微调（Shaw）"
)
     
print(f"模型创建成功: {model.name}")
print(f"模型 URL: https://replicate.com/{model.owner}/{model.name}")

然后，我们可以将数据上传到训练端点并 微调模型。

# Now use this model as the destination for your training
training = replicate.trainings.create(
    version="ostris/flux-dev-lora-trainer:4ffd32160efd92e956d39c5338a9b8fbafca58e03f791f6d8011f3e20e8ea6fa",
    input={
        "input_images": open("data.zip", "rb"),
        "steps": 1000,
    },
    destination=f"{model.owner}/{model.name}"
)
     
print(f"训练开始: {training.status}")
print(f"训练 URL: https://replicate.com/p/{training.id}")

训练 运行了 12 分钟，花费了我 1.12 美元。

使用模型

一旦训练完成，我们的模型将在 Replicate 上公开可用 (如果您在上述代码中将 visibility 设置为“私有”，则为私有)。默认会创建一个 Playground 界面，使我们能够通过 web 界面使用该模型。

以下是我的一些杰作…

微调模型输出 (采用了丰厚的头发和体态调整)。

另外，我们可以以编程方式访问微调后的模型。如果您正在构建自定义软件解决方案或 UI，这样更方便。

output = replicate.run(
    "shawhint/flux-shaw-ft:8b55b2745adc5e7f13fdee9931b2fe622fc7f12a4780e5c48d42991c19699a72",
    input={
        "prompt": "shw-tlb 的肖像照片，穿着深色西装的男人",
        "num_inference_steps": 28,
        "guidance_scale": 7.5,
        "model": "dev",
        "num_outputs": 4,
    }
)

for url in output:
    print(url)

# >> https://replicate.delivery/xezq/tr94REeVCpQtMKi2KElmyVaQBwzeTdmVhkJUVLgszjKT7EMUA/out-0.webp
# >> https://replicate.delivery/xezq/GFWxS22aHXKjBRc0eAOEfI1rueKDJqzdYM9JeajnLDKMtTwQB/out-1.webp
# >> https://replicate.delivery/xezq/WFmEIVUCO56WNxdEfr99tdln4OezSJe7fCRWcGDzLWtPtTwQB/out-2.webp
# >> https://replicate.delivery/xezq/k7f3OimVsb3iaav8jaypjxQ7cBnfvrTLZOMErUGLZilT7EMUA/out-3.webp

接下来是什么？

油管

尽管 Replicate 使这个过程变得非常简单，但它有一定的经济成本，并且需要通过他们的平台进行推理。因此，在未来的文章中，我想回顾一下如何在本地完成整个过程。

如果您有任何问题，请在评论中告诉我 :)

🗞️ 获取独家 AI 资源和项目创意 : https://the-data-entrepreneurs.kit.com/shaw

🧑‍🎓 在 6 周内通过构建来学习 AI : https://maven.com/shaw-talebi/ai-builders-bootcamp

参考文献

[1] https://replicate.com/blog/fine-tune-flux

[2] https://replicate.com/ostris/flux-dev-lora-trainer/train