Stable Diffusion 3.5 LoRA 微调指南

使用示例

以下是如何使用 StableDiffusionDataset 类的完整示例，包括数据集创建、样本查看和 DataLoader 构建：

# 创建数据集实例
dataset = StableDiffusionDataset("data.json")
# 查看数据集大小
print(f"数据集包含 {len(dataset)} 个样本")
# 获取单个样本
image, caption = dataset[0]
print(f"图像维度：{image.shape}")
print(f"文本描述：{caption}")
# 创建 DataLoader 用于批量训练
from torch.utils.data import DataLoader
dataloader = DataLoader(
    dataset,
    batch_size=4,  # 批次大小，可根据 GPU 内存调整
    shuffle=True,  # 训练时打乱数据，增加随机性
    num_workers=2,  # 并行加载进程数，加速数据加载
    pin_memory=True  # 启用内存锁定，加速数据传输到 GPU
)

2. LoRA 微调原理

LoRA 是一种参数高效微调技术，其核心思想是：

1. 冻结原有模型：保持预训练模型的权重不变，避免灾难性遗忘
2. 添加低秩适配器：在关键层（如注意力层）插入低秩矩阵对（A 和 B）
3. 仅训练低秩矩阵：通过少量参数更新即可实现模型定制

这种方法的优势在于：

• 训练参数仅为原有模型的 1%-5%，大幅降低内存占用
• 训练速度显著提升，减少计算资源消耗
• 微调结果易于保存和分享，单个 LoRA 权重文件通常仅几 MB
• 支持多 LoRA 权重组合使用，实现灵活的风格控制

3. 模型加载与 LoRA 配置

3.1 加载预训练模型

首先需要加载 Stable Diffusion 3.5 预训练模型：

from diffusers import StableDiffusion3Pipeline
import torch

# 加载预训练模型
model_id = "stabilityai/stable-diffusion-3.5-large"
pipeline = StableDiffusion3Pipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16,  # 使用半精度（float16）加速计算并减少内存占用
).to("cuda")  # 移至 GPU 设备

3.2 配置 LoRA 参数

LoRA 的关键参数包括：

• r：低秩矩阵的秩，控制适配器的容量（常用值：4, 8, 16, 32）
• alpha：缩放因子，控制 LoRA 对模型的影响程度
• target_modules：需要添加 LoRA 适配器的目标层

from peft import LoraConfig, get_peft_model

# 配置 LoRA 参数
lora_config = LoraConfig(
    r=16,  # 低秩矩阵的秩，r 越大，适配器容量越大，但参数也越多
    lora_alpha=32,  # 缩放因子，通常设置为 r 的 2 倍
    target_modules=['to_k',  # 注意力层的键（Key）投影层
                    'to_q',  # 注意力层的查询（Query）投影层
                    'to_v'],  # 注意力层的值（Value）投影层
    lora_dropout=0.05,  # Dropout 率，防止过拟合
    bias="none",  # 不对偏置项应用 LoRA
    task_type="TEXT_TO_IMAGE"  # 任务类型
)

# 为模型添加 LoRA 适配器
pipeline.transformer = get_peft_model(pipeline.transformer, lora_config)

# 冻结不需要训练的组件，仅训练 LoRA 适配器
pipeline.vae.requires_grad_(False)  # 冻结 VAE 编码器/解码器
pipeline.text_encoder.requires_grad_(False)  # 冻结文本编码器 1
pipeline.text_encoder_2.requires_grad_(False)  # 冻结文本编码器 2
pipeline.text_encoder_3.requires_grad_(False)  # 冻结文本编码器 3

# 打印可训练参数数量
print("可训练参数数量:", sum(p.numel() for p in pipeline.transformer.parameters() if p.requires_grad))

4. 训练循环实现

4.1 定义训练参数

import torch.nn.functional as F
from transformers import get_scheduler

# 训练参数配置
epochs = 5  # 训练轮次
batch_size = 4  # 批次大小
learning_rate = 1e-4  # 学习率
weight_decay = 1e-2  # 权重衰减，防止过拟合

# 优化器配置：仅优化可训练参数（LoRA 适配器参数）
optimizer = torch.optim.AdamW(
    params=filter(lambda p: p.requires_grad, pipeline.transformer.parameters()),
    lr=learning_rate,
    weight_decay=weight_decay
)

# 学习率调度器：使用余弦退火调度，逐步降低学习率
num_training_steps = len(dataloader) * epochs
scheduler = get_scheduler(
    name="cosine",  # 调度器类型
    optimizer=optimizer,
    num_warmup_steps=0,  # 预热步数
    num_training_steps=num_training_steps
)

