Stable Diffusion文生图模型训练入门实战（完整代码）

Stable Diffusion 1.5（SD1.5）是由Stability AI在2022年8月22日开源的文生图模型，是SD最经典也是社区最活跃的模型之一。

以SD1.5作为预训练模型，在火影忍者数据集上微调一个火影风格的文生图模型（非Lora方式），是学习SD训练的入门任务。

显存要求 22GB左右

在本文中，我们会使用SD-1.5模型在火影忍者数据集上做训练，同时使用SwanLab监控训练过程、评估模型效果。

• 代码：Github
• 实验日志过程：SD-naruto - SwanLab
• 模型：runwayml/stable-diffusion-v1-5
• 数据集：lambdalabs/naruto-blip-captions
• SwanLab：swanlab.cn

1.环境安装

本案例基于Python>=3.8，请在您的计算机上安装好Python；

另外，您的计算机上至少要有一张英伟达显卡（显存大约要求22GB左右）。

我们需要安装以下这几个Python库，在这之前，请确保你的环境内已安装了pytorch以及CUDA：

txt复制代码swanlab
diffusers
datasets
accelerate
torchvision
transformers

一键安装命令：

bash

复制代码pip install swanlab diffusers datasets accelerate torchvision transformers

本文的代码测试于diffusers0.29.0、accelerate0.30.1、datasets2.18.0、transformers4.41.2、swanlab==0.3.11，更多库版本可查看SwanLab记录的Python环境。

2.准备数据集

本案例是用的是火影忍者数据集，该数据集主要被用于训练文生图模型。

该数据集由1200条（图像、描述）对组成，左边是火影人物的图像，右边是对它的描述：

我们的训练任务，便是希望训练后的SD模型能够输入提示词，生成火影风格的图像：

---

数据集的大小大约700MB左右；数据集的下载方式有两种：

1. 如果你的网络与HuggingFace连接是通畅的，那么直接运行我下面提供的代码即可，它会直接通过HF的datasets库进行下载。
2. 如果网络存在问题，我也把它放到百度网盘（提取码: gtk8），下载naruto-blip-captions.zip到本地解压后，运行到与训练脚本同一目录下。

3.准备模型

这里我们使用HuggingFace上Runway发布的stable-diffusion-v1-5模型。

模型的下载方式同样有两种：

1. 如果你的网络与HuggingFace连接是通畅的，那么直接运行我下面提供的代码即可，它会直接通过HF的transformers库进行下载。
2. 如果网络存在问题，我也把它放到百度网盘（提取码: gtk8），下载stable-diffusion-v1-5.zip到本地解压后，运行到与训练脚本同一目录下。

4. 配置训练可视化工具

我们使用SwanLab来监控整个训练过程，并评估最终的模型效果。

如果你是第一次使用SwanLab，那么还需要去swanlab.cn上注册一个账号，在用户设置页面复制你的API Key，然后在训练开始时粘贴进去即可：

5.开始训练

由于训练的代码比较长，所以我把它放到了Github里，请Clone里面的代码：

bash

复制代码git clone https://github.com/Zeyi-Lin/Stable-Diffusion-Example.git

如果你与HuggingFace的网络连接通畅，那么直接运行训练：

bash复制代码python train_sd1-5_naruto.py \
  --use_ema \
  --resolution=512 --center_crop --random_flip \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --max_train_steps=15000 \
  --learning_rate=1e-05 \
  --max_grad_norm=1 \
  --seed=42 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --output_dir="sd-naruto-model"

如果你的模型或数据集用的是上面的网盘下载，那么你需要做下面的两件事：

第一步：将数据集和模型文件夹放到训练脚本同一目录下，文件结构如下：

txt复制代码|--- sd_config.py
|--- train_sd1-5_naruto.py
|--- stable-diffusion-v1-5
|--- naruto-blip-captions

stable-diffusion-v1-5是下载好的模型文件夹，naruto-blip-captions是下载好的数据集文件夹。

第二步：修改sd_config.py的代码，将pretrained_model_name_or_path和dataset_name的default值分别改为下面这样：

python复制代码    parser.add_argument(
        "--pretrained_model_name_or_path",
        type=str,
        default="./stable-diffusion-v1-5",
    )
    parser.add_argument(
        "--dataset_name",
        type=str,
        default="./naruto-blip-captions",
    )

然后运行启动命令即可。

---

看到下面的进度条即代表训练开始：

6. 训练结果演示

我们在SwanLab上查看最终的训练结果：

可以看到SD训练的特点是loss一直在震荡，随着epoch的增加，loss在最初下降后，后续的变化其实并不大：

我们来看看主观生成的图像，第一个epoch的图像长这样：

可以看到詹姆斯还是非常的“原生态”，迈克尔杰克逊生成的也怪怪的。。。

再看一下中间的状态：

经过比较长时间的训练后，效果就好了不少。

比较有意思的是，比尔盖茨生成出来的形象总是感觉非常邪恶。。。

详细训练过程看这里：SD-Naruto - SwanLab

至此，你已经完成了SD模型在火影忍者数据集上的训练。

7. 模型推理

训练好的模型会放到sd-naruto-model文件夹下，推理代码如下：

python复制代码from diffusers import StableDiffusionPipeline
import torch

model_id = "./sd-naruto-model"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda")

prompt = "Lebron James with a hat"
image = pipe(prompt).images[0]  
    
image.save("result.png")

读者福利：如果大家对大模型感兴趣，这套大模型学习资料一定对你有用

对于0基础小白入门：

如果你是零基础小白，想快速入门大模型是可以考虑的。

一方面是学习时间相对较短，学习内容更全面更集中。
二方面是可以根据这些资料规划好学习计划和方向。

资源分享

大模型AGI学习包

资料目录

1. 成长路线图&学习规划
2. 配套视频教程
3. 实战LLM
4. 人工智能比赛资料
5. AI人工智能必读书单
6. 面试题合集

《人工智能\大模型入门学习大礼包》，可以扫描下方二维码免费领取！

1.成长路线图&学习规划

要学习一门新的技术，作为新手一定要先学习成长路线图，方向不对，努力白费。

对于从来没有接触过网络安全的同学，我们帮你准备了详细的学习成长路线图&学习规划。可以说是最科学最系统的学习路线，大家跟着这个大的方向学习准没问题。

2.视频教程

很多朋友都不喜欢晦涩的文字，我也为大家准备了视频教程，其中一共有21个章节，每个章节都是当前板块的精华浓缩。

3.LLM

大家最喜欢也是最关心的LLM（大语言模型）

《人工智能\大模型入门学习大礼包》，可以扫描下方二维码免费领取！