从自拍到专业摄影：用 DreamBooth 训练 Flux.1

作者 : Rafay Farhan 在 DreamAI 软件公司（私人）有限公司

你是否曾经想创建看起来专业的头像，而不需要走进摄影工作室？作为一个AI爱好者和开发者，我最近使用先进的图像生成模型 Flux.1，以及 DreamBooth 的微调，来解决这个挑战。今天，我很高兴和大家分享我的旅程，并提供一些易于使用的笔记本，简化整个过程。

包含训练和测试笔记本的 GitHub 存储库可以在这里找到。

挑战：让使用 DreamBooth 训练 Flux 更加简单

在探索 AI 领域时，我注意到一个显著的空白：尽管 Flux.1 拥有令人印象深刻的能力，但并没有很多用户友好的资源展示如何用 DreamBooth 进行训练。官方文档存在于 Hugging Face 的 diffusers 存储库 中，但我希望让这项强大的技术对更广泛的社区更加易用。

因此，我创建了两个 Colab 笔记本，简化了这个过程，建立在 diffusers 团队的优秀工作基础上。这些笔记本让训练Flux.1在你自己的图像上变得前所未有的简单，并生成惊艳的结果。

数据集准备：少即是多？

在我的实验中，我使用了一个相对较小的个人照片数据集。虽然传统观念可能认为更多的数据总是更好，但真正令人着迷的是，Flux.1 就算在有限的训练数据下仍能产生令人满意的结果。

[原始数据集的 3/11 张图像]: 这些训练图像虽然不是专业的头像——它们是低质量的休闲照片和自拍，角度、光照、发型和背景不一，拍摄时间间隔几个月，但 Flux.1 仍然能够非常好地学习到我的面部特征。想象一下，如果有一个更精心整理的数据集，它能够做得多么出色！

但是，创建最佳数据集的一些建议如下：

• 目标是清晰、光线良好的面部照片
• 包含各种角度和表情
• 确保图像质量一致
• 删除任何模糊或遮挡的照片
• 考虑包括特写镜头和中景镜头

训练过程：两阶段方法

我开发的训练过程分为两个阶段，这被证明尤其有效。让我们看看关键组成部分：

首先，我们需要安装必要的依赖项：

!pip install --upgrade pip
!pip install accelerate transformers diffusers torch torch-fidelity mediapy wandb
!pip install -q xformers
!pip install scipy ftfy
!pip install -q triton prodigyopt
!pip install -q peft>=0.6.0

# 从源代码安装 diffusers
!git clone https://github.com/huggingface/diffusers
%cd diffusers
!pip install -e .
%cd examples/dreambooth
!pip install -r requirements_flux.txt

第一阶段：基础模型训练

首先，我们使用 DreamBooth 从头开始训练一个基础模型。这为模型理解你的面部特征和特性奠定了基础。下面是我们设置训练配置的方式：

# 设置训练配置
MODEL_NAME = "black-forest-labs/FLUX.1-schnell"
INSTANCE_DIR = "portrait_images"  # 存放训练图像的目录
OUTPUT_DIR = "latest-flux-your-identifier"
INSTANCE_PROMPT = "a photo of <your-identifier> person"  # 替换为你的身份标识

# 使用优化参数启动训练
!accelerate launch train_dreambooth_lora_flux.py 
  --pretrained_model_name_or_path="$MODEL_NAME" 
  --instance_data_dir="$INSTANCE_DIR" 
  --output_dir="$OUTPUT_DIR" 
  --mixed_precision="bf16" 
  --instance_prompt="$INSTANCE_PROMPT" 
  --resolution=256 
  --train_batch_size=1 
  --gradient_accumulation_steps=8 
  --learning_rate=1.0 
  --optimizer="prodigy" 
  --lr_scheduler="constant" 
  --lr_warmup_steps=0 
  --max_train_steps=500 
  --seed=42 
  --gradient_checkpointing 
  --cache_latents 
  --lora_layers="attn.to_k,attn.to_q,attn.to_v,attn.to_out.0" 
  --rank=4