4.2 执行训练循环

# 设置模型为训练模式
pipeline.transformer.train()
device = 'cuda:0'  # 指定 GPU 设备

def compute_density_for_timestep_sampling(batch_size, device):
    """
    基于正态分布的时间步采样，增加训练稳定性
    Args:
        batch_size: 批次大小
        device: 设备类型
    Returns:
        torch.Tensor: 采样的时间步权重，形状为 [batch_size]
    """
    u = torch.normal(0, 1, (batch_size,), device=device)
    u = torch.sigmoid(u)  # 将正态分布转换到 [0, 1] 区间
    return u

def get_sigmas(timesteps, n_dim, device):
    """
    获取对应时间步的噪声方差（sigmas）
    Args:
        timesteps: 时间步张量
        n_dim: 目标维度，用于广播 sigma
        device: 设备类型
    Returns:
        torch.Tensor: 噪声方差，形状为 [batch_size, 1, 1, 1]
    """
    # 将调度器的时间步和 sigmas 移动到当前设备
    scheduler_timesteps = pipeline.scheduler.timesteps.to(device)
    sigmas = pipeline.scheduler.sigmas.to(device)
    # 确保输入 timesteps 也在同一设备
    timesteps = timesteps.to(device)
    # 查找每个时间步对应的索引
    step_indices = [(scheduler_timesteps == t).nonzero().item() for t in timesteps]
    sigma = sigmas[step_indices].flatten()
    # 广播 sigma 到目标维度（适配 latent 形状）
    while len(sigma.shape) < n_dim:
        sigma = sigma.unsqueeze(-1)
    return sigma

# 开始训练循环
for epoch in range(epochs):
    print(f"Epoch {epoch+1}/{epochs}")
    # 重置累积损失
    total_loss = 0
    for step, (images, captions) in enumerate(dataloader):
        # ----------------------------
        # 1. 编码文本（CLIP + T5）
        # ----------------------------
        with torch.no_grad():
            prompt_embeds, _, pooled_prompt_embeds, _ = pipeline.encode_prompt(
                prompt=captions,
                prompt_2=captions,
                prompt_3=captions,
                device=device,
                negative_prompt='',
                negative_prompt_2='',
                negative_prompt_3='',
                do_classifier_free_guidance=True,
            )
        # ----------------------------
        # 2. 将图像编码为潜在表示（latent）
        # ----------------------------
        # 移动图像到 GPU 并转换为半精度
        images = images.to(device, dtype=torch.float16)
        with torch.no_grad():
            # 使用 VAE 编码器将图像转换为潜在表示
            vae_output = pipeline.vae.encode(images)
            latents = vae_output.latent_dist.sample()
            # 应用 VAE 配置的缩放和偏移因子
            latents = (latents - pipeline.vae.config.shift_factor) * pipeline.vae.config.scaling_factor
        # ----------------------------
        # 3. 采样时间步（带权重方案）
        # ----------------------------
        u = compute_density_for_timestep_sampling(
            batch_size=batch_size,
            device=device
        )
        # 将采样权重转换为时间步索引
        indices = (u * pipeline.scheduler.config.num_train_timesteps).long()
        timesteps = pipeline.scheduler.timesteps.to(device)[indices]
        # ----------------------------
        # 4. Flow Matching：生成带噪声的潜在表示
        # ----------------------------
        # 获取对应时间步的噪声方差
        sigmas = get_sigmas(timesteps, n_dim=latents.ndim, device=device)
        # 生成随机噪声
        noise = torch.randn_like(latents, device=device)
        # 生成中间状态：(1-sigma)*latent + sigma*noise
        # 注意：SD 3.5 的 Flow Matching 插值方向与标准相反
        # 0 时刻是图像的压缩态（latents），1 时刻是噪声
        noisy_latents = (1.0 - sigmas) * latents + sigmas * noise
        # ----------------------------
        # 5. 预测流场（model_pred）
        # ----------------------------
        # 使用混合精度训练，加速计算并减少内存占用
        with torch.autocast("cuda"):
            # 模型预测平均速度（方向：从压缩态到噪声）
            model_pred = pipeline.transformer(
                hidden_states=noisy_latents,
                timestep=timesteps,
                encoder_hidden_states=prompt_embeds,
                pooled_projections=pooled_prompt_embeds,
                return_dict=False
            )[0]
            # 计算预测的 latent：当前位置 + 速度*时间（反向）
            # 模型预测的是平均速度，乘以 (-sigma) 表示反向移动
            pred = model_pred * (-sigmas) + noisy_latents
            # 计算 MSE 损失：预测 latent 与真实 latent 的差距
            loss = F.mse_loss(pred, latents)
        # 反向传播与优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        scheduler.step()
        # 累积损失
        total_loss += loss.item()
        # 打印训练日志
        if step % 100 == 0:
            avg_loss = total_loss / (step + 1)
            print(f"Step {step}, Loss: {loss.item():.4f}, Avg Loss: {avg_loss:.4f}")
    # 打印 epoch 日志
    avg_epoch_loss = total_loss / len(dataloader)
    print(f"Epoch {epoch+1} 完成，平均损失：{avg_epoch_loss:.4f}")

