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左右;数据集的下载方式有两种:
- 如果你的网络与HuggingFace连接是通畅的,那么直接运行我下面提供的代码即可,它会直接通过HF的
datasets
库进行下载。 - 如果网络存在问题,我也把它放到百度网盘(提取码: gtk8),下载
naruto-blip-captions.zip
到本地解压后,运行到与训练脚本同一目录下。
3.准备模型
这里我们使用HuggingFace上Runway发布的stable-diffusion-v1-5模型。
模型的下载方式同样有两种:
- 如果你的网络与HuggingFace连接是通畅的,那么直接运行我下面提供的代码即可,它会直接通过HF的
transformers
库进行下载。 - 如果网络存在问题,我也把它放到百度网盘(提取码: 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学习包
资料目录
- 成长路线图&学习规划
- 配套视频教程
- 实战LLM
- 人工智能比赛资料
- AI人工智能必读书单
- 面试题合集
《人工智能\大模型入门学习大礼包》,可以扫描下方二维码免费领取!
![](https://img-blog.csdnimg.cn/img_convert/b151988a5e7f48e9dba74448b50edd4c.png)
1.成长路线图&学习规划
要学习一门新的技术,作为新手一定要先学习成长路线图,方向不对,努力白费。
对于从来没有接触过网络安全的同学,我们帮你准备了详细的学习成长路线图&学习规划。可以说是最科学最系统的学习路线,大家跟着这个大的方向学习准没问题。
2.视频教程
很多朋友都不喜欢晦涩的文字,我也为大家准备了视频教程,其中一共有21个章节,每个章节都是当前板块的精华浓缩。
3.LLM
大家最喜欢也是最关心的LLM(大语言模型)
《人工智能\大模型入门学习大礼包》,可以扫描下方二维码免费领取!
![](https://img-blog.csdnimg.cn/img_convert/b151988a5e7f48e9dba74448b50edd4c.png)