这些参数经过精心调整，以获得使用人像图像的最佳结果。

第二阶段：精炼

在评估初步结果后，我使用第一阶段的检查点进行了额外的 500 步训练。这个精炼阶段有助于模型捕捉更细致的细节，并提高输出质量。以下是设置精炼训练的方式：

# 为继续训练设置目录
OUTPUT_DIR = "latest-flux-your-identifier-v2"  # 精炼模型的新输出目录

# 复制之前训练的检查点文件
!mkdir -p $OUTPUT_DIR
!cp previous-checkpoint/pytorch_lora_weights.safetensors $OUTPUT_DIR/
!cp previous-checkpoint/optimizer.bin $OUTPUT_DIR/
!cp previous-checkpoint/scheduler.bin $OUTPUT_DIR/
!cp previous-checkpoint/random_states_0 $OUTPUT_DIR/

# 使用相同的参数启动精炼训练
!accelerate launch train_dreambooth_lora_flux.py 
  --pretrained_model_name_or_path="$MODEL_NAME" 
  --instance_data_dir="$INSTANCE_DIR" 
  --output_dir="$OUTPUT_DIR" 
  # ... (与第一阶段相同的参数)

训练笔记本包含注释部分，清楚地指示出根据你是从头开始还是精炼现有模型使用哪些代码。

测试和生成：用户友好的界面

这个项目中最激动人心的方面之一是我使用 Gradio 开发的测试界面。它提供了一种直观的方式来使用你的训练模型生成图像。这个界面允许你：

• 输入自定义提示
• 调整生成参数
• 实时查看结果
• 实验不同的风格和设置

下面是我们如何设置的：

import torch
from diffusers import FluxPipeline
import gradio as gr
import random

# 初始化模型
def load_model():
    pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.float16)
    pipe.load_lora_weights("latest-flux-rffzzgh-v2.1")
    return pipe.to("cuda")

# 生成图像函数
def generate_portrait(prompt_prefix, prompt_suffix, negative_prompt, guidance_scale, num_steps, seed):
    # 使用所需的身份标识构建完整提示
    full_prompt = f"a professional studio portrait of rffzzgh person, {prompt_prefix}, {prompt_suffix}"

    # 如果未提供，则生成随机种子
    if seed == -1:
        seed = random.randint(0, 999999)

    # 生成图像
    with torch.cuda.amp.autocast():
        image = pipe(
            full_prompt,
            negative_prompt=negative_prompt,
            guidance_scale=float(guidance_scale),
            num_inference_steps=int(num_steps),
            max_sequence_length=256,
            generator=torch.Generator("cuda").manual_seed(seed)
        ).images[0]

    return image, seed

# 创建 Gradio 界面
def create_interface():
    with gr.Blocks(title="Flux Portrait Generator") as interface:
        gr.Markdown("# Flux Portrait Generator")

        with gr.Row():
            with gr.Column():
                # 输入组件
                prompt_prefix = gr.Textbox(
                    label="提示前缀（风格、设置等）",
                    placeholder="惊艳的 4k UHD，专业的工作室照明",
                    value="惊艳的 4k UHD，专业的工作室照明"
                )

                prompt_suffix = gr.Textbox(
                    label="提示后缀（附加细节）",
                    placeholder="清晰聚焦，详细特征，高端时尚",
                    value="清晰聚焦，详细特征，高端时尚，清晰细节，专业摄影，获奖肖像，使用 Phase One 拍摄，完美构图，杂志质量"
                )

                negative_prompt = gr.Textbox(
                    label="消极提示",
                    value="模糊，低质量，扭曲，变形，丑陋，解剖结构差，比例失调，业余，过饱和，过曝，欠曝，颗粒感，噪音，失焦，重复，变异"
                )

                with gr.Row():
                    guidance_scale = gr.Slider(
                        label="引导比例",
                        minimum=0.0,
                        maximum=15.0,
                        value=7.5,
                        step=0.5
                    )
                    num_steps = gr.Slider(
                        label="步数",
                        minimum=1,
                        maximum=50,
                        value=4,
                        step=1
                    )

                seed = gr.Number(
                    label="种子（-1 为随机）",
                    value=-1,
                    precision=0
                )

                generate_btn = gr.Button("生成肖像", variant="primary")

            with gr.Column():
                # 输出组件
                output_image = gr.Image(label="生成肖像")
                used_seed = gr.Number(label="使用的种子", precision=0)

        # 处理生成
        generate_btn.click(
            fn=generate_portrait,
            inputs=[prompt_prefix, prompt_suffix, negative_prompt, guidance_scale, num_steps, seed],
            outputs=[output_image, used_seed]
        )

    return interface