4.3 关于损失计算的说明

SD 3.5 的 Flow Matching 训练与标准 Flow Matching 有两点关键不同：

1. 插值方向相反：
   • 标准 Flow Matching：0 时刻是噪声，1 时刻是图像
   • SD 3.5：0 时刻是图像的压缩态（latents），1 时刻是噪声
   • 中间状态：(1.0 - sigmas) * latents + sigmas * noise
2. 模型预测目标不同：
   • 模型预测的是平均速度（方向：从压缩态到噪声）
   • 距离：model_pred * sigmas，平均速度乘以时间就是距离
   • 预测公式：pred = model_pred * (-sigmas) + noisy_latents
   • 损失计算：MSE(pred, latents)，即中间点往回走一段距离的位置和起点的差距

这种设计能够更好地适应扩散模型的训练特性，提高生成质量和训练稳定性。

5. LoRA 权重保存与加载

5.1 保存 LoRA 权重

# 保存 LoRA 权重，仅保存可训练的 LoRA 参数
pipeline.transformer.save_pretrained("lora-sd35-finetuned")
print("LoRA 权重保存完成")

5.2 加载 LoRA 权重

from peft import PeftModel

# 加载预训练模型
pipeline = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-large",
    torch_dtype=torch.float16
).to("cuda")

# 加载 LoRA 权重到 transformer 组件
pipeline.transformer = PeftModel.from_pretrained(
    pipeline.transformer,
    "lora-sd35-finetuned"
)

# 设置模型为推理模式
pipeline.transformer.eval()
print("LoRA 权重加载完成")

6. 推理

推理方式与之前的 Stable Diffusion 3.5 开发指南完全相同，加载 LoRA 权重后可直接使用 pipeline 进行图像生成：

# 示例：使用微调后的 LoRA 生成图像
prompt = "A beautiful sunset over the mountains in my style"
generated_image = pipeline(
    prompt=prompt,
    negative_prompt="blur, low quality, distortion",
    num_inference_steps=30,
    guidance_scale=7.5
).images[0]
# 保存生成的图像
generated_image.save("generated_image.png")

7. 最佳实践

7.1 数据准备

• 数据质量：确保图像清晰（建议分辨率 ≥ 512x512），文本描述准确且详细
• 数据多样性：包含多种场景、角度和风格的图像，避免过拟合
• 数据量：建议至少准备 100-500 个样本，具体取决于任务复杂度
• 文本描述优化：使用详细的关键词，例如 'a photo of a cat, detailed fur, blue eyes, sunny day' 而非 'a cat"

7.2 训练参数调整

7.3 常见问题与解决方案

7.4 高级技巧

1. 多 LoRA 组合：同时加载多个 LoRA 权重，实现风格混合
2. LoRA 缩放：加载时调整 LoRA 权重的缩放因子，控制风格强度
3. 梯度检查点：启用 gradient_checkpointing 减少内存占用
4. 文本编码器微调：在数据量充足时，可解冻部分文本编码器层进行微调
5. 评估指标：使用 FID、CLIP 分数等指标评估生成质量

总结

本章详细介绍了使用 LoRA 技术微调 Stable Diffusion 3.5 模型的完整流程，包括：

1. 数据集准备与处理：创建图像 - 文本对数据集，实现自定义数据加载器
2. LoRA 微调原理：理解低秩适配器的工作机制和优势
3. 模型加载与配置：加载预训练模型，配置 LoRA 参数
4. 训练循环实现：实现 Flow Matching 训练逻辑，理解 SD 3.5 的特殊损失计算
5. 权重保存与加载：保存和加载 LoRA 权重，实现模型复用

通过 LoRA 微调，您可以高效地定制 Stable Diffusion 3.5 模型，使其适应特定领域或风格的图像生成需求。在实际应用中，建议根据具体任务调整参数和流程，以获得最佳效果。