使用自定义数据集微调 Stable Diffusion 3.5 LoRA 文生图模型

相关介绍

Stable Diffusion 3.5 是 Stability AI 于 2024 年 10 月 22 日正式发布的最新图像生成模型，是 Stable Diffusion 3.0 系列的升级版本。主要版本包括 Stable Diffusion 3.5 Large（80 亿参数）、Stable Diffusion 3.5 Large Turbo（速度提升）以及 Stable Diffusion 3.5 Medium（适合消费级硬件）。技术核心采用 MM-DiT（Modified Multimodal Diffusion Transformer），集成 Query-Key Normalization，简化微调过程。

优势包括图像质量提升、提示词遵循能力领先、多样化输出及商业友好（年收入不超过 100 万美元可免费商用）。局限性在于结构稳定性与 Flux 模型相比略逊一筹，且对硬件有一定需求（Large 版本需 10G 以上显卡）。

微调 Stable Diffusion 3.5 LoRA 文生图模型

下载项目

官方源代码：https://github.com/seochan99/stable-diffusion-3.5-text2image-lora.git

Linux 环境下克隆仓库：

git clone https://github.com/seochan99/stable-diffusion-3.5-text2image-lora.git

准备数据集

准备一个用于模型微调的数据集，目录结构如下：

dataset/
├── images/
│   ├── landscape_001.jpg
│   ├── portrait_001.jpg
│   └── ...
└── metadata.jsonl

metadata.jsonl 内容示例：

{"image":"images/landscape_001.jpg","caption":"a breathtaking mountain landscape with snow-capped peaks, a crystal clear lake reflecting the sky, and pine trees in the foreground"}
{"image":"images/portrait_001.jpg","caption":"a professional portrait of a young woman with curly brown hair, wearing a blue denim jacket, smiling softly, natural lighting"}

下载预训练模型

Hugging Face

import os
os.environ['HF_ENDPOINT']='https://hf-mirror.com'
from huggingface_hub import snapshot_download
snapshot_download(repo_id="stabilityai/stable-diffusion-3.5-large", local_dir="stabilityai/stable-diffusion-3.5-large")

Model Scope

from modelscope import snapshot_download
model_dir = snapshot_download('AI-ModelScope/stable-diffusion-3.5-medium', local_dir="stabilityai/stable-diffusion-3.5-medium")

进行训练

python stable-diffusion-3.5-text2image-lora-main/train_text_to_image_lora_sd35.py \
--pretrained_model_name_or_path stabilityai/stable-diffusion-3.5-medium \
--train_data_dir sd3_test/stable-diffusion-3.5-text2image-lora-main/examples/dataset/ \
--output_dir sd3_test/stable-diffusion-3.5-text2image-lora-main/outputs/sd35-lora \
--resolution 1024 --train_batch_size 1 --gradient_accumulation_steps 4 \
--num_train_epochs 50 --rank 64 --learning_rate 0.0001 \
--text_encoder_lr 1e-05 --lr_scheduler cosine --lr_warmup_steps 500 \
--lr_num_cycles 1 --mixed_precision fp16 --validation_epochs 2 \
--num_validation_images 2 --checkpointing_steps 500 --report_to tensorboard \
--dataloader_num_workers 2 --seed 42 --precondition_outputs 1 \
--random_flip --gradient_checkpointing --train_text_encoder

进行预测

python stable-diffusion-3.5-text2image-lora-main/inference.py \
--model_path stabilityai/stable-diffusion-3.5-medium \
--lora_path stable-diffusion-3.5-text2image-lora-main/outputs/sd35-lora/pytorch_lora_weights.safetensors \
--prompt "a professional portrait of a young woman with curly brown hair, wearing a blue denim jacket, smiling softly, natural lighting" \
--num_images 4 --seed 42

参考

[1] Patrick Esser et al., Scaling Rectified Flow Transformers for High-Resolution Image Synthesis. 2024. [2] https://huggingface.co/stabilityai/stable-diffusion-3.5-large [3] https://huggingface.co/stabilityai/stable-diffusion-3.5-medium