# 全局加载模型
print("加载模型中...")
pipe = load_model()
print("模型加载完成!")

# 启动界面
if __name__ == "__main__":
    interface = create_interface()
    interface.launch(share=True)  # 如果不想公开共享请设置 share=False

结果

实施细节

对于那些对技术细节感兴趣的人，这个项目利用了：

• Black Forest Labs 的 Flux.1 模型
• Hugging Face diffusers 库中的 DreamBooth 训练实现
• Gradio 用于用户界面
• Google Colab 用于易于访问的 GPU 计算

硬件要求

为了获得最佳性能，我推荐：

• GPU: NVIDIA A100（40GB）— 我是通过 Google Colab 使用这个
• VRAM: 至少 16GB（尽管 40GB 提供更好的性能）
• 磁盘空间：~10GB 用于模型存储和生成的图像

注意：虽然你可能在低端 GPU 上运行这个，但使用 A100 或类似的高端 GPU 将显著提高训练速度和稳定性。

软件设置

笔记本处理所有复杂的设置，包括：

• 安装依赖项
• 加载和配置模型
• 设置训练参数
• 管理模型检查点

结果与洞察

结果确实令人印象深刻。即使在一个适度的数据集上，模型也学会生成高质量的、工作室风格的头像，同时保持我的相貌，并增加了专业的光线和构图。这说明了将 Flux.1 的能力与 DreamBooth 微调相结合的强大效果。

开始使用

准备好自己尝试了吗？你可以在我的存储库中找到这些笔记本。以下是快速入门指南，包括一些有效的示例提示：

# 示例提示配置
prompt_prefix = "惊艳的 4k UHD，专业的工作室照明"
prompt_suffix = "清晰聚焦，详细特征，高端时尚，清晰细节，专业摄影，获奖肖像，使用 Phase One 拍摄，完美构图，杂志质量"

negative_prompt = """模糊，低质量，扭曲，变形，丑陋， 
                    解剖结构差，比例失调，业余，过饱和， 
                    过曝，欠曝，颗粒感，噪音，失焦， 
                    重复，变异"""

# 生成参数
guidance_scale = 7.5
num_inference_steps = 4
seed = 42  # 使用 -1 作为随机种子

1. 准备你的个人照片数据集 2. 在 Colab 中打开训练笔记本 3. 按照设置说明进行操作 4. 运行训练过程 5. 使用带有 Gradio 界面的测试笔记本生成图像

致谢

非常感谢 Hugging Face diffusers 团队出色的 DreamBooth 实现和文档。本项目在他们的基础上扩展，旨在使这项技术对所有人更具可访问性。

结论

这个项目展示了如何将强大的 AI 图像生成工具变得更具可访问性。不论你是开发者、艺术家，还是出于好奇想了解 AI，这些笔记本提供了一种简单明了的方式来尝试自定义模型训练。

Flux.1 和 DreamBooth 的结合为个性化图像生成提供了令人兴奋的可能性，希望这些工具能帮助其他人更容易地探索这项技术。